itsscb/sqlx
1
0
Fork 0
mirror of https://github.com/launchbadge/sqlx.git synced 2026-03-19 16:44:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(postgres): add missing type resolution for arrays by name

Browse Source
This commit is contained in:
Austin Bonander
2024-07-04 17:17:20 -07:00
parent efbf57265c
commit 16e3f1025a
19 changed files with 333 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
sqlx-postgres/Cargo.toml
View File

@@ -71,5 +71,8 @@ workspace = true
# We use JSON in the driver implementation itself so there's no reason not to enable it here.
features = ["json"]
[dev-dependencies]
sqlx.workspace = true
[target.'cfg(target_os = "windows")'.dependencies]
etcetera = "0.8.0"

30
sqlx-postgres/src/arguments.rs
View File

@@ -1,5 +1,6 @@
use std::fmt::{self, Write};
use std::ops::{Deref, DerefMut};
use std::sync::Arc;
use crate::encode::{Encode, IsNull};
use crate::error::Error;
@@ -7,6 +8,7 @@ use crate::ext::ustr::UStr;
use crate::types::Type;
use crate::{PgConnection, PgTypeInfo, Postgres};
use crate::type_info::PgArrayOf;
pub(crate) use sqlx_core::arguments::Arguments;
use sqlx_core::error::BoxDynError;
@@ -41,7 +43,12 @@ pub struct PgArgumentBuffer {
// This is done for Records and Arrays as the OID is needed well before we are in an async
// function and can just ask postgres.
//
type_holes: Vec<(usize, UStr)>, // Vec<{ offset, type_name }>
type_holes: Vec<(usize, HoleKind)>, // Vec<{ offset, type_name }>
}
enum HoleKind {
Type { name: UStr },
Array(Arc<PgArrayOf>),
}
struct Patch {
@@ -106,8 +113,11 @@ impl PgArguments {
(patch.callback)(buf, ty);
}
for (offset, name) in type_holes {
let oid = conn.fetch_type_id_by_name(name).await?;
for (offset, kind) in type_holes {
let oid = match kind {
HoleKind::Type { name } => conn.fetch_type_id_by_name(name).await?,
HoleKind::Array(array) => conn.fetch_array_type_id(array).await?,
};
buffer[*offset..(*offset + 4)].copy_from_slice(&oid.0.to_be_bytes());
}
@@ -186,7 +196,19 @@ impl PgArgumentBuffer {
let offset = self.len();
self.extend_from_slice(&0_u32.to_be_bytes());
self.type_holes.push((offset, type_name.clone()));
self.type_holes.push((
offset,
HoleKind::Type {
name: type_name.clone(),
},
));
}
pub(crate) fn patch_array_type(&mut self, array: Arc<PgArrayOf>) {
let offset = self.len();
self.extend_from_slice(&0_u32.to_be_bytes());
self.type_holes.push((offset, HoleKind::Array(array)));
}
fn snapshot(&self) -> PgArgumentBufferSnapshot {

45
sqlx-postgres/src/connection/describe.rs
View File

@@ -4,7 +4,7 @@ use crate::message::{ParameterDescription, RowDescription};
use crate::query_as::query_as;
use crate::query_scalar::{query_scalar, query_scalar_with};
use crate::statement::PgStatementMetadata;
use crate::type_info::{PgCustomType, PgType, PgTypeKind};
use crate::type_info::{PgArrayOf, PgCustomType, PgType, PgTypeKind};
use crate::types::Json;
use crate::types::Oid;
use crate::HashMap;
@@ -355,6 +355,19 @@ WHERE rngtypid = $1
})
}
pub(crate) async fn resolve_type_id(&mut self, ty: &PgType) -> Result<Oid, Error> {
if let Some(oid) = ty.try_oid() {
return Ok(oid);
}
match ty {
PgType::DeclareWithName(name) => self.fetch_type_id_by_name(name).await,
PgType::DeclareArrayOf(array) => self.fetch_array_type_id(array).await,
// `.try_oid()` should return `Some()` or it should be covered here
_ => unreachable!("(bug) OID should be resolvable for type {ty:?}"),
}
}
pub(crate) async fn fetch_type_id_by_name(&mut self, name: &str) -> Result<Oid, Error> {
if let Some(oid) = self.cache_type_oid.get(name) {
return Ok(*oid);
@@ -366,13 +379,41 @@ WHERE rngtypid = $1
.fetch_optional(&mut *self)
.await?
.ok_or_else(|| Error::TypeNotFound {
type_name: String::from(name),
type_name: name.into(),
})?;
self.cache_type_oid.insert(name.to_string().into(), oid);
Ok(oid)
}
pub(crate) async fn fetch_array_type_id(&mut self, array: &PgArrayOf) -> Result<Oid, Error> {
if let Some(oid) = self
.cache_type_oid
.get(&array.elem_name)
.and_then(|elem_oid| self.cache_elem_type_to_array.get(elem_oid))
{
return Ok(*oid);
}
// language=SQL
let (elem_oid, array_oid): (Oid, Oid) =
query_as("SELECT oid, typarray FROM pg_catalog.pg_type WHERE oid = $1::regtype::oid")
.bind(&*array.elem_name)
.fetch_optional(&mut *self)
.await?
.ok_or_else(|| Error::TypeNotFound {
type_name: array.name.to_string(),
})?;
// Avoids copying `elem_name` until necessary
self.cache_type_oid
.entry_ref(&array.elem_name)
.insert(elem_oid);
self.cache_elem_type_to_array.insert(elem_oid, array_oid);
Ok(array_oid)
}
pub(crate) async fn get_nullable_for_columns(
&mut self,
stmt_id: Oid,

1
sqlx-postgres/src/connection/establish.rs
View File

@@ -146,6 +146,7 @@ impl PgConnection {
cache_statement: StatementCache::new(options.statement_cache_capacity),
cache_type_oid: HashMap::new(),
cache_type_info: HashMap::new(),
cache_elem_type_to_array: HashMap::new(),
log_settings: options.log_settings.clone(),
})
}

7
sqlx-postgres/src/connection/executor.rs
View File

@@ -7,7 +7,6 @@ use crate::message::{
RowDescription,
};
use crate::statement::PgStatementMetadata;
use crate::type_info::PgType;
use crate::types::Oid;
use crate::{
statement::PgStatement, PgArguments, PgConnection, PgQueryResult, PgRow, PgTypeInfo,
@@ -36,11 +35,7 @@ async fn prepare(
let mut param_types = Vec::with_capacity(parameters.len());
for ty in parameters {
param_types.push(if let PgType::DeclareWithName(name) = &ty.0 {
conn.fetch_type_id_by_name(name).await?
} else {
ty.0.oid()
});
param_types.push(conn.resolve_type_id(&ty.0).await?);
}
// flush and wait until we are re-ready

1
sqlx-postgres/src/connection/mod.rs
View File

@@ -55,6 +55,7 @@ pub struct PgConnection {
// cache user-defined types by id <-> info
cache_type_info: HashMap<Oid, PgTypeInfo>,
cache_type_oid: HashMap<UStr, Oid>,
cache_elem_type_to_array: HashMap<Oid, Oid>,
// number of ReadyForQuery messages that we are currently expecting
pub(crate) pending_ready_for_query_count: usize,

182
sqlx-postgres/src/type_info.rs
View File

@@ -11,6 +11,34 @@ use crate::types::Oid;
pub(crate) use sqlx_core::type_info::TypeInfo;
/// Type information for a PostgreSQL type.
///
/// ### Note: Implementation of `==` ([`PartialEq::eq()`])
/// Because `==` on [`TypeInfo`]s has been used throughout the SQLx API as a synonym for type compatibility,
/// e.g. in the default impl of [`Type::compatible()`][sqlx_core::types::Type::compatible],
/// some concessions have been made in the implementation.
///
/// When comparing two `PgTypeInfo`s using the `==` operator ([`PartialEq::eq()`]),
/// if one was constructed with [`Self::with_oid()`] and the other with [`Self::with_name()`] or
/// [`Self::array_of()`], `==` will return `true`:
///
/// ```
/// # use sqlx::postgres::{types::Oid, PgTypeInfo};
/// // Potentially surprising result, this assert will pass:
/// assert_eq!(PgTypeInfo::with_oid(Oid(1)), PgTypeInfo::with_name("definitely_not_real"));
/// ```
///
/// Since it is not possible in this case to prove the types are _not_ compatible (because
/// both `PgTypeInfo`s need to be resolved by an active connection to know for sure)
/// and type compatibility is mainly done as a sanity check anyway,
/// it was deemed acceptable to fudge equality in this very specific case.
///
/// This also applies when querying with the text protocol (not using prepared statements,
/// e.g. [`sqlx::raw_sql()`][sqlx_core::raw_sql::raw_sql]), as the connection will be unable
/// to look up the type info like it normally does when preparing a statement: it won't know
/// what the OIDs of the output columns will be until it's in the middle of reading the result,
/// and by that time it's too late.
///
/// To compare types for exact equality, use [`Self::type_eq()`] instead.
#[derive(Debug, Clone, PartialEq)]
#[cfg_attr(feature = "offline", derive(serde::Serialize, serde::Deserialize))]
pub struct PgTypeInfo(pub(crate) PgType);
@@ -132,6 +160,8 @@ pub enum PgType {
// NOTE: Do we want to bring back type declaration by ID? It's notoriously fragile but
// someone may have a user for it
DeclareWithOid(Oid),
DeclareArrayOf(Arc<PgArrayOf>),
}
#[derive(Debug, Clone)]
@@ -155,6 +185,13 @@ pub enum PgTypeKind {
Range(PgTypeInfo),
}
#[derive(Debug)]
#[cfg_attr(feature = "offline", derive(serde::Serialize, serde::Deserialize))]
pub struct PgArrayOf {
pub(crate) elem_name: UStr,
pub(crate) name: Box<str>,
}
impl PgTypeInfo {
/// Returns the corresponding `PgTypeInfo` if the OID is a built-in type and recognized by SQLx.
pub(crate) fn try_from_oid(oid: Oid) -> Option<Self> {
@@ -233,18 +270,79 @@ impl PgTypeInfo {
///
/// The OID for the type will be fetched from Postgres on use of
/// a value of this type. The fetched OID will be cached per-connection.
///
/// ### Note: Type Names Prefixed with `_`
/// In `pg_catalog.pg_type`, Postgres prefixes a type name with `_` to denote an array of that
/// type, e.g. `int4[]` actually exists in `pg_type` as `_int4`.
///
/// Previously, it was necessary in manual [`PgHasArrayType`][crate::PgHasArrayType] impls
/// to return [`PgTypeInfo::with_name()`] with the type name prefixed with `_` to denote
/// an array type, but this would not work with schema-qualified names.
///
/// As of 0.8, [`PgTypeInfo::array_of()`] is used to declare an array type,
/// and the Postgres driver is now able to properly resolve arrays of custom types,
/// even in other schemas, which was not previously supported.
///
/// It is highly recommended to migrate existing usages to [`PgTypeInfo::array_of()`] where
/// applicable.
///
/// However, to maintain compatibility, the driver now infers any type name prefixed with `_`
/// to be an array of that type. This may introduce some breakages for types which use
/// a `_` prefix but which are not arrays.
///
/// As a workaround, type names with `_` as a prefix but which are not arrays should be wrapped
/// in quotes, e.g.:
/// ```
/// use sqlx::postgres::PgTypeInfo;
/// use sqlx::Type;
///
/// /// `CREATE TYPE "_foo" AS ENUM ('Bar', 'Baz');`
/// #[derive(sqlx::Type)]
/// // Will prevent SQLx from inferring `_foo` as an array type.
/// #[sqlx(type_name = r#""_foo""#)]
/// enum Foo {
/// Bar,
/// Baz
/// }
///
/// assert_eq!(Foo::type_info().name(), r#""_foo""#);
/// ```
pub const fn with_name(name: &'static str) -> Self {
Self(PgType::DeclareWithName(UStr::Static(name)))
}
/// Create a `PgTypeInfo` of an array from the name of its element type.
///
/// The array type OID will be fetched from Postgres on use of a value of this type.
/// The fetched OID will be cached per-connection.
pub fn array_of(elem_name: &'static str) -> Self {
// to satisfy `name()` and `display_name()`, we need to construct strings to return
Self(PgType::DeclareArrayOf(Arc::new(PgArrayOf {
elem_name: elem_name.into(),
name: format!("{elem_name}[]").into(),
})))
}
/// Create a `PgTypeInfo` from an OID.
///
/// Note that the OID for a type is very dependent on the environment. If you only ever use
/// one database or if this is an unhandled built-in type, you should be fine. Otherwise,
/// you will be better served using [`with_name`](Self::with_name).
/// you will be better served using [`Self::with_name()`].
///
/// ### Note: Interaction with `==`
/// This constructor may give surprising results with `==`.
///
/// See [the type-level docs][Self] for details.
pub const fn with_oid(oid: Oid) -> Self {
Self(PgType::DeclareWithOid(oid))
}
/// Returns `true` if `self` can be compared exactly to `other`.
///
/// Unlike `==`, this will return false if
pub fn type_eq(&self, other: &Self) -> bool {
self.eq_impl(other, false)
}
}
// DEVELOPER PRO TIP: find builtin type OIDs easily by grepping this file
@@ -464,6 +562,9 @@ impl PgType {
PgType::DeclareWithName(_) => {
return None;
}
PgType::DeclareArrayOf(_) => {
return None;
}
})
}
@@ -564,6 +665,7 @@ impl PgType {
PgType::Custom(ty) => &ty.name,
PgType::DeclareWithOid(_) => "?",
PgType::DeclareWithName(name) => name,
PgType::DeclareArrayOf(array) => &array.name,
}
}
@@ -664,6 +766,7 @@ impl PgType {
PgType::Custom(ty) => &ty.name,
PgType::DeclareWithOid(_) => "?",
PgType::DeclareWithName(name) => name,
PgType::DeclareArrayOf(array) => &array.name,
}
}
@@ -771,13 +874,16 @@ impl PgType {
PgType::DeclareWithName(name) => {
unreachable!("(bug) use of unresolved type declaration [name={name}]");
}
PgType::DeclareArrayOf(array) => {
unreachable!(
"(bug) use of unresolved type declaration [array of={}]",
array.elem_name
);
}
}
}
/// If `self` is an array type, return the type info for its element.
///
/// This method should only be called on resolved types: calling it on
/// a type that is merely declared (DeclareWithOid/Name) is a bug.
pub(crate) fn try_array_element(&self) -> Option<Cow<'_, PgTypeInfo>> {
// We explicitly match on all the `None` cases to ensure an exhaustive match.
match self {
@@ -885,14 +991,50 @@ impl PgType {
PgTypeKind::Enum(_) => None,
PgTypeKind::Range(_) => None,
},
PgType::DeclareWithOid(oid) => {
unreachable!("(bug) use of unresolved type declaration [oid={}]", oid.0);
}
PgType::DeclareWithOid(_) => None,
PgType::DeclareWithName(name) => {
unreachable!("(bug) use of unresolved type declaration [name={name}]");
// LEGACY: infer the array element name from a `_` prefix
UStr::strip_prefix(name, "_")
.map(|elem| Cow::Owned(PgTypeInfo(PgType::DeclareWithName(elem))))
}
PgType::DeclareArrayOf(array) => Some(Cow::Owned(PgTypeInfo(PgType::DeclareWithName(
array.elem_name.clone(),
)))),
}
}
/// Returns `true` if this type cannot be matched by name.
fn is_declare_with_oid(&self) -> bool {
matches!(self, Self::DeclareWithOid(_))
}
/// Compare two `PgType`s, first by OID, then by array element, then by name.
///
/// If `soft_eq` is true and `self` or `other` is `DeclareWithOid` but not both, return `true`
/// before checking names.
fn eq_impl(&self, other: &Self, soft_eq: bool) -> bool {
if let (Some(a), Some(b)) = (self.try_oid(), other.try_oid()) {
// If there are OIDs available, use OIDs to perform a direct match
return a == b;
}
if soft_eq && (self.is_declare_with_oid() || other.is_declare_with_oid()) {
// If we get to this point, one instance is `DeclareWithOid()` and the other is
// `DeclareArrayOf()` or `DeclareWithName()`, which means we can't compare the two.
//
// Since this is only likely to occur when using the text protocol where we can't
// resolve type names before executing a query, we can just opt out of typechecking.
return true;
}
if let (Some(elem_a), Some(elem_b)) = (self.try_array_element(), other.try_array_element())
{
return elem_a == elem_b;
}
// Otherwise, perform a match on the name
name_eq(self.name(), other.name())
}
}
impl TypeInfo for PgTypeInfo {
@@ -907,6 +1049,13 @@ impl TypeInfo for PgTypeInfo {
fn is_void(&self) -> bool {
matches!(self.0, PgType::Void)
}
fn type_compatible(&self, other: &Self) -> bool
where
Self: Sized,
{
self == other
}
}
impl PartialEq<PgCustomType> for PgCustomType {
@@ -1140,22 +1289,7 @@ impl Display for PgTypeInfo {
impl PartialEq<PgType> for PgType {
fn eq(&self, other: &PgType) -> bool {
if let (Some(a), Some(b)) = (self.try_oid(), other.try_oid()) {
// If there are OIDs available, use OIDs to perform a direct match
a == b
} else if matches!(
(self, other),
(PgType::DeclareWithName(_), PgType::DeclareWithOid(_))
| (PgType::DeclareWithOid(_), PgType::DeclareWithName(_))
) {
// One is a declare-with-name and the other is a declare-with-id
// This only occurs in the TEXT protocol with custom types
// Just opt-out of type checking here
true
} else {
// Otherwise, perform a match on the name
name_eq(self.name(), other.name())
}
self.eq_impl(other, true)
}
}

10
sqlx-postgres/src/types/array.rs
View File

@@ -156,11 +156,10 @@ where
T: Encode<'q, Postgres> + Type<Postgres>,
{
fn encode_by_ref(&self, buf: &mut PgArgumentBuffer) -> Result<IsNull, BoxDynError> {
let type_info = if self.is_empty() {
T::type_info()
} else {
self[0].produces().unwrap_or_else(T::type_info)
};
let type_info = self
.first()
.and_then(Encode::produces)
.unwrap_or_else(T::type_info);
buf.extend(&1_i32.to_be_bytes()); // number of dimensions
buf.extend(&0_i32.to_be_bytes()); // flags
@@ -168,6 +167,7 @@ where
// element type
match type_info.0 {
PgType::DeclareWithName(name) => buf.patch_type_by_name(&name),
PgType::DeclareArrayOf(array) => buf.patch_array_type(array),
ty => {
buf.extend(&ty.oid().0.to_be_bytes());