Add 'force_logical_auto_identity' shard transfer mode for non-blocking moves of tables without a replica identity by onurctirtir · Pull Request #8828 · citusdata/citus

Onur Tirtir (onurctirtir) · 2026-09-07T13:35:29Z

DESCRIPTION: Adds force_logical_auto_identity value to citus.shard_transfer_mode so tables that have no usable replica identity can still be moved / rebalanced with a non-blocking (logical replication) shard transfer.

Closes #6860.
(See the end of PR description)

Citus performs non-blocking shard moves with logical replication. Logical replication has to publish the user's concurrent UPDATE/DELETE on the source shard, which PostgreSQL only allows when the table can identify a row for those commands: a PRIMARY KEY, a replica-identity unique index, or REPLICA IDENTITY FULL. For a table that has none of these, force_logical does not reject the move, but every concurrent UPDATE/DELETE the user issues against the source shard fails for the duration of the transfer, because PostgreSQL refuses those commands on a table with no replica identity that is part of a publication publishing them.

force_logical_auto_identity is a new opt-in value for citus.shard_transfer_mode (and for the shard_transfer_mode argument of citus_move_shard_placement() and the rebalancer, including node draining). It is not the default: with any other mode the behavior is identical to before. When it is used, for every colocated table that has no usable replica identity Citus does the following, per shard and for the duration of the move only:

Temporarily sets REPLICA IDENTITY FULL on the source shard before creating the publication, so PostgreSQL lets the publisher replicate the user's concurrent UPDATE/DELETE. The ALTER needs an AccessExclusiveLock on the shard, so it is attempted up to three times with an increasing lock_timeout (1s, 2s, 4s).
Cost / trade-off: REPLICA IDENTITY FULL amplifies WAL on the source shard. This because, once the shard is set to REPLICA IDENTITY FULL, every UPDATE and every DELETE additionally writes the entire old row into WAL — even when the statement changed only one narrow column.
Restores the shard's replica identity from the distributed (shell) table at cleanup via a cleanup record that is registered before setting REPLICA IDENTITY FULL.
The deferred restore waits only on publications that publish UPDATE or DELETE. Restoring a shard away from REPLICA IDENTITY FULL while such a publication still includes it would break that publication, so cleanup defers. An INSERT-only or TRUNCATE-only publication the user created on the shard does not need a replica identity and therefore doesn't cause deferring the cleanup.
Builds a "replication helper index" on the destination shard early, before the subscriptions are enabled (so before the initial data copy and before the catch-up phase), so the apply worker can locate the row to update/delete with an index scan instead of a per-row sequential scan. If the table already has an index that PostgreSQL can use for the REPLICA IDENTITY FULL tuple lookup, that index is built early and then excluded from the normal post-copy CREATE INDEX phase, so it is created exactly once. An existing index qualifies only when all of the following hold: it is valid and live, it is a btree, it is not partial, it is not an index backing a constraint (PRIMARY KEY / UNIQUE / EXCLUDE), and its leftmost key is a plain, non-generated column (not an expression). If none of the indexes qualify, then a throwaway helper index is built instead, registered for cleanup before it is created, and dropped when the move finishes or fails. This is a speed-only optimization: if no suitable column exists or the index cannot be built, Citus logs a WARNING and falls back to the (correct) sequential scan.
Chooses the column for the throw-away index from the source shard's statistics. The column with the highest estimated number of distinct values among the eligible ones. A column is eligible only when it is not dropped, not a system column, not a generated column, has a default btree operator class, and is fixed-width (pg_attribute.attlen > 0). If the statistics query cannot be run or returns nothing usable, Citus emits a WARNING and falls back to the first eligible column.
- Generated columns are excluded because PostgreSQL refuses using an index covering a generated column for the REPLICA IDENTITY FULL (IsIndexUsableForReplicaIdentityFull) and would silently sequential-scan anyway.
- Variable-width columns (text, bytea, jsonb, unbounded varchar) are excluded because the helper index exists only on the destination: a later large value is accepted on the source but exceeds the btree entry size limit on the subscriber, and the subscription would then retry that change forever and the move would never finish.
A table with a column that has no equality operator is rejected up front. REPLICA IDENTITY FULL makes the subscriber match rows by comparing the whole old tuple, which needs a btree/hash equality operator for every replicated column.

Which code paths / UDFs benefit from `force_logical_auto_identity`?

Note: For tables that already have a usable replica identity, force_logical_auto_identity behaves exactly like force_logical (no extra work).

UDFs that accept a `shard_transfer_mode` argument

