Allow #[derive(Identifiable)] to work with composite primary keys

This is a tad bit magic in the macro piece. We're expecting the primary
keys *not to* have a trailing comma, so we can rely on the fact that
`(T)` is equivalent to `T`, but if there's more than one element it
would be a tuple.

Beyond that everything was quite straightforward with the ground work
that we've laid.

Fixes diesel-rs#42.

Author: sgrif

CHANGELOG.md

@@ -44,6 +44,11 @@ for Rust libraries in [RFC #1105](https://github.com/rust-lang/rfcs/blob/master/
 
 [exists]: http://docs.diesel.rs/diesel/expression/dsl/fn.sql.html
 
+* `#[derive(Identifiable)]` can be used with structs that have primary keys
+  other than `id`, as well as structs with composite primary keys. You can now
+  annotate the struct with `#[primary_key(nonstandard)]` or `#[primary_key(foo,
+  bar)]`.
+
 ### Changed
 
 * All macros with the same name as traits we can derive (e.g. `Queryable!`) have

diesel/src/macros/identifiable.rs

@@ -65,12 +65,12 @@ macro_rules! impl_Identifiable {
             body = $($body:tt)*
         ),
         found_option_with_name = primary_key,
-        value = ($primary_key_name:ident),
+        value = $primary_key_names:tt,
     ) => {
         impl_Identifiable! {
             (
                 table_name = $table_name,
-                primary_key_name = $primary_key_name,
+                primary_key_names = $primary_key_names,
             )
             $($body)*
         }
@@ -83,13 +83,14 @@ macro_rules! impl_Identifiable {
             struct_ty = $struct_ty:ty,
             lifetimes = ($($lifetimes:tt),*),
         ),
-        primary_key_field = {
+        found_fields_with_field_names,
+        fields = [$({
             field_name: $field_name:ident,
             column_name: $column_name:ident,
             field_ty: $field_ty:ty,
             field_kind: $field_kind:ident,
             $($rest:tt)*
-        },
+        })*],
     ) => {
         impl<$($lifetimes),*> $crate::associations::HasTable for $struct_ty {
             type Table = $table_name::table;
@@ -100,48 +101,28 @@ macro_rules! impl_Identifiable {
         }
 
         impl<'ident $(,$lifetimes)*> $crate::associations::Identifiable for &'ident $struct_ty {
-            type Id = &'ident $field_ty;
+            type Id = ($(&'ident $field_ty),*);
 
             fn id(self) -> Self::Id {
-                &self.$field_name
+                ($(&self.$field_name),*)
             }
         }
     };
 
-    // Search for the primary key field and continue
+    // Search for the primary key fields and continue
     (
         (
             table_name = $table_name:ident,
-            primary_key_name = $primary_key_name:ident,
+            primary_key_names = $primary_key_names:tt,
             $($args:tt)*
         ),
-        fields = [{
-            field_name: $field_name:ident,
-            $($rest:tt)*
-        } $($fields:tt)*],
+        fields = $fields:tt,
     ) => {
-        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)*],
-                }
-            }
+        __diesel_fields_with_field_names! {
+            (table_name = $table_name, $($args)*),
+            callback = impl_Identifiable,
+            targets = $primary_key_names,
+            fields = $fields,
         }
     };
 
@@ -345,3 +326,32 @@ fn derive_identifiable_with_non_standard_pk_given_before_table_name() {
     assert_eq!(&"hi", foo1.id());
     assert_eq!(&"there", foo2.id());
 }
+
+#[test]
+fn derive_identifiable_with_composite_pk() {
+    use associations::Identifiable;
+
+    #[allow(missing_debug_implementations, missing_copy_implementations, dead_code)]
+    struct Foo {
+        id: i32,
+        foo_id: i32,
+        bar_id: i32,
+        foo: i32,
+    }
+
+    impl_Identifiable! {
+        #[primary_key(foo_id, bar_id)]
+        #[table_name(bars)]
+        struct Foo {
+            id: i32,
+            foo_id: i32,
+            bar_id: i32,
+            foo: i32,
+        }
+    }
+
+    let foo1 = Foo { id: 1, foo_id: 2, bar_id: 3, foo: 4 };
+    let foo2 = Foo { id: 5, foo_id: 6, bar_id: 7, foo: 8 };
+    assert_eq!((&2, &3), foo1.id());
+    assert_eq!((&6, &7), foo2.id());
+}

