Merge pull request diesel-rs#1637 from diesel-rs/sg-print-schema-custom-types

Allow specifying type imports for `diesel print-schema`

Author: sgrif

diesel_cli/src/cli.rs

@@ -134,6 +134,14 @@ pub fn build_cli() -> App<'static, 'static> {
                 .long("patch-file")
                 .takes_value(true)
                 .help("A unified diff file to be applied to the final schema"),
+        )
+        .arg(
+            Arg::with_name("import-types")
+                .long("import-types")
+                .takes_value(true)
+                .multiple(true)
+                .number_of_values(1)
+                .help("A list of types to import for every table, separated by commas"),
         );
 
     let config_arg = Arg::with_name("CONFIG_FILE")

diesel_cli/src/config.rs

@@ -50,10 +50,16 @@ pub struct PrintSchema {
     pub schema: Option<String>,
     #[serde(default)]
     pub patch_file: Option<PathBuf>,
+    #[serde(default)]
+    pub import_types: Option<Vec<String>>,
 }
 
 impl PrintSchema {
     pub fn schema_name(&self) -> Option<&str> {
         self.schema.as_ref().map(|s| &**s)
     }
+
+    pub fn import_types(&self) -> Option<&[String]> {
+        self.import_types.as_ref().map(|v| &**v)
+    }
 }

diesel_cli/src/main.rs

@@ -366,6 +366,11 @@ fn run_infer_schema(matches: &ArgMatches) -> Result<(), Box<Error>> {
         config.patch_file = Some(PathBuf::from(path));
     }
 
+    if let Some(types) = matches.values_of("import-types") {
+        let types = types.map(String::from).collect();
+        config.import_types = Some(types);
+    }
+
     run_print_schema(&database_url, &config)?;
     Ok(())
 }

diesel_cli/src/print_schema.rs

@@ -70,6 +70,7 @@ pub fn output_schema(
         tables: table_data,
         fk_constraints: foreign_keys,
         include_docs: config.with_docs,
+        import_types: config.import_types(),
     };
 
     if let Some(schema_name) = config.schema_name() {
@@ -96,7 +97,7 @@ pub fn output_schema(
     Ok(())
 }
 
-struct ModuleDefinition<'a>(&'a str, TableDefinitions);
+struct ModuleDefinition<'a>(&'a str, TableDefinitions<'a>);
 
 impl<'a> Display for ModuleDefinition<'a> {
     fn fmt(&self, f: &mut Formatter) -> fmt::Result {
@@ -110,13 +111,14 @@ impl<'a> Display for ModuleDefinition<'a> {
     }
 }
 
-struct TableDefinitions {
+struct TableDefinitions<'a> {
     tables: Vec<TableData>,
     fk_constraints: Vec<ForeignKeyConstraint>,
     include_docs: bool,
+    import_types: Option<&'a [String]>,
 }
 
-impl Display for TableDefinitions {
+impl<'a> Display for TableDefinitions<'a> {
     fn fmt(&self, f: &mut Formatter) -> fmt::Result {
         let mut is_first = true;
         for table in &self.tables {
@@ -131,6 +133,7 @@ impl Display for TableDefinitions {
                 TableDefinition {
                     table,
                     include_docs: self.include_docs,
+                    import_types: self.import_types,
                 }
             )?;
         }
@@ -161,6 +164,7 @@ impl Display for TableDefinitions {
 
 struct TableDefinition<'a> {
     table: &'a TableData,
+    import_types: Option<&'a [String]>,
     include_docs: bool,
 }
 
@@ -171,6 +175,13 @@ impl<'a> Display for TableDefinition<'a> {
             let mut out = PadAdapter::new(f);
             write!(out, "\n")?;
 
+            if let Some(types) = self.import_types {
+                for import in types {
+                    writeln!(out, "use {};", import)?;
+                }
+                write!(out, "\n")?;
+            }
+
             if self.include_docs {
                 for d in self.table.docs.lines() {
                     writeln!(out, "///{}{}", if d.is_empty() { "" } else { " " }, d)?;

diesel_cli/tests/print_schema.rs

@@ -82,6 +82,14 @@ fn print_schema_patch_file() {
     test_print_schema("print_schema_patch_file", vec!["--patch-file", &path]);
 }
 
+#[test]
+fn print_schema_custom_types() {
+    test_print_schema(
+        "print_schema_custom_types",
+        vec!["--import-types", "foo::*", "--import-types", "bar::*"],
+    );
+}
+
 #[cfg(feature = "sqlite")]
 const BACKEND: &str = "sqlite";
 #[cfg(feature = "postgres")]

diesel_cli/tests/print_schema/print_schema_custom_types/diesel.toml

@@ -0,0 +1,3 @@
+[print_schema]
+file = "src/schema.rs"
+import_types = ["foo::*", "bar::*"]

diesel_cli/tests/print_schema/print_schema_custom_types/mysql/expected.rs

@@ -0,0 +1,22 @@
+table! {
+    use foo::*;
+    use bar::*;
+
+    users1 (id) {
+        id -> Integer,
+    }
+}
+
+table! {
+    use foo::*;
+    use bar::*;
+
+    users2 (id) {
+        id -> Integer,
+    }
+}
+
+allow_tables_to_appear_in_same_query!(
+    users1,
+    users2,
+);

diesel_cli/tests/print_schema/print_schema_custom_types/mysql/schema.sql

@@ -0,0 +1,2 @@
+CREATE TABLE users1 (id INTEGER PRIMARY KEY);
+CREATE TABLE users2 (id INTEGER PRIMARY KEY);

diesel_cli/tests/print_schema/print_schema_custom_types/postgres/expected.rs

@@ -0,0 +1,22 @@
+table! {
+    use foo::*;
+    use bar::*;
+
+    users1 (id) {
+        id -> Int4,
+    }
+}
+
+table! {
+    use foo::*;
+    use bar::*;
+
+    users2 (id) {
+        id -> Int4,
+    }
+}
+
+allow_tables_to_appear_in_same_query!(
+    users1,
+    users2,
+);

diesel_cli/tests/print_schema/print_schema_custom_types/postgres/schema.sql

@@ -0,0 +1,2 @@
+CREATE TABLE users1 (id SERIAL PRIMARY KEY);
+CREATE TABLE users2 (id SERIAL PRIMARY KEY);