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

datahub/ingestion/source/s3/config.py

@@ -1,7 +1,7 @@
 import logging
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, Optional, Union
 
-import pydantic
+from pydantic import ValidationInfo, field_validator, model_validator
 from pydantic.fields import Field
 
 from datahub.configuration.common import AllowDenyPattern
@@ -12,7 +12,6 @@ from datahub.configuration.validate_field_deprecation import pydantic_field_depr
 from datahub.configuration.validate_field_rename import pydantic_renamed_field
 from datahub.ingestion.source.aws.aws_common import AwsConnectionConfig
 from datahub.ingestion.source.data_lake_common.config import PathSpecsConfigMixin
-from datahub.ingestion.source.data_lake_common.path_spec import PathSpec
 from datahub.ingestion.source.s3.datalake_profiler_config import DataLakeProfilerConfig
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
     StatefulStaleMetadataRemovalConfig,
@@ -117,69 +116,91 @@ class DataLakeSourceConfig(
             self.profiling.operation_config
         )
 
-    @pydantic.validator("path_specs", always=True)
-    def check_path_specs_and_infer_platform(
-        cls, path_specs: List[PathSpec], values: Dict
-    ) -> List[PathSpec]:
+    @field_validator("path_specs", mode="before")
+    @classmethod
+    def check_path_specs(cls, path_specs: Any, info: ValidationInfo) -> Any:
         if len(path_specs) == 0:
             raise ValueError("path_specs must not be empty")
 
-        # Check that all path specs have the same platform.
-        guessed_platforms = {
-            "s3" if path_spec.is_s3 else "file" for path_spec in path_specs
-        }
-        if len(guessed_platforms) > 1:
-            raise ValueError(
-                f"Cannot have multiple platforms in path_specs: {guessed_platforms}"
-            )
-        guessed_platform = guessed_platforms.pop()
-
-        # Ensure s3 configs aren't used for file sources.
-        if guessed_platform != "s3" and (
-            values.get("use_s3_object_tags") or values.get("use_s3_bucket_tags")
-        ):
-            raise ValueError(
-                "Cannot grab s3 object/bucket tags when platform is not s3. Remove the flag or use s3."
-            )
-
-        # Infer platform if not specified.
-        if values.get("platform") and values["platform"] != guessed_platform:
-            raise ValueError(
-                f"All path_specs belong to {guessed_platform} platform, but platform is set to {values['platform']}"
-            )
-        else:
-            logger.debug(f'Setting config "platform": {guessed_platform}')
-            values["platform"] = guessed_platform
+        # Basic validation - path specs consistency and S3 config validation is now handled in model_validator
 
         return path_specs
 
-    @pydantic.validator("platform", always=True)
-    def platform_valid(cls, platform: Any, values: dict) -> str:
-        inferred_platform = values.get("platform")  # we may have inferred it above
-        platform = platform or inferred_platform
-        if not platform:
-            raise ValueError("platform must not be empty")
-
-        if platform != "s3" and values.get("use_s3_bucket_tags"):
-            raise ValueError(
-                "Cannot grab s3 bucket tags when platform is not s3. Remove the flag or ingest from s3."
-            )
-        if platform != "s3" and values.get("use_s3_object_tags"):
-            raise ValueError(
-                "Cannot grab s3 object tags when platform is not s3. Remove the flag or ingest from s3."
-            )
-        if platform != "s3" and values.get("use_s3_content_type"):
-            raise ValueError(
-                "Cannot grab s3 object content type when platform is not s3. Remove the flag or ingest from s3."
-            )
-
-        return platform
-
-    @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[DataLakeProfilerConfig] = values.get("profiling")
+    @model_validator(mode="after")
+    def ensure_profiling_pattern_is_passed_to_profiling(self) -> "DataLakeSourceConfig":
+        profiling = self.profiling
         if profiling is not None and profiling.enabled:
