use proc_macro::TokenStream;
use quote::{format_ident, quote};
use syn::parse::Parser;
use syn::{parse_macro_input, Attribute, Field, Fields, Ident, ItemStruct, Meta, Type};
#[proc_macro_derive(
Model,
attributes(model, key, autoincrement, unique, index, has_many, belongs_to, many_to_many, dbkit)
)]
pub fn derive_model(_input: TokenStream) -> TokenStream {
TokenStream::from(quote! {
compile_error!("dbkit: use #[model] instead of #[derive(Model)]");
})
}
#[proc_macro_derive(DbEnum, attributes(dbkit))]
pub fn derive_db_enum(input: TokenStream) -> TokenStream {
let input = parse_macro_input!(input as syn::ItemEnum);
match expand_db_enum(input) {
Ok(tokens) => tokens,
Err(err) => err.to_compile_error().into(),
}
}
#[proc_macro_attribute]
pub fn model(attr: TokenStream, item: TokenStream) -> TokenStream {
let input = parse_macro_input!(item as ItemStruct);
let args = parse_macro_input!(attr with syn::punctuated::Punctuated::<Meta, syn::Token![,]>::parse_terminated);
let args = parse_model_args(args);
match expand_model(args, input) {
Ok(tokens) => tokens,
Err(err) => err.to_compile_error().into(),
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum RelationKind {
HasMany,
BelongsTo,
ManyToMany,
}
struct RelationInfo {
field: Field,
param_ident: Ident,
state_mod_ident: Ident,
child_type: Type,
kind: RelationKind,
belongs_to_key: Option<Ident>,
belongs_to_ref: Option<Ident>,
many_to_many_through: Option<Ident>,
many_to_many_left_key: Option<Ident>,
many_to_many_right_key: Option<Ident>,
}
struct ScalarFieldInfo {
field: Field,
ident: Ident,
ty: Type,
column_name: String,
is_key: bool,
is_autoincrement: bool,
}
#[derive(Default)]
struct ModelArgs {
table: Option<String>,
schema: Option<String>,
}
fn expand_model(args: ModelArgs, input: ItemStruct) -> syn::Result<TokenStream> {
if !input.generics.params.is_empty() {
return Err(syn::Error::new_spanned(
input.generics,
"dbkit: #[model] does not support generics yet",
));
}
let struct_ident = input.ident;
let model_ident = format_ident!("{}Model", struct_ident);
let insert_ident = format_ident!("{}Insert", struct_ident);
let vis = input.vis;
let table_name = args.table.unwrap_or_else(|| to_snake_case(&struct_ident.to_string()));
let schema_name = args.schema;
let mut primary_keys: Vec<(Ident, Type, String)> = Vec::new();
let mut relation_fields = Vec::new();
let mut output_fields = Vec::new();
let mut insert_fields = Vec::new();
let mut scalar_fields = Vec::new();
let struct_attrs = filter_struct_attrs(&input.attrs);
let fields = match input.fields {
Fields::Named(named) => named.named,
_ => {
return Err(syn::Error::new_spanned(
struct_ident,
"dbkit: #[model] requires a struct with named fields",
))
}
};
for field in fields {
let field_ident = field
.ident
.clone()
.ok_or_else(|| syn::Error::new_spanned(&field, "dbkit: unnamed field"))?;
let is_relation =
has_attr(&field.attrs, "has_many") || has_attr(&field.attrs, "belongs_to") || has_attr(&field.attrs, "many_to_many");
let is_key = has_attr(&field.attrs, "key");
let is_autoincrement = has_attr(&field.attrs, "autoincrement");
if is_relation {
if parse_field_column_name(&field.attrs)?.is_some() {
return Err(syn::Error::new_spanned(
&field,
"dbkit: `#[dbkit(column = \"...\")]` is only supported on scalar fields",
));
}
let (kind, child_type) = relation_type(&field)?;
let state_mod_ident = format_ident!("{}_{}_state", to_snake_case(&struct_ident.to_string()), field_ident);
let param_ident = format_ident!("{}Rel", to_camel_case(&field_ident.to_string()));
let (belongs_to_key, belongs_to_ref) = if kind == RelationKind::BelongsTo {
let (key, references) = parse_belongs_to_args(&field.attrs)?;
(Some(key), Some(references))
} else {
(None, None)
};
let (many_to_many_through, many_to_many_left_key, many_to_many_right_key) = if kind == RelationKind::ManyToMany {
let (through, left_key, right_key) = parse_many_to_many_args(&field.attrs)?;
(Some(through), Some(left_key), Some(right_key))
} else {
(None, None, None)
};
relation_fields.push(RelationInfo {
field: field.clone(),
param_ident: param_ident.clone(),
state_mod_ident,
child_type,
kind,
belongs_to_key,
belongs_to_ref,
many_to_many_through,
many_to_many_left_key,
many_to_many_right_key,
});
let cleaned_field = Field {
attrs: filter_field_attrs(&field.attrs),
ty: syn::parse_quote!(#param_ident),
..field
};
output_fields.push(cleaned_field);
continue;
}
let column_name = parse_field_column_name(&field.attrs)?.unwrap_or_else(|| field_ident.to_string());
if is_key {
primary_keys.push((field_ident.clone(), field.ty.clone(), column_name.clone()));
}
let cleaned_field = Field {
attrs: filter_field_attrs(&field.attrs),
..field.clone()
};
output_fields.push(cleaned_field.clone());
if !(is_key && is_autoincrement) {
insert_fields.push(cleaned_field.clone());
}
scalar_fields.push(ScalarFieldInfo {
field: cleaned_field,
ident: field_ident,
ty: field.ty.clone(),
column_name,
is_key,
is_autoincrement,
});
}
let table_expr = if let Some(schema) = schema_name {
quote!(::dbkit::Table::new(#table_name).with_schema(#schema))
} else {
quote!(::dbkit::Table::new(#table_name))
};
if relation_fields.iter().any(|rel| rel.kind == RelationKind::ManyToMany) && primary_keys.len() != 1 {
return Err(syn::Error::new_spanned(
struct_ident,
"dbkit: many-to-many requires exactly one #[key] on the parent model",
));
}
let generics_with_defaults = relation_fields
.iter()
.map(|rel| {
let ident = &rel.param_ident;
let state_mod = &rel.state_mod_ident;
quote!(#ident: #state_mod::State = ::dbkit::NotLoaded)
})
.collect::<Vec<_>>();
let impl_generics_params = relation_fields
.iter()
.map(|rel| {
let ident = &rel.param_ident;
let state_mod = &rel.state_mod_ident;
quote!(#ident: #state_mod::State)
})
.collect::<Vec<_>>();
let generic_idents = relation_fields.iter().map(|rel| &rel.param_ident).collect::<Vec<_>>();
let struct_generics = if generics_with_defaults.is_empty() {
quote!()
} else {
quote!(<#(#generics_with_defaults),*>)
};
let impl_generics = if impl_generics_params.is_empty() {
quote!()
} else {
quote!(<#(#impl_generics_params),*>)
};
let struct_type_args = if generic_idents.is_empty() {
quote!()
} else {
quote!(<#(#generic_idents),*>)
};
let columns = output_fields
.iter()
.filter(|field| !is_relation_field(field, &relation_fields))
.map(|field| {
let ident = field.ident.as_ref().expect("field ident");
let name = scalar_column_name(&scalar_fields, ident).expect("scalar field column name");
let ty = option_inner_type(&field.ty).unwrap_or_else(|| field.ty.clone());
quote!(pub const #ident: ::dbkit::Column<#struct_ident, #ty> = ::dbkit::Column::new(Self::TABLE, #name);)
})
.collect::<Vec<_>>();
let column_refs = output_fields
.iter()
.filter(|field| !is_relation_field(field, &relation_fields))
.map(|field| {
let ident = field.ident.as_ref().expect("field ident");
quote!(Self::#ident.as_ref())
})
.collect::<Vec<_>>();
let columns_const = quote!(
pub const COLUMNS: &'static [::dbkit::ColumnRef] = &[#(#column_refs),*];
);
let primary_key_refs = primary_keys
.iter()
.map(|(ident, _, _)| quote!(Self::#ident.as_ref()))
.collect::<Vec<_>>();
let primary_keys_const = if primary_keys.is_empty() {
quote!(
pub const PRIMARY_KEYS: &'static [::dbkit::ColumnRef] = &[];
)
} else {
quote!(pub const PRIMARY_KEYS: &'static [::dbkit::ColumnRef] = &[#(#primary_key_refs),*];)
};
let insert_values = insert_fields.iter().map(|field| {
let ident = field.ident.as_ref().expect("field ident");
quote!(insert = insert.value(Self::#ident, values.#ident);)
});
let insert_field_idents = insert_fields
.iter()
.map(|field| field.ident.as_ref().expect("field ident"))
.collect::<Vec<_>>();
let active_ident = format_ident!("{}Active", struct_ident);
let active_fields = scalar_fields.iter().map(|field| {
let ident = &field.ident;
let vis = &field.field.vis;
let ty = option_inner_type(&field.ty).unwrap_or_else(|| field.ty.clone());
quote!(#vis #ident: ::dbkit::ActiveValue<#ty>)
});
let active_from_model = scalar_fields.iter().map(|field| {
let ident = &field.ident;
if option_inner_type(&field.ty).is_some() {
quote!(#ident: ::dbkit::ActiveValue::unchanged_option(#ident))
} else {
quote!(#ident: ::dbkit::ActiveValue::unchanged(#ident))
}
});
let active_destructure = scalar_fields.iter().map(|field| field.ident.clone()).collect::<Vec<_>>();
let active_insert_steps = scalar_fields.iter().map(|field| {
let ident = &field.ident;
let name = ident.to_string();
let ty = option_inner_type(&field.ty).unwrap_or_else(|| field.ty.clone());
let is_option = option_inner_type(&field.ty).is_some();
let required = !field.is_autoincrement && !is_option;
let required_check = if required {
quote!(return Err(::dbkit::Error::Decode(format!("missing required field: {}", #name)));)
} else {
quote!()
};
quote!(
match #ident {
::dbkit::ActiveValue::Unset => {
#required_check
}
::dbkit::ActiveValue::Set(value) => {
insert = insert.value(#struct_ident::#ident, value);
}
::dbkit::ActiveValue::Unchanged(value) => {
insert = insert.value(#struct_ident::#ident, value);
}
::dbkit::ActiveValue::UnchangedNull => {
insert = insert.value(#struct_ident::#ident, None::<#ty>);
}
::dbkit::ActiveValue::Null => {
insert = insert.value(#struct_ident::#ident, None::<#ty>);
}
}
)
});
let active_insert_fn = quote!(
pub async fn insert(
self,
ex: &(impl ::dbkit::Executor + Send + Sync),
) -> Result<#struct_ident, ::dbkit::Error> {
let Self { #(#active_destructure,)* } = self;
let mut insert = ::dbkit::Insert::new(#struct_ident::TABLE);
#(#active_insert_steps)*
let insert = insert.returning_all();
let row = ::dbkit::InsertExt::one(insert, ex).await?;
row.ok_or(::dbkit::Error::NotFound)
}
);
let pk_idents = primary_keys.iter().map(|(ident, _, _)| ident.clone()).collect::<Vec<_>>();
let active_update_fn = if !primary_keys.is_empty() {
let pk_vars = primary_keys
.iter()
.enumerate()
.map(|(idx, _)| format_ident!("pk_value_{}", idx))
.collect::<Vec<_>>();
let pk_extracts = primary_keys.iter().zip(pk_vars.iter()).map(|((ident, _, _), var)| {
let pk_name = ident.to_string();
quote!(
let #var = match #ident {
::dbkit::ActiveValue::Set(value) | ::dbkit::ActiveValue::Unchanged(value) => value,
::dbkit::ActiveValue::Null | ::dbkit::ActiveValue::Unset | ::dbkit::ActiveValue::UnchangedNull => {
return Err(::dbkit::Error::Decode(format!(
"missing required field: {}",
#pk_name
)));
}
};
)
});
let pk_filters = primary_keys
.iter()
.zip(pk_vars.iter())
.map(|((ident, _, _), var)| quote!(update = update.filter(#struct_ident::#ident.eq(#var));));
let update_steps = scalar_fields.iter().filter(|field| !field.is_key).map(|field| {
let ident = &field.ident;
let ty = option_inner_type(&field.ty).unwrap_or_else(|| field.ty.clone());
quote!(
match #ident {
::dbkit::ActiveValue::Unset => {}
::dbkit::ActiveValue::Set(value) => {
update = update.set(#struct_ident::#ident, value);
any_set = true;
}
::dbkit::ActiveValue::Unchanged(_) | ::dbkit::ActiveValue::UnchangedNull => {}
::dbkit::ActiveValue::Null => {
update = update.set(#struct_ident::#ident, None::<#ty>);
any_set = true;
}
}
)
});
quote!(
pub async fn update(
self,
ex: &(impl ::dbkit::Executor + Send + Sync),
) -> Result<#struct_ident, ::dbkit::Error> {
let Self { #(#active_destructure,)* } = self;
#(#pk_extracts)*
let mut update = ::dbkit::Update::new(#struct_ident::TABLE);
let mut any_set = false;
#(#update_steps)*
if !any_set {
return Err(::dbkit::Error::Decode("no fields set for update".to_string()));
}
#(#pk_filters)*
let update = update.returning_all();
let mut rows = ::dbkit::UpdateExt::all(update, ex).await?;
rows.pop().ok_or(::dbkit::Error::NotFound)
}
)
} else {
quote!()
};
let active_delete_fn = if !primary_keys.is_empty() {
let pk_vars = primary_keys
.iter()
.enumerate()
.map(|(idx, _)| format_ident!("pk_value_{}", idx))
.collect::<Vec<_>>();
let pk_extracts = primary_keys.iter().zip(pk_vars.iter()).map(|((ident, _, _), var)| {
let pk_name = ident.to_string();
quote!(
let #var = match #ident {
::dbkit::ActiveValue::Set(value) | ::dbkit::ActiveValue::Unchanged(value) => value,
::dbkit::ActiveValue::Null | ::dbkit::ActiveValue::Unset | ::dbkit::ActiveValue::UnchangedNull => {
return Err(::dbkit::Error::Decode(format!(
"missing required field: {}",
#pk_name
)));
}
};
)
});
let pk_filters = primary_keys
.iter()
.zip(pk_vars.iter())
.map(|((ident, _, _), var)| quote!(delete = delete.filter(#struct_ident::#ident.eq(#var));));
quote!(
pub async fn delete(
self,
ex: &(impl ::dbkit::Executor + Send + Sync),
) -> Result<u64, ::dbkit::Error> {
let Self { #(#pk_idents,)* .. } = self;
#(#pk_extracts)*
let mut delete = ::dbkit::Delete::new(#struct_ident::TABLE);
#(#pk_filters)*
::dbkit::DeleteExt::execute(delete, ex).await
}
)
} else {
quote!()
};
let active_save_flag_checks = scalar_fields.iter().map(|field| {
let ident = &field.ident;
quote!(
match &#ident {
::dbkit::ActiveValue::Unchanged(_) | ::dbkit::ActiveValue::UnchangedNull => {
any_loaded = true;
}
::dbkit::ActiveValue::Set(_) | ::dbkit::ActiveValue::Null => {
any_changed = true;
}
::dbkit::ActiveValue::Unset => {}
}
)
});
let active_save_model_fields = scalar_fields.iter().map(|field| {
let ident = &field.ident;
let name = ident.to_string();
if option_inner_type(&field.ty).is_some() {
quote!(
#ident: match #ident {
::dbkit::ActiveValue::Set(value) | ::dbkit::ActiveValue::Unchanged(value) => Some(value),
::dbkit::ActiveValue::Null | ::dbkit::ActiveValue::UnchangedNull => None,
::dbkit::ActiveValue::Unset => {
return Err(::dbkit::Error::Decode(format!(
"missing required field: {}",
#name
)));
}
},
)
} else {
quote!(
#ident: match #ident {
::dbkit::ActiveValue::Set(value) | ::dbkit::ActiveValue::Unchanged(value) => value,
::dbkit::ActiveValue::Null
| ::dbkit::ActiveValue::Unset
| ::dbkit::ActiveValue::UnchangedNull => {
return Err(::dbkit::Error::Decode(format!(
"missing required field: {}",
#name
)));
}
},
)
}
});
let active_save_relation_defaults = relation_fields.iter().map(|rel| {
let ident = rel.field.ident.as_ref().expect("field ident");
quote!(#ident: Default::default(),)
});
let active_save_update_branch = if !primary_keys.is_empty() {
quote!(return Self { #(#active_destructure,)* }.update(ex).await;)
} else {
quote!(
return Err(::dbkit::Error::Decode(
"update requires primary key".to_string(),
));
)
};
let active_save_fn = quote!(
pub async fn save(
self,
ex: &(impl ::dbkit::Executor + Send + Sync),
) -> Result<#struct_ident, ::dbkit::Error> {
let Self { #(#active_destructure,)* } = self;
let mut any_loaded = false;
let mut any_changed = false;
#(#active_save_flag_checks)*
if any_loaded {
if any_changed {
#active_save_update_branch
}
let model = #struct_ident {
#(#active_save_model_fields)*
#(#active_save_relation_defaults)*
};
return Ok(model);
}
Self { #(#active_destructure,)* }.insert(ex).await
}
);
let model_delete_impl = if !primary_keys.is_empty() {
let pk_filters = primary_keys
.iter()
.map(|(ident, _, _)| quote!(delete = delete.filter(Self::#ident.eq(#ident));));
quote!(
impl #impl_generics ::dbkit::ModelDelete for #model_ident #struct_type_args {
fn delete<'e, E>(self, ex: &'e E) -> ::dbkit::executor::BoxFuture<'e, Result<u64, ::dbkit::Error>>
where
E: ::dbkit::Executor + Send + Sync + 'e,
{
let Self { #(#pk_idents,)* .. } = self;
let mut delete = ::dbkit::Delete::new(Self::TABLE);
#(#pk_filters)*
::dbkit::DeleteExt::execute(delete, ex)
}
}
)
} else {
quote!()
};
let into_active_fn = quote!(
pub fn into_active(self) -> #active_ident {
let Self { #(#active_destructure,)* .. } = self;
#active_ident {
#(#active_from_model,)*
}
}
);
let primary_key_const = if primary_keys.len() == 1 {
let (_, ty, name) = primary_keys.first().expect("primary key length checked");
Some(quote!(pub const PRIMARY_KEY: ::dbkit::Column<#struct_ident, #ty> = ::dbkit::Column::new(Self::TABLE, #name);))
} else {
None
};
let by_id_fn = if primary_keys.len() == 1 {
let (ident, ty, _) = primary_keys.first().expect("primary key length checked");
Some(quote!(
pub fn by_id(id: #ty) -> ::dbkit::Select<#struct_ident> {
Self::query().filter(Self::#ident.eq(id)).limit(1)
}
))
} else {
None
};
let any_state_ident = format_ident!("{}AnyState", struct_ident);
let relation_state_modules = relation_fields.iter().map(|rel| {
let state_mod = &rel.state_mod_ident;
let (sealed_impl, state_impl) = match rel.kind {
RelationKind::HasMany | RelationKind::ManyToMany => (
quote!(
impl<T> Sealed for Vec<T> {}
),
quote!(
impl<T> State for Vec<T> {}
),
),
RelationKind::BelongsTo => (
quote!(
impl<T> Sealed for Option<T> {}
),
quote!(
impl<T> State for Option<T> {}
),
),
};
quote!(
pub mod #state_mod {
mod sealed {
pub trait Sealed {}
impl Sealed for ::dbkit::NotLoaded {}
#sealed_impl
}
pub trait State: sealed::Sealed {}
impl State for ::dbkit::NotLoaded {}
#state_impl
}
)
});
let relation_methods = relation_fields.iter().map(|rel| {
let field_ident = rel.field.ident.as_ref().expect("field ident");
let method_ident = format_ident!("{}_loaded", field_ident);
let item_ident = format_ident!("{}Item", to_camel_case(&field_ident.to_string()));
let loaded_type: Type = match rel.kind {
RelationKind::HasMany | RelationKind::ManyToMany => syn::parse_quote!(Vec<#item_ident>),
RelationKind::BelongsTo => syn::parse_quote!(Option<#item_ident>),
};
let mut other_params = Vec::new();
let mut type_params = Vec::new();
for other in &relation_fields {
if other.field.ident == rel.field.ident {
type_params.push(quote!(#loaded_type));
} else {
let ident = &other.param_ident;
let state_mod = &other.state_mod_ident;
other_params.push(quote!(#ident: #state_mod::State));
type_params.push(quote!(#ident));
}
}
let mut impl_params = Vec::new();
impl_params.push(quote!(#item_ident));
impl_params.extend(other_params);
let impl_generics = if impl_params.is_empty() {
quote!()
} else {
quote!(<#(#impl_params),*>)
};
let type_args = if type_params.is_empty() {
quote!()
} else {
quote!(<#(#type_params),*>)
};
let (return_ty, body) = match rel.kind {
RelationKind::HasMany | RelationKind::ManyToMany => (quote!(&[#item_ident]), quote!(&self.#field_ident)),
RelationKind::BelongsTo => (quote!(Option<&#item_ident>), quote!(self.#field_ident.as_ref())),
};
quote!(
impl #impl_generics #model_ident #type_args {
pub fn #method_ident(&self) -> #return_ty {
#body
}
}
)
});
let model_value_arms = output_fields
.iter()
.filter(|field| !is_relation_field(field, &relation_fields))
.map(|field| {
let ident = field.ident.as_ref().expect("field ident");
let name = scalar_column_name(&scalar_fields, ident).expect("scalar field column name");
quote!(#name => Some(self.#ident.clone().into()),)
});
let model_value_impl = quote!(
impl #impl_generics ::dbkit::ModelValue for #model_ident #struct_type_args {
fn column_value(&self, column: ::dbkit::ColumnRef) -> Option<::dbkit::Value> {
if column.table.name != Self::TABLE.name {
return None;
}
match column.name {
#(#model_value_arms)*
_ => None,
}
}
}
);
let from_row_generics = relation_fields.iter().map(|rel| {
let ident = &rel.param_ident;
let state_mod = &rel.state_mod_ident;
quote!(#ident: #state_mod::State + Default)
});
let from_row_impl_generics = if relation_fields.is_empty() {
quote!(<'r>)
} else {
quote!(<'r, #(#from_row_generics),*>)
};
let from_row_fields = output_fields.iter().map(|field| {
let ident = field.ident.as_ref().expect("field ident");
if is_relation_field(field, &relation_fields) {
quote!(#ident: Default::default())
} else {
let name = scalar_column_name(&scalar_fields, ident).expect("scalar field column name");
quote!(#ident: ::dbkit::sqlx::Row::try_get(row, #name)?)
}
});
let from_row_impl = quote!(
impl #from_row_impl_generics ::dbkit::sqlx::FromRow<'r, ::dbkit::sqlx::postgres::PgRow>
for #model_ident #struct_type_args
{
fn from_row(row: &'r ::dbkit::sqlx::postgres::PgRow) -> Result<Self, ::dbkit::sqlx::Error> {
Ok(Self {
#(#from_row_fields,)*
})
}
}
);
let joined_from_row_fields = output_fields.iter().map(|field| {
let ident = field.ident.as_ref().expect("field ident");
if is_relation_field(field, &relation_fields) {
quote!(#ident: Default::default())
} else {
let name = scalar_column_name(&scalar_fields, ident).expect("scalar field column name");
quote!(
#ident: {
let column = format!("{}{}", prefix, #name);
::dbkit::sqlx::Row::try_get(row, column.as_str())?
}
)
}
});
let joined_pk_checks = if primary_keys.is_empty() {
if let Some(first_field) = scalar_fields.first() {
let name = &first_field.column_name;
let ty = option_inner_type(&first_field.ty).unwrap_or_else(|| first_field.ty.clone());
quote!(
let value: Option<#ty> = {
let column = format!("{}{}", prefix, #name);
::dbkit::sqlx::Row::try_get(row, column.as_str())?
};
Ok(value.is_some())
)
} else {
quote!(Ok(false))
}
} else {
let checks = primary_keys.iter().map(|(_, ty, name)| {
let ty = option_inner_type(ty).unwrap_or_else(|| ty.clone());
quote!(
let value: Option<#ty> = {
let column = format!("{}{}", prefix, #name);
::dbkit::sqlx::Row::try_get(row, column.as_str())?
};
if value.is_some() {
return Ok(true);
}
)
});
quote!(
#(#checks)*
Ok(false)
)
};
let joined_model_impl = quote!(
impl #from_row_impl_generics ::dbkit::JoinedModel for #model_ident #struct_type_args {
fn joined_columns() -> &'static [::dbkit::ColumnRef] {
Self::COLUMNS
}
fn joined_primary_keys() -> &'static [::dbkit::ColumnRef] {
Self::PRIMARY_KEYS
}
fn joined_from_row_prefixed(
row: &::dbkit::sqlx::postgres::PgRow,
prefix: &str,
) -> Result<Self, ::dbkit::sqlx::Error> {
Ok(Self {
#(#joined_from_row_fields,)*
})
}
fn joined_row_has_pk(
row: &::dbkit::sqlx::postgres::PgRow,
prefix: &str,
) -> Result<bool, ::dbkit::sqlx::Error> {
#joined_pk_checks
}
}
);
let set_relation_impls = relation_fields.iter().map(|rel| {
let field_ident = rel.field.ident.as_ref().expect("field ident");
let child_type = &rel.child_type;
let item_ident = format_ident!("{}Item", to_camel_case(&field_ident.to_string()));
let (value_ty, rel_ty) = match rel.kind {
RelationKind::HasMany => (quote!(Vec<#item_ident>), quote!(::dbkit::rel::HasMany<#struct_ident, #child_type>)),
RelationKind::ManyToMany => {
let through = rel.many_to_many_through.as_ref().expect("many-to-many through");
(
quote!(Vec<#item_ident>),
quote!(::dbkit::rel::ManyToMany<#struct_ident, #child_type, #through>),
)
}
RelationKind::BelongsTo => (
quote!(Option<#item_ident>),
quote!(::dbkit::rel::BelongsTo<#struct_ident, #child_type>),
),
};
let mut other_params = Vec::new();
let mut type_params = Vec::new();
for other in &relation_fields {
if other.field.ident == rel.field.ident {
type_params.push(value_ty.clone());
} else {
let ident = &other.param_ident;
let state_mod = &other.state_mod_ident;
other_params.push(quote!(#ident: #state_mod::State));
type_params.push(quote!(#ident));
}
}
let mut impl_params = Vec::new();
impl_params.push(quote!(#item_ident));
impl_params.extend(other_params);
let impl_generics = if impl_params.is_empty() {
quote!()
} else {
quote!(<#(#impl_params),*>)
};
let type_args = if type_params.is_empty() {
quote!()
} else {
quote!(<#(#type_params),*>)
};
quote!(
impl #impl_generics ::dbkit::SetRelation<#rel_ty, #value_ty> for #model_ident #type_args {
fn set_relation(&mut self, _rel: #rel_ty, value: #value_ty) -> Result<(), ::dbkit::Error> {
self.#field_ident = value;
Ok(())
}
}
)
});
let get_relation_impls = relation_fields.iter().map(|rel| {
let field_ident = rel.field.ident.as_ref().expect("field ident");
let child_type = &rel.child_type;
let item_ident = format_ident!("{}Item", to_camel_case(&field_ident.to_string()));
let (value_ty, rel_ty) = match rel.kind {
RelationKind::HasMany => (quote!(Vec<#item_ident>), quote!(::dbkit::rel::HasMany<#struct_ident, #child_type>)),
RelationKind::ManyToMany => {
let through = rel.many_to_many_through.as_ref().expect("many-to-many through");
(
quote!(Vec<#item_ident>),
quote!(::dbkit::rel::ManyToMany<#struct_ident, #child_type, #through>),
)
}
RelationKind::BelongsTo => (
quote!(Option<#item_ident>),
quote!(::dbkit::rel::BelongsTo<#struct_ident, #child_type>),
),
};
let mut other_params = Vec::new();
let mut type_params = Vec::new();
for other in &relation_fields {
if other.field.ident == rel.field.ident {
type_params.push(value_ty.clone());
} else {
let ident = &other.param_ident;
let state_mod = &other.state_mod_ident;
other_params.push(quote!(#ident: #state_mod::State));
type_params.push(quote!(#ident));
}
}
let mut impl_params = Vec::new();
impl_params.push(quote!(#item_ident));
impl_params.extend(other_params);
let impl_generics = if impl_params.is_empty() {
quote!()
} else {
quote!(<#(#impl_params),*>)
};
let type_args = if type_params.is_empty() {
quote!()
} else {
quote!(<#(#type_params),*>)
};
quote!(
impl #impl_generics ::dbkit::GetRelation<#rel_ty, #value_ty> for #model_ident #type_args {
fn get_relation(&self, _rel: #rel_ty) -> Option<&#value_ty> {
Some(&self.#field_ident)
}
fn get_relation_mut(&mut self, _rel: #rel_ty) -> Option<&mut #value_ty> {
Some(&mut self.#field_ident)
}
}
)
});
let load_method = quote!(
pub async fn load<Rel>(
self,
rel: Rel,
ex: &(impl ::dbkit::Executor + Send + Sync),
) -> Result<<Self as ::dbkit::LoadRelation<Rel>>::Out, ::dbkit::Error>
where
Self: ::dbkit::LoadRelation<Rel>,
{
::dbkit::LoadRelation::load_relation(self, rel, ex).await
}
);
let load_relation_impls = relation_fields.iter().map(|rel| {
let field_ident = rel.field.ident.as_ref().expect("field ident");
let child_type = &rel.child_type;
let rel_type = match rel.kind {
RelationKind::HasMany => quote!(::dbkit::rel::HasMany<#struct_ident, #child_type>),
RelationKind::BelongsTo => quote!(::dbkit::rel::BelongsTo<#struct_ident, #child_type>),
RelationKind::ManyToMany => {
let through = rel.many_to_many_through.as_ref().expect("many-to-many through");
quote!(::dbkit::rel::ManyToMany<#struct_ident, #child_type, #through>)
}
};
let loaded_type = match rel.kind {
RelationKind::HasMany | RelationKind::ManyToMany => quote!(Vec<#child_type>),
RelationKind::BelongsTo => quote!(Option<#child_type>),
};
let loader_fn = match rel.kind {
RelationKind::HasMany => quote!(::dbkit::runtime::load_selectin_has_many),
RelationKind::ManyToMany => quote!(::dbkit::runtime::load_selectin_many_to_many),
RelationKind::BelongsTo => quote!(::dbkit::runtime::load_selectin_belongs_to),
};
let mut other_params = Vec::new();
let mut type_params = Vec::new();
let mut out_params = Vec::new();
for other in &relation_fields {
if other.field.ident == rel.field.ident {
type_params.push(quote!(::dbkit::NotLoaded));
out_params.push(loaded_type.clone());
} else {
let ident = &other.param_ident;
let state_mod = &other.state_mod_ident;
other_params.push(quote!(#ident: #state_mod::State + Send + 'static));
type_params.push(quote!(#ident));
out_params.push(quote!(#ident));
}
}
let impl_generics = if other_params.is_empty() {
quote!()
} else {
quote!(<#(#other_params),*>)
};
let type_args = if type_params.is_empty() {
quote!()
} else {
quote!(<#(#type_params),*>)
};
let out_type = if out_params.is_empty() {
quote!(#model_ident)
} else {
quote!(#model_ident<#(#out_params),*>)
};
let out_construct = if out_params.is_empty() {
quote!(#model_ident)
} else {
quote!(#model_ident::<#(#out_params),*>)
};
let destructure_fields = output_fields.iter().map(|field| {
let ident = field.ident.as_ref().expect("field ident");
if ident == field_ident {
quote!(#ident: _)
} else {
quote!(#ident)
}
});
let build_fields = output_fields.iter().map(|field| {
let ident = field.ident.as_ref().expect("field ident");
if ident == field_ident {
quote!(#ident: Default::default())
} else {
quote!(#ident)
}
});
quote!(
impl #impl_generics ::dbkit::LoadRelation<#rel_type> for #model_ident #type_args {
type Out = #out_type;
fn load_relation<'e, E>(
self,
rel: #rel_type,
ex: &'e E,
) -> ::dbkit::executor::BoxFuture<'e, Result<Self::Out, ::dbkit::Error>>
where
E: ::dbkit::Executor + Send + Sync + 'e,
{
Box::pin(async move {
let Self { #(#destructure_fields,)* } = self;
let mut out = #out_construct {
#(#build_fields,)*
};
let mut rows = vec![out];
#loader_fn(ex, &mut rows, rel, &::dbkit::load::NoLoad).await?;
Ok(rows.pop().expect("loaded row"))
})
}
}
)
});
let relation_consts = relation_fields.iter().filter_map(|rel| {
let field_ident = rel.field.ident.as_ref().expect("field ident");
let child_type = &rel.child_type;
match rel.kind {
RelationKind::HasMany => Some(quote!(
pub const #field_ident: ::dbkit::rel::HasMany<#struct_ident, #child_type> =
::dbkit::rel::HasMany::new(
<#child_type as ::dbkit::rel::BelongsToSpec<#struct_ident>>::PARENT_TABLE,
<#child_type as ::dbkit::rel::BelongsToSpec<#struct_ident>>::CHILD_TABLE,
<#child_type as ::dbkit::rel::BelongsToSpec<#struct_ident>>::PARENT_KEY,
<#child_type as ::dbkit::rel::BelongsToSpec<#struct_ident>>::CHILD_KEY,
);
)),
RelationKind::BelongsTo => {
let key = rel.belongs_to_key.as_ref().expect("belongs_to key");
let references = rel.belongs_to_ref.as_ref().expect("belongs_to references");
Some(quote!(
pub const #field_ident: ::dbkit::rel::BelongsTo<#struct_ident, #child_type> =
::dbkit::rel::BelongsTo::new(
Self::TABLE,
#child_type::TABLE,
Self::#key.as_ref(),
#child_type::#references.as_ref(),
);
))
}
RelationKind::ManyToMany => {
let through = rel.many_to_many_through.as_ref().expect("many-to-many through");
let left_key = rel.many_to_many_left_key.as_ref().expect("many-to-many left_key");
let right_key = rel.many_to_many_right_key.as_ref().expect("many-to-many right_key");
let parent_pk = primary_keys.first().map(|(ident, _, _)| ident).expect("many-to-many parent pk");
Some(quote!(
pub const #field_ident: ::dbkit::rel::ManyToMany<#struct_ident, #child_type, #through> =
::dbkit::rel::ManyToMany::new(
Self::TABLE,
#child_type::TABLE,
#through::TABLE,
Self::#parent_pk.as_ref(),
#child_type::PRIMARY_KEY.as_ref(),
#through::#left_key.as_ref(),
#through::#right_key.as_ref(),
);
))
}
}
});
let belongs_to_specs = relation_fields.iter().filter_map(|rel| {
if rel.kind != RelationKind::BelongsTo {
return None;
}
let parent_type = &rel.child_type;
let key = rel.belongs_to_key.as_ref().expect("belongs_to key");
let references = rel.belongs_to_ref.as_ref().expect("belongs_to references");
Some(quote!(
impl #impl_generics ::dbkit::rel::BelongsToSpec<#parent_type> for #model_ident #struct_type_args {
const CHILD_TABLE: ::dbkit::Table = Self::TABLE;
const PARENT_TABLE: ::dbkit::Table = #parent_type::TABLE;
const CHILD_KEY: ::dbkit::ColumnRef = Self::#key.as_ref();
const PARENT_KEY: ::dbkit::ColumnRef = #parent_type::#references.as_ref();
}
))
});
let apply_load_impls = relation_fields.iter().flat_map(|rel| {
let child_type = &rel.child_type;
let rel_type = match rel.kind {
RelationKind::HasMany => quote!(::dbkit::rel::HasMany<#struct_ident, #child_type>),
RelationKind::BelongsTo => quote!(::dbkit::rel::BelongsTo<#struct_ident, #child_type>),
RelationKind::ManyToMany => {
let through = rel.many_to_many_through.as_ref().expect("many-to-many through");
quote!(::dbkit::rel::ManyToMany<#struct_ident, #child_type, #through>)
}
};
let loaded_child = quote!(<Nested as ::dbkit::load::ApplyLoad<#child_type>>::Out2);
let loaded_param = match rel.kind {
RelationKind::HasMany | RelationKind::ManyToMany => quote!(Vec<#loaded_child>),
RelationKind::BelongsTo => quote!(Option<#loaded_child>),
};
let mut out_params = Vec::new();
for other in &relation_fields {
if other.field.ident == rel.field.ident {
out_params.push(loaded_param.clone());
} else {
let ident = &other.param_ident;
out_params.push(quote!(#ident));
}
}
let model_type = if generic_idents.is_empty() {
quote!(#model_ident)
} else {
quote!(#model_ident<#(#generic_idents),*>)
};
let out_type = if out_params.is_empty() {
quote!(#model_ident)
} else {
quote!(#model_ident<#(#out_params),*>)
};
let mut apply_generics = Vec::new();
apply_generics.push(quote!(Nested));
apply_generics.extend(impl_generics_params.iter().cloned());
let apply_generics = if apply_generics.is_empty() {
quote!()
} else {
quote!(<#(#apply_generics),*>)
};
let mut items = Vec::new();
for strategy in ["SelectIn", "Joined"] {
let load_ty = if strategy == "SelectIn" {
quote!(::dbkit::load::SelectIn<#rel_type, Nested>)
} else {
quote!(::dbkit::load::Joined<#rel_type, Nested>)
};
items.push(quote!(
impl #apply_generics ::dbkit::load::ApplyLoad<#model_type> for #load_ty
where
Nested: ::dbkit::load::ApplyLoad<#child_type>,
{
type Out2 = #out_type;
}
));
}
items.into_iter()
});
let run_load_impls = relation_fields.iter().flat_map(|rel| {
let child_type = &rel.child_type;
let through = rel.many_to_many_through.as_ref();
let rel_type = match rel.kind {
RelationKind::HasMany => quote!(::dbkit::rel::HasMany<#struct_ident, #child_type>),
RelationKind::BelongsTo => quote!(::dbkit::rel::BelongsTo<#struct_ident, #child_type>),
RelationKind::ManyToMany => {
let through = through.expect("many-to-many through");
quote!(::dbkit::rel::ManyToMany<#struct_ident, #child_type, #through>)
}
};
let loaded_child = quote!(<Nested as ::dbkit::load::ApplyLoad<#child_type>>::Out2);
let loaded_param = match rel.kind {
RelationKind::HasMany | RelationKind::ManyToMany => quote!(Vec<#loaded_child>),
RelationKind::BelongsTo => quote!(Option<#loaded_child>),
};
let mut out_params = Vec::new();
for other in &relation_fields {
if other.field.ident == rel.field.ident {
out_params.push(loaded_param.clone());
} else {
let ident = &other.param_ident;
out_params.push(quote!(#ident));
}
}
let out_type = if out_params.is_empty() {
quote!(#model_ident)
} else {
quote!(#model_ident<#(#out_params),*>)
};
let mut apply_generics = Vec::new();
apply_generics.push(quote!(Nested));
for other in &relation_fields {
if other.field.ident == rel.field.ident {
continue;
}
let ident = &other.param_ident;
let state_mod = &other.state_mod_ident;
apply_generics.push(quote!(#ident: #state_mod::State + Send + 'static));
}
let apply_generics = if apply_generics.is_empty() {
quote!()
} else {
quote!(<#(#apply_generics),*>)
};
let (child_bounds, loader_fn) = match rel.kind {
RelationKind::HasMany => (
quote!(#loaded_child: ::dbkit::ModelValue + for<'r> ::dbkit::sqlx::FromRow<'r, ::dbkit::sqlx::postgres::PgRow> + Send + Unpin,),
quote!(::dbkit::runtime::load_selectin_has_many),
),
RelationKind::ManyToMany => {
let through = through.expect("many-to-many through");
(
quote!(
#loaded_child: ::dbkit::ModelValue + Clone + for<'r> ::dbkit::sqlx::FromRow<'r, ::dbkit::sqlx::postgres::PgRow> + Send + Unpin,
#through: ::dbkit::ModelValue + for<'r> ::dbkit::sqlx::FromRow<'r, ::dbkit::sqlx::postgres::PgRow> + Send + Unpin,
),
quote!(::dbkit::runtime::load_selectin_many_to_many),
)
}
RelationKind::BelongsTo => (
quote!(#loaded_child: ::dbkit::ModelValue + Clone + for<'r> ::dbkit::sqlx::FromRow<'r, ::dbkit::sqlx::postgres::PgRow> + Send + Unpin,),
quote!(::dbkit::runtime::load_selectin_belongs_to),
),
};
let joined_loader_fn = match rel.kind {
RelationKind::HasMany => quote!(::dbkit::runtime::load_joined_has_many),
RelationKind::ManyToMany => quote!(::dbkit::runtime::load_joined_many_to_many),
RelationKind::BelongsTo => quote!(::dbkit::runtime::load_joined_belongs_to),
};
let mut items = Vec::new();
for (strategy, loader) in [
("SelectIn", loader_fn),
("Joined", joined_loader_fn),
] {
let load_ty = if strategy == "SelectIn" {
quote!(::dbkit::load::SelectIn<#rel_type, Nested>)
} else {
quote!(::dbkit::load::Joined<#rel_type, Nested>)
};
let out_bound = if strategy == "SelectIn" {
quote!(::dbkit::ModelValue + ::dbkit::SetRelation<#rel_type, #loaded_param>)
} else {
quote!(::dbkit::GetRelation<#rel_type, #loaded_param>)
};
items.push(quote!(
impl #apply_generics ::dbkit::runtime::RunLoad<#out_type> for #load_ty
where
Nested: ::dbkit::load::ApplyLoad<#child_type> + ::dbkit::runtime::RunLoads<#loaded_child> + Sync,
#out_type: #out_bound,
#child_bounds
{
fn run<'e, E>(
&'e self,
ex: &'e E,
rows: &'e mut [#out_type],
) -> ::dbkit::executor::BoxFuture<'e, Result<(), ::dbkit::Error>>
where
E: ::dbkit::Executor + Send + Sync + 'e,
{
#loader(ex, rows, self.rel.clone(), &self.nested)
}
}
));
}
items.into_iter()
});
let output = quote! {
#(#struct_attrs)*
#[derive(Debug, Clone)]
#vis struct #model_ident #struct_generics {
#(#output_fields,)*
}
#vis type #struct_ident = #model_ident;
#(#relation_state_modules)*
#vis trait #any_state_ident {}
impl #impl_generics #any_state_ident for #model_ident #struct_type_args {}
impl #impl_generics #model_ident #struct_type_args {
pub const TABLE: ::dbkit::Table = #table_expr;
#(#columns)*
#columns_const
#primary_key_const
#primary_keys_const
#(#relation_consts)*
pub fn query() -> ::dbkit::Select<#struct_ident> {
::dbkit::Select::new(Self::TABLE)
}
#by_id_fn
pub fn insert(values: #insert_ident) -> ::dbkit::Insert<#struct_ident> {
let mut insert = ::dbkit::Insert::new(Self::TABLE);
#(#insert_values)*
insert
}
pub fn insert_many(values: Vec<#insert_ident>) -> ::dbkit::Insert<#struct_ident> {
let mut insert = ::dbkit::Insert::new(Self::TABLE);
for value in values {
insert = insert.row(|row| {
let mut row = row;
#(
row = row.value(Self::#insert_field_idents, value.#insert_field_idents);
)*
row
});
}
insert
}
pub fn update() -> ::dbkit::Update<#struct_ident> {
::dbkit::Update::new(Self::TABLE)
}
pub fn delete() -> ::dbkit::Delete {
::dbkit::Delete::new(Self::TABLE)
}
pub fn new_active() -> #active_ident {
#active_ident::new()
}
#into_active_fn
#load_method
}
#[derive(Debug, Clone)]
#vis struct #insert_ident {
#(#insert_fields,)*
}
#[derive(Debug, Clone, Default)]
#vis struct #active_ident {
#(#active_fields,)*
}
impl #active_ident {
pub fn new() -> Self {
Self::default()
}
#active_insert_fn
#active_update_fn
#active_delete_fn
#active_save_fn
}
#(#relation_methods)*
#model_value_impl
#from_row_impl
#joined_model_impl
#(#set_relation_impls)*
#(#get_relation_impls)*
#(#load_relation_impls)*
#(#belongs_to_specs)*
#(#apply_load_impls)*
#(#run_load_impls)*
#model_delete_impl
};
Ok(output.into())
}
fn parse_model_args(args: syn::punctuated::Punctuated<Meta, syn::Token![,]>) -> ModelArgs {
let mut out = ModelArgs::default();
for meta in args {
if let Meta::NameValue(nv) = meta {
if nv.path.is_ident("table") {
if let Some(value) = extract_lit_str(&nv.value) {
out.table = Some(value);
}
} else if nv.path.is_ident("schema") {
if let Some(value) = extract_lit_str(&nv.value) {
out.schema = Some(value);
}
}
}
}
out
}
fn parse_belongs_to_args(attrs: &[Attribute]) -> syn::Result<(Ident, Ident)> {
for attr in attrs {
if !attr.path().is_ident("belongs_to") {
continue;
}
let args = attr.parse_args_with(syn::punctuated::Punctuated::<Meta, syn::Token![,]>::parse_terminated)?;
let mut key = None;
let mut references = None;
for meta in args {
if let Meta::NameValue(nv) = meta {
if nv.path.is_ident("key") {
key = extract_ident(&nv.value);
} else if nv.path.is_ident("references") {
references = extract_ident(&nv.value);
}
}
}
if let (Some(key), Some(references)) = (key, references) {
return Ok((key, references));
}
}
Err(syn::Error::new(
proc_macro2::Span::call_site(),
"dbkit: #[belongs_to] requires key = <field> and references = <field>",
))
}
fn parse_many_to_many_args(attrs: &[Attribute]) -> syn::Result<(Ident, Ident, Ident)> {
for attr in attrs {
if !attr.path().is_ident("many_to_many") {
continue;
}
let args = attr.parse_args_with(syn::punctuated::Punctuated::<Meta, syn::Token![,]>::parse_terminated)?;
let mut through = None;
let mut left_key = None;
let mut right_key = None;
for meta in args {
if let Meta::NameValue(nv) = meta {
if nv.path.is_ident("through") {
through = extract_ident(&nv.value);
} else if nv.path.is_ident("left_key") {
left_key = extract_ident(&nv.value);
} else if nv.path.is_ident("right_key") {
right_key = extract_ident(&nv.value);
}
}
}
if let (Some(through), Some(left_key), Some(right_key)) = (through, left_key, right_key) {
return Ok((through, left_key, right_key));
}
}
Err(syn::Error::new(
proc_macro2::Span::call_site(),
"dbkit: #[many_to_many] requires through = <Model>, left_key = <field>, right_key = <field>",
))
}
fn extract_lit_str(expr: &syn::Expr) -> Option<String> {
if let syn::Expr::Lit(syn::ExprLit {
lit: syn::Lit::Str(lit), ..
}) = expr
{
Some(lit.value())
} else {
None
}
}
fn extract_ident(expr: &syn::Expr) -> Option<Ident> {
if let syn::Expr::Path(path) = expr {
path.path.get_ident().cloned()
} else {
None
}
}
fn parse_field_column_name(attrs: &[Attribute]) -> syn::Result<Option<String>> {
let mut column_name = None;
for attr in attrs {
if !attr.path().is_ident("dbkit") {
continue;
}
attr.parse_nested_meta(|meta| {
if meta.path.is_ident("column") {
if column_name.is_some() {
return Err(meta.error("dbkit: duplicate field column rename"));
}
let lit: syn::LitStr = meta.value()?.parse()?;
column_name = Some(lit.value());
return Ok(());
}
Err(meta.error("dbkit: unsupported field option; expected `column`"))
})?;
}
Ok(column_name)
}
fn scalar_column_name<'a>(fields: &'a [ScalarFieldInfo], ident: &Ident) -> Option<&'a str> {
fields
.iter()
.find(|field| field.ident == *ident)
.map(|field| field.column_name.as_str())
}
fn option_inner_type(ty: &Type) -> Option<Type> {
let path = match ty {
Type::Path(path) => path,
_ => return None,
};
let segment = path.path.segments.last()?;
if segment.ident != "Option" {
return None;
}
let args = match &segment.arguments {
syn::PathArguments::AngleBracketed(args) => args,
_ => return None,
};
let inner = args.args.first()?;
match inner {
syn::GenericArgument::Type(inner_ty) => Some(inner_ty.clone()),
_ => None,
}
}
fn has_attr(attrs: &[Attribute], name: &str) -> bool {
attrs.iter().any(|attr| attr.path().is_ident(name))
}
fn filter_struct_attrs(attrs: &[Attribute]) -> Vec<Attribute> {
let mut kept = Vec::new();
for attr in attrs {
if is_model_attr(attr) {
continue;
}
if attr.path().is_ident("derive") {
if let Ok(mut paths) = attr.parse_args_with(syn::punctuated::Punctuated::<syn::Path, syn::Token![,]>::parse_terminated) {
paths = paths
.into_iter()
.filter(|path| !path.segments.last().map(|seg| seg.ident == "Model").unwrap_or(false))
.collect();
if paths.is_empty() {
continue;
}
let new_attr = quote!(#[derive(#paths)]);
let parsed = syn::Attribute::parse_outer.parse2(new_attr).expect("derive attr");
kept.extend(parsed);
continue;
}
}
kept.push(attr.clone());
}
kept
}
fn filter_field_attrs(attrs: &[Attribute]) -> Vec<Attribute> {
attrs.iter().filter(|attr| !is_field_orm_attr(attr)).cloned().collect()
}
fn is_field_orm_attr(attr: &Attribute) -> bool {
let name = attr.path().get_ident().map(|ident| ident.to_string());
matches!(
name.as_deref(),
Some("key")
| Some("autoincrement")
| Some("unique")
| Some("index")
| Some("has_many")
| Some("belongs_to")
| Some("many_to_many")
| Some("dbkit")
)
}
fn is_model_attr(attr: &Attribute) -> bool {
attr.path().is_ident("model")
}
fn relation_type(field: &Field) -> syn::Result<(RelationKind, Type)> {
let kind = if has_attr(&field.attrs, "has_many") {
RelationKind::HasMany
} else if has_attr(&field.attrs, "belongs_to") {
RelationKind::BelongsTo
} else if has_attr(&field.attrs, "many_to_many") {
RelationKind::ManyToMany
} else {
return Err(syn::Error::new_spanned(field, "dbkit: missing relation attribute"));
};
let child_type = match &field.ty {
Type::Path(path) => {
let segment = path
.path
.segments
.last()
.ok_or_else(|| syn::Error::new_spanned(&field.ty, "dbkit: invalid type"))?;
let expected = match kind {
RelationKind::HasMany => "HasMany",
RelationKind::BelongsTo => "BelongsTo",
RelationKind::ManyToMany => "ManyToMany",
};
if segment.ident != expected {
return Err(syn::Error::new_spanned(
&segment.ident,
format!("dbkit: expected {} marker type", expected),
));
}
match &segment.arguments {
syn::PathArguments::AngleBracketed(args) => {
let ty = args.args.iter().find_map(|arg| match arg {
syn::GenericArgument::Type(ty) => Some(ty.clone()),
_ => None,
});
ty.ok_or_else(|| syn::Error::new_spanned(&segment, "dbkit: missing type"))?
}
_ => return Err(syn::Error::new_spanned(&segment.arguments, "dbkit: expected generic argument")),
}
}
_ => return Err(syn::Error::new_spanned(&field.ty, "dbkit: relation marker must be a type path")),
};
Ok((kind, child_type))
}
fn is_relation_field(field: &Field, rels: &[RelationInfo]) -> bool {
rels.iter().any(|rel| rel.field.ident == field.ident)
}
fn to_snake_case(name: &str) -> String {
let chars: Vec<char> = name.chars().collect();
let mut out = String::with_capacity(name.len() + (name.len() / 4));
for (idx, &ch) in chars.iter().enumerate() {
let prev = idx.checked_sub(1).and_then(|i| chars.get(i)).copied();
let next = chars.get(idx + 1).copied();
if ch.is_uppercase() {
let prev_is_lower_or_digit = prev.map(|p| p.is_lowercase() || p.is_ascii_digit()).unwrap_or(false);
let prev_is_upper = prev.map(|p| p.is_uppercase()).unwrap_or(false);
let next_is_lower = next.map(|n| n.is_lowercase()).unwrap_or(false);
let leading_upper_pair = idx == 1 && prev_is_upper && next_is_lower;
let needs_separator = idx > 0 && (prev_is_lower_or_digit || (prev_is_upper && next_is_lower && !leading_upper_pair));
if needs_separator && !out.ends_with('_') {
out.push('_');
}
for lower in ch.to_lowercase() {
out.push(lower);
}
continue;
}
out.push(ch);
}
out
}
fn to_camel_case(name: &str) -> String {
let mut out = String::new();
let mut uppercase_next = true;
for ch in name.chars() {
if ch == '_' {
uppercase_next = true;
continue;
}
if uppercase_next {
for up in ch.to_uppercase() {
out.push(up);
}
uppercase_next = false;
} else {
out.push(ch);
}
}
out
}
#[derive(Default)]
struct DbEnumArgs {
type_name: Option<String>,
rename_all: Option<String>,
}
#[derive(Clone, Copy)]
enum DbEnumRenameAll {
AsIs,
SnakeCase,
LowerCase,
UpperCase,
ScreamingSnakeCase,
}
fn expand_db_enum(input: syn::ItemEnum) -> syn::Result<TokenStream> {
if !input.generics.params.is_empty() {
return Err(syn::Error::new_spanned(
input.generics,
"dbkit: #[derive(DbEnum)] does not support generics",
));
}
let args = parse_db_enum_args(&input.attrs)?;
let type_name = args
.type_name
.ok_or_else(|| syn::Error::new_spanned(&input.ident, "dbkit: DbEnum requires #[dbkit(type_name = \"...\")]"))?;
let rename_rule = parse_db_enum_rename_all(args.rename_all.as_deref())?;
let enum_ident = input.ident.clone();
let mut as_db_arms = Vec::new();
let mut from_db_arms = Vec::new();
let mut expected_values = Vec::new();
let mut seen_db_names: std::collections::BTreeMap<String, syn::Ident> = std::collections::BTreeMap::new();
for variant in input.variants.iter() {
if !matches!(variant.fields, syn::Fields::Unit) {
return Err(syn::Error::new_spanned(
&variant.fields,
"dbkit: DbEnum only supports unit variants",
));
}
let variant_ident = &variant.ident;
let explicit = parse_db_enum_variant_rename(&variant.attrs)?;
let db_name = match explicit {
Some(value) => value,
None => apply_db_enum_rename_rule(&variant.ident.to_string(), rename_rule),
};
if let Some(first_variant) = seen_db_names.get(&db_name) {
return Err(syn::Error::new_spanned(
variant_ident,
format!(
"dbkit: duplicate DbEnum wire name `{}` for variants `{}` and `{}`",
db_name, first_variant, variant_ident
),
));
}
seen_db_names.insert(db_name.clone(), variant_ident.clone());
let db_name_lit = syn::LitStr::new(&db_name, variant.ident.span());
expected_values.push(db_name);
as_db_arms.push(quote!(Self::#variant_ident => #db_name_lit,));
from_db_arms.push(quote!(#db_name_lit => Ok(Self::#variant_ident),));
}
if as_db_arms.is_empty() {
return Err(syn::Error::new_spanned(enum_ident, "dbkit: DbEnum requires at least one variant"));
}
let type_name_lit = syn::LitStr::new(&type_name, proc_macro2::Span::call_site());
let expected_lit = syn::LitStr::new(&expected_values.join(", "), proc_macro2::Span::call_site());
let tokens = quote! {
impl #enum_ident {
pub const DB_TYPE_NAME: &'static str = #type_name_lit;
pub fn as_db_str(&self) -> &'static str {
match self {
#(#as_db_arms)*
}
}
}
impl ::std::str::FromStr for #enum_ident {
type Err = String;
fn from_str(value: &str) -> Result<Self, Self::Err> {
match value {
#(#from_db_arms)*
_ => Err(format!(
"dbkit: invalid value `{}` for enum {} (expected one of: {})",
value,
stringify!(#enum_ident),
#expected_lit
)),
}
}
}
impl From<#enum_ident> for ::dbkit::Value {
fn from(value: #enum_ident) -> Self {
::dbkit::Value::Enum {
type_name: #type_name_lit,
value: value.as_db_str().to_string(),
}
}
}
impl ::dbkit::sqlx::Type<::dbkit::sqlx::Postgres> for #enum_ident {
fn type_info() -> ::dbkit::sqlx::postgres::PgTypeInfo {
::dbkit::sqlx::postgres::PgTypeInfo::with_name(#type_name_lit)
}
fn compatible(ty: &::dbkit::sqlx::postgres::PgTypeInfo) -> bool {
*ty == ::dbkit::sqlx::postgres::PgTypeInfo::with_name(#type_name_lit)
|| <&str as ::dbkit::sqlx::Type<::dbkit::sqlx::Postgres>>::compatible(ty)
}
}
impl<'q> ::dbkit::sqlx::Encode<'q, ::dbkit::sqlx::Postgres> for #enum_ident {
fn encode_by_ref(
&self,
buf: &mut ::dbkit::sqlx::postgres::PgArgumentBuffer,
) -> Result<::dbkit::sqlx::encode::IsNull, ::dbkit::sqlx::error::BoxDynError> {
<&str as ::dbkit::sqlx::Encode<'q, ::dbkit::sqlx::Postgres>>::encode(self.as_db_str(), buf)
}
fn produces(&self) -> Option<::dbkit::sqlx::postgres::PgTypeInfo> {
Some(::dbkit::sqlx::postgres::PgTypeInfo::with_name(#type_name_lit))
}
fn size_hint(&self) -> usize {
self.as_db_str().len()
}
}
impl<'r> ::dbkit::sqlx::Decode<'r, ::dbkit::sqlx::Postgres> for #enum_ident {
fn decode(value: ::dbkit::sqlx::postgres::PgValueRef<'r>) -> Result<Self, ::dbkit::sqlx::error::BoxDynError> {
let value = <&str as ::dbkit::sqlx::Decode<'r, ::dbkit::sqlx::Postgres>>::decode(value)?;
<Self as ::std::str::FromStr>::from_str(value).map_err(|err| err.into())
}
}
};
Ok(TokenStream::from(tokens))
}
fn parse_db_enum_args(attrs: &[Attribute]) -> syn::Result<DbEnumArgs> {
let mut args = DbEnumArgs::default();
for attr in attrs {
if !attr.path().is_ident("dbkit") {
continue;
}
attr.parse_nested_meta(|meta| {
if meta.path.is_ident("type_name") {
let lit: syn::LitStr = meta.value()?.parse()?;
args.type_name = Some(lit.value());
return Ok(());
}
if meta.path.is_ident("rename_all") {
let lit: syn::LitStr = meta.value()?.parse()?;
args.rename_all = Some(lit.value());
return Ok(());
}
Err(meta.error("dbkit: unsupported DbEnum option; expected `type_name` or `rename_all`"))
})?;
}
Ok(args)
}
fn parse_db_enum_variant_rename(attrs: &[Attribute]) -> syn::Result<Option<String>> {
let mut rename = None;
for attr in attrs {
if !attr.path().is_ident("dbkit") {
continue;
}
attr.parse_nested_meta(|meta| {
if meta.path.is_ident("rename") {
let lit: syn::LitStr = meta.value()?.parse()?;
rename = Some(lit.value());
return Ok(());
}
Err(meta.error("dbkit: unsupported DbEnum variant option; expected `rename`"))
})?;
}
Ok(rename)
}
fn parse_db_enum_rename_all(value: Option<&str>) -> syn::Result<DbEnumRenameAll> {
match value {
None => Ok(DbEnumRenameAll::AsIs),
Some("snake_case") => Ok(DbEnumRenameAll::SnakeCase),
Some("lowercase") => Ok(DbEnumRenameAll::LowerCase),
Some("UPPERCASE") => Ok(DbEnumRenameAll::UpperCase),
Some("SCREAMING_SNAKE_CASE") => Ok(DbEnumRenameAll::ScreamingSnakeCase),
Some(other) => Err(syn::Error::new(
proc_macro2::Span::call_site(),
format!(
"dbkit: unsupported rename_all strategy `{}` for DbEnum; supported values: snake_case, lowercase, UPPERCASE, SCREAMING_SNAKE_CASE",
other
),
)),
}
}
fn apply_db_enum_rename_rule(value: &str, rule: DbEnumRenameAll) -> String {
match rule {
DbEnumRenameAll::AsIs => value.to_string(),
DbEnumRenameAll::SnakeCase => to_snake_case(value),
DbEnumRenameAll::LowerCase => value.to_lowercase(),
DbEnumRenameAll::UpperCase => value.to_uppercase(),
DbEnumRenameAll::ScreamingSnakeCase => to_snake_case(value).to_uppercase(),
}
}
#[cfg(test)]
mod tests {
use super::{apply_db_enum_rename_rule, to_snake_case, DbEnumRenameAll};
#[test]
fn snake_case_respects_acronym_word_boundaries() {
assert_eq!(to_snake_case("HTTPWebhook"), "http_webhook");
assert_eq!(to_snake_case("OAuthToken"), "oauth_token");
assert_eq!(to_snake_case("XMLHttpRequest"), "xml_http_request");
assert_eq!(to_snake_case("WebhookHTTP"), "webhook_http");
}
#[test]
fn screaming_snake_case_respects_acronym_word_boundaries() {
assert_eq!(
apply_db_enum_rename_rule("HTTPWebhook", DbEnumRenameAll::ScreamingSnakeCase),
"HTTP_WEBHOOK"
);
assert_eq!(
apply_db_enum_rename_rule("XMLHttpRequest", DbEnumRenameAll::ScreamingSnakeCase),
"XML_HTTP_REQUEST"
);
}
}