UDF	Path	Benefits?	Mechanism
`citus_move_shard_placement` / `citus_move_shard_placement_with_nodeid`	move	✅ Yes	`LookupShardTransferMode` → `TransferShards` (`useAutoIdentityLogicalReplication`)
`citus_copy_shard_placement` / `citus_copy_shard_placement_with_nodeid`	copy	✅ Yes	Same `TransferShards` path; both COPY and MOVE set the flag
`master_copy_shard_placement` (deprecated alias)	copy	✅ Yes	Shares the `citus_copy_shard_placement` body
`citus_internal_copy_single_shard_placement` (internal)	copy	✅ Yes	Same `TransferShards` path
`citus_schema_move` / `citus_schema_move_with_nodeid`	move (wrapper)	✅ Yes	Forwards the mode via `DirectFunctionCall6(citus_move_shard_placement, …)`
`rebalance_table_shards`	rebalance	✅ Yes	Fans the mode out to individual moves
`citus_rebalance_start`	rebalance (background)	✅ Yes	Persists the mode; the background rebalancer moves with it
`replicate_table_shards`	replicate (copy)	✅ Yes	`LookupShardTransferMode` → copies
`citus_drain_node` / `master_drain_node` (deprecated alias)	drain (moves)	✅ Yes	Drain forwards the mode to moves
`replicate_reference_tables`	reference-table copy	✅ Yes	`LookupShardTransferMode` → copy
`citus_split_shard_by_split_points`	split	❌ Errors out	Has its own parser that rejects the mode (`ERRCODE_FEATURE_NOT_SUPPORTED`) — splits don't go through `LogicallyReplicateShards`
`isolate_tenant_to_new_shard` / `citus_isolate_tenant_to_new_shard`	tenant isolation (split-based)	❌ Errors out	Same split parser: "not supported for shard splits or tenant isolation"

Implicit / internal paths (no user-supplied mode)

Note that these paths could benefit from force_logical_auto_identity just as much — they are simply out of scope for this first phase:

Code path	What it does	Benefits?	Why
Reference-table provisioning inside `create_distributed_table`	`EnsureReferenceTablesExistOnAllNodesExtended(TRANSFER_MODE_FORCE_LOGICAL)`	❌ No	Mode is hardcoded to plain `force_logical`; the user's choice never reaches here
Node activation / `citus_add_node`, `create_reference_table`, COPY finalize (`stage_protocol`), `isolate_shards`, `citus_add_local_table_to_metadata`	`EnsureReferenceTablesExistOnAllNodes(void)` → `TRANSFER_MODE_BLOCK_WRITES`	❌ No	This path is `block_writes` by design; it doesn't use logical replication

Scope of `shard_transfer_mode=auto`

We also plan to let auto mode opt into this behavior on its own in a future phase, so eligible tables get non-blocking moves without the user having to pick force_logical_auto_identity explicitly.

…g moves of tables without a replica identity DESCRIPTION: Adds the `force_logical_auto_identity` value to `citus.shard_transfer_mode` so tables that have no usable replica identity can still be moved / rebalanced with a non-blocking (logical replication) shard transfer. Citus performs non-blocking shard moves with logical replication. Logical replication has to publish the user's concurrent `UPDATE`/`DELETE` on the source shard, which PostgreSQL only allows when the table can identify a row for those commands: a `PRIMARY KEY`, a replica-identity unique index, or `REPLICA IDENTITY FULL`. A table that has none of these cannot be moved with `force_logical` — a concurrent `UPDATE`/`DELETE` during the transfer would error out — so today it can only be moved with `block_writes`, which blocks writes for the whole move. That is exactly what a user draining a node wants to avoid. `force_logical_auto_identity` is a new opt-in value for `citus.shard_transfer_mode` (and for the `shard_transfer_mode` argument of `citus_move_shard_placement()` and the rebalancer, including node draining). It is **not** the default: with any other mode the behavior is identical to before. When it is used, for every colocated table that has no usable replica identity Citus does the following, per shard and for the duration of the move only: - **Temporarily sets `REPLICA IDENTITY FULL` on the source shard** before creating the publication, so PostgreSQL lets the publisher replicate the user's concurrent `UPDATE`/`DELETE`. The original replica identity is recorded as a cleanup object and faithfully restored when the move finishes — and also on failure/rollback via the deferred cleanup, so a crashed move never leaves a shard stuck on `REPLICA IDENTITY FULL`. - **Builds a "replication helper index" on the destination shard early**, before the initial data copy, so the apply worker can locate the row to update/delete with an index scan instead of a per-row sequential scan (`REPLICA IDENTITY FULL` otherwise makes the subscriber scan the whole shard for every `UPDATE`/`DELETE` during catch-up). If the table already has a usable btree index, that index is built early and then excluded from the normal post-copy `CREATE INDEX` phase so it is created exactly once. If the table has none, a throwaway helper index is built and dropped when the move finishes. This is a **speed-only** optimization: if no suitable column exists or the index cannot be built, Citus logs a `WARNING` and falls back to the (correct) sequential scan. - **Chooses the helper-index column by name** from the source shard's statistics — the most selective/distinct comparable column — rather than by attribute number, so the choice is stable across shard placements. - **A table with a column that has no equality operator is rejected up front.** `REPLICA IDENTITY FULL` makes the subscriber match rows by comparing the whole old tuple, which needs a btree/hash equality operator for every replicated column. A table that would be set to `FULL` and has, e.g., a `json` column is rejected before the move with a clear error that points at `force_logical` or `block_writes`. Tables that already have a replica identity / primary key are published through that index and are not checked. - **Shard splits and shard isolation in this mode are rejected** with a proper error code. The mode currently relaxes only shard moves / rebalancing / node draining. - **Nothing changes for the other modes.** `auto`, `force_logical`, and `block_writes` behave exactly as before; `force_logical_auto_identity` is purely additive and opt-in.

