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

great_expectations_cloud/agent/agent.py

@@ -2,6 +2,8 @@ from __future__ import annotations
 
 import asyncio
 import logging
+import os
+import signal
 import traceback
 import warnings
 from collections import defaultdict
@@ -9,24 +11,45 @@ from concurrent.futures import Future
 from concurrent.futures.thread import ThreadPoolExecutor
 from functools import partial
 from importlib.metadata import version as metadata_version
-from typing import TYPE_CHECKING, Any, Dict, Final
-
-from great_expectations import get_context  # type: ignore[attr-defined] # TODO: fix this
-from great_expectations.compatibility import pydantic
-from great_expectations.compatibility.pydantic import AmqpDsn, AnyUrl
+from typing import TYPE_CHECKING, Any, Callable, Final, Literal
+from urllib.parse import urljoin, urlparse
+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 packaging.version import Version
-from pika.exceptions import AuthenticationError, ProbableAuthenticationError
-from tenacity import after_log, retry, retry_if_exception_type, stop_after_attempt, wait_exponential
+from great_expectations.data_context.types.base import ProgressBarsConfig
+from pika.adapters.utils.connection_workflow import AMQPConnectorException
+from pika.exceptions import (
+    AMQPConnectionError,
+    AMQPError,
+    AuthenticationError,
+    ChannelError,
+    ProbableAuthenticationError,
+)
+from pydantic import v1 as pydantic_v1
+from pydantic.v1 import AmqpDsn, AnyUrl
+from tenacity import (
+    after_log,
+    retry,
+    retry_if_exception_type,
+    stop_after_attempt,
+    wait_random_exponential,
+)
 
 from great_expectations_cloud.agent.config import (
     GxAgentEnvVars,
     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,
+    GXAgentUnrecoverableConnectionError,
 )
 from great_expectations_cloud.agent.message_service.asyncio_rabbit_mq_client import (
     AsyncRabbitMQClient,
@@ -39,16 +62,20 @@ from great_expectations_cloud.agent.message_service.subscriber import (
     SubscriberError,
 )
 from great_expectations_cloud.agent.models import (
-    AgentBaseModel,
+    AgentBaseExtraForbid,
+    CreateScheduledJobAndSetJobStarted,
+    CreateScheduledJobAndSetJobStartedRequest,
+    DomainContext,
     JobCompleted,
     JobStarted,
     JobStatus,
+    ScheduledEventBase,
     UnknownEvent,
+    UpdateJobStatusRequest,
     build_failed_job_completed_status,
 )
 
 if TYPE_CHECKING:
-    import requests
     from great_expectations.data_context import CloudDataContext
     from typing_extensions import Self
 
@@ -56,11 +83,11 @@ if TYPE_CHECKING:
 
 LOGGER: Final[logging.Logger] = logging.getLogger(__name__)
 # TODO Set in log dict
-LOGGER.setLevel(logging.INFO)
-HandlerMap = Dict[str, OnMessageCallback]
+LOGGER.setLevel(logging.DEBUG)
+HandlerMap = dict[str, OnMessageCallback]
 
 
-class GXAgentConfig(AgentBaseModel):
+class GXAgentConfig(AgentBaseExtraForbid):
     """GXAgent configuration.
     Attributes:
         queue: name of queue
@@ -69,10 +96,33 @@ class GXAgentConfig(AgentBaseModel):
 
     queue: str
     connection_string: AmqpDsn
-    # pydantic will coerce this string to AnyUrl type
-    gx_cloud_base_url: AnyUrl = CLOUD_DEFAULT_BASE_URL
+    gx_cloud_base_url: AnyUrl = AnyUrl(url=CLOUD_DEFAULT_BASE_URL, scheme="https")
     gx_cloud_organization_id: str
     gx_cloud_access_token: str
+    enable_progress_bars: bool = True
+
+
+def orjson_dumps(v: Any, *, default: Callable[[Any], Any] | None) -> str:
+    # orjson.dumps returns bytes, to match standard json.dumps we need to decode
+    # Typing using example from https://github.com/ijl/orjson?tab=readme-ov-file#serialize
+    return orjson.dumps(
+        v,
+        default=default,
+    ).decode()
+
+
+def orjson_loads(v: bytes | bytearray | memoryview | str) -> Any:
+    # Typing using example from https://github.com/ijl/orjson?tab=readme-ov-file#deserialize
+    return orjson.loads(v)
+
+
+class Payload(AgentBaseExtraForbid):
+    data: dict[str, Any]
+
+    class Config:
+        extra = "forbid"
+        json_dumps = orjson_dumps
+        json_loads = orjson_loads
 
 
 class GXAgent:
@@ -88,19 +138,17 @@ class GXAgent:
     _PYPI_GREAT_EXPECTATIONS_PACKAGE_NAME = "great_expectations"
 
     def __init__(self: Self):
-        agent_version: str = self.get_current_gx_agent_version()
-        print(f"GX Agent version: {agent_version}")
-        print("Initializing the GX Agent.")
-        self._set_http_session_headers()
-        self._config = self._get_config()
-        print("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)
+        self._config = self._create_config()
 
-        print("DataContext is ready.")
+        agent_version: str = self.get_current_gx_agent_version()
+        great_expectations_version: str = self._get_current_great_expectations_version()
+        LOGGER.info(
+            "Initializing GX Agent.",
+            extra={
+                "agent_version": agent_version,
+                "great_expectations_version": great_expectations_version,
+            },
+        )
 
         # 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
@@ -110,45 +158,63 @@ class GXAgent:
         self._current_task: Future[Any] | None = None
         self._redeliver_msg_task: asyncio.Task[Any] | None = None
         self._correlation_ids: defaultdict[str, int] = defaultdict(lambda: 0)
+        self._listen_tries = 0
 
     def run(self) -> None:
         """Open a connection to GX Cloud."""
 
-        print("Opening connection to GX Cloud.")
+        LOGGER.debug("Opening connection to GX Cloud.")
+        self._listen_tries = 0
         self._listen()
-        print("The connection to GX Cloud has been closed.")
+        LOGGER.debug("The connection to GX Cloud has been closed.")
 
     # ZEL-505: A race condition can occur if two or more agents are started at the same time
     #          due to the generation of passwords for rabbitMQ queues. This can be mitigated
     #          by adding a delay and retrying the connection. Retrying with new credentials
     #          requires calling get_config again, which handles the password generation.
     @retry(
-        retry=retry_if_exception_type((AuthenticationError, ProbableAuthenticationError)),
-        wait=wait_exponential(multiplier=1, min=1, max=10),
+        retry=retry_if_exception_type(
+            (AuthenticationError, ProbableAuthenticationError, AMQPError, ChannelError)
+        ),
+        wait=wait_random_exponential(multiplier=1, min=1, max=10),
         stop=stop_after_attempt(3),
         after=after_log(LOGGER, logging.DEBUG),
     )
     def _listen(self) -> None:
         """Manage connection lifecycle."""
         subscriber = None
+        # force refresh if we're retrying
+        force_creds_refresh = self._listen_tries > 0
+        self._listen_tries += 1
+
+        config = self._get_config(force_refresh=force_creds_refresh)
+
         try:
-            client = AsyncRabbitMQClient(url=str(self._config.connection_string))
+            client = AsyncRabbitMQClient(url=str(config.connection_string))
             subscriber = Subscriber(client=client)
-            print("The GX Agent is ready.")
+            LOGGER.info("The GX Agent is ready.")
             # Open a connection until encountering a shutdown event
             subscriber.consume(
-                queue=self._config.queue,
+                queue=config.queue,
                 on_message=self._handle_event_as_thread_enter,
             )
         except KeyboardInterrupt:
-            print("Received request to shut down.")
+            LOGGER.debug("Received request to shut down.")
         except (SubscriberError, ClientError):
-            print("The connection to GX Cloud has encountered an error.")
-        except (AuthenticationError, ProbableAuthenticationError):
-            # Retry with new credentials
-            self._config = self._get_config()
+            LOGGER.exception("The connection to GX Cloud has encountered an error.")
+        except GXAgentUnrecoverableConnectionError:
+            LOGGER.exception("The connection to GX Cloud has encountered an unrecoverable error.")
+            os.kill(os.getpid(), signal.SIGTERM)
+        except (
+            AuthenticationError,
+            ProbableAuthenticationError,
+            AMQPConnectorException,
+            AMQPConnectionError,
+        ):
             # Raise to use the retry decorator to handle the retry logic
+            LOGGER.exception("Failed authentication to MQ.")
             raise
+
         finally:
             if subscriber is not None:
                 subscriber.close()
@@ -158,6 +224,11 @@ class GXAgent:
         version: str = metadata_version(cls._PYPI_GX_AGENT_PACKAGE_NAME)
         return version
 
+    @classmethod
+    def _get_current_great_expectations_version(cls) -> str:
+        version: str = metadata_version(cls._PYPI_GREAT_EXPECTATIONS_PACKAGE_NAME)
+        return version
+
     def _handle_event_as_thread_enter(self, event_context: EventContext) -> None:
         """Schedule _handle_event to run in a thread.
 
@@ -172,17 +243,28 @@ class GXAgent:
             event_context.processed_with_failures()
             return
         elif self._can_accept_new_task() is not True:
+            LOGGER.warning(
+                "Cannot accept new task, redelivering.",
+                extra={
+                    "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,
+                },
+            )
             # request that this message is redelivered later
             loop = asyncio.get_event_loop()
             # store a reference the task to ensure it isn't garbage collected
             self._redeliver_msg_task = loop.create_task(event_context.redeliver_message())
             return
 
-        # ensure that great_expectations.http requests to GX Cloud include the job_id/correlation_id
-        self._set_http_session_headers(correlation_id=event_context.correlation_id)
-
-        # send this message to a thread for processing
-        self._current_task = self._executor.submit(self._handle_event, event_context=event_context)
+        self._current_task = self._executor.submit(
+            self._handle_event,
+            event_context=event_context,
+        )
 
         if self._current_task is not None:
             # add a callback for when the thread exits and pass it the event context
@@ -191,6 +273,44 @@ class GXAgent:
             )
             self._current_task.add_done_callback(on_exit_callback)
 
+    def get_data_context(self, event_context: EventContext) -> CloudDataContext:
+        """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."""
+        return UUID(self._get_config().gx_cloud_organization_id)
+
+    def get_auth_key(self) -> str:
+        """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
+
     def _handle_event(self, event_context: EventContext) -> ActionResult:
         """Pass events to EventHandler.
 
@@ -201,18 +321,51 @@ class GXAgent:
             event_context: event with related properties and actions.
         """
         # warning:  this method will not be executed in the main thread
-        self._update_status(job_id=event_context.correlation_id, status=JobStarted())
-        print(f"Starting job {event_context.event.type} ({event_context.correlation_id}) ")
+
+        data_context = self.get_data_context(event_context=event_context)
+        # ensure that great_expectations.http requests to GX Cloud include the job_id/correlation_id
+        self._set_http_session_headers(
+            correlation_id=event_context.correlation_id, data_context=data_context
+        )
+
+        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, workspace_id)
+        else:
+            self._update_status(
+                correlation_id=event_context.correlation_id,
+                status=JobStarted(),
+                org_id=org_id,
+                workspace_id=workspace_id,
+            )
         LOGGER.info(
             "Starting job",
             extra={
                 "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,
             },
         )
-        handler = EventHandler(context=self._context)
+
+        self._set_sentry_tags(event_context)
+
+        handler = EventHandler(context=data_context)
         # This method might raise an exception. Allow it and handle in _handle_event_as_thread_exit
-        result = handler.handle_event(event=event_context.event, id=event_context.correlation_id)
+        result = handler.handle_event(
+            event=event_context.event,
+            id=event_context.correlation_id,
+            base_url=base_url,
+            auth_key=auth_key,
+            domain_context=DomainContext(organization_id=org_id, workspace_id=workspace_id),
+        )
         return result
 
     def _handle_event_as_thread_exit(
@@ -226,6 +379,9 @@ 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()
         if error is None:
@@ -236,24 +392,39 @@ class GXAgent:
                     success=False,
                     created_resources=[],
                     error_stack_trace="The version of the GX Agent you are using does not support this functionality. Please upgrade to the most recent image tagged with `stable`.",
+                    processed_by=self._get_processed_by(),
                 )
                 LOGGER.error(
                     "Job completed with error. Ensure agent is up-to-date.",
                     extra={
                         "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,
                     },
                 )
             else:
                 status = JobCompleted(
                     success=True,
                     created_resources=result.created_resources,
+                    processed_by=self._get_processed_by(),
                 )
                 LOGGER.info(
                     "Completed job",
                     extra={
                         "event_type": event_context.event.type,
                         "correlation_id": event_context.correlation_id,
+                        "job_duration": (
+                            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,
                     },
                 )
         else:
@@ -264,15 +435,44 @@ class GXAgent:
                 extra={
                     "event_type": event_context.event.type,
                     "correlation_id": event_context.correlation_id,
+                    "organization_id": str(org_id),
+                    "workspace_id": str(workspace_id),
                 },
             )
 
-        self._update_status(job_id=event_context.correlation_id, status=status)
+        try:
+            self._update_status(
+                correlation_id=event_context.correlation_id,
+                status=status,
+                org_id=org_id,
+                workspace_id=workspace_id,
+            )
+        except Exception:
+            LOGGER.exception(
+                "Error updating status, removing message from queue",
+                extra={
+                    "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
+            # If the status update fails, remove the message from the queue
+            # Otherwise, it would attempt to handle the error again via this done callback
+            event_context.processed_with_failures()
+            self._current_task = None
+            # Return so we don't also ack as processed successfully
+            return
 
         # ack message and cleanup resources
         event_context.processed_successfully()
         self._current_task = None
 
+    def _get_processed_by(self) -> Literal["agent", "runner"]:
+        """Return the name of the service that processed the event."""
+        return "runner" if self._get_config().queue == "gx-runner" else "agent"
+
     def _can_accept_new_task(self) -> bool:
         """Are we currently processing a task or are we free to take a new one?"""
         return self._current_task is None or self._current_task.done()
@@ -292,28 +492,33 @@ class GXAgent:
             self._correlation_ids.clear()
         return should_reject
 
+    def _get_config(self, force_refresh: bool = False) -> GXAgentConfig:
+        if force_refresh:
+            self._config = self._create_config()
+        return self._config
+
     @classmethod
-    def _get_config(cls) -> GXAgentConfig:
+    def _create_config(cls) -> GXAgentConfig:
         """Construct GXAgentConfig."""
 
         # ensure we have all required env variables, and provide a useful error if not
 
         try:
             env_vars = GxAgentEnvVars()
-        except pydantic.ValidationError as validation_err:
+        except pydantic_v1.ValidationError as validation_err:
             raise GXAgentConfigError(
                 generate_config_validation_error_text(validation_err)
             ) from validation_err
 
         # obtain the broker url and queue name from Cloud
-        agent_sessions_url = (
-            f"{env_vars.gx_cloud_base_url}/organizations/"
-            f"{env_vars.gx_cloud_organization_id}/agent-sessions"
+        agent_sessions_url = urljoin(
+            env_vars.gx_cloud_base_url,
+            f"/api/v1/organizations/{env_vars.gx_cloud_organization_id}/agent-sessions",
         )
 
         session = create_session(access_token=env_vars.gx_cloud_access_token)
-
         response = session.post(agent_sessions_url)
+        session.close()
         if response.ok is not True:
             raise GXAgentError(  # noqa: TRY003 # TODO: use AuthenticationError
                 "Unable to authenticate to GX Cloud. Please check your credentials."
@@ -323,6 +528,19 @@ class GXAgent:
         queue = json_response["queue"]
         connection_string = json_response["connection_string"]
 
+        # if overrides are set, we update the connection string. This is useful for local development to set the host
+        # to localhost, for example.
+        parsed = urlparse(connection_string)
+        if env_vars.amqp_host_override:
+            netloc = (
+                f"{parsed.username}:{parsed.password}@{env_vars.amqp_host_override}:{parsed.port}"
+            )
+            parsed = parsed._replace(netloc=netloc)  # documented in urllib docs
+        if env_vars.amqp_port_override:
+            netloc = f"{parsed.username}:{parsed.password}@{parsed.hostname}:{env_vars.amqp_port_override}"
+            parsed = parsed._replace(netloc=netloc)  # documented in urllib docs
+        connection_string = parsed.geturl()
+
         try:
             # pydantic will coerce the url to the correct type
             return GXAgentConfig(
@@ -331,71 +549,192 @@ class GXAgent:
                 gx_cloud_base_url=env_vars.gx_cloud_base_url,
                 gx_cloud_organization_id=env_vars.gx_cloud_organization_id,
                 gx_cloud_access_token=env_vars.gx_cloud_access_token,
+                enable_progress_bars=env_vars.enable_progress_bars,
             )
-        except pydantic.ValidationError as validation_err:
+        except pydantic_v1.ValidationError as validation_err:
             raise GXAgentConfigError(
                 generate_config_validation_error_text(validation_err)
             ) from validation_err
 
-    def _update_status(self, job_id: str, status: JobStatus) -> None:
+    def _configure_progress_bars(self, data_context: CloudDataContext) -> None:
+        progress_bars_enabled = self._get_config().enable_progress_bars
+
+        try:
+            data_context.variables.progress_bars = ProgressBarsConfig(
+                globally=progress_bars_enabled,
+                metric_calculations=progress_bars_enabled,
+            )
+            data_context.variables.save()
+        except Exception:
+            # Progress bars are not critical, so log and continue
+            # This is a known issue with FastAPI mercury V1 API for data-context-variables
+            LOGGER.warning(
+                "Failed to {set} progress bars".format(
+                    set="enable" if progress_bars_enabled else "disable"
+                )
+            )
+
+    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:
-            job_id: job identifier, also known as correlation_id
-            status: pydantic model encapsulating the current status
+            correlation_id: job identifier
+            status: pydantic model encapsulating the current status.
         """
