great-expectations-cloud 20250902.0.dev1__tar.gz → 20250903.0.dev0__tar.gz

{great_expectations_cloud-20250902.0.dev1 → great_expectations_cloud-20250903.0.dev0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: great_expectations_cloud
-Version: 20250902.0.dev1
+Version: 20250903.0.dev0
 Summary: Great Expectations Cloud
 License: Proprietary
 Author: The Great Expectations Team

{great_expectations_cloud-20250902.0.dev1 → great_expectations_cloud-20250903.0.dev0}/great_expectations_cloud/agent/agent.py

@@ -16,6 +16,8 @@ from uuid import UUID
 
 import orjson
 import requests
+from great_expectations import __version__
+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
@@ -150,19 +152,7 @@ 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)
+        # DataContext is created per job in get_data_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
@@ -263,6 +253,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,
@@ -287,8 +278,19 @@ 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")
+            # Validate presence of workspace_id (required), even if not passed directly to get_context
+            if getattr(event_context.event, "workspace_id", None) is None:
+                raise GXAgentError()
+
+            context: CloudDataContext = get_context(
+                cloud_mode=True,
+                user_agent_str=self.user_agent_str,
+            )
+            self._configure_progress_bars(data_context=context)
+        return context
 
     def get_organization_id(self, event_context: EventContext) -> UUID:
         """Helper method to get the organization ID. Overridden in GX-Runner."""
@@ -298,6 +300,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
@@ -320,14 +329,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",
@@ -335,6 +348,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,
@@ -366,6 +380,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()
@@ -385,6 +400,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,
@@ -405,6 +421,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,
@@ -419,12 +436,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(
@@ -433,6 +454,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
@@ -552,7 +574,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:
@@ -565,11 +589,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()
@@ -580,6 +605,7 @@ class GXAgent:
                     "correlation_id": correlation_id,
                     "status": str(status),
                     "organization_id": str(org_id),
+                    "workspace_id": str(workspace_id),
                 },
             )
             GXAgent._log_http_error(
@@ -587,7 +613,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.
 
@@ -609,13 +635,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",
@@ -636,6 +663,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(
@@ -686,8 +714,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-20250902.0.dev1 → great_expectations_cloud-20250903.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-20250902.0.dev1 → great_expectations_cloud-20250903.0.dev0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "great_expectations_cloud"
-version = "20250902.0.dev1"
+version = "20250903.0.dev0"
 
 description = "Great Expectations Cloud"
 authors = ["The Great Expectations Team <team@greatexpectations.io>"]