feat: wire up batch publishing

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** &mdash; 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) &mdash; 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) &mdash; 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** &mdash; 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.
 

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"

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"},
+    )

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",

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",

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)

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,
         )

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,