great-expectations-cloud 20240523.0.dev0__py3-none-any.whl → 20251124.0.dev1__py3-none-any.whl

great_expectations_cloud/agent/actions/run_checkpoint.py

@@ -1,5 +1,8 @@
 from __future__ import annotations
 
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
 from typing_extensions import override
 
 from great_expectations_cloud.agent.actions.agent_action import (
@@ -10,42 +13,86 @@ from great_expectations_cloud.agent.event_handler import register_event_action
 from great_expectations_cloud.agent.models import (
     CreatedResource,
     RunCheckpointEvent,
+    RunScheduledCheckpointEvent,
+    RunWindowCheckpointEvent,
 )
 
+if TYPE_CHECKING:
+    from great_expectations.data_context import CloudDataContext
+    from great_expectations.datasource.fluent.interfaces import DataAsset, Datasource
+
 
 class RunCheckpointAction(AgentAction[RunCheckpointEvent]):
     @override
     def run(self, event: RunCheckpointEvent, id: str) -> ActionResult:
-        # TODO: move connection testing into OSS; there isn't really a reason it can't be done there
-        for datasource_name, data_asset_names in event.datasource_names_to_asset_names.items():
-            datasource = self._context.get_datasource(datasource_name)
-            datasource.test_connection(test_assets=False)  # raises `TestConnectionError` on failure
-            for data_asset_name in (
-                data_asset_names
-            ):  # only test connection for assets that are validated in checkpoint
-                asset = datasource.get_asset(data_asset_name)
-                asset.test_connection()  # raises `TestConnectionError` on failure
-        checkpoint_run_result = self._context.run_checkpoint(
-            ge_cloud_id=event.checkpoint_id,
-            batch_request={"options": event.splitter_options} if event.splitter_options else None,
-        )
+        return run_checkpoint(self._context, event, id)
+
+
+class MissingCheckpointNameError(ValueError):
+    """Property checkpoint_name is required but not present."""
+
+
+# using a dataclass because we don't want the pydantic behavior of copying objects
+@dataclass
+class DataSourceAssets:
+    data_source: Datasource[Any, Any]
+    assets_by_name: dict[str, DataAsset[Any, Any]]
 
-        validation_results = checkpoint_run_result.run_results
-        created_resources = []
-        for key in validation_results.keys():
-            created_resource = CreatedResource(
-                resource_id=validation_results[key]["actions_results"]["store_validation_result"][
-                    "id"
-                ],
-                type="SuiteValidationResult",
+
+def run_checkpoint(
+    context: CloudDataContext,
+    event: RunCheckpointEvent | RunScheduledCheckpointEvent | RunWindowCheckpointEvent,
+    id: str,
+    expectation_parameters: dict[str, Any] | None = None,
+) -> ActionResult:
+    # 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
+
+    checkpoint = context.checkpoints.get(name=event.checkpoint_name)
+
+    # 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
+    # source/asset once
+    data_sources_assets_by_data_source_name: dict[str, DataSourceAssets] = {}
+    for vd in checkpoint.validation_definitions:
+        ds = vd.data_source
+        ds_name = ds.name
+        # create assets by name dict
+        if ds_name not in data_sources_assets_by_data_source_name:
+            data_sources_assets_by_data_source_name[ds_name] = DataSourceAssets(
+                data_source=ds, assets_by_name={}
             )
-            created_resources.append(created_resource)
+        data_sources_assets_by_data_source_name[ds_name].assets_by_name[vd.asset.name] = vd.asset
 
-        return ActionResult(
-            id=id,
-            type=event.type,
-            created_resources=created_resources,
+    for data_sources_assets in data_sources_assets_by_data_source_name.values():
+        data_source = data_sources_assets.data_source
+        data_source.test_connection(test_assets=False)  # raises `TestConnectionError` on failure
+        for data_asset in data_sources_assets.assets_by_name.values():
+            data_asset.test_connection()  # raises `TestConnectionError` on failure
+
+    checkpoint_run_result = checkpoint.run(
+        batch_parameters=event.splitter_options, expectation_parameters=expectation_parameters
+    )
+
+    validation_results = checkpoint_run_result.run_results
+    created_resources = []
+    for key in validation_results.keys():
+        suite_validation_result = validation_results[key]
+        if suite_validation_result.id is None:
+            raise RuntimeError(f"SuiteValidationResult.id is None for key: {key}")  # noqa: TRY003
+        created_resource = CreatedResource(
+            resource_id=suite_validation_result.id,
+            type="SuiteValidationResult",
         )
+        created_resources.append(created_resource)
+
+    return ActionResult(
+        id=id,
+        type=event.type,
+        created_resources=created_resources,
+    )
 
 
-register_event_action("0", RunCheckpointEvent, RunCheckpointAction)
+register_event_action("1", RunCheckpointEvent, RunCheckpointAction)

great_expectations_cloud/agent/actions/run_metric_list_action.py

@@ -20,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,
 )
 
@@ -29,13 +30,18 @@ if TYPE_CHECKING:
 
 
 class MetricListAction(AgentAction[RunMetricsListEvent]):
-    def __init__(
+    def __init__(  # noqa: PLR0913  # Refactor opportunity
         self,
         context: CloudDataContext,
+        base_url: str,
+        domain_context: DomainContext,
+        auth_key: str,
         metric_repository: MetricRepository | None = None,
         batch_inspector: BatchInspector | None = None,
     ):
-        super().__init__(context=context)
+        super().__init__(
+            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)
         )
@@ -45,7 +51,7 @@ class MetricListAction(AgentAction[RunMetricsListEvent]):
 
     @override
     def run(self, event: RunMetricsListEvent, id: str) -> ActionResult:
-        datasource = self._context.get_datasource(event.datasource_name)
+        datasource = self._context.data_sources.get(event.datasource_name)
         data_asset = datasource.get_asset(event.data_asset_name)
         data_asset.test_connection()  # raises `TestConnectionError` on failure
 
@@ -54,7 +60,7 @@ class MetricListAction(AgentAction[RunMetricsListEvent]):
         metric_run = self._batch_inspector.compute_metric_list_run(
             data_asset_id=data_asset.id,
             batch_request=batch_request,
-            metric_list=event.metric_names,
+            metric_list=list(event.metric_names),
         )
 
         metric_run_id = self._metric_repository.add_metric_run(metric_run)
@@ -78,4 +84,4 @@ class MetricListAction(AgentAction[RunMetricsListEvent]):
             )
 
 
-register_event_action("0", RunMetricsListEvent, MetricListAction)
+register_event_action("1", RunMetricsListEvent, MetricListAction)

great_expectations_cloud/agent/actions/run_scheduled_checkpoint.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+from urllib.parse import urljoin
+
+from great_expectations.core.http import create_session
+from great_expectations.exceptions import GXCloudError
+from typing_extensions import override
+
+from great_expectations_cloud.agent.actions.agent_action import (
+    ActionResult,
+    AgentAction,
+)
+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
+
+if TYPE_CHECKING:
+    from great_expectations.data_context import CloudDataContext
+
+
+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._domain_context.organization_id}/workspaces/{self._domain_context.workspace_id}/checkpoints/{event.checkpoint_id}/expectation-parameters",
+        )
+        return run_scheduled_checkpoint(
+            context=self._context,
+            event=event,
+            id=id,
+            auth_key=self._auth_key,
+            url=expectation_parameters_url,
+        )
+
+
+def run_scheduled_checkpoint(
+    context: CloudDataContext, event: RunScheduledCheckpointEvent, id: str, auth_key: str, url: str
+) -> ActionResult:
+    with create_session(access_token=auth_key) as session:
+        response = session.get(url=url)
+
+    if not response.ok:
+        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()
+
+    try:
+        expectation_parameters = (
+            data["data"]["expectation_parameters"]
+            if len(data["data"]["expectation_parameters"]) > 0
+            else None
+        )
+    except KeyError as e:
+        raise GXCloudError(
+            message="Malformed response received from GX Cloud",
+            response=response,
+        ) from e
+
+    return run_checkpoint(context, event, id, expectation_parameters=expectation_parameters)
+
+
+register_event_action("1", RunScheduledCheckpointEvent, RunScheduledCheckpointAction)

