Merge pull request diesel-rs#2853 from flisky/master

weiznich · web-flow · commit c2b4b62c332f · 2021-12-14T10:01:50.000Z
support `serialize_as` for `AsChangeset`
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -48,7 +48,7 @@ for Rust libraries in [RFC #1105](https://github.com/rust-lang/rfcs/blob/master/
 * Added support for PostgreSQL's `SIMILAR TO` and `NOT SIMILAR TO`.
 
 * Added `#[diesel(serialize_as)]` analogous to `#[diesel(deserialize_as)]`. This allows
-  customization of the serialization behaviour of `Insertable` structs.
+  customization of the serialization behaviour of `Insertable` and `AsChangeset` structs.
 
 * Added support for `GROUP BY` clauses
 
diff --git a/diesel_derives/src/as_changeset.rs b/diesel_derives/src/as_changeset.rs
@@ -1,6 +1,7 @@
 use proc_macro2::TokenStream;
-use syn::{DeriveInput, Path};
+use syn::{DeriveInput, Expr, Path, Type};
 
+use attrs::AttributeSpanWrapper;
 use field::Field;
 use model::Model;
 use util::{inner_of_option_ty, is_option_ty, wrap_in_dummy_mod};
@@ -37,58 +38,105 @@ pub fn derive(item: DeriveInput) -> TokenStream {
     impl_generics.params.push(parse_quote!('update));
     let (impl_generics, _, _) = impl_generics.split_for_impl();
 
-    let ref_changeset_ty = fields_for_update.iter().map(|field| {
-        field_changeset_ty(
-            field,
-            &table_name,
-            treat_none_as_null,
-            Some(quote!(&'update)),
-        )
-    });
-    let ref_changeset_expr = fields_for_update
-        .iter()
-        .map(|field| field_changeset_expr(field, &table_name, treat_none_as_null, Some(quote!(&))));
-
-    let direct_changeset_ty = fields_for_update
-        .iter()
-        .map(|field| field_changeset_ty(field, &table_name, treat_none_as_null, None));
-    let direct_changeset_expr = fields_for_update
-        .iter()
-        .map(|field| field_changeset_expr(field, &table_name, treat_none_as_null, None));
-
-    wrap_in_dummy_mod(quote!(
-        use diesel::query_builder::AsChangeset;
-        use diesel::prelude::*;
-
-        impl #impl_generics AsChangeset for &'update #struct_name #ty_generics
-        #where_clause
-        {
-            type Target = #table_name::table;
-            type Changeset = <(#(#ref_changeset_ty,)*) as AsChangeset>::Changeset;
-
-            fn as_changeset(self) -> Self::Changeset {
-                (#(#ref_changeset_expr,)*).as_changeset()
+    let mut generate_borrowed_changeset = true;
+
+    let mut direct_field_ty = Vec::with_capacity(fields_for_update.len());
+    let mut direct_field_assign = Vec::with_capacity(fields_for_update.len());
+    let mut ref_field_ty = Vec::with_capacity(fields_for_update.len());
+    let mut ref_field_assign = Vec::with_capacity(fields_for_update.len());
+
+    for field in fields_for_update {
+        match field.serialize_as.as_ref() {
+            Some(AttributeSpanWrapper { item: ty, .. }) => {
+                direct_field_ty.push(field_changeset_ty_serialize_as(
+                    field,
+                    table_name,
+                    ty,
+                    treat_none_as_null,
+                ));
+                direct_field_assign.push(field_changeset_expr_serialize_as(
+                    field,
+                    table_name,
+                    ty,
+                    treat_none_as_null,
+                ));
+
+                generate_borrowed_changeset = false; // as soon as we hit one field with #[diesel(serialize_as)] there is no point in generating the impl of AsChangeset for borrowed structs
+            }
+            None => {
+                direct_field_ty.push(field_changeset_ty(
+                    field,
+                    table_name,
+                    None,
+                    treat_none_as_null,
+                ));
+                direct_field_assign.push(field_changeset_expr(
+                    field,
+                    table_name,
+                    None,
+                    treat_none_as_null,
+                ));
+                ref_field_ty.push(field_changeset_ty(
+                    field,
+                    table_name,
+                    Some(quote!(&'update)),
+                    treat_none_as_null,
+                ));
+                ref_field_assign.push(field_changeset_expr(
+                    field,
+                    table_name,
+                    Some(quote!(&)),
+                    treat_none_as_null,
+                ));
             }
         }
+    }
 
+    let changeset_owned = quote! {
         impl #impl_generics AsChangeset for #struct_name #ty_generics
         #where_clause
         {
             type Target = #table_name::table;
-            type Changeset = <(#(#direct_changeset_ty,)*) as AsChangeset>::Changeset;
+            type Changeset = <(#(#direct_field_ty,)*) as AsChangeset>::Changeset;
 
             fn as_changeset(self) -> Self::Changeset {
-                (#(#direct_changeset_expr,)*).as_changeset()
+                (#(#direct_field_assign,)*).as_changeset()
             }
         }
+    };
+
+    let changeset_borrowed = if generate_borrowed_changeset {
+        quote! {
+            impl #impl_generics AsChangeset for &'update #struct_name #ty_generics
+            #where_clause
+            {
+                type Target = #table_name::table;
+                type Changeset = <(#(#ref_field_ty,)*) as AsChangeset>::Changeset;
+
+                fn as_changeset(self) -> Self::Changeset {
+                    (#(#ref_field_assign,)*).as_changeset()
+                }
+            }
+        }
+    } else {
+        quote! {}
+    };
+
+    wrap_in_dummy_mod(quote!(
+        use diesel::query_builder::AsChangeset;
+        use diesel::prelude::*;
+
+        #changeset_owned
+
+        #changeset_borrowed
     ))
 }
 
 fn field_changeset_ty(
     field: &Field,
     table_name: &Path,
-    treat_none_as_null: bool,
     lifetime: Option<TokenStream>,
+    treat_none_as_null: bool,
 ) -> TokenStream {
     let column_name = field.column_name();
     if !treat_none_as_null && is_option_ty(&field.ty) {
@@ -103,8 +151,8 @@ fn field_changeset_ty(
 fn field_changeset_expr(
     field: &Field,
     table_name: &Path,
-    treat_none_as_null: bool,
     lifetime: Option<TokenStream>,
+    treat_none_as_null: bool,
 ) -> TokenStream {
     let field_name = &field.name;
     let column_name = field.column_name();
@@ -118,3 +166,34 @@ fn field_changeset_expr(
         quote!(#table_name::#column_name.eq(#lifetime self.#field_name))
     }
 }
+
+fn field_changeset_ty_serialize_as(
+    field: &Field,
+    table_name: &Path,
+    ty: &Type,
+    treat_none_as_null: bool,
+) -> TokenStream {
+    let column_name = field.column_name();
+    if !treat_none_as_null && is_option_ty(&field.ty) {
+        let inner_ty = inner_of_option_ty(ty);
+        quote!(std::option::Option<diesel::dsl::Eq<#table_name::#column_name, #inner_ty>>)
+    } else {
+        quote!(diesel::dsl::Eq<#table_name::#column_name, #ty>)
+    }
+}
+
+fn field_changeset_expr_serialize_as(
+    field: &Field,
+    table_name: &Path,
+    ty: &Type,
+    treat_none_as_null: bool,
+) -> TokenStream {
+    let field_name = &field.name;
+    let column_name = field.column_name();
+    let column: Expr = parse_quote!(#table_name::#column_name);
+    if !treat_none_as_null && is_option_ty(&field.ty) {
+        quote!(self.#field_name.map(|x| #column.eq(::std::convert::Into::<#ty>::into(x))))
+    } else {
+        quote!(#column.eq(::std::convert::Into::<#ty>::into(self.#field_name)))
+    }
+}
diff --git a/diesel_derives/src/lib.rs b/diesel_derives/src/lib.rs
@@ -62,6 +62,17 @@ mod valid_grouping;
 /// from the name of the corresponding column, you can annotate the field with
 /// `#[diesel(column_name = some_column_name)]`.
 ///
+/// To provide custom serialization behavior for a field, you can use
+/// `#[diesel(serialize_as = SomeType)]`. If this attribute is present, Diesel
+/// will call `.into` on the corresponding field and serialize the instance of `SomeType`,
+/// rather than the actual field on your struct. This can be used to add custom behavior for a
+/// single field, or use types that are otherwise unsupported by Diesel.
+/// Normally, Diesel produces two implementations of the `AsChangeset` trait for your
+/// struct using this derive: one for an owned version and one for a borrowed version.
+/// Using `#[diesel(serialize_as)]` implies a conversion using `.into` which consumes the underlying value.
+/// Hence, once you use `#[diesel(serialize_as)]`, Diesel can no longer insert borrowed
+/// versions of your struct.
+///
 /// By default, any `Option` fields on the struct are skipped if their value is
 /// `None`. If you would like to assign `NULL` to the field instead, you can
 /// annotate your struct with `#[diesel(treat_none_as_null = true)]`.
@@ -87,6 +98,10 @@ mod valid_grouping;
 /// * `#[diesel(column_name = some_column_name)]`, overrides the column name
 ///    of the current field to `some_column_name`. By default the field
 ///    name is used as column name.
+/// * `#[diesel(serialize_as = SomeType)]`, instead of serializing the actual
+///    field type, Diesel will convert the field into `SomeType` using `.into` and
+///    serialize that instead. By default this derive will serialize directly using
+///    the actual field type.
 #[proc_macro_error]
 #[proc_macro_derive(
     AsChangeset,
diff --git a/diesel_derives/tests/as_changeset.rs b/diesel_derives/tests/as_changeset.rs
@@ -218,6 +218,57 @@ fn with_explicit_column_names() {
     assert_eq!(Ok(expected), actual);
 }
 
+#[test]
+fn with_serialize_as() {
+    #[derive(Debug, FromSqlRow, AsExpression)]
+    #[diesel(sql_type = sql_types::Text)]
+    struct UppercaseString(pub String);
+
+    impl Into<UppercaseString> for String {
+        fn into(self) -> UppercaseString {
+            UppercaseString(self.to_uppercase())
+        }
+    }
+
+    impl<DB> serialize::ToSql<sql_types::Text, DB> for UppercaseString
+    where
+        DB: backend::Backend,
+        String: serialize::ToSql<sql_types::Text, DB>,
+    {
+        fn to_sql<W: std::io::Write>(
+            &self,
+            out: &mut serialize::Output<W, DB>,
+        ) -> serialize::Result {
+            self.0.to_sql(out)
+        }
+    }
+
+    #[derive(AsChangeset)]
+    struct User {
+        #[diesel(serialize_as = UppercaseString)]
+        name: String,
+        #[diesel(serialize_as = UppercaseString)]
+        hair_color: Option<String>,
+    }
+
+    let connection = &mut connection_with_sean_and_tess_in_users_table();
+
+    update(users::table.find(1))
+        .set(User {
+            name: String::from("Jim"),
+            hair_color: Some(String::from("blue")),
+        })
+        .execute(connection)
+        .unwrap();
+
+    let expected = vec![
+        (1, String::from("JIM"), Some(String::from("BLUE"))),
+        (2, String::from("Tess"), Some(String::from("brown"))),
+    ];
+    let actual = users::table.order(users::id).load(connection);
+    assert_eq!(Ok(expected), actual);
+}
+
 #[test]
 fn tuple_struct() {
     #[derive(AsChangeset)]
diff --git a/diesel_tests/tests/serialize_as.rs b/diesel_tests/tests/serialize_as.rs
@@ -31,8 +31,16 @@ struct InsertableUser {
     name: String,
 }
 
+#[derive(Clone, Debug, AsChangeset, Identifiable)]
+#[diesel(table_name = users)]
+struct ChangeUser {
+    id: i32,
+    #[diesel(serialize_as = UppercaseString)]
+    name: String,
+}
+
 #[test]
-fn serialization_can_be_customized() {
+fn insert_serialization_can_be_customized() {
     use crate::schema::users::dsl::*;
     let connection = &mut connection();
 
@@ -50,3 +58,31 @@ fn serialization_can_be_customized() {
         users.select(name).first(connection)
     );
 }
+
+#[test]
+fn update_serialization_can_be_customized() {
+    use crate::schema::users::dsl::*;
+    let connection = &mut connection();
+
+    let user = InsertableUser {
+        name: "thomas".to_string(),
+    };
+    diesel::insert_into(users)
+        .values(user)
+        .execute(connection)
+        .unwrap();
+
+    let user = ChangeUser {
+        id: users.select(id).first(connection).unwrap(),
+        name: "eizinger".to_string(),
+    };
+    diesel::update(&user)
+        .set(user.clone())
+        .execute(connection)
+        .unwrap();
+
+    assert_eq!(
+        Ok("EIZINGER".to_string()),
+        users.select(name).first(connection)
+    );
+}
