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

datahub/sdk/mlmodelgroup.py

@@ -0,0 +1,240 @@
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Dict, List, Optional, Sequence, Type, Union
+
+from typing_extensions import Self
+
+from datahub.emitter.mce_builder import DEFAULT_ENV
+from datahub.metadata.schema_classes import (
+    AspectBag,
+    MLModelGroupPropertiesClass,
+)
+from datahub.metadata.urns import DataProcessInstanceUrn, MlModelGroupUrn, Urn
+from datahub.sdk._shared import (
+    DomainInputType,
+    HasDomain,
+    HasInstitutionalMemory,
+    HasOwnership,
+    HasPlatformInstance,
+    HasStructuredProperties,
+    HasTags,
+    HasTerms,
+    LinksInputType,
+    OwnersInputType,
+    StructuredPropertyInputType,
+    TagsInputType,
+    TermsInputType,
+    make_time_stamp,
+    parse_time_stamp,
+)
+from datahub.sdk.entity import Entity, ExtraAspectsType
+
+
+class MLModelGroup(
+    HasPlatformInstance,
+    HasOwnership,
+    HasInstitutionalMemory,
+    HasTags,
+    HasTerms,
+    HasDomain,
+    HasStructuredProperties,
+    Entity,
+):
+    __slots__ = ()
+
+    @classmethod
+    def get_urn_type(cls) -> Type[MlModelGroupUrn]:
+        return MlModelGroupUrn
+
+    def __init__(
+        self,
+        id: str,
+        platform: str,
+        name: Optional[str] = "",
+        platform_instance: Optional[str] = None,
+        env: str = DEFAULT_ENV,
+        # Model group properties
+        description: Optional[str] = None,
+        display_name: Optional[str] = None,
+        external_url: Optional[str] = None,
+        custom_properties: Optional[Dict[str, str]] = None,
+        created: Optional[datetime] = None,
+        last_modified: Optional[datetime] = None,
+        # Standard aspects
+        owners: Optional[OwnersInputType] = None,
+        links: Optional[LinksInputType] = None,
+        tags: Optional[TagsInputType] = None,
+        terms: Optional[TermsInputType] = None,
+        domain: Optional[DomainInputType] = None,
+        training_jobs: Optional[Sequence[Union[str, DataProcessInstanceUrn]]] = None,
+        downstream_jobs: Optional[Sequence[Union[str, DataProcessInstanceUrn]]] = None,
+        structured_properties: Optional[StructuredPropertyInputType] = None,
+        extra_aspects: ExtraAspectsType = None,
+    ):
+        urn = MlModelGroupUrn(platform=platform, name=id, env=env)
+        super().__init__(urn)
+        self._set_extra_aspects(extra_aspects)
+
+        self._set_platform_instance(urn.platform, platform_instance)
+
+        # Set MLModelGroupProperties aspect
+        self._ensure_model_group_props(name=display_name or name)
+
+        if description is not None:
+            self.set_description(description)
+        if external_url is not None:
+            self.set_external_url(external_url)
+        if custom_properties is not None:
+            self.set_custom_properties(custom_properties)
+        if created is not None:
+            self.set_created(created)
+        if last_modified is not None:
+            self.set_last_modified(last_modified)
+
+        # Standard aspects
+        if owners is not None:
+            self.set_owners(owners)
+        if links is not None:
+            self.set_links(links)
+        if tags is not None:
+            self.set_tags(tags)
+        if terms is not None:
+            self.set_terms(terms)
+        if domain is not None:
+            self.set_domain(domain)
+
+        # ML model group specific aspects
+        if training_jobs is not None:
+            self.set_training_jobs(training_jobs)
+        if downstream_jobs is not None:
+            self.set_downstream_jobs(downstream_jobs)
+        if structured_properties is not None:
+            for key, value in structured_properties.items():
+                self.set_structured_property(property_urn=key, values=value)
+
+    @classmethod
+    def _new_from_graph(cls, urn: Urn, current_aspects: AspectBag) -> Self:
+        assert isinstance(urn, MlModelGroupUrn)
+        entity = cls(
+            platform=urn.platform,
+            id=urn.name,
+            env=urn.env,
+        )
+        return entity._init_from_graph(current_aspects)
+
+    @property
+    def urn(self) -> MlModelGroupUrn:
+        return self._urn  # type: ignore
+
+    def _ensure_model_group_props(
+        self, *, name: Optional[str] = None
+    ) -> MLModelGroupPropertiesClass:
+        if name is not None:
+            return self._setdefault_aspect(MLModelGroupPropertiesClass(name=name))
+
+        props = self._get_aspect(MLModelGroupPropertiesClass)
+        if props is None:
+            # If we need properties but they don't exist and no name was provided
+            return self._setdefault_aspect(
+                MLModelGroupPropertiesClass(name=self.urn.name)
+            )
+        return props
+
+    @property
+    def name(self) -> Optional[str]:
+        return self._ensure_model_group_props().name
+
+    def set_name(self, display_name: str) -> None:
+        self._ensure_model_group_props().name = display_name
+
+    @property
+    def description(self) -> Optional[str]:
+        return self._ensure_model_group_props().description
+
+    def set_description(self, description: str) -> None:
+        self._ensure_model_group_props().description = description
+
+    @property
+    def external_url(self) -> Optional[str]:
+        return self._ensure_model_group_props().externalUrl
+
+    def set_external_url(self, external_url: str) -> None:
+        self._ensure_model_group_props().externalUrl = external_url
+
+    @property
+    def custom_properties(self) -> Optional[Dict[str, str]]:
+        return self._ensure_model_group_props().customProperties
+
+    def set_custom_properties(self, custom_properties: Dict[str, str]) -> None:
+        self._ensure_model_group_props().customProperties = custom_properties
+
+    @property
+    def created(self) -> Optional[datetime]:
+        return parse_time_stamp(self._ensure_model_group_props().created)
+
+    def set_created(self, created: datetime) -> None:
+        self._ensure_model_group_props().created = make_time_stamp(created)
+
+    @property
+    def last_modified(self) -> Optional[datetime]:
+        return parse_time_stamp(self._ensure_model_group_props().lastModified)
+
+    def set_last_modified(self, last_modified: datetime) -> None:
+        self._ensure_model_group_props().lastModified = make_time_stamp(last_modified)
+
+    @property
+    def training_jobs(self) -> Optional[List[str]]:
+        return self._ensure_model_group_props().trainingJobs
+
+    def set_training_jobs(
+        self, training_jobs: Sequence[Union[str, DataProcessInstanceUrn]]
+    ) -> None:
+        self._ensure_model_group_props().trainingJobs = [
+            str(job) for job in training_jobs
+        ]
+
+    def add_training_job(
+        self, training_job: Union[str, DataProcessInstanceUrn]
+    ) -> None:
+        props = self._ensure_model_group_props()
+        if props.trainingJobs is None:
+            props.trainingJobs = []
+        props.trainingJobs.append(str(training_job))
+
+    def remove_training_job(
+        self, training_job: Union[str, DataProcessInstanceUrn]
+    ) -> None:
+        props = self._ensure_model_group_props()
+        if props.trainingJobs is not None:
+            props.trainingJobs = [
+                job for job in props.trainingJobs if job != str(training_job)
+            ]
+
+    @property
+    def downstream_jobs(self) -> Optional[List[str]]:
+        return self._ensure_model_group_props().downstreamJobs
+
+    def set_downstream_jobs(
+        self, downstream_jobs: Sequence[Union[str, DataProcessInstanceUrn]]
+    ) -> None:
+        self._ensure_model_group_props().downstreamJobs = [
+            str(job) for job in downstream_jobs
+        ]
+
+    def add_downstream_job(
+        self, downstream_job: Union[str, DataProcessInstanceUrn]
+    ) -> None:
+        props = self._ensure_model_group_props()
+        if props.downstreamJobs is None:
+            props.downstreamJobs = []
+        props.downstreamJobs.append(str(downstream_job))
+
+    def remove_downstream_job(
+        self, downstream_job: Union[str, DataProcessInstanceUrn]
+    ) -> None:
+        props = self._ensure_model_group_props()
+        if props.downstreamJobs is not None:
+            props.downstreamJobs = [
+                job for job in props.downstreamJobs if job != str(downstream_job)
+            ]

