Add print-schema option allow_tables_in_same_query

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 
 

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")

diesel_cli/src/config.rs

@@ -166,6 +166,8 @@ impl Config {
                 let with_docs_with_indices = get_values_with_indices::<bool>(matches, "with-docs")?;
                 let with_docs_config_with_indices =
                     get_values_with_indices::<String>(matches, "with-docs-config")?;
+                let allow_tables_in_same_query_with_indices =
+                    get_values_with_indices::<String>(matches, "allow-tables-in-same-query")?;
                 let patch_file_with_indices =
                     get_values_with_indices::<PathBuf>(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::<String>("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::<String>("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,

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<W: IoWrite>(
     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<TableData>,
     fk_constraints: Vec<ForeignKeyConstraint>,
     with_docs: DocConfig,
+    allow_tables_in_same_query: AllowTablesInSameQuery,
     import_types: Option<&'a [String]>,
     custom_types_for_tables: Option<CustomTypesForTables>,
 }
@@ -574,16 +595,33 @@ 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()]
+            }
+        };
+        is_first = true;
+        for tables in table_groups {
+            if tables.len() < 2 {
+                continue;
+            }
+
+            if is_first {
+                writeln!(f)?;
+                is_first = false;
+            }
+            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)?;
+                for table in tables {
+                    if table.rust_name == table.sql_name {
+                        writeln!(out, "{},", table.sql_name)?;
                     } else {
-                        writeln!(out, "{},", table.name.rust_name)?;
+                        writeln!(out, "{},", table.rust_name)?;
                     }
                 }
             }
@@ -594,6 +632,64 @@ 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<Vec<&'a TableName>> {
+    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.contains(related_name) {
+                    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.
+        components.push(component);
+    }
+
+    components
+}
+
 struct TableDefinition<'a> {
     table: &'a TableData,
     with_docs: DocConfig,
@@ -832,6 +928,7 @@ impl DocConfig {
         "no-doc-comments",
     ];
 }
+
 impl<'de> Deserialize<'de> for DocConfig {
     fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
     where
@@ -881,7 +978,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<Self, Self::Err> {
         Ok(match s {
@@ -899,3 +997,21 @@ impl std::str::FromStr for DocConfig {
         })
     }
 }
+
+impl str::FromStr for AllowTablesInSameQuery {
+    type Err = &'static str;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        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`",
+                )
+            }
+        })
+    }
+}

diesel_cli/tests/print_schema.rs

@@ -347,6 +347,15 @@ fn print_schema_comments_dont_fallback_on_generated() {
     )
 }
 
+#[test]
+#[cfg(any(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![])

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"

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,);

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);