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

datahub/ingestion/source/data_lake_common/path_spec.py

@@ -3,11 +3,11 @@ import logging
 import os
 import re
 from enum import Enum
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import List, Optional, Tuple, Union
 
 import parse
-import pydantic
 from cached_property import cached_property
+from pydantic import field_validator, model_validator
 from pydantic.fields import Field
 from wcmatch import pathlib
 
@@ -65,7 +65,8 @@ class SortKey(ConfigModel):
         description="The date format to use when sorting. This is used to parse the date from the key. The format should follow the java [SimpleDateFormat](https://docs.oracle.com/javase/8/docs/api/java/text/SimpleDateFormat.html) format.",
     )
 
-    @pydantic.validator("date_format", always=True)
+    @field_validator("date_format", mode="before")
+    @classmethod
     def convert_date_format_to_python_format(cls, v: Optional[str]) -> Optional[str]:
         if v is None:
             return None
@@ -86,7 +87,7 @@ class PathSpec(ConfigModel):
         arbitrary_types_allowed = True
 
     include: str = Field(
-        description="Path to table. Name variable `{table}` is used to mark the folder with dataset. In absence of `{table}`, file level dataset will be created. Check below examples for more details."
+        description="Path to table. Name variable `{table}` is used to mark the folder with dataset. In absence of `{table}`, file level dataset will be created. Check below examples for more details.",
     )
     exclude: Optional[List[str]] = Field(
         [],
@@ -260,20 +261,80 @@ class PathSpec(ConfigModel):
     ) -> Union[None, parse.Result, parse.Match]:
         return self.compiled_folder_include.parse(path)
 
-    @pydantic.root_validator(skip_on_failure=True)
-    def validate_no_double_stars(cls, values: Dict) -> Dict:
-        if "include" not in values:
-            return values
+    @model_validator(mode="after")
+    def validate_path_spec_comprehensive(self):
+        """
+        Comprehensive model validator that handles multiple interdependent validations.
+
+        Consolidates related validation logic to avoid order dependencies between multiple
+        model validators and ensures reliable cross-field validation. This approach is
+        preferred over multiple separate validators when:
+
+        1. Validations depend on multiple fields (e.g., sample_files depends on include)
+        2. One validation modifies a field that another validation checks
+        3. Field validators can't reliably access other field values or defaults
+        4. Order of execution between validators is important but undefined
+
+        By combining related validations, we ensure they execute in the correct sequence
+        and have access to all field values after Pydantic has processed defaults.
+        """
+        # Handle autodetect_partitions logic first
+        if self.autodetect_partitions:
+            include = self.include
+            if include.endswith("/"):
+                include = include[:-1]
+            if include.endswith("{table}"):
+                self.include = include + "/**"
+                # Allow double stars when we add them for autodetect_partitions
+                self.allow_double_stars = True
+
+        # Handle table_name logic
+        if self.table_name is None and "{table}" in self.include:
+            self.table_name = "{table}"
+        elif self.table_name is not None:
+            parsable_include = PathSpec.get_parsable_include(self.include)
+            compiled_include = parse.compile(parsable_include)
+            if not all(
+                x in compiled_include.named_fields
+                for x in parse.compile(self.table_name).named_fields
+            ):
+                raise ValueError(
+                    f"Not all named variables used in path_spec.table_name {self.table_name} are specified in path_spec.include {self.include}"
+                )
+
+        # Handle sample_files logic - turn off sampling for non-cloud URIs
+        is_s3 = is_s3_uri(self.include)
+        is_gcs = is_gcs_uri(self.include)
+        is_abs = is_abs_uri(self.include)
+        if not is_s3 and not is_gcs and not is_abs:
+            # Sampling only makes sense on s3 and gcs currently
+            self.sample_files = False
+
+        # Validate double stars
+        if "**" in self.include and not self.allow_double_stars:
+            raise ValueError("path_spec.include cannot contain '**'")
 
