acryl-datahub 1.0.0rc18__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/configuration/env_vars.py

@@ -0,0 +1,331 @@
+# ABOUTME: Central registry for all environment variables used in metadata-ingestion.
+# ABOUTME: All environment variable reads should go through this module for discoverability and maintainability.
+
+import os
+from typing import Optional
+
+# ============================================================================
+# Core DataHub Configuration
+# ============================================================================
+
+
+def get_gms_url() -> Optional[str]:
+    """Complete GMS URL (takes precedence over separate host/port)."""
+    return os.getenv("DATAHUB_GMS_URL")
+
+
+def get_gms_host() -> Optional[str]:
+    """GMS host (fallback for URL, deprecated)."""
+    return os.getenv("DATAHUB_GMS_HOST")
+
+
+def get_gms_port() -> Optional[str]:
+    """GMS port number."""
+    return os.getenv("DATAHUB_GMS_PORT")
+
+
+def get_gms_protocol() -> str:
+    """Protocol for GMS connection (http/https)."""
+    return os.getenv("DATAHUB_GMS_PROTOCOL", "http")
+
+
+def get_gms_token() -> Optional[str]:
+    """Authentication token for GMS."""
+    return os.getenv("DATAHUB_GMS_TOKEN")
+
+
+def get_system_client_id() -> Optional[str]:
+    """System client ID for OAuth/auth."""
+    return os.getenv("DATAHUB_SYSTEM_CLIENT_ID")
+
+
+def get_system_client_secret() -> Optional[str]:
+    """System client secret for OAuth/auth."""
+    return os.getenv("DATAHUB_SYSTEM_CLIENT_SECRET")
+
+
+def get_skip_config() -> bool:
+    """Skip loading config file (forces env variables)."""
+    return os.getenv("DATAHUB_SKIP_CONFIG", "").lower() == "true"
+
+
+def get_gms_base_path() -> str:
+    """Base path for GMS API endpoints."""
+    return os.getenv("DATAHUB_GMS_BASE_PATH", "")
+
+
+# ============================================================================
+# REST Emitter Configuration
+# ============================================================================
+
+
+def get_rest_emitter_default_retry_max_times() -> str:
+    """Max retry attempts for failed requests."""
+    return os.getenv("DATAHUB_REST_EMITTER_DEFAULT_RETRY_MAX_TIMES", "4")
+
+
+def get_rest_emitter_batch_max_payload_bytes() -> int:
+    """Maximum payload size in bytes for batch operations."""
+    return int(
+        os.getenv("DATAHUB_REST_EMITTER_BATCH_MAX_PAYLOAD_BYTES", str(15 * 1024 * 1024))
+    )
+
+
+def get_rest_emitter_batch_max_payload_length() -> int:
+    """Maximum number of MCPs per batch."""
+    return int(os.getenv("DATAHUB_REST_EMITTER_BATCH_MAX_PAYLOAD_LENGTH", "200"))
+
+
+def get_emit_mode() -> Optional[str]:
+    """Emission mode (SYNC_PRIMARY, SYNC_WAIT, ASYNC, ASYNC_WAIT)."""
+    return os.getenv("DATAHUB_EMIT_MODE")
+
+
+def get_rest_emitter_default_endpoint() -> Optional[str]:
+    """REST endpoint type (RESTLI or OPENAPI)."""
+    return os.getenv("DATAHUB_REST_EMITTER_DEFAULT_ENDPOINT")
+
+
+def get_emitter_trace() -> bool:
+    """Enable detailed emitter tracing."""
+    return os.getenv("DATAHUB_EMITTER_TRACE", "").lower() == "true"
+
+
+# ============================================================================
+# REST Sink Configuration
+# ============================================================================
+
+
+def get_rest_sink_default_max_threads() -> int:
+    """Max thread pool size for async operations."""
+    return int(os.getenv("DATAHUB_REST_SINK_DEFAULT_MAX_THREADS", "15"))
+
+
+def get_rest_sink_default_mode() -> Optional[str]:
+    """Sink mode (SYNC, ASYNC, ASYNC_BATCH)."""
+    return os.getenv("DATAHUB_REST_SINK_DEFAULT_MODE")
+
+
+# ============================================================================
+# Telemetry & Monitoring
+# ============================================================================
+
+
+def get_telemetry_timeout() -> str:
+    """Telemetry timeout in seconds."""
+    return os.getenv("DATAHUB_TELEMETRY_TIMEOUT", "10")
+
+
+def get_sentry_dsn() -> Optional[str]:
+    """Sentry error tracking DSN."""
+    return os.getenv("SENTRY_DSN")
+
+
+def get_sentry_environment() -> str:
+    """Sentry environment (dev/prod)."""
+    return os.getenv("SENTRY_ENVIRONMENT", "dev")
+
+
+# ============================================================================
+# Logging & Debug Configuration
+# ============================================================================
+
+
+def get_suppress_logging_manager() -> Optional[str]:
+    """Suppress DataHub logging manager initialization."""
+    return os.getenv("DATAHUB_SUPPRESS_LOGGING_MANAGER")
+
+
+def get_no_color() -> bool:
+    """Disable colored logging output."""
+    return os.getenv("NO_COLOR", "").lower() == "true"
+
+
+def get_test_mode() -> Optional[str]:
+    """Indicates running in test context."""
+    return os.getenv("DATAHUB_TEST_MODE")
+
+
+def get_debug() -> bool:
+    """Enable debug mode."""
+    return os.getenv("DATAHUB_DEBUG", "").lower() == "true"
+
+
+# ============================================================================
+# Data Processing Configuration
+# ============================================================================
+
+
+def get_sql_agg_query_log() -> str:
+    """SQL aggregator query logging level."""
+    return os.getenv("DATAHUB_SQL_AGG_QUERY_LOG", "DISABLED")
+
+
+def get_dataset_urn_to_lower() -> str:
+    """Convert dataset URNs to lowercase."""
+    return os.getenv("DATAHUB_DATASET_URN_TO_LOWER", "false")
+
+
+# ============================================================================
+# Integration-Specific Configuration
+# ============================================================================
+
+
+def get_kafka_schema_registry_url() -> Optional[str]:
+    """Kafka schema registry URL."""
+    return os.getenv("KAFKA_SCHEMAREGISTRY_URL")
+
+
+def get_spark_version() -> Optional[str]:
+    """Spark version (for S3 source)."""
+    return os.getenv("SPARK_VERSION")
+
+
+def get_bigquery_schema_parallelism() -> int:
+    """Parallelism level for BigQuery schema extraction."""
+    return int(os.getenv("DATAHUB_BIGQUERY_SCHEMA_PARALLELISM", "20"))
+
+
+def get_snowflake_schema_parallelism() -> int:
+    """Parallelism level for Snowflake schema extraction."""
+    return int(os.getenv("DATAHUB_SNOWFLAKE_SCHEMA_PARALLELISM", "20"))
+
+
+def get_powerbi_m_query_parse_timeout() -> int:
+    """Timeout for PowerBI M query parsing."""
+    return int(os.getenv("DATAHUB_POWERBI_M_QUERY_PARSE_TIMEOUT", "60"))
+
+
+def get_trace_powerbi_mquery_parser() -> bool:
+    """Enable PowerBI M query parser tracing."""
+    return os.getenv("DATAHUB_TRACE_POWERBI_MQUERY_PARSER", "").lower() == "true"
+
+
+def get_lookml_git_test_ssh_key() -> Optional[str]:
+    """SSH key for LookML Git tests."""
+    return os.getenv("DATAHUB_LOOKML_GIT_TEST_SSH_KEY")
+
+
+# ============================================================================
+# AWS/Cloud Configuration
+# ============================================================================
+
+
+def get_aws_lambda_function_name() -> Optional[str]:
+    """Indicates running in AWS Lambda."""
+    return os.getenv("AWS_LAMBDA_FUNCTION_NAME")
+
+
+def get_aws_execution_env() -> Optional[str]:
+    """AWS execution environment."""
+    return os.getenv("AWS_EXECUTION_ENV")
+
+
+def get_aws_web_identity_token_file() -> Optional[str]:
+    """OIDC token file path."""
+    return os.getenv("AWS_WEB_IDENTITY_TOKEN_FILE")
+
+
+def get_aws_role_arn() -> Optional[str]:
+    """AWS role ARN for OIDC."""
+    return os.getenv("AWS_ROLE_ARN")
+
+
+def get_aws_app_runner_service_id() -> Optional[str]:
+    """AWS App Runner service ID."""
+    return os.getenv("AWS_APP_RUNNER_SERVICE_ID")
+
+
+def get_ecs_container_metadata_uri_v4() -> Optional[str]:
+    """ECS metadata endpoint v4."""
+    return os.getenv("ECS_CONTAINER_METADATA_URI_V4")
+
+
+def get_ecs_container_metadata_uri() -> Optional[str]:
+    """ECS metadata endpoint v3."""
+    return os.getenv("ECS_CONTAINER_METADATA_URI")
+
+
+def get_elastic_beanstalk_environment_name() -> Optional[str]:
+    """Elastic Beanstalk environment."""
+    return os.getenv("ELASTIC_BEANSTALK_ENVIRONMENT_NAME")
+
+
+# ============================================================================
+# Docker & Local Development
+# ============================================================================
+
+
+def get_compose_project_name() -> str:
+    """Docker Compose project name."""
+    return os.getenv("DATAHUB_COMPOSE_PROJECT_NAME", "datahub")
+
+
+def get_docker_compose_base() -> Optional[str]:
+    """Base path for Docker Compose files."""
+    return os.getenv("DOCKER_COMPOSE_BASE")
+
+
+def get_datahub_version() -> Optional[str]:
+    """DataHub version (set during docker init)."""
+    return os.getenv("DATAHUB_VERSION")
+
+
+def get_mapped_mysql_port() -> Optional[str]:
+    """MySQL port mapping (set during docker init)."""
+    return os.getenv("DATAHUB_MAPPED_MYSQL_PORT")
+
+
+def get_mapped_kafka_broker_port() -> Optional[str]:
+    """Kafka broker port mapping (set during docker init)."""
+    return os.getenv("DATAHUB_MAPPED_KAFKA_BROKER_PORT")
+
+
+def get_mapped_elastic_port() -> Optional[str]:
+    """Elasticsearch port mapping (set during docker init)."""
+    return os.getenv("DATAHUB_MAPPED_ELASTIC_PORT")
+
+
+def get_metadata_service_auth_enabled() -> str:
+    """Enable/disable auth in Docker."""
+    return os.getenv("METADATA_SERVICE_AUTH_ENABLED", "false")
+
+
+def get_ui_ingestion_default_cli_version() -> Optional[str]:
+    """CLI version for UI ingestion (set during init)."""
+    return os.getenv("UI_INGESTION_DEFAULT_CLI_VERSION")
+
+
+# ============================================================================
+# Utility & Helper Configuration
+# ============================================================================
+
+
+def get_datahub_component() -> str:
+    """Component name for user agent tracking."""
+    return os.getenv("DATAHUB_COMPONENT", "datahub")
+
+
+def get_force_local_quickstart_mapping() -> str:
+    """Force local quickstart mapping file."""
+    return os.getenv("FORCE_LOCAL_QUICKSTART_MAPPING", "")
+
+
+def get_dataproduct_external_url() -> Optional[str]:
+    """External URL for data products."""
+    return os.getenv("DATAHUB_DATAPRODUCT_EXTERNAL_URL")
+
+
+def get_override_sqlite_version_req() -> str:
+    """Override SQLite version requirement."""
+    return os.getenv("OVERRIDE_SQLITE_VERSION_REQ", "")
+
+
+def get_update_entity_registry() -> str:
+    """Update entity registry during tests."""
+    return os.getenv("UPDATE_ENTITY_REGISTRY", "false")
+
+
+def get_ci() -> Optional[str]:
+    """Indicates running in CI environment."""
+    return os.getenv("CI")

