Merge pull request diesel-rs#532 from diesel-rs/sg-custom-pk-identifiable

sgrif · web-flow · commit 47d1196c4416 · 2016-12-08T12:13:48.000-05:00
Allow custom primary keys to be specified for `Identifiable`
diff --git a/diesel/src/macros/identifiable.rs b/diesel/src/macros/identifiable.rs
@@ -28,50 +28,68 @@
 /// ```
 #[macro_export]
 macro_rules! impl_Identifiable {
-    // Extract table name from meta item
+    // Find table name in options
     (
         $(())*
-        #[table_name($table_name:ident)]
-        $($rest:tt)*
+        $(#[$option_name:ident $option_value:tt])*
+        $(pub)* struct $($rest:tt)*
     ) => {
-        impl_Identifiable! {
-            (table_name = $table_name,)
-            $($rest)*
+        __diesel_find_option_with_name! {
+            (options = [$(#[$option_name $option_value])*], body = $($rest)*),
+            option_name = table_name,
+            args = [$(#[$option_name $option_value])*],
+            callback = impl_Identifiable,
         }
     };
 
-    // Strip meta items that aren't table name
+    // Next, find primary key name in options
     (
-        $args:tt
-        #[$ignore:meta]
-        $($rest:tt)*
+        (options = $options:tt, $($args:tt)*),
+        found_option_with_name = table_name,
+        value = ($table_name:ident),
     ) => {
-        impl_Identifiable!($args $($rest)*);
+        __diesel_find_option_with_name! {
+            (options = $options, table_name = $table_name, $($args)*),
+            option_name = primary_key,
+            args = $options,
+            callback = impl_Identifiable,
+            default = (id),
+        }
     };
 
-    // Strip pub (if present) and struct from definition
-    // After this step, we will go to the step at the bottom.
+    // Options extracted, deal with the struct body (bottom two branches)
     (
-        $args:tt
-        $(pub)* struct $($body:tt)*
+        (
+            options = $ignore:tt,
+            table_name = $table_name:ident,
+            body = $($body:tt)*
+        ),
+        found_option_with_name = primary_key,
+        value = ($primary_key_name:ident),
     ) => {
-        impl_Identifiable!($args $($body)*);
+        impl_Identifiable! {
+            (
+                table_name = $table_name,
+                primary_key_name = $primary_key_name,
+            )
+            $($body)*
+        }
     };
 
-    // We found the `id` field, return the final impl
+    // We found the primary key field, return the final impl
     (
         (
             table_name = $table_name:ident,
             struct_ty = $struct_ty:ty,
             lifetimes = ($($lifetimes:tt),*),
         ),
-        fields = [{
-            field_name: id,
+        primary_key_field = {
+            field_name: $field_name:ident,
             column_name: $column_name:ident,
             field_ty: $field_ty:ty,
             field_kind: $field_kind:ident,
             $($rest:tt)*
-        } $($fields:tt)*],
+        },
     ) => {
         impl<$($lifetimes),*> $crate::associations::HasTable for $struct_ty {
             type Table = $table_name::table;
@@ -85,25 +103,45 @@ macro_rules! impl_Identifiable {
             type Id = &'ident $field_ty;
 
             fn id(self) -> Self::Id {
-                &self.id
+                &self.$field_name
             }
         }
     };
 
-    // Search for the `id` field and continue
+    // Search for the primary key field and continue
     (
-        $args:tt,
+        (
+            table_name = $table_name:ident,
+            primary_key_name = $primary_key_name:ident,
+            $($args:tt)*
+        ),
         fields = [{
             field_name: $field_name:ident,
-            column_name: $column_name:ident,
-            field_ty: $field_ty:ty,
-            field_kind: $field_kind:ident,
             $($rest:tt)*
         } $($fields:tt)*],
     ) => {
-        impl_Identifiable! {
-            $args,
-            fields = [$($fields)*],
+        static_cond! {
+            if $primary_key_name == $field_name {
+                impl_Identifiable! {
+                    (
+                        table_name = $table_name,
+                        $($args)*
+                    ),
+                    primary_key_field = {
+                        field_name: $field_name,
+                        $($rest)*
+                    },
+                }
+            } else {
+                impl_Identifiable! {
+                    (
+                        table_name = $table_name,
+                        primary_key_name = $primary_key_name,
+                        $($args)*
+                    ),
+                    fields = [$($fields)*],
+                }
+            }
         }
     };
 
@@ -253,3 +291,57 @@ fn derive_identifiable_on_struct_with_lifetime() {
     assert_eq!(&"hi", foo1.id());
     assert_eq!(&"there", foo2.id());
 }
+
+#[test]
+fn derive_identifiable_with_non_standard_pk() {
+    use associations::Identifiable;
+
+    #[allow(missing_debug_implementations, missing_copy_implementations, dead_code)]
+    struct Foo<'a> {
+        id: i32,
+        foo_id: &'a str,
+        foo: i32,
+    }
+
+    impl_Identifiable! {
+        #[table_name(bars)]
+        #[primary_key(foo_id)]
+        struct Foo<'a> {
+            id: i32,
+            foo_id: &'a str,
+            foo: i32,
+        }
+    }
+
+    let foo1 = Foo { id: 1, foo_id: "hi", foo: 2 };
+    let foo2 = Foo { id: 2, foo_id: "there", foo: 3 };
+    assert_eq!(&"hi", foo1.id());
+    assert_eq!(&"there", foo2.id());
+}
+
+#[test]
+fn derive_identifiable_with_non_standard_pk_given_before_table_name() {
+    use associations::Identifiable;
+
+    #[allow(missing_debug_implementations, missing_copy_implementations, dead_code)]
+    struct Foo<'a> {
+        id: i32,
+        foo_id: &'a str,
+        foo: i32,
+    }
+
+    impl_Identifiable! {
+        #[primary_key(foo_id)]
+        #[table_name(bars)]
+        struct Foo<'a> {
+            id: i32,
+            foo_id: &'a str,
+            foo: i32,
+        }
+    }
+
+    let foo1 = Foo { id: 1, foo_id: "hi", foo: 2 };
+    let foo2 = Foo { id: 2, foo_id: "there", foo: 3 };
+    assert_eq!(&"hi", foo1.id());
+    assert_eq!(&"there", foo2.id());
+}
diff --git a/diesel/src/macros/mod.rs b/diesel/src/macros/mod.rs
@@ -289,7 +289,7 @@ macro_rules! table_body {
             pub const all_columns: ($($column_name,)+) = ($($column_name,)+);
 
             #[allow(non_camel_case_types, missing_debug_implementations)]
-            #[derive(Clone, Copy)]
+            #[derive(Debug, Clone, Copy)]
             pub struct table;
 
             impl table {
diff --git a/diesel/src/macros/parse.rs b/diesel/src/macros/parse.rs
@@ -332,3 +332,67 @@ macro_rules! __diesel_field_with_column_name {
         )*
     };
 }
+
+#[doc(hidden)]
+#[macro_export]
+macro_rules! __diesel_find_option_with_name {
+    // When no default is given, just a simple search will work
+    (
+        $headers:tt,
+        option_name = $target_option_name:ident,
+        args = [$(#[$option_name:ident $option_value:tt])*],
+        callback = $callback:ident,
+    ) => { $(
+        static_cond! {
+            if $target_option_name == $option_name {
+                $callback! {
+                    $headers,
+                    found_option_with_name = $target_option_name,
+                    value = $option_value,
+                }
+            }
+        }
+    )* };
+
+    // When a default is given, we need to recurse so we know when we're done
+    (
+        $headers:tt,
+        option_name = $target_option_name:ident,
+        args = [#[$option_name:ident $option_value:tt] $($rest:tt)*],
+        callback = $callback:ident,
+        default = $default:tt,
+    ) => {
+        static_cond! {
+            if $target_option_name == $option_name {
+                $callback! {
+                    $headers,
+                    found_option_with_name = $target_option_name,
+                    value = $option_value,
+                }
+            } else {
+                __diesel_find_option_with_name! {
+                    $headers,
+                    option_name = $target_option_name,
+                    args = [$($rest)*],
+                    callback = $callback,
+                    default = $default,
+                }
+            }
+        }
+    };
+
+    // Default was given, and we didn't find a value
+    (
+        $headers:tt,
+        option_name = $target_option_name:ident,
+        args = [],
+        callback = $callback:ident,
+        default = $default:tt,
+    ) => {
+        $callback! {
+            $headers,
+            found_option_with_name = $target_option_name,
+            value = $default,
+        }
+    };
+}
diff --git a/diesel_codegen/src/identifiable.rs b/diesel_codegen/src/identifiable.rs
@@ -8,14 +8,16 @@ pub fn derive_identifiable(item: syn::MacroInput) -> Tokens {
     let table_name = model.table_name();
     let struct_ty = &model.ty;
     let lifetimes = model.generics.lifetimes;
+    let primary_key_name = model.primary_key_name;
     let fields = model.attrs;
-    if !fields.iter().any(|f| f.field_name == Some(syn::Ident::new("id"))) {
-        panic!("Could not find a field named `id` on `{}`", &model.name);
+    if !fields.iter().any(|f| f.field_name.as_ref() == Some(&primary_key_name)) {
+        panic!("Could not find a field named `{}` on `{}`", primary_key_name, &model.name);
     }
 
     quote!(impl_Identifiable! {
         (
             table_name = #table_name,
+            primary_key_name = #primary_key_name,
             struct_ty = #struct_ty,
             lifetimes = (#(#lifetimes),*),
         ),
diff --git a/diesel_codegen/src/lib.rs b/diesel_codegen/src/lib.rs
@@ -38,7 +38,7 @@ pub fn derive_queryable(input: TokenStream) -> TokenStream {
     expand_derive(input, queryable::derive_queryable)
 }
 
-#[proc_macro_derive(Identifiable, attributes(table_name))]
+#[proc_macro_derive(Identifiable, attributes(table_name, primary_key))]
 pub fn derive_identifiable(input: TokenStream) -> TokenStream {
     expand_derive(input, identifiable::derive_identifiable)
 }
diff --git a/diesel_codegen/src/model.rs b/diesel_codegen/src/model.rs
@@ -1,13 +1,14 @@
 use syn;
 
 use attr::Attr;
-use util::{struct_ty, str_value_of_attr_with_name};
+use util::*;
 
 pub struct Model {
     pub ty: syn::Ty,
     pub attrs: Vec<Attr>,
     pub name: syn::Ident,
     pub generics: syn::Generics,
+    pub primary_key_name: syn::Ident,
     table_name_from_annotation: Option<syn::Ident>,
 }
 
@@ -22,6 +23,9 @@ impl Model {
         let ty = struct_ty(item.ident.clone(), &item.generics);
         let name = item.ident.clone();
         let generics = item.generics.clone();
+        let primary_key_name = ident_value_of_attr_with_name(&item.attrs, "primary_key")
+            .map(Clone::clone)
+            .unwrap_or(syn::Ident::new("id"));
         let table_name_from_annotation = str_value_of_attr_with_name(
             &item.attrs, "table_name").map(syn::Ident::new);
 
@@ -30,6 +34,7 @@ impl Model {
             attrs: attrs,
             name: name,
             generics: generics,
+            primary_key_name: primary_key_name,
             table_name_from_annotation: table_name_from_annotation,
         })
     }
diff --git a/diesel_codegen_syntex/src/identifiable.rs b/diesel_codegen_syntex/src/identifiable.rs
@@ -1,7 +1,6 @@
 use syntax::ast;
 use syntax::codemap::Span;
 use syntax::ext::base::{Annotatable, ExtCtxt};
-use syntax::parse::token::str_to_ident;
 
 use model::Model;
 use util::lifetime_list_tokens;
@@ -17,18 +16,20 @@ pub fn expand_derive_identifiable(
         let table_name = model.table_name();
         let struct_ty = &model.ty;
         let lifetimes = lifetime_list_tokens(&model.generics.lifetimes, span);
+        let primary_key_name = model.primary_key_name;
         let fields = model.field_tokens_for_stable_macro(cx);
-        if model.attr_named(str_to_ident("id")).is_some() {
+        if model.attr_named(primary_key_name).is_some() {
             push(Annotatable::Item(quote_item!(cx, impl_Identifiable! {
                 (
                     table_name = $table_name,
+                    primary_key_name = $primary_key_name,
                     struct_ty = $struct_ty,
                     lifetimes = ($lifetimes),
                 ),
                 fields = [$fields],
             }).unwrap()));
         } else {
-            cx.span_err(span, &format!("Could not find a field named `id` on `{}`", model.name));
+            cx.span_err(span, &format!("Could not find a field named `{}` on `{}`", primary_key_name, model.name));
         }
     }
 }
diff --git a/diesel_codegen_syntex/src/model.rs b/diesel_codegen_syntex/src/model.rs
diff --git a/diesel_codegen_syntex/src/util.rs b/diesel_codegen_syntex/src/util.rs
diff --git a/diesel_tests/tests/annotations.rs b/diesel_tests/tests/annotations.rs

Original file line number	Diff line number	Diff line change
`@@ -38,7 +38,7 @@ pub fn derive_queryable(input: TokenStream) -> TokenStream {`
`38`	`38`	`expand_derive(input, queryable::derive_queryable)`
`39`	`39`	`}`
`40`	`40`
`41`		`-#[proc_macro_derive(Identifiable, attributes(table_name))]`
	`41`	`+#[proc_macro_derive(Identifiable, attributes(table_name, primary_key))]`
`42`	`42`	`pub fn derive_identifiable(input: TokenStream) -> TokenStream {`
`43`	`43`	`expand_derive(input, identifiable::derive_identifiable)`
`44`	`44`	`}`