great_expectations_cloud/agent/actions/run_window_checkpoint.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+from urllib.parse import urljoin
+
+from great_expectations.core.http import create_session
+from great_expectations.exceptions import GXCloudError
+from typing_extensions import override
+
+from great_expectations_cloud.agent.actions.agent_action import (
+    ActionResult,
+    AgentAction,
+)
+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
+
+if TYPE_CHECKING:
+    from great_expectations.data_context import CloudDataContext
+
+
+class RunWindowCheckpointAction(AgentAction[RunWindowCheckpointEvent]):
+    @override
+    def run(self, event: RunWindowCheckpointEvent, id: str) -> ActionResult:
+        expectation_parameters_url = urljoin(
+            base=self._base_url,
+            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,
+            event,
+            id,
+            auth_key=self._auth_key,
+            url=expectation_parameters_url,
+        )
+
+
+register_event_action("1", RunWindowCheckpointEvent, RunWindowCheckpointAction)
+
+
+def run_window_checkpoint(
+    context: CloudDataContext,
+    event: RunWindowCheckpointEvent,
+    id: str,
+    auth_key: str,
+    url: str,
+) -> ActionResult:
+    with create_session(access_token=auth_key) as session:
+        response = session.get(url=url)
+
+    if not response.ok:
+        raise GXCloudError(
+            message=f"RunWindowCheckpointAction 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()
+    try:
+        expectation_parameters = data["data"]["expectation_parameters"]
+    except KeyError as e:
+        raise GXCloudError(
+            message="Malformed response received from GX Cloud",
+            response=response,
+        ) from e
+
+    return run_checkpoint(context, event, id, expectation_parameters=expectation_parameters)

great_expectations_cloud/agent/actions/utils.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from great_expectations.datasource.fluent import SnowflakeDatasource, SQLDatasource
+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)
+        # See https://github.com/snowflakedb/snowflake-sqlalchemy/blob/e78319725d4b96ea205ef1264b744c65eb37853d/src/snowflake/sqlalchemy/snowdialect.py#L731
+        # Explicitly passing the schema to the inspector to get the table and view names
+        # 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_))
+        asset_names = tables + views
+    else:
+        tables = list(inspector.get_table_names())
+        views = list(inspector.get_view_names())
+        asset_names = 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