feat(ingest): add fast query fingerprinting (#10619)

datahub-project · Jul 16, 2024 · 3a9af18 · 3a9af18
1 parent 63b3ba7
commit 3a9af18
Show file tree

Hide file tree

Showing 8 changed files with 218 additions and 76 deletions.
diff --git a/metadata-ingestion/src/datahub/entrypoints.py b/metadata-ingestion/src/datahub/entrypoints.py
@@ -66,7 +66,7 @@
     "--log-file",
     type=click.Path(dir_okay=False),
     default=None,
-    help="Enable debug logging.",
+    help="Write debug-level logs to a file.",
 )
 @click.version_option(
     version=datahub_package.nice_version_name(),

diff --git a/metadata-ingestion/src/datahub/ingestion/glossary/classification_mixin.py b/metadata-ingestion/src/datahub/ingestion/glossary/classification_mixin.py
@@ -300,6 +300,8 @@ def classification_workunit_processor(
     table_name = ".".join(table_id)
     if not classification_handler.is_classification_enabled_for_table(table_name):
         yield from table_wu_generator
+        return
+
     for wu in table_wu_generator:
         maybe_schema_metadata = wu.get_aspect_of_type(SchemaMetadata)
         if (

diff --git a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_v2.py b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_v2.py
@@ -257,6 +257,7 @@ def __init__(self, ctx: PipelineContext, config: SnowflakeV2Config):
                 graph=self.ctx.graph,
                 generate_usage_statistics=False,
                 generate_operations=False,
+                format_queries=self.config.format_sql_queries,
             )
             self.report.sql_aggregator = self.aggregator.report
 

diff --git a/metadata-ingestion/src/datahub/sql_parsing/sql_parsing_aggregator.py b/metadata-ingestion/src/datahub/sql_parsing/sql_parsing_aggregator.py
@@ -160,6 +160,8 @@ class SqlAggregatorReport(Report):
     # SQL parsing (over all invocations).
     num_sql_parsed: int = 0
     sql_parsing_timer: PerfTimer = dataclasses.field(default_factory=PerfTimer)
+    sql_fingerprinting_timer: PerfTimer = dataclasses.field(default_factory=PerfTimer)
+    sql_formatting_timer: PerfTimer = dataclasses.field(default_factory=PerfTimer)
 
     # Other lineage loading metrics.
     num_known_query_lineage: int = 0
@@ -381,7 +383,8 @@ def _initialize_schema_resolver_from_graph(self, graph: DataHubGraph) -> None:
 
     def _maybe_format_query(self, query: str) -> str:
         if self.format_queries:
-            return try_format_query(query, self.platform.platform_name)
+            with self.report.sql_formatting_timer:
+                return try_format_query(query, self.platform.platform_name)
         return query
 
     def add_known_query_lineage(
@@ -405,9 +408,12 @@ def add_known_query_lineage(
         self.report.num_known_query_lineage += 1
 
         # Generate a fingerprint for the query.
-        query_fingerprint = get_query_fingerprint(
-            known_query_lineage.query_text, platform=self.platform.platform_name
-        )
+        with self.report.sql_fingerprinting_timer:
+            query_fingerprint = get_query_fingerprint(
+                known_query_lineage.query_text,
+                platform=self.platform.platform_name,
+                fast=True,
+            )
         formatted_query = self._maybe_format_query(known_query_lineage.query_text)
 
         # Register the query.

diff --git a/metadata-ingestion/src/datahub/sql_parsing/sqlglot_utils.py b/metadata-ingestion/src/datahub/sql_parsing/sqlglot_utils.py
@@ -1,6 +1,7 @@
 import functools
 import hashlib
 import logging
+import re
 from typing import Dict, Iterable, Optional, Tuple, Union
 
 import sqlglot
@@ -109,6 +110,80 @@ def _expression_to_string(
     return expression.sql(dialect=get_dialect(platform))
 
 
+_BASIC_NORMALIZATION_RULES = {
+    # Remove /* */ comments.
+    re.compile(r"/\*.*?\*/", re.DOTALL): "",
+    # Remove -- comments.
+    re.compile(r"--.*$"): "",
+    # Replace all runs of whitespace with a single space.
+    re.compile(r"\s+"): " ",
+    # Remove leading and trailing whitespace and trailing semicolons.
+    re.compile(r"^\s+|[\s;]+$"): "",
+    # Replace anything that looks like a number with a placeholder.
+    re.compile(r"\b\d+\b"): "?",
+    # Replace anything that looks like a string with a placeholder.
+    re.compile(r"'[^']*'"): "?",
+    # Replace sequences of IN/VALUES with a single placeholder.
+    re.compile(r"\b(IN|VALUES)\s*\(\?(?:, \?)*\)", re.IGNORECASE): r"\1 (?)",
+    # Normalize parenthesis spacing.
+    re.compile(r"\( "): "(",
+    re.compile(r" \)"): ")",
+}
+_TABLE_NAME_NORMALIZATION_RULES = {
+    # Replace UUID-like strings with a placeholder (both - and _ variants).
+    re.compile(
+        r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}",
+        re.IGNORECASE,
+    ): "00000000-0000-0000-0000-000000000000",
+    re.compile(
+        r"[0-9a-f]{8}_[0-9a-f]{4}_[0-9a-f]{4}_[0-9a-f]{4}_[0-9a-f]{12}",
+        re.IGNORECASE,
+    ): "00000000_0000_0000_0000_000000000000",
+    # GE temporary table names (prefix + 8 digits of a UUIDv4)
+    re.compile(
+        r"\b(ge_tmp_|ge_temp_|gx_temp_)[0-9a-f]{8}\b", re.IGNORECASE
+    ): r"\1abcdefgh",
+    # Date-suffixed table names (e.g. _20210101)
+    re.compile(r"\b(\w+)(19|20)\d{4}\b"): r"\1YYYYMM",
+    re.compile(r"\b(\w+)(19|20)\d{6}\b"): r"\1YYYYMMDD",
+    re.compile(r"\b(\w+)(19|20)\d{8}\b"): r"\1YYYYMMDDHH",
+    re.compile(r"\b(\w+)(19|20)\d{10}\b"): r"\1YYYYMMDDHHMM",
+}
+
+
+def generalize_query_fast(
+    expression: sqlglot.exp.ExpOrStr,
+    dialect: DialectOrStr,
+    change_table_names: bool = False,
+) -> str:
+    """Variant of `generalize_query` that only does basic normalization.
+
+    Args:
+        expression: The SQL query to generalize.
+        dialect: The SQL dialect to use.
+        change_table_names: If True, replace table names with placeholders. Note
+            that this should only be used for query filtering purposes, as it
+            violates the general assumption that the queries with the same fingerprint
+            have the same lineage/usage/etc.
+
+    Returns:
+        The generalized SQL query.
+    """
+
+    if isinstance(expression, sqlglot.exp.Expression):
+        expression = expression.sql(dialect=get_dialect(dialect))
+    query_text = expression
+
+    REGEX_REPLACEMENTS = {
+        **_BASIC_NORMALIZATION_RULES,
+        **(_TABLE_NAME_NORMALIZATION_RULES if change_table_names else {}),
+    }
+
+    for pattern, replacement in REGEX_REPLACEMENTS.items():
+        query_text = pattern.sub(replacement, query_text)
+    return query_text
+
+
 def generalize_query(expression: sqlglot.exp.ExpOrStr, dialect: DialectOrStr) -> str:
     """
     Generalize/normalize a SQL query.
@@ -172,11 +247,14 @@ def generate_hash(text: str) -> str:
 
 
 def get_query_fingerprint_debug(
-    expression: sqlglot.exp.ExpOrStr, platform: DialectOrStr
+    expression: sqlglot.exp.ExpOrStr, platform: DialectOrStr, fast: bool = False
 ) -> Tuple[str, Optional[str]]:
     try:
-        dialect = get_dialect(platform)
-        expression_sql = generalize_query(expression, dialect=dialect)
+        if not fast:
+            dialect = get_dialect(platform)
+            expression_sql = generalize_query(expression, dialect=dialect)
+        else:
+            expression_sql = generalize_query_fast(expression, dialect=platform)
     except (ValueError, sqlglot.errors.SqlglotError) as e:
         if not isinstance(expression, str):
             raise
@@ -193,7 +271,7 @@ def get_query_fingerprint_debug(
 
 
 def get_query_fingerprint(
-    expression: sqlglot.exp.ExpOrStr, platform: DialectOrStr
+    expression: sqlglot.exp.ExpOrStr, platform: DialectOrStr, fast: bool = False
 ) -> str:
     """Get a fingerprint for a SQL query.
 
@@ -215,7 +293,7 @@ def get_query_fingerprint(
         The fingerprint for the SQL query.
     """
 
-    return get_query_fingerprint_debug(expression, platform)[0]
+    return get_query_fingerprint_debug(expression, platform, fast=fast)[0]
 
 
 @functools.lru_cache(maxsize=FORMAT_QUERY_CACHE_SIZE)

diff --git a/metadata-ingestion/tests/integration/snowflake/test_snowflake.py b/metadata-ingestion/tests/integration/snowflake/test_snowflake.py
@@ -119,6 +119,7 @@ def test_snowflake_basic(pytestconfig, tmp_path, mock_time, mock_datahub_graph):
                         include_table_lineage=True,
                         include_view_lineage=True,
                         include_usage_stats=True,
+                        format_sql_queries=True,
                         validate_upstreams_against_patterns=False,
                         include_operational_stats=True,
                         email_as_user_identifier=True,
@@ -213,6 +214,7 @@ def test_snowflake_private_link(pytestconfig, tmp_path, mock_time, mock_datahub_
                         include_views=True,
                         include_view_lineage=True,
                         include_usage_stats=False,
+                        format_sql_queries=True,
                         incremental_lineage=False,
                         include_operational_stats=False,
                         platform_instance="instance1",

diff --git a/...ata-ingestion/tests/unit/sql_parsing/aggregator_goldens/test_add_known_query_lineage.json b/...ata-ingestion/tests/unit/sql_parsing/aggregator_goldens/test_add_known_query_lineage.json
@@ -18,7 +18,7 @@
                     },
                     "dataset": "urn:li:dataset:(urn:li:dataPlatform:redshift,dev.public.bar,PROD)",
                     "type": "TRANSFORMED",
-                    "query": "urn:li:query:02e2ec36678bea2a8c4c855fed5255d087cfeb2710d326e95fd9b48a9c4fc0ae"
+                    "query": "urn:li:query:6ed1d12fbf2ccc8138ceec08cc35b981030d6d004bfad9743c7afd84260fa63f"
                 }
             ],
             "fineGrainedLineages": [
@@ -32,7 +32,7 @@
                         "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:redshift,dev.public.foo,PROD),a)"
                     ],
                     "confidenceScore": 1.0,
-                    "query": "urn:li:query:02e2ec36678bea2a8c4c855fed5255d087cfeb2710d326e95fd9b48a9c4fc0ae"
+                    "query": "urn:li:query:6ed1d12fbf2ccc8138ceec08cc35b981030d6d004bfad9743c7afd84260fa63f"
                 },
                 {
                     "upstreamType": "FIELD_SET",
@@ -44,7 +44,7 @@
                         "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:redshift,dev.public.foo,PROD),b)"
                     ],
                     "confidenceScore": 1.0,
-                    "query": "urn:li:query:02e2ec36678bea2a8c4c855fed5255d087cfeb2710d326e95fd9b48a9c4fc0ae"
+                    "query": "urn:li:query:6ed1d12fbf2ccc8138ceec08cc35b981030d6d004bfad9743c7afd84260fa63f"
                 },
                 {
                     "upstreamType": "FIELD_SET",
@@ -56,15 +56,15 @@
                         "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:redshift,dev.public.foo,PROD),c)"
                     ],
                     "confidenceScore": 1.0,
-                    "query": "urn:li:query:02e2ec36678bea2a8c4c855fed5255d087cfeb2710d326e95fd9b48a9c4fc0ae"
+                    "query": "urn:li:query:6ed1d12fbf2ccc8138ceec08cc35b981030d6d004bfad9743c7afd84260fa63f"
                 }
             ]
         }
     }
 },
 {
     "entityType": "query",
-    "entityUrn": "urn:li:query:02e2ec36678bea2a8c4c855fed5255d087cfeb2710d326e95fd9b48a9c4fc0ae",
+    "entityUrn": "urn:li:query:6ed1d12fbf2ccc8138ceec08cc35b981030d6d004bfad9743c7afd84260fa63f",
     "changeType": "UPSERT",
     "aspectName": "queryProperties",
     "aspect": {
@@ -85,9 +85,29 @@
         }
     }
 },
+{
+    "entityType": "dataset",
+    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:redshift,dev.public.foo,PROD)",
+    "changeType": "UPSERT",
+    "aspectName": "operation",
+    "aspect": {
+        "json": {
+            "timestampMillis": 1707182625000,
+            "partitionSpec": {
+                "type": "FULL_TABLE",
+                "partition": "FULL_TABLE_SNAPSHOT"
+            },
+            "operationType": "INSERT",
+            "customProperties": {
+                "query_urn": "urn:li:query:6ed1d12fbf2ccc8138ceec08cc35b981030d6d004bfad9743c7afd84260fa63f"
+            },
+            "lastUpdatedTimestamp": 20000
+        }
+    }
+},
 {
     "entityType": "query",
-    "entityUrn": "urn:li:query:02e2ec36678bea2a8c4c855fed5255d087cfeb2710d326e95fd9b48a9c4fc0ae",
+    "entityUrn": "urn:li:query:6ed1d12fbf2ccc8138ceec08cc35b981030d6d004bfad9743c7afd84260fa63f",
     "changeType": "UPSERT",
     "aspectName": "querySubjects",
     "aspect": {
@@ -105,33 +125,13 @@
 },
 {
     "entityType": "query",
-    "entityUrn": "urn:li:query:02e2ec36678bea2a8c4c855fed5255d087cfeb2710d326e95fd9b48a9c4fc0ae",
+    "entityUrn": "urn:li:query:6ed1d12fbf2ccc8138ceec08cc35b981030d6d004bfad9743c7afd84260fa63f",
     "changeType": "UPSERT",
     "aspectName": "dataPlatformInstance",
     "aspect": {
         "json": {
             "platform": "urn:li:dataPlatform:redshift"
         }
     }
-},
-{
-    "entityType": "dataset",
-    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:redshift,dev.public.foo,PROD)",
-    "changeType": "UPSERT",
-    "aspectName": "operation",
-    "aspect": {
-        "json": {
-            "timestampMillis": 1707182625000,
-            "partitionSpec": {
-                "type": "FULL_TABLE",
-                "partition": "FULL_TABLE_SNAPSHOT"
-            },
-            "operationType": "INSERT",
-            "customProperties": {
-                "query_urn": "urn:li:query:02e2ec36678bea2a8c4c855fed5255d087cfeb2710d326e95fd9b48a9c4fc0ae"
-            },
-            "lastUpdatedTimestamp": 20000
-        }
-    }
 }
 ]