+        # Validate file extension
+        include_ext = os.path.splitext(self.include)[1].strip(".")
+        if not include_ext:
+            include_ext = (
+                "*"  # if no extension is provided, we assume all files are allowed
+            )
         if (
-            values.get("include")
-            and "**" in values["include"]
-            and not values.get("allow_double_stars")
+            include_ext not in self.file_types
+            and include_ext not in ["*", ""]
+            and not self.default_extension
+            and include_ext not in SUPPORTED_COMPRESSIONS
         ):
-            raise ValueError("path_spec.include cannot contain '**'")
-        return values
+            raise ValueError(
+                f"file type specified ({include_ext}) in path_spec.include is not in specified file "
+                f'types. Please select one from {self.file_types} or specify ".*" to allow all types'
+            )
 
-    @pydantic.validator("file_types", always=True)
+        return self
+
+    @field_validator("file_types", mode="before")
+    @classmethod
     def validate_file_types(cls, v: Optional[List[str]]) -> List[str]:
         if v is None:
             return SUPPORTED_FILE_TYPES
@@ -285,50 +346,24 @@ class PathSpec(ConfigModel):
                     )
             return v
 
-    @pydantic.validator("default_extension")
-    def validate_default_extension(cls, v):
+    @field_validator("default_extension", mode="after")
+    @classmethod
+    def validate_default_extension(cls, v: Optional[str]) -> Optional[str]:
         if v is not None and v not in SUPPORTED_FILE_TYPES:
             raise ValueError(
                 f"default extension {v} not in supported default file extension. Please specify one from {SUPPORTED_FILE_TYPES}"
             )
         return v
 
-    @pydantic.validator("sample_files", always=True)
-    def turn_off_sampling_for_non_s3(cls, v, values):
-        is_s3 = is_s3_uri(values.get("include") or "")
-        is_gcs = is_gcs_uri(values.get("include") or "")
-        is_abs = is_abs_uri(values.get("include") or "")
-        if not is_s3 and not is_gcs and not is_abs:
-            # Sampling only makes sense on s3 and gcs currently
-            v = False
-        return v
-
-    @pydantic.validator("exclude", each_item=True)
-    def no_named_fields_in_exclude(cls, v: str) -> str:
-        if len(parse.compile(v).named_fields) != 0:
-            raise ValueError(
-                f"path_spec.exclude {v} should not contain any named variables"
-            )
-        return v
-
-    @pydantic.validator("table_name", always=True)
-    def table_name_in_include(cls, v, values):
-        if "include" not in values:
-            return v
-
-        parsable_include = PathSpec.get_parsable_include(values["include"])
-        compiled_include = parse.compile(parsable_include)
-
+    @field_validator("exclude", mode="after")
+    @classmethod
+    def no_named_fields_in_exclude(cls, v: Optional[List[str]]) -> Optional[List[str]]:
         if v is None:
-            if "{table}" in values["include"]:
-                v = "{table}"
-        else:
-            if not all(
-                x in compiled_include.named_fields
-                for x in parse.compile(v).named_fields
-            ):
+            return v
+        for item in v:
+            if len(parse.compile(item).named_fields) != 0:
                 raise ValueError(
-                    f"Not all named variables used in path_spec.table_name {v} are specified in path_spec.include {values['include']}"
+                    f"path_spec.exclude {item} should not contain any named variables"
                 )
         return v
 
@@ -479,45 +514,6 @@ class PathSpec(ConfigModel):
         logger.debug(f"Setting _glob_include: {glob_include}")
         return glob_include
 
