Skip to content

Commit

Permalink
Merge pull request #3451 from weiznich/fix/3446
Browse files Browse the repository at this point in the history 
Fix #3446
  • Loading branch information
@weiznich
weiznich authored Dec 21, 2022
2 parents f7e632e + 1e9cc9c commit 0fe6167
Show file tree
Hide file tree
Showing 20 changed files with 417 additions and 143 deletions.
6 changes: 3 additions & 3 deletions diesel_cli/src/infer_schema_internals/data_structures.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -119,9 +119,9 @@ where
pub struct ForeignKeyConstraint {
pub child_table: TableName,
pub parent_table: TableName,
pub foreign_key: String,
pub foreign_key_rust_name: String,
pub primary_key: String,
pub foreign_key_columns: Vec<String>,
pub foreign_key_columns_rust: Vec<String>,
pub primary_key_columns: Vec<String>,
}

impl ForeignKeyConstraint {
Expand Down
3 changes: 2 additions & 1 deletion diesel_cli/src/infer_schema_internals/foreign_keys.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,12 +30,13 @@ pub fn remove_unsafe_foreign_keys_for_codegen(
.filter(|fk| fk.parent_table != fk.child_table)
.filter(|fk| safe_tables.contains(&fk.parent_table))
.filter(|fk| safe_tables.contains(&fk.child_table))
.filter(|fk| fk.foreign_key_columns.len() == 1)
.filter(|fk| {
let pk_columns = get_primary_keys(&mut conn, &fk.parent_table).expect(&format!(
"Error loading primary keys for `{}`",
fk.parent_table
));
pk_columns.len() == 1 && pk_columns[0] == fk.primary_key
pk_columns.len() == 1 && Some(&pk_columns[0]) == fk.primary_key_columns.get(0)
})
.filter(|fk| !duplicates.contains(&fk.ordered_tables()))
.cloned()
Expand Down
10 changes: 5 additions & 5 deletions diesel_cli/src/infer_schema_internals/inference.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -222,8 +222,7 @@ pub fn load_foreign_key_constraints(
}
#[cfg(feature = "postgres")]
InferConnection::Pg(ref mut c) => {
super::information_schema::load_foreign_key_constraints(c, schema_name)
.map_err(Into::into)
super::pg::load_foreign_key_constraints(c, schema_name).map_err(Into::into)
}
#[cfg(feature = "mysql")]
InferConnection::Mysql(ref mut c) => {
Expand All @@ -234,9 +233,10 @@ pub fn load_foreign_key_constraints(
constraints.map(|mut ct| {
ct.sort();
ct.iter_mut().for_each(|foreign_key_constraint| {
if is_reserved_name(&foreign_key_constraint.foreign_key_rust_name) {
foreign_key_constraint.foreign_key_rust_name =
format!("{}_", foreign_key_constraint.foreign_key_rust_name);
for name in &mut foreign_key_constraint.foreign_key_columns_rust {
if is_reserved_name(name) {
*name = format!("{name}_");
}
}
});
ct
Expand Down
104 changes: 0 additions & 104 deletions diesel_cli/src/infer_schema_internals/information_schema.rs
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,6 @@
use std::borrow::Cow;
use std::error::Error;

#[cfg(feature = "postgres")]
use super::ForeignKeyConstraint;
use diesel::backend::Backend;
use diesel::connection::LoadConnection;
use diesel::deserialize::FromSql;
Expand Down Expand Up @@ -186,66 +184,6 @@ where
.collect())
}

#[allow(clippy::similar_names)]
#[cfg(feature = "postgres")]
pub fn load_foreign_key_constraints(
connection: &mut PgConnection,
schema_name: Option<&str>,
) -> QueryResult<Vec<ForeignKeyConstraint>> {
use self::information_schema::key_column_usage as kcu;
use self::information_schema::referential_constraints as rc;
use self::information_schema::table_constraints as tc;

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 (mut foreign_key_table, foreign_key_column) = kcu::table
.filter(kcu::constraint_schema.eq(&foreign_key_schema))
.filter(kcu::constraint_name.eq(&foreign_key_name))
.select(((kcu::table_name, kcu::table_schema), kcu::column_name))
.first::<(TableName, String)>(connection)?;
let (mut primary_key_table, primary_key_column) = kcu::table
.filter(kcu::constraint_schema.nullable().eq(primary_key_schema))
.filter(kcu::constraint_name.nullable().eq(primary_key_name))
.select(((kcu::table_name, kcu::table_schema), kcu::column_name))
.first::<(TableName, _)>(connection)?;

foreign_key_table.strip_schema_if_matches(&default_schema);
primary_key_table.strip_schema_if_matches(&default_schema);

Ok(ForeignKeyConstraint {
child_table: foreign_key_table,
parent_table: primary_key_table,
foreign_key: foreign_key_column.clone(),
foreign_key_rust_name: foreign_key_column,
primary_key: primary_key_column,
})
},
)
.filter(|e| !matches!(e, Err(NotFound)))
.collect()
}

#[cfg(all(test, feature = "postgres"))]
mod tests {
extern crate dotenvy;
Expand Down Expand Up @@ -400,46 +338,4 @@ mod tests {
get_primary_keys(&mut connection, &table_2).unwrap()
);
}

#[test]
fn get_foreign_keys_loads_foreign_keys() {
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_one INTEGER NOT NULL REFERENCES test_schema.table_1)",
).execute(&mut connection)
.unwrap();
diesel::sql_query(
"CREATE TABLE test_schema.table_3 (id SERIAL PRIMARY KEY, fk_two INTEGER NOT NULL REFERENCES test_schema.table_2)",
).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.clone(),
parent_table: table_1,
foreign_key: "fk_one".into(),
foreign_key_rust_name: "fk_one".into(),
primary_key: "id".into(),
};
let fk_two = ForeignKeyConstraint {
child_table: table_3,
parent_table: table_2,
foreign_key: "fk_two".into(),
foreign_key_rust_name: "fk_two".into(),
primary_key: "id".into(),
};
assert_eq!(
Ok(vec![fk_one, fk_two]),
load_foreign_key_constraints(&mut connection, Some("test_schema"))
);
}
}
24 changes: 19 additions & 5 deletions diesel_cli/src/infer_schema_internals/mysql.rs
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
use diesel::mysql::{Mysql, MysqlConnection};
use diesel::*;
use heck::ToUpperCamelCase;
use std::collections::HashMap;
use std::{borrow::Cow, error::Error};

