acryl-datahub-cloud 0.3.10rc4__py3-none-any.whl → 0.3.16.1rc0__py3-none-any.whl

acryl_datahub_cloud/sdk/assertion_client/freshness.py

@@ -0,0 +1,371 @@
+from __future__ import annotations
+
+import logging
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Optional, Union
+
+from acryl_datahub_cloud.sdk.assertion.assertion_base import (
+    AssertionMode,
+    FreshnessAssertion,
+)
+from acryl_datahub_cloud.sdk.assertion_client.helpers import (
+    DEFAULT_CREATED_BY,
+    _merge_field,
+    _print_experimental_warning,
+    retrieve_assertion_and_monitor_by_urn,
+)
+from acryl_datahub_cloud.sdk.assertion_input.assertion_input import (
+    AssertionIncidentBehaviorInputTypes,
+    DetectionMechanismInputTypes,
+    TimeWindowSizeInputTypes,
+)
+from acryl_datahub_cloud.sdk.assertion_input.freshness_assertion_input import (
+    FreshnessAssertionScheduleCheckType,
+    _FreshnessAssertionInput,
+)
+from acryl_datahub_cloud.sdk.entities.assertion import Assertion, TagsInputType
+from acryl_datahub_cloud.sdk.entities.monitor import Monitor
+from acryl_datahub_cloud.sdk.errors import SDKUsageError
+from datahub.metadata import schema_classes as models
+from datahub.metadata.urns import AssertionUrn, CorpUserUrn, DatasetUrn, MonitorUrn
+
+if TYPE_CHECKING:
+    from datahub.sdk.main_client import DataHubClient
+
+logger = logging.getLogger(__name__)
+
+
+class FreshnessAssertionClient:
+    """Client for managing freshness assertions."""
+
+    def __init__(self, client: "DataHubClient"):
+        self.client = client
+
+    def sync_freshness_assertion(
+        self,
+        *,
+        dataset_urn: Union[str, DatasetUrn],
+        urn: Optional[Union[str, AssertionUrn]] = None,
+        display_name: Optional[str] = None,
+        enabled: Optional[bool] = None,
+        detection_mechanism: DetectionMechanismInputTypes = None,
+        incident_behavior: Optional[AssertionIncidentBehaviorInputTypes] = None,
+        tags: Optional[TagsInputType] = None,
+        updated_by: Optional[Union[str, CorpUserUrn]] = None,
+        freshness_schedule_check_type: Optional[
+            Union[
+                str,
+                FreshnessAssertionScheduleCheckType,
+                models.FreshnessAssertionScheduleTypeClass,
+            ]
+        ] = None,
+        schedule: Optional[Union[str, models.CronScheduleClass]] = None,
+        lookback_window: Optional[TimeWindowSizeInputTypes] = None,
+    ) -> FreshnessAssertion:
+        _print_experimental_warning()
+        now_utc = datetime.now(timezone.utc)
+
+        if updated_by is None:
+            logger.warning(
+                f"updated_by is not set, using {DEFAULT_CREATED_BY} as a placeholder"
+            )
+            updated_by = DEFAULT_CREATED_BY
+
+        # 1. Retrieve and merge the assertion input with any existing assertion and monitor entities,
+        # or build a new assertion input if the assertion does not exist:
+        assertion_input = self._retrieve_and_merge_freshness_assertion_and_monitor(
+            dataset_urn=dataset_urn,
+            urn=urn,
+            display_name=display_name,
+            enabled=enabled,
+            detection_mechanism=detection_mechanism,
+            incident_behavior=incident_behavior,
+            tags=tags,
+            updated_by=updated_by,
+            now_utc=now_utc,
+            schedule=schedule,
+            freshness_schedule_check_type=freshness_schedule_check_type,
+            lookback_window=lookback_window,
+        )
+
+        # 2. Upsert the assertion and monitor entities:
+        assertion_entity, monitor_entity = (
+            assertion_input.to_assertion_and_monitor_entities()
+        )
+        # If assertion upsert fails, we won't try to upsert the monitor
+        self.client.entities.upsert(assertion_entity)
+        # TODO: Wrap monitor upsert in a try-except and delete the assertion if monitor upsert fails (once delete is implemented https://linear.app/acryl-data/issue/OBS-1350/add-delete-method-to-entity-clientpy)
+        # try:
+        self.client.entities.upsert(monitor_entity)
+        # except Exception as e:
+        #     logger.error(f"Error upserting monitor: {e}")
+        #     self.client.entities.delete(assertion_entity)
+        #     raise e
+
+        return FreshnessAssertion._from_entities(assertion_entity, monitor_entity)
+
+    def _retrieve_and_merge_freshness_assertion_and_monitor(
+        self,
+        dataset_urn: Union[str, DatasetUrn],
+        urn: Optional[Union[str, AssertionUrn]],
+        display_name: Optional[str],
+        enabled: Optional[bool],
+        detection_mechanism: DetectionMechanismInputTypes,
+        incident_behavior: Optional[AssertionIncidentBehaviorInputTypes],
+        tags: Optional[TagsInputType],
+        updated_by: Union[str, CorpUserUrn],
+        now_utc: datetime,
+        schedule: Optional[Union[str, models.CronScheduleClass]],
+        freshness_schedule_check_type: Optional[
+            Union[str, models.FreshnessAssertionScheduleTypeClass]
+        ] = None,
+        lookback_window: Optional[TimeWindowSizeInputTypes] = None,
+    ) -> _FreshnessAssertionInput:
+        # 1. If urn is not provided, build and return a new assertion input directly
+        if urn is None:
+            logger.info("URN is not set, building a new assertion input")
+            return _FreshnessAssertionInput(
+                urn=None,
+                entity_client=self.client.entities,
+                dataset_urn=dataset_urn,
+                display_name=display_name,
+                enabled=enabled if enabled is not None else True,
+                detection_mechanism=detection_mechanism,
+                freshness_schedule_check_type=freshness_schedule_check_type,
+                lookback_window=lookback_window,
+                incident_behavior=incident_behavior,
+                tags=tags,
+                created_by=updated_by,
+                created_at=now_utc,
+                updated_by=updated_by,
+                updated_at=now_utc,
+                schedule=schedule,
+            )
+
+        # 2. Build initial assertion input for validation
+        assertion_input = _FreshnessAssertionInput(
+            urn=urn,
+            entity_client=self.client.entities,
+            dataset_urn=dataset_urn,
+            display_name=display_name,
+            detection_mechanism=detection_mechanism,
+            freshness_schedule_check_type=freshness_schedule_check_type,
+            lookback_window=lookback_window,
+            incident_behavior=incident_behavior,
+            tags=tags,
+            created_by=updated_by,  # This will be overridden by the actual created_by
+            created_at=now_utc,  # This will be overridden by the actual created_at
+            updated_by=updated_by,
+            updated_at=now_utc,
+            schedule=schedule,
+        )
+
+        # 3. Retrieve any existing assertion and monitor entities:
+        maybe_assertion_entity, monitor_urn, maybe_monitor_entity = (
+            self._retrieve_assertion_and_monitor(assertion_input)
+        )
+
+        # 4.1 If the assertion and monitor entities exist, create an assertion object from them:
+        if maybe_assertion_entity and maybe_monitor_entity:
+            existing_assertion = FreshnessAssertion._from_entities(
+                maybe_assertion_entity, maybe_monitor_entity
+            )
+        # 4.2 If the assertion exists but the monitor does not, create a placeholder monitor entity to be able to create the assertion:
+        elif maybe_assertion_entity and not maybe_monitor_entity:
+            monitor_mode = (
+                "ACTIVE" if enabled else "INACTIVE" if enabled is not None else "ACTIVE"
+            )
+            existing_assertion = FreshnessAssertion._from_entities(
+                maybe_assertion_entity,
+                Monitor(id=monitor_urn, info=("ASSERTION", monitor_mode)),
+            )
+        # 4.3 If the assertion does not exist, build and return a new assertion input:
+        elif not maybe_assertion_entity:
+            logger.info(
+                f"No existing assertion entity found for assertion urn {urn}, building a new assertion input"
+            )
+            return _FreshnessAssertionInput(
+                urn=urn,
+                entity_client=self.client.entities,
+                dataset_urn=dataset_urn,
+                display_name=display_name,
+                enabled=enabled if enabled is not None else True,
+                detection_mechanism=detection_mechanism,
+                freshness_schedule_check_type=freshness_schedule_check_type,
+                lookback_window=lookback_window,
+                incident_behavior=incident_behavior,
+                tags=tags,
+                created_by=updated_by,
+                created_at=now_utc,
+                updated_by=updated_by,
+                updated_at=now_utc,
+                schedule=schedule,
+            )
+
+        # 5. Check for any issues e.g. different dataset urns
+        if (
+            existing_assertion
+            and hasattr(existing_assertion, "dataset_urn")
+            and existing_assertion.dataset_urn != assertion_input.dataset_urn
+        ):
+            raise SDKUsageError(
+                f"Dataset URN mismatch, existing assertion: {existing_assertion.dataset_urn} != new assertion: {dataset_urn}"
+            )
+
+        # 6. Merge the existing assertion with the validated input:
+        merged_assertion_input = self._merge_freshness_input(
+            dataset_urn=dataset_urn,
+            urn=urn,
+            display_name=display_name,
+            enabled=enabled,
+            detection_mechanism=detection_mechanism,
+            incident_behavior=incident_behavior,
+            tags=tags,
+            now_utc=now_utc,
+            assertion_input=assertion_input,
+            maybe_assertion_entity=maybe_assertion_entity,
+            maybe_monitor_entity=maybe_monitor_entity,
+            existing_assertion=existing_assertion,
+            schedule=schedule,
+            freshness_schedule_check_type=freshness_schedule_check_type,
+            lookback_window=lookback_window,
+        )
+
+        return merged_assertion_input
+
+    def _retrieve_assertion_and_monitor(
+        self,
+        assertion_input: _FreshnessAssertionInput,
+    ) -> tuple[Optional[Assertion], MonitorUrn, Optional[Monitor]]:
+        """Retrieve the assertion and monitor entities from the DataHub instance.
+
+        Args:
+            assertion_input: The validated input to the function.
+
+        Returns:
+            The assertion and monitor entities.
+        """
+        assert assertion_input.urn is not None, "URN is required"
+        return retrieve_assertion_and_monitor_by_urn(
+            self.client, assertion_input.urn, assertion_input.dataset_urn
+        )
+
+    def _merge_freshness_input(
+        self,
+        dataset_urn: Union[str, DatasetUrn],
+        urn: Union[str, AssertionUrn],
+        display_name: Optional[str],
+        enabled: Optional[bool],
+        detection_mechanism: DetectionMechanismInputTypes,
+        incident_behavior: Optional[AssertionIncidentBehaviorInputTypes],
+        tags: Optional[TagsInputType],
+        now_utc: datetime,
+        assertion_input: _FreshnessAssertionInput,
+        maybe_assertion_entity: Optional[Assertion],
+        maybe_monitor_entity: Optional[Monitor],
+        existing_assertion: FreshnessAssertion,
+        schedule: Optional[Union[str, models.CronScheduleClass]],
+        freshness_schedule_check_type: Optional[
+            Union[str, models.FreshnessAssertionScheduleTypeClass]
+        ] = None,
+        lookback_window: Optional[TimeWindowSizeInputTypes] = None,
+    ) -> _FreshnessAssertionInput:
+        """Merge the input with the existing assertion and monitor entities.
+
+        Args:
+            dataset_urn: The urn of the dataset to be monitored.
+            urn: The urn of the assertion.
+            display_name: The display name of the assertion.
+            enabled: Whether the assertion is enabled.
+            incident_behavior: The incident behavior to be applied to the assertion.
+            tags: The tags to be applied to the assertion.
+            now_utc: The current UTC time from when the function is called.
+            assertion_input: The validated input to the function.
+            maybe_assertion_entity: The existing assertion entity from the DataHub instance.
+            maybe_monitor_entity: The existing monitor entity from the DataHub instance.
+            existing_assertion: The existing assertion from the DataHub instance.
+            schedule: The schedule to be applied to the assertion.
+            freshness_schedule_check_type: The freshness schedule check type to be applied to the assertion.
+            lookback_window: The lookback window to be applied to the assertion.
+
+        Returns:
+            The merged assertion input.
+        """
+        merged_assertion_input = _FreshnessAssertionInput(
+            urn=urn,
+            entity_client=self.client.entities,
+            dataset_urn=dataset_urn,
+            display_name=_merge_field(
+                display_name,
+                "display_name",
+                assertion_input,
+                existing_assertion,
+                maybe_assertion_entity.description if maybe_assertion_entity else None,
+            ),
+            enabled=_merge_field(
+                enabled,
+                "enabled",
+                assertion_input,
+                existing_assertion,
+                existing_assertion.mode == AssertionMode.ACTIVE
+                if existing_assertion
+                else None,
+            ),
+            schedule=_merge_field(
+                schedule,
+                "schedule",
+                assertion_input,
+                existing_assertion,
+                existing_assertion.schedule if existing_assertion else None,
+            ),
+            freshness_schedule_check_type=_merge_field(
+                freshness_schedule_check_type,
+                "freshness_schedule_check_type",
+                assertion_input,
+                existing_assertion,
+                existing_assertion._freshness_schedule_check_type
+                if existing_assertion
+                else None,
+            ),
+            lookback_window=_merge_field(
+                lookback_window,
+                "lookback_window",
+                assertion_input,
+                existing_assertion,
+                existing_assertion.lookback_window if existing_assertion else None,
+            ),
+            detection_mechanism=_merge_field(
+                detection_mechanism,
+                "detection_mechanism",
+                assertion_input,
+                existing_assertion,
+                FreshnessAssertion._get_detection_mechanism(
+                    maybe_assertion_entity, maybe_monitor_entity, default=None
+                )
+                if maybe_assertion_entity and maybe_monitor_entity
+                else None,
+            ),
+            incident_behavior=_merge_field(
+                incident_behavior,
+                "incident_behavior",
+                assertion_input,
+                existing_assertion,
+                FreshnessAssertion._get_incident_behavior(maybe_assertion_entity)
+                if maybe_assertion_entity
+                else None,
+            ),
+            tags=_merge_field(
+                tags,
+                "tags",
+                assertion_input,
+                existing_assertion,
+                maybe_assertion_entity.tags if maybe_assertion_entity else None,
+            ),
+            created_by=existing_assertion.created_by
+            or DEFAULT_CREATED_BY,  # Override with the existing assertion's created_by or the default created_by if not set
+            created_at=existing_assertion.created_at
+            or now_utc,  # Override with the existing assertion's created_at or now if not set
+            updated_by=assertion_input.updated_by,  # Override with the input's updated_by
+            updated_at=assertion_input.updated_at,  # Override with the input's updated_at (now)
+        )
+        return merged_assertion_input

