@@ -0,0 +1,122 @@
1+use super::{common::ForeignKey, Action, MigrationContext};
2+use crate::{
3+ db::{Conn, Transaction},
4+ schema::Schema,
5+};
6+use anyhow::Context;
7+use serde::{Deserialize, Serialize};
8+
9+#[derive(Serialize, Deserialize, Debug)]
10+pub struct AddForeignKey {
11+ pub table: String,
12+ foreign_key: ForeignKey,
13+}
14+
15+#[typetag::serde(name = "add_foreign_key")]
16+impl Action for AddForeignKey {
17+ fn describe(&self) -> String {
18+ format!(
19+ "Adding foreign key from table \"{}\" to \"{}\"",
20+ self.table, self.foreign_key.referenced_table
21+ )
22+ }
23+
24+ fn run(
25+ &self,
26+ ctx: &MigrationContext,
27+ db: &mut dyn Conn,
28+ schema: &Schema,
29+ ) -> anyhow::Result<()> {
30+ let table = schema.get_table(db, &self.table)?;
31+ let referenced_table = schema.get_table(db, &self.foreign_key.referenced_table)?;
32+
33+ // Add quotes around all column names
34+ let columns: Vec<String> = table
35+ .real_column_names(&self.foreign_key.columns)
36+ .map(|col| format!("\"{}\"", col))
37+ .collect();
38+ let referenced_columns: Vec<String> = referenced_table
39+ .real_column_names(&self.foreign_key.referenced_columns)
40+ .map(|col| format!("\"{}\"", col))
41+ .collect();
42+
43+ // Create foreign key but set is as NOT VALID.
44+ // This means the foreign key will be enforced for inserts and updates
45+ // but the existing data won't be checked, that would cause a long-lived lock.
46+ db.run(&format!(
47+ r#"
48+ ALTER TABLE "{table}"
49+ ADD CONSTRAINT {constraint_name}
50+ FOREIGN KEY ({columns})
51+ REFERENCES {referenced_table} ({referenced_columns})
52+ NOT VALID
53+ "#,
54+ table = table.real_name,
55+ constraint_name = self.temp_constraint_name(ctx),
56+ columns = columns.join(", "),
57+ referenced_table = referenced_table.real_name,
58+ referenced_columns = referenced_columns.join(", "),
59+ ))
60+ .context("failed to create foreign key")?;
61+
62+ db.run(&format!(
63+ r#"
64+ ALTER TABLE "{table}"
65+ VALIDATE CONSTRAINT "{constraint_name}"
66+ "#,
67+ table = table.real_name,
68+ constraint_name = self.temp_constraint_name(ctx),
69+ ))
70+ .context("failed to validate foreign key")?;
71+
72+ Ok(())
73+ }
74+
75+ fn complete<'a>(
76+ &self,
77+ ctx: &MigrationContext,
78+ db: &'a mut dyn Conn,
79+ ) -> anyhow::Result<Option<Transaction<'a>>> {
80+ db.run(&format!(
81+ r#"
82+ ALTER TABLE {table}
83+ RENAME CONSTRAINT {temp_constraint_name} TO {constraint_name}
84+ "#,
85+ table = self.table,
86+ temp_constraint_name = self.temp_constraint_name(ctx),
87+ constraint_name = self.final_constraint_name(),
88+ ))
89+ .context("failed to rename temporary constraint")?;
90+ Ok(None)
91+ }
92+
93+ fn update_schema(&self, _ctx: &MigrationContext, _schema: &mut Schema) {}
94+
95+ fn abort(&self, ctx: &MigrationContext, db: &mut dyn Conn) -> anyhow::Result<()> {
96+ db.run(&format!(
97+ r#"
98+ ALTER TABLE "{table}"
99+ DROP CONSTRAINT IF EXISTS "{constraint_name}"
100+ "#,
101+ table = self.table,
102+ constraint_name = self.temp_constraint_name(ctx),
103+ ))
104+ .context("failed to validate foreign key")?;
105+
106+ Ok(())
107+ }
108+}
109+
110+impl AddForeignKey {
111+ fn temp_constraint_name(&self, ctx: &MigrationContext) -> String {
112+ format!("{}_temp_fkey", ctx.prefix())
113+ }
114+
115+ fn final_constraint_name(&self) -> String {
116+ format!(
117+ "{table}_{columns}_fkey",
118+ table = self.table,
119+ columns = self.foreign_key.columns.join("_")
120+ )
121+ }
122+}