acryl-datahub 1.2.0.4rc4__py3-none-any.whl → 1.2.0.5rc1__py3-none-any.whl

datahub/ingestion/source/sql/sql_common.py

@@ -27,6 +27,7 @@ from sqlalchemy.exc import ProgrammingError
 from sqlalchemy.sql import sqltypes as types
 from sqlalchemy.types import TypeDecorator, TypeEngine
 
+from datahub.configuration.common import AllowDenyPattern
 from datahub.emitter.mce_builder import (
     make_data_platform_urn,
     make_dataplatform_instance_urn,
@@ -71,6 +72,11 @@ from datahub.ingestion.source.sql.sql_utils import (
 from datahub.ingestion.source.sql.sqlalchemy_data_reader import (
     SqlAlchemyTableDataReader,
 )
+from datahub.ingestion.source.sql.stored_procedures.base import (
+    BaseProcedure,
+    generate_procedure_container_workunits,
+    generate_procedure_workunits,
+)
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
     StaleEntityRemovalHandler,
 )
@@ -531,6 +537,24 @@ class SQLAlchemySource(StatefulIngestionSourceBase, TestableSource):
         if self.config.include_views:
             yield from self.loop_views(inspector, schema, self.config)
 
+        if getattr(self.config, "include_stored_procedures", False):
+            try:
+                yield from self.loop_stored_procedures(inspector, schema, self.config)
+            except NotImplementedError as e:
+                self.report.warning(
+                    title="Stored procedures not supported",
+                    message="The current SQL dialect does not support stored procedures.",
+                    context=f"{database}.{schema}",
+                    exc=e,
+                )
+            except Exception as e:
+                self.report.failure(
+                    title="Failed to list stored procedures for schema",
+                    message="An error occurred while listing procedures for the schema.",
+                    context=f"{database}.{schema}",
+                    exc=e,
+                )
+
     def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:
         return [
             *super().get_workunit_processors(),
@@ -1437,3 +1461,116 @@ class SQLAlchemySource(StatefulIngestionSourceBase, TestableSource):
 
     def get_report(self):
         return self.report
+
+    def loop_stored_procedures(
+        self,
+        inspector: Inspector,
+        schema: str,
+        config: Union[SQLCommonConfig, Type[SQLCommonConfig]],
+    ) -> Iterable[MetadataWorkUnit]:
+        """
+        Loop schema data for get stored procedures as dataJob-s.
+        """
+        db_name = self.get_db_name(inspector)
+
+        procedures = self.fetch_procedures_for_schema(inspector, schema, db_name)
+        if procedures:
+            yield from self._process_procedures(procedures, db_name, schema)
+
+    def fetch_procedures_for_schema(
+        self, inspector: Inspector, schema: str, db_name: str
+    ) -> List[BaseProcedure]:
+        try:
+            raw_procedures: List[BaseProcedure] = self.get_procedures_for_schema(
+                inspector, schema, db_name
+            )
+            procedures: List[BaseProcedure] = []
+            for procedure in raw_procedures:
+                procedure_qualified_name = self.get_identifier(
+                    schema=schema,
+                    entity=procedure.name,
+                    inspector=inspector,
+                )
+
+                procedure_pattern = getattr(
+                    self.config, "procedure_pattern", AllowDenyPattern.allow_all()
+                )
+                if not procedure_pattern.allowed(procedure_qualified_name):
+                    self.report.report_dropped(procedure_qualified_name)
+                else:
+                    procedures.append(procedure)
+            return procedures
+        except NotImplementedError:
+            raise
+        except Exception as e:
+            self.report.warning(
+                title="Failed to get procedures for schema",
+                message="An error occurred while fetching procedures for the schema.",
+                context=f"{db_name}.{schema}",
+                exc=e,
+            )
+            return []
+
+    def get_procedures_for_schema(
+        self, inspector: Inspector, schema: str, db_name: str
+    ) -> List[BaseProcedure]:
+        raise NotImplementedError(
+            "Subclasses must implement the 'get_procedures_for_schema' method."
+        )
+
+    def _process_procedures(
+        self,
+        procedures: List[BaseProcedure],
+        db_name: str,
+        schema: str,
+    ) -> Iterable[MetadataWorkUnit]:
+        if procedures:
+            yield from generate_procedure_container_workunits(
+                database_key=gen_database_key(
+                    database=db_name,
+                    platform=self.platform,
+                    platform_instance=self.config.platform_instance,
+                    env=self.config.env,
+                ),
+                schema_key=gen_schema_key(
+                    db_name=db_name,
+                    schema=schema,
+                    platform=self.platform,
+                    platform_instance=self.config.platform_instance,
+                    env=self.config.env,
+                ),
+            )
+        for procedure in procedures:
+            yield from self._process_procedure(procedure, schema, db_name)
+
+    def _process_procedure(
+        self,
+        procedure: BaseProcedure,
+        schema: str,
+        db_name: str,
+    ) -> Iterable[MetadataWorkUnit]:
+        try:
+            yield from generate_procedure_workunits(
+                procedure=procedure,
+                database_key=gen_database_key(
+                    database=db_name,
+                    platform=self.platform,
+                    platform_instance=self.config.platform_instance,
+                    env=self.config.env,
+                ),
+                schema_key=gen_schema_key(
+                    db_name=db_name,
+                    schema=schema,
+                    platform=self.platform,
+                    platform_instance=self.config.platform_instance,
+                    env=self.config.env,
+                ),
+                schema_resolver=self.get_schema_resolver(),
+            )
+        except Exception as e:
+            self.report.warning(
+                title="Failed to emit stored procedure",
+                message="An error occurred while emitting stored procedure",
+                context=procedure.name,
+                exc=e,
+            )

datahub/ingestion/source/superset.py

@@ -8,6 +8,7 @@ from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 
 import dateutil.parser as dp
 import requests
+import sqlglot
 from pydantic import BaseModel
 from pydantic.class_validators import root_validator, validator
 from pydantic.fields import Field
@@ -75,6 +76,7 @@ from datahub.metadata.com.linkedin.pegasus2avro.schema import (
     SchemaFieldDataType,
     SchemaMetadata,
     StringTypeClass,
+    TimeTypeClass,
 )
 from datahub.metadata.schema_classes import (
     AuditStampClass,
@@ -131,8 +133,11 @@ FIELD_TYPE_MAPPING = {
     "STRING": StringTypeClass,
     "FLOAT": NumberTypeClass,
     "DATETIME": DateTypeClass,
+    "TIMESTAMP": TimeTypeClass,
     "BOOLEAN": BooleanTypeClass,
     "SQL": StringTypeClass,
+    "NUMERIC": NumberTypeClass,
+    "TEXT": StringTypeClass,
 }
 
 
@@ -633,74 +638,130 @@ class SupersetSource(StatefulIngestionSourceBase):
 
         return input_fields
 
-    def construct_chart_cll(
-        self,
-        chart_data: dict,
-        datasource_urn: Union[str, None],
-        datasource_id: Union[Any, int],
-    ) -> List[InputField]:
-        column_data: List[Union[str, dict]] = chart_data.get("form_data", {}).get(
-            "all_columns", []
-        )
+    def _extract_columns_from_sql(self, sql_expr: Optional[str]) -> List[str]:
+        if not sql_expr:
+            return []
 
-        # the second field represents whether its a SQL expression,
-        # false being just regular column and true being SQL col
-        chart_column_data: List[Tuple[str, bool]] = [
-            (column, False)
-            if isinstance(column, str)
-            else (column.get("label", ""), True)
-            for column in column_data
-        ]
+        try:
+            parsed_expr = sqlglot.parse_one(sql_expr)
 
-        dataset_columns: List[Tuple[str, str, str]] = []
+            column_refs = set()
+            for node in parsed_expr.walk():
+                if isinstance(node, sqlglot.exp.Column):
+                    column_name = node.name
+                    column_refs.add(column_name)
 
-        # parses the superset dataset's column info, to build type and description info
-        if datasource_id:
-            dataset_info = self.get_dataset_info(datasource_id).get("result", {})
-            dataset_column_info = dataset_info.get("columns", [])
-            dataset_metric_info = dataset_info.get("metrics", [])
+            return list(column_refs)
+        except Exception as e:
+            self.report.warning(f"Failed to parse SQL expression '{sql_expr}': {e}")
+            return []
 
-            for column in dataset_column_info:
-                col_name = column.get("column_name", "")
-                col_type = column.get("type", "")
-                col_description = column.get("description", "")
+    def _process_column_item(
+        self, item: Union[str, dict], unique_columns: Dict[str, bool]
+    ) -> None:
+        """Process a single column item and add to unique_columns."""
 
-                # if missing column name or column type, cannot construct the column,
-                # so we skip this column, missing description is fine
-                if col_name == "" or col_type == "":
-                    logger.info(f"could not construct column lineage for {column}")
-                    continue
+        def add_column(col_name: str, is_sql: bool) -> None:
+            if not col_name:
+                return
+            # Always set to False if any non-SQL seen, else keep as is_sql
+            unique_columns[col_name] = unique_columns.get(col_name, True) and is_sql
+
+        if isinstance(item, str):
+            add_column(item, False)
+        elif isinstance(item, dict):
+            if item.get("expressionType") == "SIMPLE":
+                # For metrics with SIMPLE expression type
+                add_column(item.get("column", {}).get("column_name", ""), False)
+            elif item.get("expressionType") == "SQL":
+                sql_expr = item.get("sqlExpression")
+                column_refs = self._extract_columns_from_sql(sql_expr)
+                for col in column_refs:
+                    add_column(col, False)
+                if not column_refs:
+                    add_column(item.get("label", ""), True)
+
+    def _collect_all_unique_columns(self, form_data: dict) -> Dict[str, bool]:
+        """Collect all unique column names from form_data, distinguishing SQL vs non-SQL."""
+        unique_columns: Dict[str, bool] = {}
+
+        # Process regular columns
+        for column in form_data.get("all_columns", []):
+            self._process_column_item(column, unique_columns)
+
+        # Process metrics
+        # For charts with a single metric, the metric is stored in the form_data as a string in the 'metric' key
+        # For charts with multiple metrics, the metrics are stored in the form_data as a list of strings in the 'metrics' key
+        if "metric" in form_data:
+            metrics_data = [form_data.get("metric")]
+        else:
+            metrics_data = form_data.get("metrics", [])
 
-                dataset_columns.append((col_name, col_type, col_description))
+        for metric in metrics_data:
+            if metric is not None:
+                self._process_column_item(metric, unique_columns)
 
-            for metric in dataset_metric_info:
-                metric_name = metric.get("metric_name", "")
-                metric_type = metric.get("metric_type", "")
-                metric_description = metric.get("description", "")
+        # Process group by columns
+        for group in form_data.get("groupby", []):
+            self._process_column_item(group, unique_columns)
 
-                if metric_name == "" or metric_type == "":
-                    logger.info(f"could not construct metric lineage for {metric}")
-                    continue
+        # Process x-axis columns
+        x_axis_data = form_data.get("x_axis")
+        if x_axis_data is not None:
+            self._process_column_item(x_axis_data, unique_columns)
 
-                dataset_columns.append((metric_name, metric_type, metric_description))
-        else:
-            # if no datasource id, cannot build cll, just return
+        return unique_columns
+
+    def _fetch_dataset_columns(
+        self, datasource_id: Union[Any, int]
+    ) -> List[Tuple[str, str, str]]:
+        """Fetch dataset columns and metrics from Superset API."""
+        if not datasource_id:
             logger.warning(
                 "no datasource id was found, cannot build column level lineage"
             )
             return []
 
+        dataset_info = self.get_dataset_info(datasource_id).get("result", {})
+        dataset_column_info = dataset_info.get("columns", [])
+        dataset_metric_info = dataset_info.get("metrics", [])
+
+        dataset_columns: List[Tuple[str, str, str]] = []
+        for column in dataset_column_info:
+            col_name = column.get("column_name", "")
+            col_type = column.get("type", "")
+            col_description = column.get("description", "")
+
+            if col_name == "" or col_type == "":
+                logger.info(f"could not construct column lineage for {column}")
+                continue
+
+            dataset_columns.append((col_name, col_type, col_description))
+
+        for metric in dataset_metric_info:
+            metric_name = metric.get("metric_name", "")
+            metric_type = metric.get("metric_type", "")
+            metric_description = metric.get("description", "")
+
+            if metric_name == "" or metric_type == "":
+                logger.info(f"could not construct metric lineage for {metric}")
+                continue
+
+            dataset_columns.append((metric_name, metric_type, metric_description))
+
+        return dataset_columns
+
+    def _match_chart_columns_with_dataset(
+        self,
+        unique_chart_columns: Dict[str, bool],
+        dataset_columns: List[Tuple[str, str, str]],
+    ) -> List[Tuple[str, str, str]]:
+        """Match chart columns with dataset columns, preserving SQL/non-SQL status."""
         chart_columns: List[Tuple[str, str, str]] = []
-        for chart_col in chart_column_data:
-            chart_col_name, is_sql = chart_col
+
+        for chart_col_name, is_sql in unique_chart_columns.items():
             if is_sql:
-                chart_columns.append(
-                    (
-                        chart_col_name,
-                        "SQL",
-                        "",
-                    )
-                )
+                chart_columns.append((chart_col_name, "SQL", ""))
                 continue
 
             # find matching upstream column
@@ -711,13 +772,36 @@ class SupersetSource(StatefulIngestionSourceBase):
                 if dataset_col_name == chart_col_name:
                     chart_columns.append(
                         (chart_col_name, dataset_col_type, dataset_col_description)
-                    )  # column name, column type, description
+                    )
                     break
-
-            # if no matching upstream column was found
-            if len(chart_columns) == 0 or chart_columns[-1][0] != chart_col_name:
+            else:
                 chart_columns.append((chart_col_name, "", ""))
 
+        return chart_columns
+
+    def construct_chart_cll(
+        self,
+        chart_data: dict,
+        datasource_urn: Union[str, None],
+        datasource_id: Union[Any, int],
+    ) -> List[InputField]:
+        """Construct column-level lineage for a chart."""
+        form_data = chart_data.get("form_data", {})
+
+        # Extract and process all columns in one go
+        unique_columns = self._collect_all_unique_columns(form_data)
+
+        # Fetch dataset columns
+        dataset_columns = self._fetch_dataset_columns(datasource_id)
+        if not dataset_columns:
+            return []
+
+        # Match chart columns with dataset columns
+        chart_columns = self._match_chart_columns_with_dataset(
+            unique_columns, dataset_columns
+        )
+
+        # Build input fields
         return self.build_input_fields(chart_columns, datasource_urn)
 
     def construct_chart_from_chart_data(

datahub/ingestion/source/unity/config.py

@@ -275,6 +275,17 @@ class UnityCatalogSourceConfig(
         hidden_from_docs=True,
     )
 
+    databricks_api_page_size: int = pydantic.Field(
+        default=0,
+        ge=0,
+        description=(
+            "Page size for Databricks API calls when listing resources (catalogs, schemas, tables, etc.). "
+            "When set to 0 (default), uses server-side configured page length (recommended). "
+            "When set to a positive value, the page length is the minimum of this value and the server configured value. "
+            "Must be a non-negative integer."
+        ),
+    )
+
     include_usage_statistics: bool = Field(
         default=True,
         description="Generate usage statistics.",

datahub/ingestion/source/unity/connection_test.py

@@ -19,6 +19,7 @@ class UnityCatalogConnectionTest:
             self.config.token,
             self.config.profiling.warehouse_id,
             report=self.report,
+            databricks_api_page_size=self.config.databricks_api_page_size,
         )
 
     def get_connection_test(self) -> TestConnectionReport:

datahub/ingestion/source/unity/platform_resource_repository.py

@@ -0,0 +1,19 @@
+import logging
+
+from datahub.api.entities.external.external_entities import (
+    PlatformResourceRepository,
+)
+from datahub.ingestion.source.unity.tag_entities import (
+    UnityCatalogTagPlatformResource,
+    UnityCatalogTagPlatformResourceId,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class UnityCatalogPlatformResourceRepository(
+    PlatformResourceRepository[
+        UnityCatalogTagPlatformResourceId, UnityCatalogTagPlatformResource
+    ]
+):
+    """Unity Catalog-specific platform resource repository with tag-related operations."""

datahub/ingestion/source/unity/proxy.py

@@ -141,6 +141,7 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
         report: UnityCatalogReport,
         hive_metastore_proxy: Optional[HiveMetastoreProxy] = None,
         lineage_data_source: LineageDataSource = LineageDataSource.AUTO,
+        databricks_api_page_size: int = 0,
     ):
         self._workspace_client = WorkspaceClient(
             host=workspace_url,
@@ -152,6 +153,7 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
         self.report = report
         self.hive_metastore_proxy = hive_metastore_proxy
         self.lineage_data_source = lineage_data_source
+        self.databricks_api_page_size = databricks_api_page_size
         self._sql_connection_params = {
             "server_hostname": self._workspace_client.config.host.replace(
                 "https://", ""
@@ -161,7 +163,11 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
         }
 
     def check_basic_connectivity(self) -> bool:
-        return bool(self._workspace_client.catalogs.list(include_browse=True))
+        return bool(
+            self._workspace_client.catalogs.list(
+                include_browse=True, max_results=self.databricks_api_page_size
+            )
+        )
 
     def assigned_metastore(self) -> Optional[Metastore]:
         response = self._workspace_client.metastores.summary()
@@ -171,7 +177,9 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
         if self.hive_metastore_proxy:
             yield self.hive_metastore_proxy.hive_metastore_catalog(metastore)
 
-        response = self._workspace_client.catalogs.list(include_browse=True)
+        response = self._workspace_client.catalogs.list(
+            include_browse=True, max_results=self.databricks_api_page_size
+        )
         if not response:
             logger.info("Catalogs not found")
             return
@@ -203,7 +211,9 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
             yield from self.hive_metastore_proxy.hive_metastore_schemas(catalog)
             return
         response = self._workspace_client.schemas.list(
-            catalog_name=catalog.name, include_browse=True
+            catalog_name=catalog.name,
+            include_browse=True,
+            max_results=self.databricks_api_page_size,
         )
         if not response:
             logger.info(f"Schemas not found for catalog {catalog.id}")
@@ -225,6 +235,7 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
                 catalog_name=schema.catalog.name,
                 schema_name=schema.name,
                 include_browse=True,
+                max_results=self.databricks_api_page_size,
             )
             if not response:
                 logger.info(f"Tables not found for schema {schema.id}")
@@ -257,7 +268,10 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
         return group_list
 
     def workspace_notebooks(self) -> Iterable[Notebook]:
-        for obj in self._workspace_client.workspace.list("/", recursive=True):
+        workspace_objects_iter = self._workspace_client.workspace.list(
+            "/", recursive=True, max_results=self.databricks_api_page_size
+        )
+        for obj in workspace_objects_iter:
             if obj.object_type == ObjectType.NOTEBOOK and obj.object_id and obj.path:
                 yield Notebook(
                     id=obj.object_id,
@@ -299,7 +313,6 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
     def _query_history(
         self,
         filter_by: QueryFilterWithStatementTypes,
-        max_results: int = 1000,
         include_metrics: bool = False,
     ) -> Iterable[QueryInfo]:
         """Manual implementation of the query_history.list() endpoint.
@@ -311,9 +324,10 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
         """
         method = "GET"
         path = "/api/2.0/sql/history/queries"
+
         body: Dict[str, Any] = {
             "include_metrics": include_metrics,
-            "max_results": max_results,  # Max batch size
+            "max_results": self.databricks_api_page_size,  # Max batch size
         }
 
         response: dict = self._workspace_client.api_client.do(  # type: ignore

datahub/ingestion/source/unity/report.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from typing import Optional, Tuple
+from typing import TYPE_CHECKING, Optional, Tuple
 
 from datahub.ingestion.api.report import EntityFilterReport, Report
 from datahub.ingestion.source.sql.sql_report import SQLSourceReport
@@ -7,6 +7,11 @@ from datahub.ingestion.source_report.ingestion_stage import IngestionStageReport
 from datahub.utilities.lossy_collections import LossyDict, LossyList
 from datahub.utilities.perf_timer import PerfTimer
 
+if TYPE_CHECKING:
+    from datahub.ingestion.source.unity.platform_resource_repository import (
+        UnityCatalogPlatformResourceRepository,
+    )
+
 
 @dataclass
 class UnityCatalogUsagePerfReport(Report):
@@ -61,3 +66,6 @@ class UnityCatalogReport(IngestionStageReport, SQLSourceReport):
     num_tables_missing_name: int = 0
     num_columns_missing_name: int = 0
     num_queries_missing_info: int = 0
+
+    # Platform resource repository for automatic cache statistics via SupportsAsObj
+    tag_urn_resolver_cache: Optional["UnityCatalogPlatformResourceRepository"] = None