datahub/sdk/search_client.py

@@ -2,36 +2,107 @@ from __future__ import annotations
 
 from typing import (
     TYPE_CHECKING,
-    Dict,
     Iterable,
     List,
     Optional,
+    Tuple,
+    Type,
+    TypeVar,
 )
 
-from datahub.ingestion.graph.filters import RawSearchFilterRule
+from datahub.ingestion.graph.filters import RawSearchFilter, RemovedStatusFilter
 from datahub.metadata.urns import Urn
-from datahub.sdk.search_filters import Filter
+from datahub.sdk.search_filters import (
+    Filter,
+    FilterDsl,
+    _EntityTypeFilter,
+    _OrFilters,
+    _StatusFilter,
+)
+from datahub.utilities.ordered_set import OrderedSet
 
 if TYPE_CHECKING:
     from datahub.sdk.main_client import DataHubClient
 
 
+_FilterType = TypeVar("_FilterType", bound=Filter)
+
+
+def _typed_dfs(
+    filter: Optional[_FilterType], type: Type[_FilterType]
+) -> Optional[List[_FilterType]]:
+    if filter is None:
+        return None
+
+    found: Optional[List[_FilterType]] = None
+    for f in filter.dfs():
+        if isinstance(f, type):
+            if found is None:
+                found = []
+            found.append(f)
+    return found
+
+
 def compile_filters(
     filter: Optional[Filter],
-) -> Optional[List[Dict[str, List[RawSearchFilterRule]]]]:
+) -> Tuple[Optional[List[str]], RawSearchFilter]:
     # TODO: Not every filter type is supported for every entity type.
     # If we can detect issues with the filters at compile time, we should
     # raise an error.
 
