From 750235ac39b293220f3e28877a15d92dc9cf5cbf Mon Sep 17 00:00:00 2001
From: Andrey Kutejko <andy128k@gmail.com>
Date: Sun, 18 Sep 2022 16:05:41 +0200
Subject: [PATCH] Escape Rust keywords used in table names

---
 sea-orm-codegen/src/entity/base_entity.rs | 8 +++++---
 sea-orm-codegen/src/entity/relation.rs    | 7 ++++++-
 sea-orm-codegen/src/entity/writer.rs      | 7 +++++--
 3 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/sea-orm-codegen/src/entity/base_entity.rs b/sea-orm-codegen/src/entity/base_entity.rs
index a745b1f74..0858c63cc 100644
--- 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
index 57722e506..a4c8af653 100644
--- 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
index 19722e183..4ef232363 100644
--- 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;
         }