codecov · 2026-09-07T13:50:21Z

Codecov Report

❌ Patch coverage is 86.53846% with 42 lines in your changes missing coverage. Please review.
✅ Project coverage is 88.80%. Comparing base (f24c6fc) to head (52d93ef).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #8828      +/-   ##
==========================================
- Coverage   88.82%   88.80%   -0.02%     
==========================================
  Files         290      290              
  Lines       65501    65794     +293     
  Branches     8261     8304      +43     
==========================================
+ Hits        58179    58428     +249     
- Misses       4937     4966      +29     
- Partials     2385     2400      +15

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

…ion exists Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Section 8b of shard_move_without_replica_identity exercised the background rebalancer with a whole-cluster citus_rebalance_start(). In the shared enterprise_schedule the cluster is not clean at that point: earlier tests (e.g. multi_create_table) leave append-table placements whose node group has since been removed, and many unrelated tables are imbalanced. A whole-cluster rebalance therefore either errors out ("there is a shard placement in node group N but there are no nodes in that group") or tries to move dozens of unrelated shards over logical replication. The section only passed in an isolated run where those leftovers do not exist. Remove section 8b from the enterprise test (the foreground rebalancer path is still covered by section 8a) and instead cover the background rebalancer's handling of force_logical_auto_identity in the background_rebalance test, which runs against a clean cluster, alongside the existing force_logical and block_writes cases. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

The previous commit added a force_logical_auto_identity BEGIN/ROLLBACK case to background_rebalance. The background rebalancer schedules work asynchronously, and the extra citus_rebalance_start() call perturbs job scheduling enough that the test fails on the first of the flaky runner's repeated runs ("A rebalance is already running") while later repeats pass. Revert background_rebalance to its previous state. Also correct the section 8 comment in shard_move_without_replica_identity to no longer claim the background rebalancer path is covered elsewhere; it simply is not exercised from this shared schedule. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

…al_auto_identity, and cover citus_rebalance_start The per-move shard transfer path already rejects a table that force_logical_auto_identity would set to REPLICA IDENTITY FULL but that has a column with no equality operator (ErrorIfTablesCannotUseReplicaIdentityFull). The rebalancer scheduling path did not run that check for this mode, though: it only ran the analogous VerifyTablesHaveReplicaIdentity check for the auto mode. As a result rebalance_table_shards(shard_transfer_mode => 'force_logical_auto_identity') would fail only once it reached the move, and citus_rebalance_start(shard_transfer_mode => 'force_logical_auto_identity') was worse: it returned successfully, scheduled a background job, and that job then failed and retried with backoff (~30 minutes) before giving up. Run ErrorIfTablesCannotUseReplicaIdentityFull in both the foreground (RebalanceTableShards) and background (RebalanceTableShardsBackground) scheduling paths for force_logical_auto_identity, right next to the existing auto-mode replica-identity check, so the rebalancer rejects such a table up front with a clear error instead of failing mid-transfer or scheduling a doomed background job. To do this ErrorIfTablesCannotUseReplicaIdentityFull is no longer static and is declared in shard_transfer.h. Add coverage to the background_rebalance test (which runs against a clean, two-worker cluster) for citus_rebalance_start with force_logical_auto_identity: one table that cannot be moved (no replica identity plus a json column, so it is rejected synchronously) and one table that can be moved (no replica identity but all columns comparable, so the background rebalance succeeds, the data is preserved, and the source replica identity is restored). The leftover distributed table from earlier in the test is dropped first so the whole-cluster rebalance only has to move the tables created in this section. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

The previous commit added citus_rebalance_start coverage to the existing background_rebalance test. That test manipulates the coordinator, worker_3 and reference tables, and its expected output encodes a very specific cluster state whose first-run vs repeated-run behavior is fragile. Adding any extra rebalance operation to it perturbs that state and makes the flakyness job (which runs the whole file with --repeat 8) fail on the first repeat. Revert the background_rebalance change. Add a new, self-contained test file instead: background_rebalance_no_replica_identity. It runs against a clean two-worker cluster (it does not touch the coordinator, worker_3 or reference tables) and creates and drops its own tables, so a whole-cluster citus_rebalance_start only ever has to move the tables created here and the test is idempotent across repeated runs. It covers, through the background rebalancer: - a table force_logical_auto_identity cannot rescue (no replica identity plus a json column with no equality operator): citus_rebalance_start rejects it up front, before scheduling any background job. - a table force_logical_auto_identity can move (no replica identity but all columns comparable): the background rebalance succeeds, the data is preserved, and the source replica identity is restored to 'd'. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

The files were committed with Windows CRLF line endings, which the style check rejects. Convert them to LF to match the rest of the tree. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

