edda-framework 0.9.1__py3-none-any.whl → 0.11.0__py3-none-any.whl

edda/storage/pg_notify.py

@@ -0,0 +1,325 @@
+"""PostgreSQL LISTEN/NOTIFY implementation using asyncpg.
+
+This module provides a dedicated listener for PostgreSQL's LISTEN/NOTIFY
+mechanism, enabling near-instant notification delivery for workflow events.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+from collections.abc import Awaitable, Callable
+from contextlib import suppress
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+NotifyCallback = Callable[[str], Awaitable[None]]
+
+
+class PostgresNotifyListener:
+    """PostgreSQL LISTEN/NOTIFY listener using asyncpg.
+
+    This class maintains a dedicated connection for LISTEN/NOTIFY operations.
+    It provides:
+    - Automatic reconnection on connection loss
+    - Channel subscription management
+    - Callback dispatch for notifications
+
+    Example:
+        >>> listener = PostgresNotifyListener(dsn="postgresql://localhost/db")
+        >>> await listener.start()
+        >>> await listener.subscribe("my_channel", handle_notification)
+        >>> # ... later
+        >>> await listener.stop()
+    """
+
+    def __init__(
+        self,
+        dsn: str,
+        reconnect_interval: float = 5.0,
+        max_reconnect_attempts: int | None = None,
+    ) -> None:
+        """Initialize the PostgreSQL notify listener.
+
+        Args:
+            dsn: PostgreSQL connection string (postgresql://user:pass@host/db).
+            reconnect_interval: Seconds to wait between reconnection attempts.
+            max_reconnect_attempts: Maximum number of reconnection attempts.
+                                   None means unlimited.
+        """
+        self._dsn = dsn
+        self._reconnect_interval = reconnect_interval
+        self._max_reconnect_attempts = max_reconnect_attempts
+
+        self._connection: Any = None  # asyncpg.Connection
+        self._callbacks: dict[str, list[NotifyCallback]] = {}
+        self._channel_handlers: dict[str, Callable[..., None]] = {}
+        self._running = False
+        self._reconnect_task: asyncio.Task[None] | None = None
+        self._lock = asyncio.Lock()
+
+    async def start(self) -> None:
+        """Start the notification listener.
+
+        Establishes the connection and begins listening for notifications.
+        Starts the automatic reconnection task.
+
+        Raises:
+            ImportError: If asyncpg is not installed.
+        """
+        if self._running:
+            logger.warning("PostgresNotifyListener already running")
+            return
+
+        self._running = True
+        await self._establish_connection()
+
+        # Start reconnection monitor
+        self._reconnect_task = asyncio.create_task(self._reconnect_loop())
+        logger.info("PostgresNotifyListener started")
+
+    async def stop(self) -> None:
+        """Stop the notification listener.
+
+        Closes the connection and stops the reconnection task.
+        """
+        self._running = False
+
+        # Cancel reconnection task
+        if self._reconnect_task is not None:
+            self._reconnect_task.cancel()
+            with suppress(asyncio.CancelledError):
+                await self._reconnect_task
+            self._reconnect_task = None
+
+        # Close connection
+        await self._close_connection()
+        self._callbacks.clear()
+        logger.info("PostgresNotifyListener stopped")
+
+    async def subscribe(self, channel: str, callback: NotifyCallback) -> None:
+        """Subscribe to notifications on a channel.
+
+        Args:
+            channel: The PostgreSQL channel name to listen on.
+            callback: Async function called when a notification arrives.
+
+        Note:
+            Channel names must be valid PostgreSQL identifiers (max 63 chars).
+            Multiple callbacks can be registered for the same channel.
+        """
+        async with self._lock:
+            is_new_channel = channel not in self._callbacks
+
+            if channel not in self._callbacks:
+                self._callbacks[channel] = []
+            self._callbacks[channel].append(callback)
+
+            # Register listener if this is a new channel and we're connected
+            if is_new_channel and self._connection is not None:
+                try:
+                    await self._connection.add_listener(
+                        channel, self._create_notification_handler(channel)
+                    )
+                    logger.debug(f"Subscribed to channel: {channel}")
+                except Exception as e:
+                    logger.error(f"Failed to LISTEN on channel {channel}: {e}")
+
+    async def unsubscribe(self, channel: str) -> None:
+        """Unsubscribe from notifications on a channel.
+
+        Args:
+            channel: The PostgreSQL channel name to stop listening on.
+        """
+        async with self._lock:
+            if channel in self._callbacks:
+                del self._callbacks[channel]
+
+                # Remove listener if we're connected
+                if self._connection is not None:
+                    try:
+                        await self._connection.remove_listener(
+                            channel, self._create_notification_handler(channel)
+                        )
+                        logger.debug(f"Unsubscribed from channel: {channel}")
+                    except Exception as e:
+                        logger.error(f"Failed to UNLISTEN on channel {channel}: {e}")
+
+    async def notify(self, channel: str, payload: str) -> None:
+        """Send a notification on a channel.
+
+        Args:
+            channel: The PostgreSQL channel name.
+            payload: The payload string (max ~7500 bytes recommended).
+
+        Note:
+            This uses the existing connection pool from SQLAlchemy,
+            not the dedicated listener connection.
+        """
+        if self._connection is None:
+            logger.warning("Cannot send NOTIFY: not connected")
+            return
+
+        try:
+            # Use pg_notify function to properly escape the payload
+            await self._connection.execute("SELECT pg_notify($1, $2)", channel, payload)
+        except Exception as e:
+            logger.warning(f"Failed to send NOTIFY on channel {channel}: {e}")
+
+    @property
+    def is_connected(self) -> bool:
+        """Check if the listener is currently connected."""
+        return self._connection is not None and not self._connection.is_closed()
+
+    async def _establish_connection(self) -> None:
+        """Establish connection to PostgreSQL."""
+        try:
+            import asyncpg
+        except ImportError as e:
+            raise ImportError(
+                "asyncpg is required for PostgreSQL LISTEN/NOTIFY support. "
+                "Install it with: pip install edda[postgres-notify]"
+            ) from e
+
+        try:
+            self._connection = await asyncpg.connect(self._dsn)
+
+            # Re-subscribe to all channels (this also registers listeners)
+            await self._resubscribe_all()
+
+            logger.info("PostgresNotifyListener connected to database")
+        except Exception as e:
+            logger.error(f"Failed to connect to PostgreSQL: {e}")
+            self._connection = None
+            raise
+
+    async def _close_connection(self) -> None:
+        """Close the database connection."""
+        if self._connection is not None:
+            try:
+                await self._connection.close()
+            except Exception as e:
+                logger.warning(f"Error closing connection: {e}")
+            finally:
+                self._connection = None
+
+    async def _resubscribe_all(self) -> None:
+        """Re-subscribe to all registered channels after reconnection."""
+        if self._connection is None:
+            return
+
+        for channel in self._callbacks:
+            try:
+                # Register listener for this channel
+                await self._connection.add_listener(
+                    channel, self._create_notification_handler(channel)
+                )
+                logger.debug(f"Re-subscribed to channel: {channel}")
+            except Exception as e:
+                logger.error(f"Failed to re-subscribe to channel {channel}: {e}")
+
+    def _create_notification_handler(self, channel: str) -> Callable[..., None]:
+        """Create or retrieve a notification handler for a channel.
+
+        Args:
+            channel: The channel name.
+
+        Returns:
+            A handler function that can be passed to add_listener/remove_listener.
+        """
+        if channel not in self._channel_handlers:
+
+            def handler(_connection: Any, _pid: int, ch: str, payload: str) -> None:
+                """Handle incoming notification from PostgreSQL."""
+                callbacks = self._callbacks.get(ch, [])
+                for callback in callbacks:
+                    asyncio.create_task(self._safe_callback(callback, payload, ch))
+
+            self._channel_handlers[channel] = handler
+
+        return self._channel_handlers[channel]
+
+    async def _safe_callback(self, callback: NotifyCallback, payload: str, channel: str) -> None:
+        """Execute callback with error handling."""
+        try:
+            await callback(payload)
+        except Exception as e:
+            logger.error(
+                f"Error in notification callback for channel {channel}: {e}",
+                exc_info=True,
+            )
+
+    async def _reconnect_loop(self) -> None:
+        """Monitor connection and reconnect on failure."""
+        attempt = 0
+
+        with suppress(asyncio.CancelledError):
+            while self._running:
+                await asyncio.sleep(1)  # Check every second
+
+                if self._connection is None or self._connection.is_closed():
+                    attempt += 1
+                    if (
+                        self._max_reconnect_attempts is not None
+                        and attempt > self._max_reconnect_attempts
+                    ):
+                        logger.error(
+                            f"Max reconnection attempts ({self._max_reconnect_attempts}) "
+                            "exceeded, giving up"
+                        )
+                        break
+
+                    logger.info(
+                        f"Connection lost, attempting reconnection " f"(attempt {attempt})..."
+                    )
+
+                    try:
+                        await self._close_connection()
+                        await self._establish_connection()
+                        attempt = 0  # Reset on success
+                        logger.info("Reconnection successful")
+                    except Exception as e:
+                        logger.error(
+                            f"Reconnection failed: {e}, " f"retrying in {self._reconnect_interval}s"
+                        )
+                        await asyncio.sleep(self._reconnect_interval)
+
+
+def get_notify_channel_for_message(channel: str) -> str:
+    """Convert Edda channel name to PostgreSQL NOTIFY channel.
+
+    Uses a hash to ensure valid PostgreSQL identifier (max 63 chars).
+
+    Args:
+        channel: The Edda channel name.
+
+    Returns:
+        PostgreSQL-safe channel name.
+    """
+    import hashlib
+
+    h = hashlib.sha256(channel.encode()).hexdigest()[:16]
+    return f"edda_msg_{h}"
+
+
+def make_notify_payload(data: dict[str, Any]) -> str:
+    """Create JSON payload for NOTIFY.
+
+    Args:
+        data: Dictionary to serialize as JSON.
+
+    Returns:
+        JSON string (kept under 7500 bytes for PostgreSQL safety).
+    """
+    payload = json.dumps(data, separators=(",", ":"))  # Compact JSON
+    if len(payload) > 7500:
+        logger.warning(
+            f"NOTIFY payload exceeds recommended size " f"({len(payload)} > 7500 bytes), truncating"
+        )
+        # For safety, just include essential fields
+        minimal_data = {k: v for k, v in data.items() if k in ("wf_id", "ts")}
+        payload = json.dumps(minimal_data, separators=(",", ":"))
+    return payload

edda/storage/protocol.py

@@ -262,6 +262,7 @@ class StorageProtocol(Protocol):
         instance_id_filter: str | None = None,
         started_after: datetime | None = None,
         started_before: datetime | None = None,
+        input_filters: dict[str, Any] | None = None,
     ) -> dict[str, Any]:
         """
         List workflow instances with cursor-based pagination and filtering.
@@ -277,6 +278,9 @@ class StorageProtocol(Protocol):
             instance_id_filter: Optional instance ID filter (partial match, case-insensitive)
             started_after: Filter instances started after this datetime (inclusive)
             started_before: Filter instances started before this datetime (inclusive)
+            input_filters: Filter by input data values. Keys are JSON paths
+                (e.g., "order_id" or "customer.email"), values are expected
+                values (exact match). All filters are AND-combined.
 
         Returns:
             Dictionary containing:
@@ -860,7 +864,7 @@ class StorageProtocol(Protocol):
     # Workflow Resumption Methods
     # -------------------------------------------------------------------------
 
-    async def find_resumable_workflows(self) -> list[dict[str, Any]]:
+    async def find_resumable_workflows(self, limit: int | None = None) -> list[dict[str, Any]]:
         """
         Find workflows that are ready to be resumed.
 
@@ -873,6 +877,10 @@ class StorageProtocol(Protocol):
         This allows immediate resumption after message delivery rather than
         waiting for the stale lock cleanup cycle (60+ seconds).
 
+        Args:
+            limit: Optional maximum number of workflows to return.
+                   If None, returns all resumable workflows.
+
         Returns:
             List of resumable workflows.
             Each item contains: instance_id, workflow_name

edda/storage/sqlalchemy_storage.py

@@ -8,6 +8,7 @@ and transactional outbox pattern.
 
 import json
 import logging
+import re
 from collections.abc import AsyncIterator, Awaitable, Callable
 from contextlib import asynccontextmanager
 from contextvars import ContextVar
@@ -511,14 +512,43 @@ class SQLAlchemyStorage:
     - Automatic transaction management via @activity decorator
     """
 
-    def __init__(self, engine: AsyncEngine):
+    def __init__(
+        self,
+        engine: AsyncEngine,
+        notify_listener: Any | None = None,
+    ):
         """
         Initialize SQLAlchemy storage.
 
         Args:
             engine: SQLAlchemy AsyncEngine instance
+            notify_listener: Optional notify listener for PostgreSQL LISTEN/NOTIFY.
+                            If provided and PostgreSQL is used, NOTIFY messages
+                            will be sent after key operations.
         """
         self.engine = engine
+        self._notify_listener = notify_listener
+
+    @property
+    def _is_postgresql(self) -> bool:
+        """Check if the database is PostgreSQL."""
+        return self.engine.dialect.name == "postgresql"
+
+    @property
+    def _notify_enabled(self) -> bool:
+        """Check if NOTIFY is enabled (PostgreSQL with listener)."""
+        return self._is_postgresql and self._notify_listener is not None
+
+    def set_notify_listener(self, listener: Any) -> None:
+        """Set the notify listener after initialization.
+
+        This allows setting the listener after EddaApp creates the storage,
+        useful for dependency injection patterns.
+
+        Args:
+            listener: NotifyProtocol implementation (PostgresNotifyListener or NoopNotifyListener)
+        """
+        self._notify_listener = listener
 
     async def initialize(self) -> None:
         """Initialize database connection and create tables.
@@ -544,6 +574,39 @@ class SQLAlchemyStorage:
         """Close database connection."""
         await self.engine.dispose()
 
+    async def _send_notify(
+        self,
+        channel: str,
+        payload: dict[str, Any],
+    ) -> None:
+        """Send PostgreSQL NOTIFY message.
+
+        This method sends a notification on the specified channel with the given
+        payload. It's a no-op if NOTIFY is not enabled (non-PostgreSQL or no listener).
+
+        Args:
+            channel: PostgreSQL NOTIFY channel name (max 63 chars).
+            payload: Dictionary to serialize as JSON payload (max ~7500 bytes).
+        """
+        if not self._notify_enabled:
+            return
+
+        try:
+            import json as json_module
+
+            payload_str = json_module.dumps(payload, separators=(",", ":"))
+
+            # Use a separate connection for NOTIFY to avoid transaction issues
+            async with self.engine.connect() as conn:
+                await conn.execute(
+                    text("SELECT pg_notify(:channel, :payload)"),
+                    {"channel": channel, "payload": payload_str},
+                )
+                await conn.commit()
+        except Exception as e:
+            # Log but don't fail - polling will catch it as backup
+            logger.warning(f"Failed to send NOTIFY on channel {channel}: {e}")
+
     async def _initialize_schema_version(self) -> None:
         """Initialize schema version for a fresh database."""
         async with AsyncSession(self.engine) as session:
@@ -849,6 +912,60 @@ class SQLAlchemyStorage:
             # PostgreSQL/MySQL: column is already timezone-aware
             return column
 
+    def _validate_json_path(self, json_path: str) -> bool:
+        """
+        Validate JSON path to prevent SQL injection.
+
+        Args:
+            json_path: JSON path string (e.g., "order_id" or "customer.email")
+
+        Returns:
+            True if valid, False otherwise
+        """
+        # Only allow alphanumeric characters, dots, and underscores
+        # Must start with letter or underscore
+        return bool(re.match(r"^[a-zA-Z_][a-zA-Z0-9_.]*$", json_path))
+
+    def _build_json_extract_expr(self, column: Any, json_path: str) -> Any:
+        """
+        Build a database-agnostic JSON extraction expression.
+
+        Args:
+            column: SQLAlchemy column containing JSON text
+            json_path: Dot-notation path (e.g., "order_id" or "customer.email")
+
+        Returns:
+            SQLAlchemy expression that extracts the value at json_path
+
+        Raises:
+            ValueError: If json_path is invalid or dialect is unsupported
+        """
+        if not self._validate_json_path(json_path):
+            raise ValueError(f"Invalid JSON path: {json_path}")
+
+        full_path = f"$.{json_path}"
+        dialect = self.engine.dialect.name
+
+        if dialect == "sqlite":
+            # SQLite: json_extract(column, '$.key')
+            return func.json_extract(column, full_path)
+        elif dialect == "postgresql":
+            # PostgreSQL: For nested paths, use #>> operator with array path
+            # For simple paths, use ->> operator
+            # Since we're dealing with Text column, we need to cast first
+            if "." in json_path:
+                # Nested: (column)::json #>> '{customer,email}'
+                path_array = "{" + json_path.replace(".", ",") + "}"
+                return text(f"(input_data)::json #>> '{path_array}'")
+            else:
+                # Simple: (column)::json->>'key'
+                return text(f"(input_data)::json->>'{json_path}'")
+        elif dialect == "mysql":
+            # MySQL: JSON_UNQUOTE(JSON_EXTRACT(column, '$.key'))
+            return func.JSON_UNQUOTE(func.JSON_EXTRACT(column, full_path))
+        else:
+            raise ValueError(f"Unsupported database dialect: {dialect}")
+
     # -------------------------------------------------------------------------
     # Transaction Management Methods
     # -------------------------------------------------------------------------
@@ -1200,6 +1317,7 @@ class SQLAlchemyStorage:
         instance_id_filter: str | None = None,
         started_after: datetime | None = None,
         started_before: datetime | None = None,
+        input_filters: dict[str, Any] | None = None,
     ) -> dict[str, Any]:
         """List workflow instances with cursor-based pagination and filtering."""
         session = self._get_session_for_operation()
@@ -1292,6 +1410,42 @@ class SQLAlchemyStorage:
                         started_before_comparable = started_before
                     stmt = stmt.where(started_at_comparable <= started_before_comparable)
 
+            # Apply input data filters (JSON field matching)
+            if input_filters:
+                for json_path, expected_value in input_filters.items():
+                    dialect = self.engine.dialect.name
+                    if dialect == "postgresql":
+                        # PostgreSQL: use text() for the entire condition
+                        if not self._validate_json_path(json_path):
+                            raise ValueError(f"Invalid JSON path: {json_path}")
+                        if "." in json_path:
+                            path_array = "{" + json_path.replace(".", ",") + "}"
+                            json_sql = f"(input_data)::json #>> '{path_array}'"
+                        else:
+                            json_sql = f"(input_data)::json->>'{json_path}'"
+                        if expected_value is None:
+                            stmt = stmt.where(text(f"({json_sql} IS NULL OR {json_sql} = 'null')"))
+                        else:
+                            # Escape single quotes in value
+                            safe_value = str(expected_value).replace("'", "''")
+                            stmt = stmt.where(text(f"{json_sql} = '{safe_value}'"))
+                    else:
+                        # SQLite and MySQL: use func-based approach
+                        json_expr = self._build_json_extract_expr(
+                            WorkflowInstance.input_data, json_path
+                        )
+                        if expected_value is None:
+                            stmt = stmt.where(or_(json_expr.is_(None), json_expr == "null"))
+                        elif isinstance(expected_value, bool):
+                            stmt = stmt.where(json_expr == str(expected_value).lower())
+                        elif isinstance(expected_value, (int, float)):
+                            if dialect == "sqlite":
+                                stmt = stmt.where(json_expr == expected_value)
+                            else:
+                                stmt = stmt.where(json_expr == str(expected_value))
+                        else:
+                            stmt = stmt.where(json_expr == str(expected_value))
+
             # Fetch limit+1 to determine if there are more pages
             stmt = stmt.limit(limit + 1)
 
@@ -2096,6 +2250,12 @@ class SQLAlchemyStorage:
             session.add(event)
             await self._commit_if_not_in_transaction(session)
 
+        # Send NOTIFY for new outbox event
+        await self._send_notify(
+            "edda_outbox_pending",
+            {"evt_id": event_id, "evt_type": event_type},
+        )
+
     async def get_pending_outbox_events(self, limit: int = 10) -> list[dict[str, Any]]:
         """
         Get pending/failed outbox events for publishing (with row-level locking).
@@ -2719,29 +2879,34 @@ class SQLAlchemyStorage:
     # Workflow Resumption Methods
     # -------------------------------------------------------------------------
 
-    async def find_resumable_workflows(self) -> list[dict[str, Any]]:
+    async def find_resumable_workflows(self, limit: int | None = None) -> list[dict[str, Any]]:
         """
         Find workflows that are ready to be resumed.
 
         Returns workflows with status='running' that don't have an active lock.
         Used for immediate resumption after message delivery.
 
+        Args:
+            limit: Optional maximum number of workflows to return.
+                   If None, returns all resumable workflows.
+
         Returns:
             List of resumable workflows with instance_id and workflow_name.
         """
         session = self._get_session_for_operation()
         async with self._session_scope(session) as session:
-            result = await session.execute(
-                select(
-                    WorkflowInstance.instance_id,
-                    WorkflowInstance.workflow_name,
-                ).where(
-                    and_(
-                        WorkflowInstance.status == "running",
-                        WorkflowInstance.locked_by.is_(None),
-                    )
-                )
-            )
+            query = select(
+                WorkflowInstance.instance_id,
+                WorkflowInstance.workflow_name,
+            ).where(
+                and_(
+                    WorkflowInstance.status == "running",
+                    WorkflowInstance.locked_by.is_(None),
+                )
+            )
+            if limit is not None:
+                query = query.limit(limit)
+            result = await session.execute(query)
             return [
                 {
                     "instance_id": row.instance_id,
@@ -2837,6 +3002,15 @@ class SQLAlchemyStorage:
             session.add(msg)
             await self._commit_if_not_in_transaction(session)
 
+        # Send NOTIFY for message published (channel-specific)
+        import hashlib
+
+        channel_hash = hashlib.sha256(channel.encode()).hexdigest()[:16]
+        await self._send_notify(
+            f"edda_msg_{channel_hash}",
+            {"ch": channel, "msg_id": message_id},
+        )
+
         return message_id
 
     async def subscribe_to_channel(
@@ -3395,6 +3569,12 @@ class SQLAlchemyStorage:
 
                     await session.commit()
 
+                # Send NOTIFY for workflow resumable
+                await self._send_notify(
+                    "edda_workflow_resumable",
+                    {"wf_id": instance_id, "wf_name": workflow_name},
+                )
+
                 return {
                     "instance_id": instance_id,
                     "workflow_name": workflow_name,