diff --git a/pgrx-examples/custom_types/src/alignment.rs b/pgrx-examples/custom_types/src/alignment.rs new file mode 100644 index 0000000000..b874e2d190 --- /dev/null +++ b/pgrx-examples/custom_types/src/alignment.rs @@ -0,0 +1,54 @@ +//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)] +#[pgrx(alignment = "on")] +pub struct AlignedTo4Bytes { + v1: u32, + v2: [u32; 3], +} + +#[derive(PostgresType, Serialize, Deserialize)] +#[pgrx(alignment = "on")] +pub struct AlignedTo8Bytes { + v1: u64, + v2: [u64; 3], +} + +#[derive(PostgresType, Serialize, Deserialize)] +#[pgrx(alignment = "off")] +pub struct NotAlignedTo8Bytes { + 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::text FROM pg_type WHERE typname = 'alignedto4bytes'"#).unwrap().unwrap(); + + assert!(val == "i"); + + let val = Spi::get_one::(r#"SELECT typalign::text FROM pg_type WHERE typname = 'alignedto8bytes'"#).unwrap().unwrap(); + + assert!(val == "d"); + + let val = Spi::get_one::(r#"SELECT typalign::text FROM pg_type WHERE typname = 'notalignedto8bytes'"#).unwrap().unwrap(); + + assert!(val == "i"); + } +} 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..73cafcab3b 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 = "")`: Derive Postgres alignment from Rust type. One of `"on"`, or `"off"`. * `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..a0fab3ed88 100644 --- a/pgrx-sql-entity-graph/src/postgres_type/entity.rs +++ b/pgrx-sql-entity-graph/src/postgres_type/entity.rs @@ -16,13 +16,81 @@ */ 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 { + On, + Off, +} + +const INVALID_ATTR_CONTENT: &str = r#"expected `#[pgrx(alignment = align)]`, where `align` is "on", or "off""#; + +impl ToTokens for Alignment { + fn to_tokens(&self, tokens: &mut TokenStream) { + let value = match self { + Alignment::On => format_ident!("On"), + Alignment::Off => format_ident!("Off"), + }; + 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() { + "on" => Ok(Some(Self::On)), + "off" => Ok(Some(Self::Off)), + _ => 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) + } + + pub fn from_attributes(attrs: &[Attribute]) -> Result { + for attr in attrs { + if attr.path().is_ident("pgrx") { + if let Some(v) = Self::from_attribute(attr)? { + return Ok(v) + } + } + } + Ok(Self::Off) + } +} -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 +105,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 +151,7 @@ impl ToSql for PostgresTypeEntity { out_fn, out_fn_module_path, in_fn, + alignment, .. }) = item_node else { @@ -155,6 +225,21 @@ impl ToSql for PostgresTypeEntity { schema = context.schema_prefix_for(&self_index), ); + let alignment = alignment.map(|alignment| { + assert!(alignment.is_power_of_two()); + let alignment = match alignment { + 1 => "char", + 2 => "int2", + 4 => "int4", + 8 | _ => "double", + }; + format!( + ",\n\ + \tALIGNMENT = {}", + alignment + ) + }).unwrap_or_default(); + let materialized_type = format! { "\n\ -- {file}:{line}\n\ @@ -163,7 +248,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..664e47b971 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: Alignment, } impl PostgresTypeDerive { @@ -64,11 +66,12 @@ impl PostgresTypeDerive { in_fn: Ident, out_fn: Ident, to_sql_config: ToSqlConfig, + alignment: Alignment, ) -> 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,11 @@ impl ToEntityGraphTokens for PostgresTypeDerive { let to_sql_config = &self.to_sql_config; + let alignment = match &self.alignment { + Alignment::On => quote! { Some(::std::mem::align_of::<#name>()) }, + Alignment::Off => 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 +200,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 +216,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) } }