Added 'extend_existing' to Sqla Table object

docs/astro/sql/operators/drop.rst

@@ -8,7 +8,7 @@ When to Use drop operator
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 The drop operator is used to delete tables from the database if they exist. It can be used on both Temporary as well as Persistent Tables.
 
-.. literalinclude:: ../../../../example_dags/example_snowflake_partial_table_with_append.py
+.. literalinclude:: ../../../../example_dags/example_sqlite_load_transform.py
    :language: python
    :start-after: [START drop_table_example]
    :end-before: [END drop_table_example]

example_dags/example_snowflake_partial_table_with_append.py

@@ -5,15 +5,7 @@
 from airflow.decorators import dag
 
 from astro.files import File
-from astro.sql import (
-    append,
-    cleanup,
-    dataframe,
-    drop_table,
-    load_file,
-    run_raw_sql,
-    transform,
-)
+from astro.sql import append, cleanup, dataframe, load_file, run_raw_sql, transform
 from astro.sql.table import Metadata, Table
 
 """
@@ -66,10 +58,10 @@ def filter_data(homes_long: Table):
 
 
 @run_raw_sql
-def create_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 homes_reporting (
+    CREATE TABLE IF NOT EXISTS {{table}} (
       sell number,
       list number,
       variable varchar,
@@ -85,7 +77,6 @@ def example_snowflake_partial_table_with_append():
     homes_data1 = load_file(
         input_file=File(path=FILE_PATH + "homes.csv"),
         output_table=Table(
-            name="homes",
             conn_id=SNOWFLAKE_CONN_ID,
             metadata=Metadata(
                 database=os.getenv("SNOWFLAKE_DATABASE"),
@@ -97,7 +88,6 @@ def example_snowflake_partial_table_with_append():
     homes_data2 = load_file(
         input_file=File(path=FILE_PATH + "homes2.csv"),
         output_table=Table(
-            name="homes2",
             conn_id=SNOWFLAKE_CONN_ID,
             metadata=Metadata(
                 database=os.getenv("SNOWFLAKE_DATABASE"),
@@ -119,31 +109,24 @@ def example_snowflake_partial_table_with_append():
 
     filtered_data = filter_data(
         homes_long=transformed_data,
-        output_table=Table(name="expensive_homes_long"),
+        output_table=Table(),
+    )
+    homes_reporting = Table(conn_id=SNOWFLAKE_CONN_ID)
+    create_results_table = create_table(
+        table=homes_reporting, conn_id=SNOWFLAKE_CONN_ID
     )
-
-    create_results_table = create_table(conn_id=SNOWFLAKE_CONN_ID)
 
     # Append transformed & filtered data to reporting table
     # Dependency is inferred by passing the previous `filtered_data` task to `append_table` param
     # [START append_example_with_columns_list]
     record_results = append(
         source_table=filtered_data,
-        target_table=Table(name="homes_reporting", conn_id=SNOWFLAKE_CONN_ID),
+        target_table=homes_reporting,
         columns=["sell", "list", "variable", "value"],
     )
     # [END append_example_with_columns_list]
     record_results.set_upstream(create_results_table)
 
-    # We truncate this table only to avoid wasting Snowflake resources
-    # Why? Between 2022-03-25 and 2022-04-11 it accumulated 301G (89 million rows) because
-    # this example DAG used to append rows without deleting them
-    # [START drop_table_example]
-    truncate_results = drop_table(
-        table=Table(name="homes_reporting", conn_id=SNOWFLAKE_CONN_ID)
-    )
-    # [END drop_table_example]
-    truncate_results.set_upstream(record_results)
     cleanup()
 
 

example_dags/example_sqlite_load_transform.py

@@ -1,16 +1,18 @@
+import time
 from datetime import datetime
 
 from airflow import DAG
 
 from astro import sql as aql
 from astro.files import File
+from astro.sql import drop_table
 from astro.sql.table import Table
 
 START_DATE = datetime(2000, 1, 1)
 
 
 @aql.transform()
-def top_five_animations(input_table: Table):
+def get_top_five_animations(input_table: Table):
     return """
         SELECT Title, Rating
         FROM {{input_table}}
@@ -20,6 +22,8 @@ def top_five_animations(input_table: Table):
     """
 
 
+imdb_movies_name = "imdb_movies" + str(int(time.time()))
+
 with DAG(
     "example_sqlite_load_transform",
     schedule_interval=None,
@@ -32,14 +36,21 @@ def top_five_animations(input_table: Table):
             path="https://raw.githubusercontent.com/astronomer/astro-sdk/main/tests/data/imdb.csv"
         ),
         task_id="load_csv",
-        output_table=Table(name="imdb_movies", conn_id="sqlite_default"),
+        output_table=Table(name=imdb_movies_name, conn_id="sqlite_default"),
     )
 
-    top_five_animations(
+    top_five_animations = get_top_five_animations(
         input_table=imdb_movies,
         output_table=Table(
             name="top_animation",
             conn_id="sqlite_default",
         ),
     )
+    # Note - Using persistent table just to showcase drop_table operator.
+    # [START drop_table_example]
+    truncate_results = drop_table(
+        table=Table(name=imdb_movies_name, conn_id="sqlite_default")
+    )
+    # [END drop_table_example]
+    truncate_results.set_upstream(top_five_animations)
     aql.cleanup()

src/astro/databases/base.py

@@ -477,7 +477,10 @@ def get_sqla_table(self, table: Table) -> SqlaTable:
         :param table: Astro Table to be converted to SQLAlchemy table instance
         """
         return SqlaTable(
-            table.name, table.sqlalchemy_metadata, autoload_with=self.sqlalchemy_engine
+            table.name,
+            table.sqlalchemy_metadata,
+            autoload_with=self.sqlalchemy_engine,
+            extend_existing=True,
         )
 
     # ---------------------------------------------------------