@@ -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,15 @@ impl RemoveColumn {
3333 self.column
3434 )
3535 }
36+
37+ fn null_constraint_trigger_name(&self, ctx: &MigrationContext) -> String {
38+ format!(
39+ "{}_remove_column_{}_{}_nn",
40+ ctx.prefix(),
41+ self.table,
42+ self.column
43+ )
44+ }
3645}
3746
3847#[typetag::serde(name = "remove_column")]
@@ -51,6 +60,9 @@ impl Action for RemoveColumn {
5160 schema: &Schema,
5261 ) -> anyhow::Result<()> {
5362 let table = schema.get_table(db, &self.table)?;
63+ let column = table
64+ .get_column(&self.column)
65+ .ok_or_else(|| anyhow!("no such column {} exists", self.column))?;
5466
5567 // Add down trigger
5668 if let Some(down) = &self.down {
@@ -122,6 +134,62 @@ impl Action for RemoveColumn {
122134 })
123135 .collect();
124136
137+ let maybe_null_check = if !column.nullable {
138+ // Replace NOT NULL constraint with a constraint trigger that only triggers on the old schema.
139+ // As we are using a complex down function, we must remove the NOT NULL check for the new schema.
140+ // NOT NULL is not checked at the end of a transaction, but immediately upon update.
141+ let query = format!(
142+ r#"
143+ CREATE OR REPLACE FUNCTION {trigger_name}()
144+ RETURNS TRIGGER AS $$
145+ BEGIN
146+ IF NOT reshape.is_new_schema() THEN
147+ IF NEW.{column} IS NULL THEN
148+ RAISE EXCEPTION '{column} can not be null';
149+ END IF;
150+ END IF;
151+ RETURN NEW;
152+ END
153+ $$ language 'plpgsql';
154+
155+ DROP TRIGGER IF EXISTS "{trigger_name}" ON "{table}";
156+
157+ CREATE CONSTRAINT TRIGGER "{trigger_name}"
158+ AFTER INSERT OR UPDATE
159+ ON "{table}"
160+ FOR EACH ROW
161+ EXECUTE PROCEDURE {trigger_name}();
162+ "#,
163+ table = self.table,
164+ trigger_name = self.null_constraint_trigger_name(ctx),
165+ column = self.column,
166+ );
167+ db.run(&query)
168+ .context("failed to create null constraint trigger")?;
169+
170+ db.run(&format!(
171+ r#"
172+ ALTER TABLE {table}
173+ ALTER COLUMN {column}
174+ DROP NOT NULL
175+ "#,
176+ table = self.table,
177+ column = self.column
178+ ))
179+ .context("failed to remove column not null constraint")?;
180+
181+ format!(
182+ r#"
183+ IF {value} IS NULL THEN
184+ RAISE EXCEPTION '{column_name} can not be null';
185+ END IF;
186+ "#,
187+ column_name = self.column,
188+ )
189+ } else {
190+ "".to_string()
191+ };
192+
125193 // TODO: If column is NOT NULL, remove the constraint and perform a NULL check in some other way.
126194 // Otherwise, it's not possible to update the NOT NULL column from another table even in the same transaction.
127195 // Either make the NULL check in here or maybe use a constraint trigger: https://www.postgresql.org/docs/9.0/sql-createconstraint.html.
@@ -135,6 +203,7 @@ impl Action for RemoveColumn {
135203 DECLARE
136204 {declarations}
137205 BEGIN
206+ {maybe_null_check}
138207 UPDATE "migration_{existing_schema_name}"."{changed_table}" "{changed_table}"
139208 SET "{column_name}" = {value}
140209 WHERE {where};
@@ -157,6 +226,8 @@ impl Action for RemoveColumn {
157226 }
158227 }
159228
229+ if !column.nullable {}
230+
160231 Ok(())
161232 }
162233
@@ -194,10 +265,14 @@ impl Action for RemoveColumn {
194265
195266 DROP TRIGGER IF EXISTS "{trigger_name}" ON "{trigger_table}";
196267 DROP FUNCTION IF EXISTS "{trigger_name}";
268+
269+ DROP TRIGGER IF EXISTS "{null_trigger_name}" ON "{table}";
270+ DROP FUNCTION IF EXISTS "{null_trigger_name}";
197271 "#,
198272 table = self.table,
199273 column = self.column,
200274 trigger_name = self.trigger_name(ctx),
275+ null_trigger_name = self.null_constraint_trigger_name(ctx),
201276 );
202277 db.run(&query)
203278 .context("failed to drop column and down trigger")?;
@@ -214,6 +289,18 @@ impl Action for RemoveColumn {
214289 }
215290
216291 fn abort(&self, ctx: &MigrationContext, db: &mut dyn Conn) -> anyhow::Result<()> {
292+ // We might have removed the NOT NULL check so we reinstate it
293+ db.run(&format!(
294+ r#"
295+ ALTER TABLE {table}
296+ ALTER COLUMN {column}
297+ SET NOT NULL
298+ "#,
299+ table = self.table,
300+ column = self.column
301+ ))
302+ .context("failed to reinstate column not null constraint")?;
303+
217304 // Remove function and trigger
218305 let trigger_table = match &self.down {
219306 Some(Transformation::Update {
@@ -227,8 +314,13 @@ impl Action for RemoveColumn {
227314 r#"
228315 DROP TRIGGER IF EXISTS "{trigger_name}" ON "{trigger_table}";
229316 DROP FUNCTION IF EXISTS "{trigger_name}";
317+
318+ DROP TRIGGER IF EXISTS "{null_trigger_name}" ON "{table}";
319+ DROP FUNCTION IF EXISTS "{null_trigger_name}";
230320 "#,
321+ table = self.table,
231322 trigger_name = self.trigger_name(ctx),
323+ null_trigger_name = self.null_constraint_trigger_name(ctx),
232324 ))
233325 .context("failed to drop down trigger")?;
234326