From 45d27d6f0e7c46d5f8b72166c4093cbe3dd508d1 Mon Sep 17 00:00:00 2001 From: Arseniy Popov Date: Sun, 19 Jul 2026 19:15:23 +0300 Subject: [PATCH 1/2] feat: wire up batch publishing --- docs/docs/sqlbroker/design.md | 4 +- docs/docs/sqlbroker/tutorial.md | 10 +- docs/docs_src/sqlbroker/publish_batch.py | 42 ++++ faststream_sqlbroker/__init__.py | 4 + faststream_sqlbroker/sqlbroker/__init__.py | 8 +- faststream_sqlbroker/sqlbroker/client.py | 13 +- .../sqlbroker/publisher/producer.py | 30 +-- .../sqlbroker/publisher/usecase.py | 26 ++ faststream_sqlbroker/sqlbroker/response.py | 98 +++++++- pyproject.toml | 2 +- tests/docs/sqlbroker/test_publish_batch.py | 72 ++++++ tests/test_publish.py | 237 +++++++++++++++--- uv.lock | 2 +- 13 files changed, 473 insertions(+), 75 deletions(-) create mode 100644 docs/docs_src/sqlbroker/publish_batch.py create mode 100644 tests/docs/sqlbroker/test_publish_batch.py diff --git a/docs/docs/sqlbroker/design.md b/docs/docs/sqlbroker/design.md index 3952b2a..5c380fa 100644 --- a/docs/docs/sqlbroker/design.md +++ b/docs/docs/sqlbroker/design.md @@ -15,7 +15,7 @@ search: ## Message Lifecycle -A published message starts out as `PENDING`. When it is acquired by a worker, it is marked as `PROCESSING`, which prevents it from being acquired by other worker processes. If the maximum allowed number of deliveries is configured and exceeded, the message is marked as `FAILED`. If not, the message is processed. If processing didn't raise an exception or if the message was manually [Acked](../sqlbroker/tutorial.md#ack){.internal-link} in the handler, it is marked as `COMPLETED`. If the message was manually [Nacked](../sqlbroker/tutorial.md#nack){.internal-link} or [Rejected](../sqlbroker/tutorial.md#reject){.internal-link} or if processing raised an exception and [`AckPolicy`](../getting-started/acknowledgement.md){.internal-link} was set to `REJECT_ON_ERROR` or `NACK_ON_ERROR`, the message is [Nacked](../sqlbroker/tutorial.md#nack){.internal-link} or [Rejected](../sqlbroker/tutorial.md#reject){.internal-link}. [Rejected](../sqlbroker/tutorial.md#reject){.internal-link} messages are marked as `FAILED`. For [Nacked](../sqlbroker/tutorial.md#nack){.internal-link} messages, the retry policy determines if the message is allowed to be retried. If retry is allowed, the message is marked as `RETRYABLE`. If not, the message is marked as `FAILED`. +A published message starts out as `PENDING`. When it is acquired by a worker, it is marked as `PROCESSING`, which prevents it from being acquired by other worker processes. If the maximum allowed number of deliveries is configured and exceeded, the message is marked as `FAILED`. If not, the message is processed. If processing didn't raise an exception or if the message was manually [Acked](../sqlbroker/tutorial.md#ack){.internal-link} in the handler, it is marked as `COMPLETED`. If the message was manually [Nacked](../sqlbroker/tutorial.md#nack){.internal-link} or [Rejected](../sqlbroker/tutorial.md#reject){.internal-link} or if processing raised an exception and [`AckPolicy`](../sqlbroker/tutorial.md#acknowledgements){.internal-link} was set to `REJECT_ON_ERROR` or `NACK_ON_ERROR`, the message is [Nacked](../sqlbroker/tutorial.md#nack){.internal-link} or [Rejected](../sqlbroker/tutorial.md#reject){.internal-link}. [Rejected](../sqlbroker/tutorial.md#reject){.internal-link} messages are marked as `FAILED`. For [Nacked](../sqlbroker/tutorial.md#nack){.internal-link} messages, the retry policy determines if the message is allowed to be retried. If retry is allowed, the message is marked as `RETRYABLE`. If not, the message is marked as `FAILED`. `PENDING`, `PROCESSING`, and `RETRYABLE` messages reside in the main table. On status change, `COMPLETED` and `FAILED` messages are removed from the main table and, depending on `retain_in_archive_on_ack` and `retain_in_archive_on_reject`, copied to the archive table. @@ -39,7 +39,7 @@ On start, the subscriber spawns four types of concurrent loops: **1. Fetch loop** — Periodically fetches batches of `PENDING` or `RETRYABLE` messages from the database, simultaneously updating them: marking as `PROCESSING`, setting `acquired_at` to now, and incrementing `deliveries_count`. Only messages with `next_attempt_at <= now` are fetched, ordered by `next_attempt_at`. The fetched messages are placed into an internal queue. The fetch limit is the minimum of `fetch_batch_size`, the free acquired-but-not-yet-processed capacity (`fetch_batch_size * max_not_processed_factor` minus currently unprocessed messages), and the free acquired-but-not-yet-persisted capacity (`fetch_batch_size * max_not_persisted_factor` minus currently unpersisted messages). If the last fetch was "full" (returned as many messages as the limit), the next fetch happens after `min_fetch_interval`; otherwise after `max_fetch_interval`. -**2. Worker loops** (`max_workers` concurrent instances) — Each worker takes a message from the internal queue and first checks if `max_deliveries` has been exceeded; if so, the message is [Rejected](../sqlbroker/tutorial.md#reject){.internal-link} without processing. Otherwise, processing proceeds. Depending on the processing result, [`AckPolicy`](../getting-started/acknowledgement.md){.internal-link}, and manual [Ack](../sqlbroker/tutorial.md#ack){.internal-link}/[Nack](../sqlbroker/tutorial.md#nack){.internal-link}/[Reject](../sqlbroker/tutorial.md#reject){.internal-link}, the message is [Acked](../sqlbroker/tutorial.md#ack){.internal-link}, [Nacked](../sqlbroker/tutorial.md#nack){.internal-link}, or [Rejected](../sqlbroker/tutorial.md#reject){.internal-link}. For [Nacked](../sqlbroker/tutorial.md#nack){.internal-link} messages, the `retry_strategy` is consulted to determine if and when the message might be retried. If allowed to be retried, the message is marked as `RETRYABLE`; otherwise as `FAILED`. [Acked](../sqlbroker/tutorial.md#ack){.internal-link} messages are marked as `COMPLETED` and rejected messages are marked as `FAILED`. The message is then buffered for flushing. +**2. Worker loops** (`max_workers` concurrent instances) — Each worker takes a message from the internal queue and first checks if `max_deliveries` has been exceeded; if so, the message is [Rejected](../sqlbroker/tutorial.md#reject){.internal-link} without processing. Otherwise, processing proceeds. Depending on the processing result, [`AckPolicy`](../sqlbroker/tutorial.md#acknowledgements){.internal-link}, and manual [Ack](../sqlbroker/tutorial.md#ack){.internal-link}/[Nack](../sqlbroker/tutorial.md#nack){.internal-link}/[Reject](../sqlbroker/tutorial.md#reject){.internal-link}, the message is [Acked](../sqlbroker/tutorial.md#ack){.internal-link}, [Nacked](../sqlbroker/tutorial.md#nack){.internal-link}, or [Rejected](../sqlbroker/tutorial.md#reject){.internal-link}. For [Nacked](../sqlbroker/tutorial.md#nack){.internal-link} messages, the `retry_strategy` is consulted to determine if and when the message might be retried. If allowed to be retried, the message is marked as `RETRYABLE`; otherwise as `FAILED`. [Acked](../sqlbroker/tutorial.md#ack){.internal-link} messages are marked as `COMPLETED` and rejected messages are marked as `FAILED`. The message is then buffered for flushing. **3. Flush loop** — Periodically flushes the buffered message state changes to the database. `COMPLETED` and `FAILED` messages are removed from the primary table and, depending on `retain_in_archive_on_ack` and `retain_in_archive_on_reject`, copied to the archive table. The state of `RETRYABLE` messages is updated in the primary table. diff --git a/docs/docs/sqlbroker/tutorial.md b/docs/docs/sqlbroker/tutorial.md index 04c7ec7..66935f1 100644 --- a/docs/docs/sqlbroker/tutorial.md +++ b/docs/docs/sqlbroker/tutorial.md @@ -66,7 +66,7 @@ The `COMPETING_CONSUMERS` variant (version `1`) uses up to two tables — `messa {!> docs_src/sqlbroker/publish.py [ln:1-16]!} ``` -The broker's and publisher's (see [publishing](../getting-started/publishing/index.md){.internal-link}) `.publish()` methods accept: +The broker's and publisher's (see [publishing](../getting-started/publishing/index.md){.external-link target="_blank"}) `.publish()` methods accept: - **`message`** — The message body. - **`queue`** (default: `""`) — The target queue name. @@ -88,6 +88,14 @@ When `connection` is provided, the message insert participates in the same datab {!> docs_src/sqlbroker/publish.py [ln:24-30]!} ``` +### Batch publishing + +The broker's and publisher's `.publish_batch()` methods insert all messages in a single SQL statement. They accept the same arguments as [`.publish()`](#publishing){.internal-link}, applied to every message in the batch. Wrap an individual payload in `SqlBrokerPublishMessage` to override its `queue`, `headers`, `correlation_id`, or `next_attempt_at`. + +```python linenums="1" +{!> docs_src/sqlbroker/publish_batch.py [ln:16-42]!} +``` + ## Subscribing ```python linenums="1" diff --git a/docs/docs_src/sqlbroker/publish_batch.py b/docs/docs_src/sqlbroker/publish_batch.py new file mode 100644 index 0000000..2e5b75c --- /dev/null +++ b/docs/docs_src/sqlbroker/publish_batch.py @@ -0,0 +1,42 @@ +from datetime import datetime, timedelta, timezone + +from sqlalchemy.ext.asyncio import create_async_engine + +from faststream import FastStream +from faststream_sqlbroker import SqlBroker, SqlBrokerPublishMessage + +engine = create_async_engine("postgresql+asyncpg://user:pass@localhost/mydb") +broker = SqlBroker(engine=engine) +app = FastStream(broker) + +publisher_sqlbroker = broker.publisher() + +@app.after_startup +async def publish_batch_examples(): + await broker.publish_batch( + "Hello, SqlBroker!", + "Another message", + queue="my_queue", + ) + + await publisher_sqlbroker.publish_batch( + "Hello, SqlBroker!", + "Another message", + queue="my_queue", + ) + + await broker.publish_batch( + SqlBrokerPublishMessage( + "Order placed", + queue="orders", + headers={"x-source": "checkout"}, + correlation_id="order-1", + ), + SqlBrokerPublishMessage( + "Retry later", + next_attempt_at=datetime.now(timezone.utc) + timedelta(minutes=5), + ), + "Uses batch defaults", + queue="my_queue", + headers={"x-default": "batch"}, + ) diff --git a/faststream_sqlbroker/__init__.py b/faststream_sqlbroker/__init__.py index 2d5fee5..27b9270 100644 --- a/faststream_sqlbroker/__init__.py +++ b/faststream_sqlbroker/__init__.py @@ -3,7 +3,9 @@ SqlBrokerCompetingConsumersSchemaVersion, SqlBrokerMessage, SqlBrokerPublishCommand, + SqlBrokerPublishMessage, SqlBrokerPublisher, + SqlBrokerResponse, SqlBrokerRoute, SqlBrokerRouter, SqlBrokerSchemaConfig, @@ -35,7 +37,9 @@ "SqlBrokerCompetingConsumersSchemaVersion", "SqlBrokerMessage", "SqlBrokerPublishCommand", + "SqlBrokerPublishMessage", "SqlBrokerPublisher", + "SqlBrokerResponse", "SqlBrokerRoute", "SqlBrokerRouter", "SqlBrokerSchemaConfig", diff --git a/faststream_sqlbroker/sqlbroker/__init__.py b/faststream_sqlbroker/sqlbroker/__init__.py index 4ead2e5..0e5fe50 100644 --- a/faststream_sqlbroker/sqlbroker/__init__.py +++ b/faststream_sqlbroker/sqlbroker/__init__.py @@ -3,7 +3,11 @@ try: from .annotations import SqlBrokerMessage from .broker import SqlBroker, SqlBrokerPublisher, SqlBrokerRoute, SqlBrokerRouter - from .response import SqlBrokerPublishCommand + from .response import ( + SqlBrokerPublishCommand, + SqlBrokerPublishMessage, + SqlBrokerResponse, + ) from .schema import ( SqlBrokerCompetingConsumersSchemaVersion, SqlBrokerSchemaConfig, @@ -24,7 +28,9 @@ "SqlBrokerCompetingConsumersSchemaVersion", "SqlBrokerMessage", "SqlBrokerPublishCommand", + "SqlBrokerPublishMessage", "SqlBrokerPublisher", + "SqlBrokerResponse", "SqlBrokerRoute", "SqlBrokerRouter", "SqlBrokerSchemaConfig", diff --git a/faststream_sqlbroker/sqlbroker/client.py b/faststream_sqlbroker/sqlbroker/client.py index ab4dec5..c612811 100644 --- a/faststream_sqlbroker/sqlbroker/client.py +++ b/faststream_sqlbroker/sqlbroker/client.py @@ -109,24 +109,23 @@ async def enqueue( async def enqueue_batch( self, - items: Sequence[tuple[bytes, dict[str, str]]], + items: Sequence[tuple[bytes, str, dict[str, str], datetime | None]], *, - queue: str, - next_attempt_at: datetime | None = None, connection: AsyncConnection | None = None, ) -> None: if not items: return - if next_attempt_at: + if any(next_attempt_at is not None for _, _, _, next_attempt_at in items): + default_next_attempt_at = datetime.now(timezone.utc).replace(tzinfo=None) values = [ { "queue": queue, "payload": payload, "headers": headers, - "next_attempt_at": next_attempt_at, + "next_attempt_at": next_attempt_at or default_next_attempt_at, } - for payload, headers in items + for payload, queue, headers, next_attempt_at in items ] else: values = [ @@ -135,7 +134,7 @@ async def enqueue_batch( "payload": payload, "headers": headers, } - for payload, headers in items + for payload, queue, headers, _ in items ] stmt = insert(self._message_table).values(values) diff --git a/faststream_sqlbroker/sqlbroker/publisher/producer.py b/faststream_sqlbroker/sqlbroker/publisher/producer.py index da94dd2..c1cae4f 100644 --- a/faststream_sqlbroker/sqlbroker/publisher/producer.py +++ b/faststream_sqlbroker/sqlbroker/publisher/producer.py @@ -12,6 +12,8 @@ from faststream_sqlbroker.sqlbroker.response import SqlBrokerPublishCommand if TYPE_CHECKING: + from datetime import datetime + from fast_depends.library.serializer import SerializerProto from faststream._internal.types import AsyncCallable, CustomCallable @@ -35,9 +37,8 @@ async def request(self, cmd: "SqlBrokerPublishCommand") -> None: msg = "SqlBroker doesn't support synchronous requests." raise FeatureNotSupportedException(msg) - async def publish_batch(self, cmd: "SqlBrokerPublishCommand") -> None: - msg = "SqlBroker doesn't support publishing in batches." - raise FeatureNotSupportedException(msg) + @abstractmethod + async def publish_batch(self, cmd: "SqlBrokerPublishCommand") -> None: ... class SqlBrokerProducer(SqlBrokerProducerProto): @@ -76,33 +77,34 @@ async def publish(self, cmd: "SqlBrokerPublishCommand") -> None: headers_to_send = { **({"content-type": content_type} if content_type else {}), - **cmd.headers_to_publish(), + **cmd.headers_to_publish_for(0), } await cast("SqlBrokerBaseClient", self.config.client).enqueue( payload=payload, - queue=cmd.destination, + queue=cmd.queue_for(0), headers=headers_to_send, - next_attempt_at=cmd.next_attempt_at, + next_attempt_at=cmd.next_attempt_at_for(0), connection=cmd.connection, ) @override async def publish_batch(self, cmd: "SqlBrokerPublishCommand") -> None: - base_headers = cmd.headers_to_publish() - - items: list[tuple[bytes, dict[str, str]]] = [] - for body in cmd.batch_bodies: + items: list[tuple[bytes, str, dict[str, str], datetime | None]] = [] + for index, body in enumerate(cmd.batch_bodies): payload, content_type = encode_message(body, self.serializer) headers = { **({"content-type": content_type} if content_type else {}), - **base_headers, + **cmd.headers_to_publish_for(index), } - items.append((payload, headers)) + items.append(( + payload, + cmd.queue_for(index), + headers, + cmd.next_attempt_at_for(index), + )) await cast("SqlBrokerBaseClient", self.config.client).enqueue_batch( items, - queue=cmd.destination, - next_attempt_at=cmd.next_attempt_at, connection=cmd.connection, ) diff --git a/faststream_sqlbroker/sqlbroker/publisher/usecase.py b/faststream_sqlbroker/sqlbroker/publisher/usecase.py index 589448c..f4c2bd7 100644 --- a/faststream_sqlbroker/sqlbroker/publisher/usecase.py +++ b/faststream_sqlbroker/sqlbroker/publisher/usecase.py @@ -59,6 +59,32 @@ async def publish( _extra_middlewares=(), ) + async def publish_batch( + self, + *messages: "SendableMessage", + queue: str = "", + headers: dict[str, str] | None = None, + next_attempt_at: datetime | None = None, + connection: AsyncConnection | None = None, + correlation_id: str | None = None, + ) -> None: + if not messages: + return + + cmd = SqlBrokerPublishCommand( + *messages, + queue=queue or self.queue, + headers=self.headers | (headers or {}), + next_attempt_at=next_attempt_at, + connection=connection, + ) + + await self._basic_publish_batch( + cmd, + producer=self._outer_config.producer, + _extra_middlewares=(), + ) + @override async def _publish( self, diff --git a/faststream_sqlbroker/sqlbroker/response.py b/faststream_sqlbroker/sqlbroker/response.py index 5b9c40f..ca05f2e 100644 --- a/faststream_sqlbroker/sqlbroker/response.py +++ b/faststream_sqlbroker/sqlbroker/response.py @@ -2,7 +2,7 @@ from typing import TYPE_CHECKING, Union from faststream.response.publish_type import PublishType -from faststream.response.response import BatchPublishCommand, PublishCommand +from faststream.response.response import BatchPublishCommand, PublishCommand, Response from sqlalchemy.ext.asyncio import AsyncConnection from faststream_sqlbroker.sqlbroker.exceptions import DatetimeMissingTimezoneException @@ -11,6 +11,50 @@ from faststream._internal.basic_types import SendableMessage +def _normalize_next_attempt_at(value: datetime | None) -> datetime | None: + if value is None: + return None + + _validate_next_attempt_at(value) + return value.astimezone(timezone.utc).replace(tzinfo=None) + + +def _validate_next_attempt_at(value: datetime | None) -> None: + if value is not None and value.tzinfo is None: + raise DatetimeMissingTimezoneException + + +class SqlBrokerResponse(Response): + """An outgoing SQL broker message with optional per-message arguments.""" + + def __init__( + self, + body: "SendableMessage", + *, + queue: str | None = None, + headers: dict[str, str] | None = None, + correlation_id: str | None = None, + next_attempt_at: datetime | None = None, + ) -> None: + super().__init__( + body=body, + headers=headers, + correlation_id=correlation_id, + ) + self.queue = queue + _validate_next_attempt_at(next_attempt_at) + self.next_attempt_at = next_attempt_at + + def as_publish_command(self) -> "SqlBrokerPublishCommand": + return SqlBrokerPublishCommand( + self.body, + queue=self.queue or "", + headers=self.headers, + correlation_id=self.correlation_id, + next_attempt_at=self.next_attempt_at, + ) + + class SqlBrokerPublishCommand(BatchPublishCommand): def __init__( self, @@ -23,9 +67,6 @@ def __init__( next_attempt_at: datetime | None = None, connection: AsyncConnection | None = None, ) -> None: - if next_attempt_at and next_attempt_at.tzinfo is None: - raise DatetimeMissingTimezoneException - super().__init__( message, *messages, @@ -34,10 +75,19 @@ def __init__( correlation_id=correlation_id, _publish_type=PublishType.PUBLISH, ) - self.next_attempt_at = next_attempt_at - self._convert_timezone_to_utc() + self.next_attempt_at = _normalize_next_attempt_at(next_attempt_at) self.connection = connection + self._per_message_args = tuple( + body if isinstance(body, SqlBrokerResponse) else None + for body in self.batch_bodies + ) + if any(self._per_message_args): + self.batch_bodies = tuple( + body.body if isinstance(body, SqlBrokerResponse) else body + for body in self.batch_bodies + ) + @classmethod def from_cmd( cls, @@ -63,8 +113,34 @@ def headers_to_publish(self) -> dict[str, str]: return headers | (self.headers or {}) - def _convert_timezone_to_utc(self) -> None: - if self.next_attempt_at: - self.next_attempt_at = self.next_attempt_at.astimezone(timezone.utc).replace( - tzinfo=None - ) + def queue_for(self, index: int) -> str: + args = self._args_for(index) + if args is not None and args.queue is not None: + return args.queue + return self.destination + + def headers_to_publish_for(self, index: int) -> dict[str, str]: + headers = self.headers_to_publish() + args = self._args_for(index) + if args is None: + return headers + + if args.correlation_id: + headers["correlation_id"] = args.correlation_id + + return headers | args.headers + + def next_attempt_at_for(self, index: int) -> datetime | None: + args = self._args_for(index) + if args is not None and args.next_attempt_at is not None: + return _normalize_next_attempt_at(args.next_attempt_at) + return self.next_attempt_at + + def _args_for(self, index: int) -> SqlBrokerResponse | None: + try: + return self._per_message_args[index] + except IndexError: + return None + + +SqlBrokerPublishMessage = SqlBrokerResponse diff --git a/pyproject.toml b/pyproject.toml index 266b33b..35f1f85 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ authors = [ { name = "Arseniy Popov", email = "arseniypopov@gmail.com" }, ] requires-python = ">=3.10" -version = "0.1.0a8" +version = "0.1.0a9" dependencies = [ "faststream>=0.7.0rc1", diff --git a/tests/docs/sqlbroker/test_publish_batch.py b/tests/docs/sqlbroker/test_publish_batch.py new file mode 100644 index 0000000..24f6147 --- /dev/null +++ b/tests/docs/sqlbroker/test_publish_batch.py @@ -0,0 +1,72 @@ +from __future__ import annotations + +import importlib +import json +import sys +from typing import TYPE_CHECKING + +import pytest +import sqlalchemy.ext.asyncio as sa_asyncio +from faststream import TestApp +from sqlalchemy import text + +if TYPE_CHECKING: + from collections.abc import Iterator + + from sqlalchemy.ext.asyncio import AsyncEngine + +MODULE = "docs.docs_src.sqlbroker.publish_batch" + + +@pytest.fixture() +def publish_batch_module( + engine: AsyncEngine, + recreate_tables: None, + monkeypatch: pytest.MonkeyPatch, +) -> Iterator[object]: + # The docs example builds its engine from a hard-coded Postgres URL at + # import time. Swap in the real test engine (parametrised over + # postgres/mysql/sqlite) so the documented code runs against a live DB. + monkeypatch.setattr(sa_asyncio, "create_async_engine", lambda *a, **k: engine) + sys.modules.pop(MODULE, None) + try: + yield importlib.import_module(MODULE) + finally: + sys.modules.pop(MODULE, None) + + +@pytest.mark.connected() +@pytest.mark.slow() +@pytest.mark.asyncio() +async def test_publish_batch(publish_batch_module: object, engine: AsyncEngine) -> None: + # `publish_batch_examples` runs in `@app.after_startup`, so starting the + # app exercises the documented `publish_batch(...)` calls. + async with TestApp(publish_batch_module.app): + pass + + async with engine.connect() as conn: + result = await conn.execute( + text("SELECT queue, headers, next_attempt_at FROM message ORDER BY id") + ) + rows = result.all() + + assert len(rows) == 7 + assert [row.queue for row in rows] == [ + "my_queue", + "my_queue", + "my_queue", + "my_queue", + "orders", + "my_queue", + "my_queue", + ] + + def as_headers(value: object) -> dict[str, str]: + if isinstance(value, dict): + return value + return json.loads(value) # type: ignore[arg-type] + + assert as_headers(rows[4].headers)["correlation_id"] == "order-1" + assert as_headers(rows[4].headers)["x-source"] == "checkout" + assert rows[5].next_attempt_at is not None + assert as_headers(rows[6].headers)["x-default"] == "batch" diff --git a/tests/test_publish.py b/tests/test_publish.py index 12bf127..a01fe68 100644 --- a/tests/test_publish.py +++ b/tests/test_publish.py @@ -5,7 +5,7 @@ from sqlalchemy import event, text from sqlalchemy.ext.asyncio import AsyncEngine -from faststream_sqlbroker.sqlbroker import SqlBroker +from faststream_sqlbroker.sqlbroker import SqlBroker, SqlBrokerPublishMessage from faststream_sqlbroker.sqlbroker.exceptions import DatetimeMissingTimezoneException from tests.brokers.base.publish import BrokerPublishTestcase from tests.helpers import as_datetime @@ -96,6 +96,62 @@ async def test_publish_with_next_attempt_at_converts_timezone_to_utc( second=0, ) + @pytest.mark.asyncio() + @pytest.mark.parametrize("mode", ("publish", "publisher")) + async def test_publish_uses_per_message_args( + self, engine: AsyncEngine, mode: str, broker: SqlBroker + ) -> None: + next_attempt_at = datetime.datetime( + year=2026, + month=1, + day=2, + hour=12, + minute=0, + second=0, + tzinfo=datetime.timezone(datetime.timedelta(hours=3), "MSC"), + ) + message = SqlBrokerPublishMessage( + {"message": "hello1"}, + queue="override-queue", + headers={"x-shared": "override"}, + correlation_id="override-correlation", + next_attempt_at=next_attempt_at, + ) + + match mode: + case "publish": + await broker.publish( + message, + queue="default-queue", + headers={"x-default": "value", "x-shared": "default"}, + ) + case "publisher": + publisher = broker.publisher("default-queue") + await publisher.publish( + message, + headers={"x-default": "value", "x-shared": "default"}, + ) + + async with engine.connect() as conn: + result = await conn.execute( + text("SELECT queue, headers, next_attempt_at FROM message") + ) + row = result.one() + + assert row.queue == "override-queue" + headers = ( + row.headers if isinstance(row.headers, dict) else json.loads(row.headers) + ) + assert headers == { + "content-type": "application/json", + "correlation_id": "override-correlation", + "x-default": "value", + "x-shared": "override", + } + assert as_datetime(row.next_attempt_at) == datetime.datetime( # noqa: DTZ001 + year=2026, month=1, day=2, hour=9, minute=0, second=0 + ) + @pytest.mark.connected() class TestPublishTransaction(SqlBrokerTestcaseConfig): @@ -158,18 +214,35 @@ async def test_publish_in_transaction_rollback( assert len(result.all()) == 0 +async def _do_publish_batch( + broker: SqlBroker, + mode: str, + *messages: object, + queue: str = "batch-queue", + **kwargs: object, +) -> None: + match mode: + case "broker": + await broker.publish_batch(*messages, queue=queue, **kwargs) + case "publisher": + publisher = broker.publisher(queue) + await publisher.publish_batch(*messages, **kwargs) + + @pytest.mark.connected() @pytest.mark.slow() +@pytest.mark.parametrize("mode", ("broker", "publisher")) class TestPublishBatch(SqlBrokerTestcaseConfig): @pytest.mark.asyncio() async def test_publish_batch_inserts_all_messages( - self, engine: AsyncEngine, broker: SqlBroker + self, engine: AsyncEngine, broker: SqlBroker, mode: str ) -> None: - await broker.publish_batch( + await _do_publish_batch( + broker, + mode, {"message": "hello1"}, {"message": "hello2"}, {"message": "hello3"}, - queue="batch-queue", ) async with engine.connect() as conn: @@ -189,7 +262,7 @@ async def test_publish_batch_inserts_all_messages( @pytest.mark.asyncio() async def test_publish_batch_uses_single_sql_statement( - self, engine: AsyncEngine, broker: SqlBroker + self, engine: AsyncEngine, broker: SqlBroker, mode: str ) -> None: inserts: list[str] = [] @@ -199,13 +272,7 @@ def _capture(conn, cursor, statement, parameters, context, executemany) -> None: event.listen(engine.sync_engine, "before_cursor_execute", _capture) try: - await broker.publish_batch( - b"a", - b"b", - b"c", - b"d", - queue="batch-queue", - ) + await _do_publish_batch(broker, mode, b"a", b"b", b"c", b"d") finally: event.remove(engine.sync_engine, "before_cursor_execute", _capture) @@ -217,9 +284,9 @@ def _capture(conn, cursor, statement, parameters, context, executemany) -> None: @pytest.mark.asyncio() async def test_publish_batch_empty_is_noop( - self, engine: AsyncEngine, broker: SqlBroker + self, engine: AsyncEngine, broker: SqlBroker, mode: str ) -> None: - await broker.publish_batch(queue="batch-queue") + await _do_publish_batch(broker, mode) async with engine.connect() as conn: result = await conn.execute(text("SELECT COUNT(*) FROM message")) @@ -227,12 +294,13 @@ async def test_publish_batch_empty_is_noop( @pytest.mark.asyncio() async def test_publish_batch_stores_headers_per_message( - self, engine: AsyncEngine, broker: SqlBroker + self, engine: AsyncEngine, broker: SqlBroker, mode: str ) -> None: - await broker.publish_batch( + await _do_publish_batch( + broker, + mode, {"message": "hello1"}, b"raw", - queue="batch-queue", headers={"x-custom": "value"}, ) @@ -250,9 +318,109 @@ async def test_publish_batch_stores_headers_per_message( } assert headers[1] == {"x-custom": "value"} + @pytest.mark.asyncio() + async def test_publish_batch_uses_per_message_args( + self, engine: AsyncEngine, broker: SqlBroker, mode: str + ) -> None: + default_next_attempt_at = datetime.datetime( + year=2026, + month=1, + day=1, + hour=12, + minute=0, + second=0, + tzinfo=datetime.timezone(datetime.timedelta(hours=3), "MSC"), + ) + first_next_attempt_at = datetime.datetime( + year=2026, + month=1, + day=2, + hour=12, + minute=0, + second=0, + tzinfo=datetime.timezone(datetime.timedelta(hours=3), "MSC"), + ) + second_next_attempt_at = datetime.datetime( + year=2026, + month=1, + day=3, + hour=12, + minute=0, + second=0, + tzinfo=datetime.timezone(datetime.timedelta(hours=3), "MSC"), + ) + + await _do_publish_batch( + broker, + mode, + SqlBrokerPublishMessage( + {"message": "hello1"}, + queue="first-queue", + headers={"x-shared": "first"}, + correlation_id="first-correlation", + next_attempt_at=first_next_attempt_at, + ), + SqlBrokerPublishMessage( + {"message": "hello2"}, + queue="second-queue", + headers={"x-shared": "second"}, + correlation_id="second-correlation", + next_attempt_at=second_next_attempt_at, + ), + {"message": "hello3"}, + headers={"x-default": "value", "x-shared": "default"}, + next_attempt_at=default_next_attempt_at, + ) + + async with engine.connect() as conn: + result = await conn.execute( + text("SELECT queue, headers, next_attempt_at FROM message ORDER BY id") + ) + rows = result.all() + + assert [row.queue for row in rows] == [ + "first-queue", + "second-queue", + "batch-queue", + ] + headers = [ + row.headers if isinstance(row.headers, dict) else json.loads(row.headers) + for row in rows + ] + assert headers == [ + { + "content-type": "application/json", + "correlation_id": "first-correlation", + "x-default": "value", + "x-shared": "first", + }, + { + "content-type": "application/json", + "correlation_id": "second-correlation", + "x-default": "value", + "x-shared": "second", + }, + { + "content-type": "application/json", + "x-default": "value", + "x-shared": "default", + }, + ] + assert [as_datetime(row.next_attempt_at) for row in rows] == [ + datetime.datetime( # noqa: DTZ001 + year=2026, month=1, day=2, hour=9, minute=0, second=0 + ), + datetime.datetime( # noqa: DTZ001 + year=2026, month=1, day=3, hour=9, minute=0, second=0 + ), + datetime.datetime( # noqa: DTZ001 + year=2026, month=1, day=1, hour=9, minute=0, second=0 + ), + ] + @pytest.mark.asyncio() async def test_publish_batch_with_next_attempt_at( - self, engine: AsyncEngine, broker: SqlBroker + self, engine: AsyncEngine, broker: SqlBroker, mode: str ) -> None: next_attempt_at = datetime.datetime( year=2026, @@ -264,10 +432,11 @@ async def test_publish_batch_with_next_attempt_at( tzinfo=datetime.timezone(datetime.timedelta(hours=3), "MSC"), ) - await broker.publish_batch( + await _do_publish_batch( + broker, + mode, {"message": "hello1"}, {"message": "hello2"}, - queue="batch-queue", next_attempt_at=next_attempt_at, ) @@ -284,29 +453,28 @@ async def test_publish_batch_with_next_attempt_at( assert as_datetime(row.next_attempt_at) == expected @pytest.mark.asyncio() - async def test_publish_batch_without_timezone_raises(self, broker: SqlBroker) -> None: + async def test_publish_batch_without_timezone_raises( + self, broker: SqlBroker, mode: str + ) -> None: with pytest.raises(DatetimeMissingTimezoneException): - await broker.publish_batch( + await _do_publish_batch( + broker, + mode, {"message": "hello1"}, {"message": "hello2"}, - queue="batch-queue", next_attempt_at=datetime.datetime.now(), # noqa: DTZ005 ) @pytest.mark.connected() +@pytest.mark.parametrize("mode", ("broker", "publisher")) class TestPublishBatchTransaction(SqlBrokerTestcaseConfig): @pytest.mark.asyncio() async def test_publish_batch_in_transaction( - self, engine: AsyncEngine, broker: SqlBroker + self, engine: AsyncEngine, broker: SqlBroker, mode: str ) -> None: async with engine.begin() as conn: - await broker.publish_batch( - b"a", - b"b", - queue="batch-queue", - connection=conn, - ) + await _do_publish_batch(broker, mode, b"a", b"b", connection=conn) async with engine.connect() as conn: result = await conn.execute(text("SELECT COUNT(*) FROM message")) @@ -314,15 +482,10 @@ async def test_publish_batch_in_transaction( @pytest.mark.asyncio() async def test_publish_batch_in_transaction_rollback( - self, engine: AsyncEngine, broker: SqlBroker + self, engine: AsyncEngine, broker: SqlBroker, mode: str ) -> None: async with engine.begin() as conn: - await broker.publish_batch( - b"a", - b"b", - queue="batch-queue", - connection=conn, - ) + await _do_publish_batch(broker, mode, b"a", b"b", connection=conn) await conn.rollback() async with engine.connect() as conn: diff --git a/uv.lock b/uv.lock index 9d9c25e..ba7f2c0 100644 --- a/uv.lock +++ b/uv.lock @@ -717,7 +717,7 @@ wheels = [ [[package]] name = "faststream-sqlbroker" -version = "0.1.0a8" +version = "0.1.0a9" source = { editable = "." } dependencies = [ { name = "faststream" }, From e7e2ec5bb5a355450389b069c9065b1d60ac62bd Mon Sep 17 00:00:00 2001 From: Arseniy Popov Date: Sun, 19 Jul 2026 19:33:18 +0300 Subject: [PATCH 2/2] feat: wire up batch publishing --- faststream_sqlbroker/sqlbroker/response.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/faststream_sqlbroker/sqlbroker/response.py b/faststream_sqlbroker/sqlbroker/response.py index ca05f2e..bbeb56f 100644 --- a/faststream_sqlbroker/sqlbroker/response.py +++ b/faststream_sqlbroker/sqlbroker/response.py @@ -56,6 +56,17 @@ def as_publish_command(self) -> "SqlBrokerPublishCommand": class SqlBrokerPublishCommand(BatchPublishCommand): + """TODO: per-message args aren't very native in FastStream. + + Per-message overrides (queue/headers/correlation_id/next_attempt_at) are + carried on ``SqlBrokerResponse`` bodies and looked up positionally against + ``batch_bodies`` at publish time. This is fragile: anything that reorders or + filters ``batch_bodies`` after construction (e.g. a publish middleware) + desyncs the overrides from their bodies. This mirrors FastStream's own Kafka + per-message keys and is awaiting a more comprehensive upstream fix — see + https://github.com/ag2ai/faststream/issues/2943. + """ + def __init__( self, message: "SendableMessage",