Use shell environment variables in migrations using Go's text/template parsing

docker-compose.yml

@@ -15,6 +15,7 @@ services:
       MYSQL_TEST_URL: mysql://root:root@mysql/dbmate_test
       POSTGRES_TEST_URL: postgres://postgres:postgres@postgres/dbmate_test?sslmode=disable
       SQLITE_TEST_URL: sqlite3:/tmp/dbmate_test.sqlite3
+      YABBA_DABBA_DOO: Yabba dabba doo!
 
   dbmate:
     build:

pkg/dbmate/db.go

@@ -1,6 +1,7 @@
 package dbmate
 
 import (
+	"bytes"
 	"database/sql"
 	"errors"
 	"fmt"
@@ -11,6 +12,8 @@ import (
 	"path/filepath"
 	"regexp"
 	"sort"
+	"strings"
+	"text/template"
 	"time"
 
 	"github.com/amacneil/dbmate/v2/pkg/dbutil"
@@ -329,7 +332,12 @@ func (db *DB) Migrate() error {
 
 		execMigration := func(tx dbutil.Transaction) error {
 			// run actual migration
-			result, err := tx.Exec(parsed.Up)
+			upScript, err := db.resolveRefs(parsed.Up, parsed.UpOptions.EnvVars())
+			if err != nil {
+				return err
+			}
+
+			result, err := tx.Exec(upScript)
 			if err != nil {
 				return err
 			} else if db.Verbose {
@@ -361,6 +369,38 @@ func (db *DB) Migrate() error {
 	return nil
 }
 
+func (db *DB) resolveRefs(snippet string, envVars []string) (string, error) {
+	if envVars == nil {
+		return snippet, nil
+	}
+
+	envMap := db.getEnvMap()
+	model := make(map[string]string, len(envVars))
+	for _, envVar := range envVars {
+		model[envVar] = envMap[envVar]
+	}
+
+	template := template.Must(template.New("tmpl").Parse(snippet))
+
+	var buffer bytes.Buffer
+	if err := template.Execute(&buffer, model); err != nil {
+		return "", err
+	}
+
+	return buffer.String(), nil
+}
+
+func (db *DB) getEnvMap() map[string]string {
+	envMap := make(map[string]string)
+
+	for _, envVar := range os.Environ() {
+		entry := strings.SplitN(envVar, "=", 2)
+		envMap[entry[0]] = entry[1]
+	}
+
+	return envMap
+}
+
 func (db *DB) printVerbose(result sql.Result) {
 	lastInsertID, err := result.LastInsertId()
 	if err == nil {
@@ -491,7 +531,13 @@ func (db *DB) Rollback() error {
 
 	execMigration := func(tx dbutil.Transaction) error {
 		// rollback migration
-		result, err := tx.Exec(parsed.Down)
+		downScript, err := db.resolveRefs(parsed.Down, parsed.DownOptions.EnvVars())
+		if err != nil {
+			return err
+		}
+
+		result, err := tx.Exec(downScript)
+
 		if err != nil {
 			return err
 		} else if db.Verbose {

pkg/dbmate/db_test.go

@@ -237,7 +237,7 @@ func TestWaitBeforeVerbose(t *testing.T) {
 		`Applying: 20151129054053_test_migration.sql
 Rows affected: 1
 Applying: 20200227231541_test_posts.sql
-Rows affected: 0`)
+Rows affected: 1`)
 	require.Contains(t, output,
 		`Rolling back: 20200227231541_test_posts.sql
 Rows affected: 0`)
@@ -315,6 +315,11 @@ func TestUp(t *testing.T) {
 			err = sqlDB.QueryRow("select count(*) from users").Scan(&count)
 			require.NoError(t, err)
 			require.Equal(t, 1, count)
+
+			var fromEnvVar string
+			err = sqlDB.QueryRow("select name from posts where id=1").Scan(&fromEnvVar)
+			require.NoError(t, err)
+			require.Equal(t, "Yabba dabba doo!", fromEnvVar)
 		})
 	}
 }

pkg/dbmate/migration.go

@@ -48,6 +48,7 @@ type ParsedMigration struct {
 // ParsedMigrationOptions is an interface for accessing migration options
 type ParsedMigrationOptions interface {
 	Transaction() bool
+	EnvVars() []string
 }
 
 type migrationOptions map[string]string
@@ -58,6 +59,21 @@ func (m migrationOptions) Transaction() bool {
 	return m["transaction"] != "false"
 }
 
+// EnvVars returns the list of env vars enabled to templating for this migration
+// Defaults to empty list.
+func (m migrationOptions) EnvVars() []string {
+	result := make([]string, 0)
+	entry := m["env"]
+
+	if entry != "" {
+		// decode CSV encoded var names
+		varNames := strings.Split(entry, ",")
+		// add to the slice
+		result = append(result, varNames...)
+	}
+	return result
+}
+
 var (
 	upRegExp              = regexp.MustCompile(`(?m)^--\s*migrate:up(\s*$|\s+\S+)`)
 	downRegExp            = regexp.MustCompile(`(?m)^--\s*migrate:down(\s*$|\s+\S+)$`)
@@ -143,7 +159,18 @@ func parseMigrationOptions(contents string) ParsedMigrationOptions {
 
 		// if the syntax is well-formed, then store the key and value pair in options
 		if len(pair) == 2 {
-			options[pair[0]] = pair[1]
+			optKey := pair[0]
+			optValue := pair[1]
+			entry := options[optKey]
+			if entry != "" { // "env" entry already used
+				varNames := strings.Split(entry, ",")
+				// add new element to the slice
+				varNames = append(varNames, optValue)
+				// keep collected values
+				options[optKey] = strings.Join(varNames, ",")
+			} else { // first "env" entry
+				options[optKey] = optValue
+			}
 		}
 	}
 

pkg/dbmate/migration_test.go

@@ -65,9 +65,11 @@ drop table users;
 
 		require.Equal(t, "--migrate:up\ncreate table users (id serial, name text);\n\n", parsed.Up)
 		require.Equal(t, true, parsed.UpOptions.Transaction())
+		require.Equal(t, []string{}, parsed.UpOptions.EnvVars())
 
 		require.Equal(t, "--migrate:down\ndrop table users;\n", parsed.Down)
 		require.Equal(t, true, parsed.DownOptions.Transaction())
+		require.Equal(t, []string{}, parsed.UpOptions.EnvVars())
 	})
 
 	t.Run("require up before down", func(t *testing.T) {
@@ -100,6 +102,40 @@ ALTER TYPE colors ADD VALUE 'orange' AFTER 'red';
 		require.Equal(t, false, parsed.DownOptions.Transaction())
 	})
 
+	t.Run("support activating env vars", func(t *testing.T) {
+		migration := `-- migrate:up env:THE_ROLE env:THE_PASSWORD
+create role {{ .THE_ROLE }} login password {{ .THE_PASSWORD }};
+-- migrate:down env:THE_ROLE
+drop role {{ .THE_ROLE }};
+`
+
+		parsed, err := parseMigrationContents(migration)
+		require.Nil(t, err)
+
+		require.Equal(t, "-- migrate:up env:THE_ROLE env:THE_PASSWORD\ncreate role {{ .THE_ROLE }} login password {{ .THE_PASSWORD }};\n", parsed.Up)
+		require.Equal(t, []string{"THE_ROLE", "THE_PASSWORD"}, parsed.UpOptions.EnvVars())
+
+		require.Equal(t, "-- migrate:down env:THE_ROLE\ndrop role {{ .THE_ROLE }};\n", parsed.Down)
+		require.Equal(t, []string{"THE_ROLE"}, parsed.DownOptions.EnvVars())
+	})
+
+	t.Run("support activating env vars", func(t *testing.T) {
+		migration := `-- migrate:up env:THE_ROLE env:THE_PASSWORD
+create role {{ .THE_ROLE }} login password {{ .THE_PASSWORD }};
+-- migrate:down env:THE_ROLE
+drop role {{ .THE_ROLE }};
+`
+
+		parsed, err := parseMigrationContents(migration)
+		require.Nil(t, err)
+
+		require.Equal(t, "-- migrate:up env:THE_ROLE env:THE_PASSWORD\ncreate role {{ .THE_ROLE }} login password {{ .THE_PASSWORD }};\n", parsed.Up)
+		require.Equal(t, []string{"THE_ROLE", "THE_PASSWORD"}, parsed.UpOptions.EnvVars())
+
+		require.Equal(t, "-- migrate:down env:THE_ROLE\ndrop role {{ .THE_ROLE }};\n", parsed.Down)
+		require.Equal(t, []string{"THE_ROLE"}, parsed.DownOptions.EnvVars())
+	})
+
 	t.Run("require migrate blocks", func(t *testing.T) {
 		migration := `
 ALTER TABLE users

testdata/db/migrations/20200227231541_test_posts.sql

@@ -1,8 +1,9 @@
--- migrate:up
+-- migrate:up env:YABBA_DABBA_DOO
 create table posts (
   id integer,
   name varchar(255)
 );
+insert into posts (id, name) values (1, '{{ .YABBA_DABBA_DOO }}');
 
 -- migrate:down
 drop table posts;