openmetadata-ingestion 1.9.12.0__py3-none-any.whl → 1.9.13.0__py3-none-any.whl

metadata/ingestion/source/database/snowflake/usage.py

@@ -11,11 +11,22 @@
 """
 Snowflake usage module
 """
+import traceback
+from datetime import timedelta
+from typing import Iterable
+
+from metadata.generated.schema.type.basic import DateTime
+from metadata.generated.schema.type.tableQuery import TableQueries, TableQuery
+from metadata.ingestion.lineage.masker import mask_query
 from metadata.ingestion.source.database.snowflake.queries import SNOWFLAKE_SQL_STATEMENT
 from metadata.ingestion.source.database.snowflake.query_parser import (
+    SNOWFLAKE_QUERY_BATCH_SIZE,
     SnowflakeQueryParserSource,
 )
 from metadata.ingestion.source.database.usage_source import UsageSource
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
 
 
 class SnowflakeUsageSource(SnowflakeQueryParserSource, UsageSource):
@@ -41,3 +52,89 @@ class SnowflakeUsageSource(SnowflakeQueryParserSource, UsageSource):
         "INSERT",
         "MERGE",
     ]
+
+    def yield_table_queries(self) -> Iterable[TableQuery]:
+        """
+        Given an Engine, iterate over the day range and
+        query the results with pagination
+        """
+        daydiff = self.end - self.start
+        for days in range(daydiff.days):
+            logger.info(
+                f"Scanning query logs for {(self.start + timedelta(days=days)).date()} - "
+                f"{(self.start + timedelta(days=days + 1)).date()}"
+            )
+            query = None
+            offset = 0
+            total_fetched = 0
+            max_results = self.source_config.resultLimit
+            try:
+                for engine in self.get_engine():
+                    while total_fetched < max_results:
+                        batch_size = min(
+                            SNOWFLAKE_QUERY_BATCH_SIZE, max_results - total_fetched
+                        )
+                        query = self.get_sql_statement(
+                            start_time=self.start + timedelta(days=days),
+                            end_time=self.start + timedelta(days=days + 1),
+                            offset=offset,
+                            limit=batch_size,
+                        )
+                        with engine.connect() as conn:
+                            rows = conn.execute(query)
+                            queries = []
+                            row_count = 0
+                            for row in rows:
+                                row = dict(row)
+                                row_count += 1
+                                try:
+                                    row.update({k.lower(): v for k, v in row.items()})
+                                    logger.debug(f"Processing row: {row}")
+                                    query_type = row.get("query_type")
+                                    query_text = self.format_query(row["query_text"])
+                                    queries.append(
+                                        TableQuery(
+                                            query=query_text,
+                                            query_type=query_type,
+                                            exclude_usage=self.check_life_cycle_query(
+                                                query_type=query_type,
+                                                query_text=query_text,
+                                            ),
+                                            dialect=self.dialect.value,
+                                            userName=row["user_name"],
+                                            startTime=str(row["start_time"]),
+                                            endTime=str(row["end_time"]),
+                                            analysisDate=DateTime(row["start_time"]),
+                                            aborted=self.get_aborted_status(row),
+                                            databaseName=self.get_database_name(row),
+                                            duration=row.get("duration"),
+                                            serviceName=self.config.serviceName,
+                                            databaseSchema=self.get_schema_name(row),
+                                            cost=row.get("cost"),
+                                        )
+                                    )
+                                except Exception as exc:
+                                    logger.debug(traceback.format_exc())
+                                    logger.warning(
+                                        f"Unexpected exception processing row [{row}]: {exc}"
+                                    )
+                        if queries:
+                            yield TableQueries(queries=queries)
+                        total_fetched += row_count
+                        if row_count < batch_size:
+                            break
+                        offset += batch_size
+                        logger.info(
+                            f"Fetching next page with offset {offset} (fetched {total_fetched}/{max_results}) "
+                            f"for {(self.start + timedelta(days=days)).date()}"
+                        )
+            except Exception as exc:
+                if query:
+                    logger.debug(
+                        (
+                            f"###### USAGE QUERY #######\n{mask_query(query, self.dialect.value) or query}"
+                            "\n##########################"
+                        )
+                    )
+                logger.debug(traceback.format_exc())
+                logger.error(f"Source usage processing error: {exc}")

metadata/ingestion/source/database/sql_column_handler.py

@@ -268,9 +268,17 @@ class SqlColumnHandlerMixin:
 
         table_columns = []
 
