diff --git a/integration/rust/tests/integration/cross_shard_oid_drift.rs b/integration/rust/tests/integration/cross_shard_oid_drift.rs index eef84028c..a0cfd055f 100644 --- a/integration/rust/tests/integration/cross_shard_oid_drift.rs +++ b/integration/rust/tests/integration/cross_shard_oid_drift.rs @@ -9,6 +9,13 @@ struct Composite { b: String, } +#[derive(sqlx::Type, Debug, Clone, PartialEq)] +#[sqlx(type_name = "test_oid_drift_later_composite")] +struct LaterComposite { + a: String, + b: String, +} + #[tokio::test] async fn test_oid_drift() { let conn = connections_sqlx().await.pop().unwrap(); @@ -74,3 +81,85 @@ async fn test_oid_drift() { admin.execute("RELOAD").await.unwrap(); } + +/// A type created after PgDog loaded its OID mappings (here: with DDL reloads +/// disabled) is resolved on first use, so the very first query using it +/// already gets the canonical OID. +#[tokio::test] +async fn test_oid_drift_type_created_later() { + let conn = connections_sqlx().await.pop().unwrap(); + let admin = admin_sqlx().await; + + admin + .execute("SET canonicalize_type_information TO true") + .await + .unwrap(); + admin + .execute("SET reload_schema_on_ddl TO false") + .await + .unwrap(); + // Make sure the mappings are loaded before the type exists. + admin.execute("RELOAD").await.unwrap(); + conn.execute("SELECT 1").await.unwrap(); + + conn.execute("DROP TABLE IF EXISTS test_oid_drift_later") + .await + .unwrap(); + conn.execute("DROP TYPE IF EXISTS test_oid_drift_later_composite CASCADE") + .await + .unwrap(); + // Intentionally cause the OID of the type to differ between shards + conn.execute("/* pgdog_shard: 1 */ CREATE SEQUENCE foo; DROP SEQUENCE foo;") + .await + .unwrap(); + conn.execute("CREATE TYPE test_oid_drift_later_composite AS (a text, b text)") + .await + .unwrap(); + conn.execute( + "CREATE TABLE test_oid_drift_later (customer_id BIGINT, composite test_oid_drift_later_composite)", + ) + .await + .unwrap(); + + let expected_oid: Oid = sqlx::query_scalar( + "SELECT oid FROM pg_type WHERE typname = 'test_oid_drift_later_composite'", + ) + .fetch_one(&conn) + .await + .unwrap(); + + // The client learned the type's OID from shard 0 and sends it in Parse; + // shards where the OID differs must accept it on the first try. + let composite = LaterComposite { + a: String::from("a"), + b: String::from("b"), + }; + for i in 1..=20_i64 { + sqlx::query("INSERT INTO test_oid_drift_later VALUES ($1, $2)") + .bind(i) + .bind(&composite) + .execute(&conn) + .await + .unwrap(); + } + + // Reads from every shard, including the ones where the OID differs, + // must not fail even once. + for customer_id in 1..=20_i64 { + let rows = sqlx::query("SELECT composite FROM test_oid_drift_later WHERE customer_id = $1") + .bind(customer_id) + .fetch_all(&conn) + .await + .unwrap(); + assert_eq!(rows.len(), 1, "customer {customer_id}"); + assert_eq!( + rows[0].column(0).type_info().oid(), + Some(expected_oid), + "customer {customer_id}" + ); + let decoded: LaterComposite = rows[0].get(0); + assert_eq!(decoded, composite); + } + + admin.execute("RELOAD").await.unwrap(); +} diff --git a/pgdog/src/backend/error.rs b/pgdog/src/backend/error.rs index 6ebd312f0..9cd4ee9e8 100644 --- a/pgdog/src/backend/error.rs +++ b/pgdog/src/backend/error.rs @@ -24,6 +24,9 @@ pub(crate) enum Error { #[error("unexpected transaction status: {0}")] UnexpectedTransactionStatus(char), + #[error("type information is not available")] + TypeInfoUnavailable, + #[error("{0}")] ConnectionError(Box), @@ -116,9 +119,6 @@ pub(crate) 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/pool/cluster/schema_loader.rs b/pgdog/src/backend/pool/cluster/schema_loader.rs index 79f4dd254..25794848f 100644 --- a/pgdog/src/backend/pool/cluster/schema_loader.rs +++ b/pgdog/src/backend/pool/cluster/schema_loader.rs @@ -35,11 +35,7 @@ impl SchemaLoader for FromServer { tasks::spawn("load canonical oids", async move { loop { let result = tasks::shutdown_signal() - .run_until_cancelled(async { - canonical_oids - .load(&mut *shard.primary_or_replica(&Default::default()).await?) - .await - }) + .run_until_cancelled(async { canonical_oids.load(&shard).await }) .await; match result { diff --git a/pgdog/src/backend/pool/connection/binding.rs b/pgdog/src/backend/pool/connection/binding.rs index 4c32e8868..0999b5a02 100644 --- a/pgdog/src/backend/pool/connection/binding.rs +++ b/pgdog/src/backend/pool/connection/binding.rs @@ -468,6 +468,23 @@ impl Binding { } } + /// Any held server returned a type the canonical OID + /// mappings don't know about. Resets the flags. + pub(crate) fn take_oids_stale(&mut self) -> bool { + match self { + Binding::Direct(server, ..) => server.take_oids_stale(), + Binding::MultiShard(servers, _) => { + // Reset every server's flag, not just the first stale one. + let mut stale = false; + for server in servers.iter_mut() { + stale |= server.take_oids_stale(); + } + stale + } + _ => false, + } + } + pub(super) fn dirty(&mut self) { match self { Binding::Direct(server, ..) => server.mark_dirty(true), diff --git a/pgdog/src/backend/pool/mod.rs b/pgdog/src/backend/pool/mod.rs index da0bd3897..6a3203406 100644 --- a/pgdog/src/backend/pool/mod.rs +++ b/pgdog/src/backend/pool/mod.rs @@ -42,7 +42,7 @@ pub(crate) use password::Password; pub(crate) use pool_impl::Pool; pub(crate) use request::Request; pub(crate) use role::PoolRole; -pub(crate) use shard::{CanonicalOids, Oids, Shard}; +pub(crate) use shard::{CanonicalOids, OidMappings, Oids, Shard}; pub(crate) use state::State; pub(crate) use stats::Stats; diff --git a/pgdog/src/backend/pool/shard/mod.rs b/pgdog/src/backend/pool/shard/mod.rs index 341481612..671289e77 100644 --- a/pgdog/src/backend/pool/shard/mod.rs +++ b/pgdog/src/backend/pool/shard/mod.rs @@ -29,7 +29,7 @@ pub(crate) mod role_detector; use failover_signal::{FailoverSignal, FailoverSignalWatcher}; use monitor::*; -pub(crate) use oids::{CanonicalOids, Oids}; +pub(crate) use oids::{CanonicalOids, OidMappings, Oids}; use role_detector::*; #[cfg_attr(test, derive(Default))] diff --git a/pgdog/src/backend/pool/shard/oids.rs b/pgdog/src/backend/pool/shard/oids.rs index 4805716ca..c4795b6bc 100644 --- a/pgdog/src/backend/pool/shard/oids.rs +++ b/pgdog/src/backend/pool/shard/oids.rs @@ -1,18 +1,42 @@ -use super::{Request, Shard}; +//! Canonical type OID mappings. +//! +//! Types created with `CREATE TYPE` (or by extensions) get a different OID +//! on each shard. Clients cache type information by OID, so PgDog presents +//! shard 0's OIDs to clients and translates them on the way to and from +//! the other shards. +//! +//! Types created after the mappings were loaded are detected the first time +//! they appear in a message, and the mappings are refreshed from the shards +//! before the message is forwarded. + +use super::{Request, Shard, ShardInner}; use crate::{ backend::{Error, Server}, net::DataRow, sync::SetOnceCell, }; -use std::collections::HashMap; -use std::sync::Arc; -use tracing::info; +use parking_lot::{RwLock, RwLockReadGuard}; +use std::collections::{HashMap, HashSet}; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, OnceLock, Weak}; +use tokio::sync::Mutex; +use tracing::{info, warn}; + +/// OIDs below this are assigned by Postgres at bootstrap/initdb +/// and are assumed to be identical across shards. +pub(crate) const FIRST_USER_OID: u32 = 10000; #[derive(Debug)] /// The mapping from a shards type OID to a canonical one pub(crate) struct Oids { canonical_oids: Arc, - mappings: SetOnceCell, + mappings: SetOnceCell>, + /// The shard these mappings belong to, used to refresh them. + shard: OnceLock>, + /// Serialize refreshes. + refresh_lock: Mutex<()>, + /// A type couldn't be resolved and a full reload was requested. + stale: AtomicBool, } impl Oids { @@ -20,46 +44,77 @@ impl Oids { Arc::new(Self { canonical_oids: Arc::clone(canonical_oids), mappings: Default::default(), + shard: OnceLock::new(), + refresh_lock: Mutex::new(()), + stale: AtomicBool::new(false), }) } - pub(crate) async fn load(&self, shard: &Shard) -> Result<&OidMappings, Error> { + pub(crate) async fn load(&self, shard: &Shard) -> Result<(), Error> { + let _ = self.shard.set(Arc::downgrade(&shard.inner)); + 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 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 { - 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, - }) - }) + .get_or_try_init(|| async { Ok(RwLock::new(self.fetch(shard).await?)) }) .await + .map(|_| ()) + } + + /// Fetch the shard's types and build the mappings against the canonical set. + async fn fetch(&self, shard: &Shard) -> Result { + 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.read(); + let mappings = OidMappings::build(oids, &canonical, |type_name| { + warn!( + "type {} on shard {} [{}] doesn't exist on shard 0, its OID won't be canonicalized", + type_name, + shard.number(), + server_addr, + ) + }); + + info!( + "loaded type info for {} types on shard {} [{}]", + mappings.canonical_to_shard.len(), + shard.number(), + server_addr, + ); + + Ok(mappings) + } + + /// Reload the canonical set and this shard's types. + /// + /// Called when a type unknown to the mappings shows up, e.g. after + /// `CREATE TYPE` ran on another PgDog or directly on the database. + /// `still_unknown` is re-checked under the refresh lock, so concurrent + /// detections of the same type result in a single refresh. + pub(crate) async fn refresh( + &self, + still_unknown: impl Fn(&OidMappings) -> bool, + ) -> Result<(), Error> { + let _lock = self.refresh_lock.lock().await; + + let mappings = self.mappings.get().ok_or(Error::TypeInfoUnavailable)?; + if !still_unknown(&mappings.read()) { + return Ok(()); + } + + let shard = self + .shard + .get() + .and_then(Weak::upgrade) + .map(|inner| Shard { inner }) + .ok_or(Error::TypeInfoUnavailable)?; + + self.canonical_oids.refresh().await?; + let refreshed = self.fetch(&shard).await?; + *mappings.write() = refreshed; + + Ok(()) } pub(crate) async fn wait(&self) { @@ -72,20 +127,50 @@ impl Oids { } /// Get the mappings. Returns `None` if no mappings have been loaded - pub(crate) fn get(&self) -> Option<&OidMappings> { - self.mappings.get() + pub(crate) fn get(&self) -> Option> { + self.mappings.get().map(RwLock::read) + } + + /// Stop trying to resolve these OIDs; they have no counterpart on the other side. + pub(crate) fn mark_unresolvable(&self, oids: impl IntoIterator) { + if let Some(mappings) = self.mappings.get() { + mappings.write().unresolvable.extend(oids); + } + } + + /// Record that a type couldn't be resolved, requesting a full reload. + /// + /// Returns `true` only the first time, so a single reload is requested. + pub(crate) fn mark_stale(&self) -> bool { + !self.stale.swap(true, Ordering::AcqRel) } #[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 { + let this = Self { canonical_oids: Default::default(), - mappings: SetOnceCell::from(OidMappings { - canonical_to_shard, - shard_to_canonical, - }), - }) + mappings: SetOnceCell::from(RwLock::new(OidMappings::default())), + shard: OnceLock::new(), + refresh_lock: Mutex::new(()), + stale: AtomicBool::new(false), + }; + this.set_canonical(canonical_to_shard); + Arc::new(this) + } + + /// Replace the mappings, simulating a refresh. + #[cfg(test)] + pub(crate) fn set_canonical(&self, canonical_to_shard: HashMap) { + let shard_to_canonical: HashMap<_, _> = + canonical_to_shard.iter().map(|(&k, &v)| (v, k)).collect(); + let mappings = OidMappings { + known: shard_to_canonical.keys().copied().collect(), + canonical_to_shard, + shard_to_canonical, + loaded: true, + ..Default::default() + }; + *self.mappings.get().expect("test mappings are set").write() = mappings; } } @@ -93,7 +178,10 @@ impl Default for Oids { fn default() -> Self { Self { canonical_oids: Default::default(), - mappings: SetOnceCell::from(OidMappings::default()), + mappings: SetOnceCell::from(RwLock::new(OidMappings::default())), + shard: OnceLock::new(), + refresh_lock: Mutex::new(()), + stale: AtomicBool::new(false), } } } @@ -102,30 +190,120 @@ impl Default for Oids { pub(crate) struct OidMappings { pub(crate) canonical_to_shard: HashMap, pub(crate) shard_to_canonical: HashMap, + /// Every user-defined type OID present on the shard when the mappings were loaded. + known: HashSet, + /// OIDs we tried to resolve and couldn't, e.g. the type exists on one shard only. + unresolvable: HashSet, + /// Mappings were actually loaded from the shard, as opposed to skipped. + loaded: bool, +} + +impl OidMappings { + /// Build the mappings from the shard's types and the canonical set. + /// + /// Types that don't exist on the canonical shard are reported to `missing` + /// and left unmapped, instead of failing the whole shard. + fn build( + oids: impl Iterator, + canonical: &HashMap, + mut missing: impl FnMut(&str), + ) -> Self { + let mut mappings = Self { + loaded: true, + ..Default::default() + }; + + for (type_name, oid) in oids { + mappings.known.insert(oid); + let Some(&canonical) = canonical.get(&type_name) else { + missing(&type_name); + continue; + }; + if canonical == oid { + continue; + } + mappings.canonical_to_shard.insert(canonical, oid); + mappings.shard_to_canonical.insert(oid, canonical); + } + + debug_assert_eq!( + mappings.canonical_to_shard.len(), + mappings.shard_to_canonical.len() + ); + + mappings + } + + /// The mappings were loaded from the shard and can be used + /// to detect types created since. + pub(crate) fn loaded(&self) -> bool { + self.loaded + } + + fn candidate(&self, oid: u32) -> bool { + self.loaded && oid >= FIRST_USER_OID && !self.unresolvable.contains(&oid) + } + + /// An OID sent by the shard (e.g. in RowDescription) that belongs to a type + /// which didn't exist on the shard when the mappings were loaded. + pub(crate) fn is_unknown_shard_oid(&self, oid: u32) -> bool { + self.candidate(oid) && !self.known.contains(&oid) + } + + /// An OID sent by the client (e.g. in Parse) that neither maps to a type + /// on this shard nor exists on it unchanged. + pub(crate) fn is_unknown_canonical_oid(&self, oid: u32) -> bool { + self.candidate(oid) + && !self.canonical_to_shard.contains_key(&oid) + && !self.known.contains(&oid) + } } #[derive(Debug, Default)] pub(crate) struct CanonicalOids { - oids: SetOnceCell>, + oids: SetOnceCell>>, + /// The canonical shard, used to refresh the set. + shard: OnceLock>, } impl CanonicalOids { - pub(crate) async fn load(&self, server: &mut Server) -> Result<(), Error> { + pub(crate) async fn load(&self, shard: &Shard) -> Result<(), Error> { + let _ = self.shard.set(Arc::downgrade(&shard.inner)); + self.oids - .get_or_try_init(|| async { Ok(load_oids(server).await?.collect()) }) + .get_or_try_init(|| async { Ok(RwLock::new(Self::fetch(shard).await?)) }) .await .map(|_| ()) } + + async fn fetch(shard: &Shard) -> Result, Error> { + let mut server = shard.primary_or_replica(&Request::default()).await?; + Ok(load_oids(&mut server).await?.collect()) + } + + /// Reload the canonical set from the canonical shard. + async fn refresh(&self) -> Result<(), Error> { + let oids = self.oids.get().ok_or(Error::TypeInfoUnavailable)?; + let shard = self + .shard + .get() + .and_then(Weak::upgrade) + .map(|inner| Shard { inner }) + .ok_or(Error::TypeInfoUnavailable)?; + + *oids.write() = Self::fetch(&shard).await?; + + Ok(()) + } } 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::( - "SELECT nspname || '.' || typname, pg_type.oid FROM pg_type INNER JOIN pg_namespace ON typnamespace = pg_namespace.oid WHERE pg_type.oid >= 10000", - ) + .fetch_all::(&format!( + "SELECT nspname || '.' || typname, pg_type.oid FROM pg_type INNER JOIN pg_namespace ON typnamespace = pg_namespace.oid WHERE pg_type.oid >= {FIRST_USER_OID}", + )) .await? .into_iter() .map(|row| { @@ -135,3 +313,111 @@ async fn load_oids( ) })) } + +#[cfg(test)] +mod test { + use super::*; + + fn canonical() -> HashMap { + [ + ("public.mood", 16400), + ("public._mood", 16399), + ("public.same", 16500), + ] + .into_iter() + .map(|(name, oid)| (name.to_owned(), oid)) + .collect() + } + + fn build(shard: &[(&str, u32)], missing: &mut Vec) -> OidMappings { + OidMappings::build( + shard.iter().map(|(name, oid)| (name.to_string(), *oid)), + &canonical(), + |name| missing.push(name.to_owned()), + ) + } + + #[test] + fn test_build_maps_drifted_types_only() { + let mut missing = vec![]; + let mappings = build( + &[ + ("public.mood", 17000), + ("public._mood", 16999), + ("public.same", 16500), + ], + &mut missing, + ); + + assert!(missing.is_empty()); + assert_eq!(mappings.shard_to_canonical[&17000], 16400); + assert_eq!(mappings.shard_to_canonical[&16999], 16399); + assert_eq!(mappings.canonical_to_shard[&16400], 17000); + assert!(!mappings.shard_to_canonical.contains_key(&16500)); + assert!(mappings.loaded()); + } + + #[test] + fn test_build_skips_types_missing_on_canonical_shard() { + let mut missing = vec![]; + let mappings = build( + &[("public.mood", 17000), ("public.only_here", 18000)], + &mut missing, + ); + + assert_eq!(missing, vec!["public.only_here"]); + assert_eq!(mappings.shard_to_canonical.len(), 1); + // Known to the shard, so it won't trigger a refresh. + assert!(!mappings.is_unknown_shard_oid(18000)); + } + + #[test] + fn test_unknown_shard_oid() { + let mappings = build( + &[("public.mood", 17000), ("public.same", 16500)], + &mut vec![], + ); + + assert!(!mappings.is_unknown_shard_oid(17000)); + assert!(!mappings.is_unknown_shard_oid(16500)); + assert!( + !mappings.is_unknown_shard_oid(25), + "built-in types are never unknown" + ); + assert!(mappings.is_unknown_shard_oid(17001)); + + assert!( + !OidMappings::default().is_unknown_shard_oid(17001), + "skipped mappings can't detect anything" + ); + } + + #[test] + fn test_unknown_canonical_oid() { + let mappings = build( + &[("public.mood", 17000), ("public.same", 16500)], + &mut vec![], + ); + + assert!(!mappings.is_unknown_canonical_oid(16400), "mapped"); + assert!(!mappings.is_unknown_canonical_oid(16500), "same on both"); + assert!(!mappings.is_unknown_canonical_oid(25)); + assert!(mappings.is_unknown_canonical_oid(16401)); + } + + #[test] + fn test_unresolvable() { + let oids = Oids::from_canonical([(16400, 17000)].into_iter().collect()); + assert!(oids.get().unwrap().is_unknown_shard_oid(18000)); + oids.mark_unresolvable([18000]); + assert!(!oids.get().unwrap().is_unknown_shard_oid(18000)); + assert!(!oids.get().unwrap().is_unknown_canonical_oid(18000)); + } + + #[test] + fn test_mark_stale_once() { + let oids = Oids::from_canonical([(16400, 17000)].into_iter().collect()); + assert!(oids.mark_stale()); + assert!(!oids.mark_stale()); + } +} diff --git a/pgdog/src/backend/prepared_statements.rs b/pgdog/src/backend/prepared_statements.rs index 14a93270d..2cf3747d9 100644 --- a/pgdog/src/backend/prepared_statements.rs +++ b/pgdog/src/backend/prepared_statements.rs @@ -107,9 +107,37 @@ pub(crate) struct PreparedStatements { config: PreparedStatementsConfig, memory_used: usize, oids: Arc, + /// A message from the server is waiting for the OID mappings to be refreshed. + pending_oids: Option, + /// A type couldn't be resolved and the config should be reloaded. + oids_stale: bool, server_state: State, } +/// A server message that references types unknown to the OID mappings, +/// held back until the mappings are refreshed. +#[derive(Debug)] +enum PendingOidRewrite { + RowDescription { + /// Statement the RowDescription belongs to, if we're caching it. + describe: Option, + unknown: Vec, + }, + ParameterDescription { + unknown: Vec, + }, +} + +impl PendingOidRewrite { + fn unknown(&self) -> &[u32] { + match self { + Self::RowDescription { unknown, .. } | Self::ParameterDescription { unknown } => { + unknown + } + } + } +} + #[cfg(test)] impl Default for PreparedStatements { fn default() -> Self { @@ -129,6 +157,8 @@ impl PreparedStatements { config: PreparedStatementsConfig::default(), memory_used: 0, oids, + pending_oids: None, + oids_stale: false, server_state: State::Idle, } } @@ -406,13 +436,7 @@ impl PreparedStatements { } 'T' => { - let maybe_row_description = self.parse_and_rewrite_row_description(message)?; - if let Some(describe) = self.describes.pop_front() { - let row_description = maybe_row_description - .map(Ok) - .unwrap_or_else(|| RowDescription::from_bytes(message.payload()))?; - self.add_row_description(&describe, row_description); - }; + self.handle_row_description(message)?; } // No data for DELETEs @@ -448,7 +472,7 @@ impl PreparedStatements { } 't' => { - self.rewrite_parameter_description_data_types(message)?; + self.handle_parameter_description(message)?; } _ => (), @@ -626,6 +650,12 @@ impl PreparedStatements { self.oids = Arc::clone(oids) } + /// The server returned a type the OID mappings don't know about, + /// so they need to be reloaded. Resets the flag. + pub(crate) fn take_oids_stale(&mut self) -> bool { + std::mem::take(&mut self.oids_stale) + } + fn rewrite_parse_data_types(&self, parse: &mut Parse) -> bool { let Some(mappings) = self.oids.get() else { return false; @@ -633,42 +663,195 @@ impl PreparedStatements { parse.rewrite_data_types(&mappings.canonical_to_shard) } - /// 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; + /// Canonical OIDs the client is about to send in the given message + /// that this shard's mappings don't know about. + /// + /// The server would reject them, so they have to be resolved first. + pub(crate) fn unknown_canonical_oids(&self, message: &ProtocolMessage) -> Vec { + let Some(mappings) = self.oids.get().filter(|mappings| mappings.loaded()) else { + return vec![]; + }; + + let unknown = |parse: &Parse| -> Vec { + parse + .data_types() + .filter(|&oid| mappings.is_unknown_canonical_oid(oid)) + .collect() + }; - 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()); + match message { + ProtocolMessage::Parse(parse) => unknown(parse), + // Statements not yet prepared on this connection get their + // Parse prepended from the global cache. + ProtocolMessage::Bind(bind) if !self.local_cache.contains(bind.statement()) => self + .global_cache + .read() + .rewritten_parse(bind.statement()) + .map(|parse| unknown(&parse)) + .unwrap_or_default(), + ProtocolMessage::Describe(describe) + if describe.is_statement() && !self.local_cache.contains(describe.statement()) => + { + self.global_cache + .read() + .rewritten_parse(describe.statement()) + .map(|parse| unknown(&parse)) + .unwrap_or_default() } - Ok(Some(row_description)) - } else { - Ok(None) + _ => vec![], } } - fn rewrite_parameter_description_data_types(&self, message: &mut Message) -> Result<(), Error> { - let Some(mappings) = self.oids.get() else { + /// Rewrite a RowDescription to have the canonical set of OIDs and cache it + /// for the statement being described, if any. + /// + /// If it references types the mappings don't know about, hold it back + /// (see [`Self::pending_unknown_oids`]) until they are refreshed. + fn handle_row_description(&mut self, message: &mut Message) -> Result<(), Error> { + let describe = self.describes.pop_front(); + + // RowDescription is emitted during cluster startup, so we can't + // require OIDs to be loaded. + let mappings = self + .oids + .get() + .filter(|mappings| mappings.loaded() || !mappings.shard_to_canonical.is_empty()); + + let Some(mappings) = mappings else { + if let Some(describe) = describe { + let row_description = RowDescription::from_bytes(message.payload())?; + self.add_row_description(&describe, row_description); + } return Ok(()); }; - let mappings = &mappings.shard_to_canonical; - if mappings.is_empty() { + + let mut row_description = RowDescription::from_bytes(message.payload())?; + let unknown: Vec = row_description + .iter() + .map(|field| field.type_oid as u32) + .filter(|&oid| mappings.is_unknown_shard_oid(oid)) + .collect(); + + if !unknown.is_empty() { + drop(mappings); + self.pending_oids = Some(PendingOidRewrite::RowDescription { describe, unknown }); return Ok(()); } + if row_description.rewrite_data_types(&mappings.shard_to_canonical) { + message.replace_payload(row_description.to_bytes()); + } + drop(mappings); + + if let Some(describe) = describe { + self.add_row_description(&describe, row_description); + } + + Ok(()) + } + + /// Rewrite a ParameterDescription to have the canonical set of OIDs, + /// holding it back if it references unknown types. + fn handle_parameter_description(&mut self, message: &mut Message) -> Result<(), Error> { + let Some(mappings) = self.oids.get().filter(|mappings| mappings.loaded()) else { + return Ok(()); + }; + let mut parameter_description = ParameterDescription::from_bytes(message.payload())?; - parameter_description.rewrite_data_types(mappings); - message.replace_payload(parameter_description.to_bytes()); + let unknown: Vec = parameter_description + .data_types() + .filter(|&oid| mappings.is_unknown_shard_oid(oid)) + .collect(); + + if !unknown.is_empty() { + drop(mappings); + self.pending_oids = Some(PendingOidRewrite::ParameterDescription { unknown }); + return Ok(()); + } + + if parameter_description.rewrite_data_types(&mappings.shard_to_canonical) { + message.replace_payload(parameter_description.to_bytes()); + } + + Ok(()) + } + + /// OIDs in the last message from the server that the mappings don't know about. + /// The message must not be forwarded until [`Self::finish_oid_rewrite`] is called. + pub(crate) fn pending_unknown_oids(&self) -> Option<&[u32]> { + self.pending_oids.as_ref().map(PendingOidRewrite::unknown) + } + + /// Forward the held back message as is. Used for PgDog's own queries, + /// which don't need canonical OIDs. + pub(crate) fn discard_pending_oids(&mut self) { + self.pending_oids = None; + } + + /// Rewrite the held back message with the refreshed mappings. + /// + /// Types that still can't be resolved are given up on, and a config + /// reload is requested as a last resort. + pub(crate) fn finish_oid_rewrite(&mut self, message: &mut Message) -> Result<(), Error> { + let Some(pending) = self.pending_oids.take() else { + return Ok(()); + }; + + let Some(mappings) = self.oids.get() else { + return Ok(()); + }; + + match pending { + PendingOidRewrite::RowDescription { describe, .. } => { + let mut row_description = RowDescription::from_bytes(message.payload())?; + let unresolved: Vec = row_description + .iter() + .map(|field| field.type_oid as u32) + .filter(|&oid| mappings.is_unknown_shard_oid(oid)) + .collect(); + if row_description.rewrite_data_types(&mappings.shard_to_canonical) { + message.replace_payload(row_description.to_bytes()); + } + drop(mappings); + self.give_up_on(unresolved); + if let Some(describe) = describe { + self.add_row_description(&describe, row_description); + } + } + + PendingOidRewrite::ParameterDescription { .. } => { + let mut parameter_description = + ParameterDescription::from_bytes(message.payload())?; + let unresolved: Vec = parameter_description + .data_types() + .filter(|&oid| mappings.is_unknown_shard_oid(oid)) + .collect(); + if parameter_description.rewrite_data_types(&mappings.shard_to_canonical) { + message.replace_payload(parameter_description.to_bytes()); + } + drop(mappings); + self.give_up_on(unresolved); + } + } + Ok(()) } + + /// Stop trying to resolve these OIDs and request a config reload instead. + pub(crate) fn give_up_on(&mut self, unresolved: Vec) { + if unresolved.is_empty() { + return; + } + self.oids.mark_unresolvable(unresolved); + if self.oids.mark_stale() { + self.oids_stale = true; + } + } + + /// The OID mappings shared by all connections to this shard. + pub(crate) fn oids(&self) -> &Arc { + &self.oids + } } #[cfg(test)] @@ -1319,6 +1502,150 @@ pub(crate) mod test { ); } + fn row_description_with_oid(type_oid: i32) -> RowDescription { + RowDescription::new(&[crate::net::messages::Field { + name: "mood".into(), + table_oid: 0, + column: 0, + type_oid, + type_size: 4, + type_modifier: -1, + format: 0, + }]) + } + + /// Describe a statement that's already prepared on the connection and + /// forward the ParameterDescription, leaving the RowDescription expected. + fn describe_prepared(ps: &mut PreparedStatements, name: &str) { + ps.prepared(name); + ps.handle(&ProtocolMessage::Describe(Describe::new_statement(name))) + .unwrap(); + let mut params = Message::new(ParameterDescription::empty().to_bytes()); + assert!(ps.forward(&mut params).unwrap()); + } + + fn type_oid(message: &Message) -> i32 { + RowDescription::from_bytes(message.payload()) + .unwrap() + .field(0) + .unwrap() + .type_oid + } + + #[test] + fn row_description_rewrites_known_shard_oid() { + let mut ps = new_extended(); + ps.oids = Oids::from_canonical([(16400, 17000)].into_iter().collect()); + let name = insert_global("known_oid", "SELECT mood FROM t"); + describe_prepared(&mut ps, &name); + + let mut message = Message::new(row_description_with_oid(17000).to_bytes()); + assert!(ps.forward(&mut message).unwrap()); + assert_eq!(type_oid(&message), 16400); + assert!(ps.pending_unknown_oids().is_none()); + } + + #[test] + fn row_description_with_unknown_oid_is_held_until_refresh() { + let mut ps = new_extended(); + ps.oids = Oids::from_canonical([(16400, 17000)].into_iter().collect()); + let name = insert_global("unknown_oid", "SELECT mood FROM t"); + describe_prepared(&mut ps, &name); + + // A type created after the mappings were loaded. + let mut message = Message::new(row_description_with_oid(17001).to_bytes()); + ps.forward(&mut message).unwrap(); + assert_eq!(ps.pending_unknown_oids(), Some(&[17001][..])); + assert_eq!(type_oid(&message), 17001, "not rewritten yet"); + assert!( + FrontendPreparedStatements::global() + .read() + .row_description(&name) + .is_none(), + "not cached with the shard's OID" + ); + + // Refresh happened (simulated), the message can be finished. + ps.oids + .set_canonical([(16400, 17000), (16401, 17001)].into_iter().collect()); + ps.finish_oid_rewrite(&mut message).unwrap(); + assert!(ps.pending_unknown_oids().is_none()); + assert_eq!(type_oid(&message), 16401); + assert_eq!( + FrontendPreparedStatements::global() + .read() + .row_description(&name) + .unwrap() + .field(0) + .unwrap() + .type_oid, + 16401, + "cached with the canonical OID" + ); + assert!(!ps.take_oids_stale()); + } + + #[test] + fn row_description_with_unresolvable_oid_requests_reload_once() { + let mut ps = new_extended(); + ps.oids = Oids::from_canonical([(16400, 17000)].into_iter().collect()); + let name = insert_global("unresolvable_oid", "SELECT mood FROM t"); + describe_prepared(&mut ps, &name); + + let mut message = Message::new(row_description_with_oid(17001).to_bytes()); + ps.forward(&mut message).unwrap(); + assert!(ps.pending_unknown_oids().is_some()); + + // Refresh didn't help. + ps.finish_oid_rewrite(&mut message).unwrap(); + assert_eq!(type_oid(&message), 17001, "forwarded as is"); + assert!(ps.take_oids_stale(), "config reload requested"); + assert!(!ps.take_oids_stale(), "flag was reset"); + + // Not held back again. + describe_prepared(&mut ps, &name); + let mut message = Message::new(row_description_with_oid(17001).to_bytes()); + ps.forward(&mut message).unwrap(); + assert!(ps.pending_unknown_oids().is_none()); + assert!(!ps.take_oids_stale(), "reload requested only once"); + } + + #[test] + fn unknown_canonical_oids_in_parse_and_bind() { + let mut ps = new_extended(); + ps.oids = Oids::from_canonical([(16400, 17000)].into_iter().collect()); + + let parse = Parse::named("stmt1", "SELECT $1, $2, $3"); + let mapped_same_unknown = parse.with_data_types(&[16400, 17000, 16401]); + assert_eq!( + ps.unknown_canonical_oids(&ProtocolMessage::Parse(mapped_same_unknown)), + vec![16401] + ); + + let builtins = parse.with_data_types(&[25, 20, 16]); + assert!( + ps.unknown_canonical_oids(&ProtocolMessage::Parse(builtins)) + .is_empty() + ); + + // Bind for a statement not yet prepared on this connection checks + // the Parse that will be prepended. + let cached = Parse::named("unknown_bind", "SELECT $1").with_data_types(&[16401]); + let (_, name) = FrontendPreparedStatements::global().write().insert(&cached); + let bind = Bind::new_statement(&name); + assert_eq!( + ps.unknown_canonical_oids(&ProtocolMessage::Bind(bind.clone())), + vec![16401] + ); + + // Already prepared: the Parse was sent (and checked) before. + ps.prepared(&name); + assert!( + ps.unknown_canonical_oids(&ProtocolMessage::Bind(bind)) + .is_empty() + ); + } + // ------------------------------------------------------- // Simple query is unaffected by mode // ------------------------------------------------------- diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index 981b14979..c06a59012 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -14,7 +14,7 @@ use tracing::{debug, error, info, trace, warn}; use super::{ ConnectReason, DisconnectReason, Error, Oids, PreparedStatements, ServerOptions, Stats, - pool::Address, + pool::{Address, OidMappings}, prepared_statements::{HandleResult, Prepare}, }; use crate::{ @@ -42,7 +42,10 @@ use crate::{ tls::UpstreamTlsSettings, }, }; -use crate::{net::tweak, state::State}; +use crate::{net::tweak, state::State, tasks, util::safe_timeout}; + +/// How long a message is held back while type OID mappings are refreshed. +const OID_REFRESH_TIMEOUT: Duration = Duration::from_secs(5); /// A request executed on a server connection: simple-protocol queries, /// or an extended-protocol message batch ending in a Sync. @@ -148,6 +151,8 @@ pub(crate) struct Server { sending_request: bool, pooler_mode: PoolerMode, stream_buffer: MessageBuffer, + /// Message held back while the type OID mappings are refreshed. + pending_message: Option, disconnect_reason: Option, password_attempts: usize, /// Per-connection lifetime cap. When `Some`, the connection @@ -436,6 +441,7 @@ impl Server { sending_request: false, pooler_mode: PoolerMode::Transaction, stream_buffer: MessageBuffer::new(config.config.memory.message_buffer, None), + pending_message: None, disconnect_reason: None, password_attempts: 1, // This is going to be changed by parent caller. max_age: None, @@ -471,6 +477,7 @@ impl Server { } for message in client_request.messages.iter() { + self.resolve_unknown_canonical_oids(message).await; self.send_one(message).await?; } self.flush().await?; @@ -484,6 +491,47 @@ impl Server { Ok(()) } + /// Send PgDog's own messages to the server. Unlike [`Self::send`], this + /// doesn't resolve unknown type OIDs, so it can be used while resolving them. + async fn send_internal(&mut self, messages: &[ProtocolMessage]) -> Result<(), Error> { + self.sending_request = true; + self.stats.state(State::Active); + + for message in messages { + self.send_one(message).await?; + } + self.flush().await?; + + self.sending_request = false; + self.stats.state(State::ReceivingData); + + Ok(()) + } + + /// The client is using types this shard's OID mappings don't know about. + /// Refresh them before the statement is sent with the wrong OIDs. + async fn resolve_unknown_canonical_oids(&mut self, message: &ProtocolMessage) { + let unknown = self.prepared_statements.unknown_canonical_oids(message); + if unknown.is_empty() { + return; + } + + self.refresh_oids(unknown.clone(), OidMappings::is_unknown_canonical_oid) + .await; + + let unresolved = unknown + .into_iter() + .filter(|&oid| { + self.prepared_statements + .oids() + .get() + .map(|mappings| mappings.is_unknown_canonical_oid(oid)) + .unwrap_or(false) + }) + .collect(); + self.prepared_statements.give_up_on(unresolved); + } + /// Send one message to the server but don't flush the buffer, /// accelerating bulk transfers. pub(crate) async fn send_one(&mut self, message: &ProtocolMessage) -> Result<(), Error> { @@ -567,7 +615,34 @@ impl Server { /// This method is cancel-safe. /// pub(crate) async fn read(&mut self) -> Result { - let message = loop { + // A previous read was cancelled while refreshing OID mappings. + if let Some(message) = self.pending_message.take() { + let message = self.resolve_pending_oids(message).await?; + return self.received(message); + } + + let message = self.read_message().await?; + + let message = if self.prepared_statements.pending_unknown_oids().is_some() { + self.resolve_pending_oids(message).await? + } else { + message + }; + + self.received(message) + } + + /// Read a message in response to PgDog's own queries. Unlike [`Self::read`], + /// this doesn't resolve unknown type OIDs, so it can be used while resolving them. + async fn read_internal(&mut self) -> Result { + let message = self.read_message().await?; + self.prepared_statements.discard_pending_oids(); + self.received(message) + } + + /// Read the next message from the stream, or a simulated one. + async fn read_message(&mut self) -> Result { + Ok(loop { if let Some(message) = self.prepared_statements.state_mut().get_simulated() { // INVARIANT: omni dedup in multi_shard relies on this being process-unique; // never substitute a non-unique value here. @@ -606,8 +681,74 @@ impl Server { return Err(err.into()); } } - }; + }) + } + + /// The message references types unknown to the OID mappings. Refresh them + /// from the shards and rewrite the message before anyone else sees it. + /// + /// The message is parked in `self.pending_message` while we wait, so a + /// cancelled read doesn't lose it. + async fn resolve_pending_oids(&mut self, message: Message) -> Result { + self.pending_message = Some(message); + + if let Some(unknown) = self.prepared_statements.pending_unknown_oids() { + let unknown = unknown.to_vec(); + self.refresh_oids(unknown, OidMappings::is_unknown_shard_oid) + .await; + } + + let mut message = self + .pending_message + .take() + .expect("pending message set above"); + self.prepared_statements.finish_oid_rewrite(&mut message)?; + Ok(message) + } + + /// Refresh the OID mappings so the given OIDs become known. + /// Errors are logged, not returned: the caller falls back to a config reload. + async fn refresh_oids(&self, unknown: Vec, is_unknown: fn(&OidMappings, u32) -> bool) { + let oids = Arc::clone(self.prepared_statements.oids()); + let still_unknown = unknown.clone(); + + // Spawned: refreshing checks out a connection, which may open one, + // which sends messages through `send_one`. Awaiting it inline would + // make this future recursive. + let refresh = tasks::spawn("refresh type oids", async move { + oids.refresh(move |mappings| still_unknown.iter().any(|&oid| is_unknown(mappings, oid))) + .await + }); + + match safe_timeout(OID_REFRESH_TIMEOUT, refresh).await { + Ok(Ok(Ok(()))) => debug!( + "refreshed type info for unknown type oid(s) {:?} [{}]", + unknown, + self.addr() + ), + Ok(Ok(Err(err))) => warn!( + "failed to refresh type info for unknown type oid(s) {:?} [{}]: {}", + unknown, + self.addr(), + err + ), + Ok(Err(err)) => warn!( + "failed to refresh type info for unknown type oid(s) {:?} [{}]: {}", + unknown, + self.addr(), + err + ), + Err(_) => warn!( + "timed out refreshing type info for unknown type oid(s) {:?} [{}]", + unknown, + self.addr() + ), + } + } + + /// Account for a message received from the server and track transaction state. + fn received(&mut self, message: Message) -> Result { self.stats.receive(message.len(), message.code() as u8); match message.code() { @@ -784,6 +925,12 @@ impl Server { self.changed_params.clear(); } + /// The server returned a type the canonical OID mappings don't know about. + /// Resets the flag. + pub(crate) fn take_oids_stale(&mut self) -> bool { + self.prepared_statements.take_oids_stale() + } + /// We can disconnect from this server. /// /// There are no more expected messages from the server connection @@ -905,11 +1052,11 @@ impl Server { let mut messages = vec![]; let expected = request.expected; - self.send(&request.messages.into()).await?; + self.send_internal(&request.messages).await?; let mut zs = 0; while zs < expected { - let message = self.read().await?; + let message = self.read_internal().await?; if message.code() == 'Z' { zs += 1; } @@ -1004,14 +1151,14 @@ impl Server { /// attempting to return the connection into a synchronized state. pub(super) async fn drain(&mut self) -> Result<(), Error> { while self.has_more_messages() { - self.read().await?; + self.read_internal().await?; } if !self.in_sync() { - self.send(&vec![ProtocolMessage::Sync(Sync)].into()).await?; + self.send_internal(&[ProtocolMessage::Sync(Sync)]).await?; while !self.in_sync() { - self.read().await?; + self.read_internal().await?; } } @@ -1351,6 +1498,7 @@ pub(crate) mod test { id, params: Parameters::default(), changed_params: Parameters::default(), + pending_message: None, client_params: Parameters::default(), stats: Stats::connect( id, diff --git a/pgdog/src/frontend/client/query_engine/query.rs b/pgdog/src/frontend/client/query_engine/query.rs index aa88e8f99..d2832b7e7 100644 --- a/pgdog/src/frontend/client/query_engine/query.rs +++ b/pgdog/src/frontend/client/query_engine/query.rs @@ -330,6 +330,7 @@ impl QueryEngine { ) -> Result<(), Error> { if self.backend.done() { let changed_params = self.backend.changed_params(); + let oids_stale = self.backend.take_oids_stale(); // Release the connection back into the pool before flushing data to client. // Flushing can take a minute and we don't want to block the connection from being reused. @@ -350,6 +351,20 @@ impl QueryEngine { self.backend.cluster()?.identifier(), ); schema_changed()?; + } else if oids_stale + && self + .backend + .cluster() + .map(|cluster| cluster.is_canonicalizing_oids()) + .unwrap_or_default() + { + // A type was created behind our back (e.g. DDL that ran + // on another PgDog or directly on the database). + info!( + "unknown type detected, reloading config [{}]", + self.backend.cluster()?.identifier(), + ); + schema_changed()?; } self.router.reset(); diff --git a/pgdog/src/frontend/router/parser/query/ddl.rs b/pgdog/src/frontend/router/parser/query/ddl.rs index 17982084b..29099f260 100644 --- a/pgdog/src/frontend/router/parser/query/ddl.rs +++ b/pgdog/src/frontend/router/parser/query/ddl.rs @@ -1,4 +1,5 @@ use crate::frontend::client::query_engine::TempTableChange; +use pg_raw_parse::list::CastNodeList; use pg_raw_parse::raw::OnCommitAction::ONCOMMIT_DROP; use std::ffi::c_char; @@ -74,6 +75,13 @@ impl QueryParser { } } + // Dropped types leave stale entries in the canonical OID mappings. + ObjectType::OBJECT_TYPE + | ObjectType::OBJECT_DOMAIN + | ObjectType::OBJECT_EXTENSION => { + schema_changed = true; + } + _ => (), }, @@ -110,15 +118,36 @@ impl QueryParser { } } + // Creating a type assigns it a fresh OID on every shard, so the + // canonical OID mappings have to be reloaded. Node::CreateEnumStmt(stmt) => { - let table = Table::try_from(stmt.type_name()).ok(); - if let Some(table) = table { - shard = schema - .schemas - .get(table.schema()) - .map(|schema| schema.shard().into()) - .unwrap_or(Shard::All); - } + schema_changed = true; + shard = Self::shard_ddl_type_name(stmt.type_name(), schema); + } + + Node::CreateRangeStmt(stmt) => { + schema_changed = true; + shard = Self::shard_ddl_type_name(stmt.type_name(), schema); + } + + Node::CreateDomainStmt(stmt) => { + schema_changed = true; + shard = Self::shard_ddl_type_name(stmt.domainname(), schema); + } + + Node::CompositeTypeStmt(stmt) => { + schema_changed = true; + shard = Self::shard_ddl_table(stmt.typevar(), schema)?.unwrap_or(Shard::All); + } + + Node::DefineStmt(stmt) if stmt.kind == ObjectType::OBJECT_TYPE => { + schema_changed = true; + shard = Self::shard_ddl_type_name(stmt.defnames(), schema); + } + + // Extensions typically create types. + Node::CreateExtensionStmt(_) => { + schema_changed = true; } Node::AlterOwnerStmt(stmt) => { @@ -226,6 +255,19 @@ impl QueryParser { )) } + /// Shard for a type created in a (possibly sharded) schema, + /// given the qualified name of the type. + fn shard_ddl_type_name( + type_name: &CastNodeList, + schema: &ShardingSchema, + ) -> Shard { + Table::try_from(type_name) + .ok() + .and_then(|table| schema.schemas.get(table.schema())) + .map(|schema| schema.shard().into()) + .unwrap_or(Shard::All) + } + pub(super) fn shard_ddl_table( range_var: Option<&nodes::RangeVar>, schema: &ShardingSchema, @@ -459,14 +501,14 @@ mod test { fn test_create_enum_sharded() { let command = parse_stmt("CREATE TYPE shard_1.mood AS ENUM ('sad', 'ok', 'happy')"); assert_eq!(command.route().shard(), &Shard::Direct(1)); - assert!(!command.route().is_schema_changed()); + assert!(command.route().is_schema_changed()); } #[test] fn test_create_enum_unsharded() { let command = parse_stmt("CREATE TYPE public.mood AS ENUM ('sad', 'ok', 'happy')"); assert_eq!(command.route().shard(), &Shard::All); - assert!(!command.route().is_schema_changed()); + assert!(command.route().is_schema_changed()); } #[test] @@ -581,4 +623,47 @@ mod test { assert_eq!(command.route().shard(), &Shard::All); assert!(!command.route().is_schema_changed()); } + + #[test] + fn test_create_type_changes_schema() { + for query in [ + "CREATE TYPE mood AS ENUM ('sad', 'ok', 'happy')", + "CREATE TYPE complex AS (r double precision, i double precision)", + "CREATE TYPE floatrange AS RANGE (subtype = float8)", + "CREATE TYPE box_t (INPUT = box_in, OUTPUT = box_out)", + "CREATE DOMAIN posint AS integer CHECK (VALUE > 0)", + "CREATE EXTENSION IF NOT EXISTS vector", + "DROP TYPE mood", + "DROP TYPE IF EXISTS mood, complex CASCADE", + "DROP DOMAIN posint", + "DROP EXTENSION vector", + ] { + let command = parse_stmt(query); + assert_eq!(command.route().shard(), &Shard::All, "{query}"); + assert!(command.route().is_schema_changed(), "{query}"); + } + } + + #[test] + fn test_create_type_sharded_schema() { + for query in [ + "CREATE TYPE shard_1.mood AS ENUM ('sad', 'ok', 'happy')", + "CREATE TYPE shard_1.complex AS (r double precision, i double precision)", + "CREATE TYPE shard_1.floatrange AS RANGE (subtype = float8)", + "CREATE TYPE shard_1.box_t (INPUT = box_in, OUTPUT = box_out)", + "CREATE DOMAIN shard_1.posint AS integer", + ] { + let command = parse_stmt(query); + assert_eq!(command.route().shard(), &Shard::Direct(1), "{query}"); + assert!(command.route().is_schema_changed(), "{query}"); + } + } + + #[test] + fn test_alter_type_add_value_keeps_schema() { + // Adding a value doesn't change the type's OID. + let command = parse_stmt("ALTER TYPE mood ADD VALUE 'ecstatic'"); + assert_eq!(command.route().shard(), &Shard::All); + assert!(!command.route().is_schema_changed()); + } } diff --git a/pgdog/src/net/messages/parameter_description.rs b/pgdog/src/net/messages/parameter_description.rs index f329e2c38..7f9bece12 100644 --- a/pgdog/src/net/messages/parameter_description.rs +++ b/pgdog/src/net/messages/parameter_description.rs @@ -44,12 +44,22 @@ impl ParameterDescription { Self { params: Vec::new() } } - pub(crate) fn rewrite_data_types(&mut self, mapping: &HashMap) { + /// Parameter data type OIDs. + pub(crate) fn data_types(&self) -> impl Iterator + '_ { + self.params.iter().map(|¶m| param as u32) + } + + /// Replaces the data types of each parameter using the given mapping. + /// Returns whether any changes actually occurred. + pub(crate) fn rewrite_data_types(&mut self, mapping: &HashMap) -> bool { + let mut changed = false; for param in &mut self.params { if let Some(&canonical) = mapping.get(&(*param as u32)) { + changed = true; *param = canonical as i32; } } + changed } } diff --git a/pgdog/src/net/messages/parse.rs b/pgdog/src/net/messages/parse.rs index aec7ad895..6482929be 100644 --- a/pgdog/src/net/messages/parse.rs +++ b/pgdog/src/net/messages/parse.rs @@ -112,6 +112,13 @@ impl Parse { self.data_types.clone() } + /// Parameter data type OIDs, as sent by the client. + pub(crate) fn data_types(&self) -> impl Iterator + '_ { + let mut bytes = self.data_types.clone(); + let num = bytes.get_u16(); + (0..num).map(move |_| bytes.get_u32()) + } + /// Update the SQL for this prepared statement. pub(crate) fn set_query(&mut self, query: &str) { self.query = c_string_bytes(query);