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

datahub/upgrade/upgrade.py

@@ -7,13 +7,15 @@ 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__
 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__)
 
@@ -26,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):
@@ -42,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__
@@ -50,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:
@@ -109,7 +126,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,19 +141,22 @@ 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(
     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 = 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)
@@ -148,17 +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.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
+        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"}
@@ -173,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(
@@ -216,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
@@ -225,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:
@@ -257,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:
@@ -293,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:
@@ -314,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
@@ -379,7 +431,8 @@ def _maybe_print_upgrade_message(
                 + click.style(
                     f"➡️  Upgrade via \"pip install 'acryl-datahub=={version_stats.server.current.version}'\"",
                     fg="cyan",
-                )
+                ),
+                err=True,
             )
     elif client_server_compat == 0 and encourage_cli_upgrade:
         with contextlib.suppress(Exception):
@@ -389,7 +442,8 @@ def _maybe_print_upgrade_message(
                 + click.style(
                     f"You seem to be running an old version of datahub cli: {current_version} {get_days(current_release_date)}. Latest version is {latest_version} {get_days(latest_release_date)}.\nUpgrade via \"pip install -U 'acryl-datahub'\"",
                     fg="cyan",
-                )
+                ),
+                err=True,
             )
     elif encourage_quickstart_upgrade:
         try:
@@ -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")
@@ -250,7 +249,7 @@ class FileBackedDict(MutableMapping[str, _VT], Closeable, Generic[_VT]):
                 rowid INTEGER PRIMARY KEY AUTOINCREMENT,
                 key TEXT UNIQUE,
                 value BLOB
-                {"".join(f", {column_name} BLOB" for column_name in self.extra_columns.keys())}
+                {"".join(f", {column_name} BLOB" for column_name in self.extra_columns)}
             )"""
         )
 
@@ -267,7 +266,7 @@ class FileBackedDict(MutableMapping[str, _VT], Closeable, Generic[_VT]):
         if self.indexes_created:
             return
         # The key column will automatically be indexed, but we need indexes for the extra columns.
-        for column_name in self.extra_columns.keys():
+        for column_name in self.extra_columns:
             self._conn.execute(
                 f"CREATE INDEX {self.tablename}_{column_name} ON {self.tablename} ({column_name})"
             )
@@ -305,12 +304,12 @@ class FileBackedDict(MutableMapping[str, _VT], Closeable, Generic[_VT]):
                 f"""INSERT INTO {self.tablename} (
                     key,
                     value
-                    {"".join(f", {column_name}" for column_name in self.extra_columns.keys())}
+                    {"".join(f", {column_name}" for column_name in self.extra_columns)}
                 )
                 VALUES ({", ".join(["?"] * (2 + len(self.extra_columns)))})
                 ON CONFLICT (key) DO UPDATE SET
                     value = excluded.value
-                    {"".join(f", {column_name} = excluded.{column_name}" for column_name in self.extra_columns.keys())}
+                    {"".join(f", {column_name} = excluded.{column_name}" for column_name in self.extra_columns)}
                 """,
                 items_to_write,
             )
@@ -321,7 +320,7 @@ class FileBackedDict(MutableMapping[str, _VT], Closeable, Generic[_VT]):
                         f"""INSERT INTO {self.tablename} (
                             key,
                             value
-                            {"".join(f", {column_name}" for column_name in self.extra_columns.keys())}
+                            {"".join(f", {column_name}" for column_name in self.extra_columns)}
                         )
                         VALUES ({", ".join(["?"] * (2 + len(self.extra_columns)))})""",
                         item,
@@ -330,7 +329,7 @@ class FileBackedDict(MutableMapping[str, _VT], Closeable, Generic[_VT]):
                     self._conn.execute(
                         f"""UPDATE {self.tablename} SET
                             value = ?
-                            {"".join(f", {column_name} = ?" for column_name in self.extra_columns.keys())}
+                            {"".join(f", {column_name} = ?" for column_name in self.extra_columns)}
                         WHERE key = ?""",
                         (*item[1:], item[0]),
                     )

datahub/utilities/hive_schema_to_avro.py

@@ -155,7 +155,7 @@ class HiveColumnToAvroConverter:
 
     @staticmethod
     def _parse_basic_datatype_string(s: str) -> Dict[str, object]:
-        if s in HiveColumnToAvroConverter._PRIVIMITE_HIVE_TYPE_TO_AVRO_TYPE.keys():
+        if s in HiveColumnToAvroConverter._PRIVIMITE_HIVE_TYPE_TO_AVRO_TYPE:
             return {
                 "type": HiveColumnToAvroConverter._PRIVIMITE_HIVE_TYPE_TO_AVRO_TYPE[s],
                 "native_data_type": s,
@@ -218,7 +218,7 @@ class HiveColumnToAvroConverter:
         buf = ""
         level = 0
         for c in s:
-            if c in HiveColumnToAvroConverter._BRACKETS.keys():
+            if c in HiveColumnToAvroConverter._BRACKETS:
                 level += 1
                 buf += c
             elif c in HiveColumnToAvroConverter._BRACKETS.values():

datahub/utilities/ingest_utils.py

@@ -32,10 +32,10 @@ def deploy_source_vars(
     name: Optional[str],
     config: str,
     urn: Optional[str],
-    executor_id: str,
+    executor_id: Optional[str],
     cli_version: Optional[str],
     schedule: Optional[str],
-    time_zone: str,
+    time_zone: Optional[str],
     extra_pip: Optional[str],
     debug: bool = False,
 ) -> dict:

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:
@@ -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]"
 
@@ -178,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__()
@@ -188,13 +201,23 @@ 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:
     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)
@@ -212,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)
 
@@ -230,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}"