diff --git a/diesel_cli/src/infer_schema_internals/load_foreign_keys.sql b/diesel_cli/src/infer_schema_internals/load_foreign_keys.sql new file mode 100644 index 000000000000..c71b04275b68 --- /dev/null +++ b/diesel_cli/src/infer_schema_internals/load_foreign_keys.sql @@ -0,0 +1,19 @@ +SELECT + sch.nspname AS "self_schema", + tbl.relname AS "self_table", + ARRAY_AGG(col.attname ORDER BY u.attposition) AS "self_columns", + f_sch.nspname AS "foreign_schema", + f_tbl.relname AS "foreign_table", + ARRAY_AGG(f_col.attname ORDER BY f_u.attposition) AS "foreign_columns" +FROM pg_constraint c + LEFT JOIN LATERAL UNNEST(c.conkey) WITH ORDINALITY AS u(attnum, attposition) ON TRUE + LEFT JOIN LATERAL UNNEST(c.confkey) WITH ORDINALITY AS f_u(attnum, attposition) ON f_u.attposition = u.attposition + JOIN pg_class tbl ON tbl.oid = c.conrelid + JOIN pg_namespace sch ON sch.oid = tbl.relnamespace + LEFT JOIN pg_attribute col ON (col.attrelid = tbl.oid AND col.attnum = u.attnum) + LEFT JOIN pg_class f_tbl ON f_tbl.oid = c.confrelid + LEFT JOIN pg_namespace f_sch ON f_sch.oid = f_tbl.relnamespace + LEFT JOIN pg_attribute f_col ON (f_col.attrelid = f_tbl.oid AND f_col.attnum = f_u.attnum) +WHERE sch.nspname = $1 and f_sch.nspname = $1 and c.contype = 'f' +GROUP BY "self_schema", "self_table", "foreign_schema", "foreign_table" +ORDER BY "self_schema", "self_table"; diff --git a/diesel_cli/src/infer_schema_internals/pg.rs b/diesel_cli/src/infer_schema_internals/pg.rs index bb0635979aa6..03a6cc18a303 100644 --- a/diesel_cli/src/infer_schema_internals/pg.rs +++ b/diesel_cli/src/infer_schema_internals/pg.rs @@ -2,14 +2,13 @@ use super::data_structures::*; use super::information_schema::DefaultSchema; use super::TableName; use crate::print_schema::ColumnSorting; -use diesel::{ - deserialize::{self, FromStaticSqlRow, Queryable}, - dsl::AsExprOf, - expression::AsExpression, - pg::Pg, - prelude::*, - sql_types, -}; +use diesel::connection::DefaultLoadingMode; +use diesel::deserialize::{self, FromStaticSqlRow, Queryable}; +use diesel::dsl::AsExprOf; +use diesel::expression::AsExpression; +use diesel::pg::Pg; +use diesel::prelude::*; +use diesel::sql_types::{self, Array, Text}; use heck::ToUpperCamelCase; use std::borrow::Cow; use std::error::Error; @@ -173,72 +172,48 @@ pub fn load_foreign_key_constraints( connection: &mut PgConnection, schema_name: Option<&str>, ) -> QueryResult> { - use super::information_schema::information_schema::key_column_usage as kcu; - use super::information_schema::information_schema::referential_constraints as rc; - use super::information_schema::information_schema::table_constraints as tc; + #[derive(QueryableByName)] + struct ForeignKeyList { + #[diesel(sql_type = Text)] + self_schema: String, + #[diesel(sql_type = Text)] + self_table: String, + #[diesel(sql_type = Array)] + self_columns: Vec, + #[diesel(sql_type = Text)] + foreign_schema: String, + #[diesel(sql_type = Text)] + foreign_table: String, + #[diesel(sql_type = Array)] + foreign_columns: Vec, + } let default_schema = Pg::default_schema(connection)?; let schema_name = schema_name.unwrap_or(&default_schema); - let constraint_names = tc::table - .filter(tc::constraint_type.eq("FOREIGN KEY")) - .filter(tc::table_schema.eq(schema_name)) - .inner_join( - rc::table.on(tc::constraint_schema - .eq(rc::constraint_schema) - .and(tc::constraint_name.eq(rc::constraint_name))), - ) - .select(( - rc::constraint_schema, - rc::constraint_name, - rc::unique_constraint_schema, - rc::unique_constraint_name, - )) - .load::<(String, String, Option, Option)>(connection)?; - - constraint_names - .into_iter() - .map( - |(foreign_key_schema, foreign_key_name, primary_key_schema, primary_key_name)| { - let foreign_key = kcu::table - .filter(kcu::constraint_schema.eq(&foreign_key_schema)) - .filter(kcu::constraint_name.eq(&foreign_key_name)) - .group_by((kcu::table_name, kcu::table_schema)) - .select(( - kcu::table_name, - kcu::table_schema, - array_agg(kcu::column_name), - )) - .first::<(String, String, Vec)>(connection)?; - let primary_key = kcu::table - .filter(kcu::constraint_schema.nullable().eq(primary_key_schema)) - .filter(kcu::constraint_name.nullable().eq(primary_key_name)) - .group_by((kcu::table_name, kcu::table_schema)) - .select(( - kcu::table_name, - kcu::table_schema, - array_agg(kcu::column_name), - )) - .first::<(String, String, Vec)>(connection)?; - - let mut primary_key_table = TableName::new(primary_key.0, primary_key.1); - primary_key_table.strip_schema_if_matches(&default_schema); - let mut foreign_key_table = TableName::new(foreign_key.0, foreign_key.1); - foreign_key_table.strip_schema_if_matches(&default_schema); - - let primary_key_columns = primary_key.2; - let foreign_key_columns = foreign_key.2; - - Ok(ForeignKeyConstraint { - child_table: foreign_key_table, - parent_table: primary_key_table, - foreign_key_columns_rust: foreign_key_columns.clone(), - foreign_key_columns, - primary_key_columns, - }) - }, - ) - .filter(|e| !matches!(e, Err(diesel::result::Error::NotFound))) + diesel::sql_query(include_str!("load_foreign_keys.sql")) + .bind::(schema_name) + .load_iter::(connection)? + .map(|f| { + let f = f?; + let mut child_table = TableName::new(f.self_table, f.self_schema); + child_table.strip_schema_if_matches(&default_schema); + let mut parent_table = TableName::new(f.foreign_table, f.foreign_schema); + parent_table.strip_schema_if_matches(&default_schema); + + let foreign_key_columns_rust = f + .self_columns + .iter() + .map(|s| super::inference::rust_name_for_sql_name(s)) + .collect(); + Ok(ForeignKeyConstraint { + child_table, + parent_table, + foreign_key_columns: f.self_columns, + foreign_key_columns_rust, + primary_key_columns: f.foreign_columns, + }) + }) .collect() } @@ -380,4 +355,54 @@ mod test { load_foreign_key_constraints(&mut connection, Some("test_schema")) ); } + + #[test] + fn get_foreign_keys_loads_foreign_keys_with_same_name() { + let mut connection = connection(); + + diesel::sql_query("CREATE SCHEMA test_schema") + .execute(&mut connection) + .unwrap(); + diesel::sql_query("CREATE TABLE test_schema.table_1 (id SERIAL PRIMARY KEY)") + .execute(&mut connection) + .unwrap(); + diesel::sql_query( + "CREATE TABLE test_schema.table_2 (\ + id SERIAL PRIMARY KEY,\ + fk_id INTEGER NOT NULL,\ + CONSTRAINT fk FOREIGN KEY (fk_id) REFERENCES test_schema.table_1 (id))", + ) + .execute(&mut connection) + .unwrap(); + diesel::sql_query( + "CREATE TABLE test_schema.table_3 (\ + id SERIAL PRIMARY KEY,\ + fk_id INTEGER NOT NULL,\ + CONSTRAINT fk FOREIGN KEY (fk_id) REFERENCES test_schema.table_1 (id))", + ) + .execute(&mut connection) + .unwrap(); + + let table_1 = TableName::new("table_1", "test_schema"); + let table_2 = TableName::new("table_2", "test_schema"); + let table_3 = TableName::new("table_3", "test_schema"); + let fk_one = ForeignKeyConstraint { + child_table: table_2, + parent_table: table_1.clone(), + foreign_key_columns: vec!["fk_id".into()], + foreign_key_columns_rust: vec!["fk_id".into()], + primary_key_columns: vec!["id".into()], + }; + let fk_two = ForeignKeyConstraint { + child_table: table_3, + parent_table: table_1, + foreign_key_columns: vec!["fk_id".into()], + foreign_key_columns_rust: vec!["fk_id".into()], + primary_key_columns: vec!["id".into()], + }; + assert_eq!( + Ok(vec![fk_one, fk_two]), + load_foreign_key_constraints(&mut connection, Some("test_schema")) + ); + } } diff --git a/diesel_cli/tests/support/project_builder.rs b/diesel_cli/tests/support/project_builder.rs index e1eeada3bff1..c1199b1aba17 100644 --- a/diesel_cli/tests/support/project_builder.rs +++ b/diesel_cli/tests/support/project_builder.rs @@ -102,8 +102,8 @@ impl Project { use std::env; dotenv().ok(); - let mut db_url = - url::Url::parse(&env::var_os(var).unwrap().into_string().unwrap()).unwrap(); + let var_os = env::var(var).unwrap(); + let mut db_url = url::Url::parse(&var_os).unwrap(); db_url.set_path(&format!("/diesel_{}", &self.name)); db_url }