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

datahub/ingestion/source/snowflake/snowflake_queries.py

@@ -44,6 +44,11 @@ from datahub.ingestion.source.snowflake.snowflake_utils import (
     SnowflakeIdentifierBuilder,
     SnowflakeStructuredReportMixin,
 )
+from datahub.ingestion.source.snowflake.stored_proc_lineage import (
+    StoredProcCall,
+    StoredProcLineageReport,
+    StoredProcLineageTracker,
+)
 from datahub.ingestion.source.usage.usage_common import BaseUsageConfig
 from datahub.metadata.urns import CorpUserUrn
 from datahub.sql_parsing.schema_resolver import SchemaResolver
@@ -130,6 +135,7 @@ class SnowflakeQueriesExtractorReport(Report):
     aggregator_generate_timer: PerfTimer = dataclasses.field(default_factory=PerfTimer)
 
     sql_aggregator: Optional[SqlAggregatorReport] = None
+    stored_proc_lineage: Optional[StoredProcLineageReport] = None
 
     num_ddl_queries_dropped: int = 0
     num_stream_queries_observed: int = 0
@@ -261,6 +267,7 @@ class SnowflakeQueriesExtractor(SnowflakeStructuredReportMixin, Closeable):
                 TableRename,
                 TableSwap,
                 ObservedQuery,
+                StoredProcCall,
             ]
         ] = self._exit_stack.enter_context(FileBackedList(shared_connection))
 
@@ -277,12 +284,34 @@ class SnowflakeQueriesExtractor(SnowflakeStructuredReportMixin, Closeable):
                 for entry in self.fetch_query_log(users):
                     queries.append(entry)
 
+        stored_proc_tracker: StoredProcLineageTracker = self._exit_stack.enter_context(
+            StoredProcLineageTracker(
+                platform=self.identifiers.platform,
+                shared_connection=shared_connection,
+            )
+        )
+        self.report.stored_proc_lineage = stored_proc_tracker.report
+
         with self.report.audit_log_load_timer:
             for i, query in enumerate(queries):
                 if i % 1000 == 0:
                     logger.info(f"Added {i} query log entries to SQL aggregator")
 
-                self.aggregator.add(query)
+                if isinstance(query, StoredProcCall):
+                    stored_proc_tracker.add_stored_proc_call(query)
+                    continue
+
+                if not (
+                    isinstance(query, PreparsedQuery)
+                    and stored_proc_tracker.add_related_query(query)
+                ):
+                    # Only add to aggregator if it's not part of a stored procedure.
+                    self.aggregator.add(query)
+
+            # Generate and add stored procedure lineage entries.
+            for lineage_entry in stored_proc_tracker.build_merged_lineage_entries():
+                # TODO: Make this the lowest priority lineage - so that it doesn't override other lineage entries.
+                self.aggregator.add(lineage_entry)
 
         with self.report.aggregator_generate_timer:
             yield from auto_workunit(self.aggregator.gen_metadata())
@@ -342,7 +371,9 @@ class SnowflakeQueriesExtractor(SnowflakeStructuredReportMixin, Closeable):
 
     def fetch_query_log(
         self, users: UsersMapping
-    ) -> Iterable[Union[PreparsedQuery, TableRename, TableSwap, ObservedQuery]]:
+    ) -> Iterable[
+        Union[PreparsedQuery, TableRename, TableSwap, ObservedQuery, StoredProcCall]
+    ]:
         query_log_query = _build_enriched_query_log_query(
             start_time=self.config.window.start_time,
             end_time=self.config.window.end_time,
@@ -382,7 +413,9 @@ class SnowflakeQueriesExtractor(SnowflakeStructuredReportMixin, Closeable):
 
     def _parse_audit_log_row(
         self, row: Dict[str, Any], users: UsersMapping
-    ) -> Optional[Union[TableRename, TableSwap, PreparsedQuery, ObservedQuery]]:
+    ) -> Optional[
+        Union[TableRename, TableSwap, PreparsedQuery, ObservedQuery, StoredProcCall]
+    ]:
         json_fields = {
             "DIRECT_OBJECTS_ACCESSED",
             "OBJECTS_MODIFIED",
@@ -482,6 +515,17 @@ class SnowflakeQueriesExtractor(SnowflakeStructuredReportMixin, Closeable):
                 extra_info=extra_info,
             )
 
+        if snowflake_query_type == "CALL" and res["root_query_id"] is None:
+            return StoredProcCall(
+                # This is the top-level query ID that other entries will reference.
+                snowflake_root_query_id=res["query_id"],
+                query_text=query_text,
+                timestamp=timestamp,
+                user=user,
+                default_db=res["default_db"],
+                default_schema=res["default_schema"],
+            )
+
         upstreams = []
         column_usage = {}
 

datahub/ingestion/source/snowflake/snowflake_usage_v2.py

@@ -231,7 +231,10 @@ class SnowflakeUsageExtractor(SnowflakeCommonMixin, Closeable):
 
         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:
+                with (
+                    fetch_timer.pause(),
+                    self.report.usage_aggregation.result_skip_timer as skip_timer,
+                ):
                     if results.rownumber is not None and results.rownumber % 1000 == 0:
                         logger.debug(f"Processing usage row number {results.rownumber}")
                         logger.debug(self.report.usage_aggregation.as_string())
@@ -255,7 +258,10 @@ class SnowflakeUsageExtractor(SnowflakeCommonMixin, Closeable):
                             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:
+                    with (
+                        skip_timer.pause(),
+                        self.report.usage_aggregation.result_map_timer as map_timer,
+                    ):
                         wu = self.build_usage_statistics_for_dataset(
                             dataset_identifier, row
                         )

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)
+
+    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

@@ -34,6 +34,9 @@ from datahub.ingestion.source.common.subtypes import (
     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,
@@ -47,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
@@ -284,6 +292,11 @@ 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",
@@ -496,23 +509,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}",
@@ -538,6 +566,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,
@@ -563,7 +641,14 @@ class AthenaSource(SQLAlchemySource):
                 partition_keys is not None and column["name"] in partition_keys
             ),
         )
-
+        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(