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

acryl_datahub_cloud/sdk/assertion/sql_assertion.py

@@ -0,0 +1,273 @@
+import logging
+from datetime import datetime
+from typing import Optional, Union
+
+from typing_extensions import Self
+
+from acryl_datahub_cloud.sdk.assertion.assertion_base import (
+    AssertionMode,
+    _AssertionPublic,
+    _HasSchedule,
+)
+from acryl_datahub_cloud.sdk.assertion_input.assertion_input import (
+    DEFAULT_DAILY_SCHEDULE,
+    DEFAULT_DETECTION_MECHANISM,
+    AssertionIncidentBehavior,
+    _DetectionMechanismTypes,
+)
+from acryl_datahub_cloud.sdk.assertion_input.sql_assertion_input import (
+    SqlAssertionCondition,
+    SqlAssertionCriteria,
+)
+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 SDKNotYetSupportedError
+from datahub.metadata import schema_classes as models
+from datahub.metadata.urns import AssertionUrn, CorpUserUrn, DatasetUrn, TagUrn
+
+logger = logging.getLogger(__name__)
+
+
+class SqlAssertion(_AssertionPublic, _HasSchedule):
+    """
+    A class that represents a SQL assertion.
+    """
+
+    def __init__(
+        self,
+        *,
+        urn: AssertionUrn,
+        dataset_urn: DatasetUrn,
+        display_name: str,
+        mode: AssertionMode,
+        statement: str,
+        criteria: Optional[SqlAssertionCriteria],
+        schedule: models.CronScheduleClass,
+        tags: list[TagUrn],
+        incident_behavior: list[AssertionIncidentBehavior],
+        created_by: Optional[CorpUserUrn] = None,
+        created_at: Union[datetime, None] = None,
+        updated_by: Optional[CorpUserUrn] = None,
+        updated_at: Optional[datetime] = None,
+    ):
+        """
+        Initialize a SQL assertion.
+
+        Note: Values can be accessed, but not set on the assertion object.
+        To update an assertion, use the `upsert_*` method.
+        Args:
+            urn: The urn of the assertion.
+            dataset_urn: The urn of the dataset that the assertion is for.
+            display_name: The display name of the assertion.
+            mode: The mode of the assertion (active, inactive).
+            statement: The SQL statement to be used for the assertion.
+            criteria: The criteria to be used for the assertion.
+                May be None only if backend data is corrupted - this should never happen under normal
+                circumstances, but the SDK handles it gracefully to aid debugging.
+            schedule: The schedule of the assertion.
+            tags: The tags applied to the assertion.
+            incident_behavior: Whether to raise or resolve an incident when the assertion fails / passes.
+            created_by: The urn of the user that created the assertion.
+            created_at: The timestamp of when the assertion was created.
+            updated_by: The urn of the user that updated the assertion.
+            updated_at: The timestamp of when the assertion was updated.
+        """
+        # Initialize the mixins first
+        _AssertionPublic.__init__(
+            self,
+            urn=urn,
+            dataset_urn=dataset_urn,
+            display_name=display_name,
+            mode=mode,
+            tags=tags,
+            incident_behavior=incident_behavior,
+            created_by=created_by,
+            created_at=created_at,
+            updated_by=updated_by,
+            updated_at=updated_at,
+        )
+        _HasSchedule.__init__(self, schedule=schedule)
+        # Then initialize the parent class
+        self._statement = statement
+        self._criteria = criteria
+
+    @property
+    def statement(self) -> str:
+        return self._statement
+
+    @property
+    def criteria(self) -> Optional[SqlAssertionCriteria]:
+        """Return the criteria for this assertion, or None if corrupted."""
+        return self._criteria
+
+    @property
+    def criteria_condition(self) -> Union[SqlAssertionCondition, str]:
+        if self._criteria is None:
+            raise ValueError(
+                "Cannot access criteria_condition on corrupted assertion with missing criteria"
+            )
+        return self._criteria.condition
+
+    @property
+    def criteria_parameters(
+        self,
+    ) -> Union[Union[float, int], tuple[Union[float, int], Union[float, int]]]:
+        if self._criteria is None:
+            raise ValueError(
+                "Cannot access criteria_parameters on corrupted assertion with missing criteria"
+            )
+        return self._criteria.parameters
+
+    @staticmethod
+    def _get_detection_mechanism(
+        assertion: Assertion,
+        monitor: Monitor,
+        default: Optional[_DetectionMechanismTypes] = DEFAULT_DETECTION_MECHANISM,
+    ) -> Optional[_DetectionMechanismTypes]:
+        """Sql assertions do not have a detection mechanism."""
+        return None
+
+    @staticmethod
+    def _get_statement(assertion: Assertion) -> str:
+        if assertion.info is None:
+            raise SDKNotYetSupportedError(
+                f"Assertion {assertion.urn} does not have a SQL assertion info, which is not supported"
+            )
+        if isinstance(assertion.info, models.SqlAssertionInfoClass):
+            return assertion.info.statement
+        else:
+            raise SDKNotYetSupportedError(
+                f"Assertion {assertion.urn} is not a SQL assertion"
+            )
+
+    @staticmethod
+    def _get_condition_from_model_assertion_info(
+        assertion_info: models.SqlAssertionInfoClass,
+    ) -> SqlAssertionCondition:
+        """Convert stored assertion info to condition enum."""
+        # Handle value-based conditions (no change type)
+        if str(assertion_info.type) == str(models.SqlAssertionTypeClass.METRIC):
+            value_conditions = {
+                str(
+                    models.AssertionStdOperatorClass.EQUAL_TO
+                ): SqlAssertionCondition.IS_EQUAL_TO,
+                str(
+                    models.AssertionStdOperatorClass.NOT_EQUAL_TO
+                ): SqlAssertionCondition.IS_NOT_EQUAL_TO,
+                str(
+                    models.AssertionStdOperatorClass.GREATER_THAN
+                ): SqlAssertionCondition.IS_GREATER_THAN,
+                str(
+                    models.AssertionStdOperatorClass.LESS_THAN
+                ): SqlAssertionCondition.IS_LESS_THAN,
+                str(
+                    models.AssertionStdOperatorClass.BETWEEN
+                ): SqlAssertionCondition.IS_WITHIN_A_RANGE,
+            }
+            if str(assertion_info.operator) in value_conditions:
+                return value_conditions[str(assertion_info.operator)]
+
+        # Handle growth-based conditions (with change type)
+        elif str(assertion_info.type) == str(
+            models.SqlAssertionTypeClass.METRIC_CHANGE
+        ):
+            assert assertion_info.changeType is not None, (
+                "changeType must be present for METRIC_CHANGE assertions"
+            )
+
+            growth_conditions = {
+                (
+                    str(models.AssertionStdOperatorClass.LESS_THAN_OR_EQUAL_TO),
+                    str(models.AssertionValueChangeTypeClass.ABSOLUTE),
+                ): SqlAssertionCondition.GROWS_AT_MOST_ABSOLUTE,
+                (
+                    str(models.AssertionStdOperatorClass.LESS_THAN_OR_EQUAL_TO),
+                    str(models.AssertionValueChangeTypeClass.PERCENTAGE),
+                ): SqlAssertionCondition.GROWS_AT_MOST_PERCENTAGE,
+                (
+                    str(models.AssertionStdOperatorClass.GREATER_THAN_OR_EQUAL_TO),
+                    str(models.AssertionValueChangeTypeClass.ABSOLUTE),
+                ): SqlAssertionCondition.GROWS_AT_LEAST_ABSOLUTE,
+                (
+                    str(models.AssertionStdOperatorClass.GREATER_THAN_OR_EQUAL_TO),
+                    str(models.AssertionValueChangeTypeClass.PERCENTAGE),
+                ): SqlAssertionCondition.GROWS_AT_LEAST_PERCENTAGE,
+                (
+                    str(models.AssertionStdOperatorClass.BETWEEN),
+                    str(models.AssertionValueChangeTypeClass.ABSOLUTE),
+                ): SqlAssertionCondition.GROWS_WITHIN_A_RANGE_ABSOLUTE,
+                (
+                    str(models.AssertionStdOperatorClass.BETWEEN),
+                    str(models.AssertionValueChangeTypeClass.PERCENTAGE),
+                ): SqlAssertionCondition.GROWS_WITHIN_A_RANGE_PERCENTAGE,
+            }
+
+            key = (str(assertion_info.operator), str(assertion_info.changeType))
+            if key in growth_conditions:
+                return growth_conditions[key]
+
+        raise ValueError(
+            f"Unsupported combination: type={assertion_info.type}, operator={assertion_info.operator}, changeType={assertion_info.changeType}"
+        )
+
+    @staticmethod
+    def _get_criteria(assertion: Assertion) -> Optional[SqlAssertionCriteria]:
+        if assertion.info is None:
+            raise SDKNotYetSupportedError(
+                f"Assertion {assertion.urn} does not have a SQL assertion info, which is not supported"
+            )
+        if isinstance(assertion.info, models.SqlAssertionInfoClass):
+            # Return None for corrupted assertions with missing parameters or operator
+            # This allows validation to happen later with a clearer error message
+            if assertion.info.parameters is None or assertion.info.operator is None:
+                return None
+
+            parameters: Union[float, tuple[float, float]]
+            if assertion.info.parameters.value is not None:
+                parameters = float(assertion.info.parameters.value.value)
+            elif (
+                assertion.info.parameters.maxValue is not None
+                and assertion.info.parameters.minValue is not None
+            ):
+                # min and max values are in the order of min, max
+                parameters = (
+                    float(assertion.info.parameters.minValue.value),
+                    float(assertion.info.parameters.maxValue.value),
+                )
+            else:
+                # Parameters field exists but has no valid values - corrupted data
+                return None
+
+            condition = SqlAssertion._get_condition_from_model_assertion_info(
+                assertion.info
+            )
+
+            return SqlAssertionCriteria(
+                condition=condition,
+                parameters=parameters,
+            )
+        else:
+            raise SDKNotYetSupportedError(
+                f"Assertion {assertion.urn} is not a SQL assertion"
+            )
+
+    @classmethod
+    def _from_entities(cls, assertion: Assertion, monitor: Monitor) -> Self:
+        """
+        Create a SQL assertion from the assertion and monitor entities.
+        """
+        return cls(
+            urn=assertion.urn,
+            dataset_urn=assertion.dataset,
+            display_name=assertion.description or "",
+            mode=cls._get_mode(monitor),
+            statement=cls._get_statement(assertion),
+            criteria=cls._get_criteria(assertion),
+            schedule=cls._get_schedule(monitor, default=DEFAULT_DAILY_SCHEDULE),
+            tags=cls._get_tags(assertion),
+            incident_behavior=cls._get_incident_behavior(assertion),
+            created_by=cls._get_created_by(assertion),
+            created_at=cls._get_created_at(assertion),
+            updated_by=cls._get_updated_by(assertion),
+            updated_at=cls._get_updated_at(assertion),
+        )