…e background rebalancer When the background rebalancer (citus_rebalance_start) has to copy reference tables to a node that is missing them, it copies them with the same shard transfer mode. For the auto mode it already rejects reference tables without a replica identity up front (VerifyTablesHaveReplicaIdentity). The force_logical_auto_identity branch was missing the analogous check, so a reference table that this mode would set to REPLICA IDENTITY FULL but that cannot then have UPDATE and DELETE replicated (e.g. a column with no equality operator such as json) was not rejected at scheduling time; the reference table copy task was scheduled and only failed later when it executed. Run ErrorIfTablesCannotUseReplicaIdentityFull on the reference table list for force_logical_auto_identity, mirroring the existing auto-mode VerifyTablesHaveReplicaIdentity check, so such a reference table is rejected up front. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Copilot

🟡 Changes recommended

A critical cleanup-safety issue and multiple moderate correctness and coverage issues remain unresolved.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Adds opt-in force_logical_auto_identity mode for non-blocking shard transfers when tables lack usable replica identities.

Changes:

Temporarily applies and restores REPLICA IDENTITY FULL.
Builds destination helper indexes for replication catch-up.
Adds validation, cleanup, upgrade support, and regression coverage.

Review findings:

Critical (1 vote): multi_logical_replication.c:430 may register cleanup for a pre-existing same-named index and later drop that unrelated index.
Moderate (2 votes): node_protocol.c:1026 may select an unusable generated column as an index’s leftmost key.
Moderate (2 votes): shard_transfer.c:1555 skips whole-row equality validation for tables already using REPLICA IDENTITY FULL.
Moderate (2 votes): multi_logical_replication.c:488 may select generated columns for helper indexes, forcing sequential scans.
Moderate (2 votes): The downgrade script does not safely handle pending cleanup records of the new object types.
Moderate (1 vote): The interrupted-move test does not reliably wait until the intended failure window.
Nit (1 vote): Split-path hints recommend a transfer mode that splits reject.

File summaries

File	Description
`src/test/regress/sql/shard_move_without_replica_identity.sql`	Tests transfer modes, cleanup, indexes, and edge cases.
`src/test/regress/sql/background_rebalance_no_replica_identity.sql`	Tests background rebalancing.
`src/test/regress/spec/isolation_logical_replication_no_replica_identity.spec`	Tests concurrent updates and deletes.
`src/test/regress/operations_schedule`	Schedules background tests.
`src/test/regress/expected/shard_rebalancer.out`	Updates mode-hint output.
`src/test/regress/expected/shard_move_without_replica_identity.out`	Adds shard-transfer expectations.
`src/test/regress/expected/multi_tenant_isolation_nonblocking.out`	Updates mode-hint output.
`src/test/regress/expected/multi_replicate_reference_table.out`	Updates reference-copy output.
`src/test/regress/expected/multi_colocated_shard_rebalance.out`	Updates transfer output.
`src/test/regress/expected/isolation_logical_replication_no_replica_identity.out`	Adds concurrency expectations.
`src/test/regress/expected/citus_schema_move.out`	Updates schema-move output.
`src/test/regress/expected/citus_non_blocking_split_shards.out`	Updates split-mode output.
`src/test/regress/expected/background_rebalance.out`	Updates rebalancer output.
`src/test/regress/expected/background_rebalance_no_replica_identity.out`	Adds background-test expectations.
`src/test/regress/enterprise_schedule`	Schedules transfer regression tests.
`src/test/regress/enterprise_isolation_logicalrep_1_schedule`	Schedules isolation coverage.
`src/test/regress/citus_tests/test/test_shard_move_no_replica_identity_failure.py`	Tests interrupted-move cleanup.
`src/test/regress/bin/normalize.sed`	Normalizes variable tuple locations.
`src/include/distributed/shard_transfer.h`	Exposes replica-identity validation.
`src/include/distributed/shard_cleaner.h`	Adds cleanup object types.
`src/include/distributed/multi_logical_replication.h`	Propagates the new mode flag.
`src/include/distributed/coordinator_protocol.h`	Defines mode and helper-index APIs.
`src/backend/distributed/utils/reference_table_utils.c`	Serializes the new mode.
`src/backend/distributed/transaction/worker_transaction.c`	Corrects optional-command return handling.
`src/backend/distributed/sql/downgrades/citus--15.0-1--14.0-1.sql`	Handles downgrade behavior.
`src/backend/distributed/sql/citus--14.0-1--15.0-1.sql`	Adds the enum value.
`src/backend/distributed/replication/multi_logical_replication.c`	Implements temporary identities and helper indexes.
`src/backend/distributed/operations/shard_transfer.c`	Integrates validation and transfer behavior.
`src/backend/distributed/operations/shard_split.c`	Keeps automatic identity disabled for splits.
`src/backend/distributed/operations/shard_rebalancer.c`	Adds rebalancer validation.
`src/backend/distributed/operations/shard_cleaner.c`	Restores identities and removes helper indexes.
`src/backend/distributed/operations/node_protocol.c`	Selects and excludes early-built indexes.
`src/backend/distributed/operations/citus_split_shard_by_split_points.c`	Rejects the mode for split operations.

Review details

Suppressed comments (1)

src/backend/distributed/operations/shard_transfer.c:1527

