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 variable 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 occurrencies. 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 injectio 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
+ }
  }
  }