acryl-datahub-cloud 0.3.11rc0__py3-none-any.whl → 0.3.16.1rc0__py3-none-any.whl

acryl_datahub_cloud/sdk/entities/monitor.py

@@ -0,0 +1,291 @@
+import logging
+from typing import (
+    Any,
+    Dict,
+    Optional,
+    Tuple,
+    Type,
+    TypeAlias,
+    Union,
+)
+
+from typing_extensions import (
+    Self,
+    assert_never,
+)
+
+from datahub.emitter.enum_helpers import get_enum_options
+from datahub.errors import SdkUsageError
+from datahub.metadata import schema_classes as models
+from datahub.metadata.urns import (
+    AssertionUrn,
+    DatasetUrn,
+    MonitorUrn,
+    Urn,
+)
+from datahub.sdk.entity import Entity
+
+logger = logging.getLogger(__name__)
+
+
+MonitorIdentityInputType: TypeAlias = Union[
+    Tuple[DatasetUrn, str],  # (dataset urn, monitor id)
+    models.MonitorKeyClass,
+    MonitorUrn,
+    Tuple[
+        DatasetUrn, AssertionUrn
+    ],  # (dataset urn, assertion urn - monitor id is got from assertion id)
+]
+
+MonitorInfoInputType: TypeAlias = Union[
+    Tuple[  # (monitor type [ASSERTION|FRESHNESS], monitor status mode [ACTIVE|INACTIVE|PASSIVE])
+        str, str
+    ],
+    Tuple[  # (monitor type enum, monitor status mode enum)
+        models.MonitorTypeClass, models.MonitorModeClass
+    ],
+    models.MonitorInfoClass,
+]
+
+
+class Monitor(Entity):
+    """
+    Monitor entity class.
+    """
+
+    __slots__ = ()
+
+    @classmethod
+    def get_urn_type(cls) -> Type[MonitorUrn]:
+        """Get the URN type for monitors.
+
+        Returns:
+            The MonitorUrn class.
+        """
+        return MonitorUrn
+
+    def __init__(
+        self,
+        id: MonitorIdentityInputType,
+        # MonitorInfo
+        info: MonitorInfoInputType,
+        external_url: Optional[str] = None,
+        custom_properties: Optional[Dict[str, str]] = None,
+    ) -> None:
+        super().__init__(urn=Monitor._ensure_id(id=id))
+
+        self._set_info(info)
+
+        if external_url is not None:
+            # this may overwrite the value set from info#externalUrl value
+            self.set_external_url(external_url)
+        if custom_properties is not None:
+            # this may overwrite the value set from info#customPropperties value
+            self.set_custom_properties(custom_properties)
+
+    @classmethod
+    def _new_from_graph(cls, urn: Urn, current_aspects: models.AspectBag) -> Self:
+        assert isinstance(urn, MonitorUrn)
+        assert "monitorInfo" in current_aspects, "MonitorInfo is required"
+
+        entity = cls(id=urn, info=current_aspects["monitorInfo"])
+        return entity._init_from_graph(current_aspects)
+
+    @property
+    def urn(self) -> MonitorUrn:
+        assert isinstance(self._urn, MonitorUrn)
+        return self._urn
+
+    @classmethod
+    def _ensure_id(
+        cls,
+        id: MonitorIdentityInputType,
+    ) -> MonitorUrn:
+        if isinstance(id, tuple):
+            if (
+                len(id) != 2
+                or not isinstance(id[0], DatasetUrn)
+                or not isinstance(id[1], (str, AssertionUrn))
+            ):
+                raise SdkUsageError(
+                    f"Invalid monitor identity input tuple: {id}. Expected a tuple of (dataset/entity urn, monitor id)."
+                )
+            if isinstance(id[1], str):
+                # If the second element is a string, we treat it as the monitor id.
+                # This allows for flexibility in how monitor ids are specified.
+                return MonitorUrn(entity=id[0], id=id[1])
+            elif isinstance(id[1], AssertionUrn):
+                # If the second element is an AssertionUrn, we use assertion urn as the monitor id so they both are linked.
+                return MonitorUrn(entity=id[0], id=id[1].urn())
+            else:
+                assert_never(id[1])
+        elif isinstance(id, models.MonitorKeyClass):
+            # This validation may look redundant but it is not.
+            # While MonitorKey PDL model ensures entity is a DatasetUrn,
+            # generated python code for entity's type in MonitorKeyClass is not so strict and allows for any string.
+            entity: DatasetUrn
+            if isinstance(id.entity, str):
+                try:
+                    entity = DatasetUrn.from_string(id.entity)
+                except Exception as e:
+                    raise SdkUsageError(
+                        f"Invalid monitor identity input key: {id}. Invalid dataset urn for entity."
+                    ) from e
+            else:
+                if not isinstance(id.entity, DatasetUrn):
+                    raise SdkUsageError(
+                        f"Invalid monitor identity input key: {id}. Expected a DatasetUrn for entity."
+                    )
+                entity = id.entity
+            return MonitorUrn(entity=entity, id=id.id)
+        elif isinstance(id, MonitorUrn):
+            return id
+        else:
+            assert_never(id)
+
+    def _ensure_info(self) -> models.MonitorInfoClass:
+        return self._setdefault_aspect(
+            models.MonitorInfoClass(
+                # TBC: just because we need to set them!
+                type=models.MonitorTypeClass.ASSERTION,
+                status=models.MonitorStatusClass(mode=models.MonitorModeClass.ACTIVE),
+            )
+        )
+
+    @property
+    def info(self) -> models.MonitorInfoClass:
+        """Get the monitor info.
+
+        Returns:
+            The monitor info.
+        """
+        return self._ensure_info()
+
+    def _set_info(self, info: MonitorInfoInputType) -> None:
+        """Set the monitor info.
+
+        Args:
+            info: The monitor info.
+        """
+        if isinstance(info, tuple):
+            if len(info) != 2:
+                raise SdkUsageError(
+                    f"Invalid monitor info input tuple: {info}. Expected a tuple of (monitor type, monitor status)."
+                )
+            type, mode_status = info
+            if isinstance(type, str) and type not in get_enum_options(
+                models.MonitorTypeClass
+            ):
+                raise SdkUsageError(
+                    f"Invalid monitor type: {type}. Expected {get_enum_options(models.MonitorTypeClass)}."
+                )
+            if isinstance(mode_status, str) and mode_status not in get_enum_options(
+                models.MonitorModeClass
+            ):
+                raise SdkUsageError(
+                    f"Invalid monitor status: {mode_status}. Expected {get_enum_options(models.MonitorModeClass)}."
+                )
+            self._ensure_info().type = type
+            self._ensure_info().status = models.MonitorStatusClass(mode=mode_status)
+        elif isinstance(info, models.MonitorInfoClass):
+            self._set_aspect(info)
+        else:
+            assert_never(info)
+
+    @property
+    def external_url(self) -> Optional[str]:
+        """Get the external URL of the monitor.
+
+        Returns:
+            The external URL if set, None otherwise.
+        """
+        return self._ensure_info().externalUrl
+
+    def set_external_url(self, external_url: str) -> None:
+        """Set the external URL of the monitor.
+
+        Args:
+            external_url: The external URL to set.
+        """
+        self._ensure_info().externalUrl = external_url
+
+    @property
+    def custom_properties(self) -> Dict[str, str]:
+        """Get the custom properties of the monitor.
+
+        Returns:
+            Dictionary of custom properties.
+        """
+        return self._ensure_info().customProperties
+
+    def set_custom_properties(self, custom_properties: Dict[str, str]) -> None:
+        """Set the custom properties of the monitor.
+
+        Args:
+            custom_properties: Dictionary of custom properties to set.
+        """
+        self._ensure_info().customProperties = custom_properties
+
+    @property
+    def sensitivity(self) -> Optional[models.AssertionMonitorSensitivityClass]:
+        """Get the sensitivity of the monitor.
+
+        Returns:
+            The sensitivity or None if not found.
+        """
+        return _get_nested_field_for_entity_with_default(
+            self, "info.assertionMonitor.settings.adjustmentSettings.sensitivity"
+        )
+
+    @property
+    def exclusion_windows(self) -> Optional[list[models.AssertionExclusionWindowClass]]:
+        """Get the exclusion windows of the monitor.
+
+        Returns:
+            The exclusion windows or None if not found.
+        """
+        return _get_nested_field_for_entity_with_default(
+            self, "info.assertionMonitor.settings.adjustmentSettings.exclusionWindows"
+        )
+
+    @property
+    def training_data_lookback_days(self) -> Optional[int]:
+        """Get the training data lookback days of the monitor.
+
+        Returns:
+            The training data lookback days or None if not found.
+        """
+        return _get_nested_field_for_entity_with_default(
+            self,
+            "info.assertionMonitor.settings.adjustmentSettings.trainingDataLookbackWindowDays",
+        )
+
+
+def _get_nested_field_for_entity_with_default(
+    entity: Entity,
+    field_path: str,
+    default: Any = None,
+) -> Any:
+    """
+    Get a nested field from an Entity object, and warn and return default if not found.
+
+    Args:
+        entity: The entity to get the nested field from.
+        field_path: The path to the nested field.
+        default: The default value to return if the field is not found.
+    """
+    fields = field_path.split(".")
+    current = entity
+    last_valid_path = entity.entity_type_name()
+
+    for field in fields:
+        try:
+            current = getattr(current, field)
+            last_valid_path = f"{last_valid_path}.{field}"
+        except AttributeError:
+            logger.warning(
+                f"{entity.entity_type_name().capitalize()} {entity.urn} does not have an `{last_valid_path}` field, defaulting to {default}"
+            )
+            return default
+
+    return current

