acryl-datahub-cloud 0.3.12.1rc3__py3-none-any.whl → 0.3.12.2__py3-none-any.whl

acryl_datahub_cloud/sdk/assertion_input/column_metric_constants.py

@@ -0,0 +1,191 @@
+"""
+Shared constants for column metric assertions (both smart and non-smart).
+
+This module contains constants that are used by both smart and non-smart column metric assertions
+to ensure consistency and avoid duplication.
+"""
+
+from enum import Enum
+from typing import Union
+
+from datahub.metadata import schema_classes as models
+
+# Keep this in sync with the frontend in getEligibleFieldColumns
+# datahub-web-react/src/app/entityV2/shared/tabs/Dataset/Validations/assertion/builder/steps/field/utils.ts
+ALLOWED_COLUMN_TYPES_FOR_COLUMN_METRIC_ASSERTION = [
+    models.StringTypeClass(),
+    models.NumberTypeClass(),
+    models.BooleanTypeClass(),
+    models.DateTypeClass(),
+    models.TimeTypeClass(),
+    models.NullTypeClass(),
+]
+
+# Keep this in sync with FIELD_VALUES_OPERATOR_CONFIG in the frontend
+# datahub-web-react/src/app/entityV2/shared/tabs/Dataset/Validations/assertion/builder/steps/field/utils.ts
+FIELD_VALUES_OPERATOR_CONFIG = {
+    "STRING": [
+        models.AssertionStdOperatorClass.NULL,
+        models.AssertionStdOperatorClass.NOT_NULL,
+        models.AssertionStdOperatorClass.EQUAL_TO,
+        models.AssertionStdOperatorClass.IN,
+        models.AssertionStdOperatorClass.GREATER_THAN_OR_EQUAL_TO,
+        models.AssertionStdOperatorClass.REGEX_MATCH,
+        models.AssertionStdOperatorClass.GREATER_THAN,
+        models.AssertionStdOperatorClass.LESS_THAN,
+        models.AssertionStdOperatorClass.BETWEEN,
+    ],
+    "NUMBER": [
+        models.AssertionStdOperatorClass.GREATER_THAN,
+        models.AssertionStdOperatorClass.LESS_THAN,
+        models.AssertionStdOperatorClass.BETWEEN,
+        models.AssertionStdOperatorClass.NULL,
+        models.AssertionStdOperatorClass.NOT_NULL,
+        models.AssertionStdOperatorClass.EQUAL_TO,
+        models.AssertionStdOperatorClass.IN,
+        models.AssertionStdOperatorClass.GREATER_THAN_OR_EQUAL_TO,
+        models.AssertionStdOperatorClass.NOT_EQUAL_TO,
+    ],
+    "BOOLEAN": [
+        models.AssertionStdOperatorClass.IS_TRUE,
+        models.AssertionStdOperatorClass.IS_FALSE,
+        models.AssertionStdOperatorClass.NULL,
+        models.AssertionStdOperatorClass.NOT_NULL,
+    ],
+    "DATE": [
+        models.AssertionStdOperatorClass.NULL,
+        models.AssertionStdOperatorClass.NOT_NULL,
+    ],
+    "TIME": [
+        models.AssertionStdOperatorClass.NULL,
+        models.AssertionStdOperatorClass.NOT_NULL,
+    ],
+    "NULL": [
+        models.AssertionStdOperatorClass.NULL,
+        models.AssertionStdOperatorClass.NOT_NULL,
+    ],
+}
+
+# Keep this in sync with FIELD_METRIC_TYPE_CONFIG in the frontend
+# datahub-web-react/src/app/entityV2/shared/tabs/Dataset/Validations/assertion/builder/steps/field/utils.ts
+FIELD_METRIC_TYPE_CONFIG = {
+    "STRING": [
+        models.FieldMetricTypeClass.NULL_COUNT,
+        models.FieldMetricTypeClass.NULL_PERCENTAGE,
+        models.FieldMetricTypeClass.UNIQUE_COUNT,
+        models.FieldMetricTypeClass.UNIQUE_PERCENTAGE,
+        models.FieldMetricTypeClass.MAX_LENGTH,
+        models.FieldMetricTypeClass.MIN_LENGTH,
+        models.FieldMetricTypeClass.EMPTY_COUNT,
+        models.FieldMetricTypeClass.EMPTY_PERCENTAGE,
+    ],
+    "NUMBER": [
+        models.FieldMetricTypeClass.NULL_COUNT,
+        models.FieldMetricTypeClass.NULL_PERCENTAGE,
+        models.FieldMetricTypeClass.UNIQUE_COUNT,
+        models.FieldMetricTypeClass.UNIQUE_PERCENTAGE,
+        models.FieldMetricTypeClass.MAX,
+        models.FieldMetricTypeClass.MIN,
+        models.FieldMetricTypeClass.MEAN,
+        models.FieldMetricTypeClass.MEDIAN,
+        models.FieldMetricTypeClass.STDDEV,
+        models.FieldMetricTypeClass.NEGATIVE_COUNT,
+        models.FieldMetricTypeClass.NEGATIVE_PERCENTAGE,
+        models.FieldMetricTypeClass.ZERO_COUNT,
+        models.FieldMetricTypeClass.ZERO_PERCENTAGE,
+    ],
+    "BOOLEAN": [
+        models.FieldMetricTypeClass.NULL_COUNT,
+        models.FieldMetricTypeClass.NULL_PERCENTAGE,
+        models.FieldMetricTypeClass.UNIQUE_COUNT,
+        models.FieldMetricTypeClass.UNIQUE_PERCENTAGE,
+    ],
+    "DATE": [
+        models.FieldMetricTypeClass.NULL_COUNT,
+        models.FieldMetricTypeClass.NULL_PERCENTAGE,
+        models.FieldMetricTypeClass.UNIQUE_COUNT,
+        models.FieldMetricTypeClass.UNIQUE_PERCENTAGE,
+    ],
+    "TIME": [
+        models.FieldMetricTypeClass.NULL_COUNT,
+        models.FieldMetricTypeClass.NULL_PERCENTAGE,
+        models.FieldMetricTypeClass.UNIQUE_COUNT,
+        models.FieldMetricTypeClass.UNIQUE_PERCENTAGE,
+    ],
+    "NULL": [
+        models.FieldMetricTypeClass.NULL_COUNT,
+        models.FieldMetricTypeClass.NULL_PERCENTAGE,
+        models.FieldMetricTypeClass.UNIQUE_COUNT,
+        models.FieldMetricTypeClass.UNIQUE_PERCENTAGE,
+    ],
+}
+
+
+class MetricType(str, Enum):
+    """Enum for field metric types used in column metric assertions."""
+
+    NULL_COUNT = models.FieldMetricTypeClass.NULL_COUNT
+    NULL_PERCENTAGE = models.FieldMetricTypeClass.NULL_PERCENTAGE
+    UNIQUE_COUNT = models.FieldMetricTypeClass.UNIQUE_COUNT
+    UNIQUE_PERCENTAGE = models.FieldMetricTypeClass.UNIQUE_PERCENTAGE
+    MAX_LENGTH = models.FieldMetricTypeClass.MAX_LENGTH
+    MIN_LENGTH = models.FieldMetricTypeClass.MIN_LENGTH
+    EMPTY_COUNT = models.FieldMetricTypeClass.EMPTY_COUNT
+    EMPTY_PERCENTAGE = models.FieldMetricTypeClass.EMPTY_PERCENTAGE
+    MIN = models.FieldMetricTypeClass.MIN
+    MAX = models.FieldMetricTypeClass.MAX
+    MEAN = models.FieldMetricTypeClass.MEAN
+    MEDIAN = models.FieldMetricTypeClass.MEDIAN
+    STDDEV = models.FieldMetricTypeClass.STDDEV
+    NEGATIVE_COUNT = models.FieldMetricTypeClass.NEGATIVE_COUNT
+    NEGATIVE_PERCENTAGE = models.FieldMetricTypeClass.NEGATIVE_PERCENTAGE
+    ZERO_COUNT = models.FieldMetricTypeClass.ZERO_COUNT
+    ZERO_PERCENTAGE = models.FieldMetricTypeClass.ZERO_PERCENTAGE
+
+
+class OperatorType(str, Enum):
+    """Enum for assertion operators used in column metric assertions."""
+
+    EQUAL_TO = models.AssertionStdOperatorClass.EQUAL_TO
+    NOT_EQUAL_TO = models.AssertionStdOperatorClass.NOT_EQUAL_TO
+    GREATER_THAN = models.AssertionStdOperatorClass.GREATER_THAN
+    GREATER_THAN_OR_EQUAL_TO = models.AssertionStdOperatorClass.GREATER_THAN_OR_EQUAL_TO
+    LESS_THAN = models.AssertionStdOperatorClass.LESS_THAN
+    LESS_THAN_OR_EQUAL_TO = models.AssertionStdOperatorClass.LESS_THAN_OR_EQUAL_TO
+    BETWEEN = models.AssertionStdOperatorClass.BETWEEN
+    IN = models.AssertionStdOperatorClass.IN
+    NOT_IN = models.AssertionStdOperatorClass.NOT_IN
+    NULL = models.AssertionStdOperatorClass.NULL
+    NOT_NULL = models.AssertionStdOperatorClass.NOT_NULL
+    IS_TRUE = models.AssertionStdOperatorClass.IS_TRUE
+    IS_FALSE = models.AssertionStdOperatorClass.IS_FALSE
+    CONTAIN = models.AssertionStdOperatorClass.CONTAIN
+    END_WITH = models.AssertionStdOperatorClass.END_WITH
+    START_WITH = models.AssertionStdOperatorClass.START_WITH
+    REGEX_MATCH = models.AssertionStdOperatorClass.REGEX_MATCH
+
+
+class ValueType(str, Enum):
+    """Enum for assertion parameter value types."""
+
+    STRING = models.AssertionStdParameterTypeClass.STRING
+    NUMBER = models.AssertionStdParameterTypeClass.NUMBER
+    UNKNOWN = models.AssertionStdParameterTypeClass.UNKNOWN
+    # Note: LIST and SET are intentionally excluded as they are not yet supported
+    # LIST = models.AssertionStdParameterTypeClass.LIST
+    # SET = models.AssertionStdParameterTypeClass.SET
+
+
+# Type aliases
+MetricInputType = Union[MetricType, models.FieldMetricTypeClass, str]
+ValueInputType = Union[str, int, float]
+ValueTypeInputType = Union[ValueType, models.AssertionStdParameterTypeClass, str]
+RangeInputType = tuple[ValueInputType, ValueInputType]
+RangeTypeInputType = Union[
+    str,
+    tuple[str, str],
+    ValueTypeInputType,
+    tuple[ValueTypeInputType, ValueTypeInputType],
+]
+RangeTypeParsedType = tuple[ValueTypeInputType, ValueTypeInputType]
+OperatorInputType = Union[OperatorType, models.AssertionStdOperatorClass, str]

