1717import abc
1818import asyncio
1919import datetime as dt
20+ import typing
2021import uuid
2122from typing import TYPE_CHECKING
23+ from weakref import WeakKeyDictionary
2224
2325from sqlalchemy import (
2426 ARRAY ,
3032 delete ,
3133 func ,
3234 insert ,
35+ literal ,
3336 select ,
3437 text ,
3538 true ,
3639 union_all ,
3740 update ,
3841)
42+ from sqlalchemy .dialects .postgresql import insert as pg_insert
3943from sqlalchemy .ext .asyncio import AsyncSession
4044
4145from faststream_outbox import schema_validation
46+ from faststream_outbox ._scheduling import is_future_dated , resolve_next_attempt_client_side
47+ from faststream_outbox ._time import utcnow
4248from faststream_outbox .message import OutboxInnerMessage
4349from faststream_outbox .schema import (
4450 _DLQ_INJECTED_COLUMNS ,
4854
4955
5056if TYPE_CHECKING :
51- import typing
5257 from collections .abc import Mapping , Sequence
5358
5459 from sqlalchemy import CTE , ColumnElement , Insert , Table , Update
@@ -82,6 +87,41 @@ def engine(self) -> "AsyncEngine | None": ...
8287 @abc .abstractmethod
8388 def table (self ) -> "Table" : ...
8489
90+ @abc .abstractmethod
91+ async def insert (
92+ self ,
93+ session : "AsyncSession" ,
94+ * ,
95+ queue : str ,
96+ payload : bytes ,
97+ headers : dict [str , str ] | None ,
98+ activate_in : dt .timedelta | None = None ,
99+ activate_at : dt .datetime | None = None ,
100+ timer_id : str | None = None ,
101+ ) -> int | None :
102+ """Insert one encoded row on *session* and NOTIFY its queue if the row is eligible now.
103+
104+ Returns the row id, or ``None`` when a live row already holds *timer_id* on *queue*.
105+ Application code should call ``broker.publish`` instead: this skips the publish
106+ middlewares and the ``published`` metric.
107+ """
108+
109+ @abc .abstractmethod
110+ async def insert_batch (
111+ self ,
112+ session : "AsyncSession" ,
113+ * ,
114+ queue : str ,
115+ rows : "Sequence[tuple[bytes, dict[str, str] | None]]" ,
116+ activate_in : dt .timedelta | None = None ,
117+ activate_at : dt .datetime | None = None ,
118+ ) -> None :
119+ """Insert every ``(payload, headers)`` row on *session* in one statement, or none of them.
120+
121+ Application code should call ``broker.publish_batch`` instead, for the same reason as
122+ :meth:`insert`.
123+ """
124+
85125 @abc .abstractmethod
86126 async def fetch (
87127 self ,
@@ -205,6 +245,137 @@ def _build_dlq_cte(outbox: "Table", dlq: "Table") -> "Insert":
205245 return insert (dlq ).from_select ([* (col for _ , col in _DLQ_PROJECTION ), * _DLQ_INJECTED_COLUMNS ], archived )
206246
207247
248+ def _build_insert_stmts (table : "Table" , channel : str ) -> "dict[tuple[bool, bool, bool], Insert]" :
249+ """Build the single-row INSERT once per ``(timer_id set, activate_in set, notify)``; values arrive as parameters.
250+
251+ The notify variants call ``pg_notify`` in ``RETURNING``, so the NOTIFY costs no extra round trip and
252+ fires only when a row actually landed (an ``ON CONFLICT DO NOTHING`` skip returns no row).
253+ """
254+ # Server-side compute keeps timing immune to worker/DB clock skew (mirrors
255+ # client.mark_pending_with_lease).
256+ activate_in = func .now () + func .make_interval (0 , 0 , 0 , 0 , 0 , 0 , bindparam ("activate_in_seconds" , type_ = Float ))
257+ plain = insert (table )
258+ timer = pg_insert (table ).on_conflict_do_nothing (
259+ index_elements = [table .c .queue , table .c .timer_id ],
260+ index_where = table .c .timer_id .is_not (None ),
261+ )
262+ notify = func .pg_notify (literal (channel ), table .c .queue )
263+ stmts : dict [tuple [bool , bool , bool ], Insert ] = {}
264+ for has_timer , base in ((False , plain ), (True , timer )):
265+ for has_activate_in , stmt in ((False , base ), (True , base .values (next_attempt_at = activate_in ))):
266+ stmts [has_timer , has_activate_in , False ] = stmt .returning (table .c .id )
267+ stmts [has_timer , has_activate_in , True ] = stmt .returning (table .c .id , notify )
268+ return stmts
269+
270+
271+ class _OutboxInserter :
272+ """The outbox INSERT + NOTIFY on a caller's session; needs no engine, so an engine-less broker can publish."""
273+
274+ def __init__ (self , table : "Table" ) -> None :
275+ self ._table = table
276+ self ._channel = f"outbox_{ table .name } "
277+ # Per-transaction NOTIFY dedup memo (see _notify). Keyed on the innermost active
278+ # transaction object, weak-referenced so entries GC when the transaction ends -- no
279+ # explicit reset, and a rolled-back savepoint cannot suppress a later real NOTIFY.
280+ self ._notified : WeakKeyDictionary [typing .Any , set [str ]] = WeakKeyDictionary ()
281+ self ._insert_stmts = _build_insert_stmts (table , self ._channel )
282+
283+ async def insert (
284+ self ,
285+ session : "AsyncSession" ,
286+ * ,
287+ queue : str ,
288+ payload : bytes ,
289+ headers : dict [str , str ] | None ,
290+ activate_in : dt .timedelta | None = None ,
291+ activate_at : dt .datetime | None = None ,
292+ timer_id : str | None = None ,
293+ ) -> int | None :
294+ params : dict [str , typing .Any ] = {"queue" : queue , "payload" : payload , "headers" : headers }
295+ if activate_in is not None :
296+ params ["activate_in_seconds" ] = activate_in .total_seconds ()
297+ elif activate_at is not None :
298+ params ["next_attempt_at" ] = activate_at
299+ if timer_id is not None :
300+ params ["timer_id" ] = timer_id
301+ # Skip NOTIFY only when the row is genuinely future-dated. A past activate_at
302+ # (e.g. a recovered idempotency token) is immediately eligible — fire NOTIFY.
303+ # Future-dated rows skip NOTIFY: listeners can't act before the gate opens, and
304+ # polling fires them at the next tick.
305+ emitted = self ._notified_queues (session )
306+ notify = not is_future_dated (activate_in , activate_at , utcnow ()) and (emitted is None or queue not in emitted )
307+ stmt = self ._insert_stmts [timer_id is not None , activate_in is not None , notify ]
308+ result = await session .execute (stmt , params )
309+ row_id : int | None = result .scalar ()
310+ if row_id is not None and notify :
311+ # Re-read: under autobegin the transaction only exists once the INSERT has run.
312+ emitted = self ._notified_queues (session )
313+ if emitted is not None :
314+ emitted .add (queue )
315+ return row_id
316+
317+ async def insert_batch (
318+ self ,
319+ session : "AsyncSession" ,
320+ * ,
321+ queue : str ,
322+ rows : "Sequence[tuple[bytes, dict[str, str] | None]]" ,
323+ activate_in : dt .timedelta | None = None ,
324+ activate_at : dt .datetime | None = None ,
325+ ) -> None :
326+ # Client-side time for batch: executemany doesn't compose cleanly with
327+ # column-level SQL expressions, and a few-ms drift versus the DB is
328+ # harmless for user-supplied scheduling. Retries still use server time.
329+ now = utcnow ()
330+ next_at = resolve_next_attempt_client_side (activate_in , activate_at , now )
331+ values : list [dict [str , typing .Any ]] = []
332+ for payload , headers in rows :
333+ row : dict [str , typing .Any ] = {"queue" : queue , "payload" : payload , "headers" : headers }
334+ if next_at is not None :
335+ row ["next_attempt_at" ] = next_at
336+ values .append (row )
337+ await session .execute (insert (self ._table ), values )
338+ # Skip NOTIFY only when genuinely future-dated; past times are eligible.
339+ if not is_future_dated (activate_in , activate_at , now ):
340+ await self ._notify (session , queue )
341+
342+ def _notified_queues (self , session : "AsyncSession" ) -> set [str ] | None :
343+ """Return the queues already notified in *session*'s innermost transaction, or None outside one."""
344+ sync_session = session .sync_session
345+ txn = sync_session .get_nested_transaction () or sync_session .get_transaction ()
346+ if txn is None :
347+ return None
348+ return self ._notified .setdefault (txn , set ())
349+
350+ async def _notify (self , session : "AsyncSession" , queue : str ) -> None :
351+ # ``pg_notify(:channel, :payload)`` -- parameterized so channel and payload
352+ # bind cleanly (raw NOTIFY accepts only literals -- injection-prone). Runs on
353+ # the caller's session so NOTIFY commits with the row insert; rollback silently
354+ # discards it. Non-Postgres dialects ignore it.
355+ #
356+ # Dedup: Postgres already coalesces identical (channel, payload) NOTIFYs per
357+ # transaction at delivery, so a second pg_notify for the same queue in the same
358+ # transaction only wastes a round-trip. Memoize emitted queues per innermost
359+ # active transaction (get_nested_transaction() during a savepoint, else the outer
360+ # transaction) -- each a distinct, weak-referenceable object that GCs when the
361+ # transaction ends, so a rolled-back savepoint cannot suppress a later real NOTIFY.
362+ # Key on the *sync* SessionTransaction, not the async AsyncSessionTransaction proxy:
363+ # SQLAlchemy regenerates that proxy on every call and only weak-references it, so
364+ # under autobegin (no explicit `session.begin()`) it would GC between publishes and
365+ # defeat the memo. The Session holds the sync transaction strongly for its lifetime,
366+ # so the key survives across publishes whether the caller uses `session.begin()` or
367+ # autobegin. No active transaction -> emit unconditionally.
368+ emitted = self ._notified_queues (session )
369+ if emitted is not None :
370+ if queue in emitted :
371+ return
372+ emitted .add (queue )
373+ await session .execute (
374+ text ("SELECT pg_notify(:channel, :payload)" ),
375+ {"channel" : self ._channel , "payload" : queue },
376+ )
377+
378+
208379class OutboxClient (AbstractOutboxClient ):
209380 def __init__ (
210381 self ,
@@ -221,6 +392,7 @@ def __init__(
221392 self ._table = outbox_table
222393 self ._dlq_table = dlq_table
223394 self ._fetch_stmt = _build_fetch_stmt (outbox_table )
395+ self ._inserter = _OutboxInserter (outbox_table )
224396 self ._dlq_cte = None if dlq_table is None else _build_dlq_cte (outbox_table , dlq_table )
225397
226398 @property
@@ -236,6 +408,40 @@ def engine(self) -> "AsyncEngine":
236408 """
237409 return self ._engine
238410
411+ async def insert (
412+ self ,
413+ session : "AsyncSession" ,
414+ * ,
415+ queue : str ,
416+ payload : bytes ,
417+ headers : dict [str , str ] | None ,
418+ activate_in : dt .timedelta | None = None ,
419+ activate_at : dt .datetime | None = None ,
420+ timer_id : str | None = None ,
421+ ) -> int | None :
422+ return await self ._inserter .insert (
423+ session ,
424+ queue = queue ,
425+ payload = payload ,
426+ headers = headers ,
427+ activate_in = activate_in ,
428+ activate_at = activate_at ,
429+ timer_id = timer_id ,
430+ )
431+
432+ async def insert_batch (
433+ self ,
434+ session : "AsyncSession" ,
435+ * ,
436+ queue : str ,
437+ rows : "Sequence[tuple[bytes, dict[str, str] | None]]" ,
438+ activate_in : dt .timedelta | None = None ,
439+ activate_at : dt .datetime | None = None ,
440+ ) -> None :
441+ await self ._inserter .insert_batch (
442+ session , queue = queue , rows = rows , activate_in = activate_in , activate_at = activate_at
443+ )
444+
239445 async def fetch (
240446 self ,
241447 conn : "AsyncConnection | None" ,
0 commit comments