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

datahub/ingestion/sink/datahub_rest.py

@@ -34,7 +34,7 @@ from datahub.ingestion.api.sink import (
     WriteCallback,
 )
 from datahub.ingestion.api.workunit import MetadataWorkUnit
-from datahub.ingestion.graph.client import DatahubClientConfig
+from datahub.ingestion.graph.config import ClientMode, DatahubClientConfig
 from datahub.metadata.com.linkedin.pegasus2avro.mxe import (
     MetadataChangeEvent,
     MetadataChangeProposal,
@@ -140,11 +140,7 @@ class DatahubRestSink(Sink[DatahubRestSinkConfig, DataHubRestSinkReport]):
                 f"💥 Failed to connect to DataHub with {repr(self.emitter)}"
             ) from exc
 
-        self.report.gms_version = (
-            gms_config.get("versions", {})
-            .get("acryldata/datahub", {})
-            .get("version", None)
-        )
+        self.report.gms_version = gms_config.service_version
         self.report.mode = self.config.mode
         self.report.max_threads = self.config.max_threads
         logger.debug("Setting env variables to override config")
@@ -180,6 +176,8 @@ class DatahubRestSink(Sink[DatahubRestSinkConfig, DataHubRestSinkReport]):
             disable_ssl_verification=config.disable_ssl_verification,
             openapi_ingestion=config.endpoint == RestSinkEndpoint.OPENAPI,
             default_trace_mode=config.default_trace_mode == RestTraceMode.ENABLED,
+            client_mode=config.client_mode,
+            datahub_component=config.datahub_component,
         )
 
     @property
@@ -190,6 +188,7 @@ class DatahubRestSink(Sink[DatahubRestSinkConfig, DataHubRestSinkReport]):
         # https://github.com/psf/requests/issues/1871#issuecomment-32751346
         thread_local = self._emitter_thread_local
         if not hasattr(thread_local, "emitter"):
+            self.config.client_mode = ClientMode.INGESTION
             thread_local.emitter = DatahubRestSink._make_emitter(self.config)
         return thread_local.emitter
 

datahub/ingestion/source/apply/datahub_apply.py

@@ -18,6 +18,7 @@ from datahub.ingestion.api.source import MetadataWorkUnitProcessor, Source, Sour
 from datahub.ingestion.api.source_helpers import auto_workunit_reporter
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.graph.client import DataHubGraph, get_default_graph
+from datahub.ingestion.graph.config import ClientMode
 from datahub.metadata.schema_classes import (
     DomainsClass,
     GlossaryTermAssociationClass,
@@ -48,7 +49,7 @@ def apply_association_to_container(
     """
     urns: List[str] = [container_urn]
     if not graph:
-        graph = get_default_graph()
+        graph = get_default_graph(ClientMode.INGESTION)
     logger.info(f"Using {graph}")
     urns.extend(
         graph.get_urns_by_filter(

datahub/ingestion/source/ge_data_profiler.py

@@ -51,6 +51,7 @@ from typing_extensions import Concatenate, ParamSpec
 from datahub.emitter import mce_builder
 from datahub.emitter.mce_builder import get_sys_time
 from datahub.ingestion.graph.client import get_default_graph
+from datahub.ingestion.graph.config import ClientMode
 from datahub.ingestion.source.ge_profiling_config import GEProfilingConfig
 from datahub.ingestion.source.profiling.common import (
     Cardinality,
@@ -1569,7 +1570,7 @@ def _get_columns_to_ignore_sampling(
         name=dataset_name, platform=platform, env=env
     )
 
-    datahub_graph = get_default_graph()
+    datahub_graph = get_default_graph(ClientMode.INGESTION)
 
     dataset_tags = datahub_graph.get_tags(dataset_urn)
     if dataset_tags:

datahub/ingestion/source/metadata/lineage.py

@@ -36,6 +36,7 @@ from datahub.ingestion.api.source_helpers import (
 )
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.graph.client import get_default_graph
+from datahub.ingestion.graph.config import ClientMode
 from datahub.metadata.com.linkedin.pegasus2avro.dataset import (
     FineGrainedLineageDownstreamType,
     FineGrainedLineageUpstreamType,
@@ -210,7 +211,7 @@ def _get_lineage_mcp(
 
     # extract the old lineage and save it for the new mcp
     if preserve_upstream:
-        client = get_default_graph()
+        client = get_default_graph(ClientMode.INGESTION)
 
         old_upstream_lineage = get_aspects_for_entity(
             client._session,

datahub/ingestion/source/state_provider/datahub_ingestion_checkpointing_provider.py

@@ -10,7 +10,8 @@ from datahub.ingestion.api.ingestion_job_checkpointing_provider_base import (
     IngestionCheckpointingProviderConfig,
     JobId,
 )
-from datahub.ingestion.graph.client import DatahubClientConfig, DataHubGraph
+from datahub.ingestion.graph.client import DataHubGraph
+from datahub.ingestion.graph.config import DatahubClientConfig
 from datahub.metadata.schema_classes import DatahubIngestionCheckpointClass
 
 logger = logging.getLogger(__name__)

datahub/integrations/assertion/common.py

@@ -3,6 +3,7 @@ from typing import List, Optional, Tuple, TypedDict
 
 from datahub.api.entities.assertion.assertion import BaseEntityAssertion
 from datahub.ingestion.graph.client import get_default_graph
+from datahub.ingestion.graph.config import ClientMode
 from datahub.metadata.com.linkedin.pegasus2avro.dataset import DatasetProperties
 from datahub.metadata.com.linkedin.pegasus2avro.schema import SchemaMetadata
 from datahub.utilities.urns.urn import Urn
@@ -15,7 +16,7 @@ class ColumnDict(TypedDict):
 
 @lru_cache
 def get_qualified_name_from_datahub(urn: str) -> Optional[str]:
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         props: Optional[DatasetProperties] = graph.get_aspect(urn, DatasetProperties)
         if props is not None:
             return props.qualifiedName
@@ -24,7 +25,7 @@ def get_qualified_name_from_datahub(urn: str) -> Optional[str]:
 
 @lru_cache
 def get_schema_from_datahub(urn: str) -> Optional[List[ColumnDict]]:
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.INGESTION) as graph:
         schema: Optional[SchemaMetadata] = graph.get_aspect(urn, SchemaMetadata)
         if schema is not None:
             return [

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/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/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)