acryl-datahub 1.0.0.3rc9__py3-none-any.whl → 1.0.0.3rc11__py3-none-any.whl

datahub/metadata/schema_classes.py

@@ -20,16 +20,16 @@ from datahub.utilities._custom_package_loader import get_custom_models_package
 _custom_package_path = get_custom_models_package()
 
 if TYPE_CHECKING or not _custom_package_path:
-    from ._schema_classes import *
+    from ._internal_schema_classes import *
 
     # Required explicitly because __all__ doesn't include _ prefixed names.
-    from ._schema_classes import __SCHEMA_TYPES
+    from ._internal_schema_classes import __SCHEMA_TYPES
 
     if IS_SPHINX_BUILD:
         # Set __module__ to the current module so that Sphinx will document the
         # classes as belonging to this module instead of the custom package.
         for _cls in list(globals().values()):
-            if hasattr(_cls, "__module__") and "datahub.metadata._schema_classes" in _cls.__module__:
+            if hasattr(_cls, "__module__") and "datahub.metadata._internal_schema_classes" in _cls.__module__:
                 _cls.__module__ = __name__
 else:
     _custom_package = importlib.import_module(_custom_package_path)

datahub/metadata/schemas/__init__.py

@@ -15,10 +15,10 @@ import pathlib
 def _load_schema(schema_name: str) -> str:
     return (pathlib.Path(__file__).parent / f"{schema_name}.avsc").read_text()
 
-def getMetadataChangeProposalSchema() -> str:
-    return _load_schema("MetadataChangeProposal")
-
 def getMetadataChangeEventSchema() -> str:
     return _load_schema("MetadataChangeEvent")
 
+def getMetadataChangeProposalSchema() -> str:
+    return _load_schema("MetadataChangeProposal")
+
 # fmt: on

datahub/sdk/main_client.py

@@ -4,7 +4,7 @@ from typing import Optional, overload
 
 from datahub.errors import SdkUsageError
 from datahub.ingestion.graph.client import DataHubGraph, get_default_graph
-from datahub.ingestion.graph.config import DatahubClientConfig
+from datahub.ingestion.graph.config import ClientMode, DatahubClientConfig
 from datahub.sdk.entity_client import EntityClient
 from datahub.sdk.lineage_client import LineageClient
 from datahub.sdk.resolver_client import ResolverClient
@@ -84,7 +84,7 @@ class DataHubClient:
         # Inspired by the DockerClient.from_env() method.
         # TODO: This one also reads from ~/.datahubenv, so the "from_env" name might be a bit confusing.
         # That file is part of the "environment", but is not a traditional "env variable".
-        graph = get_default_graph()
+        graph = get_default_graph(ClientMode.SDK)
 
         return cls(graph=graph)
 

datahub/secret/datahub_secret_store.py

@@ -3,7 +3,8 @@ from typing import Any, Dict, List, Optional, Union
 
 from pydantic import BaseModel, validator
 
-from datahub.ingestion.graph.client import DatahubClientConfig, DataHubGraph
+from datahub.ingestion.graph.client import DataHubGraph
+from datahub.ingestion.graph.config import DatahubClientConfig
 from datahub.secret.datahub_secrets_client import DataHubSecretsClient
 from datahub.secret.secret_store import SecretStore
 

datahub/telemetry/telemetry.py

