Merge pull request #87 from sgrif/sg-table-name

Allow structs to be annotated with `#[table_name="foo"]`

diesel_codegen/src/attr.rs

@@ -1,8 +1,8 @@
 use syntax::ast;
-use syntax::attr::AttrMetaMethods;
 use syntax::ext::base::ExtCtxt;
 use syntax::ptr::P;
-use syntax::parse::token::str_to_ident;
+
+use util::str_value_of_attr_with_name;
 
 #[derive(Debug, PartialEq, Eq)]
 pub struct Attr {
@@ -14,19 +14,8 @@ pub struct Attr {
 impl Attr {
     pub fn from_struct_field(cx: &mut ExtCtxt, field: &ast::StructField) -> Option<Self> {
         let field_name = field.node.ident();
-        let column_name = field.node.attrs.iter().filter_map(|attr| {
-            if attr.check_name("column_name") {
-                attr.value_str().map(|name| {
-                    str_to_ident(&name)
-                }).or_else(|| {
-                    cx.span_err(attr.span(),
-                        r#"`column_name` must be in the form `#[column_name="something"]`"#);
-                    None
-                })
-            } else {
-                None
-            }
-        }).nth(0);
+        let column_name =
+            str_value_of_attr_with_name(cx, &field.node.attrs, "column_name");
         let ty = field.node.ty.clone();
 
         match (column_name, field_name) {

diesel_codegen/src/lib.rs

@@ -24,6 +24,8 @@ include!(concat!(env!("OUT_DIR"), "/lib.rs"));
 #[cfg(not(feature = "with-syntex"))]
 include!("lib.in.rs");
 
+mod util;
+
 #[cfg(feature = "with-syntex")]
 pub fn register(reg: &mut syntex::Registry) {
     reg.add_attr("feature(custom_derive)");
@@ -36,6 +38,8 @@ pub fn register(reg: &mut syntex::Registry) {
     reg.add_decorator("belongs_to", associations::expand_belongs_to);
     reg.add_macro("infer_table_from_schema", schema_inference::expand_load_table);
     reg.add_macro("infer_schema", schema_inference::expand_infer_schema);
+
+    reg.add_post_expansion_pass(util::strip_attributes);
 }
 
 #[cfg_attr(not(feature = "with-syntex"), plugin_registrar)]

diesel_codegen/src/model.rs

@@ -5,11 +5,13 @@ use syntax::ptr::P;
 use syntax::parse::token::str_to_ident;
 
 use attr::Attr;
+use util::str_value_of_attr_with_name;
 
 pub struct Model {
     pub ty: P<ast::Ty>,
     pub attrs: Vec<Attr>,
     pub name: ast::Ident,
+    table_name_from_annotation: Option<ast::Ident>,
 }
 
 impl Model {
@@ -19,6 +21,8 @@ impl Model {
         annotatable: &Annotatable,
     ) -> Option<Self> {
         if let Annotatable::Item(ref item) = *annotatable {
+            let table_name_from_annotation =
+                str_value_of_attr_with_name(cx, &item.attrs, "table_name");
             Attr::from_item(cx, item).map(|(generics, attrs)| {
                 let ty = builder.ty().path()
                     .segment(item.ident).with_generics(generics.clone())
@@ -27,6 +31,7 @@ impl Model {
                     ty: ty,
                     attrs: attrs,
                     name: item.ident,
+                    table_name_from_annotation: table_name_from_annotation,
                 }
             })
         } else {
@@ -39,7 +44,9 @@ impl Model {
     }
 
     pub fn table_name(&self) -> ast::Ident {
-        str_to_ident(&infer_table_name(&self.name.name.as_str()))
+        self.table_name_from_annotation.unwrap_or_else(|| {
+            str_to_ident(&infer_table_name(&self.name.name.as_str()))
+        })
     }
 
     pub fn attr_named(&self, name: ast::Ident) -> &Attr {

diesel_codegen/src/util.rs

@@ -0,0 +1,51 @@
+use syntax::ast;
+use syntax::attr::AttrMetaMethods;
+use syntax::ext::base::ExtCtxt;
+use syntax::parse::token::str_to_ident;
+
+fn str_value_of_attr(
+    cx: &mut ExtCtxt,
+    attr: &ast::Attribute,
+    name: &str,
+) -> Option<ast::Ident> {
+    attr.value_str().map(|value| {
+        str_to_ident(&value)
+    }).or_else(|| {
+        cx.span_err(attr.span(),
+            &format!(r#"`{}` must be in the form `#[{}="something"]`"#, name, name));
+        None
+    })
+}
+
+pub fn str_value_of_attr_with_name(
+    cx: &mut ExtCtxt,
+    attrs: &[ast::Attribute],
+    name: &str,
+) -> Option<ast::Ident> {
+    attrs.iter()
+        .find(|a| a.check_name(name))
+        .and_then(|a| str_value_of_attr(cx, &a, name))
+}
+
+#[cfg(feature = "with-syntex")]
+pub fn strip_attributes(krate: ast::Crate) -> ast::Crate {
+    use syntax::fold;
+
+    struct StripAttributeFolder;
+
+    impl fold::Folder for StripAttributeFolder {
+        fn fold_attribute(&mut self, attr: ast::Attribute) -> Option<ast::Attribute> {
+            if attr.check_name("table_name") {
+                None
+            } else {
+                Some(attr)
+            }
+        }
+
+        fn fold_mac(&mut self, mac: ast::Mac) -> ast::Mac {
+            fold::noop_fold_mac(mac, self)
+        }
+    }
+
+    fold::Folder::fold_crate(&mut StripAttributeFolder, krate)
+}

diesel_tests/tests/associations.rs

@@ -24,3 +24,18 @@ fn one_to_many_returns_query_source_for_association() {
     let found_posts: Vec<_> = Post::belonging_to(&tess).load(&connection).unwrap().collect();
     assert_eq!(tess_posts, found_posts);
 }
+
+#[test]
+fn association_where_struct_name_doesnt_match_table_name() {
+    let connection = connection_with_sean_and_tess_in_users_table();
+
+    let sean = find_user_by_name("Sean", &connection);
+    let post: Post = insert(&sean.new_post("Hello", None)).into(posts::table)
+        .get_result(&connection).unwrap();
+    insert(&NewComment(post.id, "comment")).into(comments::table)
+        .execute(&connection).unwrap();
+
+    let comment_text = SpecialComment::belonging_to(&post).select(comments::text)
+        .first::<String>(&connection);
+    assert_eq!(Ok("comment".into()), comment_text);
+}

diesel_tests/tests/schema.rs

@@ -59,6 +59,14 @@ pub struct Comment {
     text: String,
 }
 
+#[derive(PartialEq, Eq, Debug, Clone, Queriable)]
+#[belongs_to(post)]
+#[table_name="comments"]
+pub struct SpecialComment {
+    id: i32,
+    post_id: i32,
+}
+
 infer_schema!(dotenv!("DATABASE_URL"));
 numeric_expr!(users::id);