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

datahub/ingestion/source/sql/postgres.py

@@ -36,14 +36,8 @@ from datahub.ingestion.source.sql.sql_common import (
     register_custom_type,
 )
 from datahub.ingestion.source.sql.sql_config import BasicSQLAlchemyConfig
-from datahub.ingestion.source.sql.sql_utils import (
-    gen_database_key,
-    gen_schema_key,
-)
 from datahub.ingestion.source.sql.stored_procedures.base import (
     BaseProcedure,
-    generate_procedure_container_workunits,
-    generate_procedure_workunits,
 )
 from datahub.metadata.com.linkedin.pegasus2avro.schema import (
     ArrayTypeClass,
@@ -132,10 +126,12 @@ class PostgresConfig(BasePostgresConfig):
             "Note: this is not used if `database` or `sqlalchemy_uri` are provided."
         ),
     )
+
     include_stored_procedures: bool = Field(
         default=True,
         description="Include ingest of stored procedures.",
     )
+
     procedure_pattern: AllowDenyPattern = Field(
         default=AllowDenyPattern.allow_all(),
         description="Regex patterns for stored procedures to filter in ingestion."
@@ -310,73 +306,6 @@ class PostgresSource(SQLAlchemySource):
         except Exception as e:
             logger.error(f"failed to fetch profile metadata: {e}")
 
-    def get_schema_level_workunits(
-        self,
-        inspector: Inspector,
-        schema: str,
-        database: str,
-    ) -> Iterable[Union[MetadataWorkUnit, SqlWorkUnit]]:
-        yield from super().get_schema_level_workunits(
-            inspector=inspector,
-            schema=schema,
-            database=database,
-        )
-
-        if self.config.include_stored_procedures:
-            try:
-                yield from self.loop_stored_procedures(inspector, schema, self.config)
-            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 loop_stored_procedures(
-        self,
-        inspector: Inspector,
-        schema: str,
-        config: PostgresConfig,
-    ) -> 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,
-                )
-
-                if not self.config.procedure_pattern.allowed(procedure_qualified_name):
-                    self.report.report_dropped(procedure_qualified_name)
-                else:
-                    procedures.append(procedure)
-            return procedures
-        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]:
@@ -401,10 +330,9 @@ class PostgresSource(SQLAlchemySource):
                         pg_language l ON l.oid = p.prolang
                     WHERE
                         p.prokind = 'p'
-                        AND n.nspname = '"""
-                + schema
-                + """';
-                    """
+                        AND n.nspname = %s;
+                """,
+                (schema,),
             )
 
             procedure_rows = list(procedures)
@@ -423,60 +351,3 @@ class PostgresSource(SQLAlchemySource):
                     )
                 )
             return base_procedures
-
-    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/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."""