great-expectations-cloud 20250912.0.dev3__tar.gz → 20250915.0.dev0__tar.gz

{great_expectations_cloud-20250912.0.dev3 → great_expectations_cloud-20250915.0.dev0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: great_expectations_cloud
-Version: 20250912.0.dev3
+Version: 20250915.0.dev0
 Summary: Great Expectations Cloud
 License: Proprietary
 Author: The Great Expectations Team

{great_expectations_cloud-20250912.0.dev3 → great_expectations_cloud-20250915.0.dev0}/great_expectations_cloud/agent/actions/draft_datasource_config_action.py

@@ -40,7 +40,9 @@ class DraftDatasourceConfigAction(AgentAction[DraftDatasourceConfigEvent]):
     def check_draft_datasource_config(
         self, event: DraftDatasourceConfigEvent, id: str
     ) -> ActionResult:
-        draft_config = self.get_draft_config(config_id=event.config_id)
+        draft_config = self.get_draft_config(
+            config_id=event.config_id, workspace_id=event.workspace_id
+        )
         datasource_type = draft_config.get("type", None)
         if datasource_type is None:
             raise TypeError(  # noqa: TRY003 # one off error
@@ -58,7 +60,9 @@ class DraftDatasourceConfigAction(AgentAction[DraftDatasourceConfigEvent]):
 
         if isinstance(datasource, SQLDatasource):
             asset_names = get_asset_names(datasource)
-            self._update_asset_names_list(config_id=event.config_id, asset_names=asset_names)
+            self._update_asset_names_list(
+                config_id=event.config_id, workspace_id=event.workspace_id, asset_names=asset_names
+            )
 
         return ActionResult(
             id=id,
@@ -66,11 +70,13 @@ class DraftDatasourceConfigAction(AgentAction[DraftDatasourceConfigEvent]):
             created_resources=[],
         )
 
-    def _update_asset_names_list(self, config_id: UUID, asset_names: list[str]) -> None:
+    def _update_asset_names_list(
+        self, config_id: UUID, workspace_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._organization_id}/draft-table-names/{config_id}",
+                url=f"/api/v1/organizations/{self._organization_id}/workspaces/{workspace_id}/draft-table-names/{config_id}",
             )
             response = session.put(
                 url=url,
@@ -84,10 +90,10 @@ class DraftDatasourceConfigAction(AgentAction[DraftDatasourceConfigEvent]):
                 f"={config_id}.",
             )
 
-    def get_draft_config(self, config_id: UUID) -> dict[str, Any]:
+    def get_draft_config(self, config_id: UUID, workspace_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._organization_id}/workspaces/{workspace_id}/draft-datasources/{config_id}",
         )
         with create_session(access_token=self._auth_key) as session:
             response = session.get(resource_url)

{great_expectations_cloud-20250912.0.dev3 → great_expectations_cloud-20250915.0.dev0}/great_expectations_cloud/agent/actions/generate_data_quality_check_expectations_action.py

@@ -10,7 +10,10 @@ 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,
@@ -49,9 +52,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
 
