Escape Rust keywords used in table names

SeaQL · Sep 18, 2022 · 750235a · 750235a
1 parent 4e51b88
commit 750235a
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 6 deletions.
diff --git a/sea-orm-codegen/src/entity/base_entity.rs b/sea-orm-codegen/src/entity/base_entity.rs
@@ -1,4 +1,6 @@
-use crate::{Column, ConjunctRelation, DateTimeCrate, PrimaryKey, Relation};
+use crate::{
+    util::escape_rust_keyword, Column, ConjunctRelation, DateTimeCrate, PrimaryKey, Relation,
+};
 use heck::{CamelCase, SnakeCase};
 use proc_macro2::{Ident, TokenStream};
 use quote::format_ident;
@@ -22,11 +24,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/relation.rs b/sea-orm-codegen/src/entity/relation.rs
@@ -3,6 +3,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,
@@ -40,7 +42,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};
@@ -543,7 +543,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;
         }