-    @pydantic.root_validator(skip_on_failure=True)
-    @staticmethod
-    def validate_path_spec(values: Dict) -> Dict[str, Any]:
-        # validate that main fields are populated
-        required_fields = ["include", "file_types", "default_extension"]
-        for f in required_fields:
-            if f not in values:
-                logger.debug(
-                    f"Failed to validate because {f} wasn't populated correctly"
-                )
-                return values
-
-        if values["include"] and values["autodetect_partitions"]:
-            include = values["include"]
-            if include.endswith("/"):
-                include = include[:-1]
-
-            if include.endswith("{table}"):
-                values["include"] = include + "/**"
-
-        include_ext = os.path.splitext(values["include"])[1].strip(".")
-        if not include_ext:
-            include_ext = (
-                "*"  # if no extension is provided, we assume all files are allowed
-            )
-
-        if (
-            include_ext not in values["file_types"]
-            and include_ext not in ["*", ""]
-            and not values["default_extension"]
-            and include_ext not in SUPPORTED_COMPRESSIONS
-        ):
-            raise ValueError(
-                f"file type specified ({include_ext}) in path_spec.include is not in specified file "
-                f'types. Please select one from {values.get("file_types")} or specify ".*" to allow all types'
-            )
-
-        return values
-
     def _extract_table_name(self, named_vars: dict) -> str:
         if self.table_name is None:
             raise ValueError("path_spec.table_name is not set")

datahub/ingestion/source/datahub/config.py

@@ -2,7 +2,7 @@ import os
 from typing import Optional, Set
 
 import pydantic
-from pydantic import Field, root_validator
+from pydantic import Field, model_validator
 
 from datahub.configuration.common import AllowDenyPattern, HiddenFromDocs
 from datahub.configuration.kafka import KafkaConsumerConnectionConfig
@@ -132,20 +132,20 @@ class DataHubSourceConfig(StatefulIngestionConfigBase):
         default=True, description="Copy system metadata from the source system"
     )
 
-    @root_validator(skip_on_failure=True)
-    def check_ingesting_data(cls, values):
+    @model_validator(mode="after")
+    def check_ingesting_data(self):
         if (
-            not values.get("database_connection")
-            and not values.get("kafka_connection")
-            and not values.get("pull_from_datahub_api")
+            not self.database_connection
+            and not self.kafka_connection
+            and not self.pull_from_datahub_api
         ):
             raise ValueError(
                 "Your current config will not ingest any data."
                 " Please specify at least one of `database_connection` or `kafka_connection`, ideally both."
             )
-        return values
+        return self
 
-    @pydantic.validator("database_connection")
+    @pydantic.field_validator("database_connection")
     def validate_mysql_scheme(
         cls, v: SQLAlchemyConnectionConfig
     ) -> SQLAlchemyConnectionConfig:

datahub/ingestion/source/datahub/datahub_source.py

@@ -62,7 +62,7 @@ class DataHubSource(StatefulIngestionSourceBase):
 
     @classmethod
     def create(cls, config_dict: Dict, ctx: PipelineContext) -> "DataHubSource":
-        config: DataHubSourceConfig = DataHubSourceConfig.parse_obj(config_dict)
+        config: DataHubSourceConfig = DataHubSourceConfig.model_validate(config_dict)
         return cls(config, ctx)
 
     def get_report(self) -> SourceReport:

datahub/ingestion/source/dbt/dbt_cloud.py

@@ -1,4 +1,5 @@
 import logging
+from copy import deepcopy
 from datetime import datetime
 from json import JSONDecodeError
 from typing import Dict, List, Literal, Optional, Tuple
@@ -6,7 +7,7 @@ from urllib.parse import urlparse
 
 import dateutil.parser
 import requests
-from pydantic import Field, root_validator
+from pydantic import Field, model_validator
 
 from datahub.ingestion.api.decorators import (
     SourceCapability,
@@ -68,8 +69,13 @@ class DBTCloudConfig(DBTCommonConfig):
         description='Where should the "View in dbt" link point to - either the "Explore" UI or the dbt Cloud IDE',
     )
 
-    @root_validator(pre=True)
+    @model_validator(mode="before")
+    @classmethod
     def set_metadata_endpoint(cls, values: dict) -> dict:
+        # In-place update of the input dict would cause state contamination.
+        # So a deepcopy is performed first.
+        values = deepcopy(values)
+
         if values.get("access_url") and not values.get("metadata_endpoint"):
             metadata_endpoint = infer_metadata_endpoint(values["access_url"])
             if metadata_endpoint is None:
@@ -271,7 +277,7 @@ class DBTCloudSource(DBTSourceBase, TestableSource):
 
     @classmethod
     def create(cls, config_dict, ctx):
-        config = DBTCloudConfig.parse_obj(config_dict)
+        config = DBTCloudConfig.model_validate(config_dict)
         return cls(config, ctx)
 
     @staticmethod

datahub/ingestion/source/dbt/dbt_common.py

@@ -1,6 +1,7 @@
 import logging
 import re
 from abc import abstractmethod
+from copy import deepcopy
 from dataclasses import dataclass, field
 from datetime import datetime
 from enum import auto
@@ -8,7 +9,7 @@ from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union
 
 import more_itertools
 import pydantic
-from pydantic import root_validator, validator
+from pydantic import field_validator, model_validator
 from pydantic.fields import Field
 
 from datahub.api.entities.dataprocess.dataprocess_instance import (
@@ -194,22 +195,26 @@ class DBTEntitiesEnabled(ConfigModel):
         "Only supported with dbt core.",
     )
 
-    @root_validator(skip_on_failure=True)
-    def process_only_directive(cls, values):
+    @model_validator(mode="after")
+    def process_only_directive(self) -> "DBTEntitiesEnabled":
         # Checks that at most one is set to ONLY, and then sets the others to NO.
-
-        only_values = [k for k in values if values.get(k) == EmitDirective.ONLY]
+        only_values = [
+            k for k, v in self.model_dump().items() if v == EmitDirective.ONLY
+        ]
         if len(only_values) > 1:
             raise ValueError(
                 f"Cannot have more than 1 type of entity emission set to ONLY. Found {only_values}"
             )
 
         if len(only_values) == 1:
-            for k in values:
-                values[k] = EmitDirective.NO
-            values[only_values[0]] = EmitDirective.YES
+            # Set all fields to NO first
+            for field_name in self.model_dump():
+                setattr(self, field_name, EmitDirective.NO)
 
-        return values
+            # Set the ONLY one to YES
+            setattr(self, only_values[0], EmitDirective.YES)
+
+        return self
 
     def _node_type_allow_map(self):
         # Node type comes from dbt's node types.
@@ -412,7 +417,8 @@ class DBTCommonConfig(
         "This ensures that lineage is generated reliably, but will lose any documentation associated only with the source.",
     )
 
