acryl-datahub 1.1.0.5rc7__py3-none-any.whl → 1.1.0.5rc9__py3-none-any.whl

datahub/ingestion/source/snowflake/snowflake_config.py

@@ -1,6 +1,7 @@
 import logging
 from collections import defaultdict
 from dataclasses import dataclass
+from enum import Enum
 from typing import Dict, List, Optional, Set
 
 import pydantic
@@ -53,6 +54,11 @@ DEFAULT_TEMP_TABLES_PATTERNS = [
 ]
 
 
+class QueryDedupStrategyType(Enum):
+    STANDARD = "STANDARD"
+    NONE = "NONE"
+
+
 class TagOption(StrEnum):
     with_lineage = "with_lineage"
     without_lineage = "without_lineage"
@@ -248,6 +254,11 @@ class SnowflakeV2Config(
         "This is useful if you have a large number of schemas and want to avoid bulk fetching the schema for each table/view.",
     )
 
+    query_dedup_strategy: QueryDedupStrategyType = Field(
+        default=QueryDedupStrategyType.STANDARD,
+        description=f"Experimental: Choose the strategy for query deduplication (default value is appropriate for most use-cases; make sure you understand performance implications before changing it). Allowed values are: {', '.join([s.name for s in QueryDedupStrategyType])}",
+    )
+
     _check_role_grants_removed = pydantic_removed_field("check_role_grants")
     _provision_role_removed = pydantic_removed_field("provision_role")
 

datahub/ingestion/source/snowflake/snowflake_queries.py

@@ -28,6 +28,7 @@ from datahub.ingestion.graph.client import DataHubGraph
 from datahub.ingestion.source.snowflake.constants import SnowflakeObjectDomain
 from datahub.ingestion.source.snowflake.snowflake_config import (
     DEFAULT_TEMP_TABLES_PATTERNS,
+    QueryDedupStrategyType,
     SnowflakeFilterConfig,
     SnowflakeIdentifierConfig,
 )
@@ -118,6 +119,8 @@ class SnowflakeQueriesExtractorConfig(ConfigModel):
     include_query_usage_statistics: bool = True
     include_operations: bool = True
 
+    query_dedup_strategy: QueryDedupStrategyType = QueryDedupStrategyType.STANDARD
+
 
 class SnowflakeQueriesSourceConfig(
     SnowflakeQueriesExtractorConfig, SnowflakeIdentifierConfig, SnowflakeFilterConfig
@@ -374,12 +377,13 @@ class SnowflakeQueriesExtractor(SnowflakeStructuredReportMixin, Closeable):
     ) -> Iterable[
         Union[PreparsedQuery, TableRename, TableSwap, ObservedQuery, StoredProcCall]
     ]:
-        query_log_query = _build_enriched_query_log_query(
+        query_log_query = QueryLogQueryBuilder(
             start_time=self.config.window.start_time,
             end_time=self.config.window.end_time,
             bucket_duration=self.config.window.bucket_duration,
             deny_usernames=self.config.pushdown_deny_usernames,
-        )
+            dedup_strategy=self.config.query_dedup_strategy,
+        ).build_enriched_query_log_query()
 
         with self.structured_reporter.report_exc(
             "Error fetching query log from Snowflake"
@@ -710,63 +714,101 @@ class SnowflakeQueriesSource(Source):
         super().close()
 
 
-# Make sure we don't try to generate too much info for a single query.
-_MAX_TABLES_PER_QUERY = 20
-
-
-def _build_enriched_query_log_query(
-    start_time: datetime,
-    end_time: datetime,
-    bucket_duration: BucketDuration,
-    deny_usernames: Optional[List[str]],
-) -> str:
-    start_time_millis = int(start_time.timestamp() * 1000)
-    end_time_millis = int(end_time.timestamp() * 1000)
-
-    users_filter = "TRUE"
-    if deny_usernames:
-        user_not_in = ",".join(f"'{user.upper()}'" for user in deny_usernames)
-        users_filter = f"user_name NOT IN ({user_not_in})"
-
-    time_bucket_size = bucket_duration.value
-    assert time_bucket_size in ("HOUR", "DAY", "MONTH")
+class QueryLogQueryBuilder:
+    def __init__(
+        self,
+        start_time: datetime,
+        end_time: datetime,
+        bucket_duration: BucketDuration,
+        deny_usernames: Optional[List[str]],
+        max_tables_per_query: int = 20,
+        dedup_strategy: QueryDedupStrategyType = QueryDedupStrategyType.STANDARD,
+    ):
+        self.start_time = start_time
+        self.end_time = end_time
+        self.start_time_millis = int(start_time.timestamp() * 1000)
+        self.end_time_millis = int(end_time.timestamp() * 1000)
+        self.max_tables_per_query = max_tables_per_query
+        self.dedup_strategy = dedup_strategy
+
+        self.users_filter = "TRUE"
+        if deny_usernames:
+            user_not_in = ",".join(f"'{user.upper()}'" for user in deny_usernames)
+            self.users_filter = f"user_name NOT IN ({user_not_in})"
+
+        self.time_bucket_size = bucket_duration.value
+        assert self.time_bucket_size in ("HOUR", "DAY", "MONTH")
+
+    def _query_fingerprinted_queries(self):
+        if self.dedup_strategy == QueryDedupStrategyType.STANDARD:
+            secondary_fingerprint_sql = """
+    CASE 
+        WHEN CONTAINS(query_history.query_text, '-- Hex query metadata:')
+        -- Extract project id and hash it
+        THEN CAST(HASH(
+            REGEXP_SUBSTR(query_history.query_text, '"project_id"\\\\s*:\\\\s*"([^"]+)"', 1, 1, 'e', 1),
+            REGEXP_SUBSTR(query_history.query_text, '"context"\\\\s*:\\\\s*"([^"]+)"', 1, 1, 'e', 1)
+        ) AS VARCHAR)
+        ELSE NULL 
+    END"""
+        elif self.dedup_strategy == QueryDedupStrategyType.NONE:
+            secondary_fingerprint_sql = "NULL"
+        else:
+            raise NotImplementedError(
+                f"Strategy {self.dedup_strategy} is not implemented by the QueryLogQueryBuilder"
+            )
+        return f"""
+SELECT *,
+    -- TODO: Generate better fingerprints for each query by pushing down regex logic.
+    query_history.query_parameterized_hash as query_fingerprint,
+    -- Optional and additional hash to be used for query deduplication and final query identity
+    {secondary_fingerprint_sql} as query_secondary_fingerprint
+FROM
+    snowflake.account_usage.query_history
+WHERE
+    query_history.start_time >= to_timestamp_ltz({self.start_time_millis}, 3) -- {self.start_time.isoformat()}
+    AND query_history.start_time < to_timestamp_ltz({self.end_time_millis}, 3) -- {self.end_time.isoformat()}
+    AND execution_status = 'SUCCESS'
+    AND {self.users_filter}"""
+
+    def _query_deduplicated_queries(self):
+        if self.dedup_strategy == QueryDedupStrategyType.STANDARD:
+            return f"""
+SELECT
+    *,
+    DATE_TRUNC(
+        {self.time_bucket_size},
+        CONVERT_TIMEZONE('UTC', start_time)
+    ) AS bucket_start_time,
+    COUNT(*) OVER (PARTITION BY bucket_start_time, query_fingerprint, query_secondary_fingerprint) AS query_count,
+FROM
+    fingerprinted_queries
+QUALIFY
+    ROW_NUMBER() OVER (PARTITION BY bucket_start_time, query_fingerprint, query_secondary_fingerprint ORDER BY start_time DESC) = 1"""
+        elif self.dedup_strategy == QueryDedupStrategyType.NONE:
+            return f"""
+SELECT
+    *,
+    DATE_TRUNC(
+        {self.time_bucket_size},
+        CONVERT_TIMEZONE('UTC', start_time)
+    ) AS bucket_start_time,
+    1 AS query_count,
+FROM
+            fingerprinted_queries"""
+        else:
+            raise NotImplementedError(
+                f"Strategy {self.dedup_strategy} is not implemented by the QueryLogQueryBuilder"
+            )
 
-    return f"""\
+    def build_enriched_query_log_query(self) -> str:
+        return f"""\
 WITH
 fingerprinted_queries as (
-    SELECT *,
-        -- TODO: Generate better fingerprints for each query by pushing down regex logic.
-        query_history.query_parameterized_hash as query_fingerprint,
-        -- Optional and additional hash to be used for query deduplication and final query identity
-        CASE 
-            WHEN CONTAINS(query_history.query_text, '-- Hex query metadata:')
-            -- Extract project id and hash it
-            THEN CAST(HASH(
-                REGEXP_SUBSTR(query_history.query_text, '"project_id"\\\\s*:\\\\s*"([^"]+)"', 1, 1, 'e', 1),
-                REGEXP_SUBSTR(query_history.query_text, '"context"\\\\s*:\\\\s*"([^"]+)"', 1, 1, 'e', 1)
-            ) AS VARCHAR)
-            ELSE NULL 
-        END as query_secondary_fingerprint
-    FROM
-        snowflake.account_usage.query_history
-    WHERE
-        query_history.start_time >= to_timestamp_ltz({start_time_millis}, 3) -- {start_time.isoformat()}
-        AND query_history.start_time < to_timestamp_ltz({end_time_millis}, 3) -- {end_time.isoformat()}
-        AND execution_status = 'SUCCESS'
-        AND {users_filter}
+{self._query_fingerprinted_queries()}
 )
 , deduplicated_queries as (
-    SELECT
-        *,
-        DATE_TRUNC(
-            {time_bucket_size},
-            CONVERT_TIMEZONE('UTC', start_time)
-        ) AS bucket_start_time,
-        COUNT(*) OVER (PARTITION BY bucket_start_time, query_fingerprint, query_secondary_fingerprint) AS query_count,
-    FROM
-        fingerprinted_queries
-    QUALIFY
-        ROW_NUMBER() OVER (PARTITION BY bucket_start_time, query_fingerprint, query_secondary_fingerprint ORDER BY start_time DESC) = 1
+{self._query_deduplicated_queries()}
 )
 , raw_access_history AS (
     SELECT
@@ -780,9 +822,9 @@ fingerprinted_queries as (
     FROM
         snowflake.account_usage.access_history
     WHERE
-        query_start_time >= to_timestamp_ltz({start_time_millis}, 3) -- {start_time.isoformat()}
-        AND query_start_time < to_timestamp_ltz({end_time_millis}, 3) -- {end_time.isoformat()}
-        AND {users_filter}
+        query_start_time >= to_timestamp_ltz({self.start_time_millis}, 3) -- {self.start_time.isoformat()}
+        AND query_start_time < to_timestamp_ltz({self.end_time_millis}, 3) -- {self.end_time.isoformat()}
+        AND {self.users_filter}
         AND query_id IN (
             SELECT query_id FROM deduplicated_queries
         )
@@ -795,7 +837,7 @@ fingerprinted_queries as (
         query_start_time,
         ARRAY_SLICE(
             FILTER(direct_objects_accessed, o -> o:objectDomain IN {SnowflakeQuery.ACCESS_HISTORY_TABLE_VIEW_DOMAINS_FILTER}),
-            0, {_MAX_TABLES_PER_QUERY}
+            0, {self.max_tables_per_query}
         ) as direct_objects_accessed,
         -- TODO: Drop the columns.baseSources subfield.
         FILTER(objects_modified, o -> o:objectDomain IN {SnowflakeQuery.ACCESS_HISTORY_TABLE_VIEW_DOMAINS_FILTER}) as objects_modified,

datahub/ingestion/source/snowflake/snowflake_v2.py

@@ -32,6 +32,7 @@ from datahub.ingestion.api.source import (
 )
 from datahub.ingestion.api.source_helpers import auto_workunit
 from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.ingestion.source.common.subtypes import SourceCapabilityModifier
 from datahub.ingestion.source.snowflake.constants import (
     GENERIC_PERMISSION_ERROR_KEY,
     SnowflakeEdition,
@@ -97,7 +98,14 @@ logger: logging.Logger = logging.getLogger(__name__)
 @support_status(SupportStatus.CERTIFIED)
 @capability(SourceCapability.PLATFORM_INSTANCE, "Enabled by default")
 @capability(SourceCapability.DOMAINS, "Supported via the `domain` config field")
-@capability(SourceCapability.CONTAINERS, "Enabled by default")
+@capability(
+    SourceCapability.CONTAINERS,
+    "Enabled by default",
+    subtype_modifier=[
+        SourceCapabilityModifier.DATABASE,
+        SourceCapabilityModifier.SCHEMA,
+    ],
+)
 @capability(SourceCapability.SCHEMA_METADATA, "Enabled by default")
 @capability(
     SourceCapability.DATA_PROFILING,
@@ -577,6 +585,7 @@ class SnowflakeV2Source(
 
                 queries_extractor = SnowflakeQueriesExtractor(
                     connection=self.connection,
+                    # TODO: this should be its own section in main recipe
                     config=SnowflakeQueriesExtractorConfig(
                         window=BaseTimeWindowConfig(
                             start_time=self.config.start_time,
@@ -591,6 +600,7 @@ class SnowflakeV2Source(
                         include_query_usage_statistics=self.config.include_query_usage_statistics,
                         user_email_pattern=self.config.user_email_pattern,
                         pushdown_deny_usernames=self.config.pushdown_deny_usernames,
+                        query_dedup_strategy=self.config.query_dedup_strategy,
                     ),
                     structured_report=self.report,
                     filters=self.filters,

datahub/ingestion/source/snowflake/stored_proc_lineage.py

@@ -65,7 +65,7 @@ class StoredProcLineageTracker(Closeable):
         # { root_query_id -> StoredProcExecutionLineage }
         self._stored_proc_execution_lineage: FileBackedDict[
             StoredProcExecutionLineage
-        ] = FileBackedDict(shared_connection)
+        ] = FileBackedDict(shared_connection, tablename="stored_proc_lineage")
 
     def add_stored_proc_call(self, call: StoredProcCall) -> None:
         """Add a stored procedure call to track."""

datahub/ingestion/source/sql/hive_metastore.py

@@ -52,7 +52,6 @@ from datahub.metadata.com.linkedin.pegasus2avro.metadata.snapshot import Dataset
 from datahub.metadata.com.linkedin.pegasus2avro.mxe import MetadataChangeEvent
 from datahub.metadata.com.linkedin.pegasus2avro.schema import SchemaField
 from datahub.metadata.schema_classes import (
-    ChangeTypeClass,
     DatasetPropertiesClass,
     SubTypesClass,
     ViewPropertiesClass,
@@ -601,10 +600,7 @@ class HiveMetastoreSource(SQLAlchemySource):
                 yield dpi_aspect
 
             yield MetadataChangeProposalWrapper(
-                entityType="dataset",
-                changeType=ChangeTypeClass.UPSERT,
                 entityUrn=dataset_urn,
-                aspectName="subTypes",
                 aspect=SubTypesClass(typeNames=[self.table_subtype]),
             ).as_workunit()
 
@@ -810,10 +806,7 @@ class HiveMetastoreSource(SQLAlchemySource):
 
             # Add views subtype
             yield MetadataChangeProposalWrapper(
-                entityType="dataset",
-                changeType=ChangeTypeClass.UPSERT,
                 entityUrn=dataset_urn,
-                aspectName="subTypes",
                 aspect=SubTypesClass(typeNames=[self.view_subtype]),
             ).as_workunit()
 
@@ -824,10 +817,7 @@ class HiveMetastoreSource(SQLAlchemySource):
                 viewLogic=dataset.view_definition if dataset.view_definition else "",
             )
             yield MetadataChangeProposalWrapper(
-                entityType="dataset",
-                changeType=ChangeTypeClass.UPSERT,
                 entityUrn=dataset_urn,
-                aspectName="viewProperties",
                 aspect=view_properties_aspect,
             ).as_workunit()
 

datahub/ingestion/source/sql/sql_common.py

@@ -292,6 +292,10 @@ class ProfileMetadata:
     SourceCapability.CONTAINERS,
     "Enabled by default",
     supported=True,
+    subtype_modifier=[
+        SourceCapabilityModifier.DATABASE,
+        SourceCapabilityModifier.SCHEMA,
+    ],
 )
 @capability(
     SourceCapability.DESCRIPTIONS,
@@ -589,6 +593,10 @@ class SQLAlchemySource(StatefulIngestionSourceBase, TestableSource):
                 )
 
         # Generate workunit for aggregated SQL parsing results
+        yield from self._generate_aggregator_workunits()
+
+    def _generate_aggregator_workunits(self) -> Iterable[MetadataWorkUnit]:
+        """Generate work units from SQL parsing aggregator. Can be overridden by subclasses."""
         for mcp in self.aggregator.gen_metadata():
             yield mcp.as_workunit()