From 6816e86f4d3af520d05a4de6f1dab36c51651294 Mon Sep 17 00:00:00 2001
From: Horu <73709188+HigherOrderLogic@users.noreply.github.com>
Date: Sun, 25 Sep 2022 09:31:26 +0700
Subject: [PATCH] [cli] Add `-l`/`--lib` flag (#953)

* [cli] Add `-l`/`--lib` flag

* [cli] Change function name to reflect functionality
---
 sea-orm-cli/src/cli.rs               |  9 +++++++++
 sea-orm-cli/src/commands/generate.rs |  2 ++
 sea-orm-codegen/src/entity/writer.rs | 15 ++++++++++++---
 3 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/sea-orm-cli/src/cli.rs b/sea-orm-cli/src/cli.rs
index 36414268..001a9b9b 100644
--- a/sea-orm-cli/src/cli.rs
+++ b/sea-orm-cli/src/cli.rs
@@ -191,6 +191,15 @@ pub enum GenerateSubcommands {
             help = "The datetime crate to use for generating entities."
         )]
         date_time_crate: DateTimeCrate,
+
+        #[clap(
+            action,
+            long,
+            short = 'l',
+            default_value = "false",
+            help = "Generate index file as `lib.rs` instead of `mod.rs`."
+        )]
+        lib: bool,
     },
 }
 
diff --git a/sea-orm-cli/src/commands/generate.rs b/sea-orm-cli/src/commands/generate.rs
index 0cea31e4..1686e878 100644
--- a/sea-orm-cli/src/commands/generate.rs
+++ b/sea-orm-cli/src/commands/generate.rs
@@ -26,6 +26,7 @@ pub async fn run_generate_command(
             with_serde,
             with_copy_enums,
             date_time_crate,
+            lib,
         } => {
             if verbose {
                 let _ = tracing_subscriber::fmt()
@@ -171,6 +172,7 @@ pub async fn run_generate_command(
                 with_copy_enums,
                 date_time_crate.into(),
                 schema_name,
+                lib,
             );
             let output = EntityTransformer::transform(table_stmts)?.generate(&writer_context);
 
diff --git a/sea-orm-codegen/src/entity/writer.rs b/sea-orm-codegen/src/entity/writer.rs
index f9d6c754..be6954ca 100644
--- a/sea-orm-codegen/src/entity/writer.rs
+++ b/sea-orm-codegen/src/entity/writer.rs
@@ -42,6 +42,7 @@ pub struct EntityWriterContext {
     pub(crate) with_copy_enums: bool,
     pub(crate) date_time_crate: DateTimeCrate,
     pub(crate) schema_name: Option<String>,
+    pub(crate) lib: bool,
 }
 
 impl WithSerde {
@@ -101,6 +102,7 @@ impl EntityWriterContext {
         with_copy_enums: bool,
         date_time_crate: DateTimeCrate,
         schema_name: Option<String>,
+        lib: bool,
     ) -> Self {
         Self {
             expanded_format,
@@ -108,6 +110,7 @@ impl EntityWriterContext {
             with_copy_enums,
             date_time_crate,
             schema_name,
+            lib,
         }
     }
 }
@@ -116,7 +119,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_mod());
+        files.push(self.write_index_file(context.lib));
         files.push(self.write_prelude());
         if !self.enums.is_empty() {
             files.push(
@@ -168,7 +171,7 @@ impl EntityWriter {
             .collect()
     }
 
-    pub fn write_mod(&self) -> OutputFile {
+    pub fn write_index_file(&self, lib: 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();
@@ -188,8 +191,14 @@ impl EntityWriter {
                 }],
             );
         }
+
+        let file_name = match lib {
+            true => "lib.rs".to_owned(),
+            false => "mod.rs".to_owned(),
+        };
+
         OutputFile {
-            name: "mod.rs".to_owned(),
+            name: file_name,
             content: lines.join("\n"),
         }
     }