From f2952fb57ce2ac8040be2678f559c0c057a253ba Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Wed, 22 Jul 2026 08:07:08 -0600 Subject: [PATCH 01/21] Handle custom types with differing OIDs across shards At the wire protocol level, type information is entirely handled by OID, not type name. An OID is a unique identifier that PG assigns for virtually every schema related entity. They are extremely dependent on the order that plugins are loaded, DDL is run, etc. Because of this it's quite likely that a user or extension defined type will not have the same OID across multiple shards. To solve this, we need to do three things: - Ensure clients receive a single canonical answer when they query for type information - Rewrite any messages being sent to the server to ensure they have the correct OIDs for a given type for that sever - Rewrite any messages being sent to the client to ensure they have the canonical OIDs for a given type This commit's solution to ensuring the client gets canoincal type information is straightforward, but not ideal. We look for any select statements that could be loading type information (referencing the relevant tables in `pg_catalog`, or doing certain casts), and route them to shard 0. This means that we have a single point of failure for schema information. Ideally we would be responding to these queries ourselves rather than sending them to the server, but the breadth of different ways different client libraries will load this type information would mean we essentially have to implement our own shitty SQL server. It's possible, and probably necessary in the long run, but outside of the scope of this specific fix. It's not clear that there'd be any reason one shard would be better than another for the canonical information, but if there is one, we could probably surface this as a configuration option. But since there's no clear reason for it, we can just do the reasonable default. For rewriting messages, we need to keep track of the mappings for each shard on our end. This was actually somewhat tricky to make work, as schema loading happens on `Shard`, but the message rewriting needs to happen all the way down in `Server`, and these are not types that share an API bounary. To rectify this, we use a new primitive, which is the async equivalent of `std::sync::OnceLock`. It ensures that only one writer will ever attempt to write to it, while also allowing waiters to wait for a value to be set without attempting to initialize themselves. With this primitive, all we need to do is stick it in an `Arc` to have it shared in as many disparate places as we need. We use the `SchemaCache` that was previously introduced to ensure we only have one copy of this mapping per database, regardless of how many user/clusters we have, but its use is relatively minimal. Since this sits in a very hot path, I've done as much as I could to avoid excess allocations or loops if they're not strictly necessary. I'm assuming that the majority of users won't have custom types that vary between shards, so this change should ideally be free for those users, short of the query to determine that there are no mappings at startup. Tests were a little bit trickier. Resolving the OIDs happens at cluster launch, and the code that needs to do the rewriting will never run without the cluster launching and loading first. But many of our tests were constructing clusters without launching them. I worked around this by making `Oids::default` return an instance that is already resolved to an empty set of mappings. Since the normal construction requires canonical type information to be passed in, we don't need to worry about production code accidentally calling this incorrectly. --- integration/rust/Cargo.toml | 3 + .../integration/cross_shard_oid_drift.rs | 54 +++ integration/rust/tests/integration/mod.rs | 1 + pgdog/benches/comment_parser.rs | 13 +- pgdog/src/admin/probe.rs | 1 + pgdog/src/backend/error.rs | 3 + pgdog/src/backend/mod.rs | 4 +- pgdog/src/backend/pool/cluster.rs | 21 +- .../src/backend/pool/cluster/schema_loader.rs | 37 ++ .../backend/pool/connection/binding_test.rs | 44 ++- pgdog/src/backend/pool/guard.rs | 11 +- pgdog/src/backend/pool/inner.rs | 2 + pgdog/src/backend/pool/lb/mod.rs | 12 +- pgdog/src/backend/pool/lb/test.rs | 79 ++++- pgdog/src/backend/pool/lsn_monitor.rs | 11 +- pgdog/src/backend/pool/mod.rs | 2 +- pgdog/src/backend/pool/monitor.rs | 54 +-- pgdog/src/backend/pool/pool_impl.rs | 10 +- pgdog/src/backend/pool/shard/mod.rs | 24 +- pgdog/src/backend/pool/shard/monitor.rs | 26 +- pgdog/src/backend/pool/shard/oids.rs | 132 +++++++ pgdog/src/backend/pool/test/mod.rs | 327 +++++++++++------- pgdog/src/backend/pool/waiting.rs | 23 +- pgdog/src/backend/prepared_statements.rs | 115 ++++-- .../replication/logical/publisher/slot.rs | 2 + .../replication/logical/subscriber/tests.rs | 18 + pgdog/src/backend/schema/cache/mod.rs | 26 +- pgdog/src/backend/server.rs | 49 ++- .../prepared_statements/global_cache.rs | 4 +- pgdog/src/frontend/router/parser/query/mod.rs | 31 +- pgdog/src/frontend/router/parser/route.rs | 9 + pgdog/src/lib.rs | 1 + pgdog/src/net/messages/mod.rs | 4 + .../src/net/messages/parameter_description.rs | 9 + pgdog/src/net/messages/parse.rs | 108 +++--- pgdog/src/net/messages/row_description.rs | 14 +- pgdog/src/sync.rs | 96 +++++ 37 files changed, 1042 insertions(+), 338 deletions(-) create mode 100644 integration/rust/tests/integration/cross_shard_oid_drift.rs create mode 100644 pgdog/src/backend/pool/shard/oids.rs create mode 100644 pgdog/src/sync.rs diff --git a/integration/rust/Cargo.toml b/integration/rust/Cargo.toml index 3cbc0ecb7..427b729b1 100644 --- a/integration/rust/Cargo.toml +++ b/integration/rust/Cargo.toml @@ -6,6 +6,9 @@ edition = "2024" [lib] test = true +[features] +new_parser = [] + [dependencies] tokio-postgres = {version = "0.7.13", features = ["with-uuid-1"]} sqlx = { version = "0.8.6", features = ["postgres", "runtime-tokio", "tls-native-tls", "bigdecimal", "chrono", "json", "rust_decimal"]} diff --git a/integration/rust/tests/integration/cross_shard_oid_drift.rs b/integration/rust/tests/integration/cross_shard_oid_drift.rs new file mode 100644 index 000000000..b68434542 --- /dev/null +++ b/integration/rust/tests/integration/cross_shard_oid_drift.rs @@ -0,0 +1,54 @@ +#![cfg(feature = "new_parser")] +use crate::setup::{admin_sqlx, connections_sqlx}; +use sqlx::Executor; + +#[derive(sqlx::Type, Debug, Clone, PartialEq)] +#[sqlx(type_name = "test_oid_drift_composite")] +struct Composite { + a: String, + b: String, +} + +#[tokio::test] +async fn test_oid_drift() { + let conn = connections_sqlx().await.pop().unwrap(); + + // Intentionally cause the OID of the type to differ between shards + conn.execute("/* pgdog_shard: 0 */ CREATE SEQUENCE foo; DROP SEQUENCE foo;") + .await + .unwrap(); + conn.execute("DROP TYPE IF EXISTS test_oid_drift_composite CASCADE") + .await + .unwrap(); + conn.execute("CREATE TYPE test_oid_drift_composite AS (a text, b text)") + .await + .unwrap(); + conn.execute("DROP TABLE IF EXISTS test_oid_drift") + .await + .unwrap(); + conn.execute( + "CREATE TABLE test_oid_drift (customer_id BIGINT, composite test_oid_drift_composite)", + ) + .await + .unwrap(); + admin_sqlx().await.execute("RELOAD").await.unwrap(); + + let composite = Composite { + a: String::from("a"), + b: String::from("b"), + }; + for i in 1..=20 { + sqlx::query("INSERT INTO test_oid_drift VALUES ($1, $2)") + .bind(i) + .bind(&composite) + .execute(&conn) + .await + .unwrap(); + } + + let rows: Vec = sqlx::query_scalar("SELECT composite FROM test_oid_drift") + .fetch_all(&conn) + .await + .unwrap(); + assert_eq!(rows, vec![composite; 20]); +} diff --git a/integration/rust/tests/integration/mod.rs b/integration/rust/tests/integration/mod.rs index 55e70c4cb..3983863c7 100644 --- a/integration/rust/tests/integration/mod.rs +++ b/integration/rust/tests/integration/mod.rs @@ -10,6 +10,7 @@ pub mod client_ids; pub mod connection_recovery; pub mod copy; pub mod cross_shard_disabled; +mod cross_shard_oid_drift; pub mod distinct; pub mod explain; pub mod fake_transactions; diff --git a/pgdog/benches/comment_parser.rs b/pgdog/benches/comment_parser.rs index d5086ca5a..22afa9ed7 100644 --- a/pgdog/benches/comment_parser.rs +++ b/pgdog/benches/comment_parser.rs @@ -1,7 +1,6 @@ use brunch::{Bench, benches}; #[cfg(not(feature = "new_parser"))] use pg_query::scan_raw; -use pgdog::backend::ShardingSchema; use pgdog::frontend::router::parser::comment::parse_edge_comment; const QUERY_WITH_LEADING: &str = @@ -13,20 +12,20 @@ const QUERY_NO_COMMENT: &str = "SELECT * FROM users WHERE id = $1 AND name = $2" #[cfg(feature = "new_parser")] benches!( Bench::new("parse_edge_comment(leading)") - .run(|| parse_edge_comment(QUERY_WITH_LEADING, &ShardingSchema::default())), + .run(|| parse_edge_comment(QUERY_WITH_LEADING, &Default::default())), Bench::new("parse_edge_comment(trailing)") - .run(|| parse_edge_comment(QUERY_WITH_TRAILING, &ShardingSchema::default())), + .run(|| parse_edge_comment(QUERY_WITH_TRAILING, &Default::default())), Bench::new("parse_edge_comment(no comment)") - .run(|| parse_edge_comment(QUERY_NO_COMMENT, &ShardingSchema::default())), + .run(|| parse_edge_comment(QUERY_NO_COMMENT, &Default::default())), ); #[cfg(not(feature = "new_parser"))] benches!( Bench::new("parse_edge_comment(leading)") - .run(|| parse_edge_comment(QUERY_WITH_LEADING, &ShardingSchema::default())), + .run(|| parse_edge_comment(QUERY_WITH_LEADING, &Default::default())), Bench::new("parse_edge_comment(trailing)") - .run(|| parse_edge_comment(QUERY_WITH_TRAILING, &ShardingSchema::default())), + .run(|| parse_edge_comment(QUERY_WITH_TRAILING, &Default::default())), Bench::new("parse_edge_comment(no comment)") - .run(|| parse_edge_comment(QUERY_NO_COMMENT, &ShardingSchema::default())), + .run(|| parse_edge_comment(QUERY_NO_COMMENT, &Default::default())), Bench::new("scan_raw(leading)").run(|| scan_raw(QUERY_WITH_LEADING)), Bench::new("scan_raw(trailing)").run(|| scan_raw(QUERY_WITH_TRAILING)), Bench::new("scan_raw(no comment)").run(|| scan_raw(QUERY_NO_COMMENT)), diff --git a/pgdog/src/admin/probe.rs b/pgdog/src/admin/probe.rs index 1fe397552..4cc55bf8d 100644 --- a/pgdog/src/admin/probe.rs +++ b/pgdog/src/admin/probe.rs @@ -32,6 +32,7 @@ impl Command for Probe { &Address::try_from(self.url.clone()).map_err(|_| Error::InvalidAddress)?, ServerOptions::default(), ConnectReason::Probe, + Default::default(), ), ) .await? diff --git a/pgdog/src/backend/error.rs b/pgdog/src/backend/error.rs index 2514ca6c1..6b344c820 100644 --- a/pgdog/src/backend/error.rs +++ b/pgdog/src/backend/error.rs @@ -143,6 +143,9 @@ pub enum Error { #[error("cannot ignore response for message type: {0}")] UnsupportedHandleIgnore(char), + + #[error("missing canonical oid for type {0}")] + MissingCanonicalOid(String), } impl From for Error { diff --git a/pgdog/src/backend/mod.rs b/pgdog/src/backend/mod.rs index a4f051759..e90b36e0b 100644 --- a/pgdog/src/backend/mod.rs +++ b/pgdog/src/backend/mod.rs @@ -21,7 +21,9 @@ pub mod validation; pub use connect_reason::ConnectReason; pub use disconnect_reason::DisconnectReason; pub use error::Error; -pub use pool::{Cluster, ClusterShardConfig, LoadBalancer, Pool, Shard, ShardingSchema}; +pub(crate) use pool::{ + CanonicalOids, Cluster, ClusterShardConfig, OidMappings, Oids, Pool, Shard, ShardingSchema, +}; pub use prepared_statements::PreparedStatements; pub use protocol::*; pub use pub_sub::{PubSubClient, PubSubListener}; diff --git a/pgdog/src/backend/pool/cluster.rs b/pgdog/src/backend/pool/cluster.rs index f66654c95..85f5e71af 100644 --- a/pgdog/src/backend/pool/cluster.rs +++ b/pgdog/src/backend/pool/cluster.rs @@ -23,7 +23,9 @@ use crate::{ net::{Query, messages::FrontendPid}, }; -use super::{Address, Config, Error, Guard, MirrorStats, Request, Shard, ShardConfig}; +use super::{ + Address, CanonicalOids, Config, Error, Guard, MirrorStats, Request, Shard, ShardConfig, +}; use crate::config::LoadBalancingStrategy; use launch::Readiness; @@ -85,6 +87,10 @@ pub struct Cluster { tls_client_certificate_required: bool, #[debug(skip)] schema_loader: Box, +<<<<<<< HEAD +======= + canonical_oids: Arc, +>>>>>>> 303395b9 (WIP: Handle OID drift between shards) } /// Sharding configuration from the cluster. @@ -299,6 +305,7 @@ impl Cluster { user: user.to_owned(), database: name.to_owned(), }); + let canonical_oids = schema_cache.canonical_oids(name); Self { identifier: identifier.clone(), @@ -357,6 +364,10 @@ impl Cluster { identity: identity.clone(), tls_client_certificate_required, schema_loader: Box::new(schema_loader::FromServer), +<<<<<<< HEAD +======= + canonical_oids, +>>>>>>> 303395b9 (WIP: Handle OID drift between shards) } } @@ -411,7 +422,7 @@ impl Cluster { } /// Get all shards. - pub fn shards(&self) -> &[Shard] { + pub(crate) fn shards(&self) -> &[Shard] { &self.shards } @@ -907,6 +918,12 @@ mod test { cluster } + pub(crate) fn force_dummy_schema(&self) { + for shard in &self.shards { + shard.schema_not_needed(); + } + } + pub(crate) fn set_read_write_strategy(&mut self, rw_strategy: ReadWriteStrategy) { self.rw_strategy = rw_strategy; } diff --git a/pgdog/src/backend/pool/cluster/schema_loader.rs b/pgdog/src/backend/pool/cluster/schema_loader.rs index 21c7f244f..38a3d8d96 100644 --- a/pgdog/src/backend/pool/cluster/schema_loader.rs +++ b/pgdog/src/backend/pool/cluster/schema_loader.rs @@ -3,6 +3,7 @@ use crate::backend::pool::ee::schema_changed_hook; use crate::tasks; use crate::util::safe_sleep; use dyn_clone::DynClone; +use std::sync::Arc; use std::time::Duration; use tokio::select; use tracing::error; @@ -25,6 +26,42 @@ impl SchemaLoader for FromServer { return; } + // For now we treat shard 0 as the canonical OID source + if let Some(shard) = cluster.shards().first() { + let canonical_oids = Arc::clone(&cluster.canonical_oids); + let shard = shard.clone(); + tasks::spawn("load canonical oids", async move { + // FIXME: This shutdown signal/retry if error logic seems like + // it might be worth abstracting + let shutdown = tasks::shutdown_signal(); + + loop { + let loader = async { + canonical_oids + .load(&mut *shard.primary_or_replica(&Default::default()).await?) + .await + }; + let result = select! { + _ = shutdown.cancelled() => break, + result = loader => { result }, + }; + + match result { + Ok(_) => break, + Err(err) => { + if shard.online() { + error!("error loading canonical type information: {err}"); + sleep(Duration::from_millis(100)).await; + } else { + // Cluster is shutting down + break; + } + } + } + } + }); + } + for shard in cluster.shards() { let identifier = cluster.identifier(); let shard = shard.clone(); diff --git a/pgdog/src/backend/pool/connection/binding_test.rs b/pgdog/src/backend/pool/connection/binding_test.rs index b9f541a78..ec8ac5798 100644 --- a/pgdog/src/backend/pool/connection/binding_test.rs +++ b/pgdog/src/backend/pool/connection/binding_test.rs @@ -26,20 +26,29 @@ mod tests { let server3 = Box::new(test_server().await); // Create pools for each server using their addresses - let pool1 = Pool::new(&PoolConfig { - address: server1.addr().clone(), - config: crate::backend::pool::Config::default(), - }); + let pool1 = Pool::new( + &PoolConfig { + address: server1.addr().clone(), + config: crate::backend::pool::Config::default(), + }, + Default::default(), + ); - let pool2 = Pool::new(&PoolConfig { - address: server2.addr().clone(), - config: crate::backend::pool::Config::default(), - }); + let pool2 = Pool::new( + &PoolConfig { + address: server2.addr().clone(), + config: crate::backend::pool::Config::default(), + }, + Default::default(), + ); - let pool3 = Pool::new(&PoolConfig { - address: server3.addr().clone(), - config: crate::backend::pool::Config::default(), - }); + let pool3 = Pool::new( + &PoolConfig { + address: server3.addr().clone(), + config: crate::backend::pool::Config::default(), + }, + Default::default(), + ); let now = Instant::now(); let guards = vec![ @@ -68,10 +77,13 @@ mod tests { // Create a Direct binding instead of MultiShard let server = Box::new(test_server().await); - let pool = Pool::new(&PoolConfig { - address: server.addr().clone(), - config: crate::backend::pool::Config::default(), - }); + let pool = Pool::new( + &PoolConfig { + address: server.addr().clone(), + config: crate::backend::pool::Config::default(), + }, + Default::default(), + ); let guard = crate::backend::pool::Guard::new(pool, server, Instant::now()); let mut binding = Binding::Direct(guard, 0); diff --git a/pgdog/src/backend/pool/guard.rs b/pgdog/src/backend/pool/guard.rs index 03d27cf76..7b0f99919 100644 --- a/pgdog/src/backend/pool/guard.rs +++ b/pgdog/src/backend/pool/guard.rs @@ -348,10 +348,13 @@ mod test { }, }; - let pool = Pool::new(&PoolConfig { - address: Address::new_test(), - config, - }); + let pool = Pool::new( + &PoolConfig { + address: Address::new_test(), + config, + }, + Default::default(), + ); pool.launch(); { diff --git a/pgdog/src/backend/pool/inner.rs b/pgdog/src/backend/pool/inner.rs index a44004f8b..0d3bdc4ed 100644 --- a/pgdog/src/backend/pool/inner.rs +++ b/pgdog/src/backend/pool/inner.rs @@ -324,6 +324,7 @@ impl Inner { for conn in idle.iter_mut() { conn.stats_mut().set_pool_id(destination.id()); + conn.replace_oids(&destination.inner().oids) } (idle, taken) @@ -352,6 +353,7 @@ impl Inner { if moved.id() != self.id { server.stats_mut().set_pool_id(moved.id()); server.stats().update(); + server.replace_oids(&moved.inner().oids); moved.lock().maybe_check_in(server, now, stats, true)?; return Ok(result); } diff --git a/pgdog/src/backend/pool/lb/mod.rs b/pgdog/src/backend/pool/lb/mod.rs index 91e39a198..a192be6f8 100644 --- a/pgdog/src/backend/pool/lb/mod.rs +++ b/pgdog/src/backend/pool/lb/mod.rs @@ -18,7 +18,7 @@ use crate::{ net::Parameters, }; -use super::{Error, Guard, Pool, PoolConfig, Request}; +use super::{Error, Guard, Oids, Pool, PoolConfig, Request}; use crate::util::safe_timeout; pub mod ban; @@ -91,11 +91,12 @@ pub struct LoadBalancer { impl LoadBalancer { /// Create new replicas pools. - pub fn new( + pub(crate) fn new( primary: &Option, addrs: &[PoolConfig], lb_strategy: LoadBalancingStrategy, rw_split: ReadWriteSplit, + oids: Arc, ) -> LoadBalancer { let checkout_timeout = primary .as_ref() @@ -111,7 +112,12 @@ impl LoadBalancer { let mut targets: Vec<_> = addrs .iter() - .map(|config| Target::new(Pool::new(config), config.address.configured_role)) + .map(|config| { + Target::new( + Pool::new(config, Arc::clone(&oids)), + config.address.configured_role, + ) + }) .collect(); let primary_target = primary diff --git a/pgdog/src/backend/pool/lb/test.rs b/pgdog/src/backend/pool/lb/test.rs index f0d0dfb1b..0a8148aca 100644 --- a/pgdog/src/backend/pool/lb/test.rs +++ b/pgdog/src/backend/pool/lb/test.rs @@ -41,6 +41,7 @@ fn setup_test_replicas() -> LoadBalancer { &[pool_config1, pool_config2], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); replicas.launch(); replicas @@ -62,13 +63,14 @@ fn set_lsn_stats(target: &Target, replica: bool, lsn: i64) { async fn test_include_primary_if_replica_banned_only_primary() { let mut primary = create_test_pool_config("127.0.0.1", 5432); primary.address.configured_role = Role::Primary; - let pool = Pool::new(&primary); + let pool = Pool::new(&primary, Default::default()); let lb = LoadBalancer::new( &Some(pool), &[], LoadBalancingStrategy::default(), ReadWriteSplit::IncludePrimaryIfReplicaBanned, + Default::default(), ); lb.launch(); @@ -202,7 +204,7 @@ async fn test_pools_with_roles_and_bans() { #[tokio::test] async fn test_primary_pool_banning() { let primary_config = create_test_pool_config("127.0.0.1", 5432); - let primary_pool = Pool::new(&primary_config); + let primary_pool = Pool::new(&primary_config, Default::default()); primary_pool.launch(); let replica_configs = [create_test_pool_config("localhost", 5432)]; @@ -212,6 +214,7 @@ async fn test_primary_pool_banning() { &replica_configs, LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); replicas.launch(); @@ -355,7 +358,7 @@ async fn test_monitor_automatic_ban_expiration() { #[tokio::test] async fn test_read_write_split_exclude_primary() { let primary_config = create_test_pool_config("127.0.0.1", 5432); - let primary_pool = Pool::new(&primary_config); + let primary_pool = Pool::new(&primary_config, Default::default()); primary_pool.launch(); let replica_configs = [ @@ -368,6 +371,7 @@ async fn test_read_write_split_exclude_primary() { &replica_configs, LoadBalancingStrategy::Random, ReadWriteSplit::ExcludePrimary, + Default::default(), ); replicas.launch(); @@ -394,7 +398,7 @@ async fn test_read_write_split_exclude_primary() { #[tokio::test] async fn test_read_write_split_include_primary() { let primary_config = create_test_pool_config("127.0.0.1", 5432); - let primary_pool = Pool::new(&primary_config); + let primary_pool = Pool::new(&primary_config, Default::default()); primary_pool.launch(); let replica_configs = [create_test_pool_config("localhost", 5432)]; @@ -404,6 +408,7 @@ async fn test_read_write_split_include_primary() { &replica_configs, LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); replicas.launch(); @@ -443,7 +448,10 @@ async fn test_read_write_split_include_primary() { #[tokio::test] async fn test_prefer_primary_optin_read_honors_read_write_split() { async fn used_ids(split: ReadWriteSplit) -> (HashSet, u64) { - let primary_pool = Pool::new(&create_test_pool_config("127.0.0.1", 5432)); + let primary_pool = Pool::new( + &create_test_pool_config("127.0.0.1", 5432), + Default::default(), + ); primary_pool.launch(); let primary_id = primary_pool.id(); @@ -457,6 +465,7 @@ async fn test_prefer_primary_optin_read_honors_read_write_split() { &replica_configs, LoadBalancingStrategy::RoundRobin, split, + Default::default(), ); lb.launch(); @@ -494,7 +503,7 @@ async fn test_prefer_primary_optin_read_honors_read_write_split() { #[tokio::test] async fn test_read_write_split_exclude_primary_no_replicas() { let primary_config = create_test_pool_config("127.0.0.1", 5432); - let primary_pool = Pool::new(&primary_config); + let primary_pool = Pool::new(&primary_config, Default::default()); primary_pool.launch(); let replica_configs = []; @@ -504,6 +513,7 @@ async fn test_read_write_split_exclude_primary_no_replicas() { &replica_configs, LoadBalancingStrategy::RoundRobin, ReadWriteSplit::ExcludePrimary, + Default::default(), ); replicas.launch(); @@ -540,6 +550,7 @@ async fn test_read_write_split_exclude_primary_no_primary() { &replica_configs, LoadBalancingStrategy::Random, ReadWriteSplit::ExcludePrimary, + Default::default(), ); replicas.launch(); @@ -570,6 +581,7 @@ async fn test_read_write_split_include_primary_no_primary() { &replica_configs, LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); replicas.launch(); @@ -591,7 +603,7 @@ async fn test_read_write_split_include_primary_no_primary() { #[tokio::test] async fn test_read_write_split_with_banned_primary() { let primary_config = create_test_pool_config("127.0.0.1", 5432); - let primary_pool = Pool::new(&primary_config); + let primary_pool = Pool::new(&primary_config, Default::default()); primary_pool.launch(); let replica_configs = [create_test_pool_config("localhost", 5432)]; @@ -601,6 +613,7 @@ async fn test_read_write_split_with_banned_primary() { &replica_configs, LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); replicas.launch(); @@ -631,7 +644,7 @@ async fn test_read_write_split_with_banned_primary() { #[tokio::test] async fn test_read_write_split_with_banned_replicas() { let primary_config = create_test_pool_config("127.0.0.1", 5432); - let primary_pool = Pool::new(&primary_config); + let primary_pool = Pool::new(&primary_config, Default::default()); primary_pool.launch(); let replica_configs = [create_test_pool_config("localhost", 5432)]; @@ -641,6 +654,7 @@ async fn test_read_write_split_with_banned_replicas() { &replica_configs, LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); replicas.launch(); @@ -671,7 +685,7 @@ async fn test_read_write_split_with_banned_replicas() { #[tokio::test] async fn test_prefer_primary_with_banned_replicas_falls_back_to_primary() { let primary_config = create_test_pool_config("127.0.0.1", 5432); - let primary_pool = Pool::new(&primary_config); + let primary_pool = Pool::new(&primary_config, Default::default()); primary_pool.launch(); let replica_configs = [create_test_pool_config("localhost", 5432)]; @@ -681,6 +695,7 @@ async fn test_prefer_primary_with_banned_replicas_falls_back_to_primary() { &replica_configs, LoadBalancingStrategy::Random, ReadWriteSplit::PreferPrimary, + Default::default(), ); replicas.launch(); @@ -706,7 +721,7 @@ async fn test_prefer_primary_with_banned_replicas_falls_back_to_primary() { #[tokio::test] async fn test_read_write_split_exclude_primary_with_round_robin() { let primary_config = create_test_pool_config("127.0.0.1", 5432); - let primary_pool = Pool::new(&primary_config); + let primary_pool = Pool::new(&primary_config, Default::default()); primary_pool.launch(); let replica_configs = [ @@ -719,6 +734,7 @@ async fn test_read_write_split_exclude_primary_with_round_robin() { &replica_configs, LoadBalancingStrategy::RoundRobin, ReadWriteSplit::ExcludePrimary, + Default::default(), ); replicas.launch(); @@ -764,6 +780,7 @@ async fn test_monitor_shuts_down_on_notify() { &[pool_config1, pool_config2], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); replicas @@ -826,6 +843,7 @@ async fn test_monitor_does_not_ban_single_target() { &[pool_config], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); replicas.launch(); @@ -898,6 +916,7 @@ async fn test_monitor_does_not_ban_with_zero_ban_timeout() { &[pool_config1, pool_config2], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); replicas.launch(); @@ -944,7 +963,7 @@ async fn test_monitor_health_state_race() { #[tokio::test] async fn test_include_primary_if_replica_banned_no_bans() { let primary_config = create_test_pool_config("127.0.0.1", 5432); - let primary_pool = Pool::new(&primary_config); + let primary_pool = Pool::new(&primary_config, Default::default()); primary_pool.launch(); let replica_configs = [create_test_pool_config("localhost", 5432)]; @@ -954,6 +973,7 @@ async fn test_include_primary_if_replica_banned_no_bans() { &replica_configs, LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimaryIfReplicaBanned, + Default::default(), ); replicas.launch(); @@ -980,7 +1000,7 @@ async fn test_include_primary_if_replica_banned_no_bans() { #[tokio::test] async fn test_include_primary_if_replica_banned_with_ban() { let primary_config = create_test_pool_config("127.0.0.1", 5432); - let primary_pool = Pool::new(&primary_config); + let primary_pool = Pool::new(&primary_config, Default::default()); primary_pool.launch(); let replica_configs = [create_test_pool_config("localhost", 5432)]; @@ -990,6 +1010,7 @@ async fn test_include_primary_if_replica_banned_with_ban() { &replica_configs, LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimaryIfReplicaBanned, + Default::default(), ); replicas.launch(); @@ -1029,7 +1050,7 @@ async fn test_has_replicas_with_replicas() { #[tokio::test] async fn test_has_replicas_with_primary_and_replicas() { let primary_config = create_test_pool_config("127.0.0.1", 5432); - let primary_pool = Pool::new(&primary_config); + let primary_pool = Pool::new(&primary_config, Default::default()); primary_pool.launch(); let replica_configs = [create_test_pool_config("localhost", 5432)]; @@ -1039,6 +1060,7 @@ async fn test_has_replicas_with_primary_and_replicas() { &replica_configs, LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); lb.launch(); @@ -1050,7 +1072,7 @@ async fn test_has_replicas_with_primary_and_replicas() { #[tokio::test] async fn test_has_replicas_primary_only() { let primary_config = create_test_pool_config("127.0.0.1", 5432); - let primary_pool = Pool::new(&primary_config); + let primary_pool = Pool::new(&primary_config, Default::default()); primary_pool.launch(); let lb = LoadBalancer::new( @@ -1058,6 +1080,7 @@ async fn test_has_replicas_primary_only() { &[], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); lb.launch(); @@ -1073,6 +1096,7 @@ async fn test_has_replicas_empty() { &[], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); assert!(!lb.has_replicas()); @@ -1119,6 +1143,7 @@ async fn test_can_move_conns_to_same_config() { &[pool_config1.clone(), pool_config2.clone()], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); let lb2 = LoadBalancer::new( @@ -1126,6 +1151,7 @@ async fn test_can_move_conns_to_same_config() { &[pool_config1, pool_config2], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); assert!(lb1.can_move_conns_to(&lb2)); @@ -1143,6 +1169,7 @@ async fn test_can_move_conns_to_with_removed_replica() { &[pool_config1.clone(), pool_config2], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); let lb2 = LoadBalancer::new( @@ -1150,6 +1177,7 @@ async fn test_can_move_conns_to_with_removed_replica() { &[pool_config1], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); assert!(!lb1.can_move_conns_to(&lb2)); @@ -1167,6 +1195,7 @@ async fn test_can_move_conns_to_with_added_replica() { std::slice::from_ref(&pool_config1), LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); let lb_new = LoadBalancer::new( @@ -1174,6 +1203,7 @@ async fn test_can_move_conns_to_with_added_replica() { &[pool_config1, pool_config2], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); assert!(lb_old.can_move_conns_to(&lb_new)); @@ -1193,6 +1223,7 @@ async fn test_move_conns_to_with_added_replica_matches_by_address() { std::slice::from_ref(&pool_config1), LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); lb_old.launch(); @@ -1201,6 +1232,7 @@ async fn test_move_conns_to_with_added_replica_matches_by_address() { &[pool_config1.clone(), pool_config2.clone()], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); lb_new.launch(); @@ -1237,12 +1269,13 @@ async fn test_redetect_roles_marks_added_auto_target_replica_when_primary_unchan let mut existing_replica_config = create_test_pool_config("localhost", 5432); existing_replica_config.address.configured_role = Role::Auto; - let old_primary = Pool::new(&primary_config); + let old_primary = Pool::new(&primary_config, Default::default()); let lb_old = LoadBalancer::new( &Some(old_primary), std::slice::from_ref(&existing_replica_config), LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); set_lsn_stats(&lb_old.targets[0], true, 100); @@ -1253,12 +1286,13 @@ async fn test_redetect_roles_marks_added_auto_target_replica_when_primary_unchan let mut added_replica_config = create_test_pool_config("localhost", 5433); added_replica_config.address.configured_role = Role::Auto; - let new_primary = Pool::new(&primary_config); + let new_primary = Pool::new(&primary_config, Default::default()); let lb_new = LoadBalancer::new( &Some(new_primary), &[existing_replica_config, added_replica_config], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); lb_old.move_conns_to(&lb_new).unwrap(); @@ -1295,6 +1329,7 @@ async fn test_redetect_roles_leaves_auto_targets_pending_when_stats_are_invalid( &[config1, config2], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); assert!(lb.targets.iter().all(|target| target.role() == Role::Auto)); @@ -1323,6 +1358,7 @@ async fn test_redetect_roles_marks_auto_targets_replicas_when_all_valid_targets_ &[config1, config2], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); set_lsn_stats(&lb.targets[0], true, 100); @@ -1356,6 +1392,7 @@ async fn test_can_move_conns_to_different_addresses() { &[pool_config1, pool_config2], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); let lb2 = LoadBalancer::new( @@ -1363,6 +1400,7 @@ async fn test_can_move_conns_to_different_addresses() { &[pool_config3.clone(), pool_config3], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); assert!(!lb1.can_move_conns_to(&lb2)); @@ -1439,6 +1477,7 @@ async fn test_weighted_round_robin_smooth_distribution() { &[pool_config1, pool_config2], LoadBalancingStrategy::WeightedRoundRobin, ReadWriteSplit::IncludePrimary, + Default::default(), ); lb.launch(); @@ -1473,6 +1512,7 @@ async fn test_weighted_round_robin_equal_weights() { &[pool_config1, pool_config2], LoadBalancingStrategy::WeightedRoundRobin, ReadWriteSplit::IncludePrimary, + Default::default(), ); lb.launch(); @@ -1504,6 +1544,7 @@ async fn test_weighted_round_robin_zero_weight_never_selected() { &[pool_config1, pool_config2], LoadBalancingStrategy::WeightedRoundRobin, ReadWriteSplit::IncludePrimary, + Default::default(), ); lb.launch(); @@ -1532,6 +1573,7 @@ async fn test_weighted_round_robin_proportional_distribution() { &[pool_config1, pool_config2], LoadBalancingStrategy::WeightedRoundRobin, ReadWriteSplit::IncludePrimary, + Default::default(), ); lb.launch(); @@ -1563,6 +1605,7 @@ async fn test_least_active_connections_prefers_pool_with_fewer_checked_out() { &[pool_config1, pool_config2], LoadBalancingStrategy::LeastActiveConnections, ReadWriteSplit::IncludePrimary, + Default::default(), ); replicas.launch(); @@ -1599,6 +1642,7 @@ fn setup_test_replicas_no_launch() -> LoadBalancer { &[pool_config1, pool_config2], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ) } @@ -1799,6 +1843,7 @@ fn test_ban_check_does_not_ban_single_target() { &[pool_config], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); // Don't launch - we're unit testing ban_check @@ -1864,6 +1909,7 @@ fn test_ban_check_does_not_ban_with_zero_ban_timeout() { &[pool_config1, pool_config2], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); // Set target as unhealthy @@ -2171,6 +2217,7 @@ async fn test_params_returns_all_replicas_down_when_empty() { &[], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); let request = Request::default(); diff --git a/pgdog/src/backend/pool/lsn_monitor.rs b/pgdog/src/backend/pool/lsn_monitor.rs index 891d14607..0f5ed0cb5 100644 --- a/pgdog/src/backend/pool/lsn_monitor.rs +++ b/pgdog/src/backend/pool/lsn_monitor.rs @@ -449,10 +449,13 @@ mod test { }, }; - let pool = Pool::new(&PoolConfig { - address: Address::new_test(), - config, - }); + let pool = Pool::new( + &PoolConfig { + address: Address::new_test(), + config, + }, + Default::default(), + ); pool.launch(); // Saturate the pool by holding its only connection. diff --git a/pgdog/src/backend/pool/mod.rs b/pgdog/src/backend/pool/mod.rs index 5550a0e68..fdd8aea11 100644 --- a/pgdog/src/backend/pool/mod.rs +++ b/pgdog/src/backend/pool/mod.rs @@ -40,7 +40,7 @@ pub use monitor::Monitor; pub use password::Password; pub use pool_impl::Pool; pub use request::Request; -pub use shard::Shard; +pub(crate) use shard::{CanonicalOids, OidMappings, Oids, Shard}; pub use state::State; pub use stats::Stats; diff --git a/pgdog/src/backend/pool/monitor.rs b/pgdog/src/backend/pool/monitor.rs index c6d36bb04..a1086bfdd 100644 --- a/pgdog/src/backend/pool/monitor.rs +++ b/pgdog/src/backend/pool/monitor.rs @@ -43,6 +43,7 @@ //! The loop exits when the pool shuts down (e.g. on config reload), preventing //! refresh tasks from leaking across reloads. +use std::sync::Arc; use std::time::Duration; use super::{Error, Guard, Healtcheck, Pool, Request}; @@ -444,7 +445,12 @@ impl Monitor { for attempt in 0..connect_attempts { match safe_timeout( connect_timeout, - Server::connect(pool.addr(), options.clone(), reason), + Server::connect( + pool.addr(), + options.clone(), + reason, + Arc::clone(&pool.inner().oids), + ), ) .await { @@ -558,17 +564,20 @@ mod test { }, }; - let pool = Pool::new(&PoolConfig { - address: Address { - host: "127.0.0.1".into(), - port: 1, - database_name: "pgdog".into(), - user: "pgdog".into(), - passwords: vec!["pgdog".into()], - ..Default::default() + let pool = Pool::new( + &PoolConfig { + address: Address { + host: "127.0.0.1".into(), + port: 1, + database_name: "pgdog".into(), + user: "pgdog".into(), + passwords: vec!["pgdog".into()], + ..Default::default() + }, + config, }, - config, - }); + Default::default(), + ); pool.launch(); pool.inner().health.toggle(true); @@ -592,17 +601,20 @@ mod test { }, }; - let pool = Pool::new(&PoolConfig { - address: Address { - host: "127.0.0.1".into(), - port: 5432, - database_name: "pgdog".into(), - user: "pgdog".into(), - passwords: vec!["pgdog".into()], - ..Default::default() + let pool = Pool::new( + &PoolConfig { + address: Address { + host: "127.0.0.1".into(), + port: 5432, + database_name: "pgdog".into(), + user: "pgdog".into(), + passwords: vec!["pgdog".into()], + ..Default::default() + }, + config, }, - config, - }); + Default::default(), + ); pool.launch(); let initial_total = pool.lock().total(); diff --git a/pgdog/src/backend/pool/pool_impl.rs b/pgdog/src/backend/pool/pool_impl.rs index 1b161e22d..12bcc75f7 100644 --- a/pgdog/src/backend/pool/pool_impl.rs +++ b/pgdog/src/backend/pool/pool_impl.rs @@ -20,8 +20,8 @@ use crate::net::{Parameter, Parameters}; use super::inner::CheckInResult; use super::{ - Address, Comms, Config, Error, Guard, Healtcheck, Inner, Monitor, PoolConfig, Request, State, - Waiting, + Address, Comms, Config, Error, Guard, Healtcheck, Inner, Monitor, Oids, PoolConfig, Request, + State, Waiting, lb::TargetHealth, lsn_monitor::{LsnMonitor, ReplicaLag}, }; @@ -48,6 +48,7 @@ pub(crate) struct InnerSync { pub(super) params: OnceCell, pub(super) lsn_stats: RwLock, pub(super) lsn_role_change: Notify, + pub(super) oids: Arc, } impl std::fmt::Debug for Pool { @@ -60,7 +61,7 @@ impl std::fmt::Debug for Pool { impl Pool { /// Create new connection pool. - pub fn new(config: &PoolConfig) -> Self { + pub(crate) fn new(config: &PoolConfig, oids: Arc) -> Self { let id = next_pool_id(); Self { inner: Arc::new(InnerSync { @@ -73,6 +74,7 @@ impl Pool { params: OnceCell::new(), lsn_stats: RwLock::new(LsnStats::default()), lsn_role_change: Notify::new(), + oids, }), } } @@ -85,7 +87,7 @@ impl Pool { config: Config::default(), }; - Self::new(&config) + Self::new(&config, Default::default()) } pub(crate) fn inner(&self) -> &InnerSync { diff --git a/pgdog/src/backend/pool/shard/mod.rs b/pgdog/src/backend/pool/shard/mod.rs index 0373d0ba5..af2c291f2 100644 --- a/pgdog/src/backend/pool/shard/mod.rs +++ b/pgdog/src/backend/pool/shard/mod.rs @@ -22,9 +22,11 @@ use crate::net::messages::FrontendPid; use super::{Error, Guard, LoadBalancer, Pool, PoolConfig, Request}; pub mod monitor; +mod oids; pub mod role_detector; use monitor::*; +pub(crate) use oids::{CanonicalOids, OidMappings, Oids}; use role_detector::*; #[cfg_attr(test, derive(Default))] @@ -53,7 +55,7 @@ pub(super) struct ShardConfig<'a> { /// /// Includes a primary and replicas. #[derive(Clone, Debug)] -pub struct Shard { +pub(crate) struct Shard { inner: Arc, } @@ -134,6 +136,7 @@ impl Shard { return Ok(false); } + self.load_oids().await?; // This is syncrhonized by database/shard number, so this prevents // a thundering herd with 100s of users, for example, all fetching // the same schema. @@ -160,16 +163,23 @@ impl Shard { Ok(schema) } + async fn load_oids(&self) -> Result<(), crate::backend::Error> { + self.oids.load(self).await?; + Ok(()) + } + /// Set the schema to its default value. /// We don't need it for this shard. pub(super) fn schema_not_needed(&self) { let _ = self.schema.set(Schema::default()); + self.oids.skip_load(); } /// Wait for the shard to load the schema. /// If the schema is loaded already, this returns immediately. pub(super) async fn wait_schema_loaded(&self) { self.schema.wait().await; + self.oids.wait().await; } /// Check that the shard LB targets are all launched. @@ -321,8 +331,9 @@ impl Deref for Shard { /// Shard connection pools /// and internal state. -#[derive(Default, Debug)] -pub struct ShardInner { +#[cfg_attr(test, derive(Default))] +#[derive(Debug)] +pub(crate) struct ShardInner { number: usize, lb: LoadBalancer, comms: Arc, @@ -331,6 +342,7 @@ pub struct ShardInner { schema: SetOnce, pub_sub_enabled: bool, schema_cache: SchemaCache, + oids: Arc, } impl ShardInner { @@ -346,8 +358,9 @@ impl ShardInner { pub_sub_enabled, schema_cache, } = shard; - let primary = primary.map(Pool::new); - let lb = LoadBalancer::new(&primary, replicas, lb_strategy, rw_split); + let oids = schema_cache.oids(&identifier.database, number); + let primary = primary.map(|config| Pool::new(config, Arc::clone(&oids))); + let lb = LoadBalancer::new(&primary, replicas, lb_strategy, rw_split, Arc::clone(&oids)); let comms = Arc::new(ShardComms { shutdown: CancellationToken::new(), lsn_check_interval, @@ -362,6 +375,7 @@ impl ShardInner { schema: SetOnce::new(), pub_sub_enabled, schema_cache, + oids, } } } diff --git a/pgdog/src/backend/pool/shard/monitor.rs b/pgdog/src/backend/pool/shard/monitor.rs index af3df05e0..86f769bde 100644 --- a/pgdog/src/backend/pool/shard/monitor.rs +++ b/pgdog/src/backend/pool/shard/monitor.rs @@ -249,17 +249,23 @@ mod test { #[test] fn test_update_replica_lag_assigns_primary_minus_replica_to_replica_pool() { - let primary = Pool::new(&PoolConfig { - address: Address::new_test(), - config: Config::default(), - }); - let replica = Pool::new(&PoolConfig { - address: Address { - configured_role: Role::Replica, - ..Address::new_test() + let primary = Pool::new( + &PoolConfig { + address: Address::new_test(), + config: Config::default(), }, - config: Config::default(), - }); + Default::default(), + ); + let replica = Pool::new( + &PoolConfig { + address: Address { + configured_role: Role::Replica, + ..Address::new_test() + }, + config: Config::default(), + }, + Default::default(), + ); set_pool_lsn_stats(&primary, false, 200, "2026-07-01 13:33:10.000000+00"); set_pool_lsn_stats(&replica, true, 100, "2026-07-01 13:33:00.000000+00"); diff --git a/pgdog/src/backend/pool/shard/oids.rs b/pgdog/src/backend/pool/shard/oids.rs new file mode 100644 index 000000000..90eb8ffbc --- /dev/null +++ b/pgdog/src/backend/pool/shard/oids.rs @@ -0,0 +1,132 @@ +use super::{Request, Shard}; +use crate::{ + backend::{Error, Server}, + net::DataRow, + sync::SetOnceCell, +}; +use std::collections::HashMap; +use std::sync::Arc; +use tracing::info; + +#[derive(Debug)] +/// The mapping from a shards type OID to a canonical one +pub(crate) struct Oids { + canonical_oids: Arc, + mappings: SetOnceCell, +} + +impl Oids { + pub(crate) fn new(canonical_oids: &Arc) -> Arc { + Arc::new(Self { + canonical_oids: Arc::clone(canonical_oids), + mappings: Default::default(), + }) + } + + pub(crate) async fn load(&self, shard: &Shard) -> Result<&OidMappings, Error> { + self.mappings + .get_or_try_init(|| async { + let mut server = shard.primary_or_replica(&Request::default()).await?; + let oids = load_oids(&mut server).await?; + let canonical = self.canonical_oids.oids.wait().await; + let mut canonical_to_shard = HashMap::new(); + let mut shard_to_canonical = HashMap::new(); + for (type_name, oid) in oids { + let canonical = canonical + .get(&type_name) + .copied() + .ok_or(Error::MissingCanonicalOid(type_name))?; + if canonical == oid { + continue; + } + canonical_to_shard.insert(canonical, oid); + shard_to_canonical.insert(oid, canonical); + } + + debug_assert_eq!(canonical_to_shard.len(), shard_to_canonical.len()); + info!( + "loaded type info for {} types on shard {} [{}]", + canonical_to_shard.len(), + shard.number(), + server.addr(), + ); + + Ok(OidMappings { + canonical_to_shard, + shard_to_canonical, + }) + }) + .await + } + + pub(crate) async fn wait(&self) { + self.mappings.wait().await; + } + + /// Sets this to an empty mapping if it has not already been loaded + pub(crate) fn skip_load(&self) { + let _ = self.mappings.set(Default::default()); + } + + /// Get the mappings. Returns `None` if no mappings have been loaded + pub(crate) fn get(&self) -> Option<&OidMappings> { + self.mappings.get() + } + + #[cfg(test)] + pub(crate) fn from_canonical(canonical_to_shard: HashMap) -> Arc { + let shard_to_canonical = canonical_to_shard.iter().map(|(&k, &v)| (v, k)).collect(); + Arc::new(Self { + canonical_oids: Default::default(), + mappings: SetOnceCell::from(OidMappings { + canonical_to_shard, + shard_to_canonical, + }), + }) + } +} + +impl Default for Oids { + fn default() -> Self { + Self { + canonical_oids: Default::default(), + mappings: SetOnceCell::from(OidMappings::default()), + } + } +} + +#[derive(Debug, Default)] +pub(crate) struct OidMappings { + pub(crate) canonical_to_shard: HashMap, + pub(crate) shard_to_canonical: HashMap, +} + +#[derive(Debug, Default)] +pub(crate) struct CanonicalOids { + oids: SetOnceCell>, +} + +impl CanonicalOids { + pub(crate) async fn load(&self, server: &mut Server) -> Result<(), Error> { + self.oids + .get_or_try_init(|| async { Ok(load_oids(server).await?.collect()) }) + .await + .map(|_| ()) + } +} + +async fn load_oids(server: &mut Server) -> Result, Error> { + // OIDs < 10,000 are reserved for PG's internal use and are assumed to be stable + Ok(server + .fetch_all::( + "SELECT DISTINCT ON (typname) typname, pg_type.oid FROM pg_type WHERE pg_type.oid >= 10000", + ) + .await? + .into_iter() + .map(|row| { + ( + row.get_text(0).expect("selected 2 columns"), + row.get_int(1, true).expect("selected 2 columns") as u32, + ) + })) +} diff --git a/pgdog/src/backend/pool/test/mod.rs b/pgdog/src/backend/pool/test/mod.rs index 35af9706d..aacb9cff8 100644 --- a/pgdog/src/backend/pool/test/mod.rs +++ b/pgdog/src/backend/pool/test/mod.rs @@ -28,17 +28,20 @@ pub fn pool() -> Pool { }, }; - let pool = Pool::new(&PoolConfig { - address: Address { - host: "127.0.0.1".into(), - port: 5432, - database_name: "pgdog".into(), - user: "pgdog".into(), - passwords: vec!["pgdog".into()], - ..Default::default() + let pool = Pool::new( + &PoolConfig { + address: Address { + host: "127.0.0.1".into(), + port: 5432, + database_name: "pgdog".into(), + user: "pgdog".into(), + passwords: vec!["pgdog".into()], + ..Default::default() + }, + config, }, - config, - }); + Default::default(), + ); pool.launch(); pool } @@ -53,17 +56,20 @@ pub fn pool_with_prepared_capacity(capacity: usize) -> Pool { }, }; - let pool = Pool::new(&PoolConfig { - address: Address { - host: "127.0.0.1".into(), - port: 5432, - database_name: "pgdog".into(), - user: "pgdog".into(), - passwords: vec!["pgdog".into()], - ..Default::default() + let pool = Pool::new( + &PoolConfig { + address: Address { + host: "127.0.0.1".into(), + port: 5432, + database_name: "pgdog".into(), + user: "pgdog".into(), + passwords: vec!["pgdog".into()], + ..Default::default() + }, + config, }, - config, - }); + Default::default(), + ); pool.launch(); pool } @@ -326,17 +332,20 @@ async fn test_server_force_close_discards_connection() { }, }; - let pool = Pool::new(&PoolConfig { - address: Address { - host: "127.0.0.1".into(), - port: 5432, - database_name: "pgdog".into(), - user: "pgdog".into(), - passwords: vec!["pgdog".into()], - ..Default::default() + let pool = Pool::new( + &PoolConfig { + address: Address { + host: "127.0.0.1".into(), + port: 5432, + database_name: "pgdog".into(), + user: "pgdog".into(), + passwords: vec!["pgdog".into()], + ..Default::default() + }, + config, }, - config, - }); + Default::default(), + ); pool.launch(); let mut conn = pool.get(&Request::default()).await.unwrap(); @@ -497,17 +506,20 @@ async fn test_idle_healthcheck_loop() { }, }; - let pool = Pool::new(&PoolConfig { - address: Address { - host: "127.0.0.1".into(), - port: 5432, - database_name: "pgdog".into(), - user: "pgdog".into(), - passwords: vec!["pgdog".into()], - ..Default::default() + let pool = Pool::new( + &PoolConfig { + address: Address { + host: "127.0.0.1".into(), + port: 5432, + database_name: "pgdog".into(), + user: "pgdog".into(), + passwords: vec!["pgdog".into()], + ..Default::default() + }, + config, }, - config, - }); + Default::default(), + ); pool.launch(); let initial_healthchecks = pool.state().stats.counts.healthchecks; @@ -545,17 +557,20 @@ async fn test_idle_healthcheck_loop_disabled_with_zero_interval() { }, }; - let pool = Pool::new(&PoolConfig { - address: Address { - host: "127.0.0.1".into(), - port: 1, - database_name: "pgdog".into(), - user: "pgdog".into(), - passwords: vec!["pgdog".into()], - ..Default::default() + let pool = Pool::new( + &PoolConfig { + address: Address { + host: "127.0.0.1".into(), + port: 1, + database_name: "pgdog".into(), + user: "pgdog".into(), + passwords: vec!["pgdog".into()], + ..Default::default() + }, + config, }, - config, - }); + Default::default(), + ); pool.launch(); let initial_healthchecks = pool.state().stats.counts.healthchecks; @@ -581,10 +596,13 @@ async fn test_checkout_timeout() { }, }; - let pool = Pool::new(&PoolConfig { - address: Address::new_test(), - config, - }); + let pool = Pool::new( + &PoolConfig { + address: Address::new_test(), + config, + }, + Default::default(), + ); pool.launch(); // Hold the only connection @@ -609,30 +627,36 @@ async fn test_move_conns_to() { }, }; - let source = Pool::new(&PoolConfig { - address: Address { - host: "127.0.0.1".into(), - port: 5432, - database_name: "pgdog".into(), - user: "pgdog".into(), - passwords: vec!["pgdog".into()], - ..Default::default() + let source = Pool::new( + &PoolConfig { + address: Address { + host: "127.0.0.1".into(), + port: 5432, + database_name: "pgdog".into(), + user: "pgdog".into(), + passwords: vec!["pgdog".into()], + ..Default::default() + }, + config, }, - config, - }); + Default::default(), + ); source.launch(); - let destination = Pool::new(&PoolConfig { - address: Address { - host: "127.0.0.1".into(), - port: 5432, - database_name: "pgdog".into(), - user: "pgdog".into(), - passwords: vec!["pgdog".into()], - ..Default::default() + let destination = Pool::new( + &PoolConfig { + address: Address { + host: "127.0.0.1".into(), + port: 5432, + database_name: "pgdog".into(), + user: "pgdog".into(), + passwords: vec!["pgdog".into()], + ..Default::default() + }, + config, }, - config, - }); + Default::default(), + ); let conn1 = source.get(&Request::default()).await.unwrap(); let conn2 = source.get(&Request::default()).await.unwrap(); @@ -678,16 +702,22 @@ async fn test_move_conns_all_idle() { }, }; - let source = Pool::new(&PoolConfig { - address: Address::new_test(), - config, - }); + let source = Pool::new( + &PoolConfig { + address: Address::new_test(), + config, + }, + Default::default(), + ); source.launch(); - let destination = Pool::new(&PoolConfig { - address: Address::new_test(), - config, - }); + let destination = Pool::new( + &PoolConfig { + address: Address::new_test(), + config, + }, + Default::default(), + ); // Check out and return 3 connections so they become idle. let c1 = source.get(&Request::default()).await.unwrap(); @@ -725,16 +755,22 @@ async fn test_move_conns_all_checked_out() { }, }; - let source = Pool::new(&PoolConfig { - address: Address::new_test(), - config, - }); + let source = Pool::new( + &PoolConfig { + address: Address::new_test(), + config, + }, + Default::default(), + ); source.launch(); - let destination = Pool::new(&PoolConfig { - address: Address::new_test(), - config, - }); + let destination = Pool::new( + &PoolConfig { + address: Address::new_test(), + config, + }, + Default::default(), + ); let c1 = source.get(&Request::default()).await.unwrap(); let c2 = source.get(&Request::default()).await.unwrap(); @@ -780,16 +816,22 @@ async fn test_move_conns_destination_serves_after_launch() { }, }; - let source = Pool::new(&PoolConfig { - address: Address::new_test(), - config, - }); + let source = Pool::new( + &PoolConfig { + address: Address::new_test(), + config, + }, + Default::default(), + ); source.launch(); - let destination = Pool::new(&PoolConfig { - address: Address::new_test(), - config, - }); + let destination = Pool::new( + &PoolConfig { + address: Address::new_test(), + config, + }, + Default::default(), + ); // Create one idle connection. let c1 = source.get(&Request::default()).await.unwrap(); @@ -822,17 +864,20 @@ fn auth_pool(passwords: Vec) -> Pool { }, }; - Pool::new(&PoolConfig { - address: Address { - host: "127.0.0.1".into(), - port: 5432, - database_name: "pgdog".into(), - user: "pgdog".into(), - passwords, - ..Default::default() + Pool::new( + &PoolConfig { + address: Address { + host: "127.0.0.1".into(), + port: 5432, + database_name: "pgdog".into(), + user: "pgdog".into(), + passwords, + ..Default::default() + }, + config, }, - config, - }) + Default::default(), + ) } #[tokio::test] @@ -961,10 +1006,13 @@ async fn test_lsn_monitor() { }, }; - let pool = Pool::new(&PoolConfig { - address: Address::new_test(), - config, - }); + let pool = Pool::new( + &PoolConfig { + address: Address::new_test(), + config, + }, + Default::default(), + ); let initial_stats = pool.lsn_stats(); assert!(!initial_stats.valid()); @@ -1020,10 +1068,13 @@ async fn test_token_refresh_loop_primes_cache_on_cold_start() { let expiry = SystemTime::now() + Duration::from_millis(200); TokenCache::global().set(&addr, "initial-token".into(), expiry); - let pool = Pool::new(&PoolConfig { - address: addr.clone(), - config, - }); + let pool = Pool::new( + &PoolConfig { + address: addr.clone(), + config, + }, + Default::default(), + ); pool.launch(); // Cache must be populated immediately. @@ -1059,10 +1110,13 @@ async fn test_token_refresh_loop_refreshes_before_expiry() { let expiry = SystemTime::now() + Duration::from_secs(3600); TokenCache::global().set(&addr, "long-lived-token".into(), expiry); - let pool = Pool::new(&PoolConfig { - address: addr.clone(), - config, - }); + let pool = Pool::new( + &PoolConfig { + address: addr.clone(), + config, + }, + Default::default(), + ); pool.launch(); sleep(Duration::from_millis(100)).await; @@ -1104,10 +1158,13 @@ async fn test_token_refresh_loop_evicts_on_failed_refresh() { let expiry = SystemTime::now() + Duration::from_secs(10); TokenCache::global().set(&addr, "stale-token".into(), expiry); - let pool = Pool::new(&PoolConfig { - address: addr.clone(), - config, - }); + let pool = Pool::new( + &PoolConfig { + address: addr.clone(), + config, + }, + Default::default(), + ); pool.launch(); // Give the refresh loop time to fire and fail. @@ -1146,10 +1203,13 @@ async fn test_token_refresh_loop_not_spawned_for_password_auth() { // Poison the cache to detect any unexpected writes. TokenCache::global().evict(&addr); - let pool = Pool::new(&PoolConfig { - address: addr.clone(), - config, - }); + let pool = Pool::new( + &PoolConfig { + address: addr.clone(), + config, + }, + Default::default(), + ); pool.launch(); sleep(Duration::from_millis(100)).await; @@ -1186,10 +1246,13 @@ async fn test_token_refresh_loop_stops_on_shutdown() { let expiry = SystemTime::now() + Duration::from_secs(3600); TokenCache::global().set(&addr, "token".into(), expiry); - let pool = Pool::new(&PoolConfig { - address: addr.clone(), - config, - }); + let pool = Pool::new( + &PoolConfig { + address: addr.clone(), + config, + }, + Default::default(), + ); pool.launch(); sleep(Duration::from_millis(50)).await; diff --git a/pgdog/src/backend/pool/waiting.rs b/pgdog/src/backend/pool/waiting.rs index 8b02aad2e..0a586e574 100644 --- a/pgdog/src/backend/pool/waiting.rs +++ b/pgdog/src/backend/pool/waiting.rs @@ -150,17 +150,20 @@ mod tests { }, }; - let pool = Pool::new(&crate::backend::pool::PoolConfig { - address: crate::backend::pool::Address { - host: "127.0.0.1".into(), - port: 5432, - database_name: "pgdog".into(), - user: "pgdog".into(), - passwords: vec!["pgdog".into()], - ..Default::default() + let pool = Pool::new( + &crate::backend::pool::PoolConfig { + address: crate::backend::pool::Address { + host: "127.0.0.1".into(), + port: 5432, + database_name: "pgdog".into(), + user: "pgdog".into(), + passwords: vec!["pgdog".into()], + ..Default::default() + }, + config, }, - config, - }); + Default::default(), + ); pool.launch(); sleep(Duration::from_millis(100)).await; diff --git a/pgdog/src/backend/prepared_statements.rs b/pgdog/src/backend/prepared_statements.rs index 40f720af9..e362dbead 100644 --- a/pgdog/src/backend/prepared_statements.rs +++ b/pgdog/src/backend/prepared_statements.rs @@ -6,13 +6,13 @@ use crate::{ net::{ Close, CloseComplete, FromBytes, Message, ParseComplete, Protocol, ProtocolMessage, ToBytes, - messages::{RowDescription, parse::Parse}, + messages::{ParameterDescription, RowDescription, parse::Parse}, }, }; use parking_lot::RwLock; use pgdog_config::PreparedStatements as PreparedStatementsLevel; -use super::Error; +use super::{Error, OidMappings, Oids}; use super::{ protocol::{ProtocolState, state::Action}, state::ExecutionCode, @@ -24,7 +24,7 @@ fn str_mem(s: &str) -> usize { s.len() + std::mem::size_of::() } -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq)] pub enum HandleResult { Forward, Drop, @@ -53,17 +53,19 @@ pub struct PreparedStatements { capacity: usize, memory_used: usize, level: PreparedStatementsLevel, + oids: Arc, } +#[cfg(test)] impl Default for PreparedStatements { fn default() -> Self { - Self::new() + Self::new(Default::default()) } } impl PreparedStatements { /// New server prepared statements. - pub fn new() -> Self { + pub(crate) fn new(oids: Arc) -> Self { Self { global_cache: frontend::PreparedStatements::global(), local_cache: LruCache::unbounded(), @@ -73,6 +75,7 @@ impl PreparedStatements { capacity: usize::MAX, memory_used: 0, level: PreparedStatementsLevel::default(), + oids, } } @@ -208,24 +211,28 @@ impl PreparedStatements { } ProtocolMessage::Parse(parse) => { + let mut parse = parse.clone(); + let mut rewritten = self.rewrite_parse_data_types(&mut parse); + if !parse.anonymous() { if self.contains(parse.name()) { self.state.add_simulated(ParseComplete.message()?); return Ok(HandleResult::Drop); } else { - self.state.add('1'); self.parses.push_back(parse.name().to_string()); } // The client is sending named prepared statements, // but we're in ExtendedAnonymous mode so we rewrite // them to anonymous to avoid storing them in Postgres. if self.level.rewrite_anonymous() { - let mut parse = parse.clone(); parse.anonymize(); - return Ok(HandleResult::Rewrite(ProtocolMessage::Parse(parse))); + rewritten = true; } - } else { - self.state.add('1'); + } + + self.state.add('1'); + if rewritten { + return Ok(HandleResult::Rewrite(ProtocolMessage::Parse(parse))); } } @@ -272,7 +279,7 @@ impl PreparedStatements { } /// Should we forward the message to the client. - pub fn forward(&mut self, message: &Message) -> Result { + pub(crate) fn forward(&mut self, message: &mut Message) -> Result { let code = message.code(); let action = self.state.action(code)?; @@ -288,10 +295,7 @@ impl PreparedStatements { 'T' => { if let Some(describe) = self.describes.pop_front() { - self.add_row_description( - &describe, - &RowDescription::from_bytes(message.to_bytes())?, - ); + self.add_row_description(&describe, message)?; }; } @@ -315,6 +319,10 @@ impl PreparedStatements { self.state.action(code)?; } + 't' => { + self.rewrite_parameter_description_data_types(message)?; + } + _ => (), } @@ -383,21 +391,25 @@ impl PreparedStatements { /// Get the Parse message stored in the global prepared statements /// cache for this statement. pub(crate) fn parse(&self, name: &str) -> Option { - self.global_cache.read().rewritten_parse(name) - } - - /// Get the globally stored RowDescription for this prepared statement, - /// if any. - pub fn row_description(&self, name: &str) -> Option { - self.global_cache.read().row_description(name) + self.global_cache + .read() + .rewritten_parse(name) + .map(|mut parse| { + self.rewrite_parse_data_types(&mut parse); + parse + }) } /// Handle a Describe message, storing the RowDescription for the /// statement in the global cache. - fn add_row_description(&self, name: &str, row_description: &RowDescription) { + fn add_row_description(&self, name: &str, message: &mut Message) -> Result<(), Error> { + let mut row_description = RowDescription::from_bytes(message.payload())?; + self.rewrite_row_description_data_types(&mut row_description); + message.replace_payload(row_description.to_bytes()); self.global_cache .write() .insert_row_description(name, row_description); + Ok(()) } /// Remove statement from local cache. @@ -459,6 +471,36 @@ impl PreparedStatements { close } + + pub(crate) fn replace_oids(&mut self, oids: &Arc) { + self.oids = Arc::clone(oids) + } + + fn rewrite_parse_data_types(&self, parse: &mut Parse) -> bool { + parse.rewrite_data_types(&self.oid_mappings().canonical_to_shard) + } + + fn rewrite_row_description_data_types(&self, row_description: &mut RowDescription) { + row_description.rewrite_data_types(&self.oid_mappings().shard_to_canonical); + } + + fn rewrite_parameter_description_data_types(&self, message: &mut Message) -> Result<(), Error> { + let mappings = &self.oid_mappings().shard_to_canonical; + if mappings.is_empty() { + return Ok(()); + } + + let mut parameter_description = ParameterDescription::from_bytes(message.payload())?; + parameter_description.rewrite_data_types(&self.oid_mappings().shard_to_canonical); + message.replace_payload(parameter_description.to_bytes()); + Ok(()) + } + + fn oid_mappings(&self) -> &OidMappings { + self.oids + .get() + .expect("cluster startup always loads OID information or sets an empty mapping") + } } #[cfg(test)] @@ -473,14 +515,14 @@ mod test { /// Build a PreparedStatements instance configured for ExtendedAnonymous mode. fn new_extended_anonymous() -> PreparedStatements { - let mut ps = PreparedStatements::new(); + let mut ps = PreparedStatements::default(); ps.set_prepared_statements_level(PreparedStatementsLevel::ExtendedAnonymous); ps } /// Build a PreparedStatements instance configured for Extended (default) mode. fn new_extended() -> PreparedStatements { - let mut ps = PreparedStatements::new(); + let mut ps = PreparedStatements::default(); ps.set_prepared_statements_level(PreparedStatementsLevel::Extended); ps } @@ -709,8 +751,8 @@ mod test { // Simulate a ReadyForQuery message. // First we need to add a 'Z' to the state so we can action it. ps.state.add('Z'); - let rfq = Message::new(ReadyForQuery::idle().to_bytes()); - ps.forward(&rfq).unwrap(); + let mut rfq = Message::new(ReadyForQuery::idle().to_bytes()); + ps.forward(&mut rfq).unwrap(); // In extended_anonymous mode, cache should be cleared after done. assert_eq!(ps.len(), 0, "local cache should be cleared after RFQ"); @@ -724,8 +766,8 @@ mod test { assert_eq!(ps.len(), 2); ps.state.add('Z'); - let rfq = Message::new(ReadyForQuery::idle().to_bytes()); - ps.forward(&rfq).unwrap(); + let mut rfq = Message::new(ReadyForQuery::idle().to_bytes()); + ps.forward(&mut rfq).unwrap(); // In extended mode, cache should be preserved. assert_eq!( @@ -807,6 +849,21 @@ mod test { assert!(matches!(result, HandleResult::Forward)); } + #[test] + fn parse_rewrites_if_oids_change() { + let mut ps = new_extended(); + ps.oids = Oids::from_canonical([(10000, 10001)].into_iter().collect()); + + let parse = Parse::named("stmt1", "SELECT $1, $2"); + let client_parse = parse.with_data_types(&[10000, 10002]); + let result = ps.handle(&ProtocolMessage::Parse(client_parse)).unwrap(); + let expected = parse.with_data_types(&[10001, 10002]); + assert_eq!( + result, + HandleResult::Rewrite(ProtocolMessage::Parse(expected)) + ); + } + // ------------------------------------------------------- // Simple query is unaffected by mode // ------------------------------------------------------- diff --git a/pgdog/src/backend/replication/logical/publisher/slot.rs b/pgdog/src/backend/replication/logical/publisher/slot.rs index d34e11a12..b9fb0bc20 100644 --- a/pgdog/src/backend/replication/logical/publisher/slot.rs +++ b/pgdog/src/backend/replication/logical/publisher/slot.rs @@ -110,6 +110,7 @@ impl ReplicationSlot { &self.address, ServerOptions::new_replication(), ConnectReason::Replication, + Default::default(), ) .await?, ); @@ -127,6 +128,7 @@ impl ReplicationSlot { &self.address, ServerOptions::default(), ConnectReason::Replication, + Default::default(), ) .await?, ); diff --git a/pgdog/src/backend/replication/logical/subscriber/tests.rs b/pgdog/src/backend/replication/logical/subscriber/tests.rs index 544f2474a..c046a4986 100644 --- a/pgdog/src/backend/replication/logical/subscriber/tests.rs +++ b/pgdog/src/backend/replication/logical/subscriber/tests.rs @@ -250,12 +250,14 @@ fn x_update(u: XLogUpdate) -> CopyData { fn make_subscriber() -> StreamSubscriber { let cluster = Cluster::new_test(&config()); + cluster.force_dummy_schema(); let tables = vec![make_sharded_table(), make_sharded_test_b_table()]; StreamSubscriber::new(&cluster, &tables, OmniOwnership::test()) } fn make_subscriber_with_tables(tables: Vec) -> StreamSubscriber { let cluster = Cluster::new_test(&config()); + cluster.force_dummy_schema(); StreamSubscriber::new(&cluster, &tables, OmniOwnership::test()) } @@ -264,11 +266,13 @@ fn make_subscriber_with_tables_two_databases( partition: OmniOwnership, ) -> StreamSubscriber { let cluster = Cluster::new_test_two_databases(&config()); + cluster.force_dummy_schema(); StreamSubscriber::new(&cluster, &tables, partition) } fn make_subscriber_single_shard() -> StreamSubscriber { let cluster = Cluster::new_test_single_shard(&config()); + cluster.force_dummy_schema(); let tables = vec![make_sharded_table(), make_sharded_test_b_table()]; StreamSubscriber::new(&cluster, &tables, OmniOwnership::test()) } @@ -610,6 +614,7 @@ async fn partition_leaves_share_destination() { leaf_b.table.parent_name = "sharded".to_string(); let cluster = Cluster::new_test_single_shard(&config()); + cluster.force_dummy_schema(); let mut sub = StreamSubscriber::new(&cluster, &[leaf_a, leaf_b], OmniOwnership::test()); let mut verify = test_server().await; sub.connect().await.unwrap(); @@ -1546,6 +1551,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()); + cluster.force_dummy_schema(); let mut sub = StreamSubscriber::new( &cluster, &[make_replica_identity_nothing_table()], @@ -1584,6 +1590,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()); + cluster.force_dummy_schema(); let mut sub = StreamSubscriber::new( &cluster, &[make_full_identity_omni_table()], @@ -1627,6 +1634,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()); + cluster.force_dummy_schema(); let mut sub = StreamSubscriber::new( &cluster, &[make_full_identity_sharded_table()], @@ -1657,6 +1665,7 @@ async fn full_identity_insert_sharded() { #[tokio::test] async fn full_identity_update_fast_path() { let cluster = Cluster::new_test_single_shard(&config()); + cluster.force_dummy_schema(); let mut sub = StreamSubscriber::new( &cluster, &[make_full_identity_sharded_table()], @@ -1716,6 +1725,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()); + cluster.force_dummy_schema(); let mut sub = StreamSubscriber::new( &cluster, &[make_full_identity_sharded_table()], @@ -1780,6 +1790,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()); + cluster.force_dummy_schema(); let mut sub = StreamSubscriber::new( &cluster, &[make_full_identity_sharded_table()], @@ -1841,6 +1852,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()); + cluster.force_dummy_schema(); let mut sub = StreamSubscriber::new( &cluster, &[make_full_identity_sharded_table()], @@ -1887,6 +1899,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()); + cluster.force_dummy_schema(); let mut sub = StreamSubscriber::new( &cluster, &[make_full_identity_sharded_table()], @@ -1929,6 +1942,7 @@ async fn full_identity_delete() { #[tokio::test] async fn full_identity_insert_omni_dedup() { let cluster = Cluster::new_test_single_shard(&config()); + cluster.force_dummy_schema(); let mut sub = StreamSubscriber::new( &cluster, &[make_full_identity_omni_dedup_table()], @@ -1992,6 +2006,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()); + cluster.force_dummy_schema(); let mut sub = StreamSubscriber::new( &cluster, &[make_full_identity_dup_rows_table()], @@ -2062,6 +2077,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()); + cluster.force_dummy_schema(); let mut sub = StreamSubscriber::new( &cluster, &[make_full_identity_dup_rows_table()], @@ -2133,6 +2149,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()); + cluster.force_dummy_schema(); let mut sub = StreamSubscriber::new( &cluster, &[make_full_identity_dup_rows_table()], @@ -2198,6 +2215,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()); + cluster.force_dummy_schema(); let mut sub = StreamSubscriber::new( &cluster, &[make_full_identity_dup_rows_table()], diff --git a/pgdog/src/backend/schema/cache/mod.rs b/pgdog/src/backend/schema/cache/mod.rs index 3bfc62295..795beab59 100644 --- a/pgdog/src/backend/schema/cache/mod.rs +++ b/pgdog/src/backend/schema/cache/mod.rs @@ -5,7 +5,7 @@ use dashmap::DashMap; use std::sync::Arc; use tokio::sync::Mutex; -use crate::backend::{Schema, Shard}; +use crate::backend::{CanonicalOids, Oids, Schema, Shard}; type Entry = Arc>; @@ -13,7 +13,12 @@ type Entry = Arc>; #[derive(Debug, Default, Clone)] pub(crate) struct SchemaCache { // Database => shard => Schema - cache: Arc>>, + cache: Arc>, + /// The canonical mapping of type names to OID + /// A cluster's canonical mapping of type names to OID + canonical_oids: Arc>>, + /// Each database, shard pair's mappings to the canonical type OIDs + shard_oids: Arc>>, } impl SchemaCache { @@ -26,9 +31,7 @@ impl SchemaCache { // This is synchronized. let entry = self .cache - .entry(shard.identifier().database.clone()) - .or_default() - .entry(shard.number()) + .entry((shard.identifier().database.clone(), shard.number())) .or_default() .clone(); @@ -46,4 +49,17 @@ impl SchemaCache { Ok(schema) } + + pub(crate) fn canonical_oids(&self, database: &str) -> Arc { + Arc::clone(&self.canonical_oids.entry(database.to_owned()).or_default()) + } + + pub(crate) fn oids(&self, database: &str, shard_number: usize) -> Arc { + Arc::clone( + &self + .shard_oids + .entry((database.to_owned(), shard_number)) + .or_insert_with(|| Oids::new(&self.canonical_oids(database))), + ) + } } diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index aa715aa81..e98e108ee 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -1,6 +1,6 @@ //! PostgreSQL server connection. -use std::{ops::Deref, time::Duration}; +use std::{ops::Deref, sync::Arc, time::Duration}; use bytes::{BufMut, BytesMut}; use rustls_pki_types::ServerName; @@ -13,7 +13,7 @@ use tokio::{ use tracing::{debug, error, info, trace, warn}; use super::{ - ConnectReason, DisconnectReason, Error, PreparedStatements, ServerOptions, Stats, + ConnectReason, DisconnectReason, Error, Oids, PreparedStatements, ServerOptions, Stats, pool::Address, prepared_statements::HandleResult, }; use crate::{ @@ -96,10 +96,11 @@ impl MemoryUsage for Server { impl Server { /// Create new PostgreSQL server connection. - pub async fn connect( + pub(crate) async fn connect( addr: &Address, options: ServerOptions, connect_reason: ConnectReason, + oids: Arc, ) -> Result { let (user, auth_secrets) = addr.auth_credentials().await?; let total = auth_secrets.len(); @@ -110,6 +111,7 @@ impl Server { options.clone(), connect_reason, &auth_secret, + Arc::clone(&oids), ) .await { @@ -151,6 +153,7 @@ impl Server { options: ServerOptions, connect_reason: ConnectReason, auth_secret: &super::pool::Password, + oids: Arc, ) -> Result { debug!("=> {}", addr); let stream = TcpStream::connect(addr.addr().await?).await?; @@ -341,7 +344,7 @@ impl Server { params, changed_params: Parameters::default(), client_params: Parameters::default(), - prepared_statements: PreparedStatements::new(), + prepared_statements: PreparedStatements::new(oids), dirty: false, streaming: false, schema_changed: false, @@ -477,8 +480,8 @@ impl Server { Ok(message) => { // INVARIANT: omni dedup in multi_shard relies on this being process-unique; // never substitute a non-unique value here. - let message = message.stream(self.streaming).backend(self.id); - match self.prepared_statements.forward(&message) { + let mut message = message.stream(self.streaming).backend(self.id); + match self.prepared_statements.forward(&mut message) { Ok(forward) => { if forward { break message; @@ -1174,6 +1177,10 @@ impl Server { }, } } + + pub(crate) fn replace_oids(&mut self, oids: &Arc) { + self.prepared_statements.replace_oids(oids); + } } impl Drop for Server { @@ -1247,7 +1254,7 @@ pub mod test { &ServerOptions::default(), &Memory::default(), ), - prepared_statements: super::PreparedStatements::new(), + prepared_statements: super::PreparedStatements::default(), addr, dirty: false, streaming: false, @@ -1277,11 +1284,12 @@ pub mod test { } } - pub async fn test_server() -> Server { + pub(crate) async fn test_server() -> Server { Server::connect( &Address::new_test(), ServerOptions::default(), ConnectReason::Other, + Default::default(), ) .await .unwrap() @@ -1298,6 +1306,7 @@ pub mod test { }, ServerOptions::default(), ConnectReason::Other, + Default::default(), ) .await .unwrap() @@ -1308,6 +1317,7 @@ pub mod test { &Address::new_test(), ServerOptions::new_replication(), ConnectReason::Replication, + Default::default(), ) .await .unwrap() @@ -1366,7 +1376,13 @@ pub mod test { expected_secret, SystemTime::now() + Duration::from_secs(3600), ); - let result = Server::connect(&addr, ServerOptions::default(), ConnectReason::Other).await; + let result = Server::connect( + &addr, + ServerOptions::default(), + ConnectReason::Other, + Default::default(), + ) + .await; TokenCache::global().evict(&addr); let server = result.unwrap(); @@ -1426,7 +1442,13 @@ pub mod test { expected_secret, SystemTime::now() + Duration::from_secs(3600), ); - let result = Server::connect(&addr, ServerOptions::default(), ConnectReason::Other).await; + let result = Server::connect( + &addr, + ServerOptions::default(), + ConnectReason::Other, + Default::default(), + ) + .await; TokenCache::global().evict(&addr); let server = result.unwrap(); @@ -1600,13 +1622,6 @@ pub mod test { assert_eq!(c, msg.code()); } - // RowDescription saved. - let global = server.prepared_statements.parse(&name).unwrap(); - server - .prepared_statements - .row_description(global.name()) - .unwrap(); - server .send( &vec![ diff --git a/pgdog/src/frontend/prepared_statements/global_cache.rs b/pgdog/src/frontend/prepared_statements/global_cache.rs index f61d388a7..50eed73ba 100644 --- a/pgdog/src/frontend/prepared_statements/global_cache.rs +++ b/pgdog/src/frontend/prepared_statements/global_cache.rs @@ -224,11 +224,11 @@ impl GlobalCache { /// Client sent a Describe for a prepared statement and received a RowDescription. /// We record the RowDescription for later use by the results decoder. - pub fn insert_row_description(&mut self, name: &str, row_description: &RowDescription) { + pub fn insert_row_description(&mut self, name: &str, row_description: RowDescription) { if let Some(ref mut entry) = self.names.get_mut(name) && entry.row_description.is_none() { - entry.row_description = Some(row_description.clone()); + entry.row_description = Some(row_description); } } diff --git a/pgdog/src/frontend/router/parser/query/mod.rs b/pgdog/src/frontend/router/parser/query/mod.rs index 6f76db565..a5b340cd4 100644 --- a/pgdog/src/frontend/router/parser/query/mod.rs +++ b/pgdog/src/frontend/router/parser/query/mod.rs @@ -332,7 +332,16 @@ impl QueryParser { return Ok(Command::Deallocate); } - Node::SelectStmt(stmt) => self.select(&statement, stmt, context), + Node::SelectStmt(stmt) => { + if references_pg_type(stmt) { + // Shard 0 is considered the canonical source for now + return Ok(Command::Query(Route::read( + ShardWithPriority::new_override_canonical_schema_info(Shard::Direct(0)), + ))); + } else { + self.select(&statement, stmt, context) + } + } Node::CopyStmt(stmt) => Self::copy(stmt, context), @@ -975,5 +984,25 @@ cfg_select! { _ => {} } +#[cfg(feature = "new_parser")] +fn references_pg_type(stmt: &nodes::SelectStmt) -> bool { + use pg_raw_parse::walk::{self, Recurse}; + use std::ops::ControlFlow; + + walk::walk_manual(stmt.into(), |node| match node { + Node::RangeVar(rv) if rv.relname() == Some("pg_type") => ControlFlow::Break(true), + // atttypid references pg_type.oid + Node::RangeVar(rv) if rv.relname() == Some("pg_attribute") => ControlFlow::Break(true), + Node::TypeCast(tc) + if let Some(tn) = tc.type_name() + && tn.names().iter().filter_map(|s| s.sval()).eq(["regtype"]) => + { + ControlFlow::Break(true) + } + _ => Recurse::yes(), + }) + .unwrap_or_default() +} + #[cfg(test)] mod test; diff --git a/pgdog/src/frontend/router/parser/route.rs b/pgdog/src/frontend/router/parser/route.rs index 180c8547a..f1344ab59 100644 --- a/pgdog/src/frontend/router/parser/route.rs +++ b/pgdog/src/frontend/router/parser/route.rs @@ -420,6 +420,7 @@ pub enum OverrideReason { OnlyOneShard, RewriteUpdate, CrossShardFunction, + CanonicalSchemaInfo, } #[derive(Debug, Clone, PartialEq, Eq, Ord, PartialOrd)] @@ -508,6 +509,14 @@ impl ShardWithPriority { } } + #[cfg(feature = "new_parser")] + pub(crate) fn new_override_canonical_schema_info(shard: Shard) -> Self { + Self { + shard, + source: ShardSource::Override(OverrideReason::CanonicalSchemaInfo), + } + } + pub fn new_default_unset(shard: Shard) -> Self { Self { shard, diff --git a/pgdog/src/lib.rs b/pgdog/src/lib.rs index c1011ac13..dbba68f63 100644 --- a/pgdog/src/lib.rs +++ b/pgdog/src/lib.rs @@ -18,6 +18,7 @@ pub mod plugin; pub mod sighup; pub mod state; pub mod stats; +pub(crate) mod sync; pub mod tasks; #[cfg(test)] pub mod test_utils; diff --git a/pgdog/src/net/messages/mod.rs b/pgdog/src/net/messages/mod.rs index e4d83b878..733fd3e75 100644 --- a/pgdog/src/net/messages/mod.rs +++ b/pgdog/src/net/messages/mod.rs @@ -288,6 +288,10 @@ impl Message { pub fn transaction_error(&self) -> bool { self.code() == 'Z' && self.payload[5] as char == 'E' } + + pub fn replace_payload(&mut self, bytes: Bytes) { + self.payload = bytes; + } } /// Check that the message we received is what we expected. diff --git a/pgdog/src/net/messages/parameter_description.rs b/pgdog/src/net/messages/parameter_description.rs index cedc4ebc2..8229cadd6 100644 --- a/pgdog/src/net/messages/parameter_description.rs +++ b/pgdog/src/net/messages/parameter_description.rs @@ -1,5 +1,6 @@ use super::code; use super::prelude::*; +use std::collections::HashMap; #[derive(Debug, Clone, Default)] pub struct ParameterDescription { @@ -47,6 +48,14 @@ impl ParameterDescription { pub fn from_params(params: Vec) -> Self { Self { params } } + + pub(crate) fn rewrite_data_types(&mut self, mapping: &HashMap) { + for param in &mut self.params { + if let Some(&canonical) = mapping.get(&(*param as u32)) { + *param = canonical as i32; + } + } + } } #[cfg(test)] diff --git a/pgdog/src/net/messages/parse.rs b/pgdog/src/net/messages/parse.rs index e02ea9789..10f3c69cf 100644 --- a/pgdog/src/net/messages/parse.rs +++ b/pgdog/src/net/messages/parse.rs @@ -1,8 +1,10 @@ //! Parse (F) message. use crate::net::c_string_buf_len; +use bytes::BytesMut; +use std::collections::HashMap; use std::fmt::Debug; use std::io::Cursor; -use std::mem::size_of; +use std::mem; use std::str::from_utf8; use std::str::from_utf8_unchecked; @@ -101,13 +103,6 @@ impl Parse { } } - pub fn data_types(&self) -> DataTypesIter<'_> { - DataTypesIter { - data_types: &self.data_types, - offset: 0, - } - } - pub fn data_types_ref(&self) -> Bytes { self.data_types.clone() } @@ -117,34 +112,44 @@ impl Parse { self.query = Bytes::from(query.to_string() + "\0"); self.original = None; } -} -#[derive(Debug)] -pub struct DataTypesIter<'a> { - data_types: &'a Bytes, - offset: usize, -} + /// Rewrite the data types of this message using the given mapping. + /// Returns whether any data types were changed. + pub(crate) fn rewrite_data_types(&mut self, mapping: &HashMap) -> bool { + if mapping.is_empty() { + return false; + } -impl DataTypesIter<'_> { - pub fn len(&self) -> usize { - (self.data_types.len() - size_of::()) / size_of::() + let mut rewritten = false; + // FIXME: This will always copy the bytes, we should only copy on write + let mut bytes = Cursor::new(BytesMut::from(mem::take(&mut self.data_types))); + for _ in 0..bytes.get_u16() { + let canonical_oid = bytes.get_u32(); + if let Some(&shard_oid) = mapping.get(&canonical_oid) { + // Note: We can't use BufMut here, as that writes to the end of the + // buffer, which will allocate. We just want to write to the + // existing bytes. + let pos = bytes.position() as usize; + let oid_bytes = &mut bytes.get_mut()[(pos - 4)..pos]; + oid_bytes.copy_from_slice(&shard_oid.to_be_bytes()); + rewritten = true; + } + } + self.data_types = bytes.into_inner().freeze(); + rewritten } -} -impl Iterator for DataTypesIter<'_> { - type Item = i32; - - fn next(&mut self) -> Option { - let pos = self.offset * size_of::() + size_of::(); - self.offset += 1; - let mut cursor = Cursor::new(self.data_types); - cursor.advance(pos); - - if cursor.remaining() >= size_of::() { - Some(cursor.get_i32()) - } else { - None + #[cfg(test)] + pub(crate) fn with_data_types(&self, data_types: &[u32]) -> Self { + let mut bytes = BytesMut::new(); + bytes.put_u16(data_types.len() as _); + for &oid in data_types { + bytes.put_u32(oid); } + + let mut this = self.clone(); + this.data_types = bytes.freeze(); + this } } @@ -210,8 +215,6 @@ impl Protocol for Parse { #[cfg(test)] mod test { - use bytes::BytesMut; - use super::*; #[test] @@ -223,19 +226,7 @@ mod test { #[test] fn test_parse_from_bytes() { - let mut parse = Parse::named("__pgdog_1", "SELECT * FROM users"); - let mut data_types = BytesMut::new(); - data_types.put_i16(3); - data_types.put_i32(1); - data_types.put_i32(2); - data_types.put_i32(3); - parse.data_types = data_types.freeze(); - - let iter = parse.data_types(); - assert_eq!(iter.len(), 3); - for (i, v) in iter.enumerate() { - assert_eq!(i as i32 + 1, v); - } + let parse = Parse::named("__pgdog_1", "SELECT * FROM users"); assert_eq!(parse.name(), "__pgdog_1"); assert_eq!(parse.query(), "SELECT * FROM users"); @@ -252,8 +243,31 @@ mod test { let parse = Parse::from_bytes(b.freeze()).unwrap(); assert_eq!(parse.name(), "__pgdog_1"); assert_eq!(parse.query(), "SELECT * FROM users"); - assert_eq!(parse.data_types().len(), 0); + assert_eq!(parse.data_types_ref().get_i16(), 0); assert!(Parse::new_anonymous("SELECT 1").anonymous()); } + + #[test] + fn test_parse_rewrite_oids() { + let mut parse = Parse::named("", ""); + let mut data_types = BytesMut::new(); + data_types.put_u16(3); + data_types.put_u32(10_001); + data_types.put_u32(10_011); + data_types.put_u32(10_091); + parse.data_types = data_types.freeze(); + + let mapping = [(10_001, 10_001), (10_011, 10_002), (10_091, 10_003)] + .into_iter() + .collect(); + parse.rewrite_data_types(&mapping); + + let mut expected = BytesMut::new(); + expected.put_u16(3); + expected.put_u32(10_001); + expected.put_u32(10_002); + expected.put_u32(10_003); + assert_eq!(parse.data_types, expected.freeze()); + } } diff --git a/pgdog/src/net/messages/row_description.rs b/pgdog/src/net/messages/row_description.rs index 9e370a4e4..b032115f5 100644 --- a/pgdog/src/net/messages/row_description.rs +++ b/pgdog/src/net/messages/row_description.rs @@ -1,6 +1,6 @@ //! RowDescription (B) message. -use std::collections::BTreeSet; +use std::collections::{BTreeSet, HashMap}; use std::ops::Deref; use std::sync::Arc; @@ -266,6 +266,18 @@ impl RowDescription { true } + + pub(crate) fn rewrite_data_types(&mut self, mapping: &HashMap) { + if mapping.is_empty() { + return; + } + + for field in Arc::make_mut(&mut self.fields) { + if let Some(&canonical) = mapping.get(&(field.type_oid as u32)) { + field.type_oid = canonical as i32; + } + } + } } impl Deref for RowDescription { diff --git a/pgdog/src/sync.rs b/pgdog/src/sync.rs new file mode 100644 index 000000000..ffd9fe563 --- /dev/null +++ b/pgdog/src/sync.rs @@ -0,0 +1,96 @@ +use std::fmt; +use tokio::sync::{Semaphore, SemaphorePermit, SetError, SetOnce, TryAcquireError}; + +/// A combination of `SetOnce` and `OnceCell`. +/// It allows only a single writer to run, with no contention on future reads, +/// while also allowing callers to wait on the value. +pub(crate) struct SetOnceCell { + value: SetOnce, + semaphore: Semaphore, +} + +impl SetOnceCell { + /// Equivalent to [`tokio::sync::OnceCell::get_or_try_init`] + pub(crate) async fn get_or_try_init(&self, f: F) -> Result<&T, E> + where + F: FnOnce() -> Fut, + Fut: Future>, + { + if let Some(val) = self.get() { + Ok(val) + } else { + // Ensure only a single writer attempts to write. + // Err indicates that another task successfully initialized this + if let Ok(permit) = self.semaphore.acquire().await { + self.set_value(f().await?, permit); + } + + Ok(self.get().expect("always initialized")) + } + } + + /// Equivalent to [`SetOnce::wait`] + pub(crate) async fn wait(&self) -> &T { + self.value.wait().await + } + + /// Equivalent to [`SetOnce::get`] + pub(crate) fn get(&self) -> Option<&T> { + self.value.get() + } + + /// Equivalent to [`tokio::sync::OnceCell::set`] + pub(crate) fn set(&self, value: T) -> Result<(), SetError> { + if self.value.initialized() { + return Err(SetError::AlreadyInitializedError(value)); + } + + // Ensure no task is currently attempting to initialize this value + match self.semaphore.try_acquire() { + Ok(permit) => { + self.set_value(value, permit); + Ok(()) + } + Err(TryAcquireError::NoPermits) => Err(SetError::InitializingError(value)), + Err(TryAcquireError::Closed) => Err(SetError::AlreadyInitializedError(value)), + } + } + + /// Set the value and close the semaphore + fn set_value(&self, value: T, permit: SemaphorePermit) { + let Ok(_) = self.value.set(value) else { + panic!("set_value called when already initialized"); + }; + self.semaphore.close(); + permit.forget(); + } +} + +impl Default for SetOnceCell { + fn default() -> Self { + Self { + value: Default::default(), + semaphore: Semaphore::new(1), + } + } +} + +impl fmt::Debug for SetOnceCell +where + SetOnce: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_tuple("SetOnceCell").field(&self.value).finish() + } +} + +impl From for SetOnceCell { + fn from(value: T) -> Self { + let semaphore = Semaphore::new(0); + semaphore.close(); + Self { + value: SetOnce::from(value), + semaphore, + } + } +} From a48c38f699f36797b0a854eca0a5df555d7c63d5 Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Thu, 30 Jul 2026 14:22:10 -0600 Subject: [PATCH 02/21] Fix merge conflict, load all schema info in parallel --- pgdog/src/backend/pool/cluster.rs | 6 ------ pgdog/src/backend/pool/shard/mod.rs | 9 ++------- 2 files changed, 2 insertions(+), 13 deletions(-) diff --git a/pgdog/src/backend/pool/cluster.rs b/pgdog/src/backend/pool/cluster.rs index 85f5e71af..9793b737f 100644 --- a/pgdog/src/backend/pool/cluster.rs +++ b/pgdog/src/backend/pool/cluster.rs @@ -87,10 +87,7 @@ pub struct Cluster { tls_client_certificate_required: bool, #[debug(skip)] schema_loader: Box, -<<<<<<< HEAD -======= canonical_oids: Arc, ->>>>>>> 303395b9 (WIP: Handle OID drift between shards) } /// Sharding configuration from the cluster. @@ -364,10 +361,7 @@ impl Cluster { identity: identity.clone(), tls_client_certificate_required, schema_loader: Box::new(schema_loader::FromServer), -<<<<<<< HEAD -======= canonical_oids, ->>>>>>> 303395b9 (WIP: Handle OID drift between shards) } } diff --git a/pgdog/src/backend/pool/shard/mod.rs b/pgdog/src/backend/pool/shard/mod.rs index af2c291f2..e87b661cf 100644 --- a/pgdog/src/backend/pool/shard/mod.rs +++ b/pgdog/src/backend/pool/shard/mod.rs @@ -1,6 +1,7 @@ //! A shard is a collection of replicas and an optional primary. use arc_swap::ArcSwap; +use futures::try_join; use std::ops::Deref; use std::sync::Arc; use std::time::Duration; @@ -136,11 +137,10 @@ impl Shard { return Ok(false); } - self.load_oids().await?; // This is syncrhonized by database/shard number, so this prevents // a thundering herd with 100s of users, for example, all fetching // the same schema. - let schema = self.schema_cache.get(self).await?; + let (_, schema) = try_join!(self.oids.load(self), self.schema_cache.get(self))?; self.schema.set(schema).expect("schema was not initialized"); Ok(true) @@ -163,11 +163,6 @@ impl Shard { Ok(schema) } - async fn load_oids(&self) -> Result<(), crate::backend::Error> { - self.oids.load(self).await?; - Ok(()) - } - /// Set the schema to its default value. /// We don't need it for this shard. pub(super) fn schema_not_needed(&self) { From 6e1997db20a65e321ea90906b5e372d423ee1700 Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Thu, 30 Jul 2026 14:27:47 -0600 Subject: [PATCH 03/21] Don't copy `RowDescription` bytes if we don't need to --- pgdog/src/backend/prepared_statements.rs | 9 +++++---- pgdog/src/net/messages/row_description.rs | 10 ++++++++-- 2 files changed, 13 insertions(+), 6 deletions(-) diff --git a/pgdog/src/backend/prepared_statements.rs b/pgdog/src/backend/prepared_statements.rs index e362dbead..e10401455 100644 --- a/pgdog/src/backend/prepared_statements.rs +++ b/pgdog/src/backend/prepared_statements.rs @@ -404,8 +404,9 @@ impl PreparedStatements { /// statement in the global cache. fn add_row_description(&self, name: &str, message: &mut Message) -> Result<(), Error> { let mut row_description = RowDescription::from_bytes(message.payload())?; - self.rewrite_row_description_data_types(&mut row_description); - message.replace_payload(row_description.to_bytes()); + if self.rewrite_row_description_data_types(&mut row_description) { + message.replace_payload(row_description.to_bytes()); + } self.global_cache .write() .insert_row_description(name, row_description); @@ -480,8 +481,8 @@ impl PreparedStatements { parse.rewrite_data_types(&self.oid_mappings().canonical_to_shard) } - fn rewrite_row_description_data_types(&self, row_description: &mut RowDescription) { - row_description.rewrite_data_types(&self.oid_mappings().shard_to_canonical); + fn rewrite_row_description_data_types(&self, row_description: &mut RowDescription) -> bool { + row_description.rewrite_data_types(&self.oid_mappings().shard_to_canonical) } fn rewrite_parameter_description_data_types(&self, message: &mut Message) -> Result<(), Error> { diff --git a/pgdog/src/net/messages/row_description.rs b/pgdog/src/net/messages/row_description.rs index b032115f5..6a7c1ab94 100644 --- a/pgdog/src/net/messages/row_description.rs +++ b/pgdog/src/net/messages/row_description.rs @@ -267,16 +267,22 @@ impl RowDescription { true } - pub(crate) fn rewrite_data_types(&mut self, mapping: &HashMap) { + /// Replaces the data types of each field using the given mapping. + /// Returns whether any changes actually occurred. + pub(crate) fn rewrite_data_types(&mut self, mapping: &HashMap) -> bool { if mapping.is_empty() { - return; + return false; } + let mut changed = false; for field in Arc::make_mut(&mut self.fields) { if let Some(&canonical) = mapping.get(&(field.type_oid as u32)) { + changed = true; field.type_oid = canonical as i32; } } + + changed } } From 45fc50412b272ef4f72e852502d3f140ac0ba129 Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Thu, 30 Jul 2026 14:47:05 -0600 Subject: [PATCH 04/21] Backfill more unit tests --- pgdog/src/frontend/router/parser/query/mod.rs | 11 ++++++++- .../router/parser/query/test/test_select.rs | 23 ++++++++++++++++++- pgdog/src/net/messages/parse.rs | 23 +++++++------------ 3 files changed, 40 insertions(+), 17 deletions(-) diff --git a/pgdog/src/frontend/router/parser/query/mod.rs b/pgdog/src/frontend/router/parser/query/mod.rs index a5b340cd4..1b8d7708b 100644 --- a/pgdog/src/frontend/router/parser/query/mod.rs +++ b/pgdog/src/frontend/router/parser/query/mod.rs @@ -995,7 +995,16 @@ fn references_pg_type(stmt: &nodes::SelectStmt) -> bool { Node::RangeVar(rv) if rv.relname() == Some("pg_attribute") => ControlFlow::Break(true), Node::TypeCast(tc) if let Some(tn) = tc.type_name() - && tn.names().iter().filter_map(|s| s.sval()).eq(["regtype"]) => + && tn.names().iter().map(|s| s.sval()).eq([Some("regtype")]) => + { + ControlFlow::Break(true) + } + Node::FuncCall(fc) + if fc + .funcname() + .iter() + .map(Node::as_str) + .eq([Some("to_regclass")]) => { ControlFlow::Break(true) } diff --git a/pgdog/src/frontend/router/parser/query/test/test_select.rs b/pgdog/src/frontend/router/parser/query/test/test_select.rs index 9eea2eb8e..7301d7114 100644 --- a/pgdog/src/frontend/router/parser/query/test/test_select.rs +++ b/pgdog/src/frontend/router/parser/query/test/test_select.rs @@ -259,7 +259,7 @@ fn test_system_catalog_sharded() { ); let command = test.execute(vec![ - Query::new("SELECT * FROM pg_type WHERE typname = 'int4'").into(), + Query::new("SELECT * FROM pg_class WHERE relname = $1").into(), ]); assert_eq!( command.route().shard(), @@ -268,6 +268,27 @@ fn test_system_catalog_sharded() { ); assert!(!command.route().is_omnisharded()); + let command = test.execute(vec![Query::new("SELECT * FROM pg_type").into()]); + assert_eq!( + command.route().shard(), + &Shard::Direct(0), + "pg_type queries should go to shard 0", + ); + + let command = test.execute(vec![Query::new("SELECT $1::regclass").into()]); + assert_eq!( + command.route().shard(), + &Shard::Direct(0), + "regclass casts should go to shard 0", + ); + + let command = test.execute(vec![Query::new("SELECT to_regclass($1)").into()]); + assert_eq!( + command.route().shard(), + &Shard::Direct(0), + "to_regclass should go to shard 0", + ); + // Reset to default let mut updated = config().deref().clone(); updated.config.general.system_catalogs = SystemCatalogsBehavior::default(); diff --git a/pgdog/src/net/messages/parse.rs b/pgdog/src/net/messages/parse.rs index 10f3c69cf..c38becc73 100644 --- a/pgdog/src/net/messages/parse.rs +++ b/pgdog/src/net/messages/parse.rs @@ -250,24 +250,17 @@ mod test { #[test] fn test_parse_rewrite_oids() { - let mut parse = Parse::named("", ""); - let mut data_types = BytesMut::new(); - data_types.put_u16(3); - data_types.put_u32(10_001); - data_types.put_u32(10_011); - data_types.put_u32(10_091); - parse.data_types = data_types.freeze(); + let mut parse = Parse::named("", "").with_data_types(&[10_001, 10_011, 10_091]); let mapping = [(10_001, 10_001), (10_011, 10_002), (10_091, 10_003)] .into_iter() .collect(); - parse.rewrite_data_types(&mapping); - - let mut expected = BytesMut::new(); - expected.put_u16(3); - expected.put_u32(10_001); - expected.put_u32(10_002); - expected.put_u32(10_003); - assert_eq!(parse.data_types, expected.freeze()); + assert!(parse.rewrite_data_types(&mapping)); + + let expected = Parse::named("", "").with_data_types(&[10_001, 10_002, 10_003]); + assert_eq!(parse, expected); + + let mut parse_without_custom_types = Parse::named("", "").with_data_types(&[1]); + assert!(!parse_without_custom_types.rewrite_data_types(&mapping)); } } From 2c74be06e89e48a53a046ffc506e654b98bd951c Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Thu, 30 Jul 2026 14:56:07 -0600 Subject: [PATCH 05/21] Canonical routing only happens with the new parser --- .../router/parser/query/test/test_select.rs | 43 ++++++++++--------- 1 file changed, 23 insertions(+), 20 deletions(-) diff --git a/pgdog/src/frontend/router/parser/query/test/test_select.rs b/pgdog/src/frontend/router/parser/query/test/test_select.rs index 7301d7114..aa81527b2 100644 --- a/pgdog/src/frontend/router/parser/query/test/test_select.rs +++ b/pgdog/src/frontend/router/parser/query/test/test_select.rs @@ -268,26 +268,29 @@ fn test_system_catalog_sharded() { ); assert!(!command.route().is_omnisharded()); - let command = test.execute(vec![Query::new("SELECT * FROM pg_type").into()]); - assert_eq!( - command.route().shard(), - &Shard::Direct(0), - "pg_type queries should go to shard 0", - ); - - let command = test.execute(vec![Query::new("SELECT $1::regclass").into()]); - assert_eq!( - command.route().shard(), - &Shard::Direct(0), - "regclass casts should go to shard 0", - ); - - let command = test.execute(vec![Query::new("SELECT to_regclass($1)").into()]); - assert_eq!( - command.route().shard(), - &Shard::Direct(0), - "to_regclass should go to shard 0", - ); + #[cfg(feature = "new_parser")] + { + let command = test.execute(vec![Query::new("SELECT * FROM pg_type").into()]); + assert_eq!( + command.route().shard(), + &Shard::Direct(0), + "pg_type queries should go to shard 0", + ); + + let command = test.execute(vec![Query::new("SELECT $1::regclass").into()]); + assert_eq!( + command.route().shard(), + &Shard::Direct(0), + "regclass casts should go to shard 0", + ); + + let command = test.execute(vec![Query::new("SELECT to_regclass($1)").into()]); + assert_eq!( + command.route().shard(), + &Shard::Direct(0), + "to_regclass should go to shard 0", + ); + } // Reset to default let mut updated = config().deref().clone(); From 1f618dc0310756d18125d94849f049d9146a41ef Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Fri, 31 Jul 2026 13:30:50 -0600 Subject: [PATCH 06/21] Reduce the amount of test churn --- .../backend/pool/connection/binding_test.rs | 44 +-- pgdog/src/backend/pool/guard.rs | 11 +- pgdog/src/backend/pool/lb/mod.rs | 2 +- pgdog/src/backend/pool/lb/test.rs | 35 +- pgdog/src/backend/pool/lsn_monitor.rs | 11 +- pgdog/src/backend/pool/monitor.rs | 46 ++- pgdog/src/backend/pool/pool_impl.rs | 9 +- pgdog/src/backend/pool/shard/mod.rs | 2 +- pgdog/src/backend/pool/shard/monitor.rs | 26 +- pgdog/src/backend/pool/test/mod.rs | 327 +++++++----------- pgdog/src/backend/pool/waiting.rs | 23 +- 11 files changed, 221 insertions(+), 315 deletions(-) diff --git a/pgdog/src/backend/pool/connection/binding_test.rs b/pgdog/src/backend/pool/connection/binding_test.rs index ec8ac5798..b9f541a78 100644 --- a/pgdog/src/backend/pool/connection/binding_test.rs +++ b/pgdog/src/backend/pool/connection/binding_test.rs @@ -26,29 +26,20 @@ mod tests { let server3 = Box::new(test_server().await); // Create pools for each server using their addresses - let pool1 = Pool::new( - &PoolConfig { - address: server1.addr().clone(), - config: crate::backend::pool::Config::default(), - }, - Default::default(), - ); + let pool1 = Pool::new(&PoolConfig { + address: server1.addr().clone(), + config: crate::backend::pool::Config::default(), + }); - let pool2 = Pool::new( - &PoolConfig { - address: server2.addr().clone(), - config: crate::backend::pool::Config::default(), - }, - Default::default(), - ); + let pool2 = Pool::new(&PoolConfig { + address: server2.addr().clone(), + config: crate::backend::pool::Config::default(), + }); - let pool3 = Pool::new( - &PoolConfig { - address: server3.addr().clone(), - config: crate::backend::pool::Config::default(), - }, - Default::default(), - ); + let pool3 = Pool::new(&PoolConfig { + address: server3.addr().clone(), + config: crate::backend::pool::Config::default(), + }); let now = Instant::now(); let guards = vec![ @@ -77,13 +68,10 @@ mod tests { // Create a Direct binding instead of MultiShard let server = Box::new(test_server().await); - let pool = Pool::new( - &PoolConfig { - address: server.addr().clone(), - config: crate::backend::pool::Config::default(), - }, - Default::default(), - ); + let pool = Pool::new(&PoolConfig { + address: server.addr().clone(), + config: crate::backend::pool::Config::default(), + }); let guard = crate::backend::pool::Guard::new(pool, server, Instant::now()); let mut binding = Binding::Direct(guard, 0); diff --git a/pgdog/src/backend/pool/guard.rs b/pgdog/src/backend/pool/guard.rs index 7b0f99919..03d27cf76 100644 --- a/pgdog/src/backend/pool/guard.rs +++ b/pgdog/src/backend/pool/guard.rs @@ -348,13 +348,10 @@ mod test { }, }; - let pool = Pool::new( - &PoolConfig { - address: Address::new_test(), - config, - }, - Default::default(), - ); + let pool = Pool::new(&PoolConfig { + address: Address::new_test(), + config, + }); pool.launch(); { diff --git a/pgdog/src/backend/pool/lb/mod.rs b/pgdog/src/backend/pool/lb/mod.rs index a192be6f8..7c37068f2 100644 --- a/pgdog/src/backend/pool/lb/mod.rs +++ b/pgdog/src/backend/pool/lb/mod.rs @@ -114,7 +114,7 @@ impl LoadBalancer { .iter() .map(|config| { Target::new( - Pool::new(config, Arc::clone(&oids)), + Pool::with_oid_mapping(config, Arc::clone(&oids)), config.address.configured_role, ) }) diff --git a/pgdog/src/backend/pool/lb/test.rs b/pgdog/src/backend/pool/lb/test.rs index 0a8148aca..35bbf99dd 100644 --- a/pgdog/src/backend/pool/lb/test.rs +++ b/pgdog/src/backend/pool/lb/test.rs @@ -63,7 +63,7 @@ fn set_lsn_stats(target: &Target, replica: bool, lsn: i64) { async fn test_include_primary_if_replica_banned_only_primary() { let mut primary = create_test_pool_config("127.0.0.1", 5432); primary.address.configured_role = Role::Primary; - let pool = Pool::new(&primary, Default::default()); + let pool = Pool::new(&primary); let lb = LoadBalancer::new( &Some(pool), @@ -204,7 +204,7 @@ async fn test_pools_with_roles_and_bans() { #[tokio::test] async fn test_primary_pool_banning() { let primary_config = create_test_pool_config("127.0.0.1", 5432); - let primary_pool = Pool::new(&primary_config, Default::default()); + let primary_pool = Pool::new(&primary_config); primary_pool.launch(); let replica_configs = [create_test_pool_config("localhost", 5432)]; @@ -358,7 +358,7 @@ async fn test_monitor_automatic_ban_expiration() { #[tokio::test] async fn test_read_write_split_exclude_primary() { let primary_config = create_test_pool_config("127.0.0.1", 5432); - let primary_pool = Pool::new(&primary_config, Default::default()); + let primary_pool = Pool::new(&primary_config); primary_pool.launch(); let replica_configs = [ @@ -398,7 +398,7 @@ async fn test_read_write_split_exclude_primary() { #[tokio::test] async fn test_read_write_split_include_primary() { let primary_config = create_test_pool_config("127.0.0.1", 5432); - let primary_pool = Pool::new(&primary_config, Default::default()); + let primary_pool = Pool::new(&primary_config); primary_pool.launch(); let replica_configs = [create_test_pool_config("localhost", 5432)]; @@ -448,10 +448,7 @@ async fn test_read_write_split_include_primary() { #[tokio::test] async fn test_prefer_primary_optin_read_honors_read_write_split() { async fn used_ids(split: ReadWriteSplit) -> (HashSet, u64) { - let primary_pool = Pool::new( - &create_test_pool_config("127.0.0.1", 5432), - Default::default(), - ); + let primary_pool = Pool::new(&create_test_pool_config("127.0.0.1", 5432)); primary_pool.launch(); let primary_id = primary_pool.id(); @@ -503,7 +500,7 @@ async fn test_prefer_primary_optin_read_honors_read_write_split() { #[tokio::test] async fn test_read_write_split_exclude_primary_no_replicas() { let primary_config = create_test_pool_config("127.0.0.1", 5432); - let primary_pool = Pool::new(&primary_config, Default::default()); + let primary_pool = Pool::new(&primary_config); primary_pool.launch(); let replica_configs = []; @@ -603,7 +600,7 @@ async fn test_read_write_split_include_primary_no_primary() { #[tokio::test] async fn test_read_write_split_with_banned_primary() { let primary_config = create_test_pool_config("127.0.0.1", 5432); - let primary_pool = Pool::new(&primary_config, Default::default()); + let primary_pool = Pool::new(&primary_config); primary_pool.launch(); let replica_configs = [create_test_pool_config("localhost", 5432)]; @@ -644,7 +641,7 @@ async fn test_read_write_split_with_banned_primary() { #[tokio::test] async fn test_read_write_split_with_banned_replicas() { let primary_config = create_test_pool_config("127.0.0.1", 5432); - let primary_pool = Pool::new(&primary_config, Default::default()); + let primary_pool = Pool::new(&primary_config); primary_pool.launch(); let replica_configs = [create_test_pool_config("localhost", 5432)]; @@ -685,7 +682,7 @@ async fn test_read_write_split_with_banned_replicas() { #[tokio::test] async fn test_prefer_primary_with_banned_replicas_falls_back_to_primary() { let primary_config = create_test_pool_config("127.0.0.1", 5432); - let primary_pool = Pool::new(&primary_config, Default::default()); + let primary_pool = Pool::new(&primary_config); primary_pool.launch(); let replica_configs = [create_test_pool_config("localhost", 5432)]; @@ -721,7 +718,7 @@ async fn test_prefer_primary_with_banned_replicas_falls_back_to_primary() { #[tokio::test] async fn test_read_write_split_exclude_primary_with_round_robin() { let primary_config = create_test_pool_config("127.0.0.1", 5432); - let primary_pool = Pool::new(&primary_config, Default::default()); + let primary_pool = Pool::new(&primary_config); primary_pool.launch(); let replica_configs = [ @@ -963,7 +960,7 @@ async fn test_monitor_health_state_race() { #[tokio::test] async fn test_include_primary_if_replica_banned_no_bans() { let primary_config = create_test_pool_config("127.0.0.1", 5432); - let primary_pool = Pool::new(&primary_config, Default::default()); + let primary_pool = Pool::new(&primary_config); primary_pool.launch(); let replica_configs = [create_test_pool_config("localhost", 5432)]; @@ -1000,7 +997,7 @@ async fn test_include_primary_if_replica_banned_no_bans() { #[tokio::test] async fn test_include_primary_if_replica_banned_with_ban() { let primary_config = create_test_pool_config("127.0.0.1", 5432); - let primary_pool = Pool::new(&primary_config, Default::default()); + let primary_pool = Pool::new(&primary_config); primary_pool.launch(); let replica_configs = [create_test_pool_config("localhost", 5432)]; @@ -1050,7 +1047,7 @@ async fn test_has_replicas_with_replicas() { #[tokio::test] async fn test_has_replicas_with_primary_and_replicas() { let primary_config = create_test_pool_config("127.0.0.1", 5432); - let primary_pool = Pool::new(&primary_config, Default::default()); + let primary_pool = Pool::new(&primary_config); primary_pool.launch(); let replica_configs = [create_test_pool_config("localhost", 5432)]; @@ -1072,7 +1069,7 @@ async fn test_has_replicas_with_primary_and_replicas() { #[tokio::test] async fn test_has_replicas_primary_only() { let primary_config = create_test_pool_config("127.0.0.1", 5432); - let primary_pool = Pool::new(&primary_config, Default::default()); + let primary_pool = Pool::new(&primary_config); primary_pool.launch(); let lb = LoadBalancer::new( @@ -1269,7 +1266,7 @@ async fn test_redetect_roles_marks_added_auto_target_replica_when_primary_unchan let mut existing_replica_config = create_test_pool_config("localhost", 5432); existing_replica_config.address.configured_role = Role::Auto; - let old_primary = Pool::new(&primary_config, Default::default()); + let old_primary = Pool::new(&primary_config); let lb_old = LoadBalancer::new( &Some(old_primary), std::slice::from_ref(&existing_replica_config), @@ -1286,7 +1283,7 @@ async fn test_redetect_roles_marks_added_auto_target_replica_when_primary_unchan let mut added_replica_config = create_test_pool_config("localhost", 5433); added_replica_config.address.configured_role = Role::Auto; - let new_primary = Pool::new(&primary_config, Default::default()); + let new_primary = Pool::new(&primary_config); let lb_new = LoadBalancer::new( &Some(new_primary), &[existing_replica_config, added_replica_config], diff --git a/pgdog/src/backend/pool/lsn_monitor.rs b/pgdog/src/backend/pool/lsn_monitor.rs index 0f5ed0cb5..891d14607 100644 --- a/pgdog/src/backend/pool/lsn_monitor.rs +++ b/pgdog/src/backend/pool/lsn_monitor.rs @@ -449,13 +449,10 @@ mod test { }, }; - let pool = Pool::new( - &PoolConfig { - address: Address::new_test(), - config, - }, - Default::default(), - ); + let pool = Pool::new(&PoolConfig { + address: Address::new_test(), + config, + }); pool.launch(); // Saturate the pool by holding its only connection. diff --git a/pgdog/src/backend/pool/monitor.rs b/pgdog/src/backend/pool/monitor.rs index a1086bfdd..38e318726 100644 --- a/pgdog/src/backend/pool/monitor.rs +++ b/pgdog/src/backend/pool/monitor.rs @@ -564,20 +564,17 @@ mod test { }, }; - let pool = Pool::new( - &PoolConfig { - address: Address { - host: "127.0.0.1".into(), - port: 1, - database_name: "pgdog".into(), - user: "pgdog".into(), - passwords: vec!["pgdog".into()], - ..Default::default() - }, - config, + let pool = Pool::new(&PoolConfig { + address: Address { + host: "127.0.0.1".into(), + port: 1, + database_name: "pgdog".into(), + user: "pgdog".into(), + passwords: vec!["pgdog".into()], + ..Default::default() }, - Default::default(), - ); + config, + }); pool.launch(); pool.inner().health.toggle(true); @@ -601,20 +598,17 @@ mod test { }, }; - let pool = Pool::new( - &PoolConfig { - address: Address { - host: "127.0.0.1".into(), - port: 5432, - database_name: "pgdog".into(), - user: "pgdog".into(), - passwords: vec!["pgdog".into()], - ..Default::default() - }, - config, + let pool = Pool::new(&PoolConfig { + address: Address { + host: "127.0.0.1".into(), + port: 5432, + database_name: "pgdog".into(), + user: "pgdog".into(), + passwords: vec!["pgdog".into()], + ..Default::default() }, - Default::default(), - ); + config, + }); pool.launch(); let initial_total = pool.lock().total(); diff --git a/pgdog/src/backend/pool/pool_impl.rs b/pgdog/src/backend/pool/pool_impl.rs index 12bcc75f7..62fd13310 100644 --- a/pgdog/src/backend/pool/pool_impl.rs +++ b/pgdog/src/backend/pool/pool_impl.rs @@ -60,8 +60,13 @@ impl std::fmt::Debug for Pool { } impl Pool { + #[cfg(test)] + pub(crate) fn new(config: &PoolConfig) -> Self { + Self::with_oid_mapping(config, Default::default()) + } + /// Create new connection pool. - pub(crate) fn new(config: &PoolConfig, oids: Arc) -> Self { + pub(crate) fn with_oid_mapping(config: &PoolConfig, oids: Arc) -> Self { let id = next_pool_id(); Self { inner: Arc::new(InnerSync { @@ -87,7 +92,7 @@ impl Pool { config: Config::default(), }; - Self::new(&config, Default::default()) + Self::new(&config) } pub(crate) fn inner(&self) -> &InnerSync { diff --git a/pgdog/src/backend/pool/shard/mod.rs b/pgdog/src/backend/pool/shard/mod.rs index e87b661cf..b78722f73 100644 --- a/pgdog/src/backend/pool/shard/mod.rs +++ b/pgdog/src/backend/pool/shard/mod.rs @@ -354,7 +354,7 @@ impl ShardInner { schema_cache, } = shard; let oids = schema_cache.oids(&identifier.database, number); - let primary = primary.map(|config| Pool::new(config, Arc::clone(&oids))); + let primary = primary.map(|config| Pool::with_oid_mapping(config, Arc::clone(&oids))); let lb = LoadBalancer::new(&primary, replicas, lb_strategy, rw_split, Arc::clone(&oids)); let comms = Arc::new(ShardComms { shutdown: CancellationToken::new(), diff --git a/pgdog/src/backend/pool/shard/monitor.rs b/pgdog/src/backend/pool/shard/monitor.rs index 86f769bde..af3df05e0 100644 --- a/pgdog/src/backend/pool/shard/monitor.rs +++ b/pgdog/src/backend/pool/shard/monitor.rs @@ -249,23 +249,17 @@ mod test { #[test] fn test_update_replica_lag_assigns_primary_minus_replica_to_replica_pool() { - let primary = Pool::new( - &PoolConfig { - address: Address::new_test(), - config: Config::default(), - }, - Default::default(), - ); - let replica = Pool::new( - &PoolConfig { - address: Address { - configured_role: Role::Replica, - ..Address::new_test() - }, - config: Config::default(), + let primary = Pool::new(&PoolConfig { + address: Address::new_test(), + config: Config::default(), + }); + let replica = Pool::new(&PoolConfig { + address: Address { + configured_role: Role::Replica, + ..Address::new_test() }, - Default::default(), - ); + config: Config::default(), + }); set_pool_lsn_stats(&primary, false, 200, "2026-07-01 13:33:10.000000+00"); set_pool_lsn_stats(&replica, true, 100, "2026-07-01 13:33:00.000000+00"); diff --git a/pgdog/src/backend/pool/test/mod.rs b/pgdog/src/backend/pool/test/mod.rs index aacb9cff8..35af9706d 100644 --- a/pgdog/src/backend/pool/test/mod.rs +++ b/pgdog/src/backend/pool/test/mod.rs @@ -28,20 +28,17 @@ pub fn pool() -> Pool { }, }; - let pool = Pool::new( - &PoolConfig { - address: Address { - host: "127.0.0.1".into(), - port: 5432, - database_name: "pgdog".into(), - user: "pgdog".into(), - passwords: vec!["pgdog".into()], - ..Default::default() - }, - config, + let pool = Pool::new(&PoolConfig { + address: Address { + host: "127.0.0.1".into(), + port: 5432, + database_name: "pgdog".into(), + user: "pgdog".into(), + passwords: vec!["pgdog".into()], + ..Default::default() }, - Default::default(), - ); + config, + }); pool.launch(); pool } @@ -56,20 +53,17 @@ pub fn pool_with_prepared_capacity(capacity: usize) -> Pool { }, }; - let pool = Pool::new( - &PoolConfig { - address: Address { - host: "127.0.0.1".into(), - port: 5432, - database_name: "pgdog".into(), - user: "pgdog".into(), - passwords: vec!["pgdog".into()], - ..Default::default() - }, - config, + let pool = Pool::new(&PoolConfig { + address: Address { + host: "127.0.0.1".into(), + port: 5432, + database_name: "pgdog".into(), + user: "pgdog".into(), + passwords: vec!["pgdog".into()], + ..Default::default() }, - Default::default(), - ); + config, + }); pool.launch(); pool } @@ -332,20 +326,17 @@ async fn test_server_force_close_discards_connection() { }, }; - let pool = Pool::new( - &PoolConfig { - address: Address { - host: "127.0.0.1".into(), - port: 5432, - database_name: "pgdog".into(), - user: "pgdog".into(), - passwords: vec!["pgdog".into()], - ..Default::default() - }, - config, + let pool = Pool::new(&PoolConfig { + address: Address { + host: "127.0.0.1".into(), + port: 5432, + database_name: "pgdog".into(), + user: "pgdog".into(), + passwords: vec!["pgdog".into()], + ..Default::default() }, - Default::default(), - ); + config, + }); pool.launch(); let mut conn = pool.get(&Request::default()).await.unwrap(); @@ -506,20 +497,17 @@ async fn test_idle_healthcheck_loop() { }, }; - let pool = Pool::new( - &PoolConfig { - address: Address { - host: "127.0.0.1".into(), - port: 5432, - database_name: "pgdog".into(), - user: "pgdog".into(), - passwords: vec!["pgdog".into()], - ..Default::default() - }, - config, + let pool = Pool::new(&PoolConfig { + address: Address { + host: "127.0.0.1".into(), + port: 5432, + database_name: "pgdog".into(), + user: "pgdog".into(), + passwords: vec!["pgdog".into()], + ..Default::default() }, - Default::default(), - ); + config, + }); pool.launch(); let initial_healthchecks = pool.state().stats.counts.healthchecks; @@ -557,20 +545,17 @@ async fn test_idle_healthcheck_loop_disabled_with_zero_interval() { }, }; - let pool = Pool::new( - &PoolConfig { - address: Address { - host: "127.0.0.1".into(), - port: 1, - database_name: "pgdog".into(), - user: "pgdog".into(), - passwords: vec!["pgdog".into()], - ..Default::default() - }, - config, + let pool = Pool::new(&PoolConfig { + address: Address { + host: "127.0.0.1".into(), + port: 1, + database_name: "pgdog".into(), + user: "pgdog".into(), + passwords: vec!["pgdog".into()], + ..Default::default() }, - Default::default(), - ); + config, + }); pool.launch(); let initial_healthchecks = pool.state().stats.counts.healthchecks; @@ -596,13 +581,10 @@ async fn test_checkout_timeout() { }, }; - let pool = Pool::new( - &PoolConfig { - address: Address::new_test(), - config, - }, - Default::default(), - ); + let pool = Pool::new(&PoolConfig { + address: Address::new_test(), + config, + }); pool.launch(); // Hold the only connection @@ -627,36 +609,30 @@ async fn test_move_conns_to() { }, }; - let source = Pool::new( - &PoolConfig { - address: Address { - host: "127.0.0.1".into(), - port: 5432, - database_name: "pgdog".into(), - user: "pgdog".into(), - passwords: vec!["pgdog".into()], - ..Default::default() - }, - config, + let source = Pool::new(&PoolConfig { + address: Address { + host: "127.0.0.1".into(), + port: 5432, + database_name: "pgdog".into(), + user: "pgdog".into(), + passwords: vec!["pgdog".into()], + ..Default::default() }, - Default::default(), - ); + config, + }); source.launch(); - let destination = Pool::new( - &PoolConfig { - address: Address { - host: "127.0.0.1".into(), - port: 5432, - database_name: "pgdog".into(), - user: "pgdog".into(), - passwords: vec!["pgdog".into()], - ..Default::default() - }, - config, + let destination = Pool::new(&PoolConfig { + address: Address { + host: "127.0.0.1".into(), + port: 5432, + database_name: "pgdog".into(), + user: "pgdog".into(), + passwords: vec!["pgdog".into()], + ..Default::default() }, - Default::default(), - ); + config, + }); let conn1 = source.get(&Request::default()).await.unwrap(); let conn2 = source.get(&Request::default()).await.unwrap(); @@ -702,22 +678,16 @@ async fn test_move_conns_all_idle() { }, }; - let source = Pool::new( - &PoolConfig { - address: Address::new_test(), - config, - }, - Default::default(), - ); + let source = Pool::new(&PoolConfig { + address: Address::new_test(), + config, + }); source.launch(); - let destination = Pool::new( - &PoolConfig { - address: Address::new_test(), - config, - }, - Default::default(), - ); + let destination = Pool::new(&PoolConfig { + address: Address::new_test(), + config, + }); // Check out and return 3 connections so they become idle. let c1 = source.get(&Request::default()).await.unwrap(); @@ -755,22 +725,16 @@ async fn test_move_conns_all_checked_out() { }, }; - let source = Pool::new( - &PoolConfig { - address: Address::new_test(), - config, - }, - Default::default(), - ); + let source = Pool::new(&PoolConfig { + address: Address::new_test(), + config, + }); source.launch(); - let destination = Pool::new( - &PoolConfig { - address: Address::new_test(), - config, - }, - Default::default(), - ); + let destination = Pool::new(&PoolConfig { + address: Address::new_test(), + config, + }); let c1 = source.get(&Request::default()).await.unwrap(); let c2 = source.get(&Request::default()).await.unwrap(); @@ -816,22 +780,16 @@ async fn test_move_conns_destination_serves_after_launch() { }, }; - let source = Pool::new( - &PoolConfig { - address: Address::new_test(), - config, - }, - Default::default(), - ); + let source = Pool::new(&PoolConfig { + address: Address::new_test(), + config, + }); source.launch(); - let destination = Pool::new( - &PoolConfig { - address: Address::new_test(), - config, - }, - Default::default(), - ); + let destination = Pool::new(&PoolConfig { + address: Address::new_test(), + config, + }); // Create one idle connection. let c1 = source.get(&Request::default()).await.unwrap(); @@ -864,20 +822,17 @@ fn auth_pool(passwords: Vec) -> Pool { }, }; - Pool::new( - &PoolConfig { - address: Address { - host: "127.0.0.1".into(), - port: 5432, - database_name: "pgdog".into(), - user: "pgdog".into(), - passwords, - ..Default::default() - }, - config, + Pool::new(&PoolConfig { + address: Address { + host: "127.0.0.1".into(), + port: 5432, + database_name: "pgdog".into(), + user: "pgdog".into(), + passwords, + ..Default::default() }, - Default::default(), - ) + config, + }) } #[tokio::test] @@ -1006,13 +961,10 @@ async fn test_lsn_monitor() { }, }; - let pool = Pool::new( - &PoolConfig { - address: Address::new_test(), - config, - }, - Default::default(), - ); + let pool = Pool::new(&PoolConfig { + address: Address::new_test(), + config, + }); let initial_stats = pool.lsn_stats(); assert!(!initial_stats.valid()); @@ -1068,13 +1020,10 @@ async fn test_token_refresh_loop_primes_cache_on_cold_start() { let expiry = SystemTime::now() + Duration::from_millis(200); TokenCache::global().set(&addr, "initial-token".into(), expiry); - let pool = Pool::new( - &PoolConfig { - address: addr.clone(), - config, - }, - Default::default(), - ); + let pool = Pool::new(&PoolConfig { + address: addr.clone(), + config, + }); pool.launch(); // Cache must be populated immediately. @@ -1110,13 +1059,10 @@ async fn test_token_refresh_loop_refreshes_before_expiry() { let expiry = SystemTime::now() + Duration::from_secs(3600); TokenCache::global().set(&addr, "long-lived-token".into(), expiry); - let pool = Pool::new( - &PoolConfig { - address: addr.clone(), - config, - }, - Default::default(), - ); + let pool = Pool::new(&PoolConfig { + address: addr.clone(), + config, + }); pool.launch(); sleep(Duration::from_millis(100)).await; @@ -1158,13 +1104,10 @@ async fn test_token_refresh_loop_evicts_on_failed_refresh() { let expiry = SystemTime::now() + Duration::from_secs(10); TokenCache::global().set(&addr, "stale-token".into(), expiry); - let pool = Pool::new( - &PoolConfig { - address: addr.clone(), - config, - }, - Default::default(), - ); + let pool = Pool::new(&PoolConfig { + address: addr.clone(), + config, + }); pool.launch(); // Give the refresh loop time to fire and fail. @@ -1203,13 +1146,10 @@ async fn test_token_refresh_loop_not_spawned_for_password_auth() { // Poison the cache to detect any unexpected writes. TokenCache::global().evict(&addr); - let pool = Pool::new( - &PoolConfig { - address: addr.clone(), - config, - }, - Default::default(), - ); + let pool = Pool::new(&PoolConfig { + address: addr.clone(), + config, + }); pool.launch(); sleep(Duration::from_millis(100)).await; @@ -1246,13 +1186,10 @@ async fn test_token_refresh_loop_stops_on_shutdown() { let expiry = SystemTime::now() + Duration::from_secs(3600); TokenCache::global().set(&addr, "token".into(), expiry); - let pool = Pool::new( - &PoolConfig { - address: addr.clone(), - config, - }, - Default::default(), - ); + let pool = Pool::new(&PoolConfig { + address: addr.clone(), + config, + }); pool.launch(); sleep(Duration::from_millis(50)).await; diff --git a/pgdog/src/backend/pool/waiting.rs b/pgdog/src/backend/pool/waiting.rs index 0a586e574..8b02aad2e 100644 --- a/pgdog/src/backend/pool/waiting.rs +++ b/pgdog/src/backend/pool/waiting.rs @@ -150,20 +150,17 @@ mod tests { }, }; - let pool = Pool::new( - &crate::backend::pool::PoolConfig { - address: crate::backend::pool::Address { - host: "127.0.0.1".into(), - port: 5432, - database_name: "pgdog".into(), - user: "pgdog".into(), - passwords: vec!["pgdog".into()], - ..Default::default() - }, - config, + let pool = Pool::new(&crate::backend::pool::PoolConfig { + address: crate::backend::pool::Address { + host: "127.0.0.1".into(), + port: 5432, + database_name: "pgdog".into(), + user: "pgdog".into(), + passwords: vec!["pgdog".into()], + ..Default::default() }, - Default::default(), - ); + config, + }); pool.launch(); sleep(Duration::from_millis(100)).await; From 575d88008c5c124a827e9418f78299a0a9ccdb7d Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Mon, 3 Aug 2026 10:01:52 -0600 Subject: [PATCH 07/21] Reset original bytes when rewriting Parse OIDs --- pgdog/src/net/messages/parse.rs | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/pgdog/src/net/messages/parse.rs b/pgdog/src/net/messages/parse.rs index c38becc73..0bb8fb36d 100644 --- a/pgdog/src/net/messages/parse.rs +++ b/pgdog/src/net/messages/parse.rs @@ -133,6 +133,7 @@ impl Parse { let oid_bytes = &mut bytes.get_mut()[(pos - 4)..pos]; oid_bytes.copy_from_slice(&shard_oid.to_be_bytes()); rewritten = true; + self.original = None; } } self.data_types = bytes.into_inner().freeze(); @@ -263,4 +264,13 @@ mod test { let mut parse_without_custom_types = Parse::named("", "").with_data_types(&[1]); assert!(!parse_without_custom_types.rewrite_data_types(&mapping)); } + + #[test] + fn test_rewriting_oids_clears_original_bytes() { + let mut parse = Parse::named("", "").with_data_types(&[10_001]); + parse.original = Some(Bytes::new()); + let mapping = [(10_001, 10_002)].into_iter().collect(); + assert!(parse.rewrite_data_types(&mapping)); + assert!(!parse.to_bytes().is_empty()); + } } From d2b3401c11324027e7f1564924acdd2103df9236 Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Mon, 3 Aug 2026 10:21:50 -0600 Subject: [PATCH 08/21] Add an integration test for OID drift on simple query protocol --- .../integration/cross_shard_oid_drift.rs | 61 ++++++++++++++++++- 1 file changed, 60 insertions(+), 1 deletion(-) diff --git a/integration/rust/tests/integration/cross_shard_oid_drift.rs b/integration/rust/tests/integration/cross_shard_oid_drift.rs index b68434542..bce7c4041 100644 --- a/integration/rust/tests/integration/cross_shard_oid_drift.rs +++ b/integration/rust/tests/integration/cross_shard_oid_drift.rs @@ -1,6 +1,7 @@ #![cfg(feature = "new_parser")] -use crate::setup::{admin_sqlx, connections_sqlx}; +use crate::setup::{admin_sqlx, admin_tokio, connections_sqlx, connections_tokio}; use sqlx::Executor; +use tokio_postgres::SimpleQueryMessage; #[derive(sqlx::Type, Debug, Clone, PartialEq)] #[sqlx(type_name = "test_oid_drift_composite")] @@ -52,3 +53,61 @@ async fn test_oid_drift() { .unwrap(); assert_eq!(rows, vec![composite; 20]); } + +#[tokio::test] +async fn test_oid_drift_simple_protocol() { + let conn = connections_tokio().await.pop().unwrap(); + + // Intentionally cause the OID of the type to differ between shards + conn.simple_query("/* pgdog_shard: 0 */ CREATE SEQUENCE foo; DROP SEQUENCE foo;") + .await + .unwrap(); + conn.simple_query("DROP TYPE IF EXISTS test_oid_drift_simple_composite CASCADE") + .await + .unwrap(); + conn.simple_query("CREATE TYPE test_oid_drift_simple_composite AS (a text, b text)") + .await + .unwrap(); + conn.simple_query("DROP TABLE IF EXISTS test_oid_drift_simple") + .await + .unwrap(); + conn.simple_query( + "CREATE TABLE test_oid_drift_simple (customer_id BIGINT, composite test_oid_drift_simple_composite)", + ) + .await + .unwrap(); + admin_tokio().await.simple_query("RELOAD").await.unwrap(); + + let composite = Composite { + a: String::from("b"), + b: String::from("c"), + }; + for i in 1..=20 { + conn.simple_query(&format!( + "INSERT INTO test_oid_drift_simple VALUES ({}, ROW('{}', '{}')::test_oid_drift_simple_composite)", + i, + composite.a, + composite.b, + )) + .await + .unwrap(); + } + + let messages = conn + .simple_query("SELECT composite FROM test_oid_drift_simple") + .await + .unwrap(); + assert_eq!(messages.len(), 22); + let rows = messages + .iter() + .skip(1) // RowDescription + .take(20) // CommandComplete at end + .map(|message| { + let SimpleQueryMessage::Row(row) = message else { + panic!("not a DataRow: {:?}", message); + }; + row.get(0).unwrap() + }) + .collect::>(); + assert_eq!(rows, vec!["(b,c)"; 20]); +} From 933f1b3435337347da5e69f80c96fabe17279bdd Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Mon, 3 Aug 2026 10:28:16 -0600 Subject: [PATCH 09/21] s/sleep/safe_sleep --- pgdog/src/backend/pool/cluster/schema_loader.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgdog/src/backend/pool/cluster/schema_loader.rs b/pgdog/src/backend/pool/cluster/schema_loader.rs index 38a3d8d96..8154fea09 100644 --- a/pgdog/src/backend/pool/cluster/schema_loader.rs +++ b/pgdog/src/backend/pool/cluster/schema_loader.rs @@ -51,7 +51,7 @@ impl SchemaLoader for FromServer { Err(err) => { if shard.online() { error!("error loading canonical type information: {err}"); - sleep(Duration::from_millis(100)).await; + safe_sleep(Duration::from_millis(100)).await; } else { // Cluster is shutting down break; From 751ffb618fef5a2490d987ab3961956796fb1c46 Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Mon, 3 Aug 2026 10:49:42 -0600 Subject: [PATCH 10/21] regclass is not regtype Not sure why I got these mixed up in my head, regtype casts are what is equivalent to selecting from `pg_type` --- pgdog/src/frontend/router/parser/query/mod.rs | 2 +- .../src/frontend/router/parser/query/test/test_select.rs | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/pgdog/src/frontend/router/parser/query/mod.rs b/pgdog/src/frontend/router/parser/query/mod.rs index 1b8d7708b..c31c5f6ff 100644 --- a/pgdog/src/frontend/router/parser/query/mod.rs +++ b/pgdog/src/frontend/router/parser/query/mod.rs @@ -1004,7 +1004,7 @@ fn references_pg_type(stmt: &nodes::SelectStmt) -> bool { .funcname() .iter() .map(Node::as_str) - .eq([Some("to_regclass")]) => + .eq([Some("to_regtype")]) => { ControlFlow::Break(true) } diff --git a/pgdog/src/frontend/router/parser/query/test/test_select.rs b/pgdog/src/frontend/router/parser/query/test/test_select.rs index aa81527b2..b1ad26cc5 100644 --- a/pgdog/src/frontend/router/parser/query/test/test_select.rs +++ b/pgdog/src/frontend/router/parser/query/test/test_select.rs @@ -277,18 +277,18 @@ fn test_system_catalog_sharded() { "pg_type queries should go to shard 0", ); - let command = test.execute(vec![Query::new("SELECT $1::regclass").into()]); + let command = test.execute(vec![Query::new("SELECT $1::regtype").into()]); assert_eq!( command.route().shard(), &Shard::Direct(0), - "regclass casts should go to shard 0", + "regtype casts should go to shard 0", ); - let command = test.execute(vec![Query::new("SELECT to_regclass($1)").into()]); + let command = test.execute(vec![Query::new("SELECT to_regtype($1)").into()]); assert_eq!( command.route().shard(), &Shard::Direct(0), - "to_regclass should go to shard 0", + "to_regtype should go to shard 0", ); } From 8170856917698d1f9b336bb0982107fd9ebe88f7 Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Mon, 3 Aug 2026 11:37:50 -0600 Subject: [PATCH 11/21] Don't error when schema information is queried in direct-to-shard This isn't ideal. We can't just say "we don't perform OID mapping for direct-to-shard transactions", as type information (including OIDs) is typically cached client-side for the entire lifetime of the connection, and that same conneciton could be used to hit a different shard in the very next transaction. If we supported lazily connecting to shards as we need them, then we could make an exception for schema based sharding or transactions routed based on SET specifically to route schema queries to shard 0. But for now the best we can do is warn, and hope that the query isn't looking for any OIDs that will differ across shards. This warning should rarely fire, if ever. Queries to `pg_type` are typically handled by the client library, not the user, and they are typically run immediately upon connecting. For those queries to be in a direct-to-shard transaction they would need to be run in a transaction where search_path or pg_shard are immediately set. The test that hit this issue was a migration library, which I suspect is the only client pattern that could hit this. And since DDL shouldn't be affected by OID information, I don't believe this will end up causing any trouble for those who do receive the warning --- .../client/query_engine/route_query.rs | 14 +++++++++---- pgdog/src/frontend/router/parser/route.rs | 20 +++++++++++++++++++ 2 files changed, 30 insertions(+), 4 deletions(-) diff --git a/pgdog/src/frontend/client/query_engine/route_query.rs b/pgdog/src/frontend/client/query_engine/route_query.rs index 9f5902d40..dd64380a2 100644 --- a/pgdog/src/frontend/client/query_engine/route_query.rs +++ b/pgdog/src/frontend/client/query_engine/route_query.rs @@ -1,5 +1,5 @@ use pgdog_config::PoolerMode; -use tracing::trace; +use tracing::{trace, warn}; use crate::backend::Cluster; use crate::util::safe_timeout; @@ -116,9 +116,15 @@ impl QueryEngine { } if Self::is_shard_switch(command, &self.backend) { - self.error_response(context, ErrorResponse::direct_shard_mismatch()) - .await?; - return Ok(false); + if command.route().is_canonical_schema() { + warn!( + "type information queried in direct-to-shard transaction, cannot route to canonical source" + ); + } else { + self.error_response(context, ErrorResponse::direct_shard_mismatch()) + .await?; + return Ok(false); + } } } } diff --git a/pgdog/src/frontend/router/parser/route.rs b/pgdog/src/frontend/router/parser/route.rs index f1344ab59..70b0a5d65 100644 --- a/pgdog/src/frontend/router/parser/route.rs +++ b/pgdog/src/frontend/router/parser/route.rs @@ -375,6 +375,13 @@ impl Route { pub(crate) fn set_rewrite_plan(&mut self, plan: AggregateRewritePlan) { self.rewrite_plan = plan; } + + pub(crate) fn is_canonical_schema(&self) -> bool { + self.shard_with_priority() + .source() + .override_reason() + .is_some_and(|reason| reason.is_canonical_schema_info()) + } } /// Shard source. @@ -401,6 +408,13 @@ impl ShardSource { pub fn is_round_robin(&self) -> bool { matches!(self, Self::RoundRobin(_)) } + + pub(crate) fn override_reason(&self) -> Option<&OverrideReason> { + match self { + Self::Override(reason) => Some(reason), + _ => None, + } + } } #[derive(Debug, Clone, PartialEq, Eq, Ord, PartialOrd)] @@ -423,6 +437,12 @@ pub enum OverrideReason { CanonicalSchemaInfo, } +impl OverrideReason { + pub(crate) fn is_canonical_schema_info(&self) -> bool { + matches!(self, OverrideReason::CanonicalSchemaInfo) + } +} + #[derive(Debug, Clone, PartialEq, Eq, Ord, PartialOrd)] pub enum TableReason { Omni, From b81ad83e98820b0f587e9a0baa90ce2f96337feb Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Mon, 3 Aug 2026 12:59:45 -0600 Subject: [PATCH 12/21] Revert "Don't error when schema information is queried in direct-to-shard" This reverts commit a62bb2c76a2217a059fb575da7fcb35ccca49429. When I wrote my reasoning in that commit I didn't realize that this was only happening because of the way we are monkey patching alembic. I assume this is based on a pattern used by a customer, the real solution is likely just to have this be config option that can be disabled in that case --- .../client/query_engine/route_query.rs | 14 ++++--------- pgdog/src/frontend/router/parser/route.rs | 20 ------------------- 2 files changed, 4 insertions(+), 30 deletions(-) diff --git a/pgdog/src/frontend/client/query_engine/route_query.rs b/pgdog/src/frontend/client/query_engine/route_query.rs index dd64380a2..9f5902d40 100644 --- a/pgdog/src/frontend/client/query_engine/route_query.rs +++ b/pgdog/src/frontend/client/query_engine/route_query.rs @@ -1,5 +1,5 @@ use pgdog_config::PoolerMode; -use tracing::{trace, warn}; +use tracing::trace; use crate::backend::Cluster; use crate::util::safe_timeout; @@ -116,15 +116,9 @@ impl QueryEngine { } if Self::is_shard_switch(command, &self.backend) { - if command.route().is_canonical_schema() { - warn!( - "type information queried in direct-to-shard transaction, cannot route to canonical source" - ); - } else { - self.error_response(context, ErrorResponse::direct_shard_mismatch()) - .await?; - return Ok(false); - } + self.error_response(context, ErrorResponse::direct_shard_mismatch()) + .await?; + return Ok(false); } } } diff --git a/pgdog/src/frontend/router/parser/route.rs b/pgdog/src/frontend/router/parser/route.rs index 70b0a5d65..f1344ab59 100644 --- a/pgdog/src/frontend/router/parser/route.rs +++ b/pgdog/src/frontend/router/parser/route.rs @@ -375,13 +375,6 @@ impl Route { pub(crate) fn set_rewrite_plan(&mut self, plan: AggregateRewritePlan) { self.rewrite_plan = plan; } - - pub(crate) fn is_canonical_schema(&self) -> bool { - self.shard_with_priority() - .source() - .override_reason() - .is_some_and(|reason| reason.is_canonical_schema_info()) - } } /// Shard source. @@ -408,13 +401,6 @@ impl ShardSource { pub fn is_round_robin(&self) -> bool { matches!(self, Self::RoundRobin(_)) } - - pub(crate) fn override_reason(&self) -> Option<&OverrideReason> { - match self { - Self::Override(reason) => Some(reason), - _ => None, - } - } } #[derive(Debug, Clone, PartialEq, Eq, Ord, PartialOrd)] @@ -437,12 +423,6 @@ pub enum OverrideReason { CanonicalSchemaInfo, } -impl OverrideReason { - pub(crate) fn is_canonical_schema_info(&self) -> bool { - matches!(self, OverrideReason::CanonicalSchemaInfo) - } -} - #[derive(Debug, Clone, PartialEq, Eq, Ord, PartialOrd)] pub enum TableReason { Omni, From fbfe50f5632a522fa2a34ffea10009154999f528 Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Mon, 3 Aug 2026 13:54:29 -0600 Subject: [PATCH 13/21] Include the type namespace in type oid lookup I opted to concatenate strings instead of selecting two columns to avoid the PITA of borrowing a tuple key for a hash map lookup --- pgdog/src/backend/pool/shard/oids.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgdog/src/backend/pool/shard/oids.rs b/pgdog/src/backend/pool/shard/oids.rs index 90eb8ffbc..403a8084a 100644 --- a/pgdog/src/backend/pool/shard/oids.rs +++ b/pgdog/src/backend/pool/shard/oids.rs @@ -119,7 +119,7 @@ async fn load_oids(server: &mut Server) -> Result( - "SELECT DISTINCT ON (typname) typname, pg_type.oid FROM pg_type WHERE pg_type.oid >= 10000", + "SELECT nspname || '.' || typname, pg_type.oid FROM pg_type INNER JOIN pg_namespace ON typnamespace = pg_namespace.oid WHERE pg_type.oid >= 10000", ) .await? .into_iter() From 691f7104b446f9caffb37830605d23b8e9aa8ab5 Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Mon, 3 Aug 2026 13:59:22 -0600 Subject: [PATCH 14/21] Use CancellationToken::run_until_cancelled --- .../src/backend/pool/cluster/schema_loader.rs | 24 +++++++------------ 1 file changed, 9 insertions(+), 15 deletions(-) diff --git a/pgdog/src/backend/pool/cluster/schema_loader.rs b/pgdog/src/backend/pool/cluster/schema_loader.rs index 8154fea09..ad5bcbc5c 100644 --- a/pgdog/src/backend/pool/cluster/schema_loader.rs +++ b/pgdog/src/backend/pool/cluster/schema_loader.rs @@ -31,24 +31,18 @@ impl SchemaLoader for FromServer { let canonical_oids = Arc::clone(&cluster.canonical_oids); let shard = shard.clone(); tasks::spawn("load canonical oids", async move { - // FIXME: This shutdown signal/retry if error logic seems like - // it might be worth abstracting - let shutdown = tasks::shutdown_signal(); - loop { - let loader = async { - canonical_oids - .load(&mut *shard.primary_or_replica(&Default::default()).await?) - .await - }; - let result = select! { - _ = shutdown.cancelled() => break, - result = loader => { result }, - }; + let result = tasks::shutdown_signal() + .run_until_cancelled(async { + canonical_oids + .load(&mut *shard.primary_or_replica(&Default::default()).await?) + .await + }) + .await; match result { - Ok(_) => break, - Err(err) => { + Some(Ok(_)) | None => break, + Some(Err(err)) => { if shard.online() { error!("error loading canonical type information: {err}"); safe_sleep(Duration::from_millis(100)).await; From 8c2b60c16a144cab3ca516e903aa1e11ee914fbb Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Mon, 3 Aug 2026 14:40:31 -0600 Subject: [PATCH 15/21] Rewrite `RowDescription` messages universally Previously we were rewriting them when being inserted into the global cache, which only occurs in the extended protocol. We now always attempt to rewrite them, regardless of whether they are going into a cache or not. This introduces a little bit of trickiness. First, we expect the most common case to be that there are no mappings at all. When using the simple protocol, we previously were never parsing `RowDescription` messages. If the OID map is empty, we want to continue to avoid performing that parsing. The second bit of trickiness is that this rewrite will now be called as part of loading the OIDs themselves. So we can no longer expect them to always be set for this function in particular, and need to return an empty set instead. --- .../integration/cross_shard_oid_drift.rs | 68 ++++--------------- pgdog/src/backend/prepared_statements.rs | 35 +++++++--- pgdog/src/net/messages/row_description.rs | 4 -- 3 files changed, 39 insertions(+), 68 deletions(-) diff --git a/integration/rust/tests/integration/cross_shard_oid_drift.rs b/integration/rust/tests/integration/cross_shard_oid_drift.rs index bce7c4041..fbf2d37a6 100644 --- a/integration/rust/tests/integration/cross_shard_oid_drift.rs +++ b/integration/rust/tests/integration/cross_shard_oid_drift.rs @@ -1,7 +1,7 @@ #![cfg(feature = "new_parser")] -use crate::setup::{admin_sqlx, admin_tokio, connections_sqlx, connections_tokio}; -use sqlx::Executor; -use tokio_postgres::SimpleQueryMessage; +use crate::setup::{admin_sqlx, connections_sqlx}; +use sqlx::postgres::types::Oid; +use sqlx::{Column, Executor, Row}; #[derive(sqlx::Type, Debug, Clone, PartialEq)] #[sqlx(type_name = "test_oid_drift_composite")] @@ -51,63 +51,21 @@ async fn test_oid_drift() { .fetch_all(&conn) .await .unwrap(); - assert_eq!(rows, vec![composite; 20]); -} - -#[tokio::test] -async fn test_oid_drift_simple_protocol() { - let conn = connections_tokio().await.pop().unwrap(); + assert_eq!(rows, vec![composite.clone(); 20]); - // Intentionally cause the OID of the type to differ between shards - conn.simple_query("/* pgdog_shard: 0 */ CREATE SEQUENCE foo; DROP SEQUENCE foo;") - .await - .unwrap(); - conn.simple_query("DROP TYPE IF EXISTS test_oid_drift_simple_composite CASCADE") - .await - .unwrap(); - conn.simple_query("CREATE TYPE test_oid_drift_simple_composite AS (a text, b text)") - .await - .unwrap(); - conn.simple_query("DROP TABLE IF EXISTS test_oid_drift_simple") + let simple_rows = conn + .fetch_all("SELECT composite FROM test_oid_drift") .await .unwrap(); - conn.simple_query( - "CREATE TABLE test_oid_drift_simple (customer_id BIGINT, composite test_oid_drift_simple_composite)", - ) - .await - .unwrap(); - admin_tokio().await.simple_query("RELOAD").await.unwrap(); - let composite = Composite { - a: String::from("b"), - b: String::from("c"), - }; - for i in 1..=20 { - conn.simple_query(&format!( - "INSERT INTO test_oid_drift_simple VALUES ({}, ROW('{}', '{}')::test_oid_drift_simple_composite)", - i, - composite.a, - composite.b, - )) + let expected_oid: Oid = + sqlx::query_scalar("SELECT oid FROM pg_type WHERE typname = 'test_oid_drift_composite'") + .fetch_one(&conn) .await .unwrap(); - } + let given_oid = simple_rows.first().unwrap().column(0).type_info().oid(); + assert_eq!(given_oid, Some(expected_oid)); - let messages = conn - .simple_query("SELECT composite FROM test_oid_drift_simple") - .await - .unwrap(); - assert_eq!(messages.len(), 22); - let rows = messages - .iter() - .skip(1) // RowDescription - .take(20) // CommandComplete at end - .map(|message| { - let SimpleQueryMessage::Row(row) = message else { - panic!("not a DataRow: {:?}", message); - }; - row.get(0).unwrap() - }) - .collect::>(); - assert_eq!(rows, vec!["(b,c)"; 20]); + let simple_data: Vec = simple_rows.into_iter().map(|row| row.get(0)).collect(); + assert_eq!(simple_data, vec![composite; 20]); } diff --git a/pgdog/src/backend/prepared_statements.rs b/pgdog/src/backend/prepared_statements.rs index e10401455..9fe747ca0 100644 --- a/pgdog/src/backend/prepared_statements.rs +++ b/pgdog/src/backend/prepared_statements.rs @@ -294,8 +294,12 @@ impl PreparedStatements { } 'T' => { + let maybe_row_description = self.parse_and_rewrite_row_description(message)?; if let Some(describe) = self.describes.pop_front() { - self.add_row_description(&describe, message)?; + let row_description = maybe_row_description + .map(Ok) + .unwrap_or_else(|| RowDescription::from_bytes(message.payload()))?; + self.add_row_description(&describe, row_description); }; } @@ -402,15 +406,10 @@ impl PreparedStatements { /// Handle a Describe message, storing the RowDescription for the /// statement in the global cache. - fn add_row_description(&self, name: &str, message: &mut Message) -> Result<(), Error> { - let mut row_description = RowDescription::from_bytes(message.payload())?; - if self.rewrite_row_description_data_types(&mut row_description) { - message.replace_payload(row_description.to_bytes()); - } + fn add_row_description(&self, name: &str, row_description: RowDescription) { self.global_cache .write() .insert_row_description(name, row_description); - Ok(()) } /// Remove statement from local cache. @@ -481,8 +480,26 @@ impl PreparedStatements { parse.rewrite_data_types(&self.oid_mappings().canonical_to_shard) } - fn rewrite_row_description_data_types(&self, row_description: &mut RowDescription) -> bool { - row_description.rewrite_data_types(&self.oid_mappings().shard_to_canonical) + /// Rewrite the given RowDescription Message to have the canonical set of + /// OIDs. Returns the parsed RowDescription if parsing occurred + fn parse_and_rewrite_row_description( + &self, + message: &mut Message, + ) -> Result, Error> { + // RowDescription is emitted during cluster startup, so we can't + // require OIDs to be loaded. + let empty_mapping = Default::default(); + let mappings = &self.oids.get().unwrap_or(&empty_mapping).shard_to_canonical; + + if !mappings.is_empty() { + let mut row_description = RowDescription::from_bytes(message.payload())?; + if row_description.rewrite_data_types(mappings) { + message.replace_payload(row_description.to_bytes()); + } + Ok(Some(row_description)) + } else { + Ok(None) + } } fn rewrite_parameter_description_data_types(&self, message: &mut Message) -> Result<(), Error> { diff --git a/pgdog/src/net/messages/row_description.rs b/pgdog/src/net/messages/row_description.rs index 6a7c1ab94..4b8c119f0 100644 --- a/pgdog/src/net/messages/row_description.rs +++ b/pgdog/src/net/messages/row_description.rs @@ -270,10 +270,6 @@ impl RowDescription { /// Replaces the data types of each field using the given mapping. /// Returns whether any changes actually occurred. pub(crate) fn rewrite_data_types(&mut self, mapping: &HashMap) -> bool { - if mapping.is_empty() { - return false; - } - let mut changed = false; for field in Arc::make_mut(&mut self.fields) { if let Some(&canonical) = mapping.get(&(field.type_oid as u32)) { From a5e0382d476536c24278a650bb9228eb375b5e14 Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Mon, 3 Aug 2026 15:00:22 -0600 Subject: [PATCH 16/21] Don't panic if OIDs aren't loaded I had hoped not to silently ignore OIDs not being loaded, as panicking could potentially catch bugs in the future. But now that we have legitimate reasons that we need to allow OIDs to not be loaded, this panic is no longer worth the cost to the rest of the code base --- pgdog/src/backend/mod.rs | 2 +- pgdog/src/backend/pool/cluster.rs | 6 ------ pgdog/src/backend/pool/mod.rs | 2 +- pgdog/src/backend/pool/shard/mod.rs | 2 +- pgdog/src/backend/prepared_statements.rs | 20 +++++++++---------- .../replication/logical/subscriber/tests.rs | 18 ----------------- 6 files changed, 13 insertions(+), 37 deletions(-) diff --git a/pgdog/src/backend/mod.rs b/pgdog/src/backend/mod.rs index e90b36e0b..608fceb81 100644 --- a/pgdog/src/backend/mod.rs +++ b/pgdog/src/backend/mod.rs @@ -22,7 +22,7 @@ pub use connect_reason::ConnectReason; pub use disconnect_reason::DisconnectReason; pub use error::Error; pub(crate) use pool::{ - CanonicalOids, Cluster, ClusterShardConfig, OidMappings, Oids, Pool, Shard, ShardingSchema, + CanonicalOids, Cluster, ClusterShardConfig, Oids, Pool, Shard, ShardingSchema, }; pub use prepared_statements::PreparedStatements; pub use protocol::*; diff --git a/pgdog/src/backend/pool/cluster.rs b/pgdog/src/backend/pool/cluster.rs index 9793b737f..7c4812f26 100644 --- a/pgdog/src/backend/pool/cluster.rs +++ b/pgdog/src/backend/pool/cluster.rs @@ -912,12 +912,6 @@ mod test { cluster } - pub(crate) fn force_dummy_schema(&self) { - for shard in &self.shards { - shard.schema_not_needed(); - } - } - pub(crate) fn set_read_write_strategy(&mut self, rw_strategy: ReadWriteStrategy) { self.rw_strategy = rw_strategy; } diff --git a/pgdog/src/backend/pool/mod.rs b/pgdog/src/backend/pool/mod.rs index fdd8aea11..ada2d6e7e 100644 --- a/pgdog/src/backend/pool/mod.rs +++ b/pgdog/src/backend/pool/mod.rs @@ -40,7 +40,7 @@ pub use monitor::Monitor; pub use password::Password; pub use pool_impl::Pool; pub use request::Request; -pub(crate) use shard::{CanonicalOids, OidMappings, Oids, Shard}; +pub(crate) use shard::{CanonicalOids, Oids, Shard}; pub use state::State; pub use stats::Stats; diff --git a/pgdog/src/backend/pool/shard/mod.rs b/pgdog/src/backend/pool/shard/mod.rs index b78722f73..b8ba21a33 100644 --- a/pgdog/src/backend/pool/shard/mod.rs +++ b/pgdog/src/backend/pool/shard/mod.rs @@ -27,7 +27,7 @@ mod oids; pub mod role_detector; use monitor::*; -pub(crate) use oids::{CanonicalOids, OidMappings, Oids}; +pub(crate) use oids::{CanonicalOids, Oids}; use role_detector::*; #[cfg_attr(test, derive(Default))] diff --git a/pgdog/src/backend/prepared_statements.rs b/pgdog/src/backend/prepared_statements.rs index 9fe747ca0..dbb264f0e 100644 --- a/pgdog/src/backend/prepared_statements.rs +++ b/pgdog/src/backend/prepared_statements.rs @@ -12,7 +12,7 @@ use crate::{ use parking_lot::RwLock; use pgdog_config::PreparedStatements as PreparedStatementsLevel; -use super::{Error, OidMappings, Oids}; +use super::{Error, Oids}; use super::{ protocol::{ProtocolState, state::Action}, state::ExecutionCode, @@ -477,7 +477,10 @@ impl PreparedStatements { } fn rewrite_parse_data_types(&self, parse: &mut Parse) -> bool { - parse.rewrite_data_types(&self.oid_mappings().canonical_to_shard) + let Some(mappings) = self.oids.get() else { + return false; + }; + parse.rewrite_data_types(&mappings.canonical_to_shard) } /// Rewrite the given RowDescription Message to have the canonical set of @@ -503,22 +506,19 @@ impl PreparedStatements { } fn rewrite_parameter_description_data_types(&self, message: &mut Message) -> Result<(), Error> { - let mappings = &self.oid_mappings().shard_to_canonical; + let Some(mappings) = self.oids.get() else { + return Ok(()); + }; + let mappings = &mappings.shard_to_canonical; if mappings.is_empty() { return Ok(()); } let mut parameter_description = ParameterDescription::from_bytes(message.payload())?; - parameter_description.rewrite_data_types(&self.oid_mappings().shard_to_canonical); + parameter_description.rewrite_data_types(mappings); message.replace_payload(parameter_description.to_bytes()); Ok(()) } - - fn oid_mappings(&self) -> &OidMappings { - self.oids - .get() - .expect("cluster startup always loads OID information or sets an empty mapping") - } } #[cfg(test)] diff --git a/pgdog/src/backend/replication/logical/subscriber/tests.rs b/pgdog/src/backend/replication/logical/subscriber/tests.rs index c046a4986..544f2474a 100644 --- a/pgdog/src/backend/replication/logical/subscriber/tests.rs +++ b/pgdog/src/backend/replication/logical/subscriber/tests.rs @@ -250,14 +250,12 @@ fn x_update(u: XLogUpdate) -> CopyData { fn make_subscriber() -> StreamSubscriber { let cluster = Cluster::new_test(&config()); - cluster.force_dummy_schema(); let tables = vec![make_sharded_table(), make_sharded_test_b_table()]; StreamSubscriber::new(&cluster, &tables, OmniOwnership::test()) } fn make_subscriber_with_tables(tables: Vec
) -> StreamSubscriber { let cluster = Cluster::new_test(&config()); - cluster.force_dummy_schema(); StreamSubscriber::new(&cluster, &tables, OmniOwnership::test()) } @@ -266,13 +264,11 @@ fn make_subscriber_with_tables_two_databases( partition: OmniOwnership, ) -> StreamSubscriber { let cluster = Cluster::new_test_two_databases(&config()); - cluster.force_dummy_schema(); StreamSubscriber::new(&cluster, &tables, partition) } fn make_subscriber_single_shard() -> StreamSubscriber { let cluster = Cluster::new_test_single_shard(&config()); - cluster.force_dummy_schema(); let tables = vec![make_sharded_table(), make_sharded_test_b_table()]; StreamSubscriber::new(&cluster, &tables, OmniOwnership::test()) } @@ -614,7 +610,6 @@ async fn partition_leaves_share_destination() { leaf_b.table.parent_name = "sharded".to_string(); let cluster = Cluster::new_test_single_shard(&config()); - cluster.force_dummy_schema(); let mut sub = StreamSubscriber::new(&cluster, &[leaf_a, leaf_b], OmniOwnership::test()); let mut verify = test_server().await; sub.connect().await.unwrap(); @@ -1551,7 +1546,6 @@ 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()); - cluster.force_dummy_schema(); let mut sub = StreamSubscriber::new( &cluster, &[make_replica_identity_nothing_table()], @@ -1590,7 +1584,6 @@ 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()); - cluster.force_dummy_schema(); let mut sub = StreamSubscriber::new( &cluster, &[make_full_identity_omni_table()], @@ -1634,7 +1627,6 @@ 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()); - cluster.force_dummy_schema(); let mut sub = StreamSubscriber::new( &cluster, &[make_full_identity_sharded_table()], @@ -1665,7 +1657,6 @@ async fn full_identity_insert_sharded() { #[tokio::test] async fn full_identity_update_fast_path() { let cluster = Cluster::new_test_single_shard(&config()); - cluster.force_dummy_schema(); let mut sub = StreamSubscriber::new( &cluster, &[make_full_identity_sharded_table()], @@ -1725,7 +1716,6 @@ async fn full_identity_update_fast_path() { #[tokio::test] async fn full_identity_update_slow_path() { let cluster = Cluster::new_test_single_shard(&config()); - cluster.force_dummy_schema(); let mut sub = StreamSubscriber::new( &cluster, &[make_full_identity_sharded_table()], @@ -1790,7 +1780,6 @@ 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()); - cluster.force_dummy_schema(); let mut sub = StreamSubscriber::new( &cluster, &[make_full_identity_sharded_table()], @@ -1852,7 +1841,6 @@ 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()); - cluster.force_dummy_schema(); let mut sub = StreamSubscriber::new( &cluster, &[make_full_identity_sharded_table()], @@ -1899,7 +1887,6 @@ async fn full_identity_update_all_toasted_is_noop() { #[tokio::test] async fn full_identity_delete() { let cluster = Cluster::new_test_single_shard(&config()); - cluster.force_dummy_schema(); let mut sub = StreamSubscriber::new( &cluster, &[make_full_identity_sharded_table()], @@ -1942,7 +1929,6 @@ async fn full_identity_delete() { #[tokio::test] async fn full_identity_insert_omni_dedup() { let cluster = Cluster::new_test_single_shard(&config()); - cluster.force_dummy_schema(); let mut sub = StreamSubscriber::new( &cluster, &[make_full_identity_omni_dedup_table()], @@ -2006,7 +1992,6 @@ async fn full_identity_insert_omni_dedup() { #[tokio::test] async fn full_identity_update_duplicate_rows() { let cluster = Cluster::new_test_single_shard(&config()); - cluster.force_dummy_schema(); let mut sub = StreamSubscriber::new( &cluster, &[make_full_identity_dup_rows_table()], @@ -2077,7 +2062,6 @@ async fn full_identity_update_duplicate_rows() { #[tokio::test] async fn full_identity_delete_duplicate_rows() { let cluster = Cluster::new_test_single_shard(&config()); - cluster.force_dummy_schema(); let mut sub = StreamSubscriber::new( &cluster, &[make_full_identity_dup_rows_table()], @@ -2149,7 +2133,6 @@ 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()); - cluster.force_dummy_schema(); let mut sub = StreamSubscriber::new( &cluster, &[make_full_identity_dup_rows_table()], @@ -2215,7 +2198,6 @@ 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()); - cluster.force_dummy_schema(); let mut sub = StreamSubscriber::new( &cluster, &[make_full_identity_dup_rows_table()], From 014ce3d5b9e8cc0425fc0b5be317a5017a6f2a5e Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Mon, 3 Aug 2026 15:45:37 -0600 Subject: [PATCH 17/21] Change OID canonicalization to a config option This turns off the canonicalization behavior by default. We expect this feature to only be useful for a subset of our users, and there are certain usage patterns that cannot support OID canonicalization at all (as demonstrated in `integration/python/albemic/test_migration.py`). As such, this behavior is disabled by default, and users who are using sharding and have extension or user defined types can enable the feature explicitly. --- .../rust/tests/integration/cross_shard_oid_drift.rs | 8 +++++++- pgdog-config/src/general.rs | 13 +++++++++++++ pgdog/src/admin/set.rs | 4 ++++ pgdog/src/backend/pool/cluster.rs | 11 +++++++++-- pgdog/src/backend/pool/cluster/schema_loader.rs | 10 ++++++++-- pgdog/src/backend/pool/shard/mod.rs | 5 +++++ pgdog/src/frontend/router/parser/context.rs | 4 ++++ pgdog/src/frontend/router/parser/query/mod.rs | 2 +- 8 files changed, 51 insertions(+), 6 deletions(-) diff --git a/integration/rust/tests/integration/cross_shard_oid_drift.rs b/integration/rust/tests/integration/cross_shard_oid_drift.rs index fbf2d37a6..0c3aeb76a 100644 --- a/integration/rust/tests/integration/cross_shard_oid_drift.rs +++ b/integration/rust/tests/integration/cross_shard_oid_drift.rs @@ -13,6 +13,7 @@ struct Composite { #[tokio::test] async fn test_oid_drift() { let conn = connections_sqlx().await.pop().unwrap(); + let admin = admin_sqlx().await; // Intentionally cause the OID of the type to differ between shards conn.execute("/* pgdog_shard: 0 */ CREATE SEQUENCE foo; DROP SEQUENCE foo;") @@ -32,7 +33,10 @@ async fn test_oid_drift() { ) .await .unwrap(); - admin_sqlx().await.execute("RELOAD").await.unwrap(); + admin + .execute("SET canonicalize_type_information TO true") + .await + .unwrap(); let composite = Composite { a: String::from("a"), @@ -68,4 +72,6 @@ async fn test_oid_drift() { let simple_data: Vec = simple_rows.into_iter().map(|row| row.get(0)).collect(); assert_eq!(simple_data, vec![composite; 20]); + + admin.execute("RELOAD").await.unwrap(); } diff --git a/pgdog-config/src/general.rs b/pgdog-config/src/general.rs index f82d21b4b..947839cd8 100644 --- a/pgdog-config/src/general.rs +++ b/pgdog-config/src/general.rs @@ -767,6 +767,14 @@ pub struct General { #[serde(default = "General::load_schema")] pub load_schema: LoadSchema, + /// Controls whether PgDog maps each shard's type information to a canonical source. + /// + /// _Default:_ `false` + /// + /// + #[serde(default = "General::canonicalize_type_information")] + pub canonicalize_type_information: bool, + /// Replication lag threshold (in bytes) at which PgDog will pause traffic automatically during a traffic cutover. /// /// _Default:_ `1000000` @@ -915,6 +923,7 @@ impl Default for General { resharding_replication_retry_min_delay: Self::resharding_replication_retry_min_delay(), reload_schema_on_ddl: Self::reload_schema_on_ddl(), load_schema: Self::load_schema(), + canonicalize_type_information: Self::canonicalize_type_information(), cutover_replication_lag_threshold: Self::cutover_replication_lag_threshold(), cutover_traffic_stop_threshold: Self::cutover_traffic_stop_threshold(), cutover_last_transaction_delay: Self::cutover_last_transaction_delay(), @@ -1327,6 +1336,10 @@ impl General { Self::env_enum_or_default("PGDOG_LOAD_SCHEMA") } + fn canonicalize_type_information() -> bool { + Self::env_or_default("PGDOG_CANONICALIZE_TYPE_INFORMATION", false) + } + pub fn mirror_queue() -> usize { Self::env_or_default("PGDOG_MIRROR_QUEUE", 128) } diff --git a/pgdog/src/admin/set.rs b/pgdog/src/admin/set.rs index 28859c079..cae78e626 100644 --- a/pgdog/src/admin/set.rs +++ b/pgdog/src/admin/set.rs @@ -211,6 +211,10 @@ impl Command for Set { config.config.general.connect_timeout = self.value.parse()?; } + "canonicalize_type_information" => { + config.config.general.canonicalize_type_information = Self::from_json(&self.value)?; + } + _ => return Ok(vec![]), } diff --git a/pgdog/src/backend/pool/cluster.rs b/pgdog/src/backend/pool/cluster.rs index 7c4812f26..807c52ed7 100644 --- a/pgdog/src/backend/pool/cluster.rs +++ b/pgdog/src/backend/pool/cluster.rs @@ -87,7 +87,7 @@ pub struct Cluster { tls_client_certificate_required: bool, #[debug(skip)] schema_loader: Box, - canonical_oids: Arc, + canonical_oids: Option>, } /// Sharding configuration from the cluster. @@ -177,6 +177,7 @@ pub struct ClusterConfig<'a> { identity: &'a Option, tls_client_certificate_required: bool, schema_cache: SchemaCache, + canonicalize_oids: bool, } impl<'a> ClusterConfig<'a> { @@ -248,6 +249,7 @@ impl<'a> ClusterConfig<'a> { identity: &user.identity, tls_client_certificate_required: user.tls_client_certificate_required.unwrap_or(true), schema_cache, + canonicalize_oids: general.canonicalize_type_information, } } } @@ -296,13 +298,14 @@ impl Cluster { identity, tls_client_certificate_required, schema_cache, + canonicalize_oids, } = config; let identifier = Arc::new(DatabaseUser { user: user.to_owned(), database: name.to_owned(), }); - let canonical_oids = schema_cache.canonical_oids(name); + let canonical_oids = canonicalize_oids.then(|| schema_cache.canonical_oids(name)); Self { identifier: identifier.clone(), @@ -683,6 +686,10 @@ impl Cluster { Ok(()) } + + pub(crate) fn is_canonicalizing_oids(&self) -> bool { + self.canonical_oids.is_some() + } } #[cfg(test)] diff --git a/pgdog/src/backend/pool/cluster/schema_loader.rs b/pgdog/src/backend/pool/cluster/schema_loader.rs index ad5bcbc5c..4919f6b20 100644 --- a/pgdog/src/backend/pool/cluster/schema_loader.rs +++ b/pgdog/src/backend/pool/cluster/schema_loader.rs @@ -27,8 +27,10 @@ impl SchemaLoader for FromServer { } // For now we treat shard 0 as the canonical OID source - if let Some(shard) = cluster.shards().first() { - let canonical_oids = Arc::clone(&cluster.canonical_oids); + if let Some(shard) = cluster.shards().first() + && let Some(canonical_oids) = dbg!(&cluster.canonical_oids) + { + let canonical_oids = Arc::clone(canonical_oids); let shard = shard.clone(); tasks::spawn("load canonical oids", async move { loop { @@ -54,6 +56,10 @@ impl SchemaLoader for FromServer { } } }); + } else { + for shard in cluster.shards() { + shard.skip_loading_oids() + } } for shard in cluster.shards() { diff --git a/pgdog/src/backend/pool/shard/mod.rs b/pgdog/src/backend/pool/shard/mod.rs index b8ba21a33..79c812e62 100644 --- a/pgdog/src/backend/pool/shard/mod.rs +++ b/pgdog/src/backend/pool/shard/mod.rs @@ -167,6 +167,11 @@ impl Shard { /// We don't need it for this shard. pub(super) fn schema_not_needed(&self) { let _ = self.schema.set(Schema::default()); + self.skip_loading_oids() + } + + /// Skip loading this shard's type information + pub(super) fn skip_loading_oids(&self) { self.oids.skip_load(); } diff --git a/pgdog/src/frontend/router/parser/context.rs b/pgdog/src/frontend/router/parser/context.rs index 650c891d6..89443faff 100644 --- a/pgdog/src/frontend/router/parser/context.rs +++ b/pgdog/src/frontend/router/parser/context.rs @@ -113,4 +113,8 @@ impl<'a> QueryParserContext<'a> { pub(super) fn is_session_mode(&self) -> bool { self.router_context.cluster.pooler_mode() == crate::config::PoolerMode::Session } + + pub(super) fn is_canonicalizing_oids(&self) -> bool { + self.router_context.cluster.is_canonicalizing_oids() + } } diff --git a/pgdog/src/frontend/router/parser/query/mod.rs b/pgdog/src/frontend/router/parser/query/mod.rs index c31c5f6ff..71e9f12d1 100644 --- a/pgdog/src/frontend/router/parser/query/mod.rs +++ b/pgdog/src/frontend/router/parser/query/mod.rs @@ -333,7 +333,7 @@ impl QueryParser { } Node::SelectStmt(stmt) => { - if references_pg_type(stmt) { + if context.is_canonicalizing_oids() && references_pg_type(stmt) { // Shard 0 is considered the canonical source for now return Ok(Command::Query(Route::read( ShardWithPriority::new_override_canonical_schema_info(Shard::Direct(0)), From 22eb27d420f5dfdf52ce93220afcd45fbd9fcb17 Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Tue, 4 Aug 2026 10:04:23 -0600 Subject: [PATCH 18/21] Final bits of cleanup --- .schema/pgdog.schema.json | 6 ++++++ pgdog/src/backend/pool/cluster.rs | 6 ++++++ pgdog/src/frontend/router/parser/context.rs | 1 + .../src/frontend/router/parser/query/test/test_select.rs | 9 ++------- 4 files changed, 15 insertions(+), 7 deletions(-) diff --git a/.schema/pgdog.schema.json b/.schema/pgdog.schema.json index 745cb39d2..00344aa37 100644 --- a/.schema/pgdog.schema.json +++ b/.schema/pgdog.schema.json @@ -31,6 +31,7 @@ "ban_timeout": 300000, "broadcast_address": null, "broadcast_port": 6433, + "canonicalize_type_information": false, "checkout_timeout": 5000, "client_connection_recovery": "drop", "client_idle_in_transaction_timeout": 9223372036854775807, @@ -634,6 +635,11 @@ "maximum": 65535, "minimum": 0 }, + "canonicalize_type_information": { + "description": "Controls whether PgDog maps each shard's type information to a canonical source.\n\n_Default:_ `false`\n\n", + "type": "boolean", + "default": false + }, "checkout_timeout": { "description": "Maximum amount of time a client is allowed to wait for a connection from the pool.\n\n_Default:_ `5000`\n\n", "type": "integer", diff --git a/pgdog/src/backend/pool/cluster.rs b/pgdog/src/backend/pool/cluster.rs index 807c52ed7..5d6f29e0f 100644 --- a/pgdog/src/backend/pool/cluster.rs +++ b/pgdog/src/backend/pool/cluster.rs @@ -687,6 +687,7 @@ impl Cluster { Ok(()) } + #[cfg(feature = "new_parser")] pub(crate) fn is_canonicalizing_oids(&self) -> bool { self.canonical_oids.is_some() } @@ -825,6 +826,11 @@ mod test { 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), + canonical_oids: config + .config + .general + .canonicalize_type_information + .then(|| Default::default()), ..Default::default() } } diff --git a/pgdog/src/frontend/router/parser/context.rs b/pgdog/src/frontend/router/parser/context.rs index 89443faff..250dbb166 100644 --- a/pgdog/src/frontend/router/parser/context.rs +++ b/pgdog/src/frontend/router/parser/context.rs @@ -114,6 +114,7 @@ impl<'a> QueryParserContext<'a> { self.router_context.cluster.pooler_mode() == crate::config::PoolerMode::Session } + #[cfg(feature = "new_parser")] pub(super) fn is_canonicalizing_oids(&self) -> bool { self.router_context.cluster.is_canonicalizing_oids() } diff --git a/pgdog/src/frontend/router/parser/query/test/test_select.rs b/pgdog/src/frontend/router/parser/query/test/test_select.rs index b1ad26cc5..d06c1c1ea 100644 --- a/pgdog/src/frontend/router/parser/query/test/test_select.rs +++ b/pgdog/src/frontend/router/parser/query/test/test_select.rs @@ -247,9 +247,9 @@ fn test_system_catalog_sharded() { let mut updated = config().deref().clone(); updated.config.general.system_catalogs = SystemCatalogsBehavior::Sharded; - config::set(updated).unwrap(); + updated.config.general.canonicalize_type_information = true; - let mut test = QueryParserTest::new_with_config(&config()); + let mut test = QueryParserTest::new_with_config(&updated); let command = test.execute(vec![Query::new("SELECT * FROM pg_class").into()]); assert_eq!( @@ -291,11 +291,6 @@ fn test_system_catalog_sharded() { "to_regtype should go to shard 0", ); } - - // Reset to default - let mut updated = config().deref().clone(); - updated.config.general.system_catalogs = SystemCatalogsBehavior::default(); - config::set(updated).unwrap(); } #[test] From 50639618dcb00801188a7726b135e5bff31dbff7 Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Tue, 4 Aug 2026 12:23:14 -0600 Subject: [PATCH 19/21] Make clippy happy --- pgdog/src/backend/pool/cluster.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgdog/src/backend/pool/cluster.rs b/pgdog/src/backend/pool/cluster.rs index 5d6f29e0f..4839433c8 100644 --- a/pgdog/src/backend/pool/cluster.rs +++ b/pgdog/src/backend/pool/cluster.rs @@ -830,7 +830,7 @@ mod test { .config .general .canonicalize_type_information - .then(|| Default::default()), + .then(Default::default), ..Default::default() } } From 4e7abeda45888ab96cab314ec64e70a7209556a3 Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Tue, 4 Aug 2026 12:54:32 -0600 Subject: [PATCH 20/21] Don't hold server while awaiting canonical OIDs In theory this could deadlock if the pool size were 1 and the first attempt to load the canonical OIDs failed for some reason --- pgdog/src/backend/pool/cluster/schema_loader.rs | 2 +- pgdog/src/backend/pool/shard/oids.rs | 16 +++++++++++----- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/pgdog/src/backend/pool/cluster/schema_loader.rs b/pgdog/src/backend/pool/cluster/schema_loader.rs index 4919f6b20..79f4dd254 100644 --- a/pgdog/src/backend/pool/cluster/schema_loader.rs +++ b/pgdog/src/backend/pool/cluster/schema_loader.rs @@ -28,7 +28,7 @@ impl SchemaLoader for FromServer { // For now we treat shard 0 as the canonical OID source if let Some(shard) = cluster.shards().first() - && let Some(canonical_oids) = dbg!(&cluster.canonical_oids) + && let Some(canonical_oids) = &cluster.canonical_oids { let canonical_oids = Arc::clone(canonical_oids); let shard = shard.clone(); diff --git a/pgdog/src/backend/pool/shard/oids.rs b/pgdog/src/backend/pool/shard/oids.rs index 403a8084a..282bd1d56 100644 --- a/pgdog/src/backend/pool/shard/oids.rs +++ b/pgdog/src/backend/pool/shard/oids.rs @@ -4,6 +4,7 @@ use crate::{ net::DataRow, sync::SetOnceCell, }; +use futures::{prelude::*, try_join}; use std::collections::HashMap; use std::sync::Arc; use tracing::info; @@ -26,9 +27,12 @@ impl Oids { pub(crate) async fn load(&self, shard: &Shard) -> Result<&OidMappings, Error> { self.mappings .get_or_try_init(|| async { - let mut server = shard.primary_or_replica(&Request::default()).await?; - let oids = load_oids(&mut server).await?; - let canonical = self.canonical_oids.oids.wait().await; + let oids = async { + let mut server = shard.primary_or_replica(&Request::default()).await?; + Ok::<_, Error>((server.addr().clone(), load_oids(&mut server).await?)) + }; + let canonical = self.canonical_oids.oids.wait().map(Ok); + let ((server_addr, oids), canonical) = try_join!(oids, canonical)?; let mut canonical_to_shard = HashMap::new(); let mut shard_to_canonical = HashMap::new(); for (type_name, oid) in oids { @@ -48,7 +52,7 @@ impl Oids { "loaded type info for {} types on shard {} [{}]", canonical_to_shard.len(), shard.number(), - server.addr(), + server_addr, ); Ok(OidMappings { @@ -115,7 +119,9 @@ impl CanonicalOids { } } -async fn load_oids(server: &mut Server) -> Result, Error> { +async fn load_oids( + server: &mut Server, +) -> Result + use<>, Error> { // OIDs < 10,000 are reserved for PG's internal use and are assumed to be stable Ok(server .fetch_all::( From 4ef9db1695f65aa6441db7ccf50b581343ceb1e9 Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Tue, 4 Aug 2026 13:07:16 -0600 Subject: [PATCH 21/21] Remove the fanciness This caused an obscure lifetime related compiler error that appears to come from deep within the async/await plumbing that I can't make sense of --- pgdog/src/backend/pool/shard/oids.rs | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/pgdog/src/backend/pool/shard/oids.rs b/pgdog/src/backend/pool/shard/oids.rs index 282bd1d56..4805716ca 100644 --- a/pgdog/src/backend/pool/shard/oids.rs +++ b/pgdog/src/backend/pool/shard/oids.rs @@ -4,7 +4,6 @@ use crate::{ net::DataRow, sync::SetOnceCell, }; -use futures::{prelude::*, try_join}; use std::collections::HashMap; use std::sync::Arc; use tracing::info; @@ -27,12 +26,12 @@ impl Oids { pub(crate) async fn load(&self, shard: &Shard) -> Result<&OidMappings, Error> { self.mappings .get_or_try_init(|| async { - let oids = async { - let mut server = shard.primary_or_replica(&Request::default()).await?; - Ok::<_, Error>((server.addr().clone(), load_oids(&mut server).await?)) - }; - let canonical = self.canonical_oids.oids.wait().map(Ok); - let ((server_addr, oids), canonical) = try_join!(oids, canonical)?; + let mut server = shard.primary_or_replica(&Request::default()).await?; + let oids = load_oids(&mut server).await?; + let server_addr = server.addr().clone(); + drop(server); + + let canonical = self.canonical_oids.oids.wait().await; let mut canonical_to_shard = HashMap::new(); let mut shard_to_canonical = HashMap::new(); for (type_name, oid) in oids {