remoulade 7.0.0.dev0__tar.gz → 7.0.0.dev2__tar.gz

{remoulade-7.0.0.dev0 → remoulade-7.0.0.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: remoulade
-Version: 7.0.0.dev0
+Version: 7.0.0.dev2
 Summary: Background Processing for Python 3.
 Author-email: Wiremind <dev@wiremind.io>
 License-Expression: LGPL-3.0-or-later

remoulade-7.0.0.dev2/remoulade/api/main.py

@@ -0,0 +1,29 @@
+"""This file describe the API to get the state of messages"""
+
+from flask import Flask
+from marshmallow import ValidationError
+from werkzeug.exceptions import HTTPException
+
+from remoulade.errors import RemouladeError
+
+from .apispec import add_swagger
+
+app = Flask(__name__)
+
+
+@app.errorhandler(RemouladeError)
+def remoulade_exception(e):
+    return {"error": str(e)}, 500
+
+
+@app.errorhandler(HTTPException)
+def http_exception(e):
+    return {"error": str(e)}, e.code
+
+
+@app.errorhandler(ValidationError)
+def validation_error(e):
+    return {"error": e.normalized_messages()}, 400
+
+
+add_swagger(app, {"main": []})

{remoulade-7.0.0.dev0 → remoulade-7.0.0.dev2}/remoulade/brokers/__init__.py

@@ -15,9 +15,27 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
+import warnings
+
 from .local import LocalBroker
-from .pgmq import PgmqBroker
-from .rabbitmq import RabbitmqBroker
 from .stub import StubBroker
 
 __all__ = ["LocalBroker", "PgmqBroker", "RabbitmqBroker", "StubBroker"]
+
+try:
+    from .pgmq import PgmqBroker
+except ImportError:  # pragma: no cover
+    warnings.warn(
+        "PgmqBroker is not available.  Run `pip install remoulade[pgmq]` to add support for that broker.",
+        ImportWarning,
+        stacklevel=2,
+    )
+
+try:
+    from .rabbitmq import RabbitmqBroker
+except ImportError:  # pragma: no cover
+    warnings.warn(
+        "RabbitmqBroker is not available.  Run `pip install remoulade[rabbitmq]` to add support for that broker.",
+        ImportWarning,
+        stacklevel=2,
+    )

{remoulade-7.0.0.dev0 → remoulade-7.0.0.dev2}/remoulade/brokers/pgmq.py

@@ -1,6 +1,6 @@
 # This file is a part of Remoulade.
 #
-# Copyright (C) 2017,2018 CLEARTYPE SRL <bogdan@cleartype.io>
+# Copyright (C) 2017,2018 WIREMIND SAS <dev@wiremind.fr>
 #
 # Remoulade is free software; you can redistribute it and/or modify it
 # under the terms of the GNU Lesser General Public License as published by
@@ -15,40 +15,31 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 import json
+import time
 from collections import deque
 from collections.abc import Iterable, Iterator
 from contextlib import contextmanager
 from datetime import UTC, datetime, timedelta
-from threading import Event, Thread, local
+from threading import Event, Lock, Thread, local
 from typing import TYPE_CHECKING, Any
 
 import psycopg
 from pgmq import SQLAlchemyPGMQueue
+from pgmq.messages import Message as PgmqQueueMessage
 from psycopg import sql as psycopg_sql
-from pydantic import BaseModel, ConfigDict, ValidationError
-from sqlalchemy import bindparam, create_engine, text
+from sqlalchemy import Column, Integer, MetaData, Table, bindparam, create_engine, func, select, text
 from sqlalchemy.dialects.postgresql import JSONB
-from sqlalchemy.engine import make_url
 from sqlalchemy.orm import sessionmaker as sqlalchemy_sessionmaker
 
 from ..broker import Broker, Consumer, MessageProxy
-from ..errors import QueueNotFound, UnsupportedMessageEncoding
+from ..errors import QueueJoinTimeout, QueueNotFound, UnsupportedMessageEncoding
 from ..message import Message
 
 if TYPE_CHECKING:
     from ..middleware import Middleware
 
-
 PgmqPayload = dict[str, Any]
 
-
-class _PgmqQueueMessage(BaseModel):
-    model_config = ConfigDict(extra="ignore", from_attributes=True)
-
-    msg_id: int
-    message: PgmqPayload
-
-
 DELAYED_SEND_SQL = text(
     """
     SELECT *
@@ -75,8 +66,31 @@ class PgmqBroker(Broker):
         archive_partition_interval: int | str = "1 day",
         archive_retention_interval: int | str = "7 days",
         listen_notify_enabled: bool = True,
+        visibility_timeout: int = 30,
+        heartbeat_interval: float = 10.0,
     ) -> None:
+        """Initialize a PostgreSQL-backed broker using the PGMQ extension.
+
+        Parameters:
+          url(str): PostgreSQL URL in plain format (`postgresql://...`), used both by SQLAlchemy and psycopg.
+          middleware(list[Middleware] | None): Middleware stack applied to this broker.
+          group_transaction(bool): If True, wraps group and pipeline operations in a single transaction.
+          partition_archive_on_queue_init(bool): If True, creates partitioned queues when declaring queues.
+          archive_partition_interval(int | str): Partition interval passed to PGMQ when creating partitioned queues.
+          archive_retention_interval(int | str): Retention interval passed to PGMQ for archive partitions.
+          listen_notify_enabled(bool): If True, enables LISTEN/NOTIFY to wake consumers when new messages arrive.
+          visibility_timeout(int): Message visibility timeout in seconds after read; must be greater than 0.
+          heartbeat_interval(float): Heartbeat interval in seconds used to extend in-flight message visibility;
+            must be greater than 0 and lower than visibility_timeout.
+
+        Raises:
+          ValueError: If visibility_timeout or heartbeat_interval values are invalid.
+        """
         super().__init__(middleware=middleware)
+        if visibility_timeout <= 0:
+            raise ValueError("visibility_timeout must be greater than 0")
+        if heartbeat_interval <= 0 or heartbeat_interval >= visibility_timeout:
+            raise ValueError("heartbeat_interval must be greater than 0 and lower than visibility_timeout")
 
         self.url = url
         self.state = local()
@@ -85,12 +99,14 @@ class PgmqBroker(Broker):
         self.archive_partition_interval = archive_partition_interval
         self.archive_retention_interval = archive_retention_interval
         self.listen_notify_enabled = listen_notify_enabled
+        self.visibility_timeout = visibility_timeout
+        self.heartbeat_interval = heartbeat_interval
 
         self.engine = create_engine(self.url, pool_pre_ping=True)
         self.sessionmaker = sqlalchemy_sessionmaker(bind=self.engine)
         self.supports_native_delay = True
 
-        self.client = SQLAlchemyPGMQueue(engine=self.engine, init_extension=False)
+        self.client = SQLAlchemyPGMQueue(engine=self.engine, init_extension=True, vt=self.visibility_timeout)
 
     @contextmanager
     def tx(self) -> Iterator[None]:
@@ -107,10 +123,6 @@ class PgmqBroker(Broker):
     def _current_connection(self) -> Any | None:
         return getattr(self.state, "transaction_connection", None)
 
-    def _build_listener_conninfo(self) -> str:
-        # psycopg expects a plain PostgreSQL URL and rejects SQLAlchemy-style driver suffixes.
-        return make_url(self.url).set(drivername="postgresql").render_as_string(hide_password=False)
-
     def _try_enable_notify(self, queue_name: str) -> None:
         if not self.listen_notify_enabled:
             return
@@ -208,26 +220,92 @@ class PgmqBroker(Broker):
         for queue_name in self.queues:
             self.flush(queue_name)
 
-    def join(self, queue_name: str, *, timeout: int | None = None) -> None:
-        raise NotImplementedError("PgmqBroker.join() will be implemented in jalon 3.")
+    def _count_enqueued_messages(self, queue_name: str) -> int:
+        queue_table = Table(f"q_{queue_name}", MetaData(), Column("msg_id", Integer), schema="pgmq")
+        count_query = select(func.count()).select_from(queue_table)
+
+        with self.sessionmaker.begin() as session:
+            return session.execute(count_query).scalar_one()
+
+    def join(
+        self,
+        queue_name: str,
+        min_successes: int = 10,
+        idle_time: int = 100,
+        *,
+        timeout: int | None = None,
+    ) -> None:
+        """Wait for all the messages on the given queue to be processed.
+
+        This method checks the full PGMQ queue table and therefore waits for
+        all states: visible messages, invisible in-flight messages and native
+        delayed messages.
+
+        Parameters:
+          queue_name(str): The queue to wait on.
+          min_successes(int): The minimum number of times the queue should be
+            observed as empty.
+          idle_time(int): The number of milliseconds to wait between checks.
+          timeout(Optional[int]): The max amount of time, in milliseconds, to
+            wait on this queue.
+
+        Raises:
+          QueueNotFound: If the given queue was never declared.
+          QueueJoinTimeout: When the timeout elapses.
+        """
+        if queue_name not in self.queues:
+            raise QueueNotFound(queue_name)
+
+        deadline = time.monotonic() + timeout / 1000 if timeout is not None else None
+        successes = 0
+
+        while successes < min_successes:
+            if deadline and time.monotonic() >= deadline:
+                raise QueueJoinTimeout(queue_name)
+
+            total_messages = self._count_enqueued_messages(queue_name)
+            if total_messages == 0:
+                successes += 1
+                if successes >= min_successes:
+                    return
+            else:
+                successes = 0
+
+            time.sleep(idle_time / 1000)
 
 
 class _PgmqConsumer(Consumer):
     def __init__(self, broker: PgmqBroker, *, queue_name: str, prefetch: int, timeout: int) -> None:
+        """Initialize a consumer for a PGMQ queue.
+
+        Parameters:
+          broker(PgmqBroker): Broker instance that owns the queue and database client.
+          queue_name(str): Name of the declared queue to consume from.
+          prefetch(int): Maximum number of messages fetched per read call; values lower than 1 are coerced to 1.
+          timeout(int): Idle wait timeout in milliseconds when polling for messages; values lower than 0 are coerced to
+          0. A value of 0 performs non-blocking reads.
+        """
         self.broker = broker
         self.client = broker.client
         self.queue_name = queue_name
         self.prefetch = max(prefetch, 1)
         self.timeout = max(timeout, 0)
-        self.messages: deque[_PgmqQueueMessage] = deque()
+        self.visibility_timeout = broker.visibility_timeout
+        self.heartbeat_interval = broker.heartbeat_interval
+        self.messages: deque[PgmqQueueMessage] = deque()
         self._notify_event = Event()
         self._listener_stop = Event()
         self._listener_thread: Thread | None = None
         self._listener_connection: psycopg.Connection[Any] | None = None
         self._listener_available = False
+        self._heartbeat_stop = Event()
+        self._heartbeat_thread: Thread | None = None
+        self._heartbeat_msg_ids_lock = Lock()
+        self._heartbeat_msg_ids: set[int] = set()
 
         if self.broker.listen_notify_enabled:
             self._start_listener()
+        self._start_heartbeat()
 
     @property
     def max_poll_seconds(self) -> int:
@@ -241,38 +319,66 @@ class _PgmqConsumer(Consumer):
     def wait_timeout_seconds(self) -> float:
         return self.timeout / 1000 if self.timeout > 0 else 0
 
-    def _normalize_messages(self, messages: Any | list[Any] | None) -> list[_PgmqQueueMessage]:
+    def _normalize_messages(self, messages: PgmqQueueMessage | list[PgmqQueueMessage] | None) -> list[PgmqQueueMessage]:
         if messages is None:
             return []
-        raw_messages = messages if isinstance(messages, list) else [messages]
-        normalized_messages: list[_PgmqQueueMessage] = []
-        for message in raw_messages:
-            try:
-                normalized_messages.append(_PgmqQueueMessage.model_validate(message, from_attributes=True))
-            except ValidationError as exc:
-                raise UnsupportedMessageEncoding(
-                    "PGMQ messages must expose 'msg_id' and a JSONB-serializable object in 'message'."
-                ) from exc
-        return normalized_messages
+        return [messages] if not isinstance(messages, list) else messages
 
-    def _read_immediate(self) -> list[_PgmqQueueMessage]:
-        return self._normalize_messages(self.client.read(self.queue_name, qty=self.prefetch))
+    def _read_immediate(self) -> list[PgmqQueueMessage]:
+        return self._normalize_messages(
+            self.client.read(
+                self.queue_name,
+                vt=self.visibility_timeout,
+                qty=self.prefetch,
+            )
+        )
 
-    def _read_with_poll(self) -> list[_PgmqQueueMessage]:
+    def _read_with_poll(self) -> list[PgmqQueueMessage]:
         return self._normalize_messages(
             self.client.read_with_poll(
                 self.queue_name,
+                vt=self.visibility_timeout,
                 qty=self.prefetch,
                 max_poll_seconds=self.max_poll_seconds,
                 poll_interval_ms=self.poll_interval_ms,
             )
         )
 
+    def _start_heartbeat(self) -> None:
+        self._heartbeat_thread = Thread(
+            target=self._run_heartbeat,
+            name=f"pgmq-heartbeat-{self.queue_name}",
+            daemon=True,
+        )
+        self._heartbeat_thread.start()
+
+    def _run_heartbeat(self) -> None:
+        while not self._heartbeat_stop.wait(self.heartbeat_interval):
+            with self._heartbeat_msg_ids_lock:
+                msg_ids = list(self._heartbeat_msg_ids)
+            if not msg_ids:
+                continue
+            try:
+                self.client.set_vt(self.queue_name, msg_ids, self.visibility_timeout)
+            except Exception:
+                self.broker.logger.warning(
+                    "Failed to extend visibility timeout heartbeat for queue %s.",
+                    self.queue_name,
+                    exc_info=True,
+                )
+
+    def _register_heartbeat_msg(self, message: "_PgmqMessage") -> None:
+        with self._heartbeat_msg_ids_lock:
+            self._heartbeat_msg_ids.add(message._pgmq_message.msg_id)
+
+    def _unregister_heartbeat_msg_id(self, msg_id: int) -> None:
+        with self._heartbeat_msg_ids_lock:
+            self._heartbeat_msg_ids.discard(msg_id)
+
     def _start_listener(self) -> None:
         channel = f"pgmq.q_{self.queue_name}.INSERT"
         try:
-            conninfo = self.broker._build_listener_conninfo()
-            self._listener_connection = psycopg.connect(conninfo, autocommit=True)
+            self._listener_connection = psycopg.connect(self.broker.url, autocommit=True)
             self._listener_connection.execute(psycopg_sql.SQL("LISTEN {}").format(psycopg_sql.Identifier(channel)))
             self._listener_available = True
         except Exception:
@@ -316,21 +422,30 @@ class _PgmqConsumer(Consumer):
     def ack(self, message: "MessageProxy") -> None:
         if not isinstance(message, _PgmqMessage):
             raise ValueError("It must be a PgmqMessage")
-        self.client.delete(self.queue_name, message._pgmq_message.msg_id)
+        self._unregister_heartbeat_msg_id(message._pgmq_message.msg_id)
+        self.client.archive(self.queue_name, message._pgmq_message.msg_id)
 
     def nack(self, message: "MessageProxy") -> None:
         if not isinstance(message, _PgmqMessage):
             raise ValueError("It must be a PgmqMessage")
+        self._unregister_heartbeat_msg_id(message._pgmq_message.msg_id)
         self.client.archive(self.queue_name, message._pgmq_message.msg_id)
 
     def requeue(self, messages: Iterable["MessageProxy"]) -> None:
         msg_ids = [message._pgmq_message.msg_id for message in messages if isinstance(message, _PgmqMessage)]
         if msg_ids:
+            for msg_id in msg_ids:
+                self._unregister_heartbeat_msg_id(msg_id)
             self.client.set_vt(self.queue_name, msg_ids, 0)
 
+    def _build_message(self, pgmq_message: PgmqQueueMessage) -> "_PgmqMessage":
+        message = _PgmqMessage(pgmq_message)
+        self._register_heartbeat_msg(message)
+        return message
+
     def __next__(self) -> "_PgmqMessage | None":
         if self.messages:
-            return _PgmqMessage(self.messages.popleft())
+            return self._build_message(self.messages.popleft())
 
         if self._listener_available:
             self._notify_event.clear()
@@ -345,10 +460,11 @@ class _PgmqConsumer(Consumer):
             return None
 
         self.messages.extend(messages)
-        return _PgmqMessage(self.messages.popleft())
+        return self._build_message(self.messages.popleft())
 
     def close(self) -> None:
         self._listener_stop.set()
+        self._heartbeat_stop.set()
         self._notify_event.set()
         if self._listener_connection is not None:
             try:
@@ -359,17 +475,21 @@ class _PgmqConsumer(Consumer):
                 self._listener_connection = None
         if self._listener_thread is not None:
             self._listener_thread.join(timeout=1)
+        if self._heartbeat_thread is not None:
+            self._heartbeat_thread.join(timeout=1)
         return
 
 
 class _PgmqMessage(MessageProxy):
-    def __init__(self, pgmq_message: _PgmqQueueMessage) -> None:
+    def __init__(self, pgmq_message: PgmqQueueMessage) -> None:
         payload = pgmq_message.message
         if not isinstance(payload, dict):
             raise UnsupportedMessageEncoding("PGMQ messages must contain JSON objects.")
 
         try:
-            message: Message = Message(**payload)
+            # Re-run the global message decoder so custom encoders (e.g. PydanticEncoder)
+            # can rehydrate actor args/kwargs to their typed schemas.
+            message = Message.decode(json.dumps(payload, separators=(",", ":")).encode("utf-8"))
         except TypeError as exc:
             raise UnsupportedMessageEncoding("PGMQ message payload is not a valid Remoulade message envelope.") from exc
 

{remoulade-7.0.0.dev0 → remoulade-7.0.0.dev2}/remoulade.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: remoulade
-Version: 7.0.0.dev0
+Version: 7.0.0.dev2
 Summary: Background Processing for Python 3.
 Author-email: Wiremind <dev@wiremind.io>
 License-Expression: LGPL-3.0-or-later

remoulade-7.0.0.dev0/remoulade/api/main.py

@@ -1,175 +0,0 @@
-"""This file describe the API to get the state of messages"""
-
-import sys
-from typing import Any, TypedDict
-
-from flask import Flask
-from flask_apispec import marshal_with
-from marshmallow import Schema, ValidationError, fields, validate, validates_schema
-from werkzeug.exceptions import HTTPException
-
-import remoulade
-from remoulade import get_broker
-from remoulade.errors import NoResultBackend, NoStateBackend, RemouladeError
-from remoulade.result import Result
-from remoulade.results import ResultMissing
-
-from .apispec import add_swagger, validate_schema
-from .scheduler import scheduler_bp, scheduler_routes
-from .state import messages_bp, messages_routes
-
-app = Flask(__name__)
-app.register_blueprint(scheduler_bp)
-app.register_blueprint(messages_bp)
-
-
-class MessageSchema(Schema):
-    """
-    Class to validate post data in /messages
-    """
-
-    actor_name = fields.Str(validate=validate.Length(min=1), required=True)
-    args = fields.List(fields.Raw(), allow_none=True)
-    kwargs = fields.Dict(allow_none=True)
-    options = fields.Dict(allow_none=True)
-    delay = fields.Float(validate=validate.Range(min=1), allow_none=True)
-
-    @validates_schema
-    def validate_actor_name(self, data, **kwargs):
-        actor_name = data.get("actor_name")
-        if actor_name not in remoulade.get_broker().actors:
-            raise ValidationError(f"No actor named {actor_name} exists")
-
-
-class ResponseSchema(Schema):
-    result = fields.Raw(allow_none=True)
-    error = fields.Str(allow_none=True)
-
-
-class ArgsSchema(Schema):
-    name = fields.Str()
-    type = fields.Str()
-    default = fields.Str(allow_none=True)
-
-
-class ActorSchema(Schema):
-    name = fields.Str()
-    queue_name = fields.Str()
-    alternative_queues = fields.List(fields.Str())
-    priority = fields.Int()
-    args = fields.List(fields.Nested(ArgsSchema))
-
-
-class ActorResponseSchema(Schema):
-    result = fields.List(fields.Nested(ActorSchema))
-
-
-class OptionsResponseSchema(Schema):
-    options = fields.List(fields.Str())
-
-
-@app.route("/messages/cancel/<message_id>", methods=["POST"])
-@marshal_with(ResponseSchema)
-def cancel_message(message_id):
-    broker = get_broker()
-    cancel_backend = broker.get_cancel_backend()
-    try:
-        # If a single message in a composition is canceled, we cancel the whole composition
-        state_backend = broker.get_state_backend()
-        states_count = state_backend.get_states_count(selected_composition_ids=[message_id])
-        state = state_backend.get_state(message_id)
-        if states_count == 0 and state is None:
-            raise ValidationError(f"This message id or composition id {message_id} does not exist.")
-        if state and state.composition_id:
-            message_id = state.composition_id
-    except NoStateBackend:
-        pass
-    cancel_backend.cancel([message_id])
-    return {"result": "ok"}
-
-
-@app.route("/messages/requeue/<message_id>", methods=["POST"])
-@marshal_with(ResponseSchema)
-def requeue_message(message_id):
-    broker = get_broker()
-    backend = broker.get_state_backend()
-    state = backend.get_state(message_id)
-    if state is None:
-        raise ValidationError(f"No message with id {message_id} exists")
-    actor = broker.get_actor(state.actor_name)
-    payload = {"args": state.args, "kwargs": state.kwargs}
-    pipe_target = state.options.get("pipe_target")
-    if pipe_target is None:
-        actor.send_with_options(**payload, **state.options)
-        return {"result": "ok"}
-    else:
-        return {"error": "requeue message in a pipeline not supported"}, 400
-
-
-@app.route("/messages/result/<message_id>")
-@marshal_with(ResponseSchema)
-def get_results(message_id):
-    from ..message import get_encoder
-
-    max_size = 1e4
-    try:
-        result = Result[Any](message_id=message_id).get()
-        encoded_result = get_encoder().encode(result).decode("utf-8")
-        size_result = sys.getsizeof(encoded_result)
-        if size_result >= max_size:
-            encoded_result = f"The result is too big {size_result / 1e6}M"
-        return {"result": encoded_result}
-    except ResultMissing:
-        return {"result": "result is missing"}
-    except NoResultBackend:
-        return {"result": "no result backend"}
-    except (UnicodeDecodeError, TypeError):
-        return {"result": "non serializable result"}
-
-
-@app.route("/messages", methods=["POST"])
-@marshal_with(ResponseSchema)
-@validate_schema(MessageSchema)
-def enqueue_message(**kwargs):
-    actor = get_broker().get_actor(kwargs.pop("actor_name"))
-    options = kwargs.pop("options") or {}
-    actor.send_with_options(**kwargs, **options)
-    return {"result": "ok"}
-
-
-class GroupMessagesT(TypedDict):
-    group_id: str
-    messages: list[dict]
-
-
-@app.route("/actors")
-@marshal_with(ActorResponseSchema)
-def get_actors():
-    return {"result": [actor.as_dict() for actor in get_broker().actors.values()]}
-
-
-@app.route("/options")
-@marshal_with(OptionsResponseSchema)
-def get_options():
-    broker = get_broker()
-    return {"options": list(broker.actor_options)}
-
-
-@app.errorhandler(RemouladeError)
-def remoulade_exception(e):
-    return {"error": str(e)}, 500
-
-
-@app.errorhandler(HTTPException)
-def http_exception(e):
-    return {"error": str(e)}, e.code
-
-
-@app.errorhandler(ValidationError)
-def validation_error(e):
-    return {"error": e.normalized_messages()}, 400
-
-
-routes = [cancel_message, requeue_message, get_results, enqueue_message, get_actors, get_options]
-
-add_swagger(app, {"main": routes, "scheduler": scheduler_routes, "messages": messages_routes})