fix(ingest/snowflake): add more reporting for usage aggregation, handle lineage errors

metadata-ingestion/setup.py

@@ -99,7 +99,7 @@
 sqlglot_lib = {
     # Using an Acryl fork of sqlglot.
     # https://github.com/tobymao/sqlglot/compare/main...hsheth2:sqlglot:hsheth?expand=1
-    "acryl-sqlglot==23.11.2.dev2",
+    "acryl-sqlglot[rs]==23.11.2.dev2",
 }
 
 classification_lib = {

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

@@ -240,6 +240,7 @@ def get_known_query_lineage(
 
         downstream_table_urn = self.dataset_urn_builder(dataset_name)
 
+        # TODO: Can KnownQueryLineageInfo support unknown query lineage i.e. query = None
         known_lineage = KnownQueryLineageInfo(
             query_text=query.query_text,
             downstream=downstream_table_urn,
@@ -369,7 +370,9 @@ def _fetch_upstream_lineages_for_tables(self) -> Iterable[UpstreamLineageEdge]:
         )
         try:
             for db_row in self.query(query):
-                yield UpstreamLineageEdge.parse_obj(db_row)
+                edge = self._process_upstream_lineage_row(db_row)
+                if edge:
+                    yield edge
         except Exception as e:
             if isinstance(e, SnowflakePermissionError):
                 error_msg = "Failed to get table/view to table lineage. Please grant imported privileges on SNOWFLAKE database. "
@@ -382,6 +385,19 @@ def _fetch_upstream_lineages_for_tables(self) -> Iterable[UpstreamLineageEdge]:
                 )
             self.report_status(TABLE_LINEAGE, False)
 
+    def _process_upstream_lineage_row(
+        self, db_row: dict
+    ) -> Optional[UpstreamLineageEdge]:
+        try:
+            return UpstreamLineageEdge.parse_obj(db_row)
+        except Exception as e:
+            self.report.num_upstream_lineage_edge_parsing_failed += 1
+            self.report_warning(
+                f"Parsing lineage edge failed due to error {e}",
+                db_row.get("DOWNSTREAM_TABLE_NAME") or "",
+            )
+            return None
+
     def map_query_result_upstreams(
         self, upstream_tables: Optional[List[UpstreamTableNode]], query_id: str
     ) -> List[UrnStr]:

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

@@ -558,13 +558,20 @@ def usage_per_object_per_time_bucket_for_time_window(
         include_top_n_queries: bool,
         email_domain: Optional[str],
         email_filter: AllowDenyPattern,
+        table_deny_pattern: List[str] = DEFAULT_TABLES_DENY_LIST,
     ) -> str:
         if not include_top_n_queries:
             top_n_queries = 0
         assert (
             time_bucket_size == BucketDuration.DAY
             or time_bucket_size == BucketDuration.HOUR
         )
+
+        temp_table_filter = create_deny_regex_sql_filter(
+            table_deny_pattern,
+            ["object_name"],
+        )
+
         objects_column = (
             "BASE_OBJECTS_ACCESSED" if use_base_objects else "DIRECT_OBJECTS_ACCESSED"
         )
@@ -604,6 +611,7 @@ def usage_per_object_per_time_bucket_for_time_window(
                 )
                 t,
                 lateral flatten(input => t.{objects_column}) object
+            {("where " + temp_table_filter) if temp_table_filter else ""}
         )
         ,
         field_access_history AS

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

@@ -2,6 +2,7 @@
 from datetime import datetime
 from typing import Dict, List, MutableSet, Optional
 
+from datahub.ingestion.api.report import Report
 from datahub.ingestion.glossary.classification_mixin import ClassificationReportMixin
 from datahub.ingestion.source.snowflake.constants import SnowflakeEdition
 from datahub.ingestion.source.sql.sql_generic_profiler import ProfilingSqlReport
@@ -11,6 +12,19 @@
 from datahub.ingestion.source_report.ingestion_stage import IngestionStageReport
 from datahub.ingestion.source_report.time_window import BaseTimeWindowReport
 from datahub.sql_parsing.sql_parsing_aggregator import SqlAggregatorReport
