support flatten attribute in FromRow macro (#1959)

* support flatten attribute in FromRow macro

* added docs for flatten FromRow attribute

Author: TheoOiry

sqlx-core/src/from_row.rs

@@ -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>;
 }

sqlx-macros/src/derives/attributes.rs

@@ -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> {
@@ -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
@@ -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"),
@@ -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(

sqlx-macros/src/derives/row.rs

@@ -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,
 };
 
@@ -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);
 

tests/postgres/derives.rs

@@ -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(())
+}