This shared hint is also emitted by automatic shard splits (shard_split.c:462-464), but LookupSplitMode now rejects force_logical_auto_identity. The resulting split error recommends a mode that immediately errors; make the hint operation-aware or keep this option out of split-path hints.

						errhint("If you wish to continue without a replica "
								"identity set the shard_transfer_mode to "
								"'force_logical_auto_identity', 'force_logical' or "
								"'block_writes'.")));

Files reviewed: 32/33 changed files
Comments generated: 6
Review effort level: Balanced

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Muhammad Usama (codeforall) · 2026-09-14T17:37:54Z

+ * on the source shard, so we bound it with a short lock_timeout (see below) so that
+ * a busy shard fails the transfer fast.
+ */
+static void


I think it's worth to mention in the comment that "REPLICA IDENTITY FULL also generates much more WAL, and may be that also needs to be documented in user facing docs so that users shoild be aware of this when deciding to use this

sure, will mention that in the docs

ChooseReplicationHelperIndexToBuildEarly picks one existing index to build early on the destination shard of a force_logical_auto_identity move, so the subscriber can use an index scan (instead of a sequential scan) while it applies the changes that accumulate during catch-up. It required the leftmost index field to be a plain column (not an expression), but did not require that column to be non-generated. A generated column is not sent by the publisher over logical replication, so it is absent from the subscriber's attribute map, and PostgreSQL refuses an index whose leftmost key is generated for the REPLICA IDENTITY FULL tuple lookup (see IsIndexUsableForReplicaIdentityFull in the PostgreSQL source). Building such an index early therefore provides no catch-up acceleration. Require the leftmost column to be non-generated (attgenerated == '\0'). When the only candidate index has a generated leftmost column, the move now falls back to a throwaway helper index on a plain column, exactly like the expression-index case. Add a regression test (a no-replica-identity table whose only index has a generated leftmost column) that asserts, via log_remote_commands, that the move builds a citus_ri_helper index on a plain column instead of building the generated-column index early. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

…ns in the preflight ErrorIfTablesCannotUseReplicaIdentityFull is the preflight that rejects a force_logical_auto_identity move of a table that would be set to REPLICA IDENTITY FULL but that PostgreSQL then cannot apply UPDATE/DELETE to on the subscriber (e.g. a column with no equality operator such as json). It skipped every table for which RelationCanPublishAllModifications() returns true. That function also returns true for a table the user has already set to REPLICA IDENTITY FULL, so such a table bypassed the check entirely. But an already-FULL table still needs every replicated column to be comparable: the subscriber matches rows by comparing the whole old tuple (tuples_equal). An already-FULL table with, e.g., a json column would therefore pass the preflight and only fail later, when catch-up applies a concurrent UPDATE/DELETE. Skip a table here only when it is a partitioned parent (holds no data of its own) or when it has a usable replica identity / primary key index (GetRelationIdentityOrPK), i.e. when it is not published via REPLICA IDENTITY FULL. This still checks an already-FULL table, while preserving the partitioned-parent skip that RelationCanPublishAllModifications() provided (the leaf partitions are separate entries in the colocated list and are checked on their own). Add a regression test: an already-FULL table with a json column is now rejected up front with the equality-operator error instead of being allowed. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

When a force_logical_auto_identity move needs to build a temporary helper index on the destination shard (the table has no usable existing index), ChooseHelperIndexColumn picks the most selective column that can back a btree index. It skipped dropped and system columns and required a default btree operator class, but it did not skip generated columns. The publisher does not send generated columns over logical replication, so PostgreSQL cannot use a helper index built on one for the REPLICA IDENTITY FULL tuple lookup on the subscriber (see IsIndexUsableForReplicaIdentityFull in the PostgreSQL source); the subscriber silently falls back to a sequential scan, so the CREATE INDEX / DROP INDEX of the helper is wasted. If statistics rank a generated column highest, it was chosen. Exclude generated columns (attgenerated != '\0') from the helper-index candidates. Add a regression test: a no-index table whose generated column has the lowest attribute number (so it was chosen before this fix) now builds the helper index on a plain column instead, verified via log_remote_commands. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

PrepareReplicaIdentitiesForPublication looked up the source worker with FindWorkerNode and immediately dereferenced it (worker->groupId) when registering the replica-identity cleanup record. FindWorkerNode can return NULL (e.g. if the node was concurrently removed), which would crash. Use FindWorkerNodeOrError so we fail with a clear error instead. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

