acryl-datahub 1.2.0.10rc3__py3-none-any.whl → 1.2.0.10rc4__py3-none-any.whl

datahub/ingestion/source/sql/cockroachdb.py

@@ -1,6 +1,6 @@
 from pydantic.fields import Field
 
-from datahub.configuration.common import AllowDenyPattern
+from datahub.configuration.common import AllowDenyPattern, HiddenFromDocs
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import (
     SourceCapability,
@@ -14,8 +14,10 @@ from datahub.ingestion.source.sql.postgres import PostgresConfig, PostgresSource
 
 
 class CockroachDBConfig(PostgresConfig):
-    scheme = Field(default="cockroachdb+psycopg2", description="database scheme")
-    schema_pattern = Field(
+    scheme: HiddenFromDocs[str] = Field(
+        default="cockroachdb+psycopg2", description="database scheme"
+    )
+    schema_pattern: AllowDenyPattern = Field(
         default=AllowDenyPattern(deny=["information_schema", "crdb_internal"])
     )
 

datahub/ingestion/source/sql/druid.py

@@ -6,7 +6,7 @@ from pydantic.fields import Field
 from pydruid.db.sqlalchemy import DruidDialect
 from sqlalchemy.exc import ResourceClosedError
 
-from datahub.configuration.common import AllowDenyPattern
+from datahub.configuration.common import AllowDenyPattern, HiddenFromDocs
 from datahub.ingestion.api.decorators import (
     SourceCapability,
     SupportStatus,
@@ -34,7 +34,7 @@ DruidDialect.get_table_names = get_table_names
 
 class DruidConfig(BasicSQLAlchemyConfig):
     # defaults
-    scheme: str = "druid"
+    scheme: HiddenFromDocs[str] = "druid"
     schema_pattern: AllowDenyPattern = Field(
         default=AllowDenyPattern(deny=["^(lookup|sysgit|view).*"]),
         description="regex patterns for schemas to filter in ingestion.",

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.class_validators import validator
+from pydantic import validator
 from pydantic.fields import Field
 
 # This import verifies that the dependencies are available.
@@ -14,6 +14,7 @@ from pyhive import hive  # noqa: F401
 from pyhive.sqlalchemy_hive import HiveDate, HiveDecimal, HiveDialect, HiveTimestamp
 from sqlalchemy.engine.reflection import Inspector
 
+from datahub.configuration.common import HiddenFromDocs
 from datahub.emitter.mce_builder import (
     make_data_platform_urn,
     make_dataplatform_instance_urn,
@@ -651,10 +652,10 @@ HiveDialect.get_view_definition = get_view_definition_patched
 
 class HiveConfig(TwoTierSQLAlchemyConfig):
     # defaults
-    scheme: str = Field(default="hive", hidden_from_docs=True)
+    scheme: HiddenFromDocs[str] = Field(default="hive")
 
     # Overriding as table location lineage is richer implementation here than with include_table_location_lineage
-    include_table_location_lineage: bool = Field(default=False, hidden_from_docs=True)
+    include_table_location_lineage: HiddenFromDocs[bool] = Field(default=False)
 
     emit_storage_lineage: bool = Field(
         default=False,

datahub/ingestion/source/sql/hive_metastore.py

@@ -1,17 +1,15 @@
 import base64
+import dataclasses
 import json
 import logging
 from collections import namedtuple
 from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 
-from pydantic.dataclasses import dataclass
-from pydantic.fields import Field
-
-# This import verifies that the dependencies are available.
+from pydantic import Field
 from sqlalchemy import create_engine, text
 from sqlalchemy.engine.reflection import Inspector
 
-from datahub.configuration.common import AllowDenyPattern
+from datahub.configuration.common import AllowDenyPattern, HiddenFromDocs
 from datahub.emitter.mce_builder import make_dataset_urn_with_platform_instance
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.common import PipelineContext
@@ -73,7 +71,7 @@ class HiveMetastoreConfigMode(StrEnum):
     trino = "trino"
 
 
-@dataclass
+@dataclasses.dataclass
 class ViewDataset:
     dataset_name: str
     schema_name: str
@@ -99,7 +97,7 @@ class HiveMetastore(BasicSQLAlchemyConfig):
         default="localhost:3306",
         description="Host URL and port to connect to. Example: localhost:3306",
     )
-    scheme: str = Field(default="mysql+pymysql", description="", hidden_from_docs=True)
+    scheme: HiddenFromDocs[str] = Field(default="mysql+pymysql")
 
     database_pattern: AllowDenyPattern = Field(
         default=AllowDenyPattern.allow_all(),
@@ -123,8 +121,8 @@ class HiveMetastore(BasicSQLAlchemyConfig):
         description="Dataset Subtype name to be 'Table' or 'View' Valid options: ['True', 'False']",
     )
 
-    include_view_lineage: bool = Field(
-        default=False, description="", hidden_from_docs=True
+    include_view_lineage: HiddenFromDocs[bool] = Field(
+        default=False,
     )
 
     include_catalog_name_in_ids: bool = Field(

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

@@ -13,7 +13,7 @@ from sqlalchemy.exc import ProgrammingError, ResourceClosedError
 from sqlalchemy.sql import quoted_name
 
 import datahub.metadata.schema_classes as models
-from datahub.configuration.common import AllowDenyPattern
+from datahub.configuration.common import AllowDenyPattern, HiddenFromDocs
 from datahub.configuration.pattern_utils import UUID_REGEX
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.common import PipelineContext
@@ -75,7 +75,7 @@ DEFAULT_TEMP_TABLES_PATTERNS = [
 class SQLServerConfig(BasicSQLAlchemyConfig):
     # defaults
     host_port: str = Field(default="localhost:1433", description="MSSQL host URL.")
-    scheme: str = Field(default="mssql+pytds", description="", hidden_from_docs=True)
+    scheme: HiddenFromDocs[str] = Field(default="mssql+pytds")
 
     # TODO: rename to include_procedures ?
     include_stored_procedures: bool = Field(

datahub/ingestion/source/sql/mysql.py

@@ -9,7 +9,7 @@ from sqlalchemy.dialects.mysql import BIT, base
 from sqlalchemy.dialects.mysql.enumerated import SET
 from sqlalchemy.engine.reflection import Inspector
 
-from datahub.configuration.common import AllowDenyPattern
+from datahub.configuration.common import AllowDenyPattern, HiddenFromDocs
 from datahub.ingestion.api.decorators import (
     SourceCapability,
     SupportStatus,
@@ -57,7 +57,7 @@ base.ischema_names["decimal128"] = DECIMAL128
 class MySQLConnectionConfig(SQLAlchemyConnectionConfig):
     # defaults
     host_port: str = Field(default="localhost:3306", description="MySQL host URL.")
-    scheme: str = "mysql+pymysql"
+    scheme: HiddenFromDocs[str] = "mysql+pymysql"
 
 
 class MySQLConfig(MySQLConnectionConfig, TwoTierSQLAlchemyConfig):

datahub/ingestion/source/sql/oracle.py

@@ -110,10 +110,10 @@ class OracleConfig(BasicSQLAlchemyConfig):
         return v
 
     @pydantic.validator("data_dictionary_mode")
-    def check_data_dictionary_mode(cls, values):
-        if values not in ("ALL", "DBA"):
+    def check_data_dictionary_mode(cls, value):
+        if value not in ("ALL", "DBA"):
             raise ValueError("Specify one of data dictionary views mode: 'ALL', 'DBA'.")
-        return values
+        return value
 
     @pydantic.validator("thick_mode_lib_dir", always=True)
     def check_thick_mode_lib_dir(cls, v, values):

datahub/ingestion/source/sql/presto.py

@@ -8,6 +8,7 @@ from sqlalchemy import exc, sql
 from sqlalchemy.engine import reflection
 from sqlalchemy.engine.base import Engine
 
+from datahub.configuration.common import HiddenFromDocs
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import (
     SourceCapability,
@@ -87,7 +88,7 @@ PrestoDialect._get_full_table = _get_full_table
 
 class PrestoConfig(TrinoConfig):
     # defaults
-    scheme: str = Field(default="presto", description="", hidden_from_docs=True)
+    scheme: HiddenFromDocs[str] = Field(default="presto")
 
 
 @platform_name("Presto", doc_order=1)

datahub/ingestion/source/sql/teradata.py

@@ -468,23 +468,23 @@ class TeradataConfig(BaseTeradataConfig, BaseTimeWindowConfig):
         ),
     )
 
-    database_pattern = Field(
+    database_pattern: AllowDenyPattern = Field(
         default=AllowDenyPattern(deny=EXCLUDED_DATABASES),
         description="Regex patterns for databases to filter in ingestion.",
     )
-    include_table_lineage = Field(
+    include_table_lineage: bool = Field(
         default=False,
         description="Whether to include table lineage in the ingestion. "
         "This requires to have the table lineage feature enabled.",
     )
 
-    include_view_lineage = Field(
+    include_view_lineage: bool = Field(
         default=True,
         description="Whether to include view lineage in the ingestion. "
         "This requires to have the view lineage feature enabled.",
     )
 
-    include_queries = Field(
+    include_queries: bool = Field(
         default=True,
         description="Whether to generate query entities for SQL queries. "
         "Query entities provide metadata about individual SQL queries including "

datahub/ingestion/source/sql/trino.py

@@ -18,6 +18,7 @@ from sqlalchemy.types import TypeEngine
 from trino.sqlalchemy import datatype
 from trino.sqlalchemy.dialect import TrinoDialect
 
+from datahub.configuration.common import HiddenFromDocs
 from datahub.configuration.source_common import (
     EnvConfigMixin,
     PlatformInstanceConfigMixin,
@@ -222,7 +223,7 @@ class ConnectorDetail(PlatformInstanceConfigMixin, EnvConfigMixin):
 
 class TrinoConfig(BasicSQLAlchemyConfig):
     # defaults
-    scheme: str = Field(default="trino", description="", hidden_from_docs=True)
+    scheme: HiddenFromDocs[str] = Field(default="trino")
     database: str = Field(description="database (catalog)")
 
     catalog_to_connector_details: Dict[str, ConnectorDetail] = Field(

datahub/ingestion/source/sql/two_tier_sql_source.py

@@ -7,7 +7,7 @@ from sqlalchemy import create_engine, inspect
 from sqlalchemy.engine import URL
 from sqlalchemy.engine.reflection import Inspector
 
-from datahub.configuration.common import AllowDenyPattern
+from datahub.configuration.common import AllowDenyPattern, HiddenFromDocs
 from datahub.configuration.validate_field_rename import pydantic_renamed_field
 from datahub.emitter.mcp_builder import ContainerKey
 from datahub.ingestion.api.workunit import MetadataWorkUnit
@@ -27,11 +27,10 @@ class TwoTierSQLAlchemyConfig(BasicSQLAlchemyConfig):
         default=AllowDenyPattern.allow_all(),
         description="Regex patterns for databases to filter in ingestion.",
     )
-    schema_pattern: AllowDenyPattern = Field(
+    schema_pattern: HiddenFromDocs[AllowDenyPattern] = Field(
         # The superclass contains a `schema_pattern` field, so we need this here
         # to override the documentation.
         default=AllowDenyPattern.allow_all(),
-        hidden_from_docs=True,
         description="Deprecated in favour of database_pattern.",
     )
 

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.class_validators import validator
+from pydantic import validator
 from vertica_sqlalchemy_dialect.base import VerticaInspector
 
 from datahub.configuration.common import AllowDenyPattern

datahub/ingestion/source/sql_queries.py

@@ -8,6 +8,7 @@ from typing import ClassVar, Iterable, List, Optional, Union
 
 from pydantic import BaseModel, Field, validator
 
+from datahub.configuration.common import HiddenFromDocs
 from datahub.configuration.datetimes import parse_user_datetime
 from datahub.configuration.source_common import (
     EnvConfigMixin,
@@ -66,22 +67,21 @@ class SqlQueriesSourceConfig(
         default=BaseUsageConfig(),
     )
 
-    use_schema_resolver: bool = Field(
+    use_schema_resolver: HiddenFromDocs[bool] = Field(
+        True,
         description="Read SchemaMetadata aspects from DataHub to aid in SQL parsing. Turn off only for testing.",
-        default=True,
-        hidden_from_docs=True,
     )
     default_db: Optional[str] = Field(
+        None,
         description="The default database to use for unqualified table names",
-        default=None,
     )
     default_schema: Optional[str] = Field(
+        None,
         description="The default schema to use for unqualified table names",
-        default=None,
     )
     override_dialect: Optional[str] = Field(
+        None,
         description="The SQL dialect to use when parsing queries. Overrides automatic dialect detection.",
-        default=None,
     )
 
 

datahub/ingestion/source/state/checkpoint.py

@@ -68,7 +68,11 @@ class CheckpointStateBase(ConfigModel):
 
     @staticmethod
     def _to_bytes_utf8(model: ConfigModel) -> bytes:
-        return model.json(exclude={"version", "serde"}).encode("utf-8")
+        pydantic_json = model.model_dump_json(exclude={"version", "serde"})
+        # We decode and re-encode so that Python's default whitespace is included.
+        # This is purely to keep tests consistent as we migrate to pydantic v2,
+        # and can be removed once we're fully migrated.
+        return json.dumps(json.loads(pydantic_json)).encode("utf-8")
 
     @staticmethod
     def _to_bytes_base85_json(

datahub/ingestion/source/state/entity_removal_state.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Iterable, List, Tuple, Type
+from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Tuple, Type
 
 import pydantic
 
@@ -8,13 +8,16 @@ from datahub.utilities.checkpoint_state_util import CheckpointStateUtil
 from datahub.utilities.dedup_list import deduplicate_list
 from datahub.utilities.urns.urn import guess_entity_type
 
+if TYPE_CHECKING:
+    from pydantic.deprecated.class_validators import V1RootValidator
+
 STATEFUL_INGESTION_IGNORED_ENTITY_TYPES = {
     "dataProcessInstance",
     "query",
 }
 
 
-def pydantic_state_migrator(mapping: Dict[str, str]) -> classmethod:
+def pydantic_state_migrator(mapping: Dict[str, str]) -> "V1RootValidator":
     # mapping would be something like:
     # {
     #    'encoded_view_urns': 'dataset',

datahub/ingestion/source/state/stateful_ingestion_base.py

@@ -10,6 +10,7 @@ from datahub.configuration.common import (
     ConfigModel,
     ConfigurationError,
     DynamicTypedConfig,
+    HiddenFromDocs,
 )
 from datahub.configuration.pydantic_migration_helpers import GenericModel
 from datahub.configuration.time_window_config import BaseTimeWindowConfig
@@ -55,25 +56,21 @@ class StatefulIngestionConfig(ConfigModel):
         description="Whether or not to enable stateful ingest. "
         "Default: True if a pipeline_name is set and either a datahub-rest sink or `datahub_api` is specified, otherwise False",
     )
-    max_checkpoint_state_size: pydantic.PositiveInt = Field(
+    max_checkpoint_state_size: HiddenFromDocs[pydantic.PositiveInt] = Field(
         default=2**24,  # 16 MB
         description="The maximum size of the checkpoint state in bytes. Default is 16MB",
-        hidden_from_docs=True,
     )
-    state_provider: Optional[DynamicTypedStateProviderConfig] = Field(
+    state_provider: HiddenFromDocs[Optional[DynamicTypedStateProviderConfig]] = Field(
         default=None,
         description="The ingestion state provider configuration.",
-        hidden_from_docs=True,
     )
-    ignore_old_state: bool = Field(
+    ignore_old_state: HiddenFromDocs[bool] = Field(
         default=False,
         description="If set to True, ignores the previous checkpoint state.",
-        hidden_from_docs=True,
     )
-    ignore_new_state: bool = Field(
+    ignore_new_state: HiddenFromDocs[bool] = Field(
         default=False,
         description="If set to True, ignores the current checkpoint state.",
-        hidden_from_docs=True,
     )
 
     @pydantic.root_validator(skip_on_failure=True)

datahub/ingestion/source/superset.py

@@ -9,8 +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
-from pydantic.class_validators import root_validator, validator
+from pydantic import BaseModel, root_validator, validator
 from pydantic.fields import Field
 from requests.adapters import HTTPAdapter
 from urllib3.util.retry import Retry

datahub/ingestion/source/tableau/tableau.py

@@ -3,6 +3,7 @@ import logging
 import re
 import time
 from collections import OrderedDict, defaultdict
+from copy import deepcopy
 from dataclasses import dataclass, field as dataclass_field
 from datetime import datetime, timedelta, timezone
 from functools import lru_cache
@@ -474,6 +475,13 @@ class TableauPageSizeConfig(ConfigModel):
         return self.database_table_page_size or self.page_size
 
 
+_IngestHiddenAssetsOptionsType = Literal["worksheet", "dashboard"]
+_IngestHiddenAssetsOptions: List[_IngestHiddenAssetsOptionsType] = [
+    "worksheet",
+    "dashboard",
+]
+
+
 class TableauConfig(
     DatasetLineageProviderConfigBase,
     StatefulIngestionConfigBase,
@@ -625,8 +633,8 @@ class TableauConfig(
         description="Configuration settings for ingesting Tableau groups and their capabilities as custom properties.",
     )
 
-    ingest_hidden_assets: Union[List[Literal["worksheet", "dashboard"]], bool] = Field(
-        default=["worksheet", "dashboard"],
+    ingest_hidden_assets: Union[List[_IngestHiddenAssetsOptionsType], bool] = Field(
+        _IngestHiddenAssetsOptions,
         description=(
             "When enabled, hidden worksheets and dashboards are ingested into Datahub."
             " If a dashboard or worksheet is hidden in Tableau the luid is blank."
@@ -648,6 +656,11 @@ class TableauConfig(
     # pre = True because we want to take some decision before pydantic initialize the configuration to default values
     @root_validator(pre=True)
     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.
+        # So a copy is performed first.
+        values = deepcopy(values)
+
         projects = values.get("projects")
         project_pattern = values.get("project_pattern")
         project_path_pattern = values.get("project_path_pattern")
@@ -659,6 +672,7 @@ class TableauConfig(
             values["project_pattern"] = AllowDenyPattern(
                 allow=[f"^{prj}$" for prj in projects]
             )
+            values.pop("projects")
         elif (project_pattern or project_path_pattern) and projects:
             raise ValueError(
                 "projects is deprecated. Please use project_path_pattern only."
@@ -670,7 +684,7 @@ class TableauConfig(
 
         return values
 
-    @root_validator()
+    @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")

datahub/ingestion/source/unity/config.py

@@ -8,7 +8,12 @@ import pydantic
 from pydantic import Field
 from typing_extensions import Literal
 
-from datahub.configuration.common import AllowDenyPattern, ConfigEnum, ConfigModel
+from datahub.configuration.common import (
+    AllowDenyPattern,
+    ConfigEnum,
+    ConfigModel,
+    HiddenFromDocs,
+)
 from datahub.configuration.source_common import (
     DatasetSourceConfigMixin,
     LowerCaseDatasetUrnConfigMixin,
@@ -285,10 +290,9 @@ class UnityCatalogSourceConfig(
         description="Limit the number of columns to get column level lineage. ",
     )
 
-    lineage_max_workers: int = pydantic.Field(
+    lineage_max_workers: HiddenFromDocs[int] = pydantic.Field(
         default=5 * (os.cpu_count() or 4),
         description="Number of worker threads to use for column lineage thread pool executor. Set to 1 to disable.",
-        hidden_from_docs=True,
     )
 
     databricks_api_page_size: int = pydantic.Field(

datahub/ingestion/source/usage/usage_common.py

@@ -18,7 +18,7 @@ import pydantic
 from pydantic.fields import Field
 
 import datahub.emitter.mce_builder as builder
-from datahub.configuration.common import AllowDenyPattern
+from datahub.configuration.common import AllowDenyPattern, HiddenFromDocs
 from datahub.configuration.time_window_config import (
     BaseTimeWindowConfig,
     BucketDuration,
@@ -194,13 +194,13 @@ class GenericAggregatedDataset(Generic[ResourceType]):
 
 
 class BaseUsageConfig(BaseTimeWindowConfig):
-    queries_character_limit: int = Field(
+    queries_character_limit: HiddenFromDocs[int] = Field(
+        # Hidden since we don't want to encourage people to break elasticsearch.
         default=DEFAULT_QUERIES_CHARACTER_LIMIT,
         description=(
             "Total character limit for all queries in a single usage aspect."
             " Queries will be truncated to length `queries_character_limit / top_n_queries`."
         ),
-        hidden_from_docs=True,  # Don't want to encourage people to break elasticsearch
     )
 
     top_n_queries: pydantic.PositiveInt = Field(

datahub/ingestion/source_config/pulsar.py

@@ -2,6 +2,7 @@ import re
 from typing import Dict, List, Optional, Union
 from urllib.parse import urlparse
 
+import pydantic
 from pydantic import Field, validator
 
 from datahub.configuration.common import AllowDenyPattern
@@ -121,7 +122,8 @@ class PulsarSourceConfig(
             )
         return client_secret
 
-    @validator("web_service_url")
+    @pydantic.field_validator("web_service_url", mode="after")
+    @classmethod
     def web_service_url_scheme_host_port(cls, val: str) -> str:
         # Tokenize the web url
         url = urlparse(val)

datahub/sdk/search_filters.py

@@ -39,13 +39,7 @@ _OrFilters = List[_AndSearchFilterRule]
 
 
 class _BaseFilter(ConfigModel):
-    class Config:
-        # We can't wrap this in a TYPE_CHECKING block because the pydantic plugin
-        # doesn't recognize it properly. So unfortunately we'll need to live
-        # with the deprecation warning w/ pydantic v2.
-        allow_population_by_field_name = True
-        if PYDANTIC_VERSION_2:
-            populate_by_name = True
+    model_config = pydantic.ConfigDict(populate_by_name=True)
 
     @abc.abstractmethod
     def compile(self) -> _OrFilters: ...