Merge pull request diesel-rs#3725 from gmanninglive/fix/3717

Fix/3717 filter table on `diesel migration --diff-schema`

diesel_cli/src/cli.rs

@@ -98,6 +98,8 @@ pub fn build_cli() -> Command {
                 )
                 .arg(
                     Arg::new("MIGRATION_NAME")
+                        .index(1)
+                        .num_args(1)
                         .help("The name of the migration to create.")
                         .required(true),
                 )
@@ -145,6 +147,29 @@ pub fn build_cli() -> Command {
                         .default_missing_value("NOT_SET")
                         .num_args(0..=1)
                         .require_equals(true),
+                )
+                .arg(
+                    Arg::new("table-name")
+                        .index(2)
+                        .num_args(1..)
+                        .action(clap::ArgAction::Append)
+                        .help("Table names to filter."),
+                )
+                .arg(
+                    Arg::new("only-tables")
+                        .short('o')
+                        .long("only-tables")
+                        .action(ArgAction::SetTrue)
+                        .help("Only include tables from table-name that matches regexp.")
+                        .conflicts_with("except-tables"),
+                )
+                .arg(
+                    Arg::new("except-tables")
+                        .short('e')
+                        .long("except-tables")
+                        .action(ArgAction::SetTrue)
+                        .help("Exclude tables from table-name that matches regex.")
+                        .conflicts_with("only-tables"),
                 ),
         )
         .subcommand_required(true)
@@ -198,7 +223,7 @@ pub fn build_cli() -> Command {
                 .index(1)
                 .num_args(1..)
                 .action(clap::ArgAction::Append)
-                .help("Table names to filter (default only-tables if not empty)."),
+                .help("Table names to filter."),
         )
         .arg(
             Arg::new("only-tables")

diesel_cli/src/config.rs

@@ -1,11 +1,14 @@
 use clap::ArgMatches;
-use serde::Deserialize;
-use std::env;
+use serde::de::{self, MapAccess, Visitor};
+use serde::{Deserialize, Deserializer};
+use serde_regex::Serde as RegexWrapper;
 use std::error::Error;
 use std::fs;
 use std::path::{Path, PathBuf};
+use std::{env, fmt};
 
 use super::find_project_root;
+use crate::infer_schema_internals::TableName;
 use crate::print_schema;
 use crate::print_schema::ColumnSorting;
 
@@ -46,6 +49,26 @@ impl Config {
             migration.set_relative_path_base(base);
         }
     }
+
+    pub fn set_filter(
+        mut self,
+        matches: &ArgMatches,
+    ) -> Result<Self, Box<dyn Error + Send + Sync + 'static>> {
+        let table_names = matches
+            .get_many::<String>("table-name")
+            .unwrap_or_default()
+            .map(|table_name_regex| regex::Regex::new(table_name_regex).map(Into::into))
+            .collect::<Result<Vec<Regex>, _>>()
+            .map_err(|e| format!("invalid argument for table filtering regex: {e}"))?;
+
+        if matches.get_flag("only-tables") {
+            self.print_schema.filter = Filtering::OnlyTables(table_names)
+        } else if matches.get_flag("except-tables") {
+            self.print_schema.filter = Filtering::ExceptTables(table_names)
+        }
+
+        Ok(self)
+    }
 }
 
 #[derive(Default, Deserialize)]
@@ -56,7 +79,7 @@ pub struct PrintSchema {
     #[serde(default)]
     pub with_docs: print_schema::DocConfig,
     #[serde(default)]
-    pub filter: print_schema::Filtering,
+    pub filter: Filtering,
     #[serde(default)]
     pub column_sorting: ColumnSorting,
     #[serde(default)]
@@ -129,3 +152,86 @@ impl MigrationsDirectory {
         }
     }
 }
