From a500aaa379efd19a9c7ca809f8986ef250b5cd6a Mon Sep 17 00:00:00 2001
From: meskill <8974488+meskill@users.noreply.github.com>
Date: Fri, 4 Sep 2026 10:10:15 +0000
Subject: [PATCH 1/5] test(resharding): add test for table shape change (omni
 <-> sharded)

---
 integration/resharding/dev.sh      | 46 ++++++++++++++++++++++++++++++
 integration/resharding/pgbench.sql | 26 +++++++++++++++++
 integration/resharding/pgdog.toml  | 12 ++++++++
 integration/resharding/schema.sql  | 16 +++++++++++
 4 files changed, 100 insertions(+)

diff --git a/integration/resharding/dev.sh b/integration/resharding/dev.sh
index ca27fbf82..3c5087bdc 100644
--- a/integration/resharding/dev.sh
+++ b/integration/resharding/dev.sh
@@ -97,6 +97,8 @@ replace_copy_with_replicate projects name
 replace_copy_with_replicate tasks title
 replace_copy_with_replicate task_comments body
 replace_copy_with_replicate settings name
+replace_copy_with_replicate sharded_to_omni name
+replace_copy_with_replicate omni_to_sharded name
 
 # REPLICATION SENTINEL — must be the last DML issued against the source.
 # pgbench uses random(1, 1_000_000_000), so id=0 is reserved for this purpose.
@@ -150,6 +152,8 @@ wait_for_no_copy_rows projects name
 wait_for_no_copy_rows tasks title
 wait_for_no_copy_rows task_comments body
 wait_for_no_copy_rows settings name
+wait_for_no_copy_rows sharded_to_omni name
+wait_for_no_copy_rows omni_to_sharded name
 
 
 # pg_count PORT TABLE — row count via a direct postgres connection (bypasses pgdog).
@@ -193,11 +197,53 @@ check_omni_each_shard() {
     echo "OK omni ${table}: ${source_count} rows on each shard"
 }
 
+check_sharded_source_to_omni_destination() {
+    local table="$1"
+    local source_count dest0_count dest1_count
+
+    source_count=$(psql -d source -tAc "SELECT COUNT(*) FROM ${table}")
+    dest0_count=$(pg_count 15434 "${table}")
+    dest1_count=$(pg_count 15435 "${table}")
+
+    if [ "${source_count}" -ne "${dest0_count}" ] || [ "${source_count}" -ne "${dest1_count}" ]; then
+        echo "MISMATCH sharded->omni ${table}: source=${source_count} dest-0(15434)=${dest0_count} dest-1(15435)=${dest1_count} (expected ${source_count} on each)"
+        exit 1
+    fi
+
+    echo "OK sharded->omni ${table}: ${source_count} rows on each destination shard"
+}
+
+check_omni_source_to_sharded_destination() {
+    local table="$1"
+    local source0_count source1_count dest0_count dest1_count dest_total
+
+    source0_count=$(pg_count 15432 "${table}")
+    source1_count=$(pg_count 15433 "${table}")
+    if [ "${source0_count}" -ne "${source1_count}" ]; then
+        echo "MISMATCH omni->sharded ${table}: source shards disagree source-0(15432)=${source0_count} source-1(15433)=${source1_count}"
+        exit 1
+    fi
+
+    dest0_count=$(pg_count 15434 "${table}")
+    dest1_count=$(pg_count 15435 "${table}")
+    dest_total=$((dest0_count + dest1_count))
+
+    if [ "${source0_count}" -ne "${dest_total}" ]; then
+        echo "MISMATCH omni->sharded ${table}: source=${source0_count} dest total=${dest_total} (dest-0=${dest0_count} dest-1=${dest1_count})"
+        exit 1
+    fi
+
+    echo "OK omni->sharded ${table}: ${source0_count} rows split ${dest0_count}/${dest1_count}"
+}
+
+
 check_row_count_matches tenants
 check_row_count_matches accounts
 check_row_count_matches projects
 check_row_count_matches tasks
 check_row_count_matches task_comments
 check_omni_each_shard settings
+check_sharded_source_to_omni_destination sharded_to_omni
+check_omni_source_to_sharded_destination omni_to_sharded
 
 cleanup
diff --git a/integration/resharding/pgbench.sql b/integration/resharding/pgbench.sql
index c1b066756..59123f8f6 100644
--- a/integration/resharding/pgbench.sql
+++ b/integration/resharding/pgbench.sql
@@ -56,6 +56,24 @@ VALUES (
 )
 ON CONFLICT (id) DO NOTHING;
 
+INSERT INTO sharded_to_omni (id, org_id, name, value)
+VALUES (
+    :id_seed,
+    :id_seed,
+    'sharded-to-omni-' || :id_seed || '-copy',
+    'value-' || :id_seed || '-copy'
+)
+ON CONFLICT (id) DO NOTHING;
+
+INSERT INTO omni_to_sharded (id, org_id, name, value)
+VALUES (
+    :id_seed,
+    :id_seed,
+    'omni-to-sharded-' || :id_seed || '-copy',
+    'value-' || :id_seed || '-copy'
+)
+ON CONFLICT (id) DO NOTHING;
+
 UPDATE tenants
 SET name = 'tenant-' || :id_seed || '-replicate'
 WHERE id = :id_seed;
@@ -79,3 +97,11 @@ WHERE id = :id_seed;
 UPDATE settings
 SET name = 'setting-' || :id_seed || '-replicate'
 WHERE id = :id_seed;
+
+UPDATE sharded_to_omni
+SET name = 'sharded-to-omni-' || :id_seed || '-replicate'
+WHERE id = :id_seed;
+
+UPDATE omni_to_sharded
+SET name = 'omni-to-sharded-' || :id_seed || '-replicate'
+WHERE id = :id_seed;
diff --git a/integration/resharding/pgdog.toml b/integration/resharding/pgdog.toml
index bb5b78d01..fcfc6c0f1 100644
--- a/integration/resharding/pgdog.toml
+++ b/integration/resharding/pgdog.toml
@@ -51,6 +51,18 @@ database = "destination"
 column = "tenant_id"
 data_type = "bigint"
 
+[[sharded_tables]]
+database = "source"
+name = "sharded_to_omni"
+column = "org_id"
+data_type = "bigint"
+
+[[sharded_tables]]
+database = "destination"
+name = "omni_to_sharded"
+column = "org_id"
+data_type = "bigint"
+
 [admin]
 password = "pgdog"
 user = "pgdog"
diff --git a/integration/resharding/schema.sql b/integration/resharding/schema.sql
index 003c3f482..105198d05 100644
--- a/integration/resharding/schema.sql
+++ b/integration/resharding/schema.sql
@@ -62,6 +62,22 @@ CREATE TABLE settings (
     created_at timestamptz NOT NULL DEFAULT NOW()
 );
 
+CREATE TABLE sharded_to_omni (
+    id BIGINT PRIMARY KEY,
+    org_id BIGINT NOT NULL,
+    name VARCHAR NOT NULL,
+    value VARCHAR NOT NULL,
+    created_at timestamptz NOT NULL DEFAULT NOW()
+);
+
+CREATE TABLE omni_to_sharded (
+    id BIGINT PRIMARY KEY,
+    org_id BIGINT NOT NULL,
+    name VARCHAR NOT NULL,
+    value VARCHAR NOT NULL,
+    created_at timestamptz NOT NULL DEFAULT NOW()
+);
+
 CREATE INDEX idx_accounts_tenant_id ON accounts (tenant_id);
 CREATE INDEX idx_projects_tenant_id ON projects (tenant_id);
 CREATE INDEX idx_projects_owner ON projects (tenant_id, owner_account_id);

From 8faa314f9cfff06c1405db2845b3533ddc02e8f6 Mon Sep 17 00:00:00 2001
From: meskill <8974488+meskill@users.noreply.github.com>
Date: Thu, 3 Sep 2026 15:44:08 +0000
Subject: [PATCH 2/5] refactor(resharding): extract tables sync to separate
 function

---
 pgdog/src/backend/pool/cluster.rs             |   5 +-
 .../src/backend/replication/logical/error.rs  |   4 +-
 pgdog/src/backend/replication/logical/mod.rs  |   1 +
 .../replication/logical/publisher/mod.rs      |  55 ++++--
 .../logical/publisher/publisher_impl.rs       | 174 ++---------------
 .../replication/logical/publisher/table.rs    |   8 +-
 .../replication/logical/tables_sync.rs        | 177 ++++++++++++++++++
 pgdog/src/backend/schema/mod.rs               |   2 +-
 pgdog/src/frontend/client/test/test_client.rs |   2 +-
 pgdog/src/frontend/router/parser/context.rs   |   2 +-
 .../router/parser/rewrite/statement/update.rs |  36 ++--
 11 files changed, 272 insertions(+), 194 deletions(-)
 create mode 100644 pgdog/src/backend/replication/logical/tables_sync.rs

diff --git a/pgdog/src/backend/pool/cluster.rs b/pgdog/src/backend/pool/cluster.rs
index 53fd035f0..5f33d36f2 100644
--- a/pgdog/src/backend/pool/cluster.rs
+++ b/pgdog/src/backend/pool/cluster.rs
@@ -10,7 +10,6 @@ use std::{sync::Arc, time::Duration};
 
 use crate::backend::schema::SchemaCache;
 use crate::backend::server::ServerRequest;
