great-expectations-cloud 20250902.0.dev1__py3-none-any.whl → 20250903.0.dev0__py3-none-any.whl

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/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.dist-info → great_expectations_cloud-20250903.0.dev0.dist-info}/METADATA

@@ -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.dist-info → great_expectations_cloud-20250903.0.dev0.dist-info}/RECORD

@@ -11,7 +11,7 @@ great_expectations_cloud/agent/actions/run_scheduled_checkpoint.py,sha256=XMvv59
 great_expectations_cloud/agent/actions/run_window_checkpoint.py,sha256=khFmTSZ6W-MrInswgEHGrQAzVt-aelycurAMPbcLcLU,2264
 great_expectations_cloud/agent/actions/unknown.py,sha256=mtWw9tDZqGZSiUWj7PtIlLFJ1dM-7AHBX3SO16-u2EM,739
 great_expectations_cloud/agent/actions/utils.py,sha256=DuUy6Ntsin0RIdgG2n1K85WJUATKHsqcCgtXKjr7scQ,1269
-great_expectations_cloud/agent/agent.py,sha256=Gd8gaamnh7Gw4uiW2-sVxbMaWdeNibHQR6LDryI_uVY,29761
+great_expectations_cloud/agent/agent.py,sha256=bybMcl0DwYrnKVG5GYsqPW_ub826Wl7FRTN-W7biZd0,31021
 great_expectations_cloud/agent/agent_warnings.py,sha256=9-xl_AI2V9Py4o7KzFOQjG3lYx-vZ36fq4w2iiPNiUw,362
 great_expectations_cloud/agent/cli.py,sha256=a_HmPxBMlVD59BEmkZnlbOOAFlezVMx9djZ2XIW-3W0,2885
 great_expectations_cloud/agent/config.py,sha256=c1BOr-TrxZnciWNRuu4atGtfRh-XSmwIS0osDCzQa04,1348
@@ -21,14 +21,14 @@ great_expectations_cloud/agent/exceptions.py,sha256=XIDBVSmBfFpVQA5i9rZqEtORIdi5
 great_expectations_cloud/agent/message_service/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 great_expectations_cloud/agent/message_service/asyncio_rabbit_mq_client.py,sha256=B2EwgG_Qdm1s_xkxbGPQxdRBkDSuhkKyT6rGFpXuqQ0,12391
 great_expectations_cloud/agent/message_service/subscriber.py,sha256=K8szy9uM1MUBkMaMZt9o5ExnjqEn1hw6zYaH7pQkjoM,5975
-great_expectations_cloud/agent/models.py,sha256=yZEir-ZShLbyJ8DZxMYNAZmSzbpOSH5wGRKyPo1LcJU,9984
+great_expectations_cloud/agent/models.py,sha256=ea3d6eul8ECvsIubVSEJzeJbX6X9M40oEhaWhj4yuVI,10007
 great_expectations_cloud/agent/run.py,sha256=V33RLoB1PFmJ0h0RfHG4SB5lN_Za8tW2Dua6GUpN9yY,639
 great_expectations_cloud/agent/utils.py,sha256=3OvdcXeK1gk2oJgqG4jPvBRwlMCn8LioULW3YgRtj98,2950
 great_expectations_cloud/logging/README.md,sha256=vbwU689x8SkGjzoBYQzZOzAvh28fR0RCa1XY5WD-Dgs,1762
 great_expectations_cloud/logging/logging_cfg.py,sha256=W6mlm4_Z2bjzM5TuKmFg_WZor2XoJm4DAoLGaf2O__I,6579
 great_expectations_cloud/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-great_expectations_cloud-20250902.0.dev1.dist-info/LICENSE,sha256=_JJnoX6N_OkrAwlCRizCwil0tIjDAy2TG3GiJ50sM6k,2084
-great_expectations_cloud-20250902.0.dev1.dist-info/METADATA,sha256=lfPWV6KQi7QxYYifj3n0BkFeEKSxrummmkMp0Uq48_Y,12352
-great_expectations_cloud-20250902.0.dev1.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
-great_expectations_cloud-20250902.0.dev1.dist-info/entry_points.txt,sha256=ofJgdeS2gSzxXLyCAjfNhIaN1wmSyR7EAMs5qhVaXE4,68
-great_expectations_cloud-20250902.0.dev1.dist-info/RECORD,,
+great_expectations_cloud-20250903.0.dev0.dist-info/LICENSE,sha256=_JJnoX6N_OkrAwlCRizCwil0tIjDAy2TG3GiJ50sM6k,2084
+great_expectations_cloud-20250903.0.dev0.dist-info/METADATA,sha256=49yDtQB7jfbuCAVMXQfRWD7xgm7e77TGmrnFwg1QEyk,12352
+great_expectations_cloud-20250903.0.dev0.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
+great_expectations_cloud-20250903.0.dev0.dist-info/entry_points.txt,sha256=ofJgdeS2gSzxXLyCAjfNhIaN1wmSyR7EAMs5qhVaXE4,68
+great_expectations_cloud-20250903.0.dev0.dist-info/RECORD,,