+from datahub.utilities.perf_timer import PerfTimer
+
+
+@dataclass
+class SnowflakeUsageAggregationReport(Report):
+    query_secs: float = -1
+    query_row_count: int = -1
+    result_fetch_timer: PerfTimer = field(default_factory=PerfTimer)
+    result_skip_timer: PerfTimer = field(default_factory=PerfTimer)
+    result_map_timer: PerfTimer = field(default_factory=PerfTimer)
+    users_map_timer: PerfTimer = field(default_factory=PerfTimer)
+    queries_map_timer: PerfTimer = field(default_factory=PerfTimer)
+    fields_map_timer: PerfTimer = field(default_factory=PerfTimer)
 
 
 @dataclass
@@ -31,6 +45,10 @@ class SnowflakeUsageReport:
     usage_end_time: Optional[datetime] = None
     stateful_usage_ingestion_enabled: bool = False
 
+    usage_aggregation: SnowflakeUsageAggregationReport = (
+        SnowflakeUsageAggregationReport()
+    )
+
 
 @dataclass
 class SnowflakeReport(ProfilingSqlReport, BaseTimeWindowReport):
@@ -83,12 +101,10 @@ class SnowflakeV2Report(
     include_technical_schema: bool = False
     include_column_lineage: bool = False
 
-    usage_aggregation_query_secs: float = -1
     table_lineage_query_secs: float = -1
-    # view_lineage_parse_secs: float = -1
-    # view_upstream_lineage_query_secs: float = -1
-    # view_downstream_lineage_query_secs: float = -1
     external_lineage_queries_secs: float = -1
+    num_tables_with_known_upstreams: int = 0
+    num_upstream_lineage_edge_parsing_failed: int = 0
 
     # Reports how many times we reset in-memory `functools.lru_cache` caches of data,
     # which occurs when we occur a different database / schema.
@@ -115,14 +131,6 @@ class SnowflakeV2Report(
 
     edition: Optional[SnowflakeEdition] = None
 
-    # num_tables_with_external_upstreams_only: int = 0
-    num_tables_with_known_upstreams: int = 0
-    # num_views_with_upstreams: int = 0
-
-    # num_view_definitions_parsed: int = 0
-    # num_view_definitions_failed_parsing: int = 0
-    # num_view_definitions_failed_column_parsing: int = 0
-
     def report_entity_scanned(self, name: str, ent_type: str = "table") -> None:
         """
         Entity could be a view or a table or a schema or a database

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

@@ -36,8 +36,9 @@
 )
 from datahub.metadata.com.linkedin.pegasus2avro.timeseries import TimeWindowSize
 from datahub.metadata.schema_classes import OperationClass, OperationTypeClass
+from datahub.sql_parsing.sqlglot_utils import try_format_query
 from datahub.utilities.perf_timer import PerfTimer
-from datahub.utilities.sql_formatter import format_sql_query, trim_query
+from datahub.utilities.sql_formatter import trim_query
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -216,6 +217,7 @@ def _get_workunits_internal(
                         include_top_n_queries=self.config.include_top_n_queries,
                         email_domain=self.config.email_domain,
                         email_filter=self.config.user_email_pattern,
+                        table_deny_pattern=self.config.temporary_tables_pattern,
                     ),
                 )
             except Exception as e:
@@ -227,29 +229,46 @@ def _get_workunits_internal(
                 self.report_status(USAGE_EXTRACTION_USAGE_AGGREGATION, False)
                 return
 
-            self.report.usage_aggregation_query_secs = timer.elapsed_seconds()
+            self.report.usage_aggregation.query_secs = timer.elapsed_seconds()
+            self.report.usage_aggregation.query_row_count = results.rowcount
 
-        for row in results:
-            if not self._is_dataset_pattern_allowed(
-                row["OBJECT_NAME"],
-                row["OBJECT_DOMAIN"],
-            ):
-                continue
-
-            dataset_identifier = self.get_dataset_identifier_from_qualified_name(
-                row["OBJECT_NAME"]
-            )
-            if dataset_identifier not in discovered_datasets:
-                logger.debug(
-                    f"Skipping usage for table {dataset_identifier}, as table schema is not accessible or not allowed by recipe."
-                )
-                continue
+        with self.report.usage_aggregation.result_fetch_timer as fetch_timer:
+            for row in results:
+                with fetch_timer.pause(), self.report.usage_aggregation.result_skip_timer as skip_timer:
+                    if results.rownumber is not None and results.rownumber % 100 == 0:
+                        logger.debug(f"Processing usage row number {results.rownumber}")
+                        logger.debug(self.report.usage_aggregation.as_string())
 
-            yield from self.build_usage_statistics_for_dataset(dataset_identifier, row)
+                    if not self._is_dataset_pattern_allowed(
+                        row["OBJECT_NAME"],
+                        row["OBJECT_DOMAIN"],
+                    ):
+                        logger.debug(
+                            f"Skipping usage for {row['OBJECT_DOMAIN']} {row['OBJECT_NAME']}, as table is not allowed by recipe."
+                        )
+                        continue
+
+                    dataset_identifier = (
+                        self.get_dataset_identifier_from_qualified_name(
+                            row["OBJECT_NAME"]
+                        )
+                    )
+                    if dataset_identifier not in discovered_datasets:
+                        logger.debug(
+                            f"Skipping usage for {row['OBJECT_DOMAIN']} {dataset_identifier}, as table is not accessible."
+                        )
+                        continue
+                    with skip_timer.pause(), self.report.usage_aggregation.result_map_timer as map_timer:
+                        wu = self.build_usage_statistics_for_dataset(
+                            dataset_identifier, row
+                        )
+                        if wu:
+                            with map_timer.pause():
+                                yield wu
 
     def build_usage_statistics_for_dataset(
         self, dataset_identifier: str, row: dict
-    ) -> Iterable[MetadataWorkUnit]:
+    ) -> Optional[MetadataWorkUnit]:
         try:
             stats = DatasetUsageStatistics(
                 timestampMillis=int(row["BUCKET_START_TIME"].timestamp() * 1000),
@@ -258,18 +277,15 @@ def build_usage_statistics_for_dataset(
                 ),
                 totalSqlQueries=row["TOTAL_QUERIES"],
                 uniqueUserCount=row["TOTAL_USERS"],
-                topSqlQueries=self._map_top_sql_queries(
-                    json.loads(row["TOP_SQL_QUERIES"])
-                )
-                if self.config.include_top_n_queries
-                else None,
-                userCounts=self._map_user_counts(
-                    json.loads(row["USER_COUNTS"]),
+                topSqlQueries=(
+                    self._map_top_sql_queries(row["TOP_SQL_QUERIES"])
+                    if self.config.include_top_n_queries
+                    else None
                 ),
-                fieldCounts=self._map_field_counts(json.loads(row["FIELD_COUNTS"])),
+                userCounts=self._map_user_counts(row["USER_COUNTS"]),
+                fieldCounts=self._map_field_counts(row["FIELD_COUNTS"]),
             )
-
-            yield MetadataChangeProposalWrapper(
+            return MetadataChangeProposalWrapper(
                 entityUrn=self.dataset_urn_builder(dataset_identifier), aspect=stats
             ).as_workunit()
         except Exception as e:
@@ -281,61 +297,79 @@ def build_usage_statistics_for_dataset(
                 "Failed to parse usage statistics for dataset", dataset_identifier
             )
 
-    def _map_top_sql_queries(self, top_sql_queries: Dict) -> List[str]:
-        budget_per_query: int = int(
-            self.config.queries_character_limit / self.config.top_n_queries
-        )
-        return sorted(
-            [
-                trim_query(format_sql_query(query), budget_per_query)
-                if self.config.format_sql_queries
-                else trim_query(query, budget_per_query)
-                for query in top_sql_queries
-            ]
-        )
+        return None
+
+    def _map_top_sql_queries(self, top_sql_queries_str: str) -> List[str]:
+        with self.report.usage_aggregation.queries_map_timer:
+            top_sql_queries = json.loads(top_sql_queries_str)
+            budget_per_query: int = int(
+                self.config.queries_character_limit / self.config.top_n_queries
+            )
+            return sorted(
+                [
+                    (
+                        trim_query(
+                            try_format_query(query, self.platform), budget_per_query
+                        )
+                        if self.config.format_sql_queries
+                        else trim_query(query, budget_per_query)
+                    )
+                    for query in top_sql_queries
+                ]
+            )
 
     def _map_user_counts(
         self,
-        user_counts: Dict,
+        user_counts_str: str,
     ) -> List[DatasetUserUsageCounts]:
-        filtered_user_counts = []
-        for user_count in user_counts:
-            user_email = user_count.get("email")
-            if not user_email and self.config.email_domain and user_count["user_name"]:
-                user_email = "{0}@{1}".format(
-                    user_count["user_name"], self.config.email_domain
-                ).lower()
-            if not user_email or not self.config.user_email_pattern.allowed(user_email):
-                continue
-
-            filtered_user_counts.append(
-                DatasetUserUsageCounts(
-                    user=make_user_urn(
-                        self.get_user_identifier(
-                            user_count["user_name"],
-                            user_email,
-                            self.config.email_as_user_identifier,
-                        )
-                    ),
-                    count=user_count["total"],
-                    # NOTE: Generated emails may be incorrect, as email may be different than
-                    # username@email_domain
-                    userEmail=user_email,
+        with self.report.usage_aggregation.users_map_timer:
+            user_counts = json.loads(user_counts_str)
+            filtered_user_counts = []
+            for user_count in user_counts:
+                user_email = user_count.get("email")
+                if (
+                    not user_email
+                    and self.config.email_domain
+                    and user_count["user_name"]
+                ):
+                    user_email = "{0}@{1}".format(
+                        user_count["user_name"], self.config.email_domain
+                    ).lower()
+                if not user_email or not self.config.user_email_pattern.allowed(
+                    user_email
+                ):
+                    continue
+
+                filtered_user_counts.append(
+                    DatasetUserUsageCounts(
+                        user=make_user_urn(
+                            self.get_user_identifier(
+                                user_count["user_name"],
+                                user_email,
+                                self.config.email_as_user_identifier,
+                            )
+                        ),
+                        count=user_count["total"],
+                        # NOTE: Generated emails may be incorrect, as email may be different than
+                        # username@email_domain
+                        userEmail=user_email,
+                    )
                 )
+            return sorted(filtered_user_counts, key=lambda v: v.user)
+
+    def _map_field_counts(self, field_counts_str: str) -> List[DatasetFieldUsageCounts]:
+        with self.report.usage_aggregation.fields_map_timer:
+            field_counts = json.loads(field_counts_str)
+            return sorted(
+                [
+                    DatasetFieldUsageCounts(
+                        fieldPath=self.snowflake_identifier(field_count["col"]),
+                        count=field_count["total"],
+                    )
+                    for field_count in field_counts
+                ],
+                key=lambda v: v.fieldPath,
             )
-        return sorted(filtered_user_counts, key=lambda v: v.user)
-
-    def _map_field_counts(self, field_counts: Dict) -> List[DatasetFieldUsageCounts]:
-        return sorted(
-            [
-                DatasetFieldUsageCounts(
-                    fieldPath=self.snowflake_identifier(field_count["col"]),
-                    count=field_count["total"],
-                )
-                for field_count in field_counts
-            ],
-            key=lambda v: v.fieldPath,
-        )
 
     def _get_snowflake_history(self) -> Iterable[SnowflakeJoinedAccessEvent]:
         logger.info("Getting access history")
@@ -438,9 +472,11 @@ def _get_operation_aspect_work_unit(
                     lastUpdatedTimestamp=last_updated_timestamp,
                     actor=user_urn,
                     operationType=operation_type,
-                    customOperationType=query_type
-                    if operation_type is OperationTypeClass.CUSTOM
-                    else None,
+                    customOperationType=(
+                        query_type
+                        if operation_type is OperationTypeClass.CUSTOM
+                        else None
+                    ),
                 )
                 mcp = MetadataChangeProposalWrapper(
                     entityUrn=self.dataset_urn_builder(dataset_identifier),