acryl-datahub-cloud 0.3.12.4rc2__py3-none-any.whl → 0.3.13__py3-none-any.whl

acryl_datahub_cloud/sdk/assertion_input/smart_volume_assertion_input.py

@@ -0,0 +1,191 @@
+"""
+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_HOURLY_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,
+        *,
+        # Required fields
+        dataset_urn: Union[str, DatasetUrn],
+        entity_client: EntityClient,  # Needed to get the schema field spec for the detection mechanism if needed
+        # Optional fields
+        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_HOURLY_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

acryl_datahub_cloud/sdk/assertions_client.py

@@ -27,8 +27,6 @@ from acryl_datahub_cloud.sdk.assertion_input.assertion_input import (
     InferenceSensitivity,
     TimeWindowSizeInputTypes,
     _AssertionInput,
-    _SmartFreshnessAssertionInput,
-    _SmartVolumeAssertionInput,
 )
 from acryl_datahub_cloud.sdk.assertion_input.column_metric_assertion_input import (
     ColumnMetricAssertionParameters,
@@ -45,6 +43,12 @@ from acryl_datahub_cloud.sdk.assertion_input.freshness_assertion_input import (
 from acryl_datahub_cloud.sdk.assertion_input.smart_column_metric_assertion_input import (
     _SmartColumnMetricAssertionInput,
 )
+from acryl_datahub_cloud.sdk.assertion_input.smart_freshness_assertion_input import (
+    _SmartFreshnessAssertionInput,
+)
+from acryl_datahub_cloud.sdk.assertion_input.smart_volume_assertion_input import (
+    _SmartVolumeAssertionInput,
+)
 from acryl_datahub_cloud.sdk.assertion_input.sql_assertion_input import (
     SqlAssertionCondition,
     SqlAssertionCriteria,

{acryl_datahub_cloud-0.3.12.4rc2.dist-info → acryl_datahub_cloud-0.3.13.dist-info}/METADATA

@@ -1,6 +1,7 @@
 Metadata-Version: 2.4
 Name: acryl-datahub-cloud
-Version: 0.3.12.4rc2
+Version: 0.3.13
+Requires-Python: >=3.10
 Requires-Dist: avro-gen3==0.7.16
 Requires-Dist: acryl-datahub
 Requires-Dist: croniter
@@ -8,92 +9,93 @@ Requires-Dist: pytz
 Requires-Dist: types-croniter
 Requires-Dist: tzlocal
 Provides-Extra: datahub-lineage-features
-Requires-Dist: pandas; extra == "datahub-lineage-features"
 Requires-Dist: pydantic<2; extra == "datahub-lineage-features"
-Requires-Dist: tenacity; extra == "datahub-lineage-features"
-Requires-Dist: opensearch-py==2.4.2; extra == "datahub-lineage-features"
 Requires-Dist: duckdb; extra == "datahub-lineage-features"
 Requires-Dist: pyarrow; extra == "datahub-lineage-features"
+Requires-Dist: tenacity; extra == "datahub-lineage-features"
+Requires-Dist: opensearch-py==2.4.2; extra == "datahub-lineage-features"
+Requires-Dist: pandas; extra == "datahub-lineage-features"
 Provides-Extra: datahub-reporting-forms
-Requires-Dist: pandas; extra == "datahub-reporting-forms"
+Requires-Dist: boto3; extra == "datahub-reporting-forms"
 Requires-Dist: pydantic<2; extra == "datahub-reporting-forms"
 Requires-Dist: duckdb; extra == "datahub-reporting-forms"
-Requires-Dist: termcolor==2.5.0; extra == "datahub-reporting-forms"
 Requires-Dist: pyarrow; extra == "datahub-reporting-forms"
-Requires-Dist: boto3; extra == "datahub-reporting-forms"
+Requires-Dist: pandas; extra == "datahub-reporting-forms"
+Requires-Dist: termcolor==2.5.0; extra == "datahub-reporting-forms"
 Provides-Extra: datahub-reporting-extract-graph
-Requires-Dist: pandas; extra == "datahub-reporting-extract-graph"
+Requires-Dist: boto3; extra == "datahub-reporting-extract-graph"
 Requires-Dist: pydantic<2; extra == "datahub-reporting-extract-graph"
-Requires-Dist: opensearch-py==2.4.2; extra == "datahub-reporting-extract-graph"
 Requires-Dist: duckdb; extra == "datahub-reporting-extract-graph"
 Requires-Dist: pyarrow; extra == "datahub-reporting-extract-graph"
-Requires-Dist: boto3; extra == "datahub-reporting-extract-graph"
+Requires-Dist: opensearch-py==2.4.2; extra == "datahub-reporting-extract-graph"
+Requires-Dist: pandas; extra == "datahub-reporting-extract-graph"
 Provides-Extra: datahub-reporting-extract-sql
-Requires-Dist: pandas; extra == "datahub-reporting-extract-sql"
+Requires-Dist: boto3; extra == "datahub-reporting-extract-sql"
 Requires-Dist: pydantic<2; extra == "datahub-reporting-extract-sql"
 Requires-Dist: duckdb; extra == "datahub-reporting-extract-sql"
 Requires-Dist: pyarrow; extra == "datahub-reporting-extract-sql"
-Requires-Dist: boto3; extra == "datahub-reporting-extract-sql"
+Requires-Dist: pandas; extra == "datahub-reporting-extract-sql"
 Provides-Extra: datahub-usage-reporting
-Requires-Dist: polars==1.30.0; extra == "datahub-usage-reporting"
-Requires-Dist: opensearch-py==2.4.2; extra == "datahub-usage-reporting"
+Requires-Dist: boto3; extra == "datahub-usage-reporting"
+Requires-Dist: pyarrow<=18.0.0; extra == "datahub-usage-reporting"
+Requires-Dist: pydantic<2; extra == "datahub-usage-reporting"
+Requires-Dist: duckdb; extra == "datahub-usage-reporting"
 Requires-Dist: scipy<=1.14.1; extra == "datahub-usage-reporting"
+Requires-Dist: polars==1.30.0; extra == "datahub-usage-reporting"
+Requires-Dist: pandas; extra == "datahub-usage-reporting"
 Requires-Dist: numpy<2; extra == "datahub-usage-reporting"
 Requires-Dist: pyarrow; extra == "datahub-usage-reporting"
+Requires-Dist: opensearch-py==2.4.2; extra == "datahub-usage-reporting"
 Requires-Dist: elasticsearch==7.13.4; extra == "datahub-usage-reporting"
-Requires-Dist: pandas; extra == "datahub-usage-reporting"
-Requires-Dist: pydantic<2; extra == "datahub-usage-reporting"
-Requires-Dist: duckdb; extra == "datahub-usage-reporting"
-Requires-Dist: pyarrow<=18.0.0; extra == "datahub-usage-reporting"
 Requires-Dist: termcolor==2.5.0; extra == "datahub-usage-reporting"
-Requires-Dist: boto3; extra == "datahub-usage-reporting"
 Provides-Extra: datahub-metadata-sharing
 Requires-Dist: tenacity; extra == "datahub-metadata-sharing"
 Provides-Extra: datahub-action-request-owner
 Requires-Dist: tenacity; extra == "datahub-action-request-owner"
 Provides-Extra: acryl-cs-issues
-Requires-Dist: openai; extra == "acryl-cs-issues"
 Requires-Dist: slack-sdk; extra == "acryl-cs-issues"
-Requires-Dist: zenpy; extra == "acryl-cs-issues"
 Requires-Dist: jinja2; extra == "acryl-cs-issues"
+Requires-Dist: zenpy; extra == "acryl-cs-issues"
+Requires-Dist: openai; extra == "acryl-cs-issues"
 Provides-Extra: datahub-forms-notifications
 Requires-Dist: tenacity; extra == "datahub-forms-notifications"
 Provides-Extra: all
+Requires-Dist: boto3; extra == "all"
+Requires-Dist: scipy<=1.14.1; extra == "all"
+Requires-Dist: pandas; extra == "all"
+Requires-Dist: pyarrow; extra == "all"
+Requires-Dist: zenpy; extra == "all"
 Requires-Dist: tenacity; extra == "all"
 Requires-Dist: opensearch-py==2.4.2; extra == "all"
-Requires-Dist: jinja2; extra == "all"
-Requires-Dist: scipy<=1.14.1; extra == "all"
-Requires-Dist: openai; extra == "all"
+Requires-Dist: termcolor==2.5.0; extra == "all"
+Requires-Dist: pyarrow<=18.0.0; extra == "all"
 Requires-Dist: pydantic<2; extra == "all"
 Requires-Dist: duckdb; extra == "all"
-Requires-Dist: pyarrow<=18.0.0; extra == "all"
-Requires-Dist: termcolor==2.5.0; extra == "all"
+Requires-Dist: jinja2; extra == "all"
+Requires-Dist: openai; extra == "all"
 Requires-Dist: polars==1.30.0; extra == "all"
-Requires-Dist: slack-sdk; extra == "all"
 Requires-Dist: numpy<2; extra == "all"
-Requires-Dist: zenpy; extra == "all"
-Requires-Dist: pyarrow; extra == "all"
+Requires-Dist: slack-sdk; extra == "all"
 Requires-Dist: elasticsearch==7.13.4; extra == "all"
-Requires-Dist: pandas; extra == "all"
-Requires-Dist: boto3; extra == "all"
 Provides-Extra: dev
-Requires-Dist: polars==1.30.0; extra == "dev"
-Requires-Dist: tenacity; extra == "dev"
-Requires-Dist: opensearch-py==2.4.2; extra == "dev"
-Requires-Dist: jinja2; extra == "dev"
-Requires-Dist: scipy<=1.14.1; extra == "dev"
+Requires-Dist: boto3; extra == "dev"
 Requires-Dist: slack-sdk; extra == "dev"
+Requires-Dist: pyarrow<=18.0.0; extra == "dev"
+Requires-Dist: pydantic<2; extra == "dev"
+Requires-Dist: duckdb; extra == "dev"
+Requires-Dist: scipy<=1.14.1; extra == "dev"
+Requires-Dist: jinja2; extra == "dev"
+Requires-Dist: openai; extra == "dev"
+Requires-Dist: polars==1.30.0; extra == "dev"
+Requires-Dist: pandas; extra == "dev"
+Requires-Dist: acryl-datahub[dev]; extra == "dev"
 Requires-Dist: numpy<2; extra == "dev"
-Requires-Dist: zenpy; extra == "dev"
 Requires-Dist: pyarrow; extra == "dev"
+Requires-Dist: zenpy; extra == "dev"
+Requires-Dist: tenacity; extra == "dev"
+Requires-Dist: opensearch-py==2.4.2; extra == "dev"
 Requires-Dist: elasticsearch==7.13.4; extra == "dev"
-Requires-Dist: pandas; extra == "dev"
-Requires-Dist: acryl-datahub[dev]; extra == "dev"
-Requires-Dist: pydantic<2; extra == "dev"
-Requires-Dist: openai; extra == "dev"
-Requires-Dist: duckdb; extra == "dev"
-Requires-Dist: pyarrow<=18.0.0; extra == "dev"
 Requires-Dist: termcolor==2.5.0; extra == "dev"
-Requires-Dist: boto3; extra == "dev"
 Dynamic: provides-extra
 Dynamic: requires-dist
+Dynamic: requires-python