-    @validator("target_platform")
+    @field_validator("target_platform", mode="after")
+    @classmethod
     def validate_target_platform_value(cls, target_platform: str) -> str:
         if target_platform.lower() == DBT_PLATFORM:
             raise ValueError(
@@ -421,15 +427,21 @@ class DBTCommonConfig(
             )
         return target_platform
 
-    @root_validator(pre=True)
+    @model_validator(mode="before")
+    @classmethod
     def set_convert_column_urns_to_lowercase_default_for_snowflake(
         cls, values: dict
     ) -> dict:
+        # In-place update of the input dict would cause state contamination.
+        # So a deepcopy is performed first.
+        values = deepcopy(values)
+
         if values.get("target_platform", "").lower() == "snowflake":
             values.setdefault("convert_column_urns_to_lowercase", True)
         return values
 
-    @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(
@@ -439,10 +451,9 @@ class DBTCommonConfig(
             )
         return write_semantics
 
-    @validator("meta_mapping")
-    def meta_mapping_validator(
-        cls, meta_mapping: Dict[str, Any], values: Dict, **kwargs: Any
-    ) -> Dict[str, Any]:
+    @field_validator("meta_mapping", mode="after")
+    @classmethod
+    def meta_mapping_validator(cls, meta_mapping: Dict[str, Any]) -> Dict[str, Any]:
         for k, v in meta_mapping.items():
             if "match" not in v:
                 raise ValueError(
@@ -458,44 +469,35 @@ class DBTCommonConfig(
                     mce_builder.validate_ownership_type(owner_category)
         return meta_mapping
 
-    @validator("include_column_lineage")
-    def validate_include_column_lineage(
-        cls, include_column_lineage: bool, values: Dict
-    ) -> bool:
-        if include_column_lineage and not values.get("infer_dbt_schemas"):
+    @model_validator(mode="after")
+    def validate_include_column_lineage(self) -> "DBTCommonConfig":
+        if self.include_column_lineage and not self.infer_dbt_schemas:
             raise ValueError(
                 "`infer_dbt_schemas` must be enabled to use `include_column_lineage`"
             )
 
-        return include_column_lineage
-
-    @validator("skip_sources_in_lineage", always=True)
-    def validate_skip_sources_in_lineage(
-        cls, skip_sources_in_lineage: bool, values: Dict
-    ) -> bool:
-        entities_enabled: Optional[DBTEntitiesEnabled] = values.get("entities_enabled")
-        prefer_sql_parser_lineage: Optional[bool] = values.get(
-            "prefer_sql_parser_lineage"
-        )
+        return self
 
-        if prefer_sql_parser_lineage and not skip_sources_in_lineage:
+    @model_validator(mode="after")
+    def validate_skip_sources_in_lineage(self) -> "DBTCommonConfig":
+        if self.prefer_sql_parser_lineage and not self.skip_sources_in_lineage:
             raise ValueError(
                 "`prefer_sql_parser_lineage` requires that `skip_sources_in_lineage` is enabled."
             )
 
         if (
-            skip_sources_in_lineage
-            and entities_enabled
-            and entities_enabled.sources == EmitDirective.YES
+            self.skip_sources_in_lineage
+            and self.entities_enabled
+            and self.entities_enabled.sources == EmitDirective.YES
             # When `prefer_sql_parser_lineage` is enabled, it's ok to have `skip_sources_in_lineage` enabled
             # without also disabling sources.
-            and not prefer_sql_parser_lineage
+            and not self.prefer_sql_parser_lineage
         ):
             raise ValueError(
                 "When `skip_sources_in_lineage` is enabled, `entities_enabled.sources` must be set to NO."
             )
 
-        return skip_sources_in_lineage
+        return self
 
 
 @dataclass

datahub/ingestion/source/dbt/dbt_core.py

@@ -9,7 +9,7 @@ from urllib.parse import urlparse
 import dateutil.parser
 import requests
 from packaging import version
-from pydantic import BaseModel, Field, validator
+from pydantic import BaseModel, Field, model_validator
 
 from datahub.configuration.git import GitReference
 from datahub.configuration.validate_field_rename import pydantic_renamed_field
@@ -99,26 +99,24 @@ class DBTCoreConfig(DBTCommonConfig):
 
     _github_info_deprecated = pydantic_renamed_field("github_info", "git_info")
 
-    @validator("aws_connection", always=True)
-    def aws_connection_needed_if_s3_uris_present(
-        cls, aws_connection: Optional[AwsConnectionConfig], values: Dict, **kwargs: Any
-    ) -> Optional[AwsConnectionConfig]:
+    @model_validator(mode="after")
+    def aws_connection_needed_if_s3_uris_present(self) -> "DBTCoreConfig":
         # first check if there are fields that contain s3 uris
         uris = [
-            values.get(f)
+            getattr(self, f, None)
             for f in [
                 "manifest_path",
                 "catalog_path",
                 "sources_path",
             ]
-        ] + values.get("run_results_paths", [])
+        ] + (self.run_results_paths or [])
         s3_uris = [uri for uri in uris if is_s3_uri(uri or "")]
 
-        if s3_uris and aws_connection is None:
+        if s3_uris and self.aws_connection is None:
             raise ValueError(
                 f"Please provide aws_connection configuration, since s3 uris have been provided {s3_uris}"
             )
-        return aws_connection
+        return self
 
 
 def get_columns(
@@ -426,13 +424,13 @@ def load_run_results(
         )
         return all_nodes
 
-    dbt_metadata = DBTRunMetadata.parse_obj(test_results_json.get("metadata", {}))
+    dbt_metadata = DBTRunMetadata.model_validate(test_results_json.get("metadata", {}))
 
     all_nodes_map: Dict[str, DBTNode] = {x.dbt_name: x for x in all_nodes}
 
     results = test_results_json.get("results", [])
     for result in results:
-        run_result = DBTRunResult.parse_obj(result)
+        run_result = DBTRunResult.model_validate(result)
         id = run_result.unique_id
 
         if id.startswith("test."):
@@ -477,7 +475,7 @@ class DBTCoreSource(DBTSourceBase, TestableSource):
 
     @classmethod
     def create(cls, config_dict, ctx):
-        config = DBTCoreConfig.parse_obj(config_dict)
+        config = DBTCoreConfig.model_validate(config_dict)
         return cls(config, ctx)
 
     @staticmethod

datahub/ingestion/source/debug/datahub_debug.py

@@ -46,7 +46,7 @@ class DataHubDebugSource(Source):
 
     @classmethod
     def create(cls, config_dict, ctx):
-        config = DataHubDebugSourceConfig.parse_obj(config_dict)
+        config = DataHubDebugSourceConfig.model_validate(config_dict)
         return cls(ctx, config)
 
     def perform_dns_probe(self, url: str) -> None:

datahub/ingestion/source/delta_lake/config.py

@@ -1,9 +1,8 @@
 import logging
 from typing import Optional
 
-import pydantic
 from cached_property import cached_property
-from pydantic import Field
+from pydantic import Field, field_validator
 from typing_extensions import Literal
 
 from datahub.configuration.common import AllowDenyPattern, ConfigModel
@@ -98,8 +97,11 @@ class DeltaLakeSourceConfig(
 
         return complete_path
 
-    @pydantic.validator("version_history_lookback")
-    def negative_version_history_implies_no_limit(cls, v):
+    @field_validator("version_history_lookback", mode="after")
+    @classmethod
+    def negative_version_history_implies_no_limit(
+        cls, v: Optional[int]
+    ) -> Optional[int]:
         if v and v < 0:
             return None
         return v

datahub/ingestion/source/dremio/dremio_config.py

@@ -2,7 +2,7 @@ import os
 from typing import List, Literal, Optional
 
 import certifi
-from pydantic import Field, validator
+from pydantic import Field, ValidationInfo, field_validator
 
 from datahub.configuration.common import AllowDenyPattern, ConfigModel, HiddenFromDocs
 from datahub.configuration.source_common import (
@@ -78,8 +78,9 @@ class DremioConnectionConfig(ConfigModel):
         description="ID of Dremio Cloud Project. Found in Project Settings in the Dremio Cloud UI",
     )
 
-    @validator("authentication_method")
-    def validate_auth_method(cls, value):
+    @field_validator("authentication_method", mode="after")
+    @classmethod
+    def validate_auth_method(cls, value: str) -> str:
         allowed_methods = ["password", "PAT"]
         if value not in allowed_methods:
             raise ValueError(
@@ -87,9 +88,12 @@ class DremioConnectionConfig(ConfigModel):
             )
         return value
 
-    @validator("password")
-    def validate_password(cls, value, values):
-        if values.get("authentication_method") == "PAT" and not value:
+    @field_validator("password", mode="after")
+    @classmethod
+    def validate_password(
+        cls, value: Optional[str], info: ValidationInfo
+    ) -> Optional[str]:
+        if info.data.get("authentication_method") == "PAT" and not value:
             raise ValueError(
                 "Password (Personal Access Token) is required when using PAT authentication",
             )