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..271451a 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,36 @@ 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 + `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. 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 +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 +132,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) | `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 | @@ -117,7 +148,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/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 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..282d8f2 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. 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 @@ -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 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. @@ -80,14 +88,19 @@ 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 + 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), 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 @@ -103,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 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..410c9a8 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 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 + 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..9bfa99d 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 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 + 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..26cb90f 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,50 @@ 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._started = 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 + 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: + self._stop.set() + await self._writing + + # ── partition_data ────────────────────────────────────────────────────────────── @@ -718,18 +763,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 +776,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 +1131,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..767488e 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,114 @@ 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._started = 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 + 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: + 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/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 6e8ed59..fbdb8b5 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() @@ -2113,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 ddb4eff..cda607a 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() @@ -2123,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 ────────────────────────────────────────────────────────────