@@ -99,6 +100,8 @@ class GenerateDataQualityCheckExpectationsAction(
 
     @override
     def run(self, event: GenerateDataQualityCheckExpectationsEvent, id: str) -> ActionResult:
+        self._workspace_id = event.workspace_id
+
         created_resources: list[CreatedResource] = []
         assets_with_errors: list[str] = []
         selected_dqis: Sequence[DataQualityIssues] = event.selected_data_quality_issues or []
@@ -222,7 +225,7 @@ class GenerateDataQualityCheckExpectationsAction(
         """
         url = urljoin(
             base=self._base_url,
-            url=f"/api/v1/organizations/{self._organization_id}/expectations/",
+            url=f"/api/v1/organizations/{self._organization_id}/workspaces/{self._workspace_id}/expectations/",
         )
         with create_session(access_token=self._auth_key) as session:
             response = session.get(
@@ -524,7 +527,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._organization_id}/workspaces/{self._workspace_id}/expectations/{asset_id}",
         )
 
         expectation_payload = expectation.configuration.to_json_dict()

{great_expectations_cloud-20250912.0.dev3 → great_expectations_cloud-20250915.0.dev0}/great_expectations_cloud/agent/actions/list_asset_names.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from urllib.parse import urljoin
+from uuid import UUID
 
 from great_expectations.core.http import create_session
 from great_expectations.datasource.fluent import SQLDatasource
@@ -13,9 +14,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 +33,9 @@ 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),
+            workspace_id=event.workspace_id,
+            asset_names=asset_names,
         )
 
         return ActionResult(
@@ -43,11 +44,13 @@ class ListAssetNamesAction(AgentAction[ListAssetNamesEvent]):
             created_resources=[],
         )
 
-    def _add_or_update_asset_names_list(self, datasource_id: str, asset_names: list[str]) -> None:
+    def _add_or_update_asset_names_list(
+        self, datasource_id: str, workspace_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._organization_id}/table-names/{datasource_id}",
+                url=f"/api/v1/organizations/{self._organization_id}/workspaces/{workspace_id}/table-names/{datasource_id}",
             )
             response = session.put(
                 url=url,

{great_expectations_cloud-20250912.0.dev3 → great_expectations_cloud-20250915.0.dev0}/great_expectations_cloud/agent/actions/run_scheduled_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 (
-    RunScheduledCheckpointEvent,
-)
+from great_expectations_cloud.agent.models import RunScheduledCheckpointEvent
 
 if TYPE_CHECKING:
     from great_expectations.data_context import CloudDataContext
@@ -26,7 +24,7 @@ class RunScheduledCheckpointAction(AgentAction[RunScheduledCheckpointEvent]):
     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._organization_id}/workspaces/{event.workspace_id}/checkpoints/{event.checkpoint_id}/expectation-parameters",
         )
         return run_scheduled_checkpoint(
             context=self._context,

{great_expectations_cloud-20250912.0.dev3 → great_expectations_cloud-20250915.0.dev0}/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._organization_id}/workspaces/{event.workspace_id}/checkpoints/{event.checkpoint_id}/expectation-parameters",
         )
         return run_window_checkpoint(
             self._context,

{great_expectations_cloud-20250912.0.dev3 → great_expectations_cloud-20250915.0.dev0}/great_expectations_cloud/agent/agent.py

@@ -17,13 +17,12 @@ from uuid import UUID
 
 import orjson
 import requests
+from great_expectations import __version__, get_context
+from great_expectations.core import http
 from great_expectations.core.http import create_session
 from great_expectations.data_context.cloud_constants import CLOUD_DEFAULT_BASE_URL
-from great_expectations.data_context.data_context.context_factory import get_context
 from great_expectations.data_context.types.base import ProgressBarsConfig
-from pika.adapters.utils.connection_workflow import (
-    AMQPConnectorException,
-)
+from pika.adapters.utils.connection_workflow import AMQPConnectorException
 from pika.exceptions import (
     AMQPConnectionError,
     AMQPError,
@@ -46,9 +45,7 @@ from great_expectations_cloud.agent.config import (
     generate_config_validation_error_text,
 )
 from great_expectations_cloud.agent.constants import USER_AGENT_HEADER, HeaderName
-from great_expectations_cloud.agent.event_handler import (
-    EventHandler,
-)
+from great_expectations_cloud.agent.event_handler import EventHandler
 from great_expectations_cloud.agent.exceptions import (
     GXAgentConfigError,
     GXAgentError,
@@ -151,19 +148,6 @@ class GXAgent:
                 "great_expectations_version": great_expectations_version,
             },
         )
-        LOGGER.debug("Loading a DataContext - this might take a moment.")
-
-        with warnings.catch_warnings():
-            # suppress warnings about GX version
-            warnings.filterwarnings("ignore", message="You are using great_expectations version")
-            self._context: CloudDataContext = get_context(
-                cloud_mode=True,
-                user_agent_str=self.user_agent_str,
-            )
-            self._configure_progress_bars(data_context=self._context)
-        LOGGER.debug("DataContext is ready.")
-
-        self._set_http_session_headers(data_context=self._context)
 
         # Create a thread pool with a single worker, so we can run long-lived
         # GX processes and maintain our connection to the broker. Note that
@@ -264,6 +248,7 @@ class GXAgent:
                     "event_type": event_context.event.type,
                     "correlation_id": event_context.correlation_id,
                     "organization_id": self.get_organization_id(event_context),
+                    "workspace_id": str(self.get_workspace_id(event_context)),
                     "schedule_id": event_context.event.schedule_id
                     if isinstance(event_context.event, ScheduledEventBase)
                     else None,
@@ -288,8 +273,23 @@ class GXAgent:
             self._current_task.add_done_callback(on_exit_callback)
 
     def get_data_context(self, event_context: EventContext) -> CloudDataContext:
-        """Helper method to get a DataContext Agent. Overridden in GX-Runner."""
-        return self._context
+        """Create a new CloudDataContext for each job using the event's workspace_id."""
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", message="You are using great_expectations version")
+            workspace_id = self.get_workspace_id(event_context)
+
+            LOGGER.debug("Loading a DataContext - this might take a moment.")
+
+            context: CloudDataContext = get_context(
+                cloud_mode=True,
+                user_agent_str=self.user_agent_str,
+                cloud_workspace_id=str(workspace_id),
+            )
+            self._configure_progress_bars(data_context=context)
+
+            LOGGER.debug("DataContext is ready.")
+
+        return context
 
     def get_organization_id(self, event_context: EventContext) -> UUID:
         """Helper method to get the organization ID. Overridden in GX-Runner."""
@@ -299,6 +299,13 @@ class GXAgent:
         """Helper method to get the auth key. Overridden in GX-Runner."""
         return self._get_config().gx_cloud_access_token
 
+    def get_workspace_id(self, event_context: EventContext) -> UUID:
+        """Helper method to get the workspace ID from the event."""
+        workspace_id: UUID | None = getattr(event_context.event, "workspace_id", None)
+        if workspace_id is None:
+            raise GXAgentError()
+        return workspace_id
+
     def _set_sentry_tags(self, even_context: EventContext) -> None:
         """Used by GX-Runner to set tags for Sentry logging. No-op in the Agent."""
         pass
@@ -321,14 +328,18 @@ class GXAgent:
         )
 
         org_id = self.get_organization_id(event_context)
+        workspace_id = self.get_workspace_id(event_context)
         base_url = self._get_config().gx_cloud_base_url
         auth_key = self.get_auth_key()
 
         if isinstance(event_context.event, ScheduledEventBase):
-            self._create_scheduled_job_and_set_started(event_context, org_id)
+            self._create_scheduled_job_and_set_started(event_context, org_id, workspace_id)
         else:
             self._update_status(
-                correlation_id=event_context.correlation_id, status=JobStarted(), org_id=org_id
+                correlation_id=event_context.correlation_id,
+                status=JobStarted(),
+                org_id=org_id,
+                workspace_id=workspace_id,
             )
         LOGGER.info(
             "Starting job",
@@ -336,6 +347,7 @@ class GXAgent:
                 "event_type": event_context.event.type,
                 "correlation_id": event_context.correlation_id,
                 "organization_id": str(org_id),
+                "workspace_id": str(workspace_id),
                 "schedule_id": event_context.event.schedule_id
                 if isinstance(event_context.event, ScheduledEventBase)
                 else None,
@@ -367,6 +379,7 @@ class GXAgent:
         # warning:  this method will not be executed in the main thread
 
         org_id = self.get_organization_id(event_context)
+        workspace_id = self.get_workspace_id(event_context)
 
         # get results or errors from the thread
         error = future.exception()
@@ -386,6 +399,7 @@ class GXAgent:
                         "event_type": event_context.event.type,
                         "id": event_context.correlation_id,
                         "organization_id": str(org_id),
+                        "workspace_id": str(workspace_id),
                         "schedule_id": event_context.event.schedule_id
                         if isinstance(event_context.event, ScheduledEventBase)
                         else None,
@@ -406,6 +420,7 @@ class GXAgent:
                             result.job_duration.total_seconds() if result.job_duration else None
                         ),
                         "organization_id": str(org_id),
+                        "workspace_id": str(workspace_id),
                         "schedule_id": event_context.event.schedule_id
                         if isinstance(event_context.event, ScheduledEventBase)
                         else None,
@@ -420,12 +435,16 @@ class GXAgent:
                     "event_type": event_context.event.type,
                     "correlation_id": event_context.correlation_id,
                     "organization_id": str(org_id),
+                    "workspace_id": str(workspace_id),
                 },
             )
 
         try:
             self._update_status(
-                correlation_id=event_context.correlation_id, status=status, org_id=org_id
+                correlation_id=event_context.correlation_id,
+                status=status,
+                org_id=org_id,
+                workspace_id=workspace_id,
             )
         except Exception:
             LOGGER.exception(
@@ -434,6 +453,7 @@ class GXAgent:
                     "correlation_id": event_context.correlation_id,
                     "status": str(status),
                     "organization_id": str(org_id),
+                    "workspace_id": str(workspace_id),
                 },
             )
             # We do not want to cause an infinite loop of errors
