Skip to content

Commit

Permalink
feat: allow specifying type alignment through #[pgrx(alignment = ...)]
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 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]
}
```
  • Loading branch information
JamesGuthrie committed Nov 12, 2024
1 parent 7b68a18 commit 84e2e70
Show file tree
Hide file tree
Showing 6 changed files with 151 additions and 4 deletions.
42 changes: 42 additions & 0 deletions pgrx-examples/custom_types/src/alignment.rs
Original file line number Diff line number Diff line change
@@ -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. <[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)]
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::<bool>(r#"SELECT typalign = 'i' FROM pg_type WHERE typname = 'alignedto4bytes'"#).unwrap().unwrap();

assert!(val == true);

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

assert!(val == true);
}
}
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>")`: 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(
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
95 changes: 93 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,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<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() {
"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<Option<Self>, 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 {
Expand All @@ -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<Alignment>,
}

impl TypeMatch for PostgresTypeEntity {
Expand Down Expand Up @@ -82,6 +155,7 @@ impl ToSql for PostgresTypeEntity {
out_fn,
out_fn_module_path,
in_fn,
alignment,
..
}) = item_node
else {
Expand Down Expand Up @@ -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\
Expand All @@ -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),
Expand Down
14 changes: 12 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: Option<Alignment>,
}

impl PostgresTypeDerive {
Expand All @@ -64,11 +66,12 @@ impl PostgresTypeDerive {
in_fn: Ident,
out_fn: Ident,
to_sql_config: ToSqlConfig,
alignment: Option<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,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> {
Expand Down Expand Up @@ -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)
}
Expand All @@ -205,6 +214,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 84e2e70

Please sign in to comment.