@@ -3,7 +3,7 @@ use crate::{
33 db::{Conn, Transaction},
44 schema::Schema,
55};
6−use anyhow::{bail, Context};
6+use anyhow::{anyhow, bail, Context};
77use serde::{Deserialize, Serialize};
88
99#[derive(Serialize, Deserialize, Debug)]
@@ -33,6 +33,33 @@ impl RemoveColumn {
3333 self.column
3434 )
3535 }
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+ }
3663}
3764
3865#[typetag::serde(name = "remove_column")]
@@ -51,6 +78,9 @@ impl Action for RemoveColumn {
5178 schema: &Schema,
5279 ) -> anyhow::Result<()> {
5380 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))?;
5484
5585 // Add down trigger
5686 if let Some(down) = &self.down {
@@ -109,22 +139,76 @@ impl Action for RemoveColumn {
109139
110140 let from_table = schema.get_table(db, &from_table)?;
111141
112− let declarations: Vec<String> = from_table
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 $$
151+ BEGIN
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
113200 .columns
114201 .iter()
115202 .map(|column| {
116203 format!(
117− "{alias} public.{table}.{real_name}%TYPE := NEW.{real_name};",
118− table = from_table.real_name,
204+ "NEW.{real_name} AS {alias}",
119205 alias = column.name,
120206 real_name = column.real_name,
121207 )
122208 })
123− .collect();
209+ .collect::<Vec<String>>()
210+ .join(", ");
124211
125− // TODO: If column is NOT NULL, remove the constraint and perform a NULL check in some other way.
126− // Otherwise, it's not possible to update the NOT NULL column from another table even in the same transaction.
127− // Either make the NULL check in here or maybe use a constraint trigger: https://www.postgresql.org/docs/9.0/sql-createconstraint.html.
128212 let query = format!(
129213 r#"
130214 CREATE OR REPLACE FUNCTION {trigger_name}()
@@ -133,11 +217,21 @@ impl Action for RemoveColumn {
133217 BEGIN
134218 IF reshape.is_new_schema() THEN
135219 DECLARE
136− {declarations}
220+ {from_table} record;
137221 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+
138230 UPDATE "migration_{existing_schema_name}"."{changed_table}" "{changed_table}"
139231 SET "{column_name}" = {value}
140232 WHERE {where};
233+
234+ perform set_config('reshape.disable_triggers', '', TRUE);
141235 END;
142236 END IF;
143237 RETURN NEW;
@@ -148,12 +242,80 @@ impl Action for RemoveColumn {
148242 CREATE TRIGGER "{trigger_name}" BEFORE UPDATE OR INSERT ON "{from_table_real}" FOR EACH ROW EXECUTE PROCEDURE {trigger_name}();
149243 "#,
150244 changed_table = self.table,
245+ from_table = from_table.name,
151246 from_table_real = from_table.real_name,
152247 column_name = self.column,
153248 trigger_name = self.trigger_name(ctx),
154− declarations = declarations.join("\n"),
155249 );
156250 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")?;
157319 }
158320 }
159321
@@ -179,25 +341,20 @@ impl Action for RemoveColumn {
179341 }
180342
181343 // Remove column, function and trigger
182− let trigger_table = match &self.down {
183− Some(Transformation::Update {
184− table,
185− value: _,
186− r#where: _,
187− }) => table,
188− _ => &self.table,
189− };
190344 let query = format!(
191345 r#"
192346 ALTER TABLE "{table}"
193347 DROP COLUMN IF EXISTS "{column}";
194348
195− DROP TRIGGER IF EXISTS "{trigger_name}" ON "{trigger_table}";
196− 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;
197352 "#,
198353 table = self.table,
199354 column = self.column,
200355 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),
201358 );
202359 db.run(&query)
203360 .context("failed to drop column and down trigger")?;
@@ -214,21 +371,81 @@ impl Action for RemoveColumn {
214371 }
215372
216373 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+
217439 // Remove function and trigger
218− let trigger_table = match &self.down {
219− Some(Transformation::Update {
220− table,
221− value: _,
222− r#where: _,
223− }) => table,
224− _ => &self.table,
225− };
226440 db.run(&format!(
227441 r#"
228− DROP TRIGGER IF EXISTS "{trigger_name}" ON "{trigger_table}";
229− 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;
230445 "#,
231446 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),
232449 ))
233450 .context("failed to drop down trigger")?;
234451