acryl-datahub 1.3.1__py3-none-any.whl → 1.3.1.1__py3-none-any.whl

datahub/ingestion/source/unity/proxy.py

@@ -3,6 +3,7 @@ Manage the communication with DataBricks Server and provide equivalent dataclass
 """
 
 import dataclasses
+import json
 import logging
 import os
 from concurrent.futures import ThreadPoolExecutor
@@ -11,6 +12,7 @@ from typing import Any, Dict, Iterable, List, Optional, Sequence, Union, cast
 from unittest.mock import patch
 
 import cachetools
+import yaml
 from cachetools import cached
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.service.catalog import (
@@ -23,7 +25,11 @@ from databricks.sdk.service.catalog import (
     SchemaInfo,
     TableInfo,
 )
+from databricks.sdk.service.files import DownloadResponse, FilesAPI
 from databricks.sdk.service.iam import ServicePrincipal as DatabricksServicePrincipal
+from databricks.sdk.service.ml import (
+    ExperimentsAPI,
+)
 from databricks.sdk.service.sql import (
     QueryFilter,
     QueryInfo,
@@ -38,6 +44,7 @@ from typing_extensions import assert_never
 from datahub._version import nice_version_name
 from datahub.api.entities.external.unity_catalog_external_entites import UnityCatalogTag
 from datahub.emitter.mce_builder import parse_ts_millis
+from datahub.ingestion.source.unity.azure_auth_config import AzureAuthConfig
 from datahub.ingestion.source.unity.config import (
     LineageDataSource,
     UsageDataSource,
@@ -54,6 +61,8 @@ from datahub.ingestion.source.unity.proxy_types import (
     ExternalTableReference,
     Metastore,
     Model,
+    ModelRunDetails,
+    ModelSignature,
     ModelVersion,
     Notebook,
     NotebookReference,
@@ -155,30 +164,44 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
     _workspace_url: str
     report: UnityCatalogReport
     warehouse_id: str
+    _experiments_api: ExperimentsAPI
+    _files_api: FilesAPI
 
     def __init__(
         self,
         workspace_url: str,
-        personal_access_token: str,
         warehouse_id: Optional[str],
         report: UnityCatalogReport,
         hive_metastore_proxy: Optional[HiveMetastoreProxy] = None,
         lineage_data_source: LineageDataSource = LineageDataSource.AUTO,
         usage_data_source: UsageDataSource = UsageDataSource.AUTO,
         databricks_api_page_size: int = 0,
+        personal_access_token: Optional[str] = None,
+        azure_auth: Optional[AzureAuthConfig] = None,
     ):
-        self._workspace_client = WorkspaceClient(
-            host=workspace_url,
-            token=personal_access_token,
-            product="datahub",
-            product_version=nice_version_name(),
-        )
+        if azure_auth:
+            self._workspace_client = WorkspaceClient(
+                host=workspace_url,
+                azure_tenant_id=azure_auth.tenant_id,
+                azure_client_id=azure_auth.client_id,
+                azure_client_secret=azure_auth.client_secret.get_secret_value(),
+                product="datahub",
+                product_version=nice_version_name(),
+            )
+        else:
+            self._workspace_client = WorkspaceClient(
+                host=workspace_url,
+                token=personal_access_token,
+                product="datahub",
+                product_version=nice_version_name(),
+            )
         self.warehouse_id = warehouse_id or ""
         self.report = report
         self.hive_metastore_proxy = hive_metastore_proxy
         self.lineage_data_source = lineage_data_source
         self.usage_data_source = usage_data_source
         self.databricks_api_page_size = databricks_api_page_size
+        self._workspace_url = workspace_url
         self._sql_connection_params = {
             "server_hostname": self._workspace_client.config.host.replace(
                 "https://", ""
@@ -187,6 +210,179 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
             "access_token": self._workspace_client.config.token,
             "user_agent_entry": "datahub",
         }
+        # Initialize MLflow APIs
+        self._experiments_api = ExperimentsAPI(self._workspace_client.api_client)
+        self._files_api = FilesAPI(self._workspace_client.api_client)
+
+    def get_run_details(self, run_id: str) -> Optional[ModelRunDetails]:
+        """
+        Get comprehensive details from an MLflow run.
+
+        Args:
+            run_id: The MLflow run ID
+
+        Returns:
+            ModelRunDetails object with comprehensive run information
+        """
+        try:
+            run_response = self._experiments_api.get_run(run_id)
+            run = run_response.run
+
+            if (
+                not run
+                or not run.info
+                or not run.info.run_id
+                or not run.info.experiment_id
+            ):
+                return None
+
+            # Extract metrics
+            metrics: Dict[str, Any] = {}
+            if run.data and run.data.metrics:
+                for metric in run.data.metrics:
+                    if metric.key is not None:
+                        metrics[metric.key] = metric.value
+
+            # Extract parameters
+            parameters: Dict[str, Any] = {}
+            if run.data and run.data.params:
+                for param in run.data.params:
+                    if param.key is not None:
+                        parameters[param.key] = param.value
+
+            # Extract tags
+            tags: Dict[str, str] = {}
+            if run.data and run.data.tags:
+                for tag in run.data.tags:
+                    if tag.key is not None and tag.value is not None:
+                        tags[tag.key] = tag.value
+
+            return ModelRunDetails(
+                run_id=run.info.run_id,
+                experiment_id=run.info.experiment_id,
+                status=run.info.status.value if run.info.status else None,
+                start_time=parse_ts_millis(run.info.start_time),
+                end_time=parse_ts_millis(run.info.end_time),
+                user_id=run.info.user_id,
+                metrics=metrics,
+                parameters=parameters,
+                tags=tags,
+            )
+        except Exception as e:
+            logger.warning(
+                f"Unable to get run details for MLflow experiment, run-id: {run_id}",
+                exc_info=True,
+            )
+            self.report.report_warning(
+                title="Unable to get run details for MLflow experiment",
+                message="Error while getting run details for MLflow experiment",
+                context=f"run-id: {run_id}",
+                exc=e,
+            )
+            return None
+
+    def _extract_signature_from_files_api(
+        self, model_version: ModelVersionInfo
+    ) -> Optional[ModelSignature]:
+        """
+        Extract signature from MLmodel file using Databricks FilesAPI.
+        Uses the API endpoint: /api/2.0/fs/files/Models/{catalog}/{schema}/{model}/{version}/MLmodel
+
+        Args:
+            model_version: Unity Catalog ModelVersionInfo object with catalog_name, schema_name, model_name, version
+
+        Returns:
+            ModelSignature if found, None otherwise
+        """
+        try:
+            # Construct file path for FilesAPI
+            # The correct path format is: /Models/{catalog}/{schema}/{model}/{version}/MLmodel
+            file_path = (
+                f"/Models/{model_version.catalog_name}/{model_version.schema_name}/"
+                f"{model_version.model_name}/{model_version.version}/MLmodel"
+            )
+
+            logger.debug(f"Downloading MLmodel from FilesAPI: {file_path}")
+
+            # Download the file using FilesAPI
+            download_response: DownloadResponse = self._files_api.download(
+                file_path=file_path
+            )
+
+            # Read the file content
+            # DownloadResponse.contents is a BinaryIO object
+            if download_response and download_response.contents:
+                content_stream = download_response.contents
+
+                # Read from the binary stream
+                if content_stream:
+                    mlmodel_content: str = content_stream.read().decode("utf-8")
+
+                    logger.debug(
+                        f"MLmodel file contents from FilesAPI ({file_path}):\n{mlmodel_content}"
+                    )
+
+                    # Parse YAML content
+                    mlmodel_data = yaml.safe_load(mlmodel_content)
+
+                    # Extract signature from MLmodel YAML
+                    if mlmodel_data and "signature" in mlmodel_data:
+                        signature_raw = mlmodel_data["signature"]
+
+                        # Signature inputs and outputs are stored as JSON strings in the YAML
+                        # Parse them into proper dict/list format
+                        signature_data = {}
+                        if "inputs" in signature_raw:
+                            try:
+                                signature_data["inputs"] = json.loads(
+                                    signature_raw["inputs"]
+                                )
+                            except (json.JSONDecodeError, TypeError) as e:
+                                logger.debug(f"Failed to parse inputs JSON: {e}")
+
+                        if "outputs" in signature_raw:
+                            try:
+                                signature_data["outputs"] = json.loads(
+                                    signature_raw["outputs"]
+                                )
+                            except (json.JSONDecodeError, TypeError) as e:
+                                logger.debug(f"Failed to parse outputs JSON: {e}")
+
+                        if "params" in signature_raw:
+                            try:
+                                signature_data["params"] = json.loads(
+                                    signature_raw["params"]
+                                )
+                            except (json.JSONDecodeError, TypeError) as e:
+                                logger.debug(f"Failed to parse params JSON: {e}")
+
+                        return ModelSignature(
+                            inputs=signature_data.get("inputs"),
+                            outputs=signature_data.get("outputs"),
+                            parameters=signature_data.get("params"),
+                        )
+                    else:
+                        logger.debug(
+                            f"No signature found in MLmodel data from {file_path}"
+                        )
+                        return None
+
+            return None
+
+        except Exception as e:
+            model_name = getattr(model_version, "model_name", "unknown")
+            version_num = getattr(model_version, "version", "unknown")
+            self.report.report_warning(
+                title="Unable to extract signature from MLmodel file",
+                message="Error while extracting signature from MLmodel file",
+                context=f"model-name: {model_name}, model-version: {version_num}",
+                exc=e,
+            )
+            logger.warning(
+                f"Unable to extract signature from MLmodel file, model-name: {model_name}, model-version: {version_num}",
+                exc_info=True,
+            )
+            return None
 
     def check_basic_connectivity(self) -> bool:
         return bool(
@@ -1019,6 +1215,17 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
             for alias in obj.aliases:
                 if alias.alias_name:
                     aliases.append(alias.alias_name)
+
+        run_details: Optional[ModelRunDetails] = None
+        # Fetch run details if run_id exists
+        if obj.run_id:
+            run_details = self.get_run_details(obj.run_id)
+
+        # Extract signature separately from Files API
+        signature: Optional[ModelSignature] = self._extract_signature_from_files_api(
+            obj
+        )
+
         return ModelVersion(
             id=f"{model.id}_{obj.version}",
             name=f"{model.name}_{obj.version}",
@@ -1029,6 +1236,8 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
             created_at=parse_ts_millis(obj.created_at),
             updated_at=parse_ts_millis(obj.updated_at),
             created_by=obj.created_by,
+            run_details=run_details,
+            signature=signature,
         )
 
     def _create_service_principal(

datahub/ingestion/source/unity/proxy_types.py

@@ -339,8 +339,75 @@ class Notebook:
         )
 
 
+@dataclass
+class ModelSignature:
+    """
+    Represents the model signature with input and output schemas extracted from MLflow.
+
+    In Unity Catalog, model signatures define the expected input/output formats for ML models.
+    Model signature is stored in the MLmodel YAML file.
+
+    Attributes:
+        inputs: List of input schema specifications, each containing name, type, dtype, shape
+        outputs: List of output schema specifications, each containing name, type, dtype, shape
+        parameters: List of model parameters
+    """
+
+    inputs: Optional[List[Dict[str, str]]]
+    outputs: Optional[List[Dict[str, str]]]
+    parameters: Optional[List[Dict[str, str]]]
+
+
+@dataclass
+class ModelRunDetails:
+    """
+    Represents comprehensive details from an MLflow run associated with a Unity Catalog model version.
+
+    In Unity Catalog, each model version is linked to an MLflow run via run_id. This dataclass
+    contains all the metadata extracted from that MLflow run, including metrics, parameters,
+    and tags.
+
+    Attributes:
+        run_id: MLflow run ID
+        experiment_id: MLflow experiment ID
+        status: Run status (e.g., "FINISHED", "RUNNING")
+        start_time: Run start timestamp (milliseconds since epoch)
+        end_time: Run end timestamp (milliseconds since epoch)
+        user_id: User who initiated the run
+        metrics: Training metrics (e.g., accuracy, loss)
+        parameters: Hyperparameters used for training
+        tags: Run tags/metadata
+    """
+
+    run_id: str
+    experiment_id: str
+    status: Optional[str]
+    start_time: Optional[datetime]
+    end_time: Optional[datetime]
+    user_id: Optional[str]
+    metrics: Optional[Dict[str, str]]
+    parameters: Optional[Dict[str, str]]
+    tags: Optional[Dict[str, str]]
+
+
 @dataclass
 class Model:
+    """
+    Represents a Unity Catalog registered ML model (model group).
+
+    In Unity Catalog, a registered model is a collection of model versions.
+    This dataclass corresponds to a Unity Catalog RegisteredModelInfo.
+
+    Attributes:
+        id: Full qualified name (e.g., "catalog.schema.model_name")
+        name: Model name without catalog/schema prefix
+        schema_name: Schema name containing the model
+        catalog_name: Catalog name containing the model
+        description: Model description/comment
+        created_at: Model creation timestamp
+        updated_at: Last update timestamp
+    """
+
     id: str
     name: str
     schema_name: str
@@ -352,6 +419,28 @@ class Model:
 
 @dataclass
 class ModelVersion:
+    """
+    Represents a specific version of a Unity Catalog registered ML model.
+
+    In Unity Catalog, each model version is linked to an MLflow run (via run_id).
+    This dataclass corresponds to a Unity Catalog ModelVersionInfo.
+
+    Attributes:
+        id: Unique identifier combining model ID and version (e.g., "catalog.schema.model_1")
+        name: Versioned model name
+        model: Reference to the parent Model (model group)
+        version: Version number as string
+        aliases: List of aliases (e.g., ["prod", "latest"])
+        description: Version description/comment
+        created_at: Version creation timestamp
+        updated_at: Last update timestamp
+        created_by: User who created this version
+        run_details: Comprehensive MLflow run details (metrics, parameters, tags)
+                     extracted from the MLflow run linked to this model version.
+        signature: Model signature extracted from the MLmodel file via Files API.
+                   Contains input/output schema specifications and parameters.
+    """
+
     id: str
     name: str
     model: Model
@@ -361,3 +450,5 @@ class ModelVersion:
     created_at: Optional[datetime]
     updated_at: Optional[datetime]
     created_by: Optional[str]
+    run_details: Optional["ModelRunDetails"]
+    signature: Optional["ModelSignature"]

datahub/ingestion/source/unity/source.py

@@ -1,7 +1,9 @@
+import dataclasses
+import json
 import logging
 import re
 import time
-from typing import Dict, Iterable, List, Optional, Set, Tuple, Union
+from typing import Dict, Iterable, List, Optional, Set, Tuple, Union, cast
 from urllib.parse import urljoin
 
 from datahub.api.entities.external.unity_catalog_external_entites import UnityCatalogTag
@@ -209,13 +211,14 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
 
         self.unity_catalog_api_proxy = UnityCatalogApiProxy(
             config.workspace_url,
-            config.token,
             config.warehouse_id,
             report=self.report,
             hive_metastore_proxy=self.hive_metastore_proxy,
             lineage_data_source=config.lineage_data_source,
             usage_data_source=config.usage_data_source,
             databricks_api_page_size=config.databricks_api_page_size,
+            personal_access_token=config.token if config.token else None,
+            azure_auth=config.azure_auth if config.azure_auth else None,
         )
 
         self.external_url_base = urljoin(self.config.workspace_url, "/explore/data")
@@ -317,7 +320,7 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
 
     @classmethod
     def create(cls, config_dict, ctx):
-        config = UnityCatalogSourceConfig.parse_obj(config_dict)
+        config = UnityCatalogSourceConfig.model_validate(config_dict)
         return cls(ctx=ctx, config=config)
 
     def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:
@@ -741,6 +744,17 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
                     created=TimeStampClass(time=created_time, actor=created_actor),
                 )
             )
+        custom_properties = {}
+        if ml_model_version.signature:
+            for key, value in dataclasses.asdict(ml_model_version.signature).items():
+                if value:
+                    custom_properties[f"signature.{key}"] = json.dumps(value)
+
+        if ml_model_version.run_details:
+            if ml_model_version.run_details.tags:
+                for key, value in ml_model_version.run_details.tags.items():
+                    if value:
+                        custom_properties[key] = json.dumps(value)
 
         ml_model = MLModel(
             id=ml_model_version.id,
@@ -751,6 +765,18 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
             model_group=ml_model_urn,
             platform=self.platform,
             last_modified=ml_model_version.updated_at,
+            training_metrics=cast(
+                Optional[Dict[str, Optional[str]]], ml_model_version.run_details.metrics
+            )
+            if ml_model_version.run_details and ml_model_version.run_details.metrics
+            else None,
+            hyper_params=cast(
+                Optional[Dict[str, Optional[str]]],
+                ml_model_version.run_details.parameters,
+            )
+            if ml_model_version.run_details and ml_model_version.run_details.parameters
+            else None,
+            custom_properties=custom_properties if custom_properties else None,
             extra_aspects=extra_aspects,
         )
 

datahub/ingestion/source/usage/clickhouse_usage.py

@@ -115,7 +115,7 @@ class ClickHouseUsageSource(Source):
 
     @classmethod
     def create(cls, config_dict, ctx):
-        config = ClickHouseUsageConfig.parse_obj(config_dict)
+        config = ClickHouseUsageConfig.model_validate(config_dict)
         return cls(ctx, config)
 
     def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:

datahub/ingestion/source/usage/starburst_trino_usage.py

@@ -133,7 +133,7 @@ class TrinoUsageSource(Source):
 
     @classmethod
     def create(cls, config_dict, ctx):
-        config = TrinoUsageConfig.parse_obj(config_dict)
+        config = TrinoUsageConfig.model_validate(config_dict)
         return cls(ctx, config)
 
     def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:

datahub/ingestion/source/usage/usage_common.py

@@ -15,6 +15,7 @@ from typing import (
 )
 
 import pydantic
+from pydantic import ValidationInfo, field_validator
 from pydantic.fields import Field
 
 import datahub.emitter.mce_builder as builder
@@ -226,10 +227,11 @@ class BaseUsageConfig(BaseTimeWindowConfig):
         default=True, description="Whether to ingest the top_n_queries."
     )
 
-    @pydantic.validator("top_n_queries")
-    def ensure_top_n_queries_is_not_too_big(cls, v: int, values: dict) -> int:
+    @field_validator("top_n_queries", mode="after")
+    @classmethod
+    def ensure_top_n_queries_is_not_too_big(cls, v: int, info: ValidationInfo) -> int:
         minimum_query_size = 20
-
+        values = info.data
         max_queries = int(values["queries_character_limit"] / minimum_query_size)
         if v > max_queries:
             raise ValueError(

datahub/ingestion/source_config/csv_enricher.py

@@ -1,6 +1,5 @@
-from typing import Any, Dict
-
 import pydantic
+from pydantic import field_validator
 
 from datahub.configuration.common import ConfigModel
 
@@ -21,7 +20,8 @@ class CSVEnricherConfig(ConfigModel):
         description="Delimiter to use when parsing array fields (tags, terms and owners)",
     )
 
-    @pydantic.validator("write_semantics")
+    @field_validator("write_semantics", mode="after")
+    @classmethod
     def validate_write_semantics(cls, write_semantics: str) -> str:
         if write_semantics.lower() not in {"patch", "override"}:
             raise ValueError(
@@ -31,9 +31,10 @@ class CSVEnricherConfig(ConfigModel):
             )
         return write_semantics
 
-    @pydantic.validator("array_delimiter")
-    def validator_diff(cls, array_delimiter: str, values: Dict[str, Any]) -> str:
-        if array_delimiter == values["delimiter"]:
+    @field_validator("array_delimiter", mode="after")
+    @classmethod
+    def validator_diff(cls, array_delimiter: str, info: pydantic.ValidationInfo) -> str:
+        if array_delimiter == info.data["delimiter"]:
             raise ValueError(
                 "array_delimiter and delimiter are the same. Please choose different delimiters."
             )

datahub/ingestion/source_config/operation_config.py

@@ -3,7 +3,7 @@ import logging
 from typing import Any, Dict, Optional
 
 import cachetools
-import pydantic
+from pydantic import field_validator, model_validator
 from pydantic.fields import Field
 
 from datahub.configuration.common import ConfigModel
@@ -26,7 +26,8 @@ class OperationConfig(ConfigModel):
         description="Number between 1 to 31 for date of month (both inclusive). If not specified, defaults to Nothing and this field does not take affect.",
     )
 
-    @pydantic.root_validator(pre=True)
+    @model_validator(mode="before")
+    @classmethod
     def lower_freq_configs_are_set(cls, values: Dict[str, Any]) -> Dict[str, Any]:
         lower_freq_profile_enabled = values.get("lower_freq_profile_enabled")
         profile_day_of_week = values.get("profile_day_of_week")
@@ -41,7 +42,8 @@ class OperationConfig(ConfigModel):
             )
         return values
 
-    @pydantic.validator("profile_day_of_week")
+    @field_validator("profile_day_of_week", mode="after")
+    @classmethod
     def validate_profile_day_of_week(cls, v: Optional[int]) -> Optional[int]:
         profile_day_of_week = v
         if profile_day_of_week is None:
@@ -52,7 +54,8 @@ class OperationConfig(ConfigModel):
             )
         return profile_day_of_week
 
-    @pydantic.validator("profile_date_of_month")
+    @field_validator("profile_date_of_month", mode="after")
+    @classmethod
     def validate_profile_date_of_month(cls, v: Optional[int]) -> Optional[int]:
         profile_date_of_month = v
         if profile_date_of_month is None:

datahub/ingestion/source_config/pulsar.py

@@ -3,7 +3,7 @@ from typing import Dict, List, Optional, Union
 from urllib.parse import urlparse
 
 import pydantic
-from pydantic import Field, validator
+from pydantic import Field, model_validator
 
 from datahub.configuration.common import AllowDenyPattern
 from datahub.configuration.source_common import (
@@ -100,27 +100,23 @@ class PulsarSourceConfig(
         default_factory=dict, description="Placeholder for OpenId discovery document"
     )
 
-    @validator("token")
-    def ensure_only_issuer_or_token(
-        cls, token: Optional[str], values: Dict[str, Optional[str]]
-    ) -> Optional[str]:
-        if token is not None and values.get("issuer_url") is not None:
+    @model_validator(mode="after")
+    def ensure_only_issuer_or_token(self) -> "PulsarSourceConfig":
+        if self.token is not None and self.issuer_url is not None:
             raise ValueError(
                 "Expected only one authentication method, either issuer_url or token."
             )
-        return token
-
-    @validator("client_secret", always=True)
-    def ensure_client_id_and_secret_for_issuer_url(
-        cls, client_secret: Optional[str], values: Dict[str, Optional[str]]
-    ) -> Optional[str]:
-        if values.get("issuer_url") is not None and (
-            client_secret is None or values.get("client_id") is None
+        return self
+
+    @model_validator(mode="after")
+    def ensure_client_id_and_secret_for_issuer_url(self) -> "PulsarSourceConfig":
+        if self.issuer_url is not None and (
+            self.client_secret is None or self.client_id is None
         ):
             raise ValueError(
                 "Missing configuration: client_id and client_secret are mandatory when issuer_url is set."
             )
-        return client_secret
+        return self
 
     @pydantic.field_validator("web_service_url", mode="after")
     @classmethod

datahub/ingestion/transformer/add_dataset_browse_path.py

@@ -32,7 +32,7 @@ class AddDatasetBrowsePathTransformer(DatasetBrowsePathsTransformer):
     def create(
         cls, config_dict: dict, ctx: PipelineContext
     ) -> "AddDatasetBrowsePathTransformer":
-        config = AddDatasetBrowsePathConfig.parse_obj(config_dict)
+        config = AddDatasetBrowsePathConfig.model_validate(config_dict)
         return cls(config, ctx)
 
     @staticmethod