Skip to content

Commit

Permalink
support flatten attribute in FromRow macro (#1959)
Browse files Browse the repository at this point in the history
* support flatten attribute in FromRow macro

* added docs for flatten FromRow attribute
  • Loading branch information
TheoOiry authored Jul 12, 2022
1 parent bc3e705 commit 7cdb68b
Show file tree
Hide file tree
Showing 4 changed files with 117 additions and 36 deletions.
30 changes: 30 additions & 0 deletions sqlx-core/src/from_row.rs
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,36 @@ use crate::row::Row;
/// will set the value of the field `location` to the default value of `Option<String>`,
/// which is `None`.
///
/// ### `flatten`
///
/// If you want to handle a field that implements [`FromRow`],
/// you can use the `flatten` attribute to specify that you want
/// it to use [`FromRow`] for parsing rather than the usual method.
/// For example:
///
/// ```rust,ignore
/// #[derive(sqlx::FromRow)]
/// struct Address {
/// country: String,
/// city: String,
/// road: String,
/// }
///
/// #[derive(sqlx::FromRow)]
/// struct User {
/// id: i32,
/// name: String,
/// #[sqlx(flatten)]
/// address: Address,
/// }
/// ```
/// Given a query such as:
///
/// ```sql
/// SELECT id, name, country, city, road FROM users;
/// ```
///
/// This field is compatible with the `default` attribute.
pub trait FromRow<'r, R: Row>: Sized {
fn from_row(row: &'r R) -> Result<Self, Error>;
}
Expand Down
9 changes: 8 additions & 1 deletion sqlx-macros/src/derives/attributes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,7 @@ pub struct SqlxContainerAttributes {
pub struct SqlxChildAttributes {
pub rename: Option<String>,
pub default: bool,
pub flatten: bool,
}

pub fn parse_container_attributes(input: &[Attribute]) -> syn::Result<SqlxContainerAttributes> {
Expand Down Expand Up @@ -177,6 +178,7 @@ pub fn parse_container_attributes(input: &[Attribute]) -> syn::Result<SqlxContai
pub fn parse_child_attributes(input: &[Attribute]) -> syn::Result<SqlxChildAttributes> {
let mut rename = None;
let mut default = false;
let mut flatten = false;

for attr in input.iter().filter(|a| a.path.is_ident("sqlx")) {
let meta = attr
Expand All @@ -193,6 +195,7 @@ pub fn parse_child_attributes(input: &[Attribute]) -> syn::Result<SqlxChildAttri
..
}) if path.is_ident("rename") => try_set!(rename, val.value(), value),
Meta::Path(path) if path.is_ident("default") => default = true,
Meta::Path(path) if path.is_ident("flatten") => flatten = true,
u => fail!(u, "unexpected attribute"),
},
u => fail!(u, "unexpected attribute"),
Expand All @@ -201,7 +204,11 @@ pub fn parse_child_attributes(input: &[Attribute]) -> syn::Result<SqlxChildAttri
}
}

Ok(SqlxChildAttributes { rename, default })
Ok(SqlxChildAttributes {
rename,
default,
flatten,
})
}

pub fn check_transparent_attributes(
Expand Down
73 changes: 38 additions & 35 deletions sqlx-macros/src/derives/row.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use proc_macro2::{Span, TokenStream};
use quote::quote;
use syn::{
parse_quote, punctuated::Punctuated, token::Comma, Data, DataStruct, DeriveInput, Field,
parse_quote, punctuated::Punctuated, token::Comma, Data, DataStruct, DeriveInput, Expr, Field,
Fields, FieldsNamed, FieldsUnnamed, Lifetime, Stmt,
};

Expand Down Expand Up @@ -63,46 +63,49 @@ fn expand_derive_from_row_struct(

predicates.push(parse_quote!(&#lifetime ::std::primitive::str: ::sqlx::ColumnIndex<R>));

for field in fields {
let ty = &field.ty;

predicates.push(parse_quote!(#ty: ::sqlx::decode::Decode<#lifetime, R::Database>));
predicates.push(parse_quote!(#ty: ::sqlx::types::Type<R::Database>));
}

let (impl_generics, _, where_clause) = generics.split_for_impl();

let container_attributes = parse_container_attributes(&input.attrs)?;

let reads = fields.iter().filter_map(|field| -> Option<Stmt> {
let id = &field.ident.as_ref()?;
let attributes = parse_child_attributes(&field.attrs).unwrap();
let id_s = attributes
.rename
.or_else(|| Some(id.to_string().trim_start_matches("r#").to_owned()))
.map(|s| match container_attributes.rename_all {
Some(pattern) => rename_all(&s, pattern),
None => s,
})
.unwrap();

let ty = &field.ty;

if attributes.default {
Some(
parse_quote!(let #id: #ty = row.try_get(#id_s).or_else(|e| match e {
let reads: Vec<Stmt> = fields
.iter()
.filter_map(|field| -> Option<Stmt> {
let id = &field.ident.as_ref()?;
let attributes = parse_child_attributes(&field.attrs).unwrap();
let ty = &field.ty;

let expr: Expr = if attributes.flatten {
predicates.push(parse_quote!(#ty: ::sqlx::FromRow<#lifetime, R>));
parse_quote!(#ty::from_row(row))
} else {
predicates.push(parse_quote!(#ty: ::sqlx::decode::Decode<#lifetime, R::Database>));
predicates.push(parse_quote!(#ty: ::sqlx::types::Type<R::Database>));

let id_s = attributes
.rename
.or_else(|| Some(id.to_string().trim_start_matches("r#").to_owned()))
.map(|s| match container_attributes.rename_all {
Some(pattern) => rename_all(&s, pattern),
None => s,
})
.unwrap();
parse_quote!(row.try_get(#id_s))
};

if attributes.default {
Some(parse_quote!(let #id: #ty = #expr.or_else(|e| match e {
::sqlx::Error::ColumnNotFound(_) => {
::std::result::Result::Ok(Default::default())
},
e => ::std::result::Result::Err(e)
})?;),
)
} else {
Some(parse_quote!(
let #id: #ty = row.try_get(#id_s)?;
))
}
});
})?;))
} else {
Some(parse_quote!(
let #id: #ty = #expr?;
))
}
})
.collect();

let (impl_generics, _, where_clause) = generics.split_for_impl();

let names = fields.iter().map(|field| &field.ident);

Expand Down
41 changes: 41 additions & 0 deletions tests/postgres/derives.rs
Original file line number Diff line number Diff line change
Expand Up @@ -573,3 +573,44 @@ async fn test_default() -> anyhow::Result<()> {

Ok(())
}

#[cfg(feature = "macros")]
#[sqlx_macros::test]
async fn test_flatten() -> anyhow::Result<()> {
#[derive(Debug, Default, sqlx::FromRow)]
struct AccountDefault {
default: Option<i32>,
}

#[derive(Debug, sqlx::FromRow)]
struct UserInfo {
name: String,
surname: String,
}

#[derive(Debug, sqlx::FromRow)]
struct AccountKeyword {
id: i32,
#[sqlx(flatten)]
info: UserInfo,
#[sqlx(default)]
#[sqlx(flatten)]
default: AccountDefault,
}

let mut conn = new::<Postgres>().await?;

let account: AccountKeyword = sqlx::query_as(
r#"SELECT * from (VALUES (1, 'foo', 'bar')) accounts("id", "name", "surname")"#,
)
.fetch_one(&mut conn)
.await?;
println!("{:?}", account);

assert_eq!(1, account.id);
assert_eq!("foo", account.info.name);
assert_eq!("bar", account.info.surname);
assert_eq!(None, account.default.default);

Ok(())
}

0 comments on commit 7cdb68b

Please sign in to comment.