From ba96917a18be72fb32a0609b0643ce827175c9c5 Mon Sep 17 00:00:00 2001
From: Jordy Ruiter <jordy.ruiter@gmail.com>
Date: Sun, 3 Oct 2021 15:45:07 +0900
Subject: [PATCH] Add tables option to CLI generate entity

---
 sea-orm-cli/src/cli.rs  |  9 +++++++++
 sea-orm-cli/src/main.rs | 10 ++++++++++
 2 files changed, 19 insertions(+)

diff --git a/sea-orm-cli/src/cli.rs b/sea-orm-cli/src/cli.rs
index 100c747c..8fe9c4a4 100644
--- a/sea-orm-cli/src/cli.rs
+++ b/sea-orm-cli/src/cli.rs
@@ -41,6 +41,15 @@ pub fn build_cli() -> App<'static, 'static> {
                         .help("Generate entity file for hidden tables (i.e. table name starts with an underscore)")
                         .takes_value(false),
                 )
+                .arg(
+                    Arg::with_name("TABLES")
+                        .long("tables")
+                        .short("t")
+                        .use_delimiter(true)
+                        .help("Generate entity file for specified tables only (comma seperated)")
+                        .takes_value(true)
+                        .conflicts_with("INCLUDE_HIDDEN_TABLES"),
+                )
                 .arg(
                     Arg::with_name("EXPANDED_FORMAT")
                         .long("expanded-format")
diff --git a/sea-orm-cli/src/main.rs b/sea-orm-cli/src/main.rs
index 527e1264..c318a510 100644
--- a/sea-orm-cli/src/main.rs
+++ b/sea-orm-cli/src/main.rs
@@ -26,7 +26,15 @@ async fn run_generate_command(matches: &ArgMatches<'_>) -> Result<(), Box<dyn Er
             let url = args.value_of("DATABASE_URL").unwrap();
             let output_dir = args.value_of("OUTPUT_DIR").unwrap();
             let include_hidden_tables = args.is_present("INCLUDE_HIDDEN_TABLES");
+            let tables = args.values_of("TABLES").unwrap_or_default().collect::<Vec<_>>();
             let expanded_format = args.is_present("EXPANDED_FORMAT");
+            let filter_tables = |table: &str| -> bool {
+                if tables.len() > 0 {
+                    return tables.contains(&table);
+                }
+                
+                true
+            };
             let filter_hidden_tables = |table: &str| -> bool {
                 if include_hidden_tables {
                     true
@@ -53,6 +61,7 @@ async fn run_generate_command(matches: &ArgMatches<'_>) -> Result<(), Box<dyn Er
                 schema
                     .tables
                     .into_iter()
+                    .filter(|schema| filter_tables(&schema.info.name))
                     .filter(|schema| filter_hidden_tables(&schema.info.name))
                     .map(|schema| schema.write())
                     .collect()
@@ -67,6 +76,7 @@ async fn run_generate_command(matches: &ArgMatches<'_>) -> Result<(), Box<dyn Er
                 schema
                     .tables
                     .into_iter()
+                    .filter(|schema| filter_tables(&schema.info.name))
                     .filter(|schema| filter_hidden_tables(&schema.info.name))
                     .map(|schema| schema.write())
                     .collect()