passes dropped tables schemas to filter, adjust athena filter

Author: rudolfix

dlt/common/destination/client.py

@@ -44,6 +44,7 @@
     C_DLT_LOAD_ID,
     TLoaderReplaceStrategy,
     TTableFormat,
+    TTableSchema,
 )
 from dlt.common.destination.capabilities import DestinationCapabilitiesContext
 from dlt.common.destination.exceptions import (
@@ -686,10 +687,11 @@ def should_truncate_table_before_load_on_staging_destination(self, table_name: s
         """
         pass
 
-    def should_drop_table_on_staging_destination(self, table_name: str) -> bool:
-        """Tells if table should be dropped on staging destination (regular dataset) in addition to dropping the table on
+    def should_drop_table_on_staging_destination(self, dropped_table: TTableSchema) -> bool:
+        """Tells if `dropped_table` should be dropped on staging destination (regular dataset) in addition to dropping the table on
         final destination. This stays False for all the destinations except Athena, non-iceberg where staging destination
         holds actual data which needs to be deleted.
+        Note that `dropped_table` may not longer be present in schema. It is present only if it got recreated.
         """
         return False
 

dlt/destinations/impl/athena/athena.py

@@ -11,6 +11,8 @@
     TYPE_CHECKING,
 )
 
+from dlt.common.schema.exceptions import TableNotFound
+
 if TYPE_CHECKING:
     from mypy_boto3_lakeformation import LakeFormationClient
     from mypy_boto3_lakeformation.type_defs import (
@@ -35,6 +37,7 @@
     TColumnType,
     TSchemaTables,
     TSortOrder,
+    TTableSchema,
 )
 from dlt.common.destination import DestinationCapabilitiesContext, PreparedTableSchema
 from dlt.common.destination.client import FollowupJobRequest, SupportsStagingDestination, LoadJob
@@ -406,10 +409,15 @@ def should_truncate_table_before_load_on_staging_destination(self, table_name: s
             return True
         return False
 
-    def should_drop_table_on_staging_destination(self, table_name: str) -> bool:
+    def should_drop_table_on_staging_destination(self, dropped_table: TTableSchema) -> bool:
         # in Athena we must drop table in glue and then we must drop data in staging if table is not iceberg
-        table = self.prepare_load_table(table_name)
-        return not self._is_iceberg_table(table)
+        try:
+            existing_table = self.prepare_load_table(dropped_table["name"])
+            # do not drop data if new iceberg table got created - storage is handled by Athena
+            return not self._is_iceberg_table(existing_table)
+        except TableNotFound:
+            # table got dropped and is not recreated - drop staging destination
+            return True
 
     def should_load_data_to_staging_dataset_on_staging_destination(self, table_name: str) -> bool:
         """iceberg table data goes into staging on staging destination"""

dlt/load/utils.py

@@ -69,7 +69,7 @@ def init_client(
     expected_update: TSchemaTables,
     truncate_filter: Callable[[str], bool],
     load_staging_filter: Callable[[str], bool],
-    drop_staging_filter: Callable[[str], bool],
+    drop_staging_filter: Callable[[TTableSchema], bool],
     drop_tables: Optional[List[TTableSchema]] = None,
     truncate_tables: Optional[List[TTableSchema]] = None,
 ) -> TSchemaTables:
@@ -113,9 +113,10 @@ def init_client(
         )
     )
 
-    # get tables to drop
+    # get tables to drop, note that drop_tables are not in schema and come from the package
+    # state
     drop_table_names = (
-        {table["name"] for table in drop_tables if drop_staging_filter(table["name"])}
+        {table["name"] for table in drop_tables if drop_staging_filter(table)}
         if drop_tables
         else set()
     )

tests/load/pipeline/test_refresh_modes.py

@@ -118,11 +118,12 @@ def test_refresh_drop_sources(
     if not in_source:
         data = list(data.selected_resources.values())
 
-    # First run pipeline so destination so tables are created
+    # first run pipeline so destination so tables are created
     info = pipeline.run(data, refresh="drop_sources", **destination_config.run_kwargs)
     assert_load_info(info)
+    assert table_exists(pipeline, "some_data_3")
 
-    # Second run of pipeline with only selected resources
+    # second run of pipeline with only selected resources
     if with_wipe:
         pipeline._wipe_working_folder()
         pipeline = destination_config.setup_pipeline(pipeline_name, dataset_name=dataset_name)
@@ -144,16 +145,16 @@ def test_refresh_drop_sources(
         "some_data_2",
     }
 
-    # No "name" column should exist as table was dropped and re-created without it
+    # no "name" column should exist as table was dropped and re-created without it
     assert_only_table_columns(pipeline, "some_data_1", ["id"])
     data = load_tables_to_dicts(pipeline, "some_data_1")["some_data_1"]
     result = sorted([row["id"] for row in data])
-    # Only rows from second run should exist
+    # only rows from second run should exist
     assert result == [3, 4]
 
-    # Confirm resource tables not selected on second run got dropped
+    # confirm resource tables not selected on second run got dropped
     assert not table_exists(pipeline, "some_data_3")
-    # Loaded state is wiped
+    # loaded state is wiped
     with pipeline.destination_client() as dest_client:
         destination_state = load_pipeline_state_from_destination(
             pipeline.pipeline_name, dest_client  # type: ignore[arg-type]