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

great_expectations_cloud/agent/cli.py

@@ -91,12 +91,12 @@ def main() -> None:
     )
 
     if args.version:
-        from great_expectations_cloud.agent import get_version
+        from great_expectations_cloud.agent import get_version  # noqa: PLC0415
 
         print(f"GX Agent version: {get_version()}")
         return
 
-    from great_expectations_cloud.agent import run_agent
+    from great_expectations_cloud.agent import run_agent  # noqa: PLC0415
 
     run_agent()
 

great_expectations_cloud/agent/config.py

@@ -1,19 +1,33 @@
 from __future__ import annotations
 
-from great_expectations.compatibility.pydantic import AnyUrl, BaseSettings, ValidationError
+from typing import Optional
+
 from great_expectations.data_context.cloud_constants import CLOUD_DEFAULT_BASE_URL
+from pydantic.v1 import AnyUrl, BaseSettings, ValidationError
 
 
-class GxAgentEnvVars(BaseSettings):  # type: ignore[misc] # BaseSettings is has Any type
-    # pydantic will coerce this string to AnyUrl type
-    gx_cloud_base_url: AnyUrl = CLOUD_DEFAULT_BASE_URL
+class GxAgentEnvVars(BaseSettings):
+    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
+
+    amqp_host_override: Optional[str] = None  # noqa: UP045 # pipe not working with 3.9
+    amqp_port_override: Optional[int] = None  # noqa: UP045 # pipe not working with 3.9
+
+    def __init__(self, **overrides: str | AnyUrl) -> None:
+        """
+        Custom __init__ to prevent type error when relying on environment variables.
+
+        TODO:once mypy fully support annoting **kwargs with a Unpack[TypedDict], we should do that.
+        https://peps.python.org/pep-0692/
+        """
+        super().__init__(**overrides)
 
 
 def generate_config_validation_error_text(validation_error: ValidationError) -> str:
     missing_variables = ", ".join(
-        [validation_error["loc"][0] for validation_error in validation_error.errors()]
+        [str(validation_error["loc"][0]) for validation_error in validation_error.errors()]
     )
     error_text = f"Missing or badly formed environment variable(s). Make sure to set the following environment variable(s): {missing_variables}"
     return error_text

great_expectations_cloud/agent/event_handler.py

@@ -2,18 +2,24 @@ from __future__ import annotations
 
 import logging
 from collections import defaultdict
+from datetime import datetime, timezone
 from json import JSONDecodeError
 from typing import TYPE_CHECKING, Any, Final
+from uuid import UUID
 
 import great_expectations as gx
-from great_expectations.compatibility import pydantic
 from packaging.version import Version
 from packaging.version import parse as parse_version
+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,
+    get_event_union,
 )
 
 if TYPE_CHECKING:
@@ -41,7 +47,7 @@ _EVENT_ACTION_MAP: dict[str, dict[str, type[AgentAction[Any]]]] = defaultdict(di
 
 
 def register_event_action(
-    version: str, event_type: type[Event], action_class: type[AgentAction[Any]]
+    version: str, event_type: EventType, action_class: type[AgentAction[Any]]
 ) -> None:
     """Register an event type to an action class."""
     if version in _EVENT_ACTION_MAP and event_type.__name__ in _EVENT_ACTION_MAP[version]:
@@ -61,34 +67,65 @@ class EventHandler:
     def __init__(self, context: CloudDataContext) -> None:
         self._context = context
 
-    def get_event_action(self, event: Event) -> AgentAction[Any]:
+    def get_event_action(
+        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, domain_context.organization_id):
+            # Making message more generic
+            raise GXAgentError("Unable to process job. Invalid input.")  # noqa: TRY003
+
         action_map = _EVENT_ACTION_MAP.get(_GX_MAJOR_VERSION)
         if action_map is None:
             raise NoVersionImplementationError(version=_GX_MAJOR_VERSION)
         action_class = action_map.get(_get_event_name(event))
         if action_class is None:
             action_class = UnknownEventAction
-        return action_class(context=self._context)
-
-    def handle_event(self, event: Event, id: str) -> ActionResult:
+        return action_class(
+            context=self._context,
+            base_url=base_url,
+            domain_context=domain_context,
+            auth_key=auth_key,
+        )
+
+    def handle_event(
+        self, event: Event, id: str, base_url: str, auth_key: str, domain_context: DomainContext
+    ) -> ActionResult:
         """Transform an Event into an ActionResult."""
-        action = self.get_event_action(event=event)
+        start_time = datetime.now(tz=timezone.utc)
+        action = self.get_event_action(
+            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)
+        end_time = datetime.now(tz=timezone.utc)
+        action_result.job_duration = end_time - start_time
         return action_result
 
     @classmethod
     def parse_event_from(cls, msg_body: bytes) -> Event:
+        event_union = get_event_union()
         try:
-            event: Event = pydantic.parse_raw_as(Event, msg_body)
-        except (pydantic.ValidationError, JSONDecodeError):
+            event: Event = pydantic_v1.parse_raw_as(event_union, msg_body)
+        except (pydantic_v1.ValidationError, JSONDecodeError):
             # Log as bytes
             LOGGER.exception("Unable to parse event type", extra={"msg_body": f"{msg_body!r}"})
             return UnknownEvent()
 
         return event
 
+    @staticmethod
+    def _check_event_organization_id(event: Event, organization_id: UUID) -> bool:
+        """Check if the organization_id in the event matches the given organization_id.
+
+        This prevents processing events that are not intended for the current organization, and potentially
+        leaking sensitive information across organizations.
+        """
+        if hasattr(event, "organization_id") and event.organization_id != organization_id:
+            return False
+        return True
+
 
 class EventError(Exception): ...
 
@@ -121,12 +158,12 @@ def _get_major_version(version: str) -> str:
     return str(parsed.major)
 
 
-version = gx.__version__  # type: ignore[attr-defined] # TODO: fix this
+version = gx.__version__
 _GX_MAJOR_VERSION = _get_major_version(str(version))
 
 
 def _get_event_name(event: Event) -> str:
     try:
-        return str(event.__name__)
+        return str(event.__name__)  # type: ignore[union-attr] # FIXME
     except AttributeError:
-        return str(event.__class__.__name__)
+        return event.__class__.__name__

great_expectations_cloud/agent/exceptions.py

@@ -32,3 +32,12 @@ class ErrorCode(str, enum.Enum):
 
     GENERIC_UNHANDLED_ERROR = "generic-unhandled-error"
     WRONG_USERNAME_OR_PASSWORD = "wrong-username-or-password"  # noqa: S105  # Not a hardcoded password
+
+
+class GXAgentError(Exception): ...
+
+
+class GXAgentConfigError(GXAgentError): ...
+
+
+class GXAgentUnrecoverableConnectionError(GXAgentError): ...

great_expectations_cloud/agent/message_service/asyncio_rabbit_mq_client.py

@@ -1,19 +1,26 @@
 from __future__ import annotations
 
 import asyncio
+import logging
 import ssl
 from asyncio import AbstractEventLoop
 from dataclasses import dataclass
 from functools import partial
-from typing import TYPE_CHECKING, Callable, Protocol
+from typing import TYPE_CHECKING, Callable, Final, Protocol
 
 import pika
 from pika.adapters.asyncio_connection import AsyncioConnection
+from pika.exceptions import ChannelClosed, ConnectionClosed
+
+from great_expectations_cloud.agent.exceptions import GXAgentUnrecoverableConnectionError
 
 if TYPE_CHECKING:
     from pika.channel import Channel
     from pika.spec import Basic, BasicProperties
 
+LOGGER: Final[logging.Logger] = logging.getLogger(__name__)
+LOGGER.setLevel(logging.DEBUG)
+
 
 @dataclass(frozen=True)
 class OnMessagePayload:
@@ -43,6 +50,7 @@ class AsyncRabbitMQClient:
         self._closing = False
         self._consumer_tag = None
         self._consuming = False
+        self._is_unrecoverable = False
 
     def run(self, queue: str, on_message: OnMessageFn) -> None:
         """Run an async connection to RabbitMQ.
@@ -67,21 +75,27 @@ class AsyncRabbitMQClient:
         )
         self._connection = connection
         connection.ioloop.run_forever()
+        if self._is_unrecoverable:
+            raise GXAgentUnrecoverableConnectionError(  # noqa: TRY003
+                "AsyncRabbitMQClient has encountered an unrecoverable error."
+            )
 
     def stop(self) -> None:
         """Close the connection to RabbitMQ."""
         if self._connection is None:
             return
+        LOGGER.debug("Shutting down the connection to RabbitMQ.")
         if not self._closing:
             self._closing = True
-            if self._consuming:
-                self._stop_consuming()
-                self._connection.ioloop.run_forever()
-            else:
-                self._connection.ioloop.stop()
+
+        if self._consuming:
+            self._stop_consuming()
+        self._connection.ioloop.stop()
+        LOGGER.debug("The connection to RabbitMQ has been shut down.")
 
     def reset(self) -> None:
         """Reset client to allow a restart."""
+        LOGGER.debug("Resetting client")
         self.should_reconnect = False
         self.was_consuming = False
 
@@ -148,7 +162,7 @@ class AsyncRabbitMQClient:
             nack = partial(channel.basic_nack, delivery_tag=delivery_tag, requeue=requeue)
             loop.call_soon_threadsafe(callback=nack)
 
-    def _callback_handler(  # noqa: PLR0913
+    def _callback_handler(
         self,
         channel: Channel,
         method_frame: Basic.Deliver,
@@ -167,66 +181,118 @@ class AsyncRabbitMQClient:
 
     def _start_consuming(self, queue: str, on_message: OnMessageFn, channel: Channel) -> None:
         """Consume from a channel with the on_message callback."""
+        LOGGER.debug("Issuing consumer-related RPC commands")
         channel.add_on_cancel_callback(self._on_consumer_canceled)
+        # set RabbitMQ prefetch count to equal the max_threads value in the GX Agent's ThreadPoolExecutor
+        channel.basic_qos(prefetch_count=1)
         self._consumer_tag = channel.basic_consume(queue=queue, on_message_callback=on_message)
 
     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",
+                extra={
+                    "method_frame": method_frame,
+                },
+            )
             self._channel.close()
 
     def _reconnect(self) -> None:
         """Prepare the client to reconnect."""
+        LOGGER.debug("Preparing client to reconnect")
         self.should_reconnect = True
         self.stop()
 
     def _stop_consuming(self) -> None:
         """Cancel the channel, if it exists."""
         if self._channel is not None:
+            LOGGER.debug("Sending a Basic.Cancel RPC command to RabbitMQ")
             self._channel.basic_cancel(self._consumer_tag, callback=self._on_cancel_ok)
 
-    def _on_cancel_ok(self, method_frame: Basic.CancelOk) -> None:
+    def _on_cancel_ok(self, _unused_frame: Basic.CancelOk) -> None:
         """Callback invoked after broker confirms cancel."""
         self._consuming = False
         if self._channel is not None:
+            LOGGER.debug("RabbitMQ acknowledged the cancellation of the consumer")
             self._channel.close()
 
     def _on_connection_open(
         self, connection: AsyncioConnection, queue: str, on_message: OnMessageFn
     ) -> None:
         """Callback invoked after the broker opens the connection."""
+        LOGGER.debug("Connection to RabbitMQ has been opened")
         on_channel_open = partial(self._on_channel_open, queue=queue, on_message=on_message)
         connection.channel(on_open_callback=on_channel_open)
 
-    def _on_connection_open_error(self, connection: AsyncioConnection, reason: str) -> None:
+    def _on_connection_open_error(
+        self, _unused_connection: AsyncioConnection, reason: pika.Exception
+    ) -> None:
         """Callback invoked when there is an error while opening connection."""
         self._reconnect()
+        self._log_pika_exception("Connection open failed", reason)
 
-    def _on_connection_closed(self, connection: AsyncioConnection, reason: str) -> None:
+    def _on_connection_closed(
+        self, connection: AsyncioConnection, _unused_reason: pika.Exception
+    ) -> None:
         """Callback invoked after the broker closes the connection"""
+        LOGGER.debug("Connection to RabbitMQ has been closed")
         self._channel = None
+        self._is_unrecoverable = True
         if self._closing:
             connection.ioloop.stop()
         else:
             self._reconnect()
 
-    def _close_connection(self) -> None:
+    def _close_connection(self, reason: pika.Exception) -> None:
         """Close the connection to the broker."""
         self._consuming = False
         if self._connection is None or self._connection.is_closing or self._connection.is_closed:
+            LOGGER.debug("Connection to RabbitMQ is closing or is already closed")
             pass
         else:
-            self._connection.close()
+            LOGGER.debug("Closing connection to RabbitMQ")
+
+            if isinstance(reason, (ConnectionClosed, ChannelClosed)):
+                reply_code = reason.reply_code
+                reply_text = reason.reply_text
+            else:
+                reply_code = 999  # arbitrary value, not in the list of AMQP reply codes: https://www.rabbitmq.com/amqp-0-9-1-reference#constants
+                reply_text = str(reason)
+            self._connection.close(reply_code=reply_code, reply_text=reply_text)
+
+    def _log_pika_exception(
+        self, message: str, reason: pika.Exception, extra: dict[str, str] | None = None
+    ) -> None:
+        """Log a pika exception. Extra is key-value pairs to include in the log message."""
+        if not extra:
+            extra = {}
+        if isinstance(reason, (ConnectionClosed, ChannelClosed)):
+            default_extra: dict[str, str] = {
+                "reply_code": str(reason.reply_code),
+                "reply_text": str(reason.reply_text),
+            }
+            LOGGER.error(
+                message,
+                # mypy not happy with dict | dict, so we use the dict constructor
+                extra={**default_extra, **extra},
+            )
+        else:
+            default_extra = {"reason": str(reason)}
+            # mypy not happy with dict | dict, so we use the dict constructor
+            LOGGER.error(message, extra={**default_extra, **extra})
 
     def _on_channel_open(self, channel: Channel, queue: str, on_message: OnMessageFn) -> None:
         """Callback invoked after the broker opens the channel."""
+        LOGGER.debug("Channel opened")
         self._channel = channel
         channel.add_on_close_callback(self._on_channel_closed)
         self._start_consuming(queue=queue, on_message=on_message, channel=channel)
 
-    def _on_channel_closed(self, channel: Channel, reason: str) -> None:
+    def _on_channel_closed(self, channel: Channel, reason: ChannelClosed) -> None:
         """Callback invoked after the broker closes the channel."""
-        self._close_connection()
+        self._log_pika_exception("Channel closed", reason, extra={"channel": channel})
+        self._close_connection(reason)
 
     def _build_client_parameters(self, url: str) -> pika.URLParameters:
         """Configure parameters used to connect to the broker."""

great_expectations_cloud/agent/message_service/subscriber.py

@@ -2,9 +2,10 @@ from __future__ import annotations
 
 import asyncio
 import time
+from collections.abc import Coroutine
 from dataclasses import dataclass
 from functools import partial
-from typing import TYPE_CHECKING, Callable, Coroutine, Protocol
+from typing import TYPE_CHECKING, Callable, Protocol
 
 from pika.exceptions import (
     AMQPError,
@@ -13,6 +14,7 @@ from pika.exceptions import (
 )
 
 from great_expectations_cloud.agent.event_handler import EventHandler
+from great_expectations_cloud.agent.exceptions import GXAgentUnrecoverableConnectionError
 
 if TYPE_CHECKING:
     from great_expectations_cloud.agent.message_service.asyncio_rabbit_mq_client import (
@@ -95,11 +97,13 @@ class Subscriber:
                 self.client.stop()
                 reconnect_delay = self._get_reconnect_delay()
                 time.sleep(reconnect_delay)  # todo: update this blocking call to asyncio.sleep
+                raise
             except KeyboardInterrupt as e:
                 self.client.stop()
                 raise KeyboardInterrupt from e
-            if self.client.should_reconnect:
-                self.client.reset()
+            except GXAgentUnrecoverableConnectionError:
+                self.client.stop()
+                raise
             else:
                 break  # exit
 
@@ -159,8 +163,7 @@ class Subscriber:
             self._reconnect_delay = 0
         else:
             self._reconnect_delay += 1
-        if self._reconnect_delay > 30:  # noqa: PLR2004
-            self._reconnect_delay = 30
+        self._reconnect_delay = min(self._reconnect_delay, 30)
         return self._reconnect_delay
 
     def close(self) -> None: