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

README.md

@@ -323,8 +323,40 @@ Writing: ./db/schema.sql
 
 dbmate supports options passed to a migration block in the form of `key:value` pairs. List of supported options:
 
+- `env`
 - `transaction`
 
+**env**
+
+`env` is useful if you need to inject some environment variables into the SQL. Only environment variables explicitly listed will be available for templating:
+
+```sql
+-- migrate:up env:ADMIN_PASSWORD
+create role 'adminuser' login password '{{ .ADMIN_PASSWORD }}';
+```
+
+`env` supports multiple occurrences. If a referenced variable is not available the command will return an error:
+
+```sh
+$ dbmate up
+Applying: ...
+Error: template: tmpl:35:51: executing "tmpl" at <.ADMIN_PASSWORD>: map has no entry for key "ADMIN_PASSWORD"
+```
+
+Variables injection leverages Go's [template](https://pkg.go.dev/text/template) features. For example you can provide default values in order to prevent errors:
+
+```sql
+-- migrate:up env:THE_ROLE env:THE_PASSWORD
+create role '{{ or (index . "THE_ROLE") "adminuser" }}' login password '{{ .THE_PASSWORD }}';
+```
+
+The `js` function may help to mitigate risks of SQL injection:
+
+```sql
+-- migrate:up env:THE_ROLE env:THE_PASSWORD
+create role '{{ js .THE_ROLE }}' login password '{{ js .THE_PASSWORD }}';
+```
+
 **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:

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

@@ -13,6 +13,7 @@ import (
 	"sort"
 	"time"
 
+	"github.com/amacneil/dbmate/v2/pkg/dbmate/internal"
 	"github.com/amacneil/dbmate/v2/pkg/dbutil"
 )
 
@@ -329,7 +330,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 +367,11 @@ func (db *DB) Migrate() error {
 	return nil
 }
 
+func (db *DB) resolveRefs(snippet string, envVars []string) (string, error) {
+	envMap := internal.GetEnvMap()
+	return internal.ResolveRefs(snippet, envVars, envMap)
+}
+
 func (db *DB) printVerbose(result sql.Result) {
 	lastInsertID, err := result.LastInsertId()
 	if err == nil {
@@ -491,7 +502,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,18 @@ func TestUp(t *testing.T) {
 			err = sqlDB.QueryRow("select count(*) from users").Scan(&count)
 			require.NoError(t, err)
 			require.Equal(t, 1, count)
+
+			// check the value from env var
+			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)
+
+			// check the value from template default (env var missing)
+			var fromDefault string
+			err = sqlDB.QueryRow("select name from posts where id=2").Scan(&fromDefault)
+			require.NoError(t, err)
+			require.Equal(t, "Dino", fromDefault)
 		})
 	}
 }

pkg/dbmate/internal/envvars.go

@@ -0,0 +1,42 @@
+package internal
+
+import (
+	"bytes"
+	"os"
+	"strings"
+	"text/template"
+)
+
+func 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 ResolveRefs(snippet string, envVars []string, envMap map[string]string) (string, error) {
+	if envVars == nil {
+		return snippet, nil
+	}
+
+	model := make(map[string]string, len(envVars))
+	for _, envVar := range envVars {
+		varValue, exists := envMap[envVar]
+		if exists {
+			model[envVar] = varValue
+		}
+	}
+
+	template := template.Must(template.New("tmpl").Option("missingkey=error").Parse(snippet))
+
+	var buffer bytes.Buffer
+	if err := template.Execute(&buffer, model); err != nil {
+		return "", err
+	}
+
+	return buffer.String(), nil
+}

pkg/dbmate/internal/envvars_test.go