acryl_datahub_cloud/sdk/assertion_input/freshness_assertion_input.py

@@ -1,4 +1,5 @@
 from datetime import datetime
+from enum import Enum
 from typing import Optional, Union
 
 from acryl_datahub_cloud.sdk.assertion_input.assertion_input import (
@@ -32,6 +33,48 @@ from datahub.metadata.urns import AssertionUrn, CorpUserUrn, DatasetUrn
 from datahub.sdk.entity_client import EntityClient
 
 
+class FreshnessAssertionScheduleCheckType(str, Enum):
+    FIXED_INTERVAL = "FIXED_INTERVAL"
+    SINCE_THE_LAST_CHECK = "SINCE_THE_LAST_CHECK"
+
+
+DEFAULT_FRESHNESS_SCHEDULE_CHECK_TYPE = (
+    FreshnessAssertionScheduleCheckType.SINCE_THE_LAST_CHECK
+)
+
+
+def _parse_freshness_schedule_check_type(
+    schedule_check_type: Optional[
+        Union[
+            str,
+            FreshnessAssertionScheduleCheckType,
+            models.FreshnessAssertionScheduleTypeClass,
+        ]
+    ],
+) -> FreshnessAssertionScheduleCheckType:
+    """Parse the freshness schedule check type."""
+    if isinstance(schedule_check_type, FreshnessAssertionScheduleCheckType):
+        return schedule_check_type
+    if isinstance(schedule_check_type, models.FreshnessAssertionScheduleTypeClass):
+        return FreshnessAssertionScheduleCheckType(
+            _try_parse_and_validate_schema_classes_enum(
+                schedule_check_type, models.FreshnessAssertionScheduleTypeClass
+            )
+        )
+    if not schedule_check_type:
+        return DEFAULT_FRESHNESS_SCHEDULE_CHECK_TYPE
+
+    # Make string comparison case-insensitive
+    if isinstance(schedule_check_type, str):
+        schedule_check_type_upper = schedule_check_type.upper()
+        for member in FreshnessAssertionScheduleCheckType:
+            if member.value.upper() == schedule_check_type_upper:
+                return member
+        # If no match found, fall back to original behavior for error
+
+    return FreshnessAssertionScheduleCheckType(schedule_check_type)
+
+
 class _FreshnessAssertionInput(_AssertionInput, _HasFreshnessFeatures):
     def _assertion_type(self) -> str:
         """Get the assertion type."""
@@ -56,7 +99,11 @@ class _FreshnessAssertionInput(_AssertionInput, _HasFreshnessFeatures):
         updated_by: Union[str, CorpUserUrn],
         updated_at: datetime,
         freshness_schedule_check_type: Optional[
-            Union[str, models.FreshnessAssertionScheduleTypeClass]
+            Union[
+                str,
+                FreshnessAssertionScheduleCheckType,
+                models.FreshnessAssertionScheduleTypeClass,
+            ]
         ] = None,
         lookback_window: Optional[TimeWindowSizeInputTypes] = None,
     ):
