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

datahub/ingestion/source/state/stateful_ingestion_base.py

@@ -3,7 +3,7 @@ from dataclasses import dataclass
 from typing import Any, Dict, Generic, Optional, Type, TypeVar
 
 import pydantic
-from pydantic import root_validator
+from pydantic import model_validator
 from pydantic.fields import Field
 
 from datahub.configuration.common import (
@@ -73,14 +73,14 @@ class StatefulIngestionConfig(ConfigModel):
         description="If set to True, ignores the current checkpoint state.",
     )
 
-    @pydantic.root_validator(skip_on_failure=True)
-    def validate_config(cls, values: Dict[str, Any]) -> Dict[str, Any]:
-        if values.get("enabled"):
-            if values.get("state_provider") is None:
-                values["state_provider"] = DynamicTypedStateProviderConfig(
+    @model_validator(mode="after")
+    def validate_config(self) -> "StatefulIngestionConfig":
+        if self.enabled:
+            if self.state_provider is None:
+                self.state_provider = DynamicTypedStateProviderConfig(
                     type="datahub", config={}
                 )
-        return values
+        return self
 
 
 CustomConfig = TypeVar("CustomConfig", bound=StatefulIngestionConfig)
@@ -110,17 +110,19 @@ class StatefulLineageConfigMixin(ConfigModel):
         "store_last_lineage_extraction_timestamp", "enable_stateful_lineage_ingestion"
     )
 
-    @root_validator(skip_on_failure=True)
-    def lineage_stateful_option_validator(cls, values: Dict) -> Dict:
-        sti = values.get("stateful_ingestion")
-        if not sti or not sti.enabled:
-            if values.get("enable_stateful_lineage_ingestion"):
-                logger.warning(
-                    "Stateful ingestion is disabled, disabling enable_stateful_lineage_ingestion config option as well"
-                )
-                values["enable_stateful_lineage_ingestion"] = False
-
-        return values
+    @model_validator(mode="after")
+    def lineage_stateful_option_validator(self) -> "StatefulLineageConfigMixin":
+        try:
+            sti = getattr(self, "stateful_ingestion", None)
+            if not sti or not getattr(sti, "enabled", False):
+                if getattr(self, "enable_stateful_lineage_ingestion", False):
+                    logger.warning(
+                        "Stateful ingestion is disabled, disabling enable_stateful_lineage_ingestion config option as well"
+                    )
+                    self.enable_stateful_lineage_ingestion = False
+        except (AttributeError, RecursionError) as e:
+            logger.debug(f"Skipping stateful lineage validation due to: {e}")
+        return self
 
 
 class StatefulProfilingConfigMixin(ConfigModel):
@@ -135,16 +137,19 @@ class StatefulProfilingConfigMixin(ConfigModel):
         "store_last_profiling_timestamps", "enable_stateful_profiling"
     )
 
-    @root_validator(skip_on_failure=True)
-    def profiling_stateful_option_validator(cls, values: Dict) -> Dict:
-        sti = values.get("stateful_ingestion")
-        if not sti or not sti.enabled:
-            if values.get("enable_stateful_profiling"):
-                logger.warning(
-                    "Stateful ingestion is disabled, disabling enable_stateful_profiling config option as well"
-                )
-                values["enable_stateful_profiling"] = False
-        return values
+    @model_validator(mode="after")
+    def profiling_stateful_option_validator(self) -> "StatefulProfilingConfigMixin":
+        try:
+            sti = getattr(self, "stateful_ingestion", None)
+            if not sti or not getattr(sti, "enabled", False):
+                if getattr(self, "enable_stateful_profiling", False):
+                    logger.warning(
+                        "Stateful ingestion is disabled, disabling enable_stateful_profiling config option as well"
+                    )
+                    self.enable_stateful_profiling = False
+        except (AttributeError, RecursionError) as e:
+            logger.debug(f"Skipping stateful profiling validation due to: {e}")
+        return self
 
 
 class StatefulUsageConfigMixin(BaseTimeWindowConfig):
@@ -161,16 +166,21 @@ class StatefulUsageConfigMixin(BaseTimeWindowConfig):
         "store_last_usage_extraction_timestamp", "enable_stateful_usage_ingestion"
     )
 
