acryl-datahub 1.1.0.5rc7__py3-none-any.whl → 1.1.0.5rc9__py3-none-any.whl

datahub/upgrade/upgrade.py

@@ -32,6 +32,7 @@ class ServerVersionStats(BaseModel):
     current: VersionStats
     latest: Optional[VersionStats] = None
     current_server_type: Optional[str] = None
+    current_server_default_cli_version: Optional[VersionStats] = None
 
 
 class ClientVersionStats(BaseModel):
@@ -44,7 +45,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 +53,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 +133,7 @@ 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]]:
     import aiohttp
 
     server_config: Optional[RestServiceConfig] = None
@@ -151,12 +153,13 @@ 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
     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
         if server_type == "quickstart" and commit_hash:
             async with aiohttp.ClientSession(
                 headers={"Accept": "application/vnd.github.v3+json"}
@@ -171,7 +174,12 @@ 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,
+    )
 
 
 def retrieve_version_stats(
@@ -214,6 +222,7 @@ async def _retrieve_version_stats(
     (
         current_server_type,
         current_server_version,
+        current_server_default_cli_version,
         current_server_release_date,
     ) = results[2]
 
@@ -223,6 +232,14 @@ async def _retrieve_version_stats(
             current=VersionStats(
                 version=current_server_version, release_date=current_server_release_date
             ),
+            current_server_default_cli_version=(
+                VersionStats(
+                    version=Version(current_server_default_cli_version),
+                    release_date=None,
+                )
+                if current_server_default_cli_version
+                else None
+            ),
             latest=(
                 VersionStats(version=last_server_version, release_date=last_server_date)
                 if last_server_version
@@ -255,21 +272,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 +301,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:
@@ -429,6 +460,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/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:
         """