feat(ingest): snowflake - support deletion of containers in snowflake…

… stateful ingestion (#6180)
datahub-project · Oct 17, 2022 · 98ba36c · 98ba36c
1 parent 4abe0d2
commit 98ba36c
Showing 1 changed file with 16 additions and 0 deletions.
diff --git a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_v2.py b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_v2.py
@@ -7,6 +7,7 @@
 from snowflake.connector import SnowflakeConnection
 
 from datahub.emitter.mce_builder import (
+    make_container_urn,
     make_data_platform_urn,
     make_dataplatform_instance_urn,
     make_dataset_urn,
@@ -789,6 +790,14 @@ def add_table_to_schema_container(
             container_key=schema_container_key,
             dataset_urn=dataset_urn,
         )
+
+        self.stale_entity_removal_handler.add_entity_to_state(
+            type="container",
+            urn=make_container_urn(
+                guid=schema_container_key.guid(),
+            ),
+        )
+
         for wu in container_workunits:
             self.report.report_workunit(wu)
             yield wu
@@ -838,6 +847,13 @@ def gen_database_containers(
             domain_urn=domain_urn,
         )
 
+        self.stale_entity_removal_handler.add_entity_to_state(
+            type="container",
+            urn=make_container_urn(
+                guid=database_container_key.guid(),
+            ),
+        )
+
         for wu in container_workunits:
             self.report.report_workunit(wu)
             yield wu