@@ -3,14 +3,25 @@ use crate::{
33 db::{Conn, Transaction},
44 schema::Schema,
55};
6−use anyhow::Context;
6+use anyhow::{anyhow, bail, Context};
77use serde::{Deserialize, Serialize};
88
99#[derive(Serialize, Deserialize, Debug)]
1010pub struct RemoveColumn {
1111 pub table: String,
1212 pub column: String,
13− pub down: Option<String>,
13+ pub down: Option<Transformation>,
14+}
15+
16+#[derive(Serialize, Deserialize, Debug)]
17+#[serde(untagged)]
18+pub enum Transformation {
19+ Simple(String),
20+ Update {
21+ table: String,
22+ value: String,
23+ r#where: String,
24+ },
1425}
1526
1627impl RemoveColumn {
@@ -22,6 +33,33 @@ impl RemoveColumn {
2233 self.column
2334 )
2435 }
36+
37+ fn reverse_trigger_name(&self, ctx: &MigrationContext) -> String {
38+ format!(
39+ "{}_remove_column_{}_{}_rev",
40+ ctx.prefix(),
41+ self.table,
42+ self.column
43+ )
44+ }
45+
46+ fn not_null_constraint_trigger_name(&self, ctx: &MigrationContext) -> String {
47+ format!(
48+ "{}_remove_column_{}_{}_nn",
49+ ctx.prefix(),
50+ self.table,
51+ self.column
52+ )
53+ }
54+
55+ fn not_null_constraint_name(&self, ctx: &MigrationContext) -> String {
56+ format!(
57+ "{}_add_column_not_null_{}_{}",
58+ ctx.prefix(),
59+ self.table,
60+ self.column
61+ )
62+ }
2563}
2664
2765#[typetag::serde(name = "remove_column")]
@@ -39,10 +77,13 @@ impl Action for RemoveColumn {
3977 db: &mut dyn Conn,
4078 schema: &Schema,
4179 ) -> anyhow::Result<()> {
80+ let table = schema.get_table(db, &self.table)?;
81+ let column = table
82+ .get_column(&self.column)
83+ .ok_or_else(|| anyhow!("no such column {} exists", self.column))?;
84+
4285 // Add down trigger
4386 if let Some(down) = &self.down {
44− let table = schema.get_table(db, &self.table)?;
45−
4687 let declarations: Vec<String> = table
4788 .columns
4889 .iter()
@@ -56,32 +97,226 @@ impl Action for RemoveColumn {
5697 })
5798 .collect();
5899
59− let query = format!(
60− r#"
61− CREATE OR REPLACE FUNCTION {trigger_name}()
62− RETURNS TRIGGER AS $$
63− BEGIN
64− IF reshape.is_new_schema() THEN
65− DECLARE
66− {declarations}
100+ if let Transformation::Simple(down) = down {
101+ let query = format!(
102+ r#"
103+ CREATE OR REPLACE FUNCTION {trigger_name}()
104+ RETURNS TRIGGER AS $$
105+ BEGIN
106+ IF reshape.is_new_schema() THEN
107+ DECLARE
108+ {declarations}
109+ BEGIN
110+ NEW.{column_name} = {down};
111+ END;
112+ END IF;
113+ RETURN NEW;
114+ END
115+ $$ language 'plpgsql';
116+
117+ DROP TRIGGER IF EXISTS "{trigger_name}" ON "{table}";
118+ CREATE TRIGGER "{trigger_name}" BEFORE UPDATE OR INSERT ON "{table}" FOR EACH ROW EXECUTE PROCEDURE {trigger_name}();
119+ "#,
120+ column_name = self.column,
121+ trigger_name = self.trigger_name(ctx),
122+ down = down,
123+ table = self.table,
124+ declarations = declarations.join("\n"),
125+ );
126+ db.run(&query).context("failed to create down trigger")?;
127+ }
128+
129+ if let Transformation::Update {
130+ table: from_table,
131+ value,
132+ r#where,
133+ } = down
134+ {
135+ let existing_schema_name = match &ctx.existing_schema_name {
136+ Some(name) => name,
137+ None => bail!("can't use update without previous migration"),
138+ };
139+
140+ let from_table = schema.get_table(db, &from_table)?;
141+
142+ let maybe_null_check = if !column.nullable {
143+ // Replace NOT NULL constraint with a constraint trigger that only triggers on the old schema.
144+ // We will add a null check to the down function on the new schema below as well to cover both cases.
145+ // As we are using a complex down function, we must remove the NOT NULL check for the new schema.
146+ // NOT NULL is not checked at the end of a transaction, but immediately upon update.
147+ let query = format!(
148+ r#"
149+ CREATE OR REPLACE FUNCTION {trigger_name}()
150+ RETURNS TRIGGER AS $$
67151 BEGIN
68− NEW.{column_name} = {down};
69− END;
70− END IF;
71− RETURN NEW;
72− END
73− $$ language 'plpgsql';
74−
75− DROP TRIGGER IF EXISTS "{trigger_name}" ON "{table}";
76− CREATE TRIGGER "{trigger_name}" BEFORE UPDATE OR INSERT ON "{table}" FOR EACH ROW EXECUTE PROCEDURE {trigger_name}();
77− "#,
78− column_name = self.column,
79− trigger_name = self.trigger_name(ctx),
80− down = down,
81− table = self.table,
82− declarations = declarations.join("\n"),
83− );
84− db.run(&query).context("failed to create down trigger")?;
152+ IF NOT reshape.is_new_schema() THEN
153+ IF NEW.{column} IS NULL THEN
154+ RAISE EXCEPTION '{column} can not be null';
155+ END IF;
156+ END IF;
157+ RETURN NEW;
158+ END
159+ $$ language 'plpgsql';
160+
161+ DROP TRIGGER IF EXISTS "{trigger_name}" ON "{table}";
162+
163+ CREATE CONSTRAINT TRIGGER "{trigger_name}"
164+ AFTER INSERT OR UPDATE
165+ ON "{table}"
166+ FOR EACH ROW
167+ EXECUTE PROCEDURE {trigger_name}();
168+ "#,
169+ table = self.table,
170+ trigger_name = self.not_null_constraint_trigger_name(ctx),
171+ column = self.column,
172+ );
173+ db.run(&query)
174+ .context("failed to create null constraint trigger")?;
175+
176+ db.run(&format!(
177+ r#"
178+ ALTER TABLE {table}
179+ ALTER COLUMN {column}
180+ DROP NOT NULL
181+ "#,
182+ table = self.table,
183+ column = self.column
184+ ))
185+ .context("failed to remove column not null constraint")?;
186+
187+ format!(
188+ r#"
189+ IF {value} IS NULL THEN
190+ RAISE EXCEPTION '{column_name} can not be null';
191+ END IF;
192+ "#,
193+ column_name = self.column,
194+ )
195+ } else {
196+ "".to_string()
197+ };
198+
199+ let into_variables = from_table
200+ .columns
201+ .iter()
202+ .map(|column| {
203+ format!(
204+ "NEW.{real_name} AS {alias}",
205+ alias = column.name,
206+ real_name = column.real_name,
207+ )
208+ })
209+ .collect::<Vec<String>>()
210+ .join(", ");
211+
212+ let query = format!(
213+ r#"
214+ CREATE OR REPLACE FUNCTION {trigger_name}()
215+ RETURNS TRIGGER AS $$
216+ #variable_conflict use_variable
217+ BEGIN
218+ IF reshape.is_new_schema() THEN
219+ DECLARE
220+ {from_table} record;
221+ BEGIN
222+ SELECT {into_variables}
223+ INTO {from_table};
224+
225+ {maybe_null_check}
226+
227+ -- Don't trigger reverse trigger when making this update
228+ perform set_config('reshape.disable_triggers', 'TRUE', TRUE);
229+
230+ UPDATE "migration_{existing_schema_name}"."{changed_table}" "{changed_table}"
231+ SET "{column_name}" = {value}
232+ WHERE {where};
233+
234+ perform set_config('reshape.disable_triggers', '', TRUE);
235+ END;
236+ END IF;
237+ RETURN NEW;
238+ END
239+ $$ language 'plpgsql';
240+
241+ DROP TRIGGER IF EXISTS "{trigger_name}" ON "{from_table_real}";
242+ CREATE TRIGGER "{trigger_name}" BEFORE UPDATE OR INSERT ON "{from_table_real}" FOR EACH ROW EXECUTE PROCEDURE {trigger_name}();
243+ "#,
244+ changed_table = self.table,
245+ from_table = from_table.name,
246+ from_table_real = from_table.real_name,
247+ column_name = self.column,
248+ trigger_name = self.trigger_name(ctx),
249+ );
250+ db.run(&query).context("failed to create down trigger")?;
251+
252+ let changed_into_variables = table
253+ .columns
254+ .iter()
255+ .map(|column| {
256+ format!(
257+ "NEW.{real_name} AS {alias}",
258+ alias = column.name,
259+ real_name = column.real_name,
260+ )
261+ })
262+ .collect::<Vec<String>>()
263+ .join(", ");
264+
265+ let from_table_columns = from_table
266+ .columns
267+ .iter()
268+ .map(|column| format!("{} as {}", column.real_name, column.name))
269+ .collect::<Vec<String>>()
270+ .join(", ");
271+
272+ let query = format!(
273+ r#"
274+ CREATE OR REPLACE FUNCTION {trigger_name}()
275+ RETURNS TRIGGER AS $$
276+ #variable_conflict use_variable
277+ BEGIN
278+ IF reshape.is_new_schema() AND NOT current_setting('reshape.disable_triggers', TRUE) = 'TRUE' THEN
279+ DECLARE
280+ {changed_table} record;
281+ __temp_row record;
282+ BEGIN
283+ SELECT {changed_into_variables}
284+ INTO {changed_table};
285+
286+ SELECT *
287+ INTO __temp_row
288+ FROM (
289+ SELECT {from_table_columns}
290+ FROM public.{from_table_real}
291+ ) {from_table}
292+ WHERE {where};
293+
294+ DECLARE
295+ {from_table} record;
296+ BEGIN
297+ {from_table} := __temp_row;
298+ NEW.{column_name_real} = {value};
299+ END;
300+ END;
301+ END IF;
302+ RETURN NEW;
303+ END
304+ $$ language 'plpgsql';
305+
306+ DROP TRIGGER IF EXISTS "{trigger_name}" ON "{changed_table_real}";
307+ CREATE TRIGGER "{trigger_name}" BEFORE UPDATE OR INSERT ON "{changed_table_real}" FOR EACH ROW EXECUTE PROCEDURE {trigger_name}();
308+ "#,
309+ changed_table = table.name,
310+ changed_table_real = table.real_name,
311+ from_table = from_table.name,
312+ from_table_real = from_table.real_name,
313+ column_name_real = column.real_name,
314+ trigger_name = self.reverse_trigger_name(ctx),
315+ // declarations = declarations.join("\n"),
316+ );
317+ db.run(&query)
318+ .context("failed to create reverse down trigger")?;
319+ }
85320 }
86321
87322 Ok(())
@@ -111,12 +346,15 @@ impl Action for RemoveColumn {
111346 ALTER TABLE "{table}"
112347 DROP COLUMN IF EXISTS "{column}";
113348
114− DROP TRIGGER IF EXISTS "{trigger_name}" ON "{table}";
115− DROP FUNCTION IF EXISTS "{trigger_name}";
349+ DROP FUNCTION IF EXISTS "{trigger_name}" CASCADE;
350+ DROP FUNCTION IF EXISTS "{reverse_trigger_name}" CASCADE;
351+ DROP FUNCTION IF EXISTS "{null_trigger_name}" CASCADE;
116352 "#,
117353 table = self.table,
118354 column = self.column,
119355 trigger_name = self.trigger_name(ctx),
356+ reverse_trigger_name = self.reverse_trigger_name(ctx),
357+ null_trigger_name = self.not_null_constraint_trigger_name(ctx),
120358 );
121359 db.run(&query)
122360 .context("failed to drop column and down trigger")?;
@@ -133,14 +371,81 @@ impl Action for RemoveColumn {
133371 }
134372
135373 fn abort(&self, ctx: &MigrationContext, db: &mut dyn Conn) -> anyhow::Result<()> {
374+ // We might have temporaily removed the NOT NULL check and have to reinstate it
375+ let has_not_null_function = !db
376+ .query_with_params(
377+ "
378+ SELECT routine_name
379+ FROM information_schema.routines
380+ WHERE routine_schema = 'public'
381+ AND routine_name = $1
382+ ",
383+ &[&self.not_null_constraint_trigger_name(ctx)],
384+ )
385+ .context("failed to get any NOT NULL function")?
386+ .is_empty();
387+
388+ if has_not_null_function {
389+ // Make column NOT NULL again without taking any long lived locks with a temporary constraint
390+ let query = format!(
391+ r#"
392+ ALTER TABLE "{table}"
393+ ADD CONSTRAINT "{constraint_name}"
394+ CHECK ("{column}" IS NOT NULL) NOT VALID
395+ "#,
396+ table = self.table,
397+ constraint_name = self.not_null_constraint_name(ctx),
398+ column = self.column,
399+ );
400+ db.run(&query)
401+ .context("failed to add NOT NULL constraint")?;
402+
403+ let query = format!(
404+ r#"
405+ ALTER TABLE "{table}"
406+ VALIDATE CONSTRAINT "{constraint_name}"
407+ "#,
408+ table = self.table,
409+ constraint_name = self.not_null_constraint_name(ctx),
410+ );
411+ db.run(&query)
412+ .context("failed to validate NOT NULL constraint")?;
413+
414+ // This ALTER TABLE call will not require any exclusive locks as it can use the validated constraint from above
415+ db.run(&format!(
416+ r#"
417+ ALTER TABLE {table}
418+ ALTER COLUMN {column}
419+ SET NOT NULL
420+ "#,
421+ table = self.table,
422+ column = self.column
423+ ))
424+ .context("failed to reinstate column NOT NULL")?;
425+
426+ // Drop the temporary constraint
427+ let query = format!(
428+ r#"
429+ ALTER TABLE "{table}"
430+ DROP CONSTRAINT "{constraint_name}"
431+ "#,
432+ table = self.table,
433+ constraint_name = self.not_null_constraint_name(ctx),
434+ );
435+ db.run(&query)
436+ .context("failed to drop NOT NULL constraint")?;
437+ }
438+
136439 // Remove function and trigger
137440 db.run(&format!(
138441 r#"
139− DROP TRIGGER IF EXISTS "{trigger_name}" ON "{table}";
140− DROP FUNCTION IF EXISTS "{trigger_name}";
442+ DROP FUNCTION IF EXISTS "{trigger_name}" CASCADE;
443+ DROP FUNCTION IF EXISTS "{reverse_trigger_name}" CASCADE;
444+ DROP FUNCTION IF EXISTS "{null_trigger_name}" CASCADE;
141445 "#,
142− table = self.table,
143446 trigger_name = self.trigger_name(ctx),
447+ reverse_trigger_name = self.reverse_trigger_name(ctx),
448+ null_trigger_name = self.not_null_constraint_trigger_name(ctx),
144449 ))
145450 .context("failed to drop down trigger")?;
146451