great-expectations-cloud 20250811.1.dev0__py3-none-any.whl → 20260113.0.dev1__py3-none-any.whl

great_expectations_cloud/agent/agent.py

@@ -2,13 +2,20 @@ from __future__ import annotations
 
 import asyncio
 import logging
+import os
+import resource
+import signal
+import socket
 import sys
+import threading
+import time
 import traceback
 import warnings
 from collections import defaultdict
 from concurrent.futures import Future
 from concurrent.futures.thread import ThreadPoolExecutor
 from functools import partial
+from http import HTTPStatus
 from importlib.metadata import version as metadata_version
 from typing import TYPE_CHECKING, Any, Callable, Final, Literal
 from urllib.parse import urljoin, urlparse
@@ -16,13 +23,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,
@@ -45,9 +51,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,
@@ -67,6 +71,7 @@ from great_expectations_cloud.agent.models import (
     AgentBaseExtraForbid,
     CreateScheduledJobAndSetJobStarted,
     CreateScheduledJobAndSetJobStartedRequest,
+    DomainContext,
     JobCompleted,
     JobStarted,
     JobStatus,
@@ -138,6 +143,9 @@ class GXAgent:
     _PYPI_GX_AGENT_PACKAGE_NAME = "great_expectations_cloud"
     _PYPI_GREAT_EXPECTATIONS_PACKAGE_NAME = "great_expectations"
 
+    # Heartbeat interval in seconds (log progress every 60 seconds during job processing)
+    _HEARTBEAT_INTERVAL_SECONDS = 60
+
     def __init__(self: Self):
         self._config = self._create_config()
 
@@ -150,19 +158,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
@@ -174,6 +169,15 @@ class GXAgent:
         self._correlation_ids: defaultdict[str, int] = defaultdict(lambda: 0)
         self._listen_tries = 0
 
+        # Heartbeat tracking
+        self._heartbeat_stop_event: threading.Event | None = None
+        self._heartbeat_thread: threading.Thread | None = None
+        self._current_job_correlation_id: str | None = None
+        self._current_job_start_time: float | None = None
+
+        # Install signal handlers for graceful shutdown logging
+        self._install_signal_handlers()
+
     def run(self) -> None:
         """Open a connection to GX Cloud."""
 
@@ -218,7 +222,7 @@ class GXAgent:
             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.")
-            sys.exit(1)
+            os.kill(os.getpid(), signal.SIGTERM)
         except (
             AuthenticationError,
             ProbableAuthenticationError,
@@ -233,6 +237,98 @@ class GXAgent:
             if subscriber is not None:
                 subscriber.close()
 
+    def _install_signal_handlers(self) -> None:
+        """Install signal handlers to log when the process receives shutdown signals."""
+        original_sigterm = signal.getsignal(signal.SIGTERM)
+        original_sigint = signal.getsignal(signal.SIGINT)
+
+        def sigterm_handler(signum: int, frame: Any) -> None:
+            self._log_signal_received("SIGTERM", signum)
+            if callable(original_sigterm):
+                original_sigterm(signum, frame)
+            elif original_sigterm == signal.SIG_DFL:
+                raise SystemExit(128 + signum)
+
+        def sigint_handler(signum: int, frame: Any) -> None:
+            self._log_signal_received("SIGINT", signum)
+            if callable(original_sigint):
+                original_sigint(signum, frame)
+            elif original_sigint == signal.SIG_DFL:
+                raise KeyboardInterrupt
+
+        signal.signal(signal.SIGTERM, sigterm_handler)
+        signal.signal(signal.SIGINT, sigint_handler)
+
+    def _log_signal_received(self, signal_name: str, signum: int) -> None:
+        """Log when a shutdown signal is received, including current job info."""
+        memory_mb = self._get_memory_usage_mb()
+        LOGGER.warning(
+            f"Received {signal_name} signal - shutting down",
+            extra={
+                "signal": signal_name,
+                "signal_number": signum,
+                "hostname": socket.gethostname(),
+                "current_job_correlation_id": self._current_job_correlation_id,
+                "job_elapsed_seconds": (
+                    time.time() - self._current_job_start_time
+                    if self._current_job_start_time
+                    else None
+                ),
+                "memory_usage_mb": memory_mb,
+                "has_active_task": self._current_task is not None and not self._current_task.done(),
+            },
+        )
+
+    def _get_memory_usage_mb(self) -> float:
+        """Get current memory usage in MB using resource module."""
+        # ru_maxrss is in KB on Linux, bytes on macOS
+        usage = resource.getrusage(resource.RUSAGE_SELF)
+        # On macOS, ru_maxrss is in bytes; on Linux, it's in KB
+        if sys.platform == "darwin":
+            return usage.ru_maxrss / (1024 * 1024)
+        return usage.ru_maxrss / 1024
+
+    def _start_heartbeat(self, correlation_id: str, org_id: UUID, workspace_id: UUID) -> None:
+        """Start a background thread that logs periodic heartbeats during job processing."""
+        self._current_job_correlation_id = correlation_id
+        self._current_job_start_time = time.time()
+        self._heartbeat_stop_event = threading.Event()
+
+        def heartbeat_loop() -> None:
+            stop_event = self._heartbeat_stop_event
+            if stop_event is None:
+                return
+            while not stop_event.wait(timeout=self._HEARTBEAT_INTERVAL_SECONDS):
+                if stop_event.is_set():
+                    break
+                elapsed = time.time() - (self._current_job_start_time or time.time())
+                memory_mb = self._get_memory_usage_mb()
+                LOGGER.info(
+                    "Job heartbeat - still processing",
+                    extra={
+                        "correlation_id": correlation_id,
+                        "organization_id": str(org_id),
+                        "workspace_id": str(workspace_id),
+                        "hostname": socket.gethostname(),
+                        "elapsed_seconds": round(elapsed, 1),
+                        "memory_usage_mb": round(memory_mb, 1),
+                    },
+                )
+
+        self._heartbeat_thread = threading.Thread(target=heartbeat_loop, daemon=True)
+        self._heartbeat_thread.start()
+
+    def _stop_heartbeat(self) -> None:
+        """Stop the heartbeat thread."""
+        if self._heartbeat_stop_event:
+            self._heartbeat_stop_event.set()
+        if self._heartbeat_thread and self._heartbeat_thread.is_alive():
+            self._heartbeat_thread.join(timeout=2)
+        self._heartbeat_thread = None
+        self._heartbeat_stop_event = None
+        self._current_job_correlation_id = None
+        self._current_job_start_time = None
+
     @classmethod
     def get_current_gx_agent_version(cls) -> str:
         version: str = metadata_version(cls._PYPI_GX_AGENT_PACKAGE_NAME)
@@ -252,8 +348,26 @@ class GXAgent:
         Args:
             event_context: An Event with related properties and actions.
         """
+        # Track how many times this correlation_id has been seen BY THIS POD (for local diagnostics)
+        # Note: event_context.redelivered is set by RabbitMQ and indicates cross-pod redelivery
+        local_delivery_count = self._correlation_ids.get(event_context.correlation_id, 0)
+
         if self._reject_correlation_id(event_context.correlation_id) is True:
-            # this event has been redelivered too many times - remove it from circulation
+            # this event has been redelivered too many times to THIS pod - remove it from circulation
+            LOGGER.error(
+                "Message redelivered too many times to this pod, removing from queue",
+                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,
+                    "local_delivery_count": local_delivery_count,
+                    "redelivered": event_context.redelivered,
+                },
+            )
             event_context.processed_with_failures()
             return
         elif self._can_accept_new_task() is not True:
@@ -263,9 +377,11 @@ 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,
+                    "redelivered": event_context.redelivered,
                 },
             )
             # request that this message is redelivered later
@@ -274,6 +390,21 @@ class GXAgent:
             self._redeliver_msg_task = loop.create_task(event_context.redeliver_message())
             return
 
+        if event_context.redelivered:
+            LOGGER.warning(
+                "Accepting redelivered message - another consumer failed to acknowledge",
+                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,
+                    "redelivered": event_context.redelivered,
+                },
+            )
+
         self._current_task = self._executor.submit(
             self._handle_event,
             event_context=event_context,
@@ -287,8 +418,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."""
@@ -298,6 +444,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,27 +473,38 @@ 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,
             )
+        memory_mb = self._get_memory_usage_mb()
         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,
+                "hostname": socket.gethostname(),
+                "redelivered": event_context.redelivered,
+                "memory_usage_mb": round(memory_mb, 1),
             },
         )
 
