Skip to content

Commit

Permalink
feat: auto-generate type alignment (#[pgrx(alignment = "on")])
Browse files Browse the repository at this point in the history
Postgres allows for types to specify their alignment in the `CREATE
TYPE` statement. This change adds the ability to derive Postgres'
alignment configuration from the type's Rust alignment
(`std::mem::align_of::<T>()`). This functionality is opt-in through
the `#[pgrx(alignment = "on")]` attribute:

```
#[derive(PostgresType)]
#[pgrx(alignment = "on")]
struct AlignedTo8Bytes {
  v1: u64,
  v2: [u64; 3]
}
```
  • Loading branch information
JamesGuthrie committed Nov 15, 2024
1 parent 7b68a18 commit 71b591e
Show file tree
Hide file tree
Showing 6 changed files with 159 additions and 4 deletions.
54 changes: 54 additions & 0 deletions pgrx-examples/custom_types/src/alignment.rs
Original file line number Diff line number Diff line change
@@ -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. <[email protected]>
//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::<String>(r#"SELECT typalign::text FROM pg_type WHERE typname = 'alignedto4bytes'"#).unwrap().unwrap();

assert!(val == "i");

let val = Spi::get_one::<String>(r#"SELECT typalign::text FROM pg_type WHERE typname = 'alignedto8bytes'"#).unwrap().unwrap();

assert!(val == "d");

let val = Spi::get_one::<String>(r#"SELECT typalign::text FROM pg_type WHERE typname = 'notalignedto8bytes'"#).unwrap().unwrap();

assert!(val == "i");
}
}
1 change: 1 addition & 0 deletions pgrx-examples/custom_types/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
1 change: 1 addition & 0 deletions pgrx-macros/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 = "<align>")`: Derive Postgres alignment from Rust type. One of `"on"`, or `"off"`.
* `sql`: Same arguments as [`#[pgrx(sql = ..)]`](macro@pgrx).
*/
#[proc_macro_derive(
Expand Down
2 changes: 2 additions & 0 deletions pgrx-sql-entity-graph/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,8 @@ pub trait SqlGraphIdentifier {
fn line(&self) -> Option<u32>;
}

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 {
Expand Down
89 changes: 87 additions & 2 deletions pgrx-sql-entity-graph/src/postgres_type/entity.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<Option<Self>, 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::<PgrxAttribute>()?;
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<Self, syn::Error> {
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 {
Expand All @@ -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<usize>,
}

impl TypeMatch for PostgresTypeEntity {
Expand Down Expand Up @@ -82,6 +151,7 @@ impl ToSql for PostgresTypeEntity {
out_fn,
out_fn_module_path,
in_fn,
alignment,
..
}) = item_node
else {
Expand Down Expand Up @@ -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\
Expand All @@ -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),
Expand Down
16 changes: 14 additions & 2 deletions pgrx-sql-entity-graph/src/postgres_type/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -55,6 +56,7 @@ pub struct PostgresTypeDerive {
in_fn: Ident,
out_fn: Ident,
to_sql_config: ToSqlConfig,
alignment: Alignment,
}

impl PostgresTypeDerive {
Expand All @@ -64,11 +66,12 @@ impl PostgresTypeDerive {
in_fn: Ident,
out_fn: Ident,
to_sql_config: ToSqlConfig,
alignment: Alignment,
) -> Result<CodeEnrichment<Self>, 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(
Expand All @@ -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,
)
}
}
Expand Down Expand Up @@ -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> {
Expand Down Expand Up @@ -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)
}
Expand All @@ -205,6 +216,7 @@ impl Parse for CodeEnrichment<PostgresTypeDerive> {
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)
}
}

0 comments on commit 71b591e

Please sign in to comment.