Skip to content

Commit

Permalink
Merge pull request #3650 from tp-woven/talp/issue-3646
Browse files Browse the repository at this point in the history 
add test reproducing #3646
  • Loading branch information
@weiznich
weiznich authored Jun 8, 2023
2 parents c5a9f7b + 0cfa3d7 commit a4e6131
Show file tree
Hide file tree
Showing 3 changed files with 116 additions and 72 deletions.
19 changes: 19 additions & 0 deletions diesel_cli/src/infer_schema_internals/load_foreign_keys.sql
View file
Original file line number Diff line number Diff line change
@@ -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";
165 changes: 95 additions & 70 deletions diesel_cli/src/infer_schema_internals/pg.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -173,72 +172,48 @@ pub fn load_foreign_key_constraints(
connection: &mut PgConnection,
schema_name: Option<&str>,
) -> QueryResult<Vec<ForeignKeyConstraint>> {
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<Text>)]
self_columns: Vec<String>,
#[diesel(sql_type = Text)]
foreign_schema: String,
#[diesel(sql_type = Text)]
foreign_table: String,
#[diesel(sql_type = Array<Text>)]
foreign_columns: Vec<String>,
}

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<String>, Option<String>)>(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<String>)>(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<String>)>(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::<Text, _>(schema_name)
.load_iter::<ForeignKeyList, DefaultLoadingMode>(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()
}

Expand Down Expand Up @@ -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"))
);
}
}
4 changes: 2 additions & 2 deletions diesel_cli/tests/support/project_builder.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down

0 comments on commit a4e6131

Please sign in to comment.