Handle NOT NULL for column removal with complex down

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")?;
 

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")