diff --git a/integration/copy_data/data_sync/pgbench.sql b/integration/copy_data/data_sync/pgbench.sql index ca13fa3c7..6e8a8319c 100644 --- a/integration/copy_data/data_sync/pgbench.sql +++ b/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) diff --git a/integration/copy_data/data_sync/run.sh b/integration/copy_data/data_sync/run.sh index 5897af2a0..840984de8 100755 --- a/integration/copy_data/data_sync/run.sh +++ b/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. diff --git a/integration/copy_data/setup.sql b/integration/copy_data/setup.sql index f9894e61b..6ccbd6de6 100644 --- a/integration/copy_data/setup.sql +++ b/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); diff --git a/pgdog/src/backend/pool/monitor.rs b/pgdog/src/backend/pool/monitor.rs index 0bd7d89f8..465dbbcb8 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(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(); diff --git a/pgdog/src/backend/pool/pool_impl.rs b/pgdog/src/backend/pool/pool_impl.rs index a3a800f4c..57daf58dd 100644 --- a/pgdog/src/backend/pool/pool_impl.rs +++ b/pgdog/src/backend/pool/pool_impl.rs @@ -435,61 +435,45 @@ impl Pool { &self.inner.config } + pub(crate) fn oids(&self) -> &Arc { + &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. diff --git a/pgdog/src/backend/replication/logical/error.rs b/pgdog/src/backend/replication/logical/error.rs index 7cdf69efb..6e0124e6e 100644 --- a/pgdog/src/backend/replication/logical/error.rs +++ b/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, + }, + #[error("pool: {0}")] Pool(#[from] crate::backend::pool::Error), diff --git a/pgdog/src/backend/replication/logical/publisher/queries.rs b/pgdog/src/backend/replication/logical/publisher/queries.rs index 4d199ae0c..a377a4948 100644 --- a/pgdog/src/backend/replication/logical/publisher/queries.rs +++ b/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, BackendError> { + let versions: Vec = 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 diff --git a/pgdog/src/backend/replication/logical/schema_sync.rs b/pgdog/src/backend/replication/logical/schema_sync.rs index 67fe82510..d84bdf56b 100644 --- a/pgdog/src/backend/replication/logical/schema_sync.rs +++ b/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: + /// . + /// PostgreSQL 18 adds support: + /// . + /// Keep normal constraint validation for older destinations. pub(crate) async fn dump(&self) -> Result, 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?, )) } diff --git a/pgdog/src/backend/replication/logical/subscriber/copy.rs b/pgdog/src/backend/replication/logical/subscriber/copy.rs index bb2c96c01..6d168a617 100644 --- a/pgdog/src/backend/replication/logical/subscriber/copy.rs +++ b/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; diff --git a/pgdog/src/backend/replication/logical/subscriber/mod.rs b/pgdog/src/backend/replication/logical/subscriber/mod.rs index 45e55fc3f..249fd5b16 100644 --- a/pgdog/src/backend/replication/logical/subscriber/mod.rs +++ b/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 { + 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()), + } +} diff --git a/pgdog/src/backend/replication/logical/subscriber/stream.rs b/pgdog/src/backend/replication/logical/subscriber/stream.rs index 82587e1f3..903382060 100644 --- a/pgdog/src/backend/replication/logical/subscriber/stream.rs +++ b/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 = 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. diff --git a/pgdog/src/backend/replication/logical/subscriber/tests.rs b/pgdog/src/backend/replication/logical/subscriber/tests.rs index a8d3dda5d..6ecd1f89c 100644 --- a/pgdog/src/backend/replication/logical/subscriber/tests.rs +++ b/pgdog/src/backend/replication/logical/subscriber/tests.rs @@ -1,11 +1,12 @@ use bytes::Bytes; +use pgdog_config::User; use pgdog_postgres_types::Oid; use rand::Rng; use crate::{ backend::{ - Server, - pool::cluster::Cluster, + Error as BackendError, Server, + pool::{Address, ClusterConfig, ClusterShardConfig, PoolConfig, cluster::Cluster}, replication::logical::publisher::{ Lsn, PublicationTable, PublicationTableColumn, ReplicaIdentity, Table, }, @@ -29,7 +30,103 @@ use crate::{ }, }; -use super::stream::StreamSubscriber; +use super::{Error, connect_primary, stream::StreamSubscriber}; + +fn resharding_test_cluster(user: &str) -> Cluster { + let user = User::new(user, "pgdog", "pgdog"); + let shards = [ClusterShardConfig { + primary: Some(PoolConfig { + address: Address { + user: user.name.clone(), + passwords: vec!["pgdog".into()], + ..Address::new_test() + }, + config: Default::default(), + }), + replicas: vec![], + }]; + Cluster::new(ClusterConfig::new( + &config().config, + &user, + &shards, + Default::default(), + Default::default(), + Default::default(), + Default::default(), + )) +} + +#[tokio::test] +async fn resharding_connection_rejects_unprivileged_user() -> Result<(), Box> +{ + let user = format!("resharding_denied_{}", random_id()); + let mut admin = test_server().await; + admin + .execute_checked(format!( + "CREATE ROLE {user} LOGIN NOSUPERUSER NOINHERIT PASSWORD 'pgdog'" + )) + .await?; + let result = async { + admin + .execute_checked(format!("GRANT CONNECT ON DATABASE pgdog TO {user}")) + .await?; + let cluster = resharding_test_cluster(&user); + let result = connect_primary(cluster.shards().first().ok_or(Error::MissingData)?) + .await + .map(drop); + Ok::<_, Box>(result) + } + .await; + admin + .execute_checked(format!( + "REVOKE CONNECT ON DATABASE pgdog FROM {user}; DROP ROLE {user}" + )) + .await?; + let error = result?.expect_err("unprivileged user must not connect for resharding"); + assert!(!error.is_retryable()); + assert!(matches!( + error, + Error::ReshardingPermissionDenied { source, .. } + if matches!(source.as_ref(), BackendError::ConnectionError(response) if response.code == "42501") + )); + Ok(()) +} + +#[tokio::test] +async fn resharding_connection_accepts_privileged_user() -> Result<(), Box> { + let user = format!("resharding_allowed_{}", random_id()); + let mut admin = test_server().await; + admin + .execute_checked(format!( + "CREATE ROLE {user} LOGIN NOSUPERUSER NOINHERIT PASSWORD 'pgdog'" + )) + .await?; + let result = async { + admin + .execute_checked(format!( + "GRANT CONNECT ON DATABASE pgdog TO {user}; \ + GRANT SET ON PARAMETER session_replication_role TO {user}" + )) + .await?; + let cluster = resharding_test_cluster(&user); + let mut server = + connect_primary(cluster.shards().first().ok_or(Error::MissingData)?).await?; + let users: Vec = server.fetch_all("SELECT current_user").await?; + let roles: Vec = server.fetch_all("SHOW session_replication_role").await?; + Ok::<_, Box>((users, roles)) + } + .await; + admin + .execute_checked(format!( + "REVOKE SET ON PARAMETER session_replication_role FROM {user}; \ + REVOKE CONNECT ON DATABASE pgdog FROM {user}; DROP ROLE {user}" + )) + .await?; + let (users, roles) = result?; + assert_eq!(users, [user]); + assert_eq!(roles, ["replica"]); + Ok(()) +} fn random_id() -> String { rand::rng() diff --git a/pgdog/src/backend/replication/mod.rs b/pgdog/src/backend/replication/mod.rs index 3acd43360..1186e6acb 100644 --- a/pgdog/src/backend/replication/mod.rs +++ b/pgdog/src/backend/replication/mod.rs @@ -7,3 +7,5 @@ pub(crate) use error::Error; pub(crate) use logical::*; pub(crate) use sharded_schema::*; pub(crate) use sharded_tables::ShardedTables; +#[cfg(test)] +mod tests; diff --git a/pgdog/src/backend/replication/tests.rs b/pgdog/src/backend/replication/tests.rs new file mode 100644 index 000000000..efe487e5b --- /dev/null +++ b/pgdog/src/backend/replication/tests.rs @@ -0,0 +1,661 @@ +use std::time::Duration; + +use pgdog_config::{ConfigAndUsers, Database, ShardedTableConfig, User}; +use tokio_util::sync::CancellationToken; + +use super::logical::{Error, data_sync::DataSync, publisher::publisher_impl::Publisher}; +use crate::{ + api::{ + run_task, + schema_sync::{SchemaSyncPhase, SchemaSyncTask}, + task::TaskError, + }, + backend::{ + Cluster, ConnectReason, Error as BackendError, Server, ServerOptions, databases, + pool::{Address, Request}, + schema::sync::SchemaSyncError, + server::test::test_server, + }, + config::{config, set}, +}; + +async fn setup_replication_test( + admin: &mut Server, + schema: &str, + destination: &str, +) -> Result<(), Box> { + let mut test_config = ConfigAndUsers::default(); + test_config.config.replication.pg_dump_path = config().config.replication.pg_dump_path.clone(); + for database in [schema, destination] { + admin + .execute_checked(format!("CREATE DATABASE {database} TEMPLATE template0")) + .await?; + test_config.config.databases.push(Database { + name: database.into(), + host: "127.0.0.1".into(), + port: 5432, + ..Default::default() + }); + let mut user = User::new("pgdog", "pgdog", database); + user.schema_admin = true; + test_config.users.users.push(user); + for table in ["parents", "children"] { + test_config.config.sharded_tables.push(ShardedTableConfig { + database: database.into(), + name: Some(table.into()), + schema: Some(schema.into()), + column: "tenant_id".into(), + ..Default::default() + }); + } + } + set(test_config)?; + databases::init()?; + Ok(()) +} + +async fn replicate_until_caught_up( + publisher: &mut Publisher, + source: &Cluster, + destination: &Cluster, + slot_name: &str, +) -> Result<(), Box> { + let mut server = source.primary(0, &Request::default()).await?; + let target: Vec = server + .fetch_all("SELECT pg_current_wal_lsn()::text") + .await?; + let target = target.first().ok_or(Error::MissingData)?; + let query = format!( + "SELECT 1::bigint FROM pg_replication_slots \ + WHERE slot_name = '{slot_name}_0' \ + AND confirmed_flush_lsn >= '{target}'::pg_lsn" + ); + let mut waiter = publisher.replicate(source, destination).await?; + let caught_up = tokio::select! { + result = waiter.wait() => { + result?; + return Err(Error::MissingData.into()); + } + result = tokio::time::timeout(Duration::from_secs(10), async { + loop { + let caught_up: Vec = server.fetch_all(&query).await?; + if caught_up == [1] { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + Ok::<_, Box>(()) + }) => result, + }; + waiter.stop(); + waiter.wait().await?; + caught_up??; + Ok(()) +} + +async fn cleanup_replication_test( + publisher: &mut Publisher, + admin: &mut Server, + original_config: &ConfigAndUsers, + test_databases: [&str; 2], +) -> Result<(), Box> { + let cleanup = Box::pin(async { + publisher.cleanup().await?; + for database in test_databases { + let slots: Vec = admin + .fetch_all(format!( + "SELECT slot_name FROM pg_replication_slots WHERE database = '{database}'" + )) + .await?; + if slots.is_empty() { + continue; + } + let mut server = Server::connect( + &Address { + database_name: database.into(), + ..Address::new_test() + }, + ServerOptions::new_replication(), + ConnectReason::Resharding, + Default::default(), + ) + .await?; + for slot in slots { + match tokio::time::timeout( + Duration::from_secs(10), + server.execute_checked(format!("DROP_REPLICATION_SLOT {slot} WAIT")), + ) + .await? + { + Ok(_) => {} + Err(BackendError::ExecutionError(error)) if error.code == "42704" => {} + Err(error) => return Err(error.into()), + } + } + } + Ok::<_, Box>(()) + }) + .await; + set(original_config.clone())?; + databases::reload_from_existing()?; + for database in test_databases { + admin + .execute_checked(format!("DROP DATABASE IF EXISTS {database} WITH (FORCE)")) + .await?; + } + cleanup +} + +// Verify the case when the data related to fk update happened during tables +// copy. We can hit the constraint violations when between the related +// tables copies the updates to fk rows happened. Since we copy the table +// with specific snapshot at the moment of copy start and we doesn't start +// copying at the same moment, we can get the snapshots related to different +// moment of times and they could diverge in cross relations. +// Hitting the error +// `insert or update on table \"children\" violates foreign key constraint \"children_parent_fk\"", +// detail: Some("Key (parent_id)=(1) is not present in table \"parents\"` +#[tokio::test] +async fn test_replication_fk_conflicts_after_delete_during_copy() +-> Result<(), Box> { + let schema = "fk_replication_test".to_owned(); + let destination = format!("{schema}_dest"); + let parent = format!("{schema}.parents"); + let child = format!("{schema}.children"); + let original_config = config(); + let mut admin = test_server().await; + let mut publisher = Publisher::new(&schema, schema.clone()); + let result = async { + setup_replication_test(&mut admin, &schema, &destination).await?; + let source = databases::databases().schema_owner(&schema)?; + let mut source_server = source.primary(0, &Request::default()).await?; + + // create the tables tied by fk and put some data there + source_server + .execute_checked(format!( + "CREATE SCHEMA {schema}; \ + CREATE TABLE {parent} (id BIGINT PRIMARY KEY, tenant_id BIGINT NOT NULL); \ + CREATE TABLE {child} (id BIGINT PRIMARY KEY, tenant_id BIGINT NOT NULL, \ + parent_id BIGINT, CONSTRAINT children_parent_fk \ + FOREIGN KEY (parent_id) REFERENCES {parent}(id)); \ + INSERT INTO {parent} VALUES (1, 1), (2, 1), (3, 1); \ + INSERT INTO {child} VALUES (1, 1, 1), (2, 1, 2), (3, 1, 3); \ + CREATE PUBLICATION {schema} FOR TABLE {parent}, {child}" + )) + .await?; + + // reproduce the flow ran by pgdog - schema sync, data copy, replication. + let schema_sync = SchemaSyncTask::builder() + .databases(pgdog_stats::Databases { + source: schema.clone(), + destination: destination.clone(), + }) + .publication(schema.clone()); + // copy initial schema + run_task(schema_sync.clone().phase(SchemaSyncPhase::Pre).build()).await?; + let source = databases::databases().schema_owner(&schema)?; + let dest = databases::databases().schema_owner(&destination)?; + let cancel = CancellationToken::new(); + publisher.sync_tables(true, &source).await?; + publisher.create_slots(&source, &cancel).await?; + let tables = publisher.tables.get(&0).ok_or(Error::MissingData)?; + let child_table = tables + .iter() + .find(|table| table.table.name == "children") + .ok_or(Error::MissingData)?; + let parent_table = tables + .iter() + .find(|table| table.table.name == "parents") + .ok_or(Error::MissingData)?; + let sync = DataSync { + source: &source, + dest: &dest, + format: config().config.general.resharding_copy_format, + }; + // copy the child table first, so it won't have updates we'll do during copy + let child_table = sync + .copy_table(child_table, source_server.addr(), &cancel, |_| {}) + .await?; + + // update the fk related data, so it would be present + // only in parent table snapshot + source_server + .execute_checked(format!( + "BEGIN; \ + DELETE FROM {child} WHERE id = 1; \ + DELETE FROM {parent} WHERE id = 1; \ + COMMIT" + )) + .await?; + + // add another pair, then clear the reference and delete its parent + source_server + .execute_checked(format!( + "BEGIN; \ + INSERT INTO {parent} VALUES (4, 1); \ + INSERT INTO {child} VALUES (4, 1, 4); \ + COMMIT" + )) + .await?; + source_server + .execute_checked(format!( + "BEGIN; \ + UPDATE {child} SET parent_id = NULL WHERE id = 4; \ + DELETE FROM {parent} WHERE id = 4; \ + COMMIT" + )) + .await?; + + // and now start the copy of parent table. + // that should have an updated snapshot already with the queries + // executed above. + let parent_table = sync + .copy_table(parent_table, source_server.addr(), &cancel, |_| {}) + .await?; + publisher.post_data_sync([(0, vec![child_table, parent_table])].into()); + drop(source_server); + run_task(schema_sync.clone().phase(SchemaSyncPhase::Post).build()).await?; + + // run the replication and wait for all data to be copied + replicate_until_caught_up(&mut publisher, &source, &dest, &schema).await?; + run_task(schema_sync.phase(SchemaSyncPhase::Cutover).build()).await?; + Ok::<_, Box>(dest) + } + .await; + + let validation = async { + let dest = result?; + let mut server = dest.primary(0, &Request::default()).await?; + let parents: Vec = server + .fetch_all(format!("SELECT id FROM {parent} ORDER BY id")) + .await?; + let children: Vec = server + .fetch_all(format!("SELECT id FROM {child} ORDER BY id")) + .await?; + let parent_ids: Vec = server + .fetch_all(format!( + "SELECT COALESCE(parent_id::text, 'null') FROM {child} ORDER BY id" + )) + .await?; + Ok::<_, Box>((parents, children, parent_ids)) + } + .await; + + cleanup_replication_test( + &mut publisher, + &mut admin, + &original_config, + [&schema, &destination], + ) + .await?; + let (parents, children, parent_ids) = validation?; + assert_eq!(parents, [2, 3]); + assert_eq!(children, [2, 3, 4]); + assert_eq!(parent_ids, ["2", "3", "null"]); + Ok(()) +} + +// Verify the case when the source tables have fk constraints +// and that during the copy this constraints doesn't fire if +// the other table data is not yet present on the destination. +#[tokio::test] +async fn test_replication_fk_constraints_after_copy_child_before_parent() +-> Result<(), Box> { + let schema = "fk_copy_test"; + let destination = "fk_copy_test_dest"; + let original_config = config(); + let mut admin = test_server().await; + let mut publisher = Publisher::new(schema, schema.into()); + let result = async { + setup_replication_test(&mut admin, schema, destination).await?; + let source = databases::databases().schema_owner(schema)?; + let mut server = source.primary(0, &Request::default()).await?; + + // create the tables tied by fk and put some data there + server + .execute_checked(format!( + "CREATE SCHEMA {schema}; \ + CREATE TABLE {schema}.parents (id BIGINT PRIMARY KEY, tenant_id BIGINT NOT NULL); \ + CREATE TABLE {schema}.children (id BIGINT PRIMARY KEY, tenant_id BIGINT NOT NULL, \ + parent_id BIGINT REFERENCES {schema}.parents(id)); \ + INSERT INTO {schema}.parents VALUES (1, 1); \ + INSERT INTO {schema}.children VALUES (1, 1, 1); \ + CREATE PUBLICATION {schema} FOR TABLE {schema}.parents, {schema}.children" + )) + .await?; + let schema_sync = SchemaSyncTask::builder() + .databases(pgdog_stats::Databases { + source: schema.into(), + destination: destination.into(), + }) + .publication(schema.to_owned()); + + // copy initial schema + run_task(schema_sync.clone().phase(SchemaSyncPhase::Pre).build()).await?; + let source = databases::databases().schema_owner(schema)?; + let dest = databases::databases().schema_owner(destination)?; + let cancel = CancellationToken::new(); + publisher.sync_tables(true, &source).await?; + publisher.create_slots(&source, &cancel).await?; + let tables = publisher.tables.get(&0).ok_or(Error::MissingData)?; + let child = tables + .iter() + .find(|table| table.table.name == "children") + .ok_or(Error::MissingData)?; + let parent = tables + .iter() + .find(|table| table.table.name == "parents") + .ok_or(Error::MissingData)?; + let sync = DataSync { + source: &source, + dest: &dest, + format: config().config.general.resharding_copy_format, + }; + + // copy the child table first, while the parent data is not yet present + let child = sync + .copy_table(child, server.addr(), &cancel, |_| {}) + .await?; + + // and now copy the parent table + let parent = sync + .copy_table(parent, server.addr(), &cancel, |_| {}) + .await?; + publisher.post_data_sync([(0, vec![child, parent])].into()); + + // add rows after copy so replication must deliver them + server + .execute_checked(format!( + "BEGIN; \ + INSERT INTO {schema}.parents VALUES (2, 1); \ + INSERT INTO {schema}.children VALUES (2, 1, 2); \ + COMMIT" + )) + .await?; + drop(server); + + run_task(schema_sync.clone().phase(SchemaSyncPhase::Post).build()).await?; + replicate_until_caught_up(&mut publisher, &source, &dest, schema).await?; + run_task(schema_sync.phase(SchemaSyncPhase::Cutover).build()).await?; + Ok::<_, Box>(dest) + } + .await; + + // check that both tables have the copied rows and the fk still points to the parent + let validation = async { + let dest = result?; + let mut server = dest.primary(0, &Request::default()).await?; + let parents: Vec = server + .fetch_all(format!( + "SELECT id || ':' || tenant_id FROM {schema}.parents ORDER BY id" + )) + .await?; + let children: Vec = server + .fetch_all(format!( + "SELECT id || ':' || tenant_id || ':' || parent_id \ + FROM {schema}.children ORDER BY id" + )) + .await?; + Ok::<_, Box>((parents, children)) + } + .await; + + cleanup_replication_test( + &mut publisher, + &mut admin, + &original_config, + [schema, destination], + ) + .await?; + let (parents, children) = validation?; + assert_eq!(parents, ["1:1", "2:1"]); + assert_eq!(children, ["1:1:1", "2:1:2"]); + Ok(()) +} + +// Almost the same like [`test_replication_fk_conflicts_after_update_during_copy`] test, but +// this one uses the user trigger instead of FK. The user trigger should not fail during +// the replication for the fame reason as FK for the mentioned test. +#[tokio::test] +async fn test_replication_copy_custom_parent_trigger() -> Result<(), Box> { + let schema = "trigger_copy_test"; + let destination = "trigger_copy_test_dest"; + let original_config = config(); + let mut admin = test_server().await; + let mut publisher = Publisher::new(schema, schema.into()); + let result = async { + setup_replication_test(&mut admin, schema, destination).await?; + let source = databases::databases().schema_owner(schema)?; + let mut server = source.primary(0, &Request::default()).await?; + // create valid source data with a trigger instead of a foreign key + server + .execute_checked(format!( + "CREATE SCHEMA {schema}; \ + CREATE TABLE {schema}.parents (id BIGINT PRIMARY KEY, tenant_id BIGINT NOT NULL); \ + CREATE TABLE {schema}.children (id BIGINT PRIMARY KEY, tenant_id BIGINT NOT NULL, \ + parent_id BIGINT); \ + CREATE FUNCTION {schema}.check_parent() RETURNS trigger LANGUAGE plpgsql AS $$ \ + BEGIN \ + IF NEW.parent_id IS NOT NULL AND NOT EXISTS \ + (SELECT 1 FROM {schema}.parents WHERE id = NEW.parent_id) THEN \ + RAISE EXCEPTION 'parent is missing' USING ERRCODE = '23503'; \ + END IF; \ + RETURN NEW; \ + END; $$; \ + CREATE TRIGGER children_parent_check BEFORE INSERT OR UPDATE \ + ON {schema}.children FOR EACH ROW EXECUTE FUNCTION {schema}.check_parent(); \ + INSERT INTO {schema}.parents VALUES (1, 1); \ + INSERT INTO {schema}.children VALUES (1, 1, 1); \ + CREATE PUBLICATION {schema} FOR TABLE {schema}.parents, {schema}.children" + )) + .await?; + + // copy the schema and trigger before copying any rows + let schema_sync = SchemaSyncTask::builder() + .databases(pgdog_stats::Databases { + source: schema.into(), + destination: destination.into(), + }) + .publication(schema.to_owned()); + run_task(schema_sync.clone().phase(SchemaSyncPhase::Pre).build()).await?; + let source = databases::databases().schema_owner(schema)?; + let dest = databases::databases().schema_owner(destination)?; + let cancel = CancellationToken::new(); + publisher.sync_tables(true, &source).await?; + publisher.create_slots(&source, &cancel).await?; + let tables = publisher.tables.get(&0).ok_or(Error::MissingData)?; + let child = tables + .iter() + .find(|table| table.table.name == "children") + .ok_or(Error::MissingData)?; + let parent = tables + .iter() + .find(|table| table.table.name == "parents") + .ok_or(Error::MissingData)?; + let sync = DataSync { + source: &source, + dest: &dest, + format: config().config.general.resharding_copy_format, + }; + // copy the child first, while its parent is still missing + let child = sync + .copy_table(child, server.addr(), &cancel, |_| {}) + .await?; + let parent = sync + .copy_table(parent, server.addr(), &cancel, |_| {}) + .await?; + publisher.post_data_sync([(0, vec![child, parent])].into()); + + // add rows after copy so replication must deliver them + server + .execute_checked(format!( + "BEGIN; \ + INSERT INTO {schema}.parents VALUES (2, 1); \ + INSERT INTO {schema}.children VALUES (2, 1, 2); \ + COMMIT" + )) + .await?; + drop(server); + run_task(schema_sync.clone().phase(SchemaSyncPhase::Post).build()).await?; + replicate_until_caught_up(&mut publisher, &source, &dest, schema).await?; + run_task(schema_sync.phase(SchemaSyncPhase::Cutover).build()).await?; + Ok::<_, Box>(dest) + } + .await; + + let validation = async { + let dest = result?; + // check the copied parent and the child's reference + let mut server = dest.primary(0, &Request::default()).await?; + let parents: Vec = server + .fetch_all(format!("SELECT id FROM {schema}.parents ORDER BY id")) + .await?; + let children: Vec = server + .fetch_all(format!( + "SELECT parent_id FROM {schema}.children ORDER BY id" + )) + .await?; + Ok::<_, Box>((parents, children)) + } + .await; + + cleanup_replication_test( + &mut publisher, + &mut admin, + &original_config, + [schema, destination], + ) + .await?; + let (parents, children) = validation?; + assert_eq!(parents, [1, 2]); + assert_eq!(children, [1, 2]); + Ok(()) +} + +// Verify that we catch some data inconsistencies after resharding +// in case we created one. It's created artificially during copy, +// since we don't know for cases when we do this wrong for now. +#[ignore = "No validation for now"] +#[tokio::test] +async fn test_replication_fk_inconsistent_check_on_cutover() +-> Result<(), Box> { + let schema = "fk_post_copy_test"; + let destination = "fk_post_copy_test_dest"; + let original_config = config(); + let mut admin = test_server().await; + let mut publisher = Publisher::new(schema, schema.into()); + let result = async { + setup_replication_test(&mut admin, schema, destination).await?; + let source = databases::databases().schema_owner(schema)?; + let mut server = source.primary(0, &Request::default()).await?; + server + .execute_checked(format!( + "CREATE SCHEMA {schema}; \ + CREATE TABLE {schema}.parents (id BIGINT PRIMARY KEY, tenant_id BIGINT NOT NULL); \ + CREATE TABLE {schema}.children (id BIGINT PRIMARY KEY, tenant_id BIGINT NOT NULL, \ + parent_id BIGINT, CONSTRAINT children_parent_fk \ + FOREIGN KEY (parent_id) REFERENCES {schema}.parents(id)); \ + INSERT INTO {schema}.parents VALUES (1, 1); \ + INSERT INTO {schema}.children VALUES (1, 1, 1); \ + CREATE PUBLICATION {schema} FOR TABLE {schema}.parents, {schema}.children" + )) + .await?; + + // copy the tables without their foreign key + let schema_sync = SchemaSyncTask::builder() + .databases(pgdog_stats::Databases { + source: schema.into(), + destination: destination.into(), + }) + .publication(schema.to_owned()); + run_task(schema_sync.clone().phase(SchemaSyncPhase::Pre).build()).await?; + let source = databases::databases().schema_owner(schema)?; + let cancel = CancellationToken::new(); + publisher.sync_tables(true, &source).await?; + publisher.create_slots(&source, &cancel).await?; + let dest = databases::databases().schema_owner(destination)?; + let tables = publisher.tables.get(&0).ok_or(Error::MissingData)?; + let child = tables + .iter() + .find(|table| table.table.name == "children") + .ok_or(Error::MissingData)?; + let parent = tables + .iter() + .find(|table| table.table.name == "parents") + .ok_or(Error::MissingData)?; + let sync = DataSync { + source: &source, + dest: &dest, + format: config().config.general.resharding_copy_format, + }; + let child = sync + .copy_table(child, server.addr(), &cancel, |_| {}) + .await?; + let mut destination_server = dest.primary(0, &Request::default()).await?; + // leave an orphan that replication cannot repair + destination_server + .execute_checked(format!("INSERT INTO {schema}.children VALUES (42, 1, 999)")) + .await?; + drop(destination_server); + let parent = sync + .copy_table(parent, server.addr(), &cancel, |_| {}) + .await?; + publisher.post_data_sync([(0, vec![child, parent])].into()); + + // add valid rows that must arrive through replication + server + .execute_checked(format!( + "BEGIN; \ + INSERT INTO {schema}.parents VALUES (2, 1); \ + INSERT INTO {schema}.children VALUES (2, 1, 2); \ + COMMIT" + )) + .await?; + drop(server); + + run_task(schema_sync.clone().phase(SchemaSyncPhase::Post).build()).await?; + + // wait for the valid source rows to arrive without repairing the orphan + replicate_until_caught_up(&mut publisher, &source, &dest, schema).await?; + + // cutover should reject the orphan left on the destination + let cutover = run_task(schema_sync.phase(SchemaSyncPhase::Cutover).build()).await; + Ok::<_, Box>((dest, cutover)) + } + .await; + + let validation = async { + let (dest, cutover) = result?; + let mut server = dest.primary(0, &Request::default()).await?; + let parents: Vec = server + .fetch_all(format!("SELECT id FROM {schema}.parents ORDER BY id")) + .await?; + let children: Vec = server + .fetch_all(format!( + "SELECT parent_id FROM {schema}.children ORDER BY id" + )) + .await?; + Ok::<_, Box>((parents, children, cutover)) + } + .await; + + cleanup_replication_test( + &mut publisher, + &mut admin, + &original_config, + [schema, destination], + ) + .await?; + let (parents, children, cutover) = validation?; + assert_eq!(parents, [1, 2]); + assert_eq!(children, [1, 2, 999]); + assert!( + matches!( + &cutover, + Err(TaskError::Failed(SchemaSyncError::Backend(BackendError::ExecutionError(error)))) + if error.code == "23503" + ), + "cutover did not reject the orphan with a foreign key error: {cutover:?}" + ); + Ok(()) +} diff --git a/pgdog/src/backend/schema/sync/pg_dump.rs b/pgdog/src/backend/schema/sync/pg_dump.rs index 7e7b0db5f..a0af0358e 100644 --- a/pgdog/src/backend/schema/sync/pg_dump.rs +++ b/pgdog/src/backend/schema/sync/pg_dump.rs @@ -143,7 +143,10 @@ impl PgDump { /// Cross-checks the publication's table list on every shard, then runs /// `pg_dump --schema-only` against shard 0 only: all shards are assumed to /// share a schema, so a mismatch only warns. - pub(crate) async fn dump(&self) -> Result { + pub(crate) async fn dump( + &self, + destination_version: i64, + ) -> Result { let mut comparison: Vec = vec![]; let addr = self .source @@ -220,6 +223,7 @@ impl PgDump { Ok(PgDumpOutput { stmts, original: cleaned, + destination_version, }) } } @@ -233,6 +237,7 @@ pub(crate) struct PgDumpOutput { /// instead of deparsing, so any change to `clean` must preserve the byte /// offsets the parser saw. original: String, + destination_version: i64, } pub(crate) use pgdog_stats::SchemaSyncStatement as Statement; @@ -579,6 +584,11 @@ impl PgDumpOutput { // Get partitioned parent column types and parent-child relationships let parent_column_types = self.partitioned_parent_column_types(&columns_to_convert); let partition_parents = self.partition_parents(); + let partitioned_tables = if self.destination_version < 180_000 { + self.partitioned_tables() + } else { + HashSet::new() + }; for stmt in self.stmts.into_iter() { let original = self @@ -693,9 +703,20 @@ impl PgDumpOutput { // FK columns referencing integer PKs are // computed from fk_columns at the end if state == SyncState::PostData { - result.push( - Statement::new(original).set_skip_if_exists(), - ); + let sql = if cons.skip_validation + || stmt.relation().is_some_and(|relation| { + partitioned_tables + .contains(&Table::from(relation)) + }) { + original.to_owned() + } else { + format!( + "{} NOT VALID", + original.trim_end().trim_end_matches(';') + ) + }; + result + .push(Statement::new(sql).set_skip_if_exists()); } } else if state == SyncState::PostData { result.push( @@ -1214,6 +1235,96 @@ ALTER TABLE child ADD CONSTRAINT child_parent_fk FOREIGN KEY (parent_id) REFEREN statements[2].sql, "ALTER TABLE parent ADD CONSTRAINT parent_pkey PRIMARY KEY (id)" ); + assert_eq!( + output.statements(SyncState::PostData).unwrap()[0].sql, + "ALTER TABLE child ADD CONSTRAINT child_parent_fk FOREIGN KEY (parent_id) REFERENCES parent(id) NOT VALID" + ); + assert!(output.statements(SyncState::Cutover).unwrap().is_empty()); + } + + #[test] + fn test_foreign_key_not_valid_preserves_options() { + let sql = "ALTER TABLE child ADD CONSTRAINT child_parent_fk \ + FOREIGN KEY (tenant_id, parent_id) REFERENCES parent(tenant_id, id) \ + MATCH FULL ON DELETE CASCADE DEFERRABLE INITIALLY DEFERRED NOT VALID;"; + let output = parse(sql); + let statements = output.statements(SyncState::PostData).unwrap(); + + assert_eq!(statements.len(), 1); + assert_eq!(statements[0].sql, sql.trim_end_matches(';')); + assert!(output.statements(SyncState::PreData).unwrap().is_empty()); + assert!(output.statements(SyncState::Cutover).unwrap().is_empty()); + } + + #[test] + fn test_foreign_key_preserves_postdata_order() { + let output = parse( + "ALTER TABLE parent ADD CONSTRAINT parent_key UNIQUE (code); + ALTER TABLE child ADD CONSTRAINT child_parent_fk + FOREIGN KEY (parent_code) REFERENCES parent(code); + ALTER TABLE child ADD CONSTRAINT child_key UNIQUE (id);", + ); + let statements = output.statements(SyncState::PostData).unwrap(); + + assert_eq!(statements.len(), 3); + assert_eq!( + statements[0].sql, + "ALTER TABLE parent ADD CONSTRAINT parent_key UNIQUE (code)" + ); + assert!(statements[1].sql.contains("CONSTRAINT child_parent_fk")); + assert!(statements[1].sql.ends_with("NOT VALID")); + assert_eq!( + statements[2].sql, + "ALTER TABLE child ADD CONSTRAINT child_key UNIQUE (id)" + ); + } + + #[test] + fn test_partitioned_foreign_key_not_valid() { + let mut output = parse( + "CREATE TABLE parent (id BIGINT PRIMARY KEY); + CREATE TABLE child (id BIGINT, parent_id BIGINT) PARTITION BY RANGE (id); + CREATE TABLE child_0 PARTITION OF child FOR VALUES FROM (0) TO (100); + ALTER TABLE child ADD CONSTRAINT child_parent_fk FOREIGN KEY (parent_id) REFERENCES parent(id);", + ); + for (version, expected) in [ + ( + 170_009, + "ALTER TABLE child ADD CONSTRAINT child_parent_fk FOREIGN KEY (parent_id) REFERENCES parent(id)", + ), + ( + 180_000, + "ALTER TABLE child ADD CONSTRAINT child_parent_fk FOREIGN KEY (parent_id) REFERENCES parent(id) NOT VALID", + ), + ] { + output.destination_version = version; + let statements = output.statements(SyncState::PostData).unwrap(); + assert_eq!(statements.len(), 1); + assert_eq!(statements[0].sql.trim(), expected); + assert!(output.statements(SyncState::Cutover).unwrap().is_empty()); + } + } + + #[test] + fn test_old_destination_foreign_keys_use_referencing_table() { + let mut output = parse( + "CREATE TABLE public.parent (id BIGINT PRIMARY KEY); + CREATE TABLE public.child (id BIGINT, parent_id BIGINT); + CREATE TABLE other.child (id BIGINT, parent_id BIGINT) PARTITION BY RANGE (id); + ALTER TABLE public.child ADD CONSTRAINT ordinary_fk FOREIGN KEY (parent_id) REFERENCES public.parent(id); + ALTER TABLE other.child ADD CONSTRAINT partitioned_fk FOREIGN KEY (parent_id) REFERENCES public.parent(id) ON DELETE CASCADE DEFERRABLE INITIALLY DEFERRED;", + ); + output.destination_version = 170_009; + let statements = output.statements(SyncState::PostData).unwrap(); + assert_eq!(statements.len(), 2); + assert_eq!( + statements[0].sql.trim(), + "ALTER TABLE public.child ADD CONSTRAINT ordinary_fk FOREIGN KEY (parent_id) REFERENCES public.parent(id) NOT VALID" + ); + assert_eq!( + statements[1].sql.trim(), + "ALTER TABLE other.child ADD CONSTRAINT partitioned_fk FOREIGN KEY (parent_id) REFERENCES public.parent(id) ON DELETE CASCADE DEFERRABLE INITIALLY DEFERRED" + ); } #[test] @@ -1483,6 +1594,7 @@ b PgDumpOutput { stmts: pg_raw_parse::parse(query).unwrap().into_inner(), original: query.to_owned(), + destination_version: 180_000, } } } diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index 132ad3441..83f287749 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -420,7 +420,7 @@ impl Server { stream: Some(stream), key, id, - stats: Stats::connect(id, addr, ¶ms, &options, &config.config.memory), + stats: Stats::connect(id, addr, ¶ms, &config.config.memory), replication_mode: options.replication_mode(), params, changed_params: Parameters::default(), @@ -1352,13 +1352,7 @@ pub(crate) mod test { params: Parameters::default(), changed_params: Parameters::default(), client_params: Parameters::default(), - stats: Stats::connect( - id, - &addr, - &Parameters::default(), - &ServerOptions::default(), - &Memory::default(), - ), + stats: Stats::connect(id, &addr, &Parameters::default(), &Memory::default()), prepared_statements: super::PreparedStatements::default(), addr, dirty: false, diff --git a/pgdog/src/backend/server_options.rs b/pgdog/src/backend/server_options.rs index a669dfa2c..d4914d90d 100644 --- a/pgdog/src/backend/server_options.rs +++ b/pgdog/src/backend/server_options.rs @@ -1,12 +1,36 @@ +use std::time::Duration; + +use super::pool::Config; + use crate::net::{Parameter, parameter::ParameterValue}; -#[derive(Debug, Clone, Default)] +#[derive(Debug, Clone)] pub(crate) struct ServerOptions { pub(crate) params: Vec, - pub(crate) pool_id: u64, +} + +impl Default for ServerOptions { + fn default() -> Self { + Self { + params: vec![ + Parameter { + name: "application_name".into(), + value: "PgDog".into(), + }, + Parameter { + name: "client_encoding".into(), + value: "utf-8".into(), + }, + ], + } + } } impl ServerOptions { + pub(crate) fn add(&mut self, parameter: Parameter) { + self.params.push(parameter); + } + pub(crate) fn replication_mode(&self) -> bool { self.params.iter().any(|p| { p.name == "replication" @@ -18,12 +42,37 @@ impl ServerOptions { } pub(crate) fn new_replication() -> Self { - Self { - params: vec![Parameter { - name: "replication".into(), - value: "database".into(), - }], - pool_id: 0, - } + let mut options = Self::default(); + options.add(Parameter { + name: "replication".into(), + value: "database".into(), + }); + options + } + + pub(crate) fn new_resharding(config: &Config) -> Self { + let mut options = Self::default(); + options.add(Parameter { + name: "session_replication_role".into(), + value: "replica".into(), + }); + options.add(Parameter { + name: "statement_timeout".into(), + value: "0".into(), + }); + // 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. + options.add(Parameter { + name: "lock_timeout".into(), + value: config + .lock_timeout + .unwrap_or(Duration::from_secs(5)) + .as_millis() + .to_string() + .into(), + }); + options } } diff --git a/pgdog/src/backend/stats.rs b/pgdog/src/backend/stats.rs index e8ea0ae82..882eb8b4d 100644 --- a/pgdog/src/backend/stats.rs +++ b/pgdog/src/backend/stats.rs @@ -10,7 +10,7 @@ pub(crate) use pgdog_stats::server::Counts; use tokio::time::Instant; use crate::{ - backend::{Pool, ServerOptions, pool::stats::MemoryStats}, + backend::{Pool, pool::stats::MemoryStats}, config::Memory, net::{ Parameters, @@ -56,11 +56,10 @@ pub(crate) struct ServerStats { } impl ServerStats { - fn new(id: BackendPid, options: &ServerOptions, config: &Memory) -> Self { + fn new(id: BackendPid, config: &Memory) -> Self { let now = Instant::now(); let inner = pgdog_stats::server::Stats { memory: *MemoryStats::new(config), - pool_id: options.pool_id, ..Default::default() }; @@ -117,10 +116,9 @@ impl Stats { id: BackendPid, addr: &Address, params: &Parameters, - options: &ServerOptions, config: &Memory, ) -> Self { - let local = ServerStats::new(id, options, config); + let local = ServerStats::new(id, config); let server = ConnectedServer { stats: local,