datahub/configuration/import_resolver.py

@@ -1,15 +1,18 @@
-from typing import TypeVar, Union
+from typing import TYPE_CHECKING, Type, TypeVar, Union
 
 import pydantic
 
 from datahub.ingestion.api.registry import import_path
 
-T = TypeVar("T")
+if TYPE_CHECKING:
+    from pydantic.deprecated.class_validators import V1Validator
 
+_T = TypeVar("_T")
 
-def _pydantic_resolver(v: Union[T, str]) -> T:
+
+def _pydantic_resolver(cls: Type, v: Union[str, _T]) -> _T:
     return import_path(v) if isinstance(v, str) else v
 
 
-def pydantic_resolve_key(field: str) -> classmethod:
+def pydantic_resolve_key(field: str) -> "V1Validator":
     return pydantic.validator(field, pre=True, allow_reuse=True)(_pydantic_resolver)

datahub/configuration/kafka.py

@@ -1,16 +1,36 @@
 from pydantic import Field, validator
 
 from datahub.configuration.common import ConfigModel, ConfigurationError
+from datahub.configuration.env_vars import (
+    get_gms_base_path,
+    get_kafka_schema_registry_url,
+)
 from datahub.configuration.kafka_consumer_config import CallableConsumerConfig
 from datahub.configuration.validate_host_port import validate_host_port
 
 
+def _get_schema_registry_url() -> str:
+    """Get schema registry URL with proper base path handling."""
+    explicit_url = get_kafka_schema_registry_url()
+    if explicit_url:
+        return explicit_url
+
+    base_path = get_gms_base_path()
+    if base_path in ("/", ""):
+        base_path = ""
+
+    return f"http://localhost:8080{base_path}/schema-registry/api/"
+
+
 class _KafkaConnectionConfig(ConfigModel):
     # bootstrap servers
     bootstrap: str = "localhost:9092"
 
     # schema registry location
-    schema_registry_url: str = "http://localhost:8080/schema-registry/api/"
+    schema_registry_url: str = Field(
+        default_factory=_get_schema_registry_url,
+        description="Schema registry URL. Can be overridden with KAFKA_SCHEMAREGISTRY_URL environment variable, or will use DATAHUB_GMS_BASE_PATH if not set.",
+    )
 
     schema_registry_config: dict = Field(
         default_factory=dict,

datahub/configuration/pydantic_migration_helpers.py

@@ -1,20 +1,13 @@
 import pydantic.version
 from packaging.version import Version
 
-PYDANTIC_VERSION_2: bool
-if Version(pydantic.version.VERSION) >= Version("2.0"):
-    PYDANTIC_VERSION_2 = True
-else:
-    PYDANTIC_VERSION_2 = False
-
+_pydantic_version = Version(pydantic.version.VERSION)
 
-# This can be used to silence deprecation warnings while we migrate.
-if PYDANTIC_VERSION_2:
-    from pydantic import PydanticDeprecatedSince20  # type: ignore
-else:
+PYDANTIC_VERSION_2 = _pydantic_version >= Version("2.0")
 
-    class PydanticDeprecatedSince20(Warning):  # type: ignore
-        pass
+# The pydantic.Discriminator type was added in v2.5.0.
+# https://docs.pydantic.dev/latest/changelog/#v250-2023-11-13
+PYDANTIC_SUPPORTS_CALLABLE_DISCRIMINATOR = _pydantic_version >= Version("2.5.0")
 
 
 if PYDANTIC_VERSION_2:
@@ -50,7 +43,7 @@ class v1_ConfigModel(v1_BaseModel):
 
 __all__ = [
     "PYDANTIC_VERSION_2",
-    "PydanticDeprecatedSince20",
+    "PYDANTIC_SUPPORTS_CALLABLE_DISCRIMINATOR",
     "GenericModel",
     "v1_ConfigModel",
     "v1_Field",

datahub/configuration/source_common.py

@@ -1,6 +1,6 @@
 from typing import Dict, Optional
 
-from pydantic import validator
+import pydantic
 from pydantic.fields import Field
 
 from datahub.configuration.common import ConfigModel
@@ -16,7 +16,7 @@ class PlatformInstanceConfigMixin(ConfigModel):
         default=None,
         description="The instance of the platform that all assets produced by this recipe belong to. "
         "This should be unique within the platform. "
-        "See https://datahubproject.io/docs/platform-instances/ for more details.",
+        "See https://docs.datahub.com/docs/platform-instances/ for more details.",
     )
 
 
@@ -30,7 +30,8 @@ class EnvConfigMixin(ConfigModel):
         description="The environment that all assets produced by this connector belong to",
     )
 
