From 0227d334b80929d25b1152f02319ee5ee2d27089 Mon Sep 17 00:00:00 2001 From: Alexey Shalaev Date: Mon, 7 Sep 2026 14:04:08 +0300 Subject: [PATCH 1/6] fix: attach the window a live writer is filling, under one lock DEFAULT reconciliation moved the window's rows out in one transaction and retried the ATTACH in another. On a table that is being written to the writer refills the window in the gap, PostgreSQL refuses the attach again, and after DEFAULT_CONFLICT_MAX_RETRIES the rows are moved back -- so the current month, the one every live table is writing into, could never be attached, and each attempt cost a full scan of DEFAULT under ACCESS EXCLUSIVE. The move and the attach now share a transaction: reconcile_and_attach takes EXCLUSIVE on the parent and ACCESS EXCLUSIVE on the partition and the DEFAULT sibling, moves what is left of the window, and attaches. Nothing can reach DEFAULT in between, and no insert is mid-routing when the partition set changes -- a writer waits at the parent and is then routed into the new partition, statement unchanged. The bulk of the window still moves first under the lighter SHARE ROW EXCLUSIVE, so the exclusive window is a tail and a scan. A window that still cannot be attached is now a default_holds_rows issue rather than the driver's exception: apply() records it and the run goes on, and partition_data returns a move issue with complete=False instead of raising out of the drain. Closes #64 --- docs/agents.md | 6 +- docs/concepts/execution.md | 39 ++++- docs/guide/extending.md | 13 +- docs/guide/partition-existing-table.md | 28 +++- pg_partsmith/aio/protocols.py | 49 +++++- pg_partsmith/aio/repositories/creator.py | 162 ++++++++++++++++-- pg_partsmith/aio/repositories/repository.py | 29 +++- pg_partsmith/aio/services/execution.py | 81 +++++++-- pg_partsmith/aio/services/migration.py | 22 ++- pg_partsmith/sync/protocols.py | 49 +++++- pg_partsmith/sync/repositories/creator.py | 163 +++++++++++++++++-- pg_partsmith/sync/repositories/repository.py | 29 +++- pg_partsmith/sync/services/execution.py | 81 +++++++-- pg_partsmith/sync/services/migration.py | 22 ++- tests/integration/aio/test_migration.py | 109 +++++++++++-- tests/integration/sync/test_concurrency.py | 110 ++++++++++++- tests/unit/sync/test_execution.py | 73 ++++++--- tests/unit/test_execution.py | 73 ++++++--- 18 files changed, 980 insertions(+), 158 deletions(-) diff --git a/docs/agents.md b/docs/agents.md index 5b00a02..4979d6d 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -365,7 +365,11 @@ lifecycle units — partitions directly under the root — never once per leaf o 10. **During `partition_data`, a window's rows are invisible through the parent** between the first batch and the attach: PostgreSQL will not attach a partition while DEFAULT still holds rows for it, so no ordering keeps them visible throughout. Rows are never - in two places, and never lost. + in two places, and never lost. The attach takes the rows that arrived during the last + batch *and* goes live in one transaction, so a window the application is writing into + is attached like any other; writers wait at the parent for that commit and are then + routed into the new partition. A window that still cannot be attached comes back as a + `move` issue with `complete=False` — `partition_data` does not raise for it. 11. **Flat and composed spellings do not mix**, and `extra="forbid"` means a misspelled keyword raises rather than being ignored. 12. **`schema=` goes in, `config.db_schema` comes out** — the field is aliased to avoid diff --git a/docs/concepts/execution.md b/docs/concepts/execution.md index 6966455..487649f 100644 --- a/docs/concepts/execution.md +++ b/docs/concepts/execution.md @@ -12,6 +12,12 @@ statements. There is no long transaction wrapping a run, because `DETACH … CON cannot run inside a transaction block at all, and because a run that creates three partitions and fails on the fourth should keep the three. +Two steps are exceptions, and both for the same reason: a lock has to span more than one +statement or the thing it protects can change underneath. A drop takes its lock, revalidates, +drains and drops in one transaction; the last step of [DEFAULT reconciliation](#default-reconciliation) +moves the window's remaining rows and attaches in one transaction. `ATTACH`, unlike +`DETACH … CONCURRENTLY`, may run inside a transaction block. + The consequence for callers: pass the service an **engine**, never a session you are using elsewhere. @@ -50,13 +56,28 @@ belong to the new window, PostgreSQL refuses the attach (`23514`). The executor: 1. moves those rows from DEFAULT into the new partition, naming columns on both sides (`ATTACH` matches by name, so physical column order may differ) and leaving rows with a - NULL trailing key where PostgreSQL routes them; -2. retries the attach. - -If the attach still fails, the rows are returned to DEFAULT rather than left in a table -no query can see. For a nested branch the moved rows are routed onward into its leaves. -A DEFAULT sibling holding rows for a hash or list member is reported -(`default_holds_rows`) rather than moved: only a RANGE window can be selected by its key. + NULL trailing key where PostgreSQL routes them. This is one statement under + `SHARE ROW EXCLUSIVE` on DEFAULT: writers wait, readers do not; +2. moves whatever arrived in the meantime and attaches, **in one transaction**, under the + locks `ATTACH` needs anyway — `SHARE UPDATE EXCLUSIVE` on the parent, `ACCESS EXCLUSIVE` + on the partition and on DEFAULT, taken in that order and taken before the move. + +Step 2 is what makes the window a live writer is inserting into attachable at all. Two +transactions cannot do it: the move commits, DEFAULT is free again, and the next insert for +the window lands in it before `ATTACH` scans — under any steady write rate that never +converges, and every attempt is a full scan of DEFAULT under `ACCESS EXCLUSIVE`. Sharing one +lock closes the gap; a writer waits for the commit and is then routed into the new partition. +Step 1 is what keeps that exclusive window short: the bulk of the month moves before the +heavy lock is taken, so what it covers is a tail and a scan. + +If the attach still fails, the rows step 1 moved are returned to DEFAULT rather than left in +a table no query can see — step 2's own move rolls back with the attach, and needs no +compensation. A window that still cannot be attached is reported as `default_holds_rows`, an +issue like any other: the run goes on, and `partition_data` returns it rather than raising. +For a nested branch the moved rows are routed onward into its leaves. A DEFAULT sibling +holding rows for a hash or list member is reported rather than moved: only a RANGE window can +be selected by its key. A **foreign** DEFAULT partition cannot be locked at all, so there the +race stays and the report is all there is. ## Detach @@ -103,6 +124,8 @@ repository). |---|---| | `CREATE TABLE … (LIKE parent)` | `ACCESS SHARE` on the parent | | `ATTACH PARTITION` | `SHARE UPDATE EXCLUSIVE` on the parent, `ACCESS EXCLUSIVE` on the child and on a DEFAULT sibling; `SHARE ROW EXCLUSIVE` on tables referencing the parent through a foreign key | +| the reconciling row move ([DEFAULT reconciliation](#default-reconciliation), step 1) | `SHARE ROW EXCLUSIVE` on the DEFAULT partition and on the child: writers of DEFAULT wait, readers do not | +| the move-and-attach ([DEFAULT reconciliation](#default-reconciliation), step 2) | the `ATTACH` row, taken up front — parent first, so a concurrent `ATTACH` queues there rather than deadlocking | | `DETACH PARTITION` (plain) | `ACCESS EXCLUSIVE` on parent, partition, and every table referencing the parent | | `DETACH PARTITION … CONCURRENTLY` | `SHARE UPDATE EXCLUSIVE` on the parent; `ACCESS EXCLUSIVE` on the partition and, in its second transaction, on referencing tables | | `DROP TABLE` of a detached table | `ACCESS EXCLUSIVE` on that table only | @@ -117,7 +140,7 @@ it takes on `op.capabilities`. | What happened | Effect on the run | |---|---| -| a topology conflict at execution time — a DEFAULT sibling holding rows for a hash bucket, a name taken by a relation with other bounds, a detach PostgreSQL refuses because rows are still referenced | recorded in `result.issues`; the run goes on | +| a topology conflict at execution time — a DEFAULT partition holding rows the attach could not take, a name taken by a relation with other bounds, a detach PostgreSQL refuses because rows are still referenced | recorded in `result.issues`; the run goes on | | a `PlanStaleError` — the relation is not the one the plan saw | recorded as an issue with `continue_on_error`, otherwise raised | | any other error — a connection drop, a permission denied, a `before_*` hook raising | aborts the run, unless `continue_on_error`, in which case it is recorded and the next operation runs | | validation or lock failure | fatal, always | diff --git a/docs/guide/extending.md b/docs/guide/extending.md index 57aa0e7..1385b75 100644 --- a/docs/guide/extending.md +++ b/docs/guide/extending.md @@ -39,6 +39,7 @@ class PartitionRepository(Protocol): async def drop_partition(self, partition_name, *, expected_oid=None) -> None: ... async def adopt_partition(self, table_name, partition_name) -> bool: ... async def reconcile_default_rows(self, *, default_partition_name, target_partition_name, key_columns, from_value, to_value, limit=None) -> int: ... + async def reconcile_and_attach(self, parent_name, partition_name, bounds, *, key_columns, default_partition_name) -> int: ... async def move_rows(self, source_name, target_name, *, limit=None) -> int: ... ``` @@ -46,13 +47,23 @@ Every method takes and returns plain domain objects (`PartitionBounds`, `Partiti `DetachMode`, `LocalLeaves`), so an implementation never needs to know how the planner works. +`reconcile_and_attach` is the one method with a transaction boundary in its contract: it +must take the window's remaining rows out of the DEFAULT partition **and** attach in a +single transaction, holding one lock across both, and roll the move back if the attach +fails. Two transactions cannot do it — a writer refills the window in the gap and +PostgreSQL refuses the attach every time. The bundled implementation locks the parent +`EXCLUSIVE` (so no insert is mid-routing when the partition set changes) and the partition +and DEFAULT sibling `ACCESS EXCLUSIVE`, then moves, then attaches. It is called only after +`attach_partition` has failed on a DEFAULT conflict and the bulk of the window has been +moved by `reconcile_default_rows`. + One thing to know about errors: the executor recognises a failed `attach_partition` by the **SQLSTATE the exception carries**, not by its type. It reads `sqlstate` or `pgcode` off the exception, or off its `orig` if it wraps one — which covers a SQLAlchemy error, a bare `psycopg.Error` and an `asyncpg.PostgresError` alike. Let the driver's exception through rather than replacing it with one of your own, and three things keep working: a lost race with another worker is treated as benign, a DEFAULT partition holding rows for the new -window triggers the reconcile-and-retry, and rows already moved out of DEFAULT are put back +window triggers the reconcile-and-attach, and rows already moved out of DEFAULT are put back if the attach ultimately fails. An exception carrying no SQLSTATE is still safe — the rows are restored and it propagates — but it cannot be recognised as a race or a conflict. diff --git a/docs/guide/partition-existing-table.md b/docs/guide/partition-existing-table.md index 967bf13..ff33151 100644 --- a/docs/guide/partition-existing-table.md +++ b/docs/guide/partition-existing-table.md @@ -48,7 +48,12 @@ await maintainer.run_maintenance_safe(config) The tick creates the current month and the two after it. As each is attached, the rows of that month move out of `events_legacy` into it — ordinary DEFAULT reconciliation. -From now on new rows land in real partitions; the old ones are still in the DEFAULT. +The current month is the one the application is writing into, and it attaches on this +first tick like the empty future ones: the bulk of its rows move first, then the rows +that arrived while that ran move *and* the partition goes live in one transaction, under +the lock `ATTACH` takes on the DEFAULT partition anyway. Inserts wait for that commit and +are then routed into the new partition. From now on new rows land in real partitions; the +old ones are still in the DEFAULT. ## 3. Drain the DEFAULT partition @@ -65,7 +70,10 @@ Each call: 3. moves the window's rows into it in batches of `batch_rows` — one `DELETE … RETURNING` / `INSERT` per batch, each committing on its own, so a row is in exactly one place at every commit point; -4. attaches the partition once nothing of that window is left in DEFAULT; +4. takes whatever landed during the last batch and attaches the partition, in one + transaction under the lock `ATTACH` needs anyway — so a window still being written to + is attached on the same pass as a quiet one, and the exclusive lock covers a tail and a + scan rather than a month; 5. moves on to the next window, until DEFAULT holds only rows no window can take (rows with a NULL key), or `max_batches` is spent. @@ -80,14 +88,18 @@ finds it, finishes it and attaches it. a partition cannot be attached while DEFAULT still holds rows for it. Run the drain in a maintenance window, or with small batches during a quiet hour and readers that can tolerate a month's rows appearing a little later. Rows already in real partitions, and - rows still in DEFAULT for other windows, stay visible throughout. + rows still in DEFAULT for other windows, stay visible throughout. The attach that ends + the window holds `ACCESS EXCLUSIVE` on the DEFAULT partition — readers and writers of + the *undrained* rows wait for it, for the length of the last batch's tail plus one + scan of DEFAULT. `partition_data` takes the table's lock, so it does not race the scheduled tick. It -refuses a window it cannot create (an unmanaged partition overlaps it), and any move an -incoming foreign key's `ON DELETE` action would corrupt, with a `move` issue and -`complete=False` rather than loop. A window whose partition already exists *detached* -with this library's marker — retention retired it, and late rows for it landed in -DEFAULT — is filled and re-attached rather than given up on. +refuses a window it cannot create (an unmanaged partition overlaps it), any move an +incoming foreign key's `ON DELETE` action would corrupt, and a DEFAULT partition it could +not clear, with a `move` issue and `complete=False` rather than loop — and never by +raising at the caller. A window whose partition already exists *detached* with this +library's marker — retention retired it, and late rows for it landed in DEFAULT — is +filled and re-attached rather than given up on. ## 4. Afterwards diff --git a/pg_partsmith/aio/protocols.py b/pg_partsmith/aio/protocols.py index 86b6052..4b36c1d 100644 --- a/pg_partsmith/aio/protocols.py +++ b/pg_partsmith/aio/protocols.py @@ -23,7 +23,15 @@ from pg_partsmith.leaves import LocalLeaves from pg_partsmith.lifecycle import DetachMode, SqlPredicate from pg_partsmith.plan import PartitionBy -from pg_partsmith.topology import ActualTree, FactKind, PartitionBounds, PartitionNode, PartitionType, RelationKind +from pg_partsmith.topology import ( + ActualTree, + FactKind, + PartitionBounds, + PartitionNode, + PartitionType, + RangeBounds, + RelationKind, +) __all__ = [ "LockManager", @@ -239,6 +247,45 @@ async def reconcile_default_rows( """ ... + async def reconcile_and_attach( + self, + parent_name: str, + partition_name: str, + bounds: RangeBounds, + *, + key_columns: tuple[str, ...], + default_partition_name: str, + expected_oid: int | None = None, + expected_parent_oid: int | None = None, + expected_default_oid: int | None = None, + ) -> int: + """Take the window's last rows out of DEFAULT and attach, in one transaction. + + The pair cannot be two transactions on a table that is being written + to: the move commits, the DEFAULT partition is free again, and a row + for the window lands in it before ``ATTACH`` scans -- which fails the + attach with the very conflict the move was clearing. An implementation + must hold one lock across both, and must roll the move back with a + failed attach. + + Args: + parent_name: Partitioned relation to attach to. + partition_name: Table to attach. + bounds: The RANGE window the partition owns, and the rows to take. + key_columns: The parent's partition key, leading column first. Rows + with a NULL in a trailing column stay in DEFAULT, where + PostgreSQL routes them. + default_partition_name: Qualified name of DEFAULT partition. + expected_oid: The catalog identity of the partition, checked the + way ``attach_partition`` checks it. + expected_parent_oid: The catalog identity of the parent. + expected_default_oid: The catalog identity the rows must come from. + + Returns: + Number of rows moved out of the DEFAULT partition. + """ + ... + async def move_rows(self, source_name: str, target_name: str, *, limit: int | None = None) -> int: """Move rows from one relation into another, whatever their keys. diff --git a/pg_partsmith/aio/repositories/creator.py b/pg_partsmith/aio/repositories/creator.py index 11ff0e6..07f8c93 100644 --- a/pg_partsmith/aio/repositories/creator.py +++ b/pg_partsmith/aio/repositories/creator.py @@ -391,21 +391,7 @@ async def reconcile_default_rows( Returns: Number of rows moved. """ - if not key_columns: - msg = "reconcile_default_rows needs the parent's partition key" - raise ValueError(msg) - - column_quoted = quote_identifier(key_columns[0]) - from_quoted = quote_literal(from_value) - to_quoted = quote_literal(to_value) - - # PostgreSQL adds an IS NOT NULL test for *every* key column to a range - # partition's constraint, so a row with a NULL trailing key value - # belongs in DEFAULT whatever its leading value is. Moving it out would - # be rejected -- and the rejection would look exactly like the DEFAULT - # conflict this call exists to clear, so the retry would never converge. - not_null = "".join(f" AND {quote_identifier(column)} IS NOT NULL" for column in key_columns[1:]) - + condition = _window_condition("reconcile_default_rows", key_columns, from_value, to_value) async with asyncio.timeout(self._ddl_timeout), self._engine.begin() as conn: # Boundary literals must be interpreted in the same timezone ATTACH uses, # otherwise a non-UTC server timezone moves the wrong row range. @@ -415,8 +401,6 @@ async def reconcile_default_rows( # Deferred foreign-key checks would otherwise escape to COMMIT, # outside the per-statement translation below. await conn.execute(text("SET CONSTRAINTS ALL IMMEDIATE")) - # All identifiers and literals are properly quoted, S608 is a false positive - condition = f"{column_quoted} >= {from_quoted} AND {column_quoted} < {to_quoted}{not_null}" return await self._move( conn, default_partition_name, @@ -427,6 +411,95 @@ async def reconcile_default_rows( expected_target_oid=expected_target_oid, ) + async def reconcile_and_attach( + self, + parent_name: str, + partition_name: str, + bounds: RangeBounds, + *, + key_columns: tuple[str, ...], + default_partition_name: str, + expected_oid: int | None = None, + expected_parent_oid: int | None = None, + expected_default_oid: int | None = None, + ) -> int: + """Take the window's last rows out of DEFAULT and attach, in one transaction. + + The reconcile-then-attach pair cannot be two transactions on a table + that is being written to: the move commits, the DEFAULT partition is + free again, and a row for the window lands in it before ``ATTACH`` + scans -- which fails the attach with the very conflict the move was + clearing. Under a steady write rate that never converges. + + So both go under one lock. The transaction takes its locks up front -- + ``EXCLUSIVE`` on the parent, ``ACCESS EXCLUSIVE`` on the partition and + on the DEFAULT sibling -- moves the rows, and attaches. Nothing can put + a row into DEFAULT between the move and the scan, because nothing has + been able to reach DEFAULT since before the move; a writer waits at the + parent and is then routed into the new partition, statement unchanged. + A failed attach rolls the move back with it, so this path needs no + compensating move-back either. + + Only the tail belongs here. The bulk of a window is moved first with + :meth:`reconcile_default_rows`, which holds the lighter ``SHARE ROW + EXCLUSIVE`` and leaves readers of DEFAULT alone; what this call takes + exclusively is the rows that arrived while that ran, plus the scan + ``ATTACH`` does anyway. + + Args: + parent_name: Partitioned relation to attach to. + partition_name: Table to attach. + bounds: The RANGE window the partition owns, and the rows to take. + key_columns: The parent's partition key, leading column first. + Rows with a NULL in a trailing column stay in DEFAULT, where + PostgreSQL routes them. + default_partition_name: Qualified name of the DEFAULT partition. + expected_oid: The identity the decision to attach was made about. + expected_parent_oid: The identity of the relation the partition + should go into. + expected_default_oid: The identity of the relation the rows should + come from. + + Returns: + Number of rows moved out of the DEFAULT partition. + + Raises: + PlanStaleError: If any of the three names is not held by the + relation the caller decided about; nothing moves, nothing + stays attached. + RowMoveRefusedError: If a foreign key's ``ON DELETE`` action would + fire on the rows as they leave the DEFAULT partition. + """ + condition = _window_condition("reconcile_and_attach", key_columns, bounds.from_value, bounds.to_value) + clause, values = _values_clause(bounds, max(1, len(key_columns))) + stmt = build_ddl_statement( + "ALTER TABLE {parent} ATTACH PARTITION {partition} " + clause, + parent=parent_name, + partition=partition_name, + **values, + ) + async with asyncio.timeout(self._ddl_timeout), self._engine.begin() as conn: + if self._ddl_timezone is not None: + await conn.execute(text(f"SET LOCAL TIME ZONE {quote_literal(self._ddl_timezone)}")) + await self._lock_for_attach(conn, parent_name, partition_name, default_partition_name) + await conn.execute(text("SET CONSTRAINTS ALL IMMEDIATE")) + moved = await self._move( + conn, + default_partition_name, + partition_name, + condition=condition, + limit=None, + expected_source_oid=expected_default_oid, + expected_target_oid=expected_oid, + ) + await self._require_oid(conn, parent_name, expected_parent_oid) + await conn.execute(stmt) + # Under ATTACH's own locks now: what these see is what went live. + await self._require_oid(conn, partition_name, expected_oid) + await self._require_oid(conn, parent_name, expected_parent_oid) + await self._clear_orphan_marker(conn, partition_name) + return moved + async def move_rows(self, source_name: str, target_name: str, *, limit: int | None = None) -> int: """Move rows from one relation into another, whatever their keys. @@ -470,6 +543,42 @@ async def _lock_for_move(self, conn: AsyncConnection, *names: str) -> None: if await relation_kind(conn, name) != "f": await conn.execute(text(f"LOCK TABLE {quote_identifier(name)} IN SHARE ROW EXCLUSIVE MODE")) + async def _lock_for_attach( + self, conn: AsyncConnection, parent_name: str, partition_name: str, default_partition_name: str + ) -> None: + """Take the move-and-attach's locks: the parent's writers out, then ATTACH's own two. + + ``EXCLUSIVE`` on the parent is one level above what ``ATTACH`` itself + takes, and it is there for the writer's sake rather than this + transaction's. An INSERT routing through the parent chooses its + partition from the set it saw when it took ROW EXCLUSIVE on the parent: + one that got that far and then queued on the DEFAULT partition's lock + would come out of the wait still aimed at DEFAULT, and be rejected by + the constraint this attach just narrowed -- a lost write, and the same + thing a plain ``ATTACH`` does to a live writer. Blocking at the parent + instead means no insert is ever mid-routing while the partition set + changes: the writer waits, re-plans against the tree the attach left, + and lands in the new partition. ``EXCLUSIVE`` does not conflict with + ``ACCESS SHARE``, so readers of the other partitions carry on. + + The other two are the ``ACCESS EXCLUSIVE`` locks ``ATTACH`` takes on + the partition and on the DEFAULT sibling it scans, taken early so the + move runs under the very lock the scan will hold. Taking the parent + first is also what keeps a concurrent ``ATTACH`` queued there instead + of deadlocking against this one over the DEFAULT partition. ``ONLY`` + keeps the parent's other partitions out of it; the two below it are + locked with their own subtrees, which is where a branch's rows are + written. + + A foreign relation cannot be locked at all, so a foreign DEFAULT + partition is left to ``ATTACH``'s own scan -- with the race that + implies, and the conflict it can still raise reported as an issue. + """ + await conn.execute(text(f"LOCK TABLE ONLY {quote_identifier(parent_name)} IN EXCLUSIVE MODE")) + for name in (partition_name, default_partition_name): + if await relation_kind(conn, name) != "f": + await conn.execute(text(f"LOCK TABLE {quote_identifier(name)} IN ACCESS EXCLUSIVE MODE")) + async def _move( self, conn: AsyncConnection, @@ -756,6 +865,25 @@ async def _has_identity_always(self, conn: AsyncConnection, table_name: str) -> return bool(result.scalar()) +def _window_condition(caller: str, key_columns: tuple[str, ...], from_value: str, to_value: str) -> str: + """The WHERE clause selecting a RANGE window's rows out of a DEFAULT partition. + + PostgreSQL adds an IS NOT NULL test for *every* key column to a range + partition's constraint, so a row with a NULL trailing key value belongs in + DEFAULT whatever its leading value is. Moving it out would be rejected -- + and the rejection would look exactly like the DEFAULT conflict the move + exists to clear, so the retry would never converge. Every identifier and + literal is quoted here, which is what makes S608 a false positive at the + call sites. + """ + if not key_columns: + msg = f"{caller} needs the parent's partition key" + raise ValueError(msg) + column = quote_identifier(key_columns[0]) + not_null = "".join(f" AND {quote_identifier(name)} IS NOT NULL" for name in key_columns[1:]) + return f"{column} >= {quote_literal(from_value)} AND {column} < {quote_literal(to_value)}{not_null}" + + def _move_statement( source: str, target: str, diff --git a/pg_partsmith/aio/repositories/repository.py b/pg_partsmith/aio/repositories/repository.py index 5320bd6..fc07c1b 100644 --- a/pg_partsmith/aio/repositories/repository.py +++ b/pg_partsmith/aio/repositories/repository.py @@ -31,7 +31,7 @@ from pg_partsmith.leaves import LocalLeaves from pg_partsmith.plan import PartitionBy - from pg_partsmith.topology import PartitionBounds + from pg_partsmith.topology import PartitionBounds, RangeBounds class PostgresPartitionRepository: @@ -206,6 +206,33 @@ async def reconcile_default_rows( expected_target_oid=expected_target_oid, ) + async def reconcile_and_attach( + self, + parent_name: str, + partition_name: str, + bounds: RangeBounds, + *, + key_columns: tuple[str, ...], + default_partition_name: str, + expected_oid: int | None = None, + expected_parent_oid: int | None = None, + expected_default_oid: int | None = None, + ) -> int: + """Take the window's last rows out of DEFAULT and attach, in one transaction. + + See :meth:`PartitionCreator.reconcile_and_attach`. + """ + return await self._creator.reconcile_and_attach( + parent_name, + partition_name, + bounds, + key_columns=key_columns, + default_partition_name=default_partition_name, + expected_oid=expected_oid, + expected_parent_oid=expected_parent_oid, + expected_default_oid=expected_default_oid, + ) + async def move_rows(self, source_name: str, target_name: str, *, limit: int | None = None) -> int: """Move rows from one relation into another. diff --git a/pg_partsmith/aio/services/execution.py b/pg_partsmith/aio/services/execution.py index 7f94589..b2af363 100644 --- a/pg_partsmith/aio/services/execution.py +++ b/pg_partsmith/aio/services/execution.py @@ -440,6 +440,17 @@ async def _attach_with_reconcile( ) -> None: """Attach a partition, moving DEFAULT rows out of the way for a RANGE window. + The first attempt is the plain attach: a parent with no DEFAULT + partition, or one holding nothing for the window, pays for nothing + else. A DEFAULT conflict turns the second attempt into + ``reconcile_and_attach``, which moves and attaches in one transaction + under the lock ``ATTACH`` takes on the DEFAULT partition anyway -- the + two cannot be separate transactions on a table that is being written + to, because the writer fills the window again in the gap between them. + The bulk of the window is moved before that, in the reconcile's own + lighter transaction, so what the exclusive lock covers is the rows + that arrived while it ran, plus the scan. + If the attach ultimately fails after rows were reconciled out of the DEFAULT partition, the moved rows are returned to DEFAULT (best effort) so they do not end up stranded in a table that is invisible through the @@ -453,17 +464,34 @@ async def _attach_with_reconcile( key_arity = max(1, len(key_columns)) reconciled_from: tuple[str, int | None] | None = None window = bounds if isinstance(bounds, RangeBounds) else None + default_partition: PartitionInfo | None = None for attempt in range(1, DEFAULT_CONFLICT_MAX_RETRIES + 1): try: - await self._repo.attach_partition( - parent_name, - partition_name, - bounds, - key_arity=key_arity, - expected_oid=expected_oid, - expected_parent_oid=expected_parent_oid, - ) + if default_partition is None or window is None: + await self._repo.attach_partition( + parent_name, + partition_name, + bounds, + key_arity=key_arity, + expected_oid=expected_oid, + expected_parent_oid=expected_parent_oid, + ) + else: + moved = await self._repo.reconcile_and_attach( + parent_name, + partition_name, + window, + key_columns=key_columns, + default_partition_name=default_partition.name, + expected_oid=expected_oid, + expected_parent_oid=expected_parent_oid, + expected_default_oid=default_partition.oid, + ) + logger.info( + "Attached with the window's last rows taken under one lock", + extra={"partition_name": partition_name, "moved_rows": moved}, + ) except (asyncio.CancelledError, KeyboardInterrupt, SystemExit): # Shielded so the compensating move-back completes even mid-cancellation. await asyncio.shield( @@ -473,6 +501,11 @@ async def _attach_with_reconcile( except (OSError, TimeoutError): await self._restore_reconciled_rows(reconciled_from, partition_name, expected_oid, window, key_columns) raise + except RowMoveRefusedError as refusal: + # The move inside ``reconcile_and_attach`` rolled back with the + # attach; the rows moved before it go back here. + await self._restore_reconciled_rows(reconciled_from, partition_name, expected_oid, window, key_columns) + raise self._rows_stuck_in_default(parent_name, partition_name, refusal.detail) from refusal except Exception as exc: # Recognised structurally, by the SQLSTATE the driver carries, # rather than by a driver's exception type: a repository built @@ -512,14 +545,19 @@ async def _attach_with_reconcile( ) from exc if attempt == DEFAULT_CONFLICT_MAX_RETRIES: - logger.exception( + # Nothing left to try: the DEFAULT partition still holds + # rows for this window and this run cannot take them. That + # is a topology finding, not a failure of the run -- the + # other partitions are maintained, and ``partition_data`` + # reports the window instead of ending on the exception. + logger.warning( "Failed to attach after reconciliation retries", extra={"partition_name": partition_name, "attempts": attempt}, ) await self._restore_reconciled_rows( reconciled_from, partition_name, expected_oid, window, key_columns ) - raise + raise self._rows_stuck_in_default(parent_name, partition_name, describe_exception(exc)) from exc default_partition = await self._metadata.get_default_partition(parent_name) if not default_partition: @@ -563,18 +601,29 @@ async def _attach_with_reconcile( except RowMoveRefusedError as refusal: # Rows of DEFAULT belong to the new partition but cannot be # moved safely; the partition stays out, the run goes on. - raise PartitionTopologyError( - parent_name, - FindingReason.DEFAULT_HOLDS_ROWS.value, - f"{parent_name} cannot gain {partition_name!r} while its DEFAULT partition holds rows that " - f"belong to it, and they cannot be moved: {refusal.detail}", - ) from refusal + raise self._rows_stuck_in_default(parent_name, partition_name, refusal.detail) from refusal if moved: reconciled_from = (default_partition.name, default_partition.oid) logger.info("Reconciliation completed", extra={"partition_name": partition_name, "moved_rows": moved}) else: return + @staticmethod + def _rows_stuck_in_default(parent_name: str, partition_name: str, detail: str) -> PartitionTopologyError: + """The DEFAULT partition holds rows for the new partition that this run could not clear. + + A finding about the shape of the tree, not a failure of the run: every + other partition is still maintained, ``apply`` records it in + ``result.issues``, and ``partition_data`` reports the window rather + than ending on the driver's exception. + """ + return PartitionTopologyError( + parent_name, + FindingReason.DEFAULT_HOLDS_ROWS.value, + f"{parent_name} cannot gain {partition_name!r} while its DEFAULT partition holds rows that belong to " + f"it, and this run could not clear them: {detail}", + ) + async def _restore_reconciled_rows( self, reconciled_from: tuple[str, int | None] | None, diff --git a/pg_partsmith/aio/services/migration.py b/pg_partsmith/aio/services/migration.py index 286582c..1921a2b 100644 --- a/pg_partsmith/aio/services/migration.py +++ b/pg_partsmith/aio/services/migration.py @@ -23,7 +23,7 @@ from pg_partsmith.boundaries import Window from pg_partsmith.constants import DEFAULT_MOVE_BATCH_ROWS from pg_partsmith.entities import MaintenanceIssue, MaintenanceIssueStep, MigrationResult -from pg_partsmith.exceptions import InvalidPartitionConfigError, RowMoveRefusedError +from pg_partsmith.exceptions import InvalidPartitionConfigError, PartitionTopologyError, RowMoveRefusedError from pg_partsmith.lifecycle import DropAfter from pg_partsmith.plan import AttachPartition, CreatePartition, DetachPartition, DropPartition, MaintenancePlan, Reason from pg_partsmith.planner import to_maintenance_issue @@ -62,8 +62,16 @@ async def partition_data( Window by window, oldest first: the partition for the oldest window still in DEFAULT is created detached (subtree included), filled from DEFAULT in batches of ``batch_rows``, and attached once DEFAULT holds - nothing more for it. A partition left detached when ``max_batches`` - runs out is picked up and finished by the next call. + nothing more for it. The attach takes whatever arrived while the + batches ran in its own transaction, under the lock it needs anyway, so + the window a live writer is inserting into goes live on the same pass + as the quiet ones. A partition left detached when ``max_batches`` runs + out is picked up and finished by the next call. + + A window that cannot be finished at all -- rows an incoming foreign + key holds down, a DEFAULT partition this run could not clear -- is + reported as a ``move`` issue with ``complete=False``; it is never + raised at the caller. Args: config: The table's configuration; its root must be a RANGE level. @@ -115,7 +123,13 @@ async def partition_data( attached = await self._executor.create_partition(config, plan, op, issues=tally.issues, fill=fill) else: attached = await self._executor.attach_partition(config, plan, op, issues=tally.issues, fill=fill) - except RowMoveRefusedError as exc: + except (RowMoveRefusedError, PartitionTopologyError) as exc: + # A window this run cannot finish -- rows a foreign key holds + # down, a DEFAULT partition it could not clear, a name taken by + # a relation with other bounds. The caller asked for a drain, + # not for an exception: what stayed behind and why is on the + # result, and the loop stops rather than re-planning the same + # window forever. tally.issue( default.name, f"rows for {boundaries.describe(window)} stay in {default.name}: {exc.detail}" ) diff --git a/pg_partsmith/sync/protocols.py b/pg_partsmith/sync/protocols.py index ae18a13..b1d19e9 100644 --- a/pg_partsmith/sync/protocols.py +++ b/pg_partsmith/sync/protocols.py @@ -23,7 +23,15 @@ from pg_partsmith.leaves import LocalLeaves from pg_partsmith.lifecycle import DetachMode, SqlPredicate from pg_partsmith.plan import PartitionBy -from pg_partsmith.topology import ActualTree, FactKind, PartitionBounds, PartitionNode, PartitionType, RelationKind +from pg_partsmith.topology import ( + ActualTree, + FactKind, + PartitionBounds, + PartitionNode, + PartitionType, + RangeBounds, + RelationKind, +) __all__ = [ "LockManager", @@ -239,6 +247,45 @@ def reconcile_default_rows( """ ... + def reconcile_and_attach( + self, + parent_name: str, + partition_name: str, + bounds: RangeBounds, + *, + key_columns: tuple[str, ...], + default_partition_name: str, + expected_oid: int | None = None, + expected_parent_oid: int | None = None, + expected_default_oid: int | None = None, + ) -> int: + """Take the window's last rows out of DEFAULT and attach, in one transaction. + + The pair cannot be two transactions on a table that is being written + to: the move commits, the DEFAULT partition is free again, and a row + for the window lands in it before ``ATTACH`` scans -- which fails the + attach with the very conflict the move was clearing. An implementation + must hold one lock across both, and must roll the move back with a + failed attach. + + Args: + parent_name: Partitioned relation to attach to. + partition_name: Table to attach. + bounds: The RANGE window the partition owns, and the rows to take. + key_columns: The parent's partition key, leading column first. Rows + with a NULL in a trailing column stay in DEFAULT, where + PostgreSQL routes them. + default_partition_name: Qualified name of DEFAULT partition. + expected_oid: The catalog identity of the partition, checked the + way ``attach_partition`` checks it. + expected_parent_oid: The catalog identity of the parent. + expected_default_oid: The catalog identity the rows must come from. + + Returns: + Number of rows moved out of the DEFAULT partition. + """ + ... + def move_rows(self, source_name: str, target_name: str, *, limit: int | None = None) -> int: """Move rows from one relation into another, whatever their keys. diff --git a/pg_partsmith/sync/repositories/creator.py b/pg_partsmith/sync/repositories/creator.py index 3787ed7..fe1d45e 100644 --- a/pg_partsmith/sync/repositories/creator.py +++ b/pg_partsmith/sync/repositories/creator.py @@ -394,21 +394,7 @@ def reconcile_default_rows( Returns: Number of rows moved. """ - if not key_columns: - msg = "reconcile_default_rows needs the parent's partition key" - raise ValueError(msg) - - column_quoted = quote_identifier(key_columns[0]) - from_quoted = quote_literal(from_value) - to_quoted = quote_literal(to_value) - - # PostgreSQL adds an IS NOT NULL test for *every* key column to a range - # partition's constraint, so a row with a NULL trailing key value - # belongs in DEFAULT whatever its leading value is. Moving it out would - # be rejected -- and the rejection would look exactly like the DEFAULT - # conflict this call exists to clear, so the retry would never converge. - not_null = "".join(f" AND {quote_identifier(column)} IS NOT NULL" for column in key_columns[1:]) - + condition = _window_condition("reconcile_default_rows", key_columns, from_value, to_value) with self._engine.begin() as conn: apply_local_statement_timeout(conn, self._ddl_timeout) # Boundary literals must be interpreted in the same timezone ATTACH uses, @@ -419,8 +405,6 @@ def reconcile_default_rows( # Deferred foreign-key checks would otherwise escape to COMMIT, # outside the per-statement translation below. conn.execute(text("SET CONSTRAINTS ALL IMMEDIATE")) - # All identifiers and literals are properly quoted, S608 is a false positive - condition = f"{column_quoted} >= {from_quoted} AND {column_quoted} < {to_quoted}{not_null}" return self._move( conn, default_partition_name, @@ -431,6 +415,96 @@ def reconcile_default_rows( expected_target_oid=expected_target_oid, ) + def reconcile_and_attach( + self, + parent_name: str, + partition_name: str, + bounds: RangeBounds, + *, + key_columns: tuple[str, ...], + default_partition_name: str, + expected_oid: int | None = None, + expected_parent_oid: int | None = None, + expected_default_oid: int | None = None, + ) -> int: + """Take the window's last rows out of DEFAULT and attach, in one transaction. + + The reconcile-then-attach pair cannot be two transactions on a table + that is being written to: the move commits, the DEFAULT partition is + free again, and a row for the window lands in it before ``ATTACH`` + scans -- which fails the attach with the very conflict the move was + clearing. Under a steady write rate that never converges. + + So both go under one lock. The transaction takes its locks up front -- + ``EXCLUSIVE`` on the parent, ``ACCESS EXCLUSIVE`` on the partition and + on the DEFAULT sibling -- moves the rows, and attaches. Nothing can put + a row into DEFAULT between the move and the scan, because nothing has + been able to reach DEFAULT since before the move; a writer waits at the + parent and is then routed into the new partition, statement unchanged. + A failed attach rolls the move back with it, so this path needs no + compensating move-back either. + + Only the tail belongs here. The bulk of a window is moved first with + :meth:`reconcile_default_rows`, which holds the lighter ``SHARE ROW + EXCLUSIVE`` and leaves readers of DEFAULT alone; what this call takes + exclusively is the rows that arrived while that ran, plus the scan + ``ATTACH`` does anyway. + + Args: + parent_name: Partitioned relation to attach to. + partition_name: Table to attach. + bounds: The RANGE window the partition owns, and the rows to take. + key_columns: The parent's partition key, leading column first. + Rows with a NULL in a trailing column stay in DEFAULT, where + PostgreSQL routes them. + default_partition_name: Qualified name of the DEFAULT partition. + expected_oid: The identity the decision to attach was made about. + expected_parent_oid: The identity of the relation the partition + should go into. + expected_default_oid: The identity of the relation the rows should + come from. + + Returns: + Number of rows moved out of the DEFAULT partition. + + Raises: + PlanStaleError: If any of the three names is not held by the + relation the caller decided about; nothing moves, nothing + stays attached. + RowMoveRefusedError: If a foreign key's ``ON DELETE`` action would + fire on the rows as they leave the DEFAULT partition. + """ + condition = _window_condition("reconcile_and_attach", key_columns, bounds.from_value, bounds.to_value) + clause, values = _values_clause(bounds, max(1, len(key_columns))) + stmt = build_ddl_statement( + "ALTER TABLE {parent} ATTACH PARTITION {partition} " + clause, + parent=parent_name, + partition=partition_name, + **values, + ) + with self._engine.begin() as conn: + apply_local_statement_timeout(conn, self._ddl_timeout) + if self._ddl_timezone is not None: + conn.execute(text(f"SET LOCAL TIME ZONE {quote_literal(self._ddl_timezone)}")) + self._lock_for_attach(conn, parent_name, partition_name, default_partition_name) + conn.execute(text("SET CONSTRAINTS ALL IMMEDIATE")) + moved = self._move( + conn, + default_partition_name, + partition_name, + condition=condition, + limit=None, + expected_source_oid=expected_default_oid, + expected_target_oid=expected_oid, + ) + self._require_oid(conn, parent_name, expected_parent_oid) + conn.execute(stmt) + # Under ATTACH's own locks now: what these see is what went live. + self._require_oid(conn, partition_name, expected_oid) + self._require_oid(conn, parent_name, expected_parent_oid) + self._clear_orphan_marker(conn, partition_name) + return moved + def move_rows(self, source_name: str, target_name: str, *, limit: int | None = None) -> int: """Move rows from one relation into another, whatever their keys. @@ -475,6 +549,42 @@ def _lock_for_move(self, conn: Connection, *names: str) -> None: if relation_kind(conn, name) != "f": conn.execute(text(f"LOCK TABLE {quote_identifier(name)} IN SHARE ROW EXCLUSIVE MODE")) + def _lock_for_attach( + self, conn: Connection, parent_name: str, partition_name: str, default_partition_name: str + ) -> None: + """Take the move-and-attach's locks: the parent's writers out, then ATTACH's own two. + + ``EXCLUSIVE`` on the parent is one level above what ``ATTACH`` itself + takes, and it is there for the writer's sake rather than this + transaction's. An INSERT routing through the parent chooses its + partition from the set it saw when it took ROW EXCLUSIVE on the parent: + one that got that far and then queued on the DEFAULT partition's lock + would come out of the wait still aimed at DEFAULT, and be rejected by + the constraint this attach just narrowed -- a lost write, and the same + thing a plain ``ATTACH`` does to a live writer. Blocking at the parent + instead means no insert is ever mid-routing while the partition set + changes: the writer waits, re-plans against the tree the attach left, + and lands in the new partition. ``EXCLUSIVE`` does not conflict with + ``ACCESS SHARE``, so readers of the other partitions carry on. + + The other two are the ``ACCESS EXCLUSIVE`` locks ``ATTACH`` takes on + the partition and on the DEFAULT sibling it scans, taken early so the + move runs under the very lock the scan will hold. Taking the parent + first is also what keeps a concurrent ``ATTACH`` queued there instead + of deadlocking against this one over the DEFAULT partition. ``ONLY`` + keeps the parent's other partitions out of it; the two below it are + locked with their own subtrees, which is where a branch's rows are + written. + + A foreign relation cannot be locked at all, so a foreign DEFAULT + partition is left to ``ATTACH``'s own scan -- with the race that + implies, and the conflict it can still raise reported as an issue. + """ + conn.execute(text(f"LOCK TABLE ONLY {quote_identifier(parent_name)} IN EXCLUSIVE MODE")) + for name in (partition_name, default_partition_name): + if relation_kind(conn, name) != "f": + conn.execute(text(f"LOCK TABLE {quote_identifier(name)} IN ACCESS EXCLUSIVE MODE")) + def _move( self, conn: Connection, @@ -755,6 +865,25 @@ def _has_identity_always(self, conn: Connection, table_name: str) -> bool: return bool(result.scalar()) +def _window_condition(caller: str, key_columns: tuple[str, ...], from_value: str, to_value: str) -> str: + """The WHERE clause selecting a RANGE window's rows out of a DEFAULT partition. + + PostgreSQL adds an IS NOT NULL test for *every* key column to a range + partition's constraint, so a row with a NULL trailing key value belongs in + DEFAULT whatever its leading value is. Moving it out would be rejected -- + and the rejection would look exactly like the DEFAULT conflict the move + exists to clear, so the retry would never converge. Every identifier and + literal is quoted here, which is what makes S608 a false positive at the + call sites. + """ + if not key_columns: + msg = f"{caller} needs the parent's partition key" + raise ValueError(msg) + column = quote_identifier(key_columns[0]) + not_null = "".join(f" AND {quote_identifier(name)} IS NOT NULL" for name in key_columns[1:]) + return f"{column} >= {quote_literal(from_value)} AND {column} < {quote_literal(to_value)}{not_null}" + + def _move_statement( source: str, target: str, diff --git a/pg_partsmith/sync/repositories/repository.py b/pg_partsmith/sync/repositories/repository.py index b734b4d..293e793 100644 --- a/pg_partsmith/sync/repositories/repository.py +++ b/pg_partsmith/sync/repositories/repository.py @@ -31,7 +31,7 @@ from pg_partsmith.leaves import LocalLeaves from pg_partsmith.plan import PartitionBy - from pg_partsmith.topology import PartitionBounds + from pg_partsmith.topology import PartitionBounds, RangeBounds class PostgresPartitionRepository: @@ -206,6 +206,33 @@ def reconcile_default_rows( expected_target_oid=expected_target_oid, ) + def reconcile_and_attach( + self, + parent_name: str, + partition_name: str, + bounds: RangeBounds, + *, + key_columns: tuple[str, ...], + default_partition_name: str, + expected_oid: int | None = None, + expected_parent_oid: int | None = None, + expected_default_oid: int | None = None, + ) -> int: + """Take the window's last rows out of DEFAULT and attach, in one transaction. + + See :meth:`PartitionCreator.reconcile_and_attach`. + """ + return self._creator.reconcile_and_attach( + parent_name, + partition_name, + bounds, + key_columns=key_columns, + default_partition_name=default_partition_name, + expected_oid=expected_oid, + expected_parent_oid=expected_parent_oid, + expected_default_oid=expected_default_oid, + ) + def move_rows(self, source_name: str, target_name: str, *, limit: int | None = None) -> int: """Move rows from one relation into another. diff --git a/pg_partsmith/sync/services/execution.py b/pg_partsmith/sync/services/execution.py index 2b6d0d2..3b8b191 100644 --- a/pg_partsmith/sync/services/execution.py +++ b/pg_partsmith/sync/services/execution.py @@ -439,6 +439,17 @@ def _attach_with_reconcile( ) -> None: """Attach a partition, moving DEFAULT rows out of the way for a RANGE window. + The first attempt is the plain attach: a parent with no DEFAULT + partition, or one holding nothing for the window, pays for nothing + else. A DEFAULT conflict turns the second attempt into + ``reconcile_and_attach``, which moves and attaches in one transaction + under the lock ``ATTACH`` takes on the DEFAULT partition anyway -- the + two cannot be separate transactions on a table that is being written + to, because the writer fills the window again in the gap between them. + The bulk of the window is moved before that, in the reconcile's own + lighter transaction, so what the exclusive lock covers is the rows + that arrived while it ran, plus the scan. + If the attach ultimately fails after rows were reconciled out of the DEFAULT partition, the moved rows are returned to DEFAULT (best effort) so they do not end up stranded in a table that is invisible through the @@ -452,17 +463,34 @@ def _attach_with_reconcile( key_arity = max(1, len(key_columns)) reconciled_from: tuple[str, int | None] | None = None window = bounds if isinstance(bounds, RangeBounds) else None + default_partition: PartitionInfo | None = None for attempt in range(1, DEFAULT_CONFLICT_MAX_RETRIES + 1): try: - self._repo.attach_partition( - parent_name, - partition_name, - bounds, - key_arity=key_arity, - expected_oid=expected_oid, - expected_parent_oid=expected_parent_oid, - ) + if default_partition is None or window is None: + self._repo.attach_partition( + parent_name, + partition_name, + bounds, + key_arity=key_arity, + expected_oid=expected_oid, + expected_parent_oid=expected_parent_oid, + ) + else: + moved = self._repo.reconcile_and_attach( + parent_name, + partition_name, + window, + key_columns=key_columns, + default_partition_name=default_partition.name, + expected_oid=expected_oid, + expected_parent_oid=expected_parent_oid, + expected_default_oid=default_partition.oid, + ) + logger.info( + "Attached with the window's last rows taken under one lock", + extra={"partition_name": partition_name, "moved_rows": moved}, + ) except (KeyboardInterrupt, SystemExit): # Shielded so the compensating move-back completes even mid-cancellation. (self._restore_reconciled_rows(reconciled_from, partition_name, expected_oid, window, key_columns)) @@ -470,6 +498,11 @@ def _attach_with_reconcile( except (OSError, TimeoutError): self._restore_reconciled_rows(reconciled_from, partition_name, expected_oid, window, key_columns) raise + except RowMoveRefusedError as refusal: + # The move inside ``reconcile_and_attach`` rolled back with the + # attach; the rows moved before it go back here. + self._restore_reconciled_rows(reconciled_from, partition_name, expected_oid, window, key_columns) + raise self._rows_stuck_in_default(parent_name, partition_name, refusal.detail) from refusal except Exception as exc: # Recognised structurally, by the SQLSTATE the driver carries, # rather than by a driver's exception type: a repository built @@ -505,12 +538,17 @@ def _attach_with_reconcile( ) from exc if attempt == DEFAULT_CONFLICT_MAX_RETRIES: - logger.exception( + # Nothing left to try: the DEFAULT partition still holds + # rows for this window and this run cannot take them. That + # is a topology finding, not a failure of the run -- the + # other partitions are maintained, and ``partition_data`` + # reports the window instead of ending on the exception. + logger.warning( "Failed to attach after reconciliation retries", extra={"partition_name": partition_name, "attempts": attempt}, ) self._restore_reconciled_rows(reconciled_from, partition_name, expected_oid, window, key_columns) - raise + raise self._rows_stuck_in_default(parent_name, partition_name, describe_exception(exc)) from exc default_partition = self._metadata.get_default_partition(parent_name) if not default_partition: @@ -552,18 +590,29 @@ def _attach_with_reconcile( except RowMoveRefusedError as refusal: # Rows of DEFAULT belong to the new partition but cannot be # moved safely; the partition stays out, the run goes on. - raise PartitionTopologyError( - parent_name, - FindingReason.DEFAULT_HOLDS_ROWS.value, - f"{parent_name} cannot gain {partition_name!r} while its DEFAULT partition holds rows that " - f"belong to it, and they cannot be moved: {refusal.detail}", - ) from refusal + raise self._rows_stuck_in_default(parent_name, partition_name, refusal.detail) from refusal if moved: reconciled_from = (default_partition.name, default_partition.oid) logger.info("Reconciliation completed", extra={"partition_name": partition_name, "moved_rows": moved}) else: return + @staticmethod + def _rows_stuck_in_default(parent_name: str, partition_name: str, detail: str) -> PartitionTopologyError: + """The DEFAULT partition holds rows for the new partition that this run could not clear. + + A finding about the shape of the tree, not a failure of the run: every + other partition is still maintained, ``apply`` records it in + ``result.issues``, and ``partition_data`` reports the window rather + than ending on the driver's exception. + """ + return PartitionTopologyError( + parent_name, + FindingReason.DEFAULT_HOLDS_ROWS.value, + f"{parent_name} cannot gain {partition_name!r} while its DEFAULT partition holds rows that belong to " + f"it, and this run could not clear them: {detail}", + ) + def _restore_reconciled_rows( self, reconciled_from: tuple[str, int | None] | None, diff --git a/pg_partsmith/sync/services/migration.py b/pg_partsmith/sync/services/migration.py index f720f52..0c22422 100644 --- a/pg_partsmith/sync/services/migration.py +++ b/pg_partsmith/sync/services/migration.py @@ -23,7 +23,7 @@ from pg_partsmith.boundaries import Window from pg_partsmith.constants import DEFAULT_MOVE_BATCH_ROWS from pg_partsmith.entities import MaintenanceIssue, MaintenanceIssueStep, MigrationResult -from pg_partsmith.exceptions import InvalidPartitionConfigError, RowMoveRefusedError +from pg_partsmith.exceptions import InvalidPartitionConfigError, PartitionTopologyError, RowMoveRefusedError from pg_partsmith.lifecycle import DropAfter from pg_partsmith.plan import AttachPartition, CreatePartition, DetachPartition, DropPartition, MaintenancePlan, Reason from pg_partsmith.planner import to_maintenance_issue @@ -62,8 +62,16 @@ def partition_data( Window by window, oldest first: the partition for the oldest window still in DEFAULT is created detached (subtree included), filled from DEFAULT in batches of ``batch_rows``, and attached once DEFAULT holds - nothing more for it. A partition left detached when ``max_batches`` - runs out is picked up and finished by the next call. + nothing more for it. The attach takes whatever arrived while the + batches ran in its own transaction, under the lock it needs anyway, so + the window a live writer is inserting into goes live on the same pass + as the quiet ones. A partition left detached when ``max_batches`` runs + out is picked up and finished by the next call. + + A window that cannot be finished at all -- rows an incoming foreign + key holds down, a DEFAULT partition this run could not clear -- is + reported as a ``move`` issue with ``complete=False``; it is never + raised at the caller. Args: config: The table's configuration; its root must be a RANGE level. @@ -115,7 +123,13 @@ def partition_data( attached = self._executor.create_partition(config, plan, op, issues=tally.issues, fill=fill) else: attached = self._executor.attach_partition(config, plan, op, issues=tally.issues, fill=fill) - except RowMoveRefusedError as exc: + except (RowMoveRefusedError, PartitionTopologyError) as exc: + # A window this run cannot finish -- rows a foreign key holds + # down, a DEFAULT partition it could not clear, a name taken by + # a relation with other bounds. The caller asked for a drain, + # not for an exception: what stayed behind and why is on the + # result, and the loop stops rather than re-planning the same + # window forever. tally.issue( default.name, f"rows for {boundaries.describe(window)} stay in {default.name}: {exc.detail}" ) diff --git a/tests/integration/aio/test_migration.py b/tests/integration/aio/test_migration.py index 682a5d2..24e7965 100644 --- a/tests/integration/aio/test_migration.py +++ b/tests/integration/aio/test_migration.py @@ -7,6 +7,7 @@ from __future__ import annotations +import asyncio from datetime import UTC, datetime from typing import TYPE_CHECKING from unittest.mock import patch @@ -93,6 +94,45 @@ async def _count(engine: AsyncEngine, table: str) -> int: return int(await scalar(engine, f'SELECT count(*) FROM "{table}"')) # noqa: S608 +async def _only_count(engine: AsyncEngine, table: str) -> int: + return int(await scalar(engine, f'SELECT count(*) FROM ONLY "{table}"')) # noqa: S608 + + +# sync-mirror: skip +class _Writer: + """The application the migration happens under: one insert into a window, over and over. + + One statement that never changes, on its own connection, going as fast as + the locks let it. While the window's rows are moved out of DEFAULT these + inserts queue on the move's lock and land the moment it commits, which is + the race the move-and-attach exists to close. + """ + + def __init__(self, engine: AsyncEngine, table: str, *, month: int) -> None: + self._engine = engine + self._sql = text( + f'INSERT INTO "{table}" (created_at, payload) ' # noqa: S608 + f"VALUES (make_timestamptz(2026, {month:d}, 15, 12, 0, 0, 'UTC'), 'live')" + ) + self._stop = asyncio.Event() + self.written = 0 + + async def _write(self) -> None: + async with self._engine.connect() as base_conn: + conn = await base_conn.execution_options(isolation_level="AUTOCOMMIT") + while not self._stop.is_set(): + await conn.execute(self._sql) + self.written += 1 + + async def __aenter__(self) -> _Writer: + self._writing = asyncio.create_task(self._write()) + return self + + async def __aexit__(self, *_exc: object) -> None: + self._stop.set() + await self._writing + + # ── partition_data ────────────────────────────────────────────────────────────── @@ -718,18 +758,11 @@ async def test__ensure_partitions__default_replaced_before_a_failing_attach__row await _default_with_rows(db_engine, events, months=(3,), per_month=4) default = f"{events}_legacy" config = monthly_config(events, create_ahead=1) - original = PostgresPartitionRepository.attach_partition - attempts: list[str] = [] async def failing( self: PostgresPartitionRepository, parent: str, name: str, bounds: object, **kwargs: object - ) -> None: - attempts.append(name) - if len(attempts) == 1: - # The real attach, which PostgreSQL refuses while DEFAULT holds the rows. - await original(self, parent, name, bounds, **kwargs) # type: ignore[arg-type] - return - # The DEFAULT changes hands while its rows are out of it. + ) -> int: + # The DEFAULT changes hands while the bulk of its rows are out of it. async with db_engine.begin() as conn: await conn.execute(text(f'ALTER TABLE "{events}" DETACH PARTITION "{default}"')) await conn.execute(text(f'ALTER TABLE "{default}" RENAME TO "{default}_hijacked"')) @@ -738,7 +771,7 @@ async def failing( raise SQLAlchemyError(msg) # Act / Assert - with patch.object(PostgresPartitionRepository, "attach_partition", failing), pytest.raises(SQLAlchemyError): + with patch.object(PostgresPartitionRepository, "reconcile_and_attach", failing), pytest.raises(SQLAlchemyError): await make_service(db_engine).ensure_partitions(config, [Period(year=2026, month=3)]) # The rows are in the partition whose identity was verified, not in the stranger @@ -1093,3 +1126,59 @@ async def test__move_rows__a_temporary_table_of_the_callers_own__is_left_alone( assert list(kept) == [999] finally: await engine.dispose() + + +# ── attach under a live writer ────────────────────────────────────────────────── +# +# The window the application is inserting into is the one every live migration +# has to attach, and moving its rows out of DEFAULT in one transaction and +# attaching in another cannot do it: the writer refills the window in the gap +# and PostgreSQL refuses the attach every time. Both movers put the last rows +# and the attach under one lock. + + +# sync-mirror: skip +async def test__ensure_partitions__writer_filling_the_window__attaches_on_the_first_pass( + db_engine: AsyncEngine, events: str +) -> None: + """A tick attaches the window a writer keeps inserting into, and loses no row doing it.""" + # Arrange: March's rows in DEFAULT, and an insert going in as fast as the locks allow + default = await _default_with_rows(db_engine, events, months=(3,), per_month=2000) + config = monthly_config(events, create_ahead=1) + march = f"{events}__2026_03" + + # Act + async with _Writer(db_engine, events, month=3) as writer: + created = await make_service(db_engine).ensure_partitions(config, [Period(year=2026, month=3)]) + + # Assert: live on the first pass, with every row in exactly one place + assert writer.written, "the writer never got an insert in" + assert [info.relname for info in created] == [march] + assert await is_attached(db_engine, march) + assert await _only_count(db_engine, default) == 0 + assert await _only_count(db_engine, march) == 2000 + writer.written + assert await _count(db_engine, events) == 2000 + writer.written + + +# sync-mirror: skip +async def test__partition_data__writer_filling_the_window__drains_it_without_raising( + db_engine: AsyncEngine, events: str +) -> None: + """The drain finishes a window under a live writer instead of ending on the attach.""" + # Arrange: enough rows for several batches, and a writer adding more the whole time + default = await _default_with_rows(db_engine, events, months=(3,), per_month=2000) + config = monthly_config(events, create_ahead=1) + march = f"{events}__2026_03" + + # Act + async with _Writer(db_engine, events, month=3) as writer: + result = await make_service(db_engine).partition_data(config, batch_rows=500) + + # Assert + assert result.complete + assert result.issues == () + assert result.partitions == (f"public.{march}",) + assert writer.written, "the writer never got an insert in" + assert await is_attached(db_engine, march) + assert await _only_count(db_engine, default) == 0 + assert await _count(db_engine, events) == 2000 + writer.written diff --git a/tests/integration/sync/test_concurrency.py b/tests/integration/sync/test_concurrency.py index 549408f..7fcd5c2 100644 --- a/tests/integration/sync/test_concurrency.py +++ b/tests/integration/sync/test_concurrency.py @@ -17,7 +17,7 @@ import pytest from sqlalchemy import text -from pg_partsmith.entities import MaintenanceResult +from pg_partsmith.entities import MaintenanceResult, Period from pg_partsmith.exceptions import LockAcquisitionError from pg_partsmith.lifecycle import DetachMode from pg_partsmith.sync.maintainer import PartitionMaintainer @@ -25,7 +25,7 @@ from pg_partsmith.sync.repositories import PostgresPartitionRepository from pg_partsmith.topology import RangeBounds from tests.integration.nested_support import MONTHLY_TABLE_DDL, monthly_config -from tests.integration.sync.support import count_ddl, make_service, make_table +from tests.integration.sync.support import count_ddl, exec_sql, is_attached, make_service, make_table, scalar if TYPE_CHECKING: from collections.abc import Generator @@ -139,3 +139,109 @@ def critical(self, event: str, **kwargs: object) -> None: ... assert len(retry_msgs) >= 1, ( "Expected at least one retry warning; drop succeeded on first attempt without contention" ) + + +# ── attach under a live writer ────────────────────────────────────────────────── +# +# The thread-based twins of the two ``# sync-mirror: skip`` tests in +# ``tests/integration/aio/test_migration.py``: the window the application is +# inserting into is the one every live migration has to attach, and moving its +# rows out of DEFAULT in one transaction and attaching in another cannot do it. + + +def _default_with_rows(engine: Engine, table: str, *, month: int, rows: int) -> str: + """The migration starting point: the old monolithic table attached as DEFAULT, full of rows.""" + default = f"{table}_legacy" + exec_sql(engine, f'CREATE TABLE "{default}" (LIKE "{table}" INCLUDING ALL)') + exec_sql( + engine, + f'INSERT INTO "{default}" (created_at, payload) ' # noqa: S608 + f"SELECT make_timestamptz(2026, :month, 1 + (g % 27), 12, 0, 0, 'UTC'), 'row ' || g " + f"FROM generate_series(1, :rows) g", + month=month, + rows=rows, + ) + exec_sql(engine, f'ALTER TABLE "{table}" ATTACH PARTITION "{default}" DEFAULT') + return default + + +def _only_count(engine: Engine, table: str) -> int: + return int(scalar(engine, f'SELECT count(*) FROM ONLY "{table}"')) # noqa: S608 + + +class _Writer: + """The application the migration happens under: one insert into a window, over and over. + + One statement that never changes, on a thread of its own, going as fast as + the locks let it. While the window's rows are moved out of DEFAULT these + inserts queue on the move's lock and land the moment it commits, which is + the race the move-and-attach exists to close. + """ + + def __init__(self, engine: Engine, table: str, *, month: int) -> None: + self._engine = engine + self._sql = text( + f'INSERT INTO "{table}" (created_at, payload) ' # noqa: S608 + f"VALUES (make_timestamptz(2026, {month:d}, 15, 12, 0, 0, 'UTC'), 'live')" + ) + self._stop = threading.Event() + self._thread = threading.Thread(target=self._write) + self.written = 0 + + def _write(self) -> None: + with self._engine.connect().execution_options(isolation_level="AUTOCOMMIT") as conn: + while not self._stop.is_set(): + conn.execute(self._sql) + self.written += 1 + + def __enter__(self) -> _Writer: + self._thread.start() + return self + + def __exit__(self, *_exc: object) -> None: + self._stop.set() + self._thread.join() + + +def test__ensure_partitions__writer_filling_the_window__attaches_on_the_first_pass( + sync_db_engine: Engine, partitioned_table: str +) -> None: + """A tick attaches the window a writer keeps inserting into, and loses no row doing it.""" + # Arrange: March's rows in DEFAULT, and an insert going in as fast as the locks allow + default = _default_with_rows(sync_db_engine, partitioned_table, month=3, rows=2000) + config = monthly_config(partitioned_table, create_ahead=1) + march = f"{partitioned_table}__2026_03" + + # Act + with _Writer(sync_db_engine, partitioned_table, month=3) as writer: + created = make_service(sync_db_engine).ensure_partitions(config, [Period(year=2026, month=3)]) + + # Assert: live on the first pass, with every row in exactly one place + assert writer.written, "the writer never got an insert in" + assert [info.relname for info in created] == [march] + assert is_attached(sync_db_engine, march) + assert _only_count(sync_db_engine, default) == 0 + assert _only_count(sync_db_engine, march) == 2000 + writer.written + + +def test__partition_data__writer_filling_the_window__drains_it_without_raising( + sync_db_engine: Engine, partitioned_table: str +) -> None: + """The drain finishes a window under a live writer instead of ending on the attach.""" + # Arrange: enough rows for several batches, and a writer adding more the whole time + default = _default_with_rows(sync_db_engine, partitioned_table, month=3, rows=2000) + config = monthly_config(partitioned_table, create_ahead=1) + march = f"{partitioned_table}__2026_03" + + # Act + with _Writer(sync_db_engine, partitioned_table, month=3) as writer: + result = make_service(sync_db_engine).partition_data(config, batch_rows=500) + + # Assert + assert result.complete + assert result.issues == () + assert result.partitions == (f"public.{march}",) + assert writer.written, "the writer never got an insert in" + assert is_attached(sync_db_engine, march) + assert _only_count(sync_db_engine, default) == 0 + assert int(scalar(sync_db_engine, f'SELECT count(*) FROM "{partitioned_table}"')) == 2000 + writer.written # noqa: S608 diff --git a/tests/unit/sync/test_execution.py b/tests/unit/sync/test_execution.py index 6e8ed59..8d2d4b3 100644 --- a/tests/unit/sync/test_execution.py +++ b/tests/unit/sync/test_execution.py @@ -115,6 +115,7 @@ def repo() -> MagicMock: repo.detach_partition = MagicMock(return_value=None) repo.drop_partition = MagicMock(return_value=0) repo.reconcile_default_rows = MagicMock(return_value=0) + repo.reconcile_and_attach = MagicMock(return_value=0) return repo @@ -676,18 +677,18 @@ def test__apply__detached_branch_with_a_different_method__finding_becomes_an_iss # ── attach: DEFAULT reconciliation ────────────────────────────────────────────── -def test__apply__default_conflict_on_range_attach__moves_rows_and_retries( +def test__apply__default_conflict_on_range_attach__moves_the_bulk_then_attaches_under_one_lock( executor: PlanExecutor, repo: MagicMock, metadata: MagicMock ) -> None: # Arrange - repo.attach_partition.side_effect = [_default_conflict(), None] + repo.attach_partition.side_effect = _default_conflict() repo.reconcile_default_rows.return_value = 5 metadata.get_default_partition.return_value = _default_partition() # Act result = executor.apply(_config(), _plan(_create_op())) - # Assert + # Assert -- the bulk moves in its own transaction, the tail and the attach share one metadata.get_default_partition.assert_called_once_with("events") repo.reconcile_default_rows.assert_called_once_with( default_partition_name="events_default", @@ -698,30 +699,41 @@ def test__apply__default_conflict_on_range_attach__moves_rows_and_retries( expected_source_oid=None, expected_target_oid=101, ) - assert repo.attach_partition.call_count == 2 + repo.reconcile_and_attach.assert_called_once_with( + "events", + "events__2024_04", + APRIL, + key_columns=("created_at",), + default_partition_name="events_default", + expected_oid=101, + expected_parent_oid=None, + expected_default_oid=None, + ) + assert repo.attach_partition.call_count == 1 assert result.created_count == 1 assert result.issues == () -def test__apply__default_conflict_retries_exhausted__restores_rows_and_raises( +def test__apply__default_conflict_survives_the_locked_attach__restores_rows_and_records_an_issue( executor: PlanExecutor, repo: MagicMock, metadata: MagicMock ) -> None: - # Arrange + # Arrange -- a DEFAULT partition nothing can clear: a foreign one, which no lock covers repo.attach_partition.side_effect = _default_conflict() + repo.reconcile_and_attach.side_effect = _default_conflict() repo.reconcile_default_rows.return_value = 5 metadata.get_default_partition.return_value = _default_partition() - logger = MagicMock() - # Act / Assert - with patch("pg_partsmith.sync.services.execution.logger", logger), pytest.raises(SQLAlchemyError): - executor.apply(_config(), _plan(_create_op())) + # Act + result = executor.apply(_config(), _plan(_create_op())) - assert repo.attach_partition.call_count == 2 + # Assert -- a topology finding, not an exception out of the run + assert [issue.step for issue in result.issues] == [MaintenanceIssueStep.CREATE] + assert "could not clear them" in result.issues[0].error + assert result.created_count == 0 assert repo.reconcile_default_rows.call_count == 2 restore = repo.reconcile_default_rows.call_args_list[-1].kwargs assert restore["default_partition_name"] == "events__2024_04" assert restore["target_partition_name"] == "events_default" - logger.exception.assert_called_once() def test__apply__default_conflict_with_nothing_moved__nothing_is_restored( @@ -729,13 +741,15 @@ def test__apply__default_conflict_with_nothing_moved__nothing_is_restored( ) -> None: # Arrange repo.attach_partition.side_effect = _default_conflict() + repo.reconcile_and_attach.side_effect = _default_conflict() repo.reconcile_default_rows.return_value = 0 metadata.get_default_partition.return_value = _default_partition() - # Act / Assert - with pytest.raises(SQLAlchemyError): - executor.apply(_config(), _plan(_create_op())) + # Act + result = executor.apply(_config(), _plan(_create_op())) + # Assert + assert "could not clear them" in result.issues[0].error assert repo.reconcile_default_rows.call_count == 1 @@ -845,7 +859,8 @@ def test__apply__attach_conflict_after_reconcile__rows_are_restored_before_the_r executor: PlanExecutor, repo: MagicMock, metadata: MagicMock ) -> None: # Arrange - repo.attach_partition.side_effect = [_default_conflict(), _sqlstate_error("42P07", "duplicate")] + repo.attach_partition.side_effect = _default_conflict() + repo.reconcile_and_attach.side_effect = _sqlstate_error("42P07", "duplicate") repo.reconcile_default_rows.return_value = 3 metadata.get_default_partition.return_value = _default_partition() metadata.is_partition_attached.return_value = True @@ -866,7 +881,8 @@ def test__apply__unrelated_database_error_on_attach__propagates_after_restoring_ executor: PlanExecutor, repo: MagicMock, metadata: MagicMock ) -> None: # Arrange - repo.attach_partition.side_effect = [_default_conflict(), _sqlstate_error("53100", "disk full")] + repo.attach_partition.side_effect = _default_conflict() + repo.reconcile_and_attach.side_effect = _sqlstate_error("53100", "disk full") repo.reconcile_default_rows.return_value = 3 metadata.get_default_partition.return_value = _default_partition() @@ -882,7 +898,8 @@ def test__apply__transport_error_on_attach__propagates_after_restoring_rows( executor: PlanExecutor, repo: MagicMock, metadata: MagicMock, error: BaseException ) -> None: # Arrange - repo.attach_partition.side_effect = [_default_conflict(), error] + repo.attach_partition.side_effect = _default_conflict() + repo.reconcile_and_attach.side_effect = error repo.reconcile_default_rows.return_value = 3 metadata.get_default_partition.return_value = _default_partition() @@ -897,7 +914,8 @@ def test__apply__restore_fails__original_attach_error_still_propagates( executor: PlanExecutor, repo: MagicMock, metadata: MagicMock ) -> None: # Arrange - repo.attach_partition.side_effect = [_default_conflict(), _sqlstate_error("53100", "disk full")] + repo.attach_partition.side_effect = _default_conflict() + repo.reconcile_and_attach.side_effect = _sqlstate_error("53100", "disk full") repo.reconcile_default_rows.side_effect = [3, SQLAlchemyError("restore failed")] metadata.get_default_partition.return_value = _default_partition() logger = MagicMock() @@ -916,7 +934,8 @@ def test__apply__interrupted_while_restoring_rows__propagates_the_interruption( executor: PlanExecutor, repo: MagicMock, metadata: MagicMock ) -> None: # Arrange -- the compensating move-back itself is interrupted - repo.attach_partition.side_effect = [_default_conflict(), _sqlstate_error("53100", "disk full")] + repo.attach_partition.side_effect = _default_conflict() + repo.reconcile_and_attach.side_effect = _sqlstate_error("53100", "disk full") repo.reconcile_default_rows.side_effect = [5, KeyboardInterrupt()] metadata.get_default_partition.return_value = _default_partition() @@ -929,7 +948,8 @@ def test__apply__interrupted_during_attach_after_reconcile__restores_rows_and_re executor: PlanExecutor, repo: MagicMock, metadata: MagicMock ) -> None: # Arrange - repo.attach_partition.side_effect = [_default_conflict(), KeyboardInterrupt()] + repo.attach_partition.side_effect = _default_conflict() + repo.reconcile_and_attach.side_effect = KeyboardInterrupt() repo.reconcile_default_rows.return_value = 5 metadata.get_default_partition.return_value = _default_partition() @@ -971,7 +991,7 @@ def test__apply__driver_error_default_conflict__moves_rows_and_retries( executor: PlanExecutor, repo: MagicMock, metadata: MagicMock ) -> None: # Arrange - repo.attach_partition.side_effect = [_driver_default_conflict(), None] + repo.attach_partition.side_effect = _driver_default_conflict() repo.reconcile_default_rows.return_value = 5 metadata.get_default_partition.return_value = _default_partition() @@ -979,7 +999,8 @@ def test__apply__driver_error_default_conflict__moves_rows_and_retries( result = executor.apply(_config(), _plan(_create_op())) # Assert - assert repo.attach_partition.call_count == 2 + assert repo.attach_partition.call_count == 1 + assert repo.reconcile_and_attach.call_count == 1 assert repo.reconcile_default_rows.call_count == 1 assert result.created_count == 1 assert result.issues == () @@ -989,7 +1010,8 @@ def test__apply__driver_error_on_attach__propagates_after_restoring_rows( executor: PlanExecutor, repo: MagicMock, metadata: MagicMock ) -> None: # Arrange - repo.attach_partition.side_effect = [_driver_default_conflict(), _DriverError("disk full", "53100")] + repo.attach_partition.side_effect = _driver_default_conflict() + repo.reconcile_and_attach.side_effect = _DriverError("disk full", "53100") repo.reconcile_default_rows.return_value = 3 metadata.get_default_partition.return_value = _default_partition() @@ -1009,7 +1031,8 @@ def test__apply__repository_error_without_a_sqlstate__propagates_after_restoring executor: PlanExecutor, repo: MagicMock, metadata: MagicMock ) -> None: # Arrange -- a stub or a caching layer failing on its own terms - repo.attach_partition.side_effect = [_driver_default_conflict(), RuntimeError("stub repository")] + repo.attach_partition.side_effect = _driver_default_conflict() + repo.reconcile_and_attach.side_effect = RuntimeError("stub repository") repo.reconcile_default_rows.return_value = 3 metadata.get_default_partition.return_value = _default_partition() diff --git a/tests/unit/test_execution.py b/tests/unit/test_execution.py index ddb4eff..df05af6 100644 --- a/tests/unit/test_execution.py +++ b/tests/unit/test_execution.py @@ -116,6 +116,7 @@ def repo() -> MagicMock: repo.detach_partition = AsyncMock(return_value=None) repo.drop_partition = AsyncMock(return_value=0) repo.reconcile_default_rows = AsyncMock(return_value=0) + repo.reconcile_and_attach = AsyncMock(return_value=0) return repo @@ -679,18 +680,18 @@ async def test__apply__detached_branch_with_a_different_method__finding_becomes_ # ── attach: DEFAULT reconciliation ────────────────────────────────────────────── -async def test__apply__default_conflict_on_range_attach__moves_rows_and_retries( +async def test__apply__default_conflict_on_range_attach__moves_the_bulk_then_attaches_under_one_lock( executor: PlanExecutor, repo: MagicMock, metadata: MagicMock ) -> None: # Arrange - repo.attach_partition.side_effect = [_default_conflict(), None] + repo.attach_partition.side_effect = _default_conflict() repo.reconcile_default_rows.return_value = 5 metadata.get_default_partition.return_value = _default_partition() # Act result = await executor.apply(_config(), _plan(_create_op())) - # Assert + # Assert -- the bulk moves in its own transaction, the tail and the attach share one metadata.get_default_partition.assert_awaited_once_with("events") repo.reconcile_default_rows.assert_awaited_once_with( default_partition_name="events_default", @@ -701,30 +702,41 @@ async def test__apply__default_conflict_on_range_attach__moves_rows_and_retries( expected_source_oid=None, expected_target_oid=101, ) - assert repo.attach_partition.await_count == 2 + repo.reconcile_and_attach.assert_awaited_once_with( + "events", + "events__2024_04", + APRIL, + key_columns=("created_at",), + default_partition_name="events_default", + expected_oid=101, + expected_parent_oid=None, + expected_default_oid=None, + ) + assert repo.attach_partition.await_count == 1 assert result.created_count == 1 assert result.issues == () -async def test__apply__default_conflict_retries_exhausted__restores_rows_and_raises( +async def test__apply__default_conflict_survives_the_locked_attach__restores_rows_and_records_an_issue( executor: PlanExecutor, repo: MagicMock, metadata: MagicMock ) -> None: - # Arrange + # Arrange -- a DEFAULT partition nothing can clear: a foreign one, which no lock covers repo.attach_partition.side_effect = _default_conflict() + repo.reconcile_and_attach.side_effect = _default_conflict() repo.reconcile_default_rows.return_value = 5 metadata.get_default_partition.return_value = _default_partition() - logger = MagicMock() - # Act / Assert - with patch("pg_partsmith.aio.services.execution.logger", logger), pytest.raises(SQLAlchemyError): - await executor.apply(_config(), _plan(_create_op())) + # Act + result = await executor.apply(_config(), _plan(_create_op())) - assert repo.attach_partition.await_count == 2 + # Assert -- a topology finding, not an exception out of the run + assert [issue.step for issue in result.issues] == [MaintenanceIssueStep.CREATE] + assert "could not clear them" in result.issues[0].error + assert result.created_count == 0 assert repo.reconcile_default_rows.await_count == 2 restore = repo.reconcile_default_rows.call_args_list[-1].kwargs assert restore["default_partition_name"] == "events__2024_04" assert restore["target_partition_name"] == "events_default" - logger.exception.assert_called_once() async def test__apply__default_conflict_with_nothing_moved__nothing_is_restored( @@ -732,13 +744,15 @@ async def test__apply__default_conflict_with_nothing_moved__nothing_is_restored( ) -> None: # Arrange repo.attach_partition.side_effect = _default_conflict() + repo.reconcile_and_attach.side_effect = _default_conflict() repo.reconcile_default_rows.return_value = 0 metadata.get_default_partition.return_value = _default_partition() - # Act / Assert - with pytest.raises(SQLAlchemyError): - await executor.apply(_config(), _plan(_create_op())) + # Act + result = await executor.apply(_config(), _plan(_create_op())) + # Assert + assert "could not clear them" in result.issues[0].error assert repo.reconcile_default_rows.await_count == 1 @@ -848,7 +862,8 @@ async def test__apply__attach_conflict_after_reconcile__rows_are_restored_before executor: PlanExecutor, repo: MagicMock, metadata: MagicMock ) -> None: # Arrange - repo.attach_partition.side_effect = [_default_conflict(), _sqlstate_error("42P07", "duplicate")] + repo.attach_partition.side_effect = _default_conflict() + repo.reconcile_and_attach.side_effect = _sqlstate_error("42P07", "duplicate") repo.reconcile_default_rows.return_value = 3 metadata.get_default_partition.return_value = _default_partition() metadata.is_partition_attached.return_value = True @@ -869,7 +884,8 @@ async def test__apply__unrelated_database_error_on_attach__propagates_after_rest executor: PlanExecutor, repo: MagicMock, metadata: MagicMock ) -> None: # Arrange - repo.attach_partition.side_effect = [_default_conflict(), _sqlstate_error("53100", "disk full")] + repo.attach_partition.side_effect = _default_conflict() + repo.reconcile_and_attach.side_effect = _sqlstate_error("53100", "disk full") repo.reconcile_default_rows.return_value = 3 metadata.get_default_partition.return_value = _default_partition() @@ -885,7 +901,8 @@ async def test__apply__transport_error_on_attach__propagates_after_restoring_row executor: PlanExecutor, repo: MagicMock, metadata: MagicMock, error: BaseException ) -> None: # Arrange - repo.attach_partition.side_effect = [_default_conflict(), error] + repo.attach_partition.side_effect = _default_conflict() + repo.reconcile_and_attach.side_effect = error repo.reconcile_default_rows.return_value = 3 metadata.get_default_partition.return_value = _default_partition() @@ -900,7 +917,8 @@ async def test__apply__restore_fails__original_attach_error_still_propagates( executor: PlanExecutor, repo: MagicMock, metadata: MagicMock ) -> None: # Arrange - repo.attach_partition.side_effect = [_default_conflict(), _sqlstate_error("53100", "disk full")] + repo.attach_partition.side_effect = _default_conflict() + repo.reconcile_and_attach.side_effect = _sqlstate_error("53100", "disk full") repo.reconcile_default_rows.side_effect = [3, SQLAlchemyError("restore failed")] metadata.get_default_partition.return_value = _default_partition() logger = MagicMock() @@ -916,7 +934,8 @@ async def test__apply__cancelled_while_restoring_rows__propagates_the_cancellati executor: PlanExecutor, repo: MagicMock, metadata: MagicMock ) -> None: # Arrange -- the compensating move-back itself is cancelled - repo.attach_partition.side_effect = [_default_conflict(), _sqlstate_error("53100", "disk full")] + repo.attach_partition.side_effect = _default_conflict() + repo.reconcile_and_attach.side_effect = _sqlstate_error("53100", "disk full") repo.reconcile_default_rows.side_effect = [5, asyncio.CancelledError()] metadata.get_default_partition.return_value = _default_partition() @@ -929,7 +948,8 @@ async def test__apply__cancelled_during_attach_after_reconcile__restores_rows_an executor: PlanExecutor, repo: MagicMock, metadata: MagicMock ) -> None: # Arrange - repo.attach_partition.side_effect = [_default_conflict(), asyncio.CancelledError()] + repo.attach_partition.side_effect = _default_conflict() + repo.reconcile_and_attach.side_effect = asyncio.CancelledError() repo.reconcile_default_rows.return_value = 5 metadata.get_default_partition.return_value = _default_partition() @@ -971,7 +991,7 @@ async def test__apply__driver_error_default_conflict__moves_rows_and_retries( executor: PlanExecutor, repo: MagicMock, metadata: MagicMock ) -> None: # Arrange - repo.attach_partition.side_effect = [_driver_default_conflict(), None] + repo.attach_partition.side_effect = _driver_default_conflict() repo.reconcile_default_rows.return_value = 5 metadata.get_default_partition.return_value = _default_partition() @@ -979,7 +999,8 @@ async def test__apply__driver_error_default_conflict__moves_rows_and_retries( result = await executor.apply(_config(), _plan(_create_op())) # Assert - assert repo.attach_partition.await_count == 2 + assert repo.attach_partition.await_count == 1 + assert repo.reconcile_and_attach.await_count == 1 assert repo.reconcile_default_rows.await_count == 1 assert result.created_count == 1 assert result.issues == () @@ -989,7 +1010,8 @@ async def test__apply__driver_error_on_attach__propagates_after_restoring_rows( executor: PlanExecutor, repo: MagicMock, metadata: MagicMock ) -> None: # Arrange - repo.attach_partition.side_effect = [_driver_default_conflict(), _DriverError("disk full", "53100")] + repo.attach_partition.side_effect = _driver_default_conflict() + repo.reconcile_and_attach.side_effect = _DriverError("disk full", "53100") repo.reconcile_default_rows.return_value = 3 metadata.get_default_partition.return_value = _default_partition() @@ -1009,7 +1031,8 @@ async def test__apply__repository_error_without_a_sqlstate__propagates_after_res executor: PlanExecutor, repo: MagicMock, metadata: MagicMock ) -> None: # Arrange -- a stub or a caching layer failing on its own terms - repo.attach_partition.side_effect = [_driver_default_conflict(), RuntimeError("stub repository")] + repo.attach_partition.side_effect = _driver_default_conflict() + repo.reconcile_and_attach.side_effect = RuntimeError("stub repository") repo.reconcile_default_rows.return_value = 3 metadata.get_default_partition.return_value = _default_partition() From 245fc95dba4ab3bed7523efd7da6735c839de88d Mon Sep 17 00:00:00 2001 From: Alexey Shalaev Date: Mon, 7 Sep 2026 14:14:42 +0300 Subject: [PATCH 2/6] test: cover the locked move-and-attach and the branches around it Unit tests for reconcile_and_attach's lock order and statement order in both mirrors, for a tail move an incoming foreign key refuses inside it, and for an attach target swapped during the bulk reconcile; the docs the new lock level changes follow. --- docs/concepts/execution.md | 22 ++++-- docs/guide/partition-existing-table.md | 17 ++--- tests/integration/sync/test_migration.py | 4 ++ tests/unit/sync/test_execution.py | 44 ++++++++++++ tests/unit/sync/test_repository.py | 89 ++++++++++++++++++++++++ tests/unit/test_execution.py | 44 ++++++++++++ tests/unit/test_repository.py | 89 ++++++++++++++++++++++++ 7 files changed, 294 insertions(+), 15 deletions(-) diff --git a/docs/concepts/execution.md b/docs/concepts/execution.md index 487649f..271451a 100644 --- a/docs/concepts/execution.md +++ b/docs/concepts/execution.md @@ -58,17 +58,25 @@ belong to the new window, PostgreSQL refuses the attach (`23514`). The executor: (`ATTACH` matches by name, so physical column order may differ) and leaving rows with a NULL trailing key where PostgreSQL routes them. This is one statement under `SHARE ROW EXCLUSIVE` on DEFAULT: writers wait, readers do not; -2. moves whatever arrived in the meantime and attaches, **in one transaction**, under the - locks `ATTACH` needs anyway — `SHARE UPDATE EXCLUSIVE` on the parent, `ACCESS EXCLUSIVE` - on the partition and on DEFAULT, taken in that order and taken before the move. +2. moves whatever arrived in the meantime and attaches, **in one transaction**, under + `EXCLUSIVE` on the parent and `ACCESS EXCLUSIVE` on the partition and on DEFAULT — taken + in that order, and taken before the move. Step 2 is what makes the window a live writer is inserting into attachable at all. Two transactions cannot do it: the move commits, DEFAULT is free again, and the next insert for the window lands in it before `ATTACH` scans — under any steady write rate that never converges, and every attempt is a full scan of DEFAULT under `ACCESS EXCLUSIVE`. Sharing one -lock closes the gap; a writer waits for the commit and is then routed into the new partition. -Step 1 is what keeps that exclusive window short: the bulk of the month moves before the -heavy lock is taken, so what it covers is a tail and a scan. +lock closes the gap. Step 1 is what keeps that exclusive window short: the bulk of the month +moves before the heavy lock is taken, so what step 2 covers is a tail and a scan. + +The parent's lock is one level above what `ATTACH` itself takes, and it is there for the +writer. An INSERT picks its partition from the set it saw when it took `ROW EXCLUSIVE` on +the parent: one that got that far and then queued on DEFAULT's lock would come out of the +wait still aimed at DEFAULT and be rejected by the constraint the attach just narrowed — a +lost write, and what a plain `ATTACH` does to a live writer. Blocking at the parent instead +means no insert is ever mid-routing while the partition set changes: the writer waits, +re-plans, and lands in the new partition, its statement unchanged. `EXCLUSIVE` does not +conflict with `ACCESS SHARE`, so readers of the other partitions carry on. If the attach still fails, the rows step 1 moved are returned to DEFAULT rather than left in a table no query can see — step 2's own move rolls back with the attach, and needs no @@ -125,7 +133,7 @@ repository). | `CREATE TABLE … (LIKE parent)` | `ACCESS SHARE` on the parent | | `ATTACH PARTITION` | `SHARE UPDATE EXCLUSIVE` on the parent, `ACCESS EXCLUSIVE` on the child and on a DEFAULT sibling; `SHARE ROW EXCLUSIVE` on tables referencing the parent through a foreign key | | the reconciling row move ([DEFAULT reconciliation](#default-reconciliation), step 1) | `SHARE ROW EXCLUSIVE` on the DEFAULT partition and on the child: writers of DEFAULT wait, readers do not | -| the move-and-attach ([DEFAULT reconciliation](#default-reconciliation), step 2) | the `ATTACH` row, taken up front — parent first, so a concurrent `ATTACH` queues there rather than deadlocking | +| the move-and-attach ([DEFAULT reconciliation](#default-reconciliation), step 2) | `EXCLUSIVE` on the parent (one level above `ATTACH`'s own, so no writer is mid-routing), then `ATTACH`'s `ACCESS EXCLUSIVE` on the partition and the DEFAULT sibling — all taken before the move | | `DETACH PARTITION` (plain) | `ACCESS EXCLUSIVE` on parent, partition, and every table referencing the parent | | `DETACH PARTITION … CONCURRENTLY` | `SHARE UPDATE EXCLUSIVE` on the parent; `ACCESS EXCLUSIVE` on the partition and, in its second transaction, on referencing tables | | `DROP TABLE` of a detached table | `ACCESS EXCLUSIVE` on that table only | diff --git a/docs/guide/partition-existing-table.md b/docs/guide/partition-existing-table.md index ff33151..0c49b00 100644 --- a/docs/guide/partition-existing-table.md +++ b/docs/guide/partition-existing-table.md @@ -49,11 +49,11 @@ await maintainer.run_maintenance_safe(config) The tick creates the current month and the two after it. As each is attached, the rows of that month move out of `events_legacy` into it — ordinary DEFAULT reconciliation. The current month is the one the application is writing into, and it attaches on this -first tick like the empty future ones: the bulk of its rows move first, then the rows -that arrived while that ran move *and* the partition goes live in one transaction, under -the lock `ATTACH` takes on the DEFAULT partition anyway. Inserts wait for that commit and -are then routed into the new partition. From now on new rows land in real partitions; the -old ones are still in the DEFAULT. +first tick like the empty future ones: the bulk of its rows move first, then the rows that +arrived while that ran move *and* the partition goes live in one transaction. Inserts wait +at the parent for that commit and are then routed into the new partition, their statement +unchanged. From now on new rows land in real partitions; the old ones are still in the +DEFAULT. ## 3. Drain the DEFAULT partition @@ -89,9 +89,10 @@ finds it, finishes it and attaches it. a maintenance window, or with small batches during a quiet hour and readers that can tolerate a month's rows appearing a little later. Rows already in real partitions, and rows still in DEFAULT for other windows, stay visible throughout. The attach that ends - the window holds `ACCESS EXCLUSIVE` on the DEFAULT partition — readers and writers of - the *undrained* rows wait for it, for the length of the last batch's tail plus one - scan of DEFAULT. + a window blocks writers at the parent and holds `ACCESS EXCLUSIVE` on the DEFAULT + partition, for the length of the last batch's tail plus one scan of DEFAULT — so + inserts stall for that long rather than fail, and readers of the partitions already + drained carry on. `partition_data` takes the table's lock, so it does not race the scheduled tick. It refuses a window it cannot create (an unmanaged partition overlaps it), any move an diff --git a/tests/integration/sync/test_migration.py b/tests/integration/sync/test_migration.py index 9ddb791..77924e8 100644 --- a/tests/integration/sync/test_migration.py +++ b/tests/integration/sync/test_migration.py @@ -85,6 +85,10 @@ def _count(engine: Engine, table: str) -> int: return int(scalar(engine, f'SELECT count(*) FROM "{table}"')) # noqa: S608 +def _only_count(engine: Engine, table: str) -> int: + return int(scalar(engine, f'SELECT count(*) FROM ONLY "{table}"')) # noqa: S608 + + # ── partition_data ────────────────────────────────────────────────────────────── diff --git a/tests/unit/sync/test_execution.py b/tests/unit/sync/test_execution.py index 8d2d4b3..fbdb8b5 100644 --- a/tests/unit/sync/test_execution.py +++ b/tests/unit/sync/test_execution.py @@ -2136,6 +2136,50 @@ def test__apply__reconcile_refused_by_a_foreign_key_action__recorded_and_the_run assert result.created_count == 0 +def test__apply__attach_target_swapped_during_the_bulk_reconcile__fails_closed( + executor: PlanExecutor, repo: MagicMock, metadata: MagicMock +) -> None: + # Arrange -- the relation the rows would be moved into is no longer the planned one + repo.attach_partition.side_effect = _default_conflict() + repo.reconcile_default_rows.side_effect = PlanStaleError("events__2024_04", "the name now resolves to OID 9") + metadata.get_default_partition.return_value = _default_partition() + logger = MagicMock() + + # Act / Assert -- nothing is handed to the replacement, and nothing is restored blindly + with ( + patch("pg_partsmith.sync.services.execution.logger", logger), + pytest.raises(PlanStaleError, match="events__2024_04"), + ): + executor.apply(_config(), _plan(_create_op())) + + repo.reconcile_and_attach.assert_not_called() + assert repo.reconcile_default_rows.call_count == 1 + logger.warning.assert_called_once() + + +def test__apply__tail_move_refused_inside_the_locked_attach__rows_go_back_and_it_is_an_issue( + executor: PlanExecutor, repo: MagicMock, metadata: MagicMock +) -> None: + # Arrange -- a row that became referenced between the bulk move and the attach + repo.attach_partition.side_effect = _default_conflict() + repo.reconcile_default_rows.return_value = 4 + repo.reconcile_and_attach.side_effect = RowMoveRefusedError( + "events_default", "rows are still referenced through a foreign key" + ) + metadata.get_default_partition.return_value = _default_partition() + + # Act + result = executor.apply(_config(), _plan(_create_op())) + + # Assert -- the atomic move rolled back with the attach; the bulk one is put back here + assert [issue.step for issue in result.issues] == [MaintenanceIssueStep.CREATE] + assert "still referenced" in result.issues[0].error + assert result.created_count == 0 + restore = repo.reconcile_default_rows.call_args_list[-1].kwargs + assert restore["default_partition_name"] == "events__2024_04" + assert restore["target_partition_name"] == "events_default" + + def test__create_partition__existing_relation_vanished_before_recovery__is_stale( executor: PlanExecutor, repo: MagicMock, metadata: MagicMock ) -> None: diff --git a/tests/unit/sync/test_repository.py b/tests/unit/sync/test_repository.py index 5805584..86c9c42 100644 --- a/tests/unit/sync/test_repository.py +++ b/tests/unit/sync/test_repository.py @@ -647,6 +647,95 @@ def test__reconcile_default_rows__unknown_rowcount__reads_as_zero() -> None: assert moved == 0 +# ── reconcile_and_attach ──────────────────────────────────────────────────────── + + +def test__reconcile_and_attach__locks_the_parent_and_both_sides_then_moves_and_attaches() -> None: + # Arrange + engine, conn = _engine(_Catalog(moved_rows=7)) + repo = PostgresPartitionRepository(engine) + + # Act + moved = repo.reconcile_and_attach( + "events", + "events__2024_04", + RangeBounds(from_value="2024-04-01", to_value="2024-05-01"), + key_columns=("created_at",), + default_partition_name="events_default", + ) + + # Assert -- the parent's writers out first, then ATTACH's own two, all before the move + assert moved == 7 + assert _locks(conn) == [ + 'LOCK TABLE ONLY "events" IN EXCLUSIVE MODE', + 'LOCK TABLE "events__2024_04" IN ACCESS EXCLUSIVE MODE', + 'LOCK TABLE "events_default" IN ACCESS EXCLUSIVE MODE', + ] + statements = _statements(conn) + assert statements[0] == "SET LOCAL TIME ZONE 'UTC'" + assert statements.index(_locks(conn)[-1]) < statements.index(_move_statement_of(conn)) + assert statements.index(_move_statement_of(conn)) < statements.index(_attach_statement(conn)) + assert _attach_statement(conn) == ( + "ALTER TABLE \"events\" ATTACH PARTITION \"events__2024_04\" FOR VALUES FROM ('2024-04-01') TO ('2024-05-01')" + ) + assert statements[-1] == _MARKER_LOOKUP + + +def test__reconcile_and_attach__composite_key__pads_the_bound_and_keeps_null_keys_in_default() -> None: + # Arrange + engine, conn = _engine(_Catalog(moved_rows=1)) + repo = PostgresPartitionRepository(engine) + + # Act + repo.reconcile_and_attach( + "events", + "events__2024_04", + RangeBounds(from_value="2024-04-01", to_value="2024-05-01"), + key_columns=("created_at", "tenant_id"), + default_partition_name="events_default", + ) + + # Assert + assert _attach_statement(conn).endswith("FOR VALUES FROM ('2024-04-01', MINVALUE) TO ('2024-05-01', MINVALUE)") + assert '"tenant_id" IS NOT NULL' in _move_statement_of(conn) + + +def test__reconcile_and_attach__foreign_default_partition__is_not_locked() -> None: + # Arrange -- LOCK TABLE is refused for a foreign table, so ATTACH's own scan is all there is + engine, conn = _engine(_Catalog(relkind="f")) + repo = PostgresPartitionRepository(engine) + + # Act + repo.reconcile_and_attach( + "events", + "events__2024_04", + RangeBounds(from_value="2024-04-01", to_value="2024-05-01"), + key_columns=("created_at",), + default_partition_name="events_default", + ) + + # Assert + assert _locks(conn) == ['LOCK TABLE ONLY "events" IN EXCLUSIVE MODE'] + + +def test__reconcile_and_attach__empty_key__is_rejected_before_any_sql() -> None: + # Arrange + engine, _ = _engine() + repo = PostgresPartitionRepository(engine) + + # Act / Assert + with pytest.raises(ValueError, match="partition key"): + repo.reconcile_and_attach( + "events", + "events__2024_04", + RangeBounds(from_value="2024-04-01", to_value="2024-05-01"), + key_columns=(), + default_partition_name="events_default", + ) + + engine.begin.assert_not_called() + + # ── detach_partition ──────────────────────────────────────────────────────────── diff --git a/tests/unit/test_execution.py b/tests/unit/test_execution.py index df05af6..cda607a 100644 --- a/tests/unit/test_execution.py +++ b/tests/unit/test_execution.py @@ -2146,6 +2146,50 @@ async def test__apply__reconcile_refused_by_a_foreign_key_action__recorded_and_t assert result.created_count == 0 +async def test__apply__attach_target_swapped_during_the_bulk_reconcile__fails_closed( + executor: PlanExecutor, repo: MagicMock, metadata: MagicMock +) -> None: + # Arrange -- the relation the rows would be moved into is no longer the planned one + repo.attach_partition.side_effect = _default_conflict() + repo.reconcile_default_rows.side_effect = PlanStaleError("events__2024_04", "the name now resolves to OID 9") + metadata.get_default_partition.return_value = _default_partition() + logger = MagicMock() + + # Act / Assert -- nothing is handed to the replacement, and nothing is restored blindly + with ( + patch("pg_partsmith.aio.services.execution.logger", logger), + pytest.raises(PlanStaleError, match="events__2024_04"), + ): + await executor.apply(_config(), _plan(_create_op())) + + repo.reconcile_and_attach.assert_not_awaited() + assert repo.reconcile_default_rows.await_count == 1 + logger.warning.assert_called_once() + + +async def test__apply__tail_move_refused_inside_the_locked_attach__rows_go_back_and_it_is_an_issue( + executor: PlanExecutor, repo: MagicMock, metadata: MagicMock +) -> None: + # Arrange -- a row that became referenced between the bulk move and the attach + repo.attach_partition.side_effect = _default_conflict() + repo.reconcile_default_rows.return_value = 4 + repo.reconcile_and_attach.side_effect = RowMoveRefusedError( + "events_default", "rows are still referenced through a foreign key" + ) + metadata.get_default_partition.return_value = _default_partition() + + # Act + result = await executor.apply(_config(), _plan(_create_op())) + + # Assert -- the atomic move rolled back with the attach; the bulk one is put back here + assert [issue.step for issue in result.issues] == [MaintenanceIssueStep.CREATE] + assert "still referenced" in result.issues[0].error + assert result.created_count == 0 + restore = repo.reconcile_default_rows.call_args_list[-1].kwargs + assert restore["default_partition_name"] == "events__2024_04" + assert restore["target_partition_name"] == "events_default" + + async def test__create_partition__existing_relation_vanished_before_recovery__is_stale( executor: PlanExecutor, repo: MagicMock, metadata: MagicMock ) -> None: diff --git a/tests/unit/test_repository.py b/tests/unit/test_repository.py index 12a00bf..e6e1a24 100644 --- a/tests/unit/test_repository.py +++ b/tests/unit/test_repository.py @@ -645,6 +645,95 @@ async def test__reconcile_default_rows__unknown_rowcount__reads_as_zero() -> Non assert moved == 0 +# ── reconcile_and_attach ──────────────────────────────────────────────────────── + + +async def test__reconcile_and_attach__locks_the_parent_and_both_sides_then_moves_and_attaches() -> None: + # Arrange + engine, conn = _engine(_Catalog(moved_rows=7)) + repo = PostgresPartitionRepository(engine) + + # Act + moved = await repo.reconcile_and_attach( + "events", + "events__2024_04", + RangeBounds(from_value="2024-04-01", to_value="2024-05-01"), + key_columns=("created_at",), + default_partition_name="events_default", + ) + + # Assert -- the parent's writers out first, then ATTACH's own two, all before the move + assert moved == 7 + assert _locks(conn) == [ + 'LOCK TABLE ONLY "events" IN EXCLUSIVE MODE', + 'LOCK TABLE "events__2024_04" IN ACCESS EXCLUSIVE MODE', + 'LOCK TABLE "events_default" IN ACCESS EXCLUSIVE MODE', + ] + statements = _statements(conn) + assert statements[0] == "SET LOCAL TIME ZONE 'UTC'" + assert statements.index(_locks(conn)[-1]) < statements.index(_move_statement_of(conn)) + assert statements.index(_move_statement_of(conn)) < statements.index(_attach_statement(conn)) + assert _attach_statement(conn) == ( + "ALTER TABLE \"events\" ATTACH PARTITION \"events__2024_04\" FOR VALUES FROM ('2024-04-01') TO ('2024-05-01')" + ) + assert statements[-1] == _MARKER_LOOKUP + + +async def test__reconcile_and_attach__composite_key__pads_the_bound_and_keeps_null_keys_in_default() -> None: + # Arrange + engine, conn = _engine(_Catalog(moved_rows=1)) + repo = PostgresPartitionRepository(engine) + + # Act + await repo.reconcile_and_attach( + "events", + "events__2024_04", + RangeBounds(from_value="2024-04-01", to_value="2024-05-01"), + key_columns=("created_at", "tenant_id"), + default_partition_name="events_default", + ) + + # Assert + assert _attach_statement(conn).endswith("FOR VALUES FROM ('2024-04-01', MINVALUE) TO ('2024-05-01', MINVALUE)") + assert '"tenant_id" IS NOT NULL' in _move_statement_of(conn) + + +async def test__reconcile_and_attach__foreign_default_partition__is_not_locked() -> None: + # Arrange -- LOCK TABLE is refused for a foreign table, so ATTACH's own scan is all there is + engine, conn = _engine(_Catalog(relkind="f")) + repo = PostgresPartitionRepository(engine) + + # Act + await repo.reconcile_and_attach( + "events", + "events__2024_04", + RangeBounds(from_value="2024-04-01", to_value="2024-05-01"), + key_columns=("created_at",), + default_partition_name="events_default", + ) + + # Assert + assert _locks(conn) == ['LOCK TABLE ONLY "events" IN EXCLUSIVE MODE'] + + +async def test__reconcile_and_attach__empty_key__is_rejected_before_any_sql() -> None: + # Arrange + engine, _ = _engine() + repo = PostgresPartitionRepository(engine) + + # Act / Assert + with pytest.raises(ValueError, match="partition key"): + await repo.reconcile_and_attach( + "events", + "events__2024_04", + RangeBounds(from_value="2024-04-01", to_value="2024-05-01"), + key_columns=(), + default_partition_name="events_default", + ) + + engine.begin.assert_not_called() + + # ── detach_partition ──────────────────────────────────────────────────────────── From d7b88a1679a244b45051142e2812a180006e87ea Mon Sep 17 00:00:00 2001 From: Alexey Shalaev Date: Mon, 7 Sep 2026 14:17:02 +0300 Subject: [PATCH 3/6] docs: the legacy table still owns the serial sequence when the drain ends Following the guide's last step verbatim on a BIGSERIAL table fails: the sequence is owned by the column it was declared on, so it belongs to the renamed legacy table while every partition's default draws from it, and DROP TABLE is refused. Measured on 17, with what CASCADE does instead. --- docs/guide/partition-existing-table.md | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/docs/guide/partition-existing-table.md b/docs/guide/partition-existing-table.md index 0c49b00..5a4cba0 100644 --- a/docs/guide/partition-existing-table.md +++ b/docs/guide/partition-existing-table.md @@ -116,6 +116,21 @@ DROP TABLE events_legacy; Without a DEFAULT partition, `DETACH … CONCURRENTLY` becomes available, which is what `DetachMode.AUTO` prefers. +!!! note "A `serial` key's sequence still belongs to the old table" + `BIGSERIAL` makes the sequence **owned by** the column it was declared on, so after the + swap it belongs to `events_legacy` — while every partition's `id` default draws from it. + `DROP TABLE events_legacy` is refused for exactly that (`2BP01`, *other objects depend + on it*, listing one default per partition). Move the ownership to the live parent first: + + ```sql + ALTER SEQUENCE events_id_seq OWNED BY events.id; + ``` + + Then the drop goes through and the ids carry on where they were. Do **not** reach for + `DROP TABLE … CASCADE` here: measured on 17, it takes the sequence with the table and + every default that drew from it, and the next insert fails + `null value in column "id" … violates not-null constraint`. + ## The way back `unpartition` empties every partition into one plain table, oldest first, in the same From e4875cf770a4119d649f8fea04f3fde26efd497b Mon Sep 17 00:00:00 2001 From: Alexey Shalaev Date: Mon, 7 Sep 2026 14:18:48 +0300 Subject: [PATCH 4/6] docs: what a queued insert does when the partition set changes under it The measurement behind the move-and-attach's EXCLUSIVE lock on the parent, where the semantics page keeps the rest of the ATTACH lock findings. --- docs/design/postgresql-semantics.md | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/docs/design/postgresql-semantics.md b/docs/design/postgresql-semantics.md index 138e34b..4b59eea 100644 --- a/docs/design/postgresql-semantics.md +++ b/docs/design/postgresql-semantics.md @@ -74,6 +74,24 @@ wait is itself a period of rejected writes for that partition. Hence: create standalone with `LIKE`, attach last; never `CREATE TABLE … PARTITION OF` against a live parent; a converged tree must issue no DDL at all. +### A queued insert keeps the partition it chose (measured on 17, 2026-09-07) + +An `INSERT` through the parent picks its target from the partition set it saw when it took +`ROW EXCLUSIVE` on the parent. `ATTACH` takes only `SHARE UPDATE EXCLUSIVE` there, which does +not conflict — so an insert can route to the DEFAULT partition, queue on the `ACCESS +EXCLUSIVE` the attach holds over it, and come out of that wait still aimed at DEFAULT. The +row is then rejected by the constraint the attach has just narrowed: +`23514 new row for relation "events_legacy" violates partition constraint`. Nothing re-routes +it; the write is lost to the caller. + +Taking `EXCLUSIVE` on the parent instead — one level up, conflicting with `ROW EXCLUSIVE` but +not with `ACCESS SHARE` — makes the insert wait *before* it chooses: it re-plans against the +tree the attach left and lands in the new partition, its statement unchanged, while readers +of the other partitions carry on. That is what the move-and-attach of +[DEFAULT reconciliation](../concepts/execution.md#default-reconciliation) does; the +integration suite asserts it on PostgreSQL 15 through 18 with a writer inserting into the +window while the attach runs. + ## Foreign tables - `CREATE FOREIGN TABLE … PARTITION OF parent` and `ATTACH PARTITION` of a foreign table are From 33ab2fbefc39a53307f917797e8e4a52a1a15e6f Mon Sep 17 00:00:00 2001 From: Alexey Shalaev Date: Mon, 7 Sep 2026 14:19:32 +0300 Subject: [PATCH 5/6] docs: say which lock the move-and-attach takes, not that ATTACH takes it --- docs/guide/partition-existing-table.md | 6 +++--- pg_partsmith/aio/services/migration.py | 6 +++--- pg_partsmith/sync/services/migration.py | 6 +++--- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/docs/guide/partition-existing-table.md b/docs/guide/partition-existing-table.md index 5a4cba0..282d8f2 100644 --- a/docs/guide/partition-existing-table.md +++ b/docs/guide/partition-existing-table.md @@ -71,9 +71,9 @@ Each call: `DELETE … RETURNING` / `INSERT` per batch, each committing on its own, so a row is in exactly one place at every commit point; 4. takes whatever landed during the last batch and attaches the partition, in one - transaction under the lock `ATTACH` needs anyway — so a window still being written to - is attached on the same pass as a quiet one, and the exclusive lock covers a tail and a - scan rather than a month; + transaction and under one lock — so a window still being written to is attached on the + same pass as a quiet one, and what the lock covers is a tail and a scan rather than a + month; 5. moves on to the next window, until DEFAULT holds only rows no window can take (rows with a NULL key), or `max_batches` is spent. diff --git a/pg_partsmith/aio/services/migration.py b/pg_partsmith/aio/services/migration.py index 1921a2b..410c9a8 100644 --- a/pg_partsmith/aio/services/migration.py +++ b/pg_partsmith/aio/services/migration.py @@ -63,9 +63,9 @@ async def partition_data( still in DEFAULT is created detached (subtree included), filled from DEFAULT in batches of ``batch_rows``, and attached once DEFAULT holds nothing more for it. The attach takes whatever arrived while the - batches ran in its own transaction, under the lock it needs anyway, so - the window a live writer is inserting into goes live on the same pass - as the quiet ones. A partition left detached when ``max_batches`` runs + batches ran in its own transaction and under one lock, so the window a + live writer is inserting into goes live on the same pass as the quiet + ones. A partition left detached when ``max_batches`` runs out is picked up and finished by the next call. A window that cannot be finished at all -- rows an incoming foreign diff --git a/pg_partsmith/sync/services/migration.py b/pg_partsmith/sync/services/migration.py index 0c22422..9bfa99d 100644 --- a/pg_partsmith/sync/services/migration.py +++ b/pg_partsmith/sync/services/migration.py @@ -63,9 +63,9 @@ def partition_data( still in DEFAULT is created detached (subtree included), filled from DEFAULT in batches of ``batch_rows``, and attached once DEFAULT holds nothing more for it. The attach takes whatever arrived while the - batches ran in its own transaction, under the lock it needs anyway, so - the window a live writer is inserting into goes live on the same pass - as the quiet ones. A partition left detached when ``max_batches`` runs + batches ran in its own transaction and under one lock, so the window a + live writer is inserting into goes live on the same pass as the quiet + ones. A partition left detached when ``max_batches`` runs out is picked up and finished by the next call. A window that cannot be finished at all -- rows an incoming foreign From 53990e2849ac72759313e996804246b876b0a5d7 Mon Sep 17 00:00:00 2001 From: Alexey Shalaev Date: Mon, 7 Sep 2026 14:20:35 +0300 Subject: [PATCH 6/6] test: let the attach race a writer that is already writing The block was entered on a task that may not have had the loop yet, so on a slow runner the attach could finish before the first insert went out. --- tests/integration/aio/test_migration.py | 5 +++++ tests/integration/sync/test_concurrency.py | 5 +++++ 2 files changed, 10 insertions(+) diff --git a/tests/integration/aio/test_migration.py b/tests/integration/aio/test_migration.py index 24e7965..26cb90f 100644 --- a/tests/integration/aio/test_migration.py +++ b/tests/integration/aio/test_migration.py @@ -115,6 +115,7 @@ def __init__(self, engine: AsyncEngine, table: str, *, month: int) -> None: f"VALUES (make_timestamptz(2026, {month:d}, 15, 12, 0, 0, 'UTC'), 'live')" ) self._stop = asyncio.Event() + self._started = asyncio.Event() self.written = 0 async def _write(self) -> None: @@ -123,9 +124,13 @@ async def _write(self) -> None: while not self._stop.is_set(): await conn.execute(self._sql) self.written += 1 + self._started.set() async def __aenter__(self) -> _Writer: self._writing = asyncio.create_task(self._write()) + # The block must run against a writer that is already writing, not one + # that may not have had the loop yet. + await self._started.wait() return self async def __aexit__(self, *_exc: object) -> None: diff --git a/tests/integration/sync/test_concurrency.py b/tests/integration/sync/test_concurrency.py index 7fcd5c2..767488e 100644 --- a/tests/integration/sync/test_concurrency.py +++ b/tests/integration/sync/test_concurrency.py @@ -185,6 +185,7 @@ def __init__(self, engine: Engine, table: str, *, month: int) -> None: f"VALUES (make_timestamptz(2026, {month:d}, 15, 12, 0, 0, 'UTC'), 'live')" ) self._stop = threading.Event() + self._started = threading.Event() self._thread = threading.Thread(target=self._write) self.written = 0 @@ -193,9 +194,13 @@ def _write(self) -> None: while not self._stop.is_set(): conn.execute(self._sql) self.written += 1 + self._started.set() def __enter__(self) -> _Writer: self._thread.start() + # The block must run against a writer that is already writing, not one + # that may not have been scheduled yet. + assert self._started.wait(timeout=30), "the writer thread never got an insert in" return self def __exit__(self, *_exc: object) -> None: