Repository: amacneil/dbmate
Pull request #68 (2019-05-16 04:07:53)

Add support for options; transaction option
Addresses #46 

This PR adds support for options. The only option implemented here is `transaction` which runs the the migration outside of a transaction when `false`. This is useful for certain statements which cannot run inside of a transaction, such as the `ALTER TYPE` statement used for adding values to enums in Postgres.

I took care to make sure this didn't break any existing functionality, but please point out if I missed something.

I'd love to end-to-end test this, but not every db in the test suite supports PG's enum syntax. I'm sure we could restructure some of the tests such that we end-to-end test this, but I only have so much time I can spend on this. For now, it's unit tested and should cover all the necessary cases.


Patch:
diff --git a/README.md b/README.md
index e4e09f6..63bb976 100644
--- a/README.md
+++ b/README.md
@@ -205,6 +205,24 @@ Rolling back: 20151127184807_create_users_table.sql
 Writing: ./db/schema.sql
 ```
 
+### Migration Options
+
+dbmate supports options passed to a migration block in the form of `key:value` pairs. List of supported options:
+
+* `transaction`
+
+#### transaction
+
+`transaction` is useful if you need to run some SQL which cannot be executed from within a transaction. For example, in Postgres, you would need to disable transactions for migrations that alter an enum type to add a value:
+
+```sql
+-- migrate:up transaction:false
+ALTER TYPE colors ADD VALUE 'orange' AFTER 'red';
+ALTER TYPE colors ADD VALUE 'yellow' AFTER 'orange';
+```
+
+`transaction` will default to `true` for if your database supports it.
+
 ### Schema File
 
 When you run the `up`, `migrate`, or `rollback` commands, dbmate will automatically create a `./db/schema.sql` file containing a complete representation of your database schema. Dbmate keeps this file up to date for you, so you should not manually edit it.
diff --git a/pkg/dbmate/db.go b/pkg/dbmate/db.go
index 1e1a196..f7a786d 100644
--- a/pkg/dbmate/db.go
+++ b/pkg/dbmate/db.go
@@ -253,21 +253,29 @@ func (db *DB) Migrate() error {
 
 		fmt.Printf("Applying: %s\n", filename)
 
-		migration, err := parseMigration(filepath.Join(db.MigrationsDir, filename))
+		up, _, err := parseMigration(filepath.Join(db.MigrationsDir, filename))
 		if err != nil {
 			return err
 		}
 
-		// begin transaction
-		err = doTransaction(sqlDB, func(tx Transaction) error {
+		execMigration := func(tx Transaction) error {
 			// run actual migration
-			if _, err := tx.Exec(migration["up"]); err != nil {
+			if _, err := tx.Exec(up.Contents); err != nil {
 				return err
 			}
 
 			// record migration
 			return drv.InsertMigration(tx, ver)
-		})
+		}
+
+		if up.Options.Transaction() {
+			// begin transaction
+			err = doTransaction(sqlDB, execMigration)
+		} else {
+			// run outside of transaction
+			err = execMigration(sqlDB)
+		}
+
 		if err != nil {
 			return err
 		}
@@ -331,43 +339,6 @@ func migrationVersion(filename string) string {
 	return regexp.MustCompile(`^\d+`).FindString(filename)
 }
 
-// parseMigration reads a migration file into a map with up/down keys
-// implementation is similar to regexp.Split()
-func parseMigration(path string) (map[string]string, error) {
-	// read migration file into string
-	data, err := ioutil.ReadFile(path)
-	if err != nil {
-		return nil, err
-	}
-	contents := string(data)
-
-	// split string on our trigger comment
-	separatorRegexp := regexp.MustCompile(`(?m)^-- migrate:(.*)$`)
-	matches := separatorRegexp.FindAllStringSubmatchIndex(contents, -1)
-
-	migrations := map[string]string{}
-	direction := ""
-	beg := 0
-	end := 0
-
-	for _, match := range matches {
-		end = match[0]
-		if direction != "" {
-			// write previous direction to output map
-			migrations[direction] = contents[beg:end]
-		}
-
-		// each match records the start of a new direction
-		direction = contents[match[2]:match[3]]
-		beg = match[1]
-	}
-
-	// write final direction to output map
-	migrations[direction] = contents[beg:]
-
-	return migrations, nil
-}
-
 // Rollback rolls back the most recent migration
 func (db *DB) Rollback() error {
 	drv, sqlDB, err := db.openDatabaseForMigration()
@@ -397,21 +368,29 @@ func (db *DB) Rollback() error {
 
 	fmt.Printf("Rolling back: %s\n", filename)
 
-	migration, err := parseMigration(filepath.Join(db.MigrationsDir, filename))
+	_, down, err := parseMigration(filepath.Join(db.MigrationsDir, filename))
 	if err != nil {
 		return err
 	}
 
-	// begin transaction
-	err = doTransaction(sqlDB, func(tx Transaction) error {
+	execMigration := func(tx Transaction) error {
 		// rollback migration
-		if _, err := tx.Exec(migration["down"]); err != nil {
+		if _, err := tx.Exec(down.Contents); err != nil {
 			return err
 		}
 
 		// remove migration record
 		return drv.DeleteMigration(tx, latest)
-	})
+	}
+
+	if down.Options.Transaction() {
+		// begin transaction
+		err = doTransaction(sqlDB, execMigration)
+	} else {
+		// run outside of transaction
+		err = execMigration(sqlDB)
+	}
+
 	if err != nil {
 		return err
 	}
diff --git a/pkg/dbmate/migrations.go b/pkg/dbmate/migrations.go
new file mode 100644
index 0000000..526280d
--- /dev/null
+++ b/pkg/dbmate/migrations.go
@@ -0,0 +1,123 @@
+package dbmate
+
+import (
+	"io/ioutil"
+	"regexp"
+	"strings"
+)
+
+// MigrationOptions is an interface for accessing migration options
+type MigrationOptions interface {
+	Transaction() bool
+}
+
+type migrationOptions map[string]string
+
+// Transaction returns whether or not this migration should run in a transaction
+// Defaults to true.
+func (m migrationOptions) Transaction() bool {
+	return m["transaction"] != "false"
+}
+
+// Migration contains the migration contents and options
+type Migration struct {
+	Contents string
+	Options  MigrationOptions
+}
+
+// NewMigration constructs a Migration object
+func NewMigration() Migration {
+	return Migration{Contents: "", Options: make(migrationOptions)}
+}
+
+// parseMigration reads a migration file and returns (up Migration, down Migration, error)
+func parseMigration(path string) (Migration, Migration, error) {
+	data, err := ioutil.ReadFile(path)
+	if err != nil {
+		return NewMigration(), NewMigration(), err
+	}
+	up, down := parseMigrationContents(string(data))
+	return up, down, nil
+}
+
+var upRegExp = regexp.MustCompile(`(?m)^-- migrate:up\s+(.+)*$`)
+var downRegExp = regexp.MustCompile(`(?m)^-- migrate:down\s+(.+)*$`)
+
+// parseMigrationContents parses the string contents of a migration.
+// It will return two Migration objects, the first representing the "up"
+// block and the second representing the "down" block.
+//
+// Note that with the way this is currently defined, it is possible to
+// correctly parse a migration that does not define an "up" block or a
+// "down" block, or one that defines neither. This behavior is, in part,
+// to preserve backwards compatibility.
+func parseMigrationContents(contents string) (Migration, Migration) {
+	up := NewMigration()
+	down := NewMigration()
+
+	upMatch := upRegExp.FindStringSubmatchIndex(contents)
+	downMatch := downRegExp.FindStringSubmatchIndex(contents)
+
+	onlyDefinedUpBlock := len(upMatch) != 0 && len(downMatch) == 0
+	onlyDefinedDownBlock := len(upMatch) == 0 && len(downMatch) != 0
+
+	if onlyDefinedUpBlock {
+		up.Contents = strings.TrimSpace(contents)
+		up.Options = parseMigrationOptions(contents, upMatch[2], upMatch[3])
+	} else if onlyDefinedDownBlock {
+		down.Contents = strings.TrimSpace(contents)
+		down.Options = parseMigrationOptions(contents, downMatch[2], downMatch[3])
+	} else {
+		upStart := upMatch[0]
+		downStart := downMatch[0]
+
+		upEnd := downMatch[0]
+		downEnd := len(contents)
+
+		// If migrate:down was defined above migrate:up, correct the end indices
+		if upMatch[0] > downMatch[0] {
+			upEnd = downEnd
+			downEnd = upMatch[0]
+		}
+
+		up.Contents = strings.TrimSpace(contents[upStart:upEnd])
+		up.Options = parseMigrationOptions(contents, upMatch[2], upMatch[3])
+
+		down.Contents = strings.TrimSpace(contents[downStart:downEnd])
+		down.Options = parseMigrationOptions(contents, downMatch[2], downMatch[3])
+	}
+
+	return up, down
+}
+
+var whitespaceRegExp = regexp.MustCompile(`\s+`)
+var optionSeparatorRegExp = regexp.MustCompile(`:`)
+
+// parseMigrationOptions parses the options portion of a migration
+// block into an object that satisfies the MigrationOptions interface,
+// i.e., the 'transaction:false' piece of the following:
+//
+//     -- migrate:up transaction:false
+//     create table users (id serial, name string);
+//     -- migrate:down
+//     drop table users;
+//
+func parseMigrationOptions(contents string, begin, end int) MigrationOptions {
+	mOpts := make(migrationOptions)
+
+	if begin == -1 || end == -1 {
+		return mOpts
+	}
+
+	optionsString := strings.TrimSpace(contents[begin:end])
+
+	optionGroups := whitespaceRegExp.Split(optionsString, -1)
+	for _, group := range optionGroups {
+		pair := optionSeparatorRegExp.Split(group, -1)
+		if len(pair) == 2 {
+			mOpts[pair[0]] = pair[1]
+		}
+	}
+
+	return mOpts
+}


Test patch:
diff --git a/pkg/dbmate/migrations_test.go b/pkg/dbmate/migrations_test.go
new file mode 100644
index 0000000..bbb46c6
--- /dev/null
+++ b/pkg/dbmate/migrations_test.go
@@ -0,0 +1,52 @@
+package dbmate
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestParseMigrationContents(t *testing.T) {
+	migration := `-- migrate:up
+create table users (id serial, name text);
+-- migrate:down
+drop table users;`
+
+	up, down := parseMigrationContents(migration)
+
+	require.Equal(t, "-- migrate:up\ncreate table users (id serial, name text);", up.Contents)
+	require.Equal(t, true, up.Options.Transaction())
+
+	require.Equal(t, "-- migrate:down\ndrop table users;", down.Contents)
+	require.Equal(t, true, down.Options.Transaction())
+
+	migration = `-- migrate:down
+drop table users;
+-- migrate:up
+create table users (id serial, name text);
+`
+
+	up, down = parseMigrationContents(migration)
+
+	require.Equal(t, "-- migrate:up\ncreate table users (id serial, name text);", up.Contents)
+	require.Equal(t, true, up.Options.Transaction())
+
+	require.Equal(t, "-- migrate:down\ndrop table users;", down.Contents)
+	require.Equal(t, true, down.Options.Transaction())
+
+	// This migration would not work in Postgres if it were to
+	// run in a transaction, so we would want to disable transactions.
+	migration = `-- migrate:up transaction:false
+ALTER TYPE colors ADD VALUE 'orange' AFTER 'red';
+ALTER TYPE colors ADD VALUE 'yellow' AFTER 'orange';
+`
+
+	up, down = parseMigrationContents(migration)
+
+	require.Equal(t, "-- migrate:up transaction:false\nALTER TYPE colors ADD VALUE 'orange' AFTER 'red';\nALTER TYPE colors ADD VALUE 'yellow' AFTER 'orange';", up.Contents)
+	require.Equal(t, false, up.Options.Transaction())
+
+	require.Equal(t, "", down.Contents)
+	require.Equal(t, true, down.Options.Transaction())
+
+}