@@ -3,25 +3,14 @@ use crate::{
33 db::{Conn, Transaction},
44 schema::Schema,
55};
6−use anyhow::{anyhow, bail, Context};
6+use anyhow::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<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− },
13+ pub down: Option<String>,
2514}
2615
2716impl RemoveColumn {
@@ -33,33 +22,6 @@ impl RemoveColumn {
3322 self.column
3423 )
3524 }
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− }
6325}
6426
6527#[typetag::serde(name = "remove_column")]
@@ -77,13 +39,10 @@ impl Action for RemoveColumn {
7739 db: &mut dyn Conn,
7840 schema: &Schema,
7941 ) -> 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−
8542 // Add down trigger
8643 if let Some(down) = &self.down {
44+ let table = schema.get_table(db, &self.table)?;
45+
8746 let declarations: Vec<String> = table
8847 .columns
8948 .iter()
@@ -97,226 +56,32 @@ impl Action for RemoveColumn {
9756 })
9857 .collect();
9958
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 $$
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}
15167 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
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− }
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")?;
32085 }
32186
32287 Ok(())
@@ -346,15 +111,12 @@ impl Action for RemoveColumn {
346111 ALTER TABLE "{table}"
347112 DROP COLUMN IF EXISTS "{column}";
348113
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;
114+ DROP TRIGGER IF EXISTS "{trigger_name}" ON "{table}";
115+ DROP FUNCTION IF EXISTS "{trigger_name}";
352116 "#,
353117 table = self.table,
354118 column = self.column,
355119 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),
358120 );
359121 db.run(&query)
360122 .context("failed to drop column and down trigger")?;
@@ -371,81 +133,14 @@ impl Action for RemoveColumn {
371133 }
372134
373135 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−
439136 // Remove function and trigger
440137 db.run(&format!(
441138 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;
139+ DROP TRIGGER IF EXISTS "{trigger_name}" ON "{table}";
140+ DROP FUNCTION IF EXISTS "{trigger_name}";
445141 "#,
142+ table = self.table,
446143 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),
449144 ))
450145 .context("failed to drop down trigger")?;
451146