Repository: astronomer/astro-sdk

Issue:
CI: Fix Snowflake temp tables filling up
@utkarsharma2 @sunank200 - Can one of you add more details to this -- we need to investigate and fix the tests that is filling up our Snowflake instances with tables


Discussion:

Resolving patch:
diff --git a/python-sdk/example_dags/example_amazon_s3_snowflake_transform.py b/python-sdk/example_dags/example_amazon_s3_snowflake_transform.py
--- a/python-sdk/example_dags/example_amazon_s3_snowflake_transform.py
+++ b/python-sdk/example_dags/example_amazon_s3_snowflake_transform.py
@@ -52,21 +52,23 @@ def example_amazon_s3_snowflake_transform():
     s3_bucket = os.getenv("S3_BUCKET", "s3://tmp9")
 
     input_table_1 = Table(
-        name="ADOPTION_CENTER_1",
+        name="ADOPTION_CENTER_1_" + str(int(time.time())),
         metadata=Metadata(
             database=os.environ["SNOWFLAKE_DATABASE"],
             schema=os.environ["SNOWFLAKE_SCHEMA"],
         ),
         conn_id="snowflake_conn",
+        temp=True,
     )
     # [START metadata_example_snowflake]
     input_table_2 = Table(
-        name="ADOPTION_CENTER_2",
+        name="ADOPTION_CENTER_2_" + str(int(time.time())),
         metadata=Metadata(
             database=os.environ["SNOWFLAKE_DATABASE"],
             schema=os.environ["SNOWFLAKE_SCHEMA"],
         ),
         conn_id="snowflake_conn",
+        temp=True,
     )
     # [END metadata_example_snowflake]
 
@@ -86,16 +88,18 @@ def example_amazon_s3_snowflake_transform():
 
     cleaned_data = clean_data(combined_data)
     # [START dataframe_example_2]
+    snowflake_output_table = Table(
+        name="aggregated_adoptions_" + str(int(time.time())),
+        metadata=Metadata(
+            schema=os.environ["SNOWFLAKE_SCHEMA"],
+            database=os.environ["SNOWFLAKE_DATABASE"],
+        ),
+        conn_id="snowflake_conn",
+        temp=True,
+    )
     aggregate_data(
         cleaned_data,
-        output_table=Table(
-            name="aggregated_adoptions_" + str(int(time.time())),
-            metadata=Metadata(
-                schema=os.environ["SNOWFLAKE_SCHEMA"],
-                database=os.environ["SNOWFLAKE_DATABASE"],
-            ),
-            conn_id="snowflake_conn",
-        ),
+        output_table=snowflake_output_table,
     )
     # [END dataframe_example_2]
     aql.cleanup()
diff --git a/python-sdk/example_dags/example_snowflake_partial_table_with_append.py b/python-sdk/example_dags/example_snowflake_partial_table_with_append.py
--- a/python-sdk/example_dags/example_snowflake_partial_table_with_append.py
+++ b/python-sdk/example_dags/example_snowflake_partial_table_with_append.py
@@ -9,6 +9,7 @@
 """
 
 import os
+import time
 from datetime import datetime
 
 import pandas as pd
@@ -60,7 +61,7 @@ def filter_data(homes_long: Table):
 def create_table(table: Table):
     """Create the reporting data which will be the target of the append method"""
     return """
-    CREATE TABLE IF NOT EXISTS {{table}} (
+      CREATE OR REPLACE TABLE {{table}} (
       sell number,
       list number,
       variable varchar,
@@ -71,7 +72,7 @@ def create_table(table: Table):
 
 @dag(start_date=datetime(2021, 12, 1), schedule_interval="@daily", catchup=False)
 def example_snowflake_partial_table_with_append():
-    homes_reporting = Table(conn_id=SNOWFLAKE_CONN_ID)
+    homes_reporting = Table(name="homes_reporting_data", temp=True, conn_id=SNOWFLAKE_CONN_ID)
     create_results_table = create_table(table=homes_reporting, conn_id=SNOWFLAKE_CONN_ID)
     # [END howto_run_raw_sql_snowflake_1]
 
@@ -102,10 +103,12 @@ def example_snowflake_partial_table_with_append():
     extracted_data = extract_data(
         homes1=homes_data1,
         homes2=homes_data2,
-        output_table=Table(name="combined_homes_data"),
+        output_table=Table(name="combined_homes_data_" + str(int(time.time())), temp=True),
     )
 
-    transformed_data = transform_data(df=extracted_data, output_table=Table(name="homes_data_long"))
+    transformed_data = transform_data(
+        df=extracted_data, output_table=Table(name="homes_data_long_" + str(int(time.time())), temp=True)
+    )
 
     filtered_data = filter_data(
         homes_long=transformed_data,


Test patch:
diff --git a/python-sdk/tests/sql/operators/test_append.py b/python-sdk/tests/sql/operators/test_append.py
--- a/python-sdk/tests/sql/operators/test_append.py
+++ b/python-sdk/tests/sql/operators/test_append.py
@@ -144,6 +144,7 @@ def test_append(database_table_fixture, sample_dag, multiple_tables_fixture, app
             source_table=append_table,
         )
         validate_append(appended_table)
+        aql.cleanup()
     test_utils.run_dag(sample_dag)
 
 
diff --git a/python-sdk/tests/sql/operators/transform/test_transform.py b/python-sdk/tests/sql/operators/transform/test_transform.py
--- a/python-sdk/tests/sql/operators/transform/test_transform.py
+++ b/python-sdk/tests/sql/operators/transform/test_transform.py
@@ -260,6 +260,7 @@ def select(input_table: Table):
             return "SELECT * FROM {{input_table}} LIMIT 4;"
 
         select(input_table=homes_file, output_table=Table(conn_id="snowflake_conn_1"))
+        aql.cleanup()
     test_utils.run_dag(sample_dag)