+
+type Regex = RegexWrapper<::regex::Regex>;
+
+#[derive(Clone)]
+pub enum Filtering {
+    OnlyTables(Vec<Regex>),
+    ExceptTables(Vec<Regex>),
+    None,
+}
+
+#[allow(clippy::derivable_impls)] // that's not supported on rust 1.65
+impl Default for Filtering {
+    fn default() -> Self {
+        Filtering::None
+    }
+}
+
+impl Filtering {
+    pub fn should_ignore_table(&self, name: &TableName) -> bool {
+        use self::Filtering::*;
+
+        match *self {
+            OnlyTables(ref regexes) => !regexes.iter().any(|regex| regex.is_match(&name.sql_name)),
+            ExceptTables(ref regexes) => regexes.iter().any(|regex| regex.is_match(&name.sql_name)),
+            None => false,
+        }
+    }
+}
+
+impl<'de> Deserialize<'de> for Filtering {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        struct FilteringVisitor;
+
+        impl<'de> Visitor<'de> for FilteringVisitor {
+            type Value = Filtering;
+
+            fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
+                f.write_str("either only_tables or except_tables")
+            }
+
+            fn visit_map<V>(self, mut map: V) -> Result<Self::Value, V::Error>
+            where
+                V: MapAccess<'de>,
+            {
+                let mut only_tables = None::<Vec<Regex>>;
+                let mut except_tables = None::<Vec<Regex>>;
+                while let Some(key) = map.next_key::<String>()? {
+                    match &key as &str {
+                        "only_tables" => {
+                            if only_tables.is_some() {
+                                return Err(de::Error::duplicate_field("only_tables"));
+                            }
+                            only_tables = Some(map.next_value()?);
+                        }
+                        "except_tables" => {
+                            if except_tables.is_some() {
+                                return Err(de::Error::duplicate_field("except_tables"));
+                            }
+                            except_tables = Some(map.next_value()?);
+                        }
+                        _ => {
+                            return Err(de::Error::unknown_field(
+                                &key,
+                                &["only_tables", "except_tables"],
+                            ))
+                        }
+                    }
+                }
+                match (only_tables, except_tables) {
+                    (Some(t), None) => Ok(Filtering::OnlyTables(t)),
+                    (None, Some(t)) => Ok(Filtering::ExceptTables(t)),
+                    (None, None) => Ok(Filtering::None),
+                    _ => Err(de::Error::duplicate_field("only_tables except_tables")),
+                }
+            }
+        }
+
+        deserializer.deserialize_map(FilteringVisitor)
+    }
+}

diesel_cli/src/infer_schema_internals/inference.rs

@@ -4,6 +4,9 @@ use diesel::result::Error::NotFound;
 
 use super::data_structures::*;
 use super::table_data::*;
+
+use crate::config::Filtering;
+
 use crate::database::InferConnection;
 use crate::print_schema::{ColumnSorting, DocConfig};
 
@@ -124,6 +127,13 @@ pub fn load_table_names(
     }
 }
 