@@ -553,7 +573,9 @@ class GXAgent:
                 )
             )
 
-    def _update_status(self, correlation_id: str, status: JobStatus, org_id: UUID) -> None:
+    def _update_status(
+        self, correlation_id: str, status: JobStatus, org_id: UUID, workspace_id: UUID
+    ) -> None:
         """Update GX Cloud on the status of a job.
 
         Args:
@@ -566,11 +588,12 @@ class GXAgent:
                 "correlation_id": correlation_id,
                 "status": str(status),
                 "organization_id": str(org_id),
+                "workspace_id": str(workspace_id),
             },
         )
         agent_sessions_url = urljoin(
             self._get_config().gx_cloud_base_url,
-            f"/api/v1/organizations/{org_id}/agent-jobs/{correlation_id}",
+            f"/api/v1/organizations/{org_id}/workspaces/{workspace_id}/agent-jobs/{correlation_id}",
         )
         with create_session(access_token=self.get_auth_key()) as session:
             data = UpdateJobStatusRequest(data=status).json()
@@ -581,6 +604,7 @@ class GXAgent:
                     "correlation_id": correlation_id,
                     "status": str(status),
                     "organization_id": str(org_id),
+                    "workspace_id": str(workspace_id),
                 },
             )
             GXAgent._log_http_error(
@@ -588,7 +612,7 @@ class GXAgent:
             )
 
     def _create_scheduled_job_and_set_started(
-        self, event_context: EventContext, org_id: UUID
+        self, event_context: EventContext, org_id: UUID, workspace_id: UUID
     ) -> None:
         """Create a job in GX Cloud for scheduled events.
 