-use crate::frontend::router::sharding::ShardedTable;
 use crate::{
     backend::{
         Schema, ShardedTables, databases::User as DatabaseUser, replication::ShardedSchemas,
@@ -509,8 +508,8 @@ impl Cluster {
     }
 
     // Get sharded tables if any.
-    pub(crate) fn sharded_tables(&self) -> &[ShardedTable] {
-        self.sharded_tables.tables()
+    pub(crate) fn sharded_tables(&self) -> &ShardedTables {
+        &self.sharded_tables
     }
 
     /// Get query rewrite config.
diff --git a/pgdog/src/backend/replication/logical/error.rs b/pgdog/src/backend/replication/logical/error.rs
index ca9386db0..dfbcaf76a 100644
--- a/pgdog/src/backend/replication/logical/error.rs
+++ b/pgdog/src/backend/replication/logical/error.rs
@@ -118,8 +118,8 @@ pub(crate) enum Error {
     #[error("replication timeout")]
     ReplicationTimeout,
 
-    #[error("shard {0} has no replication tables")]
-    NoReplicationTables(usize),
+    #[error("publication \"{0}\" has no tables")]
+    EmptyPublication(String),
 
     #[error("shard {0} has no replication slot")]
     NoReplicationSlot(usize),
diff --git a/pgdog/src/backend/replication/logical/mod.rs b/pgdog/src/backend/replication/logical/mod.rs
index 9951fbfd5..d81bb673e 100644
--- a/pgdog/src/backend/replication/logical/mod.rs
+++ b/pgdog/src/backend/replication/logical/mod.rs
@@ -6,6 +6,7 @@ pub(crate) mod publisher;
 pub(crate) mod schema_sync;
 pub(crate) mod status;
 pub(crate) mod subscriber;
+pub(crate) mod tables_sync;
 
 pub(crate) use copy_statement::CopyStatement;
 pub(crate) use error::*;
diff --git a/pgdog/src/backend/replication/logical/publisher/mod.rs b/pgdog/src/backend/replication/logical/publisher/mod.rs
index 539cda2b7..6f6e029e6 100644
--- a/pgdog/src/backend/replication/logical/publisher/mod.rs
+++ b/pgdog/src/backend/replication/logical/publisher/mod.rs
@@ -21,22 +21,19 @@ pub(crate) mod test {
 
     pub(crate) struct PublicationTest {
         pub(crate) server: Server,
+        pub(crate) publication: String,
+        pub(crate) tables: Vec<String>,
     }
 
     impl PublicationTest {
         pub(crate) async fn cleanup(&mut self) {
-            self.server
-                .execute("DROP PUBLICATION IF EXISTS publication_test")
-                .await
-                .unwrap();
-            self.server
-                .execute("DROP TABLE IF EXISTS publication_test_two")
-                .await
-                .unwrap();
-            self.server
-                .execute("DROP TABLE IF EXISTS publication_test_one")
-                .await
-                .unwrap();
+            let drop_publication = format!("DROP PUBLICATION IF EXISTS {}", self.publication);
+            self.server.execute(drop_publication).await.unwrap();
+
+            for table in self.tables.iter().rev() {
+                let drop_table = format!("DROP TABLE IF EXISTS {}", table);
+                self.server.execute(drop_table).await.unwrap();
+            }
         }
     }
 
@@ -69,6 +66,38 @@ pub(crate) mod test {
             .unwrap();
         server.execute("CREATE PUBLICATION publication_test FOR TABLE publication_test_one, publication_test_two").await.unwrap();
 
-        PublicationTest { server }
+        PublicationTest {
+            server,
+            publication: "publication_test".into(),
+            tables: vec!["publication_test_one".into(), "publication_test_two".into()],
+        }
+    }
+
+    pub(crate) async fn setup_publication_tables(
+        publication: &str,
+        tables: &[&str],
+    ) -> PublicationTest {
+        let mut test = PublicationTest {
+            server: test_replication_server().await,
+            publication: publication.to_owned(),
+            tables: tables.iter().map(|table| table.to_string()).collect(),
+        };
+
+        test.cleanup().await;
+
+        for table in &test.tables {
+            let create_table =
+                format!("CREATE TABLE {} (id BIGINT PRIMARY KEY, value TEXT)", table);
+            test.server.execute(create_table).await.unwrap();
+        }
+
+        let create_publication = format!(
+            "CREATE PUBLICATION {} FOR TABLE {}",
+            test.publication,
+            test.tables.join(", ")
+        );
+        test.server.execute(create_publication).await.unwrap();
+
+        test
     }
 }
diff --git a/pgdog/src/backend/replication/logical/publisher/publisher_impl.rs b/pgdog/src/backend/replication/logical/publisher/publisher_impl.rs
index 8be44945c..a3dcaffc8 100644
--- a/pgdog/src/backend/replication/logical/publisher/publisher_impl.rs
+++ b/pgdog/src/backend/replication/logical/publisher/publisher_impl.rs
@@ -18,6 +18,7 @@ use crate::backend::replication::logical::subscriber::omni_ownership::OmniOwners
 use crate::backend::replication::logical::subscriber::stream::StreamSubscriber;
 use crate::backend::replication::publisher::Lsn;
 use crate::backend::replication::publisher::progress::Progress;
+use crate::backend::replication::tables_sync::tables_sync;
 use crate::backend::replication::{
     logical::publisher::ReplicationData, publisher::ParallelSyncManager,
 };
@@ -27,21 +28,6 @@ use crate::net::replication::ReplicationMeta;
 use crate::tasks;
 use crate::util::{safe_interval, safe_sleep};
 
-fn merge_table_lsns(
-    tables: Vec<Table>,
-    existing_lsns: Option<&HashMap<(String, String), Lsn>>,
-) -> Vec<Table> {
-    tables
-        .into_iter()
-        .map(|mut table| {
-            if let Some(lsn) = existing_lsns.and_then(|tables| tables.get(&table.key())) {
-                table.lsn = *lsn;
-            }
-            table
-        })
-        .collect()
-}
-
 #[derive(Debug, Default)]
 pub(crate) struct Publisher {
     /// Name of the publication.
@@ -70,76 +56,31 @@ impl Publisher {
         }
     }
 
-    fn distribute_omnisharded_tables(
-        &mut self,
-        omnisharded: HashMap<(String, String), Table>,
-        source: &Cluster,
-    ) {
-        let shard_count = source.shards().len();
-        // Downstream paths (e.g. Publisher::replicate) iterate every shard and
-        // require a (possibly empty) entry in `self.tables` for each one.
-        for number in 0..shard_count {
-            self.tables.entry(number).or_default();
-        }
-        for (shard_index, table) in omnisharded.into_values().enumerate() {
-            let shard = shard_index % shard_count;
-            if let Some(tables) = self.tables.get_mut(&shard) {
-                tables.push(table);
-            }
-        }
-    }
-
     /// Synchronize tables for all shards.
     pub(crate) async fn sync_tables(
         &mut self,
         data_sync: bool,
         source: &Cluster,
-        dest: &Cluster,
     ) -> Result<(), Error> {
-        let sharding_tables = dest.sharding_schema().tables;
-        let existing_lsns: HashMap<usize, HashMap<(String, String), Lsn>> = self
-            .tables
-            .iter()
-            .map(|(shard, tables)| {
-                (
-                    *shard,
-                    tables
-                        .iter()
-                        .map(|table| (table.key(), table.lsn))
-                        .collect::<HashMap<_, _>>(),
-                )
-            })
-            .collect();
-
-        // Omnisharded tables are split evenly between shards
-        // during copy to avoid duplicate key errors.
-        let mut omnisharded = HashMap::new();
-
-        for (number, shard) in source.shards().iter().enumerate() {
-            // Load tables from publication.
-            let mut primary = shard.primary(&Request::default()).await?;
-            let tables = Table::load(&self.publication, &mut primary).await?;
+        let mut tables = tables_sync(source, source.sharded_tables(), &self.publication).await?;
 
-            // For data sync, split omni tables evenly between shards.
-            if data_sync {
+        if !data_sync {
+            // fill out lsns from the existing tables if any
+            // TODO: make it explicit before running replication after copy_data task
+            for (shard, tables) in &mut tables {
                 for table in tables {
-                    let omni = !table.is_sharded(&sharding_tables);
-                    if omni {
-                        omnisharded.insert(table.key(), table);
-                    } else {
-                        let entry = self.tables.entry(number).or_insert(vec![]);
-                        entry.push(table);
+                    let existing = self
+                        .tables
+                        .get(shard)
+                        .and_then(|tables| tables.iter().find(|t| table.key_ref() == t.key_ref()));
+                    if let Some(existing) = existing {
+                        table.lsn = existing.lsn;
                     }
                 }
-            } else {
-                // For replication, process changes from all shards.
-                let tables = merge_table_lsns(tables, existing_lsns.get(&number));
-                self.tables.insert(number, tables);
             }
         }
 
-        // Distribute omni tables roughly equally between all shards.
-        self.distribute_omnisharded_tables(omnisharded, source);
+        self.tables = tables;
 
         Ok(())
     }
@@ -195,7 +136,7 @@ impl Publisher {
         let stop = CancellationToken::new();
 
         // Synchronize tables from publication.
-        self.sync_tables(false, source, dest).await?;
+        self.sync_tables(false, source).await?;
 
         // Create replication slots if we haven't already.
         if self.slots.is_empty() {
@@ -209,7 +150,8 @@ impl Publisher {
             let tables = self
                 .tables
                 .get(&number)
-                .ok_or(Error::NoReplicationTables(number))?;
+                .map(Vec::as_slice)
+                .unwrap_or_default();
             // Handles the logical replication stream messages.
             // Each subscriber owns a partition of destination shards for omni-table DML
             // (dest_shard % n_sources == source_shard), preventing cross-subscriber deadlocks.
@@ -376,7 +318,7 @@ impl Publisher {
         require_replica_identity: bool,
     ) -> Result<(), Error> {
         // Fetch schema and column metadata first — valid() depends on it.
-        self.sync_tables(true, source, dest).await?;
+        self.sync_tables(true, source).await?;
 
         // Validate replica identity up front, before a potentially multi-hour
         // copy. Only streaming consumes it (to build the per-row UPDATE/DELETE
@@ -406,11 +348,7 @@ impl Publisher {
         let mut handles = FuturesUnordered::new();
 
         for (number, shard) in source.shards().iter().enumerate() {
-            let tables = self
-                .tables
-                .get(&number)
-                .ok_or(Error::NoReplicationTables(number))?
-                .clone();
+            let tables = self.tables.get(&number).cloned().unwrap_or_default();
 
             info!(
                 "table sync starting for {} tables, shard={}",
@@ -529,85 +467,9 @@ impl Waiter {
 #[cfg(test)]
 mod test {
     use super::*;
-    use crate::backend::replication::logical::publisher::{
-        PublicationTable, PublicationTableColumn, ReplicaIdentity,
-    };
     use crate::backend::server::test::test_replication_server;
     use crate::config::config;
 
-    fn make_table(schema: &str, name: &str, lsn: i64) -> Table {
-        Table {
-            publication: "test".to_string(),
-            table: PublicationTable {
-                schema: schema.to_string(),
-                name: name.to_string(),
-                attributes: String::new(),
-                parent_schema: String::new(),
-                parent_name: String::new(),
-            },
-            identity: ReplicaIdentity {
-                oid: pgdog_postgres_types::Oid(1),
-                identity: String::new(),
-                kind: String::new(),
-            },
-            columns: vec![PublicationTableColumn {
-                oid: 1,
-                name: "tenant_id".to_string(),
-                type_oid: pgdog_postgres_types::Oid(20),
-                identity: true,
-            }],
-            lsn: Lsn::from_i64(lsn),
-        }
-    }
-
-    #[test]
-    fn merge_table_lsns_preserves_existing_offsets() {
-        let existing = HashMap::from([(
-            ("copy_data".to_string(), "users".to_string()),
-            Lsn::from_i64(123),
-        )]);
-
-        let merged = merge_table_lsns(vec![make_table("copy_data", "users", 0)], Some(&existing));
-
-        assert_eq!(merged[0].lsn, Lsn::from_i64(123));
-    }
-
-    #[test]
-    fn merge_table_lsns_leaves_unknown_tables_unset() {
-        let existing = HashMap::from([(
-            ("copy_data".to_string(), "users".to_string()),
-            Lsn::from_i64(123),
-        )]);
-
-        let merged = merge_table_lsns(vec![make_table("copy_data", "orders", 0)], Some(&existing));
-
-        assert_eq!(merged[0].lsn, Lsn::default());
-    }
-
-    #[test]
-    fn distribute_omnisharded_tables_initializes_missing_shards() {
-        let config = config();
-        let cluster = Cluster::new_test(&config);
-        let mut publisher = Publisher::new("test", "slot".into());
-        let table = make_table("public", "omni_only", 0);
-
-        publisher.distribute_omnisharded_tables(HashMap::from([(table.key(), table)]), &cluster);
-
-        assert!(
-            publisher.tables.contains_key(&0),
-            "omni-only publications should initialize shard 0 even when no sharded tables exist"
-        );
-        assert!(
-            publisher.tables.contains_key(&1),
-            "data_sync iterates every shard and needs an entry for shard 1 even if it is empty"
-        );
-        assert_eq!(
-            publisher.tables.values().map(Vec::len).sum::<usize>(),
-            1,
-            "the omnisharded table should still be assigned exactly once"
-        );
-    }
-
     /// Tables without a primary key or replica identity index must be rejected
     /// before the copy starts, not after. Validates that `data_sync` returns
     /// `TableValidation` carrying one entry per bad table and leaves no replication slots behind.
diff --git a/pgdog/src/backend/replication/logical/publisher/table.rs b/pgdog/src/backend/replication/logical/publisher/table.rs
index 481b376d2..7a7496716 100644
--- a/pgdog/src/backend/replication/logical/publisher/table.rs
+++ b/pgdog/src/backend/replication/logical/publisher/table.rs
@@ -25,7 +25,7 @@ use tokio_util::sync::CancellationToken;
 
 use tracing::info;
 
-#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+#[derive(Debug, Clone)]
 pub(crate) struct Table {
     /// Name of the table publication.
     pub(crate) publication: String,
@@ -159,10 +159,14 @@ impl Table {
     }
 
     /// Key used for duplicate check.
-    pub(super) fn key(&self) -> (String, String) {
+    pub(crate) fn key(&self) -> (String, String) {
         (self.table.schema.clone(), self.table.name.clone())
     }
 
+    pub(crate) fn key_ref(&self) -> (&String, &String) {
+        (&self.table.schema, &self.table.name)
+    }
+
     /// Check that the table supports replication.
     ///
     /// - FULL (`"f"`): valid — identity comes from `update.old`/`delete.old`, not column metadata.
diff --git a/pgdog/src/backend/replication/logical/tables_sync.rs b/pgdog/src/backend/replication/logical/tables_sync.rs
new file mode 100644
index 000000000..366252768
--- /dev/null
+++ b/pgdog/src/backend/replication/logical/tables_sync.rs
@@ -0,0 +1,177 @@
+use std::collections::{HashMap, HashSet};
+
+use super::Error;
+use super::publisher::Table;
+use crate::backend::pool::Request;
+use crate::backend::{Cluster, ShardedTables};
+
+/// Fetch the info about tables on shard and distribute the set of tables
+/// between the shards.
+///
+/// # Invariants
+///
+/// - the omni-sharded tables should be present on all the shards and should have
+///   identical content, this is due the fact that the omni-sharded table is
+///   copied/replicated only from single shard (random) and only this shard's
+///   data will be used to update the destination shards.
+pub(crate) async fn tables_sync(
+    source: &Cluster,
+    // this should be [`ShardedTables`] configuration from the source cluster
+    sharded_tables: &ShardedTables,
+    publication: &str,
+) -> Result<HashMap<usize, Vec<Table>>, Error> {
+    let mut result: HashMap<usize, Vec<Table>> = HashMap::new();
+
+    let shards_count = source.shards().len();
+    let mut omnisharded = HashSet::new();
+
+    for shard in source.shards() {
+        let mut primary = shard.primary(&Request::default()).await?;
+        let tables = Table::load(publication, &mut primary).await?;
+
+        for table in tables {
+            let shard_index = if table.is_sharded(sharded_tables) {
+                // if table is sharded on source, then we push this table to this shard
+                // so it'll be copied/replicated from this shard changes
+                shard.number()
+            } else {
+                // if the table is omnisharded, then check if we already saw it
+                // and if not, move it to only one shard
+                if !omnisharded.insert(table.key()) {
+                    continue;
+                }
+                (omnisharded.len() - 1) % shards_count
+            };
+
+            result.entry(shard_index).or_default().push(table);
+        }
+    }
+
+    if result.is_empty() {
+        return Err(Error::EmptyPublication(publication.to_owned()));
+    }
+
+    Ok(result)
+}
+
+#[cfg(test)]
+mod test {
+    use crate::backend::replication::logical::publisher::test::{
+        PublicationTest, setup_publication_tables,
+    };
+    use crate::config::{DataType, Hasher, config};
+    use crate::frontend::router::sharding::ShardedTable;
+
+    use super::*;
+
+    fn sharded_tables(names: &[&str]) -> ShardedTables {
+        names
+            .iter()
+            .map(|name| ShardedTable {
+                database: "pgdog".into(),
+                name: Some((*name).into()),
+                column: "id".into(),
+                primary: true,
+                data_type: DataType::Bigint,
+                centroid_probes: 1,
+                hasher: Hasher::Postgres,
+                ..Default::default()
+            })
+            .collect::<Vec<_>>()
+            .as_slice()
+            .into()
+    }
+
+    async fn distribute(publication: &PublicationTest, sharded: &[&str]) -> Vec<Vec<String>> {
+        let source = Cluster::new_test(&config());
+        source.launch();
+
+        let shards = source.shards().len();
+        let distribution = tables_sync(&source, &sharded_tables(sharded), &publication.publication)
+            .await
+            .unwrap();
+
+        source.shutdown();
+
+        (0..shards)
+            .map(|shard| {
+                let mut names = distribution
+                    .get(&shard)
+                    .map(|tables| {
+                        tables
+                            .iter()
+                            .map(|table| table.table.name.clone())
+                            .collect::<Vec<_>>()
+                    })
+                    .unwrap_or_default();
+                names.sort();
+                names
+            })
+            .collect()
+    }
+
+    #[tokio::test]
+    async fn test_sharded_tables_go_to_every_shard() {
+        crate::logger();
+
+        let tables = ["tables_sync_sharded_one", "tables_sync_sharded_two"];
+        let mut publication = setup_publication_tables("tables_sync_sharded_pub", &tables).await;
+
+        let distribution = distribute(&publication, &tables).await;
+
+        assert_eq!(distribution.len(), 2);
+        assert_eq!(distribution[0], tables);
+        assert_eq!(distribution[1], tables);
+
+        publication.cleanup().await;
+    }
+
+    #[tokio::test]
+    async fn test_omnisharded_tables_are_distributed() {
+        crate::logger();
+
+        let tables = [
+            "tables_sync_omni_a",
+            "tables_sync_omni_b",
+            "tables_sync_omni_c",
+            "tables_sync_omni_d",
+        ];
+        let mut publication = setup_publication_tables("tables_sync_omni_pub", &tables).await;
+
+        let distribution = distribute(&publication, &[]).await;
+
+        assert_eq!(distribution.len(), 2);
+        assert_eq!(
+            distribution[0],
+            ["tables_sync_omni_a", "tables_sync_omni_c"]
+        );
+        assert_eq!(
+            distribution[1],
+            ["tables_sync_omni_b", "tables_sync_omni_d"]
+        );
+
+        publication.cleanup().await;
+    }
+
+    #[tokio::test]
+    async fn test_mixed_tables_distribution() {
+        crate::logger();
+
+        let sharded = "tables_sync_mixed_sharded";
+        let omni = [
+            "tables_sync_mixed_omni_a",
+            "tables_sync_mixed_omni_b",
+            "tables_sync_mixed_omni_c",
+        ];
+        let tables = [omni[0], omni[1], omni[2], sharded];
+        let mut publication = setup_publication_tables("tables_sync_mixed_pub", &tables).await;
+
+        let distribution = distribute(&publication, &[sharded]).await;
+
+        assert_eq!(distribution.len(), 2);
+        assert_eq!(distribution[0], [omni[0], omni[2], sharded]);
+        assert_eq!(distribution[1], [omni[1], sharded]);
+
+        publication.cleanup().await;
+    }
+}
diff --git a/pgdog/src/backend/schema/mod.rs b/pgdog/src/backend/schema/mod.rs
index 4e74de636..3263214a6 100644
--- a/pgdog/src/backend/schema/mod.rs
+++ b/pgdog/src/backend/schema/mod.rs
@@ -175,7 +175,7 @@ impl Schema {
     /// Install PgDog-specific functions and triggers.
     pub(crate) async fn install(cluster: &Cluster) -> Result<(), Error> {
         let shards = cluster.shards();
-        let sharded_tables = cluster.sharded_tables();
+        let sharded_tables = cluster.sharded_tables().tables();
 
         if sharded_tables.is_empty() {
             return Ok(());
diff --git a/pgdog/src/frontend/client/test/test_client.rs b/pgdog/src/frontend/client/test/test_client.rs
index 11fb6214e..9612e25cf 100644
--- a/pgdog/src/frontend/client/test/test_client.rs
+++ b/pgdog/src/frontend/client/test/test_client.rs
@@ -275,7 +275,7 @@ impl TestClient {
     pub(crate) fn shard_for_id(&mut self, id: i64) -> Shard {
         let cluster = self.engine.backend().cluster().unwrap();
 
-        ContextBuilder::new(cluster.sharded_tables().first().unwrap())
+        ContextBuilder::new(cluster.sharded_tables().tables().first().unwrap())
             .data(id)
             .shards(cluster.shards().len())
             .build()
diff --git a/pgdog/src/frontend/router/parser/context.rs b/pgdog/src/frontend/router/parser/context.rs
index 994763280..76a16774b 100644
--- a/pgdog/src/frontend/router/parser/context.rs
+++ b/pgdog/src/frontend/router/parser/context.rs
@@ -65,7 +65,7 @@ impl<'a> QueryParserContext<'a> {
         let mut shards_calculator = ShardsWithPriority::default();
         let mut bare_key_lookups = Vec::new();
 
-        let sharded_tables = !router_context.cluster.sharded_tables().is_empty();
+        let sharded_tables = !router_context.cluster.sharded_tables().tables().is_empty();
         let sharding_schema = router_context.cluster.sharding_schema();
 
         router_context.parameter_hints.compute_shard(
diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/update.rs b/pgdog/src/frontend/router/parser/rewrite/statement/update.rs
index c1972c710..bc8e76b27 100644
--- a/pgdog/src/frontend/router/parser/rewrite/statement/update.rs
+++ b/pgdog/src/frontend/router/parser/rewrite/statement/update.rs
@@ -5,6 +5,7 @@ use pg_raw_parse::make::{owned, try_owned};
 use pg_raw_parse::{DeparseResult, Node, NodeMut, Owned, deparse, nodes, walk};
 use pgdog_config::RewriteMode;
 
+use crate::backend::ShardedTables;
 use crate::{
     frontend::{
         BufferedQuery, ClientRequest,
@@ -95,11 +96,11 @@ impl Deref for ShardingKeyUpdate {
 impl ShardingKeyUpdate {
     pub(crate) fn sharded_table<'a>(
         &self,
-        sharded_tables: &'a [ShardedTable],
+        sharded_tables: &'a ShardedTables,
     ) -> Option<&'a ShardedTable> {
         let table = self.target_table();
 
-        sharded_tables.iter().find(|sharded| {
+        sharded_tables.tables().iter().find(|sharded| {
             if let Some(name) = sharded.name.as_ref()
                 && !table.name_match(name)
             {
@@ -493,9 +494,8 @@ mod test {
         assert_eq!(result.delete.params, indexset![2]);
 
         let schema = default_schema();
-        let tables = schema.tables.tables();
         assert_eq!(result.target_table().name, "sharded");
-        assert_eq!(result.sharded_table(tables).unwrap().column, "id");
+        assert_eq!(result.sharded_table(&schema.tables).unwrap().column, "id");
     }
 
     #[test]
@@ -552,23 +552,29 @@ mod test {
             "DELETE FROM sharded WHERE email = $1 RETURNING *"
         );
 
-        assert!(result.sharded_table(&[]).is_none());
+        assert!(result.sharded_table(&ShardedTables::default()).is_none());
         assert!(
             result
-                .sharded_table(&[ShardedTable {
-                    name: Some("other".into()),
-                    column: "id".into(),
-                    ..Default::default()
-                }])
+                .sharded_table(&ShardedTables::from(
+                    [ShardedTable {
+                        name: Some("other".into()),
+                        column: "id".into(),
+                        ..Default::default()
+                    }]
+                    .as_slice()
+                ))
                 .is_none()
         );
         assert!(
             result
-                .sharded_table(&[ShardedTable {
-                    name: Some("sharded".into()),
-                    column: "user_id".into(),
-                    ..Default::default()
-                }])
+                .sharded_table(&ShardedTables::from(
+                    [ShardedTable {
+                        name: Some("sharded".into()),
+                        column: "user_id".into(),
+                        ..Default::default()
+                    }]
+                    .as_slice()
+                ))
                 .is_none()
         );
     }

From 5ac7fc7296e68fcfab1b72446206f68d56585e31 Mon Sep 17 00:00:00 2001
From: meskill <8974488+meskill@users.noreply.github.com>
Date: Fri, 4 Sep 2026 11:53:46 +0000
Subject: [PATCH 3/5] refactor(resharding): drop OmniOwnership

the behaviour is replaced with shared tables_sync implementation that distributes omni tables between source shards at initialization.
---
 integration/resharding/users.toml             |   1 -
 .../logical/publisher/publisher_impl.rs       |  13 +-
 .../replication/logical/subscriber/mod.rs     |   1 -
 .../logical/subscriber/omni_ownership.rs      | 117 ------------------
 .../replication/logical/subscriber/stream.rs  |  20 +--
 .../replication/logical/subscriber/tests.rs   | 102 +++------------
 6 files changed, 29 insertions(+), 225 deletions(-)
 delete mode 100644 pgdog/src/backend/replication/logical/subscriber/omni_ownership.rs

diff --git a/integration/resharding/users.toml b/integration/resharding/users.toml
index fa324f5d2..67142d309 100644
--- a/integration/resharding/users.toml
+++ b/integration/resharding/users.toml
@@ -9,4 +9,3 @@ database = "destination"
 name = "pgdog"
 password = "pgdog"
 schema_admin = true
-lock_timeout = 100
diff --git a/pgdog/src/backend/replication/logical/publisher/publisher_impl.rs b/pgdog/src/backend/replication/logical/publisher/publisher_impl.rs
index a3dcaffc8..d7897669a 100644
--- a/pgdog/src/backend/replication/logical/publisher/publisher_impl.rs
+++ b/pgdog/src/backend/replication/logical/publisher/publisher_impl.rs
@@ -14,7 +14,6 @@ use tracing::{debug, info, warn};
 use super::super::{Error, ensure_validation, publisher::Table};
 use super::ReplicationSlot;
 
-use crate::backend::replication::logical::subscriber::omni_ownership::OmniOwnership;
 use crate::backend::replication::logical::subscriber::stream::StreamSubscriber;
 use crate::backend::replication::publisher::Lsn;
 use crate::backend::replication::publisher::progress::Progress;
@@ -143,7 +142,6 @@ impl Publisher {
             Box::pin(self.create_slots(source, &stop)).await?;
         }
 
-        let n_sources = source.shards().len();
         for (number, _) in source.shards().iter().enumerate() {
             // Use table offsets from data sync
             // or from loading them above.
@@ -152,11 +150,8 @@ impl Publisher {
                 .get(&number)
                 .map(Vec::as_slice)
                 .unwrap_or_default();
-            // Handles the logical replication stream messages.
-            // Each subscriber owns a partition of destination shards for omni-table DML
-            // (dest_shard % n_sources == source_shard), preventing cross-subscriber deadlocks.
-            let mut stream =
-                StreamSubscriber::new(dest, tables, OmniOwnership::new(number, n_sources));
+
+            let mut stream = StreamSubscriber::new(dest, tables);
 
             // Take ownership of the slot for replication.
             let mut slot = self
@@ -675,7 +670,7 @@ mod test {
         let cfg = config();
         let cluster = Cluster::new_test(&cfg);
         cluster.launch();
-        let mut stream = StreamSubscriber::new(&cluster, &[], OmniOwnership::test());
+        let mut stream = StreamSubscriber::new(&cluster, &[]);
         stream.connect().await.unwrap();
 
         let result = stream.handle(begin_copy_data(1)).await;
@@ -695,7 +690,7 @@ mod test {
         let cfg = config();
         let cluster = Cluster::new_test(&cfg);
         cluster.launch();
-        let mut stream = StreamSubscriber::new(&cluster, &[], OmniOwnership::test());
+        let mut stream = StreamSubscriber::new(&cluster, &[]);
         stream.connect().await.unwrap();
 
         let result = stream.handle(commit_copy_data(1)).await;
diff --git a/pgdog/src/backend/replication/logical/subscriber/mod.rs b/pgdog/src/backend/replication/logical/subscriber/mod.rs
index 6291a10bb..45e55fc3f 100644
--- a/pgdog/src/backend/replication/logical/subscriber/mod.rs
+++ b/pgdog/src/backend/replication/logical/subscriber/mod.rs
@@ -1,6 +1,5 @@
 pub(crate) mod context;
 pub(crate) mod copy;
-pub(crate) mod omni_ownership;
 pub(crate) mod parallel_connection;
 pub(crate) mod pipeline;
 pub(crate) mod stream;
diff --git a/pgdog/src/backend/replication/logical/subscriber/omni_ownership.rs b/pgdog/src/backend/replication/logical/subscriber/omni_ownership.rs
deleted file mode 100644
index 1529c4836..000000000
--- a/pgdog/src/backend/replication/logical/subscriber/omni_ownership.rs
+++ /dev/null
@@ -1,117 +0,0 @@
-/// Controls which destination shards a subscriber writes to for omni (unsharded) tables.
-///
-/// Partitions destinations via `dest_shard % n_sources == source_shard` so that each
-/// subscriber owns a disjoint subset, preventing cross-subscriber row-lock deadlocks.
-#[derive(Debug, Clone, Copy)]
-pub(crate) struct OmniOwnership {
-    source_shard: usize,
-    n_sources: usize,
-}
-
-impl OmniOwnership {
-    pub(crate) fn new(source_shard: usize, n_sources: usize) -> Self {
-        debug_assert!(
-            n_sources == 0 || source_shard < n_sources,
-            "source_shard ({source_shard}) must be < n_sources ({n_sources})"
-        );
-        Self {
-            source_shard,
-            n_sources,
-        }
-    }
-
-    /// Returns true if this subscriber should write omni-table DML to `dest_shard`.
-    pub(crate) fn owns(&self, dest_shard: usize) -> bool {
-        if self.n_sources <= 1 {
-            return true;
-        }
-        dest_shard % self.n_sources == self.source_shard
-    }
-}
-
-#[cfg(test)]
-impl OmniOwnership {
-    pub(crate) fn test() -> Self {
-        Self::new(0, 1)
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn omni_single_source_owns_all_dests() {
-        let p = OmniOwnership::new(0, 1);
-        assert!(p.owns(0));
-        assert!(p.owns(1));
-        assert!(p.owns(2));
-        assert!(p.owns(7));
-    }
-
-    #[test]
-    fn omni_zero_sources_owns_all_dests() {
-        let p = OmniOwnership::new(0, 0);
-        assert!(p.owns(0));
-        assert!(p.owns(1));
-        assert!(p.owns(3));
-    }
-
-    #[test]
-    fn omni_equal_sources_and_dests() {
-        // n_sources == n_dests == 3: strict 1:1, each source owns only its own index.
-        assert!(OmniOwnership::new(0, 3).owns(0));
-        assert!(OmniOwnership::new(1, 3).owns(1));
-        assert!(OmniOwnership::new(2, 3).owns(2));
-
-        assert!(!OmniOwnership::new(1, 3).owns(0));
-        assert!(!OmniOwnership::new(2, 3).owns(0));
-        assert!(!OmniOwnership::new(0, 3).owns(1));
-        assert!(!OmniOwnership::new(2, 3).owns(1));
-        assert!(!OmniOwnership::new(0, 3).owns(2));
-        assert!(!OmniOwnership::new(1, 3).owns(2));
-    }
-
-    #[test]
-    fn omni_fewer_sources_than_dests() {
-        // n_sources=2, n_dests=4: sub-0 owns even dests, sub-1 owns odd dests.
-        let p0 = OmniOwnership::new(0, 2);
-        assert!(p0.owns(0));
-        assert!(p0.owns(2));
-        assert!(!p0.owns(1));
-        assert!(!p0.owns(3));
-
-        let p1 = OmniOwnership::new(1, 2);
-        assert!(p1.owns(1));
-        assert!(p1.owns(3));
-        assert!(!p1.owns(0));
-        assert!(!p1.owns(2));
-    }
-
-    #[test]
-    fn omni_more_sources_than_dests_all_dests_covered() {
-        // n_sources=5, n_dests=3: subs 0-2 each own their matching dest exclusively.
-        assert!(OmniOwnership::new(0, 5).owns(0));
-        assert!(OmniOwnership::new(1, 5).owns(1));
-        assert!(OmniOwnership::new(2, 5).owns(2));
-
-        assert!(!OmniOwnership::new(1, 5).owns(0));
-        assert!(!OmniOwnership::new(2, 5).owns(0));
-        assert!(!OmniOwnership::new(0, 5).owns(1));
-        assert!(!OmniOwnership::new(2, 5).owns(1));
-        assert!(!OmniOwnership::new(0, 5).owns(2));
-        assert!(!OmniOwnership::new(1, 5).owns(2));
-    }
-
-    #[test]
-    fn omni_more_sources_than_dests_excess_sources_idle() {
-        // n_sources=5, n_dests=3: subs 3 and 4 own no destinations.
-        assert!(!OmniOwnership::new(3, 5).owns(0));
-        assert!(!OmniOwnership::new(3, 5).owns(1));
-        assert!(!OmniOwnership::new(3, 5).owns(2));
-
-        assert!(!OmniOwnership::new(4, 5).owns(0));
-        assert!(!OmniOwnership::new(4, 5).owns(1));
-        assert!(!OmniOwnership::new(4, 5).owns(2));
-    }
-}
diff --git a/pgdog/src/backend/replication/logical/subscriber/stream.rs b/pgdog/src/backend/replication/logical/subscriber/stream.rs
index dc054d59d..e0caf7983 100644
--- a/pgdog/src/backend/replication/logical/subscriber/stream.rs
+++ b/pgdog/src/backend/replication/logical/subscriber/stream.rs
@@ -20,7 +20,6 @@ use super::super::{
 };
 use super::PipelinedConnection;
 use super::StreamContext;
-use super::omni_ownership::OmniOwnership;
 use crate::net::messages::replication::logical::tuple_data::{Identifier, TupleData};
 use crate::net::messages::replication::logical::update::Update as XLogUpdate;
 use crate::{
@@ -130,13 +129,10 @@ pub(crate) struct StreamSubscriber {
 
     // Bytes sharded
     bytes_sharded: usize,
-
-    // Determines which destination shards this subscriber owns for omni tables.
-    partition: OmniOwnership,
 }
 
 impl StreamSubscriber {
-    pub(crate) fn new(cluster: &Cluster, tables: &[Table], partition: OmniOwnership) -> Self {
+    pub(crate) fn new(cluster: &Cluster, tables: &[Table]) -> Self {
         let cluster = cluster.logical_stream();
         Self {
             cluster,
@@ -163,7 +159,6 @@ impl StreamSubscriber {
             lsn_changed: true,
             in_transaction: false,
             keys: HashMap::default(),
-            partition,
         }
     }
 
@@ -246,7 +241,6 @@ impl StreamSubscriber {
             }
         }
 
-        let partition = self.partition;
         let n_conns = self.connections.len();
         let is_direct = val.is_direct();
 
@@ -256,13 +250,9 @@ impl StreamSubscriber {
         let mut pending: Option<usize> = None;
         for shard in 0..n_conns {
             let target = match val {
-                // With a single destination shard the router collapses Shard::All
-                // to Direct(0), bypassing the partition ownership check.  Apply
-                // partition.owns() for all variants when there is only one connection
-                // so that omni-table writes are still partitioned across subscribers.
-                Shard::Direct(direct) if n_conns > 1 => shard == *direct,
-                Shard::Multi(multi) if n_conns > 1 => multi.contains(&shard),
-                _ => partition.owns(shard),
+                Shard::Direct(direct) => shard == *direct,
+                Shard::Multi(multi) => multi.contains(&shard),
+                _ => true,
             };
             if !target {
                 continue;
@@ -1028,7 +1018,7 @@ mod tests {
 
     fn make_subscriber() -> StreamSubscriber {
         let cluster = Cluster::new_test(&config());
-        StreamSubscriber::new(&cluster, &[], OmniOwnership::test())
+        StreamSubscriber::new(&cluster, &[])
     }
 
     #[test]
diff --git a/pgdog/src/backend/replication/logical/subscriber/tests.rs b/pgdog/src/backend/replication/logical/subscriber/tests.rs
index 1c271ee77..f4e9b5358 100644
--- a/pgdog/src/backend/replication/logical/subscriber/tests.rs
+++ b/pgdog/src/backend/replication/logical/subscriber/tests.rs
@@ -29,7 +29,6 @@ use crate::{
     },
 };
 
-use super::omni_ownership::OmniOwnership;
 use super::stream::StreamSubscriber;
 
 fn random_id() -> String {
@@ -247,26 +246,23 @@ fn x_update(u: XLogUpdate) -> CopyData {
 fn make_subscriber() -> StreamSubscriber {
     let cluster = Cluster::new_test(&config());
     let tables = vec![make_sharded_table(), make_sharded_test_b_table()];
-    StreamSubscriber::new(&cluster, &tables, OmniOwnership::test())
+    StreamSubscriber::new(&cluster, &tables)
 }
 
 fn make_subscriber_with_tables(tables: Vec<Table>) -> StreamSubscriber {
     let cluster = Cluster::new_test(&config());
-    StreamSubscriber::new(&cluster, &tables, OmniOwnership::test())
+    StreamSubscriber::new(&cluster, &tables)
 }
 
-fn make_subscriber_with_tables_two_databases(
-    tables: Vec<Table>,
-    partition: OmniOwnership,
-) -> StreamSubscriber {
+fn make_subscriber_with_tables_two_databases(tables: Vec<Table>) -> StreamSubscriber {
     let cluster = Cluster::new_test_two_databases(&config());
-    StreamSubscriber::new(&cluster, &tables, partition)
+    StreamSubscriber::new(&cluster, &tables)
 }
 
 fn make_subscriber_single_shard() -> StreamSubscriber {
     let cluster = Cluster::new_test_single_shard(&config());
     let tables = vec![make_sharded_table(), make_sharded_test_b_table()];
-    StreamSubscriber::new(&cluster, &tables, OmniOwnership::test())
+    StreamSubscriber::new(&cluster, &tables)
 }
 
 /// Count rows matching the given `WHERE` predicate using a separate connection.
@@ -576,7 +572,7 @@ async fn partition_leaves_share_destination() {
     leaf_b.table.parent_name = "sharded".to_string();
 
     let cluster = Cluster::new_test_single_shard(&config());
-    let mut sub = StreamSubscriber::new(&cluster, &[leaf_a, leaf_b], OmniOwnership::test());
+    let mut sub = StreamSubscriber::new(&cluster, &[leaf_a, leaf_b]);
     let mut verify = test_server().await;
     sub.connect().await.unwrap();
 
@@ -1507,11 +1503,7 @@ fn omni_insert_copy_data(oid: Oid, a: &str, b: &str) -> CopyData {
 #[tokio::test]
 async fn full_identity_nothing_rejected() {
     let cluster = Cluster::new_test_single_shard(&config());
-    let mut sub = StreamSubscriber::new(
-        &cluster,
-        &[make_replica_identity_nothing_table()],
-        OmniOwnership::test(),
-    );
+    let mut sub = StreamSubscriber::new(&cluster, &[make_replica_identity_nothing_table()]);
     sub.connect().await.unwrap();
 
     let oid = Oid(16390);
@@ -1545,11 +1537,7 @@ async fn full_identity_nothing_rejected() {
 #[tokio::test]
 async fn full_identity_omni_no_unique_index_rejected() {
     let cluster = Cluster::new_test_single_shard(&config());
-    let mut sub = StreamSubscriber::new(
-        &cluster,
-        &[make_full_identity_omni_table()],
-        OmniOwnership::test(),
-    );
+    let mut sub = StreamSubscriber::new(&cluster, &[make_full_identity_omni_table()]);
 
     // Enforce precondition: the table must exist but have no qualifying unique index.
     // A stale unique index from a prior run would make tables_missing_unique_index() return empty,
@@ -1588,11 +1576,7 @@ async fn full_identity_omni_no_unique_index_rejected() {
 #[tokio::test]
 async fn full_identity_insert_sharded() {
     let cluster = Cluster::new_test_single_shard(&config());
-    let mut sub = StreamSubscriber::new(
-        &cluster,
-        &[make_full_identity_sharded_table()],
-        OmniOwnership::test(),
-    );
+    let mut sub = StreamSubscriber::new(&cluster, &[make_full_identity_sharded_table()]);
     let mut verify = test_server().await;
     sub.connect().await.unwrap();
 
@@ -1618,11 +1602,7 @@ async fn full_identity_insert_sharded() {
 #[tokio::test]
 async fn full_identity_update_fast_path() {
     let cluster = Cluster::new_test_single_shard(&config());
-    let mut sub = StreamSubscriber::new(
-        &cluster,
-        &[make_full_identity_sharded_table()],
-        OmniOwnership::test(),
-    );
+    let mut sub = StreamSubscriber::new(&cluster, &[make_full_identity_sharded_table()]);
     let mut verify = test_server().await;
     sub.connect().await.unwrap();
 
@@ -1677,11 +1657,7 @@ async fn full_identity_update_fast_path() {
 #[tokio::test]
 async fn full_identity_update_slow_path() {
     let cluster = Cluster::new_test_single_shard(&config());
-    let mut sub = StreamSubscriber::new(
-        &cluster,
-        &[make_full_identity_sharded_table()],
-        OmniOwnership::test(),
-    );
+    let mut sub = StreamSubscriber::new(&cluster, &[make_full_identity_sharded_table()]);
     let mut verify = test_server().await;
     sub.connect().await.unwrap();
 
@@ -1741,11 +1717,7 @@ async fn full_identity_update_slow_path() {
 #[tokio::test]
 async fn full_identity_update_slow_path_realistic_old_tuple() {
     let cluster = Cluster::new_test_single_shard(&config());
-    let mut sub = StreamSubscriber::new(
-        &cluster,
-        &[make_full_identity_sharded_table()],
-        OmniOwnership::test(),
-    );
+    let mut sub = StreamSubscriber::new(&cluster, &[make_full_identity_sharded_table()]);
     let mut verify = test_server().await;
     sub.connect().await.unwrap();
 
@@ -1802,11 +1774,7 @@ async fn full_identity_update_slow_path_realistic_old_tuple() {
 #[tokio::test]
 async fn full_identity_update_all_toasted_is_noop() {
     let cluster = Cluster::new_test_single_shard(&config());
-    let mut sub = StreamSubscriber::new(
-        &cluster,
-        &[make_full_identity_sharded_table()],
-        OmniOwnership::test(),
-    );
+    let mut sub = StreamSubscriber::new(&cluster, &[make_full_identity_sharded_table()]);
     let mut verify = test_server().await;
     sub.connect().await.unwrap();
 
@@ -1848,11 +1816,7 @@ async fn full_identity_update_all_toasted_is_noop() {
 #[tokio::test]
 async fn full_identity_delete() {
     let cluster = Cluster::new_test_single_shard(&config());
-    let mut sub = StreamSubscriber::new(
-        &cluster,
-        &[make_full_identity_sharded_table()],
-        OmniOwnership::test(),
-    );
+    let mut sub = StreamSubscriber::new(&cluster, &[make_full_identity_sharded_table()]);
     let mut verify = test_server().await;
     sub.connect().await.unwrap();
 
@@ -1890,11 +1854,7 @@ async fn full_identity_delete() {
 #[tokio::test]
 async fn full_identity_insert_omni_dedup() {
     let cluster = Cluster::new_test_single_shard(&config());
-    let mut sub = StreamSubscriber::new(
-        &cluster,
-        &[make_full_identity_omni_dedup_table()],
-        OmniOwnership::test(),
-    );
+    let mut sub = StreamSubscriber::new(&cluster, &[make_full_identity_omni_dedup_table()]);
     let mut verify = test_server().await;
 
     // Ensure destination table exists with unique index before relation() runs.
@@ -1953,11 +1913,7 @@ async fn full_identity_insert_omni_dedup() {
 #[tokio::test]
 async fn full_identity_update_duplicate_rows() {
     let cluster = Cluster::new_test_single_shard(&config());
-    let mut sub = StreamSubscriber::new(
-        &cluster,
-        &[make_full_identity_dup_rows_table()],
-        OmniOwnership::test(),
-    );
+    let mut sub = StreamSubscriber::new(&cluster, &[make_full_identity_dup_rows_table()]);
     let mut verify = test_server().await;
 
     ensure_table(&mut verify, "public.full_dup_rows").await;
@@ -2023,11 +1979,7 @@ async fn full_identity_update_duplicate_rows() {
 #[tokio::test]
 async fn full_identity_delete_duplicate_rows() {
     let cluster = Cluster::new_test_single_shard(&config());
-    let mut sub = StreamSubscriber::new(
-        &cluster,
-        &[make_full_identity_dup_rows_table()],
-        OmniOwnership::test(),
-    );
+    let mut sub = StreamSubscriber::new(&cluster, &[make_full_identity_dup_rows_table()]);
     let mut verify = test_server().await;
 
     ensure_table(&mut verify, "public.full_dup_rows").await;
@@ -2094,11 +2046,7 @@ async fn full_identity_delete_duplicate_rows() {
 #[tokio::test]
 async fn full_identity_update_matches_null_column() {
     let cluster = Cluster::new_test_single_shard(&config());
-    let mut sub = StreamSubscriber::new(
-        &cluster,
-        &[make_full_identity_dup_rows_table()],
-        OmniOwnership::test(),
-    );
+    let mut sub = StreamSubscriber::new(&cluster, &[make_full_identity_dup_rows_table()]);
     let mut verify = test_server().await;
 
     // full_dup_rows has no NOT NULL on value — we can seed a NULL row.
@@ -2159,11 +2107,7 @@ async fn full_identity_update_matches_null_column() {
 #[tokio::test]
 async fn full_identity_delete_matches_null_column() {
     let cluster = Cluster::new_test_single_shard(&config());
-    let mut sub = StreamSubscriber::new(
-        &cluster,
-        &[make_full_identity_dup_rows_table()],
-        OmniOwnership::test(),
-    );
+    let mut sub = StreamSubscriber::new(&cluster, &[make_full_identity_dup_rows_table()]);
     let mut verify = test_server().await;
 
     ensure_table(&mut verify, "public.full_dup_rows").await;
@@ -2334,13 +2278,7 @@ async fn cross_subscriber_omni_deadlock_two_databases() {
         let (id1, id2) = (id1.clone(), id2.clone());
         let barrier = Arc::clone(&barrier);
         tokio::spawn(async move {
-            // Each subscriber owns a disjoint subset of destination shards:
-            // sub-0 → dest-0, sub-1 → dest-1 (dest_shard % 2 == sub_idx).
-            // This is the destination-partitioned apply fix for the cross-subscriber deadlock.
-            let mut sub = make_subscriber_with_tables_two_databases(
-                vec![make_settings_table()],
-                OmniOwnership::new(sub_idx, 2),
-            );
+            let mut sub = make_subscriber_with_tables_two_databases(vec![make_settings_table()]);
             sub.connect().await.unwrap();
             // Distinct LSN ranges so neither subscriber's LSN gating skips the other's events.
             let mut lsn = 100_000i64 + (sub_idx as i64) * 1_000_000;

From 84c4886caddcc64e4d88c7ea96dbd39c08e03986 Mon Sep 17 00:00:00 2001
From: meskill <8974488+meskill@users.noreply.github.com>
Date: Fri, 4 Sep 2026 16:09:56 +0000
Subject: [PATCH 4/5] fix(resharding): force lock_timeout during resharding

---
 pgdog/src/backend/connect_reason.rs           | 24 +++----------------
 pgdog/src/backend/pool/monitor.rs             |  2 +-
 pgdog/src/backend/pool/pool_impl.rs           | 12 ++++++++--
 .../replication/logical/publisher/slot.rs     |  4 ++--
 .../replication/logical/subscriber/copy.rs    |  2 +-
 .../replication/logical/subscriber/stream.rs  |  2 +-
 pgdog/src/backend/server.rs                   |  2 +-
 7 files changed, 19 insertions(+), 29 deletions(-)

diff --git a/pgdog/src/backend/connect_reason.rs b/pgdog/src/backend/connect_reason.rs
index 29a73318e..764ac6473 100644
--- a/pgdog/src/backend/connect_reason.rs
+++ b/pgdog/src/backend/connect_reason.rs
@@ -1,31 +1,13 @@
-use std::fmt::Display;
-
-#[derive(Debug, Clone, Copy, Default, PartialEq)]
+#[derive(Debug, Display, Clone, Copy, Default, PartialEq)]
+#[display(rename_all = "snake_case")]
 pub(crate) enum ConnectReason {
     LsnCheck,
     BelowMin,
     ClientWaiting,
-    Replication,
+    Resharding,
     PubSub,
     Probe,
     Healthcheck,
     #[default]
     Other,
 }
-
-impl Display for ConnectReason {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        let reason = match self {
-            Self::LsnCheck => "lsn check",
-            Self::BelowMin => "min",
-            Self::ClientWaiting => "client",
-            Self::Replication => "replication",
-            Self::PubSub => "pub/sub",
-            Self::Probe => "probe",
-            Self::Healthcheck => "healthcheck",
-            Self::Other => "other",
-        };
-
-        write!(f, "{}", reason)
-    }
-}
diff --git a/pgdog/src/backend/pool/monitor.rs b/pgdog/src/backend/pool/monitor.rs
index d2068c352..0bd7d89f8 100644
--- a/pgdog/src/backend/pool/monitor.rs
+++ b/pgdog/src/backend/pool/monitor.rs
@@ -434,7 +434,7 @@ impl Monitor {
         let connect_timeout = pool.config().connect_timeout;
         let connect_attempts = pool.config().connect_attempts;
         let connect_attempt_delay = pool.config().connect_attempt_delay;
-        let options = pool.server_options();
+        let options = pool.server_options(reason);
 
         let mut error = Error::ServerError;
         let now = Instant::now();
diff --git a/pgdog/src/backend/pool/pool_impl.rs b/pgdog/src/backend/pool/pool_impl.rs
index 1c351dd74..ef8f0a883 100644
--- a/pgdog/src/backend/pool/pool_impl.rs
+++ b/pgdog/src/backend/pool/pool_impl.rs
@@ -431,7 +431,7 @@ impl Pool {
     }
 
     /// Get startup parameters for new server connections.
-    pub(super) fn server_options(&self) -> ServerOptions {
+    pub(super) fn server_options(&self, reason: ConnectReason) -> ServerOptions {
         let mut params = vec![
             Parameter {
                 name: "application_name".into(),
@@ -445,6 +445,14 @@ impl Pool {
 
         let config = self.inner.config;
 
+        let lock_timeout = config
+            .lock_timeout
+            // Enforce some lock_timeout during resharding to prevent possible deadlocks.
+            // This should be mostly avoided by pgdog, but in case some invariants are not met,
+            // the resharding could deadlock and with timeout we'll probably retry the update
+            // and either succeed or fail explicitly.
+            .or(matches!(reason, ConnectReason::Resharding).then_some(Duration::from_secs(5)));
+
         if let Some(statement_timeout) = config.statement_timeout {
             params.push(Parameter {
                 name: "statement_timeout".into(),
@@ -452,7 +460,7 @@ impl Pool {
             });
         }
 
-        if let Some(lock_timeout) = config.lock_timeout {
+        if let Some(lock_timeout) = lock_timeout {
             params.push(Parameter {
                 name: "lock_timeout".into(),
                 value: lock_timeout.as_millis().to_string().into(),
diff --git a/pgdog/src/backend/replication/logical/publisher/slot.rs b/pgdog/src/backend/replication/logical/publisher/slot.rs
index 683d6ed9f..d59c619ea 100644
--- a/pgdog/src/backend/replication/logical/publisher/slot.rs
+++ b/pgdog/src/backend/replication/logical/publisher/slot.rs
@@ -107,7 +107,7 @@ impl ReplicationSlot {
             Box::pin(Server::connect(
                 &self.address,
                 ServerOptions::new_replication(),
-                ConnectReason::Replication,
+                ConnectReason::Resharding,
                 Default::default(),
             ))
             .await?,
@@ -125,7 +125,7 @@ impl ReplicationSlot {
                 Server::connect(
                     &self.address,
                     ServerOptions::default(),
-                    ConnectReason::Replication,
+                    ConnectReason::Resharding,
                     Default::default(),
                 )
                 .await?,
diff --git a/pgdog/src/backend/replication/logical/subscriber/copy.rs b/pgdog/src/backend/replication/logical/subscriber/copy.rs
index 3a35ecf94..656664c70 100644
--- a/pgdog/src/backend/replication/logical/subscriber/copy.rs
+++ b/pgdog/src/backend/replication/logical/subscriber/copy.rs
@@ -81,7 +81,7 @@ impl CopySubscriber {
                 .find(|(role, _)| role == &Role::Primary)
                 .ok_or(Error::NoPrimary)?
                 .1
-                .standalone(ConnectReason::Replication)
+                .standalone(ConnectReason::Resharding)
                 .await?;
             servers.push(ParallelConnection::new(primary)?);
         }
diff --git a/pgdog/src/backend/replication/logical/subscriber/stream.rs b/pgdog/src/backend/replication/logical/subscriber/stream.rs
index e0caf7983..82587e1f3 100644
--- a/pgdog/src/backend/replication/logical/subscriber/stream.rs
+++ b/pgdog/src/backend/replication/logical/subscriber/stream.rs
@@ -178,7 +178,7 @@ impl StreamSubscriber {
                 .find(|(r, _)| r == &Role::Primary)
                 .ok_or(Error::NoPrimary)?
                 .1
-                .standalone(ConnectReason::Replication)
+                .standalone(ConnectReason::Resharding)
                 .await?;
             conns.push(primary);
         }
diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs
index 7af7d1460..b1161e168 100644
--- a/pgdog/src/backend/server.rs
+++ b/pgdog/src/backend/server.rs
@@ -1420,7 +1420,7 @@ pub(crate) mod test {
         Server::connect(
             &Address::new_test(),
             ServerOptions::new_replication(),
-            ConnectReason::Replication,
+            ConnectReason::Resharding,
             Default::default(),
         )
         .await

From 5d951730d67aac976809c899e76481798ff20653 Mon Sep 17 00:00:00 2001
From: meskill <8974488+meskill@users.noreply.github.com>
Date: Fri, 4 Sep 2026 15:36:15 +0000
Subject: [PATCH 5/5] test(resharding): drop very specific deadlock test

this test constructs very specific environment and flow, that is not exactly how the whole flow runs.
---
 docs/RESHARDING.md                            |  11 +
 pgdog/src/backend/pool/cluster.rs             |  29 ---
 .../replication/logical/subscriber/tests.rs   | 201 +-----------------
 pgdog/src/backend/server.rs                   |  17 --
 4 files changed, 12 insertions(+), 246 deletions(-)

diff --git a/docs/RESHARDING.md b/docs/RESHARDING.md
index acd76656c..cd8193265 100644
--- a/docs/RESHARDING.md
+++ b/docs/RESHARDING.md
@@ -143,6 +143,17 @@ Two behaviours are specific to the resharding context:
   COPY. Messages at or below that LSN are skipped; the row is already on the destination.
 - **Omnisharded tables** (`statements.omni = true`): upsert is broadcast to all shards
   simultaneously rather than routed to a single shard.
+- **Table ownership** ([`tables_sync()`](../pgdog/src/backend/replication/logical/tables_sync.rs)):
+  a table that is *sharded on the source* is copied and replayed from every source shard.
+  A table that is *omnisharded on the source* is copied and replayed from one source shard
+  only, chosen by publication order, because every source shard holds the same rows.
+- **Destination row contention**: a table that is sharded on the source and omnisharded on
+  the destination is replayed by every subscriber, and every subscriber writes to every
+  destination shard. Two subscribers therefore write the same destination row whenever one
+  key reaches two source shards, for example after a sharding-key update. Two subscribers
+  can then lock the same rows on two destinations in opposite order. No Postgres instance
+  sees the whole cycle, so no instance reports a deadlock. Set `lock_timeout` on the
+  destination user so a blocked apply is cancelled and retried by `Publisher::replicate()`.
 ---
 
 ### Cutover phases
diff --git a/pgdog/src/backend/pool/cluster.rs b/pgdog/src/backend/pool/cluster.rs
index 5f33d36f2..12303a3f2 100644
--- a/pgdog/src/backend/pool/cluster.rs
+++ b/pgdog/src/backend/pool/cluster.rs
@@ -915,35 +915,6 @@ mod test {
             cluster
         }
 
-        /// Two shards targeting different databases on the same server.
-        /// Gives separate lock namespaces without needing two Postgres instances.
-        pub(crate) fn new_test_two_databases(config: &ConfigAndUsers) -> Cluster {
-            let mut cluster = Self::new_test(config);
-            let shard1 = cluster.shards.last_mut().unwrap();
-            *shard1 = Shard::new(ShardConfig {
-                number: 1,
-                primary: Some(&PoolConfig {
-                    address: Address {
-                        database_name: "pgdog1".into(),
-                        ..Address::new_test()
-                    },
-                    config: Config::default(),
-                }),
-                replicas: &[PoolConfig {
-                    address: Address {
-                        database_name: "pgdog1".into(),
-                        configured_role: Role::Replica,
-                        ..Address::new_test()
-                    },
-                    config: Config::default(),
-                }],
-                identifier: cluster.identifier.clone(),
-                lsn_check_interval: Duration::MAX,
-                ..Default::default()
-            });
-            cluster
-        }
-
         pub(crate) fn new_test_single_primary(config: &ConfigAndUsers) -> Cluster {
             let identifier = Arc::new(DatabaseUser {
                 user: "pgdog".into(),
diff --git a/pgdog/src/backend/replication/logical/subscriber/tests.rs b/pgdog/src/backend/replication/logical/subscriber/tests.rs
index f4e9b5358..a8d3dda5d 100644
--- a/pgdog/src/backend/replication/logical/subscriber/tests.rs
+++ b/pgdog/src/backend/replication/logical/subscriber/tests.rs
@@ -9,7 +9,7 @@ use crate::{
         replication::logical::publisher::{
             Lsn, PublicationTable, PublicationTableColumn, ReplicaIdentity, Table,
         },
-        server::test::{test_server, test_server_pgdog1_db},
+        server::test::test_server,
     },
     config::config,
     net::{
@@ -254,11 +254,6 @@ fn make_subscriber_with_tables(tables: Vec<Table>) -> StreamSubscriber {
     StreamSubscriber::new(&cluster, &tables)
 }
 
-fn make_subscriber_with_tables_two_databases(tables: Vec<Table>) -> StreamSubscriber {
-    let cluster = Cluster::new_test_two_databases(&config());
-    StreamSubscriber::new(&cluster, &tables)
-}
-
 fn make_subscriber_single_shard() -> StreamSubscriber {
     let cluster = Cluster::new_test_single_shard(&config());
     let tables = vec![make_sharded_table(), make_sharded_test_b_table()];
@@ -2155,197 +2150,3 @@ async fn full_identity_delete_matches_null_column() {
         "FULL identity DELETE must match NULL via IS NOT DISTINCT FROM"
     );
 }
-
-// ── Omni-table fan-out tests ─────────────────────────────────────────────────
-
-fn make_settings_table() -> Table {
-    Table {
-        publication: "test".to_string(),
-        table: PublicationTable {
-            schema: "public".to_string(),
-            name: "settings".to_string(),
-            attributes: "".to_string(),
-            parent_schema: "".to_string(),
-            parent_name: "".to_string(),
-        },
-        identity: ReplicaIdentity {
-            oid: Oid(4),
-            identity: "".to_string(),
-            kind: "".to_string(),
-        },
-        columns: vec![
-            PublicationTableColumn {
-                oid: 4,
-                name: "id".to_string(),
-                type_oid: Oid(20), // bigint
-                identity: true,
-            },
-            PublicationTableColumn {
-                oid: 4,
-                name: "name".to_string(),
-                type_oid: Oid(25), // text
-                identity: false,
-            },
-            PublicationTableColumn {
-                oid: 4,
-                name: "value".to_string(),
-                type_oid: Oid(25), // text
-                identity: false,
-            },
-        ],
-        lsn: Lsn::default(),
-    }
-}
-
-fn settings_relation(oid: Oid) -> Relation {
-    Relation {
-        oid,
-        namespace: "public".to_string(),
-        name: "settings".to_string(),
-        replica_identity: 100,
-        columns: vec![
-            RelColumn {
-                flag: 1,
-                name: "id".to_string(),
-                oid: Oid(20),
-                type_modifier: -1,
-            },
-            RelColumn {
-                flag: 0,
-                name: "name".to_string(),
-                oid: Oid(25),
-                type_modifier: -1,
-            },
-            RelColumn {
-                flag: 0,
-                name: "value".to_string(),
-                oid: Oid(25),
-                type_modifier: -1,
-            },
-        ],
-    }
-}
-
-fn settings_relation_copy_data(oid: Oid) -> CopyData {
-    xlog_copy_data(settings_relation(oid).to_bytes())
-}
-
-/// WAL UPDATE for settings(id, name, value) — full new tuple, no toasted columns.
-fn settings_update_copy_data(oid: Oid, id: &str, name: &str, value: &str) -> CopyData {
-    xlog_copy_data(
-        XLogUpdate {
-            oid,
-            identity: UpdateIdentity::Nothing,
-            new: TupleData {
-                columns: vec![text_column(id), text_column(name), text_column(value)],
-            },
-        }
-        .to_bytes(),
-    )
-}
-
-/// Two subscribers race on the same omni-table rows, reproducing the cross-destination deadlock.
-#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
-async fn cross_subscriber_omni_deadlock_two_databases() {
-    use std::sync::Arc;
-    use std::time::Duration;
-    use tokio::sync::Barrier;
-    use tokio::time::{sleep, timeout};
-
-    let oid = Oid(16393);
-    let id1 = random_id();
-    let id2 = random_id();
-
-    let mut pg0 = test_server().await;
-    let mut pg1 = test_server_pgdog1_db().await;
-    for db in [&mut pg0, &mut pg1] {
-        cleanup(db, "public.settings", &[&id1, &id2]).await;
-        for (id, name) in [(&id1, "seed1"), (&id2, "seed2")] {
-            db.execute(format!(
-                "INSERT INTO public.settings (id, name, value) VALUES ({}, '{}', 'v')",
-                id, name,
-            ))
-            .await
-            .unwrap();
-        }
-    }
-    drop(pg0);
-    drop(pg1);
-
-    let barrier = Arc::new(Barrier::new(2));
-
-    let spawn_sub = |sub_idx: usize| {
-        let (id1, id2) = (id1.clone(), id2.clone());
-        let barrier = Arc::clone(&barrier);
-        tokio::spawn(async move {
-            let mut sub = make_subscriber_with_tables_two_databases(vec![make_settings_table()]);
-            sub.connect().await.unwrap();
-            // Distinct LSN ranges so neither subscriber's LSN gating skips the other's events.
-            let mut lsn = 100_000i64 + (sub_idx as i64) * 1_000_000;
-            for round in 1..=20usize {
-                sub.handle(begin_copy_data(lsn)).await.unwrap();
-                sub.handle(settings_relation_copy_data(oid)).await.unwrap();
-                barrier.wait().await;
-                sub.handle(settings_update_copy_data(
-                    oid,
-                    &id1,
-                    &format!("r{round}-{id1}"),
-                    "v",
-                ))
-                .await
-                .expect("update id1");
-                sub.handle(settings_update_copy_data(
-                    oid,
-                    &id2,
-                    &format!("r{round}-{id2}"),
-                    "v",
-                ))
-                .await
-                .expect("update id2");
-                sub.handle(commit_copy_data(lsn + 100))
-                    .await
-                    .expect("commit");
-                lsn += 200;
-            }
-        })
-    };
-
-    let h0 = spawn_sub(0);
-    let h1 = spawn_sub(1);
-    let abort0 = h0.abort_handle();
-    let abort1 = h1.abort_handle();
-
-    let result = timeout(Duration::from_secs(10), async { tokio::join!(h0, h1) }).await;
-
-    abort0.abort();
-    abort1.abort();
-    sleep(Duration::from_millis(200)).await;
-
-    let mut pg0 = test_server().await;
-    let mut pg1 = test_server_pgdog1_db().await;
-
-    match result {
-        Err(_elapsed) => {
-            cleanup(&mut pg0, "public.settings", &[&id1, &id2]).await;
-            cleanup(&mut pg1, "public.settings", &[&id1, &id2]).await;
-            panic!("cross-subscriber omni deadlock: both subscribers hung");
-        }
-        Ok((r0, r1)) => {
-            r0.expect("sub-0 failed");
-            r1.expect("sub-1 failed");
-            for db in [&mut pg0, &mut pg1] {
-                for id in [&id1, &id2] {
-                    let count = count_where(
-                        db,
-                        "public.settings",
-                        &format!("id = {id} AND name = 'r20-{id}'"),
-                    )
-                    .await;
-                    assert_eq!(count, 1, "row {id} missing on destination");
-                }
-            }
-            cleanup(&mut pg0, "public.settings", &[&id1, &id2]).await;
-            cleanup(&mut pg1, "public.settings", &[&id1, &id2]).await;
-        }
-    }
-}
diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs
index b1161e168..b18c99630 100644
--- a/pgdog/src/backend/server.rs
+++ b/pgdog/src/backend/server.rs
@@ -1399,23 +1399,6 @@ pub(crate) mod test {
         .unwrap()
     }
 
-    /// Connect to the `pgdog1` database on the test server.
-    /// Used by tests that need a second, distinct database so that
-    /// row locks on the two databases do not share a lock namespace.
-    pub(crate) async fn test_server_pgdog1_db() -> Server {
-        Server::connect(
-            &Address {
-                database_name: "pgdog1".into(),
-                ..Address::new_test()
-            },
-            ServerOptions::default(),
-            ConnectReason::Other,
-            Default::default(),
-        )
-        .await
-        .unwrap()
-    }
-
     pub(crate) async fn test_replication_server() -> Server {
         Server::connect(
             &Address::new_test(),