acryl-datahub 1.1.1rc4__py3-none-any.whl → 1.2.0.1__py3-none-any.whl

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,
@@ -118,7 +126,7 @@ logger: logging.Logger = logging.getLogger(__name__)
 )
 @capability(
     SourceCapability.DELETION_DETECTION,
-    "Optionally enabled via `stateful_ingestion.remove_stale_metadata`",
+    "Enabled by default via stateful ingestion",
     supported=True,
 )
 @capability(
@@ -131,6 +139,7 @@ logger: logging.Logger = logging.getLogger(__name__)
     "Optionally enabled via `classification.enabled`",
     supported=True,
 )
+@capability(SourceCapability.TEST_CONNECTION, "Enabled by default")
 class SnowflakeV2Source(
     SnowflakeCommonMixin,
     StatefulIngestionSourceBase,
@@ -311,6 +320,7 @@ class SnowflakeV2Source(
                 SourceCapability.PLATFORM_INSTANCE,
                 SourceCapability.DOMAINS,
                 SourceCapability.DELETION_DETECTION,
+                SourceCapability.TEST_CONNECTION,
             )
         ]
 
@@ -575,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,
@@ -589,6 +600,9 @@ 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,
+                        push_down_database_pattern_access_history=self.config.push_down_database_pattern_access_history,
+                        additional_database_names_allowlist=self.config.additional_database_names_allowlist,
                     ),
                     structured_report=self.report,
                     filters=self.filters,

datahub/ingestion/source/snowflake/stored_proc_lineage.py

@@ -0,0 +1,143 @@
+import dataclasses
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Any, Iterable, List, Optional
+
+from datahub.ingestion.api.closeable import Closeable
+from datahub.metadata.urns import CorpUserUrn
+from datahub.sql_parsing.sql_parsing_aggregator import (
+    PreparsedQuery,
+    UrnStr,
+)
+from datahub.sql_parsing.sqlglot_utils import get_query_fingerprint
+from datahub.utilities.file_backed_collections import FileBackedDict
+
+
+@dataclasses.dataclass
+class StoredProcCall:
+    snowflake_root_query_id: str
+
+    # Query text will typically be something like:
+    # "CALL SALES_FORECASTING.CUSTOMER_ANALYSIS_PROC();"
+    query_text: str
+
+    timestamp: datetime
+    user: CorpUserUrn
+    default_db: str
+    default_schema: str
+
+
+@dataclass
+class StoredProcExecutionLineage:
+    call: StoredProcCall
+
+    inputs: List[UrnStr]
+    outputs: List[UrnStr]
+
+
+@dataclass
+class StoredProcLineageReport:
+    num_stored_proc_calls: int = 0
+    num_related_queries: int = 0
+    num_related_queries_without_proc_call: int = 0
+
+    # Incremented at generation/build time.
+    num_stored_proc_lineage_entries: int = 0
+    num_stored_proc_calls_with_no_inputs: int = 0
+    num_stored_proc_calls_with_no_outputs: int = 0
+
+
+class StoredProcLineageTracker(Closeable):
+    """
+    Tracks table-level lineage for Snowflake stored procedures.
+
+    Stored procedures in Snowflake trigger multiple SQL queries during execution.
+    Snowflake assigns each stored procedure call a unique query_id and uses this as the
+    root_query_id for all subsequent queries executed within that procedure. This allows
+    us to trace which queries belong to a specific stored procedure execution and build
+    table-level lineage by aggregating inputs/outputs from all related queries.
+    """
+
+    def __init__(self, platform: str, shared_connection: Optional[Any] = None):
+        self.platform = platform
+        self.report = StoredProcLineageReport()
+
+        # { root_query_id -> StoredProcExecutionLineage }
+        self._stored_proc_execution_lineage: FileBackedDict[
+            StoredProcExecutionLineage
+        ] = FileBackedDict(shared_connection, tablename="stored_proc_lineage")
+
+    def add_stored_proc_call(self, call: StoredProcCall) -> None:
+        """Add a stored procedure call to track."""
+        self._stored_proc_execution_lineage[call.snowflake_root_query_id] = (
+            StoredProcExecutionLineage(
+                call=call,
+                # Will be populated by subsequent queries.
+                inputs=[],
+                outputs=[],
+            )
+        )
+        self.report.num_stored_proc_calls += 1
+
+    def add_related_query(self, query: PreparsedQuery) -> bool:
+        """Add a query that might be related to a stored procedure execution.
+
+        Returns True if the query was added to a stored procedure execution, False otherwise.
+        """
+        snowflake_root_query_id = (query.extra_info or {}).get(
+            "snowflake_root_query_id"
+        )
+
+        if snowflake_root_query_id:
+            if snowflake_root_query_id not in self._stored_proc_execution_lineage:
+                self.report.num_related_queries_without_proc_call += 1
+                return False
+
+            stored_proc_execution = self._stored_proc_execution_lineage.for_mutation(
+                snowflake_root_query_id
+            )
+            stored_proc_execution.inputs.extend(query.upstreams)
+            if query.downstream is not None:
+                stored_proc_execution.outputs.append(query.downstream)
+            self.report.num_related_queries += 1
+            return True
+
+        return False
+
+    def build_merged_lineage_entries(self) -> Iterable[PreparsedQuery]:
+        # For stored procedures, we can only get table-level lineage from the audit log.
+        # We represent these as PreparsedQuery objects for now. Eventually we'll want to
+        # create dataJobInputOutput lineage instead.
+
+        for stored_proc_execution in self._stored_proc_execution_lineage.values():
+            if not stored_proc_execution.inputs:
+                self.report.num_stored_proc_calls_with_no_inputs += 1
+                continue
+
+            if not stored_proc_execution.outputs:
+                self.report.num_stored_proc_calls_with_no_outputs += 1
+                # Still continue to generate lineage for cases where we have inputs but no outputs
+
+            for downstream in stored_proc_execution.outputs:
+                stored_proc_query_id = get_query_fingerprint(
+                    stored_proc_execution.call.query_text,
+                    self.platform,
+                    fast=True,
+                    secondary_id=downstream,
+                )
+
+                lineage_entry = PreparsedQuery(
+                    query_id=stored_proc_query_id,
+                    query_text=stored_proc_execution.call.query_text,
+                    upstreams=stored_proc_execution.inputs,
+                    downstream=downstream,
+                    query_count=0,
+                    user=stored_proc_execution.call.user,
+                    timestamp=stored_proc_execution.call.timestamp,
+                )
+
+                self.report.num_stored_proc_lineage_entries += 1
+                yield lineage_entry
+
+    def close(self) -> None:
+        self._stored_proc_execution_lineage.close()

