@@ -34,16 +34,7 @@ impl RemoveColumn {
3434 )
3535 }
3636
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 {
37+ fn null_constraint_trigger_name(&self, ctx: &MigrationContext) -> String {
4738 format!(
4839 "{}_remove_column_{}_{}_nn",
4940 ctx.prefix(),
@@ -51,15 +42,6 @@ impl RemoveColumn {
5142 self.column
5243 )
5344 }
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− }
6345}
6446
6547#[typetag::serde(name = "remove_column")]
@@ -139,9 +121,21 @@ impl Action for RemoveColumn {
139121
140122 let from_table = schema.get_table(db, &from_table)?;
141123
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+
142137 let maybe_null_check = if !column.nullable {
143138 // 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.
145139 // As we are using a complex down function, we must remove the NOT NULL check for the new schema.
146140 // NOT NULL is not checked at the end of a transaction, but immediately upon update.
147141 let query = format!(
@@ -167,7 +161,7 @@ impl Action for RemoveColumn {
167161 EXECUTE PROCEDURE {trigger_name}();
168162 "#,
169163 table = self.table,
170− trigger_name = self.not_null_constraint_trigger_name(ctx),
164+ trigger_name = self.null_constraint_trigger_name(ctx),
171165 column = self.column,
172166 );
173167 db.run(&query)
@@ -196,19 +190,6 @@ impl Action for RemoveColumn {
196190 "".to_string()
197191 };
198192
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−
212193 let query = format!(
213194 r#"
214195 CREATE OR REPLACE FUNCTION {trigger_name}()
@@ -217,21 +198,12 @@ impl Action for RemoveColumn {
217198 BEGIN
218199 IF reshape.is_new_schema() THEN
219200 DECLARE
220− {from_table} record;
201+ {declarations}
221202 BEGIN
222− SELECT {into_variables}
223− INTO {from_table};
224−
225203 {maybe_null_check}
226−
227− -- Don't trigger reverse trigger when making this update
228− perform set_config('reshape.disable_triggers', 'TRUE', TRUE);
229−
230204 UPDATE "migration_{existing_schema_name}"."{changed_table}" "{changed_table}"
231205 SET "{column_name}" = {value}
232206 WHERE {where};
233−
234− perform set_config('reshape.disable_triggers', '', TRUE);
235207 END;
236208 END IF;
237209 RETURN NEW;
@@ -242,83 +214,17 @@ impl Action for RemoveColumn {
242214 CREATE TRIGGER "{trigger_name}" BEFORE UPDATE OR INSERT ON "{from_table_real}" FOR EACH ROW EXECUTE PROCEDURE {trigger_name}();
243215 "#,
244216 changed_table = self.table,
245− from_table = from_table.name,
246217 from_table_real = from_table.real_name,
247218 column_name = self.column,
248219 trigger_name = self.trigger_name(ctx),
220+ declarations = declarations.join("\n"),
249221 );
250222 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")?;
319223 }
320224 }
321225
226+ if !column.nullable {}
227+
322228 Ok(())
323229 }
324230
@@ -341,20 +247,29 @@ impl Action for RemoveColumn {
341247 }
342248
343249 // 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+ };
344258 let query = format!(
345259 r#"
346260 ALTER TABLE "{table}"
347261 DROP COLUMN IF EXISTS "{column}";
348262
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;
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}";
352268 "#,
353269 table = self.table,
354270 column = self.column,
355271 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),
272+ null_trigger_name = self.null_constraint_trigger_name(ctx),
358273 );
359274 db.run(&query)
360275 .context("failed to drop column and down trigger")?;
@@ -371,81 +286,38 @@ impl Action for RemoveColumn {
371286 }
372287
373288 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− }
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")?;
438300
439301 // 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+ };
440310 db.run(&format!(
441311 r#"
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;
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}";
445317 "#,
318+ table = self.table,
446319 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),
320+ null_trigger_name = self.null_constraint_trigger_name(ctx),
449321 ))
450322 .context("failed to drop down trigger")?;
451323