-            profiling._allow_deny_patterns = values["profile_patterns"]
-        return values
+            profiling._allow_deny_patterns = self.profile_patterns
+        return self
+
+    @model_validator(mode="after")
+    def validate_platform_and_config_consistency(self) -> "DataLakeSourceConfig":
+        """Infer platform from path_specs and validate config consistency."""
+        # Track whether platform was explicitly provided
+        platform_was_explicit = bool(self.platform)
+
+        # Infer platform from path_specs if not explicitly set
+        if not self.platform and self.path_specs:
+            guessed_platforms = set()
+            for path_spec in self.path_specs:
+                if (
+                    hasattr(path_spec, "include")
+                    and path_spec.include
+                    and path_spec.include.startswith("s3://")
+                ):
+                    guessed_platforms.add("s3")
+                else:
+                    guessed_platforms.add("file")
+
+            # Ensure all path specs belong to the same platform
+            if len(guessed_platforms) > 1:
+                raise ValueError(
+                    f"Cannot have multiple platforms in path_specs: {guessed_platforms}"
+                )
+
+            if guessed_platforms:
+                guessed_platform = guessed_platforms.pop()
+                logger.debug(f"Inferred platform: {guessed_platform}")
+                self.platform = guessed_platform
+            else:
+                self.platform = "file"
+        elif not self.platform:
+            self.platform = "file"
+
+        # Validate platform consistency only when platform was inferred (not explicitly set)
+        # This allows sources like GCS to set platform="gcs" with s3:// URIs for correct container subtypes
+        if not platform_was_explicit and self.platform and self.path_specs:
+            expected_platforms = set()
+            for path_spec in self.path_specs:
+                if (
+                    hasattr(path_spec, "include")
+                    and path_spec.include
+                    and path_spec.include.startswith("s3://")
+                ):
+                    expected_platforms.add("s3")
+                else:
+                    expected_platforms.add("file")
+
+            if len(expected_platforms) == 1:
+                expected_platform = expected_platforms.pop()
+                if self.platform != expected_platform:
+                    raise ValueError(
+                        f"All path_specs belong to {expected_platform} platform, but platform was inferred as {self.platform}"
+                    )
+
+        # Validate S3-specific configurations
+        if self.platform != "s3":
+            if self.use_s3_bucket_tags:
+                raise ValueError(
+                    "Cannot grab s3 bucket tags when platform is not s3. Remove the flag or ingest from s3."
+                )
+            if self.use_s3_object_tags:
+                raise ValueError(
+                    "Cannot grab s3 object tags when platform is not s3. Remove the flag or ingest from s3."
+                )
+            if self.use_s3_content_type:
+                raise ValueError(
+                    "Cannot grab s3 object content type when platform is not s3. Remove the flag or ingest from s3."
+                )
+
+        return self

datahub/ingestion/source/s3/datalake_profiler_config.py

@@ -1,6 +1,7 @@
-from typing import Any, Dict, Optional
+from typing import Optional
 
 import pydantic
+from pydantic import model_validator
 from pydantic.fields import Field
 
 from datahub.configuration import ConfigModel
@@ -72,21 +73,18 @@ class DataLakeProfilerConfig(ConfigModel):
         description="Whether to profile for the sample values for all columns.",
     )
 
-    @pydantic.root_validator(skip_on_failure=True)
-    def ensure_field_level_settings_are_normalized(
-        cls: "DataLakeProfilerConfig", values: Dict[str, Any]
-    ) -> Dict[str, Any]:
-        max_num_fields_to_profile_key = "max_number_of_fields_to_profile"
-        max_num_fields_to_profile = values.get(max_num_fields_to_profile_key)
+    @model_validator(mode="after")
+    def ensure_field_level_settings_are_normalized(self) -> "DataLakeProfilerConfig":
+        max_num_fields_to_profile = self.max_number_of_fields_to_profile
 
         # Disable all field-level metrics.
-        if values.get("profile_table_level_only"):
-            for field_level_metric in cls.__fields__:
-                if field_level_metric.startswith("include_field_"):
-                    values.setdefault(field_level_metric, False)
+        if self.profile_table_level_only:
+            for field_name in self.__fields__:
+                if field_name.startswith("include_field_"):
+                    setattr(self, field_name, False)
 
             assert max_num_fields_to_profile is None, (
-                f"{max_num_fields_to_profile_key} should be set to None"
+                "max_number_of_fields_to_profile should be set to None"
             )
 
-        return values
+        return self

datahub/ingestion/source/s3/source.py

