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

datahub/upgrade/upgrade.py

@@ -7,7 +7,7 @@ from typing import Any, Callable, Optional, Tuple, TypeVar
 
 import click
 import humanfriendly
-from packaging.version import Version
+from packaging.version import InvalidVersion, Version
 from pydantic import BaseModel
 
 from datahub._version import __version__
@@ -28,10 +28,24 @@ class VersionStats(BaseModel, arbitrary_types_allowed=True):
     release_date: Optional[datetime] = None
 
 
+def _safe_version_stats(version_string: str) -> Optional[VersionStats]:
+    """
+    Safely create a VersionStats object from a version string.
+    Returns None if the version string is invalid.
+    """
+    try:
+        return VersionStats(version=Version(version_string), release_date=None)
+    except InvalidVersion:
+        log.warning(f"Invalid version format received: {version_string!r}")
+        return None
+
+
 class ServerVersionStats(BaseModel):
     current: VersionStats
     latest: Optional[VersionStats] = None
     current_server_type: Optional[str] = None
+    current_server_default_cli_version: Optional[VersionStats] = None
+    is_cloud_server: Optional[bool] = None
 
 
 class ClientVersionStats(BaseModel):
@@ -44,7 +58,7 @@ class DataHubVersionStats(BaseModel):
     client: ClientVersionStats
 
 
-async def get_client_version_stats():
+async def get_client_version_stats() -> ClientVersionStats:
     import aiohttp
 
     current_version_string = __version__
@@ -52,6 +66,7 @@ async def get_client_version_stats():
     client_version_stats: ClientVersionStats = ClientVersionStats(
         current=VersionStats(version=current_version, release_date=None), latest=None
     )
+
     async with aiohttp.ClientSession() as session:
         pypi_url = "https://pypi.org/pypi/acryl_datahub/json"
         async with session.get(pypi_url) as resp:
@@ -131,7 +146,9 @@ async def get_server_config(gms_url: str, token: Optional[str]) -> RestServiceCo
 
 async def get_server_version_stats(
     server: Optional[DataHubGraph] = None,
-) -> Tuple[Optional[str], Optional[Version], Optional[datetime]]:
+) -> Tuple[
+    Optional[str], Optional[Version], Optional[str], Optional[datetime], Optional[bool]
+]:
     import aiohttp
 
     server_config: Optional[RestServiceConfig] = None
@@ -151,12 +168,15 @@ async def get_server_version_stats(
 
     server_type = None
     server_version: Optional[Version] = None
+    current_server_default_cli_version = None
     current_server_release_date = None
+    is_cloud_server: Optional[bool] = None
     if server_config:
         server_version_string = server_config.service_version
         commit_hash = server_config.commit_hash
         server_type = server_config.server_type
-
+        current_server_default_cli_version = server_config.default_cli_version
+        is_cloud_server = server_config.is_datahub_cloud
         if server_type == "quickstart" and commit_hash:
             async with aiohttp.ClientSession(
                 headers={"Accept": "application/vnd.github.v3+json"}
@@ -171,7 +191,13 @@ async def get_server_version_stats(
         if server_version_string and server_version_string.startswith("v"):
             server_version = Version(server_version_string[1:])
 
-    return (server_type, server_version, current_server_release_date)
+    return (
+        server_type,
+        server_version,
+        current_server_default_cli_version,
+        current_server_release_date,
+        is_cloud_server,
+    )
 
 
 def retrieve_version_stats(
@@ -214,7 +240,9 @@ async def _retrieve_version_stats(
     (
         current_server_type,
         current_server_version,
+        current_server_default_cli_version,
         current_server_release_date,
+        is_cloud_server,
     ) = results[2]
 
     server_version_stats = None
@@ -223,12 +251,18 @@ async def _retrieve_version_stats(
             current=VersionStats(
                 version=current_server_version, release_date=current_server_release_date
             ),
+            current_server_default_cli_version=(
+                _safe_version_stats(current_server_default_cli_version)
+                if current_server_default_cli_version
+                else None
+            ),
             latest=(
                 VersionStats(version=last_server_version, release_date=last_server_date)
                 if last_server_version
                 else None
             ),
             current_server_type=current_server_type,
+            is_cloud_server=is_cloud_server,
         )
 
     if client_version_stats and server_version_stats:
@@ -255,21 +289,14 @@ def valid_client_version(version: Version) -> bool:
     """Only version strings like 0.4.5 and 0.6.7.8 are valid. 0.8.6.7rc1 is not"""
     if version.is_prerelease or version.is_postrelease or version.is_devrelease:
         return False
-    if version.major == 0 and version.minor in [8, 9, 10, 11]:
-        return True
-
-    return False
+    return True
 
 
 def valid_server_version(version: Version) -> bool:
     """Only version strings like 0.8.x, 0.9.x or 0.10.x are valid. 0.1.x is not"""
     if version.is_prerelease or version.is_postrelease or version.is_devrelease:
         return False
-
-    if version.major == 0 and version.minor in [8, 9, 10]:
-        return True
-
-    return False
+    return True
 
 
 def is_client_server_compatible(client: VersionStats, server: VersionStats) -> int:
@@ -291,6 +318,27 @@ def is_client_server_compatible(client: VersionStats, server: VersionStats) -> i
         return server.version.micro - client.version.micro
 
 
+def is_server_default_cli_ahead(version_stats: DataHubVersionStats) -> bool:
+    """
+    Check if the server default CLI version is ahead of the current CLI version.
+    Returns True if server default CLI is newer and both versions are valid.
+    """
+    if not version_stats.server.current_server_default_cli_version:
+        return False
+
+    current_cli = version_stats.client.current
+    server_default_cli = version_stats.server.current_server_default_cli_version
+
+    is_valid_client_version = valid_client_version(current_cli.version)
+    is_valid_server_version = valid_client_version(server_default_cli.version)
+
+    if not (is_valid_client_version and is_valid_server_version):
+        return False
+
+    compatibility_result = is_client_server_compatible(current_cli, server_default_cli)
+    return compatibility_result > 0
+
+
 def _maybe_print_upgrade_message(
     version_stats: Optional[DataHubVersionStats],
 ) -> None:
@@ -312,9 +360,15 @@ def _maybe_print_upgrade_message(
             if version_stats.client.latest
             else None
         )
-        client_server_compat = is_client_server_compatible(
-            version_stats.client.current, version_stats.server.current
-        )
+        client_server_compat = 0
+        # Skip version compatibility checks for cloud servers (serverEnv="cloud")
+        # Cloud servers use different versioning schemes between server and CLI
+        is_cloud = version_stats.server.is_cloud_server
+
+        if not is_cloud:
+            client_server_compat = is_client_server_compatible(
+                version_stats.client.current, version_stats.server.current
+            )
 
         if latest_release_date and current_release_date:
             assert version_stats.client.latest
@@ -429,6 +483,8 @@ def check_upgrade_post(
 
 
 def check_upgrade(func: Callable[..., T]) -> Callable[..., T]:
+    log.debug(f"Checking upgrade for {func.__module__}.{func.__name__}")
+
     @wraps(func)
     def async_wrapper(*args: Any, **kwargs: Any) -> Any:
         with PerfTimer() as timer:

datahub/utilities/file_backed_collections.py

@@ -1,7 +1,6 @@
 import collections
 import gzip
 import logging
-import os
 import pathlib
 import pickle
 import shutil
@@ -28,18 +27,18 @@ from typing import (
     Union,
 )
 
+from datahub.configuration.env_vars import get_override_sqlite_version_req
 from datahub.ingestion.api.closeable import Closeable
 from datahub.utilities.sentinels import Unset, unset
 
 logger: logging.Logger = logging.getLogger(__name__)
 
-OVERRIDE_SQLITE_VERSION_REQUIREMENT_STR = (
-    os.environ.get("OVERRIDE_SQLITE_VERSION_REQ") or ""
-)
-OVERRIDE_SQLITE_VERSION_REQUIREMENT = (
-    OVERRIDE_SQLITE_VERSION_REQUIREMENT_STR
-    and OVERRIDE_SQLITE_VERSION_REQUIREMENT_STR.lower() != "false"
-)
+
+def _get_sqlite_version_override() -> bool:
+    """Check if SQLite version requirement should be overridden at runtime."""
+    override_str = get_override_sqlite_version_req()
+    return bool(override_str and override_str.lower() != "false")
+
 
 _DEFAULT_FILE_NAME = "sqlite.db"
 _DEFAULT_TABLE_NAME = "data"
@@ -231,7 +230,7 @@ class FileBackedDict(MutableMapping[str, _VT], Closeable, Generic[_VT]):
             # We use the ON CONFLICT clause to implement UPSERTs with sqlite.
             # This was added in 3.24.0 from 2018-06-04.
             # See https://www.sqlite.org/lang_conflict.html
-            if OVERRIDE_SQLITE_VERSION_REQUIREMENT:
+            if _get_sqlite_version_override():
                 self._use_sqlite_on_conflict = False
             else:
                 raise RuntimeError("SQLite version 3.24.0 or later is required")

datahub/utilities/is_pytest.py

@@ -1,6 +1,7 @@
-import os
 import sys
 
+from datahub.configuration.env_vars import get_test_mode
+
 
 def is_pytest_running() -> bool:
-    return "pytest" in sys.modules and os.environ.get("DATAHUB_TEST_MODE") == "1"
+    return "pytest" in sys.modules and get_test_mode() == "1"

datahub/utilities/logging_manager.py

@@ -15,13 +15,13 @@ import collections
 import contextlib
 import itertools
 import logging
-import os
 import pathlib
 import sys
 from typing import Deque, Iterator, Optional
 
 import click
 
+from datahub.configuration.env_vars import get_no_color, get_suppress_logging_manager
 from datahub.utilities.tee_io import TeeIO
 
 BASE_LOGGING_FORMAT = (
@@ -38,7 +38,7 @@ IN_MEMORY_LOG_BUFFER_SIZE = 2000  # lines
 IN_MEMORY_LOG_BUFFER_MAX_LINE_LENGTH = 2000  # characters
 
 
-NO_COLOR = os.environ.get("NO_COLOR", False)
+NO_COLOR = get_no_color()
 
 
 def extract_name_from_filename(filename: str, fallback_name: str) -> str:
@@ -179,6 +179,18 @@ class _LogBuffer:
         return text
 
 
+class _ResilientStreamHandler(logging.StreamHandler):
+    """StreamHandler that gracefully handles closed streams."""
+
+    def emit(self, record: logging.LogRecord) -> None:
+        try:
+            super().emit(record)
+        except (ValueError, OSError):
+            # Stream was closed (e.g., during pytest teardown)
+            # Silently ignore to prevent test failures
+            pass
+
+
 class _BufferLogHandler(logging.Handler):
     def __init__(self, storage: _LogBuffer) -> None:
         super().__init__()
@@ -201,7 +213,11 @@ class _BufferLogHandler(logging.Handler):
 def _remove_all_handlers(logger: logging.Logger) -> None:
     for handler in logger.handlers[:]:
         logger.removeHandler(handler)
-        handler.close()
+        try:
+            handler.close()
+        except (ValueError, OSError):
+            # Handler stream may already be closed (e.g., during pytest teardown)
+            pass
 
 
 _log_buffer = _LogBuffer(maxlen=IN_MEMORY_LOG_BUFFER_SIZE)
@@ -219,14 +235,14 @@ _default_formatter = logging.Formatter(BASE_LOGGING_FORMAT)
 def configure_logging(debug: bool, log_file: Optional[str] = None) -> Iterator[None]:
     _log_buffer.clear()
 
-    if os.environ.get("DATAHUB_SUPPRESS_LOGGING_MANAGER") == "1":
+    if get_suppress_logging_manager() == "1":
         # If we're running in pytest, we don't want to configure logging.
         yield
         return
 
     with contextlib.ExitStack() as stack:
         # Create stdout handler.
-        stream_handler = logging.StreamHandler()
+        stream_handler = _ResilientStreamHandler()
         stream_handler.addFilter(_DatahubLogFilter(debug=debug))
         stream_handler.setFormatter(_stream_formatter)
 
@@ -237,7 +253,7 @@ def configure_logging(debug: bool, log_file: Optional[str] = None) -> Iterator[N
             tee = TeeIO(sys.stdout, file)
             stack.enter_context(contextlib.redirect_stdout(tee))  # type: ignore
 
-            file_handler = logging.StreamHandler(file)
+            file_handler = _ResilientStreamHandler(file)
             file_handler.addFilter(_DatahubLogFilter(debug=True))
             file_handler.setFormatter(_default_formatter)
         else:

datahub/utilities/mapping.py

@@ -83,7 +83,7 @@ class Constants:
     MATCH = "match"
     USER_OWNER = "user"
     GROUP_OWNER = "group"
-    OPERAND_DATATYPE_SUPPORTED = [int, bool, str, float]
+    OPERAND_DATATYPE_SUPPORTED = [int, bool, str, float, list]
     TAG_PARTITION_KEY = "PARTITION_KEY"
     TAG_DIST_KEY = "DIST_KEY"
     TAG_SORT_KEY = "SORT_KEY"
@@ -455,7 +455,34 @@ class OperationProcessor:
         # function to check if a match clause is satisfied to a value.
         if not any(
             isinstance(raw_props_value, t) for t in Constants.OPERAND_DATATYPE_SUPPORTED
-        ) or not isinstance(raw_props_value, type(match_clause)):
+        ):
+            return None
+
+        # Handle list values by checking if any item in the list matches
+        if isinstance(raw_props_value, list):
+            # For lists, we need to find at least one matching item
+            # Return a match with the concatenated values of all matching items
+            matching_items = []
+            for item in raw_props_value:
+                if isinstance(item, str):
+                    match = re.match(match_clause, item)
+                    if match:
+                        matching_items.append(item)
+                elif isinstance(match_clause, type(item)):
+                    match = re.match(str(match_clause), str(item))
+                    if match:
+                        matching_items.append(str(item))
+
+            if matching_items:
+                # Create a synthetic match object with all matching items joined
+                combined_value = ",".join(matching_items)
+                return re.match(
+                    ".*", combined_value
+                )  # Always matches, returns combined value
+            return None
+
+        # Handle scalar values (existing logic)
+        elif not isinstance(raw_props_value, type(match_clause)):
             return None
         elif isinstance(raw_props_value, str):
             return re.match(match_clause, raw_props_value)

datahub/utilities/sample_data.py

@@ -1,12 +1,13 @@
-import os
 import pathlib
 import tempfile
 
 import requests
 
-DOCKER_COMPOSE_BASE = os.getenv(
-    "DOCKER_COMPOSE_BASE",
-    "https://raw.githubusercontent.com/datahub-project/datahub/master",
+from datahub.configuration.env_vars import get_docker_compose_base
+
+DOCKER_COMPOSE_BASE = (
+    get_docker_compose_base()
+    or "https://raw.githubusercontent.com/datahub-project/datahub/master"
 )
 BOOTSTRAP_MCES_FILE = "metadata-ingestion/examples/mce_files/bootstrap_mce.json"
 BOOTSTRAP_MCES_URL = f"{DOCKER_COMPOSE_BASE}/{BOOTSTRAP_MCES_FILE}"

datahub/utilities/server_config_util.py

@@ -183,6 +183,14 @@ class RestServiceConfig:
         managed_ingestion = self.raw_config.get("managedIngestion") or {}
         return managed_ingestion.get("enabled", False)
 
+    @property
+    def default_cli_version(self) -> Optional[str]:
+        """
+        Get the default CLI version.
+        """
+        managed_ingestion = self.raw_config.get("managedIngestion") or {}
+        return managed_ingestion.get("defaultCliVersion")
+
     @property
     def is_datahub_cloud(self) -> bool:
         """
@@ -234,7 +242,8 @@ class RestServiceConfig:
 
         # Check if this is a config-based feature
         if feature in config_based_features:
-            return config_based_features[feature]()
+            result = config_based_features[feature]()
+            return bool(result) if result is not None else False
 
         # For environment-based features, determine requirements based on cloud vs. non-cloud
         deployment_type = "cloud" if self.is_datahub_cloud else "core"

datahub/utilities/sqlalchemy_query_combiner.py

@@ -272,8 +272,11 @@ class SQLAlchemyQueryCombiner:
                     self.report.uncombined_queries_issued += 1
                     return _sa_execute_underlying_method(conn, query, *args, **kwargs)
 
-        with _sa_execute_method_patching_lock, unittest.mock.patch(
-            "sqlalchemy.engine.Connection.execute", _sa_execute_fake
+        with (
+            _sa_execute_method_patching_lock,
+            unittest.mock.patch(
+                "sqlalchemy.engine.Connection.execute", _sa_execute_fake
+            ),
         ):
             yield self
 

datahub/utilities/stats_collections.py

@@ -56,3 +56,7 @@ class TopKDict(DefaultDict[_KT, _VT]):
 
 def int_top_k_dict() -> TopKDict[str, int]:
     return TopKDict(int)
+
+
+def float_top_k_dict() -> TopKDict[str, float]:
+    return TopKDict(float)

datahub/utilities/urns/urn.py

@@ -1,8 +1,47 @@
-from datahub.metadata.urns import Urn
+from typing import Optional
 
-__all__ = ["Urn", "guess_entity_type"]
+from datahub.metadata.urns import (
+    DataPlatformUrn,
+    Urn,
+)
+
+__all__ = ["Urn", "guess_entity_type", "guess_platform_name"]
 
 
 def guess_entity_type(urn: str) -> str:
     assert urn.startswith("urn:li:"), "urns must start with urn:li:"
     return urn.split(":")[2]
+
+
+def guess_platform_name(urn: str) -> Optional[str]:
+    """Extract platform from URN using a mapping dictionary."""
+    urn_obj = Urn.from_string(urn)
+
+    try:
+        platform = None
+        try:
+            platform = urn_obj.platform  # type: ignore[attr-defined]
+            platform_name = DataPlatformUrn.from_string(
+                platform
+            ).get_entity_id_as_string()
+            return platform_name
+        except AttributeError:
+            pass
+        try:
+            return urn_obj.orchestrator  # type: ignore[attr-defined]
+        except AttributeError:
+            pass
+        try:
+            return urn_obj.dashboard_tool  # type: ignore[attr-defined]
+        except AttributeError:
+            pass
+        try:
+            return urn_obj.ml_model_tool  # type: ignore[attr-defined]
+        except AttributeError:
+            pass
+
+        if platform is None:
+            return None
+    except AttributeError:
+        pass
+    return None