From 84e2e70676510bd64d888a7700cce0d0a18bf1a6 Mon Sep 17 00:00:00 2001 From: James Guthrie Date: Tue, 12 Nov 2024 10:08:35 +0100 Subject: [PATCH] feat: allow specifying type alignment through #[pgrx(alignment = ...)] Postgres allows for types to specify their alignment in the `CREATE TYPE` statement. This change adds the ability for the alignment of a type to be declared directly on the type, as follows: ``` #[derive(PostgresType)] #[pgrx(alignment = "double")] struct AlignedTo8Bytes { v1: u64, v2: [u64; 3] } ``` --- pgrx-examples/custom_types/src/alignment.rs | 42 ++++++++ pgrx-examples/custom_types/src/lib.rs | 1 + pgrx-macros/src/lib.rs | 1 + pgrx-sql-entity-graph/src/lib.rs | 2 + .../src/postgres_type/entity.rs | 95 ++++++++++++++++++- .../src/postgres_type/mod.rs | 14 ++- 6 files changed, 151 insertions(+), 4 deletions(-) create mode 100644 pgrx-examples/custom_types/src/alignment.rs diff --git a/pgrx-examples/custom_types/src/alignment.rs b/pgrx-examples/custom_types/src/alignment.rs new file mode 100644 index 0000000000..9b89856622 --- /dev/null +++ b/pgrx-examples/custom_types/src/alignment.rs @@ -0,0 +1,42 @@ +//LICENSE Portions Copyright 2019-2021 ZomboDB, LLC. +//LICENSE +//LICENSE Portions Copyright 2021-2023 Technology Concepts & Design, Inc. +//LICENSE +//LICENSE Portions Copyright 2023-2023 PgCentral Foundation, Inc. +//LICENSE +//LICENSE All rights reserved. +//LICENSE +//LICENSE Use of this source code is governed by the MIT license that can be found in the LICENSE file. +use pgrx::prelude::*; +use serde::*; + +#[derive(PostgresType, Serialize, Deserialize)] +pub struct AlignedTo4Bytes { + v1: u64, + v2: [u64; 3], +} + +#[derive(PostgresType, Serialize, Deserialize)] +#[pgrx(alignment = "double")] +pub struct AlignedTo8Bytes { + v1: u64, + v2: [u64; 3], +} + +#[cfg(any(test, feature = "pg_test"))] +#[pg_schema] +mod tests { + use pgrx::prelude::*; + + #[cfg(not(feature = "no-schema-generation"))] + #[pg_test] + fn test_alignment_is_correct() { + let val = Spi::get_one::(r#"SELECT typalign = 'i' FROM pg_type WHERE typname = 'alignedto4bytes'"#).unwrap().unwrap(); + + assert!(val == true); + + let val = Spi::get_one::(r#"SELECT typalign = 'd' FROM pg_type WHERE typname = 'alignedto8bytes'"#).unwrap().unwrap(); + + assert!(val == true); + } +} diff --git a/pgrx-examples/custom_types/src/lib.rs b/pgrx-examples/custom_types/src/lib.rs index 3a750fca92..ab22ca878b 100644 --- a/pgrx-examples/custom_types/src/lib.rs +++ b/pgrx-examples/custom_types/src/lib.rs @@ -7,6 +7,7 @@ //LICENSE All rights reserved. //LICENSE //LICENSE Use of this source code is governed by the MIT license that can be found in the LICENSE file. +mod alignment; mod complex; mod fixed_size; mod generic_enum; diff --git a/pgrx-macros/src/lib.rs b/pgrx-macros/src/lib.rs index 0d8e0919bf..bfa5fb0e85 100644 --- a/pgrx-macros/src/lib.rs +++ b/pgrx-macros/src/lib.rs @@ -781,6 +781,7 @@ Optionally accepts the following attributes: * `inoutfuncs(some_in_fn, some_out_fn)`: Define custom in/out functions for the type. * `pgvarlena_inoutfuncs(some_in_fn, some_out_fn)`: Define custom in/out functions for the `PgVarlena` of this type. +* `pgrx(alignment = "")`: Define alignment for this type. One of `char`, `int2`, `int4`, or `double` (see https://www.postgresql.org/docs/current/sql-createtype.html). * `sql`: Same arguments as [`#[pgrx(sql = ..)]`](macro@pgrx). */ #[proc_macro_derive( diff --git a/pgrx-sql-entity-graph/src/lib.rs b/pgrx-sql-entity-graph/src/lib.rs index 47114eee5e..8cb9cf4253 100644 --- a/pgrx-sql-entity-graph/src/lib.rs +++ b/pgrx-sql-entity-graph/src/lib.rs @@ -94,6 +94,8 @@ pub trait SqlGraphIdentifier { fn line(&self) -> Option; } +pub use postgres_type::Alignment; + /// An entity corresponding to some SQL required by the extension. #[derive(Debug, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)] pub enum SqlGraphEntity { diff --git a/pgrx-sql-entity-graph/src/postgres_type/entity.rs b/pgrx-sql-entity-graph/src/postgres_type/entity.rs index 28c6434754..80ba26a92a 100644 --- a/pgrx-sql-entity-graph/src/postgres_type/entity.rs +++ b/pgrx-sql-entity-graph/src/postgres_type/entity.rs @@ -16,13 +16,85 @@ */ use crate::mapping::RustSqlMapping; +use crate::pgrx_attribute::{ArgValue, PgrxArg, PgrxAttribute}; use crate::pgrx_sql::PgrxSql; use crate::to_sql::entity::ToSqlConfigEntity; use crate::to_sql::ToSql; use crate::{SqlGraphEntity, SqlGraphIdentifier, TypeMatch}; +use eyre::eyre; +use proc_macro2::TokenStream; +use quote::{format_ident, quote, ToTokens, TokenStreamExt}; use std::collections::BTreeSet; +use syn::spanned::Spanned; +use syn::{AttrStyle, Attribute, Lit}; + +#[derive(Debug, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)] +pub enum Alignment { + Char, + Int2, + Int4, + Double, +} + +const INVALID_ATTR_CONTENT: &str = r#"expected `#[pgrx(alignment = align)]`, where `align` is "char", "int2", "int4", or "double""#; + +impl ToTokens for Alignment { + fn to_tokens(&self, tokens: &mut TokenStream) { + let value = match self { + Alignment::Char => format_ident!("Char"), + Alignment::Int2 => format_ident!("Int2"), + Alignment::Int4 => format_ident!("Int4"), + Alignment::Double => format_ident!("Double"), + }; + let quoted = quote! { + ::pgrx::pgrx_sql_entity_graph::Alignment::#value + }; + tokens.append_all(quoted); + } +} + +impl Alignment { + pub fn from_attribute(attr: &Attribute) -> Result, syn::Error> { + if attr.style != AttrStyle::Outer { + return Err(syn::Error::new( + attr.span(), + "#[pgrx(alignment = ..)] is only valid in an outer context", + )); + } + + let attr = attr.parse_args::()?; + for arg in attr.args.iter() { + let PgrxArg::NameValue(ref nv) = arg; + if !nv.path.is_ident("alignment") { + continue; + } + + return match nv.value { + ArgValue::Lit(Lit::Str(ref s)) => match s.value().as_ref() { + "char" => Ok(Some(Self::Char)), + "int2" => Ok(Some(Self::Int2)), + "int4" => Ok(Some(Self::Int4)), + "double" => Ok(Some(Self::Double)), + _ => Err(syn::Error::new(s.span(), INVALID_ATTR_CONTENT)), + }, + ArgValue::Path(ref p) => Err(syn::Error::new(p.span(), INVALID_ATTR_CONTENT)), + ArgValue::Lit(ref l) => Err(syn::Error::new(l.span(), INVALID_ATTR_CONTENT)), + }; + } + + Ok(None) + } + + /// Used to parse a generator config from a set of item attributes + pub fn from_attributes(attrs: &[Attribute]) -> Result, syn::Error> { + if let Some(attr) = attrs.iter().find(|attr| attr.path().is_ident("pgrx")) { + Self::from_attribute(attr) + } else { + Ok(None) + } + } +} -use eyre::eyre; /// The output of a [`PostgresType`](crate::postgres_type::PostgresTypeDerive) from `quote::ToTokens::to_tokens`. #[derive(Debug, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)] pub struct PostgresTypeEntity { @@ -37,6 +109,7 @@ pub struct PostgresTypeEntity { pub out_fn: &'static str, pub out_fn_module_path: String, pub to_sql_config: ToSqlConfigEntity, + pub alignment: Option, } impl TypeMatch for PostgresTypeEntity { @@ -82,6 +155,7 @@ impl ToSql for PostgresTypeEntity { out_fn, out_fn_module_path, in_fn, + alignment, .. }) = item_node else { @@ -155,6 +229,23 @@ impl ToSql for PostgresTypeEntity { schema = context.schema_prefix_for(&self_index), ); + let alignment = alignment + .as_ref() + .map(|a| { + let alignment = match *a { + Alignment::Char => "char", + Alignment::Int2 => "int2", + Alignment::Int4 => "int4", + Alignment::Double => "double", + }; + format!( + ",\n\ + \tALIGNMENT = {}", + alignment + ) + }) + .unwrap_or_default(); + let materialized_type = format! { "\n\ -- {file}:{line}\n\ @@ -163,7 +254,7 @@ impl ToSql for PostgresTypeEntity { \tINTERNALLENGTH = variable,\n\ \tINPUT = {schema_prefix_in_fn}{in_fn}, /* {in_fn_path} */\n\ \tOUTPUT = {schema_prefix_out_fn}{out_fn}, /* {out_fn_path} */\n\ - \tSTORAGE = extended\n\ + \tSTORAGE = extended{alignment}\n\ );\ ", schema = context.schema_prefix_for(&self_index), diff --git a/pgrx-sql-entity-graph/src/postgres_type/mod.rs b/pgrx-sql-entity-graph/src/postgres_type/mod.rs index e563c8f5ec..fd0e9ff51a 100644 --- a/pgrx-sql-entity-graph/src/postgres_type/mod.rs +++ b/pgrx-sql-entity-graph/src/postgres_type/mod.rs @@ -23,6 +23,7 @@ use quote::{format_ident, quote}; use syn::parse::{Parse, ParseStream}; use syn::{DeriveInput, Generics, ItemStruct, Lifetime, LifetimeParam}; +pub use crate::postgres_type::entity::Alignment; use crate::{CodeEnrichment, ToSqlConfig}; /// A parsed `#[derive(PostgresType)]` item. @@ -55,6 +56,7 @@ pub struct PostgresTypeDerive { in_fn: Ident, out_fn: Ident, to_sql_config: ToSqlConfig, + alignment: Option, } impl PostgresTypeDerive { @@ -64,11 +66,12 @@ impl PostgresTypeDerive { in_fn: Ident, out_fn: Ident, to_sql_config: ToSqlConfig, + alignment: Option, ) -> Result, syn::Error> { if !to_sql_config.overrides_default() { crate::ident_is_acceptable_to_postgres(&name)?; } - Ok(CodeEnrichment(Self { generics, name, in_fn, out_fn, to_sql_config })) + Ok(CodeEnrichment(Self { generics, name, in_fn, out_fn, to_sql_config, alignment })) } pub fn from_derive_input( @@ -90,12 +93,14 @@ impl PostgresTypeDerive { &format!("{}_out", derive_input.ident).to_lowercase(), derive_input.ident.span(), ); + let alignment = Alignment::from_attributes(derive_input.attrs.as_slice())?; Self::new( derive_input.ident, derive_input.generics, funcname_in, funcname_out, to_sql_config, + alignment, ) } } @@ -129,6 +134,9 @@ impl ToEntityGraphTokens for PostgresTypeDerive { let to_sql_config = &self.to_sql_config; + let alignment = + &self.alignment.as_ref().map(|a| quote! {Some(#a)}).unwrap_or(quote! {None}); + quote! { unsafe impl #impl_generics ::pgrx::pgrx_sql_entity_graph::metadata::SqlTranslatable for #name #ty_generics #where_clauses { fn argument_sql() -> core::result::Result<::pgrx::pgrx_sql_entity_graph::metadata::SqlMapping, ::pgrx::pgrx_sql_entity_graph::metadata::ArgumentError> { @@ -190,6 +198,7 @@ impl ToEntityGraphTokens for PostgresTypeDerive { path_items.join("::") }, to_sql_config: #to_sql_config, + alignment: #alignment, }; ::pgrx::pgrx_sql_entity_graph::SqlGraphEntity::Type(submission) } @@ -205,6 +214,7 @@ impl Parse for CodeEnrichment { let to_sql_config = ToSqlConfig::from_attributes(attrs.as_slice())?.unwrap_or_default(); let in_fn = Ident::new(&format!("{}_in", ident).to_lowercase(), ident.span()); let out_fn = Ident::new(&format!("{}_out", ident).to_lowercase(), ident.span()); - PostgresTypeDerive::new(ident, generics, in_fn, out_fn, to_sql_config) + let alignment = Alignment::from_attributes(attrs.as_slice())?; + PostgresTypeDerive::new(ident, generics, in_fn, out_fn, to_sql_config, alignment) } }