+        self._start_heartbeat(event_context.correlation_id, org_id, workspace_id)
+
         self._set_sentry_tags(event_context)
 
         handler = EventHandler(context=data_context)
@@ -350,7 +514,7 @@ class GXAgent:
             id=event_context.correlation_id,
             base_url=base_url,
             auth_key=auth_key,
-            organization_id=org_id,
+            domain_context=DomainContext(organization_id=org_id, workspace_id=workspace_id),
         )
         return result
 
@@ -365,10 +529,26 @@ class GXAgent:
         """
         # warning:  this method will not be executed in the main thread
 
+        self._stop_heartbeat()
+
         org_id = self.get_organization_id(event_context)
+        workspace_id = self.get_workspace_id(event_context)
+
+        memory_mb = self._get_memory_usage_mb()
+        LOGGER.debug(
+            "Job thread exiting",
+            extra={
+                "correlation_id": event_context.correlation_id,
+                "hostname": socket.gethostname(),
+                "has_exception": future.exception() is not None,
+                "cancelled": future.cancelled(),
+                "memory_usage_mb": round(memory_mb, 1),
+            },
+        )
 
         # get results or errors from the thread
         error = future.exception()
+
         if error is None:
             result: ActionResult = future.result()
 
@@ -385,6 +565,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,26 +586,35 @@ 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,
+                        "hostname": socket.gethostname(),
                     },
                 )
         else:
             status = build_failed_job_completed_status(error)
             LOGGER.info(traceback.format_exc())
-            LOGGER.info(
+            LOGGER.warning(
                 "Job completed with error",
                 extra={
                     "event_type": event_context.event.type,
                     "correlation_id": event_context.correlation_id,
                     "organization_id": str(org_id),
+                    "workspace_id": str(workspace_id),
+                    "hostname": socket.gethostname(),
+                    "error_type": type(error).__name__,
+                    "error_message": str(error)[:500],  # Truncate to avoid huge logs
                 },
             )
 
         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 +623,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 +743,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 +758,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 +774,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 +782,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 +804,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",
@@ -629,6 +825,31 @@ class GXAgent:
         with create_session(access_token=self.get_auth_key()) as session:
             payload = CreateScheduledJobAndSetJobStartedRequest(data=data).json()
             response = session.post(agent_sessions_url, data=payload)
+
+            if response.status_code == HTTPStatus.BAD_REQUEST:
+                try:
+                    response_body = response.json()
+                except Exception:
+                    response_body = response.text
+                LOGGER.warning(
+                    "Job already exists - this message was likely redelivered by RabbitMQ "
+                    "after another runner already claimed it. Continuing to process anyway "
+                    "as a safety measure in case the original runner failed.",
+                    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),
+                        "response_status": response.status_code,
+                        "response_body": response_body,
+                    },
+                )
+                # Note: We intentionally continue processing instead of NACKing.
+                # This ensures job completion even if the first runner fails.
+                # TODO: Once we add inProgress timeout in Mercury, we can
+                # safely NACK here to prevent duplicate processing.
+
             LOGGER.info(
                 "Created scheduled job and set started",
                 extra={
@@ -636,6 +857,8 @@ 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),
+                    "response_status": response.status_code,
                 },
             )
             GXAgent._log_http_error(
@@ -658,7 +881,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())
@@ -686,8 +911,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/event_handler.py

@@ -15,6 +15,7 @@ from pydantic import v1 as pydantic_v1
 from great_expectations_cloud.agent.actions.unknown import UnknownEventAction
 from great_expectations_cloud.agent.exceptions import GXAgentError
 from great_expectations_cloud.agent.models import (
+    DomainContext,
     Event,
     EventType,
     UnknownEvent,
@@ -67,11 +68,11 @@ class EventHandler:
         self._context = context
 
     def get_event_action(
-        self, event: Event, base_url: str, auth_key: str, organization_id: UUID
+        self, event: Event, base_url: str, auth_key: str, domain_context: DomainContext
     ) -> AgentAction[Any]:
         """Get the action that should be run for the given event."""
 
-        if not self._check_event_organization_id(event, organization_id):
+        if not self._check_event_organization_id(event, domain_context.organization_id):
             # Making message more generic
             raise GXAgentError("Unable to process job. Invalid input.")  # noqa: TRY003
 
@@ -84,17 +85,17 @@ class EventHandler:
         return action_class(
             context=self._context,
             base_url=base_url,
-            organization_id=organization_id,
+            domain_context=domain_context,
             auth_key=auth_key,
         )
 
-    def handle_event(  # Refactor opportunity
-        self, event: Event, id: str, base_url: str, auth_key: str, organization_id: UUID
+    def handle_event(
+        self, event: Event, id: str, base_url: str, auth_key: str, domain_context: DomainContext
     ) -> ActionResult:
-        start_time = datetime.now(tz=timezone.utc)
         """Transform an Event into an ActionResult."""
+        start_time = datetime.now(tz=timezone.utc)
         action = self.get_event_action(
-            event=event, base_url=base_url, auth_key=auth_key, organization_id=organization_id
+            event=event, base_url=base_url, auth_key=auth_key, domain_context=domain_context
         )
         LOGGER.info(f"Handling event: {event.type} -> {action.__class__.__name__}")
         action_result = action.run(event=event, id=id)

great_expectations_cloud/agent/message_service/asyncio_rabbit_mq_client.py

@@ -27,6 +27,7 @@ class OnMessagePayload:
     correlation_id: str
     delivery_tag: int
     body: bytes
+    redelivered: bool = False  # Set by RabbitMQ when message is redelivered
 
 
 class OnMessageFn(Protocol):
@@ -174,8 +175,12 @@ class AsyncRabbitMQClient:
         # param on_message is provided by the caller as an argument to AsyncRabbitMQClient.run
         correlation_id = header_frame.correlation_id
         delivery_tag = method_frame.delivery_tag
+        redelivered = method_frame.redelivered  # RabbitMQ sets this flag on redelivery
         payload = OnMessagePayload(
-            correlation_id=correlation_id, delivery_tag=delivery_tag, body=body
+            correlation_id=correlation_id,
+            delivery_tag=delivery_tag,
+            body=body,
+            redelivered=redelivered,
         )
         return on_message(payload)
 
@@ -190,10 +195,13 @@ class AsyncRabbitMQClient:
     def _on_consumer_canceled(self, method_frame: Basic.Cancel) -> None:
         """Callback invoked when the broker cancels the client's connection."""
         if self._channel is not None:
-            LOGGER.info(
-                "Consumer was cancelled remotely, shutting down",
+            LOGGER.warning(
+                "Consumer was cancelled remotely by RabbitMQ - this may indicate DAT timeout",
                 extra={
-                    "method_frame": method_frame,
+                    "consumer_tag": method_frame.consumer_tag
+                    if hasattr(method_frame, "consumer_tag")
+                    else None,
+                    "was_consuming": self.was_consuming,
                 },
             )
             self._channel.close()
@@ -232,11 +240,28 @@ class AsyncRabbitMQClient:
         self._reconnect()
         self._log_pika_exception("Connection open failed", reason)
 
-    def _on_connection_closed(
-        self, connection: AsyncioConnection, _unused_reason: pika.Exception
-    ) -> None:
+    def _on_connection_closed(self, connection: AsyncioConnection, reason: pika.Exception) -> None:
         """Callback invoked after the broker closes the connection"""
-        LOGGER.debug("Connection to RabbitMQ has been closed")
+        if isinstance(reason, (ConnectionClosed, ChannelClosed)):
+            LOGGER.warning(
+                "Connection to RabbitMQ has been closed",
+                extra={
+                    "reply_code": reason.reply_code,
+                    "reply_text": reason.reply_text,
+                    "was_consuming": self.was_consuming,
+                    "is_closing": self._closing,
+                },
+            )
+        else:
+            LOGGER.warning(
+                "Connection to RabbitMQ has been closed",
+                extra={
+                    "reason": str(reason),
+                    "reason_type": type(reason).__name__,
+                    "was_consuming": self.was_consuming,
+                    "is_closing": self._closing,
+                },
+            )
         self._channel = None
         self._is_unrecoverable = True
         if self._closing: