acryl-datahub 1.2.0.1rc1__py3-none-any.whl → 1.2.0.2rc2__py3-none-any.whl

datahub/sdk/search_filters.py

@@ -29,7 +29,7 @@ from datahub.ingestion.graph.filters import (
     _get_status_filter,
 )
 from datahub.metadata.schema_classes import EntityTypeName
-from datahub.metadata.urns import DataPlatformUrn, DomainUrn
+from datahub.metadata.urns import ContainerUrn, DataPlatformUrn, DomainUrn
 
 _AndSearchFilterRule = TypedDict(
     "_AndSearchFilterRule", {"and": List[SearchFilterRule]}
@@ -81,7 +81,7 @@ class _EntityTypeFilter(_BaseFilter):
     ENTITY_TYPE_FIELD: ClassVar[str] = "_entityType"
 
     entity_type: List[str] = pydantic.Field(
-        description="The entity type to filter on. Can be 'dataset', 'chart', 'dashboard', 'corpuser', etc.",
+        description="The entity type to filter on. Can be 'dataset', 'chart', 'dashboard', 'corpuser', 'dataProduct', etc.",
     )
 
     def _build_rule(self) -> SearchFilterRule:
@@ -174,6 +174,39 @@ class _DomainFilter(_BaseFilter):
         return [{"and": [self._build_rule()]}]
 
 
+class _ContainerFilter(_BaseFilter):
+    container: List[str]
+    direct_descendants_only: bool = pydantic.Field(
+        default=False,
+        description="If true, only entities that are direct descendants of the container will be returned.",
+    )
+
+    @pydantic.validator("container", each_item=True)
+    def validate_container(cls, v: str) -> str:
+        return str(ContainerUrn.from_string(v))
+
+    @classmethod
+    def _field_discriminator(cls) -> str:
+        return "container"
+
+    def _build_rule(self) -> SearchFilterRule:
+        if self.direct_descendants_only:
+            return SearchFilterRule(
+                field="container",
+                condition="EQUAL",
+                values=self.container,
+            )
+        else:
+            return SearchFilterRule(
+                field="browsePathV2",
+                condition="CONTAIN",
+                values=self.container,
+            )
+
+    def compile(self) -> _OrFilters:
+        return [{"and": [self._build_rule()]}]
+
+
 class _EnvFilter(_BaseFilter):
     # Note that not all entity types have an env (e.g. dashboards / charts).
     # If the env filter is specified, these will be excluded.
@@ -342,6 +375,8 @@ def _filter_discriminator(v: Any) -> Optional[str]:
     keys = list(v.keys())
     if len(keys) == 1:
         return keys[0]
+    elif set(keys).issuperset({"container"}):
+        return _ContainerFilter._field_discriminator()
     elif set(keys).issuperset({"field", "condition"}):
         return _CustomCondition._field_discriminator()
 
@@ -360,6 +395,7 @@ if TYPE_CHECKING or not PYDANTIC_SUPPORTS_CALLABLE_DISCRIMINATOR:
         _StatusFilter,
         _PlatformFilter,
         _DomainFilter,
+        _ContainerFilter,
         _EnvFilter,
         _CustomCondition,
     ]
@@ -385,6 +421,7 @@ else:
             Annotated[_StatusFilter, Tag(_StatusFilter._field_discriminator())],
             Annotated[_PlatformFilter, Tag(_PlatformFilter._field_discriminator())],
             Annotated[_DomainFilter, Tag(_DomainFilter._field_discriminator())],
+            Annotated[_ContainerFilter, Tag(_ContainerFilter._field_discriminator())],
             Annotated[_EnvFilter, Tag(_EnvFilter._field_discriminator())],
             Annotated[_CustomCondition, Tag(_CustomCondition._field_discriminator())],
         ],
@@ -468,6 +505,18 @@ class FilterDsl:
     def domain(domain: Union[str, Sequence[str]], /) -> _DomainFilter:
         return _DomainFilter(domain=[domain] if isinstance(domain, str) else domain)
 
+    @staticmethod
+    def container(
+        container: Union[str, Sequence[str]],
+        /,
+        *,
+        direct_descendants_only: bool = False,
+    ) -> _ContainerFilter:
+        return _ContainerFilter(
+            container=[container] if isinstance(container, str) else container,
+            direct_descendants_only=direct_descendants_only,
+        )
+
     @staticmethod
     def env(env: Union[str, Sequence[str]], /) -> _EnvFilter:
         return _EnvFilter(env=[env] if isinstance(env, str) else env)

datahub/sql_parsing/sql_parsing_aggregator.py

@@ -190,6 +190,7 @@ class QueryMetadata:
             source=models.QuerySourceClass.SYSTEM,
             created=self.make_created_audit_stamp(),
             lastModified=self.make_last_modified_audit_stamp(),
+            origin=self.origin.urn() if self.origin else None,
         )
 
 

datahub/upgrade/upgrade.py

@@ -352,9 +352,11 @@ 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
+        if version_stats.server.current_server_type != "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