acryl_datahub_cloud/sdk/entities/subscription.py

@@ -0,0 +1,100 @@
+from typing import (
+    Type,
+    Union,
+)
+
+from typing_extensions import (
+    Self,
+    assert_never,
+)
+
+from datahub.emitter.mcp_builder import DatahubKey
+from datahub.metadata import schema_classes as models
+from datahub.metadata.urns import (
+    SubscriptionUrn,
+    Urn,
+)
+from datahub.sdk.entity import Entity
+
+
+class SubscriptionKey(DatahubKey):
+    """
+    Key class for generating stable subscription identifiers.
+
+    The main goal is to have stable IDs that are deterministic based on the
+    entity and actor URNs, which helps prevent duplicate subscriptions during
+    eventual consistency scenarios when multiple subscription requests happen
+    in quick succession.
+
+    This implementation matches the behavior expected in the backend when a
+    subscription is created, ensuring consistent ID generation between the
+    Python SDK and Java backend services.
+    """
+
+    entity_urn: str
+    actor_urn: str
+
+
+class Subscription(Entity):
+    """
+    Subscription entity class.
+    """
+
+    __slots__ = ()
+
+    @classmethod
+    def get_urn_type(cls) -> Type[SubscriptionUrn]:
+        """Get the URN type for subscriptions.
+
+        Returns:
+            The SubscriptionUrn class.
+        """
+        return SubscriptionUrn
+
+    def __init__(
+        self,
+        # SubscriptionInfo
+        info: models.SubscriptionInfoClass,
+        # Identity
+        id: Union[str, SubscriptionUrn],
+    ):
+        """
+        Initialize the Subscription entity.
+        """
+        super().__init__(urn=Subscription._ensure_id(id))
+
+        self._set_info(info)
+
+    @classmethod
+    def _new_from_graph(cls, urn: Urn, current_aspects: models.AspectBag) -> Self:
+        assert isinstance(urn, SubscriptionUrn)
+        assert "subscriptionInfo" in current_aspects, "SubscriptionInfo is required"
+        subscription_info = current_aspects["subscriptionInfo"]
+        entity = cls(id=urn, info=subscription_info)
+        return entity._init_from_graph(current_aspects)
+
+    @property
+    def urn(self) -> SubscriptionUrn:
+        assert isinstance(self._urn, SubscriptionUrn)
+        return self._urn
+
+    @classmethod
+    def _ensure_id(cls, id: Union[str, SubscriptionUrn]) -> SubscriptionUrn:
+        if isinstance(id, str):
+            return SubscriptionUrn.from_string(id)
+        elif isinstance(id, SubscriptionUrn):
+            return id
+        else:
+            assert_never(id)
+
+    def _set_info(self, info: models.SubscriptionInfoClass) -> None:
+        self._set_aspect(info)
+
+    @property
+    def info(self) -> models.SubscriptionInfoClass:
+        """
+        Get the SubscriptionInfo aspect of the subscription.
+        """
+        subscription_info = self._get_aspect(models.SubscriptionInfoClass)
+        assert subscription_info is not None
+        return subscription_info

