implement DeriveIden in sea-orm only (#1740)

* WIP, implementing Iden * completed implementation for DeriveIden and added basic test cases * added feature flag to prevent sea-orm::sea-query::DeriveIden from crashing when sea-query is not used * fixed doc test and adjusted test case * enable `sea-query-derive`'s `sea-orm` feature * Bump `sea-query-derive` to v0.4 * Update Cargo.toml * Update Cargo.toml * adjusted test cases and updated so that iden attribute will not be snake cased * Update Cargo.toml * Update main.rs --------- Co-authored-by: Billy Chan <[email protected]> Co-authored-by: Chris Tsang <[email protected]>
SeaQL · Jul 13, 2023 · 866025a · 866025a
1 parent b4c1a69
commit 866025a
Show file tree

Hide file tree

Showing 9 changed files with 271 additions and 7 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -34,7 +34,7 @@ tracing = { version = "0.1", default-features = false, features = ["attributes",
 rust_decimal = { version = "1", default-features = false, optional = true }
 bigdecimal = { version = "0.3", default-features = false, optional = true }
 sea-orm-macros = { version = "0.12.0-rc.4", path = "sea-orm-macros", default-features = false, features = ["strum"] }
-sea-query = { version = "0.29.0-rc.2", features = ["thread-safe", "hashable-value"] }
+sea-query = { version = "0.29.0", features = ["thread-safe", "hashable-value"] }
 sea-query-binder = { version = "0.4.0-rc.2", default-features = false, optional = true }
 strum = { version = "0.25", default-features = false }
 serde = { version = "1.0", default-features = false }

diff --git a/issues/1473/Cargo.toml b/issues/1473/Cargo.toml
@@ -0,0 +1,13 @@
+[workspace]
+# A separate workspace
+
+[package]
+name = "sea-orm-issues-1473"
+version = "0.1.0"
+edition = "2021"
+publish = false
+
+[dependencies.sea-orm]
+path = "../../"
+default-features = false
+features = ["macros", "runtime-tokio-native-tls"]
diff --git a/issues/1473/src/main.rs b/issues/1473/src/main.rs
@@ -0,0 +1,17 @@
+use sea_orm::{sea_query::{self, Iden}};
+
+#[derive(Iden)]
+enum Character {
+    Table,
+    Id,
+}
+
+#[derive(Iden)]
+struct Glyph;
+
+fn main() {
+    assert_eq!(Character::Table.to_string(), "character");
+    assert_eq!(Character::Id.to_string(), "id");
+
+    assert_eq!(Glyph.to_string(), "glyph");
+}
diff --git a/sea-orm-macros/src/derives/derive_iden.rs b/sea-orm-macros/src/derives/derive_iden.rs
@@ -0,0 +1,146 @@
+use heck::ToSnakeCase;
+use proc_macro2::{self, TokenStream};
+use quote::{quote, quote_spanned};
+use syn::{
+    punctuated::Punctuated, DataEnum, DataStruct, DeriveInput, Expr, Fields, LitStr, Variant,
+};
+
+fn must_be_valid_iden(name: &str) -> bool {
+    // can only begin with [a-z_]
+    name.chars()
+        .take(1)
+        .all(|c| c == '_' || c.is_ascii_alphabetic())
+        && name.chars().all(|c| c == '_' || c.is_ascii_alphanumeric())
+}
+
+fn impl_iden_for_unit_struct(
+    ident: &proc_macro2::Ident,
+    new_iden: &str,
+) -> proc_macro2::TokenStream {
+    let prepare = if must_be_valid_iden(new_iden) {
+        quote! {
+            fn prepare(&self, s: &mut dyn ::std::fmt::Write, q: sea_orm::sea_query::Quote) {
+                write!(s, "{}", q.left()).unwrap();
+                self.unquoted(s);
+                write!(s, "{}", q.right()).unwrap();
+            }
+        }
+    } else {
+        quote! {}
+    };
+    quote! {
+        impl sea_orm::sea_query::Iden for #ident {
+            #prepare
+
+            fn unquoted(&self, s: &mut dyn ::std::fmt::Write) {
+                write!(s, #new_iden).unwrap();
+            }
+        }
+    }
+}
+
+fn impl_iden_for_enum(
+    ident: &proc_macro2::Ident,
+    variants: Punctuated<Variant, syn::token::Comma>,
+) -> proc_macro2::TokenStream {
+    let variants = variants.iter();
+    let mut all_valid = true;
+
+    let match_pair: Vec<TokenStream> = variants
+        .map(|v| {
+            let var_ident = &v.ident;
+            let mut var_name = var_ident.to_string().to_snake_case();
+            v.attrs
+                .iter()
+                .filter(|attr| attr.path().is_ident("sea_orm"))
+                .try_for_each(|attr| {
+                    attr.parse_nested_meta(|meta| {
+                        if meta.path.is_ident("iden") {
+                            let litstr: LitStr = meta.value()?.parse()?;
+                            var_name = litstr.value();
+                            all_valid &= must_be_valid_iden(var_name.as_str());
+                        } else {
+                            // Reads the value expression to advance the parse stream.
+                            // Some parameters do not have any value,
+                            // so ignoring an error occurred here.
+                            let _: Option<Expr> = meta.value().and_then(|v| v.parse()).ok();
+                        }
+                        Ok(())
+                    })
+                })
+                .expect("something something");
+            quote! { Self::#var_ident => write!(s, "{}", #var_name).unwrap() }
+        })
+        .collect();
+
+    let match_arms: TokenStream = quote! { #(#match_pair),* };
+
+    let prepare = if all_valid {
+        quote! {
+            fn prepare(&self, s: &mut dyn ::std::fmt::Write, q: sea_orm::sea_query::Quote) {
+                write!(s, "{}", q.left()).unwrap();
+                self.unquoted(s);
+                write!(s, "{}", q.right()).unwrap();
+            }
+        }
+    } else {
+        quote! {}
+    };
+
+    quote! {
+        impl sea_orm::sea_query::Iden for #ident {
+            #prepare
+
+            fn unquoted(&self, s: &mut dyn ::std::fmt::Write) {
+                match self {
+                    #match_arms
+                };
+            }
+        }
+    }
+}
+
+pub fn expand_derive_iden(input: DeriveInput) -> syn::Result<TokenStream> {
+    let DeriveInput { ident, data, .. } = input;
+
+    let mut new_iden: TokenStream = ident.to_string().to_snake_case().parse().unwrap();
+    input
+        .attrs
+        .iter()
+        .filter(|attr| attr.path().is_ident("sea_orm"))
+        .try_for_each(|attr| {
+            attr.parse_nested_meta(|meta| {
+                if meta.path.is_ident("iden") {
+                    let litstr: LitStr = meta.value()?.parse()?;
+                    new_iden = syn::parse_str::<TokenStream>(&litstr.value())?;
+                } else {
+                    // Reads the value expression to advance the parse stream.
+                    // Some parameters do not have any value,
+                    // so ignoring an error occurred here.
+                    let _: Option<Expr> = meta.value().and_then(|v| v.parse()).ok();
+                }
+                Ok(())
+            })
+        })?;
+
+    // Currently we only support enums and unit structs
+    match data {
+        syn::Data::Enum(DataEnum { variants, .. }) => {
+            if variants.is_empty() {
+                Ok(TokenStream::new())
+            } else {
+                Ok(impl_iden_for_enum(&ident, variants))
+            }
+        }
+        syn::Data::Struct(DataStruct {
+            fields: Fields::Unit,
+            ..
+        }) => Ok(impl_iden_for_unit_struct(
+            &ident,
+            new_iden.to_string().as_str(),
+        )),
+        _ => Ok(quote_spanned! {
+            ident.span() => compile_error!("you can only derive DeriveIden on unit struct or enum");
+        }),
+    }
+}
diff --git a/sea-orm-macros/src/derives/mod.rs b/sea-orm-macros/src/derives/mod.rs
@@ -4,6 +4,7 @@ mod active_model;
 mod active_model_behavior;
 mod attributes;
 mod column;
+mod derive_iden;
 mod entity;
 mod entity_model;
 mod from_query_result;
@@ -24,6 +25,7 @@ pub use active_enum_display::*;
 pub use active_model::*;
 pub use active_model_behavior::*;
 pub use column::*;
+pub use derive_iden::*;
 pub use entity::*;
 pub use entity_model::*;
 pub use from_query_result::*;

diff --git a/sea-orm-macros/src/lib.rs b/sea-orm-macros/src/lib.rs
@@ -852,3 +852,45 @@ pub fn derive_active_enum_display(input: TokenStream) -> TokenStream {
         Err(e) => e.to_compile_error().into(),
     }
 }
+
+/// The DeriveIden derive macro will implement `sea_orm::sea_query::Iden` for simplify Iden implementation.
+///
+/// ## Usage
+///
+/// ```rust
+/// use sea_orm::DeriveIden;
+///
+/// #[derive(DeriveIden)]
+/// pub enum Class {
+///     Id,
+///     Title,
+///     Text,
+/// }
+///
+/// #[derive(DeriveIden)]
+/// struct Glyph;
+/// ```
+///
+/// You can use iden = "" to customize the name
+/// ```
+/// use sea_orm::DeriveIden;
+///
+/// #[derive(DeriveIden)]
+/// pub enum Class {
+///     Id,
+///     #[sea_orm(iden = "turtle")]
+///     Title,
+///     #[sea_orm(iden = "TeXt")]
+///     Text,
+/// }
+/// ```
+#[cfg(feature = "derive")]
+#[proc_macro_derive(DeriveIden, attributes(sea_orm))]
+pub fn derive_iden(input: TokenStream) -> TokenStream {
+    let derive_input = parse_macro_input!(input as DeriveInput);
+
+    match derives::expand_derive_iden(derive_input) {
+        Ok(token_stream) => token_stream.into(),
+        Err(e) => e.to_compile_error().into(),
+    }
+}
diff --git a/src/entity/column.rs b/src/entity/column.rs
@@ -1,4 +1,4 @@
-use crate::{EntityName, IdenStatic, IntoSimpleExpr, Iterable};
+use crate::{self as sea_orm, EntityName, IdenStatic, IntoSimpleExpr, Iterable};
 use sea_query::{
     Alias, BinOper, DynIden, Expr, Iden, IntoIden, SeaRc, SelectStatement, SimpleExpr, Value,
 };

diff --git a/src/lib.rs b/src/lib.rs
@@ -349,15 +349,13 @@ pub use schema::*;
 #[cfg(feature = "macros")]
 pub use sea_orm_macros::{
     DeriveActiveEnum, DeriveActiveModel, DeriveActiveModelBehavior, DeriveColumn,
-    DeriveCustomColumn, DeriveDisplay, DeriveEntity, DeriveEntityModel, DeriveIntoActiveModel,
-    DeriveMigrationName, DeriveModel, DerivePartialModel, DerivePrimaryKey, DeriveRelatedEntity,
-    DeriveRelation, FromJsonQueryResult, FromQueryResult,
+    DeriveCustomColumn, DeriveDisplay, DeriveEntity, DeriveEntityModel, DeriveIden,
+    DeriveIntoActiveModel, DeriveMigrationName, DeriveModel, DerivePartialModel, DerivePrimaryKey,
+    DeriveRelatedEntity, DeriveRelation, FromJsonQueryResult, FromQueryResult,
 };
 
 pub use sea_query;
 pub use sea_query::Iden;
-#[cfg(feature = "macros")]
-pub use sea_query::Iden as DeriveIden;
 
 pub use sea_orm_macros::EnumIter;
 pub use strum;
diff --git a/tests/derive_iden_tests.rs b/tests/derive_iden_tests.rs
@@ -0,0 +1,46 @@
+pub mod common;
+pub use common::{features::*, setup::*, TestContext};
+use sea_orm::entity::prelude::*;
+use sea_orm_macros::DeriveIden;
+
+#[derive(DeriveIden)]
+pub enum Class {
+    Id,
+    Title,
+    Text,
+}
+
+#[derive(DeriveIden)]
+pub enum Book {
+    Id,
+    #[sea_orm(iden = "turtle")]
+    Title,
+    #[sea_orm(iden = "TeXt")]
+    Text,
+    #[sea_orm(iden = "ty_pe")]
+    Type,
+}
+
+#[derive(DeriveIden)]
+struct Glyph;
+
+#[derive(DeriveIden)]
+#[sea_orm(iden = "weRd")]
+struct Word;
+
+#[test]
+fn main() -> Result<(), DbErr> {
+    assert_eq!(Class::Id.to_string(), "id");
+    assert_eq!(Class::Title.to_string(), "title");
+    assert_eq!(Class::Text.to_string(), "text");
+
+    assert_eq!(Book::Id.to_string(), "id");
+    assert_eq!(Book::Title.to_string(), "turtle");
+    assert_eq!(Book::Text.to_string(), "TeXt");
+    assert_eq!(Book::Type.to_string(), "ty_pe");
+
+    assert_eq!(Glyph.to_string(), "glyph");
+
+    assert_eq!(Word.to_string(), "weRd");
+    Ok(())
+}