diff --git a/integration/rust/tests/integration/cross_shard_oid_drift.rs b/integration/rust/tests/integration/cross_shard_oid_drift.rs index eef84028c..49d796296 100644 --- a/integration/rust/tests/integration/cross_shard_oid_drift.rs +++ b/integration/rust/tests/integration/cross_shard_oid_drift.rs @@ -74,3 +74,78 @@ async fn test_oid_drift() { admin.execute("RELOAD").await.unwrap(); } + +#[derive(sqlx::Type, Debug, Clone, Copy, PartialEq)] +#[sqlx(type_name = "test_oid_drift_mood", rename_all = "lowercase")] +enum Mood { + Sad, + Ok, + Happy, +} + +/// Binary arrays embed the element type's OID, so arrays of custom types +/// have to be rewritten in both directions, not just the RowDescription. +#[tokio::test] +async fn test_oid_drift_arrays() { + let conn = connections_sqlx().await.pop().unwrap(); + let admin = admin_sqlx().await; + + conn.execute("DROP TABLE IF EXISTS test_oid_drift_arrays") + .await + .unwrap(); + conn.execute("DROP TYPE IF EXISTS test_oid_drift_mood 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_mood AS ENUM ('sad', 'ok', 'happy')") + .await + .unwrap(); + conn.execute( + "CREATE TABLE test_oid_drift_arrays (customer_id BIGINT, moods test_oid_drift_mood[])", + ) + .await + .unwrap(); + admin + .execute("SET canonicalize_type_information TO true") + .await + .unwrap(); + admin.execute("RELOAD").await.unwrap(); + + let canonical_oid: Oid = + sqlx::query_scalar("SELECT oid FROM pg_type WHERE typname = 'test_oid_drift_mood'") + .fetch_one(&conn) + .await + .unwrap(); + + let moods = vec![Mood::Sad, Mood::Ok, Mood::Happy]; + for i in 1..=20_i64 { + // Binary array parameter, element OID as learned from shard 0. + sqlx::query("INSERT INTO test_oid_drift_arrays VALUES ($1, $2)") + .bind(i) + .bind(&moods) + .execute(&conn) + .await + .unwrap(); + } + + for customer_id in 1..=20_i64 { + let row = sqlx::query("SELECT moods FROM test_oid_drift_arrays WHERE customer_id = $1") + .bind(customer_id) + .fetch_one(&conn) + .await + .unwrap(); + + // The element OID inside the binary array payload is shard 0's. + let raw = row.try_get_raw(0).unwrap().as_bytes().unwrap().to_vec(); + let element_oid = u32::from_be_bytes([raw[8], raw[9], raw[10], raw[11]]); + assert_eq!(element_oid, canonical_oid.0, "customer {customer_id}"); + + let decoded: Vec = row.get(0); + assert_eq!(decoded, moods); + } + + admin.execute("RELOAD").await.unwrap(); +} diff --git a/pgdog/src/backend/pool/mod.rs b/pgdog/src/backend/pool/mod.rs index da0bd3897..e38dc12f7 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, Oids, PayloadRewriter, 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..4784cdd64 100644 --- a/pgdog/src/backend/pool/shard/mod.rs +++ b/pgdog/src/backend/pool/shard/mod.rs @@ -29,7 +29,9 @@ pub(crate) mod role_detector; use failover_signal::{FailoverSignal, FailoverSignalWatcher}; use monitor::*; -pub(crate) use oids::{CanonicalOids, Oids}; +#[cfg(test)] +pub(crate) use oids::TypeKind; +pub(crate) use oids::{CanonicalOids, Oids, PayloadRewriter}; 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..79c5fbc91 100644 --- a/pgdog/src/backend/pool/shard/oids.rs +++ b/pgdog/src/backend/pool/shard/oids.rs @@ -1,3 +1,10 @@ +//! 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. + use super::{Request, Shard}; use crate::{ backend::{Error, Server}, @@ -8,6 +15,60 @@ use std::collections::HashMap; use std::sync::Arc; use tracing::info; +mod payload; +pub(crate) use payload::PayloadRewriter; + +/// What a type's binary representation looks like, as far as +/// embedded type OIDs are concerned. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum TypeKind { + /// Array: binary values carry the element type OID. + Array { element: u32 }, + /// Composite: binary values carry the OID of every field. + Composite, + /// Domain: encoded like its base type. + Domain { base: u32 }, + /// Everything else: no embedded OIDs. + Other, +} + +impl TypeKind { + /// From `pg_type` columns. + fn from_catalog(typtype: &str, typcategory: &str, typelem: u32, typbasetype: u32) -> Self { + match (typtype, typcategory) { + (_, "A") if typelem != 0 => Self::Array { element: typelem }, + ("c", _) => Self::Composite, + ("d", _) if typbasetype != 0 => Self::Domain { base: typbasetype }, + _ => Self::Other, + } + } +} + +/// A type on the shard: `schema.name`, its OID and kind. +type TypeRow = (String, u32, TypeKind); + +/// The canonical set of types, from shard 0. +#[derive(Debug, Default)] +pub(crate) struct CanonicalTypes { + by_name: HashMap, + kinds: Arc>, +} + +impl FromIterator for CanonicalTypes { + fn from_iter>(iter: I) -> Self { + let mut by_name = HashMap::new(); + let mut kinds = HashMap::new(); + for (name, oid, kind) in iter { + by_name.insert(name, oid); + kinds.insert(oid, kind); + } + Self { + by_name, + kinds: Arc::new(kinds), + } + } +} + #[derive(Debug)] /// The mapping from a shards type OID to a canonical one pub(crate) struct Oids { @@ -34,8 +95,11 @@ impl Oids { 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 mut shard_kinds = HashMap::new(); + for (type_name, oid, kind) in oids { + shard_kinds.insert(oid, kind); let canonical = canonical + .by_name .get(&type_name) .copied() .ok_or(Error::MissingCanonicalOid(type_name))?; @@ -57,6 +121,8 @@ impl Oids { Ok(OidMappings { canonical_to_shard, shard_to_canonical, + shard_kinds, + canonical_kinds: Arc::clone(&canonical.kinds), }) }) .await @@ -78,12 +144,24 @@ impl Oids { #[cfg(test)] pub(crate) fn from_canonical(canonical_to_shard: HashMap) -> Arc { + Self::from_canonical_with_kinds(canonical_to_shard, HashMap::new(), HashMap::new()) + } + + /// Mappings plus the type kinds on both sides, simulating what's loaded from `pg_type`. + #[cfg(test)] + pub(crate) fn from_canonical_with_kinds( + canonical_to_shard: HashMap, + shard_kinds: HashMap, + canonical_kinds: 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, + shard_kinds, + canonical_kinds: Arc::new(canonical_kinds), }), }) } @@ -102,11 +180,40 @@ impl Default for Oids { pub(crate) struct OidMappings { pub(crate) canonical_to_shard: HashMap, pub(crate) shard_to_canonical: HashMap, + /// Kinds of the shard's types, by shard OID. + shard_kinds: HashMap, + /// Kinds of the canonical types, by canonical OID. + canonical_kinds: Arc>, +} + +impl OidMappings { + /// Whether canonicalization has anything to do on this shard at all. + pub(crate) fn is_identity(&self) -> bool { + self.shard_to_canonical.is_empty() + } + + /// The shard's OID for a canonical type OID. + pub(crate) fn shard_oid(&self, canonical: u32) -> u32 { + self.canonical_to_shard + .get(&canonical) + .copied() + .unwrap_or(canonical) + } + + /// Rewriter for values coming from the shard (DataRow). + pub(crate) fn to_canonical(&self) -> PayloadRewriter<'_> { + PayloadRewriter::new(&self.shard_kinds, &self.shard_to_canonical) + } + + /// Rewriter for values going to the shard (Bind parameters). + pub(crate) fn to_shard(&self) -> PayloadRewriter<'_> { + PayloadRewriter::new(&self.canonical_kinds, &self.canonical_to_shard) + } } #[derive(Debug, Default)] pub(crate) struct CanonicalOids { - oids: SetOnceCell>, + oids: SetOnceCell, } impl CanonicalOids { @@ -118,20 +225,51 @@ impl CanonicalOids { } } -async fn load_oids( - server: &mut Server, -) -> Result + use<>, 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::( - "SELECT nspname || '.' || typname, pg_type.oid FROM pg_type INNER JOIN pg_namespace ON typnamespace = pg_namespace.oid WHERE pg_type.oid >= 10000", + "SELECT nspname || '.' || typname, pg_type.oid, typtype::text, typcategory::text, typelem, typbasetype \ + FROM pg_type INNER JOIN pg_namespace ON typnamespace = pg_namespace.oid \ + WHERE pg_type.oid >= 10000", ) .await? .into_iter() .map(|row| { + let name = row.get_text(0).expect("selected 6 columns"); + let oid = row.get_int(1, true).expect("selected 6 columns") as u32; + let typtype = row.get_text(2).expect("selected 6 columns"); + let typcategory = row.get_text(3).expect("selected 6 columns"); + let typelem = row.get_int(4, true).expect("selected 6 columns") as u32; + let typbasetype = row.get_int(5, true).expect("selected 6 columns") as u32; ( - row.get_text(0).expect("selected 2 columns"), - row.get_int(1, true).expect("selected 2 columns") as u32, + name, + oid, + TypeKind::from_catalog(&typtype, &typcategory, typelem, typbasetype), ) })) } + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn test_type_kind_from_catalog() { + assert_eq!( + TypeKind::from_catalog("b", "A", 16400, 0), + TypeKind::Array { element: 16400 } + ); + assert_eq!(TypeKind::from_catalog("c", "C", 0, 0), TypeKind::Composite); + assert_eq!( + TypeKind::from_catalog("d", "N", 0, 23), + TypeKind::Domain { base: 23 } + ); + assert_eq!(TypeKind::from_catalog("e", "E", 0, 0), TypeKind::Other); + // A domain over an array is category A but not itself an array. + assert_eq!( + TypeKind::from_catalog("d", "A", 16400, 16399), + TypeKind::Array { element: 16400 } + ); + } +} diff --git a/pgdog/src/backend/pool/shard/oids/payload.rs b/pgdog/src/backend/pool/shard/oids/payload.rs new file mode 100644 index 000000000..633133825 --- /dev/null +++ b/pgdog/src/backend/pool/shard/oids/payload.rs @@ -0,0 +1,358 @@ +//! Rewrite type OIDs embedded in binary-format values. +//! +//! Arrays carry their element type OID and composites (including anonymous +//! records) carry the OID of every field. Clients check those against the +//! type information they cached, so they have to be canonicalized too. + +use std::collections::HashMap; + +use bytes::{Buf, BufMut}; + +use super::TypeKind; + +/// `record`, the anonymous composite type. +const RECORD_OID: u32 = 2249; +/// `record[]`. +const RECORD_ARRAY_OID: u32 = 2287; + +/// How deep nested arrays/composites are followed before giving up. +const MAX_DEPTH: usize = 16; + +/// Malformed binary value; leave it alone. +#[derive(Debug, PartialEq)] +pub(crate) struct Malformed; + +/// Rewrites embedded OIDs of one direction (shard to canonical, or the reverse). +pub(crate) struct PayloadRewriter<'a> { + /// Kinds of the types on the source side, keyed by their OID. + kinds: &'a HashMap, + /// Source to destination OIDs. + mapping: &'a HashMap, +} + +impl<'a> PayloadRewriter<'a> { + pub(crate) fn new(kinds: &'a HashMap, mapping: &'a HashMap) -> Self { + Self { kinds, mapping } + } + + /// The kind of a type, looking through domains. + fn kind(&self, mut oid: u32) -> TypeKind { + for _ in 0..MAX_DEPTH { + match oid { + RECORD_OID => return TypeKind::Composite, + RECORD_ARRAY_OID => { + return TypeKind::Array { + element: RECORD_OID, + }; + } + _ => (), + } + match self.kinds.get(&oid) { + Some(TypeKind::Domain { base }) => oid = *base, + Some(kind) => return *kind, + None => return TypeKind::Other, + } + } + TypeKind::Other + } + + /// A binary value of this type may embed type OIDs. + pub(crate) fn needs_rewrite(&self, oid: u32) -> bool { + matches!(self.kind(oid), TypeKind::Array { .. } | TypeKind::Composite) + } + + /// Rewrite the OIDs embedded in a binary value of the given type, in place. + /// Returns whether anything changed. + pub(crate) fn rewrite(&self, oid: u32, data: &mut [u8]) -> Result { + self.rewrite_at(oid, data, 0).map(|(changed, _)| changed) + } + + /// Rewrite a value and return how many bytes it occupied. + fn rewrite_at( + &self, + oid: u32, + data: &mut [u8], + depth: usize, + ) -> Result<(bool, usize), Malformed> { + if depth > MAX_DEPTH { + return Err(Malformed); + } + + match self.kind(oid) { + TypeKind::Array { .. } => self.rewrite_array(data, depth), + TypeKind::Composite => self.rewrite_composite(data, depth), + TypeKind::Domain { .. } | TypeKind::Other => Ok((false, data.len())), + } + } + + /// Big-endian `i32` at `pos`. + fn read_i32(data: &[u8], pos: usize) -> Result { + data.get(pos..pos + 4) + .map(|mut bytes| bytes.get_i32()) + .ok_or(Malformed) + } + + /// Replace the OID stored at `pos` with its mapping, if any. + fn map_oid(&self, data: &mut [u8], pos: usize) -> Result<(u32, bool), Malformed> { + let slot = data.get_mut(pos..pos + 4).ok_or(Malformed)?; + let oid = (&slot[..]).get_u32(); + match self.mapping.get(&oid) { + Some(&mapped) => { + (&mut slot[..]).put_u32(mapped); + Ok((oid, true)) + } + None => Ok((oid, false)), + } + } + + fn rewrite_array(&self, data: &mut [u8], depth: usize) -> Result<(bool, usize), Malformed> { + if data.len() < 12 { + return Err(Malformed); + } + let ndim = Self::read_i32(data, 0)?; + if !(0..=6).contains(&ndim) { + return Err(Malformed); + } + let (element, mut changed) = self.map_oid(data, 8)?; + + let mut pos = 12; + let mut elements: usize = 1; + for _ in 0..ndim { + let size = Self::read_i32(data, pos)?; + if size < 0 { + return Err(Malformed); + } + elements = elements.checked_mul(size as usize).ok_or(Malformed)?; + pos += 8; + } + if ndim == 0 { + elements = 0; + } + + let recurse = self.needs_rewrite(element); + for _ in 0..elements { + let len = Self::read_i32(data, pos)?; + pos += 4; + if len < 0 { + continue; + } + let len = len as usize; + let value = data.get_mut(pos..pos + len).ok_or(Malformed)?; + if recurse { + changed |= self.rewrite_at(element, value, depth + 1)?.0; + } + pos += len; + } + + Ok((changed, pos)) + } + + fn rewrite_composite(&self, data: &mut [u8], depth: usize) -> Result<(bool, usize), Malformed> { + let fields = Self::read_i32(data, 0)?; + if fields < 0 { + return Err(Malformed); + } + + let mut changed = false; + let mut pos = 4; + for _ in 0..fields { + let (field_oid, mapped) = self.map_oid(data, pos)?; + changed |= mapped; + pos += 4; + let len = Self::read_i32(data, pos)?; + pos += 4; + if len < 0 { + continue; + } + let len = len as usize; + let value = data.get_mut(pos..pos + len).ok_or(Malformed)?; + if self.needs_rewrite(field_oid) { + changed |= self.rewrite_at(field_oid, value, depth + 1)?.0; + } + pos += len; + } + + Ok((changed, pos)) + } +} + +#[cfg(test)] +mod test { + use super::*; + use bytes::BytesMut; + + const MOOD: u32 = 17000; + const MOOD_ARRAY: u32 = 17001; + const PAIR: u32 = 17002; + const PAIR_ARRAY: u32 = 17003; + const POSINT: u32 = 17004; + + fn kinds() -> HashMap { + [ + (MOOD, TypeKind::Other), + (MOOD_ARRAY, TypeKind::Array { element: MOOD }), + (PAIR, TypeKind::Composite), + (PAIR_ARRAY, TypeKind::Array { element: PAIR }), + (POSINT, TypeKind::Domain { base: MOOD_ARRAY }), + ] + .into_iter() + .collect() + } + + fn mapping() -> HashMap { + [ + (MOOD, 16400), + (MOOD_ARRAY, 16401), + (PAIR, 16402), + (PAIR_ARRAY, 16403), + ] + .into_iter() + .collect() + } + + fn array(element: u32, values: &[Option<&[u8]>]) -> Vec { + let mut buf = BytesMut::new(); + buf.put_i32(1); + buf.put_i32(values.iter().any(Option::is_none) as i32); + buf.put_u32(element); + buf.put_i32(values.len() as i32); + buf.put_i32(1); + for value in values { + match value { + Some(value) => { + buf.put_i32(value.len() as i32); + buf.put_slice(value); + } + None => buf.put_i32(-1), + } + } + buf.to_vec() + } + + fn composite(fields: &[(u32, Option<&[u8]>)]) -> Vec { + let mut buf = BytesMut::new(); + buf.put_i32(fields.len() as i32); + for (oid, value) in fields { + buf.put_u32(*oid); + match value { + Some(value) => { + buf.put_i32(value.len() as i32); + buf.put_slice(value); + } + None => buf.put_i32(-1), + } + } + buf.to_vec() + } + + #[test] + fn test_needs_rewrite() { + let (kinds, mapping) = (kinds(), mapping()); + let rewriter = PayloadRewriter::new(&kinds, &mapping); + assert!(!rewriter.needs_rewrite(MOOD)); + assert!(!rewriter.needs_rewrite(25)); + assert!(rewriter.needs_rewrite(MOOD_ARRAY)); + assert!(rewriter.needs_rewrite(PAIR)); + assert!(rewriter.needs_rewrite(POSINT), "domain over an array"); + assert!(rewriter.needs_rewrite(RECORD_OID)); + assert!(rewriter.needs_rewrite(RECORD_ARRAY_OID)); + } + + #[test] + fn test_array_of_enum() { + let (kinds, mapping) = (kinds(), mapping()); + let rewriter = PayloadRewriter::new(&kinds, &mapping); + + let mut data = array(MOOD, &[Some(b"sad"), None, Some(b"happy")]); + assert_eq!(rewriter.rewrite(MOOD_ARRAY, &mut data), Ok(true)); + assert_eq!(data, array(16400, &[Some(b"sad"), None, Some(b"happy")])); + + // Through a domain. + let mut data = array(MOOD, &[Some(b"ok")]); + assert_eq!(rewriter.rewrite(POSINT, &mut data), Ok(true)); + assert_eq!(data, array(16400, &[Some(b"ok")])); + + // Builtin elements are left alone. + let mut data = array(25, &[Some(b"text")]); + assert_eq!(rewriter.rewrite(MOOD_ARRAY, &mut data), Ok(false)); + assert_eq!(data, array(25, &[Some(b"text")])); + } + + #[test] + fn test_empty_array() { + let (kinds, mapping) = (kinds(), mapping()); + let rewriter = PayloadRewriter::new(&kinds, &mapping); + + let mut buf = BytesMut::new(); + buf.put_i32(0); + buf.put_i32(0); + buf.put_u32(MOOD); + let mut data = buf.to_vec(); + assert_eq!(rewriter.rewrite(MOOD_ARRAY, &mut data), Ok(true)); + assert_eq!((&data[8..12]).get_u32(), 16400); + } + + #[test] + fn test_composite_with_nested_array() { + let (kinds, mapping) = (kinds(), mapping()); + let rewriter = PayloadRewriter::new(&kinds, &mapping); + + let moods = array(MOOD, &[Some(b"sad")]); + let mut data = composite(&[ + (25, Some(b"name")), + (MOOD, None), + (MOOD_ARRAY, Some(&moods)), + ]); + assert_eq!(rewriter.rewrite(PAIR, &mut data), Ok(true)); + + let expected_moods = array(16400, &[Some(b"sad")]); + assert_eq!( + data, + composite(&[ + (25, Some(b"name")), + (16400, None), + (16401, Some(&expected_moods)) + ]) + ); + } + + #[test] + fn test_array_of_composites_and_records() { + let (kinds, mapping) = (kinds(), mapping()); + let rewriter = PayloadRewriter::new(&kinds, &mapping); + + let pair = composite(&[(MOOD, Some(b"ok")), (20, Some(&1i64.to_be_bytes()))]); + let mut data = array(PAIR, &[Some(&pair), Some(&pair)]); + assert_eq!(rewriter.rewrite(PAIR_ARRAY, &mut data), Ok(true)); + + let expected_pair = composite(&[(16400, Some(b"ok")), (20, Some(&1i64.to_be_bytes()))]); + assert_eq!( + data, + array(16402, &[Some(&expected_pair), Some(&expected_pair)]) + ); + + // Anonymous records carry user types too. + let mut data = composite(&[(MOOD, Some(b"ok"))]); + assert_eq!(rewriter.rewrite(RECORD_OID, &mut data), Ok(true)); + assert_eq!(data, composite(&[(16400, Some(b"ok"))])); + } + + #[test] + fn test_malformed_leaves_no_partial_writes_visible() { + let (kinds, mapping) = (kinds(), mapping()); + let rewriter = PayloadRewriter::new(&kinds, &mapping); + + let mut data = vec![0, 0, 0, 1, 0, 0]; + assert_eq!(rewriter.rewrite(MOOD_ARRAY, &mut data), Err(Malformed)); + + // Element length runs past the end. + let mut data = array(MOOD, &[Some(b"sad")]); + let len = data.len(); + data[len - 4 - 3..len - 3].copy_from_slice(&100i32.to_be_bytes()); + assert_eq!(rewriter.rewrite(MOOD_ARRAY, &mut data), Err(Malformed)); + + let mut data = composite(&[(MOOD, Some(b"ok"))]); + data.truncate(10); + assert_eq!(rewriter.rewrite(PAIR, &mut data), Err(Malformed)); + } +} diff --git a/pgdog/src/backend/prepared_statements.rs b/pgdog/src/backend/prepared_statements.rs index 14a93270d..3286dc992 100644 --- a/pgdog/src/backend/prepared_statements.rs +++ b/pgdog/src/backend/prepared_statements.rs @@ -8,8 +8,9 @@ use std::{ use crate::{ frontend::{self, prepared_statements::GlobalCache}, net::{ - Close, CloseComplete, FromBytes, Message, ParseComplete, Protocol, ProtocolMessage, - ToBytes, + Bind, Close, CloseComplete, DataRow, Format, FromBytes, Message, ParseComplete, Protocol, + ProtocolMessage, ToBytes, + bind::Parameter, messages::{ParameterDescription, RowDescription, parse::Parse}, }, state::State, @@ -17,8 +18,9 @@ use crate::{ use crate::{net::ErrorResponse, util::time::deadline}; use parking_lot::RwLock; use pgdog_config::prepared_statements::PreparedStatementsConfig; +use tracing::warn; -use super::{Error, Oids}; +use super::{Error, Oids, pool::PayloadRewriter}; use super::{ protocol::{ProtocolState, state::Action}, state::ExecutionCode, @@ -107,9 +109,45 @@ pub(crate) struct PreparedStatements { config: PreparedStatementsConfig, memory_used: usize, oids: Arc, + /// Portals bound but not yet executed. + binds: VecDeque, + /// Portals being executed; DataRows belong to the front one. + executing: VecDeque, + /// Portal Describes sent and not yet answered. + portal_describes: usize, server_state: State, } +/// Bound portals never accumulate past this, even if a client +/// keeps binding named portals without executing them. +const MAX_BOUND_PORTALS: usize = 64; + +/// A portal the client bound, tracked so binary values in its +/// rows can have their embedded type OIDs canonicalized. +#[derive(Debug, Default)] +struct BoundPortal { + portal: String, + /// Global name of the statement; empty if unknown. + statement: String, + /// Result formats requested in Bind. + formats: Vec, + /// RowDescription the server sent for this portal, with the shard's OIDs. + row_description: Option, + /// Columns whose binary values embed type OIDs: `(index, shard type OID)`. + /// Computed on the first DataRow. + plan: Option>, +} + +impl BoundPortal { + fn result_format(&self, index: usize) -> Format { + match self.formats.len() { + 0 => Format::Text, + 1 => self.formats[0], + _ => self.formats.get(index).copied().unwrap_or(Format::Text), + } + } +} + #[cfg(test)] impl Default for PreparedStatements { fn default() -> Self { @@ -129,6 +167,9 @@ impl PreparedStatements { config: PreparedStatementsConfig::default(), memory_used: 0, oids, + binds: VecDeque::new(), + executing: VecDeque::new(), + portal_describes: 0, server_state: State::Idle, } } @@ -168,41 +209,9 @@ impl PreparedStatements { pub(super) fn handle(&mut self, request: &ProtocolMessage) -> Result { match request { ProtocolMessage::Bind(bind) => { - if !bind.anonymous() { - let message = self.check_prepared(bind.statement())?; - match message { - Some(mut message) => { - if message.close.is_some() { - self.state.add_ignore('3'); - } - self.state.add_ignore('1'); - self.parses.push_back(bind.statement().to_string()); - self.state.add('2'); - if self.config.level.rewrite_anonymous() { - message.anonymize(); - let mut bind = bind.clone(); - bind.anonymize(); - return Ok(HandleResult::PrependRewrite { - prepend: message, - rewrite: ProtocolMessage::Bind(bind), - }); - } else { - return Ok(HandleResult::Prepend(message)); - } - } - - None => { - self.state.add('2'); - if self.config.level.rewrite_anonymous() { - let mut bind = bind.clone(); - bind.anonymize(); - return Ok(HandleResult::Rewrite(ProtocolMessage::Bind(bind))); - } - } - } - } else { - self.state.add('2'); - } + self.bound(bind)?; + let result = self.handle_bind(bind)?; + return self.rewrite_bind_params(bind, result); } ProtocolMessage::Describe(describe) => { if !describe.anonymous() { @@ -252,14 +261,16 @@ impl PreparedStatements { } } else if describe.is_portal() { self.state.add(ExecutionCode::DescriptionOrNothing); + self.portal_describes += 1; } else if describe.is_statement() { self.state.add(ExecutionCode::DescriptionOrNothing); // t self.state.add(ExecutionCode::DescriptionOrNothing); // T } } - ProtocolMessage::Execute(_) => { + ProtocolMessage::Execute(execute) => { self.state.add(ExecutionCode::ExecutionCompleted); + self.executed(execute.portal()); } ProtocolMessage::Sync(_) => { @@ -403,6 +414,9 @@ impl PreparedStatements { // are syntactically valid. self.describes.clear(); self.parses.clear(); + self.binds.clear(); + self.executing.clear(); + self.portal_describes = 0; } 'T' => { @@ -412,18 +426,46 @@ impl PreparedStatements { .map(Ok) .unwrap_or_else(|| RowDescription::from_bytes(message.payload()))?; self.add_row_description(&describe, row_description); - }; + } else if self.portal_describes > 0 { + // Answering a portal Describe: remember the columns for its rows. + self.portal_describes -= 1; + if let Some(portal) = self.described_portal() { + portal.row_description = maybe_row_description; + } + } + } + + 'D' => { + self.rewrite_data_row(message)?; } // No data for DELETEs 'n' => { - self.describes.pop_front(); + if self.describes.pop_front().is_none() { + self.portal_describes = self.portal_describes.saturating_sub(1); + } + } + + // Portal suspended by a row limit: it stays open and the client + // will execute it again, so keep what we learned about it. + 's' => { + if let Some(portal) = self.executing.pop_front() { + self.binds.push_back(portal); + } + } + + // Empty query, nothing was executed. + 'I' => { + self.executing.pop_front(); } '1' | 'C' => { if let Some(name) = self.parses.pop_front() { self.prepared(&name); } + if code == 'C' { + self.executing.pop_front(); + } } // The close statement that is ignored and we have the parse for @@ -655,20 +697,255 @@ impl PreparedStatements { } } + /// Rewrite the ParameterDescription to canonical OIDs and cache it for the + /// statement being described, so Bind parameters of array/composite types + /// can be rewritten later. fn rewrite_parameter_description_data_types(&self, message: &mut Message) -> Result<(), Error> { - let Some(mappings) = self.oids.get() else { + let mut parameter_description = ParameterDescription::from_bytes(message.payload())?; + + if let Some(mappings) = self.oids.get() + && !mappings.is_identity() + { + parameter_description.rewrite_data_types(&mappings.shard_to_canonical); + message.replace_payload(parameter_description.to_bytes()); + } + + if let Some(describe) = self.describes.front() { + self.global_cache + .write() + .insert_parameter_description(describe, parameter_description); + } + + Ok(()) + } + + /// Upstream handling of Bind: prepare the statement first if needed, + /// and anonymize it in ExtendedAnonymous mode. + fn handle_bind(&mut self, bind: &Bind) -> Result { + if !bind.anonymous() { + let message = self.check_prepared(bind.statement())?; + match message { + Some(mut message) => { + if message.close.is_some() { + self.state.add_ignore('3'); + } + self.state.add_ignore('1'); + self.parses.push_back(bind.statement().to_string()); + self.state.add('2'); + if self.config.level.rewrite_anonymous() { + message.anonymize(); + let mut bind = bind.clone(); + bind.anonymize(); + return Ok(HandleResult::PrependRewrite { + prepend: message, + rewrite: ProtocolMessage::Bind(bind), + }); + } else { + return Ok(HandleResult::Prepend(message)); + } + } + + None => { + self.state.add('2'); + if self.config.level.rewrite_anonymous() { + let mut bind = bind.clone(); + bind.anonymize(); + return Ok(HandleResult::Rewrite(ProtocolMessage::Bind(bind))); + } + } + } + } else { + self.state.add('2'); + } + + Ok(HandleResult::Forward) + } + + /// Remember a portal the client bound, if this shard's OIDs need translating. + fn bound(&mut self, bind: &Bind) -> Result<(), Error> { + if self + .oids + .get() + .is_none_or(|mappings| mappings.is_identity()) + { + return Ok(()); + } + if self.binds.len() >= MAX_BOUND_PORTALS { + self.binds.pop_front(); + } + self.binds.push_back(BoundPortal { + portal: bind.portal()?.to_owned(), + statement: bind.statement().to_owned(), + formats: bind.result_formats().collect(), + ..Default::default() + }); + Ok(()) + } + + /// The client is executing a portal; its rows come next. + fn executed(&mut self, portal: &str) { + if let Some(index) = self.binds.iter().position(|bound| bound.portal == portal) + && let Some(bound) = self.binds.remove(index) + { + self.executing.push_back(bound); + } + } + + /// The portal a Describe(portal) response belongs to: the oldest one + /// we haven't seen a RowDescription for. + fn described_portal(&mut self) -> Option<&mut BoundPortal> { + self.executing + .iter_mut() + .chain(self.binds.iter_mut()) + .find(|portal| portal.row_description.is_none()) + } + + /// Rewrite type OIDs embedded in binary array/composite parameters from + /// canonical to this shard's, replacing the Bind in `result` if any changed. + fn rewrite_bind_params( + &self, + bind: &Bind, + result: HandleResult, + ) -> Result { + let Some(mut rewritten) = self.rewrite_params(bind)? else { + return Ok(result); + }; + + Ok(match result { + HandleResult::Forward => HandleResult::Rewrite(ProtocolMessage::Bind(rewritten)), + HandleResult::Prepend(prepend) => HandleResult::PrependRewrite { + prepend, + rewrite: ProtocolMessage::Bind(rewritten), + }, + HandleResult::Rewrite(ProtocolMessage::Bind(_)) => { + rewritten.anonymize(); + HandleResult::Rewrite(ProtocolMessage::Bind(rewritten)) + } + HandleResult::PrependRewrite { + prepend, + rewrite: ProtocolMessage::Bind(_), + } => { + rewritten.anonymize(); + HandleResult::PrependRewrite { + prepend, + rewrite: ProtocolMessage::Bind(rewritten), + } + } + other => other, + }) + } + + /// Returns the Bind with its binary array/composite parameters rewritten, if any. + /// Parameter types come from the statement's Describe response. + fn rewrite_params(&self, bind: &Bind) -> Result, Error> { + if bind.anonymous() || bind.params_raw().is_empty() { + return Ok(None); + } + let Some(mappings) = self.oids.get().filter(|mappings| !mappings.is_identity()) else { + return Ok(None); + }; + let Some(types) = self + .global_cache + .read() + .parameter_description(bind.statement()) + else { + return Ok(None); + }; + + let rewriter = mappings.to_shard(); + let mut rewritten: Option = None; + + for (index, oid) in types.data_types().enumerate() { + if !rewriter.needs_rewrite(oid) { + continue; + } + let Some(param) = bind.parameter(index)? else { + continue; + }; + if param.is_null() || param.format() != Format::Binary { + continue; + } + if let Some(data) = Self::rewrite_value(&rewriter, oid, param.data(), "parameter") { + rewritten + .get_or_insert_with(|| bind.clone()) + .set_param(index, Parameter::new(&data)); + } + } + + Ok(rewritten) + } + + /// Rewrite type OIDs embedded in binary array/composite columns + /// from this shard's to canonical. + fn rewrite_data_row(&mut self, message: &mut Message) -> Result<(), Error> { + let Some(portal) = self.executing.front_mut() else { return Ok(()); }; - let mappings = &mappings.shard_to_canonical; - if mappings.is_empty() { + let Some(mappings) = self.oids.get().filter(|mappings| !mappings.is_identity()) else { + return Ok(()); + }; + + let rewriter = mappings.to_canonical(); + let plan = match &portal.plan { + Some(plan) => plan, + None => { + // From the portal's Describe, or the statement's; both canonical. + let row_description = portal + .row_description + .clone() + .or_else(|| self.global_cache.read().row_description(&portal.statement)); + let plan = row_description + .iter() + .flat_map(|row_description| row_description.iter().enumerate()) + .filter(|(index, _)| portal.result_format(*index) == Format::Binary) + .map(|(index, field)| (index, mappings.shard_oid(field.type_oid as u32))) + .filter(|(_, oid)| rewriter.needs_rewrite(*oid)) + .collect(); + portal.plan.insert(plan) + } + }; + + if plan.is_empty() { 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 mut row = DataRow::from_bytes(message.payload())?; + let mut changed = false; + for &(index, oid) in plan { + let Some(column) = row.get_raw(index).filter(|column| !column.is_null) else { + continue; + }; + if let Some(data) = Self::rewrite_value(&rewriter, oid, column, "column") { + row.insert(index, bytes::Bytes::from(data), false); + changed = true; + } + } + + if changed { + message.replace_payload(row.to_bytes()); + } + Ok(()) } + + /// Rewrite the OIDs embedded in one binary value. Returns the new + /// bytes if anything changed; malformed values are left alone. + fn rewrite_value( + rewriter: &PayloadRewriter<'_>, + oid: u32, + value: &[u8], + what: &str, + ) -> Option> { + let mut data = value.to_vec(); + match rewriter.rewrite(oid, &mut data) { + Ok(true) => Some(data), + Ok(false) => None, + Err(_) => { + warn!("malformed binary {what} of type oid {oid}, not rewriting"); + None + } + } + } } #[cfg(test)] @@ -680,6 +957,7 @@ pub(crate) mod test { Prepare as SimplePrepare, ProtocolMessage, Query, Sync, bind::Parameter, messages::ReadyForQuery, }; + use bytes::BufMut; use pgdog_config::PreparedStatementsLevel; /// Build a PreparedStatements instance configured for ExtendedAnonymous mode. @@ -1319,6 +1597,251 @@ pub(crate) mod test { ); } + // ------------------------------------------------------- + // Embedded OIDs in binary arrays/composites + // ------------------------------------------------------- + + const MOOD: u32 = 16400; + const MOOD_ARRAY: u32 = 16401; + const SHARD_MOOD: u32 = 17000; + const SHARD_MOOD_ARRAY: u32 = 17001; + + /// Mappings for a shard where `mood` and `mood[]` have different OIDs. + fn mood_oids() -> Arc { + use crate::backend::pool::shard::TypeKind; + + Oids::from_canonical_with_kinds( + [(MOOD, SHARD_MOOD), (MOOD_ARRAY, SHARD_MOOD_ARRAY)] + .into_iter() + .collect(), + [( + SHARD_MOOD_ARRAY, + TypeKind::Array { + element: SHARD_MOOD, + }, + )] + .into_iter() + .collect(), + [(MOOD_ARRAY, TypeKind::Array { element: MOOD })] + .into_iter() + .collect(), + ) + } + + /// Binary array of one text-ish element with the given element OID. + fn mood_array(element: u32) -> Vec { + let mut buf = bytes::BytesMut::new(); + buf.put_i32(1); // ndim + buf.put_i32(0); // no nulls + buf.put_u32(element); + buf.put_i32(1); // size + buf.put_i32(1); // lower bound + buf.put_i32(3); + buf.put_slice(b"sad"); + buf.to_vec() + } + + fn array_element_oid(data: &[u8]) -> u32 { + u32::from_be_bytes([data[8], data[9], data[10], data[11]]) + } + + fn mood_array_row_description(type_oid: u32) -> RowDescription { + RowDescription::new(&[crate::net::messages::Field { + name: "moods".into(), + table_oid: 0, + column: 0, + type_oid: type_oid as i32, + type_size: -1, + type_modifier: -1, + format: 1, + }]) + } + + fn bind_complete() -> Message { + Message::new(crate::net::messages::BindComplete.to_bytes()) + } + + fn mood_row(element: u32) -> Message { + let mut row = DataRow::new(); + row.add(bytes::Bytes::from(mood_array(element))); + Message::new(row.to_bytes()) + } + + fn row_element_oid(message: &Message) -> u32 { + array_element_oid( + &DataRow::from_bytes(message.payload()) + .unwrap() + .column(0) + .unwrap(), + ) + } + + #[test] + fn bind_rewrites_binary_array_param_to_shard_oids() { + let mut ps = new_extended(); + ps.oids = mood_oids(); + let name = insert_global("array_param", "INSERT INTO t VALUES ($1)"); + FrontendPreparedStatements::global() + .write() + .insert_parameter_description( + &name, + ParameterDescription::new(vec![MOOD_ARRAY as i32]), + ); + ps.prepared(&name); + + let bind = Bind::new_params_codes( + &name, + &[Parameter::new(&mood_array(MOOD))], + &[Format::Binary], + ); + let result = ps.handle(&ProtocolMessage::Bind(bind)).unwrap(); + let HandleResult::Rewrite(ProtocolMessage::Bind(rewritten)) = result else { + panic!("expected rewritten bind, got {result:?}"); + }; + let param = rewritten.parameter(0).unwrap().unwrap(); + assert_eq!(array_element_oid(param.data()), SHARD_MOOD); + + // Text params are left alone. + let bind = Bind::new_params_codes(&name, &[Parameter::new(b"{sad}")], &[Format::Text]); + let result = ps.handle(&ProtocolMessage::Bind(bind)).unwrap(); + assert!(matches!(result, HandleResult::Forward), "{result:?}"); + } + + #[test] + fn parameter_description_is_cached_for_the_described_statement() { + let mut ps = new_extended(); + ps.oids = mood_oids(); + let name = insert_global("cache_params", "INSERT INTO t VALUES ($1)"); + ps.prepared(&name); + ps.handle(&ProtocolMessage::Describe(Describe::new_statement(&name))) + .unwrap(); + + // The shard describes the parameter with its own OID. + let mut params = + Message::new(ParameterDescription::new(vec![SHARD_MOOD_ARRAY as i32]).to_bytes()); + assert!(ps.forward(&mut params).unwrap()); + + let cached = FrontendPreparedStatements::global() + .read() + .parameter_description(&name) + .unwrap(); + assert_eq!(cached.data_types().collect::>(), vec![MOOD_ARRAY]); + } + + #[test] + fn data_row_rewrites_binary_array_using_cached_row_description() { + let mut ps = new_extended(); + ps.oids = mood_oids(); + let name = insert_global("array_rows", "SELECT moods FROM t"); + // Described earlier: cached with canonical OIDs. + FrontendPreparedStatements::global() + .write() + .insert_row_description(&name, mood_array_row_description(MOOD_ARRAY)); + ps.prepared(&name); + + let bind = Bind::new_params_codes_results(&name, &[], &[], &[1]); + ps.handle(&ProtocolMessage::Bind(bind)).unwrap(); + ps.handle(&ProtocolMessage::Execute(Execute::new())) + .unwrap(); + assert!(ps.forward(&mut bind_complete()).unwrap()); + + let mut message = mood_row(SHARD_MOOD); + assert!(ps.forward(&mut message).unwrap()); + assert_eq!(row_element_oid(&message), MOOD); + + let mut complete = Message::new(CommandComplete::from_str("SELECT 1").to_bytes()); + assert!(ps.forward(&mut complete).unwrap()); + assert!(ps.executing.is_empty()); + } + + #[test] + fn data_row_rewrites_binary_array_using_portal_row_description() { + let mut ps = new_extended(); + ps.oids = mood_oids(); + let name = insert_global("array_rows_portal", "SELECT moods FROM t"); + ps.prepared(&name); + + let bind = Bind::new_params_codes_results(&name, &[], &[], &[1]); + ps.handle(&ProtocolMessage::Bind(bind)).unwrap(); + ps.handle(&ProtocolMessage::Describe(Describe::new_portal(""))) + .unwrap(); + ps.handle(&ProtocolMessage::Execute(Execute::new())) + .unwrap(); + assert!(ps.forward(&mut bind_complete()).unwrap()); + + // The portal's RowDescription carries the shard's OIDs, canonicalized on the way out. + let mut description = Message::new(mood_array_row_description(SHARD_MOOD_ARRAY).to_bytes()); + assert!(ps.forward(&mut description).unwrap()); + assert_eq!( + RowDescription::from_bytes(description.payload()) + .unwrap() + .field(0) + .unwrap() + .type_oid, + MOOD_ARRAY as i32 + ); + + let mut message = mood_row(SHARD_MOOD); + assert!(ps.forward(&mut message).unwrap()); + assert_eq!(row_element_oid(&message), MOOD); + } + + #[test] + fn suspended_portal_keeps_rewriting_when_executed_again() { + let mut ps = new_extended(); + ps.oids = mood_oids(); + let name = insert_global("array_rows_suspended", "SELECT moods FROM t"); + FrontendPreparedStatements::global() + .write() + .insert_row_description(&name, mood_array_row_description(MOOD_ARRAY)); + ps.prepared(&name); + + let bind = Bind::new_params_codes_results(&name, &[], &[], &[1]); + ps.handle(&ProtocolMessage::Bind(bind)).unwrap(); + ps.handle(&ProtocolMessage::Execute(Execute::new())) + .unwrap(); + assert!(ps.forward(&mut bind_complete()).unwrap()); + + let mut message = mood_row(SHARD_MOOD); + assert!(ps.forward(&mut message).unwrap()); + assert_eq!(row_element_oid(&message), MOOD); + + // Row limit reached; the portal is still open. PortalSuspended: code 's', length only. + let mut suspended = Message::new(bytes::Bytes::from_static(&[b's', 0, 0, 0, 4])); + assert!(ps.forward(&mut suspended).unwrap()); + assert!(ps.executing.is_empty()); + + ps.handle(&ProtocolMessage::Execute(Execute::new())) + .unwrap(); + let mut message = mood_row(SHARD_MOOD); + assert!(ps.forward(&mut message).unwrap()); + assert_eq!(row_element_oid(&message), MOOD); + } + + #[test] + fn data_row_in_text_format_is_left_alone() { + let mut ps = new_extended(); + ps.oids = mood_oids(); + let name = insert_global("array_rows_text", "SELECT moods FROM t"); + FrontendPreparedStatements::global() + .write() + .insert_row_description(&name, mood_array_row_description(MOOD_ARRAY)); + ps.prepared(&name); + + ps.handle(&ProtocolMessage::Bind(Bind::new_statement(&name))) + .unwrap(); + ps.handle(&ProtocolMessage::Execute(Execute::new())) + .unwrap(); + assert!(ps.forward(&mut bind_complete()).unwrap()); + + let mut row = DataRow::new(); + row.add(bytes::Bytes::from_static(b"{sad}")); + let original = Message::new(row.to_bytes()); + let mut message = original.clone(); + assert!(ps.forward(&mut message).unwrap()); + assert_eq!(message.payload(), original.payload()); + } + // ------------------------------------------------------- // Simple query is unaffected by mode // ------------------------------------------------------- diff --git a/pgdog/src/frontend/prepared_statements/global_cache.rs b/pgdog/src/frontend/prepared_statements/global_cache.rs index 39269cfae..1ea9cefe7 100644 --- a/pgdog/src/frontend/prepared_statements/global_cache.rs +++ b/pgdog/src/frontend/prepared_statements/global_cache.rs @@ -2,7 +2,7 @@ use crate::{ frontend::RewritePlan, net::{ Prepare, - messages::{Parse, RowDescription}, + messages::{ParameterDescription, Parse, RowDescription}, }, stats::memory::MemoryUsage, }; @@ -73,6 +73,7 @@ impl GlobalCache { }, cache_key: cache_key.clone(), row_description: None, + parameter_description: None, }; self.insert_internal(&name, cache_key, statement); @@ -110,6 +111,7 @@ impl GlobalCache { unique_ids: rewrite_plan.unique_ids, }, row_description: None, + parameter_description: None, cache_key: cache_key.clone(), }; @@ -134,6 +136,27 @@ impl GlobalCache { } } + /// Client sent a Describe for a prepared statement and received a ParameterDescription. + /// We record it to know the types of parameters sent in Bind. + pub(crate) fn insert_parameter_description( + &mut self, + name: &str, + parameter_description: ParameterDescription, + ) { + if let Some(entry) = self.names.get_mut(name) + && entry.parameter_description.is_none() + { + entry.parameter_description = Some(parameter_description); + } + } + + /// Get the ParameterDescription for the prepared statement, if it was described. + pub(crate) fn parameter_description(&self, name: &str) -> Option { + self.names + .get(name) + .and_then(|p| p.parameter_description.clone()) + } + /// Get the Parse message for a globally unique prepared statement /// name. /// diff --git a/pgdog/src/frontend/prepared_statements/prelude.rs b/pgdog/src/frontend/prepared_statements/prelude.rs index 2cd079e5e..aa23f26b4 100644 --- a/pgdog/src/frontend/prepared_statements/prelude.rs +++ b/pgdog/src/frontend/prepared_statements/prelude.rs @@ -1,6 +1,6 @@ pub(super) use super::CacheKey; pub(super) use crate::config::config; -pub(super) use crate::net::{Parse, RowDescription}; +pub(super) use crate::net::{ParameterDescription, Parse, RowDescription}; pub(super) use crate::util::*; pub(super) use bytes::Bytes; pub(super) use std::str::from_utf8; diff --git a/pgdog/src/frontend/prepared_statements/statement.rs b/pgdog/src/frontend/prepared_statements/statement.rs index 84acbdc28..85e6fb353 100644 --- a/pgdog/src/frontend/prepared_statements/statement.rs +++ b/pgdog/src/frontend/prepared_statements/statement.rs @@ -6,6 +6,7 @@ use super::prelude::*; pub(crate) struct Statement { pub(super) stmt: StatementType, pub(super) row_description: Option, + pub(super) parameter_description: Option, pub(super) cache_key: CacheKey, } @@ -51,6 +52,11 @@ impl MemoryUsage for Statement { } else { 0 } + + self + .parameter_description + .as_ref() + .map(|params| params.memory_usage()) + .unwrap_or_default() + self.cache_key.memory_usage() } } diff --git a/pgdog/src/net/messages/bind.rs b/pgdog/src/net/messages/bind.rs index 330fa9aab..1b2a8902e 100644 --- a/pgdog/src/net/messages/bind.rs +++ b/pgdog/src/net/messages/bind.rs @@ -203,6 +203,12 @@ impl Bind { unsafe { from_utf8_unchecked(&self.statement[0..self.statement.len() - 1]) } } + /// Portal name; empty for the unnamed portal. + #[inline] + pub(crate) fn portal(&self) -> Result<&str, Error> { + Ok(from_utf8(&self.portal[0..self.portal.len() - 1])?) + } + /// Format the client asked each result column to be returned in. pub(crate) fn result_formats(&self) -> impl ExactSizeIterator + '_ { self.results.chunks_exact(2).map(|code| { diff --git a/pgdog/src/net/messages/parameter_description.rs b/pgdog/src/net/messages/parameter_description.rs index f329e2c38..cd00ccd13 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 crate::stats::memory::MemoryUsage; use std::collections::HashMap; #[derive(Debug, Clone, Default)] @@ -38,12 +39,29 @@ impl Protocol for ParameterDescription { } } +impl MemoryUsage for ParameterDescription { + #[inline] + fn memory_usage(&self) -> usize { + self.params.capacity() * std::mem::size_of::() + } +} + impl ParameterDescription { /// Create an empty parameter description. pub(crate) fn empty() -> Self { Self { params: Vec::new() } } + #[cfg(test)] + pub(crate) fn new(params: Vec) -> Self { + Self { params } + } + + /// Parameter data type OIDs. + pub(crate) fn data_types(&self) -> impl Iterator + '_ { + self.params.iter().map(|¶m| param as u32) + } + 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)) {