fix: dropping child table rows during streaming replication (#880)

For partitioned tables, we were only streaming rows for the first child
instead of all of them.

Author: levkk

integration/resharding/dev.sh

@@ -4,6 +4,18 @@ SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
 DEFAULT_BIN="${SCRIPT_DIR}/../../target/debug/pgdog"
 PGDOG_BIN=${PGDOG_BIN:-$DEFAULT_BIN}
 
+# Run in our own process group so we can kill every child on exit.
+set -m
+cleanup() {
+    trap - EXIT INT TERM
+    # Signal every process in this script's process group except ourselves.
+    pkill -TERM -P $$ 2> /dev/null || true
+    # Give children a moment to exit cleanly, then force-kill anything left.
+    sleep 1
+    pkill -KILL -P $$ 2> /dev/null || true
+}
+trap cleanup EXIT INT TERM
+
 pushd ${SCRIPT_DIR}
 docker-compose down && docker-compose up -d
 

pgdog/src/backend/replication/logical/error.rs

@@ -131,6 +131,9 @@ pub enum Error {
 
     #[error("command complete has no rows: {0}")]
     CommandCompleteNoRows(CommandComplete),
+
+    #[error("missing key in replication stream, out of sync")]
+    MissingKey,
 }
 
 impl From<ErrorResponse> for Error {

pgdog/src/backend/replication/logical/subscriber/stream.rs

@@ -131,9 +131,8 @@ pub struct StreamSubscriber {
 
     // Statements
     statements: HashMap<Oid, Statements>,
-
-    // Partitioned tables dedup.
-    partitioned_dedup: HashSet<Key>,
+    // Mapping of table keys to their oid.
+    keys: HashMap<Key, Oid>,
 
     // LSNs for each table
     table_lsns: HashMap<Oid, i64>,
@@ -171,7 +170,6 @@ impl StreamSubscriber {
             cluster,
             relations: HashMap::new(),
             statements: HashMap::new(),
-            partitioned_dedup: HashSet::new(),
             table_lsns: HashMap::new(),
             changed_tables: HashSet::new(),
             tables: tables
@@ -193,6 +191,7 @@ impl StreamSubscriber {
             in_transaction: false,
             query_parser_engine,
             missed_rows: MissedRows::default(),
+            keys: HashMap::default(),
         }
     }
 
@@ -466,73 +465,78 @@ impl StreamSubscriber {
                 name: table.table.destination_name().to_string(),
             };
 
-            if self.partitioned_dedup.contains(&dest_key) {
-                debug!("queries for table {} already prepared", dest_key);
-                return Ok(());
-            }
+            // Partition child tables target the parent on the destination shard,
+            // we don't need to prepare the same statement per child.
+            if let Some(oid) = self.keys.get(&dest_key) {
+                let statements = self.statements.get(oid).ok_or(Error::MissingKey)?;
+                self.statements.insert(relation.oid, statements.clone());
 
-            let insert = Statement::new(&table.insert(false), self.query_parser_engine)?;
-            let upsert = Statement::new(&table.insert(true), self.query_parser_engine)?;
-            let update = Statement::new(&table.update(), self.query_parser_engine)?;
-            let delete = Statement::new(&table.delete(), self.query_parser_engine)?;
+                debug!("queries for table {} already prepared", dest_key);
+            } else {
+                let insert = Statement::new(&table.insert(false), self.query_parser_engine)?;
+                let upsert = Statement::new(&table.insert(true), self.query_parser_engine)?;
+                let update = Statement::new(&table.update(), self.query_parser_engine)?;
+                let delete = Statement::new(&table.delete(), self.query_parser_engine)?;
+
+                for server in &mut self.connections {
+                    for stmt in &[&insert, &upsert, &update, &delete] {
+                        debug!("preparing \"{}\" [{}]", stmt.query(), server.addr());
+                    }
 
-            for server in &mut self.connections {
-                for stmt in &[&insert, &upsert, &update, &delete] {
-                    debug!("preparing \"{}\" [{}]", stmt.query(), server.addr());
+                    server
+                        .send(
+                            &vec![
+                                insert.parse().clone().into(),
+                                upsert.parse().clone().into(),
+                                update.parse().clone().into(),
+                                delete.parse().clone().into(),
+                                if self.in_transaction {
+                                    Flush.into()
+                                } else {
+                                    Sync.into()
+                                },
+                            ]
+                            .into(),
+                        )
+                        .await?;
                 }
 
-                server
-                    .send(
-                        &vec![
-                            insert.parse().clone().into(),
-                            upsert.parse().clone().into(),
-                            update.parse().clone().into(),
-                            delete.parse().clone().into(),
-                            if self.in_transaction {
-                                Flush.into()
-                            } else {
-                                Sync.into()
-                            },
-                        ]
-                        .into(),
-                    )
-                    .await?;
-            }
-
-            for server in &mut self.connections {
-                let num_messages = if self.in_transaction { 4 } else { 5 };
-                for _ in 0..num_messages {
-                    let msg = server.read().await?;
-                    trace!("[{}] --> {:?}", server.addr(), msg);
-
-                    match msg.code() {
-                        'E' => {
-                            return Err(Error::PgError(Box::new(ErrorResponse::from_bytes(
-                                msg.to_bytes()?,
-                            )?)))
+                for server in &mut self.connections {
+                    let num_messages = if self.in_transaction { 4 } else { 5 };
+                    for _ in 0..num_messages {
+                        let msg = server.read().await?;
+                        trace!("[{}] --> {:?}", server.addr(), msg);
+
+                        match msg.code() {
+                            'E' => {
+                                return Err(Error::PgError(Box::new(ErrorResponse::from_bytes(
+                                    msg.to_bytes()?,
+                                )?)))
+                            }
+                            'Z' => break,
+                            '1' => continue,
+                            c => return Err(Error::RelationOutOfSync(c)),
                         }
-                        'Z' => break,
-                        '1' => continue,
-                        c => return Err(Error::RelationOutOfSync(c)),
                     }
                 }
-            }
 
-            self.statements.insert(
-                relation.oid,
-                Statements {
-                    insert,
-                    upsert,
-                    update,
-                    delete,
-                    omni: !table.is_sharded(&self.cluster.sharding_schema().tables),
-                },
-            );
+                self.statements.insert(
+                    relation.oid,
+                    Statements {
+                        insert,
+                        upsert,
+                        update,
+                        delete,
+                        omni: !table.is_sharded(&self.cluster.sharding_schema().tables),
+                    },
+                );
+
+                self.keys.insert(dest_key, relation.oid);
+            }
 
             // Only record tables we expect to stream changes for.
             self.table_lsns.insert(relation.oid, table.lsn.lsn);
             self.relations.insert(relation.oid, relation);
-            self.partitioned_dedup.insert(dest_key);
         }
 
         Ok(())

pgdog/src/backend/replication/logical/subscriber/tests.rs

@@ -466,6 +466,67 @@ async fn relation_after_insert_inside_transaction() {
     cleanup(&mut verify, "public.sharded_test_b", &[&id_b]).await;
 }
 
+/// Two source tables (e.g. partition leaves) that map to the same destination
+/// must each register their oid so DML for *both* oids is applied. Regression
+/// test for the partition-dedup row-drop bug: previously the second leaf's
+/// Relation message returned early without registering its oid in `statements`,
+/// causing all subsequent inserts on that oid to be silently dropped.
+#[tokio::test]
+async fn partition_leaves_share_destination() {
+    let mut leaf_a = make_sharded_table();
+    leaf_a.table.name = "sharded_p1".to_string();
+    leaf_a.table.parent_schema = "public".to_string();
+    leaf_a.table.parent_name = "sharded".to_string();
+
+    let mut leaf_b = make_sharded_table();
+    leaf_b.table.name = "sharded_p2".to_string();
+    leaf_b.table.parent_schema = "public".to_string();
+    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], QueryParserEngine::default());
+    let mut verify = test_server().await;
+    sub.connect().await.unwrap();
+
+    let oid_a = Oid(16384);
+    let oid_b = Oid(16385);
+    let id_a = random_id();
+    let id_b = random_id();
+
+    cleanup(&mut verify, "public.sharded", &[&id_a, &id_b]).await;
+
+    // Each leaf has its own oid in the WAL stream but resolves to the same
+    // destination table via parent_schema/parent_name.
+    let mut relation_a = sharded_relation(oid_a);
+    relation_a.name = "sharded_p1".to_string();
+    let mut relation_b = sharded_relation(oid_b);
+    relation_b.name = "sharded_p2".to_string();
+
+    sub.handle(begin_copy_data(100)).await.unwrap();
+    sub.handle(xlog_copy_data(relation_a.to_bytes().unwrap()))
+        .await
+        .unwrap();
+    sub.handle(xlog_copy_data(relation_b.to_bytes().unwrap()))
+        .await
+        .unwrap();
+
+    sub.handle(insert_copy_data(oid_a, &id_a, "leaf_a"))
+        .await
+        .unwrap();
+    sub.handle(insert_copy_data(oid_b, &id_b, "leaf_b"))
+        .await
+        .unwrap();
+
+    sub.handle(commit_copy_data(200)).await.unwrap();
+
+    // Both inserts must land in the shared destination table. Before the fix,
+    // leaf_b's row would be silently dropped.
+    assert_eq!(count_row(&mut verify, "public.sharded", &id_a).await, 1);
+    assert_eq!(count_row(&mut verify, "public.sharded", &id_b).await, 1);
+
+    cleanup(&mut verify, "public.sharded", &[&id_a, &id_b]).await;
+}
+
 // ── Data flow tests ─────────────────────────────────────────────────
 
 /// Full transaction: Begin → Relation → Insert → Commit, verified in Postgres.