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

datahub/ingestion/source/sql/athena.py

@@ -16,6 +16,7 @@ from sqlalchemy.engine.reflection import Inspector
 from sqlalchemy.types import TypeEngine
 from sqlalchemy_bigquery import STRUCT
 
+from datahub.configuration.common import HiddenFromDocs
 from datahub.configuration.validate_field_rename import pydantic_renamed_field
 from datahub.emitter.mcp_builder import ContainerKey, DatabaseKey
 from datahub.ingestion.api.decorators import (
@@ -251,7 +252,7 @@ class CustomAthenaRestDialect(AthenaRestDialect):
 
 
 class AthenaConfig(SQLCommonConfig):
-    scheme: str = "awsathena+rest"
+    scheme: HiddenFromDocs[str] = "awsathena+rest"
     username: Optional[str] = pydantic.Field(
         default=None,
         description="Username credential. If not specified, detected with boto3 rules. See https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html",

datahub/ingestion/source/sql/clickhouse.py

@@ -18,6 +18,7 @@ from sqlalchemy.sql import sqltypes
 from sqlalchemy.types import BOOLEAN, DATE, DATETIME, INTEGER
 
 import datahub.emitter.mce_builder as builder
+from datahub.configuration.common import HiddenFromDocs, LaxStr
 from datahub.configuration.source_common import DatasetLineageProviderConfigBase
 from datahub.configuration.time_window_config import BaseTimeWindowConfig
 from datahub.configuration.validate_field_deprecation import pydantic_field_deprecated
@@ -128,16 +129,20 @@ class ClickHouseConfig(
 ):
     # defaults
     host_port: str = Field(default="localhost:8123", description="ClickHouse host URL.")
-    scheme: str = Field(default="clickhouse", description="", hidden_from_docs=True)
+    scheme: HiddenFromDocs[str] = Field(default="clickhouse")
     password: pydantic.SecretStr = Field(
         default=pydantic.SecretStr(""), description="password"
     )
-    secure: Optional[bool] = Field(default=None, description="")
-    protocol: Optional[str] = Field(default=None, description="")
+    secure: Optional[bool] = Field(
+        default=None, description="[deprecated] Use uri_opts instead."
+    )
+    protocol: Optional[str] = Field(
+        default=None, description="[deprecated] Use uri_opts instead."
+    )
     _deprecate_secure = pydantic_field_deprecated("secure")
     _deprecate_protocol = pydantic_field_deprecated("protocol")
 
-    uri_opts: Dict[str, str] = Field(
+    uri_opts: Dict[str, LaxStr] = Field(
         default={},
         description="The part of the URI and it's used to provide additional configuration options or parameters for the database connection.",
     )
@@ -185,9 +190,9 @@ class ClickHouseConfig(
                 "Initializing uri_opts from deprecated secure or protocol options"
             )
             values["uri_opts"] = {}
-            if secure:
-                values["uri_opts"]["secure"] = secure
-            if protocol:
+            if secure is not None:
+                values["uri_opts"]["secure"] = str(secure)
+            if protocol is not None:
                 values["uri_opts"]["protocol"] = protocol
             logger.debug(f"uri_opts: {uri_opts}")
         elif (secure or protocol) and uri_opts:

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,9 +9,10 @@ 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
 
 import datahub.emitter.mce_builder as builder
 from datahub.configuration.common import AllowDenyPattern
@@ -109,6 +110,12 @@ logger = logging.getLogger(__name__)
 
 PAGE_SIZE = 25
 
+# Retry configuration constants
+RETRY_MAX_TIMES = 3
+RETRY_STATUS_CODES = [429, 500, 502, 503, 504]
+RETRY_BACKOFF_FACTOR = 1
+RETRY_ALLOWED_METHODS = ["GET"]
+
 
 chart_type_from_viz_type = {
     "line": ChartTypeClass.LINE,
@@ -328,6 +335,19 @@ class SupersetSource(StatefulIngestionSourceBase):
         logger.debug("Got access token from superset")
 
         requests_session = requests.Session()
+
+        # Configure retry strategy for transient failures
+        retry_strategy = Retry(
+            total=RETRY_MAX_TIMES,
+            status_forcelist=RETRY_STATUS_CODES,
+            backoff_factor=RETRY_BACKOFF_FACTOR,
+            allowed_methods=RETRY_ALLOWED_METHODS,
+            raise_on_status=False,
+        )
+        adapter = HTTPAdapter(max_retries=retry_strategy)
+        requests_session.mount("http://", adapter)
+        requests_session.mount("https://", adapter)
+
         requests_session.headers.update(
             {
                 "Authorization": f"Bearer {self.access_token}",
@@ -360,8 +380,13 @@ class SupersetSource(StatefulIngestionSourceBase):
             )
 
             if response.status_code != 200:
-                logger.warning(f"Failed to get {entity_type} data: {response.text}")
-                continue
+                self.report.warning(
+                    title="Failed to fetch data from Superset API",
+                    message="Incomplete metadata extraction due to Superset API failure",
+                    context=f"Entity Type: {entity_type}, HTTP Status Code: {response.status_code}, Page: {current_page}. Response: {response.text}",
+                )
+                # we stop pagination for this entity type and we continue the overall ingestion
+                break
 
             payload = response.json()
             # Update total_items with the actual count from the response