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

acryl_datahub_cloud/sdk/assertion_input/smart_freshness_assertion_input.py

@@ -0,0 +1,218 @@
+"""
+Smart Freshness Assertion Input classes for DataHub.
+
+This module contains the _SmartFreshnessAssertionInput class that handles
+input validation and entity creation for smart freshness assertions.
+"""
+
+from datetime import datetime
+from typing import Optional, Union
+
+from acryl_datahub_cloud.sdk.assertion_input.assertion_input import (
+    DEFAULT_DAILY_SCHEDULE,
+    LAST_MODIFIED_ALLOWED_FIELD_TYPES,
+    AssertionIncidentBehaviorInputTypes,
+    DetectionMechanismInputTypes,
+    ExclusionWindowInputTypes,
+    FieldSpecType,
+    InferenceSensitivity,
+    _AssertionInput,
+    _AuditLog,
+    _DataHubOperation,
+    _HasFreshnessFeatures,
+    _HasSmartAssertionInputs,
+    _InformationSchema,
+    _LastModifiedColumn,
+)
+from acryl_datahub_cloud.sdk.entities.assertion import (
+    AssertionInfoInputType,
+    TagsInputType,
+)
+from acryl_datahub_cloud.sdk.errors import (
+    SDKNotYetSupportedError,
+    SDKUsageError,
+)
+from datahub.metadata import schema_classes as models
+from datahub.metadata.urns import AssertionUrn, CorpUserUrn, DatasetUrn
+from datahub.sdk.entity_client import EntityClient
+
+
+class _SmartFreshnessAssertionInput(
+    _AssertionInput, _HasSmartAssertionInputs, _HasFreshnessFeatures
+):
+    def __init__(
+        self,
+        *,
+        dataset_urn: Union[str, DatasetUrn],
+        entity_client: EntityClient,  # Needed to get the schema field spec for the detection mechanism if needed
+        urn: Optional[Union[str, AssertionUrn]] = None,
+        display_name: Optional[str] = None,
+        enabled: bool = True,
+        schedule: Optional[Union[str, models.CronScheduleClass]] = None,
+        detection_mechanism: DetectionMechanismInputTypes = None,
+        sensitivity: Optional[Union[str, InferenceSensitivity]] = None,
+        exclusion_windows: Optional[ExclusionWindowInputTypes] = None,
+        training_data_lookback_days: Optional[int] = None,
+        incident_behavior: Optional[AssertionIncidentBehaviorInputTypes] = None,
+        tags: Optional[TagsInputType] = None,
+        created_by: Union[str, CorpUserUrn],
+        created_at: datetime,
+        updated_by: Union[str, CorpUserUrn],
+        updated_at: datetime,
+    ):
+        _AssertionInput.__init__(
+            self,
+            dataset_urn=dataset_urn,
+            entity_client=entity_client,
+            urn=urn,
+            display_name=display_name,
+            enabled=enabled,
+            schedule=schedule
+            if schedule is not None
+            else DEFAULT_DAILY_SCHEDULE,  # Use provided schedule or default for create case
+            detection_mechanism=detection_mechanism,
+            incident_behavior=incident_behavior,
+            tags=tags,
+            source_type=models.AssertionSourceTypeClass.INFERRED,  # Smart assertions are of type inferred, not native
+            created_by=created_by,
+            created_at=created_at,
+            updated_by=updated_by,
+            updated_at=updated_at,
+        )
+        _HasSmartAssertionInputs.__init__(
+            self,
+            sensitivity=sensitivity,
+            exclusion_windows=exclusion_windows,
+            training_data_lookback_days=training_data_lookback_days,
+        )
+
+    def _assertion_type(self) -> str:
+        """Get the assertion type."""
+        return models.AssertionTypeClass.FRESHNESS
+
+    def _create_assertion_info(
+        self, filter: Optional[models.DatasetFilterClass]
+    ) -> AssertionInfoInputType:
+        """
+        Create a FreshnessAssertionInfoClass for a smart freshness assertion.
+
+        Args:
+            filter: Optional filter to apply to the assertion.
+
+        Returns:
+            A FreshnessAssertionInfoClass configured for smart freshness.
+        """
+        return models.FreshnessAssertionInfoClass(
+            type=models.FreshnessAssertionTypeClass.DATASET_CHANGE,  # Currently only dataset change is supported
+            entity=str(self.dataset_urn),
+            # schedule (optional, must be left empty for smart freshness assertions - managed by the AI inference engine)
+            filter=filter,
+        )
+
+    def _convert_schedule(self) -> models.CronScheduleClass:
+        """Create a schedule for a smart freshness assertion.
+
+        For create case, uses DEFAULT_DAILY_SCHEDULE. For update case, preserves existing schedule.
+
+        Returns:
+            A CronScheduleClass with appropriate schedule settings.
+        """
+        assert self.schedule is not None, (
+            "Schedule should never be None due to constructor logic"
+        )
+        return self.schedule
+
+    def _get_assertion_evaluation_parameters(
+        self, source_type: str, field: Optional[FieldSpecType]
+    ) -> models.AssertionEvaluationParametersClass:
+        # Ensure field is either None or FreshnessFieldSpecClass
+        freshness_field = None
+        if field is not None:
+            if not isinstance(field, models.FreshnessFieldSpecClass):
+                raise SDKUsageError(
+                    f"Expected FreshnessFieldSpecClass for freshness assertion, got {type(field).__name__}"
+                )
+            freshness_field = field
+
+        return models.AssertionEvaluationParametersClass(
+            type=models.AssertionEvaluationParametersTypeClass.DATASET_FRESHNESS,
+            datasetFreshnessParameters=models.DatasetFreshnessAssertionParametersClass(
+                sourceType=source_type, field=freshness_field
+            ),
+        )
+
+    def _convert_assertion_source_type_and_field(
+        self,
+    ) -> tuple[str, Optional[FieldSpecType]]:
+        """
+        Convert detection mechanism into source type and field specification for freshness assertions.
+
+        Returns:
+            A tuple of (source_type, field) where field may be None.
+            Note that the source_type is a string, not a models.DatasetFreshnessSourceTypeClass (or other assertion source type) since
+            the source type is not a enum in the code generated from the DatasetFreshnessSourceType enum in the PDL.
+
+        Raises:
+            SDKNotYetSupportedError: If the detection mechanism is not supported.
+            SDKUsageError: If the field (column) is not found in the dataset,
+            and the detection mechanism requires a field. Also if the field
+            is not an allowed type for the detection mechanism.
+        """
+        source_type = models.DatasetFreshnessSourceTypeClass.INFORMATION_SCHEMA
+        field = None
+
+        if isinstance(self.detection_mechanism, _LastModifiedColumn):
+            source_type = models.DatasetFreshnessSourceTypeClass.FIELD_VALUE
+            field = self._create_field_spec(
+                self.detection_mechanism.column_name,
+                LAST_MODIFIED_ALLOWED_FIELD_TYPES,
+                "last modified column",
+                models.FreshnessFieldKindClass.LAST_MODIFIED,
+                self._get_schema_field_spec,
+                self._validate_field_type,
+            )
+        elif isinstance(self.detection_mechanism, _InformationSchema):
+            source_type = models.DatasetFreshnessSourceTypeClass.INFORMATION_SCHEMA
+        elif isinstance(self.detection_mechanism, _DataHubOperation):
+            source_type = models.DatasetFreshnessSourceTypeClass.DATAHUB_OPERATION
+        elif isinstance(self.detection_mechanism, _AuditLog):
+            source_type = models.DatasetFreshnessSourceTypeClass.AUDIT_LOG
+        else:
+            raise SDKNotYetSupportedError(
+                f"Detection mechanism {self.detection_mechanism} not yet supported for smart freshness assertions"
+            )
+
+        return source_type, field
+
+    def _create_monitor_info(
+        self,
+        assertion_urn: AssertionUrn,
+        status: models.MonitorStatusClass,
+        schedule: models.CronScheduleClass,
+    ) -> 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,
+            assertionMonitor=models.AssertionMonitorClass(
+                assertions=[
+                    models.AssertionEvaluationSpecClass(
+                        assertion=str(assertion_urn),
+                        schedule=schedule,
+                        parameters=self._get_assertion_evaluation_parameters(
+                            str(source_type), field
+                        ),
+                    ),
+                ],
+                settings=models.AssertionMonitorSettingsClass(
+                    adjustmentSettings=models.AssertionAdjustmentSettingsClass(
+                        sensitivity=self._convert_sensitivity(),
+                        exclusionWindows=self._convert_exclusion_windows(),
+                        trainingDataLookbackWindowDays=self.training_data_lookback_days,
+                    ),
+                ),
+            ),
+        )

