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

datahub/ingestion/source/sql/sql_config.py

@@ -3,7 +3,7 @@ from abc import abstractmethod
 from typing import Any, Dict, Optional
 
 import pydantic
-from pydantic import Field
+from pydantic import Field, model_validator
 
 from datahub.configuration.common import AllowDenyPattern, ConfigModel
 from datahub.configuration.source_common import (
@@ -49,7 +49,8 @@ class SQLFilterConfig(ConfigModel):
         description="Regex patterns for views to filter in ingestion. Note: Defaults to table_pattern if not specified. Specify regex to match the entire view name in database.schema.view format. e.g. to match all views starting with customer in Customer database and public schema, use the regex 'Customer.public.customer.*'",
     )
 
-    @pydantic.root_validator(pre=True)
+    @model_validator(mode="before")
+    @classmethod
     def view_pattern_is_table_pattern_unless_specified(
         cls, values: Dict[str, Any]
     ) -> Dict[str, Any]:
@@ -120,11 +121,9 @@ class SQLCommonConfig(
             self.profiling.operation_config
         )
 
-    @pydantic.root_validator(skip_on_failure=True)
-    def ensure_profiling_pattern_is_passed_to_profiling(
-        cls, values: Dict[str, Any]
-    ) -> Dict[str, Any]:
-        profiling: Optional[GEProfilingConfig] = values.get("profiling")
+    @model_validator(mode="after")
+    def ensure_profiling_pattern_is_passed_to_profiling(self):
+        profiling = self.profiling
         # Note: isinstance() check is required here as unity-catalog source reuses
         # SQLCommonConfig with different profiling config than GEProfilingConfig
         if (
@@ -132,8 +131,8 @@ class SQLCommonConfig(
             and isinstance(profiling, GEProfilingConfig)
             and profiling.enabled
         ):
-            profiling._allow_deny_patterns = values["profile_pattern"]
-        return values
+            profiling._allow_deny_patterns = self.profile_pattern
+        return self
 
     @abstractmethod
     def get_sql_alchemy_url(self):

datahub/ingestion/source/sql/sql_generic.py

@@ -85,5 +85,5 @@ class SQLAlchemyGenericSource(SQLAlchemySource):
 
     @classmethod
     def create(cls, config_dict, ctx):
-        config = SQLAlchemyGenericConfig.parse_obj(config_dict)
+        config = SQLAlchemyGenericConfig.model_validate(config_dict)
         return cls(config, ctx)

datahub/ingestion/source/sql/teradata.py

@@ -860,7 +860,7 @@ ORDER by DataBaseName, TableName;
 
     @classmethod
     def create(cls, config_dict, ctx):
-        config = TeradataConfig.parse_obj(config_dict)
+        config = TeradataConfig.model_validate(config_dict)
         return cls(config, ctx)
 
     def _init_schema_resolver(self) -> SchemaResolver:

datahub/ingestion/source/sql/trino.py

@@ -413,7 +413,7 @@ class TrinoSource(SQLAlchemySource):
 
     @classmethod
     def create(cls, config_dict, ctx):
-        config = TrinoConfig.parse_obj(config_dict)
+        config = TrinoConfig.model_validate(config_dict)
         return cls(config, ctx)
 
     def get_schema_fields_for_column(

datahub/ingestion/source/sql/vertica.py

@@ -5,7 +5,7 @@ from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Set, Tupl
 
 import pydantic
 import pytest
-from pydantic import validator
+from pydantic import field_validator
 from vertica_sqlalchemy_dialect.base import VerticaInspector
 
 from datahub.configuration.common import AllowDenyPattern
@@ -105,8 +105,9 @@ class VerticaConfig(BasicSQLAlchemyConfig):
     # defaults
     scheme: str = pydantic.Field(default="vertica+vertica_python")
 
-    @validator("host_port")
-    def clean_host_port(cls, v):
+    @field_validator("host_port", mode="after")
+    @classmethod
+    def clean_host_port(cls, v: str) -> str:
         return config_clean.remove_protocol(v)
 
 
@@ -138,7 +139,7 @@ class VerticaSource(SQLAlchemySource):
 
     @classmethod
     def create(cls, config_dict: Dict, ctx: PipelineContext) -> "VerticaSource":
-        config = VerticaConfig.parse_obj(config_dict)
+        config = VerticaConfig.model_validate(config_dict)
         return cls(config, ctx)
 
     def get_workunits_internal(self) -> Iterable[Union[MetadataWorkUnit, SqlWorkUnit]]:

datahub/ingestion/source/sql_queries.py

@@ -5,10 +5,10 @@ import re
 from dataclasses import dataclass, field
 from datetime import datetime
 from functools import partial
-from typing import ClassVar, Iterable, List, Optional, Union, cast
+from typing import Any, ClassVar, Iterable, List, Optional, Union, cast
 
 import smart_open
-from pydantic import BaseModel, Field, validator
+from pydantic import BaseModel, Field, field_validator
 
 from datahub.configuration.common import HiddenFromDocs
 from datahub.configuration.datetimes import parse_user_datetime
@@ -450,19 +450,22 @@ class QueryEntry(BaseModel):
     class Config:
         arbitrary_types_allowed = True
 
-    @validator("timestamp", pre=True)
-    def parse_timestamp(cls, v):
+    @field_validator("timestamp", mode="before")
+    @classmethod
+    def parse_timestamp(cls, v: Any) -> Any:
         return None if v is None else parse_user_datetime(str(v))
 
-    @validator("user", pre=True)
-    def parse_user(cls, v):
+    @field_validator("user", mode="before")
+    @classmethod
+    def parse_user(cls, v: Any) -> Any:
         if v is None:
             return None
 
         return v if isinstance(v, CorpUserUrn) else CorpUserUrn(v)
 
-    @validator("downstream_tables", "upstream_tables", pre=True)
-    def parse_tables(cls, v):
+    @field_validator("downstream_tables", "upstream_tables", mode="before")
+    @classmethod
+    def parse_tables(cls, v: Any) -> Any:
         if not v:
             return []
 

datahub/ingestion/source/state/checkpoint.py

@@ -163,7 +163,7 @@ class Checkpoint(Generic[StateType]):
         )
         state_as_dict["version"] = checkpoint_aspect.state.formatVersion
         state_as_dict["serde"] = checkpoint_aspect.state.serde
-        return state_class.parse_obj(state_as_dict)
+        return state_class.model_validate(state_as_dict)
 
     @staticmethod
     def _from_base85_json_bytes(
@@ -179,7 +179,7 @@ class Checkpoint(Generic[StateType]):
         state_as_dict = json.loads(state_uncompressed.decode("utf-8"))
         state_as_dict["version"] = checkpoint_aspect.state.formatVersion
         state_as_dict["serde"] = checkpoint_aspect.state.serde
-        return state_class.parse_obj(state_as_dict)
+        return state_class.model_validate(state_as_dict)
 
     def to_checkpoint_aspect(
         self, max_allowed_state_size: int

datahub/ingestion/source/state/entity_removal_state.py

@@ -1,6 +1,7 @@
 from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Tuple, Type
 
 import pydantic
+from pydantic import model_validator
 
 from datahub.emitter.mce_builder import make_assertion_urn, make_container_urn
 from datahub.ingestion.source.state.checkpoint import CheckpointStateBase
@@ -59,7 +60,7 @@ def pydantic_state_migrator(mapping: Dict[str, str]) -> "V1RootValidator":
 
         return values
 
-    return pydantic.root_validator(pre=True, allow_reuse=True)(_validate_field_rename)
+    return model_validator(mode="before")(_validate_field_rename)
 
 
 class GenericCheckpointState(CheckpointStateBase):

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/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,8 @@ class UnityCatalogSourceConfig(
             )
         return workspace_url
 
-    @pydantic.validator("include_metastore")
+    @field_validator("include_metastore", mode="after")
+    @classmethod
     def include_metastore_warning(cls, v: bool) -> bool:
         if v:
             msg = (
@@ -424,60 +427,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: