diff --git a/src/migrations/remove_column.rs b/src/migrations/remove_column.rs index ff09e46..20322d6 100644 --- a/src/migrations/remove_column.rs +++ b/src/migrations/remove_column.rs @@ -3,7 +3,7 @@ use crate::{ db::{Conn, Transaction}, schema::Schema, }; -use anyhow::{bail, Context}; +use anyhow::{anyhow, bail, Context}; use serde::{Deserialize, Serialize}; #[derive(Serialize, Deserialize, Debug)] @@ -33,6 +33,15 @@ impl RemoveColumn { self.column ) } + + fn null_constraint_trigger_name(&self, ctx: &MigrationContext) -> String { + format!( + "{}_remove_column_{}_{}_nn", + ctx.prefix(), + self.table, + self.column + ) + } } #[typetag::serde(name = "remove_column")] @@ -51,6 +60,9 @@ impl Action for RemoveColumn { schema: &Schema, ) -> anyhow::Result<()> { let table = schema.get_table(db, &self.table)?; + let column = table + .get_column(&self.column) + .ok_or_else(|| anyhow!("no such column {} exists", self.column))?; // Add down trigger if let Some(down) = &self.down { @@ -122,6 +134,62 @@ impl Action for RemoveColumn { }) .collect(); + let maybe_null_check = if !column.nullable { + // Replace NOT NULL constraint with a constraint trigger that only triggers on the old schema. + // As we are using a complex down function, we must remove the NOT NULL check for the new schema. + // NOT NULL is not checked at the end of a transaction, but immediately upon update. + let query = format!( + r#" + CREATE OR REPLACE FUNCTION {trigger_name}() + RETURNS TRIGGER AS $$ + BEGIN + IF NOT reshape.is_new_schema() THEN + IF NEW.{column} IS NULL THEN + RAISE EXCEPTION '{column} can not be null'; + END IF; + END IF; + RETURN NEW; + END + $$ language 'plpgsql'; + + DROP TRIGGER IF EXISTS "{trigger_name}" ON "{table}"; + + CREATE CONSTRAINT TRIGGER "{trigger_name}" + AFTER INSERT OR UPDATE + ON "{table}" + FOR EACH ROW + EXECUTE PROCEDURE {trigger_name}(); + "#, + table = self.table, + trigger_name = self.null_constraint_trigger_name(ctx), + column = self.column, + ); + db.run(&query) + .context("failed to create null constraint trigger")?; + + db.run(&format!( + r#" + ALTER TABLE {table} + ALTER COLUMN {column} + DROP NOT NULL + "#, + table = self.table, + column = self.column + )) + .context("failed to remove column not null constraint")?; + + format!( + r#" + IF {value} IS NULL THEN + RAISE EXCEPTION '{column_name} can not be null'; + END IF; + "#, + column_name = self.column, + ) + } else { + "".to_string() + }; + // TODO: If column is NOT NULL, remove the constraint and perform a NULL check in some other way. // Otherwise, it's not possible to update the NOT NULL column from another table even in the same transaction. // 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 { DECLARE {declarations} BEGIN + {maybe_null_check} UPDATE "migration_{existing_schema_name}"."{changed_table}" "{changed_table}" SET "{column_name}" = {value} WHERE {where}; @@ -157,6 +226,8 @@ impl Action for RemoveColumn { } } + if !column.nullable {} + Ok(()) } @@ -194,10 +265,14 @@ impl Action for RemoveColumn { DROP TRIGGER IF EXISTS "{trigger_name}" ON "{trigger_table}"; DROP FUNCTION IF EXISTS "{trigger_name}"; + + DROP TRIGGER IF EXISTS "{null_trigger_name}" ON "{table}"; + DROP FUNCTION IF EXISTS "{null_trigger_name}"; "#, table = self.table, column = self.column, trigger_name = self.trigger_name(ctx), + null_trigger_name = self.null_constraint_trigger_name(ctx), ); db.run(&query) .context("failed to drop column and down trigger")?; @@ -214,6 +289,18 @@ impl Action for RemoveColumn { } fn abort(&self, ctx: &MigrationContext, db: &mut dyn Conn) -> anyhow::Result<()> { + // We might have removed the NOT NULL check so we reinstate it + db.run(&format!( + r#" + ALTER TABLE {table} + ALTER COLUMN {column} + SET NOT NULL + "#, + table = self.table, + column = self.column + )) + .context("failed to reinstate column not null constraint")?; + // Remove function and trigger let trigger_table = match &self.down { Some(Transformation::Update { @@ -227,8 +314,13 @@ impl Action for RemoveColumn { r#" DROP TRIGGER IF EXISTS "{trigger_name}" ON "{trigger_table}"; DROP FUNCTION IF EXISTS "{trigger_name}"; + + DROP TRIGGER IF EXISTS "{null_trigger_name}" ON "{table}"; + DROP FUNCTION IF EXISTS "{null_trigger_name}"; "#, + table = self.table, trigger_name = self.trigger_name(ctx), + null_trigger_name = self.null_constraint_trigger_name(ctx), )) .context("failed to drop down trigger")?; diff --git a/tests/complex.rs b/tests/complex.rs index 09fc171..661a4a1 100644 --- a/tests/complex.rs +++ b/tests/complex.rs @@ -30,10 +30,12 @@ fn extract_relation_into_new_table() { [[actions.columns]] name = "account_id" type = "INTEGER" + nullable = false [[actions.columns]] name = "account_role" type = "TEXT" + nullable = false "#, ); @@ -57,6 +59,7 @@ fn extract_relation_into_new_table() { [[actions.columns]] name = "role" type = "TEXT" + nullable = false [actions.up] table = "users" @@ -134,6 +137,13 @@ fn extract_relation_into_new_table() { == 1 ); + // Ensure NOT NULL constraint still applies to old schema + let result = old_db + .simple_query( + "INSERT INTO users (id, account_id, account_role) VALUES (2, NULL, 'developer')", + ); + assert!(result.is_err()); + // Ensure updated user role in old schema updates connection in new schema old_db .simple_query("UPDATE users SET account_role = 'admin' WHERE id = 2")