Skip to content

Commit

Permalink
chore: reduce BoxFuture's when using recursion.
Browse files Browse the repository at this point in the history
  • Loading branch information
joeydewaal committed Dec 6, 2024
1 parent 42ce24d commit 031d7b2
Showing 1 changed file with 90 additions and 105 deletions.
195 changes: 90 additions & 105 deletions sqlx-postgres/src/connection/describe.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,6 @@ use crate::types::Json;
use crate::types::Oid;
use crate::HashMap;
use crate::{PgColumn, PgConnection, PgTypeInfo};
use futures_core::future::BoxFuture;
use smallvec::SmallVec;
use sqlx_core::query_builder::QueryBuilder;
use std::sync::Arc;
Expand Down Expand Up @@ -169,7 +168,7 @@ impl PgConnection {

// fallback to asking the database directly for a type name
if should_fetch {
let info = self.fetch_type_by_oid(oid).await?;
let info = Box::pin(async { self.fetch_type_by_oid(oid).await }).await?;

// cache the type name <-> oid relationship in a paired hashmap
// so we don't come down this road again
Expand All @@ -190,74 +189,70 @@ impl PgConnection {
}
}

fn fetch_type_by_oid(&mut self, oid: Oid) -> BoxFuture<'_, Result<PgTypeInfo, Error>> {
Box::pin(async move {
let (name, typ_type, category, relation_id, element, base_type): (
String,
i8,
i8,
Oid,
Oid,
Oid,
) = query_as(
// Converting the OID to `regtype` and then `text` will give us the name that
// the type will need to be found at by search_path.
"SELECT oid::regtype::text, \
async fn fetch_type_by_oid(&mut self, oid: Oid) -> Result<PgTypeInfo, Error> {
let (name, typ_type, category, relation_id, element, base_type): (
String,
i8,
i8,
Oid,
Oid,
Oid,
) = query_as(
// Converting the OID to `regtype` and then `text` will give us the name that
// the type will need to be found at by search_path.
"SELECT oid::regtype::text, \
typtype, \
typcategory, \
typrelid, \
typelem, \
typbasetype \
FROM pg_catalog.pg_type \
WHERE oid = $1",
)
.bind(oid)
.fetch_one(&mut *self)
.await?;

let typ_type = TypType::try_from(typ_type);
let category = TypCategory::try_from(category);

match (typ_type, category) {
(Ok(TypType::Domain), _) => self.fetch_domain_by_oid(oid, base_type, name).await,

(Ok(TypType::Base), Ok(TypCategory::Array)) => {
Ok(PgTypeInfo(PgType::Custom(Arc::new(PgCustomType {
kind: PgTypeKind::Array(
self.maybe_fetch_type_info_by_oid(element, true).await?,
),
name: name.into(),
oid,
}))))
}

(Ok(TypType::Pseudo), Ok(TypCategory::Pseudo)) => {
Ok(PgTypeInfo(PgType::Custom(Arc::new(PgCustomType {
kind: PgTypeKind::Pseudo,
name: name.into(),
oid,
}))))
}
)
.bind(oid)
.fetch_one(&mut *self)
.await?;

(Ok(TypType::Range), Ok(TypCategory::Range)) => {
self.fetch_range_by_oid(oid, name).await
}
let typ_type = TypType::try_from(typ_type);
let category = TypCategory::try_from(category);

(Ok(TypType::Enum), Ok(TypCategory::Enum)) => {
self.fetch_enum_by_oid(oid, name).await
}
match (typ_type, category) {
(Ok(TypType::Domain), _) => self.fetch_domain_by_oid(oid, base_type, name).await,

(Ok(TypType::Composite), Ok(TypCategory::Composite)) => {
self.fetch_composite_by_oid(oid, relation_id, name).await
}
(Ok(TypType::Base), Ok(TypCategory::Array)) => {
Ok(PgTypeInfo(PgType::Custom(Arc::new(PgCustomType {
kind: PgTypeKind::Array(
self.maybe_fetch_type_info_by_oid(element, true).await?,
),
name: name.into(),
oid,
}))))
}

_ => Ok(PgTypeInfo(PgType::Custom(Arc::new(PgCustomType {
kind: PgTypeKind::Simple,
(Ok(TypType::Pseudo), Ok(TypCategory::Pseudo)) => {
Ok(PgTypeInfo(PgType::Custom(Arc::new(PgCustomType {
kind: PgTypeKind::Pseudo,
name: name.into(),
oid,
})))),
}))))
}

(Ok(TypType::Range), Ok(TypCategory::Range)) => {
self.fetch_range_by_oid(oid, name).await
}
})

(Ok(TypType::Enum), Ok(TypCategory::Enum)) => self.fetch_enum_by_oid(oid, name).await,

(Ok(TypType::Composite), Ok(TypCategory::Composite)) => {
self.fetch_composite_by_oid(oid, relation_id, name).await
}

_ => Ok(PgTypeInfo(PgType::Custom(Arc::new(PgCustomType {
kind: PgTypeKind::Simple,
name: name.into(),
oid,
})))),
}
}

async fn fetch_enum_by_oid(&mut self, oid: Oid, name: String) -> Result<PgTypeInfo, Error> {
Expand All @@ -280,85 +275,75 @@ ORDER BY enumsortorder
}))))
}

