acryl-datahub-cloud 0.3.12rc6__py3-none-any.whl → 0.3.12rc7__py3-none-any.whl

acryl_datahub_cloud/sdk/assertion/assertion_base.py

@@ -39,6 +39,10 @@ from acryl_datahub_cloud.sdk.assertion_input.smart_column_metric_assertion_input
 from acryl_datahub_cloud.sdk.assertion_input.sql_assertion_input import (
     SqlAssertionCriteria,
 )
+from acryl_datahub_cloud.sdk.assertion_input.volume_assertion_input import (
+    VolumeAssertionDefinition,
+    _VolumeAssertionDefinitionTypes,
+)
 from acryl_datahub_cloud.sdk.entities.assertion import Assertion
 from acryl_datahub_cloud.sdk.entities.monitor import (
     Monitor,
@@ -981,6 +985,135 @@ class SmartVolumeAssertion(_HasSchedule, _HasSmartFunctionality, _AssertionPubli
             raise SDKNotYetSupportedError(f"DatasetVolumeSourceType {source_type}")
 
 
+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,
+        definition: _VolumeAssertionDefinitionTypes,
+        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.
+            definition: The volume assertion definition (RowCountTotal or RowCountChange).
+            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._definition = definition
+
+    @property
+    def definition(self) -> _VolumeAssertionDefinitionTypes:
+        return self._definition
+
+    @staticmethod
+    def _get_volume_definition(
+        assertion: Assertion,
+    ) -> _VolumeAssertionDefinitionTypes:
+        """Get volume assertion definition from a DataHub assertion entity."""
+        return VolumeAssertionDefinition.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),
+            definition=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),
+        )
+
+
 class FreshnessAssertion(_HasSchedule, _AssertionPublic):
     """
     A class that represents a freshness assertion.
@@ -1095,30 +1228,6 @@ class FreshnessAssertion(_HasSchedule, _AssertionPublic):
                 f"Assertion {assertion.urn} is not a freshness assertion"
             )
 
-    @classmethod
-    def _from_entities(cls, assertion: Assertion, monitor: Monitor) -> Self:
-        """
-        Create a freshness 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),
-            freshness_schedule_check_type=cls._get_freshness_schedule_check_type(
-                assertion
-            ),
-            lookback_window=cls._get_lookback_window(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),
-        )
-
     @staticmethod
     def _get_detection_mechanism(
         assertion: Assertion,
@@ -1156,6 +1265,30 @@ class FreshnessAssertion(_HasSchedule, _AssertionPublic):
         else:
             raise SDKNotYetSupportedError(f"DatasetFreshnessSourceType {source_type}")
 
+    @classmethod
+    def _from_entities(cls, assertion: Assertion, monitor: Monitor) -> Self:
+        """
+        Create a freshness 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),
+            freshness_schedule_check_type=cls._get_freshness_schedule_check_type(
+                assertion
+            ),
+            lookback_window=cls._get_lookback_window(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),
+        )
+
 
 class SqlAssertion(_AssertionPublic, _HasSchedule):
     """