feat(ingest): add snowflake-queries source

metadata-ingestion/setup.py

@@ -409,6 +409,7 @@
     "salesforce": {"simple-salesforce"},
     "snowflake": snowflake_common | usage_common | sqlglot_lib,
     "snowflake-summary": snowflake_common | usage_common | sqlglot_lib,
+    "snowflake-queries": snowflake_common | usage_common | sqlglot_lib,
     "sqlalchemy": sql_common,
     "sql-queries": usage_common | sqlglot_lib,
     "slack": slack,
@@ -661,6 +662,7 @@
         "slack = datahub.ingestion.source.slack.slack:SlackSource",
         "snowflake = datahub.ingestion.source.snowflake.snowflake_v2:SnowflakeV2Source",
         "snowflake-summary = datahub.ingestion.source.snowflake.snowflake_summary:SnowflakeSummarySource",
+        "snowflake-queries = datahub.ingestion.source.snowflake.snowflake_queries:SnowflakeQueriesSource",
         "superset = datahub.ingestion.source.superset:SupersetSource",
         "tableau = datahub.ingestion.source.tableau:TableauSource",
         "openapi = datahub.ingestion.source.openapi:OpenApiSource",

metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_lineage_v2.py

@@ -30,6 +30,7 @@
 from datahub.sql_parsing.sql_parsing_aggregator import (
     ColumnLineageInfo,
     ColumnRef,
+    KnownLineageMapping,
     KnownQueryLineageInfo,
     SqlParsingAggregator,
     UrnStr,
@@ -264,13 +265,20 @@ def _populate_external_upstreams(self, discovered_tables: List[str]) -> None:
         with PerfTimer() as timer:
             self.report.num_external_table_edges_scanned = 0
 
-            self._populate_external_lineage_from_copy_history(discovered_tables)
+            for (
+                known_lineage_mapping
+            ) in self._populate_external_lineage_from_copy_history(discovered_tables):
+                self.sql_aggregator.add(known_lineage_mapping)
             logger.info(
                 "Done populating external lineage from copy history. "
                 f"Found {self.report.num_external_table_edges_scanned} external lineage edges so far."
             )
 
-            self._populate_external_lineage_from_show_query(discovered_tables)
+            for (
+                known_lineage_mapping
+            ) in self._populate_external_lineage_from_show_query(discovered_tables):
+                self.sql_aggregator.add(known_lineage_mapping)
+
             logger.info(
                 "Done populating external lineage from show external tables. "
                 f"Found {self.report.num_external_table_edges_scanned} external lineage edges so far."
@@ -282,7 +290,7 @@ def _populate_external_upstreams(self, discovered_tables: List[str]) -> None:
     # NOTE: Snowflake does not log this information to the access_history table.
     def _populate_external_lineage_from_show_query(
         self, discovered_tables: List[str]
-    ) -> None:
+    ) -> Iterable[KnownLineageMapping]:
         external_tables_query: str = SnowflakeQuery.show_external_tables()
         try:
             for db_row in self.query(external_tables_query):
@@ -293,11 +301,11 @@ def _populate_external_lineage_from_show_query(
                 if key not in discovered_tables:
                     continue
                 if db_row["location"].startswith("s3://"):
-                    self.sql_aggregator.add_known_lineage_mapping(
-                        downstream_urn=self.dataset_urn_builder(key),
+                    yield KnownLineageMapping(
                         upstream_urn=make_s3_urn_for_lineage(
                             db_row["location"], self.config.env
                         ),
+                        downstream_urn=self.dataset_urn_builder(key),
                     )
                     self.report.num_external_table_edges_scanned += 1
 
@@ -316,7 +324,7 @@ def _populate_external_lineage_from_show_query(
     # NOTE: Snowflake does not log this information to the access_history table.
     def _populate_external_lineage_from_copy_history(
         self, discovered_tables: List[str]
-    ) -> None:
+    ) -> Iterable[KnownLineageMapping]:
         query: str = SnowflakeQuery.copy_lineage_history(
             start_time_millis=int(self.start_time.timestamp() * 1000),
             end_time_millis=int(self.end_time.timestamp() * 1000),
@@ -325,7 +333,11 @@ def _populate_external_lineage_from_copy_history(
 
         try:
             for db_row in self.query(query):
-                self._process_external_lineage_result_row(db_row, discovered_tables)
+                known_lineage_mapping = self._process_external_lineage_result_row(
+                    db_row, discovered_tables
+                )
+                if known_lineage_mapping:
+                    yield known_lineage_mapping
         except Exception as e:
             if isinstance(e, SnowflakePermissionError):
                 error_msg = "Failed to get external lineage. Please grant imported privileges on SNOWFLAKE database. "
@@ -340,7 +352,7 @@ def _populate_external_lineage_from_copy_history(
 
     def _process_external_lineage_result_row(
         self, db_row: dict, discovered_tables: List[str]
-    ) -> None:
+    ) -> Optional[KnownLineageMapping]:
         # key is the down-stream table name
         key: str = self.get_dataset_identifier_from_qualified_name(
             db_row["DOWNSTREAM_TABLE_NAME"]
@@ -353,11 +365,11 @@ def _process_external_lineage_result_row(
 
             for loc in external_locations:
                 if loc.startswith("s3://"):
-                    self.sql_aggregator.add_known_lineage_mapping(
-                        downstream_urn=self.dataset_urn_builder(key),
+                    self.report.num_external_table_edges_scanned += 1
+                    return KnownLineageMapping(
                         upstream_urn=make_s3_urn_for_lineage(loc, self.config.env),
+                        downstream_urn=self.dataset_urn_builder(key),
                     )
-                    self.report.num_external_table_edges_scanned += 1
 
     def _fetch_upstream_lineages_for_tables(self) -> Iterable[UpstreamLineageEdge]:
         query: str = SnowflakeQuery.table_to_table_lineage_history_v2(