@@ -53,8 +53,11 @@ from datahub.ingestion.source.data_lake_common.data_lake_utils import (
 from datahub.ingestion.source.data_lake_common.object_store import (
     create_object_store_adapter,
 )
-from datahub.ingestion.source.data_lake_common.path_spec import FolderTraversalMethod
-from datahub.ingestion.source.s3.config import DataLakeSourceConfig, PathSpec
+from datahub.ingestion.source.data_lake_common.path_spec import (
+    FolderTraversalMethod,
+    PathSpec,
+)
+from datahub.ingestion.source.s3.config import DataLakeSourceConfig
 from datahub.ingestion.source.s3.report import DataLakeSourceReport
 from datahub.ingestion.source.schema_inference import avro, csv_tsv, json, parquet
 from datahub.ingestion.source.schema_inference.base import SchemaInferenceBase
@@ -261,7 +264,7 @@ class S3Source(StatefulIngestionSourceBase):
             )
 
         config_report = {
-            config_option: config.dict().get(config_option)
+            config_option: config.model_dump().get(config_option)
             for config_option in config_options_to_report
         }
         config_report = {
@@ -278,7 +281,7 @@ class S3Source(StatefulIngestionSourceBase):
             telemetry.telemetry_instance.ping(
                 "data_lake_profiling_config",
                 {
-                    config_flag: config.profiling.dict().get(config_flag)
+                    config_flag: config.profiling.model_dump().get(config_flag)
                     for config_flag in profiling_flags_to_report
                 },
             )
@@ -370,7 +373,7 @@ class S3Source(StatefulIngestionSourceBase):
 
     @classmethod
     def create(cls, config_dict, ctx):
-        config = DataLakeSourceConfig.parse_obj(config_dict)
+        config = DataLakeSourceConfig.model_validate(config_dict)
 
         return cls(config, ctx)
 

datahub/ingestion/source/sac/sac.py

@@ -8,7 +8,7 @@ import pyodata
 import pyodata.v2.model
 import pyodata.v2.service
 from authlib.integrations.requests_client import OAuth2Session
-from pydantic import Field, SecretStr, validator
+from pydantic import Field, SecretStr, field_validator
 from requests.adapters import HTTPAdapter
 from urllib3.util.retry import Retry
 
@@ -159,7 +159,8 @@ class SACSourceConfig(
         description="Template for generating dataset urns of consumed queries, the placeholder {query} can be used within the template for inserting the name of the query",
     )
 
-    @validator("tenant_url", "token_url")
+    @field_validator("tenant_url", "token_url", mode="after")
+    @classmethod
     def remove_trailing_slash(cls, v):
         return config_clean.remove_trailing_slashes(v)
 
@@ -209,7 +210,7 @@ class SACSource(StatefulIngestionSourceBase, TestableSource):
 
     @classmethod
     def create(cls, config_dict: dict, ctx: PipelineContext) -> "SACSource":
-        config = SACSourceConfig.parse_obj(config_dict)
+        config = SACSourceConfig.model_validate(config_dict)
         return cls(config, ctx)
 
     @staticmethod
@@ -217,7 +218,7 @@ class SACSource(StatefulIngestionSourceBase, TestableSource):
         test_report = TestConnectionReport()
 
         try:
-            config = SACSourceConfig.parse_obj(config_dict)
+            config = SACSourceConfig.model_validate(config_dict)
 
             # when creating the pyodata.Client, the metadata is automatically parsed and validated
             session, _ = SACSource.get_sac_connection(config)

datahub/ingestion/source/salesforce.py

@@ -7,7 +7,7 @@ from enum import Enum
 from typing import Any, Dict, Iterable, List, Literal, Optional, TypedDict
 
 import requests
-from pydantic import Field, validator
+from pydantic import Field, field_validator
 from simple_salesforce import Salesforce
 from simple_salesforce.exceptions import SalesforceAuthenticationFailed
 
@@ -172,7 +172,8 @@ class SalesforceConfig(
             self.profiling.operation_config
         )
 
-    @validator("instance_url")
+    @field_validator("instance_url", mode="after")
+    @classmethod
     def remove_trailing_slash(cls, v):
         return config_clean.remove_trailing_slashes(v)
 

datahub/ingestion/source/schema/json_schema.py

@@ -12,7 +12,7 @@ from urllib.parse import urlparse
 
 import jsonref
 import requests
-from pydantic import AnyHttpUrl, DirectoryPath, FilePath, validator
+from pydantic import AnyHttpUrl, DirectoryPath, FilePath, field_validator
 from pydantic.fields import Field
 
 import datahub.metadata.schema_classes as models
@@ -90,7 +90,7 @@ class JsonSchemaSourceConfig(StatefulIngestionConfigBase, DatasetSourceConfigMix
         description="Use this if URI-s need to be modified during reference resolution. Simple string match - replace capabilities are supported.",
     )
 
-    @validator("path")
+    @field_validator("path", mode="after")
     def download_http_url_to_temp_file(cls, v):
         if isinstance(v, AnyHttpUrl):
             try:

datahub/ingestion/source/sigma/data_classes.py

@@ -2,7 +2,7 @@ from copy import deepcopy
 from datetime import datetime
 from typing import Dict, List, Optional
 
-from pydantic import BaseModel, root_validator
+from pydantic import BaseModel, model_validator
 
 from datahub.emitter.mcp_builder import ContainerKey
 
@@ -22,7 +22,8 @@ class Workspace(BaseModel):
     createdAt: datetime
     updatedAt: datetime
 
-    @root_validator(pre=True)
+    @model_validator(mode="before")
+    @classmethod
     def update_values(cls, values: Dict) -> Dict:
         # Create a copy to avoid modifying the input dictionary, preventing state contamination in tests
         values = deepcopy(values)

datahub/ingestion/source/sigma/sigma.py

@@ -150,7 +150,7 @@ class SigmaSource(StatefulIngestionSourceBase, TestableSource):
 
     @classmethod
     def create(cls, config_dict, ctx):
-        config = SigmaSourceConfig.parse_obj(config_dict)
+        config = SigmaSourceConfig.model_validate(config_dict)
         return cls(config, ctx)
 
     def _gen_workbook_key(self, workbook_id: str) -> WorkbookKey:

datahub/ingestion/source/sigma/sigma_api.py

@@ -108,7 +108,7 @@ class SigmaAPI:
                 self.report.non_accessible_workspaces_count += 1
                 return None
             response.raise_for_status()
-            workspace = Workspace.parse_obj(response.json())
+            workspace = Workspace.model_validate(response.json())
             self.workspaces[workspace.workspaceId] = workspace
             return workspace
         except Exception as e:
@@ -127,7 +127,7 @@ class SigmaAPI:
                 response_dict = response.json()
                 for workspace_dict in response_dict[Constant.ENTRIES]:
                     self.workspaces[workspace_dict[Constant.WORKSPACEID]] = (
-                        Workspace.parse_obj(workspace_dict)
+                        Workspace.model_validate(workspace_dict)
                     )
                 if response_dict[Constant.NEXTPAGE]:
                     url = f"{workspace_url}&page={response_dict[Constant.NEXTPAGE]}"
@@ -197,7 +197,7 @@ class SigmaAPI:
                 response.raise_for_status()
                 response_dict = response.json()
                 for file_dict in response_dict[Constant.ENTRIES]:
-                    file = File.parse_obj(file_dict)
+                    file = File.model_validate(file_dict)
                     file.workspaceId = self.get_workspace_id_from_file_path(
                         file.parentId, file.path
                     )
@@ -225,7 +225,7 @@ class SigmaAPI:
                 response.raise_for_status()
                 response_dict = response.json()
                 for dataset_dict in response_dict[Constant.ENTRIES]:
-                    dataset = SigmaDataset.parse_obj(dataset_dict)
+                    dataset = SigmaDataset.model_validate(dataset_dict)
 
                     if dataset.datasetId not in dataset_files_metadata:
                         self.report.datasets.dropped(
@@ -354,7 +354,7 @@ class SigmaAPI:
                 element_dict[Constant.URL] = (
                     f"{workbook.url}?:nodeId={element_dict[Constant.ELEMENTID]}&:fullScreen=true"
                 )
-                element = Element.parse_obj(element_dict)
+                element = Element.model_validate(element_dict)
                 if (
                     self.config.extract_lineage
                     and self.config.workbook_lineage_pattern.allowed(workbook.name)
@@ -379,7 +379,7 @@ class SigmaAPI:
             )
             response.raise_for_status()
             for page_dict in response.json()[Constant.ENTRIES]:
-                page = Page.parse_obj(page_dict)
+                page = Page.model_validate(page_dict)
                 page.elements = self.get_page_elements(workbook, page)
                 pages.append(page)
             return pages
@@ -400,7 +400,7 @@ class SigmaAPI:
                 response.raise_for_status()
                 response_dict = response.json()
                 for workbook_dict in response_dict[Constant.ENTRIES]:
-                    workbook = Workbook.parse_obj(workbook_dict)
+                    workbook = Workbook.model_validate(workbook_dict)
 
                     if workbook.workbookId not in workbook_files_metadata:
                         # Due to a bug in the Sigma API, it seems like the /files endpoint does not

datahub/ingestion/source/slack/slack.py

@@ -260,7 +260,7 @@ class SlackSource(StatefulIngestionSourceBase):
 
     @classmethod
     def create(cls, config_dict, ctx):
-        config = SlackSourceConfig.parse_obj(config_dict)
+        config = SlackSourceConfig.model_validate(config_dict)
         return cls(ctx, config)
 
     def get_slack_client(self) -> WebClient:

datahub/ingestion/source/snaplogic/snaplogic.py

@@ -351,5 +351,5 @@ class SnaplogicSource(StatefulIngestionSourceBase):
 
     @classmethod
     def create(cls, config_dict: dict, ctx: PipelineContext) -> "SnaplogicSource":
-        config = SnaplogicConfig.parse_obj(config_dict)
+        config = SnaplogicConfig.model_validate(config_dict)
         return cls(config, ctx)

datahub/ingestion/source/snowflake/snowflake_assertion.py

@@ -91,7 +91,7 @@ class SnowflakeAssertionsHandler:
         self, result_row: dict, discovered_datasets: List[str]
     ) -> Optional[MetadataChangeProposalWrapper]:
         try:
-            result = DataQualityMonitoringResult.parse_obj(result_row)
+            result = DataQualityMonitoringResult.model_validate(result_row)
             assertion_guid = result.METRIC_NAME.split("__")[-1].lower()
             status = bool(result.VALUE)  # 1 if PASS, 0 if FAIL
             assertee = self.identifiers.get_dataset_identifier(

datahub/ingestion/source/snowflake/snowflake_config.py

@@ -5,7 +5,7 @@ from enum import Enum
 from typing import Dict, List, Optional, Set
 
 import pydantic
-from pydantic import Field, root_validator, validator
+from pydantic import Field, ValidationInfo, field_validator, model_validator
 
 from datahub.configuration.common import AllowDenyPattern, ConfigModel, HiddenFromDocs
 from datahub.configuration.pattern_utils import UUID_REGEX
@@ -122,10 +122,10 @@ class SnowflakeFilterConfig(SQLFilterConfig):
         description="Whether `schema_pattern` is matched against fully qualified schema name `<catalog>.<schema>`.",
     )
 
-    @root_validator(pre=False, skip_on_failure=True)
-    def validate_legacy_schema_pattern(cls, values: Dict) -> Dict:
-        schema_pattern: Optional[AllowDenyPattern] = values.get("schema_pattern")
-        match_fully_qualified_names = values.get("match_fully_qualified_names")
+    @model_validator(mode="after")
+    def validate_legacy_schema_pattern(self) -> "SnowflakeFilterConfig":
+        schema_pattern: Optional[AllowDenyPattern] = self.schema_pattern
+        match_fully_qualified_names = self.match_fully_qualified_names
 
         if (
             schema_pattern is not None
@@ -145,7 +145,7 @@ class SnowflakeFilterConfig(SQLFilterConfig):
             assert isinstance(schema_pattern, AllowDenyPattern)
             schema_pattern.deny.append(r".*INFORMATION_SCHEMA$")
 
-        return values
+        return self
 
 
 class SnowflakeIdentifierConfig(
@@ -391,7 +391,8 @@ class SnowflakeV2Config(
         "This may be required in the case of _eg_ temporary tables being created in a different database than the ones in the database_name patterns.",
     )
 
-    @validator("convert_urns_to_lowercase")
+    @field_validator("convert_urns_to_lowercase", mode="after")
+    @classmethod
     def validate_convert_urns_to_lowercase(cls, v):
         if not v:
             add_global_warning(
@@ -400,30 +401,31 @@ class SnowflakeV2Config(
 
         return v
 
-    @validator("include_column_lineage")
-    def validate_include_column_lineage(cls, v, values):
-        if not values.get("include_table_lineage") and v:
+    @field_validator("include_column_lineage", mode="after")
+    @classmethod
+    def validate_include_column_lineage(cls, v, info):
+        if not info.data.get("include_table_lineage") and v:
             raise ValueError(
                 "include_table_lineage must be True for include_column_lineage to be set."
             )
         return v
 
-    @root_validator(pre=False, skip_on_failure=True)
-    def validate_unsupported_configs(cls, values: Dict) -> Dict:
-        value = values.get("include_read_operational_stats")
-        if value is not None and value:
+    @model_validator(mode="after")
+    def validate_unsupported_configs(self) -> "SnowflakeV2Config":
+        if (
+            hasattr(self, "include_read_operational_stats")
+            and self.include_read_operational_stats
+        ):
             raise ValueError(
                 "include_read_operational_stats is not supported. Set `include_read_operational_stats` to False.",
             )
 
-        include_technical_schema = values.get("include_technical_schema")
-        include_profiles = (
-            values.get("profiling") is not None and values["profiling"].enabled
-        )
+        include_technical_schema = self.include_technical_schema
+        include_profiles = self.profiling is not None and self.profiling.enabled
         delete_detection_enabled = (
-            values.get("stateful_ingestion") is not None
-            and values["stateful_ingestion"].enabled
-            and values["stateful_ingestion"].remove_stale_metadata
+            self.stateful_ingestion is not None
+            and self.stateful_ingestion.enabled
+            and self.stateful_ingestion.remove_stale_metadata
         )
 
         # TODO: Allow profiling irrespective of basic schema extraction,
@@ -435,13 +437,14 @@ class SnowflakeV2Config(
                 "Cannot perform Deletion Detection or Profiling without extracting snowflake technical schema. Set `include_technical_schema` to True or disable Deletion Detection and Profiling."
             )
 
-        return values
+        return self
 
-    @validator("shares")
+    @field_validator("shares", mode="after")
+    @classmethod
     def validate_shares(
-        cls, shares: Optional[Dict[str, SnowflakeShareConfig]], values: Dict
+        cls, shares: Optional[Dict[str, SnowflakeShareConfig]], info: ValidationInfo
     ) -> Optional[Dict[str, SnowflakeShareConfig]]:
-        current_platform_instance = values.get("platform_instance")
+        current_platform_instance = info.data.get("platform_instance")
 
         if shares:
             # Check: platform_instance should be present
@@ -479,11 +482,12 @@ class SnowflakeV2Config(
 
         return shares
 
-    @root_validator(pre=False, skip_on_failure=True)
-    def validate_queries_v2_stateful_ingestion(cls, values: Dict) -> Dict:
-        if values.get("use_queries_v2"):
-            if values.get("enable_stateful_lineage_ingestion") or values.get(
-                "enable_stateful_usage_ingestion"
+    @model_validator(mode="after")
+    def validate_queries_v2_stateful_ingestion(self) -> "SnowflakeV2Config":
+        if self.use_queries_v2:
+            if (
+                self.enable_stateful_lineage_ingestion
+                or self.enable_stateful_usage_ingestion
             ):
                 logger.warning(
                     "enable_stateful_lineage_ingestion and enable_stateful_usage_ingestion are deprecated "
@@ -491,7 +495,7 @@ class SnowflakeV2Config(
                     "For queries v2, use enable_stateful_time_window instead to enable stateful ingestion "
                     "for the unified time window extraction (lineage + usage + operations + queries)."
                 )
-        return values
+        return self
 
     def outbounds(self) -> Dict[str, Set[DatabaseId]]:
         """