@@ -78,28 +125,30 @@ class _FreshnessAssertionInput(_AssertionInput, _HasFreshnessFeatures):
             updated_at=updated_at,
         )
 
-        self.freshness_schedule_check_type = (
-            _try_parse_and_validate_schema_classes_enum(
-                freshness_schedule_check_type
-                or models.FreshnessAssertionScheduleTypeClass.SINCE_THE_LAST_CHECK,
-                models.FreshnessAssertionScheduleTypeClass,
-            )
+        self.freshness_schedule_check_type = _parse_freshness_schedule_check_type(
+            freshness_schedule_check_type
         )
         self.lookback_window = (
             _try_parse_time_window_size(lookback_window) if lookback_window else None
         )
+        self._validate_schedule_check_type()
+
+    def _validate_schedule_check_type(self) -> None:
+        """Validate the schedule check type."""
+        if self.freshness_schedule_check_type is None:
+            raise SDKUsageError("Freshness schedule check type is required.")
         if (
             self.freshness_schedule_check_type
-            is models.FreshnessAssertionScheduleTypeClass.FIXED_INTERVAL
-            and lookback_window is None
+            == FreshnessAssertionScheduleCheckType.FIXED_INTERVAL
+            and self.lookback_window is None
         ):
             raise SDKUsageError(
                 "Fixed interval freshness assertions must have a lookback_window provided."
             )
         if (
             self.freshness_schedule_check_type
-            is models.FreshnessAssertionScheduleTypeClass.SINCE_THE_LAST_CHECK
-            and lookback_window is not None
+            == FreshnessAssertionScheduleCheckType.SINCE_THE_LAST_CHECK
+            and self.lookback_window is not None
         ):
             raise SDKUsageError(
                 "Since the last check freshness assertions cannot have a lookback_window provided."