@@ -352,10 +352,10 @@ class Telemetry:
             }
         else:
             return {
-                "server_type": server.server_config.get("datahub", {}).get(
+                "server_type": server.server_config.raw_config.get("datahub", {}).get(
                     "serverType", "missing"
                 ),
-                "server_version": server.server_config.get("versions", {})
+                "server_version": server.server_config.raw_config.get("versions", {})
                 .get("acryldata/datahub", {})
                 .get("version", "missing"),
                 "server_id": server.server_id or "missing",

datahub/testing/check_imports.py

@@ -9,7 +9,7 @@ def ensure_no_indirect_model_imports(dirs: List[pathlib.Path]) -> None:
     # If our needs become more complex, we should move to a proper linter.
     denied_imports = {
         "src.": "datahub.*",
-        "datahub.metadata._schema_classes": "datahub.metadata.schema_classes",
+        "datahub.metadata._internal_schema_classes": "datahub.metadata.schema_classes",
         "datahub.metadata._urns": "datahub.metadata.urns",
     }
     ignored_files = {

datahub/upgrade/upgrade.py

@@ -13,7 +13,9 @@ from pydantic import BaseModel
 from datahub._version import __version__
 from datahub.cli.config_utils import load_client_config
 from datahub.ingestion.graph.client import DataHubGraph
+from datahub.ingestion.graph.config import ClientMode
 from datahub.utilities.perf_timer import PerfTimer
+from datahub.utilities.server_config_util import RestServiceConfig
 
 log = logging.getLogger(__name__)
 
@@ -109,7 +111,7 @@ async def get_github_stats():
             return (latest_server_version, latest_server_date)
 
 
-async def get_server_config(gms_url: str, token: Optional[str]) -> dict:
+async def get_server_config(gms_url: str, token: Optional[str]) -> RestServiceConfig:
     import aiohttp
 
     headers = {
@@ -124,7 +126,7 @@ async def get_server_config(gms_url: str, token: Optional[str]) -> dict:
         config_endpoint = f"{gms_url}/config"
         async with session.get(config_endpoint, headers=headers) as dh_response:
             dh_response_json = await dh_response.json()
-            return dh_response_json
+            return RestServiceConfig(raw_config=dh_response_json)
 
 
 async def get_server_version_stats(
@@ -132,11 +134,12 @@ async def get_server_version_stats(
 ) -> Tuple[Optional[str], Optional[Version], Optional[datetime]]:
     import aiohttp
 
-    server_config = None
+    server_config: Optional[RestServiceConfig] = None
     if not server:
         try:
             # let's get the server from the cli config
             client_config = load_client_config()
+            client_config.client_mode = ClientMode.CLI
             host = client_config.server
             token = client_config.token
             server_config = await get_server_config(host, token)
@@ -150,15 +153,10 @@ async def get_server_version_stats(
     server_version: Optional[Version] = None
     current_server_release_date = None
     if server_config:
-        server_version_string = (
-            server_config.get("versions", {})
-            .get("acryldata/datahub", {})
-            .get("version")
-        )
-        commit_hash = (
-            server_config.get("versions", {}).get("acryldata/datahub", {}).get("commit")
-        )
-        server_type = server_config.get("datahub", {}).get("serverType", "unknown")
+        server_version_string = server_config.service_version
+        commit_hash = server_config.commit_hash
+        server_type = server_config.server_type
+
         if server_type == "quickstart" and commit_hash:
             async with aiohttp.ClientSession(
                 headers={"Accept": "application/vnd.github.v3+json"}

datahub/utilities/logging_manager.py

@@ -161,6 +161,7 @@ class _LogBuffer:
         self._buffer: Deque[str] = collections.deque(maxlen=maxlen)
 
     def write(self, line: str) -> None:
+        # We do not expect `line` to have a trailing newline.
         if len(line) > IN_MEMORY_LOG_BUFFER_MAX_LINE_LENGTH:
             line = line[:IN_MEMORY_LOG_BUFFER_MAX_LINE_LENGTH] + "[truncated]"
 
@@ -188,7 +189,13 @@ class _BufferLogHandler(logging.Handler):
             message = self.format(record)
         except TypeError as e:
             message = f"Error formatting log message: {e}\nMessage: {record.msg}, Args: {record.args}"
-        self._storage.write(message)
+
+        # For exception stack traces, the message is split over multiple lines,
+        # but we store it as a single string. Because we truncate based on line
+        # length, it's better for us to split it into multiple lines so that we
+        # don't lose any information on deeper stack traces.
+        for line in message.split("\n"):
+            self._storage.write(line)
 
 
 def _remove_all_handlers(logger: logging.Logger) -> None:

datahub/utilities/server_config_util.py

@@ -1,23 +1,391 @@
-from typing import Any, Dict, Optional
+import logging
+import re
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import (
+    Any,
+    Dict,
+    Optional,
+    Tuple,
+    Union,
+)
 
+import requests
+
+from datahub.configuration.common import (
+    ConfigurationError,
+)
 from datahub.telemetry.telemetry import suppress_telemetry
 
+logger = logging.getLogger(__name__)
+
 # Only to be written to for logging server related information
 global_debug: Dict[str, Any] = {}
 
 
-def set_gms_config(config: Dict) -> Any:
+def get_gms_config() -> Dict:
+    return global_debug.get("gms_config", {})
+
+
+class ServiceFeature(Enum):
+    """
+    Enum representing supported features in the REST service.
+    """
+
+    OPEN_API_SDK = "openapi_sdk"
+    API_TRACING = "api_tracing"
+    NO_CODE = "no_code"
+    STATEFUL_INGESTION = "stateful_ingestion"
+    IMPACT_ANALYSIS = "impact_analysis"
+    PATCH_CAPABLE = "patch_capable"
+    CLI_TELEMETRY = "cli_telemetry"
+    DATAHUB_CLOUD = "datahub_cloud"
+    # Add more features as needed
+
+
+_REQUIRED_VERSION_OPENAPI_TRACING = {
+    "acryl": (
+        0,
+        3,
+        11,
+        0,
+    ),  # Requires v0.3.11.0 or higher for acryl versions
+    "cloud": (0, 3, 11, 0),  # Special case for '-cloud' suffix
+    "any_suffix": (0, 3, 11, 0),  # Generic requirement for any other suffix
+    "none": (1, 0, 1, 0),  # Requirement for versions without suffix
+}
+
+
+@dataclass
+class RestServiceConfig:
+    """
+    A class to represent REST service configuration with semantic version parsing capabilities.
+    """
+
+    session: Optional[requests.Session] = None
+    url: Optional[str] = None
+    raw_config: Dict[str, Any] = field(default_factory=dict)
+    _version_cache: Optional[Tuple[int, int, int, int]] = None
+
+    def fetch_config(self) -> Dict[str, Any]:
+        """
+        Fetch configuration from the server if not already loaded.
+
+        Returns:
+            The configuration dictionary
+
+        Raises:
+            ConfigurationError: If there's an error fetching or validating the configuration
+        """
+        if not self.raw_config:
+            if self.session is None or self.url is None:
+                raise ConfigurationError(
+                    "Session and URL are required to load configuration"
+                )
+
+            response = self.session.get(self.url)
+
+            if response.status_code == 200:
+                config = response.json()
+
+                # Validate that we're connected to the correct service
+                if config.get("noCode") == "true":
+                    self.raw_config = config
+                else:
+                    raise ConfigurationError(
+                        "You seem to have connected to the frontend service instead of the GMS endpoint. "
+                        "The rest emitter should connect to DataHub GMS (usually <datahub-gms-host>:8080) or Frontend GMS API (usually <frontend>:9002/api/gms). "
+                        "For Acryl users, the endpoint should be https://<name>.acryl.io/gms"
+                    )
+            else:
+                logger.debug(
+                    f"Unable to connect to {self.url} with status_code: {response.status_code}. Response: {response.text}"
+                )
+
+                if response.status_code == 401:
+                    message = f"Unable to connect to {self.url} - got an authentication error: {response.text}."
+                else:
+                    message = f"Unable to connect to {self.url} with status_code: {response.status_code}."
+
+                message += "\nPlease check your configuration and make sure you are talking to the DataHub GMS (usually <datahub-gms-host>:8080) or Frontend GMS API (usually <frontend>:9002/api/gms)."
+                raise ConfigurationError(message)
+
+        return self.raw_config
+
+    @property
+    def config(self) -> Dict[str, Any]:
+        """
+        Get the full configuration dictionary, loading it if necessary.
+
+        Returns:
+            The configuration dictionary
+        """
+        return self.fetch_config()
+
+    @property
+    def commit_hash(self) -> Optional[str]:
+        """
+        Get the commit hash for the current version.
+
+        Returns:
+            The commit hash or None if not found
+        """
+        versions = self.config.get("versions") or {}
+        datahub_info = versions.get("acryldata/datahub") or {}
+        return datahub_info.get("commit")
+
+    @property
+    def server_type(self) -> str:
+        """
+        Get the server type.
+
+        Returns:
+            The server type or "unknown" if not found
+        """
+        datahub = self.config.get("datahub") or {}
+        return datahub.get("serverType", "unknown")
+
+    @property
+    def service_version(self) -> Optional[str]:
+        """
+        Get the raw service version string.
+
+        Returns:
+            The version string or None if not found
+        """
+        config = self.fetch_config()
+        versions = config.get("versions") or {}
+        datahub_info = versions.get("acryldata/datahub") or {}
+        return datahub_info.get("version")
+
+    def _parse_version(
+        self, version_str: Optional[str] = None
+    ) -> Tuple[int, int, int, int]:
+        """
+        Parse a semantic version string into its components, ignoring rc and suffixes.
+        Supports standard three-part versions (1.0.0) and four-part versions (1.0.0.1).
+
+        Args:
+            version_str: Version string to parse. If None, uses the service version.
+
+        Returns:
+            Tuple of (major, minor, patch, build) version numbers where build is 0 for three-part versions
+
+        Raises:
+            ValueError: If the version string cannot be parsed
+        """
+        if version_str is None:
+            version_str = self.service_version
+
+        if not version_str:
+            return (0, 0, 0, 0)
+
+        # Remove 'v' prefix if present
+        if version_str.startswith("v"):
+            version_str = version_str[1:]
+
+        # Extract the semantic version part (before any rc or suffix)
+        # This pattern will match both three-part (1.0.0) and four-part (1.0.0.1) versions
+        match = re.match(r"(\d+)\.(\d+)\.(\d+)(?:\.(\d+))?(?:rc\d+|-.*)?", version_str)
+        if not match:
+            raise ValueError(f"Invalid version format: {version_str}")
+
+        major = int(match.group(1))
+        minor = int(match.group(2))
+        patch = int(match.group(3))
+        build = (
+            int(match.group(4)) if match.group(4) else 0
+        )  # Default to 0 if not present
+
+        return (major, minor, patch, build)
+
+    @property
+    def parsed_version(self) -> Optional[Tuple[int, int, int, int]]:
+        """
+        Get the parsed semantic version of the service.
+        Uses caching for efficiency.
+
+        Returns:
+            Tuple of (major, minor, patch) version numbers
+        """
+        if self._version_cache is None:
+            self._version_cache = self._parse_version()
+        return self._version_cache
+
+    def is_version_at_least(
+        self, major: int, minor: int = 0, patch: int = 0, build: int = 0
+    ) -> bool:
+        """
+        Check if the service version is at least the specified version.
+
+        Args:
+            major: Major version to check against
+            minor: Minor version to check against
+            patch: Patch version to check against
+            build: Build version to check against (for four-part versions)
+
+        Returns:
+            True if the service version is at least the specified version
+        """
+        current_version = self.parsed_version or (0, 0, 0, 0)
+        requested_version = (major, minor, patch, build)
+
+        return current_version >= requested_version
+
+    @property
+    def is_no_code_enabled(self) -> bool:
+        """
+        Check if noCode is enabled.
+
+        Returns:
+            True if noCode is set to "true"
+        """
+        return self.config.get("noCode") == "true"
+
+    @property
+    def is_managed_ingestion_enabled(self) -> bool:
+        """
+        Check if managedIngestion is enabled.
+
+        Returns:
+            True if managedIngestion.enabled is True
+        """
+        managed_ingestion = self.config.get("managedIngestion") or {}
+        return managed_ingestion.get("enabled", False)
+
+    @property
+    def is_datahub_cloud(self) -> bool:
+        """
+        Check if DataHub Cloud is enabled.
+
+        Returns:
+            True if the server environment is not 'oss'
+        """
+        datahub_config = self.config.get("datahub") or {}
+        server_env = datahub_config.get("serverEnv")
+
+        # Return False if serverEnv is None or empty string
+        if not server_env:
+            return False
+
+        return server_env != "oss"
+
+    def supports_feature(self, feature: ServiceFeature) -> bool:
+        """
+        Determines whether a specific feature is supported based on service version.
+
+        Version categorization follows these rules:
+        1. Has '-acryl' suffix (highest priority)
+        2. Has a specific known suffix (e.g. '-other')
+        3. Has some other suffix (catchall for any suffix)
+        4. No suffix
+
+        Args:
+            feature: Feature enum value to check
+
+        Returns:
+            Boolean indicating whether the feature is supported
+        """
+        version = self.service_version
+        if not version:
+            return False
+
+        # Determine the suffix category
+        suffix_category = "none"  # Default: no suffix
+
+        if "-" in version:
+            suffix = version.split("-", 1)[1]
+
+            if suffix == "acryl":
+                suffix_category = "acryl"
+            elif suffix == "cloud":  # Example of a specific override
+                suffix_category = "cloud"
+            else:
+                suffix_category = "any_suffix"  # Catchall for any other suffix
+
+        # Define feature requirements based on version scheme
+        # This can be expanded to include more features
+        feature_requirements = {
+            ServiceFeature.OPEN_API_SDK: _REQUIRED_VERSION_OPENAPI_TRACING,
+            ServiceFeature.API_TRACING: _REQUIRED_VERSION_OPENAPI_TRACING,
+            # Additional features can be defined here
+        }
+
+        # Special handling for features that rely on config flags instead of version
+        config_based_features = {
+            ServiceFeature.NO_CODE: lambda: self.is_no_code_enabled,
+            ServiceFeature.STATEFUL_INGESTION: lambda: self.config.get(
+                "statefulIngestionCapable", False
+            )
+            is True,
+            ServiceFeature.IMPACT_ANALYSIS: lambda: self.config.get(
+                "supportsImpactAnalysis", False
+            )
+            is True,
+            ServiceFeature.PATCH_CAPABLE: lambda: self.config.get("patchCapable", False)
+            is True,
+            ServiceFeature.CLI_TELEMETRY: lambda: (
+                self.config.get("telemetry") or {}
+            ).get("enabledCli", None),
+            # Add more config-based feature checks as needed
+        }
+
+        # Check if this is a config-based feature
+        if feature in config_based_features:
+            return config_based_features[feature]()
+
+        # Check if the feature exists in our requirements dictionary
+        if feature not in feature_requirements:
+            # Unknown feature, assume not supported
+            return False
+
+        # Get version requirements for this feature and version category
+        feature_reqs = feature_requirements[feature]
+        requirements = feature_reqs.get(suffix_category)
+
+        if not requirements:
+            # Fallback to the no-suffix requirements if specific requirements aren't defined
+            requirements = feature_reqs.get(
+                "none", (99, 99, 99, 99)
+            )  # Very high version if none defined
+
+        # Check if the current version meets the requirements
+        req_major, req_minor, req_patch, req_build = requirements
+        return self.is_version_at_least(req_major, req_minor, req_patch, req_build)
+
+    def __str__(self) -> str:
+        """
+        Return a string representation of the configuration as JSON.
+
+        Returns:
+            A string representation of the configuration dictionary
+        """
+        return str(self.config)
+
+    def __repr__(self) -> str:
+        """
+        Return a representation of the object that can be used to recreate it.
+
+        Returns:
+            A string representation that can be used with pprint
+        """
+        return str(self.config)
+
+
+def set_gms_config(config: Union[Dict[str, Any], RestServiceConfig]) -> None:
     global_debug["gms_config"] = config
 
-    cli_telemtry_enabled = is_cli_telemetry_enabled()
-    if cli_telemtry_enabled is not None and not cli_telemtry_enabled:
+    config_obj = (
+        config
+        if isinstance(config, RestServiceConfig)
+        else RestServiceConfig(raw_config=config)
+    )
+
+    cli_telemetry_enabled = is_cli_telemetry_enabled(config_obj)
+    if cli_telemetry_enabled is not None and not cli_telemetry_enabled:
         # server requires telemetry to be disabled on client
         suppress_telemetry()
 
 
-def get_gms_config() -> Dict:
-    return global_debug.get("gms_config", {})
-
-
-def is_cli_telemetry_enabled() -> Optional[bool]:
-    return get_gms_config().get("telemetry", {}).get("enabledCli", None)
+def is_cli_telemetry_enabled(config: RestServiceConfig) -> bool:
+    return config.supports_feature(ServiceFeature.CLI_TELEMETRY)

datahub/utilities/sqlalchemy_query_combiner.py

@@ -272,11 +272,10 @@ class SQLAlchemyQueryCombiner:
                     self.report.uncombined_queries_issued += 1
                     return _sa_execute_underlying_method(conn, query, *args, **kwargs)
 
-        with _sa_execute_method_patching_lock:
-            with unittest.mock.patch(
-                "sqlalchemy.engine.Connection.execute", _sa_execute_fake
-            ):
-                yield self
+        with _sa_execute_method_patching_lock, unittest.mock.patch(
+            "sqlalchemy.engine.Connection.execute", _sa_execute_fake
+        ):
+            yield self
 
     def run(self, method: Callable[[], None]) -> None:
         """

datahub/utilities/urn_encoder.py

@@ -4,7 +4,7 @@ from typing import List
 # NOTE: Frontend relies on encoding these three characters. Specifically, we decode and encode schema fields for column level lineage.
 # If this changes, make appropriate changes to datahub-web-react/src/app/lineage/utils/columnLineageUtils.ts
 # We also rely on encoding these exact three characters when generating schemaField urns in our graphQL layer. Update SchemaFieldUtils if this changes.
-# Also see https://datahubproject.io/docs/what/urn/#restrictions
+# Also see https://docs.datahub.com/docs/what/urn/#restrictions
 RESERVED_CHARS = {",", "(", ")", "␟"}
 RESERVED_CHARS_EXTENDED = RESERVED_CHARS.union({"%"})