acryl_datahub_cloud/sdk/assertion_input/smart_sql_assertion_input.py

@@ -0,0 +1,181 @@
+"""
+Smart SQL Assertion Input classes for DataHub.
+
+This module contains the _SmartSqlAssertionInput class that handles
+input validation and entity creation for smart SQL assertions.
+"""
+
+from datetime import datetime
+from typing import Optional, Union
+
+from acryl_datahub_cloud.sdk.assertion_input.assertion_input import (
+    DEFAULT_EVERY_SIX_HOURS_SCHEDULE,
+    AssertionIncidentBehaviorInputTypes,
+    ExclusionWindowInputTypes,
+    FieldSpecType,
+    InferenceSensitivity,
+    _AssertionInput,
+    _HasSmartAssertionInputs,
+)
+from acryl_datahub_cloud.sdk.entities.assertion import (
+    AssertionInfoInputType,
+    TagsInputType,
+)
+from datahub.metadata import schema_classes as models
+from datahub.metadata.urns import AssertionUrn, CorpUserUrn, DatasetUrn
+from datahub.sdk.entity_client import EntityClient
+
+
+class _SmartSqlAssertionInput(_AssertionInput, _HasSmartAssertionInputs):
+    """Input class for smart SQL assertions with AI-powered inference."""
+
+    def __init__(
+        self,
+        *,
+        dataset_urn: Union[str, DatasetUrn],
+        entity_client: EntityClient,
+        statement: Optional[str],
+        created_by: Union[str, CorpUserUrn],
+        created_at: datetime,
+        updated_by: Union[str, CorpUserUrn],
+        updated_at: datetime,
+        urn: Optional[Union[str, AssertionUrn]] = None,
+        display_name: Optional[str] = None,
+        enabled: bool = True,
+        schedule: Optional[Union[str, models.CronScheduleClass]] = None,
+        sensitivity: Optional[Union[str, InferenceSensitivity]] = None,
+        exclusion_windows: Optional[ExclusionWindowInputTypes] = None,
+        training_data_lookback_days: Optional[int] = None,
+        incident_behavior: Optional[AssertionIncidentBehaviorInputTypes] = None,
+        tags: Optional[TagsInputType] = None,
+    ):
+        _AssertionInput.__init__(
+            self,
+            dataset_urn=dataset_urn,
+            entity_client=entity_client,
+            urn=urn,
+            display_name=display_name,
+            enabled=enabled,
+            schedule=schedule,
+            detection_mechanism=None,  # Smart SQL assertions don't use detection mechanisms
+            incident_behavior=incident_behavior,
+            tags=tags,
+            source_type=models.AssertionSourceTypeClass.INFERRED,  # Smart assertions are of type inferred
+            created_by=created_by,
+            created_at=created_at,
+            updated_by=updated_by,
+            updated_at=updated_at,
+        )
+        _HasSmartAssertionInputs.__init__(
+            self,
+            sensitivity=sensitivity,
+            exclusion_windows=exclusion_windows,
+            training_data_lookback_days=training_data_lookback_days,
+        )
+        self.statement = statement
+
+    def _create_assertion_info(
+        self, filter: Optional[models.DatasetFilterClass]
+    ) -> AssertionInfoInputType:
+        """
+        Create a SqlAssertionInfoClass for a smart SQL assertion.
+
+        For smart SQL assertions, we use METRIC type with a placeholder operator/parameters
+        since the actual thresholds are inferred by AI.
+
+        Args:
+            filter: Optional filter (not used for SQL assertions).
+
+        Returns:
+            A SqlAssertionInfoClass configured for smart SQL.
+        """
+
+        assert self.statement is not None, (
+            "statement is required for Smart SQL assertions"
+        )
+
+        # For smart SQL assertions, we use METRIC type with placeholder values
+        # The actual thresholds are inferred by the AI model
+        return models.SqlAssertionInfoClass(
+            type=models.SqlAssertionTypeClass.METRIC,
+            entity=str(self.dataset_urn),
+            statement=self.statement,
+            changeType=None,
+            # Use GREATER_THAN_OR_EQUAL_TO with 0 as a placeholder - AI will infer actual thresholds
+            operator=models.AssertionStdOperatorClass.GREATER_THAN_OR_EQUAL_TO,
+            parameters=models.AssertionStdParametersClass(
+                value=models.AssertionStdParameterClass(
+                    type=models.AssertionStdParameterTypeClass.NUMBER,
+                    value="0",
+                ),
+            ),
+        )
+
+    def _convert_schedule(self) -> models.CronScheduleClass:
+        """Create a schedule for a smart SQL assertion.
+
+        Returns:
+            A CronScheduleClass with appropriate schedule settings.
+        """
+        if self.schedule is None:
+            return DEFAULT_EVERY_SIX_HOURS_SCHEDULE
+
+        return models.CronScheduleClass(
+            cron=self.schedule.cron,
+            timezone=self.schedule.timezone,
+        )
+
+    def _get_assertion_evaluation_parameters(
+        self, source_type: str, field: Optional[FieldSpecType]
+    ) -> models.AssertionEvaluationParametersClass:
+        return models.AssertionEvaluationParametersClass(
+            type=models.AssertionEvaluationParametersTypeClass.DATASET_SQL,
+        )
+
+    def _convert_assertion_source_type_and_field(
+        self,
+    ) -> tuple[str, None]:
+        """SQL assertions do not have source type or field."""
+        return "None", None
+
+    def _create_filter_from_detection_mechanism(
+        self,
+    ) -> Optional[models.DatasetFilterClass]:
+        """SQL assertions do not use filters."""
+        return None
+
+    def _create_monitor_info(
+        self,
+        assertion_urn: AssertionUrn,
+        status: models.MonitorStatusClass,
+        schedule: models.CronScheduleClass,
+    ) -> models.MonitorInfoClass:
+        """
+        Create a MonitorInfoClass with all the necessary components for smart SQL assertion.
+        """
+        return models.MonitorInfoClass(
+            type=models.MonitorTypeClass.ASSERTION,
+            status=status,
+            assertionMonitor=models.AssertionMonitorClass(
+                assertions=[
+                    models.AssertionEvaluationSpecClass(
+                        assertion=str(assertion_urn),
+                        schedule=schedule,
+                        parameters=self._get_assertion_evaluation_parameters(
+                            "None", None
+                        ),
+                    ),
+                ],
+                settings=models.AssertionMonitorSettingsClass(
+                    adjustmentSettings=models.AssertionAdjustmentSettingsClass(
+                        sensitivity=self._convert_sensitivity(),
+                        exclusionWindows=self._convert_exclusion_windows(),
+                        trainingDataLookbackWindowDays=self.training_data_lookback_days,
+                    ),
+                ),
+            ),
+        )
+
+    def _assertion_type(self) -> str:
+        """Get the assertion type."""
+        return models.AssertionTypeClass.SQL

