Add update/delete support to logical replication (#603)

Author: levkk

.github/workflows/ci.yml

@@ -113,6 +113,7 @@ jobs:
           sudo -u postgres createdb $USER
           sudo -u postgres psql -c 'ALTER SYSTEM SET max_connections TO 1000;'
           sudo -u postgres psql -c 'ALTER SYSTEM SET max_prepared_transactions TO 1000;'
+          sudo -u postgres psql -c 'ALTER SYSTEM SET wal_level TO logical;'
           sudo service postgresql restart
           bash integration/setup.sh
           sudo apt update && sudo apt install -y python3-virtualenv mold
@@ -163,6 +164,8 @@ jobs:
         run: bash integration/rust/run.sh
       - name: Stop shared PgDog
         run: bash -lc 'source integration/common.sh; stop_pgdog'
+      - name: Data sync
+        run: bash integration/copy_data/dev.sh
       - name: Python
         run: bash integration/python/run.sh
       - name: Load balancer
@@ -190,6 +193,7 @@ jobs:
   plugin-unit-tests:
     runs-on: blacksmith-4vcpu-ubuntu-2404
     continue-on-error: true
+    timeout-minutes: 30
     steps:
       - uses: actions/checkout@v4
       - uses: actions-rs/toolchain@v1
@@ -212,6 +216,7 @@ jobs:
   plugin-integration-tests:
     runs-on: blacksmith-4vcpu-ubuntu-2404
     continue-on-error: true
+    timeout-minutes: 30
     steps:
       - uses: actions/checkout@v4
       - uses: actions-rs/toolchain@v1

integration/copy_data/dev.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+set -e
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+DEFAULT_BIN="${SCRIPT_DIR}/../../target/release/pgdog"
+PGDOG_BIN=${PGDOG_BIN:-$DEFAULT_BIN}
+
+export PGUSER=pgdog
+export PGDATABASE=pgdog
+export PGHOST=127.0.0.1
+export PGPORT=5432
+export PGPASSWORD=pgdog
+
+pushd ${SCRIPT_DIR}
+
+psql -f init.sql
+
+${PGDOG_BIN} schema-sync --from-database source --to-database destination --publication pgdog
+${PGDOG_BIN} data-sync --sync-only --from-database source --to-database destination --publication pgdog --replication-slot copy_data
+popd

integration/copy_data/init.sql

@@ -0,0 +1,8 @@
+\c pgdog1
+DROP SCHEMA IF EXISTS copy_data CASCADE;
+\c pgdog2
+DROP SCHEMA IF EXISTS copy_data CASCADE;
+\c pgdog
+DROP SCHEMA IF EXISTS copy_data CASCADE;
+SELECT pg_drop_replication_slot(slot_name) FROM pg_replication_slots;
+\i setup.sql

integration/copy_data/psql.sh

@@ -0,0 +1,2 @@
+#!/bin/bash
+PGPASSWORD=pgdog psql -h 127.0.0.1 -p 5432 -U pgdog $1

integration/copy_data/setup.sql

@@ -5,7 +5,8 @@ CREATE TABLE IF NOT EXISTS copy_data.users (
    tenant_id BIGINT NOT NULL,
    email VARCHAR NOT NULL,
    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
-   settings JSONB NOT NULL DEFAULT '{}'::jsonb
+   settings JSONB NOT NULL DEFAULT '{}'::jsonb,
+   PRIMARY KEY(id, tenant_id)
 ) PARTITION BY HASH(tenant_id);
 
 CREATE TABLE IF NOT EXISTS copy_data.users_0 PARTITION OF copy_data.users

pgdog-config/src/core.rs

@@ -5,7 +5,7 @@ use std::path::PathBuf;
 use tracing::{info, warn};
 
 use crate::sharding::ShardedSchema;
-use crate::{Memory, PassthoughAuth, PreparedStatements};
+use crate::{Memory, PassthoughAuth, PreparedStatements, RewriteMode};
 
 use super::database::Database;
 use super::error::Error;
@@ -326,6 +326,20 @@ impl Config {
             }
             _ => (),
         }
+
+        if !self.general.two_phase_commit {
+            if self.rewrite.enabled {
+                if self.rewrite.shard_key == RewriteMode::Rewrite {
+                    warn!("rewrite.shard_key=rewrite will apply non-atomic shard-key rewrites; enabling two_phase_commit is strongly recommended"
+                    );
+                }
+
+                if self.rewrite.split_inserts == RewriteMode::Rewrite {
+                    warn!("rewrite.split_inserts=rewrite may commit partial multi-row INSERTs; enabling two_phase_commit is strongly recommended"
+                    );
+                }
+            }
+        }
     }
 
     /// Multi-tenancy is enabled.

pgdog/src/backend/databases.rs

@@ -430,6 +430,7 @@ pub(crate) fn new_pool(
             sharded_tables,
             config.multi_tenant(),
             sharded_schemas,
+            &config.rewrite,
         );
 
         Some((

pgdog/src/backend/pool/cluster.rs

@@ -1,6 +1,7 @@
 //! A collection of replicas and a primary.
 
 use parking_lot::{Mutex, RwLock};
+use pgdog_config::{PreparedStatements, Rewrite};
 use std::sync::{
     atomic::{AtomicBool, Ordering},
     Arc,
@@ -52,6 +53,12 @@ pub struct Cluster {
     two_phase_commit: bool,
     two_phase_commit_auto: bool,
     online: Arc<AtomicBool>,
+    rewrite: Rewrite,
+    prepared_statements: PreparedStatements,
+    dry_run: bool,
+    expanded_explain: bool,
+    pub_sub_channel_size: usize,
+    query_parser_enabled: bool,
 }
 
 /// Sharding configuration from the cluster.
@@ -94,6 +101,12 @@ pub struct ClusterConfig<'a> {
     pub two_pc: bool,
     pub two_pc_auto: bool,
     pub sharded_schemas: ShardedSchemas,
+    pub rewrite: &'a Rewrite,
+    pub prepared_statements: &'a PreparedStatements,
+    pub dry_run: bool,
+    pub expanded_explain: bool,
+    pub pub_sub_channel_size: usize,
+    pub query_parser_enabled: bool,
 }
 
 impl<'a> ClusterConfig<'a> {
@@ -104,6 +117,7 @@ impl<'a> ClusterConfig<'a> {
         sharded_tables: ShardedTables,
         multi_tenant: &'a Option<MultiTenant>,
         sharded_schemas: ShardedSchemas,
+        rewrite: &'a Rewrite,
     ) -> Self {
         Self {
             name: &user.database,
@@ -126,6 +140,12 @@ impl<'a> ClusterConfig<'a> {
                 .two_phase_commit_auto
                 .unwrap_or(general.two_phase_commit_auto.unwrap_or(false)), // Disable by default.
             sharded_schemas,
+            rewrite,
+            prepared_statements: &general.prepared_statements,
+            dry_run: general.dry_run,
+            expanded_explain: general.expanded_explain,
+            pub_sub_channel_size: general.pub_sub_channel_size,
+            query_parser_enabled: general.query_parser_enabled,
         }
     }
 }
@@ -150,6 +170,12 @@ impl Cluster {
             two_pc,
             two_pc_auto,
             sharded_schemas,
+            rewrite,
+            prepared_statements,
+            dry_run,
+            expanded_explain,
+            pub_sub_channel_size,
+            query_parser_enabled,
         } = config;
 
         Self {
@@ -175,9 +201,23 @@ impl Cluster {
             two_phase_commit: two_pc && shards.len() > 1,
             two_phase_commit_auto: two_pc_auto && shards.len() > 1,
             online: Arc::new(AtomicBool::new(false)),
+            rewrite: rewrite.clone(),
+            prepared_statements: prepared_statements.clone(),
+            dry_run,
+            expanded_explain,
+            pub_sub_channel_size,
+            query_parser_enabled,
         }
     }
 
+    /// Change config to work with logical replication streaming.
+    pub fn logical_stream(&self) -> Self {
+        let mut cluster = self.clone();
+        // Disable rewrites, we are only sending valid statements.
+        cluster.rewrite.enabled = false;
+        cluster
+    }
+
     /// Get a connection to a primary of the given shard.
     pub async fn primary(&self, shard: usize, request: &Request) -> Result<Guard, Error> {
         let shard = self.shards.get(shard).ok_or(Error::NoShard(shard))?;
@@ -251,6 +291,31 @@ impl Cluster {
         self.sharded_tables.tables()
     }
 
+    /// Get query rewrite config.
+    pub fn rewrite(&self) -> &Rewrite {
+        &self.rewrite
+    }
+
+    pub fn query_parser_enabled(&self) -> bool {
+        self.query_parser_enabled
+    }
+
+    pub fn prepared_statements(&self) -> &PreparedStatements {
+        &self.prepared_statements
+    }
+
+    pub fn dry_run(&self) -> bool {
+        self.dry_run
+    }
+
+    pub fn expanded_explain(&self) -> bool {
+        self.expanded_explain
+    }
+
+    pub fn pub_sub_enabled(&self) -> bool {
+        self.pub_sub_channel_size > 0
+    }
+
     /// Find sharded column position, if the table and columns match the configuration.
     pub fn sharded_column(&self, table: &str, columns: &[&str]) -> Option<ShardedColumn> {
         self.sharded_tables.sharded_column(table, columns)
@@ -412,10 +477,12 @@ mod test {
     use std::sync::Arc;
 
     use crate::{
-        backend::pool::{Address, Config, PoolConfig},
-        backend::{Shard, ShardedTables},
+        backend::{
+            pool::{Address, Config, PoolConfig},
+            Shard, ShardedTables,
+        },
         config::{
-            DataType, Hasher, LoadBalancingStrategy, ReadWriteSplit, ReadWriteStrategy,
+            config, DataType, Hasher, LoadBalancingStrategy, ReadWriteSplit, ReadWriteStrategy,
             ShardedTable,
         },
     };
@@ -424,6 +491,7 @@ mod test {
 
     impl Cluster {
         pub fn new_test() -> Self {
+            let config = config();
             Cluster {
                 sharded_tables: ShardedTables::new(
                     vec![ShardedTable {
@@ -470,6 +538,13 @@ mod test {
                     user: "pgdog".into(),
                     database: "pgdog".into(),
                 }),
+                prepared_statements: config.config.general.prepared_statements,
+                dry_run: config.config.general.dry_run,
+                expanded_explain: config.config.general.expanded_explain,
+                query_parser_enabled: config.config.general.query_parser_enabled,
+                rewrite: config.config.rewrite.clone(),
+                two_phase_commit: config.config.general.two_phase_commit,
+                two_phase_commit_auto: config.config.general.two_phase_commit_auto.unwrap_or(false),
                 ..Default::default()
             }
         }

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

@@ -6,13 +6,16 @@ use crate::{backend::replication::publisher::PublicationTable, net::ErrorRespons
 
 #[derive(Debug, Error)]
 pub enum Error {
-    #[error("{0}")]
+    #[error("backend: {0}")]
     Backend(#[from] crate::backend::Error),
 
-    #[error("{0}")]
+    #[error("pool: {0}")]
     Pool(#[from] crate::backend::pool::Error),
 
-    #[error("{0}")]
+    #[error("router: {0}")]
+    Router(#[from] crate::frontend::router::Error),
+
+    #[error("net: {0}")]
     Net(#[from] crate::net::Error),
 
     #[error("transaction not started")]
@@ -21,6 +24,12 @@ pub enum Error {
     #[error("out of sync, got {0}")]
     OutOfSync(char),
 
+    #[error("out of sync during commit, got {0}")]
+    CommitOutOfSync(char),
+
+    #[error("out of sync during relation prepare, got {0}")]
+    RelationOutOfSync(char),
+
     #[error("missing data")]
     MissingData,
 
@@ -30,7 +39,7 @@ pub enum Error {
     #[error("copy error")]
     Copy,
 
-    #[error("{0}")]
+    #[error("pg_error: {0}")]
     PgError(Box<ErrorResponse>),
 
     #[error("table \"{0}\".\"{1}\" has no replica identity")]
@@ -39,13 +48,16 @@ pub enum Error {
     #[error("lsn decode")]
     LsnDecode,
 
+    #[error("replication slot \"{0}\" doesn't exist, but it should")]
+    MissingReplicationSlot(String),
+
     #[error("parse int")]
     ParseInt(#[from] ParseIntError),
 
     #[error("shard has no primary")]
     NoPrimary,
 
-    #[error("{0}")]
+    #[error("parser: {0}")]
     Parser(#[from] crate::frontend::router::parser::Error),
 
     #[error("not connected")]
@@ -68,6 +80,9 @@ pub enum Error {
 
     #[error("table {0} doesn't have a primary key")]
     NoPrimaryKey(PublicationTable),
+
+    #[error("router returned incorrect command")]
+    IncorrectCommand,
 }
 
 impl From<ErrorResponse> for Error {