acryl_datahub_cloud/sdk/assertion/types.py

@@ -0,0 +1,20 @@
+from typing import Union
+
+from acryl_datahub_cloud.sdk.assertion.assertion_base import (
+    FreshnessAssertion,
+    SmartFreshnessAssertion,
+    SmartVolumeAssertion,
+    SqlAssertion,
+)
+from acryl_datahub_cloud.sdk.assertion.smart_column_metric_assertion import (
+    SmartColumnMetricAssertion,
+)
+
+AssertionTypes = Union[
+    SmartFreshnessAssertion,
+    SmartVolumeAssertion,
+    FreshnessAssertion,
+    SmartColumnMetricAssertion,
+    SqlAssertion,
+    # TODO: Add other assertion types here as we add them.
+]

acryl_datahub_cloud/sdk/assertion/volume_assertion.py

@@ -0,0 +1,156 @@
+import logging
+from datetime import datetime
+from typing import Optional, Union
+
+from typing_extensions import Self
+
+from acryl_datahub_cloud.sdk.assertion.assertion_base import (
+    AssertionMode,
+    _AssertionPublic,
+    _HasSchedule,
+)
+from acryl_datahub_cloud.sdk.assertion_input.assertion_input import (
+    DEFAULT_DETECTION_MECHANISM,
+    AssertionIncidentBehavior,
+    DetectionMechanism,
+    _DetectionMechanismTypes,
+)
+from acryl_datahub_cloud.sdk.assertion_input.volume_assertion_input import (
+    VolumeAssertionCriteria,
+)
+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 SDKNotYetSupportedError
+from datahub.metadata import schema_classes as models
+from datahub.metadata.urns import AssertionUrn, CorpUserUrn, DatasetUrn, TagUrn
+
+logger = logging.getLogger(__name__)
+
+
+class VolumeAssertion(_HasSchedule, _AssertionPublic):
+    """
+    A class that represents a volume assertion.
+    """
+
+    def __init__(
+        self,
+        *,
+        urn: AssertionUrn,
+        dataset_urn: DatasetUrn,
+        display_name: str,
+        mode: AssertionMode,
+        schedule: models.CronScheduleClass,
+        criteria: VolumeAssertionCriteria,
+        tags: list[TagUrn],
+        incident_behavior: list[AssertionIncidentBehavior],
+        detection_mechanism: Optional[
+            _DetectionMechanismTypes
+        ] = DEFAULT_DETECTION_MECHANISM,
+        created_by: Optional[CorpUserUrn] = None,
+        created_at: Union[datetime, None] = None,
+        updated_by: Optional[CorpUserUrn] = None,
+        updated_at: Optional[datetime] = None,
+    ):
+        """
+        Initialize a volume assertion.
+
+        Note: Values can be accessed, but not set on the assertion object.
+        To update an assertion, use the `upsert_*` method.
+        Args:
+            urn: The urn of the assertion.
+            dataset_urn: The urn of the dataset that the assertion is for.
+            display_name: The display name of the assertion.
+            mode: The mode of the assertion (active, inactive).
+            schedule: The schedule of the assertion.
+            criteria: The volume assertion criteria.
+            tags: The tags applied to the assertion.
+            incident_behavior: Whether to raise or resolve an incident when the assertion fails / passes.
+            detection_mechanism: The detection mechanism of the assertion.
+            created_by: The urn of the user that created the assertion.
+            created_at: The timestamp of when the assertion was created.
+            updated_by: The urn of the user that updated the assertion.
+            updated_at: The timestamp of when the assertion was updated.
+        """
+        _HasSchedule.__init__(self, schedule=schedule)
+        _AssertionPublic.__init__(
+            self,
+            urn=urn,
+            dataset_urn=dataset_urn,
+            display_name=display_name,
+            mode=mode,
+            incident_behavior=incident_behavior,
+            detection_mechanism=detection_mechanism,
+            created_by=created_by,
+            created_at=created_at,
+            updated_by=updated_by,
+            updated_at=updated_at,
+            tags=tags,
+        )
+        self._criteria = criteria
+
+    @property
+    def criteria(self) -> VolumeAssertionCriteria:
+        return self._criteria
+
+    @staticmethod
+    def _get_volume_definition(
+        assertion: Assertion,
+    ) -> VolumeAssertionCriteria:
+        """Get volume assertion definition from a DataHub assertion entity."""
+        return VolumeAssertionCriteria.from_assertion(assertion)
+
+    @staticmethod
+    def _get_detection_mechanism(
+        assertion: Assertion,
+        monitor: Monitor,
+        default: Optional[_DetectionMechanismTypes] = DEFAULT_DETECTION_MECHANISM,
+    ) -> Optional[_DetectionMechanismTypes]:
+        """Get the detection mechanism for volume assertions."""
+        parameters = _AssertionPublic._get_validated_detection_context(
+            monitor,
+            assertion,
+            models.AssertionEvaluationParametersTypeClass.DATASET_VOLUME,
+            models.VolumeAssertionInfoClass,
+            default,
+        )
+        if parameters is None:
+            return default
+        if parameters.datasetVolumeParameters is None:
+            logger.warning(
+                f"Monitor does not have datasetVolumeParameters, defaulting detection mechanism to {DEFAULT_DETECTION_MECHANISM}"
+            )
+            if default is None:
+                return DEFAULT_DETECTION_MECHANISM
+            else:
+                return default
+        source_type = parameters.datasetVolumeParameters.sourceType
+        if source_type == models.DatasetVolumeSourceTypeClass.INFORMATION_SCHEMA:
+            return DetectionMechanism.INFORMATION_SCHEMA
+        elif source_type == models.DatasetVolumeSourceTypeClass.QUERY:
+            additional_filter = _AssertionPublic._get_additional_filter(assertion)
+            return DetectionMechanism.QUERY(additional_filter=additional_filter)
+        elif source_type == models.DatasetVolumeSourceTypeClass.DATAHUB_DATASET_PROFILE:
+            return DetectionMechanism.DATASET_PROFILE
+        else:
+            raise SDKNotYetSupportedError(f"DatasetVolumeSourceType {source_type}")
+
+    @classmethod
+    def _from_entities(cls, assertion: Assertion, monitor: Monitor) -> Self:
+        """
+        Create a volume assertion from the assertion and monitor entities.
+        """
+        return cls(
+            urn=assertion.urn,
+            dataset_urn=assertion.dataset,
+            display_name=assertion.description or "",
+            mode=cls._get_mode(monitor),
+            schedule=cls._get_schedule(monitor),
+            criteria=cls._get_volume_definition(assertion),
+            incident_behavior=cls._get_incident_behavior(assertion),
+            detection_mechanism=cls._get_detection_mechanism(assertion, monitor),
+            created_by=cls._get_created_by(assertion),
+            created_at=cls._get_created_at(assertion),
+            updated_by=cls._get_updated_by(assertion),
+            updated_at=cls._get_updated_at(assertion),
+            tags=cls._get_tags(assertion),
+        )