fix(resharding): use replica_session role and disable FK validation

integration/copy_data/data_sync/pgbench.sql

@@ -1,6 +1,22 @@
 \set tenant_id random(1, 20)
 \set order_amount random(1, 50000) / 100.0
 
+BEGIN;
+SELECT -txid_current()::bigint AS fk_id \gset
+INSERT INTO copy_data.users (id, tenant_id, email)
+VALUES (:fk_id, :tenant_id, 'fk_' || :fk_id || '@example.com');
+INSERT INTO copy_data.log_actions (id, tenant_id, user_id, action)
+VALUES (:fk_id, :tenant_id, :fk_id, 'fk_linked');
+COMMIT;
+
+BEGIN;
+UPDATE copy_data.log_actions
+SET user_id = NULL, action = 'fk_released'
+WHERE id = :fk_id AND tenant_id = :tenant_id;
+DELETE FROM copy_data.users
+WHERE id = :fk_id AND tenant_id = :tenant_id;
+COMMIT;
+
 -- Upsert a persistent user for this transaction. txid_current() is stable
 -- within the transaction and unique enough across the benchmark run.
 INSERT INTO copy_data.users (id, tenant_id, email, settings)

integration/copy_data/data_sync/run.sh

@@ -167,11 +167,11 @@ psql -d "${SRC_DB}" -c "UPDATE copy_data.full_identity_events SET label = 'dup_c
 # PgDog must fill body from old_full before routing (P1 fix) and before building the INSERT.
 psql -d "${SRC_DB}" -c "UPDATE copy_data.full_identity_events SET tenant_id = 3 WHERE seq = 1"
 
+stop_pgbench
 # REPLICATION SENTINEL — must be the last DML issued against the source.
 # Updating this row to 'sentinel_done' produces a WAL record that is downstream of
 # every preceding change. The poll loop below waits for it to land on the destination.
 psql -d "${SRC_DB}" -c "UPDATE copy_data.full_identity_events SET label = 'sentinel_done' WHERE seq = 999"
-stop_pgbench
 
 # Wait for the replication sentinel to land on the destination.
 # seq=999 is dedicated solely to this purpose — see setup.sql.

integration/copy_data/setup.sql

@@ -39,10 +39,15 @@ ALTER TABLE copy_data.order_items REPLICA IDENTITY FULL;
 CREATE TABLE IF NOT EXISTS copy_data.log_actions(
     id BIGSERIAL PRIMARY KEY,
     tenant_id BIGINT,
+    user_id BIGINT,
     action VARCHAR,
-    created_at TIMESTAMPTZ NOT NULL DEFAULT now()
+    created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
+    CONSTRAINT log_actions_user_fk FOREIGN KEY (user_id, tenant_id)
+        REFERENCES copy_data.users (id, tenant_id)
 );
 
+CREATE INDEX log_actions_user_idx ON copy_data.log_actions (user_id, tenant_id);
+
 CREATE TABLE copy_data.with_identity(
     id BIGINT PRIMARY KEY GENERATED ALWAYS AS identity,
     tenant_id BIGINT NOT NULL
@@ -181,6 +186,11 @@ SELECT
     ] AS action
 FROM generate_series(1, 100000);
 
+INSERT INTO copy_data.log_actions (tenant_id, user_id, action)
+SELECT tenant_id, id, 'fk_seed'
+FROM copy_data.users
+WHERE id <= 100;
+
 
 INSERT INTO copy_data.with_identity (tenant_id)
 SELECT floor(random() * 10000)::bigint FROM generate_series(1, 100000);

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(reason);
+        let options = pool.server_options();
 
         let mut error = Error::ServerError;
         let now = Instant::now();
@@ -455,6 +455,7 @@ impl Monitor {
             .await
             {
                 Ok(Ok(mut conn)) => {
+                    conn.stats_mut().set_pool_id(pool.id());
                     let elapsed = now.elapsed();
                     {
                         let mut guard = pool.lock();

pgdog/src/backend/pool/pool_impl.rs

@@ -435,61 +435,45 @@ impl Pool {
         &self.inner.config
     }
 
+    pub(crate) fn oids(&self) -> &Arc<Oids> {
+        &self.inner.oids
+    }
+
     /// Get startup parameters for new server connections.
-    pub(super) fn server_options(&self, reason: ConnectReason) -> ServerOptions {
-        let mut params = vec![
-            Parameter {
-                name: "application_name".into(),
-                value: "PgDog".into(),
-            },
-            Parameter {
-                name: "client_encoding".into(),
-                value: "utf-8".into(),
-            },
-        ];
+    pub(super) fn server_options(&self) -> ServerOptions {
+        let mut options = ServerOptions::default();
 
         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 {
+            options.add(Parameter {
                 name: "statement_timeout".into(),
                 value: statement_timeout.as_millis().to_string().into(),
             });
         }
 
-        if let Some(lock_timeout) = lock_timeout {
-            params.push(Parameter {
+        if let Some(lock_timeout) = config.lock_timeout {
+            options.add(Parameter {
                 name: "lock_timeout".into(),
                 value: lock_timeout.as_millis().to_string().into(),
             });
         }
 
         if config.replication_mode {
-            params.push(Parameter {
+            options.add(Parameter {
                 name: "replication".into(),
                 value: "database".into(),
             });
         }
 
         if config.read_only {
-            params.push(Parameter {
+            options.add(Parameter {
                 name: "default_transaction_read_only".into(),
                 value: "on".into(),
             });
         }
 
-        ServerOptions {
-            params,
-            pool_id: self.id(),
-        }
+        options
     }
 
     /// Pool state.

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

@@ -72,6 +72,18 @@ pub(crate) enum Error {
     #[error("backend: {0}")]
     Backend(#[from] crate::backend::Error),
 
+    #[error(
+        "Resharding connection denied for user \"{user}\" on database \"{database}\": {source}.
+    Check and update user permissions:
+    Resharding requires SET ON PARAMETER session_replication_role, or a superuser."
+    )]
+    ReshardingPermissionDenied {
+        user: String,
+        database: String,
+        #[source]
+        source: Box<crate::backend::Error>,
+    },
+
     #[error("pool: {0}")]
     Pool(#[from] crate::backend::pool::Error),
 

pgdog/src/backend/replication/logical/publisher/queries.rs

@@ -8,13 +8,20 @@ use std::{collections::HashSet, fmt::Display};
 use pgdog_postgres_types::Oid;
 
 use crate::{
-    backend::Server,
+    backend::{Error as BackendError, Server},
     net::{DataRow, Format},
 };
 
 use super::super::Error;
 use crate::util::sql::quote_literal;
 
+pub(crate) async fn server_version(server: &mut Server) -> Result<Option<i64>, BackendError> {
+    let versions: Vec<i64> = server
+        .fetch_all("SELECT current_setting('server_version_num')::bigint")
+        .await?;
+    Ok(versions.into_iter().next())
+}
+
 /// Get list of tables in publication.
 static TABLES: &str = "SELECT DISTINCT
   n.nspname,
@@ -280,6 +287,16 @@ mod test {
 
     use super::*;
 
+    #[tokio::test]
+    async fn server_version_returns_value() -> Result<(), BackendError> {
+        let mut server = test_server().await;
+        let version = server_version(&mut server)
+            .await?
+            .expect("PostgreSQL returns its server version");
+        assert!((16_000..1_000_000).contains(&version));
+        Ok(())
+    }
+
     #[test]
     fn test_replica_identity_decodes_oid_above_i32_max() {
         // Regression for issue #847: pg_class.oid is unsigned 32-bit, and in

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

@@ -2,6 +2,7 @@
 
 use std::sync::Arc;
 
+use futures::future::join_all;
 use pgdog_config::RewriteMode;
 use tracing::{info, warn};
 
@@ -12,6 +13,8 @@ use crate::backend::{
     schema::sync::{PgDump, SchemaSyncError, Statement, pg_dump::PgDumpOutput},
 };
 
+use super::publisher::queries::server_version;
+
 /// Sync the schema from a source database to a destination, one phase at a
 /// time.
 #[derive(Debug)]
@@ -43,10 +46,33 @@ impl SchemaSync {
         Ok(())
     }
 
-    /// Dump the source schema.
+    /// Dump the source schema for the oldest known destination version.
+    ///
+    /// PostgreSQL 17 rejects NOT VALID foreign keys on partitioned tables:
+    /// <https://www.postgresql.org/docs/17/sql-altertable.html#SQL-ALTERTABLE-DESC-ADD-TABLE-CONSTRAINT>.
+    /// PostgreSQL 18 adds support:
+    /// <https://www.postgresql.org/docs/18/release-18.html#RELEASE-18-CONSTRAINTS>.
+    /// Keep normal constraint validation for older destinations.
     pub(crate) async fn dump(&self) -> Result<Arc<PgDumpOutput>, SchemaSyncError> {
+        let destination_version = join_all((0..self.destination.shards().len()).map(
+            |shard| async move {
+                let mut primary = self
+                    .destination
+                    .primary(shard, &Request::default())
+                    .await
+                    .ok()?;
+                server_version(&mut primary).await.ok().flatten()
+            },
+        ))
+        .await
+        .into_iter()
+        .flatten()
+        .min()
+        .unwrap_or(0);
         Ok(Arc::new(
-            PgDump::new(&self.source, &self.publication).dump().await?,
+            PgDump::new(&self.source, &self.publication)
+                .dump(destination_version)
+                .await?,
         ))
     }
 

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

@@ -12,8 +12,7 @@ use crate::frontend::client::query_engine::two_pc::{
 
 use crate::frontend::router::parser::Error as ParseError;
 use crate::{
-    backend::{Cluster, ConnectReason, replication::subscriber::ParallelConnection},
-    config::Role,
+    backend::{Cluster, replication::subscriber::ParallelConnection},
     frontend::router::parser::{CopyParser, Shard},
     net::{
         CopyData, CopyDone, ErrorResponse, FromBytes, Message, Protocol, ProtocolMessage, Query,
@@ -22,6 +21,7 @@ use crate::{
 };
 
 use super::super::{CopyStatement, Error};
+use super::connect_primary;
 
 // Not really needed, but we're currently
 // sharding 3 CopyData messages at a time.
@@ -73,15 +73,8 @@ impl CopySubscriber {
     pub(crate) async fn connect(&mut self) -> Result<(), Error> {
         let mut servers = vec![];
         for shard in self.cluster.shards() {
-            let primary = shard
-                .pools_with_roles()
-                .iter()
-                .find(|(role, _)| role == &Role::Primary)
-                .ok_or(Error::NoPrimary)?
-                .1
-                .standalone(ConnectReason::Resharding)
-                .await?;
-            servers.push(ParallelConnection::new(primary)?);
+            let server = connect_primary(shard).await?;
+            servers.push(ParallelConnection::new(server)?);
         }
 
         self.connections = servers;

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

@@ -1,3 +1,12 @@
+use std::sync::Arc;
+
+use crate::{
+    backend::{ConnectReason, Error as BackendError, Server, ServerOptions, pool::Shard},
+    config::Role,
+};
+
+use super::Error;
+
 pub(crate) mod context;
 pub(crate) mod copy;
 pub(crate) mod parallel_connection;
@@ -11,3 +20,29 @@ pub(crate) use context::StreamContext;
 pub(crate) use copy::CopySubscriber;
 pub(crate) use parallel_connection::ParallelConnection;
 pub(crate) use pipeline::PipelinedConnection;
+
+async fn connect_primary(shard: &Shard) -> Result<Server, Error> {
+    let pools = shard.pools_with_roles();
+    let (_, primary) = pools
+        .iter()
+        .find(|(role, _)| role == &Role::Primary)
+        .ok_or(Error::NoPrimary)?;
+    match Box::pin(Server::connect(
+        primary.addr(),
+        ServerOptions::new_resharding(primary.config()),
+        ConnectReason::Resharding,
+        Arc::clone(primary.oids()),
+    ))
+    .await
+    {
+        Ok(server) => Ok(server),
+        Err(BackendError::ConnectionError(error)) if error.code == "42501" => {
+            Err(Error::ReshardingPermissionDenied {
+                user: primary.addr().user.clone(),
+                database: primary.addr().database_name.clone(),
+                source: Box::new(BackendError::ConnectionError(error)),
+            })
+        }
+        Err(error) => Err(error.into()),
+    }
+}

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

@@ -18,13 +18,12 @@ use super::super::publisher::{NonIdentityColumnsPresence, tables_missing_unique_
 use super::super::{
     Error, TableValidationError, TableValidationErrorKind, ensure_validation, publisher::Table,
 };
-use super::PipelinedConnection;
 use super::StreamContext;
+use super::{PipelinedConnection, connect_primary};
 use crate::net::messages::replication::logical::tuple_data::{Identifier, TupleData};
 use crate::net::messages::replication::logical::update::Update as XLogUpdate;
 use crate::{
-    backend::{Cluster, ConnectReason, Server},
-    config::Role,
+    backend::{Cluster, Server},
     frontend::router::parser::Shard,
     net::{
         Bind, CopyData, ErrorResponse, FromBytes, Parse, Protocol, Sync, ToBytes,
@@ -172,15 +171,7 @@ impl StreamSubscriber {
         let mut conns: Vec<Server> = vec![];
 
         for shard in self.cluster.shards() {
-            let primary = shard
-                .pools_with_roles()
-                .iter()
-                .find(|(r, _)| r == &Role::Primary)
-                .ok_or(Error::NoPrimary)?
-                .1
-                .standalone(ConnectReason::Resharding)
-                .await?;
-            conns.push(primary);
+            conns.push(connect_primary(shard).await?);
         }
 
         // Transaction control statements.