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

great_expectations_cloud/agent/__init__.py

@@ -2,4 +2,7 @@ from __future__ import annotations
 
 import great_expectations_cloud.agent.actions  # import actions to register them
 from great_expectations_cloud.agent.agent import GXAgent
+from great_expectations_cloud.agent.event_handler import register_event_action
+from great_expectations_cloud.agent.models import Event, EventType
 from great_expectations_cloud.agent.run import get_version, run_agent
+from great_expectations_cloud.agent.utils import triangular_interpolation

great_expectations_cloud/agent/actions/__init__.py

@@ -7,13 +7,16 @@ from great_expectations_cloud.agent.actions.agent_action import (
 )
 
 # Import all actions to register them:
-from great_expectations_cloud.agent.actions.data_assistants import (
-    RunMissingnessDataAssistantAction,
-    RunOnboardingDataAssistantAction,
-)
 from great_expectations_cloud.agent.actions.draft_datasource_config_action import (
     DraftDatasourceConfigAction,
 )
-from great_expectations_cloud.agent.actions.list_table_names import ListTableNamesAction
+from great_expectations_cloud.agent.actions.generate_data_quality_check_expectations_action import (
+    GenerateDataQualityCheckExpectationsAction,
+)
+from great_expectations_cloud.agent.actions.list_asset_names import ListAssetNamesAction
 from great_expectations_cloud.agent.actions.run_checkpoint import RunCheckpointAction
 from great_expectations_cloud.agent.actions.run_metric_list_action import MetricListAction
+from great_expectations_cloud.agent.actions.run_scheduled_checkpoint import (
+    RunScheduledCheckpointAction,
+)
+from great_expectations_cloud.agent.actions.run_window_checkpoint import RunWindowCheckpointAction

great_expectations_cloud/agent/actions/agent_action.py

@@ -1,28 +1,43 @@
 from __future__ import annotations
 
+import datetime
 from abc import abstractmethod
-from typing import TYPE_CHECKING, Generic, Sequence, TypeVar
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Generic, Optional, TypeVar, Union
 
-from great_expectations.compatibility.pydantic import BaseModel
+from pydantic.v1 import BaseModel
 
-from great_expectations_cloud.agent.models import CreatedResource, Event
+from great_expectations_cloud.agent.models import (
+    AgentBaseExtraForbid,
+    AgentBaseExtraIgnore,
+    CreatedResource,
+    DomainContext,
+)
 
 if TYPE_CHECKING:
     from great_expectations.data_context import CloudDataContext
 
 
-class ActionResult(BaseModel):  # type: ignore[misc] # BaseSettings is has Any type
+class ActionResult(BaseModel):
     id: str
     type: str
     created_resources: Sequence[CreatedResource]
+    job_duration: Optional[datetime.timedelta] = (  # noqa: UP045
+        None  # Python 3.8 doesn't support `X | Y` for type annotation
+    )
 
 
-_EventT = TypeVar("_EventT", bound=Event)
+_EventT = TypeVar("_EventT", bound=Union[AgentBaseExtraForbid, AgentBaseExtraIgnore])
 
 
 class AgentAction(Generic[_EventT]):
-    def __init__(self, context: CloudDataContext):
+    def __init__(
+        self, context: CloudDataContext, base_url: str, domain_context: DomainContext, auth_key: str
+    ):
         self._context = context
+        self._base_url = base_url
+        self._domain_context = domain_context
+        self._auth_key = auth_key
 
     @abstractmethod
     def run(self, event: _EventT, id: str) -> ActionResult: ...

great_expectations_cloud/agent/actions/draft_datasource_config_action.py

@@ -1,24 +1,25 @@
 from __future__ import annotations
 
 from typing import Any
+from urllib.parse import urljoin
 from uuid import UUID
 
-from great_expectations.compatibility import pydantic
 from great_expectations.core.http import create_session
+from great_expectations.datasource.fluent import SQLDatasource
 from great_expectations.datasource.fluent.interfaces import TestConnectionError
 from typing_extensions import override
 
 from great_expectations_cloud.agent.actions import ActionResult, AgentAction