-    if filter is None:
-        return None
+    existing_soft_deleted_filter = _typed_dfs(filter, _StatusFilter)
+    if existing_soft_deleted_filter is None:
+        soft_deleted_filter = FilterDsl.soft_deleted(
+            RemovedStatusFilter.NOT_SOFT_DELETED
+        )
+        if filter is None:
+            filter = soft_deleted_filter
+        else:
+            filter = FilterDsl.and_(filter, soft_deleted_filter)
+
+    # This should be safe - if filter were None coming in, then we would replace it
+    # with the soft-deleted filter.
+    assert filter is not None
 
     initial_filters = filter.compile()
-    return [
+
+    compiled_filters: RawSearchFilter = [
         {"and": [rule.to_raw() for rule in andClause["and"]]}
         for andClause in initial_filters
     ]
 
+    entity_types = compute_entity_types(initial_filters)
+
+    return entity_types, compiled_filters
+
+
+def compute_entity_types(
+    filters: _OrFilters,
+) -> Optional[List[str]]:
+    found_filters = False
+    found_positive_filters = False
+    entity_types: OrderedSet[str] = OrderedSet()
+    for ands in filters:
+        for clause in ands["and"]:
+            if clause.field == _EntityTypeFilter.ENTITY_TYPE_FIELD:
+                found_filters = True
+                if not clause.negated:
+                    found_positive_filters = True
+
+                entity_types.update(clause.values)
+
+    if not found_filters:
+        # If we didn't find any filters, use None so we use the default set.
+        return None
+
+    if not found_positive_filters:
+        # If we only found negated filters, then it's probably a query like
+        # "find me all entities except for dashboards". In that case, we
+        # still want to use the default set.
+        return None
+
+    return list(entity_types)
+
 
 class SearchClient:
     def __init__(self, client: DataHubClient):
@@ -41,10 +112,16 @@ class SearchClient:
         self,
         query: Optional[str] = None,
         filter: Optional[Filter] = None,
+        *,
+        skip_cache: bool = False,
     ) -> Iterable[Urn]:
         # TODO: Add better limit / pagination support.
+        types, compiled_filters = compile_filters(filter)
         for urn in self._client._graph.get_urns_by_filter(
             query=query,
-            extra_or_filters=compile_filters(filter),
+            status=None,
+            extra_or_filters=compiled_filters,
+            entity_types=types,
+            skip_cache=skip_cache,
         ):
             yield Urn.from_string(urn)