…cally test_auto_identity_move_failure_restores_source waited for the source shard to reach REPLICA IDENTITY FULL and then assumed the move was parked on the advisory lock. That ALTER commits at the very first mutation step, before the publication, replication slot and subscription are created, so the poll could return while the move was still anywhere before the park point. The subsequent "publication exists" assertion and the backend termination were therefore racy. Poll pg_stat_activity until the move backend is actually blocked on the advisory lock (wait_event_type = 'Lock', wait_event = 'advisory') before checking resources and terminating it. At that point the move is guaranteed to be past the REPLICA IDENTITY FULL switch and the publication/slot/subscription creation. Also reconcile the docstring with the lock key. The test claimed to pause "right before the initial copy", but (44000, 55152) is the AFTER-copy lock tag (see isolation_shard_rebalancer_progress.spec); the before-copy tag is (55152, 44000). Switch the lock to (55152, 44000) so the test parks where its docstring says -- right after the source is set to REPLICA IDENTITY FULL and the publication / slot / subscription are created, but before the initial copy. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Tidy up the code that decides which column a replica-identity helper index may be built on, with no behavior change: - node_protocol.c (ChooseReplicationHelperIndexToBuildEarly): shorten the comment explaining why the leftmost column must be non-generated. - shard_transfer.c (ErrorIfTablesCannotUseReplicaIdentityFull): shorten two comments about already-REPLICA-IDENTITY-FULL tables. - multi_logical_replication.c (ChooseHelperIndexColumn): fold the dropped / system-column check into the generated-column check as a single guard and shorten its comment.

The deferred replica-identity cleanup record used to encode the shard's replica identity captured at move time (as "<relreplident_char>:<shard name>"), and cleanup restored the shard to that captured value. That value can go stale. After a failed force_logical_auto_identity move the source shard is left at REPLICA IDENTITY FULL and its cleanup is deferred, holding no lock on the table, so the user can ALTER the distributed table's replica identity in that window. Restoring the value captured at move time would then clobber the shard back to the pre-move setting. Store only the shard id in the cleanup record and, at cleanup time, restore the shard to match its distributed (shell) table's CURRENT replica identity. In Citus a shard's replica identity mirrors its distributed table, and force_logical_auto_identity changes only the shard with DDL propagation off, so the shell table still holds the value the shard should have. The REPLICA IDENTITY USING INDEX case is reconstructed with the shard-qualified index name when the shell table has a usable replica identity index, and falls back to NOTHING only when that index is missing or invalid. A shell table that is itself REPLICA IDENTITY FULL leaves the shard as-is. Tests (shard_move_without_replica_identity): - SECTION 13: after a simulated failed move leaves the shard on FULL, a shell table that is REPLICA IDENTITY USING INDEX has the shard restored to USING INDEX, and the restored shard points at the shard-qualified index. - 4g: rewritten to drive the unusable-index branch through the shell table (index dropped so the shell stays 'i' with no usable index); cleanup restores the shard to NOTHING. - 4d: the synthetic cleanup record now uses the shard-id object name. - SECTION 5: pin citus.shard_replication_factor to 1 so isolate_tenant_to_new_shard reaches the shard transfer mode check.

…anup warning TryResetReplicaIdentityOutsideTransaction emits a WARNING when it restores a shard to REPLICA IDENTITY NOTHING because the shell table replica identity index is missing. The expected output for the t_i_dropped section of shard_move_without_replica_identity was not regenerated after that warning was added, so the test failed. Add the missing warning line to the expected output (no code or behavior change). Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

force_logical_auto_identity moves a table that has no replica identity by temporarily setting REPLICA IDENTITY FULL on the source shard and building a throwaway single-column btree helper index on the destination shard to speed up the logical replication catch-up. ChooseHelperIndexColumn picked the most selective eligible column, which could be a variable-width column such as text, bytea, jsonb or unbounded varchar. That is unsafe. The helper index exists only on the destination. A future INSERT of a large value is accepted on the source but the destination apply worker fails to insert it into the helper btree once it exceeds the btree entry size limit (~1/3 of a page). The subscription then retries that change forever and the move never finishes. Restrict the eligible columns to fixed-width types (pg_attribute.attlen > 0), which can never exceed the btree entry limit, so no future value can stall replication. Variable-width columns are conservatively excluded even when bounded; a table with only variable-width candidate columns simply gets no helper index and the subscriber sequential-scans during catch-up, exactly as when no eligible column exists today. Tests (shard_move_without_replica_identity): - t_unindexable (int + wide text): the wide text column is no longer eligible, so the helper is built on the fixed-width int column and the move succeeds without the previous "index row size exceeds btree maximum" build failure. - t_wide_only (text-only): new section asserting a table whose only column is variable-width gets no helper index (the move emits "no suitable column found ... falling back to a sequential scan"), still succeeds, and preserves a value far larger than the btree entry limit. - t_dropped_column: its candidate column is now int so it remains eligible under the fixed-width rule while still exercising name-based column matching. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