diesel/src/macros/parse.rs

@@ -333,6 +333,105 @@ macro_rules! __diesel_field_with_column_name {
     };
 }
 
+#[doc(hidden)]
+#[macro_export]
+macro_rules! __diesel_field_with_field_name {
+    (
+        $headers:tt,
+        callback = $callback:ident,
+        target = $target_field_name:ident,
+        fields = [$({
+            field_name: $field_name:ident,
+            $($field_info:tt)*
+        })*],
+    ) => {
+        $(
+            static_cond! {
+                if $target_field_name == $field_name {
+                    $callback! {
+                        $headers,
+                        found_field_with_field_name = $field_name,
+                        field = {
+                            field_name: $field_name,
+                            $($field_info)*
+                        },
+                    }
+                }
+            }
+        )*
+    };
+}
+
+#[doc(hidden)]
+#[macro_export]
+macro_rules! __diesel_fields_with_field_names {
+    // Entrypoint, start search
+    (
+        $headers:tt,
+        callback = $callback:ident,
+        targets = ($target_field_name:ident $(,$rest:ident)*),
+        fields = $fields:tt,
+    ) => {
+        __diesel_field_with_field_name! {
+            (
+                targets = ($($rest),*),
+                fields = $fields,
+                headers = $headers,
+                callback = $callback,
+                found_fields = [],
+            ),
+            callback = __diesel_fields_with_field_names,
+            target = $target_field_name,
+            fields = $fields,
+        }
+    };
+
+    // Found field, more to search for
+    (
+        (
+            targets = ($target_field_name:ident $(,$rest:ident)*),
+            fields = $fields:tt,
+            headers = $headers:tt,
+            callback = $callback:ident,
+            found_fields = [$($found_fields:tt)*],
+        ),
+        found_field_with_field_name = $ignore:tt,
+        field = $field:tt,
+    ) => {
+        __diesel_field_with_field_name! {
+            (
+                targets = ($($rest),*),
+                fields = $fields,
+                headers = $headers,
+                callback = $callback,
+                found_fields = [$($found_fields)* $field],
+            ),
+            callback = __diesel_fields_with_field_names,
+            target = $target_field_name,
+            fields = $fields,
+        }
+    };
+
+    // Found field, no more to search for
+    (
+        (
+            targets = (),
+            fields = $fields:tt,
+            headers = $headers:tt,
+            callback = $callback:ident,
+            found_fields = [$($found_fields:tt)*],
+        ),
+        found_field_with_field_name = $ignore:tt,
+        field = $field:tt,
+    ) => {
+        $callback! {
+            $headers,
+            found_fields_with_field_names,
+            fields = [$($found_fields)* $field],
+        }
+    }
+}
+
 #[doc(hidden)]
 #[macro_export]
 macro_rules! __diesel_find_option_with_name {

diesel_codegen/src/identifiable.rs

@@ -8,16 +8,18 @@ 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 primary_key_names = model.primary_key_names;
     let fields = model.attrs;
-    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);
+    for pk in &primary_key_names {
+        if !fields.iter().any(|f| f.field_name.as_ref() == Some(pk)) {
+            panic!("Could not find a field named `{}` on `{}`", pk, &model.name);
+        }
     }
 
     quote!(impl_Identifiable! {
         (
             table_name = #table_name,
-            primary_key_name = #primary_key_name,
+            primary_key_names = (#(#primary_key_names),*),
             struct_ty = #struct_ty,
             lifetimes = (#(#lifetimes),*),
         ),

diesel_codegen/src/model.rs

@@ -8,7 +8,7 @@ pub struct Model {
     pub attrs: Vec<Attr>,
     pub name: syn::Ident,
     pub generics: syn::Generics,
-    pub primary_key_name: syn::Ident,
+    pub primary_key_names: Vec<syn::Ident>,
     table_name_from_annotation: Option<syn::Ident>,
 }
 
@@ -23,9 +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 primary_key_names = list_value_of_attr_with_name(&item.attrs, "primary_key")
+            .map(|v| v.into_iter().map(Clone::clone).collect())
+            .unwrap_or_else(|| vec![syn::Ident::new("id")]);
         let table_name_from_annotation = str_value_of_attr_with_name(
             &item.attrs, "table_name").map(syn::Ident::new);
 
@@ -34,7 +34,7 @@ impl Model {
             attrs: attrs,
             name: name,
             generics: generics,
-            primary_key_name: primary_key_name,
+            primary_key_names: primary_key_names,
             table_name_from_annotation: table_name_from_annotation,
         })
     }

diesel_codegen/src/util.rs

@@ -35,7 +35,19 @@ pub fn ident_value_of_attr_with_name<'a>(
     attrs: &'a [Attribute],
     name: &str,
 ) -> Option<&'a Ident> {
-    attr_with_name(attrs, name).map(|attr| single_arg_value_of_attr(attr, name))
+    list_value_of_attr_with_name(attrs, name).map(|idents| {
+        if idents.len() != 1 {
+            panic!(r#"`{}` must be in the form `#[{}(something)]`"#, name, name);
+        }
+        idents[0]
+    })
+}
+
+pub fn list_value_of_attr_with_name<'a>(
+    attrs: &'a [Attribute],
+    name: &str,
+) -> Option<Vec<&'a Ident>> {
+    attr_with_name(attrs, name).map(|attr| list_value_of_attr(attr, name))
 }
 
 pub fn attr_with_name<'a>(
@@ -56,19 +68,15 @@ pub fn str_value_of_meta_item<'a>(item: &'a MetaItem, name: &str) -> &'a str {
     }
 }
 
-fn single_arg_value_of_attr<'a>(attr: &'a Attribute, name: &str) -> &'a Ident {
-    let usage_err = || panic!(r#"`{}` must be in the form `#[{}(something)]`"#, name, name);
+fn list_value_of_attr<'a>(attr: &'a Attribute, name: &str) -> Vec<&'a Ident> {
     match attr.value {
         MetaItem::List(_, ref items) => {
-            if items.len() != 1 {
-                return usage_err();
-            }
-            match items[0] {
+            items.iter().map(|item| match *item {
                 NestedMetaItem::MetaItem(MetaItem::Word(ref name)) => name,
-                _ => usage_err(),
-            }
+                _ => panic!(r#"`{}` must be in the form `#[{}(something)]`"#, name, name),
+            }).collect()
         }
-        _ => usage_err(),
+        _ => panic!(r#"`{}` must be in the form `#[{}(something)]`"#, name, name),
     }
 }
 

diesel_codegen_syntex/src/identifiable.rs

@@ -1,9 +1,10 @@
 use syntax::ast;
 use syntax::codemap::Span;
 use syntax::ext::base::{Annotatable, ExtCtxt};
+use syntax::parse::token;
 
 use model::Model;
-use util::lifetime_list_tokens;
+use util::{lifetime_list_tokens, comma_delimited_tokens};
 
 pub fn expand_derive_identifiable(
     cx: &mut ExtCtxt,
@@ -16,20 +17,25 @@ 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 primary_key_names = model.primary_key_names();
         let fields = model.field_tokens_for_stable_macro(cx);
-        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 `{}` on `{}`", primary_key_name, model.name));
+        for name in primary_key_names {
+            if model.attr_named(*name).is_none() {
+                cx.span_err(span, &format!("Could not find a field named `{}` on `{}`", name, model.name));
+                return;
+            }
         }
+
+        let primary_key_names = comma_delimited_tokens(
+            primary_key_names.into_iter().map(|n| token::Ident(*n)), span);
+        push(Annotatable::Item(quote_item!(cx, impl_Identifiable! {
+            (
+                table_name = $table_name,
+                primary_key_names = ($primary_key_names),
+                struct_ty = $struct_ty,
+                lifetimes = ($lifetimes),
+            ),
+            fields = [$fields],
+        }).unwrap()));
     }
 }