-        columns = self._get_columns_internal(
-            schema_name, table_name, db_name, inspector, table_type
-        )
+        try:
+            columns = self._get_columns_internal(
+                schema_name, table_name, db_name, inspector, table_type
+            )
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(
+                f"Unexpected exception getting columns for table [{table_name}] "
+                f"(schema: '{schema_name}', db: '{db_name}'): {exc}"
+            )
+            columns = []
 
         def process_column(column: dict):
             (

metadata/ingestion/source/database/unitycatalog/client.py

@@ -27,7 +27,7 @@ from metadata.ingestion.source.database.unitycatalog.models import (
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
-TABLE_LINEAGE_PATH = "/lineage-tracking/table-lineage/get"
+TABLE_LINEAGE_PATH = "/lineage-tracking/table-lineage"
 COLUMN_LINEAGE_PATH = "/lineage-tracking/column-lineage/get"
 TABLES_PATH = "/unity-catalog/tables"
 

metadata/ingestion/source/database/unitycatalog/lineage.py

@@ -15,6 +15,7 @@ import traceback
 from typing import Iterable, Optional
 
 from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
+from metadata.generated.schema.entity.data.container import ContainerDataModel
 from metadata.generated.schema.entity.data.database import Database
 from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
 from metadata.generated.schema.entity.data.table import Table
@@ -90,6 +91,41 @@ class UnitycatalogLineageSource(Source):
             )
         return cls(config, metadata)
 
+    def _get_data_model_column_fqn(
+        self, data_model_entity: ContainerDataModel, column: str
+    ) -> Optional[str]:
+        if not data_model_entity:
+            return None
+        for entity_column in data_model_entity.columns:
+            if entity_column.displayName.lower() == column.lower():
+                return entity_column.fullyQualifiedName.root
+        return None
+
+    def _get_container_column_lineage(
+        self, data_model_entity: ContainerDataModel, table_entity: Table
+    ) -> Optional[LineageDetails]:
+        try:
+            column_lineage = []
+            for column in table_entity.columns:
+                from_column = self._get_data_model_column_fqn(
+                    data_model_entity=data_model_entity, column=column.name.root
+                )
+                to_column = column.fullyQualifiedName.root
+                if from_column and to_column:
+                    column_lineage.append(
+                        ColumnLineage(fromColumns=[from_column], toColumn=to_column)
+                    )
+            if column_lineage:
+                return LineageDetails(
+                    columnsLineage=column_lineage,
+                    source=LineageSource.ExternalTableLineage,
+                )
+            return None
+        except Exception as exc:
+            logger.debug(f"Error computing container column lineage: {exc}")
+            logger.debug(traceback.format_exc())
+            return None
+
     def _get_lineage_details(
         self, from_table: Table, to_table: Table, databricks_table_fqn: str
     ) -> Optional[LineageDetails]:
@@ -124,16 +160,87 @@ class UnitycatalogLineageSource(Source):
             logger.debug(traceback.format_exc())
             return None
 
+    def _handle_external_location_lineage(
+        self, file_info, table: Table, is_upstream: bool
+    ) -> Iterable[Either[AddLineageRequest]]:
+        try:
+            if not file_info.storage_location:
+                logger.debug("No storage location found in fileInfo")
+                return
+
+            location_entity = self.metadata.es_search_container_by_path(
+                full_path=file_info.storage_location, fields="dataModel"
+            )
+
+            if location_entity and location_entity[0]:
+                lineage_details = None
+                if location_entity[0].dataModel:
+                    lineage_details = self._get_container_column_lineage(
+                        location_entity[0].dataModel, table
+                    )
+
+                if is_upstream:
+                    yield Either(
+                        left=None,
+                        right=AddLineageRequest(
+                            edge=EntitiesEdge(
+                                fromEntity=EntityReference(
+                                    id=location_entity[0].id,
+                                    type="container",
+                                ),
+                                toEntity=EntityReference(
+                                    id=table.id,
+                                    type="table",
+                                ),
+                                lineageDetails=lineage_details,
+                            )
+                        ),
+                    )
+                else:
+                    yield Either(
+                        left=None,
+                        right=AddLineageRequest(
+                            edge=EntitiesEdge(
+                                fromEntity=EntityReference(
+                                    id=table.id,
+                                    type="table",
+                                ),
+                                toEntity=EntityReference(
+                                    id=location_entity[0].id,
+                                    type="container",
+                                ),
+                                lineageDetails=lineage_details,
+                            )
+                        ),
+                    )
+            else:
+                logger.debug(
+                    f"Unable to find container for external location: {file_info.storage_location}"
+                )
+        except Exception as exc:
+            logger.debug(
+                f"Error while processing external location lineage for {file_info.storage_location}: {exc}"
+            )
+            logger.debug(traceback.format_exc())
+
     def _handle_upstream_table(
         self,
         table_streams: LineageTableStreams,
         table: Table,
         databricks_table_fqn: str,
     ) -> Iterable[Either[AddLineageRequest]]:
-        for upstream_table in table_streams.upstream_tables:
+        for upstream_entity in table_streams.upstreams:
             try:
-                if not upstream_table.name:
+                if upstream_entity.fileInfo:
+                    yield from self._handle_external_location_lineage(
+                        upstream_entity.fileInfo, table, is_upstream=True
+                    )
+                    continue
+
+                if not upstream_entity.tableInfo or not upstream_entity.tableInfo.name:
                     continue
+
+                upstream_table = upstream_entity.tableInfo
                 from_entity_fqn = fqn.build(
                     metadata=self.metadata,
                     entity_type=Table,
@@ -172,9 +279,69 @@ class UnitycatalogLineageSource(Source):
                     )
             except Exception:
                 logger.debug(
-                    "Error while processing lineage for "
-                    f"{upstream_table.catalog_name}.{upstream_table.schema_name}.{upstream_table.name}"
-                    f" -> {databricks_table_fqn}"
+                    "Error while processing upstream lineage for "
+                    f"{databricks_table_fqn}"
+                )
+                logger.debug(traceback.format_exc())
+
+    def _handle_downstream_table(
+        self,
+        table_streams: LineageTableStreams,
+        table: Table,
+        databricks_table_fqn: str,
+    ) -> Iterable[Either[AddLineageRequest]]:
+        for downstream_entity in table_streams.downstreams:
+            try:
+                if downstream_entity.fileInfo:
+                    yield from self._handle_external_location_lineage(
+                        downstream_entity.fileInfo, table, is_upstream=False
+                    )
+                    continue
+
+                if (
+                    not downstream_entity.tableInfo
+                    or not downstream_entity.tableInfo.name
+                ):
+                    continue
+
+                downstream_table = downstream_entity.tableInfo
+                to_entity_fqn = fqn.build(
+                    metadata=self.metadata,
+                    entity_type=Table,
+                    database_name=downstream_table.catalog_name,
+                    schema_name=downstream_table.schema_name,
+                    table_name=downstream_table.name,
+                    service_name=self.config.serviceName,
+                )
+
+                to_entity = self.metadata.get_by_name(entity=Table, fqn=to_entity_fqn)
+                if to_entity:
+                    downstream_table_fqn = f"{downstream_table.catalog_name}.{downstream_table.schema_name}.{downstream_table.name}"
+                    lineage_details = self._get_lineage_details(
+                        from_table=table,
+                        to_table=to_entity,
+                        databricks_table_fqn=downstream_table_fqn,
+                    )
+                    yield Either(
+                        left=None,
+                        right=AddLineageRequest(
+                            edge=EntitiesEdge(
+                                fromEntity=EntityReference(id=table.id, type="table"),
+                                toEntity=EntityReference(id=to_entity.id, type="table"),
+                                lineageDetails=lineage_details,
+                            )
+                        ),
+                    )
+                else:
+                    logger.debug(
+                        f"Unable to find downstream entity for "
+                        f"{databricks_table_fqn} -> "
+                        f"{downstream_table.catalog_name}.{downstream_table.schema_name}.{downstream_table.name}"
+                    )
+            except Exception:
+                logger.debug(
+                    "Error while processing downstream lineage for "
+                    f"{databricks_table_fqn}"
                 )
                 logger.debug(traceback.format_exc())
 
@@ -225,10 +392,17 @@ class UnitycatalogLineageSource(Source):
                     table_streams: LineageTableStreams = self.client.get_table_lineage(
                         databricks_table_fqn
                     )
+
+                    # Process upstream lineage
                     yield from self._handle_upstream_table(
                         table_streams, table, databricks_table_fqn
                     )
 
+                    # Process downstream lineage
+                    yield from self._handle_downstream_table(
+                        table_streams, table, databricks_table_fqn
+                    )
+
     def test_connection(self) -> None:
         test_connection_common(
             self.metadata, self.connection_obj, self.service_connection

metadata/ingestion/source/database/unitycatalog/models.py

@@ -22,6 +22,8 @@ class DatabricksTable(BaseModel):
     name: Optional[str] = None
     catalog_name: Optional[str] = None
     schema_name: Optional[str] = None
+    table_type: Optional[str] = None
+    lineage_timestamp: Optional[str] = None
 
 
 class DatabricksColumn(BaseModel):
@@ -31,9 +33,23 @@ class DatabricksColumn(BaseModel):
     table_name: Optional[str] = None
 
 
+class FileInfo(BaseModel):
+    path: Optional[str] = None
+    has_permission: Optional[bool] = None
+    securable_name: Optional[str] = None
+    storage_location: Optional[str] = None
+    securable_type: Optional[str] = None
+    lineage_timestamp: Optional[str] = None
+
+
+class LineageEntity(BaseModel):
+    tableInfo: Optional[DatabricksTable] = None
+    fileInfo: Optional[FileInfo] = None
+
+
 class LineageTableStreams(BaseModel):
-    upstream_tables: Optional[List[DatabricksTable]] = []
-    downstream_tables: Optional[List[DatabricksTable]] = []
+    upstreams: Optional[List[LineageEntity]] = []
+    downstreams: Optional[List[LineageEntity]] = []
 
 
 class LineageColumnStreams(BaseModel):