acryl_datahub_cloud/sdk/assertion_input/smart_volume_assertion_input.py

@@ -0,0 +1,189 @@
+"""
+Smart Volume Assertion Input classes for DataHub.
+
+This module contains the _SmartVolumeAssertionInput class that handles
+input validation and entity creation for smart volume assertions.
+"""
+
+from datetime import datetime
+from typing import Optional, Union
+
+from acryl_datahub_cloud.sdk.assertion_input.assertion_input import (
+    DEFAULT_DAILY_SCHEDULE,
+    AssertionIncidentBehaviorInputTypes,
+    DetectionMechanismInputTypes,
+    ExclusionWindowInputTypes,
+    FieldSpecType,
+    InferenceSensitivity,
+    _AssertionInput,
+    _DatasetProfile,
+    _HasSmartAssertionInputs,
+    _InformationSchema,
+    _Query,
+)
+from acryl_datahub_cloud.sdk.entities.assertion import (
+    AssertionInfoInputType,
+    TagsInputType,
+)
+from acryl_datahub_cloud.sdk.errors import SDKNotYetSupportedError
+from datahub.metadata import schema_classes as models
+from datahub.metadata.urns import AssertionUrn, CorpUserUrn, DatasetUrn
+from datahub.sdk.entity_client import EntityClient
+
+
+class _SmartVolumeAssertionInput(_AssertionInput, _HasSmartAssertionInputs):
+    def __init__(
+        self,
+        *,
+        dataset_urn: Union[str, DatasetUrn],
+        entity_client: EntityClient,  # Needed to get the schema field spec for the detection mechanism if needed
+        urn: Optional[Union[str, AssertionUrn]] = None,
+        display_name: Optional[str] = None,
+        enabled: bool = True,
+        schedule: Optional[Union[str, models.CronScheduleClass]] = None,
+        detection_mechanism: DetectionMechanismInputTypes = None,
+        sensitivity: Optional[Union[str, InferenceSensitivity]] = None,
+        exclusion_windows: Optional[ExclusionWindowInputTypes] = None,
+        training_data_lookback_days: Optional[int] = None,
+        incident_behavior: Optional[AssertionIncidentBehaviorInputTypes] = None,
+        tags: Optional[TagsInputType] = None,
+        created_by: Union[str, CorpUserUrn],
+        created_at: datetime,
+        updated_by: Union[str, CorpUserUrn],
+        updated_at: datetime,
+    ):
+        _AssertionInput.__init__(
+            self,
+            dataset_urn=dataset_urn,
+            entity_client=entity_client,
+            urn=urn,
+            display_name=display_name,
+            enabled=enabled,
+            schedule=schedule,
+            detection_mechanism=detection_mechanism,
+            incident_behavior=incident_behavior,
+            tags=tags,
+            source_type=models.AssertionSourceTypeClass.INFERRED,  # Smart assertions are of type inferred, not native
+            created_by=created_by,
+            created_at=created_at,
+            updated_by=updated_by,
+            updated_at=updated_at,
+        )
+        _HasSmartAssertionInputs.__init__(
+            self,
+            sensitivity=sensitivity,
+            exclusion_windows=exclusion_windows,
+            training_data_lookback_days=training_data_lookback_days,
+        )
+
+    def _create_assertion_info(
+        self, filter: Optional[models.DatasetFilterClass]
+    ) -> AssertionInfoInputType:
+        """
+        Create a VolumeAssertionInfoClass for a smart volume assertion.
+
+        Args:
+            filter: Optional filter to apply to the assertion.
+
+        Returns:
+            A VolumeAssertionInfoClass configured for smart volume.
+        """
+        return models.VolumeAssertionInfoClass(
+            type=models.VolumeAssertionTypeClass.ROW_COUNT_TOTAL,  # Currently only ROW_COUNT_TOTAL is supported for smart volume
+            entity=str(self.dataset_urn),
+            filter=filter,
+        )
+
+    def _convert_schedule(self) -> models.CronScheduleClass:
+        """Create a schedule for a smart volume assertion.
+
+        Returns:
+            A CronScheduleClass with appropriate schedule settings.
+        """
+        if self.schedule is None:
+            return DEFAULT_DAILY_SCHEDULE
+
+        return models.CronScheduleClass(
+            cron=self.schedule.cron,
+            timezone=self.schedule.timezone,
+        )
+
+    def _get_assertion_evaluation_parameters(
+        self, source_type: str, field: Optional[FieldSpecType]
+    ) -> models.AssertionEvaluationParametersClass:
+        return models.AssertionEvaluationParametersClass(
+            type=models.AssertionEvaluationParametersTypeClass.DATASET_VOLUME,
+            datasetVolumeParameters=models.DatasetVolumeAssertionParametersClass(
+                sourceType=source_type,
+            ),
+        )
+
+    def _convert_assertion_source_type_and_field(
+        self,
+    ) -> tuple[str, Optional[FieldSpecType]]:
+        """
+        Convert detection mechanism into source type and field specification for volume assertions.
+
+        Returns:
+            A tuple of (source_type, field) where field may be None.
+            Note that the source_type is a string, not a models.DatasetFreshnessSourceTypeClass (or other assertion source type) since
+            the source type is not a enum in the code generated from the DatasetFreshnessSourceType enum in the PDL.
+
+        Raises:
+            SDKNotYetSupportedError: If the detection mechanism is not supported.
+            SDKUsageError: If the field (column) is not found in the dataset,
+            and the detection mechanism requires a field. Also if the field
+            is not an allowed type for the detection mechanism.
+        """
+        source_type = models.DatasetVolumeSourceTypeClass.INFORMATION_SCHEMA
+        field = None
+
+        if isinstance(self.detection_mechanism, _Query):
+            source_type = models.DatasetVolumeSourceTypeClass.QUERY
+        elif isinstance(self.detection_mechanism, _InformationSchema):
+            source_type = models.DatasetVolumeSourceTypeClass.INFORMATION_SCHEMA
+        elif isinstance(self.detection_mechanism, _DatasetProfile):
+            source_type = models.DatasetVolumeSourceTypeClass.DATAHUB_DATASET_PROFILE
+        else:
+            raise SDKNotYetSupportedError(
+                f"Detection mechanism {self.detection_mechanism} not yet supported for smart volume assertions"
+            )
+
+        return source_type, field
+
+    def _create_monitor_info(
+        self,
+        assertion_urn: AssertionUrn,
+        status: models.MonitorStatusClass,
+        schedule: models.CronScheduleClass,
+    ) -> 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,
+            assertionMonitor=models.AssertionMonitorClass(
+                assertions=[
+                    models.AssertionEvaluationSpecClass(
+                        assertion=str(assertion_urn),
+                        schedule=schedule,
+                        parameters=self._get_assertion_evaluation_parameters(
+                            str(source_type), field
+                        ),
+                    ),
+                ],
+                settings=models.AssertionMonitorSettingsClass(
+                    adjustmentSettings=models.AssertionAdjustmentSettingsClass(
+                        sensitivity=self._convert_sensitivity(),
+                        exclusionWindows=self._convert_exclusion_windows(),
+                        trainingDataLookbackWindowDays=self.training_data_lookback_days,
+                    ),
+                ),
+            ),
+        )
+
+    def _assertion_type(self) -> str:
+        """Get the assertion type."""
+        return models.AssertionTypeClass.VOLUME