-    @root_validator(skip_on_failure=True)
-    def last_usage_extraction_stateful_option_validator(cls, values: Dict) -> Dict:
-        sti = values.get("stateful_ingestion")
-        if not sti or not sti.enabled:
-            if values.get("enable_stateful_usage_ingestion"):
-                logger.warning(
-                    "Stateful ingestion is disabled, disabling enable_stateful_usage_ingestion config option as well"
-                )
-                values["enable_stateful_usage_ingestion"] = False
-        return values
+    @model_validator(mode="after")
+    def last_usage_extraction_stateful_option_validator(
+        self,
+    ) -> "StatefulUsageConfigMixin":
+        try:
+            sti = getattr(self, "stateful_ingestion", None)
+            if not sti or not getattr(sti, "enabled", False):
+                if getattr(self, "enable_stateful_usage_ingestion", False):
+                    logger.warning(
+                        "Stateful ingestion is disabled, disabling enable_stateful_usage_ingestion config option as well"
+                    )
+                    self.enable_stateful_usage_ingestion = False
+        except (AttributeError, RecursionError) as e:
+            logger.debug(f"Skipping stateful usage validation due to: {e}")
+        return self
 
 
 class StatefulTimeWindowConfigMixin(BaseTimeWindowConfig):
@@ -185,16 +195,16 @@ class StatefulTimeWindowConfigMixin(BaseTimeWindowConfig):
         "and queries together from a single audit log and uses a unified time window.",
     )
 
-    @root_validator(skip_on_failure=True)
-    def time_window_stateful_option_validator(cls, values: Dict) -> Dict:
-        sti = values.get("stateful_ingestion")
-        if not sti or not sti.enabled:
-            if values.get("enable_stateful_time_window"):
+    @model_validator(mode="after")
+    def time_window_stateful_option_validator(self) -> "StatefulTimeWindowConfigMixin":
+        sti = getattr(self, "stateful_ingestion", None)
+        if not sti or not getattr(sti, "enabled", False):
+            if getattr(self, "enable_stateful_time_window", False):
                 logger.warning(
                     "Stateful ingestion is disabled, disabling enable_stateful_time_window config option as well"
                 )
-                values["enable_stateful_time_window"] = False
-        return values
+                self.enable_stateful_time_window = False
+        return self
 
 
 @dataclass

datahub/ingestion/source/state_provider/datahub_ingestion_checkpointing_provider.py

@@ -40,7 +40,7 @@ class DatahubIngestionCheckpointingProvider(IngestionCheckpointingProviderBase):
     def create(
         cls, config_dict: Dict[str, Any], ctx: PipelineContext
     ) -> "DatahubIngestionCheckpointingProvider":
-        config = DatahubIngestionStateProviderConfig.parse_obj(config_dict)
+        config = DatahubIngestionStateProviderConfig.model_validate(config_dict)
         if config.datahub_api is not None:
             return cls(DataHubGraph(config.datahub_api))
         elif ctx.graph:

datahub/ingestion/source/state_provider/file_ingestion_checkpointing_provider.py

@@ -32,7 +32,7 @@ class FileIngestionCheckpointingProvider(IngestionCheckpointingProviderBase):
     def create(
         cls, config_dict: Dict[str, Any], ctx: PipelineContext
     ) -> "FileIngestionCheckpointingProvider":
