acryl-datahub 1.1.1rc4__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

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
@@ -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,
@@ -72,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
@@ -376,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})"
 

datahub/emitter/rest_emitter.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import functools
 import json
 import logging
-import os
+import re
 import time
 from collections import defaultdict
 from dataclasses import dataclass
@@ -32,7 +32,6 @@ from typing_extensions import deprecated
 from datahub._version import nice_version_name
 from datahub.cli import config_utils
 from datahub.cli.cli_utils import ensure_has_system_metadata, fixup_gms_url, get_or_else
-from datahub.cli.env_utils import get_boolean_env_variable
 from datahub.configuration.common import (
     ConfigEnum,
     ConfigModel,
@@ -41,6 +40,14 @@ from datahub.configuration.common import (
     TraceTimeoutError,
     TraceValidationError,
 )
+from datahub.configuration.env_vars import (
+    get_emit_mode,
+    get_emitter_trace,
+    get_rest_emitter_batch_max_payload_bytes,
+    get_rest_emitter_batch_max_payload_length,
+    get_rest_emitter_default_endpoint,
+    get_rest_emitter_default_retry_max_times,
+)
 from datahub.emitter.generic_emitter import Emitter
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.emitter.request_helper import OpenApiRequest, make_curl_command
@@ -60,6 +67,10 @@ from datahub.metadata.com.linkedin.pegasus2avro.mxe import (
     MetadataChangeProposal,
 )
 from datahub.metadata.com.linkedin.pegasus2avro.usage import UsageAggregation
+from datahub.metadata.schema_classes import (
+    KEY_ASPECT_NAMES,
+    ChangeTypeClass,
+)
 from datahub.utilities.server_config_util import RestServiceConfig, ServiceFeature
 
 if TYPE_CHECKING:
@@ -77,11 +88,9 @@ _DEFAULT_RETRY_STATUS_CODES = [  # Additional status codes to retry on
     504,
 ]
 _DEFAULT_RETRY_METHODS = ["HEAD", "GET", "POST", "PUT", "DELETE", "OPTIONS", "TRACE"]
-_DEFAULT_RETRY_MAX_TIMES = int(
-    os.getenv("DATAHUB_REST_EMITTER_DEFAULT_RETRY_MAX_TIMES", "4")
-)
+_DEFAULT_RETRY_MAX_TIMES = int(get_rest_emitter_default_retry_max_times())
 
-_DATAHUB_EMITTER_TRACE = get_boolean_env_variable("DATAHUB_EMITTER_TRACE", False)
+_DATAHUB_EMITTER_TRACE = get_emitter_trace()
 
 _DEFAULT_CLIENT_MODE: ClientMode = ClientMode.SDK
 
@@ -90,18 +99,32 @@ TRACE_INITIAL_BACKOFF = 1.0  # Start with 1 second
 TRACE_MAX_BACKOFF = 300.0  # Cap at 5 minutes
 TRACE_BACKOFF_FACTOR = 2.0  # Double the wait time each attempt
 
-# The limit is 16mb. We will use a max of 15mb to have some space
+# The limit is 16,000,000 bytes. We will use a max of 15mb to have some space
 # for overhead like request headers.
 # This applies to pretty much all calls to GMS.
-INGEST_MAX_PAYLOAD_BYTES = 15 * 1024 * 1024
+INGEST_MAX_PAYLOAD_BYTES = get_rest_emitter_batch_max_payload_bytes()
 
 # This limit is somewhat arbitrary. All GMS endpoints will timeout
 # and return a 500 if processing takes too long. To avoid sending
 # too much to the backend and hitting a timeout, we try to limit
 # the number of MCPs we send in a batch.
-BATCH_INGEST_MAX_PAYLOAD_LENGTH = int(
-    os.getenv("DATAHUB_REST_EMITTER_BATCH_MAX_PAYLOAD_LENGTH", 200)
-)
+BATCH_INGEST_MAX_PAYLOAD_LENGTH = get_rest_emitter_batch_max_payload_length()
+
+
+def preserve_unicode_escapes(obj: Any) -> Any:
+    """Recursively convert unicode characters back to escape sequences"""
+    if isinstance(obj, dict):
+        return {k: preserve_unicode_escapes(v) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [preserve_unicode_escapes(item) for item in obj]
+    elif isinstance(obj, str):
+        # Convert non-ASCII characters back to \u escapes
+        def escape_unicode(match: Any) -> Any:
+            return f"\\u{ord(match.group(0)):04x}"
+
+        return re.sub(r"[^\x00-\x7F]", escape_unicode, obj)
+    else:
+        return obj
 
 
 class EmitMode(ConfigEnum):
@@ -124,7 +147,7 @@ class EmitMode(ConfigEnum):
 
 _DEFAULT_EMIT_MODE = pydantic.parse_obj_as(
     EmitMode,
-    os.getenv("DATAHUB_EMIT_MODE", EmitMode.SYNC_PRIMARY),
+    get_emit_mode() or EmitMode.SYNC_PRIMARY,
 )
 
 
@@ -135,7 +158,7 @@ class RestSinkEndpoint(ConfigEnum):
 
 DEFAULT_REST_EMITTER_ENDPOINT = pydantic.parse_obj_as(
     RestSinkEndpoint,
-    os.getenv("DATAHUB_REST_EMITTER_DEFAULT_ENDPOINT", RestSinkEndpoint.RESTLI),
+    get_rest_emitter_default_endpoint() or RestSinkEndpoint.RESTLI,
 )
 
 
@@ -314,6 +337,7 @@ class DataHubRestEmitter(Closeable, Emitter):
         openapi_ingestion: Optional[bool] = None,
         client_mode: Optional[ClientMode] = None,
         datahub_component: Optional[str] = None,
+        server_config_refresh_interval: Optional[int] = None,
     ):
         if not gms_server:
             raise ConfigurationError("gms server is required")
@@ -329,6 +353,8 @@ class DataHubRestEmitter(Closeable, Emitter):
         self._openapi_ingestion = (
             openapi_ingestion  # Re-evaluated after test connection
         )
+        self._server_config_refresh_interval = server_config_refresh_interval
+        self._config_fetch_time: Optional[float] = None
 
         headers = {
             "X-RestLi-Protocol-Version": "2.0.0",
@@ -398,7 +424,17 @@ class DataHubRestEmitter(Closeable, Emitter):
         Raises:
             ConfigurationError: If there's an error fetching or validating the configuration
         """
-        if not hasattr(self, "_server_config") or not self._server_config:
+
+        if (
+            not hasattr(self, "_server_config")
+            or self._server_config is None
+            or (
+                self._server_config_refresh_interval is not None
+                and self._config_fetch_time is not None
+                and (time.time() - self._config_fetch_time)
+                > self._server_config_refresh_interval
+            )
+        ):
             if self._session is None or self._gms_server is None:
                 raise ConfigurationError(
                     "Session and URL are required to load configuration"
@@ -419,6 +455,7 @@ class DataHubRestEmitter(Closeable, Emitter):
                     )
 
                 self._server_config = RestServiceConfig(raw_config=raw_config)
+                self._config_fetch_time = time.time()
                 self._post_fetch_server_config()
 
             else:
@@ -441,7 +478,7 @@ class DataHubRestEmitter(Closeable, Emitter):
         if self._openapi_ingestion is None:
             # No constructor parameter
             if (
-                not os.getenv("DATAHUB_REST_EMITTER_DEFAULT_ENDPOINT")
+                not get_rest_emitter_default_endpoint()
                 and self._session_config.client_mode == ClientMode.SDK
                 and self._server_config.supports_feature(ServiceFeature.OPEN_API_SDK)
             ):
@@ -453,6 +490,8 @@ class DataHubRestEmitter(Closeable, Emitter):
                     DEFAULT_REST_EMITTER_ENDPOINT == RestSinkEndpoint.OPENAPI
                 )
 
+    def test_connection(self) -> None:
+        self.fetch_server_config()
         logger.debug(
             f"Using {'OpenAPI' if self._openapi_ingestion else 'Restli'} for ingestion."
         )
@@ -460,12 +499,21 @@ class DataHubRestEmitter(Closeable, Emitter):
             f"{EmitMode.ASYNC_WAIT} {'IS' if self._should_trace(emit_mode=EmitMode.ASYNC_WAIT, warn=False) else 'IS NOT'} supported."
         )
 
-    def test_connection(self) -> None:
-        self.fetch_server_config()
-
     def get_server_config(self) -> dict:
         return self.server_config.raw_config
 
+    def invalidate_config_cache(self) -> None:
+        """Manually invalidate the configuration cache."""
+        if (
+            hasattr(self, "_server_config")
+            and self._server_config is not None
+            and self._server_config_refresh_interval is not None
+        ):
+            # Set fetch time to beyond TTL in the past to force refresh on next access
+            self._config_fetch_time = (
+                time.time() - self._server_config_refresh_interval - 1
+            )
+
     def to_graph(self) -> "DataHubGraph":
         from datahub.ingestion.graph.client import DataHubGraph
 
@@ -538,6 +586,11 @@ class DataHubRestEmitter(Closeable, Emitter):
             "systemMetadata": system_metadata_obj,
         }
         payload = json.dumps(snapshot)
+        if len(payload) > INGEST_MAX_PAYLOAD_BYTES:
+            logger.warning(
+                f"MCE object has size {len(payload)} that exceeds the max payload size of {INGEST_MAX_PAYLOAD_BYTES}, "
+                "so this metadata will likely fail to be emitted."
+            )
 
         self._emit_generic(url, payload)
 
@@ -584,15 +637,27 @@ class DataHubRestEmitter(Closeable, Emitter):
                     trace_data = extract_trace_data(response) if response else None
 
         else:
-            url = f"{self._gms_server}/aspects?action=ingestProposal"
+            if mcp.changeType == ChangeTypeClass.DELETE:
+                if mcp.aspectName not in KEY_ASPECT_NAMES:
+                    raise OperationalError(
+                        f"Delete not supported for non key aspect: {mcp.aspectName} for urn: "
+                        f"{mcp.entityUrn}"
+                    )
 
-            mcp_obj = pre_json_transform(mcp.to_obj())
-            payload_dict = {
-                "proposal": mcp_obj,
-                "async": "true"
-                if emit_mode in (EmitMode.ASYNC, EmitMode.ASYNC_WAIT)
-                else "false",
-            }
+                url = f"{self._gms_server}/entities?action=delete"
+                payload_dict = {
+                    "urn": mcp.entityUrn,
+                }
+            else:
+                url = f"{self._gms_server}/aspects?action=ingestProposal"
+
+                mcp_obj = preserve_unicode_escapes(pre_json_transform(mcp.to_obj()))
+                payload_dict = {
+                    "proposal": mcp_obj,
+                    "async": "true"
+                    if emit_mode in (EmitMode.ASYNC, EmitMode.ASYNC_WAIT)
+                    else "false",
+                }
 
             payload = json.dumps(payload_dict)
 
@@ -704,16 +769,24 @@ class DataHubRestEmitter(Closeable, Emitter):
         url = f"{self._gms_server}/aspects?action=ingestProposalBatch"
 
         mcp_objs = [pre_json_transform(mcp.to_obj()) for mcp in mcps]
+        if len(mcp_objs) == 0:
+            return 0
 
         # As a safety mechanism, we need to make sure we don't exceed the max payload size for GMS.
         # If we will exceed the limit, we need to break it up into chunks.
-        mcp_obj_chunks: List[List[str]] = []
-        current_chunk_size = INGEST_MAX_PAYLOAD_BYTES
+        mcp_obj_chunks: List[List[str]] = [[]]
+        current_chunk_size = 0
         for mcp_obj in mcp_objs:
+            mcp_identifier = f"{mcp_obj.get('entityUrn')}-{mcp_obj.get('aspectName')}"
             mcp_obj_size = len(json.dumps(mcp_obj))
             if _DATAHUB_EMITTER_TRACE:
                 logger.debug(
-                    f"Iterating through object with size {mcp_obj_size} (type: {mcp_obj.get('aspectName')}"
+                    f"Iterating through object ({mcp_identifier}) with size {mcp_obj_size}"
+                )
+            if mcp_obj_size > INGEST_MAX_PAYLOAD_BYTES:
+                logger.warning(
+                    f"MCP object {mcp_identifier} has size {mcp_obj_size} that exceeds the max payload size of {INGEST_MAX_PAYLOAD_BYTES}, "
+                    "so this metadata will likely fail to be emitted."
                 )
 
             if (
@@ -726,7 +799,7 @@ class DataHubRestEmitter(Closeable, Emitter):
                 current_chunk_size = 0
             mcp_obj_chunks[-1].append(mcp_obj)
             current_chunk_size += mcp_obj_size
-        if len(mcp_obj_chunks) > 0:
+        if len(mcp_obj_chunks) > 1 or _DATAHUB_EMITTER_TRACE:
             logger.debug(
                 f"Decided to send {len(mcps)} MCP batch in {len(mcp_obj_chunks)} chunks"
             )

datahub/entrypoints.py

@@ -10,6 +10,7 @@ import click
 import datahub._version as datahub_version
 from datahub.cli.check_cli import check
 from datahub.cli.cli_utils import (
+    enable_auto_decorators,
     fixup_gms_url,
     generate_access_token,
     make_shim_command,
@@ -21,6 +22,7 @@ from datahub.cli.docker_cli import docker
 from datahub.cli.env_utils import get_boolean_env_variable
 from datahub.cli.exists_cli import exists
 from datahub.cli.get_cli import get
+from datahub.cli.graphql_cli import graphql
 from datahub.cli.ingest_cli import ingest
 from datahub.cli.migrate import migrate
 from datahub.cli.put_cli import put
@@ -38,7 +40,6 @@ from datahub.cli.timeline_cli import timeline
 from datahub.configuration.common import should_show_stack_trace
 from datahub.ingestion.graph.client import get_default_graph
 from datahub.ingestion.graph.config import ClientMode
-from datahub.telemetry import telemetry
 from datahub.utilities._custom_package_loader import model_version_name
 from datahub.utilities.logging_manager import configure_logging
 from datahub.utilities.server_config_util import get_gms_config
@@ -111,7 +112,6 @@ def datahub(
     default=False,
     help="If passed will show server config. Assumes datahub init has happened.",
 )
-@telemetry.with_telemetry()
 def version(include_server: bool = False) -> None:
     """Print version number and exit."""
 
@@ -131,7 +131,6 @@ def version(include_server: bool = False) -> None:
     default=False,
     help="If passed then uses password to initialise token.",
 )
-@telemetry.with_telemetry()
 def init(use_password: bool = False) -> None:
     """Configure which datahub instance to connect to"""
 
@@ -171,6 +170,7 @@ datahub.add_command(ingest)
 datahub.add_command(delete)
 datahub.add_command(exists)
 datahub.add_command(get)
+datahub.add_command(graphql)
 datahub.add_command(put)
 datahub.add_command(state)
 datahub.add_command(telemetry_cli)
@@ -218,6 +218,9 @@ except ImportError as e:
         make_shim_command("actions", "run `pip install acryl-datahub-actions`")
     )
 
+# Adding telemetry and upgrade decorators to all commands
+enable_auto_decorators(datahub)
+
 
 def main(**kwargs):
     # We use threads in a variety of places within our CLI. The multiprocessing