feat: environment variables interpolation (#604)

Co-authored-by: Mike Fridman <mf192@icloud.com>

CHANGELOG.md

@@ -7,6 +7,26 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
 
+- Add environment variable substitution for SQL migrations. (#604)
+
+  - This feature is **disabled by default**, and can be enabled by adding an annotation to the
+    migration file:
+
+    ```sql
+    -- +goose ENVSUB ON
+    ```
+
+  - When enabled, goose will attempt to substitute environment variables in the SQL migration
+    queries until the end of the file, or until the annotation `-- +goose ENVSUB OFF` is found. For
+    example, if the environment variable `REGION` is set to `us_east_1`, the following SQL migration
+    will be substituted to `SELECT * FROM regions WHERE name = 'us_east_1';`
+
+    ```sql
+    -- +goose ENVSUB ON
+    -- +goose Up
+    SELECT * FROM regions WHERE name = '${REGION}';
+    ```
+
 ## [v3.17.0] - 2023-12-15
 
 - Standardised the MIT license (#647)

README.md

@@ -282,6 +282,48 @@ language plpgsql;
 -- +goose StatementEnd
 ```
 
+Goose supports environment variable substitution in SQL migrations through annotations. To enable
+this feature, use the `-- +goose ENVSUB ON` annotation before the queries where you want
+substitution applied. It stays active until the `-- +goose ENVSUB OFF` annotation is encountered.
+You can use these annotations multiple times within a file.
+
+This feature is disabled by default for backward compatibility with existing scripts.
+
+For `PL/pgSQL` functions or other statements where substitution is not desired, wrap the annotations
+explicitly around the relevant parts. For example, to exclude escaping the `**` characters:
+
+```sql
+-- +goose StatementBegin
+CREATE OR REPLACE FUNCTION test_func()
+RETURNS void AS $$
+-- +goose ENVSUB ON
+BEGIN
+	RAISE NOTICE '${SOME_ENV_VAR}';
+END;
+-- +goose ENVSUB OFF
+$$ LANGUAGE plpgsql;
+-- +goose StatementEnd
+```
+
+<details>
+<summary>Supported expansions (click here to expand):</summary>
+
+- `${VAR}` or $VAR - expands to the value of the environment variable `VAR`
+- `${VAR:-default}` - expands to the value of the environment variable `VAR`, or `default` if `VAR`
+  is unset or null
+- `${VAR-default}` - expands to the value of the environment variable `VAR`, or `default` if `VAR`
+  is unset
+- `${VAR?err_msg}` - expands to the value of the environment variable `VAR`, or prints `err_msg` and
+  error if `VAR` unset
+- ~~`${VAR:?err_msg}` - expands to the value of the environment variable `VAR`, or prints `err_msg`
+  and error if `VAR` unset or null.~~ **THIS IS NOT SUPPORTED**
+
+See
+[mfridman/interpolate](https://github.com/mfridman/interpolate?tab=readme-ov-file#supported-expansions)
+for more details on supported expansions.
+
+</details>
+
 ## Embedded sql migrations
 
 Go 1.16 introduced new feature: [compile-time embedding](https://pkg.go.dev/embed/) files into binary and

go.mod

@@ -6,6 +6,7 @@ require (
 	github.com/ClickHouse/clickhouse-go/v2 v2.16.0
 	github.com/go-sql-driver/mysql v1.7.1
 	github.com/jackc/pgx/v5 v5.5.1
+	github.com/mfridman/interpolate v0.0.2
 	github.com/microsoft/go-mssqldb v1.6.0
 	github.com/ory/dockertest/v3 v3.10.0
 	github.com/sethvargo/go-retry v0.2.4

go.sum

@@ -174,6 +174,8 @@ github.com/mattn/go-isatty v0.0.12/go.mod h1:cbi8OIDigv2wuxKPP5vlRcQ1OAZbq2CE4Ky
 github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
 github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
 github.com/mattn/go-sqlite3 v1.14.16 h1:yOQRA0RpS5PFz/oikGwBEqvAWhWg5ufRz4ETLjwpU1Y=
+github.com/mfridman/interpolate v0.0.2 h1:pnuTK7MQIxxFz1Gr+rjSIx9u7qVjf5VOoM/u6BbAxPY=
+github.com/mfridman/interpolate v0.0.2/go.mod h1:p+7uk6oE07mpE/Ik1b8EckO0O4ZXiGAfshKBWLUM9Xg=
 github.com/microsoft/go-mssqldb v1.6.0 h1:mM3gYdVwEPFrlg/Dvr2DNVEgYFG7L42l+dGc67NNNpc=
 github.com/microsoft/go-mssqldb v1.6.0/go.mod h1:00mDtPbeQCRGC1HwOOR5K/gr30P1NcEG0vx6Kbv2aJU=
 github.com/mitchellh/mapstructure v1.5.0 h1:jeMsZIYE/09sWLaz43PL7Gy6RuMjD2eJVyuac5Z2hdY=

internal/sqlparser/parser.go

@@ -7,8 +7,11 @@ import (
 	"fmt"
 	"io"
 	"log"
+	"os"
 	"strings"
 	"sync"
+
+	"github.com/mfridman/interpolate"
 )
 
 type Direction string
@@ -107,6 +110,7 @@ func ParseSQLMigration(r io.Reader, direction Direction, debug bool) (stmts []st
 
 	stateMachine := newStateMachine(start, debug)
 	useTx = true
+	useEnvsub := false
 
 	var buf bytes.Buffer
 	for scanner.Scan() {
@@ -171,6 +175,14 @@ func ParseSQLMigration(r io.Reader, direction Direction, debug bool) (stmts []st
 			case "+goose NO TRANSACTION":
 				useTx = false
 				continue
+
+			case "+goose ENVSUB ON":
+				useEnvsub = true
+				continue
+
+			case "+goose ENVSUB OFF":
+				useEnvsub = false
+				continue
 			}
 		}
 		// Once we've started parsing a statement the buffer is no longer empty,
@@ -187,6 +199,13 @@ func ParseSQLMigration(r io.Reader, direction Direction, debug bool) (stmts []st
 		case gooseStatementEndDown, gooseStatementEndUp:
 			// Do not include the "+goose StatementEnd" annotation in the final statement.
 		default:
+			if useEnvsub {
+				expanded, err := interpolate.Interpolate(&envWrapper{}, line)
+				if err != nil {
+					return nil, false, fmt.Errorf("variable substitution failed: %w:\n%s", err, line)
+				}
+				line = expanded
+			}
 			// Write SQL line to a buffer.
 			if _, err := buf.WriteString(line + "\n"); err != nil {
 				return nil, false, fmt.Errorf("failed to write to buf: %w", err)
@@ -266,7 +285,14 @@ func missingSemicolonError(state parserState, direction Direction, s string) err
 	)
 }
 
-// cleanupStatement trims whitespace from the given statement.
+type envWrapper struct{}
+
+var _ interpolate.Env = (*envWrapper)(nil)
+
+func (e *envWrapper) Get(key string) (string, bool) {
+	return os.LookupEnv(key)
+}
+
 func cleanupStatement(input string) string {
 	return strings.TrimSpace(input)
 }

internal/sqlparser/parser_test.go

@@ -380,9 +380,9 @@ func TestValidUp(t *testing.T) {
 	// to the parser. Then we compare the statements against the golden files.
 	// Each golden file is equivalent to one statement.
 	//
-	// ├── 01.golden.sql
-	// ├── 02.golden.sql
-	// ├── 03.golden.sql
+	// ├── 01.up.golden.sql
+	// ├── 02.up.golden.sql
+	// ├── 03.up.golden.sql
 	// └── input.sql
 	tests := []struct {
 		Name            string
@@ -401,36 +401,36 @@ func TestValidUp(t *testing.T) {
 	for _, tc := range tests {
 		path := filepath.Join("testdata", "valid-up", tc.Name)
 		t.Run(tc.Name, func(t *testing.T) {
-			testValidUp(t, path, tc.StatementsCount)
+			testValid(t, path, tc.StatementsCount, DirectionUp)
 		})
 	}
 }
 
-func testValidUp(t *testing.T, dir string, count int) {
+func testValid(t *testing.T, dir string, count int, direction Direction) {
 	t.Helper()
 
 	f, err := os.Open(filepath.Join(dir, "input.sql"))
 	check.NoError(t, err)
 	t.Cleanup(func() { f.Close() })
-	statements, _, err := ParseSQLMigration(f, DirectionUp, debug)
+	statements, _, err := ParseSQLMigration(f, direction, debug)
 	check.NoError(t, err)
 	check.Number(t, len(statements), count)
-	compareStatements(t, dir, statements)
+	compareStatements(t, dir, statements, direction)
 }
 
-func compareStatements(t *testing.T, dir string, statements []string) {
+func compareStatements(t *testing.T, dir string, statements []string, direction Direction) {
 	t.Helper()
 
-	files, err := filepath.Glob(filepath.Join(dir, "*.golden.sql"))
+	files, err := filepath.Glob(filepath.Join(dir, fmt.Sprintf("*.%s.golden.sql", direction)))
 	check.NoError(t, err)
 	if len(statements) != len(files) {
-		t.Fatalf("mismatch between parsed statements (%d) and golden files (%d), did you check in NN.golden.sql file in %q?", len(statements), len(files), dir)
+		t.Fatalf("mismatch between parsed statements (%d) and golden files (%d), did you check in NN.{up|down}.golden.sql file in %q?", len(statements), len(files), dir)
 	}
 	for _, goldenFile := range files {
 		goldenFile = filepath.Base(goldenFile)
-		before, _, ok := cut(goldenFile, ".")
+		before, _, ok := strings.Cut(goldenFile, ".")
 		if !ok {
-			t.Fatal(`failed to cut on file delimiter ".", must be of the format NN.golden.sql`)
+			t.Fatal(`failed to cut on file delimiter ".", must be of the format NN.{up|down}.golden.sql`)
 		}
 		index, err := strconv.Atoi(before)
 		check.NoError(t, err)
@@ -458,16 +458,52 @@ func compareStatements(t *testing.T, dir string, statements []string) {
 	}
 }
 
-// copied directly from strings.Cut (go1.18) to support older Go versions.
-// In the future, replace this with the upstream function.
-func cut(s, sep string) (before, after string, found bool) {
-	if i := strings.Index(s, sep); i >= 0 {
-		return s[:i], s[i+len(sep):], true
-	}
-	return s, "", false
-}
-
 func isCIEnvironment() bool {
 	ok, _ := strconv.ParseBool(os.Getenv("CI"))
 	return ok
 }
+
+func TestEnvsub(t *testing.T) {
+	// Do not run in parallel, as this test sets environment variables.
+
+	// Test valid migrations with ${var} like statements when on are substituted for the whole
+	// migration.
+	t.Setenv("GOOSE_ENV_REGION", "us_east_")
+	t.Setenv("GOOSE_ENV_SET_BUT_EMPTY_VALUE", "")
+	t.Setenv("GOOSE_ENV_NAME", "foo")
+
+	tests := []struct {
+		Name      string
+		DownCount int
+		UpCount   int
+	}{
+		{Name: "test01", UpCount: 4, DownCount: 1},
+		{Name: "test02", UpCount: 3, DownCount: 0},
+		{Name: "test03", UpCount: 1, DownCount: 0},
+	}
+	for _, tc := range tests {
+		t.Run(tc.Name, func(t *testing.T) {
+			dir := filepath.Join("testdata", "envsub", tc.Name)
+			testValid(t, dir, tc.UpCount, DirectionUp)
+			testValid(t, dir, tc.DownCount, DirectionDown)
+		})
+	}
+}
+
+func TestEnvsubError(t *testing.T) {
+	t.Parallel()
+
+	s := `
+-- +goose ENVSUB ON
+-- +goose Up
+CREATE TABLE post (
+	id int NOT NULL,
+	title text,
+	${SOME_UNSET_VAR?required env var not set} text,
+	PRIMARY KEY(id)
+);
+`
+	_, _, err := ParseSQLMigration(strings.NewReader(s), DirectionUp, debug)
+	check.HasError(t, err)
+	check.Contains(t, err.Error(), "variable substitution failed: $SOME_UNSET_VAR: required env var not set:")
+}

internal/sqlparser/testdata/envsub/test01/01.down.golden.sql

@@ -0,0 +1 @@
+DROP TABLE us_east_post;    -- 1st stmt

internal/sqlparser/testdata/envsub/test01/01.up.golden.sql

@@ -0,0 +1,6 @@
+CREATE TABLE us_east_post (
+		id int NOT NULL,
+		title text,
+		body text,
+		PRIMARY KEY(id)
+);                  -- 1st stmt

internal/sqlparser/testdata/envsub/test01/02.up.golden.sql

@@ -0,0 +1 @@
+SELECT 2;           -- 2nd stmt

internal/sqlparser/testdata/envsub/test01/03.up.golden.sql

@@ -0,0 +1 @@
+SELECT 3; SELECT 3; -- 3rd stmt

internal/sqlparser/testdata/envsub/test01/04.up.golden.sql

@@ -0,0 +1 @@
+SELECT 4;           -- 4th stmt

internal/sqlparser/testdata/envsub/test01/input.sql

@@ -0,0 +1,17 @@
+-- +goose ENVSUB ON
+-- +goose Up
+CREATE TABLE ${GOOSE_ENV_REGION}post (
+		id int NOT NULL,
+		title text,
+		body text,
+		PRIMARY KEY(id)
+);                  -- 1st stmt
+
+-- comment
+SELECT 2;           -- 2nd stmt
+SELECT 3; SELECT 3; -- 3rd stmt
+SELECT 4;           -- 4th stmt
+
+-- +goose Down
+-- comment
+DROP TABLE ${GOOSE_ENV_REGION}post;    -- 1st stmt

internal/sqlparser/testdata/envsub/test02/01.up.golden.sql

@@ -0,0 +1,8 @@
+CREATE TABLE post (
+	id int NOT NULL,
+	title text,
+	foo text,
+	footitle3 text,
+	defaulttitle4 text,
+	title5 text,
+);

internal/sqlparser/testdata/envsub/test02/02.up.golden.sql

@@ -0,0 +1,8 @@
+CREATE TABLE post (
+	id int NOT NULL,
+	title text,
+	$GOOSE_ENV_NAME text,
+	${GOOSE_ENV_NAME}title3 text,
+	${ANOTHER_VAR:-default}title4 text,
+	${GOOSE_ENV_SET_BUT_EMPTY_VALUE-default}title5 text,
+);

internal/sqlparser/testdata/envsub/test02/03.up.golden.sql

@@ -0,0 +1,6 @@
+CREATE OR REPLACE FUNCTION test_func()
+RETURNS void AS $$
+BEGIN
+	RAISE NOTICE 'foo $GOOSE_ENV_NAME $GOOSE_ENV_NAME';
+END;
+$$ LANGUAGE plpgsql;

internal/sqlparser/testdata/envsub/test02/input.sql

@@ -0,0 +1,32 @@
+-- +goose Up
+
+-- +goose ENVSUB ON
+CREATE TABLE post (
+	id int NOT NULL,
+	title text,
+	$GOOSE_ENV_NAME text,
+	${GOOSE_ENV_NAME}title3 text,
+	${ANOTHER_VAR:-default}title4 text,
+	${GOOSE_ENV_SET_BUT_EMPTY_VALUE-default}title5 text,
+);
+-- +goose ENVSUB OFF
+
+CREATE TABLE post (
+	id int NOT NULL,
+	title text,
+	$GOOSE_ENV_NAME text,
+	${GOOSE_ENV_NAME}title3 text,
+	${ANOTHER_VAR:-default}title4 text,
+	${GOOSE_ENV_SET_BUT_EMPTY_VALUE-default}title5 text,
+);
+
+-- +goose StatementBegin
+CREATE OR REPLACE FUNCTION test_func()
+RETURNS void AS $$
+-- +goose ENVSUB ON
+BEGIN
+	RAISE NOTICE '${GOOSE_ENV_NAME} \$GOOSE_ENV_NAME \$GOOSE_ENV_NAME';
+END;
+-- +goose ENVSUB OFF
+$$ LANGUAGE plpgsql;
+-- +goose StatementEnd

internal/sqlparser/testdata/envsub/test03/01.up.golden.sql

@@ -0,0 +1,8 @@
+CREATE TABLE post (
+	id int NOT NULL,
+	title text,
+	$NAME text,
+	${NAME}title3 text,
+	${ANOTHER_VAR:-default}title4 text,
+	${SET_BUT_EMPTY_VALUE-default}title5 text,
+);

internal/sqlparser/testdata/envsub/test03/input.sql

@@ -0,0 +1,9 @@
+-- +goose Up
+CREATE TABLE post (
+	id int NOT NULL,
+	title text,
+	$NAME text,
+	${NAME}title3 text,
+	${ANOTHER_VAR:-default}title4 text,
+	${SET_BUT_EMPTY_VALUE-default}title5 text,
+);