… fast force_logical_auto_identity temporarily sets REPLICA IDENTITY FULL on a source shard that has no replica identity, which needs an AccessExclusiveLock on the shard. PrepareReplicaIdentitiesForPublication ran the ALTER once under a fixed 1s lock_timeout, so a briefly-held conflicting lock (a concurrent read holding an AccessShareLock, or a session left idle in transaction after touching the shard) canceled the ALTER and failed the entire move with a bare "canceling statement due to lock timeout" that did not explain what happened. Try the ALTER up to three times with an increasing lock_timeout (1s, 2s, 4s) using the non-throwing SendOptionalCommandListToWorkerOutsideTransactionWith- Connection, which aborts and resets the remote transaction on failure and is therefore safe to call again on the same connection. A briefly-held lock no longer fails the move, while a persistently-locked shard still fails in bounded time. When every attempt fails, raise a clear error that states the fact (the alter did not succeed after N attempts) and hints at the likely cause (a concurrent session holding a conflicting lock) and that the move can be retried. The cleanup record is still registered before the attempts, so the resource cleanup framework restores the original replica identity on failure. No behavior change on the common path: the first attempt still issues the same 1s lock_timeout, so a move with no lock contention is unaffected. A deterministic regression test would require a concurrent lock holder (isolation/injection harness) and is left out of this change. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Before scheduling or running a rebalance, RebalanceTableShards and RebalanceTableShardsBackground verify replica identity of every table that would be moved: VerifyTablesHaveReplicaIdentity for TRANSFER_MODE_AUTOMATIC and ErrorIfTablesCannotUseReplicaIdentityFull for force_logical_auto_identity. Both looped over placementUpdateList, which holds one entry per shard(-group) move, and called ColocatedTableList plus the per-table check for each entry. All shards of a colocation group return the same colocated table list and the same answer, so for e.g. 64 shard-groups in one colocation group the same tables were opened and checked dozens of times. Collect the tables to check once, visiting each colocation group only once (CollectCitusTableIdsForPlacementUpdates). For each placement update it reads the shard's colocation id and relation id from a single shard id cache lookup (new metadata cache accessor RelationIdAndColocationIdForShard), skips shards whose colocation id was already seen, and otherwise appends that group's colocated table list. The caller then runs its check once over the combined list, which holds each relevant table exactly once. This is applied to both the AUTOMATIC and force_logical_auto_identity branches in both the foreground and background paths. Only distributed table shards appear in placementUpdateList, so every entry has a valid colocation id; reference tables (one shared colocation id) and Citus local tables (INVALID_COLOCATION_ID) are not moved by the rebalancer and never reach this path, so no colocation id needs special handling. Behavior is unchanged: the same errors are raised for the same inputs, only the redundant repeated work is removed, so this performance refactor adds no new regression test. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

…DELETE The deferred replica-identity cleanup would not restore a shard's replica identity while the shard was part of ANY publication. That is too broad: only a publication that publishes UPDATE or DELETE forces the shard to keep a usable replica identity. PostgreSQL rejects an UPDATE/DELETE on a table with no replica identity that is part of a publication publishing that operation, but INSERT and TRUNCATE never need one. So a user's own INSERT-only (or TRUNCATE-only) publication on the physical source shard would block the restore forever: the shard stayed on REPLICA IDENTITY FULL and the pg_dist_cleanup record was retried indefinitely, even though restoring the shard could not break anything. Narrow the guard (renamed RelationHasPublication -> RelationHasUpdateOrDeletePublication) to only defer the restore while a publication that publishes UPDATE or DELETE still includes the shard. This keeps the original safety net for our own move publication, which publishes UPDATE/DELETE and is dropped by an earlier cleanup record, while no longer false-blocking on INSERT/TRUNCATE-only publications. Add test 4d2 (the counterpart of 4d): an INSERT-only publication on the source shard no longer blocks the restore, so cleanup removes the record and restores the shard away from FULL. 4d itself (a default publication, which publishes UPDATE/DELETE) still blocks the restore, unchanged. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

PrepareReplicaIdentitiesForPublication's header comment still said the short lock_timeout makes a busy shard "fail the transfer fast". That is stale: the ALTER ... REPLICA IDENTITY FULL is now retried a few times with an increasing lock_timeout (1s/2s/4s), so a transiently-locked shard no longer fails the transfer; only a shard that stays locked fails it, in bounded time. Update the header to match the retry loop below it. Also drop the normalize.sed rule that masked the btree "index row size exceeds maximum" WARNING (its "references tuple (X,Y) in relation t_unindexable_..." detail). Now that the helper index is only built on fixed-width columns, t_unindexable builds its helper on the fixed-width int column and the move succeeds without that WARNING, so the rule matches nothing and is dead. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

CreateCatchupIndexesForLogicalReplication used to build the early existing index or the temporary replica identity helper index one shard at a time, each over a new connection. Now the per-shard loop only collects the CREATE INDEX commands, and all of them are run in parallel with ExecuteTaskListOutsideTransaction, the same way the late CREATE INDEX phase (ExecuteCreateIndexCommands) already does. The helper index must stay optional: its CREATE INDEX is wrapped in a DO block that turns an error into a WARNING on the destination node, which is propagated to the client. An existing index built early stays mandatory, because the late phase skips it. Add regression tests for both failure paths. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Muhammad Usama (codeforall)

Looks good to me, thanks for the work on this.

I have one minor comment on CollectCitusTableIdsForPlacementUpdates in shard_rebalancer.c, the handling of colocation ID 0 for append/range tables (see the inline comment). Once that's fixed, feel free to merge.

…ncer

Conflicts: src/backend/distributed/sql/citus--14.0-1--15.0-1.sql

Onur Tirtir (onurctirtir) commented Sep 7, 2026

View reviewed changes

Comment thread src/backend/distributed/operations/shard_cleaner.c Outdated

Comment thread src/backend/distributed/operations/shard_cleaner.c Outdated

