acryl-datahub 1.3.0.1rc9__py3-none-any.whl → 1.3.1.1__py3-none-any.whl

datahub/ingestion/source/unity/config.py

@@ -1,10 +1,10 @@
 import logging
 import os
 from datetime import datetime, timedelta, timezone
-from typing import Any, Dict, List, Optional, Union
+from typing import Dict, List, Optional, Union
 
 import pydantic
-from pydantic import Field
+from pydantic import Field, field_validator, model_validator
 from typing_extensions import Literal
 
 from datahub.configuration.common import (
@@ -397,13 +397,15 @@ class UnityCatalogSourceConfig(
         default=None, description="Unity Catalog Stateful Ingestion Config."
     )
 
-    @pydantic.validator("start_time")
+    @field_validator("start_time", mode="after")
+    @classmethod
     def within_thirty_days(cls, v: datetime) -> datetime:
         if (datetime.now(timezone.utc) - v).days > 30:
             raise ValueError("Query history is only maintained for 30 days.")
         return v
 
-    @pydantic.validator("workspace_url")
+    @field_validator("workspace_url", mode="after")
+    @classmethod
     def workspace_url_should_start_with_http_scheme(cls, workspace_url: str) -> str:
         if not workspace_url.lower().startswith(("http://", "https://")):
             raise ValueError(
@@ -411,7 +413,26 @@ class UnityCatalogSourceConfig(
             )
         return workspace_url
 
-    @pydantic.validator("include_metastore")
+    @model_validator(mode="before")
+    def either_token_or_azure_auth_provided(cls, values: dict) -> dict:
+        token = values.get("token")
+        azure_auth = values.get("azure_auth")
+
+        # Check if exactly one of the authentication methods is provided
+        if not token and not azure_auth:
+            raise ValueError(
+                "Either 'azure_auth' or 'token' (personal access token) must be provided in the configuration."
+            )
+
+        if token and azure_auth:
+            raise ValueError(
+                "Cannot specify both 'token' and 'azure_auth'. Please provide only one authentication method."
+            )
+
+        return values
+
+    @field_validator("include_metastore", mode="after")
+    @classmethod
     def include_metastore_warning(cls, v: bool) -> bool:
         if v:
             msg = (
@@ -424,60 +445,56 @@ class UnityCatalogSourceConfig(
             add_global_warning(msg)
         return v
 
-    @pydantic.root_validator(skip_on_failure=True)
-    def set_warehouse_id_from_profiling(cls, values: Dict[str, Any]) -> Dict[str, Any]:
-        profiling: Optional[
-            Union[UnityCatalogGEProfilerConfig, UnityCatalogAnalyzeProfilerConfig]
-        ] = values.get("profiling")
-        if not values.get("warehouse_id") and profiling and profiling.warehouse_id:
-            values["warehouse_id"] = profiling.warehouse_id
+    @model_validator(mode="after")
+    def set_warehouse_id_from_profiling(self):
+        profiling = self.profiling
+        if not self.warehouse_id and profiling and profiling.warehouse_id:
+            self.warehouse_id = profiling.warehouse_id
         if (
-            values.get("warehouse_id")
+            self.warehouse_id
             and profiling
             and profiling.warehouse_id
-            and values["warehouse_id"] != profiling.warehouse_id
+            and self.warehouse_id != profiling.warehouse_id
         ):
             raise ValueError(
                 "When `warehouse_id` is set, it must match the `warehouse_id` in `profiling`."
             )
 
-        if values.get("warehouse_id") and profiling and not profiling.warehouse_id:
-            profiling.warehouse_id = values["warehouse_id"]
+        if self.warehouse_id and profiling and not profiling.warehouse_id:
+            profiling.warehouse_id = self.warehouse_id
 
         if profiling and profiling.enabled and not profiling.warehouse_id:
             raise ValueError("warehouse_id must be set when profiling is enabled.")
 
-        return values
+        return self
 
-    @pydantic.root_validator(skip_on_failure=True)
-    def validate_lineage_data_source_with_warehouse(
-        cls, values: Dict[str, Any]
-    ) -> Dict[str, Any]:
-        lineage_data_source = values.get("lineage_data_source", LineageDataSource.AUTO)
-        warehouse_id = values.get("warehouse_id")
+    @model_validator(mode="after")
+    def validate_lineage_data_source_with_warehouse(self):
+        lineage_data_source = self.lineage_data_source or LineageDataSource.AUTO
 
-        if lineage_data_source == LineageDataSource.SYSTEM_TABLES and not warehouse_id:
+        if (
+            lineage_data_source == LineageDataSource.SYSTEM_TABLES
+            and not self.warehouse_id
+        ):
             raise ValueError(
                 f"lineage_data_source='{LineageDataSource.SYSTEM_TABLES.value}' requires warehouse_id to be set"
             )
 
-        return values
+        return self
 
-    @pydantic.root_validator(skip_on_failure=True)
-    def validate_usage_data_source_with_warehouse(
-        cls, values: Dict[str, Any]
-    ) -> Dict[str, Any]:
-        usage_data_source = values.get("usage_data_source", UsageDataSource.AUTO)
-        warehouse_id = values.get("warehouse_id")
+    @model_validator(mode="after")
+    def validate_usage_data_source_with_warehouse(self):
+        usage_data_source = self.usage_data_source or UsageDataSource.AUTO
 
-        if usage_data_source == UsageDataSource.SYSTEM_TABLES and not warehouse_id:
+        if usage_data_source == UsageDataSource.SYSTEM_TABLES and not self.warehouse_id:
             raise ValueError(
                 f"usage_data_source='{UsageDataSource.SYSTEM_TABLES.value}' requires warehouse_id to be set"
             )
 
-        return values
+        return self
 
-    @pydantic.validator("schema_pattern", always=True)
+    @field_validator("schema_pattern", mode="before")
+    @classmethod
     def schema_pattern_should__always_deny_information_schema(
         cls, v: AllowDenyPattern
     ) -> AllowDenyPattern:

datahub/ingestion/source/unity/connection.py

@@ -8,6 +8,7 @@ from pydantic import Field
 
 from datahub.configuration.common import ConfigModel
 from datahub.ingestion.source.sql.sqlalchemy_uri import make_sqlalchemy_uri
+from datahub.ingestion.source.unity.azure_auth_config import AzureAuthConfig
 
 DATABRICKS = "databricks"
 
@@ -19,7 +20,12 @@ class UnityCatalogConnectionConfig(ConfigModel):
     """
 
     scheme: str = DATABRICKS
-    token: str = pydantic.Field(description="Databricks personal access token")
+    token: Optional[str] = pydantic.Field(
+        default=None, description="Databricks personal access token"
+    )
+    azure_auth: Optional[AzureAuthConfig] = Field(
+        default=None, description="Azure configuration"
+    )
     workspace_url: str = pydantic.Field(
         description="Databricks workspace url. e.g. https://my-workspace.cloud.databricks.com"
     )

datahub/ingestion/source/unity/connection_test.py

@@ -16,10 +16,10 @@ class UnityCatalogConnectionTest:
         self.report = UnityCatalogReport()
         self.proxy = UnityCatalogApiProxy(
             self.config.workspace_url,
-            self.config.token,
             self.config.profiling.warehouse_id,
             report=self.report,
             databricks_api_page_size=self.config.databricks_api_page_size,
+            personal_access_token=self.config.token,
         )
 
     def get_connection_test(self) -> TestConnectionReport:

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]: