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

datahub/ingestion/source/snowflake/snowflake_connection.py

@@ -6,6 +6,7 @@ import pydantic
 import snowflake.connector
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import serialization
+from pydantic import field_validator, model_validator
 from snowflake.connector import SnowflakeConnection as NativeSnowflakeConnection
 from snowflake.connector.cursor import DictCursor
 from snowflake.connector.network import (
@@ -125,26 +126,28 @@ class SnowflakeConnectionConfig(ConfigModel):
 
     rename_host_port_to_account_id = pydantic_renamed_field("host_port", "account_id")  # type: ignore[pydantic-field]
 
-    @pydantic.validator("account_id")
-    def validate_account_id(cls, account_id: str, values: Dict) -> str:
+    @field_validator("account_id", mode="after")
+    @classmethod
+    def validate_account_id(cls, account_id: str, info: pydantic.ValidationInfo) -> str:
         account_id = remove_protocol(account_id)
         account_id = remove_trailing_slashes(account_id)
         # Get the domain from config, fallback to default
-        domain = values.get("snowflake_domain", DEFAULT_SNOWFLAKE_DOMAIN)
+        domain = info.data.get("snowflake_domain", DEFAULT_SNOWFLAKE_DOMAIN)
         snowflake_host_suffix = f".{domain}"
         account_id = remove_suffix(account_id, snowflake_host_suffix)
         return account_id
 
-    @pydantic.validator("authentication_type", always=True)
-    def authenticator_type_is_valid(cls, v, values):
+    @field_validator("authentication_type", mode="before")
+    @classmethod
+    def authenticator_type_is_valid(cls, v: Any, info: pydantic.ValidationInfo) -> Any:
         if v not in _VALID_AUTH_TYPES:
             raise ValueError(
                 f"unsupported authenticator type '{v}' was provided,"
                 f" use one of {list(_VALID_AUTH_TYPES.keys())}"
             )
         if (
-            values.get("private_key") is not None
-            or values.get("private_key_path") is not None
+            info.data.get("private_key") is not None
+            or info.data.get("private_key_path") is not None
         ) and v != "KEY_PAIR_AUTHENTICATOR":
             raise ValueError(
                 f"Either `private_key` and `private_key_path` is set but `authentication_type` is {v}. "
@@ -153,21 +156,22 @@ class SnowflakeConnectionConfig(ConfigModel):
         if v == "KEY_PAIR_AUTHENTICATOR":
             # If we are using key pair auth, we need the private key path and password to be set
             if (
-                values.get("private_key") is None
-                and values.get("private_key_path") is None
+                info.data.get("private_key") is None
+                and info.data.get("private_key_path") is None
             ):
                 raise ValueError(
                     f"Both `private_key` and `private_key_path` are none. "
                     f"At least one should be set when using {v} authentication"
                 )
         elif v == "OAUTH_AUTHENTICATOR":
-            cls._check_oauth_config(values.get("oauth_config"))
+            cls._check_oauth_config(info.data.get("oauth_config"))
         logger.info(f"using authenticator type '{v}'")
         return v
 
-    @pydantic.validator("token", always=True)
-    def validate_token_oauth_config(cls, v, values):
-        auth_type = values.get("authentication_type")
+    @field_validator("token", mode="before")
+    @classmethod
+    def validate_token_oauth_config(cls, v: Any, info: pydantic.ValidationInfo) -> Any:
+        auth_type = info.data.get("authentication_type")
         if auth_type == "OAUTH_AUTHENTICATOR_TOKEN":
             if not v:
                 raise ValueError("Token required for OAUTH_AUTHENTICATOR_TOKEN.")
@@ -177,6 +181,24 @@ class SnowflakeConnectionConfig(ConfigModel):
             )
         return v
 
+    @model_validator(mode="after")
+    def validate_authentication_config(self):
+        """Validate authentication configuration consistency."""
+        # Check token requirement for OAUTH_AUTHENTICATOR_TOKEN
+        if self.authentication_type == "OAUTH_AUTHENTICATOR_TOKEN":
+            if not self.token:
+                raise ValueError("Token required for OAUTH_AUTHENTICATOR_TOKEN.")
+
+        # Check private key authentication consistency
+        if self.private_key is not None or self.private_key_path is not None:
+            if self.authentication_type != "KEY_PAIR_AUTHENTICATOR":
+                raise ValueError(
+                    f"Either `private_key` and `private_key_path` is set but `authentication_type` is {self.authentication_type}. "
+                    f"Should be set to 'KEY_PAIR_AUTHENTICATOR' when using key pair authentication"
+                )
+
+        return self
+
     @staticmethod
     def _check_oauth_config(oauth_config: Optional[OAuthConfiguration]) -> None:
         if oauth_config is None:

datahub/ingestion/source/snowflake/snowflake_lineage_v2.py

@@ -14,7 +14,7 @@ from typing import (
     Type,
 )
 
-from pydantic import BaseModel, Field, validator
+from pydantic import BaseModel, Field, field_validator
 
 from datahub.configuration.datetimes import parse_absolute_time
 from datahub.ingestion.api.closeable import Closeable
@@ -70,7 +70,7 @@ def pydantic_parse_json(field: str) -> "V1Validator":
             return json.loads(v)
         return v
 
-    return validator(field, pre=True, allow_reuse=True)(_parse_from_json)
+    return field_validator(field, mode="before")(_parse_from_json)
 
 
 class UpstreamColumnNode(BaseModel):
@@ -379,7 +379,7 @@ class SnowflakeLineageExtractor(SnowflakeCommonMixin, Closeable):
                 # To avoid that causing a pydantic error we are setting it to an empty list
                 # instead of a list with an empty object
                 db_row["QUERIES"] = "[]"
-            return UpstreamLineageEdge.parse_obj(db_row)
+            return UpstreamLineageEdge.model_validate(db_row)
         except Exception as e:
             self.report.num_upstream_lineage_edge_parsing_failed += 1
             upstream_tables = db_row.get("UPSTREAM_TABLES")

datahub/ingestion/source/snowflake/snowflake_queries.py

@@ -78,6 +78,7 @@ from datahub.utilities.file_backed_collections import (
     ConnectionWrapper,
     FileBackedList,
 )
+from datahub.utilities.lossy_collections import LossyList
 from datahub.utilities.perf_timer import PerfTimer
 
 logger = logging.getLogger(__name__)
@@ -169,6 +170,10 @@ class SnowflakeQueriesExtractorReport(Report):
     num_stream_queries_observed: int = 0
     num_create_temp_view_queries_observed: int = 0
     num_users: int = 0
+    num_queries_with_empty_column_name: int = 0
+    queries_with_empty_column_name: LossyList[str] = dataclasses.field(
+        default_factory=LossyList
+    )
 
 
 @dataclass
@@ -626,9 +631,28 @@ class SnowflakeQueriesExtractor(SnowflakeStructuredReportMixin, Closeable):
 
             columns = set()
             for modified_column in obj["columns"]:
-                columns.add(
-                    self.identifiers.snowflake_identifier(modified_column["columnName"])
-                )
+                column_name = modified_column["columnName"]
+                # An empty column name in the audit log would cause an error when creating column URNs.
+                # To avoid this and still extract lineage, the raw query text is parsed as a fallback.
+                if not column_name or not column_name.strip():
+                    query_id = res["query_id"]
+                    self.report.num_queries_with_empty_column_name += 1
+                    self.report.queries_with_empty_column_name.append(query_id)
+                    logger.info(f"Query {query_id} has empty column name in audit log.")
+
+                    return ObservedQuery(
+                        query=query_text,
+                        session_id=res["session_id"],
+                        timestamp=timestamp,
+                        user=user,
+                        default_db=res["default_db"],
+                        default_schema=res["default_schema"],
+                        query_hash=get_query_fingerprint(
+                            query_text, self.identifiers.platform, fast=True
+                        ),
+                        extra_info=extra_info,
+                    )
+                columns.add(self.identifiers.snowflake_identifier(column_name))
 
             upstreams.append(dataset)
             column_usage[dataset] = columns
@@ -782,7 +806,7 @@ class SnowflakeQueriesSource(Source):
 
     @classmethod
     def create(cls, config_dict: dict, ctx: PipelineContext) -> Self:
-        config = SnowflakeQueriesSourceConfig.parse_obj(config_dict)
+        config = SnowflakeQueriesSourceConfig.model_validate(config_dict)
         return cls(ctx, config)
 
     def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:

datahub/ingestion/source/sql/athena.py

@@ -386,7 +386,7 @@ class AthenaSource(SQLAlchemySource):
 
     @classmethod
     def create(cls, config_dict, ctx):
-        config = AthenaConfig.parse_obj(config_dict)
+        config = AthenaConfig.model_validate(config_dict)
         return cls(config, ctx)
 
     # overwrite this method to allow to specify the usage of a custom dialect

datahub/ingestion/source/sql/clickhouse.py

@@ -10,6 +10,7 @@ import clickhouse_sqlalchemy.types as custom_types
 import pydantic
 from clickhouse_sqlalchemy.drivers import base
 from clickhouse_sqlalchemy.drivers.base import ClickHouseDialect
+from pydantic import model_validator
 from pydantic.fields import Field
 from sqlalchemy import create_engine, text
 from sqlalchemy.engine import reflection
@@ -175,7 +176,8 @@ class ClickHouseConfig(
         return str(url)
 
     # pre = True because we want to take some decision before pydantic initialize the configuration to default values
-    @pydantic.root_validator(pre=True)
+    @model_validator(mode="before")
+    @classmethod
     def projects_backward_compatibility(cls, values: Dict) -> Dict:
         secure = values.get("secure")
         protocol = values.get("protocol")
@@ -423,7 +425,7 @@ class ClickHouseSource(TwoTierSQLAlchemySource):
 
     @classmethod
     def create(cls, config_dict, ctx):
-        config = ClickHouseConfig.parse_obj(config_dict)
+        config = ClickHouseConfig.model_validate(config_dict)
         return cls(config, ctx)
 
     def get_workunits_internal(self) -> Iterable[Union[MetadataWorkUnit, SqlWorkUnit]]:

datahub/ingestion/source/sql/cockroachdb.py

@@ -39,5 +39,5 @@ class CockroachDBSource(PostgresSource):
 
     @classmethod
     def create(cls, config_dict, ctx):
-        config = CockroachDBConfig.parse_obj(config_dict)
+        config = CockroachDBConfig.model_validate(config_dict)
         return cls(config, ctx)

datahub/ingestion/source/sql/druid.py

@@ -77,5 +77,5 @@ class DruidSource(SQLAlchemySource):
 
     @classmethod
     def create(cls, config_dict, ctx):
-        config = DruidConfig.parse_obj(config_dict)
+        config = DruidConfig.model_validate(config_dict)
         return cls(config, ctx)

datahub/ingestion/source/sql/hana.py

@@ -36,5 +36,5 @@ class HanaSource(SQLAlchemySource):
 
     @classmethod
     def create(cls, config_dict: Dict, ctx: PipelineContext) -> "HanaSource":
-        config = HanaConfig.parse_obj(config_dict)
+        config = HanaConfig.model_validate(config_dict)
         return cls(config, ctx)

datahub/ingestion/source/sql/hive.py

@@ -6,7 +6,7 @@ from enum import Enum
 from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 from urllib.parse import urlparse
 
-from pydantic import validator
+from pydantic import field_validator
 from pydantic.fields import Field
 
 # This import verifies that the dependencies are available.
@@ -674,11 +674,13 @@ class HiveConfig(TwoTierSQLAlchemyConfig):
         description="Platform instance for the storage system",
     )
 
-    @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)
 
-    @validator("hive_storage_lineage_direction")
+    @field_validator("hive_storage_lineage_direction", mode="after")
+    @classmethod
     def _validate_direction(cls, v: str) -> str:
         """Validate the lineage direction."""
         if v.lower() not in ["upstream", "downstream"]:
@@ -725,7 +727,7 @@ class HiveSource(TwoTierSQLAlchemySource):
 
     @classmethod
     def create(cls, config_dict, ctx):
-        config = HiveConfig.parse_obj(config_dict)
+        config = HiveConfig.model_validate(config_dict)
         return cls(config, ctx)
 
     def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:

datahub/ingestion/source/sql/hive_metastore.py

@@ -351,7 +351,7 @@ class HiveMetastoreSource(SQLAlchemySource):
 
     @classmethod
     def create(cls, config_dict, ctx):
-        config = HiveMetastore.parse_obj(config_dict)
+        config = HiveMetastore.model_validate(config_dict)
         return cls(config, ctx)
 
     def gen_database_containers(

datahub/ingestion/source/sql/mssql/source.py

@@ -3,8 +3,8 @@ import re
 import urllib.parse
 from typing import Any, Dict, Iterable, List, Optional, Tuple
 
-import pydantic
 import sqlalchemy.dialects.mssql
+from pydantic import ValidationInfo, field_validator
 from pydantic.fields import Field
 from sqlalchemy import create_engine, inspect
 from sqlalchemy.engine.base import Connection
@@ -140,11 +140,18 @@ class SQLServerConfig(BasicSQLAlchemyConfig):
         description="Indicates if the SQL Server instance is running on AWS RDS. When None (default), automatic detection will be attempted using server name analysis.",
     )
 
-    @pydantic.validator("uri_args")
-    def passwords_match(cls, v, values, **kwargs):
-        if values["use_odbc"] and not values["sqlalchemy_uri"] and "driver" not in v:
+    @field_validator("uri_args", mode="after")
+    @classmethod
+    def passwords_match(
+        cls, v: Dict[str, Any], info: ValidationInfo, **kwargs: Any
+    ) -> Dict[str, Any]:
+        if (
+            info.data["use_odbc"]
+            and not info.data["sqlalchemy_uri"]
+            and "driver" not in v
+        ):
             raise ValueError("uri_args must contain a 'driver' option")
-        elif not values["use_odbc"] and v:
+        elif not info.data["use_odbc"] and v:
             raise ValueError("uri_args is not supported when ODBC is disabled")
         return v
 
@@ -314,7 +321,7 @@ class SQLServerSource(SQLAlchemySource):
 
     @classmethod
     def create(cls, config_dict: Dict, ctx: PipelineContext) -> "SQLServerSource":
-        config = SQLServerConfig.parse_obj(config_dict)
+        config = SQLServerConfig.model_validate(config_dict)
         return cls(config, ctx)
 
     # override to get table descriptions

datahub/ingestion/source/sql/mysql.py

@@ -150,7 +150,7 @@ class MySQLSource(TwoTierSQLAlchemySource):
 
     @classmethod
     def create(cls, config_dict, ctx):
-        config = MySQLConfig.parse_obj(config_dict)
+        config = MySQLConfig.model_validate(config_dict)
         return cls(config, ctx)
 
     def _setup_rds_iam_event_listener(

datahub/ingestion/source/sql/oracle.py

@@ -10,8 +10,8 @@ from typing import Any, Dict, Iterable, List, NoReturn, Optional, Tuple, Union,
 from unittest.mock import patch
 
 import oracledb
-import pydantic
 import sqlalchemy.engine
+from pydantic import ValidationInfo, field_validator
 from pydantic.fields import Field
 from sqlalchemy import event, sql
 from sqlalchemy.dialects.oracle.base import ischema_names
@@ -101,25 +101,32 @@ class OracleConfig(BasicSQLAlchemyConfig):
         "On Linux, this value is ignored, as ldconfig or LD_LIBRARY_PATH will define the location.",
     )
 
-    @pydantic.validator("service_name")
-    def check_service_name(cls, v, values):
-        if values.get("database") and v:
+    @field_validator("service_name", mode="after")
+    @classmethod
+    def check_service_name(
+        cls, v: Optional[str], info: ValidationInfo
+    ) -> Optional[str]:
+        if info.data.get("database") and v:
             raise ValueError(
                 "specify one of 'database' and 'service_name', but not both"
             )
         return v
 
-    @pydantic.validator("data_dictionary_mode")
-    def check_data_dictionary_mode(cls, value):
+    @field_validator("data_dictionary_mode", mode="after")
+    @classmethod
+    def check_data_dictionary_mode(cls, value: str) -> str:
         if value not in ("ALL", "DBA"):
             raise ValueError("Specify one of data dictionary views mode: 'ALL', 'DBA'.")
         return value
 
-    @pydantic.validator("thick_mode_lib_dir", always=True)
-    def check_thick_mode_lib_dir(cls, v, values):
+    @field_validator("thick_mode_lib_dir", mode="before")
+    @classmethod
+    def check_thick_mode_lib_dir(
+        cls, v: Optional[str], info: ValidationInfo
+    ) -> Optional[str]:
         if (
             v is None
-            and values.get("enable_thick_mode")
+            and info.data.get("enable_thick_mode")
             and (platform.system() == "Darwin" or platform.system() == "Windows")
         ):
             raise ValueError(
@@ -659,7 +666,7 @@ class OracleSource(SQLAlchemySource):
 
     @classmethod
     def create(cls, config_dict, ctx):
-        config = OracleConfig.parse_obj(config_dict)
+        config = OracleConfig.model_validate(config_dict)
         return cls(config, ctx)
 
     def get_db_name(self, inspector: Inspector) -> str:

datahub/ingestion/source/sql/postgres.py

@@ -212,7 +212,7 @@ class PostgresSource(SQLAlchemySource):
 
     @classmethod
     def create(cls, config_dict, ctx):
-        config = PostgresConfig.parse_obj(config_dict)
+        config = PostgresConfig.model_validate(config_dict)
         return cls(config, ctx)
 
     def _setup_rds_iam_event_listener(
@@ -288,7 +288,7 @@ class PostgresSource(SQLAlchemySource):
                 return {}
 
             for row in results:
-                data.append(ViewLineageEntry.parse_obj(row))
+                data.append(ViewLineageEntry.model_validate(row))
 
         lineage_elements: Dict[Tuple[str, str], List[str]] = defaultdict(list)
         # Loop over the lineages in the JSON data.

datahub/ingestion/source/sql/presto.py

@@ -115,7 +115,7 @@ class PrestoSource(TrinoSource):
 
     @classmethod
     def create(cls, config_dict, ctx):
-        config = PrestoConfig.parse_obj(config_dict)
+        config = PrestoConfig.model_validate(config_dict)
         return cls(config, ctx)
 
 

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