diff --git a/metadata-ingestion/src/datahub/upgrade/upgrade.py b/metadata-ingestion/src/datahub/upgrade/upgrade.py index adebd45f47..c5dbdda2e7 100644 --- a/metadata-ingestion/src/datahub/upgrade/upgrade.py +++ b/metadata-ingestion/src/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,6 +28,18 @@ 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 @@ -233,10 +245,7 @@ async def _retrieve_version_stats( 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, - ) + _safe_version_stats(current_server_default_cli_version) if current_server_default_cli_version else None ),