great-expectations-cloud 20250811.1.dev0__py3-none-any.whl → 20260113.0.dev1__py3-none-any.whl

great_expectations_cloud/agent/actions/agent_action.py

@@ -4,7 +4,6 @@ import datetime
 from abc import abstractmethod
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Generic, Optional, TypeVar, Union
-from uuid import UUID
 
 from pydantic.v1 import BaseModel
 
@@ -12,6 +11,7 @@ from great_expectations_cloud.agent.models import (
     AgentBaseExtraForbid,
     AgentBaseExtraIgnore,
     CreatedResource,
+    DomainContext,
 )
 
 if TYPE_CHECKING:
@@ -32,11 +32,11 @@ _EventT = TypeVar("_EventT", bound=Union[AgentBaseExtraForbid, AgentBaseExtraIgn
 
 class AgentAction(Generic[_EventT]):
     def __init__(
-        self, context: CloudDataContext, base_url: str, organization_id: UUID, auth_key: str
+        self, context: CloudDataContext, base_url: str, domain_context: DomainContext, auth_key: str
     ):
         self._context = context
         self._base_url = base_url
-        self._organization_id = organization_id
+        self._domain_context = domain_context
         self._auth_key = auth_key
 
     @abstractmethod

great_expectations_cloud/agent/actions/draft_datasource_config_action.py

@@ -70,7 +70,7 @@ class DraftDatasourceConfigAction(AgentAction[DraftDatasourceConfigEvent]):
         with create_session(access_token=self._auth_key) as session:
             url = urljoin(
                 base=self._base_url,
-                url=f"/api/v1/organizations/{self._organization_id}/draft-table-names/{config_id}",
+                url=f"/api/v1/organizations/{self._domain_context.organization_id}/workspaces/{self._domain_context.workspace_id}/draft-table-names/{config_id}",
             )
             response = session.put(
                 url=url,
@@ -87,7 +87,7 @@ class DraftDatasourceConfigAction(AgentAction[DraftDatasourceConfigEvent]):
     def get_draft_config(self, config_id: UUID) -> dict[str, Any]:
         resource_url = urljoin(
             base=self._base_url,
-            url=f"/api/v1/organizations/{self._organization_id}/draft-datasources/{config_id}",
+            url=f"/api/v1/organizations/{self._domain_context.organization_id}/workspaces/{self._domain_context.workspace_id}/draft-datasources/{config_id}",
         )
         with create_session(access_token=self._auth_key) as session:
             response = session.get(resource_url)

great_expectations_cloud/agent/actions/generate_data_quality_check_expectations_action.py

@@ -10,9 +10,13 @@ from uuid import UUID
 
 import great_expectations.expectations as gx_expectations
 from great_expectations.core.http import create_session
-from great_expectations.exceptions import GXCloudError, InvalidExpectationConfigurationError
+from great_expectations.exceptions import (
+    GXCloudError,
+    InvalidExpectationConfigurationError,
+)
 from great_expectations.expectations.metadata_types import (
     DataQualityIssues,
+    FailureSeverity,
 )
 from great_expectations.expectations.window import Offset, Window
 from great_expectations.experimental.metric_repository.batch_inspector import (
@@ -39,6 +43,7 @@ from great_expectations_cloud.agent.event_handler import register_event_action
 from great_expectations_cloud.agent.exceptions import GXAgentError
 from great_expectations_cloud.agent.models import (
     CreatedResource,
+    DomainContext,
     GenerateDataQualityCheckExpectationsEvent,
 )
 from great_expectations_cloud.agent.utils import (
@@ -48,9 +53,7 @@ from great_expectations_cloud.agent.utils import (
 )
 
 if TYPE_CHECKING:
-    from great_expectations.core.suite_parameters import (
-        SuiteParameterDict,
-    )
+    from great_expectations.core.suite_parameters import SuiteParameterDict
     from great_expectations.data_context import CloudDataContext
     from great_expectations.datasource.fluent import DataAsset
 
@@ -81,13 +84,13 @@ class GenerateDataQualityCheckExpectationsAction(
         self,
         context: CloudDataContext,
         base_url: str,
-        organization_id: UUID,
+        domain_context: DomainContext,
         auth_key: str,
         metric_repository: MetricRepository | None = None,
         batch_inspector: BatchInspector | None = None,
     ):
         super().__init__(
-            context=context, base_url=base_url, organization_id=organization_id, auth_key=auth_key
+            context=context, base_url=base_url, domain_context=domain_context, auth_key=auth_key
         )
         self._metric_repository = metric_repository or MetricRepository(
             data_store=CloudDataStore(self._context)
@@ -155,6 +158,7 @@ class GenerateDataQualityCheckExpectationsAction(
                             asset_id=data_asset.id,
                             pre_existing_completeness_change_expectations=pre_existing_completeness_change_expectations,
                             created_via=created_via,
+                            use_forecast=event.use_forecast,
                         )
                         for exp_id in completeness_change_expectation_ids:
                             created_resources.append(
@@ -201,7 +205,7 @@ class GenerateDataQualityCheckExpectationsAction(
             metric_list=[
                 MetricTypes.TABLE_COLUMNS,
                 MetricTypes.TABLE_COLUMN_TYPES,
-                MetricTypes.COLUMN_NULL_COUNT,
+                MetricTypes.COLUMN_NON_NULL_COUNT,
                 MetricTypes.TABLE_ROW_COUNT,
             ],
         )
@@ -220,7 +224,7 @@ class GenerateDataQualityCheckExpectationsAction(
         """
         url = urljoin(
             base=self._base_url,
-            url=f"/api/v1/organizations/{self._organization_id}/expectations/",
+            url=f"/api/v1/organizations/{self._domain_context.organization_id}/workspaces/{self._domain_context.workspace_id}/expectations/",
         )
         with create_session(access_token=self._auth_key) as session:
             response = session.get(
@@ -316,6 +320,7 @@ class GenerateDataQualityCheckExpectationsAction(
             strict_max=strict_max,
             min_value=min_value,
             max_value=max_value,
+            severity=FailureSeverity.WARNING,
         )
         expectation_id = self._create_expectation_for_asset(
             expectation=expectation, asset_id=asset_id, created_via=created_via
@@ -338,7 +343,8 @@ class GenerateDataQualityCheckExpectationsAction(
             raise RuntimeError("missing TABLE_COLUMNS metric")  # noqa: TRY003
 
         expectation = gx_expectations.ExpectTableColumnsToMatchSet(
-            column_set=table_columns_metric.value
+            column_set=table_columns_metric.value,
+            severity=FailureSeverity.WARNING,
         )
         expectation_id = self._create_expectation_for_asset(
             expectation=expectation, asset_id=asset_id, created_via=created_via
@@ -353,6 +359,7 @@ class GenerateDataQualityCheckExpectationsAction(
             dict[Any, Any]
         ],  # list of ExpectationConfiguration dicts
         created_via: str | None,
+        use_forecast: bool = False,
     ) -> list[UUID]:
         table_row_count = next(
             metric
@@ -367,11 +374,11 @@ class GenerateDataQualityCheckExpectationsAction(
             metric
             for metric in metric_run.metrics
             if isinstance(metric, ColumnMetric)
-            and metric.metric_name == MetricTypes.COLUMN_NULL_COUNT
+            and metric.metric_name == MetricTypes.COLUMN_NON_NULL_COUNT
         ]
 
         if not column_null_values_metric or len(column_null_values_metric) == 0:
-            raise RuntimeError("missing COLUMN_NULL_COUNT metrics")  # noqa: TRY003
+            raise RuntimeError("missing COLUMN_NON_NULL_COUNT metrics")  # noqa: TRY003
 
         expectation_ids = []
         # Single-expectation approach using ExpectColumnProportionOfNonNullValuesToBeBetween
@@ -382,7 +389,7 @@ class GenerateDataQualityCheckExpectationsAction(
         )
         for column in columns_missing_completeness_coverage:
             column_name = column.column
-            null_count = column.value
+            non_null_count = column.value
             row_count = table_row_count.value
             expectation: gx_expectations.Expectation
 
@@ -392,18 +399,42 @@ class GenerateDataQualityCheckExpectationsAction(
             max_param_name = f"{unique_id}_proportion_max"
 
             # Calculate non-null proportion
-            non_null_count = row_count - null_count if row_count > 0 else 0
             non_null_proportion = non_null_count / row_count if row_count > 0 else 0
 
-            if non_null_proportion == 0:
+            if use_forecast:
+                expectation = gx_expectations.ExpectColumnProportionOfNonNullValuesToBeBetween(
+                    windows=[
+                        Window(
+                            constraint_fn=ExpectationConstraintFunction.FORECAST,
+                            parameter_name=min_param_name,
+                            range=1,
+                            offset=Offset(positive=0.0, negative=0.0),
+                            strict=True,
+                        ),
+                        Window(
+                            constraint_fn=ExpectationConstraintFunction.FORECAST,
+                            parameter_name=max_param_name,
+                            range=1,
+                            offset=Offset(positive=0.0, negative=0.0),
+                            strict=True,
+                        ),
+                    ],
+                    column=column_name,
+                    min_value={"$PARAMETER": min_param_name},
+                    max_value={"$PARAMETER": max_param_name},
+                    severity=FailureSeverity.WARNING,
+                )
+            elif non_null_proportion == 0:
                 expectation = gx_expectations.ExpectColumnProportionOfNonNullValuesToBeBetween(
                     column=column_name,
                     max_value=0,
+                    severity=FailureSeverity.WARNING,
                 )
             elif non_null_proportion == 1:
                 expectation = gx_expectations.ExpectColumnProportionOfNonNullValuesToBeBetween(
                     column=column_name,
                     min_value=1,
+                    severity=FailureSeverity.WARNING,
                 )
             else:
                 # Use triangular interpolation to compute min/max values
@@ -435,6 +466,7 @@ class GenerateDataQualityCheckExpectationsAction(
                     column=column_name,
                     min_value={"$PARAMETER": min_param_name},
                     max_value={"$PARAMETER": max_param_name},
+                    severity=FailureSeverity.WARNING,
                 )
 
             expectation_id = self._create_expectation_for_asset(
@@ -493,7 +525,7 @@ class GenerateDataQualityCheckExpectationsAction(
     ) -> UUID:
         url = urljoin(
             base=self._base_url,
-            url=f"/api/v1/organizations/{self._organization_id}/expectations/{asset_id}",
+            url=f"/api/v1/organizations/{self._domain_context.organization_id}/workspaces/{self._domain_context.workspace_id}/expectations/{asset_id}",
         )
 
         expectation_payload = expectation.configuration.to_json_dict()
@@ -505,15 +537,6 @@ class GenerateDataQualityCheckExpectationsAction(
         expectation_type = expectation_payload.pop("type")
         expectation_payload["expectation_type"] = expectation_type
 
-        # Add failure severity to kwargs
-        if "kwargs" not in expectation_payload:
-            expectation_payload["kwargs"] = {}
-        if not isinstance(expectation_payload["kwargs"], dict):
-            raise InvalidExpectationConfigurationError(  # noqa: TRY003 # one off error
-                "Expectation configuration kwargs must be a dict."
-            )
-        expectation_payload["kwargs"]["severity"] = "warning"
-
         with create_session(access_token=self._auth_key) as session:
             response = session.post(url=url, json=expectation_payload)
 

great_expectations_cloud/agent/actions/list_asset_names.py

@@ -13,9 +13,7 @@ from great_expectations_cloud.agent.actions.agent_action import (
 )
 from great_expectations_cloud.agent.actions.utils import get_asset_names
 from great_expectations_cloud.agent.event_handler import register_event_action
-from great_expectations_cloud.agent.models import (
-    ListAssetNamesEvent,
-)
+from great_expectations_cloud.agent.models import ListAssetNamesEvent
 
 
 class ListAssetNamesAction(AgentAction[ListAssetNamesEvent]):
@@ -34,7 +32,8 @@ class ListAssetNamesAction(AgentAction[ListAssetNamesEvent]):
         asset_names = get_asset_names(datasource)
 
         self._add_or_update_asset_names_list(
-            datasource_id=str(datasource.id), asset_names=asset_names
+            datasource_id=str(datasource.id),
+            asset_names=asset_names,
         )
 
         return ActionResult(
@@ -47,7 +46,7 @@ class ListAssetNamesAction(AgentAction[ListAssetNamesEvent]):
         with create_session(access_token=self._auth_key) as session:
             url = urljoin(
                 base=self._base_url,
-                url=f"/api/v1/organizations/{self._organization_id}/table-names/{datasource_id}",
+                url=f"/api/v1/organizations/{self._domain_context.organization_id}/workspaces/{self._domain_context.workspace_id}/table-names/{datasource_id}",
             )
             response = session.put(
                 url=url,

great_expectations_cloud/agent/actions/run_checkpoint.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
 
+import logging
+import socket
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Final
 
 from typing_extensions import override
 
@@ -21,6 +23,8 @@ if TYPE_CHECKING:
     from great_expectations.data_context import CloudDataContext
     from great_expectations.datasource.fluent.interfaces import DataAsset, Datasource
 
+LOGGER: Final[logging.Logger] = logging.getLogger(__name__)
+
 
 class RunCheckpointAction(AgentAction[RunCheckpointEvent]):
     @override
@@ -45,12 +49,28 @@ def run_checkpoint(
     id: str,
     expectation_parameters: dict[str, Any] | None = None,
 ) -> ActionResult:
+    """Run a checkpoint and return the result."""
+    hostname = socket.gethostname()
+    log_extra = {
+        "correlation_id": id,
+        "checkpoint_name": event.checkpoint_name,
+        "hostname": hostname,
+    }
+
     # the checkpoint_name property on possible events is optional for backwards compatibility,
     # but this action requires it in order to run:
     if not event.checkpoint_name:
         raise MissingCheckpointNameError
 
+    LOGGER.debug("Fetching checkpoint from context", extra=log_extra)
     checkpoint = context.checkpoints.get(name=event.checkpoint_name)
+    LOGGER.debug(
+        "Checkpoint fetched successfully",
+        extra={
+            **log_extra,
+            "validation_definitions_count": len(checkpoint.validation_definitions),
+        },
+    )
 
     # only GX-managed Checkpoints are currently validated here and they contain only one validation definition, but
     # the Checkpoint does allow for multiple validation definitions so we'll be defensive and ensure we only test each
@@ -66,15 +86,48 @@ def run_checkpoint(
             )
         data_sources_assets_by_data_source_name[ds_name].assets_by_name[vd.asset.name] = vd.asset
 
-    for data_sources_assets in data_sources_assets_by_data_source_name.values():
+    # Test connections to all datasources and assets
+    for ds_name, data_sources_assets in data_sources_assets_by_data_source_name.items():
         data_source = data_sources_assets.data_source
+        LOGGER.debug(
+            "Testing datasource connection",
+            extra={**log_extra, "datasource_name": ds_name},
+        )
         data_source.test_connection(test_assets=False)  # raises `TestConnectionError` on failure
-        for data_asset in data_sources_assets.assets_by_name.values():
+        LOGGER.debug(
+            "Datasource connection successful",
+            extra={**log_extra, "datasource_name": ds_name},
+        )
+
+        for asset_name, data_asset in data_sources_assets.assets_by_name.items():
+            LOGGER.debug(
+                "Testing data asset connection",
+                extra={**log_extra, "datasource_name": ds_name, "asset_name": asset_name},
+            )
             data_asset.test_connection()  # raises `TestConnectionError` on failure
+            LOGGER.debug(
+                "Data asset connection successful",
+                extra={**log_extra, "datasource_name": ds_name, "asset_name": asset_name},
+            )
 
+    LOGGER.debug(
+        "Running checkpoint",
+        extra={
+            **log_extra,
+            "datasources_count": len(data_sources_assets_by_data_source_name),
+            "has_expectation_parameters": expectation_parameters is not None,
+        },
+    )
     checkpoint_run_result = checkpoint.run(
         batch_parameters=event.splitter_options, expectation_parameters=expectation_parameters
     )
+    LOGGER.debug(
+        "Checkpoint run completed",
+        extra={
+            **log_extra,
+            "run_results_count": len(checkpoint_run_result.run_results),
+        },
+    )
 
     validation_results = checkpoint_run_result.run_results
     created_resources = []
@@ -88,6 +141,14 @@ def run_checkpoint(
         )
         created_resources.append(created_resource)
 
+    LOGGER.debug(
+        "Checkpoint action completed successfully",
+        extra={
+            **log_extra,
+            "created_resources_count": len(created_resources),
+        },
+    )
+
     return ActionResult(
         id=id,
         type=event.type,

great_expectations_cloud/agent/actions/run_metric_list_action.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 from typing import TYPE_CHECKING
-from uuid import UUID
 
 from great_expectations.experimental.metric_repository.batch_inspector import (
     BatchInspector,
@@ -21,6 +20,7 @@ from great_expectations_cloud.agent.actions import ActionResult, AgentAction
 from great_expectations_cloud.agent.event_handler import register_event_action
 from great_expectations_cloud.agent.models import (
     CreatedResource,
+    DomainContext,
     RunMetricsListEvent,
 )
 
@@ -34,13 +34,13 @@ class MetricListAction(AgentAction[RunMetricsListEvent]):
         self,
         context: CloudDataContext,
         base_url: str,
-        organization_id: UUID,
+        domain_context: DomainContext,
         auth_key: str,
         metric_repository: MetricRepository | None = None,
         batch_inspector: BatchInspector | None = None,
     ):
         super().__init__(
-            context=context, base_url=base_url, organization_id=organization_id, auth_key=auth_key
+            context=context, base_url=base_url, domain_context=domain_context, auth_key=auth_key
         )
         self._metric_repository = metric_repository or MetricRepository(
             data_store=CloudDataStore(self._context)

great_expectations_cloud/agent/actions/run_scheduled_checkpoint.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+import logging
+import socket
+from typing import TYPE_CHECKING, Final
 from urllib.parse import urljoin
 
 from great_expectations.core.http import create_session
@@ -13,20 +15,20 @@ from great_expectations_cloud.agent.actions.agent_action import (
 )
 from great_expectations_cloud.agent.actions.run_checkpoint import run_checkpoint
 from great_expectations_cloud.agent.event_handler import register_event_action
-from great_expectations_cloud.agent.models import (
-    RunScheduledCheckpointEvent,
-)
+from great_expectations_cloud.agent.models import RunScheduledCheckpointEvent
 
 if TYPE_CHECKING:
     from great_expectations.data_context import CloudDataContext
 
+LOGGER: Final[logging.Logger] = logging.getLogger(__name__)
+
 
 class RunScheduledCheckpointAction(AgentAction[RunScheduledCheckpointEvent]):
     @override
     def run(self, event: RunScheduledCheckpointEvent, id: str) -> ActionResult:
         expectation_parameters_url = urljoin(
             base=self._base_url,
-            url=f"/api/v1/organizations/{self._organization_id}/checkpoints/{event.checkpoint_id}/expectation-parameters",
+            url=f"/api/v1/organizations/{self._domain_context.organization_id}/workspaces/{self._domain_context.workspace_id}/checkpoints/{event.checkpoint_id}/expectation-parameters",
         )
         return run_scheduled_checkpoint(
             context=self._context,
@@ -40,16 +42,32 @@ class RunScheduledCheckpointAction(AgentAction[RunScheduledCheckpointEvent]):
 def run_scheduled_checkpoint(
     context: CloudDataContext, event: RunScheduledCheckpointEvent, id: str, auth_key: str, url: str
 ) -> ActionResult:
+    """Run a scheduled checkpoint, fetching expectation parameters from GX Cloud first."""
+    hostname = socket.gethostname()
+    log_extra = {
+        "correlation_id": id,
+        "checkpoint_id": str(event.checkpoint_id),
+        "schedule_id": str(event.schedule_id),
+        "hostname": hostname,
+    }
+
+    LOGGER.debug("Fetching expectation parameters from GX Cloud", extra=log_extra)
     with create_session(access_token=auth_key) as session:
         response = session.get(url=url)
 
     if not response.ok:
+        LOGGER.error(
+            "Failed to fetch expectation parameters",
+            extra={**log_extra, "response_status": response.status_code},
+        )
         raise GXCloudError(
             message=f"RunScheduledCheckpointAction encountered an error while connecting to GX Cloud. "
             f"Unable to retrieve expectation_parameters for Checkpoint with ID={event.checkpoint_id}.",
             response=response,
         )
     data = response.json()
+    LOGGER.debug("Expectation parameters fetched successfully", extra=log_extra)
+
     try:
         expectation_parameters = (
             data["data"]["expectation_parameters"]
@@ -57,11 +75,16 @@ def run_scheduled_checkpoint(
             else None
         )
     except KeyError as e:
+        LOGGER.exception("Malformed response from GX Cloud", extra=log_extra)
         raise GXCloudError(
             message="Malformed response received from GX Cloud",
             response=response,
         ) from e
 
+    LOGGER.debug(
+        "Proceeding to run checkpoint",
+        extra={**log_extra, "has_expectation_parameters": expectation_parameters is not None},
+    )
     return run_checkpoint(context, event, id, expectation_parameters=expectation_parameters)
 
 

great_expectations_cloud/agent/actions/run_window_checkpoint.py

@@ -13,9 +13,7 @@ from great_expectations_cloud.agent.actions.agent_action import (
 )
 from great_expectations_cloud.agent.actions.run_checkpoint import run_checkpoint
 from great_expectations_cloud.agent.event_handler import register_event_action
-from great_expectations_cloud.agent.models import (
-    RunWindowCheckpointEvent,
-)
+from great_expectations_cloud.agent.models import RunWindowCheckpointEvent
 
 if TYPE_CHECKING:
     from great_expectations.data_context import CloudDataContext
@@ -26,7 +24,7 @@ class RunWindowCheckpointAction(AgentAction[RunWindowCheckpointEvent]):
     def run(self, event: RunWindowCheckpointEvent, id: str) -> ActionResult:
         expectation_parameters_url = urljoin(
             base=self._base_url,
-            url=f"/api/v1/organizations/{self._organization_id}/checkpoints/{event.checkpoint_id}/expectation-parameters",
+            url=f"/api/v1/organizations/{self._domain_context.organization_id}/workspaces/{self._domain_context.workspace_id}/checkpoints/{event.checkpoint_id}/expectation-parameters",
         )
         return run_window_checkpoint(
             self._context,

great_expectations_cloud/agent/actions/utils.py

@@ -7,10 +7,13 @@ from sqlalchemy import inspect
 
 if TYPE_CHECKING:
     from sqlalchemy.engine import Inspector
+    from sqlalchemy.sql.compiler import IdentifierPreparer
 
 
 def get_asset_names(datasource: SQLDatasource) -> list[str]:
     inspector: Inspector = inspect(datasource.get_engine())
+    identifier_preparer: IdentifierPreparer = inspector.dialect.identifier_preparer
+
     if isinstance(datasource, SnowflakeDatasource) and datasource.schema_:
         # Snowflake-SQLAlchemy uses the default_schema if no schema is provided to get_table_names
         # Or if the role does not have access to the schema (it silently fails and defaults to using default_schema)
@@ -19,8 +22,14 @@ def get_asset_names(datasource: SQLDatasource) -> list[str]:
         # Also converting to list to ensure JSON serializable
         tables = list(inspector.get_table_names(schema=datasource.schema_))
         views = list(inspector.get_view_names(schema=datasource.schema_))
-        return tables + views
+        asset_names = tables + views
+    else:
+        tables = list(inspector.get_table_names())
+        views = list(inspector.get_view_names())
+        asset_names = tables + views
 
-    tables = list(inspector.get_table_names())
-    views = list(inspector.get_view_names())
-    return tables + views
+    # the identifier preparer adds quotes when they are necessary
+    quoted_asset_names: list[str] = [
+        identifier_preparer.quote(asset_name) for asset_name in asset_names
+    ]
+    return quoted_asset_names