From 7e92876feb397fde84c95e97bcd5bd4f7176b178 Mon Sep 17 00:00:00 2001 From: Andrey Kutejko Date: Wed, 26 Oct 2022 05:22:08 +0200 Subject: [PATCH] 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 --- sea-orm-codegen/src/entity/base_entity.rs | 8 +++++--- sea-orm-codegen/src/entity/conjunct_relation.rs | 6 ++++-- sea-orm-codegen/src/entity/relation.rs | 7 ++++++- sea-orm-codegen/src/entity/writer.rs | 7 +++++-- 4 files changed, 20 insertions(+), 8 deletions(-) diff --git a/sea-orm-codegen/src/entity/base_entity.rs b/sea-orm-codegen/src/entity/base_entity.rs index 0147c6f06..d58a6529c 100644 --- 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 { diff --git a/sea-orm-codegen/src/entity/conjunct_relation.rs b/sea-orm-codegen/src/entity/conjunct_relation.rs index 05e8010a7..a156cf77a 100644 --- 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 index 185af3178..2c29eef9b 100644 --- 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 index 3eea90569..006beedab 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}; @@ -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; }