@@ -0,0 +1,112 @@
+package internal_test
+
+import (
+	"testing"
+
+	"github.com/amacneil/dbmate/v2/pkg/dbmate/internal"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestResolveRefs(t *testing.T) {
+	parseUp := `create role '{{ .THE_ROLE }}' login password '{{ .THE_PASSWORD }}';`
+	parsedUpOptsEnvVars := []string{
+		"THE_ROLE",
+		"THE_PASSWORD",
+	}
+	envMap := map[string]string{
+		"THE_ROLE":     "Barney",
+		"THE_PASSWORD": "Betty",
+	}
+
+	resolved, err := internal.ResolveRefs(parseUp, parsedUpOptsEnvVars, envMap)
+
+	require.NoError(t, err)
+	require.Equal(t, "create role 'Barney' login password 'Betty';", resolved)
+}
+
+func TestResolveRefsUsingDefaults(t *testing.T) {
+	parseUp := `create role '{{ or (index . "THE_ROLE") "Fred" }}' login password '{{ .THE_PASSWORD }}';`
+	parsedUpOptsEnvVars := []string{
+		"THE_ROLE",
+		"THE_PASSWORD",
+	}
+	envMap := map[string]string{
+		"THE_PASSWORD": "Wilma",
+	}
+
+	resolved, err := internal.ResolveRefs(parseUp, parsedUpOptsEnvVars, envMap)
+
+	require.NoError(t, err)
+	require.Equal(t, "create role 'Fred' login password 'Wilma';", resolved)
+}
+
+func TestResolveRefsIgnoringDefaults(t *testing.T) {
+	parseUp := `create role '{{ or (index . "THE_ROLE") "Fred" }}' login password '{{ .THE_PASSWORD }}';`
+	parsedUpOptsEnvVars := []string{
+		"THE_ROLE",
+		"THE_PASSWORD",
+	}
+	envMap := map[string]string{
+		"THE_ROLE":     "Dino",
+		"THE_PASSWORD": "Wilma",
+	}
+
+	resolved, err := internal.ResolveRefs(parseUp, parsedUpOptsEnvVars, envMap)
+
+	require.NoError(t, err)
+	require.Equal(t, "create role 'Dino' login password 'Wilma';", resolved)
+}
+
+func TestResolveRefsErroringOnMissingVar(t *testing.T) {
+	parseUp := `create role '{{ or (index . "THE_ROLE") "Fred" }}' login password '{{ .THE_PASSWORD }}';`
+	parsedUpOptsEnvVars := []string{
+		"THE_ROLE",
+		"THE_PASSWORD",
+	}
+	envMap := map[string]string{
+		"THE_ROLE": "Dino",
+	}
+
+	resolved, err := internal.ResolveRefs(parseUp, parsedUpOptsEnvVars, envMap)
+
+	require.Error(t, err)
+	require.Equal(t, "", resolved)
+}
+
+func TestResolveWithSqlInjection(t *testing.T) {
+	parseUp := `create role '{{ .THE_ROLE }}' login password '{{ .THE_PASSWORD }}';`
+	parsedUpOptsEnvVars := []string{
+		"THE_ROLE",
+		"THE_PASSWORD",
+	}
+	envMap := map[string]string{
+		"THE_ROLE":     "Slate'; drop table SALARY; create role 'Barney",
+		"THE_PASSWORD": "Betty",
+	}
+
+	resolved, err := internal.ResolveRefs(parseUp, parsedUpOptsEnvVars, envMap)
+
+	require.NoError(t, err)
+	// sql injection is not prevented here
+	require.Equal(t, "create role 'Slate'; drop table SALARY; create role 'Barney' login password 'Betty';", resolved)
+}
+
+func TestResolveMitigatingSqlInjection(t *testing.T) {
+	parseUp := `create role '{{ js .THE_ROLE }}' login password '{{ js .THE_PASSWORD }}';`
+	parsedUpOptsEnvVars := []string{
+		"THE_ROLE",
+		"THE_PASSWORD",
+	}
+	envMap := map[string]string{
+		// simulating naive SQL injection attempt
+		"THE_ROLE":     "Slate'; drop table SALARY; create role 'Barney",
+		"THE_PASSWORD": "Betty",
+	}
+
+	resolved, err := internal.ResolveRefs(parseUp, parsedUpOptsEnvVars, envMap)
+
+	require.NoError(t, err)
+	// sql injection mitigated
+	require.Equal(t, "create role 'Slate\\'; drop table SALARY; create role \\'Barney' login password 'Betty';", resolved)
+}

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
+			}
 		}
 	}