-        config = FileIngestionStateProviderConfig.parse_obj(config_dict)
+        config = FileIngestionStateProviderConfig.model_validate(config_dict)
         return cls(config)
 
     def get_latest_checkpoint(

datahub/ingestion/source/superset.py

@@ -9,7 +9,7 @@ from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 import dateutil.parser as dp
 import requests
 import sqlglot
-from pydantic import BaseModel, root_validator, validator
+from pydantic import BaseModel, field_validator, model_validator
 from pydantic.fields import Field
 from requests.adapters import HTTPAdapter
 from urllib3.util.retry import Retry
@@ -246,16 +246,16 @@ class SupersetConfig(
         # This is required to allow preset configs to get parsed
         extra = "allow"
 
-    @validator("connect_uri", "display_uri")
-    def remove_trailing_slash(cls, v):
+    @field_validator("connect_uri", "display_uri", mode="after")
+    @classmethod
+    def remove_trailing_slash(cls, v: str) -> str:
         return config_clean.remove_trailing_slashes(v)
 
-    @root_validator(skip_on_failure=True)
-    def default_display_uri_to_connect_uri(cls, values):
-        base = values.get("display_uri")
-        if base is None:
-            values["display_uri"] = values.get("connect_uri")
-        return values
+    @model_validator(mode="after")
+    def default_display_uri_to_connect_uri(self) -> "SupersetConfig":
+        if self.display_uri is None:
+            self.display_uri = self.connect_uri
+        return self
 
 
 def get_metric_name(metric):

datahub/ingestion/source/tableau/tableau.py

@@ -25,7 +25,7 @@ from urllib.parse import quote, urlparse
 
 import dateutil.parser as dp
 import tableauserverclient as TSC
-from pydantic import root_validator, validator
+from pydantic import field_validator, model_validator
 from pydantic.fields import Field
 from requests.adapters import HTTPAdapter
 from tableauserverclient import (
@@ -257,8 +257,9 @@ class TableauConnectionConfig(ConfigModel):
         description="When enabled, extracts column-level lineage from Tableau Datasources",
     )
 
-    @validator("connect_uri")
-    def remove_trailing_slash(cls, v):
+    @field_validator("connect_uri", mode="after")
+    @classmethod
+    def remove_trailing_slash(cls, v: str) -> str:
         return config_clean.remove_trailing_slashes(v)
 
     def get_tableau_auth(
@@ -652,8 +653,9 @@ class TableauConfig(
         "fetch_size",
     )
 
-    # pre = True because we want to take some decision before pydantic initialize the configuration to default values
-    @root_validator(pre=True)
+    # mode = "before" because we want to take some decision before pydantic initialize the configuration to default values
+    @model_validator(mode="before")
+    @classmethod
     def projects_backward_compatibility(cls, values: Dict) -> Dict:
         # In-place update of the input dict would cause state contamination. This was discovered through test failures
         # in test_hex.py where the same dict is reused.
@@ -683,27 +685,23 @@ class TableauConfig(
 
         return values
 
-    @root_validator(skip_on_failure=True)
-    def validate_config_values(cls, values: Dict) -> Dict:
-        tags_for_hidden_assets = values.get("tags_for_hidden_assets")
-        ingest_tags = values.get("ingest_tags")
+    @model_validator(mode="after")
+    def validate_config_values(self) -> "TableauConfig":
         if (
-            not ingest_tags
-            and tags_for_hidden_assets
-            and len(tags_for_hidden_assets) > 0
+            not self.ingest_tags
+            and self.tags_for_hidden_assets
+            and len(self.tags_for_hidden_assets) > 0
         ):
             raise ValueError(
                 "tags_for_hidden_assets is only allowed with ingest_tags enabled. Be aware that this will overwrite tags entered from the UI."
             )
 
-        use_email_as_username = values.get("use_email_as_username")
-        ingest_owner = values.get("ingest_owner")
-        if use_email_as_username and not ingest_owner:
+        if self.use_email_as_username and not self.ingest_owner:
             raise ValueError(
                 "use_email_as_username requires ingest_owner to be enabled."
             )
 
-        return values
+        return self
 
 
 class WorkbookKey(ContainerKey):

datahub/ingestion/source/unity/azure_auth_config.py

@@ -0,0 +1,15 @@
+from pydantic import Field, SecretStr
+
+from datahub.configuration import ConfigModel
+
+
+class AzureAuthConfig(ConfigModel):
+    client_secret: SecretStr = Field(
+        description="Azure application client secret used for authentication. This is a confidential credential that should be kept secure."
+    )
+    client_id: str = Field(
+        description="Azure application (client) ID. This is the unique identifier for the registered Azure AD application.",
+    )
+    tenant_id: str = Field(
+        description="Azure tenant (directory) ID. This identifies the Azure AD tenant where the application is registered.",
+    )

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: