Escape Rust keywords used in table names (#1052)

* Escape Rust keywords used in table names * Escape rust keyword when generating conjunct relation impl Co-authored-by: Billy Chan <ccw.billy.123@gmail.com>
SeaQL · Oct 26, 2022 · 7e92876 · 7e92876
1 parent a3fb8e8
commit 7e92876
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 8 deletions.
diff --git a/sea-orm-codegen/src/entity/base_entity.rs b/sea-orm-codegen/src/entity/base_entity.rs
@@ -4,7 +4,9 @@ use quote::format_ident;
 use quote::quote;
 use sea_query::ColumnType;
 
-use crate::{Column, ConjunctRelation, DateTimeCrate, PrimaryKey, Relation};
+use crate::{
+    util::escape_rust_keyword, Column, ConjunctRelation, DateTimeCrate, PrimaryKey, Relation,
+};
 
 #[derive(Clone, Debug)]
 pub struct Entity {
@@ -25,11 +27,11 @@ impl Entity {
     }
 
     pub fn get_table_name_snake_case_ident(&self) -> Ident {
-        format_ident!("{}", self.get_table_name_snake_case())
+        format_ident!("{}", escape_rust_keyword(self.get_table_name_snake_case()))
     }
 
     pub fn get_table_name_camel_case_ident(&self) -> Ident {
-        format_ident!("{}", self.get_table_name_camel_case())
+        format_ident!("{}", escape_rust_keyword(self.get_table_name_camel_case()))
     }
 
     pub fn get_column_names_snake_case(&self) -> Vec<Ident> {

diff --git a/sea-orm-codegen/src/entity/conjunct_relation.rs b/sea-orm-codegen/src/entity/conjunct_relation.rs
@@ -2,6 +2,8 @@ use heck::{CamelCase, SnakeCase};
 use proc_macro2::Ident;
 use quote::format_ident;
 
+use crate::util::escape_rust_keyword;
+
 #[derive(Clone, Debug)]
 pub struct ConjunctRelation {
     pub(crate) via: String,
@@ -10,11 +12,11 @@ pub struct ConjunctRelation {
 
 impl ConjunctRelation {
     pub fn get_via_snake_case(&self) -> Ident {
-        format_ident!("{}", self.via.to_snake_case())
+        format_ident!("{}", escape_rust_keyword(self.via.to_snake_case()))
     }
 
     pub fn get_to_snake_case(&self) -> Ident {
-        format_ident!("{}", self.to.to_snake_case())
+        format_ident!("{}", escape_rust_keyword(self.to.to_snake_case()))
     }
 
     pub fn get_to_camel_case(&self) -> Ident {

diff --git a/sea-orm-codegen/src/entity/relation.rs b/sea-orm-codegen/src/entity/relation.rs
@@ -4,6 +4,8 @@ use proc_macro2::{Ident, TokenStream};
 use quote::{format_ident, quote};
 use sea_query::{ForeignKeyAction, TableForeignKey};
 
+use crate::util::escape_rust_keyword;
+
 #[derive(Clone, Debug)]
 pub enum RelationType {
     HasOne,
@@ -41,7 +43,10 @@ impl Relation {
         if self.self_referencing {
             None
         } else {
-            Some(format_ident!("{}", self.ref_table.to_snake_case()))
+            Some(format_ident!(
+                "{}",
+                escape_rust_keyword(self.ref_table.to_snake_case())
+            ))
         }
     }
 

diff --git a/sea-orm-codegen/src/entity/writer.rs b/sea-orm-codegen/src/entity/writer.rs
@@ -1,4 +1,4 @@
-use crate::{ActiveEnum, Entity};
+use crate::{util::escape_rust_keyword, ActiveEnum, Entity};
 use heck::CamelCase;
 use proc_macro2::TokenStream;
 use quote::{format_ident, quote};
@@ -544,7 +544,10 @@ impl EntityWriter {
     }
 
     pub fn gen_mod(entity: &Entity) -> TokenStream {
-        let table_name_snake_case_ident = entity.get_table_name_snake_case_ident();
+        let table_name_snake_case_ident = format_ident!(
+            "{}",
+            escape_rust_keyword(entity.get_table_name_snake_case_ident())
+        );
         quote! {
             pub mod #table_name_snake_case_ident;
         }