Comment thread src/backend/distributed/transaction/worker_transaction.c

Onur Tirtir and others added 8 commits September 7, 2026 17:34

Be explicit that a failed publication check assumes publication exists

08ab19c

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Be explicit that an invalid publication check result assumes publicat…

23bd623

…ion exists Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Use LF line endings for the new background rebalancer test files

15ec19b

The files were committed with Windows CRLF line endings, which the style check rejects. Convert them to LF to match the rest of the tree. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Onur Tirtir (onurctirtir) commented Sep 8, 2026

View reviewed changes

Comment thread src/backend/distributed/operations/shard_rebalancer.c Outdated

Onur Tirtir (onurctirtir) and others added 3 commits September 8, 2026 13:36

Apply suggestion from @onurctirtir

291ac16

Merge branch 'main' into force-logical-auto-identity

2ae3305

Merge branch 'main' into force-logical-auto-identity

0d0572c

ibrahim halatci (ihalatci) requested review from Muhammad Usama (codeforall) and a balanced review from Copilot September 9, 2026 08:52

Copilot started reviewing on behalf of ibrahim halatci (ihalatci) September 9, 2026 08:52

View session

Copilot AI reviewed Sep 9, 2026

View reviewed changes

Muhammad Usama (codeforall) reviewed Sep 9, 2026

View reviewed changes

Comment thread src/backend/distributed/replication/multi_logical_replication.c Outdated

Muhammad Usama (codeforall) reviewed Sep 9, 2026

View reviewed changes

Comment thread src/backend/distributed/operations/shard_cleaner.c

Muhammad Usama (codeforall) reviewed Sep 9, 2026

View reviewed changes

Comment thread src/backend/distributed/operations/shard_cleaner.c

Muhammad Usama (codeforall) reviewed Sep 10, 2026

View reviewed changes

Comment thread src/backend/distributed/replication/multi_logical_replication.c Outdated

Muhammad Usama (codeforall) reviewed Sep 10, 2026

View reviewed changes

Comment thread src/backend/distributed/replication/multi_logical_replication.c Outdated

Muhammad Usama (codeforall) reviewed Sep 14, 2026

View reviewed changes

Comment thread src/backend/distributed/operations/shard_rebalancer.c Outdated

Muhammad Usama (codeforall) reviewed Sep 14, 2026

View reviewed changes

Comment thread src/backend/distributed/replication/multi_logical_replication.c Outdated

Muhammad Usama (codeforall) reviewed Sep 14, 2026

View reviewed changes

Onur Tirtir and others added 3 commits September 15, 2026 01:30

Onur Tirtir and others added 2 commits September 15, 2026 01:30

Onur Tirtir (onurctirtir) force-pushed the force-logical-auto-identity branch from fb6a5ca to d2aca5c Compare September 22, 2026 14:03

Onur Tirtir (onurctirtir) and others added 7 commits September 23, 2026 07:53

Onur Tirtir (onurctirtir) force-pushed the force-logical-auto-identity branch from d2aca5c to 6a793e7 Compare September 23, 2026 09:08

Onur Tirtir (onurctirtir) and others added 5 commits September 23, 2026 09:13

reindent

5a26b97

add a test

d18d6d7

Merge branch 'main' into force-logical-auto-identity

e402175

Muhammad Usama (codeforall) reviewed Oct 23, 2026

View reviewed changes

Comment thread src/backend/distributed/operations/shard_rebalancer.c

Muhammad Usama (codeforall) approved these changes Oct 23, 2026

View reviewed changes

Onur Tirtir (onurctirtir) added 3 commits October 27, 2026 12:35

Check replica identity of each non-colocated table up front in rebala…

ed6fa0d

…ncer

Merge remote-tracking branch 'origin/main' into fix8828

82c0351

Conflicts: src/backend/distributed/sql/citus--14.0-1--15.0-1.sql

Document force_logical_auto_identity in the technical README

52d93ef

Onur Tirtir (onurctirtir) enabled auto-merge (squash) October 27, 2026 12:59

Onur Tirtir (onurctirtir) merged commit 1aba449 into main Oct 27, 2026
434 of 439 checks passed

Onur Tirtir (onurctirtir) deleted the force-logical-auto-identity branch October 27, 2026 13:27

Conversation

Onur Tirtir (onurctirtir) commented Sep 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which code paths / UDFs benefit from force_logical_auto_identity?

UDFs that accept a shard_transfer_mode argument

Implicit / internal paths (no user-supplied mode)

Scope of shard_transfer_mode=auto

Uh oh!

codecov Bot commented Sep 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Suppressed comments (1)

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Muhammad Usama (codeforall) Sep 14, 2026

Choose a reason for hiding this comment

Uh oh!

Onur Tirtir (onurctirtir) Sep 23, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Muhammad Usama (codeforall) left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Onur Tirtir (onurctirtir) commented Sep 7, 2026 •

edited

Loading

Which code paths / UDFs benefit from `force_logical_auto_identity`?

UDFs that accept a `shard_transfer_mode` argument

Scope of `shard_transfer_mode=auto`

codecov Bot commented Sep 7, 2026 •

edited

Loading