1313import pytest
1414from faststream import BaseMiddleware , Context , Depends
1515from faststream .kafka import KafkaBroker , TestKafkaBroker
16- from sqlalchemy import MetaData , Table , event , func , insert , literal , select , text
16+ from sqlalchemy import MetaData , Table , case , event , func , insert , literal , select , text
1717from sqlalchemy .exc import IntegrityError
1818from sqlalchemy .ext .asyncio import AsyncEngine , AsyncSession , async_sessionmaker , create_async_engine
1919
@@ -424,13 +424,43 @@ async def test_validate_schema_fails_when_pending_index_missing(pg_engine, outbo
424424 await client .validate_schema ()
425425
426426
427- async def test_validate_schema_fails_when_lease_index_missing (pg_engine , outbox_table ) -> None :
428- """Missing the expired-lease partial index sends the fetch CTE back to seq-scan; flag it."""
429- drop_sql = f'DROP INDEX "{ outbox_table .name } _lease_idx"'
427+ async def test_validate_schema_reports_the_pre_0_17_lease_index (
428+ pg_engine : AsyncEngine ,
429+ outbox_table : Table ,
430+ ) -> None :
431+ """An outbox migrated before 0.17 has ``_lease_idx`` on ``(queue, acquired_at)`` and lacks ``_leased_idx``."""
430432 async with pg_engine .begin () as conn :
431- await conn .exec_driver_sql (drop_sql )
433+ await conn .execute (text (f'DROP INDEX "{ outbox_table .name } _leased_idx"' ))
434+ await conn .execute (
435+ text (
436+ f'CREATE INDEX "{ outbox_table .name } _lease_idx" ON { outbox_table .name } (queue, acquired_at) '
437+ "WHERE acquired_token IS NOT NULL" ,
438+ ),
439+ )
432440 client = OutboxClient (pg_engine , outbox_table )
433- with pytest .raises (RuntimeError , match = "missing index" ):
441+ with pytest .raises (RuntimeError , match = r"missing index.*leased_idx.*operations/alembic/#upgrading-to-017" ):
442+ await client .validate_schema ()
443+
444+
445+ @pytest .mark .parametrize (
446+ ("predicate" , "error" ),
447+ [("" , "is not a partial index" ), ("WHERE acquired_token IS NULL" , "has wrong partial predicate" )],
448+ )
449+ async def test_validate_schema_checks_the_leased_index_predicate (
450+ pg_engine : AsyncEngine ,
451+ outbox_table : Table ,
452+ predicate : str ,
453+ error : str ,
454+ ) -> None :
455+ """A ``_leased_idx`` rebuilt without ``acquired_token IS NOT NULL`` is reported (alembic's diff ignores it)."""
456+ idx = f"{ outbox_table .name } _leased_idx"
457+ async with pg_engine .begin () as conn :
458+ await conn .execute (text (f'DROP INDEX "{ idx } "' ))
459+ await conn .execute (
460+ text (f'CREATE INDEX "{ idx } " ON { outbox_table .name } (queue, next_attempt_at, id) { predicate } ' )
461+ )
462+ client = OutboxClient (pg_engine , outbox_table )
463+ with pytest .raises (RuntimeError , match = f"'{ idx } ' { error } " ):
434464 await client .validate_schema ()
435465
436466
@@ -616,6 +646,35 @@ def _scanned_rows(plan: dict[str, Any], table: str) -> list[float]:
616646 return found
617647
618648
649+ async def _explain_fetch (engine : AsyncEngine , table : Table , * , limit : int ) -> dict [str , Any ]:
650+ """Run ``EXPLAIN ANALYZE`` on the statement a real fetch of ``orders`` sends, rolled back."""
651+ captured : list [tuple [str , Any ]] = []
652+
653+ def capture (_conn : Any , _cursor : Any , statement : str , parameters : Any , _context : Any , _many : bool ) -> None :
654+ if "FOR UPDATE SKIP LOCKED" in statement :
655+ captured .append ((statement , parameters ))
656+
657+ client = OutboxClient (engine , table )
658+ event .listen (engine .sync_engine , "before_cursor_execute" , capture )
659+ try :
660+ async with engine .connect () as conn :
661+ await client .fetch (conn , ["orders" ], limit = limit , lease_ttl_seconds = 60.0 )
662+ finally :
663+ event .remove (engine .sync_engine , "before_cursor_execute" , capture )
664+ statement , parameters = captured [0 ]
665+
666+ async with engine .connect () as conn :
667+ raw = (await conn .get_raw_connection ()).driver_connection
668+ assert raw is not None
669+ tx = raw .transaction ()
670+ await tx .start ()
671+ try :
672+ explained = await raw .fetchval (f"EXPLAIN (ANALYZE, FORMAT JSON) { statement } " , * parameters )
673+ finally :
674+ await tx .rollback ()
675+ return explained [0 ]["Plan" ]
676+
677+
619678async def test_fetch_reads_only_the_rows_it_claims_from_a_backlog (
620679 pg_engine : AsyncEngine ,
621680 outbox_table : Table ,
@@ -633,31 +692,36 @@ async def test_fetch_reads_only_the_rows_it_claims_from_a_backlog(
633692 await conn .execute (insert (outbox_table ).from_select (["queue" , "payload" ], backlog ))
634693 await conn .execute (text (f"ANALYZE { outbox_table .name } " ))
635694
636- captured : list [ tuple [ str , Any ]] = []
695+ plan = await _explain_fetch ( pg_engine , outbox_table , limit = limit )
637696
638- def capture (_conn : Any , _cursor : Any , statement : str , parameters : Any , _context : Any , _many : bool ) -> None :
639- if "FOR UPDATE SKIP LOCKED" in statement :
640- captured .append ((statement , parameters ))
697+ assert max (_scanned_rows (plan , outbox_table .name )) <= 2 * limit
641698
642- client = OutboxClient (pg_engine , outbox_table )
643- event .listen (pg_engine .sync_engine , "before_cursor_execute" , capture )
644- try :
645- async with pg_engine .connect () as conn :
646- await client .fetch (conn , ["orders" ], limit = limit , lease_ttl_seconds = 60.0 )
647- finally :
648- event .remove (pg_engine .sync_engine , "before_cursor_execute" , capture )
649- statement , parameters = captured [0 ]
650699
651- async with pg_engine .connect () as conn :
652- raw = (await conn .get_raw_connection ()).driver_connection
653- assert raw is not None
654- tx = raw .transaction ()
655- await tx .start ()
656- try :
657- explained = await raw .fetchval (f"EXPLAIN (ANALYZE, FORMAT JSON) { statement } " , * parameters )
658- finally :
659- await tx .rollback ()
660- plan = explained [0 ]["Plan" ]
700+ async def test_fetch_reads_only_the_rows_it_claims_from_expired_leases (
701+ pg_engine : AsyncEngine ,
702+ outbox_table : Table ,
703+ ) -> None :
704+ """INVARIANT: reclaiming expired leases reads about ``limit`` rows, however many handlers hung.
705+
706+ The expired-lease branch must come off an ordered scan of ``_leased_idx``. An index keyed on
707+ ``acquired_at`` makes Postgres read and sort every expired lease on the queue.
708+ """
709+ limit = 10
710+ async with pg_engine .begin () as conn :
711+ series = func .generate_series (1 , 20000 ).column_valued ("g" )
712+ leased = series % 4 == 0
713+ backlog = select (
714+ literal ("orders" ),
715+ literal (b"x" ),
716+ func .now () - func .make_interval (0 , 0 , 0 , 0 , 0 , 0 , series ),
717+ case ((leased , func .gen_random_uuid ())),
718+ case ((leased , func .now () - func .make_interval (0 , 0 , 0 , 0 , 1 ))),
719+ )
720+ columns = ["queue" , "payload" , "next_attempt_at" , "acquired_token" , "acquired_at" ]
721+ await conn .execute (insert (outbox_table ).from_select (columns , backlog ))
722+ await conn .execute (text (f"ANALYZE { outbox_table .name } " ))
723+
724+ plan = await _explain_fetch (pg_engine , outbox_table , limit = limit )
661725
662726 assert max (_scanned_rows (plan , outbox_table .name )) <= 2 * limit
663727
0 commit comments