diff --git a/CHANGELOG.md b/CHANGELOG.md index ad4d64531bae..1b37e9de64ae 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -26,6 +26,7 @@ in a way that makes the pools suitable for use in parallel tests. * Added `Json` and `Jsonb` support for the SQLite backend. * Added a `#[diesel::declare_sql_function]` attribute macro to easily define support for multiple sql functions at once via an `extern "SQL"` block +* Support `[print_schema] allow_tables_in_same_query = "fk_related_tables"` to generate separate `allow_tables_to_appear_in_same_query!` calls containing only tables that are related through foreign keys. Without this option, or with an explicit value of `"all_tables"`, a single call is generated that contains all tables. ### Fixed diff --git a/diesel_cli/src/cli.rs b/diesel_cli/src/cli.rs index 345e3e880e20..a560289391cb 100644 --- a/diesel_cli/src/cli.rs +++ b/diesel_cli/src/cli.rs @@ -280,6 +280,14 @@ pub fn build_cli() -> Command { .action(ArgAction::Append) .value_parser(PossibleValuesParser::new(print_schema::DocConfig::VARIANTS_STR)), ) + .arg( + Arg::new("allow-tables-in-same-query") + .long("allow-tables-in-same-query") + .help("Group tables in allow_tables_to_appear_in_same_query!().") + .num_args(1) + .action(ArgAction::Append) + .value_parser(PossibleValuesParser::new(["fk_related_tables", "all_tables"])), + ) .arg( Arg::new("column-sorting") .long("column-sorting") diff --git a/diesel_cli/src/config.rs b/diesel_cli/src/config.rs index fa6a8d9ab463..b7ab9b840b03 100644 --- a/diesel_cli/src/config.rs +++ b/diesel_cli/src/config.rs @@ -166,6 +166,8 @@ impl Config { let with_docs_with_indices = get_values_with_indices::(matches, "with-docs")?; let with_docs_config_with_indices = get_values_with_indices::(matches, "with-docs-config")?; + let allow_tables_in_same_query_with_indices = + get_values_with_indices::(matches, "allow-tables-in-same-query")?; let patch_file_with_indices = get_values_with_indices::(matches, "patch-file")?; let column_sorting_with_indices = @@ -234,6 +236,20 @@ impl Config { })?; } + if let Some(allow_tables_in_same_query) = + allow_tables_in_same_query_with_indices + .clone() + .and_then(|v| v.range(boundary).nth(0).map(|v| v.1.clone())) + { + print_schema.allow_tables_in_same_query = + allow_tables_in_same_query.parse().map_err(|_| { + crate::errors::Error::UnsupportedFeature(format!( + "Invalid `allow_tables_in_same_query` config mode: \ + {allow_tables_in_same_query}" + )) + })?; + } + if let Some(sorting) = column_sorting_with_indices .clone() .and_then(|v| v.range(boundary).nth(0).map(|v| v.1.clone())) @@ -325,6 +341,18 @@ impl Config { })?; } + if let Some(allow_tables_in_same_query) = + matches.get_one::("allow-tables-in-same-query") + { + config.allow_tables_in_same_query = + allow_tables_in_same_query.parse().map_err(|_| { + crate::errors::Error::UnsupportedFeature(format!( + "Invalid `allow_tables_in_same_query` config mode: \ + {allow_tables_in_same_query}" + )) + })?; + } + if let Some(sorting) = matches.get_one::("column-sorting") { match sorting as &str { "ordinal_position" => config.column_sorting = ColumnSorting::OrdinalPosition, @@ -426,6 +454,8 @@ pub struct PrintSchema { #[serde(default)] pub with_docs: print_schema::DocConfig, #[serde(default)] + pub allow_tables_in_same_query: print_schema::AllowTablesInSameQuery, + #[serde(default)] pub filter: Filtering, #[serde(default)] pub column_sorting: ColumnSorting, diff --git a/diesel_cli/src/print_schema.rs b/diesel_cli/src/print_schema.rs index 6882bf49c002..51e102918cf0 100644 --- a/diesel_cli/src/print_schema.rs +++ b/diesel_cli/src/print_schema.rs @@ -3,10 +3,11 @@ use crate::database::{Backend, InferConnection}; use crate::infer_schema_internals::*; use serde::{Deserialize, Serialize}; -use std::collections::HashSet; +use std::collections::{BTreeSet, HashSet}; use std::fmt::{self, Display, Formatter, Write}; use std::io::Write as IoWrite; use std::process; +use std::str; const SCHEMA_HEADER: &str = "// @generated automatically by Diesel CLI.\n"; @@ -42,6 +43,24 @@ impl Default for DocConfig { } } +/// How to group tables in `allow_tables_to_appear_in_same_query!()`. +#[derive(Debug, Deserialize, Serialize, Clone, Copy)] +pub enum AllowTablesInSameQuery { + /// Group by foreign key relations + #[serde(rename = "fk_related_tables")] + FkRelatedTables, + /// List all tables in invocation + #[serde(rename = "all_tables")] + AllTables, +} + +#[allow(clippy::derivable_impls)] // that's not supported on rust 1.65 +impl Default for AllowTablesInSameQuery { + fn default() -> Self { + AllowTablesInSameQuery::AllTables + } +} + pub fn run_print_schema( connection: &mut InferConnection, config: &config::PrintSchema, @@ -237,6 +256,7 @@ pub fn output_schema( tables: table_data, fk_constraints: foreign_keys, with_docs: config.with_docs, + allow_tables_in_same_query: config.allow_tables_in_same_query, custom_types_for_tables: columns_custom_types.map(|custom_types_sorted| { CustomTypesForTables { backend, @@ -538,6 +558,7 @@ struct TableDefinitions<'a> { tables: Vec, fk_constraints: Vec, with_docs: DocConfig, + allow_tables_in_same_query: AllowTablesInSameQuery, import_types: Option<&'a [String]>, custom_types_for_tables: Option, } @@ -574,17 +595,31 @@ impl Display for TableDefinitions<'_> { writeln!(f, "{}", Joinable(foreign_key))?; } - if self.tables.len() > 1 { - write!(f, "\ndiesel::allow_tables_to_appear_in_same_query!(")?; + let table_groups = match self.allow_tables_in_same_query { + AllowTablesInSameQuery::FkRelatedTables => { + foreign_key_table_groups(&self.tables, &self.fk_constraints) + } + AllowTablesInSameQuery::AllTables => { + vec![self.tables.iter().map(|table| &table.name).collect()] + } + }; + for (table_group_index, table_group) in table_groups + .into_iter() + .filter(|table_group| table_group.len() >= 2) + .enumerate() + { + if table_group_index == 0 { + writeln!(f)?; + } + write!(f, "diesel::allow_tables_to_appear_in_same_query!(")?; { let mut out = PadAdapter::new(f); writeln!(out)?; - for table in &self.tables { - if table.name.rust_name == table.name.sql_name { - writeln!(out, "{},", table.name.sql_name)?; - } else { - writeln!(out, "{},", table.name.rust_name)?; + for (table_index, table) in table_group.into_iter().enumerate() { + if table_index != 0 { + write!(out, ", ")?; } + write!(out, "{}", table.rust_name)?; } } writeln!(f, ");")?; @@ -594,6 +629,71 @@ impl Display for TableDefinitions<'_> { } } +/// Calculates groups of tables that are related by foreign key. +/// +/// Given the graph of all tables and their foreign key relations, this returns the set of connected +/// components of that graph. +fn foreign_key_table_groups<'a>( + tables: &'a [TableData], + fk_constraints: &'a [ForeignKeyConstraint], +) -> Vec> { + let mut visited = BTreeSet::new(); + let mut components = vec![]; + + // Find connected components in table graph. For the intended purpose of this function, we treat + // the foreign key relation as being symmetrical, i.e. we are operating on the undirected graph. + // + // The algorithm is not optimized and suffers from repeated lookups in the foreign key list, but + // it should be sufficient for typical table counts from a few dozen up to a few hundred tables. + for table in tables { + let name = &table.name; + if visited.contains(name) { + // This table is already part of another connected component. + continue; + } + + visited.insert(name); + let mut component = vec![]; + let mut pending = vec![name]; + + // Start a depth-first search with the current table name, walking the foreign key relations + // in both directions. + while let Some(name) = pending.pop() { + component.push(name); + + let mut visit = |related_name: &'a TableName| { + if visited.insert(related_name) { + pending.push(related_name); + } + }; + + // Visit all remaining child tables that have this table as parent. + for foreign_key in fk_constraints.iter().filter(|fk| fk.parent_table == *name) { + visit(&foreign_key.child_table); + } + + // Visit all remaining parent tables that have this table as child. + for foreign_key in fk_constraints.iter().filter(|fk| fk.child_table == *name) { + visit(&foreign_key.parent_table); + } + } + + // The component contains all tables that are reachable in either direction from the current + // table. Sort that list by table name to ensure a stable output that does not depend on the + // algorithm's specific implementation. + component.sort(); + + components.push(component); + } + + // Sort the list of components to ensure a stable output that does not depend on the algorithm's + // specific implementation. This sorts the list of components by the name of the first tables in + // each component. + components.sort(); + + components +} + struct TableDefinition<'a> { table: &'a TableData, with_docs: DocConfig, @@ -832,6 +932,7 @@ impl DocConfig { "no-doc-comments", ]; } + impl<'de> Deserialize<'de> for DocConfig { fn deserialize(deserializer: D) -> Result where @@ -881,7 +982,8 @@ impl<'de> Deserialize<'de> for DocConfig { deserializer.deserialize_any(DocConfigVisitor) } } -impl std::str::FromStr for DocConfig { + +impl str::FromStr for DocConfig { type Err = &'static str; fn from_str(s: &str) -> Result { Ok(match s { @@ -899,3 +1001,21 @@ impl std::str::FromStr for DocConfig { }) } } + +impl str::FromStr for AllowTablesInSameQuery { + type Err = &'static str; + + fn from_str(s: &str) -> Result { + Ok(match s { + "fk_related_tables" => AllowTablesInSameQuery::FkRelatedTables, + "all_tables" => AllowTablesInSameQuery::AllTables, + _ => { + return Err( + "Unknown variant for `allow_tables_in_same_query` config, expected one of: \ + `fk_related_tables`, \ + `all_tables`", + ) + } + }) + } +} diff --git a/diesel_cli/tests/print_schema.rs b/diesel_cli/tests/print_schema.rs index d720eded3873..963f33dd5c19 100644 --- a/diesel_cli/tests/print_schema.rs +++ b/diesel_cli/tests/print_schema.rs @@ -347,6 +347,15 @@ fn print_schema_comments_dont_fallback_on_generated() { ) } +#[test] +#[cfg(feature = "postgres")] +fn print_schema_allow_tables_in_same_query() { + test_print_schema( + "print_schema_allow_tables_in_same_query", + vec!["--allow-tables-in-same-query", "fk_related_tables"], + ) +} + #[test] fn print_schema_reserved_names() { test_print_schema("print_schema_reserved_name_mitigation_issue_3404", vec![]) diff --git a/diesel_cli/tests/print_schema/print_schema_allow_tables_in_same_query/diesel.toml b/diesel_cli/tests/print_schema/print_schema_allow_tables_in_same_query/diesel.toml new file mode 100644 index 000000000000..4919dc4153e9 --- /dev/null +++ b/diesel_cli/tests/print_schema/print_schema_allow_tables_in_same_query/diesel.toml @@ -0,0 +1,3 @@ +[print_schema] +file = "src/schema.rs" +allow_tables_in_same_query = "fk_related_tables" diff --git a/diesel_cli/tests/print_schema/print_schema_allow_tables_in_same_query/postgres/expected.snap b/diesel_cli/tests/print_schema/print_schema_allow_tables_in_same_query/postgres/expected.snap new file mode 100644 index 000000000000..ae28fda95549 --- /dev/null +++ b/diesel_cli/tests/print_schema/print_schema_allow_tables_in_same_query/postgres/expected.snap @@ -0,0 +1,59 @@ +--- +source: diesel_cli/tests/print_schema.rs +assertion_line: 500 +description: "Test: print_schema_allow_tables_in_same_query" +snapshot_kind: text +--- +// @generated automatically by Diesel CLI. + +diesel::table! { + bikes (id) { + id -> Int4, + } +} + +diesel::table! { + cars (id) { + id -> Int4, + } +} + +diesel::table! { + comments (id) { + id -> Int4, + post_id -> Int4, + } +} + +diesel::table! { + posts (id) { + id -> Int4, + user_id -> Int4, + } +} + +diesel::table! { + sessions (id) { + id -> Int4, + } +} + +diesel::table! { + transactions (id) { + id -> Int4, + session_id -> Int4, + } +} + +diesel::table! { + users (id) { + id -> Int4, + } +} + +diesel::joinable!(comments -> posts (post_id)); +diesel::joinable!(posts -> users (user_id)); +diesel::joinable!(transactions -> sessions (session_id)); + +diesel::allow_tables_to_appear_in_same_query!(comments, posts, users); +diesel::allow_tables_to_appear_in_same_query!(sessions, transactions); diff --git a/diesel_cli/tests/print_schema/print_schema_allow_tables_in_same_query/postgres/schema.sql b/diesel_cli/tests/print_schema/print_schema_allow_tables_in_same_query/postgres/schema.sql new file mode 100644 index 000000000000..175f82a4b761 --- /dev/null +++ b/diesel_cli/tests/print_schema/print_schema_allow_tables_in_same_query/postgres/schema.sql @@ -0,0 +1,12 @@ +-- Three related tables. +CREATE TABLE users (id SERIAL PRIMARY KEY); +CREATE TABLE posts (id SERIAL PRIMARY KEY, user_id INTEGER NOT NULL REFERENCES users); +CREATE TABLE comments (id SERIAL PRIMARY KEY, post_id INTEGER NOT NULL REFERENCES posts); + +-- Two related tables. +CREATE TABLE sessions (id SERIAL PRIMARY KEY); +CREATE TABLE transactions (id SERIAL PRIMARY KEY, session_id INTEGER NOT NULL REFERENCES sessions); + +-- Unrelated tables. +CREATE TABLE cars (id SERIAL PRIMARY KEY); +CREATE TABLE bikes (id SERIAL PRIMARY KEY);