-from great_expectations_cloud.agent.config import (
-    GxAgentEnvVars,
-    generate_config_validation_error_text,
-)
+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.exceptions import ErrorCode, raise_with_error_code
 from great_expectations_cloud.agent.models import DraftDatasourceConfigEvent
 
 
 class DraftDatasourceConfigAction(AgentAction[DraftDatasourceConfigEvent]):
+    # TODO: New actions need to be created that are compatible with GX v1 and registered for v1.
+    #  This action is registered for v0, see register_event_action()
+
     @override
     def run(self, event: DraftDatasourceConfigEvent, id: str) -> ActionResult:
         try:
@@ -43,11 +44,10 @@ class DraftDatasourceConfigAction(AgentAction[DraftDatasourceConfigEvent]):
         datasource_type = draft_config.get("type", None)
         if datasource_type is None:
             raise TypeError(  # noqa: TRY003 # one off error
-                "The DraftDatasourceConfigAction can only be used with a "
-                "fluent-style Data Source."
+                "The DraftDatasourceConfigAction can only be used with a fluent-style Data Source."
             )
         try:
-            datasource_cls = self._context.sources.type_lookup[datasource_type]
+            datasource_cls = self._context.data_sources.type_lookup[datasource_type]
         except LookupError as exc:
             raise TypeError(  # noqa: TRY003 # one off error
                 "DraftDatasourceConfigAction received an unknown Data Source type."
@@ -55,32 +55,53 @@ class DraftDatasourceConfigAction(AgentAction[DraftDatasourceConfigEvent]):
         datasource = datasource_cls(**draft_config)
         datasource._data_context = self._context
         datasource.test_connection(test_assets=True)  # raises `TestConnectionError` on failure
-        return ActionResult(id=id, type=event.type, created_resources=[])
 
-    def get_draft_config(self, config_id: UUID) -> dict[str, Any]:
-        try:
-            config = GxAgentEnvVars()
-        except pydantic.ValidationError as validation_err:
-            raise RuntimeError(
-                generate_config_validation_error_text(validation_err)
-            ) from validation_err
-        resource_url = (
-            f"{config.gx_cloud_base_url}/organizations/"
-            f"{config.gx_cloud_organization_id}/datasources/drafts/{config_id}"
+        if isinstance(datasource, SQLDatasource):
+            asset_names = get_asset_names(datasource)
+            self._update_asset_names_list(config_id=event.config_id, asset_names=asset_names)
+
+        return ActionResult(
+            id=id,
+            type=event.type,
+            created_resources=[],
         )
-        session = create_session(access_token=config.gx_cloud_access_token)
-        response = session.get(resource_url)
+
+    def _update_asset_names_list(self, config_id: UUID, asset_names: list[str]) -> None:
+        with create_session(access_token=self._auth_key) as session:
+            url = urljoin(
+                base=self._base_url,
+                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,
+                json={"data": {"table_names": asset_names}},
+            )
         if not response.ok:
             raise RuntimeError(  # noqa: TRY003 # one off error
-                "DraftDatasourceConfigAction encountered an error while " "connecting to GX Cloud"
+                f"DraftDatasourceConfigAction encountered an error while connecting to GX Cloud. "
+                f"Unable to update "
+                f"table_names for Draft Config with ID"
+                f"={config_id}.",
             )
+
+    def get_draft_config(self, config_id: UUID) -> dict[str, Any]:
+        resource_url = urljoin(
+            base=self._base_url,
+            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)
+            if not response.ok:
+                raise RuntimeError(  # noqa: TRY003 # one off error
+                    "DraftDatasourceConfigAction encountered an error while connecting to GX Cloud"
+                )
         data = response.json()
         try:
-            return data["data"]["attributes"]["draft_config"]  # type: ignore[no-any-return]
+            return data["data"]["config"]  # type: ignore[no-any-return]
         except KeyError as e:
             raise RuntimeError(  # noqa: TRY003 # one off error
                 "Malformed response received from GX Cloud"
             ) from e
 
 
-register_event_action("0", DraftDatasourceConfigEvent, DraftDatasourceConfigAction)
+register_event_action("1", DraftDatasourceConfigEvent, DraftDatasourceConfigAction)