Use system variables in migrations

Add support to selectively enable usage of env vars within migrations.
Named env vars are enabled though "env-var:VAR_NAME" entries within
the up/down block headers., reusing the existing parsing logic
The name of activated env vars are then kept with migration options and
passed to go templating just before executing relevant migrations.

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-var"]
+
+	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-var" 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-var" 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-var:THE_ROLE env-var:THE_PASSWORD
+create role {{ .THE_ROLE }} login password {{ .THE_PASSWORD }};
+-- migrate:down env-var:THE_ROLE
+drop role {{ .THE_ROLE }};
+`
+
+		parsed, err := parseMigrationContents(migration)
+		require.Nil(t, err)
+
+		require.Equal(t, "-- migrate:up env-var:THE_ROLE env-var: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-var: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-var:THE_ROLE env-var:THE_PASSWORD
+create role {{ .THE_ROLE }} login password {{ .THE_PASSWORD }};
+-- migrate:down env-var:THE_ROLE
+drop role {{ .THE_ROLE }};
+`
+
+		parsed, err := parseMigrationContents(migration)
+		require.Nil(t, err)
+
+		require.Equal(t, "-- migrate:up env-var:THE_ROLE env-var: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-var: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-var: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;