datahub/ingestion/source/sql/athena.py

@@ -29,8 +29,14 @@ from datahub.ingestion.api.decorators import (
 from datahub.ingestion.api.source import StructuredLogLevel
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.aws.s3_util import make_s3_urn
-from datahub.ingestion.source.common.subtypes import DatasetContainerSubTypes
+from datahub.ingestion.source.common.subtypes import (
+    DatasetContainerSubTypes,
+    SourceCapabilityModifier,
+)
 from datahub.ingestion.source.ge_profiling_config import GEProfilingConfig
+from datahub.ingestion.source.sql.athena_properties_extractor import (
+    AthenaPropertiesExtractor,
+)
 from datahub.ingestion.source.sql.sql_common import (
     SQLAlchemySource,
     register_custom_type,
@@ -44,12 +50,17 @@ from datahub.ingestion.source.sql.sql_utils import (
 )
 from datahub.ingestion.source.sql.sqlalchemy_uri import make_sqlalchemy_uri
 from datahub.metadata.com.linkedin.pegasus2avro.schema import SchemaField
-from datahub.metadata.schema_classes import MapTypeClass, RecordTypeClass
+from datahub.metadata.schema_classes import (
+    ArrayTypeClass,
+    MapTypeClass,
+    RecordTypeClass,
+)
 from datahub.utilities.hive_schema_to_avro import get_avro_schema_for_hive_column
 from datahub.utilities.sqlalchemy_type_converter import (
     MapType,
     get_schema_fields_for_sqlalchemy_column,
 )
+from datahub.utilities.urns.field_paths import get_simple_field_path_from_v2_field_path
 
 try:
     from typing_extensions import override
@@ -281,12 +292,22 @@ class AthenaConfig(SQLCommonConfig):
         description="Extract partitions for tables. Partition extraction needs to run a query (`select * from table$partitions`) on the table. Disable this if you don't want to grant select permission.",
     )
 
+    extract_partitions_using_create_statements: bool = pydantic.Field(
+        default=False,
+        description="Extract partitions using the `SHOW CREATE TABLE` statement instead of querying the table's partitions directly. This needs to be enabled to extract Iceberg partitions. If extraction fails it falls back to the default partition extraction. This is experimental.",
+    )
+
     _s3_staging_dir_population = pydantic_renamed_field(
         old_name="s3_staging_dir",
         new_name="query_result_location",
         print_warning=True,
     )
 
+    emit_schema_fieldpaths_as_v1: bool = pydantic.Field(
+        default=False,
+        description="Convert simple field paths to DataHub field path v1 format. Simple column paths are those that do not contain any nested fields.",
+    )
+
     profiling: AthenaProfilingConfig = AthenaProfilingConfig()
 
     def get_sql_alchemy_url(self):
@@ -321,8 +342,18 @@ class Partitionitem:
 @capability(
     SourceCapability.DATA_PROFILING,
     "Optionally enabled via configuration. Profiling uses sql queries on whole table which can be expensive operation.",
+    subtype_modifier=[SourceCapabilityModifier.TABLE],
+)
+@capability(
+    SourceCapability.LINEAGE_COARSE,
+    "Supported for S3 tables",
+    subtype_modifier=[SourceCapabilityModifier.TABLE],
+)
+@capability(
+    SourceCapability.LINEAGE_FINE,
+    "Supported for S3 tables",
+    subtype_modifier=[SourceCapabilityModifier.TABLE],
 )
-@capability(SourceCapability.LINEAGE_COARSE, "Supported for S3 tables")
 @capability(SourceCapability.DESCRIPTIONS, "Enabled by default")
 class AthenaSource(SQLAlchemySource):
     """
@@ -483,23 +514,38 @@ class AthenaSource(SQLAlchemySource):
     def get_partitions(
         self, inspector: Inspector, schema: str, table: str
     ) -> Optional[List[str]]:
-        if not self.config.extract_partitions:
+        if (
+            not self.config.extract_partitions
+            and not self.config.extract_partitions_using_create_statements
+        ):
             return None
 
         if not self.cursor:
             return None
 
-        metadata: AthenaTableMetadata = self.cursor.get_table_metadata(
-            table_name=table, schema_name=schema
-        )
+        if self.config.extract_partitions_using_create_statements:
+            try:
+                partitions = self._get_partitions_create_table(schema, table)
+            except Exception as e:
+                logger.warning(
+                    f"Failed to get partitions from create table statement for {schema}.{table} because of {e}. Falling back to SQLAlchemy.",
+                    exc_info=True,
+                )
+
+                # If we can't get create table statement, we fall back to SQLAlchemy
+                partitions = self._get_partitions_sqlalchemy(schema, table)
+        else:
+            partitions = self._get_partitions_sqlalchemy(schema, table)
 
-        partitions = []
-        for key in metadata.partition_keys:
-            if key.name:
-                partitions.append(key.name)
         if not partitions:
             return []
 
+        if (
+            not self.config.profiling.enabled
+            or not self.config.profiling.partition_profiling_enabled
+        ):
+            return partitions
+
         with self.report.report_exc(
             message="Failed to extract partition details",
             context=f"{schema}.{table}",
@@ -525,6 +571,56 @@ class AthenaSource(SQLAlchemySource):
 
         return partitions
 
+    def _get_partitions_create_table(self, schema: str, table: str) -> List[str]:
+        assert self.cursor
+        try:
+            res = self.cursor.execute(f"SHOW CREATE TABLE `{schema}`.`{table}`")
+        except Exception as e:
+            # Athena does not support SHOW CREATE TABLE for views
+            # and will throw an error. We need to handle this case
+            # and caller needs to fallback to sqlalchemy's get partitions call.
+            logger.debug(
+                f"Failed to get table properties for {schema}.{table}: {e}",
+                exc_info=True,
+            )
+            raise e
+        rows = res.fetchall()
+
+        # Concatenate all rows into a single string with newlines
+        create_table_statement = "\n".join(row[0] for row in rows)
+
+        try:
+            athena_table_info = AthenaPropertiesExtractor.get_table_properties(
+                create_table_statement
+            )
+        except Exception as e:
+            logger.debug(
+                f"Failed to parse table properties for {schema}.{table}: {e} and statement: {create_table_statement}",
+                exc_info=True,
+            )
+            raise e
+
+        partitions = []
+        if (
+            athena_table_info.partition_info
+            and athena_table_info.partition_info.simple_columns
+        ):
+            partitions = [
+                ci.name for ci in athena_table_info.partition_info.simple_columns
+            ]
+        return partitions
+
+    def _get_partitions_sqlalchemy(self, schema: str, table: str) -> List[str]:
+        assert self.cursor
+        metadata: AthenaTableMetadata = self.cursor.get_table_metadata(
+            table_name=table, schema_name=schema
+        )
+        partitions = []
+        for key in metadata.partition_keys:
+            if key.name:
+                partitions.append(key.name)
+        return partitions
+
     # Overwrite to modify the creation of schema fields
     def get_schema_fields_for_column(
         self,
@@ -551,6 +647,18 @@ class AthenaSource(SQLAlchemySource):
             ),
         )
 
+        # Keeping it as individual check to make it more explicit and easier to understand
+        if not self.config.emit_schema_fieldpaths_as_v1:
+            return fields
+
+        if isinstance(
+            fields[0].type.type, (RecordTypeClass, MapTypeClass, ArrayTypeClass)
+        ):
+            return fields
+        else:
+            fields[0].fieldPath = get_simple_field_path_from_v2_field_path(
+                fields[0].fieldPath
+            )
         return fields
 
     def generate_partition_profiler_query(