acryl_datahub_cloud/sdk/assertion_client/helpers.py

@@ -0,0 +1,166 @@
+"""Helper utilities for assertion clients."""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any, Optional, Union
+
+from acryl_datahub_cloud.sdk.assertion.assertion_base import _AssertionPublic
+from acryl_datahub_cloud.sdk.assertion_input.assertion_input import _AssertionInput
+from acryl_datahub_cloud.sdk.entities.assertion import Assertion
+from acryl_datahub_cloud.sdk.entities.monitor import Monitor
+from acryl_datahub_cloud.sdk.errors import SDKUsageError
+from datahub.errors import ItemNotFoundError
+from datahub.metadata.urns import AssertionUrn, CorpUserUrn, DatasetUrn, MonitorUrn
+from datahub.sdk.search_filters import FilterDsl
+
+if TYPE_CHECKING:
+    from datahub.sdk.main_client import DataHubClient
+
+logger = logging.getLogger(__name__)
+
+# TODO: Replace __datahub_system with the actual datahub system user https://linear.app/acryl-data/issue/OBS-1351/auditstamp-actor-hydration-pattern-for-sdk-calls
+DEFAULT_CREATED_BY = CorpUserUrn.from_string("urn:li:corpuser:__datahub_system")
+
+
+def _merge_field(
+    input_field_value: Any,
+    input_field_name: str,
+    validated_assertion_input: _AssertionInput,
+    validated_existing_assertion: _AssertionPublic,
+    existing_entity_value: Optional[Any] = None,  # TODO: Can we do better than Any?
+) -> Any:
+    """Merge the input field value with any existing entity value or default value.
+
+    The merge logic is as follows:
+    - If the input is None, use the existing value
+    - If the input is not None, use the input value
+    - If the input is an empty list or empty string, still use the input value (falsy values can be used to unset fields)
+    - If the input is a non-empty list or non-empty string, use the input value
+    - If the input is None and the existing value is None, use the default value from _AssertionInput
+
+    Args:
+        input_field_value: The value of the field in the input e.g. passed to the function.
+        input_field_name: The name of the field in the input.
+        validated_assertion_input: The *validated* input to the function.
+        validated_existing_assertion: The *validated* existing assertion from the DataHub instance.
+        existing_entity_value: The value of the field in the existing entity from the DataHub instance, directly retrieved from the entity.
+
+        Returns:
+            The merged value of the field.
+
+    """
+    if input_field_value is None:  # Input value default
+        if existing_entity_value is not None:  # Existing entity value set
+            return existing_entity_value
+        elif (
+            getattr(validated_existing_assertion, input_field_name) is None
+        ):  # Validated existing value not set
+            return getattr(validated_assertion_input, input_field_name)
+        else:  # Validated existing value set
+            return getattr(validated_existing_assertion, input_field_name)
+    else:  # Input value set
+        return input_field_value
+
+
+def _print_experimental_warning() -> None:
+    print(
+        "Warning: The assertions client is experimental and under heavy development. Expect breaking changes."
+    )
+
+
+def _validate_required_field(
+    field_value: Optional[Any], field_name: str, context: str
+) -> None:
+    """Validate that a required field is not None and raise SDKUsageError if it is."""
+    if field_value is None:
+        raise SDKUsageError(f"{field_name} is required {context}")
+
+
+def resolve_updated_by(
+    updated_by: Optional[Union[str, CorpUserUrn]],
+) -> Union[str, CorpUserUrn]:
+    """Resolve the updated_by field, using DEFAULT_CREATED_BY if not provided.
+
+    Args:
+        updated_by: The updated_by value from the user input, or None.
+
+    Returns:
+        The resolved updated_by value.
+    """
+    if updated_by is None:
+        logger.warning(
+            f"updated_by is not set, using {DEFAULT_CREATED_BY} as a placeholder"
+        )
+        return DEFAULT_CREATED_BY
+    return updated_by
+
+
+def retrieve_assertion_and_monitor_by_urn(
+    client: "DataHubClient",
+    assertion_urn: Union[str, AssertionUrn],
+    dataset_urn: Union[str, DatasetUrn],
+) -> tuple[Optional[Assertion], MonitorUrn, Optional[Monitor]]:
+    """Retrieve the assertion and monitor entities from the DataHub instance by URN.
+
+    Args:
+        client: The DataHub client.
+        assertion_urn: The assertion URN.
+        dataset_urn: The dataset URN.
+
+    Returns:
+        A tuple of (assertion entity or None, monitor URN, monitor entity or None).
+    """
+    urn: AssertionUrn = (
+        assertion_urn
+        if isinstance(assertion_urn, AssertionUrn)
+        else AssertionUrn.from_string(assertion_urn)
+    )
+    _dataset_urn: DatasetUrn = (
+        dataset_urn
+        if isinstance(dataset_urn, DatasetUrn)
+        else DatasetUrn.from_string(dataset_urn)
+    )
+
+    # Get assertion entity
+    maybe_assertion_entity: Optional[Assertion] = None
+    try:
+        entity = client.entities.get(urn)
+        if entity is not None:
+            assert isinstance(entity, Assertion)
+            maybe_assertion_entity = entity
+    except ItemNotFoundError:
+        pass
+
+    # Get monitor entity by searching for monitors where assertionUrn equals the assertion urn
+    monitor_urn: Optional[MonitorUrn] = None
+    maybe_monitor_entity: Optional[Monitor] = None
+    try:
+        # Search for monitor entities with assertionUrn matching the assertion urn
+        monitor_filter = FilterDsl.and_(
+            FilterDsl.entity_type("monitor"),
+            FilterDsl.custom_filter("assertionUrn", "EQUAL", [str(urn)]),
+        )
+        monitor_urns = list(client.search.get_urns(filter=monitor_filter))
+
+        if monitor_urns:
+            # Log if there are multiple monitors found, because this is unexpected
+            if len(monitor_urns) > 1:
+                logger.warning(
+                    f"Multiple monitors found for assertion {urn}, which should never happen: {monitor_urns}"
+                )
+            # Use the first matching monitor
+            monitor_urn = MonitorUrn.from_string(str(monitor_urns[0]))
+            entity = client.entities.get(monitor_urn)
+            if entity is not None:
+                assert isinstance(entity, Monitor)
+                maybe_monitor_entity = entity
+    except ItemNotFoundError as e:
+        logger.debug(f"Could not find monitor for assertion {urn}: {e}")
+        pass
+
+    # If no monitor found via search, fall back to creating a new monitor with the dataset urn and assertion urn
+    if monitor_urn is None:
+        monitor_urn = Monitor._ensure_id(id=(_dataset_urn, urn))
+
+    return maybe_assertion_entity, monitor_urn, maybe_monitor_entity