Add new custom action

This action enables migrations to run custom SQL and is meant to be used
when no other actions fit. It doesn't provide any guarantees for
zero-downtime.

README.md

@@ -33,6 +33,7 @@ Designed for Postgres 12 and later.
 	- [Enums](#enums)
 		- [Create enum](#create-enum)
 		- [Remove enum](#remove-enum)
+	- [Custom](#custom)
 - [Commands and options](#commands-and-options)
 	- [`reshape migrate`](#reshape-migrate)
 	- [`reshape complete`](#reshape-complete)
@@ -531,6 +532,33 @@ type = "remove_enum"
 enum = "mood"
 ```
 
+### Custom
+
+The `custom` action lets you create a migration which runs custom SQL. It should be used with great care as it provides no guarantees of zero-downtime and will simply run whatever SQL is provided. Use other actions whenever possible as they are explicitly designed for zero downtime.
+
+There are three optional settings available which all accept SQL queries. All queries need to be idempotent, for example by using `IF NOT EXISTS` wherever available.
+
+- `start`: run when a migration is started using `reshape migrate`
+- `complete`: run when a migration is completed using `reshape complete`
+- `abort`: run when a migration is aborted using `reshape abort`
+
+*Example: enable PostGIS and pg_stat_statements extensions*
+
+```toml
+[[actions]]
+type = "custom"
+
+start = """
+	CREATE EXTENSION IF NOT EXISTS postgis;
+	CREATE EXTENSION IF NOT EXISTS pg_stat_statements;
+"""
+
+abort = """
+	DROP EXTENSION IF NOT EXISTS postgis;
+	DROP EXTENSION IF NOT EXISTS pg_stat_statements;
+"""
+```
+
 ## Commands and options
 
 ### `reshape migrate`

src/migrations/custom.rs

@@ -0,0 +1,61 @@
+use super::{Action, MigrationContext};
+use crate::{
+    db::{Conn, Transaction},
+    schema::Schema,
+};
+use serde::{Deserialize, Serialize};
+
+#[derive(Serialize, Deserialize, Debug)]
+pub struct Custom {
+    #[serde(default)]
+    pub start: Option<String>,
+
+    #[serde(default)]
+    pub complete: Option<String>,
+
+    #[serde(default)]
+    pub abort: Option<String>,
+}
+
+#[typetag::serde(name = "custom")]
+impl Action for Custom {
+    fn describe(&self) -> String {
+        "Running custom migration".to_string()
+    }
+
+    fn run(
+        &self,
+        _ctx: &MigrationContext,
+        db: &mut dyn Conn,
+        _schema: &Schema,
+    ) -> anyhow::Result<()> {
+        if let Some(start_query) = &self.start {
+            println!("Running query: {}", start_query);
+            db.run(start_query)?;
+        }
+
+        Ok(())
+    }
+
+    fn complete<'a>(
+        &self,
+        _ctx: &MigrationContext,
+        db: &'a mut dyn Conn,
+    ) -> anyhow::Result<Option<Transaction<'a>>> {
+        if let Some(complete_query) = &self.complete {
+            db.run(complete_query)?;
+        }
+
+        Ok(None)
+    }
+
+    fn update_schema(&self, _ctx: &MigrationContext, _schema: &mut Schema) {}
+
+    fn abort(&self, _ctx: &MigrationContext, db: &mut dyn Conn) -> anyhow::Result<()> {
+        if let Some(abort_query) = &self.abort {
+            db.run(abort_query)?;
+        }
+
+        Ok(())
+    }
+}

src/migrations/mod.rs

@@ -39,6 +39,9 @@ pub use create_enum::CreateEnum;
 mod remove_enum;
 pub use remove_enum::RemoveEnum;
 
+mod custom;
+pub use custom::Custom;
+
 mod add_foreign_key;
 pub use add_foreign_key::AddForeignKey;
 

tests/common.rs

@@ -12,6 +12,7 @@ pub struct Test<'a> {
     second_migration: Option<Migration>,
     expect_failure: bool,
 
+    clear_fn: Option<fn(&mut Client) -> ()>,
     after_first_fn: Option<fn(&mut Client) -> ()>,
     intermediate_fn: Option<fn(&mut Client, &mut Client) -> ()>,
     after_completion_fn: Option<fn(&mut Client) -> ()>,
@@ -36,6 +37,7 @@ impl Test<'_> {
             first_migration: None,
             second_migration: None,
             expect_failure: false,
+            clear_fn: None,
             after_first_fn: None,
             intermediate_fn: None,
             after_completion_fn: None,
@@ -54,6 +56,12 @@ impl Test<'_> {
         self
     }
 
+    #[allow(dead_code)]
+    pub fn clear(&mut self, f: fn(&mut Client) -> ()) -> &mut Self {
+        self.clear_fn = Some(f);
+        self
+    }
+
     #[allow(dead_code)]
     pub fn after_first(&mut self, f: fn(&mut Client) -> ()) -> &mut Self {
         self.after_first_fn = Some(f);
@@ -115,6 +123,10 @@ impl Test<'_> {
         print_subheading("Clearing database");
         self.reshape.remove().unwrap();
 
+        if let Some(clear_fn) = self.clear_fn {
+            clear_fn(&mut self.old_db);
+        }
+
         // Apply first migration, will automatically complete
         print_subheading("Applying first migration");
         let first_migration = self

tests/custom.rs

@@ -0,0 +1,128 @@
+mod common;
+use common::Test;
+
+#[test]
+fn custom_enable_extension() {
+    let mut test = Test::new("Custom migration");
+
+    test.clear(|db| {
+        db.simple_query(
+            "
+			DROP EXTENSION IF EXISTS bloom;
+			DROP EXTENSION IF EXISTS btree_gin;
+			DROP EXTENSION IF EXISTS btree_gist;
+            ",
+        )
+        .unwrap();
+    });
+
+    test.first_migration(
+        r#"
+		name = "empty_migration"
+
+		[[actions]]
+		type = "custom"
+		"#,
+    );
+
+    test.second_migration(
+        r#"
+		name = "enable_extensions"
+
+		[[actions]]
+		type = "custom"
+
+		start = """
+			CREATE EXTENSION IF NOT EXISTS bloom;
+			CREATE EXTENSION IF NOT EXISTS btree_gin;
+		"""
+
+        complete = "CREATE EXTENSION IF NOT EXISTS btree_gist"
+
+		abort = """
+			DROP EXTENSION IF EXISTS bloom;
+			DROP EXTENSION IF EXISTS btree_gin;
+		"""
+		"#,
+    );
+
+    test.intermediate(|db, _| {
+        let bloom_activated = !db
+            .query("SELECT * FROM pg_extension WHERE extname = 'bloom'", &[])
+            .unwrap()
+            .is_empty();
+        assert!(bloom_activated);
+
+        let btree_gin_activated = !db
+            .query(
+                "SELECT * FROM pg_extension WHERE extname = 'btree_gin'",
+                &[],
+            )
+            .unwrap()
+            .is_empty();
+        assert!(btree_gin_activated);
+
+        let btree_gist_activated = !db
+            .query(
+                "SELECT * FROM pg_extension WHERE extname = 'btree_gist'",
+                &[],
+            )
+            .unwrap()
+            .is_empty();
+        assert!(!btree_gist_activated);
+    });
+
+    test.after_completion(|db| {
+        let bloom_activated = !db
+            .query("SELECT * FROM pg_extension WHERE extname = 'bloom'", &[])
+            .unwrap()
+            .is_empty();
+        assert!(bloom_activated);
+
+        let btree_gin_activated = !db
+            .query(
+                "SELECT * FROM pg_extension WHERE extname = 'btree_gin'",
+                &[],
+            )
+            .unwrap()
+            .is_empty();
+        assert!(btree_gin_activated);
+
+        let btree_gist_activated = !db
+            .query(
+                "SELECT * FROM pg_extension WHERE extname = 'btree_gist'",
+                &[],
+            )
+            .unwrap()
+            .is_empty();
+        assert!(btree_gist_activated);
+    });
+
+    test.after_abort(|db| {
+        let bloom_activated = !db
+            .query("SELECT * FROM pg_extension WHERE extname = 'bloom'", &[])
+            .unwrap()
+            .is_empty();
+        assert!(!bloom_activated);
+
+        let btree_gin_activated = !db
+            .query(
+                "SELECT * FROM pg_extension WHERE extname = 'btree_gin'",
+                &[],
+            )
+            .unwrap()
+            .is_empty();
+        assert!(!btree_gin_activated);
+
+        let btree_gist_activated = !db
+            .query(
+                "SELECT * FROM pg_extension WHERE extname = 'btree_gist'",
+                &[],
+            )
+            .unwrap()
+            .is_empty();
+        assert!(!btree_gist_activated);
+    });
+
+    test.run();
+}