Skip to content

Commit e648424

Browse files
authored
refactor: move the outbox insert and its NOTIFY into the client (#237)
* refactor: move the outbox insert and its NOTIFY into the client AbstractOutboxClient gains insert / insert_batch. OutboxClient takes the producer's prebuilt INSERT statements and per-transaction NOTIFY memo unchanged; FakeOutboxClient stores rows with the same validation and timer_id dedup. The producer now encodes, inserts through the broker's client and emits published. A broker without an engine keeps publishing through an engine-less inserter. Part of #236. * test: keep only the insert rejection every SQLAlchemy version shares SQLAlchemy 2.0.0 binds String(255) values as $1::VARCHAR(255), and Postgres truncates an explicit cast instead of raising, so the real client does not reject over-long queue or timer_id there.
1 parent 356af94 commit e648424

8 files changed

Lines changed: 481 additions & 134 deletions

File tree

‎AGENTS.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ will **not** tell you:
3131
that — `tests/test_client_contract.py::test_delete_noop_on_token_mismatch` is the claim.
3232
- `client.py` and `testing.py` implement the same rules twice, in SQL and in Python, because one
3333
runs in the database and one in the process. They cannot share an implementation, so
34-
`tests/test_client_contract.py` couples them by behaviour — a change to either adapter's fetch,
34+
`tests/test_client_contract.py` couples them by behaviour — a change to either adapter's insert, fetch,
3535
terminal, or DLQ semantics belongs in that suite.
3636
- `schema.py` — the three partial indexes and the `<table>_lease_ck` CHECK are load-bearing, not
3737
decoration.

‎docs/usage/timers.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,7 @@ for user-supplied scheduling.
5959
## Deduplication with `timer_id`
6060

6161
`timer_id` flows into a `String(255)` column with a partial unique index
62-
on `(queue, timer_id) WHERE timer_id IS NOT NULL`. The producer switches
62+
on `(queue, timer_id) WHERE timer_id IS NOT NULL`. The insert switches
6363
to `pg_insert(...).on_conflict_do_nothing(...)` so re-publishing the same
6464
id is a silent no-op (returns `None`):
6565

‎faststream_outbox/broker.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -187,6 +187,7 @@ def __init__( # noqa: PLR0913
187187
# Serializer lives on fd_config — wire it onto the producer so encoded
188188
# bodies use the same path as the broker's own publish flow.
189189
producer.serializer = fd_config._serializer # noqa: SLF001
190+
producer.broker_config = broker_config
190191
specification = BrokerSpec(
191192
url=_spec_url(engine, outbox_table),
192193
protocol="postgresql",

‎faststream_outbox/client.py‎

Lines changed: 207 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,10 @@
1717
import abc
1818
import asyncio
1919
import datetime as dt
20+
import typing
2021
import uuid
2122
from typing import TYPE_CHECKING
23+
from weakref import WeakKeyDictionary
2224

2325
from sqlalchemy import (
2426
ARRAY,
@@ -30,15 +32,19 @@
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
3943
from sqlalchemy.ext.asyncio import AsyncSession
4044

4145
from 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
4248
from faststream_outbox.message import OutboxInnerMessage
4349
from faststream_outbox.schema import (
4450
_DLQ_INJECTED_COLUMNS,
@@ -48,7 +54,6 @@
4854

4955

5056
if 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+
208379
class 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

Comments
 (0)