diff --git a/test_codegen/codegen/src/queries/stress.rs b/test_codegen/codegen/src/queries/stress.rs index 5bcad3dc..acabde7c 100644 --- a/test_codegen/codegen/src/queries/stress.rs +++ b/test_codegen/codegen/src/queries/stress.rs @@ -930,6 +930,57 @@ pub mod sync { Ok(it) } } + pub struct SchemaNightmareCompositeQuery<'a, C: GenericClient, T, const N: usize> { + client: &'a mut C, + params: [&'a (dyn postgres_types::ToSql + Sync); N], + stmt: &'a mut crate::client::sync::Stmt, + extractor: fn(&postgres::Row) -> crate::types::schema::NightmareCompositeBorrowed, + mapper: fn(crate::types::schema::NightmareCompositeBorrowed) -> T, + } + impl<'a, C, T: 'a, const N: usize> SchemaNightmareCompositeQuery<'a, C, T, N> + where + C: GenericClient, + { + pub fn map( + self, + mapper: fn(crate::types::schema::NightmareCompositeBorrowed) -> R, + ) -> SchemaNightmareCompositeQuery<'a, C, R, N> { + SchemaNightmareCompositeQuery { + client: self.client, + params: self.params, + stmt: self.stmt, + extractor: self.extractor, + mapper, + } + } + pub fn one(self) -> Result { + let stmt = self.stmt.prepare(self.client)?; + let row = self.client.query_one(stmt, &self.params)?; + Ok((self.mapper)((self.extractor)(&row))) + } + pub fn all(self) -> Result, postgres::Error> { + self.iter()?.collect() + } + pub fn opt(self) -> Result, postgres::Error> { + let stmt = self.stmt.prepare(self.client)?; + Ok(self + .client + .query_opt(stmt, &self.params)? + .map(|row| (self.mapper)((self.extractor)(&row)))) + } + pub fn iter( + self, + ) -> Result> + 'a, postgres::Error> + { + let stmt = self.stmt.prepare(self.client)?; + let it = self + .client + .query_raw(stmt, crate::slice_iter(&self.params))? + .iterator() + .map(move |res| res.map(|row| (self.mapper)((self.extractor)(&row)))); + Ok(it) + } + } pub fn select_everything() -> SelectEverythingStmt { SelectEverythingStmt(crate::client::sync::Stmt::new( "SELECT @@ -1604,6 +1655,47 @@ FROM client.execute(stmt, &[composite]) } } + pub fn select_schema_nightmare() -> SelectSchemaNightmareStmt { + SelectSchemaNightmareStmt(crate::client::sync::Stmt::new( + "SELECT + * +FROM + schema.nightmare", + )) + } + pub struct SelectSchemaNightmareStmt(crate::client::sync::Stmt); + impl SelectSchemaNightmareStmt { + pub fn bind<'a, C: GenericClient>( + &'a mut self, + client: &'a mut C, + ) -> SchemaNightmareCompositeQuery<'a, C, crate::types::schema::NightmareComposite, 0> + { + SchemaNightmareCompositeQuery { + client, + params: [], + stmt: &mut self.0, + extractor: |row| row.get(0), + mapper: |it| it.into(), + } + } + } + pub fn insert_schema_nightmare() -> InsertSchemaNightmareStmt { + InsertSchemaNightmareStmt(crate::client::sync::Stmt::new( + "INSERT INTO schema.nightmare (composite) + VALUES ($1)", + )) + } + pub struct InsertSchemaNightmareStmt(crate::client::sync::Stmt); + impl InsertSchemaNightmareStmt { + pub fn bind<'a, C: GenericClient>( + &'a mut self, + client: &'a mut C, + composite: &'a crate::types::schema::NightmareCompositeParams<'a>, + ) -> Result { + let stmt = self.0.prepare(client)?; + client.execute(stmt, &[composite]) + } + } } pub mod async_ { use crate::client::async_::GenericClient; @@ -1883,6 +1975,61 @@ pub mod async_ { Ok(it) } } + pub struct SchemaNightmareCompositeQuery<'a, C: GenericClient, T, const N: usize> { + client: &'a C, + params: [&'a (dyn postgres_types::ToSql + Sync); N], + stmt: &'a mut crate::client::async_::Stmt, + extractor: fn(&tokio_postgres::Row) -> crate::types::schema::NightmareCompositeBorrowed, + mapper: fn(crate::types::schema::NightmareCompositeBorrowed) -> T, + } + impl<'a, C, T: 'a, const N: usize> SchemaNightmareCompositeQuery<'a, C, T, N> + where + C: GenericClient, + { + pub fn map( + self, + mapper: fn(crate::types::schema::NightmareCompositeBorrowed) -> R, + ) -> SchemaNightmareCompositeQuery<'a, C, R, N> { + SchemaNightmareCompositeQuery { + client: self.client, + params: self.params, + stmt: self.stmt, + extractor: self.extractor, + mapper, + } + } + pub async fn one(self) -> Result { + let stmt = self.stmt.prepare(self.client).await?; + let row = self.client.query_one(stmt, &self.params).await?; + Ok((self.mapper)((self.extractor)(&row))) + } + pub async fn all(self) -> Result, tokio_postgres::Error> { + self.iter().await?.try_collect().await + } + pub async fn opt(self) -> Result, tokio_postgres::Error> { + let stmt = self.stmt.prepare(self.client).await?; + Ok(self + .client + .query_opt(stmt, &self.params) + .await? + .map(|row| (self.mapper)((self.extractor)(&row)))) + } + pub async fn iter( + self, + ) -> Result< + impl futures::Stream> + 'a, + tokio_postgres::Error, + > { + let stmt = self.stmt.prepare(self.client).await?; + let it = self + .client + .query_raw(stmt, crate::slice_iter(&self.params)) + .await? + .map(move |res| res.map(|row| (self.mapper)((self.extractor)(&row)))) + .into_stream(); + Ok(it) + } + } pub fn select_everything() -> SelectEverythingStmt { SelectEverythingStmt(crate::client::async_::Stmt::new( "SELECT @@ -2569,4 +2716,45 @@ FROM client.execute(stmt, &[composite]).await } } + pub fn select_schema_nightmare() -> SelectSchemaNightmareStmt { + SelectSchemaNightmareStmt(crate::client::async_::Stmt::new( + "SELECT + * +FROM + schema.nightmare", + )) + } + pub struct SelectSchemaNightmareStmt(crate::client::async_::Stmt); + impl SelectSchemaNightmareStmt { + pub fn bind<'a, C: GenericClient>( + &'a mut self, + client: &'a C, + ) -> SchemaNightmareCompositeQuery<'a, C, crate::types::schema::NightmareComposite, 0> + { + SchemaNightmareCompositeQuery { + client, + params: [], + stmt: &mut self.0, + extractor: |row| row.get(0), + mapper: |it| it.into(), + } + } + } + pub fn insert_schema_nightmare() -> InsertSchemaNightmareStmt { + InsertSchemaNightmareStmt(crate::client::async_::Stmt::new( + "INSERT INTO schema.nightmare (composite) + VALUES ($1)", + )) + } + pub struct InsertSchemaNightmareStmt(crate::client::async_::Stmt); + impl InsertSchemaNightmareStmt { + pub async fn bind<'a, C: GenericClient>( + &'a mut self, + client: &'a C, + composite: &'a crate::types::schema::NightmareCompositeParams<'a>, + ) -> Result { + let stmt = self.0.prepare(client).await?; + client.execute(stmt, &[composite]).await + } + } } diff --git a/test_codegen/codegen/src/types.rs b/test_codegen/codegen/src/types.rs index 61f92e44..60945a93 100644 --- a/test_codegen/codegen/src/types.rs +++ b/test_codegen/codegen/src/types.rs @@ -1199,3 +1199,147 @@ impl<'a> postgres_types::FromSql<'a> for SyntaxEnum { } } } +pub mod schema { + #[derive(serde::Serialize, Debug, postgres_types :: FromSql, Clone, PartialEq)] + #[postgres(name = "nightmare_composite")] + pub struct NightmareComposite { + #[postgres(name = "custom")] + pub custom: Vec, + #[postgres(name = "spongebob")] + pub spongebob: Vec, + #[postgres(name = "domain")] + pub domain: String, + } + #[derive(Debug)] + pub struct NightmareCompositeBorrowed<'a> { + pub custom: crate::ArrayIterator<'a, super::CustomCompositeBorrowed<'a>>, + pub spongebob: crate::ArrayIterator<'a, super::SpongebobCharacter>, + pub domain: &'a str, + } + impl<'a> From> for NightmareComposite { + fn from( + NightmareCompositeBorrowed { + custom, + spongebob, + domain, + }: NightmareCompositeBorrowed<'a>, + ) -> Self { + Self { + custom: custom.map(|v| v.into()).collect(), + spongebob: spongebob.map(|v| v).collect(), + domain: domain.into(), + } + } + } + impl<'a> postgres_types::FromSql<'a> for NightmareCompositeBorrowed<'a> { + fn from_sql( + ty: &postgres_types::Type, + out: &'a [u8], + ) -> Result, Box> + { + let fields = match *ty.kind() { + postgres_types::Kind::Composite(ref fields) => fields, + _ => unreachable!(), + }; + let mut out = out; + let num_fields = postgres_types::private::read_be_i32(&mut out)?; + if num_fields as usize != fields.len() { + return std::result::Result::Err(std::convert::Into::into(format!( + "invalid field count: {} vs {}", + num_fields, + fields.len() + ))); + } + let _oid = postgres_types::private::read_be_i32(&mut out)?; + let custom = postgres_types::private::read_value(fields[0].type_(), &mut out)?; + let _oid = postgres_types::private::read_be_i32(&mut out)?; + let spongebob = postgres_types::private::read_value(fields[1].type_(), &mut out)?; + let _oid = postgres_types::private::read_be_i32(&mut out)?; + let domain = postgres_types::private::read_value(fields[2].type_(), &mut out)?; + Ok(NightmareCompositeBorrowed { + custom, + spongebob, + domain, + }) + } + fn accepts(ty: &postgres_types::Type) -> bool { + ty.name() == "nightmare_composite" && ty.schema() == "schema" + } + } + #[derive(Debug)] + pub struct NightmareCompositeParams<'a> { + pub custom: &'a [super::CustomCompositeBorrowed<'a>], + pub spongebob: &'a [super::SpongebobCharacter], + pub domain: &'a str, + } + impl<'a> postgres_types::ToSql for NightmareCompositeParams<'a> { + fn to_sql( + &self, + ty: &postgres_types::Type, + out: &mut postgres_types::private::BytesMut, + ) -> Result> { + let NightmareCompositeParams { + custom, + spongebob, + domain, + } = self; + let fields = match *ty.kind() { + postgres_types::Kind::Composite(ref fields) => fields, + _ => unreachable!(), + }; + out.extend_from_slice(&(fields.len() as i32).to_be_bytes()); + for field in fields { + out.extend_from_slice(&field.type_().oid().to_be_bytes()); + let base = out.len(); + out.extend_from_slice(&[0; 4]); + let r = match field.name() { + "custom" => postgres_types::ToSql::to_sql(custom, field.type_(), out), + "spongebob" => postgres_types::ToSql::to_sql(spongebob, field.type_(), out), + "domain" => { + postgres_types::ToSql::to_sql(&crate::Domain(domain), field.type_(), out) + } + _ => unreachable!(), + }; + let count = match r? { + postgres_types::IsNull::Yes => -1, + postgres_types::IsNull::No => { + let len = out.len() - base - 4; + if len > i32::max_value() as usize { + return Err(Into::into("value too large to transmit")); + } + len as i32 + } + }; + out[base..base + 4].copy_from_slice(&count.to_be_bytes()); + } + Ok(postgres_types::IsNull::No) + } + fn accepts(ty: &postgres_types::Type) -> bool { + if ty.name() != "nightmare_composite" { + return false; + } + match *ty.kind() { + postgres_types::Kind::Composite(ref fields) => { + if fields.len() != 3 { + return false; + } + fields.iter().all(| f | match f.name() + { + "custom" => < &'a [super::CustomCompositeBorrowed<'a>] as postgres_types :: + ToSql > :: accepts(f.type_()),"spongebob" => < &'a [super::SpongebobCharacter] as postgres_types :: + ToSql > :: accepts(f.type_()),"domain" => < crate::Domain::<&'a str> as postgres_types :: + ToSql > :: accepts(f.type_()),_ => false, + }) + } + _ => false, + } + } + fn to_sql_checked( + &self, + ty: &postgres_types::Type, + out: &mut postgres_types::private::BytesMut, + ) -> Result> { + postgres_types::__to_sql_checked(self, ty, out) + } + } +} diff --git a/test_codegen/queries/stress.sql b/test_codegen/queries/stress.sql index e7d9e9e7..da56994b 100644 --- a/test_codegen/queries/stress.sql +++ b/test_codegen/queries/stress.sql @@ -40,3 +40,14 @@ FROM INSERT INTO nightmare (composite) VALUES (:composite); + +--! select_schema_nightmare +SELECT + * +FROM + schema.nightmare; + +--! insert_schema_nightmare +INSERT INTO schema.nightmare (composite) + VALUES (:composite); + diff --git a/test_codegen/schema.sql b/test_codegen/schema.sql index 7d7fcea9..8e2e99b4 100644 --- a/test_codegen/schema.sql +++ b/test_codegen/schema.sql @@ -189,6 +189,20 @@ CREATE TABLE nightmare ( composite nightmare_composite NOT NULL ); +-- Schema + +CREATE SCHEMA schema; + +CREATE TYPE schema.nightmare_composite AS ( + custom custom_composite[], + spongebob spongebob_character[], + domain my_domain +); + +CREATE TABLE schema.nightmare ( + composite schema.nightmare_composite NOT NULL +); + -- Syntax CREATE TYPE syntax_composite AS ( diff --git a/test_codegen/src/main.rs b/test_codegen/src/main.rs index 8a8f019f..69efbb22 100644 --- a/test_codegen/src/main.rs +++ b/test_codegen/src/main.rs @@ -34,8 +34,9 @@ use codegen::{ }, stress::{ sync::{ - insert_everything, insert_everything_array, insert_nightmare, select_everything, - select_everything_array, select_nightmare, + insert_everything, insert_everything_array, insert_nightmare, + insert_schema_nightmare, select_everything, select_everything_array, + select_nightmare, select_schema_nightmare, }, Everything, EverythingArray, EverythingArrayParams, EverythingParams, }, @@ -45,12 +46,13 @@ use codegen::{ }, }, types::{ - CloneCompositeBorrowed, CopyComposite, CustomComposite, CustomCompositeBorrowed, + schema, CloneCompositeBorrowed, CopyComposite, CustomComposite, CustomCompositeBorrowed, DomainComposite, DomainCompositeParams, EnumWithDot, NamedComposite, NamedCompositeBorrowed, NamedCompositeWithDot, NightmareComposite, NightmareCompositeParams, NullityComposite, NullityCompositeParams, SpongebobCharacter, SyntaxComposite, SyntaxEnum, - }, IterSql, + }, + IterSql, }; pub fn main() { @@ -556,6 +558,30 @@ pub fn test_stress(client: &mut Client) { assert_eq!(1, insert_nightmare().bind(client, ¶ms).unwrap()); let actual = select_nightmare().bind(client).one().unwrap(); assert_eq!(expected, actual); + + // In a named schema + let expected = schema::NightmareComposite { + custom: vec![CustomComposite { + wow: "Bob".to_string(), + such_cool: 42, + nice: SpongebobCharacter::Squidward, + }], + spongebob: vec![SpongebobCharacter::Bob, SpongebobCharacter::Patrick], + domain: "Hello".to_string(), + }; + let params = schema::NightmareCompositeParams { + custom: &[CustomCompositeBorrowed { + wow: "Bob", + such_cool: 42, + nice: SpongebobCharacter::Squidward, + }], + spongebob: &[SpongebobCharacter::Bob, SpongebobCharacter::Patrick], + domain: "Hello", + }; + + assert_eq!(1, insert_schema_nightmare().bind(client, ¶ms).unwrap()); + let actual = select_schema_nightmare().bind(client).one().unwrap(); + assert_eq!(expected, actual); } // Test keyword escaping