codegen: register seaography entity module (#2403)

* codegen: register seaography entity module * codegen: register seaography active enum * Update seaography examples * fmt
SeaQL · Dec 2, 2024 · bfca886 · bfca886
1 parent 760b825
commit bfca886
Show file tree

Hide file tree

Showing 7 changed files with 61 additions and 15 deletions.
diff --git a/examples/loco_seaography/src/graphql/query_root.rs b/examples/loco_seaography/src/graphql/query_root.rs
@@ -14,11 +14,7 @@ pub fn schema(
     // Builder of Seaography query root
     let mut builder = Builder::new(&CONTEXT, database.clone());
     // Register SeaORM entities
-    seaography::register_entities!(
-        builder,
-        // List all models we want to include in the GraphQL endpoint here
-        [files, notes, users]
-    );
+    let builder = crate::models::_entities::register_entity_modules(builder);
     // Configure async GraphQL limits
     let schema = builder
         .schema_builder()

diff --git a/examples/loco_seaography/src/models/_entities/mod.rs b/examples/loco_seaography/src/models/_entities/mod.rs
@@ -5,3 +5,5 @@ pub mod prelude;
 pub mod files;
 pub mod notes;
 pub mod users;
+
+seaography::register_entity_modules!([files, notes, users]);
diff --git a/examples/react_admin/backend/src/graphql/query_root.rs b/examples/react_admin/backend/src/graphql/query_root.rs
@@ -12,13 +12,9 @@ pub fn schema(
     complexity: usize,
 ) -> Result<Schema, SchemaError> {
     // Builder of Seaography query root
-    let mut builder = Builder::new(&CONTEXT, database.clone());
+    let builder = Builder::new(&CONTEXT, database.clone());
     // Register SeaORM entities
-    seaography::register_entities!(
-        builder,
-        // List all models we want to include in the GraphQL endpoint here
-        [files, notes, users]
-    );
+    let builder = crate::models::_entities::register_entity_modules(builder);
     // Configure async GraphQL limits
     let schema = builder
         .schema_builder()

diff --git a/examples/react_admin/backend/src/models/_entities/mod.rs b/examples/react_admin/backend/src/models/_entities/mod.rs
@@ -5,3 +5,5 @@ pub mod prelude;
 pub mod files;
 pub mod notes;
 pub mod users;
+
+seaography::register_entity_modules!([files, notes, users]);
diff --git a/examples/seaography_example/graphql/src/entities/mod.rs b/examples/seaography_example/graphql/src/entities/mod.rs
@@ -6,3 +6,5 @@ pub mod baker;
 pub mod bakery;
 pub mod cake;
 pub mod cake_baker;
+
+seaography::register_entity_modules!([baker, bakery, cake, cake_baker]);
diff --git a/examples/seaography_example/graphql/src/query_root.rs b/examples/seaography_example/graphql/src/query_root.rs
@@ -11,8 +11,8 @@ pub fn schema(
     depth: Option<usize>,
     complexity: Option<usize>,
 ) -> Result<Schema, SchemaError> {
-    let mut builder = Builder::new(&CONTEXT, database.clone());
-    seaography::register_entities!(builder, [baker, bakery, cake, cake_baker,]);
+    let builder = Builder::new(&CONTEXT, database.clone());
+    let builder = crate::entities::register_entity_modules(builder);
     let schema = builder.schema_builder();
     let schema = if let Some(depth) = depth {
         schema.limit_depth(depth)

diff --git a/sea-orm-codegen/src/entity/writer.rs b/sea-orm-codegen/src/entity/writer.rs
@@ -171,7 +171,7 @@ impl EntityWriter {
     pub fn generate(self, context: &EntityWriterContext) -> WriterOutput {
         let mut files = Vec::new();
         files.extend(self.write_entities(context));
-        files.push(self.write_index_file(context.lib));
+        files.push(self.write_index_file(context.lib, context.seaography));
         files.push(self.write_prelude());
         if !self.enums.is_empty() {
             files.push(self.write_sea_orm_active_enums(
@@ -245,7 +245,7 @@ impl EntityWriter {
             .collect()
     }
 
-    pub fn write_index_file(&self, lib: bool) -> OutputFile {
+    pub fn write_index_file(&self, lib: bool, seaography: bool) -> OutputFile {
         let mut lines = Vec::new();
         Self::write_doc_comment(&mut lines);
         let code_blocks: Vec<TokenStream> = self.entities.iter().map(Self::gen_mod).collect();
@@ -266,6 +266,12 @@ impl EntityWriter {
             );
         }
 
+        if seaography {
+            lines.push("".to_owned());
+            let ts = Self::gen_seaography_entity_mod(&self.entities, &self.enums);
+            Self::write(&mut lines, vec![ts]);
+        }
+
         let file_name = match lib {
             true => "lib.rs".to_owned(),
             false => "mod.rs".to_owned(),
@@ -704,6 +710,48 @@ impl EntityWriter {
         }
     }
 
+    pub fn gen_seaography_entity_mod(
+        entities: &[Entity],
+        enums: &BTreeMap<String, ActiveEnum>,
+    ) -> TokenStream {
+        let mut ts = TokenStream::new();
+        for entity in entities {
+            let table_name_snake_case_ident = format_ident!(
+                "{}",
+                escape_rust_keyword(entity.get_table_name_snake_case_ident())
+            );
+            ts = quote! {
+                #ts
+                #table_name_snake_case_ident,
+            }
+        }
+        ts = quote! {
+            seaography::register_entity_modules!([
+                #ts
+            ]);
+        };
+
+        let mut enum_ts = TokenStream::new();
+        for active_enum in enums.values() {
+            let enum_name = &active_enum.enum_name.to_string();
+            let enum_iden = format_ident!("{}", enum_name.to_upper_camel_case());
+            enum_ts = quote! {
+                #enum_ts
+                sea_orm_active_enums::#enum_iden
+            }
+        }
+        if !enum_ts.is_empty() {
+            ts = quote! {
+                #ts
+
+                seaography::register_active_enums!([
+                    #enum_ts
+                ]);
+            };
+        }
+        ts
+    }
+
     pub fn gen_prelude_use(entity: &Entity) -> TokenStream {
         let table_name_snake_case_ident = entity.get_table_name_snake_case_ident();
         let table_name_camel_case_ident = entity.get_table_name_camel_case_ident();