use super::data_structures::*;
Expand Down Expand Up @@ -123,20 +124,33 @@ pub fn load_foreign_key_constraints(
(kcu::referenced_table_name, kcu::referenced_table_schema),
kcu::column_name,
kcu::referenced_column_name,
kcu::constraint_name,
))
.load::<(TableName, TableName, String, _)>(connection)?
.load::<(TableName, TableName, String, String, String)>(connection)?
.into_iter()
.fold(
HashMap::new(),
|mut acc, (child_table, parent_table, foreign_key, primary_key, fk_constraint_name)| {
let entry = acc
.entry(fk_constraint_name)
.or_insert_with(|| (child_table, parent_table, Vec::new(), Vec::new()));
entry.2.push(foreign_key);
entry.3.push(primary_key);
acc
},
)
.into_values()
.map(
|(mut child_table, mut parent_table, foreign_key, primary_key)| {
|(mut child_table, mut parent_table, foreign_key_columns, primary_key_columns)| {
child_table.strip_schema_if_matches(&default_schema);
parent_table.strip_schema_if_matches(&default_schema);

ForeignKeyConstraint {
child_table,
parent_table,
foreign_key: foreign_key.clone(),
foreign_key_rust_name: foreign_key,
primary_key,
primary_key_columns,
foreign_key_columns_rust: foreign_key_columns.clone(),
foreign_key_columns,
}
},
)
Expand Down
121 changes: 121 additions & 0 deletions diesel_cli/src/infer_schema_internals/pg.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -115,6 +115,85 @@ mod information_schema {
}
}

sql_function! {
#[aggregate]
fn array_agg(input: diesel::sql_types::Text) -> diesel::sql_types::Array<diesel::sql_types::Text>;
}

#[allow(clippy::similar_names)]
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;

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)))
.collect()
}

#[cfg(test)]
mod test {
extern crate dotenvy;
Expand Down Expand Up @@ -202,4 +281,46 @@ mod test {
);
assert_eq!(Ok(None), get_table_comment(&mut connection, &table_2));
}

#[test]
fn get_foreign_keys_loads_foreign_keys() {
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_one INTEGER NOT NULL REFERENCES test_schema.table_1)",
).execute(&mut connection)
.unwrap();
diesel::sql_query(
"CREATE TABLE test_schema.table_3 (id SERIAL PRIMARY KEY, fk_two INTEGER NOT NULL REFERENCES test_schema.table_2)",
).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.clone(),
parent_table: table_1,
foreign_key_columns: vec!["fk_one".into()],
foreign_key_columns_rust: vec!["fk_one".into()],
primary_key_columns: vec!["id".into()],
};
let fk_two = ForeignKeyConstraint {
child_table: table_3,
parent_table: table_2,
foreign_key_columns: vec!["fk_two".into()],
foreign_key_columns_rust: vec!["fk_two".into()],
primary_key_columns: vec!["id".into()],
};
assert_eq!(
Ok(vec![fk_one, fk_two]),
load_foreign_key_constraints(&mut connection, Some("test_schema"))
);
}
}
Loading

0 comments on commit 0fe6167

Please sign in to comment.