acryl-datahub 1.2.0__py3-none-any.whl → 1.2.0.1__py3-none-any.whl

datahub/sdk/search_filters.py

@@ -2,6 +2,8 @@ from __future__ import annotations
 
 import abc
 from typing import (
+    TYPE_CHECKING,
+    Annotated,
     Any,
     ClassVar,
     Iterator,
@@ -15,7 +17,10 @@ from typing import (
 import pydantic
 
 from datahub.configuration.common import ConfigModel
-from datahub.configuration.pydantic_migration_helpers import PYDANTIC_VERSION_2
+from datahub.configuration.pydantic_migration_helpers import (
+    PYDANTIC_SUPPORTS_CALLABLE_DISCRIMINATOR,
+    PYDANTIC_VERSION_2,
+)
 from datahub.ingestion.graph.client import flexible_entity_type_to_graphql
 from datahub.ingestion.graph.filters import (
     FilterOperator,
@@ -42,12 +47,29 @@ class _BaseFilter(ConfigModel):
             populate_by_name = True
 
     @abc.abstractmethod
-    def compile(self) -> _OrFilters:
-        pass
+    def compile(self) -> _OrFilters: ...
 
     def dfs(self) -> Iterator[_BaseFilter]:
         yield self
 
+    @classmethod
+    def _field_discriminator(cls) -> str:
+        if cls is _BaseFilter:
+            raise ValueError("Cannot get discriminator for _BaseFilter")
+        if PYDANTIC_VERSION_2:
+            fields: dict = cls.model_fields  # type: ignore
+        else:
+            fields = cls.__fields__  # type: ignore
+
+        # Assumes that there's only one field name per filter.
+        # If that's not the case, this method should be overridden.
+        if len(fields.keys()) != 1:
+            raise ValueError(
+                f"Found multiple fields that could be the discriminator for this filter: {list(fields.keys())}"
+            )
+        name, field = next(iter(fields.items()))
+        return field.alias or name  # type: ignore
+
 
 class _EntityTypeFilter(_BaseFilter):
     """Filter for specific entity types.
@@ -74,15 +96,19 @@ class _EntityTypeFilter(_BaseFilter):
 
 
 class _EntitySubtypeFilter(_BaseFilter):
-    entity_subtype: str = pydantic.Field(
+    entity_subtype: List[str] = pydantic.Field(
         description="The entity subtype to filter on. Can be 'Table', 'View', 'Source', etc. depending on the native platform's concepts.",
     )
 
+    @pydantic.validator("entity_subtype", pre=True)
+    def validate_entity_subtype(cls, v: str) -> List[str]:
+        return [v] if not isinstance(v, list) else v
+
     def _build_rule(self) -> SearchFilterRule:
         return SearchFilterRule(
             field="typeNames",
             condition="EQUAL",
-            values=[self.entity_subtype],
+            values=self.entity_subtype,
         )
 
     def compile(self) -> _OrFilters:
@@ -196,6 +222,10 @@ class _CustomCondition(_BaseFilter):
         )
         return [{"and": [rule]}]
 
+    @classmethod
+    def _field_discriminator(cls) -> str:
+        return "_custom"
+
 
 class _And(_BaseFilter):
     """Represents an AND conjunction of filters."""
@@ -302,31 +332,69 @@ class _Not(_BaseFilter):
         yield from self.not_.dfs()
 
 
-# TODO: With pydantic 2, we can use a RootModel with a
-# discriminated union to make the error messages more informative.
-Filter = Union[
-    _And,
-    _Or,
-    _Not,
-    _EntityTypeFilter,
-    _EntitySubtypeFilter,
-    _StatusFilter,
-    _PlatformFilter,
-    _DomainFilter,
-    _EnvFilter,
-    _CustomCondition,
-]
-
-
-# Required to resolve forward references to "Filter"
-if PYDANTIC_VERSION_2:
-    _And.model_rebuild()  # type: ignore
-    _Or.model_rebuild()  # type: ignore
-    _Not.model_rebuild()  # type: ignore
-else:
+def _filter_discriminator(v: Any) -> Optional[str]:
+    if isinstance(v, _BaseFilter):
+        return v._field_discriminator()
+
+    if not isinstance(v, dict):
+        return None
+
+    keys = list(v.keys())
+    if len(keys) == 1:
+        return keys[0]
+    elif set(keys).issuperset({"field", "condition"}):
+        return _CustomCondition._field_discriminator()
+
+    return None
+
+
+if TYPE_CHECKING or not PYDANTIC_SUPPORTS_CALLABLE_DISCRIMINATOR:
+    # The `not TYPE_CHECKING` bit is required to make the linter happy,
+    # since we currently only run mypy with pydantic v1.
+    Filter = Union[
+        _And,
+        _Or,
+        _Not,
+        _EntityTypeFilter,
+        _EntitySubtypeFilter,
+        _StatusFilter,
+        _PlatformFilter,
+        _DomainFilter,
+        _EnvFilter,
+        _CustomCondition,
+    ]
+
     _And.update_forward_refs()
     _Or.update_forward_refs()
     _Not.update_forward_refs()
+else:
+    from pydantic import Discriminator, Tag
+
+    # TODO: Once we're fully on pydantic 2, we can use a RootModel here.
+    # That way we'd be able to attach methods to the Filter type.
+    # e.g. replace load_filters(...) with Filter.load(...)
+    Filter = Annotated[
+        Union[
+            Annotated[_And, Tag(_And._field_discriminator())],
+            Annotated[_Or, Tag(_Or._field_discriminator())],
+            Annotated[_Not, Tag(_Not._field_discriminator())],
+            Annotated[_EntityTypeFilter, Tag(_EntityTypeFilter._field_discriminator())],
+            Annotated[
+                _EntitySubtypeFilter, Tag(_EntitySubtypeFilter._field_discriminator())
+            ],
+            Annotated[_StatusFilter, Tag(_StatusFilter._field_discriminator())],
+            Annotated[_PlatformFilter, Tag(_PlatformFilter._field_discriminator())],
+            Annotated[_DomainFilter, Tag(_DomainFilter._field_discriminator())],
+            Annotated[_EnvFilter, Tag(_EnvFilter._field_discriminator())],
+            Annotated[_CustomCondition, Tag(_CustomCondition._field_discriminator())],
+        ],
+        Discriminator(_filter_discriminator),
+    ]
+
+    # Required to resolve forward references to "Filter"
+    _And.model_rebuild()  # type: ignore
+    _Or.model_rebuild()  # type: ignore
+    _Not.model_rebuild()  # type: ignore
 
 
 def load_filters(obj: Any) -> Filter:

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