-    @validator("env")
+    @pydantic.field_validator("env", mode="after")
+    @classmethod
     def env_must_be_one_of(cls, v: str) -> str:
         if v.upper() not in ALL_ENV_TYPES:
             raise ValueError(f"env must be one of {ALL_ENV_TYPES}, found {v}")

datahub/configuration/validate_field_deprecation.py

@@ -1,11 +1,14 @@
 import warnings
-from typing import Any, Optional, Type
+from typing import TYPE_CHECKING, Any, Optional, Type
 
 import pydantic
 
 from datahub.configuration.common import ConfigurationWarning
 from datahub.utilities.global_warning_util import add_global_warning
 
+if TYPE_CHECKING:
+    from pydantic.deprecated.class_validators import V1RootValidator
+
 _unset = object()
 
 
@@ -13,7 +16,7 @@ def pydantic_field_deprecated(
     field: str,
     warn_if_value_is_not: Any = _unset,
     message: Optional[str] = None,
-) -> classmethod:
+) -> "V1RootValidator":
     if message:
         output = message
     else:

datahub/configuration/validate_field_removal.py

@@ -1,15 +1,18 @@
 import warnings
-from typing import Type
+from typing import TYPE_CHECKING, Type
 
 import pydantic
 
 from datahub.configuration.common import ConfigurationWarning
 