acryl_datahub_cloud/sdk/errors.py

@@ -0,0 +1,34 @@
+import json
+from typing import Collection, Mapping, Union
+
+
+class SDKUsageError(Exception):
+    def __init__(self, msg: str):
+        super().__init__(msg)
+
+
+class SDKUsageErrorWithExamples(SDKUsageError):
+    def __init__(
+        self,
+        msg: str,
+        examples: Union[
+            Mapping[str, Union[dict[str, str], str]], dict[str, Collection[str]]
+        ],
+    ):
+        super().__init__(msg)
+        self.msg = msg
+        self.examples = examples
+
+    def __str__(self) -> str:
+        examples_str = "\n\n".join(
+            f"Example - {key}:\n{json.dumps(value, indent=2) if isinstance(value, dict) else value}"
+            for key, value in self.examples.items()
+        )
+        return f"{self.msg}\n\n*** Examples of Accepted Values ***\n\n{examples_str}\n\n*** End of Examples ***"
+
+
+class SDKNotYetSupportedError(Exception):
+    def __init__(self, supported_item: str):
+        msg = f"This feature is not yet supported in the Python SDK: {supported_item}"
+        super().__init__(msg)
+        self.msg = msg

acryl_datahub_cloud/sdk/resolver_client.py

@@ -0,0 +1,42 @@
+from typing import List, Optional
+
+from datahub.metadata.urns import (
+    SubscriptionUrn,
+)
+from datahub.sdk.main_client import DataHubClient
+from datahub.sdk.resolver_client import ResolverClient as OSSResolverClient
+from datahub.sdk.search_filters import Filter, FilterDsl as F
+
+
+class ResolverClient(OSSResolverClient):
+    def __init__(self, client: DataHubClient):
+        super().__init__(client=client)
+
+    def subscription(
+        self,
+        *,
+        entity_urn: Optional[str] = None,
+        actor_urn: Optional[str] = None,
+        skip_cache: bool = False,
+    ) -> List[SubscriptionUrn]:
+        """Retrieve subscriptions for a given entity or actor, or both if both are given.
+        Args:
+            entity_urn: Optional URN of the entity to filter subscriptions by.
+            actor_urn: Optional URN of the actor to filter subscriptions by.
+        Returns:
+            List[SubscriptionUrn]: List of SubscriptionUrns matching the filters,
+            empty list if none.
+        """
+        filters: List[Filter] = [F.entity_type(SubscriptionUrn.ENTITY_TYPE)]
+
+        if entity_urn is not None:
+            filters.append(F.custom_filter("entityUrn", "EQUAL", [entity_urn]))
+
+        if actor_urn is not None:
+            filters.append(F.custom_filter("actorUrn", "EQUAL", [actor_urn]))
+
+        filter = F.and_(*filters)
+        subscriptions = list(
+            self._client.search.get_urns(filter=filter, skip_cache=skip_cache)
+        )
+        return [SubscriptionUrn.from_string(urn) for urn in subscriptions]