Add new create_enum action

README.md

@@ -28,6 +28,8 @@ Designed for Postgres 12 and later.
 	- [Indices](#indices)
 		- [Add index](#add-index)
 		- [Remove index](#remove-index)
+	- [Enums](#enums)
+		- [Create enum](#create-enum)
 - [Commands and options](#commands-and-options)
 	- [`reshape migrate`](#reshape-migrate)
 	- [`reshape complete`](#reshape-complete)
@@ -435,6 +437,19 @@ type = "remove_index"
 index = "name_idx"
 ```
 
+### Enums
+
+#### Create enum
+
+The `create_enum` action will create a new [enum type](https://www.postgresql.org/docs/current/datatype-enum.html) with the specified values.
+
+```toml
+[[actions]]
+type = "create_enum"
+name = "mood"
+values = ["happy", "ok", "sad"]
+```
+
 ## Commands and options
 
 ### `reshape migrate`

src/lib.rs

@@ -103,6 +103,16 @@ impl Reshape {
                 ))?;
             }
 
+            // Remove all enums
+            let enums: Vec<String> = db
+                .query("SELECT typname FROM pg_type WHERE typcategory = 'E'")?
+                .iter()
+                .map(|row| row.get("typname"))
+                .collect();
+            for enum_type in enums {
+                db.run(&format!("DROP TYPE {}", enum_type))?;
+            }
+
             // Reset state
             state.clear(db)?;
 

src/migrations/create_enum.rs

@@ -0,0 +1,83 @@
+use super::{Action, MigrationContext};
+use crate::{
+    db::{Conn, Transaction},
+    schema::Schema,
+};
+use anyhow::Context;
+use serde::{Deserialize, Serialize};
+
+#[derive(Serialize, Deserialize, Debug)]
+pub struct CreateEnum {
+    pub name: String,
+    pub values: Vec<String>,
+}
+
+#[typetag::serde(name = "create_enum")]
+impl Action for CreateEnum {
+    fn describe(&self) -> String {
+        format!("Creating enum \"{}\"", self.name)
+    }
+
+    fn run(
+        &self,
+        _ctx: &MigrationContext,
+        db: &mut dyn Conn,
+        _schema: &Schema,
+    ) -> anyhow::Result<()> {
+        // Check if enum already exists. CREATE TYPE doesn't have
+        // a IF NOT EXISTS option so we have to do it manually.
+        let enum_exists = !db
+            .query(&format!(
+                "
+                SELECT typname
+                FROM pg_catalog.pg_type
+                WHERE typcategory = 'E'
+                AND typname = '{name}'
+                ",
+                name = self.name,
+            ))?
+            .is_empty();
+        if enum_exists {
+            return Ok(());
+        }
+
+        let values_def: Vec<String> = self
+            .values
+            .iter()
+            .map(|value| format!("'{}'", value))
+            .collect();
+
+        db.run(&format!(
+            r#"
+            CREATE TYPE "{name}" AS ENUM ({values})
+            "#,
+            name = self.name,
+            values = values_def.join(", "),
+        ))
+        .context("failed to create enum")?;
+
+        Ok(())
+    }
+
+    fn complete<'a>(
+        &self,
+        _ctx: &MigrationContext,
+        _db: &'a mut dyn Conn,
+    ) -> anyhow::Result<Option<Transaction<'a>>> {
+        Ok(None)
+    }
+
+    fn update_schema(&self, _ctx: &MigrationContext, _schema: &mut Schema) {}
+
+    fn abort(&self, _ctx: &MigrationContext, db: &mut dyn Conn) -> anyhow::Result<()> {
+        db.run(&format!(
+            r#"
+            DROP TYPE IF EXISTS {name}
+            "#,
+            name = self.name,
+        ))
+        .context("failed to drop enum")?;
+
+        Ok(())
+    }
+}

src/migrations/mod.rs

@@ -33,6 +33,9 @@ pub use remove_table::RemoveTable;
 mod rename_table;
 pub use rename_table::RenameTable;
 
+mod create_enum;
+pub use create_enum::CreateEnum;
+
 #[derive(Serialize, Deserialize, Debug)]
 pub struct Migration {
     pub name: String,

tests/create_enum.rs

@@ -0,0 +1,41 @@
+mod common;
+use common::Test;
+
+#[test]
+fn create_enum() {
+    let mut test = Test::new("Create enum");
+
+    test.first_migration(
+        r#"
+		name = "create_enum_and_table"
+
+		[[actions]]
+		type = "create_enum"
+		name = "mood"
+		values = ["happy", "ok", "sad"]
+
+		[[actions]]
+		type = "create_table"
+		name = "updates"
+		primary_key = ["id"]
+
+			[[actions.columns]]
+			name = "id"
+			type = "INTEGER"
+
+			[[actions.columns]]
+			name = "status"
+			type = "mood"
+		"#,
+    );
+
+    test.after_first(|db| {
+        // Valid enum values should succeed
+        db.simple_query(
+            "INSERT INTO updates (id, status) VALUES (1, 'happy'), (2, 'ok'), (3, 'sad')",
+        )
+        .unwrap();
+    });
+
+    test.run();
+}