fn fetch_composite_by_oid(
async fn fetch_composite_by_oid(
&mut self,
oid: Oid,
relation_id: Oid,
name: String,
) -> BoxFuture<'_, Result<PgTypeInfo, Error>> {
Box::pin(async move {
let raw_fields: Vec<(String, Oid)> = query_as(
r#"
) -> Result<PgTypeInfo, Error> {
let raw_fields: Vec<(String, Oid)> = query_as(
r#"
SELECT attname, atttypid
FROM pg_catalog.pg_attribute
WHERE attrelid = $1
AND NOT attisdropped
AND attnum > 0
ORDER BY attnum
"#,
)
.bind(relation_id)
.fetch_all(&mut *self)
.await?;
)
.bind(relation_id)
.fetch_all(&mut *self)
.await?;

let mut fields = Vec::new();
let mut fields = Vec::new();

for (field_name, field_oid) in raw_fields.into_iter() {
let field_type = self.maybe_fetch_type_info_by_oid(field_oid, true).await?;
for (field_name, field_oid) in raw_fields.into_iter() {
let field_type = self.maybe_fetch_type_info_by_oid(field_oid, true).await?;

fields.push((field_name, field_type));
}
fields.push((field_name, field_type));
}

Ok(PgTypeInfo(PgType::Custom(Arc::new(PgCustomType {
oid,
name: name.into(),
kind: PgTypeKind::Composite(Arc::from(fields)),
}))))
})
Ok(PgTypeInfo(PgType::Custom(Arc::new(PgCustomType {
oid,
name: name.into(),
kind: PgTypeKind::Composite(Arc::from(fields)),
}))))
}

fn fetch_domain_by_oid(
async fn fetch_domain_by_oid(
&mut self,
oid: Oid,
base_type: Oid,
name: String,
) -> BoxFuture<'_, Result<PgTypeInfo, Error>> {
Box::pin(async move {
let base_type = self.maybe_fetch_type_info_by_oid(base_type, true).await?;
) -> Result<PgTypeInfo, Error> {
let base_type = self.maybe_fetch_type_info_by_oid(base_type, true).await?;

Ok(PgTypeInfo(PgType::Custom(Arc::new(PgCustomType {
oid,
name: name.into(),
kind: PgTypeKind::Domain(base_type),
}))))
})
Ok(PgTypeInfo(PgType::Custom(Arc::new(PgCustomType {
oid,
name: name.into(),
kind: PgTypeKind::Domain(base_type),
}))))
}

fn fetch_range_by_oid(
&mut self,
oid: Oid,
name: String,
) -> BoxFuture<'_, Result<PgTypeInfo, Error>> {
Box::pin(async move {
let element_oid: Oid = query_scalar(
r#"
async fn fetch_range_by_oid(&mut self, oid: Oid, name: String) -> Result<PgTypeInfo, Error> {
let element_oid: Oid = query_scalar(
r#"
SELECT rngsubtype
FROM pg_catalog.pg_range
WHERE rngtypid = $1
"#,
)
.bind(oid)
.fetch_one(&mut *self)
.await?;
)
.bind(oid)
.fetch_one(&mut *self)
.await?;

let element = self.maybe_fetch_type_info_by_oid(element_oid, true).await?;
let element = self.maybe_fetch_type_info_by_oid(element_oid, true).await?;

Ok(PgTypeInfo(PgType::Custom(Arc::new(PgCustomType {
kind: PgTypeKind::Range(element),
name: name.into(),
oid,
}))))
})
Ok(PgTypeInfo(PgType::Custom(Arc::new(PgCustomType {
kind: PgTypeKind::Range(element),
name: name.into(),
oid,
}))))
}

pub(crate) async fn resolve_type_id(&mut self, ty: &PgType) -> Result<Oid, Error> {
Expand Down

0 comments on commit 031d7b2

Please sign in to comment.