@@ -610,13 +634,14 @@ class GXAgent:
                 "correlation_id": str(event_context.correlation_id),
                 "event_type": str(event_context.event.type),
                 "organization_id": str(org_id),
+                "workspace_id": str(workspace_id),
                 "schedule_id": str(event_context.event.schedule_id),
             },
         )
 
         agent_sessions_url = urljoin(
             self._get_config().gx_cloud_base_url,
-            f"/api/v1/organizations/{org_id}/agent-jobs",
+            f"/api/v1/organizations/{org_id}/workspaces/{workspace_id}/agent-jobs",
         )
         data = CreateScheduledJobAndSetJobStarted(
             type="run_scheduled_checkpoint.received",
@@ -637,6 +662,7 @@ class GXAgent:
                     "event_type": str(event_context.event.type),
                     "organization_id": str(org_id),
                     "schedule_id": str(event_context.event.schedule_id),
+                    "workspace_id": str(workspace_id),
                 },
             )
             GXAgent._log_http_error(
@@ -659,7 +685,9 @@ class GXAgent:
         """
         Sets headers on all stores in the data context.
         """
-        from great_expectations.data_context.store.gx_cloud_store_backend import GXCloudStoreBackend  # noqa: I001, PLC0415
+        from great_expectations.data_context.store.gx_cloud_store_backend import (  # noqa: PLC0415
+            GXCloudStoreBackend,
+        )
 
         # OSS doesn't use the same session for all requests, so we need to set the header for each store
         stores = list(data_context.stores.values())
@@ -687,8 +715,6 @@ class GXAgent:
         Note: the Agent-Job-Id header value will be set for all GX Cloud request until this method is
         called again.
         """
-        from great_expectations import __version__  # noqa: PLC0415
-        from great_expectations.core import http  # noqa: PLC0415
 
         header_name = self.get_header_name()
         user_agent_header_value = self.user_agent_str

{great_expectations_cloud-20250912.0.dev3 → great_expectations_cloud-20250915.0.dev0}/great_expectations_cloud/agent/models.py

@@ -39,6 +39,7 @@ class AgentBaseExtraIgnore(BaseModel):
 class EventBase(AgentBaseExtraIgnore):
     type: str
     organization_id: Optional[UUID] = None  # noqa: UP045
+    workspace_id: UUID
 
 
 class ScheduledEventBase(EventBase):

{great_expectations_cloud-20250912.0.dev3 → great_expectations_cloud-20250915.0.dev0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "great_expectations_cloud"
-version = "20250912.0.dev3"
+version = "20250915.0.dev0"
 description = "Great Expectations Cloud"
 authors = ["The Great Expectations Team <team@greatexpectations.io>"]
 repository = "https://github.com/great-expectations/cloud"