+pub fn filter_table_names(table_names: Vec<TableName>, table_filter: &Filtering) -> Vec<TableName> {
+    table_names
+        .into_iter()
+        .filter(|t| !table_filter.should_ignore_table(t))
+        .collect::<_>()
+}
+
 fn get_table_comment(
     conn: &mut InferConnection,
     table: &TableName,

diesel_cli/src/main.rs

@@ -31,7 +31,6 @@ use database::InferConnection;
 use diesel::backend::Backend;
 use diesel::Connection;
 use diesel_migrations::{FileBasedMigrations, HarnessWithOutput, MigrationHarness};
-use regex::Regex;
 use std::error::Error;
 use std::fmt::Display;
 use std::io::stdout;
@@ -220,25 +219,12 @@ fn run_infer_schema(matches: &ArgMatches) -> Result<(), Box<dyn Error + Send + S
     use crate::print_schema::*;
 
     let mut conn = InferConnection::from_matches(matches);
-    let mut config = Config::read(matches)?.print_schema;
+    let mut config = Config::read(matches)?.set_filter(matches)?.print_schema;
 
     if let Some(schema_name) = matches.get_one::<String>("schema") {
         config.schema = Some(schema_name.clone())
     }
 
-    let filter = matches
-        .get_many::<String>("table-name")
-        .unwrap_or_default()
-        .map(|table_name_regex| Regex::new(table_name_regex).map(Into::into))
-        .collect::<Result<_, _>>()
-        .map_err(|e| format!("invalid argument for table filtering regex: {e}"));
-
-    if matches.get_flag("only-tables") {
-        config.filter = Filtering::OnlyTables(filter?)
-    } else if matches.get_flag("except-tables") {
-        config.filter = Filtering::ExceptTables(filter?)
-    }
-
     if matches.get_flag("with-docs") {
         config.with_docs = DocConfig::DatabaseCommentsFallbackToAutoGeneratedDocComment;
     }
@@ -282,16 +268,16 @@ fn regenerate_schema_if_file_specified(
 ) -> Result<(), Box<dyn Error + Send + Sync + 'static>> {
     use std::io::Read;
 
-    let config = Config::read(matches)?;
-    if let Some(ref path) = config.print_schema.file {
+    let config = Config::read(matches)?.print_schema;
+    if let Some(ref path) = config.file {
         let mut connection = InferConnection::from_matches(matches);
         if let Some(parent) = path.parent() {
             fs::create_dir_all(parent)?;
         }
 
         if matches.get_flag("LOCKED_SCHEMA") {
             let mut buf = Vec::new();
-            print_schema::run_print_schema(&mut connection, &config.print_schema, &mut buf)?;
+            print_schema::run_print_schema(&mut connection, &config, &mut buf)?;
 
             let mut old_buf = Vec::new();
             let mut file = fs::File::open(path)?;
@@ -309,7 +295,7 @@ fn regenerate_schema_if_file_specified(
             use std::io::Write;
 
             let mut file = fs::File::create(path)?;
-            let schema = print_schema::output_schema(&mut connection, &config.print_schema)?;
+            let schema = print_schema::output_schema(&mut connection, &config)?;
             file.write_all(schema.as_bytes())?;
         }
     }

diesel_cli/src/migrations/diff_schema.rs

@@ -13,7 +13,8 @@ use syn::visit::Visit;
 use crate::config::Config;
 use crate::database::InferConnection;
 use crate::infer_schema_internals::{
-    ColumnDefinition, ColumnType, ForeignKeyConstraint, TableData, TableName,
+    filter_table_names, load_table_names, ColumnDefinition, ColumnType, ForeignKeyConstraint,
+    TableData, TableName,
 };
 use crate::print_schema::DocConfig;
 
@@ -27,10 +28,12 @@ fn compatible_type_list() -> HashMap<&'static str, Vec<&'static str>> {
 }
 
 pub fn generate_sql_based_on_diff_schema(
-    _config: Config,
+    config: Config,
     matches: &ArgMatches,
     schema_file_path: &Path,
 ) -> Result<(String, String), Box<dyn Error + Send + Sync>> {
+    let config = config.set_filter(matches)?;
+
     let project_root = crate::find_project_root()?;
 
     let schema_path = project_root.join(schema_file_path);
@@ -43,7 +46,11 @@ pub fn generate_sql_based_on_diff_schema(
 
     tables_from_schema.visit_file(&syn_file);
     let mut conn = InferConnection::from_matches(matches);
-    let tables_from_database = crate::infer_schema_internals::load_table_names(&mut conn, None)?;
+    let tables_from_database = filter_table_names(
+        load_table_names(&mut conn, None)?,
+        &config.print_schema.filter,
+    );
+
     let foreign_keys =
         crate::infer_schema_internals::load_foreign_key_constraints(&mut conn, None)?;
     let foreign_key_map = foreign_keys.into_iter().fold(HashMap::new(), |mut acc, t| {

diesel_cli/src/migrations/mod.rs

@@ -88,7 +88,7 @@ pub(super) fn run_migration_command(
                 if let Some(diff_schema) = diff_schema {
                     self::diff_schema::generate_sql_based_on_diff_schema(
                         config,
-                        matches,
+                        args,
                         &diff_schema,
                     )?
                 } else {