@@ -34,7 +34,16 @@ impl RemoveColumn {
3434 )
3535 }
3636
37− fn null_constraint_trigger_name(&self, ctx: &MigrationContext) -> String {
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 {
3847 format!(
3948 "{}_remove_column_{}_{}_nn",
4049 ctx.prefix(),
@@ -42,6 +51,15 @@ impl RemoveColumn {
4251 self.column
4352 )
4453 }
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+ }
4563}
4664
4765#[typetag::serde(name = "remove_column")]
@@ -121,21 +139,9 @@ impl Action for RemoveColumn {
121139
122140 let from_table = schema.get_table(db, &from_table)?;
123141
124− let declarations: Vec<String> = from_table
125− .columns
126− .iter()
127− .map(|column| {
128− format!(
129− "{alias} public.{table}.{real_name}%TYPE := NEW.{real_name};",
130− table = from_table.real_name,
131− alias = column.name,
132− real_name = column.real_name,
133− )
134− })
135− .collect();
136−
137142 let maybe_null_check = if !column.nullable {
138143 // 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.
139145 // As we are using a complex down function, we must remove the NOT NULL check for the new schema.
140146 // NOT NULL is not checked at the end of a transaction, but immediately upon update.
141147 let query = format!(
@@ -161,7 +167,7 @@ impl Action for RemoveColumn {
161167 EXECUTE PROCEDURE {trigger_name}();
162168 "#,
163169 table = self.table,
164− trigger_name = self.null_constraint_trigger_name(ctx),
170+ trigger_name = self.not_null_constraint_trigger_name(ctx),
165171 column = self.column,
166172 );
167173 db.run(&query)
@@ -190,6 +196,19 @@ impl Action for RemoveColumn {
190196 "".to_string()
191197 };
192198
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+
193212 let query = format!(
194213 r#"
195214 CREATE OR REPLACE FUNCTION {trigger_name}()
@@ -198,12 +217,21 @@ impl Action for RemoveColumn {
198217 BEGIN
199218 IF reshape.is_new_schema() THEN
200219 DECLARE
201− {declarations}
220+ {from_table} record;
202221 BEGIN
222+ SELECT {into_variables}
223+ INTO {from_table};
224+
203225 {maybe_null_check}
226+
227+ -- Don't trigger reverse trigger when making this update
228+ perform set_config('reshape.disable_triggers', 'TRUE', TRUE);
229+
204230 UPDATE "migration_{existing_schema_name}"."{changed_table}" "{changed_table}"
205231 SET "{column_name}" = {value}
206232 WHERE {where};
233+
234+ perform set_config('reshape.disable_triggers', '', TRUE);
207235 END;
208236 END IF;
209237 RETURN NEW;
@@ -214,17 +242,83 @@ impl Action for RemoveColumn {
214242 CREATE TRIGGER "{trigger_name}" BEFORE UPDATE OR INSERT ON "{from_table_real}" FOR EACH ROW EXECUTE PROCEDURE {trigger_name}();
215243 "#,
216244 changed_table = self.table,
245+ from_table = from_table.name,
217246 from_table_real = from_table.real_name,
218247 column_name = self.column,
219248 trigger_name = self.trigger_name(ctx),
220− declarations = declarations.join("\n"),
221249 );
222250 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")?;
223319 }
224320 }
225321
226− if !column.nullable {}
227−
228322 Ok(())
229323 }
230324
@@ -247,29 +341,20 @@ impl Action for RemoveColumn {
247341 }
248342
249343 // Remove column, function and trigger
250− let trigger_table = match &self.down {
251− Some(Transformation::Update {
252− table,
253− value: _,
254− r#where: _,
255− }) => table,
256− _ => &self.table,
257− };
258344 let query = format!(
259345 r#"
260346 ALTER TABLE "{table}"
261347 DROP COLUMN IF EXISTS "{column}";
262348
263− DROP TRIGGER IF EXISTS "{trigger_name}" ON "{trigger_table}";
264− DROP FUNCTION IF EXISTS "{trigger_name}";
265−
266− DROP TRIGGER IF EXISTS "{null_trigger_name}" ON "{table}";
267− DROP FUNCTION IF EXISTS "{null_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;
268352 "#,
269353 table = self.table,
270354 column = self.column,
271355 trigger_name = self.trigger_name(ctx),
272− null_trigger_name = self.null_constraint_trigger_name(ctx),
356+ reverse_trigger_name = self.reverse_trigger_name(ctx),
357+ null_trigger_name = self.not_null_constraint_trigger_name(ctx),
273358 );
274359 db.run(&query)
275360 .context("failed to drop column and down trigger")?;
@@ -286,38 +371,81 @@ impl Action for RemoveColumn {
286371 }
287372
288373 fn abort(&self, ctx: &MigrationContext, db: &mut dyn Conn) -> anyhow::Result<()> {
289− // We might have removed the NOT NULL check so we reinstate it
290− db.run(&format!(
291− r#"
292− ALTER TABLE {table}
293− ALTER COLUMN {column}
294− SET NOT NULL
295− "#,
296− table = self.table,
297− column = self.column
298− ))
299− .context("failed to reinstate column not null constraint")?;
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+ }
300438
301439 // Remove function and trigger
302− let trigger_table = match &self.down {
303− Some(Transformation::Update {
304− table,
305− value: _,
306− r#where: _,
307− }) => table,
308− _ => &self.table,
309− };
310440 db.run(&format!(
311441 r#"
312− DROP TRIGGER IF EXISTS "{trigger_name}" ON "{trigger_table}";
313− DROP FUNCTION IF EXISTS "{trigger_name}";
314−
315− DROP TRIGGER IF EXISTS "{null_trigger_name}" ON "{table}";
316− DROP FUNCTION IF EXISTS "{null_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;
317445 "#,
318− table = self.table,
319446 trigger_name = self.trigger_name(ctx),
320− null_trigger_name = self.null_constraint_trigger_name(ctx),
447+ reverse_trigger_name = self.reverse_trigger_name(ctx),
448+ null_trigger_name = self.not_null_constraint_trigger_name(ctx),
321449 ))
322450 .context("failed to drop down trigger")?;
323451