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

acryl_datahub_cloud/_codegen_config.json

@@ -1,6 +1,6 @@
 {
   "name": "acryl-datahub-cloud",
-  "version": "0.3.12rc5",
+  "version": "0.3.12rc6",
   "install_requires": [
     "avro-gen3==0.7.16",
     "acryl-datahub"

acryl_datahub_cloud/datahub_forms_notifications/get_feature_flag.gql

@@ -0,0 +1,7 @@
+query getAppConfig {
+  appConfig {
+    featureFlags {
+		formsNotificationsEnabled
+	}
+  }
+}

acryl_datahub_cloud/sdk/__init__.py

@@ -2,6 +2,7 @@ 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,
@@ -13,6 +14,9 @@ from acryl_datahub_cloud.sdk.assertion_input.assertion_input import (
     InferenceSensitivity,
     TimeWindowSize,
 )
+from acryl_datahub_cloud.sdk.assertion_input.sql_assertion_input import (
+    SqlAssertionCriteria,
+)
 from acryl_datahub_cloud.sdk.assertions_client import AssertionsClient
 from acryl_datahub_cloud.sdk.resolver_client import ResolverClient
 from acryl_datahub_cloud.sdk.subscription_client import SubscriptionClient
@@ -30,4 +34,6 @@ __all__ = [
     "AssertionsClient",
     "ResolverClient",
     "SubscriptionClient",
+    "SqlAssertion",
+    "SqlAssertionCriteria",
 ]

acryl_datahub_cloud/sdk/assertion/assertion_base.py

@@ -17,6 +17,7 @@ from typing_extensions import Self
 from acryl_datahub_cloud.sdk.assertion_input.assertion_input import (
     ASSERTION_MONITOR_DEFAULT_TRAINING_LOOKBACK_WINDOW_DAYS,
     DEFAULT_DETECTION_MECHANISM,
+    DEFAULT_EVERY_SIX_HOURS_SCHEDULE,
     DEFAULT_SCHEDULE,
     DEFAULT_SENSITIVITY,
     AssertionIncidentBehavior,
@@ -35,6 +36,9 @@ from acryl_datahub_cloud.sdk.assertion_input.smart_column_metric_assertion_input
     ValueInputType,
     ValueTypeInputType,
 )
+from acryl_datahub_cloud.sdk.assertion_input.sql_assertion_input import (
+    SqlAssertionCriteria,
+)
 from acryl_datahub_cloud.sdk.entities.assertion import Assertion
 from acryl_datahub_cloud.sdk.entities.monitor import (
     Monitor,
@@ -73,7 +77,9 @@ class _HasSchedule:
         return self._schedule
 
     @staticmethod
-    def _get_schedule(monitor: Monitor) -> models.CronScheduleClass:
+    def _get_schedule(
+        monitor: Monitor, default: models.CronScheduleClass = DEFAULT_SCHEDULE
+    ) -> models.CronScheduleClass:
         """Get the schedule from the monitor."""
         assertion_evaluation_specs = _get_nested_field_for_entity_with_default(
             monitor,
@@ -81,11 +87,11 @@ class _HasSchedule:
             [],
         )
         if len(assertion_evaluation_specs) == 0:
-            return DEFAULT_SCHEDULE
+            return default
         assertion_evaluation_spec = assertion_evaluation_specs[0]
         schedule = assertion_evaluation_spec.schedule
         if schedule is None:
-            return DEFAULT_SCHEDULE
+            return default
         return schedule
 
 
@@ -335,6 +341,7 @@ class _AssertionPublic(ABC):
     Abstract base class that represents a public facing assertion and contains the common properties of all assertions.
     """
 
+    # TODO: have the individual classes self-declare this
     _SUPPORTED_WITH_FILTER_ASSERTION_TYPES = (
         models.FreshnessAssertionInfoClass,
         models.VolumeAssertionInfoClass,
@@ -1148,3 +1155,163 @@ class FreshnessAssertion(_HasSchedule, _AssertionPublic):
             raise SDKNotYetSupportedError("FILE_METADATA DatasetFreshnessSourceType")
         else:
             raise SDKNotYetSupportedError(f"DatasetFreshnessSourceType {source_type}")
+
+
+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: 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.
+            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) -> SqlAssertionCriteria:
+        return self._criteria
+
+    @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_criteria(assertion: Assertion) -> 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):
+            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:
+                raise SDKNotYetSupportedError(
+                    f"Assertion {assertion.urn} does not have a valid parameters for the SQL assertion"
+                )
+
+            return SqlAssertionCriteria(
+                type=assertion.info.type
+                if isinstance(assertion.info.type, str)
+                else str(assertion.info.type),
+                change_type=assertion.info.changeType
+                if assertion.info.changeType is None
+                else (
+                    assertion.info.changeType
+                    if isinstance(assertion.info.changeType, str)
+                    else str(assertion.info.changeType)
+                ),
+                operator=assertion.info.operator
+                if isinstance(assertion.info.operator, str)
+                else str(assertion.info.operator),
+                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_EVERY_SIX_HOURS_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

@@ -4,6 +4,7 @@ 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,
@@ -14,5 +15,6 @@ AssertionTypes = Union[
     SmartVolumeAssertion,
     FreshnessAssertion,
     SmartColumnMetricAssertion,
+    SqlAssertion,
     # TODO: Add other assertion types here as we add them.
 ]

acryl_datahub_cloud/sdk/assertion_input/assertion_input.py

@@ -142,6 +142,40 @@ class _DatasetProfile(AbstractDetectionMechanism):
     type: Literal["dataset_profile"] = "dataset_profile"
 
 
+# Operators that require a single value numeric parameter
+SINGLE_VALUE_NUMERIC_OPERATORS = [
+    models.AssertionStdOperatorClass.EQUAL_TO,
+    models.AssertionStdOperatorClass.NOT_EQUAL_TO,
+    models.AssertionStdOperatorClass.GREATER_THAN,
+    models.AssertionStdOperatorClass.LESS_THAN,
+    models.AssertionStdOperatorClass.GREATER_THAN_OR_EQUAL_TO,
+    models.AssertionStdOperatorClass.LESS_THAN_OR_EQUAL_TO,
+]
+
+# Operators that require a single value parameter
+SINGLE_VALUE_OPERATORS = [
+    models.AssertionStdOperatorClass.CONTAIN,
+    models.AssertionStdOperatorClass.END_WITH,
+    models.AssertionStdOperatorClass.START_WITH,
+    models.AssertionStdOperatorClass.REGEX_MATCH,
+    models.AssertionStdOperatorClass.IN,
+    models.AssertionStdOperatorClass.NOT_IN,
+] + SINGLE_VALUE_NUMERIC_OPERATORS
+
+# Operators that require a numeric range parameter
+RANGE_OPERATORS = [
+    models.AssertionStdOperatorClass.BETWEEN,
+]
+
+# Operators that require no parameters
+NO_PARAMETER_OPERATORS = [
+    models.AssertionStdOperatorClass.NULL,
+    models.AssertionStdOperatorClass.NOT_NULL,
+    models.AssertionStdOperatorClass.IS_TRUE,
+    models.AssertionStdOperatorClass.IS_FALSE,
+]
+
+
 # Keep these two lists in sync:
 _DETECTION_MECHANISM_CONCRETE_TYPES = (
     _InformationSchema,
@@ -1096,15 +1130,12 @@ class _AssertionInput(ABC):
         Returns:
             A Monitor entity configured with the assertion input parameters.
         """
-        source_type, field = self._convert_assertion_source_type_and_field()
         return Monitor(
             id=(self.dataset_urn, assertion_urn),
             info=self._create_monitor_info(
                 assertion_urn=assertion_urn,
                 status=self._convert_monitor_status(),
                 schedule=self._convert_schedule(),
-                source_type=source_type,
-                field=field,
             ),
         )
 
@@ -1182,8 +1213,6 @@ class _AssertionInput(ABC):
         assertion_urn: AssertionUrn,
         status: models.MonitorStatusClass,
         schedule: models.CronScheduleClass,
-        source_type: Union[str, models.DatasetFreshnessSourceTypeClass],
-        field: Optional[FieldSpecType],
     ) -> models.MonitorInfoClass:
         """
         Create a MonitorInfoClass with all the necessary components.
@@ -1191,8 +1220,6 @@ class _AssertionInput(ABC):
         Args:
             status: The monitor status.
             schedule: The monitor schedule.
-            source_type: The source type.
-            field: Optional field specification.
         Returns:
             A MonitorInfoClass configured with all the provided components.
         """
@@ -1432,12 +1459,11 @@ class _SmartFreshnessAssertionInput(
         assertion_urn: AssertionUrn,
         status: models.MonitorStatusClass,
         schedule: models.CronScheduleClass,
-        source_type: Union[str, models.DatasetFreshnessSourceTypeClass],
-        field: Optional[FieldSpecType],
     ) -> models.MonitorInfoClass:
         """
         Create a MonitorInfoClass with all the necessary components.
         """
+        source_type, field = self._convert_assertion_source_type_and_field()
         return models.MonitorInfoClass(
             type=models.MonitorTypeClass.ASSERTION,
             status=status,
@@ -1591,12 +1617,11 @@ class _SmartVolumeAssertionInput(_AssertionInput, _HasSmartAssertionInputs):
         assertion_urn: AssertionUrn,
         status: models.MonitorStatusClass,
         schedule: models.CronScheduleClass,
-        source_type: Union[str, models.DatasetFreshnessSourceTypeClass],
-        field: Optional[FieldSpecType],
     ) -> models.MonitorInfoClass:
         """
         Create a MonitorInfoClass with all the necessary components.
         """
+        source_type, field = self._convert_assertion_source_type_and_field()
         return models.MonitorInfoClass(
             type=models.MonitorTypeClass.ASSERTION,
             status=status,

acryl_datahub_cloud/sdk/assertion_input/freshness_assertion_input.py

@@ -112,12 +112,11 @@ class _FreshnessAssertionInput(_AssertionInput, _HasFreshnessFeatures):
         assertion_urn: AssertionUrn,
         status: models.MonitorStatusClass,
         schedule: models.CronScheduleClass,
-        source_type: Union[str, models.DatasetFreshnessSourceTypeClass],
-        field: Optional[FieldSpecType],
     ) -> models.MonitorInfoClass:
         """
         Create a MonitorInfoClass with all the necessary components.
         """
+        source_type, field = self._convert_assertion_source_type_and_field()
         return models.MonitorInfoClass(
             type=models.MonitorTypeClass.ASSERTION,
             status=status,

acryl_datahub_cloud/sdk/assertion_input/smart_column_metric_assertion_input.py

@@ -5,6 +5,9 @@ from typing import Optional, Union
 from acryl_datahub_cloud.sdk.assertion_input.assertion_input import (
     DEFAULT_EVERY_SIX_HOURS_SCHEDULE,
     HIGH_WATERMARK_ALLOWED_FIELD_TYPES,
+    NO_PARAMETER_OPERATORS,
+    RANGE_OPERATORS,
+    SINGLE_VALUE_OPERATORS,
     AssertionIncidentBehavior,
     AssertionInfoInputType,
     DetectionMechanismInputTypes,
@@ -85,35 +88,6 @@ FIELD_VALUES_OPERATOR_CONFIG = {
     ],
 }
 
-# Operators that require a single value parameter
-SINGLE_VALUE_OPERATORS = [
-    models.AssertionStdOperatorClass.EQUAL_TO,
-    models.AssertionStdOperatorClass.NOT_EQUAL_TO,
-    models.AssertionStdOperatorClass.GREATER_THAN,
-    models.AssertionStdOperatorClass.LESS_THAN,
-    models.AssertionStdOperatorClass.GREATER_THAN_OR_EQUAL_TO,
-    models.AssertionStdOperatorClass.LESS_THAN_OR_EQUAL_TO,
-    models.AssertionStdOperatorClass.CONTAIN,
-    models.AssertionStdOperatorClass.END_WITH,
-    models.AssertionStdOperatorClass.START_WITH,
-    models.AssertionStdOperatorClass.REGEX_MATCH,
-    models.AssertionStdOperatorClass.IN,
-    models.AssertionStdOperatorClass.NOT_IN,
-]
-
-# Operators that require a range parameter
-RANGE_OPERATORS = [
-    models.AssertionStdOperatorClass.BETWEEN,
-]
-
-# Operators that require no parameters
-NO_PARAMETER_OPERATORS = [
-    models.AssertionStdOperatorClass.NULL,
-    models.AssertionStdOperatorClass.NOT_NULL,
-    models.AssertionStdOperatorClass.IS_TRUE,
-    models.AssertionStdOperatorClass.IS_FALSE,
-]
-
 # 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 = {
@@ -442,12 +416,11 @@ class _SmartColumnMetricAssertionInput(_AssertionInput, _HasSmartAssertionInputs
         assertion_urn: AssertionUrn,
         status: models.MonitorStatusClass,
         schedule: models.CronScheduleClass,
-        source_type: Union[str, models.DatasetFreshnessSourceTypeClass],
-        field: Optional[FieldSpecType],
     ) -> models.MonitorInfoClass:
         """
         Create a MonitorInfoClass with all the necessary components.
         """
+        source_type, field = self._convert_assertion_source_type_and_field()
         return models.MonitorInfoClass(
             type=models.MonitorTypeClass.ASSERTION,
             status=status,