use chrono::{DateTime, Local, NaiveDate, NaiveDateTime, NaiveTime}; use deadpool_postgres::{ManagerConfig, Pool, PoolError, RecyclingMethod, Runtime}; use futures::{SinkExt, StreamExt}; use openssl::ssl::{SslConnector, SslFiletype, SslMethod, SslVerifyMode}; use percent_encoding::percent_decode_str; use rustls::client::danger::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier}; use rustls::client::verify_server_cert_signed_by_trust_anchor; use rustls::crypto::{verify_tls12_signature, verify_tls13_signature, CryptoProvider}; use rustls::pki_types::{CertificateDer, PrivateKeyDer, ServerName, UnixTime}; use rustls::server::ParsedCertificate; use sqlparser::ast::{SetExpr, Statement}; use sqlparser::dialect::PostgreSqlDialect; use sqlparser::parser::Parser; use std::collections::{BTreeMap, BTreeSet, HashMap}; use std::fmt::Write as _; use std::fs::File; use std::future::Future; use std::io::BufReader; use std::pin::Pin; use std::str::FromStr; use std::sync::atomic::AtomicBool; use std::sync::{Arc, Mutex, OnceLock, Weak}; use std::time::{Duration, Instant}; use tokio::task::JoinHandle; use tokio_postgres::config::SslMode; use tokio_postgres::tls::{MakeTlsConnect, TlsConnect}; use tokio_postgres::types::{FromSql, Kind, Type}; use tokio_postgres::{AsyncMessage, NoTls, Row, SimpleQueryMessage, Socket}; use tokio_util::sync::CancellationToken; use super::file_validator::validate_file_path; use crate::query::{await_stream_with_progress_timeout, DbOperationBudget, StreamProgressClock}; use crate::sql::starts_with_executable_sql_keyword; use crate::types::{ ColumnInfo, CompletionAssistantCandidate, CompletionAssistantCandidateKind, CompletionAssistantMatchMode, CompletionAssistantObjectKind, CompletionAssistantRequest, CompletionAssistantResponse, ConstraintInfo, CustomTypeDdl, CustomTypeDetails, CustomTypeDomainConstraint, CustomTypeKind, CustomTypeMember, CustomTypeProperties, DatabaseInfo, DatabaseStorageInfo, ExtensionInfo, ForeignKeyInfo, FunctionInfo, IndexInfo, ObjectInfo, ObjectStatistics, OwnerInfo, QueryMessage, QueryResult, RuleInfo, SchemaInfo, SequenceInfo, SpatialColumnBuilder, TableInfo, TriggerInfo, }; pub(crate) const GAUSSDB_COMPATIBILITY_SQL: &str = "SELECT datcompatibility FROM pg_catalog.pg_database WHERE datname = current_database()"; pub async fn gaussdb_identifier_quote(pool: &Pool) -> Option { let timeout = super::connection_timeout(); let client = checkout_postgres_client(pool, None, timeout).await.ok()?; let row = tokio::time::timeout(timeout, client.query_opt(GAUSSDB_COMPATIBILITY_SQL, &[])).await.ok()?.ok()??; let compatibility_mode = row.try_get::<_, String>(0).ok()?; gaussdb_identifier_quote_for_compatibility_mode(&compatibility_mode).map(str::to_string) } pub(crate) fn gaussdb_identifier_quote_for_compatibility_mode(compatibility_mode: &str) -> Option<&'static str> { match compatibility_mode.trim().to_ascii_uppercase().as_str() { "M" | "B" | "MYSQL" => Some("`"), "A" | "PG" | "ORA" | "POSTGRESQL" => Some("\""), _ => None, } } #[derive(Debug, Clone, PartialEq, Eq)] pub struct PostgresTablePrivilegeInfo { pub grantor: String, pub grantee: String, pub privilege_type: String, pub is_grantable: bool, pub column_name: Option, } #[derive(Debug, Clone, PartialEq, Eq)] pub struct PostgresTableAccessInfo { pub owner: String, pub owner_default_privileges: Vec, pub privileges: Vec, } #[derive(Debug, Clone, Default, PartialEq, Eq)] pub struct PostgresTablePartitionInfo { pub is_partition: bool, pub parent_schema: Option, pub parent_table: Option, pub bound: Option, pub key: Option, /// True when this relation is itself a foreign table (relkind 'f'), /// whether or not it is also a partition. A partition backed by a /// foreign server must be declared with `CREATE FOREIGN TABLE ... /// SERVER`, not `CREATE TABLE`. pub is_foreign: bool, pub foreign_server: Option, pub foreign_options: Vec<(String, String)>, } /// The state of a partition-local column DEFAULT relative to the parent's /// default for the same column (a partition with no entry in the owning /// `column_defaults` map simply inherits the parent's default unchanged). #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum PostgresColumnDefaultState { /// This partition's own default differs from (or has no counterpart in) /// the parent's default — declared via `column WITH OPTIONS DEFAULT ...`. Overridden, /// The parent has a default for this column, but this partition /// explicitly ran `ALTER TABLE ONLY ... ALTER COLUMN ... DROP DEFAULT` /// to remove its own (otherwise auto-inherited-at-creation) copy — /// replayed as a standalone `ALTER TABLE ONLY ... DROP DEFAULT;`. Dropped, } #[derive(Debug, Clone, Default, PartialEq, Eq)] pub struct PostgresTablePartitionLocalObjects { pub has_primary_key: bool, pub foreign_keys: BTreeSet, pub indexes: BTreeSet, /// CHECK constraints with a local definition on this partition, as /// reported by pg_constraint.conislocal. A merged constraint may remain /// local even when coninhcount is greater than zero. pub check_constraints: BTreeSet, /// Columns with a local override or explicit drop of the parent's /// default; a column absent from this map is purely inherited. pub column_defaults: BTreeMap, } fn pg_temporal_to_json_value(row: &Row, idx: usize) -> Option { if let Ok(v) = row.try_get::<_, DateTime>(idx) { return Some(serde_json::Value::String(format_pg_timestamptz(v))); } if let Ok(v) = row.try_get::<_, NaiveDateTime>(idx) { return Some(serde_json::Value::String(v.to_string())); } if let Ok(v) = row.try_get::<_, NaiveDate>(idx) { return Some(serde_json::Value::String(v.to_string())); } if let Ok(v) = row.try_get::<_, NaiveTime>(idx) { return Some(serde_json::Value::String(v.to_string())); } None } struct PgSystemU32(u32); impl<'a> FromSql<'a> for PgSystemU32 { fn from_sql(_: &Type, raw: &'a [u8]) -> Result> { let bytes: [u8; 4] = raw.try_into().map_err(|_| "expected 4 bytes for PostgreSQL system u32")?; Ok(Self(u32::from_be_bytes(bytes))) } fn accepts(ty: &Type) -> bool { matches!(*ty, Type::XID | Type::CID) } } /// A `FromSql` adapter that accepts any PostgreSQL type and reads its raw /// bytes as a UTF-8 string. This is used as a last-resort fallback to handle /// custom types (enums, domains, etc.) that tokio_postgres cannot map to /// built-in Rust types in the binary protocol. struct PgAnyString(String); impl<'a> FromSql<'a> for PgAnyString { fn from_sql(_: &Type, raw: &'a [u8]) -> Result> { std::str::from_utf8(raw) .map(|s| PgAnyString(s.to_string())) .map_err(|e| Box::new(e) as Box) } fn accepts(_: &Type) -> bool { true } } struct PgNumeric(String); impl<'a> FromSql<'a> for PgNumeric { fn from_sql(_: &Type, raw: &'a [u8]) -> Result> { decode_pg_numeric_bytes(raw).map(Self).ok_or_else(|| "invalid PostgreSQL numeric binary value".into()) } fn accepts(ty: &Type) -> bool { *ty == Type::NUMERIC } } fn decode_pg_numeric_bytes(raw: &[u8]) -> Option { const POSITIVE: u16 = 0x0000; const NEGATIVE: u16 = 0x4000; const NAN: u16 = 0xC000; const POSITIVE_INFINITY: u16 = 0xD000; const NEGATIVE_INFINITY: u16 = 0xF000; let digit_count = usize::from(u16::from_be_bytes(raw.get(0..2)?.try_into().ok()?)); let weight = i16::from_be_bytes(raw.get(2..4)?.try_into().ok()?); let sign = u16::from_be_bytes(raw.get(4..6)?.try_into().ok()?); let scale = usize::from(u16::from_be_bytes(raw.get(6..8)?.try_into().ok()?)); let expected_len = digit_count.checked_mul(2)?.checked_add(8)?; if raw.len() != expected_len { return None; } match sign { NAN => return Some("NaN".to_string()), POSITIVE_INFINITY => return Some("Infinity".to_string()), NEGATIVE_INFINITY => return Some("-Infinity".to_string()), POSITIVE | NEGATIVE => {} _ => return None, } let digits = raw[8..].chunks_exact(2).map(|chunk| u16::from_be_bytes([chunk[0], chunk[1]])).collect::>(); if digits.iter().any(|digit| *digit > 9999) { return None; } let integer_groups = if weight >= 0 { usize::try_from(weight).ok()?.checked_add(1)? } else { 0 }; let fractional_groups = scale.checked_add(3)? / 4; let mut text = String::with_capacity(integer_groups.saturating_mul(4).max(1).saturating_add(scale + 2)); if sign == NEGATIVE && digits.iter().any(|digit| *digit != 0) { text.push('-'); } if integer_groups == 0 { text.push('0'); } else { for group_index in 0..integer_groups { let digit = digits.get(group_index).copied().unwrap_or(0); if group_index == 0 { write!(&mut text, "{digit}").ok()?; } else { write!(&mut text, "{digit:04}").ok()?; } } } if scale > 0 { text.push('.'); let fractional_start = text.len(); for group_index in 0..fractional_groups { let digit_index = i32::from(weight) + 1 + i32::try_from(group_index).ok()?; let digit = usize::try_from(digit_index).ok().and_then(|index| digits.get(index)).copied().unwrap_or(0); write!(&mut text, "{digit:04}").ok()?; } text.truncate(fractional_start + scale); } Some(text) } /// PostgreSQL `money` is sent in the binary protocol as a signed 64-bit /// integer scaled by 100 (the smallest unit of the currency). It is not /// decoded by `rust_decimal`, so attempting to read it as `Decimal` produces /// NULL in the result grid. struct PgMoney(i64); impl<'a> FromSql<'a> for PgMoney { fn from_sql(_: &Type, raw: &'a [u8]) -> Result> { let bytes: [u8; 8] = raw.try_into().map_err(|_| "expected 8 bytes for PostgreSQL money")?; Ok(Self(i64::from_be_bytes(bytes))) } fn accepts(ty: &Type) -> bool { *ty == Type::MONEY } } fn format_pg_money(value: i64) -> String { let negative = value < 0; let absolute = (value as i128).abs(); let formatted = format!("{}.{:02}", absolute / 100, absolute % 100); if negative { format!("-{formatted}") } else { formatted } } /// A `FromSql` adapter that accepts any PostgreSQL type and returns the raw /// bytes unchanged. Used to decode custom types like pgvector whose binary /// format we handle ourselves. struct PgRawBytes(Vec); impl<'a> FromSql<'a> for PgRawBytes { fn from_sql(_: &Type, raw: &'a [u8]) -> Result> { Ok(PgRawBytes(raw.to_vec())) } fn accepts(_: &Type) -> bool { true } } #[derive(Clone, Copy, Debug, PartialEq)] struct PgPoint { x: f64, y: f64, } impl<'a> FromSql<'a> for PgPoint { fn from_sql(_: &Type, raw: &'a [u8]) -> Result> { decode_pg_point_bytes(raw).ok_or_else(|| "expected 16 bytes for PostgreSQL point".into()) } fn accepts(ty: &Type) -> bool { *ty == Type::POINT } } fn decode_pg_point_bytes(raw: &[u8]) -> Option { let raw: [u8; 16] = raw.try_into().ok()?; Some(PgPoint { x: f64::from_be_bytes(raw[0..8].try_into().ok()?), y: f64::from_be_bytes(raw[8..16].try_into().ok()?), }) } fn format_pg_float(value: f64) -> String { if value == f64::INFINITY { "Infinity".to_string() } else if value == f64::NEG_INFINITY { "-Infinity".to_string() } else { value.to_string() } } fn format_pg_point(point: PgPoint) -> String { format!("({},{})", format_pg_float(point.x), format_pg_float(point.y)) } #[derive(Clone, Copy, Debug, PartialEq, Eq)] struct PgInterval { microseconds: i64, days: i32, months: i32, } impl<'a> FromSql<'a> for PgInterval { fn from_sql(_: &Type, raw: &'a [u8]) -> Result> { decode_pg_interval_bytes(raw).ok_or_else(|| "expected 16 bytes for PostgreSQL interval".into()) } fn accepts(ty: &Type) -> bool { *ty == Type::INTERVAL } } fn decode_pg_interval_bytes(raw: &[u8]) -> Option { let raw: [u8; 16] = raw.try_into().ok()?; Some(PgInterval { microseconds: i64::from_be_bytes(raw[0..8].try_into().ok()?), days: i32::from_be_bytes(raw[8..12].try_into().ok()?), months: i32::from_be_bytes(raw[12..16].try_into().ok()?), }) } fn push_pg_interval_component(parts: &mut Vec, value: i64, singular: &str, plural: &str) { if value == 0 { return; } let unit = if value.abs() == 1 { singular } else { plural }; parts.push(format!("{value} {unit}")); } fn format_pg_interval_time(microseconds: i64) -> String { let signed_microseconds = i128::from(microseconds); let sign = if signed_microseconds < 0 { "-" } else { "" }; let absolute_microseconds = signed_microseconds.abs(); let hours = absolute_microseconds / 3_600_000_000; let minutes = absolute_microseconds / 60_000_000 % 60; let seconds = absolute_microseconds / 1_000_000 % 60; let fraction = absolute_microseconds % 1_000_000; let mut formatted = format!("{sign}{hours:02}:{minutes:02}:{seconds:02}"); if fraction != 0 { let fraction = format!("{fraction:06}"); formatted.push('.'); formatted.push_str(fraction.trim_end_matches('0')); } formatted } fn format_pg_interval(interval: PgInterval) -> String { let total_months = i64::from(interval.months); let years = total_months / 12; let months = total_months % 12; let mut parts = Vec::with_capacity(4); push_pg_interval_component(&mut parts, years, "year", "years"); push_pg_interval_component(&mut parts, months, "mon", "mons"); push_pg_interval_component(&mut parts, i64::from(interval.days), "day", "days"); parts.push(format_pg_interval_time(interval.microseconds)); parts.join(" ") } struct PgDateRange(String); impl<'a> FromSql<'a> for PgDateRange { fn from_sql(_: &Type, raw: &'a [u8]) -> Result> { decode_pg_daterange_bytes(raw).map(Self).ok_or_else(|| "invalid PostgreSQL daterange binary value".into()) } fn accepts(ty: &Type) -> bool { *ty == Type::DATE_RANGE } } const PG_RANGE_EMPTY: u8 = 0b0000_0001; const PG_RANGE_LOWER_INCLUSIVE: u8 = 0b0000_0010; const PG_RANGE_UPPER_INCLUSIVE: u8 = 0b0000_0100; const PG_RANGE_LOWER_UNBOUNDED: u8 = 0b0000_1000; const PG_RANGE_UPPER_UNBOUNDED: u8 = 0b0001_0000; const PG_RANGE_KNOWN_FLAGS: u8 = PG_RANGE_EMPTY | PG_RANGE_LOWER_INCLUSIVE | PG_RANGE_UPPER_INCLUSIVE | PG_RANGE_LOWER_UNBOUNDED | PG_RANGE_UPPER_UNBOUNDED; fn decode_pg_date_bound(raw: &[u8], cursor: &mut usize, unbounded: bool) -> Option { if unbounded { return Some(String::new()); } let length = read_i32_be(raw, cursor)?; if length != 4 { return None; } let days = read_i32_be(raw, cursor)?; match days { i32::MIN => Some("-infinity".to_string()), i32::MAX => Some("infinity".to_string()), days => NaiveDate::from_ymd_opt(2000, 1, 1)? .checked_add_signed(chrono::Duration::days(i64::from(days))) .map(|date| date.to_string()), } } fn decode_pg_daterange_bytes(raw: &[u8]) -> Option { let (&flags, _) = raw.split_first()?; if flags & !PG_RANGE_KNOWN_FLAGS != 0 { return None; } if flags == PG_RANGE_EMPTY { return (raw.len() == 1).then(|| "empty".to_string()); } if flags & PG_RANGE_EMPTY != 0 { return None; } let lower_unbounded = flags & PG_RANGE_LOWER_UNBOUNDED != 0; let upper_unbounded = flags & PG_RANGE_UPPER_UNBOUNDED != 0; if (lower_unbounded && flags & PG_RANGE_LOWER_INCLUSIVE != 0) || (upper_unbounded && flags & PG_RANGE_UPPER_INCLUSIVE != 0) { return None; } let mut cursor = 1; let lower = decode_pg_date_bound(raw, &mut cursor, lower_unbounded)?; let upper = decode_pg_date_bound(raw, &mut cursor, upper_unbounded)?; if cursor != raw.len() { return None; } let lower_delimiter = if flags & PG_RANGE_LOWER_INCLUSIVE != 0 { '[' } else { '(' }; let upper_delimiter = if flags & PG_RANGE_UPPER_INCLUSIVE != 0 { ']' } else { ')' }; Some(format!("{lower_delimiter}{lower},{upper}{upper_delimiter}")) } /// Decode pgvector binary format into a Vec. /// /// pgvector binary layout (big-endian): /// - 2 bytes: dimensions (uint16) /// - 2 bytes: unused (padding) /// - N*4 bytes: IEEE 754 f32 values fn decode_pgvector_bytes(raw: &[u8]) -> Option> { if raw.len() < 4 { return None; } let dims = u16::from_be_bytes([raw[0], raw[1]]) as usize; let expected_len = 4 + dims * 4; if raw.len() != expected_len { return None; } let floats: Vec = raw[4..].chunks_exact(4).map(|chunk| f32::from_be_bytes([chunk[0], chunk[1], chunk[2], chunk[3]])).collect(); Some(floats) } fn pg_u32_number(v: u32) -> serde_json::Value { serde_json::Value::Number(serde_json::Number::from(v)) } fn pg_system_u32_to_json(row: &Row, idx: usize) -> Option { if let Ok(v) = row.try_get::<_, u32>(idx) { return Some(pg_u32_number(v)); } row.try_get::<_, PgSystemU32>(idx).ok().map(|v| pg_u32_number(v.0)) } fn pg_optional_array_to_json( values: Vec>, map_value: impl Fn(T) -> serde_json::Value, ) -> serde_json::Value { serde_json::Value::Array( values.into_iter().map(|value| value.map(&map_value).unwrap_or(serde_json::Value::Null)).collect(), ) } fn pg_json_array_values_to_json(values: Vec>) -> serde_json::Value { pg_optional_array_to_json(values, |value| serde_json::Value::String(value.to_string())) } fn pg_float_number(v: f64) -> serde_json::Value { serde_json::Number::from_f64(v).map(serde_json::Value::Number).unwrap_or(serde_json::Value::Null) } fn decode_pg_network_address_bytes(raw: &[u8], force_cidr_output: bool) -> Option { let family = *raw.first()?; let bits = *raw.get(1)?; let is_cidr = *raw.get(2)? != 0; let addr_len = *raw.get(3)? as usize; let addr = raw.get(4..)?; if addr.len() != addr_len { return None; } let (address, host_bits) = match (family, addr_len) { (2, 4) => { let bytes: [u8; 4] = addr.try_into().ok()?; (std::net::IpAddr::V4(std::net::Ipv4Addr::from(bytes)).to_string(), 32) } (3, 16) => { let bytes: [u8; 16] = addr.try_into().ok()?; (std::net::IpAddr::V6(std::net::Ipv6Addr::from(bytes)).to_string(), 128) } _ => return None, }; if bits > host_bits { return None; } if force_cidr_output || is_cidr || bits != host_bits { Some(format!("{address}/{bits}")) } else { Some(address) } } fn decode_pg_macaddr_bytes(raw: &[u8]) -> Option { if !matches!(raw.len(), 6 | 8) { return None; } Some(raw.iter().map(|byte| format!("{byte:02x}")).collect::>().join(":")) } fn decode_pg_bit_string_bytes(raw: &[u8]) -> Option { let mut cursor = 0; let bit_len = read_i32_be(raw, &mut cursor)?; if bit_len < 0 { return None; } let bit_len = bit_len as usize; let data = raw.get(cursor..)?; if data.len() != bit_len.div_ceil(8) { return None; } let mut bits = String::with_capacity(bit_len); for index in 0..bit_len { let byte = data[index / 8]; let bit = (byte >> (7 - (index % 8))) & 1; bits.push(if bit == 1 { '1' } else { '0' }); } Some(bits) } fn pg_network_address_to_json_value(row: &Row, idx: usize, force_cidr_output: bool) -> Option { row.try_get::<_, PgRawBytes>(idx) .ok() .and_then(|raw| decode_pg_network_address_bytes(&raw.0, force_cidr_output)) .map(serde_json::Value::String) } fn pg_macaddr_to_json_value(row: &Row, idx: usize) -> Option { row.try_get::<_, PgRawBytes>(idx) .ok() .and_then(|raw| decode_pg_macaddr_bytes(&raw.0)) .map(serde_json::Value::String) } fn pg_bit_string_to_json_value(row: &Row, idx: usize) -> Option { row.try_get::<_, PgRawBytes>(idx) .ok() .and_then(|raw| decode_pg_bit_string_bytes(&raw.0)) .map(serde_json::Value::String) } fn pg_network_address_array_to_json_value(row: &Row, idx: usize, force_cidr_output: bool) -> Option { row.try_get::<_, Vec>>(idx).ok().map(|values| { pg_optional_array_to_json(values, |raw| { decode_pg_network_address_bytes(&raw.0, force_cidr_output) .map(serde_json::Value::String) .unwrap_or_else(|| super::binary_value_to_json(&raw.0)) }) }) } fn pg_macaddr_array_to_json_value(row: &Row, idx: usize) -> Option { row.try_get::<_, Vec>>(idx).ok().map(|values| { pg_optional_array_to_json(values, |raw| { decode_pg_macaddr_bytes(&raw.0) .map(serde_json::Value::String) .unwrap_or_else(|| super::binary_value_to_json(&raw.0)) }) }) } fn pg_bit_string_array_to_json_value(row: &Row, idx: usize) -> Option { row.try_get::<_, Vec>>(idx).ok().map(|values| { pg_optional_array_to_json(values, |raw| { decode_pg_bit_string_bytes(&raw.0) .map(serde_json::Value::String) .unwrap_or_else(|| super::binary_value_to_json(&raw.0)) }) }) } fn pg_array_to_json_value(row: &Row, idx: usize) -> Option { if let Ok(values) = row.try_get::<_, Vec>>(idx) { return Some(pg_json_array_values_to_json(values)); } if let Ok(values) = row.try_get::<_, Vec>>(idx) { return Some(pg_optional_array_to_json(values, serde_json::Value::String)); } if let Ok(values) = row.try_get::<_, Vec>>(idx) { return Some(pg_optional_array_to_json(values, serde_json::Value::Bool)); } if let Ok(values) = row.try_get::<_, Vec>>(idx) { return Some(pg_optional_array_to_json(values, |v| serde_json::Value::String(v.0))); } if let Ok(values) = row.try_get::<_, Vec>>(idx) { return Some(pg_optional_array_to_json(values, |v| serde_json::Value::String(v.to_string()))); } if let Ok(values) = row.try_get::<_, Vec>>>(idx) { return Some(pg_optional_array_to_json(values, |v| serde_json::Value::String(format_pg_timestamptz(v)))); } if let Ok(values) = row.try_get::<_, Vec>>(idx) { return Some(pg_optional_array_to_json(values, |v| serde_json::Value::String(v.to_string()))); } if let Ok(values) = row.try_get::<_, Vec>>(idx) { return Some(pg_optional_array_to_json(values, |v| serde_json::Value::String(v.to_string()))); } if let Ok(values) = row.try_get::<_, Vec>>(idx) { return Some(pg_optional_array_to_json(values, |v| serde_json::Value::String(v.to_string()))); } if let Ok(values) = row.try_get::<_, Vec>>(idx) { return Some(pg_optional_array_to_json(values, pg_u32_number)); } if let Ok(values) = row.try_get::<_, Vec>>(idx) { return Some(pg_optional_array_to_json(values, |v| serde_json::Value::Number(v.into()))); } if let Ok(values) = row.try_get::<_, Vec>>(idx) { return Some(pg_optional_array_to_json(values, |v| serde_json::Value::Number(v.into()))); } if let Ok(values) = row.try_get::<_, Vec>>(idx) { return Some(pg_optional_array_to_json(values, |v| serde_json::Value::Number(v.into()))); } if let Ok(values) = row.try_get::<_, Vec>>(idx) { return Some(pg_optional_array_to_json(values, super::safe_i64_to_json)); } if let Ok(values) = row.try_get::<_, Vec>>(idx) { return Some(pg_optional_array_to_json(values, |v| pg_float_number(v as f64))); } if let Ok(values) = row.try_get::<_, Vec>>(idx) { return Some(pg_optional_array_to_json(values, pg_float_number)); } if let Ok(values) = row.try_get::<_, Vec>>(idx) { return Some(pg_optional_array_to_json(values, |v| serde_json::Value::String(v.0))); } None } fn format_pg_timestamptz(value: DateTime) -> String { value.to_rfc3339() } /// 时间类型解码失败后的回退目标,与原 if 链中时间分支之后的匹配顺序一一对应。 #[derive(Clone, Copy, PartialEq, Eq, Debug)] pub(crate) enum PgTemporalFallback { /// 数组类型名(下划线开头)落到通用数组解码。 GenericArray, /// `VECTOR(...)` 形式的类型名落到 pgvector 解码。 Vector, /// 其余落到通用试探链。 Probe, } /// 每列一次的类型分类结果,避免在逐单元格路径上重复 `to_uppercase` 与字符串比较链。 #[derive(Clone, Copy, PartialEq, Eq, Debug)] pub(crate) enum PgColType { Bytea, Json, Bool, Point, Interval, DateRange, Temporal { fallback: PgTemporalFallback }, Numeric, Money, Uuid, Inet { cidr: bool }, MacAddr, BitString, TsVector, SystemU32, InetArray { cidr: bool }, MacAddrArray, BitStringArray, GenericArray, Vector, Geometry, Other, } fn pg_scalar_type_requires_text_protocol(oid: u32, col_type: PgColType) -> bool { Type::from_oid(oid).is_none() && !matches!(col_type, PgColType::Vector | PgColType::Geometry) } /// PostgreSQL `reg*` types are OID-backed, but their useful representation is /// the catalog name returned by the text protocol (for example `regtype` /// returns `integer`). Decoding their binary payload as a generic value leaves /// the raw four-byte OID on the UI boundary, which is rendered as gibberish. fn pg_type_is_reg_type(pg_type: &Type) -> bool { matches!( pg_type.oid(), oid if [ Type::REGPROC.oid(), Type::REGPROCEDURE.oid(), Type::REGOPER.oid(), Type::REGOPERATOR.oid(), Type::REGCLASS.oid(), Type::REGTYPE.oid(), Type::REGNAMESPACE.oid(), Type::REGROLE.oid(), Type::REGCOLLATION.oid(), Type::REGCONFIG.oid(), Type::REGDICTIONARY.oid(), Type::REGPROC_ARRAY.oid(), Type::REGPROCEDURE_ARRAY.oid(), Type::REGOPER_ARRAY.oid(), Type::REGOPERATOR_ARRAY.oid(), Type::REGCLASS_ARRAY.oid(), Type::REGTYPE_ARRAY.oid(), Type::REGNAMESPACE_ARRAY.oid(), Type::REGROLE_ARRAY.oid(), Type::REGCOLLATION_ARRAY.oid(), Type::REGCONFIG_ARRAY.oid(), Type::REGDICTIONARY_ARRAY.oid(), ] .contains(&oid) ) } fn pg_type_requires_text_protocol(pg_type: &Type, col_type: PgColType) -> bool { if pg_type.oid() == Type::RECORD.oid() || pg_type.oid() == Type::RECORD_ARRAY.oid() || pg_type_is_reg_type(pg_type) { return true; } match pg_type.kind() { Kind::Enum(_) => false, Kind::Array(element_type) => Type::from_oid(element_type.oid()).is_none(), Kind::Simple => pg_scalar_type_requires_text_protocol(pg_type.oid(), col_type), _ => Type::from_oid(pg_type.oid()).is_none(), } } pub(crate) fn classify_pg_type(type_name: &str) -> PgColType { let upper = type_name.to_uppercase(); if upper == "BYTEA" { return PgColType::Bytea; } if upper == "JSON" || upper == "JSONB" { return PgColType::Json; } if upper == "BOOL" { return PgColType::Bool; } if upper == "POINT" { return PgColType::Point; } if upper == "INTERVAL" { return PgColType::Interval; } if upper == "DATERANGE" { return PgColType::DateRange; } if upper.contains("TIMESTAMP") || upper == "DATE" || upper == "TIME" || upper == "TIMETZ" || upper.contains("INTERVAL") { let fallback = if upper.starts_with('_') { PgTemporalFallback::GenericArray } else if upper.starts_with("VECTOR(") { PgTemporalFallback::Vector } else { PgTemporalFallback::Probe }; return PgColType::Temporal { fallback }; } if upper == "MONEY" { return PgColType::Money; } if upper == "NUMERIC" || upper == "DECIMAL" { return PgColType::Numeric; } if upper == "UUID" { return PgColType::Uuid; } if matches!(upper.as_str(), "INET" | "CIDR") { return PgColType::Inet { cidr: upper == "CIDR" }; } if matches!(upper.as_str(), "MACADDR" | "MACADDR8") { return PgColType::MacAddr; } if matches!(upper.as_str(), "BIT" | "VARBIT") { return PgColType::BitString; } if upper == "TSVECTOR" { return PgColType::TsVector; } if matches!(upper.as_str(), "OID" | "XID" | "CID") { return PgColType::SystemU32; } if matches!(upper.as_str(), "_INET" | "_CIDR") { return PgColType::InetArray { cidr: upper == "_CIDR" }; } if matches!(upper.as_str(), "_MACADDR" | "_MACADDR8") { return PgColType::MacAddrArray; } if matches!(upper.as_str(), "_BIT" | "_VARBIT") { return PgColType::BitStringArray; } if upper.starts_with('_') { return PgColType::GenericArray; } if upper == "VECTOR" || upper.starts_with("VECTOR(") { return PgColType::Vector; } if upper == "GEOMETRY" || upper == "GEOGRAPHY" { return PgColType::Geometry; } PgColType::Other } pub(crate) fn classify_pg_column_types(column_types: &[String]) -> Vec { column_types.iter().map(|type_name| classify_pg_type(type_name)).collect() } pub(crate) fn pg_value_to_json_classified(row: &Row, idx: usize, col_type: PgColType) -> serde_json::Value { match col_type { PgColType::Bytea => row .try_get::<_, Vec>(idx) .map(|bytes| super::binary_value_to_json(&bytes)) .unwrap_or(serde_json::Value::Null), PgColType::Json => { if let Ok(v) = row.try_get::<_, serde_json::Value>(idx) { return serde_json::Value::String(v.to_string()); } if let Ok(v) = row.try_get::<_, String>(idx) { return serde_json::Value::String(v); } serde_json::Value::Null } PgColType::Bool => pg_bool_value_to_json(row, idx), PgColType::Point => row .try_get::<_, PgPoint>(idx) .map(|point| serde_json::Value::String(format_pg_point(point))) .unwrap_or(serde_json::Value::Null), PgColType::Interval => row .try_get::<_, PgInterval>(idx) .map(|interval| serde_json::Value::String(format_pg_interval(interval))) .unwrap_or_else(|_| pg_fallback_value_to_json(row, idx)), PgColType::DateRange => row .try_get::<_, PgDateRange>(idx) .map(|range| serde_json::Value::String(range.0)) .unwrap_or_else(|_| pg_fallback_value_to_json(row, idx)), PgColType::Temporal { fallback } => { if let Some(v) = pg_temporal_to_json_value(row, idx) { return v; } match fallback { PgTemporalFallback::GenericArray => pg_array_to_json_value(row, idx).unwrap_or(serde_json::Value::Null), PgTemporalFallback::Vector => pg_vector_value_to_json(row, idx), PgTemporalFallback::Probe => pg_fallback_value_to_json(row, idx), } } PgColType::Numeric => { row.try_get::<_, PgNumeric>(idx).map(|v| serde_json::Value::String(v.0)).unwrap_or(serde_json::Value::Null) } PgColType::Money => row .try_get::<_, PgMoney>(idx) .map(|v| serde_json::Value::String(format_pg_money(v.0))) .unwrap_or(serde_json::Value::Null), PgColType::Uuid => row .try_get::<_, uuid::Uuid>(idx) .map(|v| serde_json::Value::String(v.to_string())) .unwrap_or(serde_json::Value::Null), PgColType::Inet { cidr } => pg_network_address_to_json_value(row, idx, cidr).unwrap_or(serde_json::Value::Null), PgColType::MacAddr => pg_macaddr_to_json_value(row, idx).unwrap_or(serde_json::Value::Null), PgColType::BitString => pg_bit_string_to_json_value(row, idx).unwrap_or(serde_json::Value::Null), PgColType::TsVector => row .try_get::<_, PgRawBytes>(idx) .ok() .and_then(|raw| decode_tsvector_bytes(&raw.0)) .map(serde_json::Value::String) .unwrap_or(serde_json::Value::Null), PgColType::SystemU32 => pg_system_u32_to_json(row, idx).unwrap_or(serde_json::Value::Null), PgColType::InetArray { cidr } => { pg_network_address_array_to_json_value(row, idx, cidr).unwrap_or(serde_json::Value::Null) } PgColType::MacAddrArray => pg_macaddr_array_to_json_value(row, idx).unwrap_or(serde_json::Value::Null), PgColType::BitStringArray => pg_bit_string_array_to_json_value(row, idx).unwrap_or(serde_json::Value::Null), PgColType::GenericArray => pg_array_to_json_value(row, idx).unwrap_or(serde_json::Value::Null), PgColType::Vector => pg_vector_value_to_json(row, idx), PgColType::Geometry => { if let Ok(PgRawBytes(raw)) = row.try_get::<_, PgRawBytes>(idx) { return super::wkb::wkb_to_wkt(&raw) .map(serde_json::Value::String) .unwrap_or_else(|| super::binary_value_to_json(&raw)); } serde_json::Value::Null } PgColType::Other => pg_fallback_value_to_json(row, idx), } } fn pg_value_to_json_with_srid(row: &Row, idx: usize, col_type: PgColType) -> (serde_json::Value, Option) { if col_type != PgColType::Geometry { return (pg_value_to_json_classified(row, idx, col_type), None); } if let Ok(PgRawBytes(raw)) = row.try_get::<_, PgRawBytes>(idx) { return match super::wkb::decode_wkb_geometry(&raw) { Some(geometry) => (serde_json::Value::String(geometry.wkt), geometry.srid), None => (super::binary_value_to_json(&raw), None), }; } (serde_json::Value::Null, None) } fn pg_text_fallback_value(value: &str, col_type: Option) -> (serde_json::Value, Option) { let (value, srid, _) = pg_text_fallback_value_with_spatial(value, col_type); (value, srid) } fn pg_text_fallback_value_with_spatial( value: &str, col_type: Option, ) -> (serde_json::Value, Option, bool) { match col_type { Some(PgColType::Geometry) => decode_pg_text_wkb(value) .map(|geometry| (serde_json::Value::String(geometry.wkt), geometry.srid, true)) .or_else(|| split_pg_ewkt(value, false).map(|(value, srid)| (value, srid, true))) .unwrap_or_else(|| (serde_json::Value::String(value.to_string()), None, true)), Some(_) => (serde_json::Value::String(value.to_string()), None, false), None => decode_pg_text_wkb(value) .map(|geometry| (serde_json::Value::String(geometry.wkt), geometry.srid, true)) .or_else(|| split_pg_ewkt(value, true).map(|(value, srid)| (value, srid, true))) .unwrap_or_else(|| (serde_json::Value::String(value.to_string()), None, false)), } } fn decode_pg_text_wkb(value: &str) -> Option { let trimmed = value.trim(); let hex = trimmed .strip_prefix("0x") .or_else(|| trimmed.strip_prefix("0X")) .or_else(|| trimmed.strip_prefix("\\x")) .or_else(|| trimmed.strip_prefix("\\X")) .unwrap_or(trimmed); if hex.len() < 10 || !hex.len().is_multiple_of(2) || !matches!(&hex[..2], "00" | "01") || !hex.is_ascii() { return None; } let bytes = hex .as_bytes() .chunks_exact(2) .map(|pair| { let text = std::str::from_utf8(pair).ok()?; u8::from_str_radix(text, 16).ok() }) .collect::>>()?; super::wkb::decode_wkb_geometry(&bytes) } fn split_pg_ewkt(value: &str, require_recognizable_wkt: bool) -> Option<(serde_json::Value, Option)> { let rest = value.strip_prefix("SRID=")?; let (srid, wkt) = rest.split_once(';')?; let srid = srid.parse::().ok().and_then(|value| u32::try_from(value).ok())?; if require_recognizable_wkt && !is_recognizable_wkt(wkt) { return None; } Some((serde_json::Value::String(wkt.to_string()), (srid != 0).then_some(srid))) } fn is_recognizable_wkt(value: &str) -> bool { let trimmed = value.trim_start(); const TYPES: [&str; 7] = ["POINT", "LINESTRING", "POLYGON", "MULTIPOINT", "MULTILINESTRING", "MULTIPOLYGON", "GEOMETRYCOLLECTION"]; TYPES.iter().any(|geometry_type| { trimmed .get(..geometry_type.len()) .filter(|prefix| prefix.eq_ignore_ascii_case(geometry_type)) .and_then(|_| trimmed.as_bytes().get(geometry_type.len())) .is_some_and(|next| next.is_ascii_whitespace() || *next == b'(') }) } /// Serialize a pgvector `vector` component with f32 shortest round-trip decimal text. /// /// Casting through `f64` (or fixed fractional rounding) either expands binary noise or /// truncates remaining single-precision digits; formatting via `f32` display keeps the /// full float4 value that pgvector stores. fn pg_vector_element_number(v: f32) -> serde_json::Value { v.to_string().parse().map(serde_json::Value::Number).unwrap_or(serde_json::Value::Null) } fn pg_vector_value_to_json(row: &Row, idx: usize) -> serde_json::Value { if let Ok(PgRawBytes(raw)) = row.try_get::<_, PgRawBytes>(idx) { if let Some(floats) = decode_pgvector_bytes(&raw) { return serde_json::Value::Array(floats.into_iter().map(pg_vector_element_number).collect()); } } serde_json::Value::Null } fn pg_fallback_value_to_json(row: &Row, idx: usize) -> serde_json::Value { row.try_get::<_, String>(idx) .map(serde_json::Value::String) .or_else(|e| pg_system_u32_to_json(row, idx).ok_or(e)) .or_else(|_| row.try_get::<_, i64>(idx).map(super::safe_i64_to_json)) .or_else(|_| row.try_get::<_, i32>(idx).map(|v| serde_json::Value::Number(v.into()))) .or_else(|_| row.try_get::<_, i16>(idx).map(|v| serde_json::Value::Number(v.into()))) .or_else(|_| row.try_get::<_, i8>(idx).map(|v| serde_json::Value::Number(v.into()))) .or_else(|e| pg_array_to_json_value(row, idx).ok_or(e)) .or_else(|_| { row.try_get::<_, f64>(idx).map(|v| { serde_json::Number::from_f64(v).map(serde_json::Value::Number).unwrap_or(serde_json::Value::Null) }) }) .or_else(|_| { row.try_get::<_, f32>(idx).map(|v| { serde_json::Number::from_f64((v as f64 * 1_000_000.0).round() / 1_000_000.0) .map(serde_json::Value::Number) .unwrap_or(serde_json::Value::Null) }) }) .or_else(|_| row.try_get::<_, bool>(idx).map(serde_json::Value::Bool)) .or_else(|_| row.try_get::<_, uuid::Uuid>(idx).map(|v| serde_json::Value::String(v.to_string()))) .or_else(|e| pg_temporal_to_json_value(row, idx).ok_or(e)) .or_else(|_| row.try_get::<_, Vec>(idx).map(|bytes| super::binary_value_to_json(&bytes))) .or_else(|_| row.try_get::<_, PgAnyString>(idx).map(|v| serde_json::Value::String(v.0))) .or_else(|_| row.try_get::<_, PgRawBytes>(idx).map(|v| super::binary_value_to_json(&v.0))) .unwrap_or(serde_json::Value::Null) } fn decode_tsvector_bytes(raw: &[u8]) -> Option { let mut cursor = 0; let count = read_i32_be(raw, &mut cursor)?; if count < 0 { return None; } let mut entries = Vec::with_capacity(count as usize); for _ in 0..count { let start = cursor; while cursor < raw.len() && raw[cursor] != 0 { cursor += 1; } if cursor >= raw.len() { return None; } let lexeme = std::str::from_utf8(&raw[start..cursor]).ok()?; cursor += 1; let position_count = read_u16_be(raw, &mut cursor)? as usize; let mut positions = Vec::with_capacity(position_count); for _ in 0..position_count { let encoded = read_u16_be(raw, &mut cursor)?; let position = encoded & 0x3fff; let weight = match encoded >> 14 { 3 => "A", 2 => "B", 1 => "C", _ => "", }; positions.push(format!("{position}{weight}")); } let mut entry = format!("'{}'", escape_tsvector_lexeme(lexeme)); if !positions.is_empty() { entry.push(':'); entry.push_str(&positions.join(",")); } entries.push(entry); } if cursor == raw.len() { Some(entries.join(" ")) } else { None } } fn read_i32_be(raw: &[u8], cursor: &mut usize) -> Option { let bytes: [u8; 4] = raw.get(*cursor..*cursor + 4)?.try_into().ok()?; *cursor += 4; Some(i32::from_be_bytes(bytes)) } fn read_u16_be(raw: &[u8], cursor: &mut usize) -> Option { let bytes: [u8; 2] = raw.get(*cursor..*cursor + 2)?.try_into().ok()?; *cursor += 2; Some(u16::from_be_bytes(bytes)) } fn escape_tsvector_lexeme(value: &str) -> String { value.replace('\\', "\\\\").replace('\'', "''") } fn pg_error_to_string(err: tokio_postgres::Error) -> String { err.as_db_error().map(ToString::to_string).unwrap_or_else(|| err.to_string()) } /// Tries each SQL tier in `tiers` in order (most-capable first), via `run`, /// returning the first tier that succeeds. Every driver-compat query in this /// module (a "does this server have the newer catalog column" primary/compat /// split, occasionally with a further information_schema fallback) used to /// hand-roll this same try/log/combine-errors shape once per query; this is /// the shared version. /// /// If every tier fails, all of their errors are logged together at debug /// level (so a fallback firing in production is diagnosable) and the last /// tier's error is returned to the caller, since it's usually the most /// specific one for whatever the connected server actually is. async fn query_with_compat_fallback( log_context: &str, tiers: &[&'static str], mut run: F, ) -> Result where F: FnMut(&'static str) -> Fut, Fut: std::future::Future>, { let mut errors: Vec = Vec::new(); for sql in tiers { match run(sql).await { Ok(value) => return Ok(value), Err(error) => errors.push(pg_error_to_string(error)), } } log::debug!("[postgres][{log_context}:compat-failed] {}", errors.join("; ")); Err(errors.into_iter().next_back().unwrap_or_else(|| format!("[postgres][{log_context}] no SQL tiers configured"))) } fn pg_db_error_to_string(err: &tokio_postgres::error::DbError) -> String { format!("{err} (SQLSTATE {})", err.code().code()) } fn pg_error_from_sources(err: &(dyn std::error::Error + 'static)) -> Option { let mut current = Some(err); while let Some(source) = current { if let Some(pg_error) = source.downcast_ref::() { if let Some(db_error) = pg_error.as_db_error() { return Some(pg_db_error_to_string(db_error)); } } if let Some(db_error) = source.downcast_ref::() { return Some(pg_db_error_to_string(db_error)); } current = source.source(); } None } fn error_with_sources_to_string(err: &(dyn std::error::Error + 'static)) -> String { let mut messages = vec![err.to_string()]; let mut current = err.source(); while let Some(source) = current { let message = source.to_string(); if !messages.iter().any(|existing| existing == &message) { messages.push(message); } current = source.source(); } messages.join(": ") } fn pg_pool_error_to_string(err: PoolError) -> String { pg_error_from_sources(&err).unwrap_or_else(|| error_with_sources_to_string(&err)) } fn should_retry_postgres_text_query(err: &tokio_postgres::Error) -> bool { let message = err.as_db_error().map(ToString::to_string).unwrap_or_else(|| err.to_string()).to_ascii_lowercase(); should_retry_postgres_text_query_message(&message) } fn should_retry_postgres_text_query_message(message: &str) -> bool { message.contains("no binary output function") || message.contains("no binary send function") || message.contains("cannot display a value of type") } fn should_retry_postgres_stale_cache(err: &tokio_postgres::Error) -> bool { if let Some(db_error) = err.as_db_error() { return should_retry_postgres_stale_cache_fields( Some(db_error.code().code()), db_error.routine(), db_error.message(), ); } should_retry_postgres_stale_cache_fields(None, None, &err.to_string()) } fn should_retry_postgres_stale_cache_fields(sqlstate: Option<&str>, routine: Option<&str>, message: &str) -> bool { let structured_match = sqlstate == Some("0A000") && routine.is_some_and(|routine| routine.eq_ignore_ascii_case("RevalidateCachedQuery")); structured_match || message.to_ascii_lowercase().contains("cached plan must not change result type") } fn should_fallback_postgres_missing_prepared_statement(err: &tokio_postgres::Error) -> bool { if let Some(db_error) = err.as_db_error() { return should_fallback_postgres_missing_prepared_statement_fields( Some(db_error.code().code()), db_error.message(), ); } should_fallback_postgres_missing_prepared_statement_fields(None, &err.to_string()) } fn should_fallback_postgres_missing_prepared_statement_fields(sqlstate: Option<&str>, message: &str) -> bool { if sqlstate == Some("26000") { return true; } let message = message.to_ascii_lowercase(); message.contains("prepared statement") && message.contains("does not exist") } fn postgres_unnamed_statement_clients() -> &'static Mutex>> { static CLIENTS: OnceLock>>> = OnceLock::new(); CLIENTS.get_or_init(|| Mutex::new(HashMap::new())) } fn postgres_client_uses_unnamed_statements(client: &deadpool_postgres::Client) -> bool { let statement_cache = &client.statement_cache; let key = Arc::as_ptr(statement_cache) as usize; let mut clients = postgres_unnamed_statement_clients().lock().unwrap_or_else(|poisoned| poisoned.into_inner()); match clients.get(&key).and_then(Weak::upgrade) { Some(cached) if Arc::ptr_eq(&cached, statement_cache) => true, _ => { clients.remove(&key); false } } } fn mark_postgres_client_unnamed_statements(client: &deadpool_postgres::Client) { let statement_cache = &client.statement_cache; let key = Arc::as_ptr(statement_cache) as usize; let mut clients = postgres_unnamed_statement_clients().lock().unwrap_or_else(|poisoned| poisoned.into_inner()); clients.retain(|_, cached| cached.strong_count() > 0); clients.insert(key, Arc::downgrade(statement_cache)); } fn postgres_typed_params<'a>( params: &[&'a (dyn tokio_postgres::types::ToSql + Sync)], param_types: &[Type], ) -> Option> { (params.len() == param_types.len()).then(|| params.iter().copied().zip(param_types.iter().cloned()).collect()) } async fn postgres_query_unnamed( client: &deadpool_postgres::Client, sql: &str, ) -> Result { client.query_typed_raw(sql, std::iter::empty::<(&(dyn tokio_postgres::types::ToSql + Sync), Type)>()).await } async fn postgres_query_cached( client: &deadpool_postgres::Client, sql: &str, params: &[&(dyn tokio_postgres::types::ToSql + Sync)], ) -> Result, tokio_postgres::Error> { if postgres_client_uses_unnamed_statements(client) && params.is_empty() { return client.query_typed(sql, &[]).await; } let stmt = client.prepare_cached(sql).await?; if postgres_client_uses_unnamed_statements(client) { let Some(typed_params) = postgres_typed_params(params, stmt.params()) else { return client.query(&stmt, params).await; }; return client.query_typed(sql, &typed_params).await; } match client.query(&stmt, params).await { Ok(rows) => Ok(rows), Err(err) if should_fallback_postgres_missing_prepared_statement(&err) => { let Some(typed_params) = postgres_typed_params(params, stmt.params()) else { return Err(err); }; log::warn!( "[postgres][metadata:missing_prepared_statement] downgrading connection to unnamed statements: {}", pg_error_to_string(err) ); mark_postgres_client_unnamed_statements(client); client.query_typed(sql, &typed_params).await } Err(err) if should_retry_postgres_stale_cache(&err) => { // Metadata queries can be cached while a table/view definition is // changed from another session. Evict and retry once with fresh // statement/type metadata instead of surfacing PostgreSQL's stale // cached-plan error to the UI. log::warn!("[postgres][metadata:stale_cache] evicting cached statement: {}", pg_error_to_string(err)); client.statement_cache.remove(sql, &[]); client.clear_type_cache(); let stmt = client.prepare_cached(sql).await?; client.query(&stmt, params).await } Err(err) => Err(err), } } async fn postgres_query_one_cached( client: &deadpool_postgres::Client, sql: &str, params: &[&(dyn tokio_postgres::types::ToSql + Sync)], ) -> Result { if postgres_client_uses_unnamed_statements(client) && params.is_empty() { return client.query_typed_one(sql, &[]).await; } let stmt = client.prepare_cached(sql).await?; if postgres_client_uses_unnamed_statements(client) { let Some(typed_params) = postgres_typed_params(params, stmt.params()) else { return client.query_one(&stmt, params).await; }; return client.query_typed_one(sql, &typed_params).await; } match client.query_one(&stmt, params).await { Ok(row) => Ok(row), Err(err) if should_fallback_postgres_missing_prepared_statement(&err) => { let Some(typed_params) = postgres_typed_params(params, stmt.params()) else { return Err(err); }; log::warn!( "[postgres][metadata_one:missing_prepared_statement] downgrading connection to unnamed statements: {}", pg_error_to_string(err) ); mark_postgres_client_unnamed_statements(client); client.query_typed_one(sql, &typed_params).await } Err(err) if should_retry_postgres_stale_cache(&err) => { // Same stale-cache protection as postgres_query_cached, for scalar // catalog probes such as pg_proc feature detection. log::warn!("[postgres][metadata:stale_cache] evicting cached statement: {}", pg_error_to_string(err)); client.statement_cache.remove(sql, &[]); client.clear_type_cache(); let stmt = client.prepare_cached(sql).await?; client.query_one(&stmt, params).await } Err(err) => Err(err), } } #[allow(clippy::large_enum_variant)] enum PreparedSelectOutcome { Complete(Box), TextFallback { column_types: Vec, unsupported_type: String }, } struct PreparedSelectMetadata { columns: Vec, column_types: Vec, column_classes: Vec, unsupported_type: Option, } fn prepared_select_metadata(columns: &[tokio_postgres::Column]) -> PreparedSelectMetadata { let column_names: Vec = columns.iter().map(|c| c.name().to_string()).collect(); let column_types: Vec = columns.iter().map(|c| c.type_().name().to_string()).collect(); let column_classes = classify_pg_column_types(&column_types); let unsupported_type = columns.iter().zip(&column_classes).find_map(|(column, col_type)| { let pg_type = column.type_(); pg_type_requires_text_protocol(pg_type, *col_type).then(|| pg_type.name().to_string()) }); PreparedSelectMetadata { columns: column_names, column_types, column_classes, unsupported_type } } async fn prepare_select_with_metadata( client: &deadpool_postgres::Client, sql: &str, ) -> Result<(tokio_postgres::Statement, PreparedSelectMetadata), tokio_postgres::Error> { let mut stmt = client.prepare_cached(sql).await?; let mut metadata = prepared_select_metadata(stmt.columns()); if metadata.unsupported_type.is_some() { stmt = client.prepare(sql).await?; metadata = prepared_select_metadata(stmt.columns()); } Ok((stmt, metadata)) } enum PostgresSelectStreamOutcome { Binary { stream: tokio_postgres::RowStream, metadata: PreparedSelectMetadata }, TextFallback { column_types: Vec, unsupported_type: String }, } fn postgres_select_stream_outcome(stream: tokio_postgres::RowStream) -> PostgresSelectStreamOutcome { let metadata = prepared_select_metadata(stream.columns()); if let Some(unsupported_type) = metadata.unsupported_type.clone() { return PostgresSelectStreamOutcome::TextFallback { column_types: metadata.column_types, unsupported_type }; } PostgresSelectStreamOutcome::Binary { stream, metadata } } async fn prepare_unnamed_select_metadata( client: &deadpool_postgres::Client, sql: &str, ) -> Result { // query_typed_raw sends Describe and Execute together. If its result has an // unknown user-defined type, tokio-postgres then resolves that type with a // second query on the same connection before returning the RowStream. A // large result can fill the connection buffers and block that type lookup. // Describe once without executing so custom types are cached before the // actual unnamed stream starts. The stream remains unnamed to tolerate // server-side prepared statement loss in snapshot/proxy environments. let stmt = client.prepare(sql).await?; Ok(prepared_select_metadata(stmt.columns())) } async fn start_postgres_select_stream( client: &deadpool_postgres::Client, sql: &str, force_unnamed: bool, ) -> Result { if force_unnamed || postgres_client_uses_unnamed_statements(client) { let metadata = prepare_unnamed_select_metadata(client, sql).await?; if let Some(unsupported_type) = metadata.unsupported_type { return Ok(PostgresSelectStreamOutcome::TextFallback { column_types: metadata.column_types, unsupported_type, }); } return postgres_query_unnamed(client, sql).await.map(postgres_select_stream_outcome); } let (stmt, metadata) = prepare_select_with_metadata(client, sql).await?; if let Some(unsupported_type) = metadata.unsupported_type { return Ok(PostgresSelectStreamOutcome::TextFallback { column_types: metadata.column_types, unsupported_type }); } let params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> = Vec::new(); match client.query_raw(&stmt, params).await { Ok(stream) => Ok(postgres_select_stream_outcome(stream)), Err(err) if should_fallback_postgres_missing_prepared_statement(&err) => { // Bind failed before execution. Downgrade this physical connection // so later queries do not repeat the named-statement round trip. log::warn!( "[postgres][prepared_statement:missing] downgrading connection to unnamed statements: {}", pg_error_to_string(err) ); mark_postgres_client_unnamed_statements(client); postgres_query_unnamed(client, sql).await.map(postgres_select_stream_outcome) } Err(err) => Err(err), } } async fn execute_select_prepared( client: &deadpool_postgres::Client, sql: &str, start: Instant, row_limit: usize, progress_clock: Option<&StreamProgressClock>, force_unnamed: bool, ) -> Result { let stream_start = Instant::now(); let stream_outcome = start_postgres_select_stream(client, sql, force_unnamed).await?; if let Some(progress_clock) = progress_clock { progress_clock.mark(); } log::info!( "[postgres][select:stream_ready] elapsed_ms={} total_ms={}", stream_start.elapsed().as_millis(), start.elapsed().as_millis() ); let (stream, metadata) = match stream_outcome { PostgresSelectStreamOutcome::Binary { stream, metadata } => (stream, metadata), PostgresSelectStreamOutcome::TextFallback { column_types, unsupported_type } => { return Ok(PreparedSelectOutcome::TextFallback { column_types, unsupported_type }); } }; let PreparedSelectMetadata { columns, column_types, column_classes, unsupported_type: _ } = metadata; if let Some(progress_clock) = progress_clock { progress_clock.mark(); } log::info!( "[postgres][select:metadata_ready] total_ms={} column_count={}", start.elapsed().as_millis(), columns.len() ); tokio::pin!(stream); let mut result_rows: Vec> = Vec::new(); let mut spatial_values: Vec>> = Vec::new(); let mut spatial_columns = SpatialColumnBuilder::new( column_classes .iter() .enumerate() .filter_map(|(index, col_type)| (*col_type == PgColType::Geometry).then_some(index)), ); let mut truncated = false; let rows_start = Instant::now(); while let Some(row_result) = stream.next().await { if let Some(progress_clock) = progress_clock { progress_clock.mark(); } if result_rows.len() >= row_limit { truncated = true; break; } let row = row_result?; let mut values = Vec::with_capacity(row.columns().len()); let mut row_srids = vec![None; row.columns().len()]; for (i, row_srid) in row_srids.iter_mut().enumerate() { let col_type = column_classes.get(i).copied().unwrap_or(PgColType::Other); let (value, srid) = pg_value_to_json_with_srid(&row, i, col_type); if col_type == PgColType::Geometry { spatial_columns.observe(i, srid); *row_srid = srid; } values.push(value); } result_rows.push(values); spatial_values.push(row_srids); } log::info!( "[postgres][select:rows:done] elapsed_ms={} total_ms={} row_count={} truncated={}", rows_start.elapsed().as_millis(), start.elapsed().as_millis(), result_rows.len(), truncated ); let (spatial_columns, spatial_values) = spatial_columns.finish_with_values(spatial_values); Ok(PreparedSelectOutcome::Complete(Box::new(QueryResult { columns, column_types, column_sortables: Vec::new(), spatial_columns, spatial_values, rows: result_rows, affected_rows: 0, execution_time_ms: start.elapsed().as_millis(), truncated, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: Vec::new(), }))) } fn matching_pg_text_column_types(columns: &[String], prepared: Option>) -> Vec { prepared.filter(|types| types.len() == columns.len()).unwrap_or_default() } async fn execute_select_text( client: &deadpool_postgres::Client, sql: &str, start: Instant, row_limit: usize, prepared_column_types: Option>, progress_clock: Option<&StreamProgressClock>, ) -> Result { let stream = client.simple_query_raw(sql).await.map_err(pg_error_to_string)?; if let Some(progress_clock) = progress_clock { progress_clock.mark(); } tokio::pin!(stream); let mut columns: Vec = Vec::new(); let mut result_rows: Vec> = Vec::new(); let mut spatial_values: Vec>> = Vec::new(); let column_classes = prepared_column_types.as_ref().map(|types| classify_pg_column_types(types)).unwrap_or_default(); let mut spatial_columns = SpatialColumnBuilder::new( column_classes .iter() .enumerate() .filter_map(|(index, col_type)| (*col_type == PgColType::Geometry).then_some(index)), ); let mut truncated = false; while let Some(message) = stream.next().await { if let Some(progress_clock) = progress_clock { progress_clock.mark(); } match message { Ok(SimpleQueryMessage::RowDescription(cols)) => { columns = cols.iter().map(|c| c.name().to_string()).collect(); } Ok(SimpleQueryMessage::Row(row)) => { if columns.is_empty() { columns = row.columns().iter().map(|c| c.name().to_string()).collect(); } if result_rows.len() >= row_limit { truncated = true; break; } let mut values = Vec::with_capacity(row.len()); let mut row_srids = vec![None; row.len()]; for (i, row_srid) in row_srids.iter_mut().enumerate() { match row.try_get(i).map_err(pg_error_to_string)? { Some(value) => { let (decoded, srid, is_spatial) = pg_text_fallback_value_with_spatial(value, column_classes.get(i).copied()); values.push(decoded); if is_spatial { spatial_columns.observe(i, srid); *row_srid = srid; } } None => { values.push(serde_json::Value::Null); } } } result_rows.push(values); spatial_values.push(row_srids); } Err(_) if result_rows.len() >= row_limit => { truncated = true; break; } Err(err) => return Err(pg_error_to_string(err)), Ok(SimpleQueryMessage::CommandComplete(_)) => {} Ok(_) => {} } } let (spatial_columns, spatial_values) = spatial_columns.finish_with_values(spatial_values); Ok(QueryResult { column_types: matching_pg_text_column_types(&columns, prepared_column_types), columns, column_sortables: Vec::new(), spatial_columns, spatial_values, rows: result_rows, affected_rows: 0, execution_time_ms: start.elapsed().as_millis(), truncated, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: Vec::new(), }) } async fn finish_prepared_select( client: &deadpool_postgres::Client, sql: &str, start: Instant, row_limit: usize, outcome: PreparedSelectOutcome, progress_clock: Option<&StreamProgressClock>, ) -> Result { match outcome { PreparedSelectOutcome::Complete(result) => Ok(*result), PreparedSelectOutcome::TextFallback { column_types, unsupported_type } => { log::info!( "[postgres][select:text_fallback] unsupported_type={} switching_to=simple_query", unsupported_type ); execute_select_text(client, sql, start, row_limit, Some(column_types), progress_clock).await } } } pub(crate) async fn execute_select_query( client: &deadpool_postgres::Client, sql: &str, start: Instant, row_limit: usize, ) -> Result { execute_select_query_with_progress(client, sql, start, row_limit, None, false).await } pub(crate) async fn execute_select_query_unnamed( client: &deadpool_postgres::Client, sql: &str, start: Instant, row_limit: usize, ) -> Result { execute_select_query_with_progress(client, sql, start, row_limit, None, true).await } async fn execute_select_query_with_progress( client: &deadpool_postgres::Client, sql: &str, start: Instant, row_limit: usize, progress_clock: Option<&StreamProgressClock>, force_unnamed: bool, ) -> Result { match execute_select_prepared(client, sql, start, row_limit, progress_clock, force_unnamed).await { Ok(outcome) => finish_prepared_select(client, sql, start, row_limit, outcome, progress_clock).await, Err(err) if should_retry_postgres_stale_cache(&err) => { // The cached prepared statement is stale (e.g. the view or table // schema changed since the statement was prepared). Evict the // stale entry and retry with a fresh server-side prepare. log::warn!("[postgres][select:stale_cache] evicting cached statement: {}", pg_error_to_string(err)); client.statement_cache.remove(sql, &[]); match execute_select_prepared(client, sql, start, row_limit, progress_clock, force_unnamed).await { Ok(outcome) => finish_prepared_select(client, sql, start, row_limit, outcome, progress_clock).await, Err(err) if should_retry_postgres_text_query(&err) => { execute_select_text(client, sql, start, row_limit, None, progress_clock).await } Err(err) => Err(pg_error_to_string(err)), } } Err(err) if should_retry_postgres_text_query(&err) => { execute_select_text(client, sql, start, row_limit, None, progress_clock).await } Err(err) => Err(pg_error_to_string(err)), } } pub enum PostgresQueryStreamItem { Columns { columns: Vec, column_types: Vec }, Row(Vec), } enum PostgresQueryStreamError { Postgres { err: tokio_postgres::Error, emitted: bool }, TextFallback { column_types: Vec, unsupported_type: String }, Export(String), } impl PostgresQueryStreamError { fn into_string(self) -> String { match self { Self::Postgres { err, .. } => pg_error_to_string(err), Self::TextFallback { unsupported_type, .. } => { format!("PostgreSQL type {unsupported_type} requires text protocol") } Self::Export(err) => err, } } } async fn stream_select_query_prepared( client: &deadpool_postgres::Client, sql: &str, row_limit: Option, on_item: &mut impl FnMut(PostgresQueryStreamItem) -> Result<(), String>, force_unnamed: bool, ) -> Result { let stream_outcome = start_postgres_select_stream(client, sql, force_unnamed) .await .map_err(|err| PostgresQueryStreamError::Postgres { err, emitted: false })?; let (stream, metadata) = match stream_outcome { PostgresSelectStreamOutcome::Binary { stream, metadata } => (stream, metadata), PostgresSelectStreamOutcome::TextFallback { column_types, unsupported_type } => { return Err(PostgresQueryStreamError::TextFallback { column_types, unsupported_type }); } }; let PreparedSelectMetadata { columns, column_types, column_classes, unsupported_type: _ } = metadata; tokio::pin!(stream); let mut rows_streamed = 0_u64; let mut columns_emitted = false; while let Some(row_result) = stream.next().await { if row_limit.is_some_and(|limit| rows_streamed as usize >= limit) { break; } let row = row_result .map_err(|err| PostgresQueryStreamError::Postgres { err, emitted: columns_emitted || rows_streamed > 0 })?; if !columns_emitted { on_item(PostgresQueryStreamItem::Columns { columns: columns.clone(), column_types: column_types.clone() }) .map_err(PostgresQueryStreamError::Export)?; columns_emitted = true; } let values = (0..row.columns().len()) .map(|i| pg_value_to_json_classified(&row, i, column_classes.get(i).copied().unwrap_or(PgColType::Other))) .collect(); on_item(PostgresQueryStreamItem::Row(values)).map_err(PostgresQueryStreamError::Export)?; rows_streamed += 1; } if !columns_emitted { on_item(PostgresQueryStreamItem::Columns { columns, column_types }) .map_err(PostgresQueryStreamError::Export)?; } Ok(rows_streamed) } async fn stream_select_query_text( client: &deadpool_postgres::Client, sql: &str, row_limit: Option, prepared_column_types: Option>, on_item: &mut impl FnMut(PostgresQueryStreamItem) -> Result<(), String>, ) -> Result { let stream = client.simple_query_raw(sql).await.map_err(pg_error_to_string)?; tokio::pin!(stream); let mut columns: Vec = Vec::new(); let column_classes = prepared_column_types.as_ref().map(|types| classify_pg_column_types(types)).unwrap_or_default(); let mut rows_streamed = 0_u64; while let Some(message) = stream.next().await { match message.map_err(pg_error_to_string)? { SimpleQueryMessage::RowDescription(cols) => { columns = cols.iter().map(|c| c.name().to_string()).collect(); let column_types = matching_pg_text_column_types(&columns, prepared_column_types.clone()); on_item(PostgresQueryStreamItem::Columns { columns: columns.clone(), column_types })?; } SimpleQueryMessage::Row(row) => { if row_limit.is_some_and(|limit| rows_streamed as usize >= limit) { break; } if columns.is_empty() { columns = row.columns().iter().map(|c| c.name().to_string()).collect(); let column_types = matching_pg_text_column_types(&columns, prepared_column_types.clone()); on_item(PostgresQueryStreamItem::Columns { columns: columns.clone(), column_types })?; } let mut values = Vec::with_capacity(row.len()); for i in 0..row.len() { values.push(match row.try_get(i).map_err(pg_error_to_string)? { Some(value) => pg_text_fallback_value(value, column_classes.get(i).copied()).0, None => serde_json::Value::Null, }); } on_item(PostgresQueryStreamItem::Row(values))?; rows_streamed += 1; } SimpleQueryMessage::CommandComplete(_) => {} _ => {} } } Ok(rows_streamed) } #[cfg(test)] pub(crate) async fn stream_select_query_inner_unnamed( client: &deadpool_postgres::Client, sql: &str, row_limit: Option, on_item: &mut impl FnMut(PostgresQueryStreamItem) -> Result<(), String>, ) -> Result { stream_select_query_inner_with_mode(client, sql, row_limit, on_item, true).await } pub(crate) async fn stream_select_query_inner_unnamed_with_cancel( client: &deadpool_postgres::Client, sql: &str, row_limit: Option, on_item: &mut impl FnMut(PostgresQueryStreamItem) -> Result<(), String>, cancel_token: Option<&CancellationToken>, budget: &DbOperationBudget, cancel_context: Option<&PostgresCancelContext>, ) -> Result { if cancel_token.is_some_and(CancellationToken::is_cancelled) { return Err(crate::query::canceled_error()); } let pg_cancel_token = client.cancel_token(); let progress_clock = Arc::new(StreamProgressClock::new()); let progress_clock_for_stream = progress_clock.clone(); let mut on_stream_item = |item| { let row_received = matches!(&item, PostgresQueryStreamItem::Row(_)); on_item(item)?; if row_received { progress_clock_for_stream.mark(); } Ok(()) }; let timeout_error = format!("Query timed out after {} seconds", budget.query_timeout.map_or(0, |timeout| timeout.as_secs())); let stream = stream_select_query_inner_with_mode(client, sql, row_limit, &mut on_stream_item, true); tokio::pin!(stream); let result = await_stream_with_progress_timeout( stream.as_mut(), budget.query_timeout, progress_clock, cancel_token, timeout_error.clone(), ) .await; if result.as_ref().is_err_and(|error| error == &timeout_error || error == crate::query::QUERY_CANCELED) { let original_error = result.as_ref().expect_err("timeout or cancellation result").clone(); cancel_postgres_query(pg_cancel_token, cancel_context, budget.cancel_timeout).await; if tokio::time::timeout(budget.cleanup_timeout, stream.as_mut()).await.is_err() { return Err(format!( "{}; PostgreSQL stream cleanup timed out after {} seconds", original_error, budget.cleanup_timeout.as_secs() )); } } result } async fn stream_select_query_inner_with_mode( client: &deadpool_postgres::Client, sql: &str, row_limit: Option, on_item: &mut impl FnMut(PostgresQueryStreamItem) -> Result<(), String>, force_unnamed: bool, ) -> Result { match stream_select_query_prepared(client, sql, row_limit, on_item, force_unnamed).await { Ok(rows) => Ok(rows), Err(PostgresQueryStreamError::TextFallback { column_types, unsupported_type }) => { log::info!( "[postgres][stream:text_fallback] unsupported_type={} switching_to=simple_query", unsupported_type ); stream_select_query_text(client, sql, row_limit, Some(column_types), on_item).await } Err(PostgresQueryStreamError::Postgres { err, emitted: false }) if should_retry_postgres_stale_cache(&err) => { // The cached prepared statement can become stale after schema changes. // Evict and retry once, matching the normal query execution path. log::warn!("[postgres][stream:stale_cache] evicting cached statement: {}", pg_error_to_string(err)); client.statement_cache.remove(sql, &[]); match stream_select_query_prepared(client, sql, row_limit, on_item, force_unnamed).await { Ok(rows) => Ok(rows), Err(PostgresQueryStreamError::Postgres { err, emitted: false }) if should_retry_postgres_text_query(&err) => { stream_select_query_text(client, sql, row_limit, None, on_item).await } Err(PostgresQueryStreamError::TextFallback { column_types, unsupported_type }) => { log::info!( "[postgres][stream:text_fallback] unsupported_type={} switching_to=simple_query", unsupported_type ); stream_select_query_text(client, sql, row_limit, Some(column_types), on_item).await } Err(err) => Err(err.into_string()), } } Err(PostgresQueryStreamError::Postgres { err, emitted: false }) if should_retry_postgres_text_query(&err) => { stream_select_query_text(client, sql, row_limit, None, on_item).await } Err(err) => Err(err.into_string()), } } pub async fn stream_query_rows( pool: &Pool, sql: &str, max_rows: Option, cancelled: &AtomicBool, mut on_row: impl FnMut(&[serde_json::Value]) -> Result<(), String>, ) -> Result { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; match stream_query_rows_on_client(&client, sql, max_rows, cancelled, &mut on_row).await { Ok(rows) => Ok(rows), Err(error) if should_retry_postgres_text_query_message(&error.to_ascii_lowercase()) => { stream_query_rows_text_on_client(&client, sql, max_rows, cancelled, None, &mut on_row).await } Err(error) => Err(error), } } async fn stream_query_rows_on_client( client: &deadpool_postgres::Client, sql: &str, max_rows: Option, cancelled: &AtomicBool, on_row: &mut impl FnMut(&[serde_json::Value]) -> Result<(), String>, ) -> Result { let stream_outcome = start_postgres_select_stream(client, sql, false).await.map_err(pg_error_to_string)?; let (stream, metadata) = match stream_outcome { PostgresSelectStreamOutcome::Binary { stream, metadata } => (stream, metadata), PostgresSelectStreamOutcome::TextFallback { column_types, unsupported_type } => { log::info!( "[postgres][row_stream:text_fallback] unsupported_type={} switching_to=simple_query", unsupported_type ); let column_classes = classify_pg_column_types(&column_types); return stream_query_rows_text_on_client(client, sql, max_rows, cancelled, Some(&column_classes), on_row) .await; } }; let PreparedSelectMetadata { column_classes, .. } = metadata; tokio::pin!(stream); let row_limit = max_rows.unwrap_or(usize::MAX); let mut rows_exported = 0_u64; while let Some(row_result) = stream.next().await { if cancelled.load(std::sync::atomic::Ordering::SeqCst) { return Err(crate::query::canceled_error()); } if rows_exported as usize >= row_limit { break; } let row = row_result.map_err(pg_error_to_string)?; let values: Vec = (0..row.columns().len()) .map(|i| pg_value_to_json_classified(&row, i, column_classes.get(i).copied().unwrap_or(PgColType::Other))) .collect(); on_row(&values)?; rows_exported += 1; } Ok(rows_exported) } async fn stream_query_rows_text_on_client( client: &deadpool_postgres::Client, sql: &str, max_rows: Option, cancelled: &AtomicBool, column_classes: Option<&[PgColType]>, on_row: &mut impl FnMut(&[serde_json::Value]) -> Result<(), String>, ) -> Result { let stream = client.simple_query_raw(sql).await.map_err(pg_error_to_string)?; tokio::pin!(stream); let row_limit = max_rows.unwrap_or(usize::MAX); let mut rows_exported = 0_u64; while let Some(message) = stream.next().await { if cancelled.load(std::sync::atomic::Ordering::SeqCst) { return Err(crate::query::canceled_error()); } if rows_exported as usize >= row_limit { break; } let message = message.map_err(pg_error_to_string)?; if let SimpleQueryMessage::Row(row) = message { let mut values = Vec::with_capacity(row.len()); for i in 0..row.len() { values.push(match row.try_get(i).map_err(pg_error_to_string)? { Some(value) => { pg_text_fallback_value(value, column_classes.and_then(|classes| classes.get(i)).copied()).0 } None => serde_json::Value::Null, }); } on_row(&values)?; rows_exported += 1; } } Ok(rows_exported) } pub async fn connect(url: &str, fallback_timeout: Duration) -> Result { connect_with_max_connections(url, fallback_timeout, 10).await } /// Creates a PostgreSQL pool with an explicit checkout bound. /// /// Session-scoped DBX pools use a single connection so temporary tables and /// other connection-local state cannot migrate between physical clients. pub async fn connect_with_max_connections( url: &str, fallback_timeout: Duration, max_connections: usize, ) -> Result { #[cfg(all(windows, target_vendor = "win7"))] { connect_with_optional_local_timezone_with_max_connections(url, fallback_timeout, None, max_connections).await } #[cfg(not(all(windows, target_vendor = "win7")))] { let timezone = iana_time_zone::get_timezone().unwrap_or_else(|_| "UTC".to_string()); connect_with_local_timezone_with_max_connections(url, fallback_timeout, &timezone, max_connections).await } } /// Test-only thin wrappers (the library now connects through /// `connect_with_max_connections`, which carries the pool bound). #[cfg(test)] async fn connect_with_local_timezone(url: &str, fallback_timeout: Duration, timezone: &str) -> Result { connect_with_local_timezone_with_max_connections(url, fallback_timeout, timezone, 10).await } async fn connect_with_local_timezone_with_max_connections( url: &str, fallback_timeout: Duration, timezone: &str, max_connections: usize, ) -> Result { connect_with_optional_local_timezone_with_max_connections(url, fallback_timeout, Some(timezone), max_connections) .await } /// Identity of one physical backend connection for notice attribution: /// (server address, server port, backend PID). The PID alone is not unique /// across different servers, so the server's own address/port disambiguate /// (`inet_server_addr()` is NULL for Unix sockets, hence the fallback). type PostgresConnectionKey = (String, String, String); const POSTGRES_CONNECTION_IDENTITY_SQL: &str = "SELECT pg_backend_pid()::text, \ COALESCE(host(inet_server_addr()), 'unix'), \ COALESCE(inet_server_port()::text, current_setting('port'))"; fn postgres_connection_key_from_row(row: &Row) -> Option { Some((row.try_get::<_, String>(1).ok()?, row.try_get::<_, String>(2).ok()?, row.try_get::<_, String>(0).ok()?)) } /// Notice buffers for live connections, keyed by connection identity. Entries /// are weak so they disappear once the pooled connection (and its driver /// task) is dropped. type PostgresNoticeBuffers = HashMap>>>; fn postgres_notice_buffers() -> &'static Mutex { static BUFFERS: OnceLock> = OnceLock::new(); BUFFERS.get_or_init(|| Mutex::new(HashMap::new())) } /// Connection identity cache, keyed by the per-physical-connection statement /// cache pointer (same identity pattern as `postgres_single_schema_clients`). /// The `Weak` guards against pointer reuse: a new connection whose /// `StatementCache` lands on a freed entry's address fails the `ptr_eq` check /// and is treated as a miss. The value is `None` when the identity query /// failed, so it is issued at most once per physical connection — it must /// never be retried from `drain_postgres_notices`, which can run inside the /// read-only transaction used for EXPLAIN, where a failing query would abort /// the user's statement. type PostgresClientKeys = HashMap, Option)>; fn postgres_client_keys() -> &'static Mutex { static KEYS: OnceLock> = OnceLock::new(); KEYS.get_or_init(|| Mutex::new(HashMap::new())) } /// Establishes connections like deadpool's `ConfigConnectImpl`, but drives /// each connection with a task that captures `NoticeResponse` messages /// (`RAISE NOTICE`/`WARNING`, etc.) into a per-backend buffer instead of /// discarding them. Query execution drains the buffer so notices are /// attached to the `QueryResult` of the statement that raised them. struct NoticeCapturingConnect where T: MakeTlsConnect + Clone + Sync + Send + 'static, T::Stream: Sync + Send, T::TlsConnect: Sync + Send, >::Future: Send, { tls: T, } impl deadpool_postgres::Connect for NoticeCapturingConnect where T: MakeTlsConnect + Clone + Sync + Send + 'static, T::Stream: Sync + Send, T::TlsConnect: Sync + Send, >::Future: Send, { fn connect( &self, pg_config: &tokio_postgres::Config, ) -> Pin< Box), tokio_postgres::Error>> + Send + '_>, > { let tls = self.tls.clone(); let pg_config = pg_config.clone(); Box::pin(async move { let (client, mut connection) = pg_config.connect(tls).await?; // No query can complete before the connection is being driven, so // the notice buffer is handed to the driver task through a slot // that is filled once the backend PID is known. let notice_buffer = Arc::new(Mutex::new(None::>>>)); let task_buffer = Arc::clone(¬ice_buffer); let conn_task = tokio::spawn(async move { loop { match std::future::poll_fn(|cx| connection.poll_message(cx)).await { Some(Ok(AsyncMessage::Notice(error))) => { let message = QueryMessage { severity: error.severity().to_string(), message: error.message().to_string(), code: Some(error.code().code().to_string()), detail: error.detail().map(str::to_string), hint: error.hint().map(str::to_string), }; let buffer = task_buffer.lock().unwrap_or_else(|poisoned| poisoned.into_inner()).clone(); match buffer { Some(buffer) => { buffer.lock().unwrap_or_else(|poisoned| poisoned.into_inner()).push(message); } None => { log::info!("[postgres][notice] {}: {}", message.severity, message.message); } } } // LISTEN/NOTIFY messages are not surfaced anywhere. Some(Ok(_)) => {} Some(Err(err)) => { log::warn!("[postgres] connection driver error: {err}"); break; } None => break, } } }); // Best-effort: without the connection identity, notices cannot be // attributed to query results on this connection and are logged // by the driver task instead. Never fail the connection over this. if let Ok(row) = client.query_one(POSTGRES_CONNECTION_IDENTITY_SQL, &[]).await { if let Some(key) = postgres_connection_key_from_row(&row) { let buffer = Arc::new(Mutex::new(Vec::new())); let mut buffers = postgres_notice_buffers().lock().unwrap_or_else(|poisoned| poisoned.into_inner()); buffers.retain(|_, weak| weak.strong_count() > 0); buffers.insert(key, Arc::downgrade(&buffer)); drop(buffers); *notice_buffer.lock().unwrap_or_else(|poisoned| poisoned.into_inner()) = Some(buffer); } } Ok((client, conn_task)) }) } } /// Cached identity lookup. `Some(Some(key))` = resolved, `Some(None)` = /// identity query failed before (do not retry), `None` = never seen (or the /// entry belonged to a dropped connection whose address was reused). fn cached_postgres_client_key(client: &deadpool_postgres::Client) -> Option> { let cache_key = Arc::as_ptr(&client.statement_cache) as usize; let mut keys = postgres_client_keys().lock().unwrap_or_else(|poisoned| poisoned.into_inner()); match keys.get(&cache_key) { Some((cached, key)) if cached.upgrade().is_some_and(|cached| Arc::ptr_eq(&cached, &client.statement_cache)) => { Some(key.clone()) } Some(_) => { keys.remove(&cache_key); None } None => None, } } /// Best-effort identity resolution. Failures are cached as `None` so the /// identity query runs at most once per physical connection. async fn resolve_postgres_client_key(client: &deadpool_postgres::Client) -> Option { if let Some(key) = cached_postgres_client_key(client) { return key; } let key: Option = client .query_one(POSTGRES_CONNECTION_IDENTITY_SQL, &[]) .await .ok() .and_then(|row| postgres_connection_key_from_row(&row)); let mut keys = postgres_client_keys().lock().unwrap_or_else(|poisoned| poisoned.into_inner()); keys.retain(|_, (cached, _)| cached.strong_count() > 0); let cache_key = Arc::as_ptr(&client.statement_cache) as usize; keys.insert(cache_key, (Arc::downgrade(&client.statement_cache), key.clone())); key } async fn postgres_client_key(client: &deadpool_postgres::Client) -> Option { resolve_postgres_client_key(client).await } fn take_notices_for_key(key: &PostgresConnectionKey) -> Vec { let buffer = { let mut buffers = postgres_notice_buffers().lock().unwrap_or_else(|poisoned| poisoned.into_inner()); buffers.retain(|_, weak| weak.strong_count() > 0); buffers.get(key).and_then(Weak::upgrade) }; let Some(buffer) = buffer else { return Vec::new(); }; let notices = std::mem::take(&mut *buffer.lock().unwrap_or_else(|poisoned| poisoned.into_inner())); notices } async fn drain_postgres_notices(client: &deadpool_postgres::Client) -> Vec { match postgres_client_key(client).await { Some(key) => take_notices_for_key(&key), None => Vec::new(), } } #[cfg(test)] async fn connect_with_optional_local_timezone( url: &str, fallback_timeout: Duration, timezone: Option<&str>, ) -> Result { connect_with_optional_local_timezone_with_max_connections(url, fallback_timeout, timezone, 10).await } async fn connect_with_optional_local_timezone_with_max_connections( url: &str, fallback_timeout: Duration, timezone: Option<&str>, max_connections: usize, ) -> Result { let url_with_keepalive = inject_postgres_keepalive_params(url); let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); let timeout = super::parse_connect_timeout_with_fallback(url, fallback_timeout); let first_attempt = connect_postgres_pool_attempt(&url_with_keepalive, timeout, max_connections).await; let (pool, client) = match first_attempt { Err(error) if postgres_error_should_retry_without_tls(&error) => { let Some(fallback_url) = postgres_ssl_fallback_url(&url_with_keepalive) else { return Err(error); }; log::info!("PostgreSQL TLS handshake failed in sslmode=prefer; retrying without TLS"); connect_postgres_pool_attempt(&fallback_url, timeout, max_connections).await? } result => result?, }; // Creating the physical connection and applying session defaults are two // sequential network phases. Give each phase the configured connection // timeout instead of sharing one deadline that can expire during the // optional SET timezone round-trip on higher-latency tunnels. if !pg_url_has_timezone_setting(url) { if let Some(timezone) = timezone { postgres_session_setup_with_timeout(timeout, set_automatic_postgres_timezone(&client, timezone)).await?; } } drop(client); Ok(pool) } async fn connect_postgres_pool_attempt( url: &str, timeout: Duration, max_connections: usize, ) -> Result<(Pool, deadpool_postgres::Client), String> { let postgres_url = postgres_connection_url(url)?; super::with_connection_timeout("PostgreSQL", timeout, async { let pg_config = tokio_postgres::Config::from_str(&postgres_url.url) .map_err(|e| format!("Invalid PostgreSQL connection URL: {e}"))?; // Fast recycling only checks whether the connection is already closed // instead of issuing a validation query on every checkout, saving one // round-trip per query. Connections that went stale without being // observed are caught when the query runs and recovered by the // executor's ReconnectAndRetry path (see pool_error_action / do_execute // in query.rs). let mgr_config = ManagerConfig { recycling_method: RecyclingMethod::Fast }; let mgr = if postgres_url.legacy_tls { log::info!("PostgreSQL legacy TLS compatibility enabled"); let tls = postgres_openssl_connector( &postgres_url.ssl_files, postgres_url.accepts_invalid_certs, postgres_url.verifies_hostname, )?; deadpool_postgres::Manager::from_connect(pg_config.clone(), NoticeCapturingConnect { tls }, mgr_config) } else { let tls_config = postgres_tls_config( &pg_config, &postgres_url.ssl_files, postgres_url.accepts_invalid_certs, postgres_url.verifies_hostname, )?; deadpool_postgres::Manager::from_connect( pg_config.clone(), NoticeCapturingConnect { tls: tokio_postgres_rustls::MakeRustlsConnect::new(tls_config) }, mgr_config, ) }; let pool = Pool::builder(mgr) .max_size(max_connections.max(1)) .runtime(Runtime::Tokio1) .wait_timeout(Some(timeout)) .create_timeout(Some(timeout)) .recycle_timeout(Some(timeout)) .build() .map_err(|e| format!("Failed to create PostgreSQL pool: {e}"))?; // Verify connectivity. Explicit connection options are handled by // PostgreSQL during startup and must remain strict. let client = pool.get().await.map_err(|e| format!("PostgreSQL connection failed: {}", pg_pool_error_to_string(e)))?; Ok((pool, client)) }) .await } fn postgres_error_should_retry_without_tls(error: &str) -> bool { error.to_ascii_lowercase().contains("tls handshake") } fn postgres_ssl_fallback_url(url: &str) -> Option { let parsed = postgres_connection_url(url).ok()?; let pg_config = tokio_postgres::Config::from_str(&parsed.url).ok()?; if pg_config.get_ssl_mode() != SslMode::Prefer { return None; } let (base, fragment) = url.split_once('#').map_or((url, None), |(base, fragment)| (base, Some(fragment))); let (prefix, query) = base.split_once('?').map_or((base, None), |(prefix, query)| (prefix, Some(query))); let mut params = Vec::new(); let mut inserted_sslmode = false; for param in query.into_iter().flat_map(|query| query.split('&')).filter(|param| !param.is_empty()) { let is_sslmode = param.split_once('=').is_some_and(|(key, _)| key.trim().eq_ignore_ascii_case("sslmode")); if is_sslmode { if !inserted_sslmode { params.push("sslmode=disable".to_string()); inserted_sslmode = true; } } else { params.push(param.to_string()); } } if !inserted_sslmode { params.insert(0, "sslmode=disable".to_string()); } let mut fallback = format!("{prefix}?{}", params.join("&")); if let Some(fragment) = fragment { fallback.push('#'); fallback.push_str(fragment); } Some(fallback) } async fn postgres_session_setup_with_timeout(timeout: Duration, future: F) -> Result where F: Future>, { super::with_connection_timeout("PostgreSQL session setup", timeout, future).await } async fn set_automatic_postgres_timezone(client: &deadpool_postgres::Client, timezone: &str) -> Result<(), String> { let candidates = postgres_timezone_candidates(timezone); for (index, candidate) in candidates.iter().enumerate() { let sql = format!("SET timezone = '{}'", candidate.replace('\'', "''")); match client.execute(&sql, &[]).await { Ok(_) => { if *candidate != timezone { log::warn!( "PostgreSQL does not recognize local timezone '{timezone}'; using compatible alias '{candidate}'" ); } return Ok(()); } Err(error) if postgres_timezone_error_is_nonfatal(&error) => { let detail = pg_error_to_string(error); if index + 1 == candidates.len() { // A connected server may have older tzdata or only partial PostgreSQL compatibility. // Keep its session default rather than making optional local display alignment fatal. log::warn!( "PostgreSQL connected, but automatic local timezone '{timezone}' was rejected; \ keeping the server default timezone: {detail}" ); return Ok(()); } } Err(error) => { return Err(format!("PostgreSQL SET timezone failed after connecting: {}", pg_error_to_string(error))); } } } Ok(()) } fn postgres_timezone_error_is_nonfatal(error: &tokio_postgres::Error) -> bool { let Some(db_error) = error.as_db_error() else { return false; }; // SET failures reported as ordinary SQL errors are optional session setup. // FATAL/PANIC responses mean the connection itself is not safe to return. !matches!( db_error.parsed_severity(), Some(tokio_postgres::error::Severity::Fatal | tokio_postgres::error::Severity::Panic) ) && !matches!(db_error.severity().to_ascii_uppercase().as_str(), "FATAL" | "PANIC") } fn postgres_timezone_candidates(timezone: &str) -> Vec<&str> { let legacy_alias = match timezone { "Asia/Saigon" => Some("Asia/Ho_Chi_Minh"), "Asia/Ho_Chi_Minh" => Some("Asia/Saigon"), "Europe/Kyiv" => Some("Europe/Kiev"), "Europe/Kiev" => Some("Europe/Kyiv"), "Asia/Calcutta" => Some("Asia/Kolkata"), "Asia/Kolkata" => Some("Asia/Calcutta"), _ => None, }; std::iter::once(timezone).chain(legacy_alias).collect() } #[derive(Debug, Default, Clone, PartialEq, Eq)] pub struct PostgresSslFiles { pub sslcert: Option, pub sslkey: Option, pub sslrootcert: Option, } /// TLS context info, used to reconstruct the TLS connector when cancelling a query. #[derive(Debug, Clone)] pub struct PostgresCancelContext { pub ssl_files: PostgresSslFiles, pub accepts_invalid_certs: bool, pub verifies_hostname: bool, pub ssl_mode: SslMode, pub legacy_tls: bool, } /// Build a TLS cancel context from the connection URL. /// Returns None if URL parsing fails or sslmode=disable (no TLS cancel needed). pub fn build_postgres_cancel_context(url: &str) -> Option { let postgres_url = postgres_connection_url(url).ok()?; let pg_config = tokio_postgres::Config::from_str(&postgres_url.url).ok()?; if pg_config.get_ssl_mode() == SslMode::Disable { return None; } Some(PostgresCancelContext { ssl_files: postgres_url.ssl_files, accepts_invalid_certs: postgres_url.accepts_invalid_certs, verifies_hostname: postgres_url.verifies_hostname, ssl_mode: pg_config.get_ssl_mode(), legacy_tls: postgres_url.legacy_tls, }) } /// Reconstruct a TLS connector from the cancel context, used for TLS connection cancellation. fn make_rustls_connect_from_context( ctx: &PostgresCancelContext, ) -> Result { // Build a minimal pg_config solely for ssl_mode determination let mut pg_config = tokio_postgres::Config::new(); pg_config.ssl_mode(ctx.ssl_mode); let tls_config = postgres_tls_config(&pg_config, &ctx.ssl_files, ctx.accepts_invalid_certs, ctx.verifies_hostname)?; Ok(tokio_postgres_rustls::MakeRustlsConnect::new(tls_config)) } fn make_openssl_connect_from_context( ctx: &PostgresCancelContext, ) -> Result { postgres_openssl_connector(&ctx.ssl_files, ctx.accepts_invalid_certs, ctx.verifies_hostname) } #[derive(Debug, Clone, PartialEq, Eq)] struct PostgresConnectionUrl { url: String, ssl_files: PostgresSslFiles, accepts_invalid_certs: bool, verifies_hostname: bool, legacy_tls: bool, } /// Inject TCP keepalive parameters into the PostgreSQL URL (only when the user has not explicitly specified them). /// Default parameters shorten half-open connection detection time, suitable for desktop/VPN/NAT environments. fn inject_postgres_keepalive_params(url: &str) -> String { let (base, fragment) = url.split_once('#').map_or((url, ""), |(base, fragment)| (base, fragment)); let query = base.split('?').nth(1); let has_keepalives = query .map(|q| q.split('&').any(|p| p.split('=').next().is_some_and(|k| k.eq_ignore_ascii_case("keepalives")))) .unwrap_or(false); if has_keepalives { return url.to_string(); // User has explicitly configured keepalive } let separator = if base.contains('?') { "&" } else { "?" }; let injected = format!("{base}{separator}keepalives=1&keepalives_idle=30&keepalives_interval=10&keepalives_retries=3"); if fragment.is_empty() { injected } else { format!("{injected}#{fragment}") } } fn postgres_connection_url(url: &str) -> Result { let Some(query_start) = url.find('?') else { let pg_config = tokio_postgres::Config::from_str(url).map_err(|e| format!("Invalid PostgreSQL connection URL: {e}"))?; return Ok(PostgresConnectionUrl { url: url.to_string(), ssl_files: PostgresSslFiles::default(), accepts_invalid_certs: postgres_sslmode_accepts_invalid_certs(pg_config.get_ssl_mode()), verifies_hostname: false, legacy_tls: false, }); }; let prefix = &url[..query_start]; let suffix = &url[query_start + 1..]; let (query_string, fragment) = suffix.split_once('#').map_or((suffix, ""), |(query, fragment)| (query, fragment)); let mut ssl_files = PostgresSslFiles::default(); let mut kept_params = Vec::new(); let mut accepts_invalid_certs = true; let mut verifies_hostname = false; let mut legacy_tls = false; for param in query_string.split('&') { if param.is_empty() { continue; } let Some((key, value)) = param.split_once('=') else { kept_params.push(param.to_string()); continue; }; if key.eq_ignore_ascii_case("legacy_tls") { legacy_tls = matches!(value.trim().to_ascii_lowercase().as_str(), "1" | "true" | "yes" | "on"); } else if key.eq_ignore_ascii_case("sslcert") || key.eq_ignore_ascii_case("sslkey") || key.eq_ignore_ascii_case("sslrootcert") { let decoded = percent_decode_str(value) .decode_utf8() .map_err(|_| format!("Invalid URL encoding in {key}"))? .into_owned(); validate_file_path(&decoded, |_| false).map_err(|e| format!("{key}: {e}"))?; if key.eq_ignore_ascii_case("sslcert") { ssl_files.sslcert = Some(decoded); } else if key.eq_ignore_ascii_case("sslkey") { ssl_files.sslkey = Some(decoded); } else { ssl_files.sslrootcert = Some(decoded); } } else if key.eq_ignore_ascii_case("channel_binding") { // channel_binding=require fails when the server does not offer // SCRAM-SHA-256-PLUS (e.g. Neon). Normalize require→prefer so // channel binding is used when available but does not cause a // hard failure when the server doesn't support it. match value.to_ascii_lowercase().as_str() { "require" => kept_params.push("channel_binding=prefer".to_string()), _ => kept_params.push(param.to_string()), } } else if key.eq_ignore_ascii_case("sslmode") { match value.to_ascii_lowercase().as_str() { "verify-ca" => { accepts_invalid_certs = false; kept_params.push("sslmode=require".to_string()); } "verify-full" | "verify_identity" | "verify-identity" => { accepts_invalid_certs = false; verifies_hostname = true; kept_params.push("sslmode=require".to_string()); } "disable" => { accepts_invalid_certs = false; kept_params.push(param.to_string()); } "prefer" | "require" => { accepts_invalid_certs = true; kept_params.push(param.to_string()); } _ => kept_params.push(param.to_string()), } } else { kept_params.push(param.to_string()); } } let mut sanitized_url = prefix.to_string(); if !kept_params.is_empty() { sanitized_url.push('?'); sanitized_url.push_str(&kept_params.join("&")); } if !fragment.is_empty() { sanitized_url.push('#'); sanitized_url.push_str(fragment); } Ok(PostgresConnectionUrl { url: sanitized_url, ssl_files, accepts_invalid_certs, verifies_hostname, legacy_tls }) } fn postgres_openssl_connector( ssl_files: &PostgresSslFiles, accepts_invalid_certs: bool, verifies_hostname: bool, ) -> Result { let mut builder = SslConnector::builder(SslMethod::tls_client()) .map_err(|e| format!("Failed to initialize PostgreSQL legacy TLS: {e}"))?; if accepts_invalid_certs { builder.set_verify(SslVerifyMode::NONE); } else { builder.set_verify(SslVerifyMode::PEER); if let Some(path) = ssl_files.sslrootcert.as_deref() { builder .set_ca_file(path) .map_err(|e| format!("sslrootcert: failed to load CA certificate from {path}: {e}"))?; } } match (ssl_files.sslcert.as_deref(), ssl_files.sslkey.as_deref()) { (Some(cert_path), Some(key_path)) => { builder .set_certificate_chain_file(cert_path) .map_err(|e| format!("sslcert: failed to load certificate from {cert_path}: {e}"))?; builder .set_private_key_file(key_path, SslFiletype::PEM) .map_err(|e| format!("sslkey: failed to load private key from {key_path}: {e}"))?; builder.check_private_key().map_err(|e| format!("sslkey: private key does not match sslcert: {e}"))?; } (Some(_), None) => return Err("PostgreSQL sslcert requires sslkey".to_string()), (None, Some(_)) => return Err("PostgreSQL sslkey requires sslcert".to_string()), (None, None) => {} } let mut connector = postgres_openssl::MakeTlsConnector::new(builder.build()); connector.set_callback(move |config, _| { config.set_verify_hostname(verifies_hostname); Ok(()) }); Ok(connector) } fn postgres_tls_config( pg_config: &tokio_postgres::Config, ssl_files: &PostgresSslFiles, accepts_invalid_certs: bool, verifies_hostname: bool, ) -> Result { if pg_config.get_ssl_mode() != SslMode::Disable && accepts_invalid_certs { let provider = Arc::new(rustls::crypto::aws_lc_rs::default_provider()); let builder = rustls::ClientConfig::builder() .dangerous() .with_custom_certificate_verifier(Arc::new(NoPostgresCertVerification { provider })); return postgres_tls_client_auth(builder, ssl_files); } let root_store = postgres_root_cert_store(ssl_files)?; let builder = if verifies_hostname { rustls::ClientConfig::builder().with_root_certificates(root_store) } else { let provider = Arc::new(rustls::crypto::aws_lc_rs::default_provider()); rustls::ClientConfig::builder().dangerous().with_custom_certificate_verifier(Arc::new( PostgresCaOnlyCertVerification { provider, roots: Arc::new(root_store) }, )) }; postgres_tls_client_auth(builder, ssl_files) } fn postgres_root_cert_store(ssl_files: &PostgresSslFiles) -> Result { let mut root_store = rustls::RootCertStore::empty(); if let Some(path) = ssl_files.sslrootcert.as_deref() { let certs = read_postgres_pem_certs("sslrootcert", path)?; let (valid_count, _) = root_store.add_parsable_certificates(certs); if valid_count == 0 { return Err(format!("sslrootcert: no valid CA certificates found in {path}")); } } else { root_store.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned()); } Ok(root_store) } fn postgres_tls_client_auth( builder: rustls::ConfigBuilder, ssl_files: &PostgresSslFiles, ) -> Result { match (ssl_files.sslcert.as_deref(), ssl_files.sslkey.as_deref()) { (Some(cert_path), Some(key_path)) => { let certs = read_postgres_pem_certs("sslcert", cert_path)?; if certs.is_empty() { return Err(format!("sslcert: no certificates found in {cert_path}")); } let private_key = read_postgres_private_key(key_path)?; let provider = rustls::crypto::aws_lc_rs::default_provider(); let signing_key = provider .key_provider .load_private_key(private_key) .map_err(|e| format!("PostgreSQL client private key is invalid: {e}"))?; // rustls' convenience builder parses the leaf with webpki while matching keys, // which rejects otherwise usable X.509 v1 client certificates before TLS starts. let certified_key = rustls::sign::CertifiedKey::new(certs, signing_key); Ok(builder.with_client_cert_resolver(Arc::new(rustls::sign::SingleCertAndKey::from(certified_key)))) } (Some(_), None) => Err("PostgreSQL sslcert requires sslkey".to_string()), (None, Some(_)) => Err("PostgreSQL sslkey requires sslcert".to_string()), (None, None) => Ok(builder.with_no_client_auth()), } } fn read_postgres_pem_certs(label: &str, path: &str) -> Result>, String> { let file = File::open(path).map_err(|e| format!("{label}: failed to open {path}: {e}"))?; let mut reader = BufReader::new(file); rustls_pemfile::certs(&mut reader) .collect::, _>>() .map_err(|e| format!("{label}: failed to read PEM certificates from {path}: {e}")) } fn read_postgres_private_key(path: &str) -> Result, String> { let file = File::open(path).map_err(|e| format!("sslkey: failed to open {path}: {e}"))?; let mut reader = BufReader::new(file); rustls_pemfile::private_key(&mut reader) .map_err(|e| format!("sslkey: failed to read PEM private key from {path}: {e}"))? .ok_or_else(|| format!("sslkey: no private key found in {path}")) } fn postgres_sslmode_accepts_invalid_certs(ssl_mode: SslMode) -> bool { matches!(ssl_mode, SslMode::Prefer | SslMode::Require) } #[derive(Debug)] struct NoPostgresCertVerification { provider: Arc, } impl ServerCertVerifier for NoPostgresCertVerification { fn verify_server_cert( &self, _end_entity: &CertificateDer<'_>, _intermediates: &[CertificateDer<'_>], _server_name: &ServerName<'_>, _ocsp_response: &[u8], _now: UnixTime, ) -> Result { Ok(ServerCertVerified::assertion()) } fn verify_tls12_signature( &self, _message: &[u8], cert: &CertificateDer<'_>, _dss: &rustls::DigitallySignedStruct, ) -> Result { self.accept_tls_signature_for_unverified_cert(cert) } fn verify_tls13_signature( &self, _message: &[u8], cert: &CertificateDer<'_>, _dss: &rustls::DigitallySignedStruct, ) -> Result { self.accept_tls_signature_for_unverified_cert(cert) } fn supported_verify_schemes(&self) -> Vec { self.provider.signature_verification_algorithms.supported_schemes() } } impl NoPostgresCertVerification { fn accept_tls_signature_for_unverified_cert( &self, _cert: &CertificateDer<'_>, ) -> Result { // PostgreSQL sslmode=prefer/require does not authenticate the server certificate. // Avoid rustls' default signature helpers here because they parse the certificate // before chain verification and reject legacy server certificates that libpq/JDBC // still accept in these non-verifying modes. Ok(HandshakeSignatureValid::assertion()) } } #[derive(Debug)] struct PostgresCaOnlyCertVerification { provider: Arc, roots: Arc, } impl ServerCertVerifier for PostgresCaOnlyCertVerification { fn verify_server_cert( &self, end_entity: &CertificateDer<'_>, intermediates: &[CertificateDer<'_>], _server_name: &ServerName<'_>, _ocsp_response: &[u8], now: UnixTime, ) -> Result { let cert = ParsedCertificate::try_from(end_entity)?; verify_server_cert_signed_by_trust_anchor( &cert, &self.roots, intermediates, now, self.provider.signature_verification_algorithms.all, )?; Ok(ServerCertVerified::assertion()) } fn verify_tls12_signature( &self, message: &[u8], cert: &CertificateDer<'_>, dss: &rustls::DigitallySignedStruct, ) -> Result { verify_tls12_signature(message, cert, dss, &self.provider.signature_verification_algorithms) } fn verify_tls13_signature( &self, message: &[u8], cert: &CertificateDer<'_>, dss: &rustls::DigitallySignedStruct, ) -> Result { verify_tls13_signature(message, cert, dss, &self.provider.signature_verification_algorithms) } fn supported_verify_schemes(&self) -> Vec { self.provider.signature_verification_algorithms.supported_schemes() } } /// Check whether the user's connection URL already specifies a timezone via /// the `options` parameter so we don't overwrite it with the local timezone. fn pg_url_has_timezone_setting(url: &str) -> bool { let Some(query) = url.split_once('?').map(|(_, query)| query.split('#').next().unwrap_or(query)) else { return false; }; query.split('&').any(|parameter| { let (raw_key, raw_value) = parameter.split_once('=').unwrap_or((parameter, "")); let key = percent_decode_str(raw_key).decode_utf8_lossy(); if !key.eq_ignore_ascii_case("options") { return false; } let options = percent_decode_str(raw_value).decode_utf8_lossy().to_ascii_lowercase(); options.split_ascii_whitespace().any(|token| { let option = token.trim_start_matches('-'); option.starts_with("timezone=") || option.starts_with("time_zone=") }) }) } #[cfg(test)] fn validate_postgres_ssl_paths(url: &str) -> Result<(), String> { postgres_connection_url(url).map(|_| ()) } fn list_databases_sql() -> &'static str { "SELECT datname \ FROM pg_database \ WHERE datallowconn = true \ ORDER BY datname" } fn database_metadata_sql() -> &'static str { "SELECT d.datname, \ CASE \ WHEN has_database_privilege(d.datname, 'CONNECT') \ OR COALESCE(( \ SELECT pg_has_role(current_user, r.oid, 'MEMBER') \ FROM pg_roles r \ WHERE r.rolname = 'pg_read_all_stats' \ ), false) \ THEN pg_database_size(d.oid) \ ELSE NULL \ END AS size_bytes, \ obj_description(d.oid, 'pg_database') AS comment, \ pg_encoding_to_char(d.encoding) AS default_charset \ FROM pg_database d \ WHERE datallowconn = true \ ORDER BY d.datname" } fn database_storage_sql() -> &'static str { "SELECT d.datname, \ CASE \ WHEN has_database_privilege(d.datname, 'CONNECT') \ OR COALESCE(( \ SELECT pg_has_role(current_user, r.oid, 'MEMBER') \ FROM pg_roles r \ WHERE r.rolname = 'pg_read_all_stats' \ ), false) \ THEN pg_database_size(d.oid) \ ELSE NULL \ END AS size_bytes \ FROM pg_database d \ WHERE d.datallowconn = true \ AND d.datname = ANY($1::text[]) \ ORDER BY d.datname" } pub async fn list_databases(pool: &Pool) -> Result, String> { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let rows = postgres_query_cached(&client, list_databases_sql(), &[]).await.map_err(|e| e.to_string())?; Ok(rows.iter().map(|row| DatabaseInfo { name: pg_row_try_string(row, 0), ..Default::default() }).collect()) } pub async fn list_database_metadata(pool: &Pool) -> Result, String> { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let rows = postgres_query_cached(&client, database_metadata_sql(), &[]).await.map_err(|e| e.to_string())?; Ok(rows .iter() .map(|row| DatabaseInfo { name: pg_row_try_string(row, 0), size_bytes: row.try_get::<_, Option>(1).ok().flatten(), comment: row.try_get::<_, Option>(2).ok().flatten().filter(|value| !value.trim().is_empty()), default_charset: row .try_get::<_, Option>(3) .ok() .flatten() .filter(|value| !value.trim().is_empty()), ..Default::default() }) .collect()) } pub async fn list_database_storage(pool: &Pool, database_names: &[String]) -> Result, String> { if database_names.is_empty() { return Ok(Vec::new()); } let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let rows = postgres_query_cached(&client, database_storage_sql(), &[&database_names]).await.map_err(|e| e.to_string())?; Ok(rows .iter() .map(|row| DatabaseStorageInfo { name: pg_row_try_string(row, 0), size_bytes: row.try_get::<_, Option>(1).ok().flatten(), }) .collect()) } pub async fn list_tables(pool: &Pool, schema: &str) -> Result, String> { list_tables_filtered(pool, schema, None, None, None).await } pub async fn list_tables_filtered( pool: &Pool, schema: &str, filter: Option<&str>, limit: Option, offset: Option, ) -> Result, String> { list_tables_filtered_by_kind(pool, schema, filter, limit, offset, false).await } /// Lists only table-like relations, excluding views and materialized views while /// retaining server-side filtering and pagination. pub async fn list_table_objects_filtered( pool: &Pool, schema: &str, filter: Option<&str>, limit: Option, offset: Option, ) -> Result, String> { list_tables_filtered_by_kind(pool, schema, filter, limit, offset, true).await } async fn list_tables_filtered_by_kind( pool: &Pool, schema: &str, filter: Option<&str>, limit: Option, offset: Option, table_objects_only: bool, ) -> Result, String> { let schema = if schema.is_empty() { "public" } else { schema }; let filter = filter.unwrap_or("").trim(); let filter_pattern = like_contains_pattern(filter); let fuzzy_filter_pattern = if crate::sql::fuzzy_filter_enabled(filter) { like_fuzzy_pattern(filter) } else { String::new() }; let limit_param = limit.and_then(|value| i64::try_from(value).ok()); let offset_param = offset.and_then(|value| i64::try_from(value).ok()).unwrap_or(0); let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let sql = if table_objects_only { postgres_table_objects_sql(limit_param, offset_param) } else { postgres_tables_sql(limit_param, offset_param) }; let params: &[(&(dyn tokio_postgres::types::ToSql + Sync), Type)] = &[(&schema, Type::VARCHAR), (&filter_pattern, Type::VARCHAR), (&fuzzy_filter_pattern, Type::VARCHAR)]; // The pagination literals make this SQL vary by page. Use an unnamed typed // query so each load stays one round trip without growing the statement cache. let rows = client.query_typed(&sql, params).await.map_err(|e| e.to_string())?; Ok(rows .iter() .map(|row| TableInfo { name: pg_row_try_string(row, 0), table_type: pg_row_try_string(row, 1), comment: row.try_get::<_, Option>(2).ok().flatten().filter(|s| !s.is_empty()), parent_schema: row.try_get::<_, Option>(3).ok().flatten().filter(|s| !s.is_empty()), parent_name: row.try_get::<_, Option>(4).ok().flatten().filter(|s| !s.is_empty()), }) .collect()) } pub async fn completion_assistant_search( pool: &Pool, request: &CompletionAssistantRequest, ) -> Result { let schema = request.schema.as_deref().or(request.parent_schema.as_deref()); let routine_schema = schema.unwrap_or("public"); let limit = request.max_results.unwrap_or(100).clamp(1, 1000); let kinds = if request.object_kinds.is_empty() { vec![CompletionAssistantObjectKind::Table, CompletionAssistantObjectKind::View] } else { request.object_kinds.clone() }; let pattern = postgres_completion_like_pattern(&request.mask, request.match_mode.as_ref()); let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let mut candidates = Vec::new(); if kinds.iter().any(|kind| matches!(kind, CompletionAssistantObjectKind::Schema)) { for row in postgres_query_cached( &client, "SELECT nspname FROM pg_catalog.pg_namespace \ WHERE nspname NOT LIKE 'pg_%' AND nspname <> 'information_schema' \ AND ($1 = '%%' OR nspname ILIKE $1 ESCAPE '~') \ ORDER BY nspname LIMIT $2", &[&pattern, &(limit as i64)], ) .await .map_err(|e| e.to_string())? { let schema_name: String = pg_row_try_string(&row, 0); candidates.push(CompletionAssistantCandidate { name: schema_name.clone(), kind: CompletionAssistantCandidateKind::Schema, database: Some(request.database.clone()), schema: Some(schema_name), parent_schema: None, parent_name: None, comment: None, data_type: None, signature: None, }); } } if candidates.len() < limit && kinds.iter().any(CompletionAssistantObjectKind::is_table_like) { let relkinds = postgres_completion_relkinds(&kinds); let rows = postgres_query_cached( &client, postgres_completion_tables_sql(), &[&schema, &pattern, &relkinds, &((limit - candidates.len()) as i64)], ) .await .map_err(|e| e.to_string())?; for row in rows { let table_type: String = pg_row_try_string(&row, 2); candidates.push(CompletionAssistantCandidate { name: pg_row_try_string(&row, 0), kind: if table_type.contains("VIEW") { CompletionAssistantCandidateKind::View } else { CompletionAssistantCandidateKind::Table }, database: Some(request.database.clone()), schema: Some(pg_row_try_string(&row, 1)), parent_schema: row.try_get::<_, Option>(4).ok().flatten(), parent_name: row.try_get::<_, Option>(5).ok().flatten(), comment: row.try_get::<_, Option>(3).ok().flatten(), data_type: None, signature: None, }); } } if candidates.len() < limit && kinds.iter().any(CompletionAssistantObjectKind::is_routine_like) { let has_proc_prokind = postgres_proc_has_prokind(&client).await?; let rows = if has_proc_prokind { let prokinds = postgres_completion_prokinds(&kinds); postgres_query_cached( &client, postgres_completion_routines_sql(true), &[&routine_schema, &pattern, &prokinds, &((limit - candidates.len()) as i64)], ) .await .map_err(|e| e.to_string())? } else if kinds.iter().any(|kind| { matches!(kind, CompletionAssistantObjectKind::Function | CompletionAssistantObjectKind::Routine) }) { postgres_query_cached( &client, postgres_completion_routines_sql(false), &[&routine_schema, &pattern, &((limit - candidates.len()) as i64)], ) .await .map_err(|e| e.to_string())? } else { Vec::new() }; for row in rows { let routine_type: String = pg_row_try_string(&row, 2); candidates.push(CompletionAssistantCandidate { name: pg_row_try_string(&row, 0), kind: if routine_type == "PROCEDURE" { CompletionAssistantCandidateKind::Procedure } else { CompletionAssistantCandidateKind::Function }, database: Some(request.database.clone()), schema: Some(pg_row_try_string(&row, 1)), parent_schema: None, parent_name: None, comment: row.try_get::<_, Option>(3).ok().flatten(), data_type: row.try_get::<_, Option>(4).ok().flatten(), signature: row.try_get::<_, Option>(5).ok().flatten(), }); } } if candidates.len() < limit && kinds.iter().any(|kind| matches!(kind, CompletionAssistantObjectKind::Sequence)) { let rows = postgres_query_cached( &client, postgres_completion_sequences_sql(), &[&schema, &pattern, &request.case_sensitive, &((limit - candidates.len()) as i64)], ) .await .map_err(|e| e.to_string())?; for row in rows { candidates.push(CompletionAssistantCandidate { name: pg_row_try_string(&row, 0), kind: CompletionAssistantCandidateKind::Sequence, database: Some(request.database.clone()), schema: Some(pg_row_try_string(&row, 1)), parent_schema: None, parent_name: None, comment: row.try_get::<_, Option>(2).ok().flatten(), data_type: None, signature: None, }); } } if candidates.len() < limit && kinds.iter().any(|kind| matches!(kind, CompletionAssistantObjectKind::Column)) { let table = request.parent_name.as_deref().unwrap_or(""); if !table.is_empty() { // Unqualified PostgreSQL objects resolve through search_path, so column // metadata must use the same visible relation instead of assuming public. let resolved_schema = match schema { Some(schema) => Some(schema.to_string()), None => postgres_query_cached(&client, postgres_visible_table_schema_sql(), &[&table]) .await .map_err(|e| e.to_string())? .first() .map(|row| pg_row_try_string(row, 0)), }; let Some(resolved_schema) = resolved_schema else { return Ok(CompletionAssistantResponse { incomplete: false, candidates, fallback_used: false }); }; let rows = postgres_query_cached( &client, postgres_completion_columns_sql(), &[&resolved_schema, &table, &pattern, &((limit - candidates.len()) as i64)], ) .await .map_err(|e| e.to_string())?; for row in rows { candidates.push(CompletionAssistantCandidate { name: pg_row_try_string(&row, 0), kind: CompletionAssistantCandidateKind::Column, database: Some(request.database.clone()), schema: Some(resolved_schema.clone()), parent_schema: Some(resolved_schema.clone()), parent_name: Some(table.to_string()), comment: row.try_get::<_, Option>(2).ok().flatten(), data_type: Some(pg_row_try_string(&row, 1)), signature: None, }); } } } Ok(CompletionAssistantResponse { incomplete: candidates.len() >= limit, candidates, fallback_used: false }) } fn postgres_completion_tables_sql() -> &'static str { "SELECT c.relname, n.nspname, \ CASE c.relkind WHEN 'v' THEN 'VIEW' WHEN 'm' THEN 'VIEW' ELSE 'TABLE' END AS table_type, \ obj_description(c.oid) AS table_comment, \ CASE WHEN pc.relkind = 'p' THEN pn.nspname ELSE NULL END AS parent_schema, \ CASE WHEN pc.relkind = 'p' THEN pc.relname ELSE NULL END AS parent_name \ FROM pg_catalog.pg_class c \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ LEFT JOIN pg_catalog.pg_inherits i ON i.inhrelid = c.oid \ LEFT JOIN pg_catalog.pg_class pc ON pc.oid = i.inhparent \ LEFT JOIN pg_catalog.pg_namespace pn ON pn.oid = pc.relnamespace \ WHERE ($1::text IS NOT NULL AND n.nspname = $1 \ OR $1::text IS NULL AND pg_catalog.pg_table_is_visible(c.oid)) \ AND c.relkind::text = ANY($3::text[]) \ AND ($2 = '%%' OR c.relname ILIKE $2 ESCAPE '~') \ ORDER BY c.relname LIMIT $4" } fn postgres_completion_routines_sql(has_proc_prokind: bool) -> &'static str { if has_proc_prokind { return "SELECT p.proname, n.nspname, CASE p.prokind WHEN 'p' THEN 'PROCEDURE' ELSE 'FUNCTION' END, \ obj_description(p.oid) AS routine_comment, COALESCE(pg_get_function_result(p.oid), '') AS data_type, \ pg_get_function_identity_arguments(p.oid) AS signature \ FROM pg_catalog.pg_proc p \ JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \ WHERE n.nspname = $1 AND p.prokind::text = ANY($3::text[]) \ AND ($2 = '%%' OR p.proname ILIKE $2 ESCAPE '~') \ ORDER BY p.proname LIMIT $4"; } "SELECT p.proname, n.nspname, 'FUNCTION'::text, \ obj_description(p.oid) AS routine_comment, COALESCE(pg_get_function_result(p.oid), '') AS data_type, \ pg_get_function_identity_arguments(p.oid) AS signature \ FROM pg_catalog.pg_proc p \ JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \ WHERE n.nspname = $1 \ AND ($2 = '%%' OR p.proname ILIKE $2 ESCAPE '~') \ ORDER BY p.proname LIMIT $3" } fn postgres_completion_sequences_sql() -> &'static str { "SELECT c.relname, n.nspname, obj_description(c.oid) AS sequence_comment \ FROM pg_catalog.pg_class c \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ WHERE ($1::text IS NOT NULL AND n.nspname = $1 \ OR $1::text IS NULL AND pg_catalog.pg_table_is_visible(c.oid)) \ AND c.relkind = 'S' \ AND pg_catalog.has_schema_privilege(n.oid, 'USAGE') \ AND CASE WHEN c.relkind = 'S' \ THEN pg_catalog.has_sequence_privilege(c.oid, 'USAGE, SELECT, UPDATE') \ ELSE FALSE END \ AND ($2 = '%%' OR CASE WHEN $3 THEN c.relname LIKE $2 ESCAPE '~' \ ELSE c.relname ILIKE $2 ESCAPE '~' END) \ ORDER BY c.relname LIMIT $4" } fn postgres_completion_columns_sql() -> &'static str { "SELECT a.attname, pg_catalog.format_type(a.atttypid, a.atttypmod), col_description(c.oid, a.attnum) \ FROM pg_catalog.pg_attribute a \ JOIN pg_catalog.pg_class c ON c.oid = a.attrelid \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ WHERE n.nspname = $1 AND c.relname = $2 AND a.attnum > 0 AND NOT a.attisdropped \ AND ($3 = '%%' OR a.attname ILIKE $3 ESCAPE '~') \ ORDER BY a.attnum LIMIT $4" } fn postgres_visible_table_schema_sql() -> &'static str { "SELECT n.nspname FROM pg_catalog.pg_class c \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ WHERE c.relname = $1 AND pg_catalog.pg_table_is_visible(c.oid) \ LIMIT 1" } fn postgres_completion_relkinds(kinds: &[CompletionAssistantObjectKind]) -> Vec { let mut relkinds = Vec::new(); if kinds.iter().any(|kind| matches!(kind, CompletionAssistantObjectKind::Table)) { relkinds.extend(["r", "p", "f"].into_iter().map(str::to_string)); } if kinds.iter().any(|kind| matches!(kind, CompletionAssistantObjectKind::View)) { relkinds.extend(["v", "m"].into_iter().map(str::to_string)); } relkinds } fn postgres_completion_prokinds(kinds: &[CompletionAssistantObjectKind]) -> Vec { let mut prokinds = Vec::new(); if kinds .iter() .any(|kind| matches!(kind, CompletionAssistantObjectKind::Procedure | CompletionAssistantObjectKind::Routine)) { prokinds.push("p".to_string()); } if kinds .iter() .any(|kind| matches!(kind, CompletionAssistantObjectKind::Function | CompletionAssistantObjectKind::Routine)) { prokinds.push("f".to_string()); } prokinds } fn postgres_completion_like_pattern(value: &str, mode: Option<&CompletionAssistantMatchMode>) -> String { if value.trim().is_empty() || value == "%" { return "%%".to_string(); } let escaped = value.trim().replace('~', "~~").replace('%', "~%").replace('_', "~_"); match mode.unwrap_or(&CompletionAssistantMatchMode::Prefix) { CompletionAssistantMatchMode::Prefix => format!("{escaped}%"), CompletionAssistantMatchMode::Contains => format!("%{escaped}%"), } } pub async fn get_table_comment(pool: &Pool, schema: &str, table: &str) -> Result, String> { let schema = if schema.is_empty() { "public" } else { schema }; let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let rows = postgres_query_cached(&client, postgres_table_comment_sql(), &[&schema, &table]) .await .map_err(|e| e.to_string())?; Ok(rows.first().and_then(|row| row.try_get::<_, Option>(0).ok().flatten()).filter(|s| !s.is_empty())) } pub async fn get_table_partition_info( pool: &Pool, schema: &str, table: &str, ) -> Result { let schema = if schema.is_empty() { "public" } else { schema }; let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let relation_rows = postgres_query_cached(&client, postgres_table_partition_relation_sql(), &[&schema, &table]) .await .map_err(|e| e.to_string())?; let Some(relation) = relation_rows.first() else { return Ok(PostgresTablePartitionInfo::default()); }; let relkind = relation.try_get::<_, String>(0).unwrap_or_default(); let is_partition = relation.try_get::<_, bool>(1).unwrap_or(false); let is_foreign = relkind == "f"; if relkind != "p" && !is_foreign && !is_partition { return Ok(PostgresTablePartitionInfo::default()); } let params: [&(dyn tokio_postgres::types::ToSql + Sync); 2] = [&schema, &table]; let rows = query_with_compat_fallback( "get_table_partition_info", &[postgres_table_partition_info_sql(), postgres_table_partition_info_compat_sql()], |sql| postgres_query_cached(&client, sql, ¶ms), ) .await?; let Some(row) = rows.first() else { return Ok(PostgresTablePartitionInfo { is_partition, is_foreign, ..Default::default() }); }; let foreign_options = row .try_get::<_, Option>>(5) .ok() .flatten() .unwrap_or_default() .into_iter() .filter_map(|option| option.split_once('=').map(|(key, value)| (key.to_string(), value.to_string()))) .collect(); Ok(PostgresTablePartitionInfo { is_partition, parent_schema: row.try_get::<_, Option>(0).ok().flatten().filter(|value| !value.is_empty()), parent_table: row.try_get::<_, Option>(1).ok().flatten().filter(|value| !value.is_empty()), bound: row.try_get::<_, Option>(2).ok().flatten().filter(|value| !value.is_empty()), key: row.try_get::<_, Option>(3).ok().flatten().filter(|value| !value.is_empty()), is_foreign, foreign_server: row.try_get::<_, Option>(4).ok().flatten().filter(|value| !value.is_empty()), foreign_options, }) } pub async fn get_table_partition_key(pool: &Pool, schema: &str, table: &str) -> Result, String> { Ok(get_table_partition_info(pool, schema, table).await?.key) } /// Classifies one row of `postgres_table_partition_local_objects_sql` (or its /// `_for_relations` sibling, which has the same `object_kind`/`object_name`/ /// `object_type` columns plus a leading `relid`) into `entry`. Shared so a /// future local-object kind, or a fix to how a kind is recognized, can't be /// applied to only one of the two query paths. fn apply_partition_local_object_row( entry: &mut PostgresTablePartitionLocalObjects, object_kind: &str, object_name: String, object_type: &str, ) { match object_kind { "constraint" if object_type == "p" => entry.has_primary_key = true, "constraint" if object_type == "f" && !object_name.is_empty() => { entry.foreign_keys.insert(object_name); } "check" if !object_name.is_empty() => { entry.check_constraints.insert(object_name); } "index" if !object_name.is_empty() => { entry.indexes.insert(object_name); } "column_default" if !object_name.is_empty() => { let state = if object_type == "dropped" { PostgresColumnDefaultState::Dropped } else { PostgresColumnDefaultState::Overridden }; entry.column_defaults.insert(object_name, state); } _ => {} } } pub async fn get_table_partition_local_objects( pool: &Pool, schema: &str, table: &str, ) -> Result { let schema = if schema.is_empty() { "public" } else { schema }; let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let rows = postgres_query_cached(&client, postgres_table_partition_local_objects_sql(), &[&schema, &table]) .await .map_err(|e| e.to_string())?; let mut result = PostgresTablePartitionLocalObjects::default(); for row in rows { let object_kind = row.try_get::<_, String>(0).unwrap_or_default(); let object_name = row.try_get::<_, String>(1).unwrap_or_default(); let object_type = row.try_get::<_, Option>(2).ok().flatten().unwrap_or_default(); apply_partition_local_object_row(&mut result, &object_kind, object_name, &object_type); } Ok(result) } /// One relation in a partition tree, as discovered by /// `fetch_postgres_partition_tree`. #[derive(Debug, Clone)] pub struct PostgresPartitionTreeNode { pub oid: i64, pub schema: String, pub table: String, pub parent_oid: Option, pub parent_schema: Option, pub parent_table: Option, pub partition_info: PostgresTablePartitionInfo, } fn postgres_partition_tree_sql() -> &'static str { "WITH RECURSIVE tree AS ( \ SELECT c.oid::bigint AS oid, n.nspname::text AS schema, c.relname::text AS relname, \ ap.oid::bigint AS parent_oid, an.nspname::text COLLATE \"C\" AS parent_schema, ap.relname::text COLLATE \"C\" AS parent_relname, \ c.relkind::text AS relkind, ARRAY[c.oid]::bigint[] AS path \ FROM pg_catalog.pg_class c \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ LEFT JOIN pg_catalog.pg_inherits ai ON ai.inhrelid = c.oid AND c.relispartition \ LEFT JOIN pg_catalog.pg_class ap ON ap.oid = ai.inhparent \ LEFT JOIN pg_catalog.pg_namespace an ON an.oid = ap.relnamespace \ WHERE n.nspname = $1 AND c.relname = $2 AND c.relkind IN ('r','p','f') \ UNION ALL \ SELECT c.oid::bigint, n.nspname::text, c.relname::text, tree.oid, tree.schema, tree.relname, c.relkind::text, \ tree.path || c.oid::bigint \ FROM pg_catalog.pg_inherits i \ JOIN pg_catalog.pg_class c ON c.oid = i.inhrelid AND c.relispartition \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ JOIN tree ON tree.oid = i.inhparent \ WHERE NOT c.oid = ANY(tree.path) \ ) \ SELECT t.oid, t.schema, t.relname, t.parent_oid, t.parent_schema, t.parent_relname, t.relkind, \ pg_catalog.pg_get_expr(c.relpartbound, c.oid, true) AS partition_bound, \ CASE WHEN t.relkind = 'p' THEN pg_catalog.pg_get_partkeydef(c.oid) ELSE NULL END AS partition_key, \ fs.srvname AS foreign_server, \ ft.ftoptions AS foreign_options \ FROM tree t \ JOIN pg_catalog.pg_class c ON c.oid = t.oid \ LEFT JOIN pg_catalog.pg_foreign_table ft ON ft.ftrelid = c.oid \ LEFT JOIN pg_catalog.pg_foreign_server fs ON fs.oid = ft.ftserver \ ORDER BY t.oid" } /// 9.x 没有 relispartition/relpartbound/pg_get_partkeydef,主查询会用到一个 /// 不存在的列,直接报错;旧式 INHERITS 子表在 9.x 按普通表处理, /// 这里只返回请求的关系本身。 /// /// Pre-10 servers lack the declarative-partition catalog (relispartition / /// relpartbound / pg_get_partkeydef) the primary query needs; old-style /// INHERITS children are plain tables there, so just the requested relation /// is returned. fn postgres_partition_tree_compat_sql() -> &'static str { "SELECT c.oid::bigint AS oid, n.nspname::text AS schema, c.relname::text AS relname, \ NULL::bigint AS parent_oid, NULL::text AS parent_schema, NULL::text AS parent_relname, \ c.relkind::text AS relkind, \ NULL::text AS partition_bound, NULL::text AS partition_key, \ fs.srvname AS foreign_server, \ ft.ftoptions AS foreign_options \ FROM pg_catalog.pg_class c \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ LEFT JOIN pg_catalog.pg_foreign_table ft ON ft.ftrelid = c.oid \ LEFT JOIN pg_catalog.pg_foreign_server fs ON fs.oid = ft.ftserver \ WHERE n.nspname = $1 AND c.relname = $2 AND c.relkind IN ('r','p','f')" } /// The whole partition tree rooted at (schema, table) — the root itself plus /// every descendant partition at any depth — plus each node's partition info, /// in a single round trip. Used by `pg_ddl_with_partitions` instead of /// recursing per-relation (which reruns the full metadata query chain once /// per node and scales the request count linearly with tree size). pub async fn fetch_postgres_partition_tree( pool: &Pool, schema: &str, table: &str, ) -> Result, String> { let schema = if schema.is_empty() { "public" } else { schema }; let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let params: [&(dyn tokio_postgres::types::ToSql + Sync); 2] = [&schema, &table]; let rows = query_with_compat_fallback( "fetch_postgres_partition_tree", &[postgres_partition_tree_sql(), postgres_partition_tree_compat_sql()], |sql| postgres_query_cached(&client, sql, ¶ms), ) .await?; Ok(rows .into_iter() .filter_map(|row| { let oid = row.try_get::<_, i64>(0).ok()?; let node_schema = row.try_get::<_, String>(1).ok()?; let node_table = row.try_get::<_, String>(2).ok()?; let parent_oid = row.try_get::<_, Option>(3).ok().flatten(); let parent_schema = row.try_get::<_, Option>(4).ok().flatten(); let parent_table = row.try_get::<_, Option>(5).ok().flatten(); let relkind = row.try_get::<_, String>(6).unwrap_or_default(); let is_foreign = relkind == "f"; let foreign_options = row .try_get::<_, Option>>(10) .ok() .flatten() .unwrap_or_default() .into_iter() .filter_map(|option| option.split_once('=').map(|(key, value)| (key.to_string(), value.to_string()))) .collect(); let partition_info = PostgresTablePartitionInfo { is_partition: parent_oid.is_some(), parent_schema: parent_schema.clone().filter(|value| !value.is_empty()), parent_table: parent_table.clone().filter(|value| !value.is_empty()), bound: row.try_get::<_, Option>(7).ok().flatten().filter(|value| !value.is_empty()), key: row.try_get::<_, Option>(8).ok().flatten().filter(|value| !value.is_empty()), is_foreign, foreign_server: row.try_get::<_, Option>(9).ok().flatten().filter(|value| !value.is_empty()), foreign_options, }; Some(PostgresPartitionTreeNode { oid, schema: node_schema, table: node_table, parent_oid, parent_schema, parent_table, partition_info, }) }) .collect()) } /// Batched sibling of get_columns, fetching every relation through at most /// two OID-scoped catalog queries. The compatibility tier avoids PostgreSQL /// 12+ catalog fields so PostgreSQL 10-compatible servers remain bounded. pub async fn get_columns_for_relations( pool: &Pool, relations: &[(i64, String, String)], ) -> Result>, String> { let oids: Vec = relations.iter().map(|(oid, _, _)| *oid).collect(); let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let tiers = postgres_columns_for_relations_query_tiers(); query_with_compat_fallback("get_columns_for_relations", &tiers, |sql| { get_columns_for_relations_with_sql(&client, sql, &oids) }) .await } async fn get_columns_for_relations_with_sql( client: &deadpool_postgres::Client, sql: &str, oids: &[i64], ) -> Result>, tokio_postgres::Error> { let rows = postgres_query_cached(client, sql, &[&oids]).await?; let mut result: HashMap> = HashMap::new(); for row in &rows { let Ok(relid) = row.try_get::<_, i64>(0) else { continue }; result.entry(relid).or_default().push(column_info_from_row_offset(row, 1)); } Ok(result) } fn postgres_columns_for_relations_query_tiers() -> [&'static str; 2] { [postgres_columns_for_relations_sql(), postgres_columns_for_relations_compat_sql()] } // Sibling of `POSTGRES_COLUMNS_SQL`/`POSTGRES_COLUMNS_COMPAT_SQL` below (~line // 4880): same column list and detection logic (identity/serial inference, // numeric precision/scale, enum values), batched by oid instead of a single // (schema, table) pair. Kept as a separate literal rather than sharing a // fragment — `pg_class` needs the `c` alias here for the oid filter, which // pushes `information_schema.columns` to `ic` instead of the single-relation // version's `c`, and every column's position in the SELECT list is relied on // positionally by `column_info_from_row_offset`. A change to one almost // certainly needs the same change in the other. fn postgres_columns_for_relations_sql() -> &'static str { "SELECT c.oid::bigint AS relid, a.attname AS column_name, \ format_type(a.atttypid, a.atttypmod) AS full_type, \ COALESCE(ic.is_nullable = 'YES', NOT a.attnotnull) AS is_nullable, \ CASE WHEN a.attgenerated <> '' THEN NULL ELSE pg_get_expr(ad.adbin, ad.adrelid) END AS column_default, \ EXISTS ( \ SELECT 1 FROM pg_constraint co \ JOIN pg_index i ON i.indrelid = co.conrelid AND co.conindid = i.indexrelid \ WHERE co.conrelid = a.attrelid AND co.contype = 'p' \ AND a.attnum = ANY(i.indkey) \ ) AS is_pk, \ col_description(a.attrelid, a.attnum) AS column_comment, \ CASE a.attidentity \ WHEN 'd' THEN 'generated by default as identity' || CASE WHEN pseq.seqstart IS NOT NULL THEN format(' (start with %s increment by %s)', pseq.seqstart, pseq.seqincrement) ELSE '' END \ WHEN 'a' THEN 'generated always as identity' || CASE WHEN pseq.seqstart IS NOT NULL THEN format(' (start with %s increment by %s)', pseq.seqstart, pseq.seqincrement) ELSE '' END \ ELSE CASE a.attgenerated \ WHEN 's' THEN 'generated always as (' || pg_get_expr(ad.adbin, ad.adrelid) || ') stored' \ WHEN 'v' THEN 'generated always as (' || pg_get_expr(ad.adbin, ad.adrelid) || ') virtual' \ ELSE CASE WHEN a.atttypid IN (20, 21, 23) AND dep.deptype = 'a' \ AND pseq.seqrelid IS NOT NULL \ AND pg_get_expr(ad.adbin, ad.adrelid) = format('nextval(%L::regclass)', dep.objid::regclass::text) \ THEN CASE a.atttypid \ WHEN 21 THEN 'smallserial' \ WHEN 23 THEN 'serial' \ WHEN 20 THEN 'bigserial' \ END ELSE NULL END \ END \ END AS column_extra, \ CASE WHEN t.typname = 'numeric' AND a.atttypmod > 0 \ THEN ((a.atttypmod - 4) >> 16) & 65535 ELSE NULL END AS numeric_precision, \ CASE WHEN t.typname = 'numeric' AND a.atttypmod > 0 \ THEN (a.atttypmod - 4) & 65535 ELSE NULL END AS numeric_scale, \ CASE WHEN t.typname IN ('varchar', 'bpchar') AND a.atttypmod > 0 \ THEN a.atttypmod - 4 ELSE NULL END AS character_maximum_length, \ CASE WHEN enum_t.oid IS NULL THEN NULL \ ELSE COALESCE((SELECT array_to_json(array_agg(e.enumlabel ORDER BY e.enumsortorder))::text \ FROM pg_enum e WHERE e.enumtypid = enum_t.oid), '[]') END AS enum_values \ FROM pg_attribute a \ JOIN pg_class c ON c.oid = a.attrelid \ JOIN pg_namespace n ON n.oid = c.relnamespace \ JOIN pg_type t ON t.oid = a.atttypid \ LEFT JOIN pg_type enum_t ON enum_t.oid = CASE WHEN t.typtype = 'd' THEN t.typbasetype WHEN t.typtype = 'e' THEN t.oid ELSE NULL END AND enum_t.typtype = 'e' \ LEFT JOIN pg_attrdef ad ON ad.adrelid = a.attrelid AND ad.adnum = a.attnum \ LEFT JOIN LATERAL ( \ SELECT sequence_dep.objid, sequence_dep.deptype \ FROM pg_catalog.pg_depend sequence_dep \ JOIN pg_catalog.pg_class sequence_class \ ON sequence_class.oid = sequence_dep.objid AND sequence_class.relkind = 'S' \ WHERE sequence_dep.classid = 'pg_catalog.pg_class'::pg_catalog.regclass \ AND sequence_dep.objsubid = 0 \ AND sequence_dep.refclassid = 'pg_catalog.pg_class'::pg_catalog.regclass \ AND sequence_dep.refobjid = a.attrelid AND sequence_dep.refobjsubid = a.attnum \ AND ((a.attidentity <> '' AND sequence_dep.deptype = 'i') OR (a.attidentity = '' \ AND sequence_dep.deptype = 'a' AND EXISTS ( \ SELECT 1 FROM pg_catalog.pg_depend serial_default_dep \ WHERE serial_default_dep.classid = 'pg_catalog.pg_attrdef'::pg_catalog.regclass \ AND serial_default_dep.objid = ad.oid AND serial_default_dep.objsubid = 0 \ AND serial_default_dep.refclassid = 'pg_catalog.pg_class'::pg_catalog.regclass \ AND serial_default_dep.refobjid = sequence_dep.objid \ AND serial_default_dep.refobjsubid = 0 AND serial_default_dep.deptype = 'n' \ ))) \ ORDER BY sequence_dep.objid \ LIMIT 1 \ ) dep ON TRUE \ LEFT JOIN pg_sequence pseq ON pseq.seqrelid = dep.objid \ LEFT JOIN information_schema.columns ic \ ON ic.table_schema = n.nspname AND ic.table_name = c.relname AND ic.column_name = a.attname \ WHERE c.oid = ANY($1::bigint[]) \ AND a.attnum > 0 AND NOT a.attisdropped \ ORDER BY c.oid, a.attnum" } // Compat-tier sibling of `POSTGRES_COLUMNS_COMPAT_SQL` (~line 4938) — see the // note on `postgres_columns_for_relations_sql` above. fn postgres_columns_for_relations_compat_sql() -> &'static str { "SELECT c.oid::bigint AS relid, a.attname AS column_name, \ format_type(a.atttypid, a.atttypmod) AS full_type, \ COALESCE(ic.is_nullable = 'YES', NOT a.attnotnull) AS is_nullable, \ pg_get_expr(ad.adbin, ad.adrelid) AS column_default, \ EXISTS ( \ SELECT 1 FROM pg_constraint co \ JOIN pg_index i ON i.indrelid = co.conrelid AND co.conindid = i.indexrelid \ WHERE co.conrelid = a.attrelid AND co.contype = 'p' \ AND a.attnum = ANY(i.indkey) \ ) AS is_pk, \ col_description(a.attrelid, a.attnum) AS column_comment, \ CASE WHEN a.atttypid IN (20, 21, 23) AND serial_seq.oid IS NOT NULL \ AND pg_get_expr(ad.adbin, ad.adrelid) = format('nextval(%L::regclass)', serial_seq.oid::regclass::text) \ THEN CASE a.atttypid \ WHEN 21 THEN 'smallserial' \ WHEN 23 THEN 'serial' \ WHEN 20 THEN 'bigserial' \ END ELSE NULL END AS column_extra, \ CASE WHEN t.typname = 'numeric' AND a.atttypmod > 0 \ THEN ((a.atttypmod - 4) >> 16) & 65535 ELSE NULL END AS numeric_precision, \ CASE WHEN t.typname = 'numeric' AND a.atttypmod > 0 \ THEN (a.atttypmod - 4) & 65535 ELSE NULL END AS numeric_scale, \ CASE WHEN t.typname IN ('varchar', 'bpchar') AND a.atttypmod > 0 \ THEN a.atttypmod - 4 ELSE NULL END AS character_maximum_length, \ NULL::text AS enum_values \ FROM pg_attribute a \ JOIN pg_class c ON c.oid = a.attrelid \ JOIN pg_namespace n ON n.oid = c.relnamespace \ JOIN pg_type t ON t.oid = a.atttypid \ LEFT JOIN pg_attrdef ad ON ad.adrelid = a.attrelid AND ad.adnum = a.attnum \ LEFT JOIN pg_catalog.pg_class serial_seq ON serial_seq.oid = ( \ SELECT sequence_dep.objid \ FROM pg_catalog.pg_depend sequence_dep \ JOIN pg_catalog.pg_class sequence_class \ ON sequence_class.oid = sequence_dep.objid AND sequence_class.relkind = 'S' \ WHERE sequence_dep.classid = 'pg_catalog.pg_class'::pg_catalog.regclass \ AND sequence_dep.objsubid = 0 \ AND sequence_dep.refclassid = 'pg_catalog.pg_class'::pg_catalog.regclass \ AND sequence_dep.refobjid = a.attrelid AND sequence_dep.refobjsubid = a.attnum \ AND sequence_dep.deptype = 'a' AND EXISTS ( \ SELECT 1 FROM pg_catalog.pg_depend serial_default_dep \ WHERE serial_default_dep.classid = 'pg_catalog.pg_attrdef'::pg_catalog.regclass \ AND serial_default_dep.objid = ad.oid AND serial_default_dep.objsubid = 0 \ AND serial_default_dep.refclassid = 'pg_catalog.pg_class'::pg_catalog.regclass \ AND serial_default_dep.refobjid = sequence_dep.objid \ AND serial_default_dep.refobjsubid = 0 AND serial_default_dep.deptype = 'n' \ ) \ ORDER BY sequence_dep.objid \ LIMIT 1 \ ) AND serial_seq.relkind = 'S' \ LEFT JOIN information_schema.columns ic \ ON ic.table_schema = n.nspname AND ic.table_name = c.relname AND ic.column_name = a.attname \ WHERE c.oid = ANY($1::bigint[]) \ AND a.attnum > 0 AND NOT a.attisdropped \ ORDER BY c.oid, a.attnum" } /// Same field layout as `column_info_from_row`, offset by one leading `relid` /// column. fn column_info_from_row_offset(row: &Row, offset: usize) -> ColumnInfo { let full_type = row.try_get::<_, Option>(offset + 1).ok().flatten().unwrap_or_default(); ColumnInfo { name: pg_row_try_string(row, offset), data_type: full_type, is_nullable: pg_row_try_bool(row, offset + 2).unwrap_or(true), column_default: row.try_get::<_, Option>(offset + 3).ok().flatten(), is_primary_key: pg_row_try_bool(row, offset + 4).unwrap_or(false), extra: row.try_get::<_, Option>(offset + 6).ok().flatten(), comment: row.try_get::<_, Option>(offset + 5).ok().flatten(), numeric_precision: row.try_get::<_, Option>(offset + 7).ok().flatten(), numeric_scale: row.try_get::<_, Option>(offset + 8).ok().flatten(), character_maximum_length: row.try_get::<_, Option>(offset + 9).ok().flatten(), enum_values: parse_enum_values_from_row(row, offset + 10), ..Default::default() } } /// Batched sibling of list_indexes, using a modern and a PostgreSQL 10 /// compatible OID-scoped tier so the request count stays bounded. pub async fn list_indexes_for_relations( pool: &Pool, relations: &[(i64, String, String)], ) -> Result>, String> { let oids: Vec = relations.iter().map(|(oid, _, _)| *oid).collect(); let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let tiers = postgres_indexes_for_relations_query_tiers(); query_with_compat_fallback("list_indexes_for_relations", &tiers, |sql| { list_indexes_for_relations_with_sql(&client, sql, &oids) }) .await } async fn list_indexes_for_relations_with_sql( client: &deadpool_postgres::Client, sql: &str, oids: &[i64], ) -> Result>, tokio_postgres::Error> { let rows = postgres_query_cached(client, sql, &[&oids]).await?; let mut result: HashMap> = HashMap::new(); for row in &rows { let Ok(relid) = row.try_get::<_, i64>(0) else { continue }; let all_cols: Vec = row.try_get::<_, Vec>(2).unwrap_or_default(); let all_opclasses: Vec> = row.try_get::<_, Vec>>(3).unwrap_or_default(); let nkeyatts = row.try_get::<_, Option>(8).ok().flatten().unwrap_or(all_cols.len() as i16) as usize; let split_at = nkeyatts.min(all_cols.len()); let key_cols = all_cols[..split_at].to_vec(); let key_opclasses = if all_opclasses.len() == all_cols.len() { all_opclasses[..split_at].to_vec() } else { vec![None; split_at] }; let included = if split_at < all_cols.len() { all_cols[split_at..].to_vec() } else { vec![] }; let all_is_expr: Vec = row.try_get::<_, Vec>(11).unwrap_or_default(); let key_is_expression = if all_is_expr.len() == all_cols.len() { all_is_expr[..split_at].to_vec() } else { Vec::new() }; result.entry(relid).or_default().push(IndexInfo { name: pg_row_try_string(row, 1), columns: key_cols, is_unique: pg_row_try_bool(row, 4).unwrap_or(false), is_primary: pg_row_try_bool(row, 5).unwrap_or(false), filter: row.try_get::<_, Option>(6).ok().flatten(), index_type: row.try_get::<_, Option>(7).ok().flatten(), included_columns: if included.is_empty() { None } else { Some(included) }, comment: row.try_get::<_, Option>(10).ok().flatten(), key_is_expression, column_opclasses: key_opclasses, }); } Ok(result) } fn postgres_indexes_for_relations_query_tiers() -> [&'static str; 2] { [postgres_indexes_for_relations_sql(), postgres_indexes_for_relations_compat_sql()] } // Sibling of `POSTGRES_INDEXES_SQL`/`POSTGRES_INDEXES_COMPAT_SQL` (~line // 6042): same index-detection logic, batched by oid instead of a single // (schema, table) pair. Not merged into a shared fragment for the same // reason as the columns queries above — an alias would need renaming to // line up, and result columns are read positionally. fn postgres_indexes_for_relations_sql() -> &'static str { "SELECT t.oid::bigint AS relid, i.relname AS index_name, \ array_agg(COALESCE(a.attname, pg_get_indexdef(ix.indexrelid, k.n::int, false)) ORDER BY k.n) AS columns, \ array_agg(CASE WHEN oc.opcdefault THEN NULL ELSE quote_ident(opcns.nspname) || '.' || quote_ident(oc.opcname) END ORDER BY k.n) AS column_opclasses, \ (ix.indisunique AND ix.indisvalid) AS is_unique, \ ix.indisprimary AS is_primary, \ pg_get_expr(ix.indpred, ix.indrelid) AS filter_expr, \ am.amname AS index_type, \ ix.indnkeyatts AS nkeyatts, \ ix.indkey AS indkey, \ obj_description(i.oid, 'pg_class') AS index_comment, \ array_agg(a.attname IS NULL ORDER BY k.n) AS key_is_expression \ FROM pg_index ix \ JOIN pg_class t ON t.oid = ix.indrelid \ JOIN pg_class i ON i.oid = ix.indexrelid \ JOIN pg_am am ON am.oid = i.relam \ JOIN LATERAL unnest(ix.indkey, ix.indclass) WITH ORDINALITY AS k(attnum, class_oid, n) ON true \ LEFT JOIN pg_attribute a ON a.attrelid = t.oid AND a.attnum = k.attnum AND k.attnum > 0 \ LEFT JOIN pg_opclass oc ON oc.oid = k.class_oid \ LEFT JOIN pg_namespace opcns ON opcns.oid = oc.opcnamespace \ WHERE t.oid = ANY($1::bigint[]) \ GROUP BY t.oid, i.relname, i.oid, ix.indisunique, ix.indisvalid, ix.indisprimary, ix.indpred, ix.indrelid, am.amname, ix.indnkeyatts, ix.indkey \ ORDER BY t.oid, i.relname" } // Compat-tier sibling of `POSTGRES_INDEXES_COMPAT_SQL` (~line 6063) — see the // note on `postgres_indexes_for_relations_sql` above. fn postgres_indexes_for_relations_compat_sql() -> &'static str { "SELECT t.oid::bigint AS relid, i.relname AS index_name, \ ARRAY( \ SELECT COALESCE(a.attname, pg_get_indexdef(ix.indexrelid, pos.n, false)) \ FROM generate_series(1, array_length(string_to_array(ix.indkey::text, ' '), 1)) AS pos(n) \ LEFT JOIN pg_attribute a \ ON a.attrelid = t.oid \ AND a.attnum = (string_to_array(ix.indkey::text, ' '))[pos.n]::int2 \ AND a.attnum > 0 \ ORDER BY pos.n \ ) AS columns, \ ARRAY( \ SELECT CASE WHEN oc.opcdefault THEN NULL \ ELSE oc.opcname \ END \ FROM generate_series(1, array_length(string_to_array(ix.indkey::text, ' '), 1)) AS pos(n) \ LEFT JOIN pg_attribute a \ ON a.attrelid = t.oid \ AND a.attnum = (string_to_array(ix.indkey::text, ' '))[pos.n]::int2 \ AND a.attnum > 0 \ LEFT JOIN pg_opclass oc \ ON oc.oid = (string_to_array(ix.indclass::text, ' '))[pos.n]::oid \ ORDER BY pos.n \ ) AS column_opclasses, \ (ix.indisunique AND ix.indisvalid) AS is_unique, \ ix.indisprimary AS is_primary, \ pg_get_expr(ix.indpred, ix.indrelid) AS filter_expr, \ am.amname AS index_type, \ NULL::smallint AS nkeyatts, \ ix.indkey AS indkey, \ obj_description(i.oid, 'pg_class') AS index_comment, \ ARRAY( \ SELECT a.attname IS NULL \ FROM generate_series(1, array_length(string_to_array(ix.indkey::text, ' '), 1)) AS pos(n) \ LEFT JOIN pg_attribute a \ ON a.attrelid = t.oid \ AND a.attnum = (string_to_array(ix.indkey::text, ' '))[pos.n]::int2 \ AND a.attnum > 0 \ ORDER BY pos.n \ ) AS key_is_expression \ FROM pg_index ix \ JOIN pg_class t ON t.oid = ix.indrelid \ JOIN pg_class i ON i.oid = ix.indexrelid \ JOIN pg_am am ON am.oid = i.relam \ WHERE t.oid = ANY($1::bigint[]) \ ORDER BY t.oid, i.relname" } /// Batched sibling of `list_foreign_keys`, keyed by `(schema, table)`. pub async fn list_foreign_keys_for_relations( pool: &Pool, relations: &[(String, String)], ) -> Result>, String> { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let schemas: Vec<&str> = relations.iter().map(|(schema, _)| schema.as_str()).collect(); let tables: Vec<&str> = relations.iter().map(|(_, table)| table.as_str()).collect(); let tiers = postgres_foreign_keys_for_relations_query_tiers(); query_with_compat_fallback("list_foreign_keys_for_relations", &tiers, |sql| { list_foreign_keys_for_relations_with_sql(&client, sql, &schemas, &tables) }) .await } async fn list_foreign_keys_for_relations_with_sql( client: &deadpool_postgres::Client, sql: &str, schemas: &[&str], tables: &[&str], ) -> Result>, tokio_postgres::Error> { let rows = postgres_query_cached(client, sql, &[&schemas, &tables]).await?; let mut result: HashMap<(String, String), Vec> = HashMap::new(); for row in &rows { let key = (pg_row_try_string(row, 0), pg_row_try_string(row, 1)); result.entry(key).or_default().push(ForeignKeyInfo { name: pg_row_try_string(row, 2), column: pg_row_try_string(row, 3), ref_schema: Some(pg_row_try_string(row, 4)), ref_table: pg_row_try_string(row, 5), ref_column: pg_row_try_string(row, 6), on_update: postgres_fk_action_label(pg_row_try_optional_text(row, 7)), on_delete: postgres_fk_action_label(pg_row_try_optional_text(row, 8)), }); } Ok(result) } fn postgres_foreign_keys_for_relations_query_tiers() -> [&'static str; 2] { [postgres_foreign_keys_for_relations_sql(), postgres_foreign_keys_for_relations_compat_sql()] } fn postgres_foreign_keys_for_relations_sql() -> &'static str { "SELECT rel.rel_schema, rel.rel_table, \ con.conname AS constraint_name, \ a.attname AS column_name, \ ref_n.nspname AS ref_schema, \ ref_c.relname AS ref_table, \ ref_a.attname AS ref_column, \ con.confupdtype::text AS on_update_raw, \ con.confdeltype::text AS on_delete_raw \ FROM unnest($1::text[], $2::text[]) AS rel(rel_schema, rel_table) \ JOIN pg_catalog.pg_constraint con ON true \ JOIN pg_catalog.pg_class c ON c.oid = con.conrelid \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ JOIN pg_catalog.pg_class ref_c ON ref_c.oid = con.confrelid \ JOIN pg_catalog.pg_namespace ref_n ON ref_n.oid = ref_c.relnamespace \ JOIN LATERAL unnest(con.conkey) WITH ORDINALITY AS fk(attnum, ord) ON true \ JOIN pg_catalog.pg_attribute a ON a.attrelid = c.oid AND a.attnum = fk.attnum AND NOT a.attisdropped \ JOIN LATERAL unnest(con.confkey) WITH ORDINALITY AS pk(attnum, ord) ON pk.ord = fk.ord \ JOIN pg_catalog.pg_attribute ref_a ON ref_a.attrelid = ref_c.oid AND ref_a.attnum = pk.attnum AND NOT ref_a.attisdropped \ WHERE con.contype = 'f' AND n.nspname = rel.rel_schema AND c.relname = rel.rel_table \ ORDER BY rel.rel_schema, rel.rel_table, con.conname, fk.ord" } // PostgreSQL 9.3 and older only accept one array argument to unnest(). Pair // the schema/table arrays through their shared subscript so the fallback stays // one bounded query and preserves each relation tuple's position. // WITH ORDINALITY is equally unavailable before 9.4, so pair conkey/confkey // positions with generate_series over the array length plus plain subscripts — // the same pre-9.4 technique the index compat SQL relies on. fn postgres_foreign_keys_for_relations_compat_sql() -> &'static str { "SELECT rel.rel_schema, rel.rel_table, \ con.conname AS constraint_name, \ a.attname AS column_name, \ ref_n.nspname AS ref_schema, \ ref_c.relname AS ref_table, \ ref_a.attname AS ref_column, \ con.confupdtype::text AS on_update_raw, \ con.confdeltype::text AS on_delete_raw \ FROM ( \ SELECT ($1::text[])[rel.i] AS rel_schema, ($2::text[])[rel.i] AS rel_table \ FROM generate_subscripts($1::text[], 1) AS rel(i) \ ) AS rel \ JOIN pg_catalog.pg_constraint con ON true \ JOIN pg_catalog.pg_class c ON c.oid = con.conrelid \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ JOIN pg_catalog.pg_class ref_c ON ref_c.oid = con.confrelid \ JOIN pg_catalog.pg_namespace ref_n ON ref_n.oid = ref_c.relnamespace \ CROSS JOIN generate_series(1, array_length(con.conkey, 1)) AS fk(ord) \ JOIN pg_catalog.pg_attribute a ON a.attrelid = c.oid AND a.attnum = (con.conkey)[fk.ord] AND NOT a.attisdropped \ JOIN pg_catalog.pg_attribute ref_a ON ref_a.attrelid = ref_c.oid AND ref_a.attnum = (con.confkey)[fk.ord] AND NOT ref_a.attisdropped \ WHERE con.contype = 'f' AND n.nspname = rel.rel_schema AND c.relname = rel.rel_table \ ORDER BY rel.rel_schema, rel.rel_table, con.conname, fk.ord" } /// Batched sibling of `get_table_comment`. pub async fn get_table_comments_for_relations( pool: &Pool, oids: &[i64], ) -> Result>, String> { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let rows = postgres_query_cached( &client, "SELECT c.oid::bigint AS relid, obj_description(c.oid) AS table_comment \ FROM pg_catalog.pg_class c WHERE c.oid = ANY($1::bigint[])", &[&oids], ) .await .map_err(|e| e.to_string())?; let mut result: HashMap> = HashMap::new(); for row in &rows { let Ok(relid) = row.try_get::<_, i64>(0) else { continue }; let comment = row.try_get::<_, Option>(1).ok().flatten().filter(|s| !s.is_empty()); result.insert(relid, comment); } Ok(result) } /// Batched sibling of `list_trigger_definitions`, with the same pre-13 /// `tgparentid` capability probe. pub async fn list_trigger_definitions_for_relations( pool: &Pool, oids: &[i64], ) -> Result>, String> { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let sql = if postgres_trigger_has_tgparentid(&client).await? { postgres_trigger_definitions_for_relations_sql() } else { postgres_trigger_definitions_for_relations_sql_without_tgparentid() }; let rows = postgres_query_cached(&client, sql, &[&oids]).await.map_err(|e| e.to_string())?; let mut result: HashMap> = HashMap::new(); for row in &rows { let Ok(relid) = row.try_get::<_, i64>(0) else { continue }; let definition = pg_row_try_string(row, 1); if !definition.trim().is_empty() { result.entry(relid).or_default().push(definition); } } Ok(result) } fn postgres_trigger_definitions_for_relations_sql() -> &'static str { "SELECT t.tgrelid::bigint AS relid, pg_catalog.pg_get_triggerdef(t.oid, true) AS trigger_definition \ FROM pg_catalog.pg_trigger t \ WHERE t.tgrelid = ANY($1::bigint[]) AND NOT t.tgisinternal AND t.tgparentid = 0 \ ORDER BY t.tgrelid, t.tgname, t.oid" } fn postgres_trigger_definitions_for_relations_sql_without_tgparentid() -> &'static str { "SELECT t.tgrelid::bigint AS relid, pg_catalog.pg_get_triggerdef(t.oid, true) AS trigger_definition \ FROM pg_catalog.pg_trigger t \ WHERE t.tgrelid = ANY($1::bigint[]) AND NOT t.tgisinternal \ ORDER BY t.tgrelid, t.tgname, t.oid" } /// Batched sibling of `list_check_constraints`. pub async fn list_check_constraints_for_relations( pool: &Pool, oids: &[i64], ) -> Result>, String> { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let rows = postgres_query_cached( &client, "SELECT con.conrelid::bigint AS relid, con.conname, pg_catalog.pg_get_constraintdef(con.oid, true) AS definition \ FROM pg_catalog.pg_constraint con \ WHERE con.conrelid = ANY($1::bigint[]) AND con.contype = 'c' \ ORDER BY con.conrelid, con.conname", &[&oids], ) .await .map_err(|e| e.to_string())?; let mut result: HashMap> = HashMap::new(); for row in &rows { let Ok(relid) = row.try_get::<_, i64>(0) else { continue }; let name = pg_row_try_string(row, 1); let definition = pg_row_try_string(row, 2); if !name.is_empty() && !definition.is_empty() { result.entry(relid).or_default().push((name, definition)); } } Ok(result) } /// Batched sibling of `get_table_partition_local_objects`. pub async fn get_table_partition_local_objects_for_relations( pool: &Pool, oids: &[i64], ) -> Result, String> { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let tiers = postgres_table_partition_local_objects_for_relations_query_tiers(); query_with_compat_fallback("get_table_partition_local_objects_for_relations", &tiers, |sql| { get_table_partition_local_objects_for_relations_with_sql(&client, sql, oids) }) .await } async fn get_table_partition_local_objects_for_relations_with_sql( client: &deadpool_postgres::Client, sql: &str, oids: &[i64], ) -> Result, tokio_postgres::Error> { let rows = postgres_query_cached(client, sql, &[&oids]).await?; let mut result: HashMap = HashMap::new(); for row in &rows { let Ok(relid) = row.try_get::<_, i64>(0) else { continue }; let object_kind = row.try_get::<_, String>(1).unwrap_or_default(); let object_name = row.try_get::<_, String>(2).unwrap_or_default(); let object_type = row.try_get::<_, Option>(3).ok().flatten().unwrap_or_default(); apply_partition_local_object_row(result.entry(relid).or_default(), &object_kind, object_name, &object_type); } Ok(result) } fn postgres_table_partition_local_objects_for_relations_query_tiers() -> [&'static str; 2] { [ postgres_table_partition_local_objects_for_relations_sql(), postgres_table_partition_local_objects_for_relations_compat_sql(), ] } fn postgres_table_partition_local_objects_for_relations_sql() -> &'static str { "SELECT con.conrelid::bigint AS relid, 'constraint'::text AS object_kind, con.conname AS object_name, con.contype::text AS object_type \ FROM pg_catalog.pg_constraint con \ WHERE con.conrelid = ANY($1::bigint[]) AND con.contype IN ('p','f') \ AND COALESCE(NULLIF(pg_catalog.row_to_json(con)->>'conparentid', '')::oid, 0) = 0 \ UNION ALL \ SELECT con.conrelid::bigint, 'check'::text AS object_kind, con.conname AS object_name, NULL::text AS object_type \ FROM pg_catalog.pg_constraint con \ WHERE con.conrelid = ANY($1::bigint[]) AND con.contype = 'c' AND con.conislocal \ UNION ALL \ SELECT ix.indrelid::bigint, 'index'::text AS object_kind, idx.relname AS object_name, NULL::text AS object_type \ FROM pg_catalog.pg_index ix \ JOIN pg_catalog.pg_class idx ON idx.oid = ix.indexrelid \ WHERE ix.indrelid = ANY($1::bigint[]) \ AND NOT EXISTS (SELECT 1 FROM pg_catalog.pg_inherits i WHERE i.inhrelid = idx.oid) \ UNION ALL \ SELECT a.attrelid::bigint, 'column_default'::text AS object_kind, a.attname AS object_name, \ CASE WHEN ad.oid IS NOT NULL THEN 'overridden' ELSE 'dropped' END AS object_type \ FROM pg_catalog.pg_attribute a \ JOIN pg_catalog.pg_inherits i ON i.inhrelid = a.attrelid \ JOIN pg_catalog.pg_attribute pa \ ON pa.attrelid = i.inhparent AND pa.attname = a.attname AND pa.attnum > 0 AND NOT pa.attisdropped \ LEFT JOIN pg_catalog.pg_attrdef ad ON ad.adrelid = a.attrelid AND ad.adnum = a.attnum \ LEFT JOIN pg_catalog.pg_attrdef pad ON pad.adrelid = pa.attrelid AND pad.adnum = pa.attnum \ WHERE a.attrelid = ANY($1::bigint[]) AND a.attnum > 0 AND NOT a.attisdropped \ AND ( \ (ad.oid IS NOT NULL AND (pad.oid IS NULL \ OR pg_catalog.pg_get_expr(ad.adbin, ad.adrelid) \ IS DISTINCT FROM pg_catalog.pg_get_expr(pad.adbin, pad.adrelid))) \ OR (ad.oid IS NULL AND pad.oid IS NOT NULL) \ ) \ ORDER BY relid, object_kind, object_name" } // PostgreSQL 9.2 has row_to_json(), but not the JSON ->> operator used above // to probe conparentid without directly referencing that newer catalog // column. Pre-11 servers do not have parent-linked constraint clones, so the // compatibility tier can omit only that filter and keep the same batched // result shape. fn postgres_table_partition_local_objects_for_relations_compat_sql() -> &'static str { "SELECT con.conrelid::bigint AS relid, 'constraint'::text AS object_kind, con.conname AS object_name, con.contype::text AS object_type \ FROM pg_catalog.pg_constraint con \ WHERE con.conrelid = ANY($1::bigint[]) AND con.contype IN ('p','f') \ UNION ALL \ SELECT con.conrelid::bigint, 'check'::text AS object_kind, con.conname AS object_name, NULL::text AS object_type \ FROM pg_catalog.pg_constraint con \ WHERE con.conrelid = ANY($1::bigint[]) AND con.contype = 'c' AND con.conislocal \ UNION ALL \ SELECT ix.indrelid::bigint, 'index'::text AS object_kind, idx.relname AS object_name, NULL::text AS object_type \ FROM pg_catalog.pg_index ix \ JOIN pg_catalog.pg_class idx ON idx.oid = ix.indexrelid \ WHERE ix.indrelid = ANY($1::bigint[]) \ AND NOT EXISTS (SELECT 1 FROM pg_catalog.pg_inherits i WHERE i.inhrelid = idx.oid) \ UNION ALL \ SELECT a.attrelid::bigint, 'column_default'::text AS object_kind, a.attname AS object_name, \ CASE WHEN ad.oid IS NOT NULL THEN 'overridden' ELSE 'dropped' END AS object_type \ FROM pg_catalog.pg_attribute a \ JOIN pg_catalog.pg_inherits i ON i.inhrelid = a.attrelid \ JOIN pg_catalog.pg_attribute pa \ ON pa.attrelid = i.inhparent AND pa.attname = a.attname AND pa.attnum > 0 AND NOT pa.attisdropped \ LEFT JOIN pg_catalog.pg_attrdef ad ON ad.adrelid = a.attrelid AND ad.adnum = a.attnum \ LEFT JOIN pg_catalog.pg_attrdef pad ON pad.adrelid = pa.attrelid AND pad.adnum = pa.attnum \ WHERE a.attrelid = ANY($1::bigint[]) AND a.attnum > 0 AND NOT a.attisdropped \ AND ( \ (ad.oid IS NOT NULL AND (pad.oid IS NULL \ OR pg_catalog.pg_get_expr(ad.adbin, ad.adrelid) \ IS DISTINCT FROM pg_catalog.pg_get_expr(pad.adbin, pad.adrelid))) \ OR (ad.oid IS NULL AND pad.oid IS NOT NULL) \ ) \ ORDER BY relid, object_kind, object_name" } /// CHECK constraints on a table, as (name, definition) pairs (definition /// includes the leading `CHECK` keyword, e.g. `CHECK (id > 0)`). pub async fn list_check_constraints(pool: &Pool, schema: &str, table: &str) -> Result, String> { let schema = if schema.is_empty() { "public" } else { schema }; let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let rows = postgres_query_cached(&client, postgres_check_constraints_sql(), &[&schema, &table]) .await .map_err(|e| e.to_string())?; Ok(rows .into_iter() .filter_map(|row| { let name = row.try_get::<_, String>(0).ok()?; let definition = row.try_get::<_, String>(1).ok()?; (!name.is_empty() && !definition.is_empty()).then_some((name, definition)) }) .collect()) } fn postgres_check_constraints_sql() -> &'static str { "SELECT con.conname, pg_catalog.pg_get_constraintdef(con.oid, true) AS definition \ FROM pg_catalog.pg_constraint con \ JOIN pg_catalog.pg_class c ON c.oid = con.conrelid \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ WHERE n.nspname = $1 AND c.relname = $2 AND con.contype = 'c' \ ORDER BY con.conname" } /// All constraints on a table (primary key, foreign key, unique, check, /// exclude, and — on PG 18+ — not-null), as full `ConstraintInfo` records. /// `pg_constraint.contype` letters are expanded to display labels and the /// `conkey`/`confkey` attribute numbers are resolved to column names. /// PostgreSQL has no per-constraint disabled state, so `enabled` is always /// true; a `NOT VALID` constraint (e.g. an added-but-unvalidated CHECK or /// FK) reports `valid = false` and surfaces the "Not validated" badge. pub async fn list_constraints(pool: &Pool, schema: &str, table: &str) -> Result, String> { let schema = if schema.is_empty() { "public" } else { schema }; let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let rows = postgres_query_cached(&client, postgres_constraints_sql(), &[&schema, &table]) .await .map_err(|e| e.to_string())?; Ok(rows .iter() .map(|row| ConstraintInfo { name: pg_row_try_string(row, 0), constraint_type: postgres_constraint_type_label(&pg_row_try_string(row, 1)), definition: pg_row_try_string(row, 2), columns: row.try_get::<_, Vec>(3).unwrap_or_default(), ref_schema: row.try_get::<_, Option>(4).ok().flatten(), ref_table: row.try_get::<_, Option>(5).ok().flatten(), ref_columns: row.try_get::<_, Vec>(6).unwrap_or_default(), match_type: postgres_constraint_match_type(row.try_get::<_, Option>(7).ok().flatten()), on_update: postgres_fk_action_label(row.try_get::<_, Option>(8).ok().flatten()), on_delete: postgres_fk_action_label(row.try_get::<_, Option>(9).ok().flatten()), deferrable: row.try_get::<_, bool>(10).unwrap_or(false), initially_deferred: row.try_get::<_, bool>(11).unwrap_or(false), enabled: true, valid: row.try_get::<_, bool>(12).unwrap_or(true), }) .collect()) } /// OpenGauss-compatible constraint metadata without PostgreSQL's array/LATERAL /// dependencies. Older OpenGauss releases expose conkey/confkey as catalog /// vectors but may reject WITH ORDINALITY or decode them differently on the /// wire, so resolve attribute numbers in Rust. pub async fn list_opengauss_constraints(pool: &Pool, schema: &str, table: &str) -> Result, String> { let schema = if schema.is_empty() { "public" } else { schema }; let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let rows = match postgres_query_cached(&client, opengauss_constraints_sql(true), &[&schema, &table]).await { Ok(rows) => rows, Err(error) if is_missing_opengauss_constraint_column(&error, "convalidated") => { log::debug!( "[opengauss][constraints] convalidated unavailable; defaulting valid=true: {}", pg_error_to_string(error) ); postgres_query_cached(&client, opengauss_constraints_sql(false), &[&schema, &table]) .await .map_err(|e| e.to_string())? } Err(error) => return Err(error.to_string()), }; if rows.is_empty() { return Ok(Vec::new()); } let local_relation_oid = pg_row_try_u32(&rows[0], 13).ok_or("OpenGauss constraint metadata did not return a local relation OID")?; let local_attributes = opengauss_relation_attributes(&client, local_relation_oid).await?; let mut referenced_attributes: HashMap> = HashMap::new(); let mut result = Vec::with_capacity(rows.len()); for row in rows { let name = pg_row_try_string(&row, 0); let kind = pg_row_try_string(&row, 1); let column_numbers = parse_opengauss_attribute_numbers(&pg_row_try_string(&row, 3)) .map_err(|error| format!("failed to parse OpenGauss constraint {name} columns: {error}"))?; let ref_column_numbers = parse_opengauss_attribute_numbers(&pg_row_try_string(&row, 6)) .map_err(|error| format!("failed to parse OpenGauss constraint {name} referenced columns: {error}"))?; let referenced_relation_oid = pg_row_try_u32(&row, 14).filter(|oid| *oid != 0); let ref_columns = if kind == "f" { if let Some(oid) = referenced_relation_oid { if let std::collections::hash_map::Entry::Vacant(entry) = referenced_attributes.entry(oid) { entry.insert(opengauss_relation_attributes(&client, oid).await?); } map_opengauss_attribute_names(&ref_column_numbers, referenced_attributes.get(&oid).unwrap()) } else { Vec::new() } } else { Vec::new() }; result.push(ConstraintInfo { name, constraint_type: postgres_constraint_type_label(&kind), definition: pg_row_try_string(&row, 2), columns: map_opengauss_attribute_names(&column_numbers, &local_attributes), ref_schema: row.try_get::<_, Option>(4).ok().flatten(), ref_table: row.try_get::<_, Option>(5).ok().flatten(), ref_columns, match_type: postgres_constraint_match_type(pg_row_try_optional_text(&row, 7)), on_update: postgres_fk_action_label(pg_row_try_optional_text(&row, 8)), on_delete: postgres_fk_action_label(pg_row_try_optional_text(&row, 9)), deferrable: pg_row_try_bool(&row, 10).unwrap_or(false), initially_deferred: pg_row_try_bool(&row, 11).unwrap_or(false), enabled: true, valid: pg_row_try_bool(&row, 12).unwrap_or(true), }); } Ok(result) } fn opengauss_constraints_sql(include_validated: bool) -> &'static str { if include_validated { "SELECT COALESCE(con.conname, ''), \ con.contype::text, \ COALESCE(pg_catalog.pg_get_constraintdef(con.oid, true), ''), \ COALESCE(con.conkey::text, ''), \ refn.nspname, refc.relname, COALESCE(con.confkey::text, ''), \ CASE WHEN con.contype = 'f' THEN con.confmatchtype::text END, \ CASE WHEN con.contype = 'f' THEN con.confupdtype::text END, \ CASE WHEN con.contype = 'f' THEN con.confdeltype::text END, \ con.condeferrable, con.condeferred, con.convalidated, \ con.conrelid::oid, con.confrelid::oid \ FROM pg_catalog.pg_constraint con \ JOIN pg_catalog.pg_class c ON c.oid = con.conrelid \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ LEFT JOIN pg_catalog.pg_class refc ON refc.oid = con.confrelid \ LEFT JOIN pg_catalog.pg_namespace refn ON refn.oid = refc.relnamespace \ WHERE n.nspname = $1 AND c.relname = $2 AND c.relkind IN ('r','p','f') \ ORDER BY COALESCE(con.conname, '')" } else { "SELECT COALESCE(con.conname, ''), \ con.contype::text, \ COALESCE(pg_catalog.pg_get_constraintdef(con.oid, true), ''), \ COALESCE(con.conkey::text, ''), \ refn.nspname, refc.relname, COALESCE(con.confkey::text, ''), \ CASE WHEN con.contype = 'f' THEN con.confmatchtype::text END, \ CASE WHEN con.contype = 'f' THEN con.confupdtype::text END, \ CASE WHEN con.contype = 'f' THEN con.confdeltype::text END, \ con.condeferrable, con.condeferred, TRUE, \ con.conrelid::oid, con.confrelid::oid \ FROM pg_catalog.pg_constraint con \ JOIN pg_catalog.pg_class c ON c.oid = con.conrelid \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ LEFT JOIN pg_catalog.pg_class refc ON refc.oid = con.confrelid \ LEFT JOIN pg_catalog.pg_namespace refn ON refn.oid = refc.relnamespace \ WHERE n.nspname = $1 AND c.relname = $2 AND c.relkind IN ('r','p','f') \ ORDER BY COALESCE(con.conname, '')" } } async fn opengauss_relation_attributes( client: &tokio_postgres::Client, relation_oid: u32, ) -> Result, String> { let rows = client .query( "SELECT a.attnum, a.attname::text FROM pg_catalog.pg_attribute a \ WHERE a.attrelid = $1::oid AND a.attnum > 0 AND NOT a.attisdropped ORDER BY a.attnum", &[&relation_oid], ) .await .map_err(|e| e.to_string())?; Ok(rows.iter().filter_map(|row| Some((row.try_get::<_, i16>(0).ok()?, pg_row_try_string(row, 1)))).collect()) } fn parse_opengauss_attribute_numbers(raw: &str) -> Result, String> { let value = raw.trim().trim_matches(|character| matches!(character, '{' | '}' | '[' | ']')); if value.is_empty() { return Ok(Vec::new()); } value .split(|character: char| character == ',' || character.is_whitespace()) .filter(|part| !part.is_empty()) .map(|part| part.parse::().map_err(|_| format!("invalid attribute number {part:?}"))) .collect() } fn map_opengauss_attribute_names(numbers: &[i16], attributes: &HashMap) -> Vec { numbers.iter().filter_map(|number| attributes.get(number).cloned()).collect() } fn pg_row_try_u32(row: &Row, idx: usize) -> Option { row.try_get::<_, u32>(idx) .ok() .or_else(|| row.try_get::<_, i64>(idx).ok().and_then(|value| u32::try_from(value).ok())) .or_else(|| pg_row_try_string(row, idx).parse::().ok()) } fn pg_row_try_optional_text(row: &Row, idx: usize) -> Option { if let Ok(value) = row.try_get::<_, Option>(idx) { return value.filter(|value| !value.is_empty()); } row.try_get::<_, PgRawBytes>(idx) .ok() .and_then(|value| String::from_utf8(value.0).ok()) .filter(|value| !value.is_empty()) } fn is_missing_opengauss_constraint_column(error: &tokio_postgres::Error, column: &str) -> bool { error.as_db_error().is_some_and(|db_error| { db_error.code().code() == "42703" && db_error.message().to_ascii_lowercase().contains(column) }) } fn postgres_constraints_sql() -> &'static str { "SELECT con.conname, \ con.contype::text, \ pg_catalog.pg_get_constraintdef(con.oid, true) AS definition, \ COALESCE(conkey.attnames, ARRAY[]::text[]) AS columns, \ refn.nspname AS ref_schema, \ refc.relname AS ref_table, \ COALESCE(confkey.attnames, ARRAY[]::text[]) AS ref_columns, \ CASE WHEN con.contype = 'f' THEN con.confmatchtype::text END AS match_type, \ CASE WHEN con.contype = 'f' THEN con.confupdtype::text END AS on_update, \ CASE WHEN con.contype = 'f' THEN con.confdeltype::text END AS on_delete, \ con.condeferrable, \ con.condeferred, \ con.convalidated \ FROM pg_catalog.pg_constraint con \ JOIN pg_catalog.pg_class c ON c.oid = con.conrelid \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ LEFT JOIN pg_catalog.pg_class refc ON refc.oid = con.confrelid \ LEFT JOIN pg_catalog.pg_namespace refn ON refn.oid = refc.relnamespace \ LEFT JOIN LATERAL ( \ SELECT array_agg(a.attname::text ORDER BY ord.ord) AS attnames \ FROM unnest(con.conkey) WITH ORDINALITY AS ord(attnum, ord) \ JOIN pg_catalog.pg_attribute a ON a.attrelid = con.conrelid AND a.attnum = ord.attnum AND NOT a.attisdropped \ ) conkey ON true \ LEFT JOIN LATERAL ( \ SELECT array_agg(a.attname::text ORDER BY ord.ord) AS attnames \ FROM unnest(con.confkey) WITH ORDINALITY AS ord(attnum, ord) \ JOIN pg_catalog.pg_attribute a ON a.attrelid = con.confrelid AND a.attnum = ord.attnum AND NOT a.attisdropped \ ) confkey ON true \ WHERE n.nspname = $1 AND c.relname = $2 AND c.relkind IN ('r','p','f') \ ORDER BY con.conname" } /// Expand a `pg_constraint.contype` letter to a human-readable label. fn postgres_constraint_type_label(contype: &str) -> String { match contype { "c" => "CHECK".to_string(), "f" => "FOREIGN KEY".to_string(), "p" => "PRIMARY KEY".to_string(), "u" => "UNIQUE".to_string(), "t" => "CONSTRAINT TRIGGER".to_string(), "x" => "EXCLUDE".to_string(), "n" => "NOT NULL".to_string(), _ => contype.to_string(), } } /// Normalize a `confupdtype`/`confdeltype` letter to an `information_schema` /// style referential-action label. fn postgres_fk_action_label(action: Option) -> Option { action .as_deref() .and_then(|value| match value { "a" => Some("NO ACTION"), "r" => Some("RESTRICT"), "c" => Some("CASCADE"), "n" => Some("SET NULL"), "d" => Some("SET DEFAULT"), _ => None, }) .map(str::to_string) } /// Normalize a `confmatchtype` letter to a match-type label. fn postgres_constraint_match_type(match_type: Option) -> Option { match_type .as_deref() .and_then(|value| match value { "f" => Some("FULL"), "p" => Some("PARTIAL"), "s" | "u" => Some("SIMPLE"), _ => None, }) .map(str::to_string) } fn postgres_table_partition_relation_sql() -> &'static str { "SELECT c.relkind::text, \ EXISTS ( \ SELECT 1 FROM pg_catalog.pg_inherits i \ WHERE i.inhrelid = c.oid \ AND (SELECT parent.relkind FROM pg_catalog.pg_class parent WHERE parent.oid = i.inhparent) = 'p' \ ) AS is_partition \ FROM pg_catalog.pg_class c \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ WHERE n.nspname = $1 AND c.relname = $2 AND c.relkind IN ('r','p','f') \ LIMIT 1" } fn postgres_table_partition_info_sql() -> &'static str { "SELECT CASE WHEN c.relispartition THEN pn.nspname ELSE NULL END AS parent_schema, \ CASE WHEN c.relispartition THEN pc.relname ELSE NULL END AS parent_table, \ CASE WHEN c.relispartition THEN pg_catalog.pg_get_expr(c.relpartbound, c.oid, true) ELSE NULL END AS partition_bound, \ CASE WHEN c.relkind = 'p' THEN pg_catalog.pg_get_partkeydef(c.oid) ELSE NULL END AS partition_key, \ CASE WHEN c.relkind = 'f' THEN fs.srvname ELSE NULL END AS foreign_server, \ CASE WHEN c.relkind = 'f' THEN ft.ftoptions ELSE NULL END AS foreign_options \ FROM pg_catalog.pg_class c \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ LEFT JOIN pg_catalog.pg_inherits i ON i.inhrelid = c.oid \ LEFT JOIN pg_catalog.pg_class pc ON pc.oid = i.inhparent \ LEFT JOIN pg_catalog.pg_namespace pn ON pn.oid = pc.relnamespace \ LEFT JOIN pg_catalog.pg_foreign_table ft ON ft.ftrelid = c.oid \ LEFT JOIN pg_catalog.pg_foreign_server fs ON fs.oid = ft.ftserver \ WHERE n.nspname = $1 AND c.relname = $2 AND c.relkind IN ('r','p','f') \ ORDER BY i.inhseqno NULLS LAST \ LIMIT 1" } /// 9.x 没有分区父表(也没有 relpartbound/pg_get_partkeydef),只保留 /// 外部表字段,其余列给 NULL,列序与主查询保持一致。 /// /// Pre-10 servers have no partition parent (nor the 10+ columns/functions); /// only the foreign-table fields are kept, same column order as the primary. fn postgres_table_partition_info_compat_sql() -> &'static str { "SELECT NULL::text AS parent_schema, NULL::text AS parent_table, NULL::text AS partition_bound, \ NULL::text AS partition_key, \ CASE WHEN c.relkind = 'f' THEN fs.srvname ELSE NULL END AS foreign_server, \ CASE WHEN c.relkind = 'f' THEN ft.ftoptions ELSE NULL END AS foreign_options \ FROM pg_catalog.pg_class c \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ LEFT JOIN pg_catalog.pg_foreign_table ft ON ft.ftrelid = c.oid \ LEFT JOIN pg_catalog.pg_foreign_server fs ON fs.oid = ft.ftserver \ WHERE n.nspname = $1 AND c.relname = $2 AND c.relkind IN ('r','p','f') \ LIMIT 1" } fn postgres_table_partition_local_objects_sql() -> &'static str { "SELECT 'constraint'::text AS object_kind, con.conname AS object_name, con.contype::text AS object_type \ FROM pg_catalog.pg_constraint con \ JOIN pg_catalog.pg_class c ON c.oid = con.conrelid \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ WHERE n.nspname = $1 AND c.relname = $2 AND con.contype IN ('p','f') \ AND COALESCE(NULLIF(pg_catalog.row_to_json(con)->>'conparentid', '')::oid, 0) = 0 \ UNION ALL \ SELECT 'check'::text AS object_kind, con.conname AS object_name, NULL::text AS object_type \ FROM pg_catalog.pg_constraint con \ JOIN pg_catalog.pg_class c ON c.oid = con.conrelid \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ WHERE n.nspname = $1 AND c.relname = $2 AND con.contype = 'c' AND con.conislocal \ UNION ALL \ SELECT 'index'::text AS object_kind, idx.relname AS object_name, NULL::text AS object_type \ FROM pg_catalog.pg_index ix \ JOIN pg_catalog.pg_class c ON c.oid = ix.indrelid \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ JOIN pg_catalog.pg_class idx ON idx.oid = ix.indexrelid \ WHERE n.nspname = $1 AND c.relname = $2 \ AND NOT EXISTS (SELECT 1 FROM pg_catalog.pg_inherits i WHERE i.inhrelid = idx.oid) \ UNION ALL \ SELECT 'column_default'::text AS object_kind, a.attname AS object_name, \ CASE WHEN ad.oid IS NOT NULL THEN 'overridden' ELSE 'dropped' END AS object_type \ FROM pg_catalog.pg_class c \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ JOIN pg_catalog.pg_attribute a ON a.attrelid = c.oid AND a.attnum > 0 AND NOT a.attisdropped \ JOIN pg_catalog.pg_inherits i ON i.inhrelid = c.oid \ JOIN pg_catalog.pg_attribute pa \ ON pa.attrelid = i.inhparent AND pa.attname = a.attname AND pa.attnum > 0 AND NOT pa.attisdropped \ LEFT JOIN pg_catalog.pg_attrdef ad ON ad.adrelid = a.attrelid AND ad.adnum = a.attnum \ LEFT JOIN pg_catalog.pg_attrdef pad ON pad.adrelid = pa.attrelid AND pad.adnum = pa.attnum \ WHERE n.nspname = $1 AND c.relname = $2 \ AND ( \ (ad.oid IS NOT NULL AND (pad.oid IS NULL \ OR pg_catalog.pg_get_expr(ad.adbin, ad.adrelid) \ IS DISTINCT FROM pg_catalog.pg_get_expr(pad.adbin, pad.adrelid))) \ OR (ad.oid IS NULL AND pad.oid IS NOT NULL) \ ) \ ORDER BY object_kind, object_name" } fn postgres_table_comment_sql() -> &'static str { "SELECT obj_description(c.oid) AS table_comment \ FROM pg_catalog.pg_class c \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ WHERE n.nspname = $1 AND c.relname = $2 AND c.relkind IN ('r','m','f','p') \ LIMIT 1" } fn postgres_tables_sql(limit: Option, offset: i64) -> String { postgres_tables_sql_with_kind(limit, offset, false) } fn postgres_table_objects_sql(limit: Option, offset: i64) -> String { postgres_tables_sql_with_kind(limit, offset, true) } fn postgres_tables_sql_with_kind(limit: Option, offset: i64, table_objects_only: bool) -> String { // PostgreSQL-compatible servers do not agree on the inferred wire types // or accepted expression grammar for LIMIT/OFFSET parameters. These values // originate as usize and are converted to non-negative i64 literals. // Omitting an explicit ESCAPE keeps compatibility with servers that expose // only two-argument ILIKE; bound patterns use the default backslash escape. let pagination = match limit { Some(limit) => format!("LIMIT {limit} OFFSET {offset}"), None => format!("OFFSET {offset}"), }; let relation_kinds = if table_objects_only { "'r','f','p'" } else { "'r','v','m','f','p'" }; format!( "SELECT c.relname AS table_name, \ CASE c.relkind WHEN 'r' THEN 'BASE TABLE' WHEN 'v' THEN 'VIEW' \ WHEN 'm' THEN 'MATERIALIZED_VIEW' WHEN 'f' THEN 'FOREIGN TABLE' \ WHEN 'p' THEN 'BASE TABLE' END AS table_type, \ obj_description(c.oid) AS table_comment, \ CASE WHEN pc.relkind = 'p' THEN pn.nspname ELSE NULL END AS parent_schema, \ CASE WHEN pc.relkind = 'p' THEN pc.relname ELSE NULL END AS parent_name \ FROM pg_catalog.pg_class c \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ LEFT JOIN pg_catalog.pg_inherits i ON i.inhrelid = c.oid \ LEFT JOIN pg_catalog.pg_class pc ON pc.oid = i.inhparent \ LEFT JOIN pg_catalog.pg_namespace pn ON pn.oid = pc.relnamespace \ WHERE n.nspname = $1 AND c.relkind IN ({relation_kinds}) \ AND ($2 = '%%' OR c.relname ILIKE $2 OR ($3 <> '' AND c.relname ILIKE $3)) \ ORDER BY CASE WHEN pc.relkind = 'p' THEN 1 ELSE 0 END, c.relname \ {pagination}" ) } fn like_contains_pattern(value: &str) -> String { if value.is_empty() { return "%%".to_string(); } let mut pattern = String::with_capacity(value.len() + 2); pattern.push('%'); for ch in value.chars() { if ch == '\\' || ch == '%' || ch == '_' { pattern.push('\\'); } pattern.push(ch); } pattern.push('%'); pattern } fn like_fuzzy_pattern(value: &str) -> String { crate::sql::fuzzy_like_pattern_with_escape(value, |value| { let mut escaped = String::with_capacity(value.len() + 1); for ch in value.chars() { if ch == '\\' || ch == '%' || ch == '_' { escaped.push('\\'); } escaped.push(ch); } escaped }) } fn list_object_relations_sql(include_timestamps: bool) -> &'static str { if include_timestamps { return "SELECT c.relname AS object_name, \ CASE c.relkind \ WHEN 'v' THEN 'VIEW' \ WHEN 'm' THEN 'MATERIALIZED_VIEW' \ WHEN 'S' THEN 'SEQUENCE' \ ELSE 'TABLE' \ END AS object_type, \ obj_description(c.oid) AS object_comment, \ stat.creation::text AS created_at, \ COALESCE( \ CASE WHEN current_setting('track_commit_timestamp', true) = 'on' \ THEN pg_xact_commit_timestamp(c.xmin)::text END, \ stat.modification::text \ ) AS updated_at, \ CASE WHEN pc.relkind = 'p' THEN pn.nspname ELSE NULL END AS parent_schema, \ CASE WHEN pc.relkind = 'p' THEN pc.relname ELSE NULL END AS parent_name, \ NULL::text AS signature, \ CASE c.relkind WHEN 'v' THEN 1 WHEN 'm' THEN 1 WHEN 'S' THEN 4 ELSE 0 END AS sort_order \ FROM pg_catalog.pg_class c \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ LEFT JOIN pg_catalog.pg_inherits i ON i.inhrelid = c.oid \ LEFT JOIN pg_catalog.pg_class pc ON pc.oid = i.inhparent \ LEFT JOIN pg_catalog.pg_namespace pn ON pn.oid = pc.relnamespace \ LEFT JOIN LATERAL pg_stat_file( \ CASE WHEN c.relkind IN ('r','m','f','p') THEN pg_relation_filepath(c.oid) END, true \ ) stat ON true \ WHERE n.nspname = $1 AND c.relkind IN ('r','v','m','f','p','S')"; } "SELECT c.relname AS object_name, \ CASE c.relkind \ WHEN 'v' THEN 'VIEW' \ WHEN 'm' THEN 'MATERIALIZED_VIEW' \ WHEN 'S' THEN 'SEQUENCE' \ ELSE 'TABLE' \ END AS object_type, \ obj_description(c.oid) AS object_comment, \ NULL::text AS created_at, \ NULL::text AS updated_at, \ CASE WHEN pc.relkind = 'p' THEN pn.nspname ELSE NULL END AS parent_schema, \ CASE WHEN pc.relkind = 'p' THEN pc.relname ELSE NULL END AS parent_name, \ NULL::text AS signature, \ CASE c.relkind WHEN 'v' THEN 1 WHEN 'm' THEN 1 WHEN 'S' THEN 4 ELSE 0 END AS sort_order \ FROM pg_catalog.pg_class c \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ LEFT JOIN pg_catalog.pg_inherits i ON i.inhrelid = c.oid \ LEFT JOIN pg_catalog.pg_class pc ON pc.oid = i.inhparent \ LEFT JOIN pg_catalog.pg_namespace pn ON pn.oid = pc.relnamespace \ WHERE n.nspname = $1 AND c.relkind IN ('r','v','m','f','p','S')" } fn list_object_routines_sql(include_timestamps: bool, has_proc_prokind: bool, has_proc_prosp: bool) -> &'static str { if has_proc_prokind && has_proc_prosp { if include_timestamps { return "SELECT p.proname AS object_name, \ CASE WHEN p.prokind = 'p' OR p.prosp THEN 'PROCEDURE' ELSE 'FUNCTION' END AS object_type, \ obj_description(p.oid) AS object_comment, \ NULL::text AS created_at, \ CASE WHEN current_setting('track_commit_timestamp', true) = 'on' \ THEN pg_xact_commit_timestamp(p.xmin)::text END AS updated_at, \ NULL::text AS parent_schema, \ NULL::text AS parent_name, \ pg_get_function_identity_arguments(p.oid) AS signature, \ CASE WHEN p.prokind = 'p' OR p.prosp THEN 2 ELSE 3 END AS sort_order \ FROM pg_catalog.pg_proc p \ JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \ WHERE n.nspname = $1 AND (p.prokind IN ('p','f') OR p.prosp)"; } return "SELECT p.proname AS object_name, \ CASE WHEN p.prokind = 'p' OR p.prosp THEN 'PROCEDURE' ELSE 'FUNCTION' END AS object_type, \ obj_description(p.oid) AS object_comment, \ NULL::text AS created_at, \ NULL::text AS updated_at, \ NULL::text AS parent_schema, \ NULL::text AS parent_name, \ pg_get_function_identity_arguments(p.oid) AS signature, \ CASE WHEN p.prokind = 'p' OR p.prosp THEN 2 ELSE 3 END AS sort_order \ FROM pg_catalog.pg_proc p \ JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \ WHERE n.nspname = $1 AND (p.prokind IN ('p','f') OR p.prosp)"; } if has_proc_prokind { if include_timestamps { return "SELECT p.proname AS object_name, \ CASE p.prokind WHEN 'p' THEN 'PROCEDURE' ELSE 'FUNCTION' END AS object_type, \ obj_description(p.oid) AS object_comment, \ NULL::text AS created_at, \ CASE WHEN current_setting('track_commit_timestamp', true) = 'on' \ THEN pg_xact_commit_timestamp(p.xmin)::text END AS updated_at, \ NULL::text AS parent_schema, \ NULL::text AS parent_name, \ pg_get_function_identity_arguments(p.oid) AS signature, \ CASE p.prokind WHEN 'p' THEN 2 ELSE 3 END AS sort_order \ FROM pg_catalog.pg_proc p \ JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \ WHERE n.nspname = $1 AND p.prokind IN ('p','f')"; } return "SELECT p.proname AS object_name, \ CASE p.prokind WHEN 'p' THEN 'PROCEDURE' ELSE 'FUNCTION' END AS object_type, \ obj_description(p.oid) AS object_comment, \ NULL::text AS created_at, \ NULL::text AS updated_at, \ NULL::text AS parent_schema, \ NULL::text AS parent_name, \ pg_get_function_identity_arguments(p.oid) AS signature, \ CASE p.prokind WHEN 'p' THEN 2 ELSE 3 END AS sort_order \ FROM pg_catalog.pg_proc p \ JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \ WHERE n.nspname = $1 AND p.prokind IN ('p','f')"; } if has_proc_prosp { if include_timestamps { return "SELECT p.proname AS object_name, \ CASE WHEN p.prosp THEN 'PROCEDURE' ELSE 'FUNCTION' END AS object_type, \ obj_description(p.oid) AS object_comment, \ NULL::text AS created_at, \ CASE WHEN current_setting('track_commit_timestamp', true) = 'on' \ THEN pg_xact_commit_timestamp(p.xmin)::text END AS updated_at, \ NULL::text AS parent_schema, \ NULL::text AS parent_name, \ pg_get_function_identity_arguments(p.oid) AS signature, \ CASE WHEN p.prosp THEN 2 ELSE 3 END AS sort_order \ FROM pg_catalog.pg_proc p \ JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \ WHERE n.nspname = $1 AND NOT p.proisagg AND NOT p.proiswindow"; } return "SELECT p.proname AS object_name, \ CASE WHEN p.prosp THEN 'PROCEDURE' ELSE 'FUNCTION' END AS object_type, \ obj_description(p.oid) AS object_comment, \ NULL::text AS created_at, \ NULL::text AS updated_at, \ NULL::text AS parent_schema, \ NULL::text AS parent_name, \ pg_get_function_identity_arguments(p.oid) AS signature, \ CASE WHEN p.prosp THEN 2 ELSE 3 END AS sort_order \ FROM pg_catalog.pg_proc p \ JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \ WHERE n.nspname = $1 AND NOT p.proisagg AND NOT p.proiswindow"; } if include_timestamps { return "SELECT p.proname AS object_name, \ 'FUNCTION' AS object_type, \ obj_description(p.oid) AS object_comment, \ NULL::text AS created_at, \ CASE WHEN current_setting('track_commit_timestamp', true) = 'on' \ THEN pg_xact_commit_timestamp(p.xmin)::text END AS updated_at, \ NULL::text AS parent_schema, \ NULL::text AS parent_name, \ pg_get_function_identity_arguments(p.oid) AS signature, \ 3 AS sort_order \ FROM pg_catalog.pg_proc p \ JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \ WHERE n.nspname = $1 AND NOT p.proisagg AND NOT p.proiswindow"; } "SELECT p.proname AS object_name, \ 'FUNCTION' AS object_type, \ obj_description(p.oid) AS object_comment, \ NULL::text AS created_at, \ NULL::text AS updated_at, \ NULL::text AS parent_schema, \ NULL::text AS parent_name, \ pg_get_function_identity_arguments(p.oid) AS signature, \ 3 AS sort_order \ FROM pg_catalog.pg_proc p \ JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \ WHERE n.nspname = $1 AND NOT p.proisagg AND NOT p.proiswindow" } fn redshift_routine_objects_sql() -> &'static str { "SELECT function_name AS object_name, \ CASE function_type \ WHEN 'STORED PROCEDURE' THEN 'PROCEDURE' \ ELSE 'FUNCTION' \ END AS object_type, \ NULL::varchar AS object_comment, \ NULL::varchar AS created_at, \ NULL::varchar AS updated_at, \ NULL::varchar AS parent_schema, \ NULL::varchar AS parent_name, \ argument_type AS signature, \ CASE function_type WHEN 'STORED PROCEDURE' THEN 2 ELSE 3 END AS sort_order \ FROM svv_redshift_functions \ WHERE database_name = current_database() \ AND schema_name = $1 \ AND function_type IN ('STORED PROCEDURE', 'REGULAR FUNCTION') \ ORDER BY sort_order, object_name" } /// SQL for listing user-defined types in one schema. /// /// Only explicitly created types are returned: base types (b), standalone /// composite types (c), domains (d), enums (e), ranges (r) and multiranges (m). /// Relation auto-generated row types (table/view/materialized view/foreign /// table/partitioned table) are excluded via `typrelid = 0 OR relkind = 'c'`, /// and array companion types are excluded via `typelem = 0`. The type branch /// has no usable timestamp/stat columns, so created_at/updated_at are always /// NULL. Column order must match the relation and routine branches. fn list_object_custom_types_sql() -> &'static str { "SELECT t.typname AS object_name, \ 'TYPE' AS object_type, \ d.description AS object_comment, \ NULL::text AS created_at, \ NULL::text AS updated_at, \ NULL::text AS parent_schema, \ NULL::text AS parent_name, \ (t.typtype::text || ':' || CASE \ WHEN t.typtype = 'c' THEN CASE WHEN EXISTS ( \ SELECT 1 FROM pg_catalog.pg_attribute a \ WHERE a.attrelid = t.typrelid AND a.attnum > 0 AND NOT a.attisdropped \ ) THEN '1' ELSE '0' END \ WHEN t.typtype = 'e' THEN CASE WHEN EXISTS ( \ SELECT 1 FROM pg_catalog.pg_enum e WHERE e.enumtypid = t.oid \ ) THEN '1' ELSE '0' END \ ELSE '0' \ END) AS signature, \ 5 AS sort_order \ FROM pg_catalog.pg_type t \ JOIN pg_catalog.pg_namespace n ON n.oid = t.typnamespace \ LEFT JOIN pg_catalog.pg_class c ON c.oid = t.typrelid \ LEFT JOIN pg_catalog.pg_description d \ ON d.objoid = t.oid \ AND d.classoid = 'pg_catalog.pg_type'::regclass \ AND d.objsubid = 0 \ WHERE n.nspname = $1 \ AND t.typtype IN ('b', 'c', 'd', 'e', 'r', 'm') \ AND t.typisdefined \ AND t.typelem = 0 \ AND (t.typrelid = 0 OR c.relkind = 'c') \ AND n.nspname <> 'pg_catalog' \ AND n.nspname <> 'information_schema' \ AND n.nspname NOT LIKE 'pg_toast%' \ AND n.nspname NOT LIKE 'pg_temp%'" } fn is_postgres_system_schema(schema: &str) -> bool { schema == "pg_catalog" || schema == "information_schema" || schema.starts_with("pg_toast") || schema.starts_with("pg_temp") } fn list_objects_sql( include_timestamps: bool, has_proc_prokind: bool, has_proc_prosp: bool, has_function_identity_arguments: bool, include_relations: bool, include_routines: bool, include_custom_types: bool, ) -> String { let mut parts: Vec = Vec::new(); if include_relations { parts.push(list_object_relations_sql(include_timestamps).to_string()); } if include_routines { parts.push(list_object_routines_sql(include_timestamps, has_proc_prokind, has_proc_prosp).to_string()); } if include_custom_types { parts.push(list_object_custom_types_sql().to_string()); } let mut sql = parts.join(" UNION ALL "); if !parts.is_empty() { sql = format!("{sql} ORDER BY sort_order, object_name"); } if has_function_identity_arguments { sql } else { // Redshift and older PostgreSQL-compatible servers may only expose the // older formatter. It includes argument names but still distinguishes // overloads instead of making the whole schema browser unavailable. sql.replace("pg_get_function_identity_arguments(p.oid)", "pg_get_function_arguments(p.oid)") } } fn postgres_has_function_identity_arguments_sql() -> &'static str { "SELECT EXISTS ( \ SELECT 1 \ FROM pg_catalog.pg_proc p \ JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \ WHERE n.nspname = 'pg_catalog' \ AND p.proname = 'pg_get_function_identity_arguments' \ )" } async fn postgres_has_function_identity_arguments(client: &deadpool_postgres::Client) -> Result { let row = postgres_query_one_cached(client, postgres_has_function_identity_arguments_sql(), &[]) .await .map_err(|e| e.to_string())?; Ok(pg_row_try_bool(&row, 0).unwrap_or(false)) } fn postgres_proc_has_prokind_sql() -> &'static str { "SELECT EXISTS ( \ SELECT 1 \ FROM pg_catalog.pg_attribute \ WHERE attrelid = 'pg_catalog.pg_proc'::regclass \ AND attname = 'prokind' \ AND NOT attisdropped \ )" } async fn postgres_proc_has_prokind(client: &deadpool_postgres::Client) -> Result { let row = postgres_query_one_cached(client, postgres_proc_has_prokind_sql(), &[]).await.map_err(|e| e.to_string())?; Ok(pg_row_try_bool(&row, 0).unwrap_or(false)) } fn postgres_trigger_has_tgparentid_sql() -> &'static str { "SELECT EXISTS ( \ SELECT 1 \ FROM pg_catalog.pg_attribute \ WHERE attrelid = 'pg_catalog.pg_trigger'::regclass \ AND attname = 'tgparentid' \ AND NOT attisdropped \ )" } /// PostgreSQL 13 added `pg_trigger.tgparentid`, tracking a partition's /// automatically-cloned copy of its parent's trigger. Servers older than 13 /// lack the column (and never clone triggers onto partitions in the first /// place, so `tgisinternal` alone is sufficient there). async fn postgres_trigger_has_tgparentid(client: &deadpool_postgres::Client) -> Result { let row = postgres_query_one_cached(client, postgres_trigger_has_tgparentid_sql(), &[]) .await .map_err(|e| e.to_string())?; Ok(pg_row_try_bool(&row, 0).unwrap_or(false)) } fn postgres_proc_has_prosp_sql() -> &'static str { "SELECT EXISTS ( \ SELECT 1 \ FROM pg_catalog.pg_attribute \ WHERE attrelid = 'pg_catalog.pg_proc'::regclass \ AND attname = 'prosp' \ AND NOT attisdropped \ )" } async fn postgres_proc_has_prosp(client: &deadpool_postgres::Client) -> Result { let row = postgres_query_one_cached(client, postgres_proc_has_prosp_sql(), &[]).await.map_err(|e| e.to_string())?; Ok(pg_row_try_bool(&row, 0).unwrap_or(false)) } async fn list_objects_rows( client: &deadpool_postgres::Client, schema: &str, include_timestamps: bool, has_proc_prokind: bool, has_proc_prosp: bool, has_function_identity_arguments: bool, include_relations: bool, include_routines: bool, include_custom_types: bool, ) -> Result, String> { let sql = list_objects_sql( include_timestamps, has_proc_prokind, has_proc_prosp, has_function_identity_arguments, include_relations, include_routines, include_custom_types, ); if sql.is_empty() { // No branch was selected (e.g. an object_types filter the catalog // cannot serve); skip the round-trip instead of executing empty SQL. return Ok(Vec::new()); } postgres_query_cached(client, &sql, &[&schema]).await.map_err(|e| e.to_string()) } pub async fn list_objects( pool: &Pool, schema: &str, include_relations: bool, include_routines: bool, include_custom_types: bool, ) -> Result, String> { // System schemas may be visible in the schema tree, but their catalog // types are implementation details and are never custom types. let include_custom_types = include_custom_types && !is_postgres_system_schema(schema); let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; // Routine catalog probes are only needed when the routine branch runs. // Skipping them for relation/type-only requests avoids three extra // pg_proc/pg_attribute round-trips and keeps compatible catalogs that lack // prokind/prosp from breaking a plain type listing. let (has_proc_prokind, has_proc_prosp, has_function_identity_arguments) = if include_routines { let has_proc_prokind = postgres_proc_has_prokind(&client).await?; // Some GaussDB-compatible catalogs expose prosp alongside, or instead of, // PostgreSQL 11's prokind. Treat prosp as an extra procedure signal. let has_proc_prosp = postgres_proc_has_prosp(&client).await?; let has_function_identity_arguments = postgres_has_function_identity_arguments(&client).await?; (has_proc_prokind, has_proc_prosp, has_function_identity_arguments) } else { (false, false, false) }; let rows = match list_objects_rows( &client, schema, true, has_proc_prokind, has_proc_prosp, has_function_identity_arguments, include_relations, include_routines, include_custom_types, ) .await { Ok(rows) => rows, Err(primary_error) => { log::debug!("[postgres][list_objects:timestamp-fallback] primary_error={}", primary_error); match list_objects_rows( &client, schema, false, has_proc_prokind, has_proc_prosp, has_function_identity_arguments, include_relations, include_routines, include_custom_types, ) .await { Ok(rows) => rows, Err(fallback_error) => { return Err(format!("{primary_error}; timestamp fallback failed: {fallback_error}")); } } } }; Ok(object_rows_to_infos(&rows, schema)) } pub async fn list_redshift_objects( pool: &Pool, schema: &str, include_relations: bool, include_routines: bool, ) -> Result, String> { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let mut rows = Vec::new(); if include_relations { // Redshift does not support PostgreSQL's generic file, object-location, // or transaction-ID helpers, so skip the timestamp variant entirely. rows = list_objects_rows(&client, schema, false, false, false, false, true, false, false).await?; } if include_routines { let routine_rows = client .query_typed(redshift_routine_objects_sql(), &[(&schema, Type::VARCHAR)]) .await .map_err(|error| format!("Redshift routine metadata query failed: {error}"))?; rows.extend(routine_rows); } Ok(object_rows_to_infos(&rows, schema)) } fn object_rows_to_infos(rows: &[Row], schema: &str) -> Vec { rows.iter() .map(|row| { let object_type = pg_row_try_string(row, 1); let raw_signature = row.try_get::<_, Option>(7).ok().flatten(); let (signature, custom_type_kind, has_members) = if object_type == "TYPE" { let (kind, has_members) = custom_type_list_metadata(raw_signature.as_deref()); (None, kind, has_members) } else { (raw_signature, None, None) }; ObjectInfo { name: pg_row_try_string(row, 0), object_type, schema: Some(schema.to_string()), valid: None, signature, custom_type_kind, has_members, comment: row.try_get::<_, Option>(2).ok().flatten().filter(|s| !s.is_empty()), created_at: row.try_get::<_, Option>(3).ok().flatten().filter(|s| !s.is_empty()), updated_at: row.try_get::<_, Option>(4).ok().flatten().filter(|s| !s.is_empty()), parent_schema: row.try_get::<_, Option>(5).ok().flatten().filter(|s| !s.is_empty()), parent_name: row.try_get::<_, Option>(6).ok().flatten().filter(|s| !s.is_empty()), trigger: None, xugu_type_members_expandable: None, } }) .collect() } fn custom_type_list_metadata(value: Option<&str>) -> (Option, Option) { let Some((kind_code, has_members)) = value.and_then(|value| value.split_once(':')) else { return (None, None); }; let kind = custom_type_kind_for(kind_code).map(|kind| kind.as_str().to_string()); let has_members = match has_members { "1" => Some(true), "0" => Some(false), _ => None, }; (kind, has_members) } /// General `pg_type` row for a user-defined type located by schema + name. /// /// The query never interpolates user input: schema and name are bound as /// parameters. Function OIDs are joined to `pg_proc` so the DDL generator and /// the properties panel get human-readable names without extra round-trips. struct CustomTypeGeneralInfo { oid: u32, typtype: String, typisdefined: bool, typbasetype: u32, typnotnull: bool, typrelid: u32, typelem: u32, typcollation: u32, typdefaultbin: Option, typdefault: Option, typlen: i16, typbyval: bool, typalign: String, typstorage: String, typtypmod: i32, input_function: Option, output_function: Option, receive_function: Option, send_function: Option, analyze_function: Option, comment: Option, relkind: Option, collation: Option, } fn custom_type_general_info_sql() -> &'static str { "SELECT t.oid, t.typtype::text, t.typisdefined, \ t.typbasetype, t.typnotnull, t.typrelid, t.typelem, \ t.typcollation, t.typdefaultbin, t.typdefault, \ t.typlen, t.typbyval, t.typalign::text, t.typstorage::text, t.typtypmod, \ pi.proname AS input_fn, po.proname AS output_fn, \ pr.proname AS receive_fn, ps.proname AS send_fn, pa.proname AS analyze_fn, \ d.description, \ CASE WHEN t.typrelid != 0 THEN \ (SELECT c.relkind::text FROM pg_catalog.pg_class c WHERE c.oid = t.typrelid) \ END AS relkind, \ CASE WHEN cl.oid IS NULL THEN NULL ELSE quote_ident(ncl.nspname) || '.' || quote_ident(cl.collname) END \ FROM pg_catalog.pg_type t \ JOIN pg_catalog.pg_namespace n ON n.oid = t.typnamespace \ LEFT JOIN pg_catalog.pg_description d \ ON d.objoid = t.oid AND d.classoid = 'pg_catalog.pg_type'::regclass AND d.objsubid = 0 \ LEFT JOIN pg_catalog.pg_proc pi ON pi.oid = t.typinput \ LEFT JOIN pg_catalog.pg_proc po ON po.oid = t.typoutput \ LEFT JOIN pg_catalog.pg_proc pr ON pr.oid = t.typreceive \ LEFT JOIN pg_catalog.pg_proc ps ON ps.oid = t.typsend \ LEFT JOIN pg_catalog.pg_proc pa ON pa.oid = t.typanalyze \ LEFT JOIN pg_catalog.pg_collation cl ON cl.oid = t.typcollation \ LEFT JOIN pg_catalog.pg_namespace ncl ON ncl.oid = cl.collnamespace \ WHERE n.nspname = $1 AND t.typname = $2" } async fn custom_type_general_info( client: &deadpool_postgres::Client, schema: &str, name: &str, ) -> Result { let rows = postgres_query_cached(client, custom_type_general_info_sql(), &[&schema, &name]) .await .map_err(|e| format!("failed to locate custom type {schema}.{name}: {e}"))?; let row = rows.first().ok_or_else(|| format!("custom type {schema}.{name} does not exist"))?; Ok(CustomTypeGeneralInfo { oid: row.try_get::<_, u32>(0).unwrap_or(0), typtype: pg_row_try_string(row, 1), typisdefined: pg_row_try_bool(row, 2).unwrap_or(false), typbasetype: row.try_get::<_, u32>(3).unwrap_or(0), typnotnull: pg_row_try_bool(row, 4).unwrap_or(false), typrelid: row.try_get::<_, u32>(5).unwrap_or(0), typelem: row.try_get::<_, u32>(6).unwrap_or(0), typcollation: row.try_get::<_, u32>(7).unwrap_or(0), typdefaultbin: row.try_get::<_, Option>(8).ok().flatten(), typdefault: row.try_get::<_, Option>(9).ok().flatten(), typlen: row.try_get::<_, i16>(10).unwrap_or(-1), typbyval: pg_row_try_bool(row, 11).unwrap_or(false), typalign: pg_row_try_string(row, 12), typstorage: pg_row_try_string(row, 13), typtypmod: row.try_get::<_, i32>(14).unwrap_or(-1), input_function: row.try_get::<_, Option>(15).ok().flatten(), output_function: row.try_get::<_, Option>(16).ok().flatten(), receive_function: row.try_get::<_, Option>(17).ok().flatten(), send_function: row.try_get::<_, Option>(18).ok().flatten(), analyze_function: row.try_get::<_, Option>(19).ok().flatten(), comment: row.try_get::<_, Option>(20).ok().flatten(), relkind: row.try_get::<_, Option>(21).ok().flatten(), collation: row.try_get::<_, Option>(22).ok().flatten(), }) } async fn custom_type_rendered_domain_default( client: &deadpool_postgres::Client, oid: u32, ) -> Result, String> { let rows = postgres_query_cached( client, "SELECT pg_catalog.pg_get_expr(t.typdefaultbin, 0) \ FROM pg_catalog.pg_type t WHERE t.oid = $1", &[&oid], ) .await .map_err(|error| format!("failed to render domain default: {error}"))?; Ok(rows .first() .and_then(|row| row.try_get::<_, Option>(0).ok().flatten()) .filter(|value| !value.is_empty())) } fn domain_default_from_render_result(result: Result, String>) -> (Option, Option) { match result { Ok(Some(value)) if !value.is_empty() => (Some(value), None), Ok(_) => (None, Some("default value could not be rendered; the generated DDL is incomplete".to_string())), Err(error) => { (None, Some(format!("default value could not be rendered; the generated DDL is incomplete: {error}"))) } } } fn custom_type_kind_for(typtype: &str) -> Option { match typtype { "b" => Some(CustomTypeKind::Base), "c" => Some(CustomTypeKind::Composite), "d" => Some(CustomTypeKind::Domain), "e" => Some(CustomTypeKind::Enum), "r" => Some(CustomTypeKind::Range), "m" => Some(CustomTypeKind::Multirange), _ => None, } } async fn custom_type_enum_members( client: &deadpool_postgres::Client, oid: u32, ) -> Result, String> { let rows = postgres_query_cached( client, "SELECT e.enumlabel, e.enumsortorder \ FROM pg_catalog.pg_enum e \ WHERE e.enumtypid = $1 \ ORDER BY e.enumsortorder", &[&oid], ) .await .map_err(|e| format!("failed to read enum values: {e}"))?; Ok(rows .iter() .enumerate() .map(|(index, row)| CustomTypeMember { name: String::new(), data_type: String::new(), // enumsortorder is float4: ALTER TYPE ... ADD VALUE BEFORE/AFTER can // produce fractional values (1.5). Use the position after ORDER BY // instead so ordinals stay unique and stable for UI keys. ordinal: index as i32 + 1, nullable: None, default: None, comment: None, enum_value: Some(pg_row_try_string(row, 0)), }) .collect()) } async fn custom_type_composite_members( client: &deadpool_postgres::Client, typrelid: u32, ) -> Result, String> { let data_type = postgres_qualified_format_type_expression("at", "atn", "elem", "elem_n", "a.atttypid", "a.atttypmod"); let rows = postgres_query_cached( client, &format!( "SELECT a.attname, {data_type} AS data_type, \ a.attnum, NOT a.attnotnull AS nullable, a.atthasdef, \ pg_get_expr(ad.adbin, ad.adrelid) AS default_expr, \ col_description($1, a.attnum) AS comment \ FROM pg_catalog.pg_attribute a \ JOIN pg_catalog.pg_type at ON at.oid = a.atttypid \ JOIN pg_catalog.pg_namespace atn ON atn.oid = at.typnamespace \ LEFT JOIN pg_catalog.pg_type elem ON elem.oid = at.typelem \ LEFT JOIN pg_catalog.pg_namespace elem_n ON elem_n.oid = elem.typnamespace \ LEFT JOIN pg_catalog.pg_attrdef ad \ ON ad.adrelid = a.attrelid AND ad.adnum = a.attnum \ WHERE a.attrelid = $1 AND a.attnum > 0 AND NOT a.attisdropped \ ORDER BY a.attnum" ), &[&typrelid], ) .await .map_err(|e| format!("failed to read composite fields: {e}"))?; let mut members = Vec::with_capacity(rows.len()); for row in rows { let name = row.try_get::<_, String>(0).map_err(|error| format!("failed to decode composite field name: {error}"))?; let data_type = row.try_get::<_, String>(1).map_err(|error| format!("failed to decode composite field type: {error}"))?; let ordinal = row.try_get::<_, i16>(2).map_err(|error| format!("failed to decode composite field ordinal: {error}"))?; if name.is_empty() || data_type.is_empty() { return Err("composite field name or type is empty".to_string()); } let has_default = pg_row_try_bool(&row, 4).unwrap_or(false); members.push(CustomTypeMember { name, data_type, ordinal: ordinal as i32, nullable: pg_row_try_bool(&row, 3), default: if has_default { row.try_get::<_, Option>(5).ok().flatten() } else { None }, comment: row.try_get::<_, Option>(6).ok().flatten(), enum_value: None, }); } Ok(members) } /// Catalog-qualified type rendering for generated DDL. `format_type` omits a /// user type's schema when it is on the current search path, so use it only /// for pg_catalog types where it also preserves typmods. User-defined types /// and their array companions retain an explicit, safely quoted schema. fn postgres_qualified_format_type_expression( type_alias: &str, namespace_alias: &str, element_alias: &str, element_namespace_alias: &str, oid_expression: &str, typmod_expression: &str, ) -> String { format!( "CASE \ WHEN {type_alias}.typelem <> 0 AND {element_namespace_alias}.nspname <> 'pg_catalog' \ THEN quote_ident({element_namespace_alias}.nspname) || '.' || quote_ident({element_alias}.typname) || '[]' \ WHEN {namespace_alias}.nspname <> 'pg_catalog' \ THEN quote_ident({namespace_alias}.nspname) || '.' || quote_ident({type_alias}.typname) \ ELSE format_type({oid_expression}, {typmod_expression}) \ END" ) } /// Domain-only attributes. Runs a category query that some compatible /// catalogs (older openGauss/GaussDB) may reject; on failure the caller /// degrades to missing attributes with a warning instead of failing the whole /// details request. async fn custom_type_domain_attributes( client: &deadpool_postgres::Client, info: &CustomTypeGeneralInfo, properties: &mut CustomTypeProperties, ) -> Vec { let mut warnings = Vec::new(); let base_type_expression = postgres_qualified_format_type_expression("t", "n", "elem", "elem_n", "t.oid", "$2::int4"); let base_type = postgres_query_cached( client, &format!( "SELECT {base_type_expression} \ FROM pg_catalog.pg_type t \ JOIN pg_catalog.pg_namespace n ON n.oid = t.typnamespace \ LEFT JOIN pg_catalog.pg_type elem ON elem.oid = t.typelem \ LEFT JOIN pg_catalog.pg_namespace elem_n ON elem_n.oid = elem.typnamespace \ WHERE t.oid = $1" ), &[&info.typbasetype, &info.typtypmod], ) .await .ok() .and_then(|rows| rows.first().map(|row| pg_row_try_string(row, 0))); if let Some(base) = base_type.filter(|value| !value.is_empty()) { properties.base_type = Some(base); } properties.not_null = Some(info.typnotnull); properties.default = info.typdefault.clone().filter(|value| !value.is_empty()); if properties.default.is_none() && info.typdefaultbin.as_deref().is_some_and(|value| !value.is_empty()) { let (default, warning) = domain_default_from_render_result(custom_type_rendered_domain_default(client, info.oid).await); properties.default = default; if let Some(warning) = warning { warnings.push(warning); } } if info.typcollation != 0 { properties.collation = info.collation.clone(); } match postgres_query_cached( client, "SELECT c.conname, pg_get_constraintdef(c.oid, true) AS definition \ FROM pg_catalog.pg_constraint c \ WHERE c.contypid = $1 \ ORDER BY c.conname", &[&info.oid], ) .await { Ok(rows) => { properties.domain_constraints.clear(); for row in rows { let name = match row.try_get::<_, String>(0) { Ok(value) => value, Err(error) => { warnings.push(format!("domain constraints could not be decoded (name): {error}")); continue; } }; let definition = match row.try_get::<_, String>(1) { Ok(value) => value, Err(error) => { warnings.push(format!("domain constraints could not be decoded ({name}): {error}")); continue; } }; if !definition.is_empty() { properties.domain_constraints.push(CustomTypeDomainConstraint { name, definition }); } } } Err(error) => { warnings.push(format!("domain constraints could not be read: {error}")); } } warnings } /// Range/multirange attributes. Like the domain branch this is a category /// extension: failures degrade to missing attributes with a warning. The /// multirange companion name (PG 13+ `rngmultitypid`) is queried separately /// because openGauss/GaussDB kernels predate that column. async fn custom_type_range_attributes( client: &deadpool_postgres::Client, oid: u32, is_multirange: bool, properties: &mut CustomTypeProperties, ) -> Vec { // pg_range.rngtypid always stores the RANGE oid. A multirange view must // first resolve its owning range through rngmultitypid. let range_oid_clause = if is_multirange { "WHERE r.rngmultitypid = $1" } else { "WHERE r.rngtypid = $1" }; let subtype = postgres_qualified_format_type_expression("st", "stn", "elem", "elem_n", "r.rngsubtype", "NULL::integer"); let rows = match postgres_query_cached( client, &format!( "SELECT {subtype} AS subtype, \ CASE WHEN pcan.oid IS NULL THEN NULL ELSE quote_ident(ncan.nspname) || '.' || quote_ident(pcan.proname) END AS canonical_fn, \ CASE WHEN pdiff.oid IS NULL THEN NULL ELSE quote_ident(ndiff.nspname) || '.' || quote_ident(pdiff.proname) END AS subdiff_fn, \ CASE WHEN opc.oid IS NULL THEN NULL ELSE quote_ident(nopc.nspname) || '.' || quote_ident(opc.opcname) END AS subtype_opclass \ FROM pg_catalog.pg_range r \ JOIN pg_catalog.pg_type st ON st.oid = r.rngsubtype \ JOIN pg_catalog.pg_namespace stn ON stn.oid = st.typnamespace \ LEFT JOIN pg_catalog.pg_type elem ON elem.oid = st.typelem \ LEFT JOIN pg_catalog.pg_namespace elem_n ON elem_n.oid = elem.typnamespace \ LEFT JOIN pg_catalog.pg_proc pcan ON pcan.oid = r.rngcanonical \ LEFT JOIN pg_catalog.pg_proc pdiff ON pdiff.oid = r.rngsubdiff \ LEFT JOIN pg_catalog.pg_opclass opc ON opc.oid = r.rngsubopc \ LEFT JOIN pg_catalog.pg_namespace ncan ON ncan.oid = pcan.pronamespace \ LEFT JOIN pg_catalog.pg_namespace ndiff ON ndiff.oid = pdiff.pronamespace \ LEFT JOIN pg_catalog.pg_namespace nopc ON nopc.oid = opc.opcnamespace \ {range_oid_clause}" ), &[&oid], ) .await { Ok(rows) => rows, Err(error) => return vec![format!("range attributes could not be read: {error}")], }; let mut warnings = Vec::new(); if let Some(row) = rows.first() { match row.try_get::<_, Option>(0) { Ok(value) => properties.range_subtype = value.filter(|value| !value.is_empty()), Err(error) => warnings.push(format!("range subtype could not be read: {error}")), } match row.try_get::<_, Option>(1) { Ok(value) => properties.range_canonical_function = value.filter(|value| !value.is_empty()), Err(error) => warnings.push(format!("range canonical function could not be read: {error}")), } match row.try_get::<_, Option>(2) { Ok(value) => properties.range_subtype_diff_function = value.filter(|value| !value.is_empty()), Err(error) => warnings.push(format!("range subtype diff function could not be read: {error}")), } match row.try_get::<_, Option>(3) { Ok(value) => properties.range_subtype_opclass = value.filter(|value| !value.is_empty()), Err(error) => warnings.push(format!("range subtype opclass could not be read: {error}")), } } else { warnings.push("range attributes returned no rows".to_string()); } // Optional PG 13+ multirange companion; older kernels simply have no column. match postgres_query_cached( client, "SELECT mt.typname \ FROM pg_catalog.pg_range r \ JOIN pg_catalog.pg_type mt ON mt.oid = r.rngmultitypid \ WHERE r.rngtypid = $1", &[&oid], ) .await { Ok(rows) => { properties.range_multirange_name = rows.first().map(|row| pg_row_try_string(row, 0)).filter(|value| !value.is_empty()); } Err(error) => warnings.push(format!("multirange companion could not be read: {error}")), } warnings } /// Generate normalized `CREATE TYPE` text for a user-defined type. /// /// `complete` is only true when the text can be executed standalone in the /// current schema. Multiranges (auto-generated companions) and base types /// (whose I/O functions cannot be reconstructed from catalogs) are always /// marked incomplete with visible warnings. fn build_custom_type_ddl( schema: &str, name: &str, kind: CustomTypeKind, info: &CustomTypeGeneralInfo, members: &[CustomTypeMember], properties: &CustomTypeProperties, warnings: &[String], ) -> CustomTypeDdl { let qualified = format!("{}.{}", pg_quote_ident(schema), pg_quote_ident(name)); match kind { CustomTypeKind::Enum => { let values = members .iter() .map(|member| member.enum_value.as_deref().unwrap_or_default()) .map(pg_quote_literal) .collect::>() .join(", "); CustomTypeDdl { sql: format!("CREATE TYPE {qualified} AS ENUM ({values});"), complete: true, warnings: warnings.to_vec(), } } CustomTypeKind::Composite => { let fields = members .iter() .map(|member| format!("{} {}", pg_quote_ident(&member.name), member.data_type)) .collect::>() .join(",\n "); let mut sql = format!("CREATE TYPE {qualified} AS (\n {fields}\n);"); let comments = members .iter() .filter_map(|member| { member.comment.as_ref().map(|comment| { format!( "COMMENT ON COLUMN {}.{} IS {};", qualified, pg_quote_ident(&member.name), pg_quote_literal(comment) ) }) }) .collect::>() .join("\n"); if !comments.is_empty() { sql = format!("{sql}\n{comments}"); } CustomTypeDdl { sql, complete: true, warnings: warnings.to_vec() } } CustomTypeKind::Domain => { let mut warnings = warnings.to_vec(); // A domain without a resolvable base type or constraints cannot be // rebuilt faithfully; mark it incomplete so the UI never presents // the truncated text as executable source. let mut complete = true; let base = match properties.base_type.as_deref().filter(|v| !v.is_empty()) { Some(base) => base.to_string(), None => { complete = false; warnings.push("base type could not be resolved; the generated DDL is incomplete".to_string()); "unknown".to_string() } }; let mut parts = vec![format!("CREATE DOMAIN {qualified} AS {base}")]; if let Some(collation) = properties.collation.as_deref().filter(|v| !v.is_empty()) { parts.push(format!("COLLATE {}", pg_quote_catalog_identifier(collation))); } if let Some(default) = properties.default.as_deref().filter(|v| !v.is_empty()) { parts.push(format!("DEFAULT {default}")); } if properties.not_null == Some(true) { parts.push("NOT NULL".to_string()); } let incomplete_warning = warnings.iter().any(|warning| { warning.contains("domain constraints") || warning.contains("default value could not be rendered") }); if incomplete_warning { complete = false; } for constraint in &properties.domain_constraints { let body = constraint.definition.trim().trim_start_matches("CHECK").trim(); let constraint_name = if constraint.name.is_empty() { format!("{name}_check") } else { constraint.name.clone() }; parts.push(format!("CONSTRAINT {} CHECK {body}", pg_quote_ident(&constraint_name))); } CustomTypeDdl { sql: format!("{};", parts.join("\n ")), complete, warnings } } CustomTypeKind::Range => { let mut args = Vec::new(); if let Some(subtype) = properties.range_subtype.as_deref().filter(|v| !v.is_empty()) { args.push(format!("subtype = {subtype}")); } if let Some(opclass) = properties.range_subtype_opclass.as_deref().filter(|v| !v.is_empty()) { args.push(format!("subtype_opclass = {}", pg_quote_catalog_identifier(opclass))); } if let Some(value) = properties.range_canonical_function.as_deref().filter(|v| !v.is_empty()) { args.push(format!("canonical = {}", pg_quote_catalog_identifier(value))); } if let Some(value) = properties.range_subtype_diff_function.as_deref().filter(|v| !v.is_empty()) { args.push(format!("subtype_diff = {}", pg_quote_catalog_identifier(value))); } if let Some(value) = properties.range_multirange_name.as_deref().filter(|v| !v.is_empty()) { args.push(format!("multirange_type_name = {}", pg_quote_catalog_identifier(value))); } if properties.range_subtype.as_deref().is_none_or(|value| value.is_empty()) { let mut warnings = warnings.to_vec(); warnings.push("range attributes could not be resolved".to_string()); CustomTypeDdl { sql: format!("CREATE TYPE {qualified} AS RANGE (subtype = unknown);"), complete: false, warnings, } } else { CustomTypeDdl { sql: format!("CREATE TYPE {qualified} AS RANGE (\n {}\n);", args.join(",\n ")), complete: true, warnings: warnings.to_vec(), } } } CustomTypeKind::Multirange => { let mut all = warnings.to_vec(); all.push(format!( "{} is the auto-generated multirange companion of a range type; it has no standalone CREATE statement", qualified )); CustomTypeDdl { sql: String::new(), complete: false, warnings: all } } CustomTypeKind::Base => { let mut all = warnings.to_vec(); all.push(format!( "{} is a base type; its input/output functions ({}) cannot be rebuilt from catalogs", qualified, info.input_function.clone().unwrap_or_else(|| "unknown".to_string()) )); CustomTypeDdl { sql: format!("CREATE TYPE {qualified}; -- base type attributes require manual reconstruction"), complete: false, warnings: all, } } } } /// Category-independent attribute assembly shared by all kinds. fn custom_type_common_properties(info: &CustomTypeGeneralInfo) -> CustomTypeProperties { CustomTypeProperties { input_function: info.input_function.clone(), output_function: info.output_function.clone(), receive_function: info.receive_function.clone(), send_function: info.send_function.clone(), analyze_function: info.analyze_function.clone(), internallength: (info.typlen != -1 && info.typlen > 0).then_some(info.typlen as i32), passed_by_value: Some(info.typbyval), alignment: (!info.typalign.is_empty()).then(|| info.typalign.clone()), storage: (!info.typstorage.is_empty()).then(|| info.typstorage.clone()), ..Default::default() } } /// Fetch the read-only details of a user-defined type. /// /// The caller must already have routed the connection to a verified /// PostgreSQL-family database; this function validates that the target object /// is an independent user-defined type (never a relation row type, an array /// companion, an undefined type or a system-schema type). pub async fn get_custom_type_details(pool: &Pool, schema: &str, name: &str) -> Result { let schema = schema.trim(); let name = name.trim(); if schema.is_empty() || name.is_empty() { return Err("schema and type name are required".to_string()); } if is_postgres_system_schema(schema) { return Err(format!("system schema {schema} is not supported for custom type details")); } let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let info = custom_type_general_info(&client, schema, name).await?; if !info.typisdefined { return Err(format!("custom type {schema}.{name} is not fully defined")); } if info.typelem != 0 { return Err(format!("custom type {schema}.{name} is an array companion type")); } let Some(kind) = custom_type_kind_for(&info.typtype) else { return Err(format!("custom type {schema}.{name} is a pseudo type (typtype={})", info.typtype)); }; if let Some(relkind) = info.relkind.as_deref() { if relkind != "c" { return Err(format!( "{schema}.{name} is the auto-generated row type of a relation, not an independent custom type" )); } } let mut members = Vec::new(); let mut properties = custom_type_common_properties(&info); let mut warnings = Vec::new(); match kind { CustomTypeKind::Enum => { members = custom_type_enum_members(&client, info.oid).await?; } CustomTypeKind::Composite => { members = custom_type_composite_members(&client, info.typrelid).await?; } CustomTypeKind::Domain => { warnings.extend(custom_type_domain_attributes(&client, &info, &mut properties).await); } CustomTypeKind::Range => { warnings.extend(custom_type_range_attributes(&client, info.oid, false, &mut properties).await); } CustomTypeKind::Multirange => { warnings.extend(custom_type_range_attributes(&client, info.oid, true, &mut properties).await); } CustomTypeKind::Base => {} } let ddl = build_custom_type_ddl(schema, name, kind, &info, &members, &properties, &warnings); Ok(CustomTypeDetails { name: name.to_string(), schema: schema.to_string(), kind, comment: info.comment.clone(), members, properties, ddl: Some(ddl), }) } pub async fn list_object_statistics(pool: &Pool, schema: &str) -> Result, String> { let schema = if schema.is_empty() { "public" } else { schema }; let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let rows = postgres_query_cached( &client, "SELECT c.relname, \ GREATEST(c.reltuples, 0)::bigint AS estimated_rows, \ pg_catalog.pg_total_relation_size(c.oid)::bigint AS total_bytes \ FROM pg_catalog.pg_class c \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ WHERE n.nspname = $1 AND c.relkind IN ('r','m','f','p') \ ORDER BY c.relname", &[&schema], ) .await .map_err(|e| e.to_string())?; Ok(rows .iter() .map(|row| ObjectStatistics { name: pg_row_try_string(row, 0), schema: Some(schema.to_string()), estimated_rows: row.try_get::<_, i64>(1).ok(), total_bytes: row.try_get::<_, i64>(2).ok(), }) .collect()) } pub async fn list_schemas(pool: &Pool) -> Result, String> { list_schemas_with_system(pool, false).await } pub async fn list_schema_infos(pool: &Pool) -> Result, String> { list_schema_infos_with_system(pool, false).await } pub async fn list_schemas_with_system(pool: &Pool, show_system_schemas: bool) -> Result, String> { Ok(list_schema_infos_with_system(pool, show_system_schemas).await?.into_iter().map(|schema| schema.name).collect()) } const POSTGRES_SCHEMA_INFOS_SQL: &str = "SELECT n.nspname AS schema_name, d.description AS schema_comment \ FROM pg_catalog.pg_namespace n \ LEFT JOIN pg_catalog.pg_description d \ ON d.objoid = n.oid \ AND d.objsubid = 0 \ AND d.classoid = 'pg_namespace'::regclass \ ORDER BY n.nspname"; const POSTGRES_SCHEMA_INFOS_HIDE_SYSTEM_SQL: &str = "SELECT n.nspname AS schema_name, d.description AS schema_comment \ FROM pg_catalog.pg_namespace n \ LEFT JOIN pg_catalog.pg_description d \ ON d.objoid = n.oid \ AND d.objsubid = 0 \ AND d.classoid = 'pg_namespace'::regclass \ WHERE n.nspname NOT IN ('information_schema', 'pg_catalog', 'pg_toast') \ AND n.nspname NOT LIKE 'pg_toast_temp_%' \ AND n.nspname NOT LIKE 'pg_temp_%' \ ORDER BY n.nspname"; fn postgres_schema_infos_sql(show_system_schemas: bool) -> &'static str { if show_system_schemas { POSTGRES_SCHEMA_INFOS_SQL } else { POSTGRES_SCHEMA_INFOS_HIDE_SYSTEM_SQL } } pub async fn list_schema_infos_with_system(pool: &Pool, show_system_schemas: bool) -> Result, String> { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let rows = postgres_query_cached(&client, postgres_schema_infos_sql(show_system_schemas), &[]) .await .map_err(|e| e.to_string())?; Ok(rows .iter() .map(|row| SchemaInfo { name: pg_row_try_string(row, 0), comment: row.try_get::<_, Option>(1).ok().flatten(), }) .collect()) } // Sibling of `postgres_columns_for_relations_sql`/`_compat_sql` above (~line // 3086), for a single (schema, table) instead of a batch of oids — see the // note there about why these aren't merged into a shared fragment. const POSTGRES_COLUMNS_SQL: &str = "SELECT a.attname AS column_name, \ format_type(a.atttypid, a.atttypmod) AS full_type, \ COALESCE(c.is_nullable = 'YES', NOT a.attnotnull) AS is_nullable, \ CASE WHEN a.attgenerated <> '' THEN NULL ELSE pg_get_expr(ad.adbin, ad.adrelid) END AS column_default, \ EXISTS ( \ SELECT 1 FROM pg_constraint co \ JOIN pg_index i ON i.indrelid = co.conrelid AND co.conindid = i.indexrelid \ WHERE co.conrelid = a.attrelid AND co.contype = 'p' \ AND a.attnum = ANY(i.indkey) \ ) AS is_pk, \ col_description(a.attrelid, a.attnum) AS column_comment, \ CASE a.attidentity \ WHEN 'd' THEN 'generated by default as identity' || CASE WHEN pseq.seqstart IS NOT NULL THEN format(' (start with %s increment by %s)', pseq.seqstart, pseq.seqincrement) ELSE '' END \ WHEN 'a' THEN 'generated always as identity' || CASE WHEN pseq.seqstart IS NOT NULL THEN format(' (start with %s increment by %s)', pseq.seqstart, pseq.seqincrement) ELSE '' END \ ELSE CASE a.attgenerated \ WHEN 's' THEN 'generated always as (' || pg_get_expr(ad.adbin, ad.adrelid) || ') stored' \ WHEN 'v' THEN 'generated always as (' || pg_get_expr(ad.adbin, ad.adrelid) || ') virtual' \ ELSE CASE WHEN a.atttypid IN (20, 21, 23) AND dep.deptype = 'a' \ AND pseq.seqrelid IS NOT NULL \ AND pg_get_expr(ad.adbin, ad.adrelid) = format('nextval(%L::regclass)', dep.objid::regclass::text) \ THEN CASE a.atttypid \ WHEN 21 THEN 'smallserial' \ WHEN 23 THEN 'serial' \ WHEN 20 THEN 'bigserial' \ END ELSE NULL END \ END \ END AS column_extra, \ CASE WHEN t.typname = 'numeric' AND a.atttypmod > 0 \ THEN ((a.atttypmod - 4) >> 16) & 65535 ELSE NULL END AS numeric_precision, \ CASE WHEN t.typname = 'numeric' AND a.atttypmod > 0 \ THEN (a.atttypmod - 4) & 65535 ELSE NULL END AS numeric_scale, \ CASE WHEN t.typname IN ('varchar', 'bpchar') AND a.atttypmod > 0 \ THEN a.atttypmod - 4 ELSE NULL END AS character_maximum_length, \ CASE WHEN enum_t.oid IS NULL THEN NULL \ ELSE COALESCE((SELECT array_to_json(array_agg(e.enumlabel ORDER BY e.enumsortorder))::text \ FROM pg_enum e WHERE e.enumtypid = enum_t.oid), '[]') END AS enum_values \ FROM pg_attribute a \ JOIN pg_type t ON t.oid = a.atttypid \ JOIN pg_class relation_class ON relation_class.oid = a.attrelid \ JOIN pg_namespace relation_namespace ON relation_namespace.oid = relation_class.relnamespace \ LEFT JOIN pg_type enum_t ON enum_t.oid = CASE WHEN t.typtype = 'd' THEN t.typbasetype WHEN t.typtype = 'e' THEN t.oid ELSE NULL END AND enum_t.typtype = 'e' \ LEFT JOIN pg_attrdef ad ON ad.adrelid = a.attrelid AND ad.adnum = a.attnum \ LEFT JOIN LATERAL ( \ SELECT sequence_dep.objid, sequence_dep.deptype \ FROM pg_catalog.pg_depend sequence_dep \ JOIN pg_catalog.pg_class sequence_class \ ON sequence_class.oid = sequence_dep.objid AND sequence_class.relkind = 'S' \ WHERE sequence_dep.classid = 'pg_catalog.pg_class'::pg_catalog.regclass \ AND sequence_dep.objsubid = 0 \ AND sequence_dep.refclassid = 'pg_catalog.pg_class'::pg_catalog.regclass \ AND sequence_dep.refobjid = a.attrelid AND sequence_dep.refobjsubid = a.attnum \ AND ((a.attidentity <> '' AND sequence_dep.deptype = 'i') OR (a.attidentity = '' \ AND sequence_dep.deptype = 'a' AND EXISTS ( \ SELECT 1 FROM pg_catalog.pg_depend serial_default_dep \ WHERE serial_default_dep.classid = 'pg_catalog.pg_attrdef'::pg_catalog.regclass \ AND serial_default_dep.objid = ad.oid AND serial_default_dep.objsubid = 0 \ AND serial_default_dep.refclassid = 'pg_catalog.pg_class'::pg_catalog.regclass \ AND serial_default_dep.refobjid = sequence_dep.objid \ AND serial_default_dep.refobjsubid = 0 AND serial_default_dep.deptype = 'n' \ ))) \ ORDER BY sequence_dep.objid \ LIMIT 1 \ ) dep ON TRUE \ LEFT JOIN pg_sequence pseq ON pseq.seqrelid = dep.objid \ LEFT JOIN information_schema.columns c \ ON c.table_schema = relation_namespace.nspname AND c.table_name = relation_class.relname AND c.column_name = a.attname \ WHERE a.attrelid = (CASE WHEN $1 = '' THEN quote_ident($2) ELSE quote_ident($1) || '.' || quote_ident($2) END)::regclass \ AND a.attnum > 0 AND NOT a.attisdropped \ ORDER BY a.attnum"; // Compat-tier sibling of `postgres_columns_for_relations_compat_sql` (~line // 3148) — see the note on `POSTGRES_COLUMNS_SQL` above. const POSTGRES_COLUMNS_COMPAT_SQL: &str = "SELECT a.attname AS column_name, \ format_type(a.atttypid, a.atttypmod) AS full_type, \ COALESCE(c.is_nullable = 'YES', NOT a.attnotnull) AS is_nullable, \ pg_get_expr(ad.adbin, ad.adrelid) AS column_default, \ EXISTS ( \ SELECT 1 FROM pg_constraint co \ JOIN pg_index i ON i.indrelid = co.conrelid AND co.conindid = i.indexrelid \ WHERE co.conrelid = a.attrelid AND co.contype = 'p' \ AND a.attnum = ANY(i.indkey) \ ) AS is_pk, \ col_description(a.attrelid, a.attnum) AS column_comment, \ CASE WHEN a.atttypid IN (20, 21, 23) AND serial_seq.oid IS NOT NULL \ AND pg_get_expr(ad.adbin, ad.adrelid) = format('nextval(%L::regclass)', serial_seq.oid::regclass::text) \ THEN CASE a.atttypid \ WHEN 21 THEN 'smallserial' \ WHEN 23 THEN 'serial' \ WHEN 20 THEN 'bigserial' \ END ELSE NULL END AS column_extra, \ CASE WHEN t.typname = 'numeric' AND a.atttypmod > 0 \ THEN ((a.atttypmod - 4) >> 16) & 65535 ELSE NULL END AS numeric_precision, \ CASE WHEN t.typname = 'numeric' AND a.atttypmod > 0 \ THEN (a.atttypmod - 4) & 65535 ELSE NULL END AS numeric_scale, \ CASE WHEN t.typname IN ('varchar', 'bpchar') AND a.atttypmod > 0 \ THEN a.atttypmod - 4 ELSE NULL END AS character_maximum_length, \ NULL::text AS enum_values \ FROM pg_attribute a \ JOIN pg_type t ON t.oid = a.atttypid \ JOIN pg_class relation_class ON relation_class.oid = a.attrelid \ JOIN pg_namespace relation_namespace ON relation_namespace.oid = relation_class.relnamespace \ LEFT JOIN pg_attrdef ad ON ad.adrelid = a.attrelid AND ad.adnum = a.attnum \ LEFT JOIN pg_catalog.pg_class serial_seq ON serial_seq.oid = ( \ SELECT sequence_dep.objid \ FROM pg_catalog.pg_depend sequence_dep \ JOIN pg_catalog.pg_class sequence_class \ ON sequence_class.oid = sequence_dep.objid AND sequence_class.relkind = 'S' \ WHERE sequence_dep.classid = 'pg_catalog.pg_class'::pg_catalog.regclass \ AND sequence_dep.objsubid = 0 \ AND sequence_dep.refclassid = 'pg_catalog.pg_class'::pg_catalog.regclass \ AND sequence_dep.refobjid = a.attrelid AND sequence_dep.refobjsubid = a.attnum \ AND sequence_dep.deptype = 'a' AND EXISTS ( \ SELECT 1 FROM pg_catalog.pg_depend serial_default_dep \ WHERE serial_default_dep.classid = 'pg_catalog.pg_attrdef'::pg_catalog.regclass \ AND serial_default_dep.objid = ad.oid AND serial_default_dep.objsubid = 0 \ AND serial_default_dep.refclassid = 'pg_catalog.pg_class'::pg_catalog.regclass \ AND serial_default_dep.refobjid = sequence_dep.objid \ AND serial_default_dep.refobjsubid = 0 AND serial_default_dep.deptype = 'n' \ ) \ ORDER BY sequence_dep.objid \ LIMIT 1 \ ) AND serial_seq.relkind = 'S' \ LEFT JOIN information_schema.columns c \ ON c.table_schema = relation_namespace.nspname AND c.table_name = relation_class.relname AND c.column_name = a.attname \ WHERE a.attrelid = (CASE WHEN $1 = '' THEN quote_ident($2) ELSE quote_ident($1) || '.' || quote_ident($2) END)::regclass \ AND a.attnum > 0 AND NOT a.attisdropped \ ORDER BY a.attnum"; const POSTGRES_COLUMNS_INFORMATION_SCHEMA_SQL: &str = "SELECT c.column_name, \ CASE WHEN c.data_type = 'USER-DEFINED' THEN c.udt_name ELSE c.data_type END AS full_type, \ c.is_nullable = 'YES' AS is_nullable, \ c.column_default, \ EXISTS ( \ SELECT 1 FROM information_schema.table_constraints tc \ JOIN information_schema.key_column_usage kcu \ ON kcu.constraint_catalog = tc.constraint_catalog \ AND kcu.constraint_schema = tc.constraint_schema \ AND kcu.constraint_name = tc.constraint_name \ AND kcu.table_schema = tc.table_schema \ AND kcu.table_name = tc.table_name \ WHERE tc.constraint_type = 'PRIMARY KEY' \ AND tc.table_schema = c.table_schema \ AND tc.table_name = c.table_name \ AND kcu.column_name = c.column_name \ ) AS is_pk, \ NULL::text AS column_comment, \ NULL::text AS column_extra, \ CAST(c.numeric_precision AS int) AS numeric_precision, \ CAST(c.numeric_scale AS int) AS numeric_scale, \ CAST(c.character_maximum_length AS int) AS character_maximum_length, \ NULL::text AS enum_values \ FROM information_schema.columns c \ WHERE c.table_schema = $1 AND c.table_name = $2 \ ORDER BY c.ordinal_position"; fn parse_enum_values_from_row(row: &Row, index: usize) -> Option> { let raw = row.try_get::<_, Option>(index).ok().flatten()?; serde_json::from_str::>(&raw).ok() } /// Decode a boolean column to JSON, tolerating databases (e.g. GaussDB) that /// encode booleans as the ASCII bytes `t` (0x74) / `f` (0x66) in the binary /// protocol instead of the standard PostgreSQL 0x00 / 0x01. fn pg_bool_value_to_json(row: &Row, idx: usize) -> serde_json::Value { if let Some(v) = pg_row_try_bool(row, idx) { return serde_json::Value::Bool(v); } serde_json::Value::Null } /// Map raw boolean bytes to a Rust `bool`. /// /// Standard PostgreSQL binary uses `[0x00]` / `[0x01]`; GaussDB sends the ASCII /// text representation `[b't']` / `[b'f']` instead. fn decode_bool_bytes(raw: &[u8]) -> Option { match raw { [0x00] => Some(false), [0x01] => Some(true), [b't'] | [b'T'] => Some(true), [b'f'] | [b'F'] => Some(false), _ => None, } } fn decode_bool_candidates(raw: Option<&[u8]>, standard: Option) -> Option { raw.and_then(decode_bool_bytes).or(standard) } /// Read a boolean column from a PostgreSQL row, tolerating databases that /// encode booleans as integers (0/1) or text ('t'/'f') instead of the standard /// `bool` OID. Returns `None` when the column is NULL or truly unreadable. fn pg_row_try_bool(row: &Row, idx: usize) -> Option { // GaussDB encodes boolean as ASCII 't' (0x74) / 'f' (0x66) in binary. let raw = row.try_get::<_, PgRawBytes>(idx).ok(); let standard = row.try_get::<_, bool>(idx).ok(); if let Some(v) = decode_bool_candidates(raw.as_ref().map(|value| value.0.as_slice()), standard) { return Some(v); } if let Ok(v) = row.try_get::<_, i32>(idx) { return Some(v != 0); } if let Ok(v) = row.try_get::<_, i16>(idx) { return Some(v != 0); } if let Ok(Some(v)) = row.try_get::<_, Option>(idx) { match v.as_str() { "t" | "true" | "1" | "yes" | "YES" => return Some(true), "f" | "false" | "0" | "no" | "NO" => return Some(false), _ => return None, } } None } /// Read a String column from a PostgreSQL row, tolerating databases that /// return text as other types. Falls back to i64/i32/i16/bool formatting. fn pg_row_try_string(row: &Row, idx: usize) -> String { if let Ok(v) = row.try_get::<_, String>(idx) { return v; } if let Ok(v) = row.try_get::<_, i64>(idx) { return v.to_string(); } if let Ok(v) = row.try_get::<_, i32>(idx) { return v.to_string(); } if let Ok(v) = row.try_get::<_, i16>(idx) { return v.to_string(); } if let Some(v) = pg_row_try_bool(row, idx) { return v.to_string(); } String::new() } fn column_info_from_row(row: &Row) -> ColumnInfo { column_info_from_row_offset(row, 0) } async fn get_columns_with_sql( client: &deadpool_postgres::Client, sql: &str, schema: &str, table: &str, ) -> Result, tokio_postgres::Error> { let rows = postgres_query_cached(client, sql, &[&schema, &table]).await?; Ok(rows.iter().map(column_info_from_row).collect()) } pub async fn get_columns(pool: &Pool, schema: &str, table: &str) -> Result, String> { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let tiers = [POSTGRES_COLUMNS_SQL, POSTGRES_COLUMNS_COMPAT_SQL, POSTGRES_COLUMNS_INFORMATION_SCHEMA_SQL]; query_with_compat_fallback("get_columns", &tiers, |sql| get_columns_with_sql(&client, sql, schema, table)).await } fn pg_quote_literal(value: &str) -> String { format!("'{}'", value.replace('\'', "''")) } fn redshift_columns_sql(schema: &str, table: &str) -> String { format!( "SELECT c.column_name, \ CASE WHEN c.data_type = 'USER-DEFINED' THEN c.udt_name ELSE c.data_type END AS full_type, \ c.is_nullable, \ c.column_default, \ CAST(c.numeric_precision AS varchar) AS numeric_precision, \ CAST(c.numeric_scale AS varchar) AS numeric_scale, \ CAST(c.character_maximum_length AS varchar) AS character_maximum_length \ FROM information_schema.columns c \ WHERE c.table_schema = {} AND c.table_name = {} \ ORDER BY c.ordinal_position", pg_quote_literal(schema), pg_quote_literal(table) ) } fn query_result_text(row: &[serde_json::Value], index: usize) -> Option { row.get(index).and_then(|value| match value { serde_json::Value::String(value) => Some(value.clone()), serde_json::Value::Number(value) => Some(value.to_string()), serde_json::Value::Bool(value) => Some(value.to_string()), serde_json::Value::Null | serde_json::Value::Array(_) | serde_json::Value::Object(_) => None, }) } fn query_result_i32(row: &[serde_json::Value], index: usize) -> Option { query_result_text(row, index)?.parse().ok() } fn redshift_columns_from_query_result(result: QueryResult) -> Vec { result .rows .into_iter() .filter_map(|row| { Some(ColumnInfo { name: query_result_text(&row, 0)?, data_type: query_result_text(&row, 1).unwrap_or_default(), is_nullable: query_result_text(&row, 2).is_none_or(|value| value.eq_ignore_ascii_case("YES")), column_default: query_result_text(&row, 3), is_primary_key: false, extra: None, comment: None, numeric_precision: query_result_i32(&row, 4), numeric_scale: query_result_i32(&row, 5), character_maximum_length: query_result_i32(&row, 6), enum_values: None, ..Default::default() }) }) .collect() } pub async fn get_redshift_columns(pool: &Pool, schema: &str, table: &str) -> Result, String> { let schema = if schema.is_empty() { "public" } else { schema }; let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let result = execute_select_text( &client, &redshift_columns_sql(schema, table), Instant::now(), crate::query::MAX_ROWS, None, None, ) .await?; Ok(redshift_columns_from_query_result(result)) } pub(crate) fn pg_quote_ident(ident: &str) -> String { format!("\"{}\"", ident.replace('"', "\"\"")) } fn pg_quote_catalog_identifier(ident: &str) -> String { let ident = ident.trim(); if ident.starts_with('"') && ident.ends_with('"') && ident.contains("\".\"") { return ident.to_string(); } if let Some((schema, name)) = ident.rsplit_once('.') { return format!("{}.{}", pg_quote_ident(schema), pg_quote_ident(name)); } pg_quote_ident(ident) } #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) enum PostgresSearchPathContext { Query, Transaction, LocalTransaction, LocalQueryTransaction, } #[derive(Debug, Clone, PartialEq, Eq)] struct PostgresSearchPathBaseline { configured: String, first_resolved_schema: Option, has_explicit_pg_catalog: bool, } type PostgresSearchPathBaselineEntry = (Weak, PostgresSearchPathBaseline); type PostgresSearchPathBaselineCache = Mutex>; pub(crate) fn postgres_set_search_path_sql(schema: &str, context: PostgresSearchPathContext) -> String { let (scope, suffix) = match context { // Ordinary queries and exports historically fall back to public for // extensions and helper functions after checking the selected schema. PostgresSearchPathContext::Query => ("", ", pg_catalog, public"), PostgresSearchPathContext::Transaction => ("", ", pg_catalog"), PostgresSearchPathContext::LocalTransaction => (" LOCAL", ", pg_catalog"), PostgresSearchPathContext::LocalQueryTransaction => (" LOCAL", ", pg_catalog, public"), }; // PostgreSQL otherwise searches pg_catalog before every explicit path item. format!("SET{scope} search_path TO {}{suffix}", pg_quote_ident(schema)) } fn postgres_set_single_schema_search_path_sql(schema: &str, context: PostgresSearchPathContext) -> String { let scope = if matches!( context, PostgresSearchPathContext::LocalTransaction | PostgresSearchPathContext::LocalQueryTransaction ) { " LOCAL" } else { "" }; format!("SET{scope} search_path TO {}", pg_quote_ident(schema)) } fn postgres_set_preserved_search_path_sql( schema: &str, context: PostgresSearchPathContext, baseline: &PostgresSearchPathBaseline, ) -> String { let scope = if matches!( context, PostgresSearchPathContext::LocalTransaction | PostgresSearchPathContext::LocalQueryTransaction ) { " LOCAL" } else { "" }; let configured = baseline.configured.trim(); let selected_schema = pg_quote_ident(schema); let mut path = if baseline.first_resolved_schema.as_deref() == Some(schema) { configured.to_string() } else if configured.is_empty() { selected_schema.clone() } else { format!("{selected_schema}, {configured}") }; if path.is_empty() { path = selected_schema; } if !baseline.has_explicit_pg_catalog { path.push_str(", pg_catalog"); } format!("SET{scope} search_path TO {path}") } fn postgres_requires_single_schema_search_path(error: &str) -> bool { error.to_ascii_lowercase().contains("does not support search_path with multiple names") } fn postgres_single_schema_clients() -> &'static Mutex>> { static CLIENTS: OnceLock>>> = OnceLock::new(); CLIENTS.get_or_init(|| Mutex::new(HashMap::new())) } fn postgres_client_uses_single_schema_search_path(client: &deadpool_postgres::Client) -> bool { let statement_cache = &client.statement_cache; let key = Arc::as_ptr(statement_cache) as usize; let mut clients = postgres_single_schema_clients().lock().unwrap_or_else(|poisoned| poisoned.into_inner()); match clients.get(&key).and_then(Weak::upgrade) { Some(cached) if Arc::ptr_eq(&cached, statement_cache) => true, _ => { clients.remove(&key); false } } } fn mark_postgres_client_single_schema_search_path(client: &deadpool_postgres::Client) { let statement_cache = &client.statement_cache; let key = Arc::as_ptr(statement_cache) as usize; let mut clients = postgres_single_schema_clients().lock().unwrap_or_else(|poisoned| poisoned.into_inner()); clients.retain(|_, cached| cached.strong_count() > 0); clients.insert(key, Arc::downgrade(statement_cache)); } fn postgres_search_path_baselines() -> &'static PostgresSearchPathBaselineCache { static BASELINES: OnceLock = OnceLock::new(); BASELINES.get_or_init(|| Mutex::new(HashMap::new())) } fn cached_postgres_search_path_baseline(client: &deadpool_postgres::Client) -> Option { let statement_cache = &client.statement_cache; let key = Arc::as_ptr(statement_cache) as usize; let mut baselines = postgres_search_path_baselines().lock().unwrap_or_else(|poisoned| poisoned.into_inner()); match baselines.get(&key) { Some((cached, baseline)) if cached.upgrade().is_some_and(|cached| Arc::ptr_eq(&cached, statement_cache)) => { Some(baseline.clone()) } _ => { baselines.remove(&key); None } } } fn cache_postgres_search_path_baseline(client: &deadpool_postgres::Client, baseline: PostgresSearchPathBaseline) { let statement_cache = &client.statement_cache; let key = Arc::as_ptr(statement_cache) as usize; let mut baselines = postgres_search_path_baselines().lock().unwrap_or_else(|poisoned| poisoned.into_inner()); baselines.retain(|_, (cached, _)| cached.strong_count() > 0); baselines.insert(key, (Arc::downgrade(statement_cache), baseline)); } async fn postgres_search_path_baseline( client: &deadpool_postgres::Client, timeout_duration: Duration, ) -> Result { if let Some(baseline) = cached_postgres_search_path_baseline(client) { return Ok(baseline); } let row = tokio::time::timeout( timeout_duration, client.query_one( "SELECT current_setting('search_path'), (current_schemas(false))[1], \ 'pg_catalog' = ANY(current_schemas(false))", &[], ), ) .await .map_err(|_| format!("PostgreSQL schema.get timed out after {} seconds", timeout_duration.as_secs()))? .map_err(pg_error_to_string)?; let baseline = PostgresSearchPathBaseline { configured: row.try_get(0).map_err(|error| error.to_string())?, first_resolved_schema: row.try_get(1).map_err(|error| error.to_string())?, has_explicit_pg_catalog: row.try_get(2).map_err(|error| error.to_string())?, }; cache_postgres_search_path_baseline(client, baseline.clone()); Ok(baseline) } pub(crate) async fn set_postgres_search_path( client: &deadpool_postgres::Client, schema: &str, context: PostgresSearchPathContext, timeout_duration: Duration, ) -> Result { if postgres_client_uses_single_schema_search_path(client) { return execute_postgres_infra_statement( client, &postgres_set_single_schema_search_path_sql(schema, context), timeout_duration, "schema.set", ) .await; } let primary_sql = match postgres_search_path_baseline(client, timeout_duration).await { Ok(baseline) => postgres_set_preserved_search_path_sql(schema, context, &baseline), Err(error) => { log::warn!("[postgres][schema.get:error] schema={schema} error={error}; using compatibility fallback"); postgres_set_search_path_sql(schema, context) } }; match execute_postgres_infra_statement(client, &primary_sql, timeout_duration, "schema.set").await { Ok(affected) => Ok(affected), Err(primary_error) if postgres_requires_single_schema_search_path(&primary_error) => { mark_postgres_client_single_schema_search_path(client); log::info!("[postgres][schema.set:single-schema-fallback] schema={schema}"); execute_postgres_infra_statement( client, &postgres_set_single_schema_search_path_sql(schema, context), timeout_duration, "schema.set", ) .await .map_err(|fallback_error| { format!("{primary_error}; single-schema search_path fallback failed: {fallback_error}") }) } Err(error) => Err(error), } } fn query_result_row_limit(max_rows: Option) -> usize { max_rows.unwrap_or(crate::query::MAX_ROWS).max(1) } /// Returns whether a parsed DML statement produces a result set, or `None` when /// the statement is not DML. fn postgres_dml_statement_returns_rows(statement: &Statement) -> Option { match statement { Statement::Insert(insert) => Some(insert.returning.is_some()), Statement::Update(update) => Some(update.returning.is_some()), Statement::Delete(delete) => Some(delete.returning.is_some()), Statement::Merge(merge) => Some(merge.output.is_some()), _ => None, } } /// `WITH cte AS (...) UPDATE ...` parses as a query whose body is the wrapped /// DML statement; returns that statement so its `RETURNING` clause can decide. fn postgres_query_body_dml(body: &SetExpr) -> Option<&Statement> { match body { SetExpr::Insert(statement) | SetExpr::Update(statement) | SetExpr::Delete(statement) | SetExpr::Merge(statement) => Some(statement), _ => None, } } /// Returns whether PostgreSQL should execute this statement through the row /// retrieval path. DML without `RETURNING` needs `execute` for its command /// tag/affected-row count, while DML with `RETURNING` produces a result set. pub(crate) fn postgres_statement_returns_rows(sql: &str) -> bool { if starts_with_executable_sql_keyword(sql, &["SELECT", "SHOW", "EXPLAIN", "TABLE"]) { return true; } // A leading `WITH` is usually a CTE query, but PostgreSQL also allows a CTE // to be followed directly by INSERT/UPDATE/DELETE/MERGE, which only returns // rows with `RETURNING`. The prefix alone cannot decide, so parse it and // keep the row-returning assumption whenever parsing does not say otherwise. let leads_with_cte = starts_with_executable_sql_keyword(sql, &["WITH"]); let Ok(statements) = Parser::parse_sql(&PostgreSqlDialect {}, sql) else { return leads_with_cte; }; let [statement] = statements.as_slice() else { return leads_with_cte; }; if let Some(returns_rows) = postgres_dml_statement_returns_rows(statement) { return returns_rows; } match statement { Statement::Query(query) => { postgres_query_body_dml(&query.body).and_then(postgres_dml_statement_returns_rows).unwrap_or(true) } _ => leads_with_cte, } } pub async fn execute_query(pool: &Pool, sql: &str) -> Result { execute_query_with_max_rows(pool, sql, None).await } pub async fn execute_query_with_max_rows( pool: &Pool, sql: &str, max_rows: Option, ) -> Result { let start = Instant::now(); let row_limit = query_result_row_limit(max_rows); if postgres_statement_returns_rows(sql) { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; execute_select_query(&client, sql, start, row_limit).await } else { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let affected = client.execute(sql, &[]).await.map_err(pg_error_to_string)?; clear_postgres_caches_after_ddl(pool, Some(&client), sql); Ok(QueryResult { columns: vec![], column_types: Vec::new(), column_sortables: Vec::new(), spatial_columns: vec![], spatial_values: vec![], rows: vec![], affected_rows: affected, execution_time_ms: start.elapsed().as_millis(), truncated: false, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: Vec::new(), }) } } pub async fn execute_query_with_max_rows_and_cancel( pool: &Pool, sql: &str, max_rows: Option, cancel_token: Option, budget: DbOperationBudget, cancel_context: Option, prefer_text_protocol: bool, ) -> Result { let client = checkout_postgres_client(pool, cancel_token.as_ref(), budget.checkout_timeout).await?; execute_postgres_user_query( &client, sql, max_rows, cancel_token, budget.query_timeout, budget.cancel_timeout, cancel_context, prefer_text_protocol, ) .await } fn postgres_read_only_transaction_setup() -> Vec<(String, &'static str)> { vec![("BEGIN READ ONLY".to_string(), "explain_analyze.begin")] } fn postgres_read_only_transaction_cleanup_error(error: String) -> String { format!("PostgreSQL read-only transaction cleanup failed: {error}") } fn merge_postgres_operation_and_rollback_result( operation_result: Result, rollback_result: Result<(), String>, ) -> Result { match (operation_result, rollback_result) { (Ok(result), Ok(())) => Ok(result), (Err(operation_error), Ok(())) => Err(operation_error), (Ok(_), Err(rollback_error)) => Err(postgres_read_only_transaction_cleanup_error(rollback_error)), (Err(operation_error), Err(rollback_error)) => { Err(format!("{operation_error}; {}", postgres_read_only_transaction_cleanup_error(rollback_error))) } } } async fn run_postgres_operation_with_rollback( operation: Operation, rollback: Rollback, ) -> Result where Operation: FnOnce() -> OperationFuture, OperationFuture: Future>, Rollback: FnOnce() -> RollbackFuture, RollbackFuture: Future>, { let operation_result = operation().await; let rollback_result = rollback().await; merge_postgres_operation_and_rollback_result(operation_result, rollback_result) } pub async fn execute_query_in_read_only_transaction_with_rollback( pool: &Pool, schema: Option<&str>, sql: &str, max_rows: Option, cancel_token: Option, budget: DbOperationBudget, cancel_context: Option, ) -> Result { let client = checkout_postgres_client(pool, cancel_token.as_ref(), budget.checkout_timeout).await?; let setup = postgres_read_only_transaction_setup(); run_postgres_operation_with_rollback( || async { for (statement, stage) in setup { execute_postgres_infra_statement(&client, &statement, budget.recycle_timeout, stage).await?; } if let Some(schema) = schema.map(str::trim).filter(|schema| !schema.is_empty()) { set_postgres_search_path( &client, schema, PostgresSearchPathContext::LocalQueryTransaction, budget.recycle_timeout, ) .await?; } execute_postgres_user_query_with_mode( &client, sql, max_rows, cancel_token, budget.query_timeout, budget.cancel_timeout, cancel_context, false, true, ) .await }, || async { execute_postgres_infra_statement(&client, "ROLLBACK", budget.cleanup_timeout, "explain_analyze.rollback") .await .map(|_| ()) }, ) .await } pub async fn stream_select_query_with_cancel( pool: &Pool, schema: Option<&str>, setup_sql: &[String], sql: &str, max_rows: Option, cancel_token: Option, budget: DbOperationBudget, cancel_context: Option, on_item: impl FnMut(PostgresQueryStreamItem) -> Result<(), String>, ) -> Result { let start = Instant::now(); let client = checkout_postgres_client(pool, cancel_token.as_ref(), budget.checkout_timeout).await?; let mut on_item = on_item; let row_limit = max_rows.map(|limit| limit.max(1)); let schema = schema.map(str::trim).filter(|schema| !schema.is_empty()); let schema_was_set = schema.is_some_and(|_| !is_transaction_recovery_statement(sql)); if let Some(schema) = schema.filter(|_| schema_was_set) { // Match normal query execution: export may reference unqualified names // in the active schema, so the streaming path must use the same search_path. set_postgres_search_path(&client, schema, PostgresSearchPathContext::Query, budget.recycle_timeout).await?; } let setup_transaction_started = !setup_sql.is_empty(); if setup_transaction_started { execute_postgres_infra_statement(&client, "BEGIN", budget.recycle_timeout, "export_setup.begin").await?; } let query_timeout = budget.query_timeout; let timeout_error = format!("Query timed out after {} seconds", query_timeout.map_or(0, |timeout| timeout.as_secs())); let setup_result = async { for setup_statement in setup_sql { wait_postgres_query( client.cancel_token(), cancel_context.clone(), cancel_token.clone(), query_timeout, budget.cancel_timeout, async { client.batch_execute(setup_statement).await.map_err(pg_error_to_string)?; Ok(()) }, ) .await?; } Ok(()) } .await; let result = match setup_result { Ok(()) => { let pg_cancel_token = client.cancel_token(); let progress_clock = Arc::new(StreamProgressClock::new()); let progress_clock_for_stream = progress_clock.clone(); let mut on_stream_item = |item| { on_item(item)?; progress_clock_for_stream.mark(); Ok(()) }; let result = await_stream_with_progress_timeout( stream_select_query_inner_with_mode( &client, sql, row_limit, &mut on_stream_item, setup_transaction_started, ), query_timeout, progress_clock, cancel_token.as_ref(), timeout_error.clone(), ) .await; if result.as_ref().is_err_and(|error| error == &timeout_error || error == crate::query::QUERY_CANCELED) { cancel_postgres_query(pg_cancel_token, cancel_context.as_ref(), budget.cancel_timeout).await; } result } Err(error) => Err(error), }; let result = if setup_transaction_started { let rollback_result = execute_postgres_infra_statement(&client, "ROLLBACK", budget.cleanup_timeout, "export_setup.rollback") .await; match (result, rollback_result) { (Ok(rows), Ok(_)) => Ok(rows), (Err(query_err), Ok(_)) => Err(query_err), (Ok(_), Err(rollback_err)) => Err(rollback_err), (Err(query_err), Err(rollback_err)) => Err(format!("{query_err}; {rollback_err}")), } } else { result }; if schema_was_set { let reset_result = reset_postgres_search_path(&client, budget.cleanup_timeout, start).await; match (result, reset_result) { (Ok(rows), Ok(())) => Ok(rows), (Err(query_err), Ok(())) => Err(query_err), (Ok(_), Err(reset_err)) => Err(reset_err), (Err(query_err), Err(reset_err)) => Err(format!("{query_err}; {reset_err}")), } } else { result } } pub async fn execute_query_with_schema(pool: &Pool, schema: &str, sql: &str) -> Result { execute_query_with_schema_and_max_rows(pool, schema, sql, None).await } pub async fn execute_query_with_schema_and_max_rows( pool: &Pool, schema: &str, sql: &str, max_rows: Option, ) -> Result { let start = Instant::now(); let checkout_start = Instant::now(); let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; log::info!( "[postgres][execute_with_schema:pool:done] elapsed_ms={} total_ms={} schema={}", checkout_start.elapsed().as_millis(), start.elapsed().as_millis(), schema ); if is_transaction_recovery_statement(sql) { log::info!( "[postgres][execute_with_schema:skip-search-path] total_ms={} reason=transaction-recovery", start.elapsed().as_millis() ); return execute_query_with_max_rows_inner(&client, sql, max_rows, false, None, false).await; } let set_schema_start = Instant::now(); set_postgres_search_path(&client, schema, PostgresSearchPathContext::Query, super::connection_timeout()).await?; log::info!( "[postgres][execute_with_schema:set-search-path:done] elapsed_ms={} total_ms={}", set_schema_start.elapsed().as_millis(), start.elapsed().as_millis() ); let query_start = Instant::now(); let result = execute_query_with_max_rows_inner(&client, sql, max_rows, false, None, false).await; if result.is_ok() { clear_postgres_caches_after_ddl(pool, Some(&client), sql); } log::info!( "[postgres][execute_with_schema:query:done] elapsed_ms={} total_ms={} ok={}", query_start.elapsed().as_millis(), start.elapsed().as_millis(), result.is_ok() ); let reset_result = reset_postgres_search_path(&client, super::connection_timeout(), start).await; merge_postgres_query_and_reset_result(result, reset_result) } pub async fn execute_query_with_schema_and_max_rows_and_cancel( pool: &Pool, schema: &str, sql: &str, max_rows: Option, cancel_token: Option, budget: DbOperationBudget, cancel_context: Option, prefer_text_protocol: bool, ) -> Result { let start = Instant::now(); let checkout_start = Instant::now(); let client = checkout_postgres_client(pool, cancel_token.as_ref(), budget.checkout_timeout).await?; log::info!( "[postgres][execute_with_schema:pool:done] elapsed_ms={} total_ms={} schema={}", checkout_start.elapsed().as_millis(), start.elapsed().as_millis(), schema ); if is_transaction_recovery_statement(sql) { log::info!( "[postgres][execute_with_schema:skip-search-path] total_ms={} reason=transaction-recovery", start.elapsed().as_millis() ); return execute_postgres_user_query( &client, sql, max_rows, cancel_token, budget.query_timeout, budget.cancel_timeout, cancel_context, prefer_text_protocol, ) .await; } let set_schema_start = Instant::now(); set_postgres_search_path(&client, schema, PostgresSearchPathContext::Query, budget.recycle_timeout).await?; log::info!( "[postgres][execute_with_schema:set-search-path:done] elapsed_ms={} total_ms={}", set_schema_start.elapsed().as_millis(), start.elapsed().as_millis() ); let query_start = Instant::now(); let result = execute_postgres_user_query( &client, sql, max_rows, cancel_token, budget.query_timeout, budget.cancel_timeout, cancel_context, prefer_text_protocol, ) .await; if result.is_ok() { clear_postgres_caches_after_ddl(pool, Some(&client), sql); } log::info!( "[postgres][execute_with_schema:query:done] elapsed_ms={} total_ms={} ok={}", query_start.elapsed().as_millis(), start.elapsed().as_millis(), result.is_ok() ); let reset_result = reset_postgres_search_path(&client, budget.cleanup_timeout, start).await; merge_postgres_query_and_reset_result(result, reset_result) } async fn reset_postgres_search_path( client: &deadpool_postgres::Client, timeout_duration: Duration, start: Instant, ) -> Result<(), String> { let reset_start = Instant::now(); match execute_postgres_infra_statement(client, "RESET search_path", timeout_duration, "schema.reset").await { Ok(_) => { log::info!( "[postgres][execute_with_schema:reset-search-path:done] elapsed_ms={} total_ms={}", reset_start.elapsed().as_millis(), start.elapsed().as_millis() ); Ok(()) } Err(err) => { log::warn!( "[postgres][execute_with_schema:reset-search-path:error] elapsed_ms={} total_ms={} error={}", reset_start.elapsed().as_millis(), start.elapsed().as_millis(), err ); Err(postgres_schema_reset_cleanup_error(err)) } } } fn merge_postgres_query_and_reset_result( query_result: Result, reset_result: Result<(), String>, ) -> Result { match (query_result, reset_result) { (Ok(result), Ok(())) => Ok(result), (Err(query_err), Ok(())) => Err(query_err), (Ok(_), Err(reset_err)) => Err(reset_err), (Err(query_err), Err(reset_err)) => Err(format!("{query_err}; {reset_err}")), } } fn postgres_schema_reset_cleanup_error(err: String) -> String { format!("PostgreSQL schema.reset cleanup failed: {err}") } pub(crate) async fn execute_postgres_infra_statement( client: &deadpool_postgres::Client, sql: &str, timeout_duration: Duration, stage: &str, ) -> Result { tokio::time::timeout(timeout_duration, client.execute_typed(sql, &[])) .await .map_err(|_| format!("PostgreSQL {stage} timed out after {} seconds", timeout_duration.as_secs()))? .map_err(pg_error_to_string) } pub(crate) async fn wait_postgres_operation( pg_cancel_token: tokio_postgres::CancelToken, cancel_context: Option, timeout_duration: Option, cancel_timeout: Duration, future: F, ) -> Result where F: Future>, { wait_postgres_query(pg_cancel_token, cancel_context, None, timeout_duration, cancel_timeout, future).await } async fn wait_postgres_query( pg_cancel_token: tokio_postgres::CancelToken, cancel_context: Option, cancel_token: Option, timeout_duration: Option, cancel_timeout: Duration, future: F, ) -> Result where F: Future>, { match (cancel_token, timeout_duration) { (Some(token), Some(duration)) => { tokio::select! { biased; _ = token.cancelled() => { cancel_postgres_query(pg_cancel_token, cancel_context.as_ref(), cancel_timeout).await; Err(crate::query::canceled_error()) } result = tokio::time::timeout(duration, future) => match result { Ok(result) => result, Err(_) => { cancel_postgres_query(pg_cancel_token, cancel_context.as_ref(), cancel_timeout).await; Err(format!("Query timed out after {} seconds", duration.as_secs())) } }, } } (None, Some(duration)) => match tokio::time::timeout(duration, future).await { Ok(result) => result, Err(_) => { cancel_postgres_query(pg_cancel_token, cancel_context.as_ref(), cancel_timeout).await; Err(format!("Query timed out after {} seconds", duration.as_secs())) } }, (Some(token), None) => { tokio::select! { biased; _ = token.cancelled() => { cancel_postgres_query(pg_cancel_token, cancel_context.as_ref(), cancel_timeout).await; Err(crate::query::canceled_error()) } result = future => result, } } (None, None) => future.await, } } async fn execute_postgres_user_query( client: &deadpool_postgres::Client, sql: &str, max_rows: Option, cancel_token: Option, timeout_duration: Option, cancel_timeout: Duration, cancel_context: Option, prefer_text_protocol: bool, ) -> Result { execute_postgres_user_query_with_mode( client, sql, max_rows, cancel_token, timeout_duration, cancel_timeout, cancel_context, prefer_text_protocol, false, ) .await } async fn execute_postgres_user_query_with_mode( client: &deadpool_postgres::Client, sql: &str, max_rows: Option, cancel_token: Option, timeout_duration: Option, cancel_timeout: Duration, cancel_context: Option, prefer_text_protocol: bool, force_unnamed: bool, ) -> Result { let pg_cancel_token = client.cancel_token(); // Commands do not expose incremental results, so keep their timeout as a // wall-clock deadline. Row-returning queries may spend longer transferring // a bounded result through a slow tunnel; for those, the same setting is an // inactivity budget that resets only after real PostgreSQL progress. if !postgres_statement_returns_rows(sql) { return wait_postgres_query( pg_cancel_token, cancel_context, cancel_token, timeout_duration, cancel_timeout, execute_query_with_max_rows_inner(client, sql, max_rows, prefer_text_protocol, None, force_unnamed), ) .await; } let timeout_error = format!("Query timed out after {} seconds", timeout_duration.map_or(0, |timeout| timeout.as_secs())); let progress_clock = Arc::new(StreamProgressClock::new()); let result = await_stream_with_progress_timeout( execute_query_with_max_rows_inner( client, sql, max_rows, prefer_text_protocol, Some(progress_clock.clone()), force_unnamed, ), timeout_duration, progress_clock, cancel_token.as_ref(), timeout_error.clone(), ) .await; if result.as_ref().is_err_and(|error| error == &timeout_error || error == crate::query::QUERY_CANCELED) { cancel_postgres_query(pg_cancel_token, cancel_context.as_ref(), cancel_timeout).await; } result } /// PostgreSQL pool checkout with timeout and cancel token support. /// When the checkout phase is stuck, the cancel token can terminate the wait early. pub(crate) type PostgresCheckoutStage = super::PoolCheckoutStage; pub(crate) fn classify_postgres_checkout_stage(status: deadpool_postgres::Status) -> PostgresCheckoutStage { if status.waiting > 0 || (status.available == 0 && status.size >= status.max_size) { PostgresCheckoutStage::Wait } else if status.available > 0 { PostgresCheckoutStage::Recycle } else if status.size < status.max_size { PostgresCheckoutStage::Create } else { PostgresCheckoutStage::Unknown } } fn postgres_pool_error_stage(error: &PoolError) -> PostgresCheckoutStage { match error { PoolError::Timeout(deadpool_postgres::TimeoutType::Wait) => PostgresCheckoutStage::Wait, PoolError::Timeout(deadpool_postgres::TimeoutType::Create) => PostgresCheckoutStage::Create, PoolError::Timeout(deadpool_postgres::TimeoutType::Recycle) => PostgresCheckoutStage::Recycle, PoolError::Backend(_) | PoolError::PostCreateHook(_) => PostgresCheckoutStage::Create, PoolError::Closed | PoolError::NoRuntimeSpecified => PostgresCheckoutStage::Unknown, } } pub(crate) async fn checkout_postgres_client_classified( pool: &Pool, cancel_token: Option<&CancellationToken>, checkout_timeout: Duration, ) -> Result { let start = Instant::now(); let get_future = async { match tokio::time::timeout(checkout_timeout, pool.get()).await { Ok(Ok(client)) => Ok(client), Ok(Err(error)) => { let stage = postgres_pool_error_stage(&error); let timed_out = matches!(&error, PoolError::Timeout(_)); let detail = pg_pool_error_to_string(error); log::warn!( "[db:pool.checkout:error] elapsed_ms={} timeout_ms={} stage={} error={}", start.elapsed().as_millis(), checkout_timeout.as_millis(), stage.as_str(), detail ); if timed_out { Err(super::PoolCheckoutError::Timeout { database: "PostgreSQL", stage, timeout: checkout_timeout }) } else { Err(super::PoolCheckoutError::Failed { database: "PostgreSQL", stage, detail }) } } Err(_) => { let stage = classify_postgres_checkout_stage(pool.status()); log::warn!( "[db:pool.checkout:error] elapsed_ms={} timeout_ms={} stage={} error=checkout timed out", start.elapsed().as_millis(), checkout_timeout.as_millis(), stage.as_str() ); Err(super::PoolCheckoutError::Timeout { database: "PostgreSQL", stage, timeout: checkout_timeout }) } } }; match cancel_token { Some(token) => tokio::select! { biased; _ = token.cancelled() => { log::info!( "[db:pool.checkout:cancelled] elapsed_ms={} timeout_ms={}", start.elapsed().as_millis(), checkout_timeout.as_millis() ); Err(super::PoolCheckoutError::Canceled) } result = get_future => result, }, None => get_future.await, } } pub async fn checkout_postgres_client( pool: &Pool, cancel_token: Option<&CancellationToken>, checkout_timeout: Duration, ) -> Result { let checkout_timeout = effective_postgres_checkout_timeout(pool, checkout_timeout); let start = Instant::now(); let result = checkout_postgres_client_classified(pool, cancel_token, checkout_timeout) .await .map_err(|error| error.to_string()); if let Ok(client) = &result { log::debug!( "[db:pool.checkout:done] elapsed_ms={} timeout_ms={}", start.elapsed().as_millis(), checkout_timeout.as_millis() ); // Resolve the notice-attribution identity here, outside of any // transaction: a lazy first lookup from `drain_postgres_notices` // could run inside the read-only EXPLAIN transaction, where a // failing identity query would abort the user's statement. Cached // after the first checkout of each physical connection. let _ = resolve_postgres_client_key(client).await; } result } fn effective_postgres_checkout_timeout(pool: &Pool, requested_timeout: Duration) -> Duration { let configured = pool.timeouts(); [configured.wait, configured.create, configured.recycle] .into_iter() .flatten() .fold(requested_timeout, std::cmp::max) } async fn cancel_postgres_query( pg_cancel_token: tokio_postgres::CancelToken, cancel_context: Option<&PostgresCancelContext>, cancel_timeout: Duration, ) { let cancel_timeout = postgres_cancel_attempt_timeout(cancel_timeout, cancel_context); if let Some(ctx) = cancel_context { let tls_result = if ctx.legacy_tls { match make_openssl_connect_from_context(ctx) { Ok(tls) => Some(tokio::time::timeout(cancel_timeout, pg_cancel_token.cancel_query(tls)).await), Err(err) => { log::warn!("Failed to build legacy TLS connector for cancel: {err}; falling back to NoTls cancel"); None } } } else { match make_rustls_connect_from_context(ctx) { Ok(tls) => Some(tokio::time::timeout(cancel_timeout, pg_cancel_token.cancel_query(tls)).await), Err(err) => { log::warn!("Failed to build TLS connector for cancel: {err}; falling back to NoTls cancel"); None } } }; if let Some(result) = tls_result { match result { Ok(Ok(())) => return, Ok(Err(err)) => { log::warn!("Failed to send PostgreSQL TLS cancel request: {err}"); if ctx.ssl_mode != SslMode::Prefer { return; } } Err(_) => { log::warn!("Timed out sending PostgreSQL TLS cancel request ({}s)", cancel_timeout.as_secs()); if ctx.ssl_mode != SslMode::Prefer { return; } } } } } match tokio::time::timeout(cancel_timeout, pg_cancel_token.cancel_query(NoTls)).await { Ok(Ok(())) => {} Ok(Err(err)) => log::warn!("Failed to send PostgreSQL cancel request: {err}"), Err(_) => log::warn!("Timed out sending PostgreSQL cancel request ({}s)", cancel_timeout.as_secs()), } } fn postgres_cancel_attempt_timeout( cancel_timeout: Duration, _cancel_context: Option<&PostgresCancelContext>, ) -> Duration { cancel_timeout } fn is_transaction_recovery_statement(sql: &str) -> bool { starts_with_executable_sql_keyword(sql, &["ROLLBACK", "ABORT", "COMMIT", "END"]) } async fn execute_query_with_max_rows_inner( client: &deadpool_postgres::Client, sql: &str, max_rows: Option, prefer_text_protocol: bool, progress_clock: Option>, force_unnamed: bool, ) -> Result { let start = Instant::now(); let row_limit = query_result_row_limit(max_rows); // Discard stale notices from infrastructure statements (e.g. the timezone // SET issued at connect time) so only messages raised by this statement // are attached to its result. let _ = drain_postgres_notices(client).await; let result = if postgres_statement_returns_rows(sql) { if prefer_text_protocol { execute_select_text(client, sql, start, row_limit, None, progress_clock.as_deref()).await } else { execute_select_query_with_progress(client, sql, start, row_limit, progress_clock.as_deref(), force_unnamed) .await } } else { let affected = if force_unnamed { client.execute_typed(sql, &[]).await } else { client.execute(sql, &[]).await }; affected.map_err(pg_error_to_string).map(|affected| QueryResult { columns: vec![], column_types: Vec::new(), column_sortables: Vec::new(), spatial_columns: vec![], spatial_values: vec![], rows: vec![], affected_rows: affected, execution_time_ms: start.elapsed().as_millis(), truncated: false, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: Vec::new(), }) }; match result { Ok(mut result) => { result.messages = drain_postgres_notices(client).await; Ok(result) } Err(error) => { // Drop notices so an errored statement's messages cannot leak // into the next query on this pooled connection. let _ = drain_postgres_notices(client).await; Err(error) } } } // Sibling of `postgres_indexes_for_relations_sql` (~line 3288), for a single // (schema, table) instead of a batch of oids — see the note there. const POSTGRES_INDEXES_SQL: &str = "SELECT i.relname AS index_name, \ array_agg(COALESCE(a.attname, pg_get_indexdef(ix.indexrelid, k.n::int, false)) ORDER BY k.n) AS columns, \ array_agg(CASE WHEN oc.opcdefault THEN NULL ELSE quote_ident(opcns.nspname) || '.' || quote_ident(oc.opcname) END ORDER BY k.n) AS column_opclasses, \ (ix.indisunique AND ix.indisvalid) AS is_unique, \ ix.indisprimary AS is_primary, \ pg_get_expr(ix.indpred, ix.indrelid) AS filter_expr, \ am.amname AS index_type, \ ix.indnkeyatts AS nkeyatts, \ ix.indkey AS indkey, \ obj_description(i.oid, 'pg_class') AS index_comment, \ array_agg(a.attname IS NULL ORDER BY k.n) AS key_is_expression \ FROM pg_index ix \ JOIN pg_class t ON t.oid = ix.indrelid \ JOIN pg_class i ON i.oid = ix.indexrelid \ JOIN pg_namespace n ON n.oid = t.relnamespace \ JOIN pg_am am ON am.oid = i.relam \ JOIN LATERAL unnest(ix.indkey, ix.indclass) WITH ORDINALITY AS k(attnum, class_oid, n) ON true \ LEFT JOIN pg_attribute a ON a.attrelid = t.oid AND a.attnum = k.attnum AND k.attnum > 0 \ LEFT JOIN pg_opclass oc ON oc.oid = k.class_oid \ LEFT JOIN pg_namespace opcns ON opcns.oid = oc.opcnamespace \ WHERE t.oid = (CASE WHEN $1 = '' THEN quote_ident($2) ELSE quote_ident($1) || '.' || quote_ident($2) END)::regclass \ GROUP BY i.relname, i.oid, ix.indisunique, ix.indisvalid, ix.indisprimary, ix.indpred, ix.indrelid, am.amname, ix.indnkeyatts, ix.indkey \ ORDER BY i.relname"; // Compat-tier sibling of `postgres_indexes_for_relations_compat_sql` (~line // 3312) — see the note on `POSTGRES_INDEXES_SQL` above. const POSTGRES_INDEXES_COMPAT_SQL: &str = "SELECT i.relname AS index_name, \ ARRAY( \ SELECT COALESCE(a.attname, pg_get_indexdef(ix.indexrelid, pos.n, false)) \ FROM generate_series(1, array_length(string_to_array(ix.indkey::text, ' '), 1)) AS pos(n) \ LEFT JOIN pg_attribute a \ ON a.attrelid = t.oid \ AND a.attnum = (string_to_array(ix.indkey::text, ' '))[pos.n]::int2 \ AND a.attnum > 0 \ ORDER BY pos.n \ ) AS columns, \ ARRAY( \ SELECT CASE WHEN oc.opcdefault THEN NULL \ ELSE oc.opcname \ END \ FROM generate_series(1, array_length(string_to_array(ix.indkey::text, ' '), 1)) AS pos(n) \ LEFT JOIN pg_attribute a \ ON a.attrelid = t.oid \ AND a.attnum = (string_to_array(ix.indkey::text, ' '))[pos.n]::int2 \ AND a.attnum > 0 \ LEFT JOIN pg_opclass oc \ ON oc.oid = (string_to_array(ix.indclass::text, ' '))[pos.n]::oid \ ORDER BY pos.n \ ) AS column_opclasses, \ (ix.indisunique AND ix.indisvalid) AS is_unique, \ ix.indisprimary AS is_primary, \ pg_get_expr(ix.indpred, ix.indrelid) AS filter_expr, \ am.amname AS index_type, \ NULL::smallint AS nkeyatts, \ ix.indkey AS indkey, \ obj_description(i.oid, 'pg_class') AS index_comment, \ ARRAY( \ SELECT a.attname IS NULL \ FROM generate_series(1, array_length(string_to_array(ix.indkey::text, ' '), 1)) AS pos(n) \ LEFT JOIN pg_attribute a \ ON a.attrelid = t.oid \ AND a.attnum = (string_to_array(ix.indkey::text, ' '))[pos.n]::int2 \ AND a.attnum > 0 \ ORDER BY pos.n \ ) AS key_is_expression \ FROM pg_index ix \ JOIN pg_class t ON t.oid = ix.indrelid \ JOIN pg_class i ON i.oid = ix.indexrelid \ JOIN pg_namespace n ON n.oid = t.relnamespace \ JOIN pg_am am ON am.oid = i.relam \ WHERE t.oid = (CASE WHEN $1 = '' THEN quote_ident($2) ELSE quote_ident($1) || '.' || quote_ident($2) END)::regclass \ ORDER BY i.relname"; const POSTGRES_OWNERS_SQL: &str = "SELECT n.nspname, c.relname, c.relkind::text AS relkind, pg_get_userbyid(c.relowner) \ FROM pg_class c \ JOIN pg_namespace n ON n.oid = c.relnamespace \ WHERE n.nspname = $1 \ AND c.relkind IN ('r', 'v', 'm', 'S', 'f', 'p')"; const POSTGRES_TABLE_OWNER_SQL: &str = "SELECT pg_get_userbyid(c.relowner)::text, \ ARRAY(SELECT default_acl.privilege_type::text \ FROM pg_catalog.aclexplode(pg_catalog.acldefault('r', c.relowner)) default_acl \ WHERE default_acl.grantee = c.relowner \ ORDER BY default_acl.privilege_type) \ FROM pg_catalog.pg_class c \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ WHERE n.nspname = $1 AND c.relname = $2 AND c.relkind IN ('r', 'p') \ ORDER BY c.oid LIMIT 1"; const POSTGRES_TABLE_ACL_PRIVILEGES_SQL: &str = "SELECT CASE WHEN acl.grantee = 0 THEN 'PUBLIC' ELSE grantee.rolname END::text, \ acl.privilege_type::text, acl.is_grantable, pg_get_userbyid(acl.grantor)::text \ FROM pg_catalog.pg_class c \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ JOIN LATERAL pg_catalog.aclexplode(COALESCE(c.relacl, pg_catalog.acldefault('r', c.relowner))) acl ON true \ LEFT JOIN pg_catalog.pg_roles grantee ON grantee.oid = acl.grantee \ WHERE n.nspname = $1 AND c.relname = $2 AND c.relkind IN ('r', 'p') \ ORDER BY 4, 1, 2, 3"; const POSTGRES_COLUMN_ACL_PRIVILEGES_SQL: &str = "SELECT CASE WHEN acl.grantee = 0 THEN 'PUBLIC' ELSE grantee.rolname END::text, \ acl.privilege_type::text, acl.is_grantable, a.attname::text, pg_get_userbyid(acl.grantor)::text \ FROM pg_catalog.pg_class c \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ JOIN pg_catalog.pg_attribute a ON a.attrelid = c.oid AND a.attnum > 0 AND NOT a.attisdropped \ JOIN LATERAL pg_catalog.aclexplode(a.attacl) acl ON true \ LEFT JOIN pg_catalog.pg_roles grantee ON grantee.oid = acl.grantee \ WHERE n.nspname = $1 AND c.relname = $2 AND c.relkind IN ('r', 'p') \ ORDER BY 5, 1, 2, 3, 4"; pub(crate) fn postgres_owner_object_type(relkind: &str) -> &str { match relkind { "r" => "TABLE", "v" => "VIEW", "m" => "MATERIALIZED_VIEW", "S" => "SEQUENCE", "f" => "FOREIGN TABLE", "p" => "PARTITIONED TABLE", "I" => "PARTITIONED INDEX", _ => relkind, } } /// Looks up a relation's own `relkind` (e.g. `'r'`, `'v'`, `'S'`), regardless /// of whether it's a kind this driver otherwise treats as a table. Used to /// give a specific diagnostic ("it's a view, not a table") instead of a bare /// "not found" when a `(schema, table)` request turns out not to be a /// table/partition/foreign table. pub(crate) async fn postgres_relation_relkind( pool: &Pool, schema: &str, table: &str, ) -> Result, String> { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let row = client .query_opt( "SELECT c.relkind::text FROM pg_catalog.pg_class c \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ WHERE n.nspname = $1 AND c.relname = $2 LIMIT 1", &[&schema, &table], ) .await .map_err(pg_error_to_string)?; row.map(|row| row.try_get::<_, String>(0)).transpose().map_err(pg_error_to_string) } /// Compute per-column operator classes from `pg_opclass` joined via /// `pg_index.indclass`. For each key column position, the schema-qualified /// opclass (`quote_ident(nspname) || '.' || quote_ident(opcname)`) is returned /// unless the class is the type's default (`oc.opcdefault`), so DDL regeneration /// resolves the opclass regardless of `search_path`. This applies to every key /// position, including expression keys (`a.attname IS NULL`): the per-column /// `pg_get_indexdef(indexrelid, colno, pretty)` call returns only the bare /// expression text (PostgreSQL sets `attrsOnly = (colno != 0)`, so the /// opclass/COLLATE/DESC block is skipped — see `ruleutils.c`), so the opclass /// must be read from `indclass` rather than parsed out of that text. async fn list_indexes_with_sql( client: &deadpool_postgres::Client, sql: &str, schema: &str, table: &str, ) -> Result, tokio_postgres::Error> { let rows = postgres_query_cached(client, sql, &[&schema, &table]).await?; Ok(rows .iter() .map(|row| { let all_cols: Vec = row.try_get::<_, Vec>(1).unwrap_or_default(); let all_opclasses: Vec> = row.try_get::<_, Vec>>(2).unwrap_or_default(); let nkeyatts = row.try_get::<_, Option>(7).ok().flatten().unwrap_or(all_cols.len() as i16) as usize; let split_at = nkeyatts.min(all_cols.len()); let key_cols = all_cols[..split_at].to_vec(); let key_opclasses = if all_opclasses.len() == all_cols.len() { all_opclasses[..split_at].to_vec() } else { vec![None; split_at] }; let included = if split_at < all_cols.len() { all_cols[split_at..].to_vec() } else { vec![] }; // `a.attname IS NULL` at a given key position means that key part came back from // pg_get_indexdef (a functional/expression key part), not from a real column (#6295). let all_is_expr: Vec = row.try_get::<_, Vec>(10).unwrap_or_default(); let key_is_expression = if all_is_expr.len() == all_cols.len() { all_is_expr[..split_at].to_vec() } else { Vec::new() }; IndexInfo { name: pg_row_try_string(row, 0), columns: key_cols, is_unique: pg_row_try_bool(row, 3).unwrap_or(false), is_primary: pg_row_try_bool(row, 4).unwrap_or(false), filter: row.try_get::<_, Option>(5).ok().flatten(), index_type: row.try_get::<_, Option>(6).ok().flatten(), included_columns: if included.is_empty() { None } else { Some(included) }, comment: row.try_get::<_, Option>(9).ok().flatten(), key_is_expression, column_opclasses: key_opclasses, } }) .collect()) } pub async fn list_indexes(pool: &Pool, schema: &str, table: &str) -> Result, String> { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let tiers = [POSTGRES_INDEXES_SQL, POSTGRES_INDEXES_COMPAT_SQL]; query_with_compat_fallback("list_indexes", &tiers, |sql| list_indexes_with_sql(&client, sql, schema, table)).await } /// Names of same-table indexes whose `pg_index.indisvalid` is `false`. /// /// A cancelled `CREATE INDEX CONCURRENTLY` leaves an INVALID index behind with /// the requested name; until it is dropped, any retry of the same build fails /// with `relation already exists`. The structure editor consults this before /// applying a concurrent build so it can surface the leftover explicitly /// instead of failing silently. const POSTGRES_INVALID_INDEXES_SQL: &str = "SELECT idx.relname \ FROM pg_catalog.pg_index ix \ JOIN pg_catalog.pg_class idx ON idx.oid = ix.indexrelid \ JOIN pg_catalog.pg_class t ON t.oid = ix.indrelid \ JOIN pg_catalog.pg_namespace n ON n.oid = t.relnamespace \ WHERE n.nspname = $1 AND t.relname = $2 AND ix.indisvalid = false \ ORDER BY idx.relname"; pub async fn list_invalid_indexes(pool: &Pool, schema: &str, table: &str) -> Result, String> { let schema = if schema.is_empty() { "public" } else { schema }; let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let rows = postgres_query_cached(&client, POSTGRES_INVALID_INDEXES_SQL, &[&schema, &table]) .await .map_err(|e| e.to_string())?; Ok(rows.iter().map(|row| pg_row_try_string(row, 0)).collect()) } fn postgres_foreign_keys_sql() -> &'static str { "SELECT con.conname AS constraint_name, \ a.attname AS column_name, \ ref_n.nspname AS ref_schema, \ ref_c.relname AS ref_table, \ ref_a.attname AS ref_column, \ con.confupdtype::text AS on_update_raw, \ con.confdeltype::text AS on_delete_raw \ FROM pg_catalog.pg_constraint con \ JOIN pg_catalog.pg_class c ON c.oid = con.conrelid \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ JOIN pg_catalog.pg_class ref_c ON ref_c.oid = con.confrelid \ JOIN pg_catalog.pg_namespace ref_n ON ref_n.oid = ref_c.relnamespace \ JOIN LATERAL unnest(con.conkey) WITH ORDINALITY AS fk(attnum, ord) ON true \ JOIN pg_catalog.pg_attribute a ON a.attrelid = c.oid AND a.attnum = fk.attnum AND NOT a.attisdropped \ JOIN LATERAL unnest(con.confkey) WITH ORDINALITY AS pk(attnum, ord) ON pk.ord = fk.ord \ JOIN pg_catalog.pg_attribute ref_a ON ref_a.attrelid = ref_c.oid AND ref_a.attnum = pk.attnum AND NOT ref_a.attisdropped \ WHERE con.contype = 'f' AND n.nspname = $1 AND c.relname = $2 \ ORDER BY con.conname, fk.ord" } // Pre-9.4 sibling of `postgres_foreign_keys_sql`: WITH ORDINALITY requires // PostgreSQL 9.4, so pair conkey/confkey positions with generate_series over // the array length plus plain subscripts. fn postgres_foreign_keys_compat_sql() -> &'static str { "SELECT con.conname AS constraint_name, \ a.attname AS column_name, \ ref_n.nspname AS ref_schema, \ ref_c.relname AS ref_table, \ ref_a.attname AS ref_column, \ con.confupdtype::text AS on_update_raw, \ con.confdeltype::text AS on_delete_raw \ FROM pg_catalog.pg_constraint con \ JOIN pg_catalog.pg_class c ON c.oid = con.conrelid \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ JOIN pg_catalog.pg_class ref_c ON ref_c.oid = con.confrelid \ JOIN pg_catalog.pg_namespace ref_n ON ref_n.oid = ref_c.relnamespace \ CROSS JOIN generate_series(1, array_length(con.conkey, 1)) AS fk(ord) \ JOIN pg_catalog.pg_attribute a ON a.attrelid = c.oid AND a.attnum = (con.conkey)[fk.ord] AND NOT a.attisdropped \ JOIN pg_catalog.pg_attribute ref_a ON ref_a.attrelid = ref_c.oid AND ref_a.attnum = (con.confkey)[fk.ord] AND NOT ref_a.attisdropped \ WHERE con.contype = 'f' AND n.nspname = $1 AND c.relname = $2 \ ORDER BY con.conname, fk.ord" } async fn list_foreign_keys_with_sql( client: &deadpool_postgres::Client, sql: &'static str, schema: &str, table: &str, ) -> Result, tokio_postgres::Error> { let rows = postgres_query_cached(client, sql, &[&schema, &table]).await?; Ok(rows .iter() .map(|row| ForeignKeyInfo { name: pg_row_try_string(row, 0), column: pg_row_try_string(row, 1), ref_schema: Some(pg_row_try_string(row, 2)), ref_table: pg_row_try_string(row, 3), ref_column: pg_row_try_string(row, 4), on_update: postgres_fk_action_label(pg_row_try_optional_text(row, 5)), on_delete: postgres_fk_action_label(pg_row_try_optional_text(row, 6)), }) .collect()) } pub async fn list_foreign_keys(pool: &Pool, schema: &str, table: &str) -> Result, String> { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let tiers = [postgres_foreign_keys_sql(), postgres_foreign_keys_compat_sql()]; query_with_compat_fallback("list_foreign_keys", &tiers, |sql| { list_foreign_keys_with_sql(&client, sql, schema, table) }) .await } /// OpenGauss-safe foreign key metadata. Mirrors `list_opengauss_constraints`: /// older OpenGauss releases may reject WITH ORDINALITY or decode catalog arrays /// differently on the wire, so read conkey/confkey as text and resolve the /// attribute numbers in Rust. pub async fn list_opengauss_foreign_keys( pool: &Pool, schema: &str, table: &str, ) -> Result, String> { let schema = if schema.is_empty() { "public" } else { schema }; let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let rows = postgres_query_cached(&client, opengauss_foreign_keys_sql(), &[&schema, &table]) .await .map_err(|e| e.to_string())?; if rows.is_empty() { return Ok(Vec::new()); } let local_relation_oid = pg_row_try_u32(&rows[0], 7).ok_or("OpenGauss foreign key metadata did not return a local relation OID")?; let local_attributes = opengauss_relation_attributes(&client, local_relation_oid).await?; let mut referenced_attributes: HashMap> = HashMap::new(); let mut result = Vec::new(); for row in &rows { let name = pg_row_try_string(row, 0); let column_numbers = parse_opengauss_attribute_numbers(&pg_row_try_string(row, 3)) .map_err(|error| format!("failed to parse OpenGauss foreign key {name} columns: {error}"))?; let ref_column_numbers = parse_opengauss_attribute_numbers(&pg_row_try_string(row, 4)) .map_err(|error| format!("failed to parse OpenGauss foreign key {name} referenced columns: {error}"))?; let ref_schema = pg_row_try_string(row, 1); let ref_table = pg_row_try_string(row, 2); let on_update = postgres_fk_action_label(pg_row_try_optional_text(row, 5)); let on_delete = postgres_fk_action_label(pg_row_try_optional_text(row, 6)); let Some(referenced_relation_oid) = pg_row_try_u32(row, 8).filter(|oid| *oid != 0) else { continue; }; if let std::collections::hash_map::Entry::Vacant(entry) = referenced_attributes.entry(referenced_relation_oid) { entry.insert(opengauss_relation_attributes(&client, referenced_relation_oid).await?); } let ref_attributes = referenced_attributes.get(&referenced_relation_oid).unwrap(); for (fk_number, pk_number) in column_numbers.iter().zip(ref_column_numbers.iter()) { let Some(column) = local_attributes.get(fk_number) else { continue }; let Some(ref_column) = ref_attributes.get(pk_number) else { continue }; result.push(ForeignKeyInfo { name: name.clone(), column: column.clone(), ref_schema: Some(ref_schema.clone()), ref_table: ref_table.clone(), ref_column: ref_column.clone(), on_update: on_update.clone(), on_delete: on_delete.clone(), }); } } Ok(result) } fn opengauss_foreign_keys_sql() -> &'static str { "SELECT con.conname, \ ref_n.nspname, ref_c.relname, \ COALESCE(con.conkey::text, ''), COALESCE(con.confkey::text, ''), \ con.confupdtype::text, con.confdeltype::text, \ con.conrelid::oid, con.confrelid::oid \ FROM pg_catalog.pg_constraint con \ JOIN pg_catalog.pg_class c ON c.oid = con.conrelid \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ JOIN pg_catalog.pg_class ref_c ON ref_c.oid = con.confrelid \ JOIN pg_catalog.pg_namespace ref_n ON ref_n.oid = ref_c.relnamespace \ WHERE con.contype = 'f' AND n.nspname = $1 AND c.relname = $2 \ ORDER BY con.conname" } fn postgres_table_dependencies_sql() -> &'static str { // Foreign keys aren't the only ordering constraint on export/replay: a // partition must be created after its parent table exists too, so union // in `pg_inherits` partition-of edges alongside the FK edges. "SELECT child.relname AS table_name, parent.relname AS ref_table \ FROM pg_catalog.pg_constraint con \ JOIN pg_catalog.pg_class child ON child.oid = con.conrelid \ JOIN pg_catalog.pg_namespace child_schema ON child_schema.oid = child.relnamespace \ JOIN pg_catalog.pg_class parent ON parent.oid = con.confrelid \ JOIN pg_catalog.pg_namespace parent_schema ON parent_schema.oid = parent.relnamespace \ WHERE con.contype = 'f' \ AND child_schema.nspname = $1 \ AND parent_schema.nspname = $1 \ UNION \ SELECT child.relname AS table_name, parent.relname AS ref_table \ FROM pg_catalog.pg_inherits i \ JOIN pg_catalog.pg_class child ON child.oid = i.inhrelid AND child.relispartition \ JOIN pg_catalog.pg_namespace child_schema ON child_schema.oid = child.relnamespace \ JOIN pg_catalog.pg_class parent ON parent.oid = i.inhparent \ JOIN pg_catalog.pg_namespace parent_schema ON parent_schema.oid = parent.relnamespace \ WHERE child_schema.nspname = $1 \ AND parent_schema.nspname = $1 \ ORDER BY table_name, ref_table" } /// 9.x 没有 relispartition;INHERITS 子表同样要先建父表,所以保留这条边。 /// /// No relispartition pre-10; INHERITS children still need their parent /// created first, so the edge is kept without the filter. fn postgres_table_dependencies_compat_sql() -> &'static str { "SELECT child.relname AS table_name, parent.relname AS ref_table \ FROM pg_catalog.pg_constraint con \ JOIN pg_catalog.pg_class child ON child.oid = con.conrelid \ JOIN pg_catalog.pg_namespace child_schema ON child_schema.oid = child.relnamespace \ JOIN pg_catalog.pg_class parent ON parent.oid = con.confrelid \ JOIN pg_catalog.pg_namespace parent_schema ON parent_schema.oid = parent.relnamespace \ WHERE con.contype = 'f' \ AND child_schema.nspname = $1 \ AND parent_schema.nspname = $1 \ UNION \ SELECT child.relname AS table_name, parent.relname AS ref_table \ FROM pg_catalog.pg_inherits i \ JOIN pg_catalog.pg_class child ON child.oid = i.inhrelid \ JOIN pg_catalog.pg_namespace child_schema ON child_schema.oid = child.relnamespace \ JOIN pg_catalog.pg_class parent ON parent.oid = i.inhparent \ JOIN pg_catalog.pg_namespace parent_schema ON parent_schema.oid = parent.relnamespace \ WHERE child_schema.nspname = $1 \ AND parent_schema.nspname = $1 \ ORDER BY table_name, ref_table" } /// Fetch all same-schema table dependencies in one round trip. Whole-database /// exports use this instead of issuing one information_schema query per table. pub async fn list_table_dependencies(pool: &Pool, schema: &str) -> Result, String> { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let params: [&(dyn tokio_postgres::types::ToSql + Sync); 1] = [&schema]; let rows = query_with_compat_fallback( "list_table_dependencies", &[postgres_table_dependencies_sql(), postgres_table_dependencies_compat_sql()], |sql| postgres_query_cached(&client, sql, ¶ms), ) .await?; Ok(rows.iter().map(|row| (pg_row_try_string(row, 0), pg_row_try_string(row, 1))).collect()) } pub async fn list_triggers(pool: &Pool, schema: &str, table: &str) -> Result, String> { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let rows = postgres_query_cached( &client, "SELECT trigger_name, event_manipulation, action_timing \ FROM information_schema.triggers \ WHERE trigger_schema = $1 AND event_object_table = $2 \ ORDER BY trigger_name", &[&schema, &table], ) .await .map_err(|e| e.to_string())?; Ok(rows .iter() .map(|row| TriggerInfo { name: pg_row_try_string(row, 0), event: pg_row_try_string(row, 1), timing: pg_row_try_string(row, 2), level: None, condition: None, language: None, enabled: None, valid: None, comment: None, created_at: None, statement: None, }) .collect()) } // A trigger declared on a partitioned table is automatically cloned onto each // partition (tracked via pg_trigger.tgparentid since PostgreSQL 13); such // clones are not independently creatable statements — attaching/creating the // partition already recreates them — so they must be excluded here, or the // partition's own DDL would re-declare (and fail to create, or duplicate) a // trigger the parent's DDL already installs. PostgreSQL servers older than 13 // lack the tgparentid column (probed via postgres_trigger_has_tgparentid), so // use the pre-13 query there; those servers never clone triggers onto // partitions, so tgisinternal alone is sufficient there. pub async fn list_trigger_definitions(pool: &Pool, schema: &str, table: &str) -> Result, String> { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let sql = if postgres_trigger_has_tgparentid(&client).await? { postgres_trigger_definitions_sql() } else { postgres_trigger_definitions_sql_without_tgparentid() }; let rows = postgres_query_cached(&client, sql, &[&schema, &table]).await.map_err(|e| e.to_string())?; Ok(rows.iter().map(|row| pg_row_try_string(row, 0)).filter(|definition| !definition.trim().is_empty()).collect()) } fn postgres_trigger_definitions_sql() -> &'static str { "SELECT pg_catalog.pg_get_triggerdef(t.oid, true) AS trigger_definition \ FROM pg_catalog.pg_trigger t \ JOIN pg_catalog.pg_class c ON c.oid = t.tgrelid \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ WHERE n.nspname = $1 AND c.relname = $2 AND NOT t.tgisinternal AND t.tgparentid = 0 \ ORDER BY t.tgname, t.oid" } fn postgres_trigger_definitions_sql_without_tgparentid() -> &'static str { "SELECT pg_catalog.pg_get_triggerdef(t.oid, true) AS trigger_definition \ FROM pg_catalog.pg_trigger t \ JOIN pg_catalog.pg_class c ON c.oid = t.tgrelid \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ WHERE n.nspname = $1 AND c.relname = $2 AND NOT t.tgisinternal \ ORDER BY t.tgname, t.oid" } fn postgres_functions_sql(has_proc_prokind: bool) -> &'static str { if has_proc_prokind { return "SELECT p.proname, \ CASE p.prokind WHEN 'f' THEN 'FUNCTION' WHEN 'p' THEN 'PROCEDURE' ELSE 'FUNCTION' END, \ COALESCE(pg_get_function_result(p.oid), ''), \ pg_get_functiondef(p.oid), \ COALESCE(pg_get_function_arguments(p.oid), '') \ FROM pg_proc p \ JOIN pg_namespace n ON n.oid = p.pronamespace \ WHERE n.nspname = $1 AND p.prokind IN ('f', 'p') \ ORDER BY p.proname"; } // PostgreSQL 10 and older do not have pg_proc.prokind; procedures were // introduced with prokind, so the legacy path can only return functions. "SELECT p.proname, \ 'FUNCTION', \ COALESCE(pg_get_function_result(p.oid), ''), \ pg_get_functiondef(p.oid), \ COALESCE(pg_get_function_arguments(p.oid), '') \ FROM pg_proc p \ JOIN pg_namespace n ON n.oid = p.pronamespace \ WHERE n.nspname = $1 AND NOT p.proisagg AND NOT p.proiswindow \ ORDER BY p.proname" } pub async fn list_functions(pool: &Pool, schema: &str) -> Result, String> { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; // Use pg_proc + pg_get_functiondef() instead of information_schema.routines // for reliable function definition retrieval (information_schema.routines.routine_definition // is NULL for non-SQL functions like plpgsql) let has_proc_prokind = postgres_proc_has_prokind(&client).await?; let rows = postgres_query_cached(&client, postgres_functions_sql(has_proc_prokind), &[&schema]) .await .map_err(|e| e.to_string())?; Ok(rows .iter() .map(|row| { let def: String = pg_row_try_string(row, 3); // Remove schema qualification from CREATE FUNCTION statement // to avoid false differences when comparing across schemas. // Handle both "schema.name" and schema.name formats. let normalized_def = def .replace(&format!("CREATE OR REPLACE FUNCTION \"{}\".", schema), "CREATE OR REPLACE FUNCTION ") .replace(&format!("CREATE OR REPLACE FUNCTION {}.", schema), "CREATE OR REPLACE FUNCTION "); FunctionInfo { name: pg_row_try_string(row, 0), function_type: pg_row_try_string(row, 1), data_type: pg_row_try_string(row, 2), definition: normalized_def, arguments: pg_row_try_string(row, 4), } }) .collect()) } fn postgres_sequences_sql() -> &'static str { "SELECT c.relname, \ COALESCE(format_type(s.seqtypid, NULL), 'bigint'), \ COALESCE(s.seqstart::text, '1'), \ COALESCE(s.seqmin::text, '1'), \ COALESCE(s.seqmax::text, '9223372036854775807'), \ COALESCE(s.seqincrement::text, '1'), \ CASE WHEN s.seqcycle THEN 'YES' ELSE 'NO' END \ FROM pg_class c \ JOIN pg_namespace n ON n.oid = c.relnamespace \ LEFT JOIN pg_sequence s ON s.seqrelid = c.oid \ WHERE c.relkind = 'S' AND n.nspname = $1 \ ORDER BY c.relname" } fn opengauss_sequences_sql() -> &'static str { "SELECT s.sequence_name, \ COALESCE(s.data_type::text, 'bigint'), \ COALESCE(s.start_value::text, '1'), \ COALESCE(s.minimum_value::text, '1'), \ COALESCE(s.maximum_value::text, '9223372036854775807'), \ COALESCE(s.increment::text, '1'), \ COALESCE(s.cycle_option::text, 'NO') \ FROM information_schema.sequences s \ JOIN pg_namespace n ON n.nspname = s.sequence_schema \ JOIN pg_class c ON c.relnamespace = n.oid AND c.relname = s.sequence_name \ WHERE s.sequence_schema = $1 AND c.relkind IN ('S','L','z','Z') \ ORDER BY s.sequence_name" } fn postgres_sequence_last_values_sql() -> &'static str { "SELECT c.relname, pg_sequence_last_value(c.oid)::text \ FROM pg_class c \ JOIN pg_namespace n ON n.oid = c.relnamespace \ WHERE c.relkind = 'S' AND n.nspname = $1" } fn opengauss_sequence_last_values_sql() -> &'static str { "SELECT c.relname, (pg_sequence_last_value(c.oid)).last_value::text \ FROM pg_class c \ JOIN pg_namespace n ON n.oid = c.relnamespace \ WHERE c.relkind IN ('S','L','z','Z') AND n.nspname = $1" } async fn list_sequences_with_sql( pool: &Pool, schema: &str, with_last_values: bool, metadata_sql: &str, last_values_sql: &str, ) -> Result, String> { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let rows = postgres_query_cached(&client, metadata_sql, &[&schema]).await.map_err(|e| e.to_string())?; let mut sequences: Vec = rows .iter() .map(|row| SequenceInfo { name: pg_row_try_string(row, 0), data_type: pg_row_try_string(row, 1), start_value: pg_row_try_string(row, 2), min_value: pg_row_try_string(row, 3), max_value: pg_row_try_string(row, 4), increment: pg_row_try_string(row, 5), cycle: pg_row_try_string(row, 6) == "YES", last_value: None, }) .collect(); if with_last_values { if let Ok(rows) = postgres_query_cached(&client, last_values_sql, &[&schema]).await { for row in rows { let name: String = pg_row_try_string(&row, 0); if let Ok(Some(value)) = row.try_get::<_, Option>(1) { if let Some(seq) = sequences.iter_mut().find(|s| s.name == name) { seq.last_value = Some(value); } } } } } Ok(sequences) } pub async fn list_sequences(pool: &Pool, schema: &str, with_last_values: bool) -> Result, String> { // PostgreSQL 10+ stores sequence properties in pg_sequence. list_sequences_with_sql( pool, schema, with_last_values, postgres_sequences_sql(), postgres_sequence_last_values_sql(), ) .await } pub async fn list_opengauss_sequences( pool: &Pool, schema: &str, with_last_values: bool, ) -> Result, String> { // openGauss does not expose PostgreSQL 10's pg_sequence catalog. Its // information_schema view contains the portable sequence properties, while // pg_sequence_last_value returns a record rather than a scalar. list_sequences_with_sql( pool, schema, with_last_values, opengauss_sequences_sql(), opengauss_sequence_last_values_sql(), ) .await } pub async fn list_rules(pool: &Pool, schema: &str) -> Result, String> { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let rows = postgres_query_cached( &client, "SELECT schemaname, tablename, rulename, definition \ FROM pg_rules \ WHERE schemaname = $1 \ ORDER BY rulename", &[&schema], ) .await .map_err(|e| e.to_string())?; Ok(rows .iter() .map(|row| RuleInfo { name: pg_row_try_string(row, 2), table_name: pg_row_try_string(row, 1), definition: pg_row_try_string(row, 3), }) .collect()) } pub async fn list_extensions(pool: &Pool, schema: Option<&str>) -> Result, String> { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let rows = if let Some(schema) = schema.filter(|value| !value.is_empty()) { postgres_query_cached( &client, "SELECT e.extname, COALESCE(e.extversion, '') AS extversion, d.description, n.nspname \ FROM pg_catalog.pg_extension e \ JOIN pg_catalog.pg_namespace n ON n.oid = e.extnamespace \ LEFT JOIN pg_catalog.pg_description d ON d.objoid = e.oid AND d.classoid = 'pg_extension'::regclass \ WHERE n.nspname = $1 \ ORDER BY e.extname", &[&schema], ) .await } else { postgres_query_cached( &client, "SELECT e.extname, COALESCE(e.extversion, '') AS extversion, d.description, n.nspname \ FROM pg_catalog.pg_extension e \ JOIN pg_catalog.pg_namespace n ON n.oid = e.extnamespace \ LEFT JOIN pg_catalog.pg_description d ON d.objoid = e.oid AND d.classoid = 'pg_extension'::regclass \ ORDER BY n.nspname, e.extname", &[], ) .await } .map_err(|e| e.to_string())?; Ok(rows .iter() .map(|row| ExtensionInfo { name: pg_row_try_string(row, 0), version: pg_row_try_string(row, 1), comment: row.try_get::<_, Option>(2).ok().flatten().filter(|s| !s.is_empty()), schema: row.try_get::<_, Option>(3).ok().flatten().filter(|s| !s.is_empty()), }) .collect()) } fn list_extension_member_objects_sql() -> &'static str { "SELECT 'RELATION'::text AS object_kind, c.relname, ''::text AS signature \ FROM pg_catalog.pg_class c \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ WHERE n.nspname = $1 \ AND EXISTS ( \ SELECT 1 FROM pg_catalog.pg_depend d \ WHERE d.classid = 'pg_catalog.pg_class'::regclass \ AND d.objid = c.oid \ AND d.refclassid = 'pg_catalog.pg_extension'::regclass \ AND d.deptype = 'e' \ ) \ UNION ALL \ SELECT 'FUNCTION'::text, p.proname, pg_get_function_identity_arguments(p.oid) \ FROM pg_catalog.pg_proc p \ JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \ WHERE n.nspname = $1 \ AND EXISTS ( \ SELECT 1 FROM pg_catalog.pg_depend d \ WHERE d.classid = 'pg_catalog.pg_proc'::regclass \ AND d.objid = p.oid \ AND d.refclassid = 'pg_catalog.pg_extension'::regclass \ AND d.deptype = 'e' \ )" } pub async fn list_extension_member_objects(pool: &Pool, schema: &str) -> Result, String> { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let rows = match postgres_query_cached(&client, list_extension_member_objects_sql(), &[&schema]).await { Ok(rows) => rows, Err(primary_error) => { // PostgreSQL-compatible servers before the identity-argument // formatter can still be filtered using their legacy formatter. let fallback_sql = list_extension_member_objects_sql() .replace("pg_get_function_identity_arguments(p.oid)", "pg_get_function_arguments(p.oid)"); postgres_query_cached(&client, &fallback_sql, &[&schema]) .await .map_err(|fallback_error| format!("{primary_error}; legacy fallback failed: {fallback_error}"))? } }; Ok(rows .iter() .map(|row| (pg_row_try_string(row, 0), pg_row_try_string(row, 1), pg_row_try_string(row, 2))) .collect()) } pub async fn list_available_extensions(pool: &Pool) -> Result, String> { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let rows = postgres_query_cached( &client, "SELECT name, default_version, comment \ FROM pg_catalog.pg_available_extensions \ WHERE installed_version IS NULL \ ORDER BY name", &[], ) .await .map_err(|e| e.to_string())?; Ok(rows .iter() .map(|row| ExtensionInfo { name: pg_row_try_string(row, 0), version: pg_row_try_string(row, 1), comment: row.try_get::<_, Option>(2).ok().flatten().filter(|s| !s.is_empty()), schema: None, }) .collect()) } pub async fn list_owners(pool: &Pool, schema: &str) -> Result, String> { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let rows = postgres_query_cached(&client, POSTGRES_OWNERS_SQL, &[&schema]).await.map_err(|e| e.to_string())?; Ok(rows .iter() .map(|row| { let relkind: String = pg_row_try_string(row, 2); OwnerInfo { object_name: pg_row_try_string(row, 1), object_type: postgres_owner_object_type(&relkind).to_string(), owner: pg_row_try_string(row, 3), } }) .collect()) } pub async fn get_table_owner(pool: &Pool, schema: &str, table: &str) -> Result, String> { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let params: [&(dyn tokio_postgres::types::ToSql + Sync); 2] = [&schema, &table]; let rows = postgres_query_cached(&client, POSTGRES_TABLE_OWNER_SQL, ¶ms).await.map_err(pg_error_to_string)?; Ok(rows.first().map(|row| pg_row_try_string(row, 0)).filter(|owner| !owner.is_empty())) } pub async fn get_table_access(pool: &Pool, schema: &str, table: &str) -> Result { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let params: [&(dyn tokio_postgres::types::ToSql + Sync); 2] = [&schema, &table]; let owner_rows = postgres_query_cached(&client, POSTGRES_TABLE_OWNER_SQL, ¶ms).await.map_err(pg_error_to_string)?; let owner_row = owner_rows.first().ok_or_else(|| "Table owner not found".to_string())?; let owner = pg_row_try_string(owner_row, 0); if owner.is_empty() { return Err("Table owner not found".to_string()); } let owner_default_privileges = owner_row.try_get::<_, Vec>(1).unwrap_or_default(); if owner_default_privileges.is_empty() { return Err("Table owner default privileges are unavailable".to_string()); } let (table_privileges, column_privileges) = tokio::try_join!( postgres_query_cached(&client, POSTGRES_TABLE_ACL_PRIVILEGES_SQL, ¶ms), postgres_query_cached(&client, POSTGRES_COLUMN_ACL_PRIVILEGES_SQL, ¶ms), ) .map_err(pg_error_to_string)?; let privileges = table_privileges .iter() .map(|row| PostgresTablePrivilegeInfo { grantor: pg_row_try_string(row, 3), grantee: pg_row_try_string(row, 0), privilege_type: pg_row_try_string(row, 1), is_grantable: pg_row_try_bool(row, 2).unwrap_or(false), column_name: None, }) .chain(column_privileges.iter().map(|row| PostgresTablePrivilegeInfo { grantor: pg_row_try_string(row, 4), grantee: pg_row_try_string(row, 0), privilege_type: pg_row_try_string(row, 1), is_grantable: pg_row_try_bool(row, 2).unwrap_or(false), column_name: Some(pg_row_try_string(row, 3)), })) .collect::>(); if privileges.iter().any(|privilege| { privilege.grantor.is_empty() || privilege.grantee.is_empty() || privilege.privilege_type.is_empty() }) { return Err("Table ACL metadata is incomplete".to_string()); } Ok(PostgresTableAccessInfo { owner, owner_default_privileges, privileges }) } /// Execute multiple SQL statements in a single round-trip using batch_execute. /// Best for DDL scripts where per-statement affected-row counts are not needed. pub async fn execute_batch(pool: &Pool, statements: &[String]) -> Result<(), String> { let combined = statements.iter().map(|s| s.trim()).filter(|s| !s.is_empty()).collect::>().join(";\n"); if combined.is_empty() { return Ok(()); } let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; client.batch_execute(&combined).await.map_err(pg_error_to_string)?; clear_postgres_caches_after_ddl(pool, Some(&client), &combined); Ok(()) } pub async fn terminate_current_user_database_backends(pool: &Pool, database: &str) -> Result { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; client .execute( "SELECT pg_terminate_backend(pid) \ FROM pg_stat_activity \ WHERE datname = $1 \ AND pid <> pg_backend_pid() \ AND usename = current_user", &[&database], ) .await .map_err(pg_error_to_string) } fn clear_postgres_caches_after_ddl(pool: &Pool, client: Option<&deadpool_postgres::Client>, sql: &str) { if !invalidates_postgres_statement_cache(sql) { return; } pool.manager().statement_caches.clear(); if let Some(client) = client { client.clear_type_cache(); } } fn invalidates_postgres_statement_cache(sql: &str) -> bool { let trimmed = sql.trim_start(); starts_with_executable_sql_keyword( trimmed, &["ALTER", "CREATE", "DROP", "TRUNCATE", "COMMENT", "REINDEX", "VACUUM"], ) } /// Export data via COPY TO STDOUT. `sql` must be a complete COPY statement, e.g. /// `COPY table (col1, col2) TO STDOUT (FORMAT CSV, HEADER)`. /// Returns the raw COPY output bytes. pub async fn copy_out(pool: &Pool, sql: &str) -> Result, String> { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let stream = client.copy_out(sql).await.map_err(pg_error_to_string)?; tokio::pin!(stream); let mut result = Vec::new(); while let Some(chunk) = stream.next().await { result.extend_from_slice(&chunk.map_err(pg_error_to_string)?); } Ok(result) } /// Import data via COPY FROM STDIN. `sql` must be a complete COPY statement, e.g. /// `COPY table (col1, col2) FROM STDIN (FORMAT CSV)`. /// `data` is the raw input in the format specified by the COPY command. pub async fn copy_in(pool: &Pool, sql: &str, data: &[u8]) -> Result<(), String> { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let sink = client.copy_in::(sql).await.map_err(pg_error_to_string)?; let mut sink = Box::pin(sink); sink.as_mut().send(bytes::Bytes::copy_from_slice(data)).await.map_err(pg_error_to_string)?; sink.as_mut().close().await.map_err(pg_error_to_string) } #[cfg(test)] mod tests { use super::*; use openssl::asn1::Asn1Time; use openssl::bn::{BigNum, MsbOption}; use openssl::hash::MessageDigest; use openssl::pkey::PKey; use openssl::rsa::Rsa; use openssl::ssl::{Ssl, SslAcceptor, SslMethod, SslVersion}; use openssl::x509::extension::{BasicConstraints, KeyUsage, SubjectAlternativeName}; use openssl::x509::{X509NameBuilder, X509}; use std::cell::Cell; use std::pin::Pin; use std::process::Command; use std::time::Instant; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::TcpListener; use tokio_postgres::types::FromSql; fn postgres_numeric_binary(weight: i16, sign: u16, scale: u16, digits: &[u16]) -> Vec { let mut raw = Vec::with_capacity(8 + digits.len() * 2); raw.extend_from_slice(&u16::try_from(digits.len()).unwrap().to_be_bytes()); raw.extend_from_slice(&weight.to_be_bytes()); raw.extend_from_slice(&sign.to_be_bytes()); raw.extend_from_slice(&scale.to_be_bytes()); for digit in digits { raw.extend_from_slice(&digit.to_be_bytes()); } raw } fn postgres_numeric_array_binary(values: &[Option>]) -> Vec { let mut raw = Vec::new(); raw.extend_from_slice(&1_i32.to_be_bytes()); raw.extend_from_slice(&i32::from(values.iter().any(Option::is_none)).to_be_bytes()); raw.extend_from_slice(&Type::NUMERIC.oid().to_be_bytes()); raw.extend_from_slice(&i32::try_from(values.len()).unwrap().to_be_bytes()); raw.extend_from_slice(&1_i32.to_be_bytes()); for value in values { if let Some(value) = value { raw.extend_from_slice(&i32::try_from(value.len()).unwrap().to_be_bytes()); raw.extend_from_slice(value); } else { raw.extend_from_slice(&(-1_i32).to_be_bytes()); } } raw } #[test] fn postgres_numeric_binary_values_preserve_arbitrary_precision() { assert_eq!(decode_pg_numeric_bytes(&postgres_numeric_binary(0, 0x0000, 0, &[])).as_deref(), Some("0")); assert_eq!(decode_pg_numeric_bytes(&postgres_numeric_binary(0, 0x0000, 2, &[])).as_deref(), Some("0.00")); assert_eq!( decode_pg_numeric_bytes(&postgres_numeric_binary(1, 0x0000, 4, &[1, 2345, 6789])).as_deref(), Some("12345.6789") ); assert_eq!( decode_pg_numeric_bytes(&postgres_numeric_binary(-1, 0x4000, 7, &[12, 3000])).as_deref(), Some("-0.0012300") ); let tiny = decode_pg_numeric_bytes(&postgres_numeric_binary(-77, 0x0000, 307, &[10])).unwrap(); assert_eq!(tiny, format!("0.{}1", "0".repeat(306))); assert_eq!(tiny.len(), 309); let mut huge_digits = vec![9]; huge_digits.extend(std::iter::repeat_n(9999, 77)); let huge = "9".repeat(309); assert_eq!( decode_pg_numeric_bytes(&postgres_numeric_binary(77, 0x0000, 0, &huge_digits)).as_deref(), Some(huge.as_str()) ); let decoded = PgNumeric::from_sql(&Type::NUMERIC, &postgres_numeric_binary(-77, 0x0000, 307, &[10])).unwrap(); assert_eq!(decoded.0, tiny); assert!(PgNumeric::accepts(&Type::NUMERIC)); assert!(!PgNumeric::accepts(&Type::INT8)); let array_raw = postgres_numeric_array_binary(&[ Some(postgres_numeric_binary(-77, 0x0000, 307, &[10])), None, Some(postgres_numeric_binary(77, 0x0000, 0, &huge_digits)), ]); let array = Vec::>::from_sql(&Type::NUMERIC_ARRAY, &array_raw).unwrap(); assert_eq!( array.into_iter().map(|value| value.map(|value| value.0)).collect::>(), vec![Some(tiny), None, Some(huge)] ); } #[test] fn postgres_numeric_binary_values_preserve_special_values_and_reject_malformed_payloads() { assert_eq!(decode_pg_numeric_bytes(&postgres_numeric_binary(0, 0xC000, 0, &[])).as_deref(), Some("NaN")); assert_eq!(decode_pg_numeric_bytes(&postgres_numeric_binary(0, 0xD000, 0, &[])).as_deref(), Some("Infinity")); assert_eq!(decode_pg_numeric_bytes(&postgres_numeric_binary(0, 0xF000, 0, &[])).as_deref(), Some("-Infinity")); assert!(decode_pg_numeric_bytes(&[]).is_none()); assert!(decode_pg_numeric_bytes(&postgres_numeric_binary(0, 0x8000, 0, &[])).is_none()); assert!(decode_pg_numeric_bytes(&postgres_numeric_binary(0, 0x0000, 0, &[10_000])).is_none()); let mut trailing_bytes = postgres_numeric_binary(0, 0x0000, 0, &[1]); trailing_bytes.push(0); assert!(decode_pg_numeric_bytes(&trailing_bytes).is_none()); } #[test] fn postgres_money_binary_values_are_decoded_with_two_decimal_places() { assert_eq!(format_pg_money(12345), "123.45"); assert_eq!(format_pg_money(-5), "-0.05"); assert_eq!(format_pg_money(0), "0.00"); let decoded = PgMoney::from_sql(&Type::MONEY, &(-12345_i64).to_be_bytes()).unwrap(); assert_eq!(format_pg_money(decoded.0), "-123.45"); } fn postgres_error_response(message: &str) -> Vec { let mut payload = Vec::new(); payload.extend_from_slice(b"SERROR\0"); payload.extend_from_slice(b"C0A000\0"); payload.push(b'M'); payload.extend_from_slice(message.as_bytes()); payload.extend_from_slice(b"\0\0"); let mut response = vec![b'E']; response.extend_from_slice(&u32::try_from(payload.len() + 4).unwrap().to_be_bytes()); response.extend_from_slice(&payload); response.extend_from_slice(&[b'Z', 0, 0, 0, 5, b'I']); response } async fn serve_tls_handshake_failure_then_plaintext(listener: TcpListener) { let (mut tls_socket, _) = listener.accept().await.unwrap(); let mut ssl_request = [0_u8; 8]; tls_socket.read_exact(&mut ssl_request).await.unwrap(); assert_eq!(ssl_request, [0, 0, 0, 8, 4, 210, 22, 47]); tls_socket.write_all(b"S").await.unwrap(); tls_socket.write_all(&[0x15, 0x03, 0x03, 0, 2, 2, 40]).await.unwrap(); drop(tls_socket); let (mut plain_socket, _) = listener.accept().await.unwrap(); let startup_len = plain_socket.read_u32().await.unwrap(); assert!(startup_len >= 8); let mut startup = vec![0_u8; startup_len as usize - 4]; plain_socket.read_exact(&mut startup).await.unwrap(); assert_eq!(&startup[..4], &[0, 3, 0, 0]); plain_socket .write_all(&[ b'R', 0, 0, 0, 8, 0, 0, 0, 0, // AuthenticationOk b'K', 0, 0, 0, 12, 0, 0, 0, 1, 0, 0, 0, 2, // BackendKeyData b'Z', 0, 0, 0, 5, b'I', // ReadyForQuery ]) .await .unwrap(); let mut request = [0_u8; 1024]; let read = plain_socket.read(&mut request).await.unwrap(); assert!(read > 0, "client should issue the best-effort identity query"); plain_socket.write_all(&postgres_error_response("identity probe unavailable")).await.unwrap(); } fn tls12_rsa_certificate_acceptor(cipher_list: &str) -> (SslAcceptor, Vec) { let key = PKey::from_rsa(Rsa::generate(2048).unwrap()).unwrap(); let mut name = X509NameBuilder::new().unwrap(); name.append_entry_by_text("CN", "localhost").unwrap(); let name = name.build(); let mut serial = BigNum::new().unwrap(); serial.rand(64, MsbOption::MAYBE_ZERO, false).unwrap(); let serial = serial.to_asn1_integer().unwrap(); let mut certificate = X509::builder().unwrap(); certificate.set_version(2).unwrap(); certificate.set_serial_number(&serial).unwrap(); certificate.set_subject_name(&name).unwrap(); certificate.set_issuer_name(&name).unwrap(); certificate.set_pubkey(&key).unwrap(); certificate.set_not_before(Asn1Time::days_from_now(0).unwrap().as_ref()).unwrap(); certificate.set_not_after(Asn1Time::days_from_now(1).unwrap().as_ref()).unwrap(); certificate.append_extension(BasicConstraints::new().critical().ca().build().unwrap()).unwrap(); certificate .append_extension(KeyUsage::new().digital_signature().key_encipherment().key_cert_sign().build().unwrap()) .unwrap(); let subject_alt_name = SubjectAlternativeName::new().dns("localhost").build(&certificate.x509v3_context(None, None)).unwrap(); certificate.append_extension(subject_alt_name).unwrap(); certificate.sign(&key, MessageDigest::sha256()).unwrap(); let certificate = certificate.build(); let certificate_pem = certificate.to_pem().unwrap(); let mut acceptor = SslAcceptor::mozilla_intermediate(SslMethod::tls_server()).unwrap(); acceptor.set_min_proto_version(Some(SslVersion::TLS1_2)).unwrap(); acceptor.set_max_proto_version(Some(SslVersion::TLS1_2)).unwrap(); acceptor.set_cipher_list(cipher_list).unwrap(); acceptor.set_private_key(&key).unwrap(); acceptor.set_certificate(&certificate).unwrap(); acceptor.check_private_key().unwrap(); (acceptor.build(), certificate_pem) } async fn accept_tls12_with_acceptor( mut socket: tokio::net::TcpStream, acceptor: &SslAcceptor, ) -> Result, openssl::ssl::Error> { let mut ssl_request = [0_u8; 8]; socket.read_exact(&mut ssl_request).await.unwrap(); assert_eq!(ssl_request, [0, 0, 0, 8, 4, 210, 22, 47]); socket.write_all(b"S").await.unwrap(); let ssl = Ssl::new(acceptor.context()).unwrap(); let mut tls_socket = tokio_openssl::SslStream::new(ssl, socket).unwrap(); Pin::new(&mut tls_socket).accept().await?; Ok(tls_socket) } async fn accept_tls12( socket: tokio::net::TcpStream, cipher_list: &str, ) -> Result, openssl::ssl::Error> { let (acceptor, _) = tls12_rsa_certificate_acceptor(cipher_list); accept_tls12_with_acceptor(socket, &acceptor).await } async fn authenticate_tls12_postgres(tls_socket: &mut tokio_openssl::SslStream) { let startup_len = tls_socket.read_u32().await.unwrap(); assert!(startup_len >= 8); let mut startup = vec![0_u8; startup_len as usize - 4]; tls_socket.read_exact(&mut startup).await.unwrap(); assert_eq!(&startup[..4], &[0, 3, 0, 0]); tls_socket .write_all(&[ b'R', 0, 0, 0, 8, 0, 0, 0, 0, // AuthenticationOk b'K', 0, 0, 0, 12, 0, 0, 0, 1, 0, 0, 0, 2, // BackendKeyData b'Z', 0, 0, 0, 5, b'I', // ReadyForQuery ]) .await .unwrap(); let mut request = [0_u8; 1024]; let read = tls_socket.read(&mut request).await.unwrap(); assert!(read > 0, "client should issue the best-effort identity query"); tls_socket.write_all(&postgres_error_response("identity probe unavailable")).await.unwrap(); } async fn serve_tls12_postgres(listener: TcpListener, cipher_list: &'static str, expect_handshake: bool) { let (socket, _) = listener.accept().await.unwrap(); let handshake = accept_tls12(socket, cipher_list).await; if !expect_handshake { assert!(handshake.is_err(), "rustls unexpectedly negotiated a static RSA cipher suite"); return; } let mut tls_socket = handshake.unwrap(); authenticate_tls12_postgres(&mut tls_socket).await; } async fn serve_tls12_postgres_with_acceptor(listener: TcpListener, acceptor: SslAcceptor, expect_handshake: bool) { let (socket, _) = listener.accept().await.unwrap(); let handshake = accept_tls12_with_acceptor(socket, &acceptor).await; if !expect_handshake { assert!(handshake.is_err(), "TLS handshake unexpectedly succeeded"); return; } let mut tls_socket = handshake.unwrap(); authenticate_tls12_postgres(&mut tls_socket).await; } async fn serve_static_rsa_postgres_with_cancel(listener: TcpListener) { let (main_socket, _) = listener.accept().await.unwrap(); let mut main_tls = accept_tls12(main_socket, "AES128-GCM-SHA256").await.unwrap(); authenticate_tls12_postgres(&mut main_tls).await; let (cancel_socket, _) = listener.accept().await.unwrap(); let mut cancel_tls = accept_tls12(cancel_socket, "AES128-GCM-SHA256").await.unwrap(); let mut cancel_request = [0_u8; 16]; cancel_tls.read_exact(&mut cancel_request).await.unwrap(); assert_eq!(u32::from_be_bytes(cancel_request[0..4].try_into().unwrap()), 16); assert_eq!(u32::from_be_bytes(cancel_request[4..8].try_into().unwrap()), 80_877_102); assert_eq!(u32::from_be_bytes(cancel_request[8..12].try_into().unwrap()), 1); assert_eq!(u32::from_be_bytes(cancel_request[12..16].try_into().unwrap()), 2); } #[tokio::test] async fn postgres_prefer_retries_plaintext_after_tls_handshake_failure() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let port = listener.local_addr().unwrap().port(); let server = tokio::spawn(serve_tls_handshake_failure_then_plaintext(listener)); let url = format!("postgres://postgres@127.0.0.1:{port}/postgres?sslmode=prefer"); let result = connect_with_optional_local_timezone(&url, Duration::from_secs(2), None).await; assert!(result.is_ok(), "prefer mode should retry without TLS: {result:?}"); server.await.unwrap(); } #[tokio::test] async fn postgres_default_tls_rejects_static_rsa_key_exchange() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let port = listener.local_addr().unwrap().port(); let server = tokio::spawn(serve_tls12_postgres(listener, "AES128-GCM-SHA256", false)); let url = format!("postgres://postgres@127.0.0.1:{port}/postgres?sslmode=require"); let error = connect_with_optional_local_timezone(&url, Duration::from_secs(2), None) .await .expect_err("rustls must not negotiate static RSA key exchange"); assert!(error.to_ascii_lowercase().contains("tls"), "{error}"); server.await.unwrap(); } #[tokio::test] async fn postgres_legacy_tls_connects_with_static_rsa_key_exchange() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let port = listener.local_addr().unwrap().port(); let server = tokio::spawn(serve_tls12_postgres(listener, "AES128-GCM-SHA256", true)); let url = format!("postgres://postgres@127.0.0.1:{port}/postgres?sslmode=require&legacy_tls=true"); let result = connect_with_optional_local_timezone(&url, Duration::from_secs(2), None).await; assert!(result.is_ok(), "legacy TLS should negotiate static RSA: {result:?}"); server.await.unwrap(); } #[tokio::test] async fn postgres_legacy_tls_verify_ca_uses_configured_root_without_hostname_check() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let port = listener.local_addr().unwrap().port(); let (acceptor, certificate_pem) = tls12_rsa_certificate_acceptor("AES128-GCM-SHA256"); let root_dir = tempfile::tempdir().unwrap(); let root_path = root_dir.path().join("root.pem"); std::fs::write(&root_path, certificate_pem).unwrap(); let root_path_text = root_path.to_string_lossy(); let encoded_root = percent_encoding::utf8_percent_encode(root_path_text.as_ref(), percent_encoding::NON_ALPHANUMERIC); let server = tokio::spawn(serve_tls12_postgres_with_acceptor(listener, acceptor, true)); let url = format!( "postgres://postgres@127.0.0.1:{port}/postgres?sslmode=verify-ca&legacy_tls=true&sslrootcert={encoded_root}" ); let result = connect_with_optional_local_timezone(&url, Duration::from_secs(2), None).await; assert!(result.is_ok(), "verify-ca should trust the configured root without checking the host: {result:?}"); server.await.unwrap(); } #[tokio::test] async fn postgres_legacy_tls_verify_full_rejects_hostname_mismatch() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let port = listener.local_addr().unwrap().port(); let (acceptor, certificate_pem) = tls12_rsa_certificate_acceptor("AES128-GCM-SHA256"); let root_dir = tempfile::tempdir().unwrap(); let root_path = root_dir.path().join("root.pem"); std::fs::write(&root_path, certificate_pem).unwrap(); let root_path_text = root_path.to_string_lossy(); let encoded_root = percent_encoding::utf8_percent_encode(root_path_text.as_ref(), percent_encoding::NON_ALPHANUMERIC); let server = tokio::spawn(serve_tls12_postgres_with_acceptor(listener, acceptor, false)); let url = format!( "postgres://postgres@127.0.0.1:{port}/postgres?sslmode=verify-full&legacy_tls=true&sslrootcert={encoded_root}" ); let error = connect_with_optional_local_timezone(&url, Duration::from_secs(2), None) .await .expect_err("verify-full must reject a certificate for another host"); assert!(error.to_ascii_lowercase().contains("certificate"), "{error}"); server.await.unwrap(); } #[tokio::test] async fn postgres_default_tls_still_connects_with_ecdhe_rsa() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let port = listener.local_addr().unwrap().port(); let server = tokio::spawn(serve_tls12_postgres(listener, "ECDHE-RSA-AES128-GCM-SHA256", true)); let url = format!("postgres://postgres@127.0.0.1:{port}/postgres?sslmode=require"); let result = connect_with_optional_local_timezone(&url, Duration::from_secs(2), None).await; assert!(result.is_ok(), "default TLS should keep negotiating ECDHE-RSA: {result:?}"); server.await.unwrap(); } #[tokio::test] async fn postgres_legacy_tls_cancel_uses_static_rsa_connector() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let port = listener.local_addr().unwrap().port(); let server = tokio::spawn(serve_static_rsa_postgres_with_cancel(listener)); let url = format!("postgres://postgres@127.0.0.1:{port}/postgres?sslmode=require&legacy_tls=true"); let pool = connect_with_optional_local_timezone(&url, Duration::from_secs(2), None).await.unwrap(); let client = pool.get().await.unwrap(); let context = build_postgres_cancel_context(&url).unwrap(); assert!(context.legacy_tls); cancel_postgres_query(client.cancel_token(), Some(&context), Duration::from_secs(2)).await; server.await.unwrap(); } fn pg_array_binary(element_oid: u32, elements: &[Option>]) -> Vec { let mut raw = Vec::new(); raw.extend_from_slice(&1_i32.to_be_bytes()); raw.extend_from_slice(&i32::from(elements.iter().any(Option::is_none)).to_be_bytes()); raw.extend_from_slice(&element_oid.to_be_bytes()); raw.extend_from_slice(&(elements.len() as i32).to_be_bytes()); raw.extend_from_slice(&1_i32.to_be_bytes()); for element in elements { match element { Some(bytes) => { raw.extend_from_slice(&(bytes.len() as i32).to_be_bytes()); raw.extend_from_slice(bytes); } None => raw.extend_from_slice(&(-1_i32).to_be_bytes()), } } raw } fn pg_jsonb_binary(value: &[u8]) -> Vec { let mut raw = Vec::with_capacity(value.len() + 1); raw.push(1); raw.extend_from_slice(value); raw } #[test] fn gaussdb_compatibility_mode_selects_identifier_quote() { for mode in ["M", "B", "mysql", " MYSQL "] { assert_eq!(gaussdb_identifier_quote_for_compatibility_mode(mode), Some("`")); } for mode in ["A", "PG", "ora", " PostgreSQL "] { assert_eq!(gaussdb_identifier_quote_for_compatibility_mode(mode), Some("\"")); } assert_eq!(gaussdb_identifier_quote_for_compatibility_mode("C"), None); assert_eq!(gaussdb_identifier_quote_for_compatibility_mode(""), None); } fn test_query_message(message: &str) -> QueryMessage { QueryMessage { severity: "NOTICE".to_string(), message: message.to_string(), code: Some("00000".to_string()), detail: None, hint: None, } } #[test] fn postgres_connection_identity_normalizes_vendor_numeric_types_to_text() { assert!(POSTGRES_CONNECTION_IDENTITY_SQL.contains("pg_backend_pid()::text")); assert!(POSTGRES_CONNECTION_IDENTITY_SQL.contains("inet_server_port()::text")); } #[tokio::test] #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a PostgreSQL-compatible database"] async fn postgres_connection_identity_supports_compatible_servers() { let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); let pool = connect_with_local_timezone(&url, Duration::from_secs(10), "UTC") .await .expect("connect PostgreSQL-compatible database"); let client = pool.get().await.expect("checkout PostgreSQL-compatible database"); let key = postgres_client_key(&client).await.expect("resolve text connection identity"); assert!(!key.0.is_empty()); assert!(!key.1.is_empty()); assert!(!key.2.is_empty()); } #[tokio::test] #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a PostgreSQL database"] async fn postgres_connection_identity_preserves_notice_capture() { let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); let pool = connect_with_local_timezone(&url, Duration::from_secs(10), "UTC") .await .expect("connect PostgreSQL database"); let client = pool.get().await.expect("checkout PostgreSQL database"); let result = execute_query_with_max_rows_inner( &client, "DO $$ BEGIN RAISE NOTICE 'dbx notice identity regression'; END $$", None, false, None, false, ) .await .expect("execute statement with notice"); assert!(result.messages.iter().any(|message| message.message == "dbx notice identity regression")); } #[test] fn take_notices_for_key_returns_buffered_notices_and_empties_buffer() { let key = ("test-host".to_string(), "9000001".to_string(), "9000001".to_string()); let buffer = Arc::new(Mutex::new(vec![test_query_message("first"), test_query_message("second")])); postgres_notice_buffers() .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()) .insert(key.clone(), Arc::downgrade(&buffer)); let notices = take_notices_for_key(&key); assert_eq!(notices.len(), 2); assert_eq!(notices[0].message, "first"); assert_eq!(notices[1].message, "second"); assert_eq!(notices[0].severity, "NOTICE"); assert_eq!(notices[0].code.as_deref(), Some("00000")); // The buffer was drained but stays registered while the connection lives. assert!(take_notices_for_key(&key).is_empty()); buffer.lock().unwrap_or_else(|poisoned| poisoned.into_inner()).push(test_query_message("third")); let notices = take_notices_for_key(&key); assert_eq!(notices.len(), 1); assert_eq!(notices[0].message, "third"); } #[test] fn take_notices_for_key_prunes_dead_buffers_and_misses_return_empty() { let live_key = ("test-host".to_string(), "9000002".to_string(), "9000002".to_string()); let dead_key = ("test-host".to_string(), "9000003".to_string(), "9000003".to_string()); let live = Arc::new(Mutex::new(vec![test_query_message("live")])); let dead = Arc::new(Mutex::new(vec![test_query_message("dead")])); { let mut buffers = postgres_notice_buffers().lock().unwrap_or_else(|poisoned| poisoned.into_inner()); buffers.insert(live_key.clone(), Arc::downgrade(&live)); buffers.insert(dead_key.clone(), Arc::downgrade(&dead)); } drop(dead); assert!( take_notices_for_key(&("test-host".to_string(), "9000004".to_string(), "9000004".to_string())).is_empty() ); assert!(take_notices_for_key(&dead_key).is_empty()); let buffers = postgres_notice_buffers().lock().unwrap_or_else(|poisoned| poisoned.into_inner()); assert!(!buffers.contains_key(&dead_key)); assert!(buffers.contains_key(&live_key)); drop(buffers); let notices = take_notices_for_key(&live_key); assert_eq!(notices.len(), 1); assert_eq!(notices[0].message, "live"); } #[test] fn take_notices_for_key_distinguishes_same_pid_on_different_servers() { // Backend PIDs collide across servers; the (address, port, pid) key // keeps notice attribution separate. let key_a = ("server-a".to_string(), "5432".to_string(), "42".to_string()); let key_b = ("server-b".to_string(), "5432".to_string(), "42".to_string()); let buffer_a = Arc::new(Mutex::new(vec![test_query_message("from-a")])); let buffer_b = Arc::new(Mutex::new(vec![test_query_message("from-b")])); { let mut buffers = postgres_notice_buffers().lock().unwrap_or_else(|poisoned| poisoned.into_inner()); buffers.insert(key_a.clone(), Arc::downgrade(&buffer_a)); buffers.insert(key_b.clone(), Arc::downgrade(&buffer_b)); } let notices_a = take_notices_for_key(&key_a); let notices_b = take_notices_for_key(&key_b); assert_eq!(notices_a.len(), 1); assert_eq!(notices_a[0].message, "from-a"); assert_eq!(notices_b.len(), 1); assert_eq!(notices_b[0].message, "from-b"); } #[test] fn postgres_json_arrays_decode_elements_without_jsonb_version_bytes() { let json_raw = pg_array_binary( Type::JSON.oid(), &[ Some(br#"{"kind":"json"}"#.to_vec()), Some(br#""text""#.to_vec()), Some(br#"[1,true,null]"#.to_vec()), None, ], ); let jsonb_raw = pg_array_binary( Type::JSONB.oid(), &[ Some(pg_jsonb_binary(br#"{"port":10031,"type":"admin_web"}"#)), Some(pg_jsonb_binary(br#""quoted""#)), Some(pg_jsonb_binary(br#"[2,false,{"nested":true}]"#)), None, ], ); let json_values = Vec::>::from_sql(&Type::JSON_ARRAY, &json_raw).unwrap(); let jsonb_values = Vec::>::from_sql(&Type::JSONB_ARRAY, &jsonb_raw).unwrap(); assert_eq!( pg_json_array_values_to_json(json_values), serde_json::json!([r#"{"kind":"json"}"#, r#""text""#, "[1,true,null]", null]) ); let decoded = pg_json_array_values_to_json(jsonb_values); assert_eq!( decoded, serde_json::json!([ r#"{"port":10031,"type":"admin_web"}"#, r#""quoted""#, r#"[2,false,{"nested":true}]"#, null ]) ); assert!(!decoded.to_string().contains('\u{1}')); } #[test] fn postgres_json_array_decoder_is_limited_to_json_element_types() { assert!(Vec::>::accepts(&Type::JSON_ARRAY)); assert!(Vec::>::accepts(&Type::JSONB_ARRAY)); assert!(!Vec::>::accepts(&Type::TEXT_ARRAY)); assert!(!Vec::>::accepts(&Type::INT4_ARRAY)); } #[test] fn postgres_explain_analyze_uses_read_only_transaction_local_schema() { assert_eq!( postgres_read_only_transaction_setup(), vec![("BEGIN READ ONLY".to_string(), "explain_analyze.begin")] ); } #[tokio::test] async fn postgres_explain_analyze_rolls_back_after_success_and_query_failures() { let rollback_calls = Cell::new(0); let result = run_postgres_operation_with_rollback( || async { Ok::<_, String>(7) }, || async { rollback_calls.set(rollback_calls.get() + 1); Ok(()) }, ) .await; assert_eq!(result, Ok(7)); assert_eq!(rollback_calls.get(), 1); for operation_error in ["query failed", crate::query::QUERY_CANCELED, "Query timed out after 30 seconds"] { let rollback_calls = Cell::new(0); let result = run_postgres_operation_with_rollback( || async { Err::<(), _>(operation_error.to_string()) }, || async { rollback_calls.set(rollback_calls.get() + 1); Ok(()) }, ) .await; assert_eq!(result, Err(operation_error.to_string())); assert_eq!(rollback_calls.get(), 1); } } #[tokio::test] async fn postgres_explain_analyze_marks_rollback_failure_as_pool_pollution() { let result = run_postgres_operation_with_rollback( || async { Err::<(), _>("query failed".to_string()) }, || async { Err("PostgreSQL explain_analyze.rollback timed out after 3 seconds".to_string()) }, ) .await; assert_eq!( result, Err("query failed; PostgreSQL read-only transaction cleanup failed: PostgreSQL explain_analyze.rollback timed out after 3 seconds".to_string()) ); } fn pg_interval_bytes(microseconds: i64, days: i32, months: i32) -> [u8; 16] { let mut raw = [0_u8; 16]; raw[0..8].copy_from_slice(µseconds.to_be_bytes()); raw[8..12].copy_from_slice(&days.to_be_bytes()); raw[12..16].copy_from_slice(&months.to_be_bytes()); raw } #[test] fn postgres_interval_binary_decodes_and_formats_components() { let microseconds = 4 * 3_600_000_000 + 5 * 60_000_000 + 6 * 1_000_000 + 123_456; let interval = PgInterval::from_sql(&Type::INTERVAL, &pg_interval_bytes(microseconds, 3, 14)).unwrap(); assert_eq!(interval, PgInterval { microseconds, days: 3, months: 14 }); assert_eq!(format_pg_interval(interval), "1 year 2 mons 3 days 04:05:06.123456"); } #[test] fn postgres_interval_formats_negative_mixed_and_zero_values() { assert_eq!( format_pg_interval(PgInterval { microseconds: -3_723_450_000, days: -2, months: -13 }), "-1 year -1 mon -2 days -01:02:03.45" ); assert_eq!( format_pg_interval(PgInterval { microseconds: -1, days: 2, months: -1 }), "-1 mon 2 days -00:00:00.000001" ); assert_eq!(format_pg_interval(PgInterval { microseconds: 0, days: 0, months: 0 }), "00:00:00"); } #[test] fn postgres_interval_formats_now_minus_xact_start_shape() { let elapsed = PgInterval { microseconds: 123_450_000, days: 0, months: 0 }; assert_eq!(format_pg_interval(elapsed), "00:02:03.45"); } #[test] fn postgres_interval_rejects_invalid_binary_and_keeps_binary_protocol() { assert!(PgInterval::from_sql(&Type::INTERVAL, &[0; 15]).is_err()); assert_eq!(classify_pg_type("interval"), PgColType::Interval); assert_eq!(classify_pg_type("_interval"), PgColType::Temporal { fallback: PgTemporalFallback::GenericArray }); assert!(!pg_type_requires_text_protocol(&Type::INTERVAL, PgColType::Interval)); } fn pg_date_range_bytes(flags: u8, lower: Option, upper: Option) -> Vec { let mut raw = vec![flags]; for days in [lower, upper].into_iter().flatten() { raw.extend_from_slice(&4_i32.to_be_bytes()); raw.extend_from_slice(&days.to_be_bytes()); } raw } #[test] fn postgres_daterange_binary_decodes_reported_value() { let raw = pg_date_range_bytes(PG_RANGE_LOWER_INCLUSIVE, Some(9_163), Some(9_168)); let range = PgDateRange::from_sql(&Type::DATE_RANGE, &raw).unwrap(); assert_eq!(range.0, "[2025-02-01,2025-02-06)"); } #[test] fn postgres_daterange_binary_handles_empty_unbounded_and_infinite_bounds() { assert_eq!(PgDateRange::from_sql(&Type::DATE_RANGE, &[PG_RANGE_EMPTY]).unwrap().0, "empty"); assert_eq!( PgDateRange::from_sql( &Type::DATE_RANGE, &pg_date_range_bytes(PG_RANGE_UPPER_UNBOUNDED | PG_RANGE_LOWER_INCLUSIVE, Some(9_163), None), ) .unwrap() .0, "[2025-02-01,)" ); assert_eq!( PgDateRange::from_sql( &Type::DATE_RANGE, &pg_date_range_bytes(PG_RANGE_LOWER_UNBOUNDED, None, Some(9_168)), ) .unwrap() .0, "(,2025-02-06)" ); assert_eq!( PgDateRange::from_sql( &Type::DATE_RANGE, &pg_date_range_bytes(PG_RANGE_LOWER_INCLUSIVE, Some(i32::MIN), Some(i32::MAX)), ) .unwrap() .0, "[-infinity,infinity)" ); } #[test] fn postgres_daterange_rejects_malformed_binary_and_keeps_binary_protocol() { assert!(PgDateRange::from_sql(&Type::DATE_RANGE, &[]).is_err()); assert!(PgDateRange::from_sql(&Type::DATE_RANGE, &[PG_RANGE_EMPTY, 0]).is_err()); assert!(PgDateRange::from_sql( &Type::DATE_RANGE, &pg_date_range_bytes(PG_RANGE_LOWER_UNBOUNDED | PG_RANGE_LOWER_INCLUSIVE, None, Some(9_168)), ) .is_err()); assert!(PgDateRange::from_sql(&Type::DATE_RANGE, &[PG_RANGE_LOWER_INCLUSIVE, 0, 0, 0, 3, 0, 0, 0]).is_err()); assert_eq!(classify_pg_type("daterange"), PgColType::DateRange); assert_eq!(classify_pg_type("_daterange"), PgColType::GenericArray); assert!(!pg_type_requires_text_protocol(&Type::DATE_RANGE, PgColType::DateRange)); } #[test] fn postgres_custom_other_type_requires_text_protocol() { assert!(pg_scalar_type_requires_text_protocol(8_880, PgColType::Other)); assert!(pg_scalar_type_requires_text_protocol(98_765, PgColType::Other)); assert!(pg_scalar_type_requires_text_protocol(98_765, PgColType::GenericArray)); } #[test] fn postgres_record_types_require_text_protocol() { assert!(pg_type_requires_text_protocol(&Type::RECORD, PgColType::Other)); assert!(pg_type_requires_text_protocol(&Type::RECORD_ARRAY, PgColType::GenericArray)); let dynamic_record = Type::new("record".to_string(), Type::RECORD.oid(), Kind::Simple, "pg_catalog".to_string()); let dynamic_record_array = Type::new("_record".to_string(), Type::RECORD_ARRAY.oid(), Kind::Simple, "pg_catalog".to_string()); assert!(pg_type_requires_text_protocol(&dynamic_record, PgColType::Other)); assert!(pg_type_requires_text_protocol(&dynamic_record_array, PgColType::GenericArray)); } #[test] fn postgres_reg_types_use_text_protocol_for_catalog_names() { // Binary reg* values are OIDs, while PostgreSQL's text output is the // human-readable object/type name users expect to see in the grid. for pg_type in [ Type::REGPROC, Type::REGPROCEDURE, Type::REGOPER, Type::REGOPERATOR, Type::REGCLASS, Type::REGTYPE, Type::REGNAMESPACE, Type::REGROLE, Type::REGCOLLATION, Type::REGCONFIG, Type::REGDICTIONARY, Type::REGPROC_ARRAY, Type::REGPROCEDURE_ARRAY, Type::REGOPER_ARRAY, Type::REGOPERATOR_ARRAY, Type::REGCLASS_ARRAY, Type::REGTYPE_ARRAY, Type::REGNAMESPACE_ARRAY, Type::REGROLE_ARRAY, Type::REGCOLLATION_ARRAY, Type::REGCONFIG_ARRAY, Type::REGDICTIONARY_ARRAY, ] { assert!(pg_type_is_reg_type(&pg_type)); assert!(pg_type_requires_text_protocol(&pg_type, PgColType::Other)); } } #[test] fn postgres_enum_keeps_binary_protocol() { let enum_type = Type::new( "withdraw_btc_status".to_string(), 98_765, Kind::Enum(vec!["pending".to_string(), "completed".to_string()]), "risk".to_string(), ); let enum_array = Type::new("_withdraw_btc_status".to_string(), 98_766, Kind::Array(enum_type.clone()), "risk".to_string()); assert!(!pg_type_requires_text_protocol(&enum_type, PgColType::Other)); assert!(pg_type_requires_text_protocol(&enum_array, PgColType::GenericArray)); } #[test] fn postgres_builtin_or_supported_type_keeps_binary_protocol() { assert!(!pg_type_requires_text_protocol(&Type::INT4, PgColType::Other)); assert!(!pg_type_requires_text_protocol(&Type::VARCHAR, PgColType::Other)); assert!(!pg_type_requires_text_protocol(&Type::INT4_ARRAY, PgColType::GenericArray)); assert!(!pg_scalar_type_requires_text_protocol(98_765, PgColType::Vector)); assert!(!pg_scalar_type_requires_text_protocol(98_765, PgColType::Geometry)); } #[test] fn postgres_compatible_builtin_names_with_unknown_low_oids_require_text_protocol() { let compatible_date = Type::new("date".to_string(), 8_881, Kind::Simple, "pg_catalog".to_string()); let compatible_tinyint = Type::new("tinyint".to_string(), 8_882, Kind::Simple, "pg_catalog".to_string()); assert!(Type::from_oid(compatible_date.oid()).is_none()); assert!(Type::from_oid(compatible_tinyint.oid()).is_none()); assert!(pg_type_requires_text_protocol( &compatible_date, PgColType::Temporal { fallback: PgTemporalFallback::Probe } )); assert!(pg_type_requires_text_protocol(&compatible_tinyint, PgColType::Other)); } #[test] fn postgres_query_uses_text_when_any_output_type_is_unsupported() { let columns = [(Type::INT4.oid(), PgColType::Other), (98_765, PgColType::Other), (Type::TEXT.oid(), PgColType::Other)]; assert!(columns.into_iter().any(|(oid, col_type)| pg_scalar_type_requires_text_protocol(oid, col_type))); } #[test] fn postgres_text_fallback_keeps_matching_prepared_column_types() { let columns = vec!["payload".to_string(), "id".to_string()]; let types = vec!["payload_type".to_string(), "int4".to_string()]; assert_eq!(matching_pg_text_column_types(&columns, Some(types.clone())), types); } #[test] fn postgres_text_fallback_discards_misaligned_column_types() { let columns = vec!["payload".to_string(), "id".to_string()]; let types = vec!["payload_type".to_string()]; assert!(matching_pg_text_column_types(&columns, Some(types)).is_empty()); assert!(matching_pg_text_column_types(&columns, None).is_empty()); } #[test] fn postgres_query_search_path_preserves_public_after_catalog() { assert_eq!( postgres_set_search_path_sql("application", PostgresSearchPathContext::Query), "SET search_path TO \"application\", pg_catalog, public" ); } #[test] fn postgres_transaction_search_paths_prioritize_selected_schema() { assert_eq!( postgres_set_search_path_sql("application", PostgresSearchPathContext::Transaction), "SET search_path TO \"application\", pg_catalog" ); assert_eq!( postgres_set_search_path_sql("application", PostgresSearchPathContext::LocalTransaction), "SET LOCAL search_path TO \"application\", pg_catalog" ); assert_eq!( postgres_set_search_path_sql("application", PostgresSearchPathContext::LocalQueryTransaction), "SET LOCAL search_path TO \"application\", pg_catalog, public" ); } #[test] fn postgres_search_path_safely_quotes_selected_schema() { assert_eq!( postgres_set_search_path_sql("tenant\"; RESET search_path; --", PostgresSearchPathContext::Query,), "SET search_path TO \"tenant\"\"; RESET search_path; --\", pg_catalog, public" ); } #[test] fn postgres_single_schema_search_path_preserves_scope_and_quoting() { assert_eq!( postgres_set_single_schema_search_path_sql("application", PostgresSearchPathContext::Query), "SET search_path TO \"application\"" ); assert_eq!( postgres_set_single_schema_search_path_sql("application", PostgresSearchPathContext::Transaction), "SET search_path TO \"application\"" ); assert_eq!( postgres_set_single_schema_search_path_sql( "tenant\"; RESET search_path; --", PostgresSearchPathContext::LocalTransaction, ), "SET LOCAL search_path TO \"tenant\"\"; RESET search_path; --\"" ); assert_eq!( postgres_set_single_schema_search_path_sql("application", PostgresSearchPathContext::LocalQueryTransaction), "SET LOCAL search_path TO \"application\"" ); } #[test] fn postgres_search_path_keeps_the_configured_path_when_selected_schema_is_first() { let baseline = PostgresSearchPathBaseline { configured: "application, extensions, public".to_string(), first_resolved_schema: Some("application".to_string()), has_explicit_pg_catalog: false, }; assert_eq!( postgres_set_preserved_search_path_sql("application", PostgresSearchPathContext::Query, &baseline), "SET search_path TO application, extensions, public, pg_catalog" ); } #[test] fn postgres_search_path_prepends_selected_schema_without_dropping_configured_items() { let baseline = PostgresSearchPathBaseline { configured: "\"$user\", extensions, public".to_string(), first_resolved_schema: Some("public".to_string()), has_explicit_pg_catalog: false, }; assert_eq!( postgres_set_preserved_search_path_sql("application", PostgresSearchPathContext::Query, &baseline), "SET search_path TO \"application\", \"$user\", extensions, public, pg_catalog" ); assert_eq!( postgres_set_preserved_search_path_sql( "application", PostgresSearchPathContext::LocalQueryTransaction, &baseline, ), "SET LOCAL search_path TO \"application\", \"$user\", extensions, public, pg_catalog" ); } #[test] fn postgres_search_path_keeps_explicit_pg_catalog_position() { let baseline = PostgresSearchPathBaseline { configured: "extensions, pg_catalog, public".to_string(), first_resolved_schema: Some("extensions".to_string()), has_explicit_pg_catalog: true, }; assert_eq!( postgres_set_preserved_search_path_sql("application", PostgresSearchPathContext::Query, &baseline), "SET search_path TO \"application\", extensions, pg_catalog, public" ); } #[test] fn postgres_single_schema_fallback_only_matches_compatible_server_error() { assert!(postgres_requires_single_schema_search_path( "ERROR: Hologres does not support search_path with multiple names: admaterial." )); assert!(!postgres_requires_single_schema_search_path("ERROR: permission denied for schema admaterial")); } #[test] fn postgres_statement_returns_rows_for_returning_dml_only() { for sql in [ "INSERT INTO users (id) VALUES (1) RETURNING id", "UPDATE users SET name = 'Ada' RETURNING id, name", "DELETE FROM users WHERE id = 1 RETURNING id", "MERGE INTO users AS target USING updates AS source ON target.id = source.id WHEN MATCHED THEN UPDATE SET name = source.name RETURNING target.id", "WITH removed AS (DELETE FROM users WHERE id = 1 RETURNING id) SELECT * FROM removed", ] { assert!(postgres_statement_returns_rows(sql), "expected result rows for: {sql}"); } for sql in [ "INSERT INTO users (id) VALUES (1)", "UPDATE users SET name = 'Ada'", "DELETE FROM users WHERE id = 1", "MERGE INTO users AS target USING updates AS source ON target.id = source.id WHEN MATCHED THEN UPDATE SET name = source.name", "INSERT INTO users (note) VALUES ('RETURNING is text')", ] { assert!(!postgres_statement_returns_rows(sql), "expected command result for: {sql}"); } } #[test] fn cte_wrapped_dml_parses_as_query_body() { let statements = Parser::parse_sql( &PostgreSqlDialect {}, "WITH cte AS (SELECT id FROM t) UPDATE t SET x = 1 WHERE id IN (SELECT id FROM cte)", ) .expect("CTE-wrapped UPDATE should parse"); let [Statement::Query(query)] = statements.as_slice() else { panic!("expected a single Statement::Query, got: {statements:?}"); }; assert!(query.with.is_some(), "expected the CTE to be attached to the query"); assert!( matches!(postgres_query_body_dml(&query.body), Some(Statement::Update(_))), "expected an UPDATE body, got: {:?}", query.body ); } #[test] fn postgres_statement_returns_rows_for_cte_wrapped_dml() { for sql in [ "WITH cte AS (SELECT id FROM t) SELECT * FROM cte", "WITH cte AS (SELECT id FROM t) UPDATE t SET x = 1 WHERE id IN (SELECT id FROM cte) RETURNING id", "WITH cte AS (SELECT id FROM t) DELETE FROM t WHERE id IN (SELECT id FROM cte) RETURNING id", "WITH cte AS (SELECT id FROM t) INSERT INTO u (id) SELECT id FROM cte RETURNING id", // Not parseable by sqlparser: keep assuming a CTE query returns rows. "WITH RECURSIVE cte AS MATERIALIZED (SELECT 1) SELECT * FROM cte", ] { assert!(postgres_statement_returns_rows(sql), "expected result rows for: {sql}"); } for sql in [ "WITH cte AS (SELECT id FROM t) UPDATE t SET x = 1 WHERE id IN (SELECT id FROM cte)", "WITH cte AS (SELECT id FROM t) DELETE FROM t WHERE id IN (SELECT id FROM cte)", "WITH cte AS (SELECT id FROM t) INSERT INTO u (id) SELECT id FROM cte", " with cte as (select id from t)\n update t set x = 1 where id in (select id from cte) ", ] { assert!(!postgres_statement_returns_rows(sql), "expected command result for: {sql}"); } } #[test] fn database_list_does_not_collect_storage_usage() { assert!(list_databases_sql().contains("pg_database")); assert!(!list_databases_sql().contains("pg_database_size")); } #[test] fn database_storage_is_scoped_and_permission_guarded() { let sql = database_storage_sql(); assert!(sql.contains("d.datname = ANY($1::text[])")); assert!(sql.contains("has_database_privilege")); assert!(sql.contains("pg_read_all_stats")); assert!(sql.contains("pg_database_size")); assert!(sql.contains("ELSE NULL")); } #[test] fn classify_pg_type_covers_all_dispatch_branches() { assert_eq!(classify_pg_type("bytea"), PgColType::Bytea); assert_eq!(classify_pg_type("json"), PgColType::Json); assert_eq!(classify_pg_type("JSONB"), PgColType::Json); assert_eq!(classify_pg_type("bool"), PgColType::Bool); assert_eq!(classify_pg_type("point"), PgColType::Point); assert_eq!(classify_pg_type("timestamp"), PgColType::Temporal { fallback: PgTemporalFallback::Probe }); assert_eq!(classify_pg_type("timestamptz"), PgColType::Temporal { fallback: PgTemporalFallback::Probe }); assert_eq!(classify_pg_type("date"), PgColType::Temporal { fallback: PgTemporalFallback::Probe }); assert_eq!(classify_pg_type("time"), PgColType::Temporal { fallback: PgTemporalFallback::Probe }); assert_eq!(classify_pg_type("timetz"), PgColType::Temporal { fallback: PgTemporalFallback::Probe }); assert_eq!(classify_pg_type("interval"), PgColType::Interval); // 时间数组类型名在原实现中先进时间分支、解码失败后落到通用数组分支 assert_eq!(classify_pg_type("_timestamp"), PgColType::Temporal { fallback: PgTemporalFallback::GenericArray }); assert_eq!(classify_pg_type("_interval"), PgColType::Temporal { fallback: PgTemporalFallback::GenericArray }); // 同时命中时间关键字与 VECTOR( 前缀的类型名,原实现时间解码失败后走 vector 分支 assert_eq!(classify_pg_type("vector(timestamp)"), PgColType::Temporal { fallback: PgTemporalFallback::Vector }); assert_eq!(classify_pg_type("numeric"), PgColType::Numeric); assert_eq!(classify_pg_type("money"), PgColType::Money); assert_eq!(classify_pg_type("uuid"), PgColType::Uuid); assert_eq!(classify_pg_type("inet"), PgColType::Inet { cidr: false }); assert_eq!(classify_pg_type("cidr"), PgColType::Inet { cidr: true }); assert_eq!(classify_pg_type("macaddr"), PgColType::MacAddr); assert_eq!(classify_pg_type("macaddr8"), PgColType::MacAddr); assert_eq!(classify_pg_type("bit"), PgColType::BitString); assert_eq!(classify_pg_type("varbit"), PgColType::BitString); assert_eq!(classify_pg_type("tsvector"), PgColType::TsVector); assert_eq!(classify_pg_type("oid"), PgColType::SystemU32); assert_eq!(classify_pg_type("xid"), PgColType::SystemU32); assert_eq!(classify_pg_type("_inet"), PgColType::InetArray { cidr: false }); assert_eq!(classify_pg_type("_cidr"), PgColType::InetArray { cidr: true }); assert_eq!(classify_pg_type("_macaddr"), PgColType::MacAddrArray); assert_eq!(classify_pg_type("_bit"), PgColType::BitStringArray); assert_eq!(classify_pg_type("_varbit"), PgColType::BitStringArray); assert_eq!(classify_pg_type("_int4"), PgColType::GenericArray); assert_eq!(classify_pg_type("_time"), PgColType::GenericArray); assert_eq!(classify_pg_type("vector"), PgColType::Vector); assert_eq!(classify_pg_type("vector(3)"), PgColType::Vector); assert_eq!(classify_pg_type("geometry"), PgColType::Geometry); assert_eq!(classify_pg_type("geography"), PgColType::Geometry); assert_eq!(classify_pg_type("int4"), PgColType::Other); assert_eq!(classify_pg_type("varchar"), PgColType::Other); assert_eq!(classify_pg_type(""), PgColType::Other); } #[test] fn postgres_text_spatial_value_separates_srid_from_wkt() { assert_eq!( pg_text_fallback_value("SRID=4326;POINT(1 2)", Some(PgColType::Geometry)), (serde_json::json!("POINT(1 2)"), Some(4326)) ); assert_eq!( pg_text_fallback_value("SRID=0;POINT(1 2)", Some(PgColType::Geometry)), (serde_json::json!("POINT(1 2)"), None) ); assert_eq!( pg_text_fallback_value("POINT(1 2)", Some(PgColType::Geometry)), (serde_json::json!("POINT(1 2)"), None) ); } #[test] fn postgres_text_spatial_value_decodes_hex_ewkb() { let ewkb = "0101000020E6100000000000000000F03F0000000000000040"; for value in [ewkb.to_string(), format!("0x{ewkb}"), format!("\\x{ewkb}")] { assert_eq!( pg_text_fallback_value(&value, Some(PgColType::Geometry)), (serde_json::json!("POINT(1 2)"), Some(4326)) ); assert_eq!(pg_text_fallback_value(&value, None), (serde_json::json!("POINT(1 2)"), Some(4326))); } let srid_zero = "010100002000000000000000000000F03F0000000000000040"; assert_eq!(pg_text_fallback_value(srid_zero, None), (serde_json::json!("POINT(1 2)"), None)); } #[test] fn postgres_text_fallback_does_not_reinterpret_ordinary_text() { for value in ["SRID=abc;POINT(1 2)", "SRID=4326;not geometry", "0101-not-hex", "POINTLESS"] { assert_eq!(pg_text_fallback_value(value, None), (serde_json::json!(value), None)); } assert_eq!(pg_text_fallback_value("SRID=4326;point(1 2)", None), (serde_json::json!("point(1 2)"), Some(4326))); } struct DockerPostgres { name: String, port: u16, } impl DockerPostgres { fn url(&self) -> String { format!("postgres://postgres:postgres@127.0.0.1:{}/postgres?sslmode=disable", self.port) } } impl Drop for DockerPostgres { fn drop(&mut self) { let _ = Command::new("docker").args(["rm", "-f", &self.name]).status(); } } fn docker_ready() -> bool { Command::new("docker") .args(["version", "--format", "{{.Server.Version}}"]) .output() .map(|output| output.status.success()) .unwrap_or(false) } async fn start_docker_postgres() -> Option { if !docker_ready() { eprintln!("skipping docker-backed postgres test because Docker is unavailable"); return None; } let port = portpicker::pick_unused_port().expect("pick unused postgres port"); let container = DockerPostgres { name: format!("dbx-postgres-enum-{}", uuid::Uuid::new_v4()), port }; let status = Command::new("docker") .args([ "run", "-d", "--rm", "--name", &container.name, "-e", "POSTGRES_PASSWORD=postgres", "-e", "POSTGRES_USER=postgres", "-e", "POSTGRES_DB=postgres", "-p", &format!("{port}:5432"), "postgres:16-alpine", ]) .status() .expect("start docker postgres"); assert!(status.success(), "docker run postgres container should succeed"); let deadline = Instant::now() + Duration::from_secs(60); loop { match connect(&container.url(), Duration::from_secs(2)).await { Ok(pool) => { drop(pool); return Some(container); } Err(_) if Instant::now() < deadline => tokio::time::sleep(Duration::from_millis(500)).await, Err(error) => panic!("docker postgres did not become ready: {error}"), } } } #[tokio::test] async fn postgres_dml_returning_preserves_result_rows() { let Some(container) = start_docker_postgres().await else { return; }; let pool = connect(&container.url(), Duration::from_secs(5)).await.expect("connect postgres"); execute_query(&pool, "CREATE TABLE dml_returning (id integer PRIMARY KEY, name text NOT NULL)") .await .expect("create table"); let inserted = execute_query(&pool, "INSERT INTO dml_returning VALUES (1, 'alice'), (2, 'bob')") .await .expect("insert rows"); let updated = execute_query(&pool, "UPDATE dml_returning SET name = 'bob-updated' WHERE id = 2") .await .expect("update rows"); let deleted = execute_query(&pool, "DELETE FROM dml_returning WHERE id = 1").await.expect("delete row"); let unmatched = execute_query(&pool, "UPDATE dml_returning SET name = name WHERE id = 999").await.expect("update no rows"); assert_eq!(inserted.affected_rows, 2); assert_eq!(updated.affected_rows, 1); assert_eq!(deleted.affected_rows, 1); assert_eq!(unmatched.affected_rows, 0); let insert_returning = execute_query(&pool, "INSERT INTO dml_returning VALUES (3, 'carol') RETURNING id, name") .await .expect("insert returning"); let update_returning = execute_query(&pool, "UPDATE dml_returning SET name = 'carol-updated' WHERE id = 3 RETURNING id, name") .await .expect("update returning"); let delete_returning = execute_query(&pool, "DELETE FROM dml_returning WHERE id = 3 RETURNING id, name") .await .expect("delete returning"); let empty_returning = execute_query(&pool, "UPDATE dml_returning SET name = name WHERE id = 999 RETURNING id, name") .await .expect("empty returning"); for result in [&insert_returning, &update_returning, &delete_returning] { assert_eq!(result.columns, vec!["id", "name"]); assert_eq!(result.rows.len(), 1); } assert_eq!(insert_returning.rows[0], vec![serde_json::json!(3), serde_json::json!("carol")]); assert_eq!(update_returning.rows[0], vec![serde_json::json!(3), serde_json::json!("carol-updated")]); assert_eq!(delete_returning.rows[0], vec![serde_json::json!(3), serde_json::json!("carol-updated")]); assert_eq!(empty_returning.columns, vec!["id", "name"]); assert!(empty_returning.rows.is_empty()); } async fn assert_postgres_18(pool: &Pool) { let version = execute_query(pool, "SHOW server_version_num").await.expect("query PostgreSQL version"); let version_num = version.rows[0][0] .as_str() .expect("server_version_num should be text") .parse::() .expect("server_version_num should be numeric"); assert!((180_000..190_000).contains(&version_num), "expected PostgreSQL 18, got {version_num}"); } #[tokio::test] #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a writable PostgreSQL 18 database"] async fn postgres_custom_composite_result_uses_server_text_output() { let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); let pool = connect(&url, Duration::from_secs(5)).await.expect("connect postgres"); assert_postgres_18(&pool).await; let schema = format!("dbx_custom_text_{}", uuid::Uuid::new_v4().simple()); let schema_ident = pg_quote_ident(&schema); let payload_type = format!("{schema_ident}.payload"); execute_query(&pool, &format!("CREATE SCHEMA {schema_ident}")).await.expect("create schema"); let exercise = async { execute_query(&pool, &format!("CREATE TYPE {payload_type} AS (id integer, label text)")).await?; let custom = execute_query(&pool, &format!("SELECT ROW(7, 'alpha')::{payload_type} AS payload, 42::int4 AS id")) .await?; let builtin = execute_query(&pool, "SELECT 42::int4 AS id").await?; Ok::<_, String>((custom, builtin)) } .await; let cleanup = execute_query(&pool, &format!("DROP SCHEMA {schema_ident} CASCADE")).await; cleanup.expect("drop schema"); let (custom, builtin) = exercise.expect("exercise custom composite fallback"); assert_eq!(custom.columns, vec!["payload", "id"]); assert_eq!(custom.column_types, vec!["payload", "int4"]); assert_eq!(custom.rows[0][0], serde_json::Value::String("(7,alpha)".to_string())); assert_eq!(custom.rows[0][1], serde_json::Value::String("42".to_string())); assert!(!custom.rows[0][0].as_str().unwrap().chars().any(char::is_control)); assert_eq!(builtin.column_types, vec!["int4"]); assert_eq!(builtin.rows[0][0], serde_json::Value::Number(42.into())); } #[tokio::test] #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a writable PostgreSQL 18 database"] async fn postgres_custom_type_arrays_and_exports_use_server_text_output() { let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); let pool = connect(&url, Duration::from_secs(5)).await.expect("connect postgres"); assert_postgres_18(&pool).await; let schema = format!("dbx_custom_array_{}", uuid::Uuid::new_v4().simple()); let schema_ident = pg_quote_ident(&schema); let payload_type = format!("{schema_ident}.payload"); let mood_type = format!("{schema_ident}.mood"); let score_type = format!("{schema_ident}.positive_int"); let underscore_scalar_type = format!("{schema_ident}._hidden"); let vector_named_enum_type = format!("{schema_ident}.vector"); let table = format!("{schema_ident}.custom_arrays"); let select_sql = format!("SELECT payloads, moods, scores FROM {table}"); execute_query(&pool, &format!("CREATE SCHEMA {schema_ident}")).await.expect("create schema"); let exercise = async { execute_query(&pool, &format!("CREATE TYPE {payload_type} AS (id integer, label text)")).await?; execute_query(&pool, &format!("CREATE TYPE {mood_type} AS ENUM ('ready', 'done')")).await?; execute_query(&pool, &format!("CREATE DOMAIN {score_type} AS integer CHECK (VALUE > 0)")).await?; execute_query(&pool, &format!("CREATE TYPE {underscore_scalar_type} AS ENUM ('secret')")).await?; execute_query(&pool, &format!("CREATE TYPE {vector_named_enum_type} AS ENUM ('label')")).await?; execute_query( &pool, &format!( "CREATE TABLE {table} (payloads {payload_type}[], moods {mood_type}[], scores {score_type}[])" ), ) .await?; execute_query( &pool, &format!( "INSERT INTO {table} VALUES \ (ARRAY[ROW(7, 'alpha')::{payload_type}], ARRAY['ready'::{mood_type}], ARRAY[7::{score_type}])" ), ) .await?; let query = execute_query(&pool, &select_sql).await?; let underscore_scalar = execute_query(&pool, &format!("SELECT 'secret'::{underscore_scalar_type} AS hidden")).await?; let vector_named_enum = execute_query(&pool, &format!("SELECT 'label'::{vector_named_enum_type} AS label")).await?; let client = checkout_postgres_client(&pool, None, Duration::from_secs(5)).await?; let mut query_export_rows = Vec::new(); stream_select_query_inner_with_mode( &client, &select_sql, None, &mut |item| { if let PostgresQueryStreamItem::Row(row) = item { query_export_rows.push(row); } Ok(()) }, false, ) .await?; let cancelled = AtomicBool::new(false); let mut table_export_rows = Vec::new(); stream_query_rows_on_client(&client, &select_sql, None, &cancelled, &mut |row| { table_export_rows.push(row.to_vec()); Ok(()) }) .await?; drop(client); Ok::<_, String>((query, underscore_scalar, vector_named_enum, query_export_rows, table_export_rows)) } .await; let cleanup = execute_query(&pool, &format!("DROP SCHEMA {schema_ident} CASCADE")).await; cleanup.expect("drop schema"); let (query, underscore_scalar, vector_named_enum, query_export_rows, table_export_rows) = exercise.expect("exercise custom array fallbacks"); let expected = vec![ serde_json::Value::String(r#"{"(7,alpha)"}"#.to_string()), serde_json::Value::String("{ready}".to_string()), serde_json::Value::String("{7}".to_string()), ]; assert_eq!(query.rows, vec![expected.clone()]); assert_eq!(underscore_scalar.rows, vec![vec![serde_json::Value::String("secret".to_string())]]); assert_eq!(vector_named_enum.rows, vec![vec![serde_json::Value::String("label".to_string())]]); assert_eq!(query_export_rows, vec![expected.clone()]); assert_eq!(table_export_rows, vec![expected]); } #[tokio::test] #[ignore = "requires DBX_TEST_GAUSSDB_URL pointing at a writable GaussDB database"] async fn gaussdb_list_custom_types_excludes_relation_row_types_and_arrays() { let url = std::env::var("DBX_TEST_GAUSSDB_URL").expect("DBX_TEST_GAUSSDB_URL"); let pool = connect(&url, Duration::from_secs(5)).await.expect("connect gaussdb"); let schema = format!("dbx_types_{}", uuid::Uuid::new_v4().simple()); let schema_ident = pg_quote_ident(&schema); let status_type = format!("{schema_ident}.status"); let address_type = format!("{schema_ident}.address"); let orders_table = format!("{schema_ident}.orders"); execute_query(&pool, &format!("CREATE SCHEMA {schema_ident}")).await.expect("create schema"); let exercise = async { execute_query(&pool, &format!("CREATE TYPE {status_type} AS ENUM ('draft', 'published')")).await?; execute_query(&pool, &format!("CREATE TYPE {address_type} AS (city text, zip text)")).await?; execute_query(&pool, &format!("COMMENT ON TYPE {status_type} IS '订单状态'")).await?; execute_query( &pool, &format!("CREATE TABLE {orders_table} (id bigint, state {status_type}, ship_to {address_type})"), ) .await?; let custom = list_objects(&pool, &schema, false, false, true).await?; let all = list_objects(&pool, &schema, true, true, true).await?; Ok::<_, String>((custom, all)) } .await; let cleanup = execute_query(&pool, &format!("DROP SCHEMA {schema_ident} CASCADE")).await; cleanup.expect("drop schema"); let (custom, all) = exercise.expect("exercise gaussdb custom type listing"); // GaussDB supports enum and composite user types but not domains. let mut type_names: Vec<&str> = custom.iter().map(|o| o.name.as_str()).collect(); type_names.sort_unstable(); assert_eq!(type_names, vec!["address", "status"], "custom types = {custom:?}"); for object in &custom { assert_eq!(object.object_type, "TYPE"); assert_eq!(object.schema.as_deref(), Some(schema.as_str())); } let status = custom.iter().find(|o| o.name == "status").expect("status type"); assert_eq!(status.comment.as_deref(), Some("订单状态"), "type comment was lost: {custom:?}"); let all_names: Vec<&str> = all.iter().map(|o| o.name.as_str()).collect(); assert!(all_names.contains(&"orders"), "table missing from full listing: {all:?}"); assert!( all_names.iter().all(|name| !name.starts_with('_')), "auto-generated array companions leaked into the listing: {all:?}" ); } #[tokio::test] #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a writable PostgreSQL database"] async fn postgres_list_custom_types_excludes_relation_row_types_and_arrays() { let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); let pool = connect(&url, Duration::from_secs(5)).await.expect("connect postgres"); let schema = format!("dbx_types_{}", uuid::Uuid::new_v4().simple()); let schema_ident = pg_quote_ident(&schema); let status_type = format!("{schema_ident}.status"); let email_domain = format!("{schema_ident}.email"); let address_type = format!("{schema_ident}.address"); let orders_table = format!("{schema_ident}.orders"); execute_query(&pool, &format!("CREATE SCHEMA {schema_ident}")).await.expect("create schema"); let exercise = async { execute_query(&pool, &format!("CREATE TYPE {status_type} AS ENUM ('draft', 'published')")).await?; execute_query(&pool, &format!("CREATE DOMAIN {email_domain} AS text CHECK (VALUE ~ '.+@.+')")).await?; execute_query(&pool, &format!("CREATE TYPE {address_type} AS (city text, zip text)")).await?; execute_query( &pool, &format!("CREATE TABLE {orders_table} (id bigint, state {status_type}, ship_to {address_type})"), ) .await?; let custom = list_objects(&pool, &schema, false, false, true).await?; let all = list_objects(&pool, &schema, true, true, true).await?; Ok::<_, String>((custom, all)) } .await; let cleanup = execute_query(&pool, &format!("DROP SCHEMA {schema_ident} CASCADE")).await; cleanup.expect("drop schema"); let (custom, all) = exercise.expect("exercise custom type listing"); // The type-only request returns exactly the user-created enum, domain // and composite type, sorted by name. The table's auto-generated row // type (orders) and the array companions (_status, _email, _address) // must stay out. let mut type_names: Vec<&str> = custom.iter().map(|o| o.name.as_str()).collect(); type_names.sort_unstable(); assert_eq!(type_names, vec!["address", "email", "status"], "custom types = {custom:?}"); for object in &custom { assert_eq!(object.object_type, "TYPE"); assert_eq!(object.schema.as_deref(), Some(schema.as_str())); } let all_names: Vec<&str> = all.iter().map(|o| o.name.as_str()).collect(); assert!(all_names.contains(&"orders"), "table missing from full listing: {all:?}"); assert!( all_names.iter().all(|name| !name.starts_with('_')), "auto-generated array companions leaked into the listing: {all:?}" ); } #[tokio::test] #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a writable PostgreSQL database"] async fn postgres_custom_type_details_reads_members_and_ddl() { let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); let pool = connect(&url, Duration::from_secs(5)).await.expect("connect postgres"); let schema = format!("dbx_details_{}", uuid::Uuid::new_v4().simple()); let schema_ident = pg_quote_ident(&schema); let status_type = format!("{schema_ident}.status"); let email_domain = format!("{schema_ident}.email"); let address_type = format!("{schema_ident}.address"); let price_range_type = format!("{schema_ident}.price_range"); let orders_table = format!("{schema_ident}.orders"); execute_query(&pool, &format!("CREATE SCHEMA {schema_ident}")).await.expect("create schema"); let exercise = async { execute_query(&pool, &format!("CREATE TYPE {status_type} AS ENUM ('draft', 'published', '已归档')")) .await?; execute_query(&pool, &format!("CREATE DOMAIN {email_domain} AS text DEFAULT '' CHECK (VALUE <> '')")) .await?; execute_query(&pool, &format!("CREATE TYPE {address_type} AS (city text, zip numeric(6))")).await?; execute_query(&pool, &format!("COMMENT ON TYPE {address_type} IS 'shipping address'")).await?; execute_query(&pool, &format!("COMMENT ON COLUMN {address_type}.city IS 'city name'")).await?; execute_query(&pool, &format!("CREATE TYPE {price_range_type} AS RANGE (subtype = numeric)")).await?; execute_query( &pool, &format!( "CREATE TABLE {orders_table} (state {status_type}, address {address_type}, email {email_domain})" ), ) .await?; let status = get_custom_type_details(&pool, &schema, "status").await?; let email = get_custom_type_details(&pool, &schema, "email").await?; let address = get_custom_type_details(&pool, &schema, "address").await?; let price_range = get_custom_type_details(&pool, &schema, "price_range").await?; let row_type_error = get_custom_type_details(&pool, &schema, "orders").await.err(); let array_type_error = get_custom_type_details(&pool, &schema, "_status").await.err(); Ok::<_, String>((status, email, address, price_range, row_type_error, array_type_error)) } .await; let cleanup = execute_query(&pool, &format!("DROP SCHEMA {schema_ident} CASCADE")).await; cleanup.expect("drop schema"); let (status, email, address, price_range, row_type_error, array_type_error) = exercise.expect("exercise custom type details"); assert_eq!(status.kind, CustomTypeKind::Enum); let enum_values: Vec<&str> = status.members.iter().map(|m| m.enum_value.as_deref().unwrap()).collect(); assert_eq!(enum_values, vec!["draft", "published", "已归档"]); let status_ddl = status.ddl.expect("enum ddl"); assert!(status_ddl.complete); assert!(status_ddl.sql.contains("AS ENUM ('draft', 'published', '已归档')"), "{}", status_ddl.sql); assert_eq!(email.kind, CustomTypeKind::Domain); assert_eq!(email.properties.base_type.as_deref(), Some("text")); assert!(email.properties.default.is_some(), "domain default was lost"); assert!( email.properties.domain_constraints.iter().any(|c| c.definition.contains("VALUE")), "domain constraint was lost: {:?}", email.properties.domain_constraints ); let email_ddl = email.ddl.expect("domain ddl"); assert!(email_ddl.complete); assert!(email_ddl.sql.contains("CREATE DOMAIN"), "{}", email_ddl.sql); assert_eq!(address.kind, CustomTypeKind::Composite); assert_eq!(address.comment.as_deref(), Some("shipping address")); assert_eq!(address.members.len(), 2); assert_eq!(address.members[0].name, "city"); assert_eq!(address.members[0].data_type, "text"); assert_eq!(address.members[0].comment.as_deref(), Some("city name")); assert_eq!(address.members[1].name, "zip"); // format_type normalizes numeric(6) to numeric(6,0) on real servers. assert_eq!(address.members[1].data_type, "numeric(6,0)"); let address_ddl = address.ddl.expect("composite ddl"); assert!(address_ddl.complete); assert!(address_ddl.sql.contains("\"city\" text"), "{}", address_ddl.sql); assert!( address_ddl.sql.contains(&format!("COMMENT ON COLUMN \"{schema}\".\"address\".\"city\" IS 'city name'")), "{}", address_ddl.sql ); assert_eq!(price_range.kind, CustomTypeKind::Range); assert_eq!(price_range.properties.range_subtype.as_deref(), Some("numeric")); let range_ddl = price_range.ddl.expect("range ddl"); assert!(range_ddl.complete); assert!(range_ddl.sql.contains("subtype = numeric"), "{}", range_ddl.sql); // Relation row types and array companions must be rejected, never // presented as independent composite types. assert!(row_type_error.is_some(), "relation row type must be rejected"); assert!(array_type_error.is_some(), "array companion must be rejected"); } #[tokio::test] #[ignore = "requires DBX_TEST_GAUSSDB_URL pointing at a writable GaussDB database"] async fn gaussdb_custom_type_details_reads_members_and_ddl() { let url = std::env::var("DBX_TEST_GAUSSDB_URL").expect("DBX_TEST_GAUSSDB_URL"); let pool = connect(&url, Duration::from_secs(5)).await.expect("connect gaussdb"); let schema = format!("dbx_gdetails_{}", uuid::Uuid::new_v4().simple()); let schema_ident = pg_quote_ident(&schema); let status_type = format!("{schema_ident}.status"); let address_type = format!("{schema_ident}.address"); let price_range_type = format!("{schema_ident}.price_range"); let orders_table = format!("{schema_ident}.orders"); execute_query(&pool, &format!("CREATE SCHEMA {schema_ident}")).await.expect("create schema"); let exercise = async { execute_query(&pool, &format!("CREATE TYPE {status_type} AS ENUM ('draft', 'published')")).await?; execute_query(&pool, &format!("CREATE TYPE {address_type} AS (city text, zip numeric(6))")).await?; execute_query(&pool, &format!("COMMENT ON COLUMN {address_type}.city IS 'city name'")).await?; execute_query(&pool, &format!("CREATE TYPE {price_range_type} AS RANGE (subtype = numeric)")).await?; execute_query(&pool, &format!("CREATE TABLE {orders_table} (state {status_type}, address {address_type})")) .await?; let status = get_custom_type_details(&pool, &schema, "status").await?; let address = get_custom_type_details(&pool, &schema, "address").await?; let price_range = get_custom_type_details(&pool, &schema, "price_range").await?; let row_type_error = get_custom_type_details(&pool, &schema, "orders").await.err(); let array_type_error = get_custom_type_details(&pool, &schema, "_status").await.err(); Ok::<_, String>((status, address, price_range, row_type_error, array_type_error)) } .await; let cleanup = execute_query(&pool, &format!("DROP SCHEMA {schema_ident} CASCADE")).await; if let Err(error) = cleanup { // GaussDB range types own system canonical functions, so CASCADE // cleanup is subject to instance privileges; that is a server // limitation, not a details-query failure. eprintln!("[gaussdb][custom-type-details][cleanup] {error}"); } let (status, address, price_range, row_type_error, array_type_error) = exercise.expect("exercise gaussdb custom type details"); assert_eq!(status.kind, CustomTypeKind::Enum); let enum_values: Vec<&str> = status.members.iter().map(|m| m.enum_value.as_deref().unwrap()).collect(); assert_eq!(enum_values, vec!["draft", "published"]); let status_ddl = status.ddl.expect("enum ddl"); assert!(status_ddl.complete); assert!(status_ddl.sql.contains("AS ENUM ('draft', 'published')"), "{}", status_ddl.sql); assert_eq!(address.kind, CustomTypeKind::Composite); assert_eq!(address.members.len(), 2); assert_eq!(address.members[0].name, "city"); assert_eq!(address.members[0].comment.as_deref(), Some("city name")); let address_ddl = address.ddl.expect("composite ddl"); assert!(address_ddl.complete); assert!(address_ddl.sql.contains("\"city\" text"), "{}", address_ddl.sql); assert_eq!(price_range.kind, CustomTypeKind::Range); assert_eq!(price_range.properties.range_subtype.as_deref(), Some("numeric")); let range_ddl = price_range.ddl.expect("range ddl"); assert!(range_ddl.complete); assert!(range_ddl.sql.contains("subtype = numeric"), "{}", range_ddl.sql); assert!(row_type_error.is_some(), "relation row type must be rejected"); assert!(array_type_error.is_some(), "array companion must be rejected"); } #[tokio::test] #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a writable PostgreSQL 18 database"] async fn postgres_custom_type_fallback_refreshes_stale_cached_metadata() { let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); let pool_a = connect(&url, Duration::from_secs(5)).await.expect("connect postgres pool A"); let pool_b = connect(&url, Duration::from_secs(5)).await.expect("connect postgres pool B"); assert_postgres_18(&pool_a).await; let schema = format!("dbx_custom_stale_{}", uuid::Uuid::new_v4().simple()); let schema_ident = pg_quote_ident(&schema); let payload_type = format!("{schema_ident}.payload"); let view = format!("{schema_ident}.cached_payload"); let view_sql = format!("SELECT payload FROM {view}"); execute_query(&pool_a, &format!("CREATE SCHEMA {schema_ident}")).await.expect("create schema"); let exercise = async { execute_query(&pool_a, &format!("CREATE TYPE {payload_type} AS (id integer, label text)")).await?; execute_query(&pool_a, &format!("CREATE VIEW {view} AS SELECT ROW(7, 'alpha')::{payload_type} AS payload")) .await?; let custom = execute_query(&pool_a, &view_sql).await?; execute_query(&pool_b, &format!("DROP VIEW {view}")).await?; execute_query(&pool_b, &format!("CREATE VIEW {view} AS SELECT 42::int4 AS payload")).await?; let builtin = execute_query(&pool_a, &view_sql).await?; Ok::<_, String>((custom, builtin)) } .await; let cleanup = execute_query(&pool_a, &format!("DROP SCHEMA {schema_ident} CASCADE")).await; cleanup.expect("drop schema"); let (custom, builtin) = exercise.expect("exercise stale cached custom metadata"); assert_eq!(custom.column_types, vec!["payload"]); assert_eq!(custom.rows[0][0], serde_json::Value::String("(7,alpha)".to_string())); assert_eq!(builtin.column_types, vec!["int4"]); assert_eq!(builtin.rows[0][0], serde_json::Value::Number(42.into())); } #[tokio::test] #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a writable PostgreSQL 18 database"] async fn postgres_text_fallback_stops_before_late_row_error_at_limit() { let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); let pool = connect(&url, Duration::from_secs(5)).await.expect("connect postgres"); assert_postgres_18(&pool).await; let schema = format!("dbx_custom_limit_{}", uuid::Uuid::new_v4().simple()); let schema_ident = pg_quote_ident(&schema); let payload_type = format!("{schema_ident}.payload"); let fail_after_two = format!("{schema_ident}.fail_after_two"); execute_query(&pool, &format!("CREATE SCHEMA {schema_ident}")).await.expect("create schema"); let exercise = async { execute_query(&pool, &format!("CREATE TYPE {payload_type} AS (id integer)")).await?; execute_query( &pool, &format!( "CREATE FUNCTION {fail_after_two}(i integer) RETURNS integer LANGUAGE plpgsql AS $$ \ BEGIN IF i >= 2 THEN RAISE EXCEPTION 'late row failure'; END IF; RETURN i; END $$" ), ) .await?; let client = checkout_postgres_client(&pool, None, Duration::from_secs(5)).await?; let custom_sql = format!( "SELECT ROW({fail_after_two}(i))::{payload_type} AS payload \ FROM generate_series(1, 2) AS series(i)" ); let limited = execute_select_query(&client, &custom_sql, Instant::now(), 1).await; let cancelled = AtomicBool::new(false); let mut exported_rows = Vec::new(); let exported = stream_query_rows_on_client(&client, &custom_sql, Some(1), &cancelled, &mut |row| { exported_rows.push(row.to_vec()); Ok(()) }) .await; let recovery = execute_select_query(&client, "SELECT 1::int4 AS value", Instant::now(), 1).await; drop(client); Ok::<_, String>((limited, exported, exported_rows, recovery)) } .await; let cleanup = execute_query(&pool, &format!("DROP SCHEMA {schema_ident} CASCADE")).await; cleanup.expect("drop schema"); let (limited, exported, exported_rows, recovery) = exercise.expect("set up late row error query"); let limited = limited.expect("query should stop before late row error"); let exported = exported.expect("streamed export should stop before late row error"); let recovery = recovery.expect("connection should remain reusable"); assert_eq!(limited.column_types, vec!["payload"]); assert_eq!(limited.rows, vec![vec![serde_json::Value::String("(1)".to_string())]]); assert!(limited.truncated); assert_eq!(exported, 1); assert_eq!(exported_rows, vec![vec![serde_json::Value::String("(1)".to_string())]]); assert_eq!(recovery.column_types, vec!["int4"]); assert_eq!(recovery.rows[0][0], serde_json::Value::Number(1.into())); } fn state_enum_values(columns: &[ColumnInfo]) -> Option> { columns.iter().find(|column| column.name == "state").and_then(|column| column.enum_values.clone()) } // --- pg_quote_ident --- #[test] fn pg_system_u32_decodes_catalog_integer_types() { let raw = 42_u32.to_be_bytes(); assert_eq!(u32::from_sql(&Type::OID, &raw).unwrap(), 42); assert_eq!(PgSystemU32::from_sql(&Type::XID, &raw).unwrap().0, 42); assert_eq!(PgSystemU32::from_sql(&Type::CID, &raw).unwrap().0, 42); assert!(u32::accepts(&Type::OID)); assert!(PgSystemU32::accepts(&Type::XID)); assert!(PgSystemU32::accepts(&Type::CID)); assert!(!PgSystemU32::accepts(&Type::OID)); assert!(!PgSystemU32::accepts(&Type::INT4)); } #[test] fn pg_point_decodes_binary_coordinates_as_text() { let mut raw = Vec::new(); raw.extend_from_slice(&(-194.0_f64).to_be_bytes()); raw.extend_from_slice(&53.0_f64.to_be_bytes()); let point = PgPoint::from_sql(&Type::POINT, &raw).unwrap(); assert_eq!(point, PgPoint { x: -194.0, y: 53.0 }); assert_eq!(format_pg_point(point), "(-194,53)"); assert_eq!(format_pg_point(PgPoint { x: f64::NEG_INFINITY, y: f64::INFINITY }), "(-Infinity,Infinity)"); assert!(PgPoint::from_sql(&Type::POINT, &[0; 15]).is_err()); assert!(PgPoint::accepts(&Type::POINT)); assert!(!PgPoint::accepts(&Type::BYTEA)); } #[test] fn pg_any_string_accepts_all_types_and_decodes_utf8() { // Accepts any type — built-in, custom enum OIDs, domains, etc. assert!(PgAnyString::accepts(&Type::TEXT)); assert!(PgAnyString::accepts(&Type::INT4)); assert!(PgAnyString::accepts(&Type::UNKNOWN)); assert!(PgAnyString::accepts(&Type::OID)); assert!(PgAnyString::accepts(&Type::BOOL)); let label = PgAnyString::from_sql(&Type::UNKNOWN, b"pending").unwrap(); assert_eq!(label.0, "pending"); let label = PgAnyString::from_sql(&Type::UNKNOWN, b"hello world").unwrap(); assert_eq!(label.0, "hello world"); // Non-UTF-8 bytes should fail gracefully assert!(PgAnyString::from_sql(&Type::UNKNOWN, &[0xFF, 0xFE, 0xFD]).is_err()); } #[test] fn pg_raw_bytes_accepts_all_types_and_preserves_binary_payloads() { assert!(PgRawBytes::accepts(&Type::TEXT)); assert!(PgRawBytes::accepts(&Type::UNKNOWN)); assert!(PgRawBytes::accepts(&Type::OID)); let raw = PgRawBytes::from_sql(&Type::UNKNOWN, &[0x01, 0xAB, 0xFF]).unwrap(); assert_eq!(raw.0, vec![0x01, 0xAB, 0xFF]); } #[test] fn decode_bool_bytes_handles_standard_and_gaussdb_encodings() { // Standard PostgreSQL binary boolean: 0x00 / 0x01 assert_eq!(decode_bool_bytes(&[0x00]), Some(false)); assert_eq!(decode_bool_bytes(&[0x01]), Some(true)); // GaussDB binary boolean: ASCII 't' (0x74) / 'f' (0x66) assert_eq!(decode_bool_bytes(&[0x74]), Some(true)); assert_eq!(decode_bool_bytes(&[0x66]), Some(false)); assert_eq!(decode_bool_bytes(b"t"), Some(true)); assert_eq!(decode_bool_bytes(b"f"), Some(false)); assert_eq!(decode_bool_bytes(b"T"), Some(true)); assert_eq!(decode_bool_bytes(b"F"), Some(false)); // Unrecognized encodings return None assert_eq!(decode_bool_bytes(&[0x02]), None); assert_eq!(decode_bool_bytes(&[0x74, 0x66]), None); assert_eq!(decode_bool_bytes(&[]), None); } #[test] fn raw_gaussdb_boolean_takes_precedence_over_standard_decoder() { assert_eq!(decode_bool_candidates(Some(b"f"), Some(true)), Some(false)); assert_eq!(decode_bool_candidates(Some(b"t"), Some(true)), Some(true)); assert_eq!(decode_bool_candidates(Some(&[0x00]), Some(true)), Some(false)); assert_eq!(decode_bool_candidates(Some(&[0x01]), Some(false)), Some(true)); assert_eq!(decode_bool_candidates(Some(&[0x02]), Some(false)), Some(false)); } #[test] fn postgres_foreign_keys_sql_selects_referential_actions() { let sql = postgres_foreign_keys_sql(); assert!(sql.contains("pg_catalog.pg_constraint")); assert!(sql.contains("con.confupdtype::text AS on_update_raw")); assert!(sql.contains("con.confdeltype::text AS on_delete_raw")); assert!(sql.contains("confdeltype")); assert!(sql.contains("ref_n.nspname")); assert!(sql.contains("unnest(con.conkey) WITH ORDINALITY")); assert!(sql.contains("n.nspname = $1 AND c.relname = $2")); assert!(sql.contains("con.contype = 'f'")); assert!(sql.contains("NOT a.attisdropped")); assert!(sql.contains("JOIN LATERAL")); assert!(!sql.contains("information_schema")); } #[test] fn postgres_foreign_key_metadata_has_legacy_catalog_fallback() { // The compat tiers exist for pre-9.4 servers, so they must avoid WITH // ORDINALITY (a 9.4 feature) and pair conkey/confkey positions through // generate_series + plain array subscripts instead, mirroring the index // compat invariant. for compat_sql in [postgres_foreign_keys_compat_sql(), postgres_foreign_keys_for_relations_compat_sql()] { assert!(!compat_sql.contains("WITH ORDINALITY")); assert!(compat_sql.contains("generate_series")); assert!(compat_sql.contains("array_length(con.conkey, 1)")); assert!(compat_sql.contains("con.contype = 'f'")); assert!(compat_sql.contains("NOT a.attisdropped")); assert!(!compat_sql.contains("information_schema")); } // OpenGauss resolves attribute numbers in Rust, so its SQL must not // expand catalog arrays at all. let opengauss_sql = opengauss_foreign_keys_sql(); assert!(!opengauss_sql.contains("WITH ORDINALITY")); assert!(!opengauss_sql.contains("unnest")); assert!(opengauss_sql.contains("con.conkey::text")); assert!(opengauss_sql.contains("con.confkey::text")); } #[test] fn postgres_table_dependencies_sql_batches_schema_foreign_keys() { let sql = postgres_table_dependencies_sql(); let compat_sql = postgres_table_dependencies_compat_sql(); assert!(sql.contains("pg_catalog.pg_constraint")); assert!(sql.contains("con.contype = 'f'")); assert!(sql.contains("child_schema.nspname = $1")); assert!(sql.contains("parent_schema.nspname = $1")); assert!(!sql.contains("information_schema")); assert!(sql.contains("pg_catalog.pg_inherits")); assert!(sql.contains("child.relispartition")); // 兼容版面向 9.x:不能引用 relispartition,但 INHERITS 边要保留 // (旧式子表同样要先建父表)。 assert!(!compat_sql.contains("relispartition")); assert!(compat_sql.contains("pg_catalog.pg_inherits")); assert!(compat_sql.contains("con.contype = 'f'")); assert!(compat_sql.contains("ORDER BY table_name, ref_table")); } #[test] fn decodes_tsvector_binary_output() { let raw = [ 0, 0, 0, 2, b'b', b'a', b'c', b'k', b'\\', b's', b'l', b'a', b's', b'h', 0, 0, 1, 0x80, 0x03, b'o', b'\'', b'c', b'l', b'o', b'c', b'k', 0, 0, 2, 0, 1, 0xc0, 0x02, ]; assert_eq!(decode_tsvector_bytes(&raw).as_deref(), Some("'back\\\\slash':3B 'o''clock':1,2A")); } fn encode_pgvector_bytes(values: &[f32]) -> Vec { let dims = u16::try_from(values.len()).expect("dim fits u16"); let mut raw = Vec::with_capacity(4 + values.len() * 4); raw.extend_from_slice(&dims.to_be_bytes()); raw.extend_from_slice(&0u16.to_be_bytes()); for value in values { raw.extend_from_slice(&value.to_be_bytes()); } raw } #[test] fn decodes_pgvector_binary_output() { let values = [0.1f32, -2.5f32, 1.2345679e-5f32]; let decoded = decode_pgvector_bytes(&encode_pgvector_bytes(&values)).expect("decode vector"); assert_eq!(decoded, values); } #[test] fn pgvector_element_number_round_trips_full_f32_precision() { let values = [0.1f32, 0.12345679f32, 1.2345679e-5f32, -0.00012345679f32, 1.2345678f32, 1e20f32]; for value in values { let json = pg_vector_element_number(value); let text = json.to_string(); let restored: f32 = text.parse().expect("json number parses as f32"); let rounded_six = ((value as f64 * 1_000_000.0).round() / 1_000_000.0) as f32; // Display text must recover the exact stored float4 bits. assert_eq!(restored, value, "lost f32 precision for {value} -> {text}"); // Fixed 6-decimal rounding is what caused #3931; reject that path when it differs. if rounded_six != value { assert_ne!(restored, rounded_six, "still clamped to 6 decimals for {value}"); } } } #[test] fn pgvector_binary_to_json_preserves_component_precision() { let values = [0.12345679f32, 1.2345679e-5f32, -2.5f32]; let decoded = decode_pgvector_bytes(&encode_pgvector_bytes(&values)).expect("decode vector"); let json = serde_json::Value::Array(decoded.into_iter().map(pg_vector_element_number).collect()); let arr = json.as_array().expect("vector json array"); assert_eq!(arr.len(), values.len()); for (component, expected) in arr.iter().zip(values) { let restored: f32 = component.to_string().parse().expect("component parses as f32"); assert_eq!(restored, expected); } } fn decode_hex(hex: &str) -> Vec { assert_eq!(hex.len() % 2, 0, "hex input must have an even number of chars"); (0..hex.len()).step_by(2).map(|idx| u8::from_str_radix(&hex[idx..idx + 2], 16).unwrap()).collect() } #[test] fn decodes_postgres_inet_binary_output() { assert_eq!( decode_pg_network_address_bytes(&decode_hex("02200004c0a8010a"), false).as_deref(), Some("192.168.1.10") ); assert_eq!( decode_pg_network_address_bytes(&decode_hex("0310001020010db8abcd00120000000000000001"), false).as_deref(), Some("2001:db8:abcd:12::1/16") ); assert_eq!( decode_pg_network_address_bytes(&decode_hex("0340001020010db8abcd00120000000000000001"), false).as_deref(), Some("2001:db8:abcd:12::1/64") ); } #[test] fn decodes_postgres_cidr_binary_output() { assert_eq!( decode_pg_network_address_bytes(&decode_hex("02180104c0a80100"), true).as_deref(), Some("192.168.1.0/24") ); assert_eq!( decode_pg_network_address_bytes(&decode_hex("02200104c0a8010a"), true).as_deref(), Some("192.168.1.10/32") ); assert_eq!( decode_pg_network_address_bytes(&decode_hex("0380011000000000000000000000000000000001"), true).as_deref(), Some("::1/128") ); } #[test] fn rejects_invalid_postgres_network_binary_output() { assert_eq!(decode_pg_network_address_bytes(&[], false), None); assert_eq!(decode_pg_network_address_bytes(&decode_hex("04200004c0a8010a"), false), None); assert_eq!(decode_pg_network_address_bytes(&decode_hex("02210004c0a8010a"), false), None); assert_eq!(decode_pg_network_address_bytes(&decode_hex("02200004c0a801"), false), None); } #[test] fn decodes_postgres_macaddr_binary_output() { assert_eq!(decode_pg_macaddr_bytes(&decode_hex("08002b010203")).as_deref(), Some("08:00:2b:01:02:03")); assert_eq!( decode_pg_macaddr_bytes(&decode_hex("08002bfffe010203")).as_deref(), Some("08:00:2b:ff:fe:01:02:03") ); assert_eq!(decode_pg_macaddr_bytes(&decode_hex("08002b")), None); } #[test] fn decodes_postgres_bit_string_binary_output() { assert_eq!(decode_pg_bit_string_bytes(&decode_hex("00000005a8")).as_deref(), Some("10101")); assert_eq!(decode_pg_bit_string_bytes(&decode_hex("00000009a880")).as_deref(), Some("101010001")); assert_eq!(decode_pg_bit_string_bytes(&decode_hex("00000000")).as_deref(), Some("")); assert_eq!(decode_pg_bit_string_bytes(&decode_hex("00000005a8ff")), None); assert_eq!(decode_pg_bit_string_bytes(&decode_hex("ffffffff")), None); } #[test] fn ewkb_point_with_srid_formats_as_wkt() { let raw = decode_hex("0101000020E6100000C520B07268195D404E62105839F44340"); assert_eq!(super::super::wkb::wkb_to_wkt(&raw), Some("POINT(116.397 39.908)".to_string())); } #[test] fn ewkb_multi_polygon_formats_as_wkt() { let raw = decode_hex( "0106000020E610000002000000010300000001000000050000000000000000005D4000000000000044400000000000405D4000000000000044400000000000405D4000000000008044400000000000005D4000000000008044400000000000005D400000000000004440010300000001000000050000000000000000805D4000000000008043400000000000C05D4000000000008043400000000000C05D4000000000000044400000000000805D4000000000000044400000000000805D400000000000804340", ); assert_eq!( super::super::wkb::wkb_to_wkt(&raw), Some( "MULTIPOLYGON(((116 40,117 40,117 41,116 41,116 40)),((118 39,119 39,119 40,118 40,118 39)))" .to_string() ) ); } #[test] fn ewkb_geometry_collection_formats_as_wkt() { let raw = decode_hex( "0107000020E61000000200000001010000000000000000005D4000000000000044400102000000020000000000000000405D4000000000008044400000000000805D400000000000004540", ); assert_eq!( super::super::wkb::wkb_to_wkt(&raw), Some("GEOMETRYCOLLECTION(POINT(116 40),LINESTRING(117 41,118 42))".to_string()) ); } #[test] fn pg_optional_array_to_json_preserves_text_values_and_nulls() { let value = pg_optional_array_to_json( vec![Some("productManager".to_string()), None, Some("projectOwner".to_string())], serde_json::Value::String, ); assert_eq!(value, serde_json::json!(["productManager", null, "projectOwner"])); } #[test] fn pg_quote_ident_plain_identifier() { assert_eq!(pg_quote_ident("public"), "\"public\""); } #[test] fn pg_quote_ident_escapes_double_quotes() { assert_eq!(pg_quote_ident("my\"schema"), "\"my\"\"schema\""); } #[test] fn pg_quote_ident_empty_string() { assert_eq!(pg_quote_ident(""), "\"\""); } #[test] fn pg_quote_ident_special_chars() { // PostgreSQL allows many special chars in quoted identifiers let ident = "my schema with spaces"; assert_eq!(pg_quote_ident(ident), "\"my schema with spaces\""); } #[test] fn pg_quote_ident_injection_attempt() { // A malicious schema name that tries to break out of quotes let malicious = r#"public"; DROP TABLE users; --"#; let escaped = pg_quote_ident(malicious); // Double quotes should be doubled, not breaking out assert_eq!(escaped, r#""public""; DROP TABLE users; --""#); assert!(escaped.matches('"').count().is_multiple_of(2), "quote count should be even"); } #[test] fn postgres_table_access_reads_complete_catalog_acls() { assert!(POSTGRES_TABLE_OWNER_SQL.contains("acldefault('r', c.relowner)")); assert!( POSTGRES_TABLE_ACL_PRIVILEGES_SQL.contains("COALESCE(c.relacl, pg_catalog.acldefault('r', c.relowner))") ); assert!(POSTGRES_COLUMN_ACL_PRIVILEGES_SQL.contains("aclexplode(a.attacl)")); assert!(POSTGRES_TABLE_ACL_PRIVILEGES_SQL.contains("acl.grantee = 0 THEN 'PUBLIC'")); assert!(POSTGRES_COLUMN_ACL_PRIVILEGES_SQL.contains("acl.grantee = 0 THEN 'PUBLIC'")); assert!(POSTGRES_TABLE_ACL_PRIVILEGES_SQL.contains("pg_get_userbyid(acl.grantor)")); assert!(POSTGRES_COLUMN_ACL_PRIVILEGES_SQL.contains("pg_get_userbyid(acl.grantor)")); } // --- query_result_row_limit --- #[test] fn row_limit_uses_max_rows_when_present() { assert_eq!(query_result_row_limit(Some(50)), 50); } #[test] fn row_limit_falls_back_to_default() { let default = crate::query::MAX_ROWS; assert_eq!(query_result_row_limit(None), default); } #[test] fn row_limit_clamps_zero_to_one() { assert_eq!(query_result_row_limit(Some(0)), 1); } #[test] fn row_limit_allows_max_rows_override() { assert_eq!(query_result_row_limit(Some(5)), 5); } #[test] fn timestamptz_display_preserves_local_offset() { let text = format_pg_timestamptz(Local::now()); assert!(!text.ends_with("+00:00") || Local::now().offset().local_minus_utc() == 0); } // --- validate_postgres_ssl_paths --- #[test] fn ssl_validation_passes_for_clean_url() { assert!(validate_postgres_ssl_paths("postgres://localhost/db").is_ok()); } #[test] fn ssl_validation_passes_for_url_without_query() { assert!(validate_postgres_ssl_paths("host=localhost dbname=test").is_ok()); } #[test] fn ssl_validation_passes_for_irrelevant_params() { assert!(validate_postgres_ssl_paths("postgres://localhost/db?sslmode=require&connect_timeout=10").is_ok()); } #[test] fn ssl_validation_rejects_nonexistent_sslcert_path() { let result = validate_postgres_ssl_paths("postgres://localhost/db?sslcert=/nonexistent/path/cert.pem"); assert!(result.is_err()); assert!(result.unwrap_err().contains("sslcert"), "error should mention sslcert"); } #[test] fn ssl_validation_rejects_nonexistent_sslkey_path() { let result = validate_postgres_ssl_paths("postgres://localhost/db?sslkey=/nonexistent/path/key.pem"); assert!(result.is_err()); assert!(result.unwrap_err().contains("sslkey"), "error should mention sslkey"); } #[test] fn ssl_validation_rejects_nonexistent_sslrootcert_path() { let result = validate_postgres_ssl_paths("postgres://localhost/db?sslrootcert=/nonexistent/path/root.crt"); assert!(result.is_err()); assert!(result.unwrap_err().contains("sslrootcert"), "error should mention sslrootcert"); } #[test] fn ssl_validation_rejects_path_traversal_in_sslcert() { let result = validate_postgres_ssl_paths("postgres://localhost/db?sslcert=../../../etc/passwd"); assert!(result.is_err()); } #[test] fn ssl_validation_handles_url_encoded_ssl_param() { // %2F = '/', so sslcert=%2Ftmp%2Fcert.pem means sslcert=/tmp/cert.pem let result = validate_postgres_ssl_paths("postgres://localhost/db?sslcert=%2Fnonexistent%2Fcert.pem"); assert!(result.is_err()); } #[test] fn ssl_validation_handles_multiple_params() { let result = validate_postgres_ssl_paths("postgres://localhost/db?sslmode=require&sslcert=/nonexistent/cert.pem"); assert!(result.is_err()); } #[test] fn postgres_connection_url_strips_ssl_file_params_before_driver_parse() { let dir = std::env::temp_dir(); let cert = dir.join(format!("dbx-postgres-cert-{}.pem", std::process::id())); let key = dir.join(format!("dbx-postgres-key-{}.pem", std::process::id())); let root = dir.join(format!("dbx-postgres-root-{}.pem", std::process::id())); std::fs::write(&cert, "not a real cert").unwrap(); std::fs::write(&key, "not a real key").unwrap(); std::fs::write(&root, "not a real root").unwrap(); let url = format!( "postgres://localhost/db?sslmode=verify-full&sslcert={}&sslkey={}&sslrootcert={}&application_name=dbx", cert.display(), key.display(), root.display() ); let parsed = postgres_connection_url(&url).unwrap(); assert_eq!(parsed.url, "postgres://localhost/db?sslmode=require&application_name=dbx"); assert_eq!(parsed.ssl_files.sslcert.as_deref(), Some(cert.to_str().unwrap())); assert_eq!(parsed.ssl_files.sslkey.as_deref(), Some(key.to_str().unwrap())); assert_eq!(parsed.ssl_files.sslrootcert.as_deref(), Some(root.to_str().unwrap())); assert!(!parsed.accepts_invalid_certs); assert!(parsed.verifies_hostname); tokio_postgres::Config::from_str(&parsed.url).unwrap(); let _ = std::fs::remove_file(cert); let _ = std::fs::remove_file(key); let _ = std::fs::remove_file(root); } #[test] fn postgres_connection_url_keeps_verify_ca_ca_only_semantics() { let parsed = postgres_connection_url("postgres://localhost/db?sslmode=verify-ca").unwrap(); assert_eq!(parsed.url, "postgres://localhost/db?sslmode=require"); assert!(!parsed.accepts_invalid_certs); assert!(!parsed.verifies_hostname); } #[test] fn postgres_connection_url_extracts_legacy_tls_before_driver_parse() { let parsed = postgres_connection_url("postgres://localhost/db?sslmode=require&legacy_tls=true&application_name=dbx") .unwrap(); assert_eq!(parsed.url, "postgres://localhost/db?sslmode=require&application_name=dbx"); assert!(parsed.legacy_tls); tokio_postgres::Config::from_str(&parsed.url).unwrap(); let parsed = postgres_connection_url("postgres://localhost/db?legacy_tls=false").unwrap(); assert_eq!(parsed.url, "postgres://localhost/db"); assert!(!parsed.legacy_tls); } #[test] fn postgres_connection_url_normalizes_channel_binding_require_to_prefer() { let parsed = postgres_connection_url("postgres://localhost/db?sslmode=require&channel_binding=require").unwrap(); assert_eq!(parsed.url, "postgres://localhost/db?sslmode=require&channel_binding=prefer"); // The sanitized URL must be parseable by the driver tokio_postgres::Config::from_str(&parsed.url).unwrap(); } #[test] fn postgres_connection_url_keeps_channel_binding_prefer() { let parsed = postgres_connection_url("postgres://localhost/db?channel_binding=prefer").unwrap(); assert_eq!(parsed.url, "postgres://localhost/db?channel_binding=prefer"); tokio_postgres::Config::from_str(&parsed.url).unwrap(); } #[test] fn postgres_tls_rejects_unpaired_client_cert_and_key() { let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); let pg_config = tokio_postgres::Config::from_str("postgres://localhost/db?sslmode=require").unwrap(); let ssl_files = PostgresSslFiles { sslcert: Some("/tmp/client.crt".to_string()), sslkey: None, sslrootcert: None }; let error = match postgres_tls_config(&pg_config, &ssl_files, true, false) { Ok(_) => panic!("expected missing sslkey to fail"), Err(error) => error, }; assert!(error.contains("sslkey")); } #[test] fn postgres_tls_rejects_empty_client_certificate() { let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); let dir = std::env::temp_dir(); let suffix = uuid::Uuid::new_v4().simple(); let cert = dir.join(format!("dbx-postgres-empty-client-{suffix}.crt")); let key = dir.join(format!("dbx-postgres-empty-client-{suffix}.key")); std::fs::write(&cert, []).unwrap(); std::fs::write(&key, []).unwrap(); let pg_config = tokio_postgres::Config::from_str("postgres://localhost/db?sslmode=require").unwrap(); let ssl_files = PostgresSslFiles { sslcert: Some(cert.to_string_lossy().into_owned()), sslkey: Some(key.to_string_lossy().into_owned()), sslrootcert: None, }; let error = postgres_tls_config(&pg_config, &ssl_files, true, false) .expect_err("empty client certificate should fail before connecting"); assert!(error.contains("no certificates"), "{error}"); let _ = std::fs::remove_file(cert); let _ = std::fs::remove_file(key); } #[test] fn postgres_tls_rejects_malformed_private_key() { let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); let dir = std::env::temp_dir(); let suffix = uuid::Uuid::new_v4().simple(); let cert = dir.join(format!("dbx-postgres-client-{suffix}.crt")); let key = dir.join(format!("dbx-postgres-malformed-client-{suffix}.key")); std::fs::write(&cert, "-----BEGIN CERTIFICATE-----\nMAA=\n-----END CERTIFICATE-----\n").unwrap(); std::fs::write(&key, "-----BEGIN PRIVATE KEY-----\nMAA=\n-----END PRIVATE KEY-----\n").unwrap(); let pg_config = tokio_postgres::Config::from_str("postgres://localhost/db?sslmode=require").unwrap(); let ssl_files = PostgresSslFiles { sslcert: Some(cert.to_string_lossy().into_owned()), sslkey: Some(key.to_string_lossy().into_owned()), sslrootcert: None, }; let error = postgres_tls_config(&pg_config, &ssl_files, true, false) .expect_err("malformed private key should fail before connecting"); assert!(error.contains("client private key is invalid"), "{error}"); let _ = std::fs::remove_file(cert); let _ = std::fs::remove_file(key); } #[tokio::test] #[ignore = "requires DBX_TEST_POSTGRES_MTLS_URL pointing at PostgreSQL 14 with required client certificates"] async fn postgres_mtls_accepts_configured_client_identity() { let url = std::env::var("DBX_TEST_POSTGRES_MTLS_URL").expect("DBX_TEST_POSTGRES_MTLS_URL"); let pool = connect(&url, Duration::from_secs(5)).await.expect("connect with client certificate"); let client = pool.get().await.expect("get PostgreSQL mTLS client"); let row = client .query_one( "SELECT current_user, ssl, current_setting('server_version_num')::int \ FROM pg_stat_ssl WHERE pid = pg_backend_pid()", &[], ) .await .expect("query PostgreSQL TLS session"); assert_eq!(row.get::<_, String>(0), "db"); assert!(row.get::<_, bool>(1)); assert!((140_000..150_000).contains(&row.get::<_, i32>(2))); } #[test] #[ignore = "requires DBX_TEST_POSTGRES_MTLS_URL with a valid client certificate"] fn postgres_mtls_cancel_connector_reuses_client_identity() { let url = std::env::var("DBX_TEST_POSTGRES_MTLS_URL").expect("DBX_TEST_POSTGRES_MTLS_URL"); let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); let context = build_postgres_cancel_context(&url).expect("build PostgreSQL TLS cancel context"); make_rustls_connect_from_context(&context).expect("rebuild TLS connector with client certificate"); } #[tokio::test] #[ignore = "requires DBX_TEST_POSTGRES_MTLS_REJECT_URL that the TLS server must reject"] async fn postgres_mtls_rejects_invalid_tls_identity() { let url = std::env::var("DBX_TEST_POSTGRES_MTLS_REJECT_URL").expect("DBX_TEST_POSTGRES_MTLS_REJECT_URL"); let error = connect(&url, Duration::from_secs(5)).await.expect_err("invalid TLS identity must not authenticate"); assert!(error.contains("PostgreSQL connection failed"), "{error}"); } #[test] fn postgres_accept_all_tls_signature_does_not_parse_unverified_cert() { let verifier = NoPostgresCertVerification { provider: Arc::new(rustls::crypto::aws_lc_rs::default_provider()) }; let malformed_cert = CertificateDer::from(vec![0x30, 0x03, 0x02, 0x01, 0x00]); assert!(verifier.accept_tls_signature_for_unverified_cert(&malformed_cert).is_ok()); } #[test] fn inject_postgres_keepalive_params_preserves_url_fragment() { let url = "postgres://localhost/app?sslmode=require#read-only"; assert_eq!( inject_postgres_keepalive_params(url), "postgres://localhost/app?sslmode=require&keepalives=1&keepalives_idle=30&keepalives_interval=10&keepalives_retries=3#read-only" ); } #[test] fn postgres_cancel_attempt_timeout_is_single_budget() { assert_eq!(postgres_cancel_attempt_timeout(Duration::from_secs(5), None), Duration::from_secs(5)); assert_eq!( postgres_cancel_attempt_timeout( Duration::from_secs(5), Some(&PostgresCancelContext { ssl_files: PostgresSslFiles::default(), accepts_invalid_certs: true, verifies_hostname: false, ssl_mode: SslMode::Require, legacy_tls: false, }) ), Duration::from_secs(5) ); } #[test] fn postgres_checkout_timeout_respects_pool_configuration() { let manager = deadpool_postgres::Manager::new(tokio_postgres::Config::new(), NoTls); let pool = Pool::builder(manager) .runtime(Runtime::Tokio1) .wait_timeout(Some(Duration::from_secs(10))) .create_timeout(Some(Duration::from_secs(10))) .recycle_timeout(Some(Duration::from_secs(10))) .build() .expect("build postgres pool"); assert_eq!(effective_postgres_checkout_timeout(&pool, Duration::from_secs(5)), Duration::from_secs(10)); assert_eq!(effective_postgres_checkout_timeout(&pool, Duration::from_secs(15)), Duration::from_secs(15)); } #[tokio::test] async fn postgres_checkout_fault_injection_classifies_full_pool_as_wait() { assert!(tokio::time::timeout(Duration::from_millis(1), std::future::pending::<()>()).await.is_err()); assert_eq!( classify_postgres_checkout_stage(deadpool_postgres::Status { max_size: 1, size: 1, available: 0, waiting: 1 }), PostgresCheckoutStage::Wait ); } #[tokio::test] async fn postgres_checkout_fault_injection_preserves_create_and_recycle_stages() { assert!(tokio::time::timeout(Duration::from_millis(1), std::future::pending::<()>()).await.is_err()); assert_eq!( classify_postgres_checkout_stage(deadpool_postgres::Status { max_size: 2, size: 1, available: 0, waiting: 0 }), PostgresCheckoutStage::Create ); assert_eq!( classify_postgres_checkout_stage(deadpool_postgres::Status { max_size: 1, size: 1, available: 1, waiting: 0 }), PostgresCheckoutStage::Recycle ); } #[tokio::test] async fn postgres_session_setup_reports_its_own_timeout_stage() { let error = postgres_session_setup_with_timeout(Duration::from_millis(1), std::future::pending::>()) .await .expect_err("pending session setup should time out"); assert!(error.starts_with("PostgreSQL session setup connection timed out"), "{error}"); } #[tokio::test] async fn postgres_row_query_timeout_tracks_inactivity_instead_of_total_duration() { let progress_clock = Arc::new(StreamProgressClock::new()); let progress_clock_for_query = progress_clock.clone(); let future = async { for _ in 0..4 { tokio::time::sleep(Duration::from_millis(40)).await; progress_clock_for_query.mark(); } Ok::<_, String>(()) }; let result = await_stream_with_progress_timeout( future, Some(Duration::from_millis(100)), progress_clock, None, "query inactivity timeout".to_string(), ) .await; assert_eq!(result, Ok(())); } #[tokio::test] async fn postgres_row_query_timeout_still_fires_after_no_progress() { let progress_clock = Arc::new(StreamProgressClock::new()); let error = await_stream_with_progress_timeout( std::future::pending::>(), Some(Duration::from_millis(10)), progress_clock, None, "query inactivity timeout".to_string(), ) .await .expect_err("a stalled query should time out"); assert_eq!(error, "query inactivity timeout"); } #[tokio::test] async fn postgres_row_query_timeout_covers_stall_after_last_row() { let progress_clock = Arc::new(StreamProgressClock::new()); let progress_clock_for_query = progress_clock.clone(); let error = await_stream_with_progress_timeout( async move { progress_clock_for_query.mark(); std::future::pending::>().await }, Some(Duration::from_millis(10)), progress_clock, None, "query completion inactivity timeout".to_string(), ) .await .expect_err("a stream stalled after its last row should time out"); assert_eq!(error, "query completion inactivity timeout"); } #[tokio::test] async fn postgres_row_query_without_timeout_still_honors_cancellation() { let progress_clock = Arc::new(StreamProgressClock::new()); let cancel_token = CancellationToken::new(); cancel_token.cancel(); let error = await_stream_with_progress_timeout( std::future::pending::>(), None, progress_clock, Some(&cancel_token), "disabled query timeout".to_string(), ) .await .expect_err("cancellation must remain active when query timeout is disabled"); assert_eq!(error, crate::query::QUERY_CANCELED); } #[test] fn postgres_cancel_context_omits_disabled_ssl_mode() { assert!(build_postgres_cancel_context("postgres://localhost/app?sslmode=disable").is_none()); } #[test] fn postgres_tls_accepts_invalid_certs_for_require_sslmode() { let pg_config = tokio_postgres::Config::from_str("postgres://localhost/db?sslmode=require").unwrap(); assert!(postgres_sslmode_accepts_invalid_certs(pg_config.get_ssl_mode())); } #[test] fn postgres_tls_accepts_invalid_certs_for_default_prefer_sslmode() { let pg_config = tokio_postgres::Config::from_str("postgres://localhost/db").unwrap(); assert!(postgres_sslmode_accepts_invalid_certs(pg_config.get_ssl_mode())); } #[test] fn postgres_tls_keeps_verification_off_only_when_ssl_is_disabled() { let pg_config = tokio_postgres::Config::from_str("postgres://localhost/db?sslmode=disable").unwrap(); assert!(!postgres_sslmode_accepts_invalid_certs(pg_config.get_ssl_mode())); } #[test] fn postgres_prefer_fallback_url_preserves_other_options() { assert_eq!( postgres_ssl_fallback_url("postgres://localhost/db?application_name=dbx&sslmode=prefer&connect_timeout=5"), Some("postgres://localhost/db?application_name=dbx&sslmode=disable&connect_timeout=5".to_string()) ); assert_eq!( postgres_ssl_fallback_url("postgres://localhost/db?application_name=dbx"), Some("postgres://localhost/db?sslmode=disable&application_name=dbx".to_string()) ); } #[test] fn postgres_prefer_fallback_never_downgrades_strict_tls_modes() { for sslmode in ["require", "verify-ca", "verify-full", "disable"] { let url = format!("postgres://localhost/db?sslmode={sslmode}"); assert_eq!(postgres_ssl_fallback_url(&url), None, "unexpected fallback for {sslmode}"); } } #[test] fn postgres_prefer_fallback_only_retries_tls_handshake_errors() { assert!(postgres_error_should_retry_without_tls( "error performing TLS handshake: received fatal alert: HandshakeFailure" )); assert!(!postgres_error_should_retry_without_tls("password authentication failed for user postgres")); assert!(!postgres_error_should_retry_without_tls("connection timed out")); assert!(!postgres_error_should_retry_without_tls("server does not support TLS")); } // --- SQL generation --- #[test] fn postgres_tables_sql_contains_expected_columns() { let sql = postgres_tables_sql(Some(500), 0); assert!(sql.contains("table_name")); assert!(sql.contains("table_type")); assert!(sql.contains("table_comment")); assert!(sql.contains("pg_catalog.pg_inherits")); assert!(sql.contains("parent_schema")); assert!(sql.contains("parent_name")); assert!(sql.contains("pc.relkind = 'p'")); assert!(sql.contains("$1")); assert!(sql.contains("BASE TABLE")); assert!(sql.contains("VIEW")); assert!(sql.contains("MATERIALIZED_VIEW")); assert!(sql.contains("FOREIGN TABLE")); } #[test] fn postgres_table_objects_sql_excludes_views_before_pagination() { let sql = postgres_table_objects_sql(Some(101), 100); assert!(sql.contains("c.relkind IN ('r','f','p')")); assert!(sql.contains("LIMIT 101 OFFSET 100")); } #[test] fn postgres_table_comment_sql_targets_single_table() { let sql = postgres_table_comment_sql(); assert!(sql.contains("obj_description(c.oid)")); assert!(sql.contains("n.nspname = $1")); assert!(sql.contains("c.relname = $2")); assert!(sql.contains("LIMIT 1")); assert!(!sql.contains("ORDER BY")); } #[test] fn postgres_table_partition_sql_tracks_parents_bounds_and_local_objects() { let relation_sql = postgres_table_partition_relation_sql(); let info_sql = postgres_table_partition_info_sql(); let info_compat_sql = postgres_table_partition_info_compat_sql(); let local_objects_sql = postgres_table_partition_local_objects_sql(); assert!(!relation_sql.contains("row_to_json")); assert!(!relation_sql.contains("relispartition")); assert!(!relation_sql.contains("relpartbound")); assert!(relation_sql.contains("pg_catalog.pg_inherits")); assert!(relation_sql.contains("parent.oid = i.inhparent")); assert!(relation_sql.contains(") = 'p'")); // Foreign tables (relkind 'f') must be reachable here too, or a // foreign partition looked up directly (not via recursion from its // parent) would never be classified as a partition at all. assert!(relation_sql.contains("c.relkind IN ('r','p','f')")); assert!(info_sql.contains("pg_catalog.pg_get_expr(c.relpartbound, c.oid, true)")); assert!(info_sql.contains("pg_catalog.pg_get_partkeydef(c.oid)")); assert!(info_sql.contains("pg_catalog.pg_inherits")); assert!(info_sql.contains("parent_schema")); assert!(info_sql.contains("parent_table")); assert!(info_sql.contains("pg_catalog.pg_foreign_table")); assert!(info_sql.contains("pg_catalog.pg_foreign_server")); assert!(info_sql.contains("fs.srvname")); assert!(info_sql.contains("ft.ftoptions")); assert!(info_sql.contains("c.relkind IN ('r','p','f')")); // 兼容版面向 9.x:没有 relispartition/relpartbound/pg_get_partkeydef, // 但保留了与主查询一致的列序和 LIMIT 形状。 assert!(!info_compat_sql.contains("relispartition")); assert!(!info_compat_sql.contains("relpartbound")); assert!(!info_compat_sql.contains("pg_get_partkeydef")); assert!(!info_compat_sql.contains("pg_inherits")); assert!(info_compat_sql.contains("NULL::text AS parent_schema")); assert!(info_compat_sql.contains("NULL::text AS parent_table")); assert!(info_compat_sql.contains("NULL::text AS partition_bound")); assert!(info_compat_sql.contains("NULL::text AS partition_key")); assert!(info_compat_sql.contains("fs.srvname")); assert!(info_compat_sql.contains("ft.ftoptions")); assert!(info_compat_sql.contains("c.relkind IN ('r','p','f')")); assert!(info_compat_sql.contains("LIMIT 1")); assert!(local_objects_sql.contains("row_to_json(con)->>'conparentid'")); assert!(local_objects_sql.contains("con.contype IN ('p','f')")); assert!(local_objects_sql.contains("i.inhrelid = idx.oid")); assert!(local_objects_sql.contains("con.contype = 'c' AND con.conislocal")); assert!(!local_objects_sql.contains("con.coninhcount = 0")); assert!(local_objects_sql.contains("pg_catalog.pg_attrdef")); } #[test] fn postgres_partition_tree_compat_sql_collapses_to_single_relation() { let sql = postgres_partition_tree_sql(); let compat_sql = postgres_partition_tree_compat_sql(); // 主查询只支持 PostgreSQL 10+(声明式分区);兼容版不能引用任何 // 10+ 的目录列/函数。 assert!(sql.contains("c.relispartition")); assert!(sql.contains("pg_catalog.pg_get_expr(c.relpartbound")); assert!(sql.contains("pg_catalog.pg_get_partkeydef")); assert!(!compat_sql.contains("relispartition")); assert!(!compat_sql.contains("relpartbound")); assert!(!compat_sql.contains("pg_get_partkeydef")); assert!(!compat_sql.contains("pg_inherits")); assert!(!compat_sql.contains("WITH RECURSIVE")); assert!(compat_sql.contains("NULL::bigint AS parent_oid")); assert!(compat_sql.contains("NULL::text AS parent_schema")); assert!(compat_sql.contains("NULL::text AS parent_relname")); assert!(compat_sql.contains("NULL::text AS partition_bound")); assert!(compat_sql.contains("NULL::text AS partition_key")); assert!(compat_sql.contains("c.relkind::text AS relkind")); assert!(compat_sql.contains("fs.srvname AS foreign_server")); assert!(compat_sql.contains("ft.ftoptions AS foreign_options")); assert!(compat_sql.contains("n.nspname = $1 AND c.relname = $2")); assert!(compat_sql.contains("c.relkind IN ('r','p','f')")); } #[test] fn postgres_check_constraints_sql_selects_definitions_for_relation() { let sql = postgres_check_constraints_sql(); assert!(sql.contains("con.contype = 'c'")); assert!(sql.contains("pg_catalog.pg_get_constraintdef(con.oid, true)")); assert!(sql.contains("n.nspname = $1 AND c.relname = $2")); assert!(sql.contains("ORDER BY con.conname")); } #[test] fn postgres_constraints_sql_covers_all_contypes_and_referential_fields() { let sql = postgres_constraints_sql(); assert!(sql.contains("con.contype::text")); assert!(sql.contains("pg_catalog.pg_get_constraintdef(con.oid, true)")); // FK-only referential columns are conditionally selected. assert!(sql.contains("CASE WHEN con.contype = 'f' THEN con.confmatchtype::text END")); assert!(sql.contains("CASE WHEN con.contype = 'f' THEN con.confupdtype::text END")); assert!(sql.contains("CASE WHEN con.contype = 'f' THEN con.confdeltype::text END")); // conkey/confkey attribute numbers are resolved to names via laterals. assert!(sql.contains("unnest(con.conkey) WITH ORDINALITY")); assert!(sql.contains("unnest(con.confkey) WITH ORDINALITY")); // attname (pg_attribute type `name`) is explicitly cast to text so the // array decodes as text[] (matching the COALESCE fallback and Vec). assert!(sql.contains("array_agg(a.attname::text ORDER BY ord.ord)")); assert_eq!(sql.matches("attname::text").count(), 2); assert!(sql.contains("a.attisdropped")); assert!(sql.contains("n.nspname = $1 AND c.relname = $2")); assert!(sql.contains("c.relkind IN ('r','p','f')")); assert!(sql.contains("ORDER BY con.conname")); } #[test] fn opengauss_constraint_sql_avoids_array_and_ordinality_dependencies() { let sql = opengauss_constraints_sql(true); assert!(sql.contains("con.conkey::text")); assert!(sql.contains("con.confkey::text")); assert!(sql.contains("con.conrelid::oid")); assert!(sql.contains("con.confrelid::oid")); assert!(!sql.contains("unnest(")); assert!(!sql.contains("WITH ORDINALITY")); assert!(!sql.contains("LATERAL")); } #[test] fn opengauss_constraint_sql_falls_back_without_convalidated() { let sql = opengauss_constraints_sql(false); assert!(sql.contains("TRUE")); assert!(!sql.contains("con.convalidated")); } #[test] fn opengauss_constraint_attribute_vectors_are_strict() { assert_eq!(parse_opengauss_attribute_numbers("1 2 4").unwrap(), vec![1, 2, 4]); assert_eq!(parse_opengauss_attribute_numbers("{3,2}").unwrap(), vec![3, 2]); assert_eq!(parse_opengauss_attribute_numbers("[4, 5]").unwrap(), vec![4, 5]); assert!(parse_opengauss_attribute_numbers("1 bad 3").is_err()); } #[test] fn postgres_constraint_type_label_expands_contype_letters() { assert_eq!(postgres_constraint_type_label("p"), "PRIMARY KEY"); assert_eq!(postgres_constraint_type_label("f"), "FOREIGN KEY"); assert_eq!(postgres_constraint_type_label("u"), "UNIQUE"); assert_eq!(postgres_constraint_type_label("c"), "CHECK"); assert_eq!(postgres_constraint_type_label("x"), "EXCLUDE"); assert_eq!(postgres_constraint_type_label("n"), "NOT NULL"); assert_eq!(postgres_constraint_type_label("unknown"), "unknown"); } #[test] fn postgres_fk_action_label_and_match_type_normalize_letters() { assert_eq!(postgres_fk_action_label(Some("c".to_string())).as_deref(), Some("CASCADE")); assert_eq!(postgres_fk_action_label(Some("n".to_string())).as_deref(), Some("SET NULL")); assert_eq!(postgres_fk_action_label(Some("d".to_string())).as_deref(), Some("SET DEFAULT")); assert_eq!(postgres_fk_action_label(Some("a".to_string())).as_deref(), Some("NO ACTION")); assert_eq!(postgres_fk_action_label(Some("r".to_string())).as_deref(), Some("RESTRICT")); assert_eq!(postgres_fk_action_label(Some("zz".to_string())), None); assert_eq!(postgres_fk_action_label(None), None); assert_eq!(postgres_constraint_match_type(Some("f".to_string())).as_deref(), Some("FULL")); assert_eq!(postgres_constraint_match_type(Some("s".to_string())).as_deref(), Some("SIMPLE")); assert_eq!(postgres_constraint_match_type(Some("u".to_string())).as_deref(), Some("SIMPLE")); assert_eq!(postgres_constraint_match_type(Some("p".to_string())).as_deref(), Some("PARTIAL")); assert_eq!(postgres_constraint_match_type(Some("zz".to_string())), None); } #[test] fn postgres_column_metadata_reads_identity_extra() { assert!(POSTGRES_COLUMNS_SQL.contains("a.attidentity")); assert!(POSTGRES_COLUMNS_SQL.contains("pg_sequence")); assert!(POSTGRES_COLUMNS_SQL.contains("generated by default as identity")); assert!(POSTGRES_COLUMNS_SQL.contains("generated always as identity")); assert!(POSTGRES_COLUMNS_SQL.contains("COALESCE(c.is_nullable = 'YES', NOT a.attnotnull)")); assert!(POSTGRES_COLUMNS_SQL.contains("LEFT JOIN information_schema.columns")); assert!(POSTGRES_COLUMNS_SQL.contains("pg_enum")); assert!(POSTGRES_COLUMNS_SQL.contains("AS enum_values")); } #[test] fn postgres_single_relation_metadata_can_resolve_through_search_path() { for sql in [POSTGRES_COLUMNS_SQL, POSTGRES_COLUMNS_COMPAT_SQL] { assert!(sql.contains("CASE WHEN $1 = '' THEN quote_ident($2)")); assert!(sql.contains("c.table_schema = relation_namespace.nspname")); } for sql in [POSTGRES_INDEXES_SQL, POSTGRES_INDEXES_COMPAT_SQL] { assert!(sql.contains("t.oid = (CASE WHEN $1 = '' THEN quote_ident($2)")); } } #[test] fn postgres_column_metadata_marks_only_owned_integer_sequence_defaults_as_serial() { let modern_sql = [POSTGRES_COLUMNS_SQL, postgres_columns_for_relations_sql()]; let compat_sql = [POSTGRES_COLUMNS_COMPAT_SQL, postgres_columns_for_relations_compat_sql()]; for sql in modern_sql.into_iter().chain(compat_sql) { assert!(sql.contains("a.atttypid IN (20, 21, 23)")); assert!(sql.contains("WHEN 21 THEN 'smallserial'")); assert!(sql.contains("WHEN 23 THEN 'serial'")); assert!(sql.contains("WHEN 20 THEN 'bigserial'")); } for sql in modern_sql { assert!(sql.contains("LEFT JOIN LATERAL")); assert!(sql.contains("a.attidentity <> '' AND sequence_dep.deptype = 'i'")); assert!(sql.contains("a.attidentity = ''")); assert!(sql.contains("sequence_dep.deptype = 'a' AND EXISTS")); assert!(sql.contains("sequence_class.relkind = 'S'")); assert!(sql.contains("ORDER BY sequence_dep.objid")); assert!(sql.contains("LIMIT 1")); assert!(sql.contains("pseq.seqrelid IS NOT NULL")); assert!(sql.contains("serial_default_dep.objid = ad.oid")); assert!(sql.contains("serial_default_dep.refobjid = sequence_dep.objid")); assert!(sql.contains( "pg_get_expr(ad.adbin, ad.adrelid) = format('nextval(%L::regclass)', dep.objid::regclass::text)" )); } for sql in compat_sql { assert!(!sql.contains("LEFT JOIN LATERAL")); assert!(sql.contains("serial_seq.oid = (")); assert!(sql.contains("sequence_dep.deptype = 'a' AND EXISTS")); assert!(sql.contains("sequence_class.relkind = 'S'")); assert!(sql.contains("ORDER BY sequence_dep.objid")); assert!(sql.contains("LIMIT 1")); assert!(sql.contains("serial_seq.relkind = 'S'")); assert!(sql.contains("serial_default_dep.objid = ad.oid")); assert!(sql.contains("serial_default_dep.refobjid = sequence_dep.objid")); assert!(sql.contains( "pg_get_expr(ad.adbin, ad.adrelid) = format('nextval(%L::regclass)', serial_seq.oid::regclass::text)" )); } assert!(!POSTGRES_COLUMNS_INFORMATION_SCHEMA_SQL.contains("serial_dep")); assert!(POSTGRES_COLUMNS_INFORMATION_SCHEMA_SQL.contains("NULL::text AS column_extra")); } #[test] fn postgres_partition_local_checks_keep_merged_local_definitions() { for sql in [postgres_table_partition_local_objects_sql(), postgres_table_partition_local_objects_for_relations_sql()] { assert!(sql.contains("con.contype = 'c' AND con.conislocal")); assert!(!sql.contains("con.coninhcount = 0")); } } #[test] fn postgres_partition_batch_metadata_uses_bounded_compat_tiers() { let column_tiers = postgres_columns_for_relations_query_tiers(); assert_eq!(column_tiers.len(), 2); assert!(column_tiers.iter().all(|sql| sql.contains("c.oid = ANY($1::bigint[])"))); assert!(column_tiers[0].contains("a.attgenerated")); assert!(!column_tiers[1].contains("a.attgenerated")); assert!(!column_tiers[1].contains("pg_sequence")); assert!(column_tiers[1].contains("sequence_dep.deptype = 'a'")); assert!(!column_tiers[1].contains("LEFT JOIN LATERAL")); let index_tiers = postgres_indexes_for_relations_query_tiers(); assert_eq!(index_tiers.len(), 2); assert!(index_tiers.iter().all(|sql| sql.contains("t.oid = ANY($1::bigint[])"))); assert!(index_tiers[0].contains("ix.indnkeyatts")); assert!(!index_tiers[1].contains("ix.indnkeyatts")); assert!(index_tiers[1].contains("string_to_array(ix.indkey::text, ' ')")); let local_object_tiers = postgres_table_partition_local_objects_for_relations_query_tiers(); assert_eq!(local_object_tiers.len(), 2); assert!(local_object_tiers[0].contains("row_to_json(con)->>'conparentid'")); assert!(!local_object_tiers[1].contains("row_to_json")); assert!(!local_object_tiers[1].contains("conparentid")); assert!(local_object_tiers.iter().all(|sql| sql.contains("con.conrelid = ANY($1::bigint[])"))); } #[test] fn postgres_partition_batch_foreign_keys_keep_ordered_paired_compat_tier() { let tiers = postgres_foreign_keys_for_relations_query_tiers(); assert_eq!(tiers.len(), 2); let modern_sql = tiers[0]; assert_eq!(modern_sql, postgres_foreign_keys_for_relations_sql()); assert!(modern_sql.contains("unnest($1::text[], $2::text[])")); assert!(!modern_sql.contains("generate_subscripts")); assert!(modern_sql.contains("pg_catalog.pg_constraint")); assert!(!modern_sql.contains("information_schema")); assert!(modern_sql.contains("con.confupdtype::text")); assert!(modern_sql.contains("confdeltype")); assert!(modern_sql.contains("con.contype = 'f'")); assert!(modern_sql.contains("ORDER BY rel.rel_schema, rel.rel_table, con.conname, fk.ord")); let compat_sql = tiers[1]; assert_eq!(compat_sql, postgres_foreign_keys_for_relations_compat_sql()); assert!(!compat_sql.contains("unnest($1::text[]")); assert!(compat_sql.contains("generate_subscripts($1::text[], 1)")); assert!(compat_sql.contains("($1::text[])[rel.i] AS rel_schema")); assert!(compat_sql.contains("($2::text[])[rel.i] AS rel_table")); assert!(compat_sql.contains("pg_catalog.pg_constraint")); assert!(!compat_sql.contains("information_schema")); assert!(compat_sql.contains("confdeltype")); assert!(compat_sql.contains("con.contype = 'f'")); assert!(compat_sql.contains("ORDER BY rel.rel_schema, rel.rel_table, con.conname, fk.ord")); } #[test] fn postgres_column_metadata_has_opengauss_compatible_fallback() { assert!(!POSTGRES_COLUMNS_COMPAT_SQL.contains("a.attidentity")); assert!(!POSTGRES_COLUMNS_COMPAT_SQL.contains("pg_sequence")); assert!(POSTGRES_COLUMNS_COMPAT_SQL.contains("AS column_extra")); assert!(POSTGRES_COLUMNS_COMPAT_SQL.contains("col_description")); assert!(POSTGRES_COLUMNS_COMPAT_SQL.contains("COALESCE(c.is_nullable = 'YES', NOT a.attnotnull)")); assert!(POSTGRES_COLUMNS_COMPAT_SQL.contains("LEFT JOIN information_schema.columns")); assert!(POSTGRES_COLUMNS_COMPAT_SQL.contains("NULL::text AS enum_values")); assert!(!POSTGRES_COLUMNS_COMPAT_SQL.contains("pg_enum")); } #[test] fn postgres_column_metadata_has_information_schema_fallback() { assert!(POSTGRES_COLUMNS_INFORMATION_SCHEMA_SQL.contains("information_schema.columns")); assert!(POSTGRES_COLUMNS_INFORMATION_SCHEMA_SQL.contains("information_schema.table_constraints")); assert!(POSTGRES_COLUMNS_INFORMATION_SCHEMA_SQL.contains("information_schema.key_column_usage")); assert!(POSTGRES_COLUMNS_INFORMATION_SCHEMA_SQL.contains("NULL::text AS enum_values")); assert!(!POSTGRES_COLUMNS_INFORMATION_SCHEMA_SQL.contains("pg_attribute")); assert!(!POSTGRES_COLUMNS_INFORMATION_SCHEMA_SQL.contains("regclass")); } #[test] fn opengauss_sequence_metadata_uses_compatible_information_schema_view() { let sql = opengauss_sequences_sql(); assert!(sql.contains("information_schema.sequences")); assert!(sql.contains("s.sequence_schema = $1")); assert!(sql.contains("c.relkind IN ('S','L','z','Z')")); assert!(sql.contains("sequence_name")); assert!(sql.contains("start_value")); assert!(sql.contains("minimum_value")); assert!(sql.contains("maximum_value")); assert!(sql.contains("increment")); assert!(sql.contains("cycle_option")); assert!(!sql.contains("pg_sequence s")); } #[test] fn opengauss_sequence_last_values_extract_record_field_as_text() { let sql = opengauss_sequence_last_values_sql(); assert!(sql.contains("(pg_sequence_last_value(c.oid)).last_value::text")); assert!(sql.contains("c.relkind IN ('S','L','z','Z')")); assert!(sql.contains("n.nspname = $1")); } #[test] fn postgres_sequence_last_values_are_read_as_text() { assert!(postgres_sequence_last_values_sql().contains("pg_sequence_last_value(c.oid)::text")); } #[test] fn extension_member_query_filters_only_owned_relations_and_routines() { let sql = list_extension_member_objects_sql(); assert!(sql.contains("d.classid = 'pg_catalog.pg_class'::regclass")); assert!(sql.contains("d.classid = 'pg_catalog.pg_proc'::regclass")); assert!(sql.contains("d.refclassid = 'pg_catalog.pg_extension'::regclass")); assert!(sql.contains("d.deptype = 'e'")); assert!(sql.contains("pg_get_function_identity_arguments(p.oid)")); assert!(!sql.contains("d.deptype = 'x'")); } #[tokio::test] async fn postgres_column_metadata_query_returns_enum_values_against_real_postgres() { let Some(container) = start_docker_postgres().await else { return; }; let pool = connect(&container.url(), Duration::from_secs(5)).await.expect("connect postgres"); let schema = format!("dbx_enum_meta_{}", std::process::id()); let schema_ident = format!("\"{}\"", schema.replace('\"', "\"\"")); let table = format!("{schema_ident}.orders"); let type_ident = format!("{schema_ident}.\"status\""); execute_query(&pool, &format!("CREATE SCHEMA {schema_ident}")).await.expect("create schema"); execute_query(&pool, &format!("CREATE TYPE {type_ident} AS ENUM ('pending', 'active', 'archived')")) .await .expect("create enum type"); execute_query(&pool, &format!("CREATE TABLE {table} (id integer PRIMARY KEY, state {type_ident} NOT NULL)")) .await .expect("create table"); let client = checkout_postgres_client(&pool, None, crate::db::connection_timeout()).await.expect("checkout client"); let columns = get_columns_with_sql(&client, POSTGRES_COLUMNS_SQL, &schema, "orders").await.expect("primary columns"); assert_eq!( state_enum_values(&columns), Some(vec!["pending".to_string(), "active".to_string(), "archived".to_string()]) ); } #[tokio::test] #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a writable PostgreSQL database"] async fn postgres_unqualified_metadata_follows_search_path() { let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); let pool = connect(&url, Duration::from_secs(5)).await.expect("connect postgres"); let client = checkout_postgres_client(&pool, None, crate::db::connection_timeout()).await.expect("checkout client"); let suffix = std::process::id(); let schema = format!("dbx_search_path_meta_{suffix}"); let table = format!("comments_{suffix}"); let schema_ident = format!("\"{}\"", schema.replace('\"', "\"\"")); let table_ident = format!("\"{}\"", table.replace('\"', "\"\"")); client .batch_execute(&format!( "CREATE SCHEMA {schema_ident}; \ CREATE TABLE {schema_ident}.{table_ident} (id integer PRIMARY KEY, display_name text); \ COMMENT ON COLUMN {schema_ident}.{table_ident}.display_name IS '客户姓名'; \ SET search_path TO {schema_ident}, public;" )) .await .expect("create search_path metadata fixtures"); let columns = get_columns_with_sql(&client, POSTGRES_COLUMNS_SQL, "", &table) .await .expect("load columns through search_path"); assert_eq!( columns.iter().find(|column| column.name == "display_name").and_then(|column| column.comment.as_deref()), Some("客户姓名") ); let indexes = list_indexes_with_sql(&client, POSTGRES_INDEXES_SQL, "", &table) .await .expect("load indexes through search_path"); assert!(indexes.iter().any(|index| index.is_primary && index.columns == ["id"])); client .batch_execute(&format!("RESET search_path; DROP SCHEMA {schema_ident} CASCADE")) .await .expect("clean search_path metadata fixtures"); } #[tokio::test] #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a writable PostgreSQL database"] async fn postgres_column_metadata_decode_type_mismatch_uses_fallbacks() { let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); let pool = connect(&url, std::time::Duration::from_secs(5)).await.expect("connect postgres"); let client = checkout_postgres_client(&pool, None, std::time::Duration::from_secs(5)).await.expect("checkout postgres"); let row = client .query_one( "SELECT \ 1::int4 AS column_name, \ 'text'::text AS full_type, \ 'YES'::text AS is_nullable, \ NULL::text AS column_default, \ 1::int4 AS is_pk, \ NULL::text AS column_comment, \ NULL::text AS column_extra, \ NULL::int4 AS numeric_precision, \ NULL::int4 AS numeric_scale, \ NULL::int4 AS character_maximum_length", &[], ) .await .expect("query mismatched metadata row"); let info = column_info_from_row(&row); // int4 column_name should be converted to string "1" instead of panicking assert_eq!(info.name, "1"); // text 'YES' is not a standard bool, pg_row_try_bool falls back to string match assert!(info.is_nullable); // int4 1 should be interpreted as true for is_primary_key assert!(info.is_primary_key); } #[tokio::test] #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a writable PostgreSQL database"] async fn postgres_owned_serial_table_ddl_round_trips_without_external_sequences() { let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); let pool = connect(&url, Duration::from_secs(5)).await.expect("connect postgres"); let suffix = uuid::Uuid::new_v4().simple(); let schema = format!("dbx 6405 \"{suffix}"); let schema_ident = pg_quote_ident(&schema); let table_name = "order\"items"; let table_ident = pg_quote_ident(table_name); let table = format!("{schema_ident}.{table_ident}"); let custom_sequence_name = "custom\"source"; let custom_sequence = format!("{schema_ident}.{}", pg_quote_ident(custom_sequence_name)); let custom_table_name = "custom defaults"; let custom_table = format!("{schema_ident}.{}", pg_quote_ident(custom_table_name)); let owned_custom_sequence_name = "owned custom source"; let owned_custom_sequence = format!("{schema_ident}.{}", pg_quote_ident(owned_custom_sequence_name)); let owned_custom_table_name = "owned custom defaults"; let owned_custom_table = format!("{schema_ident}.{}", pg_quote_ident(owned_custom_table_name)); let dropped_default_table_name = "dropped default"; let dropped_default_table = format!("{schema_ident}.{}", pg_quote_ident(dropped_default_table_name)); let client = pool.get().await.expect("get postgres client"); client .batch_execute(&format!( "CREATE SCHEMA {schema_ident}; \ CREATE TABLE {table} (\ \"small\"\"id\" smallserial, \ \"regular id\" serial, \ \"large\"\"id\" bigserial, \ identity_id bigint GENERATED BY DEFAULT AS IDENTITY, \ base bigint DEFAULT 2, \ generated bigint GENERATED ALWAYS AS (base * 2) STORED\ ); \ CREATE SEQUENCE {custom_sequence}; \ CREATE TABLE {custom_table} (id bigint DEFAULT nextval({}::regclass)); \ CREATE TABLE {owned_custom_table} (id bigint NOT NULL); \ CREATE SEQUENCE {owned_custom_sequence} OWNED BY {owned_custom_table}.id; \ ALTER TABLE {owned_custom_table} ALTER COLUMN id \ SET DEFAULT nextval({}::regclass) + 7; \ CREATE TABLE {dropped_default_table} (id serial); \ ALTER TABLE {dropped_default_table} ALTER COLUMN id DROP DEFAULT", pg_quote_literal(&custom_sequence), pg_quote_literal(&owned_custom_sequence), )) .await .expect("create serial probe objects"); let ddl = crate::schema::pg_ddl(&pool, &schema, table_name).await.expect("read serial table ddl"); let custom_ddl = crate::schema::pg_ddl(&pool, &schema, custom_table_name).await.expect("read custom-default table ddl"); let owned_custom_ddl = crate::schema::pg_ddl(&pool, &schema, owned_custom_table_name) .await .expect("read owned custom-default table ddl"); let dropped_default_ddl = crate::schema::pg_ddl(&pool, &schema, dropped_default_table_name) .await .expect("read dropped-default table ddl"); assert!(ddl.contains("\"small\"\"id\" smallserial NOT NULL"), "ddl: {ddl}"); assert!(ddl.contains("\"regular id\" serial NOT NULL"), "ddl: {ddl}"); assert!(ddl.contains("\"large\"\"id\" bigserial NOT NULL"), "ddl: {ddl}"); assert!(ddl.contains("\"identity_id\" bigint generated by default as identity"), "ddl: {ddl}"); assert!(ddl.contains("\"generated\" bigint generated always as ((base * 2)) stored"), "ddl: {ddl}"); assert!(!ddl.contains("nextval("), "ddl: {ddl}"); assert!(custom_ddl.contains("bigint DEFAULT nextval("), "ddl: {custom_ddl}"); assert!(!custom_ddl.contains("bigserial"), "ddl: {custom_ddl}"); assert!(owned_custom_ddl.contains("bigint NOT NULL DEFAULT (nextval("), "ddl: {owned_custom_ddl}"); assert!(owned_custom_ddl.contains(" + 7)"), "ddl: {owned_custom_ddl}"); assert!(!owned_custom_ddl.contains("bigserial"), "ddl: {owned_custom_ddl}"); assert!(dropped_default_ddl.contains("\"id\" integer NOT NULL"), "ddl: {dropped_default_ddl}"); assert!(!dropped_default_ddl.contains("serial"), "ddl: {dropped_default_ddl}"); assert!(!dropped_default_ddl.contains("DEFAULT"), "ddl: {dropped_default_ddl}"); client.batch_execute(&format!("DROP TABLE {table} CASCADE; {ddl}")).await.expect("replay copied table ddl"); let inserted = client .query_one( &format!( "INSERT INTO {table} DEFAULT VALUES RETURNING \"small\"\"id\", \"regular id\", \"large\"\"id\", identity_id, generated" ), &[], ) .await .expect("insert with replayed sequence defaults"); assert_eq!(inserted.get::<_, i16>(0), 1); assert_eq!(inserted.get::<_, i32>(1), 1); assert_eq!(inserted.get::<_, i64>(2), 1); assert_eq!(inserted.get::<_, i64>(3), 1); assert_eq!(inserted.get::<_, i64>(4), 4); client.batch_execute(&format!("DROP SCHEMA {schema_ident} CASCADE")).await.expect("drop serial probe schema"); } #[tokio::test] #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a writable PostgreSQL database"] async fn postgres_column_metadata_dependency_joins_return_each_column_once() { let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); let pool = connect(&url, Duration::from_secs(5)).await.expect("connect postgres"); let suffix = uuid::Uuid::new_v4().simple(); let schema = format!("dbx 6547 \"{suffix}"); let schema_ident = pg_quote_ident(&schema); let table_name = "issue6547_repro"; let table_ident = pg_quote_ident(table_name); let table = format!("{schema_ident}.{table_ident}"); let client = pool.get().await.expect("get postgres client"); client .batch_execute(&format!( "CREATE SCHEMA {schema_ident}; \ CREATE TABLE {table} (\ id serial PRIMARY KEY, \ value text, \ small_id smallserial, \ big_id bigserial, \ ident_id bigint GENERATED BY DEFAULT AS IDENTITY, \ plain_int integer, \ vc varchar(32), \ num numeric(10, 2), \ ts timestamp\ ); \ CREATE SEQUENCE {schema_ident}.serial_decoy OWNED BY {table}.id; \ CREATE SEQUENCE {schema_ident}.identity_decoy OWNED BY {table}.ident_id; \ CREATE INDEX idx_issue6547_repro_id ON {table}(id); \ CREATE INDEX idx_issue6547_repro_big ON {table}(big_id); \ CREATE UNIQUE INDEX uniq_issue6547_repro_vc ON {table}(vc);" )) .await .expect("create repro table"); let columns = get_columns(&pool, &schema, table_name).await.expect("get columns"); // One pg_attribute row must produce exactly one ColumnInfo even when // the column has multiple sequence dependencies and several indexes. // Only the sequence used by the default, or the internal identity // sequence, is semantic metadata for the column. let mut seen: std::collections::HashSet<&str> = std::collections::HashSet::new(); for column in &columns { assert!(seen.insert(column.name.as_str()), "duplicate column {:?} in {columns:?}", column.name); } assert_eq!(columns.len(), 9, "expected 9 columns, got {columns:?}"); for (name, expected_extra) in [("id", "serial"), ("small_id", "smallserial"), ("big_id", "bigserial")] { let column = columns.iter().find(|c| c.name == name).unwrap_or_else(|| panic!("missing {name}")); assert_eq!(column.extra.as_deref(), Some(expected_extra), "{name} extra mismatch"); } let ident = columns.iter().find(|c| c.name == "ident_id").expect("ident_id column"); assert!( ident.extra.as_deref().is_some_and(|extra| extra.starts_with("generated by default as identity")), "ident_id extra: {:?}", ident.extra ); // Ordinary columns keep their plain metadata. let vc = columns.iter().find(|c| c.name == "vc").expect("vc column"); assert_eq!(vc.data_type, "character varying(32)"); let num = columns.iter().find(|c| c.name == "num").expect("num column"); assert_eq!(num.data_type, "numeric(10,2)"); assert_eq!(num.numeric_precision, Some(10)); assert_eq!(num.numeric_scale, Some(2)); let ts = columns.iter().find(|c| c.name == "ts").expect("ts column"); assert_eq!(ts.data_type, "timestamp without time zone"); // The rendered column section of the DDL must contain each column // definition exactly once. Before the dependency-join fix these // queries could return "id" (serial + PRIMARY KEY + secondary index // = several pg_depend rows) any number of times. let ddl = crate::schema::pg_ddl(&pool, &schema, table_name).await.expect("read repro table ddl"); for needle in [r#" "id" serial NOT NULL"#, r#" "value" text"#, r#" "vc" character varying(32)"#] { assert_eq!(ddl.matches(needle).count(), 1, "needle {needle:?} not unique in ddl:\n{ddl}"); } drop(client); execute_query(&pool, &format!("DROP SCHEMA {schema_ident} CASCADE")).await.expect("drop repro schema"); } #[tokio::test] #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a writable PostgreSQL database"] async fn postgres_column_metadata_compat_query_returns_each_column_once() { let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); let pool = connect(&url, Duration::from_secs(5)).await.expect("connect postgres"); let suffix = uuid::Uuid::new_v4().simple(); let schema = format!("dbx 6547 compat \"{suffix}"); let schema_ident = pg_quote_ident(&schema); let table_name = "issue6547_repro"; let table_ident = pg_quote_ident(table_name); let table = format!("{schema_ident}.{table_ident}"); let client = pool.get().await.expect("get postgres client"); client .batch_execute(&format!( "CREATE SCHEMA {schema_ident}; \ CREATE TABLE {table} (id serial PRIMARY KEY, value text); \ CREATE SEQUENCE {schema_ident}.serial_decoy OWNED BY {table}.id; \ CREATE INDEX idx_issue6547_repro_id ON {table}(id)" )) .await .expect("create compat repro table"); let columns = get_columns_with_sql(&client, POSTGRES_COLUMNS_COMPAT_SQL, &schema, table_name) .await .expect("compat columns"); let names: Vec<&str> = columns.iter().map(|c| c.name.as_str()).collect(); let mut seen: std::collections::HashSet<&str> = std::collections::HashSet::new(); for name in &names { assert!(seen.insert(*name), "duplicate column {name:?} in {names:?}"); } assert_eq!(names, vec!["id", "value"], "unexpected compat columns {names:?}"); let id = columns.iter().find(|c| c.name == "id").expect("id column"); assert_eq!(id.extra.as_deref(), Some("serial"), "compat id extra: {:?}", id.extra); drop(client); execute_query(&pool, &format!("DROP SCHEMA {schema_ident} CASCADE")).await.expect("drop compat repro schema"); } #[tokio::test] #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a writable PostgreSQL database"] async fn postgres_batched_column_metadata_returns_each_column_once() { let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); let pool = connect(&url, Duration::from_secs(5)).await.expect("connect postgres"); let suffix = uuid::Uuid::new_v4().simple(); let schema = format!("dbx 6547 batch \"{suffix}"); let schema_ident = pg_quote_ident(&schema); let table_name = "issue6547_repro"; let table_ident = pg_quote_ident(table_name); let table = format!("{schema_ident}.{table_ident}"); let client = pool.get().await.expect("get postgres client"); client .batch_execute(&format!( "CREATE SCHEMA {schema_ident}; \ CREATE TABLE {table} (id serial PRIMARY KEY, value text); \ CREATE SEQUENCE {schema_ident}.serial_decoy OWNED BY {table}.id; \ CREATE INDEX idx_issue6547_repro_id ON {table}(id)" )) .await .expect("create batch repro table"); let oid = client .query_one( "SELECT c.oid::bigint FROM pg_catalog.pg_class c \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ WHERE n.nspname = $1 AND c.relname = $2", &[&schema, &table_name], ) .await .expect("lookup repro table oid") .get::<_, i64>(0); let relations = vec![(oid, schema.clone(), table_name.to_string())]; let columns_by_oid = get_columns_for_relations(&pool, &relations).await.expect("batched columns"); let columns = columns_by_oid.get(&oid).unwrap_or_else(|| panic!("no batched columns for {oid}")); let names: Vec<&str> = columns.iter().map(|c| c.name.as_str()).collect(); let mut seen: std::collections::HashSet<&str> = std::collections::HashSet::new(); for name in &names { assert!(seen.insert(*name), "duplicate column {name:?} in {names:?}"); } assert_eq!(names, vec!["id", "value"], "unexpected batched columns {names:?}"); let id = columns.iter().find(|c| c.name == "id").expect("id column"); assert_eq!(id.extra.as_deref(), Some("serial"), "batched id extra: {:?}", id.extra); // The batched sibling must stay duplicate-free too: the partition DDL // path (pg_ddl_with_partitions) consumes exactly these results. let columns_by_oid_compat = get_columns_for_relations_with_sql(&client, postgres_columns_for_relations_compat_sql(), &[oid]) .await .expect("batched compat columns"); let names_compat: Vec<&str> = columns_by_oid_compat .get(&oid) .expect("no batched compat columns") .iter() .map(|c| c.name.as_str()) .collect(); let mut seen_compat: std::collections::HashSet<&str> = std::collections::HashSet::new(); for name in &names_compat { assert!(seen_compat.insert(*name), "duplicate column {name:?} in {names_compat:?}"); } assert_eq!(names_compat, vec!["id", "value"], "unexpected batched compat columns {names_compat:?}"); drop(client); execute_query(&pool, &format!("DROP SCHEMA {schema_ident} CASCADE")).await.expect("drop batch repro schema"); } #[tokio::test] #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a writable PostgreSQL 9.x database"] async fn postgres_legacy_table_ddl_uses_compatible_partition_probe() { let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); let pool = connect(&url, Duration::from_secs(5)).await.expect("connect postgres"); let version = execute_query(&pool, "SHOW server_version_num").await.expect("query PostgreSQL version"); let version_num = version.rows[0][0] .as_str() .expect("server_version_num should be text") .parse::() .expect("server_version_num should be numeric"); assert!((90000..100000).contains(&version_num), "expected PostgreSQL 9.x, got {version_num}"); let schema = format!("dbx_legacy_ddl_{}", uuid::Uuid::new_v4().simple()); let schema_ident = pg_quote_ident(&schema); let parent = format!("{schema_ident}.parent"); let child = format!("{schema_ident}.child"); let function = format!("{schema_ident}.fill_child_name"); let client = pool.get().await.expect("get postgres client"); client .batch_execute(&format!( "CREATE SCHEMA {schema_ident}; \ CREATE TABLE {parent} (id integer PRIMARY KEY); \ CREATE TABLE {child} (parent_id integer REFERENCES {parent}(id), name text) INHERITS ({parent}); \ ALTER TABLE {child} ADD PRIMARY KEY (id); \ CREATE INDEX child_name_idx ON {child}(name); \ CREATE FUNCTION {function}() RETURNS trigger LANGUAGE plpgsql AS $$BEGIN NEW.name := COALESCE(NEW.name, chr(120)); RETURN NEW; END$$; \ CREATE TRIGGER child_bi BEFORE INSERT ON {child} FOR EACH ROW EXECUTE PROCEDURE {function}()" )) .await .expect("create PostgreSQL 9.x DDL fixtures"); drop(client); let partition_info = get_table_partition_info(&pool, &schema, "child").await.expect("classify PostgreSQL 9.x inherited table"); let ddl = crate::schema::pg_ddl(&pool, &schema, "child").await; // 查看表 DDL(DISPLAY 路径)在 9.x 上不能因 relispartition 报错, // INHERITS 子表也不能被当成声明式分区渲染。 // The view-DDL path must not fail on the missing relispartition // column pre-10, and INHERITS children must render as plain tables. let ddl_with_partitions = crate::schema::pg_ddl_with_partitions(&pool, &schema, "child").await; let parent_ddl = crate::schema::pg_ddl_with_partitions(&pool, &schema, "parent").await; execute_query(&pool, &format!("DROP SCHEMA {schema_ident} CASCADE")) .await .expect("drop PostgreSQL 9.x DDL fixtures"); let ddl = ddl.expect("generate PostgreSQL 9.x table DDL"); let ddl_with_partitions = ddl_with_partitions.expect("generate PostgreSQL 9.x partition-tree DDL"); let parent_ddl = parent_ddl.expect("generate PostgreSQL 9.x partition-tree DDL for parent"); assert!(ddl.contains("CREATE TABLE"), "ddl: {ddl}"); assert!(ddl.contains("PRIMARY KEY"), "ddl: {ddl}"); assert!(ddl.contains("FOREIGN KEY"), "ddl: {ddl}"); assert!(ddl.contains("CREATE INDEX \"child_name_idx\""), "ddl: {ddl}"); assert!(ddl.contains("CREATE TRIGGER child_bi"), "ddl: {ddl}"); assert_eq!(partition_info, PostgresTablePartitionInfo::default()); assert!(!ddl.contains("PARTITION OF"), "ddl: {ddl}"); assert!(ddl_with_partitions.contains("CREATE TABLE"), "ddl_with_partitions: {ddl_with_partitions}"); assert!(!ddl_with_partitions.contains("PARTITION OF"), "ddl_with_partitions: {ddl_with_partitions}"); // 9.x 没有声明式分区:父表的分区树只含父表自身(子表按普通表)。 assert_eq!(parent_ddl.matches("CREATE TABLE").count(), 1, "parent_ddl: {parent_ddl}"); assert!(!parent_ddl.contains("PARTITION OF"), "parent_ddl: {parent_ddl}"); } #[tokio::test] #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a writable PostgreSQL database"] async fn postgres_legacy_batched_ddl_metadata_preserves_relation_pairs_across_compat_tiers() { let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); let pool = connect(&url, Duration::from_secs(5)).await.expect("connect postgres"); let schema = format!("dbx_batch_fkeys_{}", uuid::Uuid::new_v4().simple()); let schema_ident = pg_quote_ident(&schema); let client = pool.get().await.expect("get postgres client"); client .batch_execute(&format!( "CREATE SCHEMA {schema_ident}; \ CREATE TABLE {schema_ident}.parent_a (id integer PRIMARY KEY); \ CREATE TABLE {schema_ident}.parent_b (id integer PRIMARY KEY); \ CREATE TABLE {schema_ident}.child_a ( \ id integer PRIMARY KEY, parent_id integer, \ CONSTRAINT child_a_parent_fk FOREIGN KEY (parent_id) REFERENCES {schema_ident}.parent_a(id) \ ); \ CREATE TABLE {schema_ident}.child_b ( \ id integer PRIMARY KEY, parent_id integer, \ CONSTRAINT child_b_parent_fk FOREIGN KEY (parent_id) REFERENCES {schema_ident}.parent_b(id) \ ); \ CREATE TABLE {schema_ident}.without_fk (id integer PRIMARY KEY)" )) .await .expect("create batched foreign-key fixtures"); let version_num = client .query_one("SHOW server_version_num", &[]) .await .expect("query PostgreSQL version") .get::<_, String>(0) .parse::() .expect("server_version_num should be numeric"); let has_conparentid = client .query_one( "SELECT EXISTS ( \ SELECT 1 FROM pg_catalog.pg_attribute \ WHERE attrelid = 'pg_catalog.pg_constraint'::pg_catalog.regclass \ AND attname = 'conparentid' AND NOT attisdropped \ )", &[], ) .await .expect("probe pg_constraint.conparentid") .get::<_, bool>(0); if version_num < 110000 { assert!(!has_conparentid, "pre-11 PostgreSQL must not have parent-linked constraint clones"); } let empty = list_foreign_keys_for_relations(&pool, &[]).await.expect("empty relation batch"); assert!(empty.is_empty()); let single_relations = vec![(schema.clone(), "child_a".to_string())]; let single = list_foreign_keys_for_relations(&pool, &single_relations).await.expect("single relation batch"); let single_fk = &single[&(schema.clone(), "child_a".to_string())][0]; assert_eq!(single_fk.name, "child_a_parent_fk"); assert_eq!(single_fk.ref_table, "parent_a"); // Keep the pairs deliberately out of name order. A cross product or // independently unnested array would duplicate or misattribute keys. let relations = vec![ (schema.clone(), "child_b".to_string()), (schema.clone(), "without_fk".to_string()), (schema.clone(), "child_a".to_string()), ]; let batched = list_foreign_keys_for_relations(&pool, &relations).await.expect("multi-relation batch"); assert_eq!(batched.len(), 2, "batched keys: {batched:?}"); assert!(!batched.contains_key(&(schema.clone(), "without_fk".to_string()))); assert_eq!(batched[&(schema.clone(), "child_a".to_string())][0].ref_table, "parent_a"); assert_eq!(batched[&(schema.clone(), "child_b".to_string())][0].ref_table, "parent_b"); let schemas: Vec<&str> = relations.iter().map(|(schema, _)| schema.as_str()).collect(); let tables: Vec<&str> = relations.iter().map(|(_, table)| table.as_str()).collect(); let compat = list_foreign_keys_for_relations_with_sql( &client, postgres_foreign_keys_for_relations_compat_sql(), &schemas, &tables, ) .await .expect("subscript-paired compatibility query"); assert_eq!(compat[&(schema.clone(), "child_a".to_string())][0].ref_table, "parent_a"); assert_eq!(compat[&(schema.clone(), "child_b".to_string())][0].ref_table, "parent_b"); let modern = list_foreign_keys_for_relations_with_sql( &client, postgres_foreign_keys_for_relations_sql(), &schemas, &tables, ) .await; if version_num >= 90400 { let modern = modern.expect("multi-array unnest primary query"); assert_eq!(modern[&(schema.clone(), "child_a".to_string())][0].ref_table, "parent_a"); assert_eq!(modern[&(schema.clone(), "child_b".to_string())][0].ref_table, "parent_b"); } else { modern.expect_err("pre-9.4 server must reject multi-array unnest"); } let relation_oids: Vec = client .query( "SELECT c.oid::bigint \ FROM pg_catalog.pg_class c \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ WHERE n.nspname = $1 ORDER BY c.relname", &[&schema], ) .await .expect("query fixture relation oids") .into_iter() .map(|row| row.get(0)) .collect(); get_table_partition_local_objects_for_relations_with_sql( &client, postgres_table_partition_local_objects_for_relations_compat_sql(), &relation_oids, ) .await .expect("pre-parent-clone local-object query"); let modern_local_objects = get_table_partition_local_objects_for_relations_with_sql( &client, postgres_table_partition_local_objects_for_relations_sql(), &relation_oids, ) .await; if version_num >= 90300 { modern_local_objects.expect("JSON-extraction local-object query"); } else { modern_local_objects.expect_err("PostgreSQL 9.2 must reject the JSON ->> operator"); } drop(client); let child_ddl = crate::schema::pg_ddl_with_partitions(&pool, &schema, "child_a") .await .expect("display DDL with foreign key"); let without_fk_ddl = crate::schema::pg_ddl_with_partitions(&pool, &schema, "without_fk") .await .expect("display DDL without foreign key"); execute_query(&pool, &format!("DROP SCHEMA {schema_ident} CASCADE")) .await .expect("drop batched foreign-key fixtures"); assert!(child_ddl.contains("FOREIGN KEY"), "child ddl: {child_ddl}"); assert!(!without_fk_ddl.contains("FOREIGN KEY"), "without-fk ddl: {without_fk_ddl}"); } #[tokio::test] #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a writable PostgreSQL database"] async fn postgres_ddl_emits_not_valid_check_constraints_as_separate_replayable_alter() { let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); let pool = connect(&url, Duration::from_secs(5)).await.expect("connect postgres"); let pid = std::process::id(); let schema = format!("dbx_not_valid_check_{pid}"); let schema_ident = pg_quote_ident(&schema); let table = format!("{schema_ident}.events"); execute_query(&pool, &format!("CREATE SCHEMA {schema_ident}")).await.expect("create schema"); let client = pool.get().await.expect("get postgres client"); client .batch_execute(&format!( "CREATE TABLE {table} (id integer, amount integer); \ ALTER TABLE {table} ADD CONSTRAINT c_pos CHECK (amount > 0) NOT VALID" )) .await .expect("create table with unvalidated check constraint"); let ddl = crate::schema::pg_ddl(&pool, &schema, "events").await.expect("events ddl"); // A `NOT VALID` CHECK constraint is illegal inline in a `CREATE // TABLE` column list — Postgres only accepts it after `ALTER TABLE // ADD CONSTRAINT`. Replaying the generated DDL onto a fresh schema // exercises that this is real, executable syntax, not just a string // match. let replay_schema = format!("{schema}_replay"); let replay_schema_ident = pg_quote_ident(&replay_schema); execute_query(&pool, &format!("CREATE SCHEMA {replay_schema_ident}")).await.expect("create replay schema"); client .batch_execute(&ddl.replace(&schema, &replay_schema)) .await .unwrap_or_else(|error| panic!("replay not-valid-check ddl failed: {error:?}; ddl: {ddl}")); client .batch_execute(&format!("DROP SCHEMA {schema_ident} CASCADE; DROP SCHEMA {replay_schema_ident} CASCADE")) .await .expect("drop schemas"); assert!(!ddl.contains("NOT VALID)"), "NOT VALID must not appear inline in the column list: {ddl}"); assert!(ddl.contains("CREATE TABLE"), "ddl: {ddl}"); assert!(!ddl.contains("\"amount\" integer,"), "check constraint must not be inlined into columns: {ddl}"); assert!(ddl.contains("ALTER TABLE"), "ddl: {ddl}"); assert!( ddl.contains("ADD CONSTRAINT \"c_pos\" CHECK (amount > 0) NOT VALID;"), "expected a separate ALTER TABLE ADD CONSTRAINT ... NOT VALID statement: {ddl}" ); } #[tokio::test] #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a writable PostgreSQL database"] async fn postgres_partition_metadata_renders_replayable_children_and_subpartitions() { let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); let pool = connect(&url, Duration::from_secs(5)).await.expect("connect postgres"); let pid = std::process::id(); let schema = format!("dbx_partition_meta_{pid}"); let schema_ident = pg_quote_ident(&schema); let replay_schema = format!("{schema}_replay"); let replay_schema_ident = pg_quote_ident(&replay_schema); let parent = format!("{schema_ident}.parent"); let child = format!("{schema_ident}.child"); let default_child = format!("{schema_ident}.child_default"); let subpartition = format!("{schema_ident}.subpartition"); let inherited_parent = format!("{schema_ident}.inherited_parent"); let inherited_child = format!("{schema_ident}.inherited_child"); // Declared in "public" (not the per-test schema) so the same // qualified name keeps resolving after `schema_ident` substrings are // rewritten to `replay_schema_ident` for the replay pass below. let trigger_fn = format!("public.dbx_partition_meta_trigger_fn_{pid}"); execute_query(&pool, &format!("CREATE SCHEMA {schema_ident}")).await.expect("create schema"); let client = pool.get().await.expect("get postgres client"); client .batch_execute(&format!( "CREATE OR REPLACE FUNCTION {trigger_fn}() RETURNS trigger LANGUAGE plpgsql AS \ $$ BEGIN RETURN NEW; END $$; \ CREATE TABLE {parent} (id integer, payload text, status text DEFAULT 'pending') PARTITION BY RANGE (id); \ CREATE TABLE {child} PARTITION OF {parent} ( \ PRIMARY KEY (id), \ CONSTRAINT child_payload_not_blank CHECK (payload <> ''), \ status WITH OPTIONS DEFAULT 'child-status' \ ) FOR VALUES FROM (1) TO (10); \ CREATE INDEX child_payload_idx ON {child} (payload); \ CREATE TABLE {default_child} PARTITION OF {parent} DEFAULT; \ CREATE TABLE {subpartition} PARTITION OF {parent} FOR VALUES FROM (10) TO (20) PARTITION BY HASH (payload); \ ALTER TABLE {parent} ADD CONSTRAINT parent_id_positive CHECK (id > 0); \ CREATE TRIGGER parent_biu BEFORE INSERT ON {parent} FOR EACH ROW EXECUTE FUNCTION {trigger_fn}(); \ CREATE TABLE {inherited_parent} (id integer, bucket integer, PRIMARY KEY (id, bucket)) PARTITION BY RANGE (bucket); \ CREATE TABLE {inherited_child} PARTITION OF {inherited_parent} FOR VALUES FROM (1) TO (10)" )) .await .expect("create partitioned tables"); let parent_info = get_table_partition_info(&pool, &schema, "parent").await.expect("parent metadata"); let child_info = get_table_partition_info(&pool, &schema, "child").await.expect("child metadata"); let default_info = get_table_partition_info(&pool, &schema, "child_default").await.expect("default metadata"); let subpartition_info = get_table_partition_info(&pool, &schema, "subpartition").await.expect("subpartition metadata"); let child_local_objects = get_table_partition_local_objects(&pool, &schema, "child").await.expect("child local objects"); let inherited_child_local_objects = get_table_partition_local_objects(&pool, &schema, "inherited_child") .await .expect("inherited child local objects"); let parent_ddl = crate::schema::pg_ddl_with_partitions(&pool, &schema, "parent").await.expect("parent ddl"); let child_ddl = crate::schema::pg_ddl(&pool, &schema, "child").await.expect("child ddl"); let default_ddl = crate::schema::pg_ddl(&pool, &schema, "child_default").await.expect("default ddl"); let subpartition_ddl = crate::schema::pg_ddl(&pool, &schema, "subpartition").await.expect("subpartition ddl"); let inherited_parent_ddl = crate::schema::pg_ddl_with_partitions(&pool, &schema, "inherited_parent") .await .expect("inherited parent ddl"); let inherited_child_ddl = crate::schema::pg_ddl(&pool, &schema, "inherited_child").await.expect("inherited child ddl"); // The bare (non-recursive) parent DDL, exactly what database export // and table transfer produce per-relation — used below to give the // single-relation partition ddls (child_ddl, default_ddl, // subpartition_ddl, inherited_child_ddl) a real parent shell to // replay against, so their syntax is actually exercised by // PostgreSQL rather than only checked with string assertions. let bare_parent_ddl = crate::schema::pg_ddl(&pool, &schema, "parent").await.expect("bare parent ddl"); let bare_inherited_parent_ddl = crate::schema::pg_ddl(&pool, &schema, "inherited_parent").await.expect("bare inherited parent ddl"); // `parent_ddl`/`inherited_parent_ddl` already embed their full // partition subtree (that's the fix under test), so replaying the // per-partition ddl strings on top would recreate the same relations // twice. Only the two top-level requests need replaying; the // individually-fetched partition ddls above are replayed separately // below, on top of a bare (childless) copy of their parent. execute_query(&pool, &format!("CREATE SCHEMA {replay_schema_ident}")).await.expect("create replay schema"); for ddl in [&parent_ddl, &inherited_parent_ddl] { client .batch_execute(&ddl.replace(&schema, &replay_schema)) .await .unwrap_or_else(|error| panic!("replay partition ddl failed: {error:?}; ddl: {ddl}")); } // Single-relation partition ddls are now a real production code // path (database export/table transfer fetch each partition this // way), so replay them for real too: a fresh schema gets just the // bare parents, then each single-relation partition ddl on top. let bare_schema = format!("{schema}_bare"); let bare_schema_ident = pg_quote_ident(&bare_schema); execute_query(&pool, &format!("CREATE SCHEMA {bare_schema_ident}")).await.expect("create bare schema"); for ddl in [&bare_parent_ddl, &bare_inherited_parent_ddl] { client .batch_execute(&ddl.replace(&schema, &bare_schema)) .await .unwrap_or_else(|error| panic!("replay bare parent ddl failed: {error:?}; ddl: {ddl}")); } for ddl in [&child_ddl, &default_ddl, &subpartition_ddl, &inherited_child_ddl] { client .batch_execute(&ddl.replace(&schema, &bare_schema)) .await .unwrap_or_else(|error| panic!("replay single-relation partition ddl failed: {error:?}; ddl: {ddl}")); } client .batch_execute(&format!( "DROP SCHEMA {schema_ident} CASCADE; DROP SCHEMA {replay_schema_ident} CASCADE; \ DROP SCHEMA {bare_schema_ident} CASCADE; DROP FUNCTION IF EXISTS {trigger_fn}()" )) .await .expect("drop schemas and trigger function"); assert_eq!(parent_info.key.as_deref(), Some("RANGE (id)")); assert!(!parent_info.is_partition); assert_eq!(child_info.parent_schema.as_deref(), Some(schema.as_str())); assert_eq!(child_info.parent_table.as_deref(), Some("parent")); assert_eq!(child_info.bound.as_deref(), Some("FOR VALUES FROM (1) TO (10)")); assert!(child_local_objects.has_primary_key); assert!(child_local_objects.indexes.contains("child_payload_idx")); assert!(child_local_objects.check_constraints.contains("child_payload_not_blank")); assert!(!child_local_objects.check_constraints.contains("parent_id_positive")); assert_eq!(child_local_objects.column_defaults.get("status"), Some(&PostgresColumnDefaultState::Overridden)); assert!(!inherited_child_local_objects.has_primary_key); assert!(inherited_child_local_objects.indexes.is_empty()); assert_eq!(default_info.bound.as_deref(), Some("DEFAULT")); assert_eq!(subpartition_info.key.as_deref(), Some("HASH (payload)")); assert!(parent_ddl.contains(") PARTITION BY RANGE (id);"), "ddl: {parent_ddl}"); assert!(parent_ddl.contains("CONSTRAINT \"parent_id_positive\" CHECK (id > 0)"), "ddl: {parent_ddl}"); // The parent's trigger is automatically cloned onto every partition // by PostgreSQL itself when the partition is created/attached, so it // must appear exactly once in the combined DDL, not once per relation. assert_eq!(parent_ddl.matches("CREATE TRIGGER").count(), 1, "ddl: {parent_ddl}"); assert!(child_ddl.contains("CREATE TABLE"), "ddl: {child_ddl}"); assert!(child_ddl.contains("PARTITION OF"), "ddl: {child_ddl}"); assert!(child_ddl.contains("PRIMARY KEY (\"id\")"), "ddl: {child_ddl}"); assert!(child_ddl.contains("FOR VALUES FROM (1) TO (10);"), "ddl: {child_ddl}"); assert!(child_ddl.contains("CREATE INDEX \"child_payload_idx\""), "ddl: {child_ddl}"); assert!(child_ddl.contains("CONSTRAINT \"child_payload_not_blank\" CHECK (payload"), "ddl: {child_ddl}"); assert!(!child_ddl.contains("parent_id_positive"), "ddl: {child_ddl}"); assert!(child_ddl.contains("\"status\" WITH OPTIONS DEFAULT 'child-status'"), "ddl: {child_ddl}"); assert!( !child_ddl.contains("CREATE TRIGGER"), "child must not redeclare the parent's cloned trigger: {child_ddl}" ); assert!(default_ddl.contains(" DEFAULT;"), "ddl: {default_ddl}"); assert!( subpartition_ddl.contains("FOR VALUES FROM (10) TO (20) PARTITION BY HASH (payload);"), "ddl: {subpartition_ddl}" ); assert!(!inherited_child_ddl.contains("PRIMARY KEY"), "ddl: {inherited_child_ddl}"); assert!(!inherited_child_ddl.contains("CREATE INDEX"), "ddl: {inherited_child_ddl}"); } #[tokio::test] #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a writable PostgreSQL database with postgres_fdw available"] async fn postgres_partition_ddl_uses_foreign_table_syntax_for_foreign_partitions() { let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); let pool = connect(&url, Duration::from_secs(5)).await.expect("connect postgres"); let pid = std::process::id(); let schema = format!("dbx_partition_fdw_{pid}"); let schema_ident = pg_quote_ident(&schema); let server = format!("dbx_partition_fdw_server_{pid}"); let server_ident = pg_quote_ident(&server); let replay_schema = format!("{schema}_replay"); let replay_schema_ident = pg_quote_ident(&replay_schema); let parent = format!("{schema_ident}.parent"); let local_child = format!("{schema_ident}.local_child"); let foreign_child = format!("{schema_ident}.foreign_child"); execute_query(&pool, &format!("CREATE SCHEMA {schema_ident}")).await.expect("create schema"); let client = pool.get().await.expect("get postgres client"); client .batch_execute(&format!( "CREATE EXTENSION IF NOT EXISTS postgres_fdw; \ CREATE SERVER {server_ident} FOREIGN DATA WRAPPER postgres_fdw \ OPTIONS (host 'localhost', dbname 'postgres', port '5432'); \ CREATE USER MAPPING FOR CURRENT_USER SERVER {server_ident} \ OPTIONS (user 'postgres', password 'postgres'); \ CREATE TABLE {parent} (id integer, payload text) PARTITION BY RANGE (id); \ CREATE TABLE {local_child} PARTITION OF {parent} FOR VALUES FROM (1) TO (10); \ CREATE FOREIGN TABLE {foreign_child} PARTITION OF {parent} \ FOR VALUES FROM (10) TO (20) \ SERVER {server_ident} OPTIONS (schema_name 'public', table_name 'foreign_child_data')" )) .await .expect("create partitioned tables with a foreign partition"); let foreign_info = get_table_partition_info(&pool, &schema, "foreign_child").await.expect("foreign child info"); let local_info = get_table_partition_info(&pool, &schema, "local_child").await.expect("local child info"); let parent_ddl = crate::schema::pg_ddl_with_partitions(&pool, &schema, "parent").await.expect("parent ddl"); let foreign_ddl = crate::schema::pg_ddl(&pool, &schema, "foreign_child").await.expect("foreign child ddl"); execute_query(&pool, &format!("CREATE SCHEMA {replay_schema_ident}")).await.expect("create replay schema"); client .batch_execute(&parent_ddl.replace(&schema, &replay_schema)) .await .unwrap_or_else(|error| panic!("replay foreign partition ddl failed: {error:?}; ddl: {parent_ddl}")); client .batch_execute(&format!( "DROP SCHEMA {schema_ident} CASCADE; DROP SCHEMA {replay_schema_ident} CASCADE; \ DROP SERVER {server_ident} CASCADE" )) .await .expect("drop schemas and foreign server"); assert!(!local_info.is_foreign); assert!(foreign_info.is_foreign); assert_eq!(foreign_info.foreign_server.as_deref(), Some(server.as_str())); assert!( foreign_info.foreign_options.contains(&("schema_name".to_string(), "public".to_string())), "options: {:?}", foreign_info.foreign_options ); assert!(parent_ddl.contains("CREATE FOREIGN TABLE"), "ddl: {parent_ddl}"); assert!(foreign_ddl.starts_with("CREATE FOREIGN TABLE"), "ddl: {foreign_ddl}"); assert!(foreign_ddl.contains("PARTITION OF"), "ddl: {foreign_ddl}"); assert!(foreign_ddl.contains(&format!("SERVER {}", pg_quote_ident(&server))), "ddl: {foreign_ddl}"); assert!(foreign_ddl.contains("\"schema_name\" 'public'"), "ddl: {foreign_ddl}"); assert!(foreign_ddl.contains("\"table_name\" 'foreign_child_data'"), "ddl: {foreign_ddl}"); assert!(!foreign_ddl.contains("PRIMARY KEY"), "ddl: {foreign_ddl}"); assert!(!foreign_ddl.contains("CREATE INDEX"), "ddl: {foreign_ddl}"); } #[tokio::test] #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a writable PostgreSQL database"] async fn postgres_schema_context_prioritizes_selected_schema_and_cleans_up() { let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); let pool = connect(&url, Duration::from_secs(5)).await.expect("connect postgres"); let suffix = format!("{}_{}", std::process::id(), uuid::Uuid::new_v4().simple()); let schema = format!("dbx_issue_830_\"{suffix}"); let schema_ident = pg_quote_ident(&schema); let helper = format!("dbx_issue_830_public_{suffix}"); let helper_ident = pg_quote_ident(&helper); let initial_path = execute_query(&pool, "SHOW search_path").await.expect("read initial search_path"); let initial_path_value = initial_path.rows[0][0].as_str().expect("search_path string").to_string(); let client = pool.get().await.expect("get setup client"); client .batch_execute(&format!( "CREATE SCHEMA {schema_ident}; \ CREATE TABLE {schema_ident}.pg_settings(marker text); \ INSERT INTO {schema_ident}.pg_settings VALUES ('selected-schema'); \ CREATE FUNCTION public.{helper_ident}() RETURNS text \ LANGUAGE SQL IMMUTABLE AS $$ SELECT 'public-fallback'::text $$" )) .await .expect("create search_path fixtures"); drop(client); let query_sql = format!("SELECT marker, {helper_ident}() AS helper FROM pg_settings"); let stale_session_recovery = async { let client = pool.get().await.map_err(|error| error.to_string())?; set_postgres_search_path(&client, &schema, PostgresSearchPathContext::Query, Duration::from_secs(5)) .await?; client.execute("SET search_path TO public", &[]).await.map_err(pg_error_to_string)?; set_postgres_search_path(&client, &schema, PostgresSearchPathContext::Query, Duration::from_secs(5)) .await?; let selected: String = client .query_one("SELECT marker FROM pg_settings", &[]) .await .map_err(pg_error_to_string)? .try_get(0) .map_err(pg_error_to_string)?; client.execute("RESET search_path", &[]).await.map_err(pg_error_to_string)?; Ok::<_, String>(selected) } .await; let ordinary_result = execute_query_with_schema(&pool, &schema, &query_sql).await; let path_after_ordinary = execute_query(&pool, "SHOW search_path").await; let read_only_result = execute_query_in_read_only_transaction_with_rollback( &pool, Some(&schema), &query_sql, None, None, DbOperationBudget::with_defaults(), None, ) .await; let path_after_read_only = execute_query(&pool, "SHOW search_path").await; let mut streamed_rows = Vec::new(); let streaming_result = stream_select_query_with_cancel( &pool, Some(&schema), &[], &query_sql, None, None, DbOperationBudget::with_defaults(), None, |item| { if let PostgresQueryStreamItem::Row(row) = item { streamed_rows.push(row); } Ok(()) }, ) .await; let path_after_streaming = execute_query(&pool, "SHOW search_path").await; let transaction_cleanup = async { let client = pool.get().await.map_err(|error| error.to_string())?; client .execute(&postgres_set_search_path_sql(&schema, PostgresSearchPathContext::Transaction), &[]) .await .map_err(pg_error_to_string)?; let selected: String = client .query_one("SELECT marker FROM pg_settings", &[]) .await .map_err(pg_error_to_string)? .try_get(0) .map_err(pg_error_to_string)?; client.execute("RESET search_path", &[]).await.map_err(pg_error_to_string)?; let after_reset: String = client .query_one("SHOW search_path", &[]) .await .map_err(pg_error_to_string)? .try_get(0) .map_err(pg_error_to_string)?; client.execute("BEGIN", &[]).await.map_err(pg_error_to_string)?; client .execute(&postgres_set_search_path_sql(&schema, PostgresSearchPathContext::LocalTransaction), &[]) .await .map_err(pg_error_to_string)?; let local_selected: String = client .query_one("SELECT marker FROM pg_settings", &[]) .await .map_err(pg_error_to_string)? .try_get(0) .map_err(pg_error_to_string)?; client.execute("COMMIT", &[]).await.map_err(pg_error_to_string)?; let after_commit: String = client .query_one("SHOW search_path", &[]) .await .map_err(pg_error_to_string)? .try_get(0) .map_err(pg_error_to_string)?; Ok::<_, String>((selected, after_reset, local_selected, after_commit)) } .await; let cleanup_client = pool.get().await.expect("get cleanup client"); cleanup_client .batch_execute(&format!("DROP FUNCTION public.{helper_ident}(); DROP SCHEMA {schema_ident} CASCADE")) .await .expect("clean search_path fixtures"); let ordinary = ordinary_result.expect("ordinary schema query"); assert_eq!(stale_session_recovery.expect("recover stale session search_path"), "selected-schema"); assert_eq!( ordinary.rows, vec![vec![serde_json::json!("selected-schema"), serde_json::json!("public-fallback")]] ); assert_eq!(path_after_ordinary.expect("path after ordinary query").rows, initial_path.rows); assert_eq!( read_only_result.expect("read-only schema query").rows, vec![vec![serde_json::json!("selected-schema"), serde_json::json!("public-fallback")]] ); assert_eq!(path_after_read_only.expect("path after read-only query").rows, initial_path.rows); assert_eq!(streaming_result.expect("streaming schema query"), 1); assert_eq!( streamed_rows, vec![vec![serde_json::json!("selected-schema"), serde_json::json!("public-fallback")]] ); assert_eq!(path_after_streaming.expect("path after streaming query").rows, initial_path.rows); let (selected, after_reset, local_selected, after_commit) = transaction_cleanup.expect("transaction cleanup"); assert_eq!(selected, "selected-schema"); assert_eq!(local_selected, "selected-schema"); assert_eq!(after_reset, initial_path_value); assert_eq!(after_commit, initial_path_value); } #[test] fn postgres_index_metadata_has_legacy_catalog_fallback() { assert!(POSTGRES_INDEXES_SQL.contains("ix.indnkeyatts")); assert!(!POSTGRES_INDEXES_COMPAT_SQL.contains("ix.indnkeyatts")); assert!(POSTGRES_INDEXES_COMPAT_SQL.contains("NULL::smallint AS nkeyatts")); assert!(!POSTGRES_INDEXES_COMPAT_SQL.contains("LATERAL")); assert!(!POSTGRES_INDEXES_COMPAT_SQL.contains("WITH ORDINALITY")); assert!(POSTGRES_INDEXES_COMPAT_SQL.contains("generate_series")); assert!(POSTGRES_INDEXES_COMPAT_SQL.contains("string_to_array(ix.indkey::text, ' ')")); for sql in [ POSTGRES_INDEXES_SQL, POSTGRES_INDEXES_COMPAT_SQL, postgres_indexes_for_relations_sql(), postgres_indexes_for_relations_compat_sql(), ] { assert!(sql.contains("ix.indisunique AND ix.indisvalid")); } } #[test] fn postgres_index_metadata_tracks_expression_key_provenance() { // #6295 fix: each index key part's `a.attname IS NULL` tags whether it came from a real // column or from pg_get_indexdef, so DDL generation never has to guess from the text. assert!(POSTGRES_INDEXES_SQL.contains("array_agg(a.attname IS NULL ORDER BY k.n) AS key_is_expression")); assert!(POSTGRES_INDEXES_COMPAT_SQL.contains("SELECT a.attname IS NULL")); assert!(POSTGRES_INDEXES_COMPAT_SQL.contains("AS key_is_expression")); } #[test] fn postgres_index_metadata_schema_qualifies_opclass() { // The LATERAL (modern) index-introspection SQL returns each non-default // opclass schema-qualified (`quote_ident(nspname) || '.' || quote_ident(opcname)`) // and joins `pg_namespace` on `opcnamespace`, so DDL regeneration resolves // opclasses (e.g. `gin_trgm_ops` from `pg_trgm`) regardless of `search_path`. // The compat (pre-LATERAL) fallback keeps bare `oc.opcname` — no regression, // just no schema-qualification for the rare old-PG × non-default-schema case. for sql in [POSTGRES_INDEXES_SQL, postgres_indexes_for_relations_sql()] { assert!( sql.contains("quote_ident(opcns.nspname) || '.' || quote_ident(oc.opcname)"), "expected schema-qualified opclass in {sql}" ); assert!( sql.contains("LEFT JOIN pg_namespace opcns ON opcns.oid = oc.opcnamespace"), "expected opcnamespace join in {sql}" ); } } #[test] fn postgres_owner_metadata_casts_relkind_to_text() { assert!(POSTGRES_OWNERS_SQL.contains("c.relkind::text AS relkind")); assert!(POSTGRES_OWNERS_SQL.contains("c.relkind IN ('r', 'v', 'm', 'S', 'f', 'p')")); } #[test] fn postgres_owner_object_type_maps_relkind_codes() { assert_eq!(postgres_owner_object_type("r"), "TABLE"); assert_eq!(postgres_owner_object_type("v"), "VIEW"); assert_eq!(postgres_owner_object_type("m"), "MATERIALIZED_VIEW"); assert_eq!(postgres_owner_object_type("S"), "SEQUENCE"); assert_eq!(postgres_owner_object_type("f"), "FOREIGN TABLE"); assert_eq!(postgres_owner_object_type("p"), "PARTITIONED TABLE"); assert_eq!(postgres_owner_object_type("?"), "?"); } fn test_custom_type_info(typtype: &str) -> CustomTypeGeneralInfo { CustomTypeGeneralInfo { oid: 100, typtype: typtype.to_string(), typisdefined: true, typbasetype: 0, typnotnull: false, typrelid: 0, typelem: 0, typcollation: 0, typdefaultbin: None, typdefault: None, typlen: -1, typbyval: false, typalign: "c".to_string(), typstorage: "p".to_string(), typtypmod: -1, input_function: None, output_function: None, receive_function: None, send_function: None, analyze_function: None, comment: None, relkind: None, collation: None, } } #[test] fn custom_type_ddl_enum_quotes_values_in_order() { let info = test_custom_type_info("e"); let members = vec![ CustomTypeMember { name: String::new(), data_type: String::new(), ordinal: 1, nullable: None, default: None, comment: None, enum_value: Some("draft".to_string()), }, CustomTypeMember { name: String::new(), data_type: String::new(), ordinal: 2, nullable: None, default: None, comment: None, enum_value: Some("已归档".to_string()), }, ]; let ddl = build_custom_type_ddl( "app", "status", CustomTypeKind::Enum, &info, &members, &CustomTypeProperties::default(), &[], ); assert_eq!(ddl.sql, "CREATE TYPE \"app\".\"status\" AS ENUM ('draft', '已归档');"); assert!(ddl.complete); assert!(ddl.warnings.is_empty()); } #[test] fn custom_type_ddl_composite_lists_fields_and_comments() { let info = test_custom_type_info("c"); let members = vec![ CustomTypeMember { name: "city".to_string(), data_type: "text".to_string(), ordinal: 1, nullable: Some(true), default: None, comment: Some("city name".to_string()), enum_value: None, }, CustomTypeMember { name: "zip".to_string(), data_type: "numeric(6)".to_string(), ordinal: 2, nullable: Some(true), default: None, comment: None, enum_value: None, }, ]; let ddl = build_custom_type_ddl( "app", "address", CustomTypeKind::Composite, &info, &members, &CustomTypeProperties::default(), &[], ); assert_eq!( ddl.sql, "CREATE TYPE \"app\".\"address\" AS (\n \"city\" text,\n \"zip\" numeric(6)\n);\nCOMMENT ON COLUMN \"app\".\"address\".\"city\" IS 'city name';" ); assert!(ddl.complete); } #[test] fn custom_type_ddl_domain_combines_base_default_notnull_and_constraints() { let info = test_custom_type_info("d"); let properties = CustomTypeProperties { base_type: Some("text".to_string()), collation: Some("C".to_string()), default: Some("''::text".to_string()), not_null: Some(true), domain_constraints: vec![CustomTypeDomainConstraint { name: "email_valid".to_string(), definition: "CHECK ((VALUE <> ''::text))".to_string(), }], ..Default::default() }; let ddl = build_custom_type_ddl("app", "email", CustomTypeKind::Domain, &info, &[], &properties, &[]); assert_eq!( ddl.sql, "CREATE DOMAIN \"app\".\"email\" AS text\n COLLATE \"C\"\n DEFAULT ''::text\n NOT NULL\n CONSTRAINT \"email_valid\" CHECK ((VALUE <> ''::text));" ); assert!(ddl.complete); } #[test] fn custom_type_ddl_domain_is_incomplete_when_attributes_failed() { let info = test_custom_type_info("d"); let properties = CustomTypeProperties { base_type: Some("text".to_string()), ..Default::default() }; let ddl = build_custom_type_ddl( "app", "email", CustomTypeKind::Domain, &info, &[], &properties, &["domain constraints could not be read: x".to_string()], ); assert!(!ddl.complete, "DDL must be marked incomplete when constraints failed: {ddl:?}"); let constraint_decode_failed = build_custom_type_ddl( "app", "email", CustomTypeKind::Domain, &info, &[], &properties, &["domain constraints could not be decoded: x".to_string()], ); assert!( !constraint_decode_failed.complete, "DDL must be marked incomplete when constraints cannot be decoded: {constraint_decode_failed:?}" ); let default_failed = build_custom_type_ddl( "app", "email", CustomTypeKind::Domain, &info, &[], &properties, &["default value could not be rendered; the generated DDL is incomplete".to_string()], ); assert!(!default_failed.complete, "DDL must be marked incomplete when the default failed: {default_failed:?}"); assert!(default_failed.warnings.iter().any(|w| w.contains("default value"))); } #[test] fn custom_type_qualified_format_type_expression_qualifies_user_types() { let expression = postgres_qualified_format_type_expression("t", "n", "elem", "elem_n", "t.oid", "t.typtypmod"); assert!(expression.contains("quote_ident(n.nspname) || '.' || quote_ident(t.typname)")); assert!(expression.contains("quote_ident(elem_n.nspname) || '.' || quote_ident(elem.typname) || '[]'")); assert!(expression.contains("format_type(t.oid, t.typtypmod)")); } #[test] fn custom_type_ddl_range_uses_subtype_and_canonical() { let info = test_custom_type_info("r"); let properties = CustomTypeProperties { range_subtype: Some("numeric".to_string()), range_canonical_function: Some("numeric_range_canonical".to_string()), range_subtype_diff_function: Some("numeric_range_subdiff".to_string()), ..Default::default() }; let ddl = build_custom_type_ddl("app", "price_range", CustomTypeKind::Range, &info, &[], &properties, &[]); assert_eq!( ddl.sql, "CREATE TYPE \"app\".\"price_range\" AS RANGE (\n subtype = numeric,\n canonical = \"numeric_range_canonical\",\n subtype_diff = \"numeric_range_subdiff\"\n);" ); assert!(ddl.complete); } #[test] fn custom_type_ddl_range_without_subtype_is_incomplete() { let info = test_custom_type_info("r"); let properties = CustomTypeProperties { range_multirange_name: Some("price_multirange".to_string()), range_canonical_function: Some("range_canonical".to_string()), ..Default::default() }; let ddl = build_custom_type_ddl("app", "price_range", CustomTypeKind::Range, &info, &[], &properties, &[]); assert!(!ddl.complete, "a range without subtype cannot be reconstructed: {ddl:?}"); assert_eq!(ddl.sql, "CREATE TYPE \"app\".\"price_range\" AS RANGE (subtype = unknown);"); assert!(ddl.warnings.iter().any(|warning| warning.contains("range attributes"))); } #[test] fn custom_type_ddl_range_keeps_catalog_qualified_names() { let info = test_custom_type_info("r"); let properties = CustomTypeProperties { range_subtype: Some("numeric".to_string()), range_canonical_function: Some("\"extensions\".\"range_canonical\"".to_string()), range_subtype_diff_function: Some("\"extensions\".\"range_subdiff\"".to_string()), range_subtype_opclass: Some("\"extensions\".\"numeric_ops\"".to_string()), ..Default::default() }; let ddl = build_custom_type_ddl("app", "price_range", CustomTypeKind::Range, &info, &[], &properties, &[]); assert!(ddl.complete, "{ddl:?}"); assert!(ddl.sql.contains("subtype_opclass = \"extensions\".\"numeric_ops\"")); assert!(ddl.sql.contains("canonical = \"extensions\".\"range_canonical\"")); assert!(ddl.sql.contains("subtype_diff = \"extensions\".\"range_subdiff\"")); } #[test] fn custom_type_ddl_range_emits_custom_multirange_name() { let info = test_custom_type_info("r"); let properties = CustomTypeProperties { range_subtype: Some("numeric".to_string()), range_multirange_name: Some("price_multirange".to_string()), ..Default::default() }; let ddl = build_custom_type_ddl("app", "price_range", CustomTypeKind::Range, &info, &[], &properties, &[]); assert!(ddl.complete, "{ddl:?}"); assert!( ddl.sql.contains("multirange_type_name = \"price_multirange\""), "custom multirange name must be emitted: {}", ddl.sql ); } #[test] fn custom_type_ddl_multirange_and_base_are_marked_incomplete() { let multirange = build_custom_type_ddl( "app", "_price_range", CustomTypeKind::Multirange, &test_custom_type_info("m"), &[], &CustomTypeProperties::default(), &[], ); assert!(!multirange.complete); assert!(multirange.sql.is_empty()); assert!(multirange.warnings.iter().any(|w| w.contains("multirange companion"))); let base = build_custom_type_ddl( "app", "point2d", CustomTypeKind::Base, &test_custom_type_info("b"), &[], &CustomTypeProperties::default(), &[], ); assert!(!base.complete); assert!(base.warnings.iter().any(|w| w.contains("base type"))); } #[test] fn custom_type_ddl_escapes_identifiers_and_literals() { let info = test_custom_type_info("e"); let members = vec![CustomTypeMember { name: String::new(), data_type: String::new(), ordinal: 1, nullable: None, default: None, comment: None, enum_value: Some("it's \"quoted\"".to_string()), }]; let ddl = build_custom_type_ddl( "we\"ird", "ty\"pe", CustomTypeKind::Enum, &info, &members, &CustomTypeProperties::default(), &[], ); assert_eq!(ddl.sql, "CREATE TYPE \"we\"\"ird\".\"ty\"\"pe\" AS ENUM ('it''s \"quoted\"');"); } #[test] fn custom_type_kind_for_maps_catalog_codes() { assert_eq!(custom_type_kind_for("b"), Some(CustomTypeKind::Base)); assert_eq!(custom_type_kind_for("c"), Some(CustomTypeKind::Composite)); assert_eq!(custom_type_kind_for("d"), Some(CustomTypeKind::Domain)); assert_eq!(custom_type_kind_for("e"), Some(CustomTypeKind::Enum)); assert_eq!(custom_type_kind_for("r"), Some(CustomTypeKind::Range)); assert_eq!(custom_type_kind_for("m"), Some(CustomTypeKind::Multirange)); assert_eq!(custom_type_kind_for("p"), None); } #[test] fn custom_type_list_metadata_decodes_kind_and_member_capability() { assert_eq!(custom_type_list_metadata(Some("c:1")), (Some("composite".to_string()), Some(true))); assert_eq!(custom_type_list_metadata(Some("d:0")), (Some("domain".to_string()), Some(false))); assert_eq!(custom_type_list_metadata(None), (None, None)); assert_eq!(custom_type_list_metadata(Some("invalid")), (None, None)); } #[test] fn custom_type_general_info_sql_looks_up_by_schema_and_name() { let sql = custom_type_general_info_sql(); assert!(sql.contains("pg_catalog.pg_type")); assert!(sql.contains("pg_catalog.pg_namespace")); assert!(sql.contains("pg_catalog.pg_description")); assert!(sql.contains("pg_catalog.pg_proc")); assert!(sql.contains("pg_catalog.pg_collation")); assert!(sql.contains("pg_catalog.pg_namespace ncl")); assert!(sql.contains("d.classoid = 'pg_catalog.pg_type'::regclass")); assert!(sql.contains("n.nspname = $1 AND t.typname = $2")); assert!(!sql.contains("pg_get_expr")); } #[test] fn custom_type_domain_default_render_failure_is_degradable() { let (default, warning) = domain_default_from_render_result(Err( "function pg_get_expr(pg_node_tree, integer) does not exist".to_string(), )); assert!(default.is_none()); assert!(warning.as_deref().is_some_and(|value| value.contains("DDL is incomplete"))); } #[test] fn list_objects_sql_includes_routines() { let sql = list_objects_sql(true, true, false, true, true, true, false); assert!(sql.contains("pg_catalog.pg_class")); assert!(sql.contains("pg_catalog.pg_proc")); assert!(sql.contains("pg_catalog.pg_inherits")); assert!(sql.contains("parent_schema")); assert!(sql.contains("parent_name")); assert!(sql.contains("NULL::text AS signature")); assert!(sql.contains("pg_get_function_identity_arguments(p.oid) AS signature")); assert!(sql.contains("pc.relkind = 'p'")); assert!(sql.contains("pg_stat_file")); assert!(sql.contains("pg_xact_commit_timestamp")); assert!(sql.contains("'PROCEDURE'")); assert!(sql.contains("'FUNCTION'")); assert!(!sql.contains("pg_catalog.pg_type")); } #[test] fn list_objects_sql_includes_custom_types_when_enabled() { let sql = list_objects_sql(true, true, false, true, true, true, true); assert!(sql.contains("pg_catalog.pg_type")); assert!(sql.contains("pg_catalog.pg_class")); assert!(sql.contains("pg_catalog.pg_description")); assert!(sql.contains("t.typtype IN ('b', 'c', 'd', 'e', 'r', 'm')")); assert!(sql.contains("t.typisdefined")); assert!(sql.contains("t.typelem = 0")); assert!(sql.contains("t.typrelid = 0 OR c.relkind = 'c'")); assert!(sql.contains("FROM pg_catalog.pg_attribute a")); assert!(sql.contains("FROM pg_catalog.pg_enum e")); assert!(sql.contains("END) AS signature")); assert!(sql.contains("'TYPE' AS object_type")); assert!(sql.contains("5 AS sort_order")); assert!(sql.contains("n.nspname = $1")); assert!(sql.contains("n.nspname <> 'pg_catalog'")); assert!(sql.contains("n.nspname <> 'information_schema'")); assert!(sql.contains("n.nspname NOT LIKE 'pg_toast%'")); assert!(sql.contains("n.nspname NOT LIKE 'pg_temp%'")); } #[test] fn postgres_system_schemas_are_not_custom_type_schemas() { for schema in ["pg_catalog", "information_schema", "pg_toast", "pg_toast_temp_5", "pg_temp_5"] { assert!(is_postgres_system_schema(schema), "{schema} should be a system schema"); } assert!(!is_postgres_system_schema("public")); assert!(!is_postgres_system_schema("app")); } #[test] fn list_objects_sql_omits_custom_types_when_disabled() { let sql = list_objects_sql(true, true, false, true, true, true, false); assert!(!sql.contains("pg_type")); assert!(!sql.contains("t.typtype")); } #[test] fn list_objects_sql_type_only_skips_relations_and_routines() { let sql = list_objects_sql(true, true, false, true, false, false, true); assert!(sql.contains("pg_catalog.pg_type")); assert!(!sql.contains("FROM pg_catalog.pg_class")); assert!(!sql.contains("pg_catalog.pg_proc")); assert!(!sql.contains("pg_stat_file")); assert!(!sql.contains("pg_get_function_identity_arguments")); } #[test] fn list_objects_sql_table_only_skips_routines_and_custom_types() { let sql = list_objects_sql(true, true, false, true, true, false, false); assert!(sql.contains("pg_catalog.pg_class")); assert!(!sql.contains("pg_catalog.pg_proc")); assert!(!sql.contains("pg_catalog.pg_type")); } #[test] fn list_objects_sql_routine_only_skips_relations_and_custom_types() { let sql = list_objects_sql(true, true, false, true, false, true, false); assert!(sql.contains("pg_catalog.pg_proc")); assert!(!sql.contains("pg_catalog.pg_class")); assert!(!sql.contains("pg_catalog.pg_type")); } #[test] fn list_objects_sql_empty_scope_produces_no_sql() { let sql = list_objects_sql(true, true, false, true, false, false, false); assert!(sql.is_empty()); } #[test] fn list_objects_sql_custom_types_branch_in_both_timestamp_variants() { // The pg_type branch carries the same filters whether the timestamp // query or the timestamp fallback query is used. for sql in [ list_objects_sql(true, true, false, true, false, false, true), list_objects_sql(false, true, false, true, false, false, true), ] { assert!(sql.contains("pg_catalog.pg_type")); assert!(sql.contains("t.typtype IN ('b', 'c', 'd', 'e', 'r', 'm')")); assert!(sql.contains("t.typelem = 0")); assert!(sql.contains("t.typrelid = 0 OR c.relkind = 'c'")); assert!(sql.contains("'TYPE' AS object_type")); } } #[test] fn list_objects_sql_custom_types_keeps_comment_join_semantics() { let sql = list_objects_sql(true, true, false, true, false, false, true); assert!(sql.contains("d.objoid = t.oid")); assert!(sql.contains("d.classoid = 'pg_catalog.pg_type'::regclass")); assert!(sql.contains("d.objsubid = 0")); assert!(sql.contains("d.description AS object_comment")); } #[test] fn list_objects_sql_custom_types_uses_null_timestamps() { let sql = list_objects_sql(true, true, false, true, false, false, true); assert!(sql.contains("NULL::text AS created_at")); assert!(sql.contains("NULL::text AS updated_at")); assert!(sql.contains("NULL::text AS parent_schema")); assert!(sql.contains("NULL::text AS parent_name")); assert!(sql.contains("t.typtype::text || ':'")); assert!(sql.contains("END) AS signature")); } #[test] fn list_objects_sql_without_timestamps_omits_stat_file() { let sql = list_objects_sql(false, true, false, true, true, true, false); assert!(!sql.contains("pg_stat_file")); assert!(sql.contains("NULL::text AS created_at")); assert!(sql.contains("NULL::text AS updated_at")); } #[test] fn redshift_compatible_list_objects_sql_uses_legacy_argument_formatter() { let sql = list_objects_sql(false, false, false, false, true, true, false); assert!(sql.contains("pg_get_function_arguments(p.oid) AS signature")); assert!(!sql.contains("pg_get_function_identity_arguments")); } #[test] fn redshift_routine_objects_sql_uses_supported_catalog_view() { let sql = redshift_routine_objects_sql(); assert!(sql.contains("FROM svv_redshift_functions")); assert!(sql.contains("database_name = current_database()")); assert!(sql.contains("schema_name = $1")); assert!(sql.contains("'STORED PROCEDURE'")); assert!(sql.contains("'REGULAR FUNCTION'")); assert!(!sql.contains("pg_proc")); assert!(!sql.contains("pg_get_function")); assert!(!sql.contains("UNION")); } #[test] fn redshift_relation_objects_sql_avoids_unsupported_postgres_helpers() { let sql = list_objects_sql(false, false, false, false, true, false, false); assert!(sql.contains("pg_catalog.pg_class")); assert!(!sql.contains("pg_catalog.pg_proc")); assert!(!sql.contains("pg_stat_file")); assert!(!sql.contains("pg_relation_filepath")); assert!(!sql.contains("pg_xact_commit_timestamp")); } #[test] fn redshift_columns_sql_uses_simple_information_schema_metadata() { let sql = redshift_columns_sql("tenant's", "orders"); assert!(sql.contains("FROM information_schema.columns c")); assert!(sql.contains("c.table_schema = 'tenant''s'")); assert!(sql.contains("c.table_name = 'orders'")); assert!(!sql.contains("pg_attribute")); assert!(!sql.contains("pg_index")); assert!(!sql.contains('$')); } #[test] fn redshift_columns_from_text_result_preserves_basic_metadata() { let result = QueryResult { columns: vec![ "column_name".to_string(), "full_type".to_string(), "is_nullable".to_string(), "column_default".to_string(), "numeric_precision".to_string(), "numeric_scale".to_string(), "character_maximum_length".to_string(), ], column_types: Vec::new(), column_sortables: Vec::new(), spatial_columns: Vec::new(), spatial_values: Vec::new(), rows: vec![vec![ serde_json::json!("amount"), serde_json::json!("numeric"), serde_json::json!("NO"), serde_json::Value::Null, serde_json::json!(18), serde_json::json!(2), serde_json::Value::Null, ]], affected_rows: 0, execution_time_ms: 0, truncated: false, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: Vec::new(), }; let columns = redshift_columns_from_query_result(result); assert_eq!(columns.len(), 1); assert_eq!(columns[0].name, "amount"); assert_eq!(columns[0].data_type, "numeric"); assert!(!columns[0].is_nullable); assert_eq!(columns[0].numeric_precision, Some(18)); assert_eq!(columns[0].numeric_scale, Some(2)); assert_eq!(columns[0].character_maximum_length, None); assert!(!columns[0].is_primary_key); } #[test] fn function_identity_arguments_probe_uses_pg_proc() { let sql = postgres_has_function_identity_arguments_sql(); assert!(sql.contains("pg_catalog.pg_proc")); assert!(sql.contains("n.nspname = 'pg_catalog'")); assert!(sql.contains("p.proname = 'pg_get_function_identity_arguments'")); } #[test] fn both_list_objects_sql_variants_use_parameter() { assert!(list_objects_sql(true, true, true, true, true, true, false).contains("$1")); assert!(list_objects_sql(false, true, true, true, true, true, false).contains("$1")); assert!(list_objects_sql(true, true, false, true, true, true, false).contains("$1")); assert!(list_objects_sql(false, true, false, true, true, true, false).contains("$1")); assert!(list_objects_sql(true, false, true, true, true, true, false).contains("$1")); assert!(list_objects_sql(false, false, true, true, true, true, false).contains("$1")); assert!(list_objects_sql(true, false, false, true, true, true, false).contains("$1")); assert!(list_objects_sql(false, false, false, true, true, true, false).contains("$1")); } #[test] fn both_list_objects_sql_variants_include_pg_proc() { assert!(list_objects_sql(true, true, true, true, true, true, false).contains("pg_catalog.pg_proc")); assert!(list_objects_sql(false, true, true, true, true, true, false).contains("pg_catalog.pg_proc")); assert!(list_objects_sql(true, true, false, true, true, true, false).contains("pg_catalog.pg_proc")); assert!(list_objects_sql(false, true, false, true, true, true, false).contains("pg_catalog.pg_proc")); assert!(list_objects_sql(true, false, true, true, true, true, false).contains("pg_catalog.pg_proc")); assert!(list_objects_sql(false, false, true, true, true, true, false).contains("pg_catalog.pg_proc")); assert!(list_objects_sql(true, false, false, true, true, true, false).contains("pg_catalog.pg_proc")); assert!(list_objects_sql(false, false, false, true, true, true, false).contains("pg_catalog.pg_proc")); } #[test] fn legacy_list_objects_sql_avoids_pg11_proc_kind_column() { let sql = list_objects_sql(true, false, false, true, true, true, false); assert!(!sql.contains("p.prokind")); assert!(!sql.contains("p.prosp")); assert!(sql.contains("NOT p.proisagg")); assert!(sql.contains("NOT p.proiswindow")); assert!(sql.contains("pg_get_function_identity_arguments(p.oid) AS signature")); assert!(sql.contains("'FUNCTION' AS object_type")); assert!(!sql.contains("'PROCEDURE'")); } #[test] fn gaussdb_compatible_list_objects_sql_uses_prosp_when_prokind_is_missing() { let sql = list_objects_sql(true, false, true, true, true, true, false); assert!(!sql.contains("p.prokind")); assert!(sql.contains("CASE WHEN p.prosp THEN 'PROCEDURE' ELSE 'FUNCTION' END AS object_type")); assert!(sql.contains("CASE WHEN p.prosp THEN 2 ELSE 3 END AS sort_order")); assert!(sql.contains("NOT p.proisagg")); assert!(sql.contains("NOT p.proiswindow")); assert!(sql.contains("pg_get_function_identity_arguments(p.oid) AS signature")); } #[test] fn gaussdb_compatible_list_objects_sql_uses_prosp_with_prokind_when_available() { let sql = list_objects_sql(true, true, true, true, true, true, false); assert!( sql.contains("CASE WHEN p.prokind = 'p' OR p.prosp THEN 'PROCEDURE' ELSE 'FUNCTION' END AS object_type") ); assert!(sql.contains("CASE WHEN p.prokind = 'p' OR p.prosp THEN 2 ELSE 3 END AS sort_order")); assert!(sql.contains("p.prokind IN ('p','f') OR p.prosp")); assert!(sql.contains("pg_get_function_identity_arguments(p.oid) AS signature")); } #[test] fn postgres_functions_sql_uses_proc_kind_when_available() { let sql = postgres_functions_sql(true); assert!(sql.contains("p.prokind IN ('f', 'p')")); assert!(sql.contains("WHEN 'p' THEN 'PROCEDURE'")); assert!(!sql.contains("p.proisagg")); assert!(!sql.contains("p.proiswindow")); } #[test] fn legacy_postgres_functions_sql_avoids_proc_kind_column() { let sql = postgres_functions_sql(false); assert!(!sql.contains("p.prokind")); assert!(sql.contains("NOT p.proisagg")); assert!(sql.contains("NOT p.proiswindow")); assert!(sql.contains("'FUNCTION'")); assert!(!sql.contains("'PROCEDURE'")); } #[test] fn postgres_proc_has_prokind_sql_checks_catalog_attribute() { let sql = postgres_proc_has_prokind_sql(); assert!(sql.contains("pg_catalog.pg_attribute")); assert!(sql.contains("'pg_catalog.pg_proc'::regclass")); assert!(sql.contains("attname = 'prokind'")); } #[test] fn postgres_proc_has_prosp_sql_checks_catalog_attribute() { let sql = postgres_proc_has_prosp_sql(); assert!(sql.contains("pg_catalog.pg_attribute")); assert!(sql.contains("'pg_catalog.pg_proc'::regclass")); assert!(sql.contains("attname = 'prosp'")); } #[test] fn postgres_trigger_definitions_sql_excludes_internal_triggers() { let sql = postgres_trigger_definitions_sql(); assert!(sql.contains("pg_catalog.pg_get_triggerdef(t.oid, true) AS trigger_definition")); assert!(sql.contains("NOT t.tgisinternal")); assert!(sql.contains("t.tgparentid = 0")); assert!(sql.contains("ORDER BY t.tgname, t.oid")); } #[test] fn postgres_trigger_definitions_sql_without_tgparentid_omits_partition_clone_filter() { let sql = postgres_trigger_definitions_sql_without_tgparentid(); assert!(sql.contains("NOT t.tgisinternal")); assert!(!sql.contains("tgparentid")); } #[test] fn postgres_trigger_has_tgparentid_sql_checks_catalog_attribute() { let sql = postgres_trigger_has_tgparentid_sql(); assert!(sql.contains("pg_catalog.pg_trigger")); assert!(sql.contains("attname = 'tgparentid'")); } #[test] fn transaction_recovery_statement_detection_matches_common_postgres_commands() { assert!(is_transaction_recovery_statement("ROLLBACK")); assert!(is_transaction_recovery_statement("rollback work")); assert!(is_transaction_recovery_statement("ABORT TRANSACTION")); assert!(is_transaction_recovery_statement("commit")); assert!(is_transaction_recovery_statement("END")); } #[test] fn transaction_recovery_statement_detection_ignores_regular_queries() { assert!(!is_transaction_recovery_statement("SELECT 1")); assert!(!is_transaction_recovery_statement("BEGIN")); assert!(!is_transaction_recovery_statement("UPDATE users SET name = 'dbx'")); } #[test] fn postgres_ddl_detection_covers_schema_changing_statements() { assert!(invalidates_postgres_statement_cache("ALTER TABLE users ADD COLUMN email text")); assert!(invalidates_postgres_statement_cache(" CREATE INDEX idx_users_email ON users(email)")); assert!(invalidates_postgres_statement_cache("COMMENT ON COLUMN users.email IS 'Email'")); assert!(invalidates_postgres_statement_cache("DROP TABLE users")); assert!(invalidates_postgres_statement_cache("TRUNCATE users")); assert!(invalidates_postgres_statement_cache("REINDEX TABLE users")); assert!(invalidates_postgres_statement_cache("VACUUM users")); } #[test] fn postgres_ddl_detection_ignores_regular_dml_and_selects() { assert!(!invalidates_postgres_statement_cache("SELECT * FROM users")); assert!(!invalidates_postgres_statement_cache("UPDATE users SET name = 'Ada'")); assert!(!invalidates_postgres_statement_cache("INSERT INTO users(name) VALUES ('Ada')")); assert!(!invalidates_postgres_statement_cache("DELETE FROM users WHERE id = 1")); } #[test] fn postgres_stale_cache_retry_uses_structured_fields_for_localized_errors() { assert!(should_retry_postgres_stale_cache_fields( Some("0A000"), Some("RevalidateCachedQuery"), "已缓冲的计划不能改变结果类型", )); assert!(should_retry_postgres_stale_cache_fields(None, None, "cached plan must not change result type",)); } #[test] fn postgres_stale_cache_retry_rejects_other_feature_errors() { assert!(!should_retry_postgres_stale_cache_fields(Some("0A000"), Some("CheckFeatureSupport"), "不支持该功能",)); assert!(!should_retry_postgres_stale_cache_fields( Some("23505"), Some("RevalidateCachedQuery"), "duplicate key value violates unique constraint", )); } #[test] fn postgres_missing_prepared_statement_fallback_uses_sqlstate_and_message() { assert!(should_fallback_postgres_missing_prepared_statement_fields(Some("26000"), "本地化的预编译语句错误",)); assert!(should_fallback_postgres_missing_prepared_statement_fields( None, "prepared statement \"s63\" does not exist", )); } #[test] fn postgres_missing_prepared_statement_fallback_rejects_unrelated_errors() { assert!(!should_fallback_postgres_missing_prepared_statement_fields( Some("23505"), "duplicate key value violates unique constraint", )); assert!(!should_fallback_postgres_missing_prepared_statement_fields( None, "prepared statement result type changed", )); } #[test] fn postgres_typed_params_pairs_values_with_inferred_types() { let id = 42_i32; let name = "dbx"; let params: &[&(dyn tokio_postgres::types::ToSql + Sync)] = &[&id, &name]; let typed = postgres_typed_params(params, &[Type::INT4, Type::TEXT]).expect("matching parameter count"); assert_eq!(typed.iter().map(|(_, ty)| ty.clone()).collect::>(), vec![Type::INT4, Type::TEXT]); assert!(postgres_typed_params(params, &[Type::INT4]).is_none()); } // --- execute_batch --- #[tokio::test] async fn execute_batch_empty_statements_returns_ok() { // Empty input should not error or try to connect // We can't test with a real pool, but we can verify the empty-early-return logic // by testing that an empty Vec doesn't need a pool reference let statements: Vec = vec![]; // This test validates the early return logic at code review level // Actual execution requires a pool; we just verify the empty path exists assert!(statements.is_empty()); } #[tokio::test] async fn execute_batch_whitespace_only_is_filtered() { let statements = [" ".to_string(), "\t\n".to_string(), "".to_string()]; let combined = statements.iter().map(|s| s.trim()).filter(|s| !s.is_empty()).collect::>().join(";\n"); assert!(combined.is_empty()); } #[test] fn execute_batch_joins_with_semicolons() { let statements = ["SELECT 1".to_string(), "SELECT 2".to_string()]; let combined = statements.iter().map(|s| s.trim()).filter(|s| !s.is_empty()).collect::>().join(";\n"); assert_eq!(combined, "SELECT 1;\nSELECT 2"); } // --- SET timezone escaping --- #[test] fn timezone_single_quotes_are_doubled() { let tz = "UTC"; let escaped = tz.replace('\'', "''"); assert_eq!(escaped, "UTC"); } #[test] fn timezone_with_quote_is_escaped() { let tz = "Some'Zone"; let escaped = tz.replace('\'', "''"); assert_eq!(escaped, "Some''Zone"); } // --- pg_url_has_timezone_setting --- #[test] fn url_without_timezone_returns_false() { assert!(!pg_url_has_timezone_setting("postgres://localhost/db")); assert!(!pg_url_has_timezone_setting("postgres://localhost/db?sslmode=require")); } #[test] fn url_with_options_timezone_returns_true() { assert!(pg_url_has_timezone_setting("postgres://localhost/db?options=-c timezone=Asia/Shanghai")); } #[test] fn url_with_url_encoded_timezone_returns_true() { assert!(pg_url_has_timezone_setting("postgres://localhost/db?options=-c%20timezone%3DUTC")); } #[test] fn url_with_uppercase_timezone_returns_true() { assert!(pg_url_has_timezone_setting("postgres://localhost/db?options=--TimeZone=UTC")); } #[test] fn unrelated_timezone_text_is_not_treated_as_explicit() { assert!(!pg_url_has_timezone_setting("postgres://localhost/db?timezone=UTC")); assert!(!pg_url_has_timezone_setting( "postgres://localhost/db?application_name=timezone%3DUTC&options=-c%20search_path%3Dpublic" )); } #[test] fn postgres_timezone_candidates_include_known_tzdata_aliases() { assert_eq!(postgres_timezone_candidates("Europe/Kyiv"), vec!["Europe/Kyiv", "Europe/Kiev"]); assert_eq!(postgres_timezone_candidates("Asia/Kolkata"), vec!["Asia/Kolkata", "Asia/Calcutta"]); assert_eq!(postgres_timezone_candidates("America/New_York"), vec!["America/New_York"]); } #[tokio::test] #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a PostgreSQL database"] async fn automatic_invalid_timezone_keeps_connected_server_default() { let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); let pool = connect_with_local_timezone(&url, Duration::from_secs(5), "Invalid/DBX_Timezone") .await .expect("automatic local timezone rejection must not reject a valid connection"); let client = pool.get().await.expect("checkout postgres"); let timezone: String = client.query_one("SHOW timezone", &[]).await.unwrap().get(0); assert_ne!(timezone, "Invalid/DBX_Timezone"); } #[tokio::test] #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a PostgreSQL database"] async fn cached_queries_downgrade_connection_after_server_deallocates_statements() { let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); let pool = connect_with_local_timezone(&url, Duration::from_secs(5), "UTC") .await .expect("connect PostgreSQL database"); let client = pool.get().await.expect("checkout postgres"); let sql = "SELECT $1::int4 AS value"; let rows = postgres_query_cached(&client, sql, &[&41_i32]).await.expect("prime statement cache"); assert_eq!(rows[0].get::<_, i32>(0), 41); client.batch_execute("DEALLOCATE ALL").await.expect("drop server-side prepared statements"); let rows = postgres_query_cached(&client, sql, &[&42_i32]).await.expect("retry through unnamed typed query"); assert_eq!(rows[0].get::<_, i32>(0), 42); assert!(postgres_client_uses_unnamed_statements(&client)); client.execute_typed("DEALLOCATE ALL", &[]).await.expect("clear any remaining named statements"); let rows = postgres_query_cached(&client, sql, &[&43_i32]).await.expect("reuse inferred types without named bind"); assert_eq!(rows[0].get::<_, i32>(0), 43); let prepared_count = client .query_typed_one("SELECT count(*)::int8 FROM pg_prepared_statements", &[]) .await .expect("count server-side prepared statements") .get::<_, i64>(0); assert_eq!(prepared_count, 0, "downgraded connection must not bind or re-prepare the cached query"); } #[tokio::test] #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a writable PostgreSQL database"] async fn missing_statement_fallback_uses_current_metadata_for_empty_results() { let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); let pool = connect_with_local_timezone(&url, Duration::from_secs(5), "UTC") .await .expect("connect PostgreSQL database"); let client = pool.get().await.expect("checkout postgres"); let select_sql = "SELECT * FROM dbx_statement_metadata_test WHERE false"; client .execute_typed("CREATE TEMP TABLE dbx_statement_metadata_test (id int4)", &[]) .await .expect("create temporary test table"); let initial = execute_select_query(&client, select_sql, Instant::now(), 10).await.expect("prime statement cache"); assert_eq!(initial.columns, vec!["id"]); assert!(initial.rows.is_empty()); client .execute_typed("ALTER TABLE dbx_statement_metadata_test ADD COLUMN label text", &[]) .await .expect("change result metadata"); client.execute_typed("DEALLOCATE ALL", &[]).await.expect("drop server-side prepared statements"); let current = execute_select_query(&client, select_sql, Instant::now(), 10) .await .expect("fallback with current row description"); assert_eq!(current.columns, vec!["id", "label"]); assert_eq!(current.column_types, vec!["int4", "text"]); assert!(current.rows.is_empty()); } #[tokio::test] #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a PostgreSQL database"] async fn explicit_transaction_queries_start_unnamed_and_remain_usable() { let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); let pool = connect_with_local_timezone(&url, Duration::from_secs(5), "UTC") .await .expect("connect PostgreSQL database"); let client = pool.get().await.expect("checkout postgres"); client.execute_typed("BEGIN", &[]).await.expect("begin transaction"); client.execute_typed("DEALLOCATE ALL", &[]).await.expect("simulate statement loss"); let result = execute_select_query_unnamed(&client, "SELECT 44::int4 AS value", Instant::now(), 10) .await .expect("execute unnamed query inside transaction"); assert_eq!(result.rows[0][0], serde_json::json!(44)); client.execute_typed("SELECT 1", &[]).await.expect("transaction remains usable"); client.execute_typed("ROLLBACK", &[]).await.expect("rollback transaction"); } #[tokio::test] #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a PostgreSQL database"] async fn backup_snapshot_stream_starts_unnamed_and_remains_usable() { let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); let pool = connect_with_local_timezone(&url, Duration::from_secs(5), "UTC") .await .expect("connect PostgreSQL database"); let client = pool.get().await.expect("checkout postgres"); client .execute_typed("BEGIN TRANSACTION ISOLATION LEVEL REPEATABLE READ READ ONLY", &[]) .await .expect("begin backup snapshot"); client.execute_typed("DEALLOCATE ALL", &[]).await.expect("simulate statement loss"); let mut columns = Vec::new(); let mut rows = Vec::new(); stream_select_query_inner_unnamed(&client, "SELECT 45::int4 AS value", None, &mut |item| { match item { PostgresQueryStreamItem::Columns { columns: names, .. } => columns = names, PostgresQueryStreamItem::Row(row) => rows.push(row), } Ok(()) }) .await .expect("stream unnamed query inside backup snapshot"); assert_eq!(columns, vec!["value"]); assert_eq!(rows[0][0], serde_json::json!(45)); let prepared_count = client .query_typed_one("SELECT count(*)::int8 FROM pg_prepared_statements", &[]) .await .expect("count server-side prepared statements") .get::<_, i64>(0); assert_eq!(prepared_count, 0, "snapshot stream metadata preparation must not retain a named statement"); client.execute_typed("SELECT 1", &[]).await.expect("snapshot remains usable"); client.execute_typed("ROLLBACK", &[]).await.expect("rollback backup snapshot"); } #[tokio::test] #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a PostgreSQL database"] async fn explicit_timezone_remains_strict_and_overrides_local_timezone() { let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); let separator = if url.contains('?') { '&' } else { '?' }; let explicit_url = format!("{url}{separator}options=-c%20TimeZone%3DAsia%2FShanghai"); let pool = connect_with_local_timezone(&explicit_url, Duration::from_secs(5), "UTC") .await .expect("valid explicit timezone"); let client = pool.get().await.expect("checkout postgres"); let timezone: String = client.query_one("SHOW timezone", &[]).await.unwrap().get(0); assert_eq!(timezone, "Asia/Shanghai"); let invalid_url = format!("{url}{separator}options=-c%20TimeZone%3DInvalid%2FDBX_Timezone"); let error = connect_with_local_timezone(&invalid_url, Duration::from_secs(5), "UTC") .await .expect_err("invalid explicit timezone must remain a connection error"); assert!(error.contains("Invalid/DBX_Timezone") || error.contains("time zone"), "{error}"); } #[tokio::test] #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a PostgreSQL database"] async fn valid_automatic_timezone_is_applied_normally() { let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); let pool = connect_with_local_timezone(&url, Duration::from_secs(5), "UTC").await.expect("valid automatic timezone"); let client = pool.get().await.expect("checkout postgres"); let timezone: String = client.query_one("SHOW timezone", &[]).await.unwrap().get(0); assert_eq!(timezone, "UTC"); } #[tokio::test] #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a writable PostgreSQL-compatible database"] async fn list_tables_filtered_supports_risingwave_pagination_and_filtering() { let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); let pool = connect_with_local_timezone(&url, Duration::from_secs(5), "UTC") .await .expect("connect PostgreSQL-compatible database"); let client = pool.get().await.expect("checkout postgres"); client .batch_execute( r#"DROP TABLE IF EXISTS public."dbx_issue_5584_live_page_a"; DROP TABLE IF EXISTS public."dbx_issue_5584_live_page_b"; DROP TABLE IF EXISTS public."dbx_issue_5584_live_order_100%"; DROP TABLE IF EXISTS public."dbx_issue_5584_live_back\slash"; DROP TABLE IF EXISTS public."dbx_issue_5584_live_system_users"; CREATE TABLE public."dbx_issue_5584_live_page_a"(id int); CREATE TABLE public."dbx_issue_5584_live_page_b"(id int); CREATE TABLE public."dbx_issue_5584_live_order_100%"(id int); CREATE TABLE public."dbx_issue_5584_live_back\slash"(id int); CREATE TABLE public."dbx_issue_5584_live_system_users"(id int);"#, ) .await .expect("create live table fixtures"); drop(client); let all_tables = list_tables(&pool, "public").await.expect("expand complete table list"); let first_page = list_tables_filtered(&pool, "public", Some("dbx_issue_5584_live_page_"), Some(1), Some(0)) .await .expect("list first table page"); let second_page = list_tables_filtered(&pool, "public", Some("dbx_issue_5584_live_page_"), Some(1), Some(1)) .await .expect("list second table page"); let wildcard_match = list_tables_filtered(&pool, "public", Some("dbx_issue_5584_live_order_100%"), Some(10), Some(0)) .await .expect("filter table with wildcard characters"); let backslash_match = list_tables_filtered(&pool, "public", Some(r"dbx_issue_5584_live_back\slash"), Some(10), Some(0)) .await .expect("filter table with backslash"); let fuzzy_match = list_tables_filtered(&pool, "public", Some("i5584su"), Some(10), Some(0)) .await .expect("fuzzy filter table name"); assert!(all_tables.iter().any(|table| table.name == "dbx_issue_5584_live_page_a")); assert!(all_tables.iter().any(|table| table.name == "dbx_issue_5584_live_page_b")); assert_eq!( first_page.iter().map(|table| table.name.as_str()).collect::>(), ["dbx_issue_5584_live_page_a"] ); assert_eq!( second_page.iter().map(|table| table.name.as_str()).collect::>(), ["dbx_issue_5584_live_page_b"] ); assert_eq!( wildcard_match.iter().map(|table| table.name.as_str()).collect::>(), ["dbx_issue_5584_live_order_100%"] ); assert_eq!( backslash_match.iter().map(|table| table.name.as_str()).collect::>(), [r"dbx_issue_5584_live_back\slash"] ); assert_eq!( fuzzy_match.iter().map(|table| table.name.as_str()).collect::>(), ["dbx_issue_5584_live_system_users"] ); let client = pool.get().await.expect("checkout postgres for cleanup"); client .batch_execute( r#"DROP TABLE IF EXISTS public."dbx_issue_5584_live_page_a"; DROP TABLE IF EXISTS public."dbx_issue_5584_live_page_b"; DROP TABLE IF EXISTS public."dbx_issue_5584_live_order_100%"; DROP TABLE IF EXISTS public."dbx_issue_5584_live_back\slash"; DROP TABLE IF EXISTS public."dbx_issue_5584_live_system_users";"#, ) .await .expect("drop live table fixtures"); } #[tokio::test] #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a writable PostgreSQL database"] async fn completion_assistant_searches_visible_and_qualified_sequences() { let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); let pool = connect_with_local_timezone(&url, Duration::from_secs(5), "UTC") .await .expect("connect PostgreSQL database"); let client = pool.get().await.expect("checkout postgres"); client .batch_execute( r#"DROP SCHEMA IF EXISTS dbx_issue_4005_visible CASCADE; DROP SCHEMA IF EXISTS dbx_issue_4005_shadowed CASCADE; CREATE SCHEMA dbx_issue_4005_visible; CREATE SCHEMA dbx_issue_4005_shadowed; CREATE SEQUENCE dbx_issue_4005_visible.order_seq; CREATE SEQUENCE dbx_issue_4005_shadowed.order_seq; CREATE SEQUENCE dbx_issue_4005_visible."OrderSequence"; SET search_path = dbx_issue_4005_visible, dbx_issue_4005_shadowed, public;"#, ) .await .expect("create live sequence fixtures"); drop(client); let visible = completion_assistant_search( &pool, &CompletionAssistantRequest { connection_id: "live".into(), database: "postgres".into(), schema: None, object_kinds: vec![CompletionAssistantObjectKind::Sequence], mask: "order_".into(), case_sensitive: false, global_search: false, max_results: Some(20), search_in_comments: false, search_in_definitions: false, parent_schema: None, parent_name: None, match_mode: Some(CompletionAssistantMatchMode::Prefix), }, ) .await .expect("complete visible sequence"); assert_eq!( visible .candidates .iter() .map(|candidate| (candidate.name.as_str(), candidate.schema.as_deref())) .collect::>(), [("order_seq", Some("dbx_issue_4005_visible"))] ); let qualified = completion_assistant_search( &pool, &CompletionAssistantRequest { connection_id: "live".into(), database: "postgres".into(), schema: Some("dbx_issue_4005_visible".into()), object_kinds: vec![CompletionAssistantObjectKind::Sequence], mask: "OrderS".into(), case_sensitive: true, global_search: false, max_results: Some(20), search_in_comments: false, search_in_definitions: false, parent_schema: None, parent_name: None, match_mode: Some(CompletionAssistantMatchMode::Prefix), }, ) .await .expect("complete qualified mixed-case sequence"); assert_eq!(qualified.candidates[0].name, "OrderSequence"); let client = pool.get().await.expect("checkout postgres for cleanup"); client .batch_execute( r#"DROP SCHEMA IF EXISTS dbx_issue_4005_visible CASCADE; DROP SCHEMA IF EXISTS dbx_issue_4005_shadowed CASCADE;"#, ) .await .expect("drop live sequence fixtures"); } #[test] fn like_contains_pattern_escapes_wildcards() { assert_eq!(like_contains_pattern(""), "%%"); assert_eq!(like_contains_pattern("order_100%"), "%order\\_100\\%%"); assert_eq!(like_contains_pattern("tilde~name"), "%tilde~name%"); assert_eq!(like_contains_pattern(r"foo\bar"), r"%foo\\bar%"); } #[test] fn like_fuzzy_pattern_escapes_wildcards() { assert_eq!(like_fuzzy_pattern(""), "%%"); assert_eq!(like_fuzzy_pattern("sysu"), "%s%y%s%u%"); assert_eq!(like_fuzzy_pattern("user_%"), "%u%s%e%r%\\_%\\%%"); assert_eq!(like_fuzzy_pattern("tilde~name"), "%t%i%l%d%e%~%n%a%m%e%"); } #[test] fn postgres_tables_sql_uses_literal_pagination_without_escape_clause() { let paged_sql = postgres_tables_sql(Some(500), 200); assert!(paged_sql.contains("ILIKE $2 OR")); assert!(paged_sql.contains("$3 <> ''")); assert!(paged_sql.contains("ILIKE $3")); assert!(!paged_sql.contains("ESCAPE")); assert!(paged_sql.contains("ORDER BY CASE WHEN pc.relkind = 'p' THEN 1 ELSE 0 END, c.relname")); assert!(paged_sql.contains("LIMIT 500 OFFSET 200")); assert!(!paged_sql.contains("$4")); assert!(!paged_sql.contains("$5")); let unbounded_sql = postgres_tables_sql(None, 0); assert!(unbounded_sql.ends_with("OFFSET 0")); assert!(!unbounded_sql.contains("LIMIT")); } #[test] fn postgres_completion_like_pattern_uses_prefix_by_default() { assert_eq!(postgres_completion_like_pattern("Temp", Some(&CompletionAssistantMatchMode::Prefix)), "Temp%"); assert_eq!(postgres_completion_like_pattern("Temp", Some(&CompletionAssistantMatchMode::Contains)), "%Temp%"); assert_eq!( postgres_completion_like_pattern("order_100%", Some(&CompletionAssistantMatchMode::Prefix)), "order~_100~%%" ); } #[test] fn postgres_completion_sql_filters_before_limit() { assert!(postgres_completion_tables_sql().contains("c.relname ILIKE $2 ESCAPE '~'")); assert!(postgres_completion_tables_sql().contains("pg_catalog.pg_table_is_visible(c.oid)")); assert!(postgres_completion_tables_sql().contains("c.relkind::text = ANY($3::text[])")); assert!(postgres_completion_tables_sql().contains("ORDER BY c.relname LIMIT $4")); assert!(postgres_completion_routines_sql(true).contains("p.proname ILIKE $2 ESCAPE '~'")); assert!(postgres_completion_routines_sql(true).contains("p.prokind::text = ANY($3::text[])")); assert!( postgres_completion_routines_sql(true).contains("pg_get_function_identity_arguments(p.oid) AS signature") ); assert!(postgres_completion_routines_sql(true).contains("ORDER BY p.proname LIMIT $4")); assert!(!postgres_completion_routines_sql(false).contains("p.prokind")); assert!(postgres_completion_routines_sql(false).contains("'FUNCTION'::text")); assert!(postgres_completion_routines_sql(false).contains("ORDER BY p.proname LIMIT $3")); assert!(postgres_completion_sequences_sql().contains("c.relkind = 'S'")); assert!(postgres_completion_sequences_sql().contains("has_schema_privilege(n.oid, 'USAGE')")); assert!(postgres_completion_sequences_sql().contains("pg_catalog.pg_table_is_visible(c.oid)")); assert!(postgres_completion_sequences_sql().contains("CASE WHEN c.relkind = 'S'")); assert!(postgres_completion_sequences_sql() .contains("THEN pg_catalog.has_sequence_privilege(c.oid, 'USAGE, SELECT, UPDATE')")); assert!(postgres_completion_sequences_sql().contains("CASE WHEN $3 THEN c.relname LIKE $2 ESCAPE '~'")); assert!(postgres_completion_sequences_sql().contains("ELSE c.relname ILIKE $2 ESCAPE '~' END")); assert!(postgres_completion_sequences_sql().contains("ORDER BY c.relname LIMIT $4")); assert!(postgres_completion_columns_sql().contains("a.attname ILIKE $3 ESCAPE '~'")); assert!(postgres_visible_table_schema_sql().contains("pg_catalog.pg_table_is_visible(c.oid)")); } #[test] fn postgres_schema_info_sql_only_filters_system_schemas_when_disabled() { let hidden_sql = postgres_schema_infos_sql(false); assert!(hidden_sql.contains("information_schema")); assert!(hidden_sql.contains("pg_temp_%")); let visible_sql = postgres_schema_infos_sql(true); assert!(!visible_sql.contains("NOT IN")); assert!(!visible_sql.contains("NOT LIKE")); } }