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

…gestion
datahub-project · Oct 15, 2022 · 96b335b · 96b335b
1 parent 73fd358
commit 96b335b
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