+if TYPE_CHECKING:
+    from pydantic.deprecated.class_validators import V1RootValidator
+
 
 def pydantic_removed_field(
     field: str,
     print_warning: bool = True,
-) -> classmethod:
+) -> "V1RootValidator":
     def _validate_field_removal(cls: Type, values: dict) -> dict:
         if field in values:
             if print_warning:
@@ -21,6 +24,9 @@ def pydantic_removed_field(
             values.pop(field)
         return values
 
+    # Mark the function as handling a removed field for doc generation
+    _validate_field_removal._doc_removed_field = field  # type: ignore[attr-defined]
+
     # Hack: Pydantic maintains unique list of validators by referring its __name__.
     # https://github.com/pydantic/pydantic/blob/v1.10.9/pydantic/main.py#L264
     # This hack ensures that multiple field removals do not overwrite each other.

datahub/configuration/validate_field_rename.py

@@ -1,11 +1,14 @@
 import warnings
-from typing import Callable, Type, TypeVar
+from typing import TYPE_CHECKING, Callable, Type, TypeVar
 
 import pydantic
 
 from datahub.configuration.common import ConfigurationWarning
 from datahub.utilities.global_warning_util import add_global_warning
 
+if TYPE_CHECKING:
+    from pydantic.deprecated.class_validators import V1RootValidator
+
 _T = TypeVar("_T")
 
 
@@ -18,7 +21,7 @@ def pydantic_renamed_field(
     new_name: str,
     transform: Callable = _default_rename_transform,
     print_warning: bool = True,
-) -> classmethod:
+) -> "V1RootValidator":
     def _validate_field_rename(cls: Type, values: dict) -> dict:
         if old_name in values:
             if new_name in values:
@@ -49,6 +52,4 @@ def pydantic_renamed_field(
     # validator with pre=True gets all the values that were passed in.
     # Given that a renamed field doesn't show up in the fields list, we can't use
     # the field-level validator, even with a different field name.
-    return pydantic.root_validator(pre=True, skip_on_failure=True, allow_reuse=True)(
-        _validate_field_rename
-    )
+    return pydantic.root_validator(pre=True, allow_reuse=True)(_validate_field_rename)

datahub/configuration/validate_multiline_string.py

@@ -1,9 +1,12 @@
-from typing import Optional, Type, Union
+from typing import TYPE_CHECKING, Optional, Type, Union
 
 import pydantic
 
+if TYPE_CHECKING:
+    from pydantic.deprecated.class_validators import V1Validator
 
-def pydantic_multiline_string(field: str) -> classmethod:
+
+def pydantic_multiline_string(field: str) -> "V1Validator":
     """If the field is present and contains an escaped newline, replace it with a real newline.
 
     This makes the assumption that the field value is never supposed to have a

datahub/emitter/mce_builder.py

@@ -3,7 +3,6 @@
 import hashlib
 import json
 import logging
-import os
 import re
 import time
 from datetime import datetime, timezone
@@ -26,6 +25,7 @@ import typing_inspect
 from avrogen.dict_wrapper import DictWrapper
 from typing_extensions import assert_never
 
+from datahub.configuration.env_vars import get_dataset_urn_to_lower
 from datahub.emitter.enum_helpers import get_enum_options
 from datahub.metadata.schema_classes import (
     AssertionKeyClass,
@@ -59,6 +59,7 @@ from datahub.metadata.urns import (
     DataJobUrn,
     DataPlatformUrn,
     DatasetUrn,
+    OwnershipTypeUrn,
     TagUrn,
 )
 from datahub.utilities.urn_encoder import UrnEncoder
@@ -71,9 +72,7 @@ ALL_ENV_TYPES: Set[str] = set(get_enum_options(FabricTypeClass))
 
 DEFAULT_FLOW_CLUSTER = "prod"
 UNKNOWN_USER = "urn:li:corpuser:unknown"
-DATASET_URN_TO_LOWER: bool = (
-    os.getenv("DATAHUB_DATASET_URN_TO_LOWER", "false") == "true"
-)
+DATASET_URN_TO_LOWER: bool = get_dataset_urn_to_lower() == "true"
 
 if TYPE_CHECKING:
     from datahub.emitter.mcp_builder import DatahubKey
@@ -125,9 +124,7 @@ def parse_ts_millis(ts: Optional[float]) -> Optional[datetime]:
 
 
 def make_data_platform_urn(platform: str) -> str:
-    if platform.startswith("urn:li:dataPlatform:"):
-        return platform
-    return DataPlatformUrn.create_from_id(platform).urn()
+    return DataPlatformUrn(platform).urn()
 
 
 def make_dataset_urn(platform: str, name: str, env: str = DEFAULT_ENV) -> str:
@@ -377,6 +374,12 @@ def make_domain_urn(domain: str) -> str:
     return f"urn:li:domain:{domain}"
 
 
+def make_data_product_urn(data_product_id: str) -> str:
+    if data_product_id.startswith("urn:li:dataProduct:"):
+        return data_product_id
+    return f"urn:li:dataProduct:{data_product_id}"
+
+
 def make_ml_primary_key_urn(feature_table_name: str, primary_key_name: str) -> str:
     return f"urn:li:mlPrimaryKey:({feature_table_name},{primary_key_name})"
 
@@ -408,7 +411,8 @@ def make_ml_model_group_urn(platform: str, group_name: str, env: str) -> str:
 
 def validate_ownership_type(ownership_type: str) -> Tuple[str, Optional[str]]:
     if ownership_type.startswith("urn:li:"):
-        return OwnershipTypeClass.CUSTOM, ownership_type
+        ownership_type_urn = OwnershipTypeUrn.from_string(ownership_type)
+        return OwnershipTypeClass.CUSTOM, ownership_type_urn.urn()
     ownership_type = ownership_type.upper()
     if ownership_type in get_enum_options(OwnershipTypeClass):
         return ownership_type, None