-        LOGGER.info("Updating status", extra={"job_id": job_id, "status": str(status)})
-        agent_sessions_url = (
-            f"{self._config.gx_cloud_base_url}/organizations/{self._config.gx_cloud_organization_id}"
-            + f"/agent-jobs/{job_id}"
+        LOGGER.info(
+            "Updating status",
+            extra={
+                "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}/workspaces/{workspace_id}/agent-jobs/{correlation_id}",
         )
-        session = create_session(access_token=self._config.gx_cloud_access_token)
-        data = status.json()
-        session.patch(agent_sessions_url, data=data)
+        with create_session(access_token=self.get_auth_key()) as session:
+            data = UpdateJobStatusRequest(data=status).json()
+            response = session.patch(agent_sessions_url, data=data)
+            LOGGER.info(
+                "Status updated",
+                extra={
+                    "correlation_id": correlation_id,
+                    "status": str(status),
+                    "organization_id": str(org_id),
+                    "workspace_id": str(workspace_id),
+                },
+            )
+            GXAgent._log_http_error(
+                response, message="Status Update action had an error while connecting to GX Cloud."
+            )
+
+    def _create_scheduled_job_and_set_started(
+        self, event_context: EventContext, org_id: UUID, workspace_id: UUID
+    ) -> None:
+        """Create a job in GX Cloud for scheduled events.
+
+        This is because the scheduler + lambda create the event in the queue, and the agent consumes it. The agent then
+        sends a request to the agent-jobs endpoint to create the job in mercury to keep track of the job status.
+        Non-scheduled events by contrast create the job in mercury and the event in the queue at the same time.
 
-    def _set_http_session_headers(self, correlation_id: str | None = None) -> None:
+        Args:
+            event_context: event with related properties and actions.
         """
-        Set the the session headers for requests to GX Cloud.
+        if not isinstance(event_context.event, ScheduledEventBase):
+            raise GXAgentError(  # noqa: TRY003
+                "Unable to create a scheduled job for a non-scheduled event."
+            )
+
+        LOGGER.info(
+            "Creating scheduled job and setting started",
+            extra={
+                "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}/workspaces/{workspace_id}/agent-jobs",
+        )
+        data = CreateScheduledJobAndSetJobStarted(
+            type="run_scheduled_checkpoint.received",
+            correlation_id=UUID(event_context.correlation_id),
+            schedule_id=event_context.event.schedule_id,
+            checkpoint_id=event_context.event.checkpoint_id,
+            datasource_names_to_asset_names=event_context.event.datasource_names_to_asset_names,
+            splitter_options=event_context.event.splitter_options,
+            checkpoint_name=event_context.event.checkpoint_name,
+        )
+        with create_session(access_token=self.get_auth_key()) as session:
+            payload = CreateScheduledJobAndSetJobStartedRequest(data=data).json()
+            response = session.post(agent_sessions_url, data=payload)
+            LOGGER.info(
+                "Created scheduled job and set started",
+                extra={
+                    "correlation_id": str(event_context.correlation_id),
+                    "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(
+                response,
+                message="Create schedule job action had an error while connecting to GX Cloud.",
+            )
+
+    def get_header_name(self) -> type[HeaderName]:
+        return HeaderName
+
+    def get_user_agent_header(self) -> str:
+        return USER_AGENT_HEADER
+
+    def _get_version(self) -> str:
+        return self.get_current_gx_agent_version()
+
+    def _set_data_context_store_headers(
+        self, data_context: CloudDataContext, headers: dict[HeaderName, str]
+    ) -> None:
+        """
+        Sets headers on all stores in the data context.
+        """
+        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())
+        # some stores are treated differently
+        stores.extend([data_context._datasource_store, data_context._data_asset_store])
+        for store in stores:
+            backend = store._store_backend
+            if isinstance(backend, GXCloudStoreBackend):
+                backend._session.headers.update({str(key): value for key, value in headers.items()})
+
+    @property
+    def user_agent_str(self) -> str:
+        user_agent_header_prefix = self.get_user_agent_header()
+        agent_version = self._get_version()
+        return f"{user_agent_header_prefix}/{agent_version}"
+
+    def _set_http_session_headers(
+        self, data_context: CloudDataContext, correlation_id: str | None = None
+    ) -> None:
+        """
+        Set the session headers for requests to GX Cloud.
         In particular, set the User-Agent header to identify the GX Agent and the correlation_id as
         Agent-Job-Id if provided.
 
         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__  # type: ignore[attr-defined] # TODO: fix this
-        from great_expectations.core import http
-        from great_expectations.data_context.store.gx_cloud_store_backend import GXCloudStoreBackend
 
-        if Version(__version__) > Version(
-            "0.19"  # using 0.19 instead of 1.0 to account for pre-releases
-        ):
-            # TODO: public API should be available in v1
-            LOGGER.info(
-                "Unable to set header for requests to GX Cloud",
-                extra={
-                    "user_agent": HeaderName.USER_AGENT,
-                    "agent_job_id": HeaderName.AGENT_JOB_ID,
-                },
-            )
-            return
+        header_name = self.get_header_name()
+        user_agent_header_value = self.user_agent_str
 
-        agent_version = self.get_current_gx_agent_version()
         LOGGER.debug(
             "Setting session headers for GX Cloud",
             extra={
-                "user_agent": HeaderName.USER_AGENT,
-                "agent_version": agent_version,
-                "job_id": HeaderName.AGENT_JOB_ID,
+                "user_agent_header_name": header_name.USER_AGENT,
+                "user_agent_header_value": user_agent_header_value,
+                "correlation_id_header_name": header_name.AGENT_JOB_ID,
+                "correlation_id_header_value": correlation_id,
                 "correlation_id": correlation_id,
             },
         )
 
+        core_headers = {header_name.USER_AGENT: user_agent_header_value}
         if correlation_id:
-            # OSS doesn't use the same session for all requests, so we need to set the header for each store
-            for store in self._context.stores.values():
-                backend = store._store_backend
-                if isinstance(backend, GXCloudStoreBackend):
-                    backend._session.headers[HeaderName.AGENT_JOB_ID] = correlation_id
+            core_headers.update({header_name.AGENT_JOB_ID: correlation_id})
+        self._set_data_context_store_headers(data_context=data_context, headers=core_headers)
 
         def _update_headers_agent_patch(
             session: requests.Session, access_token: str
@@ -407,10 +746,10 @@ class GXAgent:
                 "Content-Type": "application/vnd.api+json",
                 "Authorization": f"Bearer {access_token}",
                 "Gx-Version": __version__,
-                HeaderName.USER_AGENT: f"{USER_AGENT_HEADER}/{agent_version}",
+                header_name.USER_AGENT: user_agent_header_value,
             }
             if correlation_id:
-                headers[HeaderName.AGENT_JOB_ID] = correlation_id
+                headers[header_name.AGENT_JOB_ID] = correlation_id
             session.headers.update(headers)
             return session
 
@@ -418,8 +757,12 @@ class GXAgent:
         # use a public API once it is available
         http._update_headers = _update_headers_agent_patch
 
-
-class GXAgentError(Exception): ...
-
-
-class GXAgentConfigError(GXAgentError): ...
+    @staticmethod
+    def _log_http_error(response: requests.Response, message: str) -> None:
+        """
+        Log the http error if the response is not successful.
+        """
+        try:
+            response.raise_for_status()
+        except requests.HTTPError:
+            LOGGER.exception(message, extra={"response": response})