templating: add logs around template exec

database/clickhouse/clickhouse.go

@@ -192,7 +192,13 @@ func (ch *ClickHouse) SetVersion(version int, dirty bool) error {
 		return err
 	}
 
-	query := "INSERT INTO " + ch.config.MigrationsTable + " (version, dirty, sequence) VALUES (?, ?, ?)"
+	query := fmt.Sprintf(
+		"INSERT INTO %s (version, dirty, sequence) VALUES (%d, %d, %d)",
+		ch.config.MigrationsTable,
+		version,
+		bool(dirty),
+		time.Now().UnixNano(),
+	)
 	if _, err := tx.Exec(query, version, bool(dirty), time.Now().UnixNano()); err != nil {
 		return &database.Error{OrigErr: err, Query: []byte(query)}
 	}

internal/cli/main.go

@@ -63,6 +63,7 @@ func Main(version string) {
 	helpPtr := flag.Bool("help", false, "")
 	versionPtr := flag.Bool("version", false, "")
 	verbosePtr := flag.Bool("verbose", false, "")
+	templatePtr := flag.Bool("template", false, "")
 	prefetchPtr := flag.Uint("prefetch", 10, "")
 	lockTimeoutPtr := flag.Uint("lock-timeout", 15, "")
 	pathPtr := flag.String("path", "", "")
@@ -80,6 +81,19 @@ Options:
   -database        Run migrations against this database (driver://url)
   -prefetch N      Number of migrations to load in advance before executing (default 10)
   -lock-timeout N  Allow N seconds to acquire database lock (default 15)
+  -template        Treat migration files as go text templates; making environment variables accessible in your
+                   migration files. 
+                     i.e. If you set the LOCAL_WAREHOUSE environment variable to MY_DB and have a migration file with the
+                     following contents:
+                       INSERT INTO {{.LOCAL_WAREHOUSE}}.INVENTORY.RECORDS ('foo') VALUES ('bar');
+                     it will be transformed into the following before being executed:
+                       INSERT INTO MY_DB.INVENTORY.RECORDS ('foo') VALUES ('bar');
+
+                     Note that enabling templating requires that the contents of the migration file be brought into memory 
+                     in order to perform the transformation, and streaming the file directly from the source driver into
+                     the database is not currently possible with the go templating implementation.
+
+                     See https://pkg.go.dev/text/template for more information on supported template formats.
   -verbose         Print verbose logging
   -version         Print version
   -help            Print usage
@@ -134,6 +148,9 @@ Database drivers: `+strings.Join(database.List(), ", ")+"\n", createUsage, gotoU
 		migrater.Log = log
 		migrater.PrefetchMigrations = *prefetchPtr
 		migrater.LockTimeout = time.Duration(int64(*lockTimeoutPtr)) * time.Second
+		if *templatePtr {
+			migrater.EnableTemplating = true
+		}
 
 		// handle Ctrl+c
 		signals := make(chan os.Signal, 1)

migrate.go

@@ -7,6 +7,7 @@ package migrate
 import (
 	"errors"
 	"fmt"
+	"io"
 	"os"
 	"sync"
 	"time"
@@ -80,6 +81,17 @@ type Migrate struct {
 	// LockTimeout defaults to DefaultLockTimeout,
 	// but can be set per Migrate instance.
 	LockTimeout time.Duration
+
+	// EnableTemplating treats migration files as go text templates; making environment variables accessible in your
+	// migration files.
+	// i.e. If you set the LOCAL_WAREHOUSE environment variable to MY_DB and have a migration file with the
+	// following contents:
+	//   INSERT INTO {{.LOCAL_WAREHOUSE}}.INVENTORY.RECORDS ('foo') VALUES ('bar');
+	// it will be transformed into the following before being executed:
+	//   INSERT INTO MY_DB.INVENTORY.RECORDS ('foo') VALUES ('bar');
+	//
+	// See https://pkg.go.dev/text/template for more information on template formats
+	EnableTemplating bool
 }
 
 // New returns a new Migrate instance from a source URL and a database URL.
@@ -833,45 +845,36 @@ func (m *Migrate) stop() bool {
 // specified version and targetVersion.
 func (m *Migrate) newMigration(version uint, targetVersion int) (*Migration, error) {
 	var migr *Migration
+	var r io.ReadCloser
+	var identifier string
+	var err error
 
 	if targetVersion >= int(version) {
-		r, identifier, err := m.sourceDrv.ReadUp(version)
-		if errors.Is(err, os.ErrNotExist) {
-			// create "empty" migration
-			migr, err = NewMigration(nil, "", version, targetVersion)
-			if err != nil {
-				return nil, err
-			}
+		r, identifier, err = m.sourceDrv.ReadUp(version)
+	} else {
+		r, identifier, err = m.sourceDrv.ReadDown(version)
+	}
 
-		} else if err != nil {
+	if errors.Is(err, os.ErrNotExist) {
+		// create "empty" migration
+		migr, err = NewMigration(nil, "", version, targetVersion)
+		if err != nil {
 			return nil, err
-
-		} else {
-			// create migration from up source
-			migr, err = NewMigration(r, identifier, version, targetVersion)
-			if err != nil {
-				return nil, err
-			}
 		}
 
+	} else if err != nil {
+		return nil, err
+
 	} else {
-		r, identifier, err := m.sourceDrv.ReadDown(version)
-		if errors.Is(err, os.ErrNotExist) {
-			// create "empty" migration
-			migr, err = NewMigration(nil, "", version, targetVersion)
-			if err != nil {
+		if m.EnableTemplating {
+			if r, err = applyEnvironmentTemplate(r, m.Log); err != nil {
 				return nil, err
 			}
-
-		} else if err != nil {
+		}
+		// create migration from source
+		migr, err = NewMigration(r, identifier, version, targetVersion)
+		if err != nil {
 			return nil, err
-
-		} else {
-			// create migration from down source
-			migr, err = NewMigration(r, identifier, version, targetVersion)
-			if err != nil {
-				return nil, err
-			}
 		}
 	}
 
@@ -880,7 +883,6 @@ func (m *Migrate) newMigration(version uint, targetVersion int) (*Migration, err
 	} else {
 		m.logVerbosePrintf("Scheduled %v\n", migr.LogString())
 	}
-
 	return migr, nil
 }
 

template.go

@@ -0,0 +1,69 @@
+package migrate
+
+import (
+	"fmt"
+	"io"
+	"os"
+	"strings"
+	"sync"
+	"text/template"
+)
+
+var envMapCache map[string]string
+var envMapCacheLock sync.Mutex
+
+func envMap() map[string]string {
+	// get the lock before accessing envMap to prevent concurrent reads and writes
+	envMapCacheLock.Lock()
+	defer envMapCacheLock.Unlock()
+	if envMapCache != nil {
+		return envMapCache
+	}
+	envMapCache = make(map[string]string)
+	for _, kvp := range os.Environ() {
+		kvParts := strings.SplitN(kvp, "=", 2)
+		envMapCache[kvParts[0]] = kvParts[1]
+	}
+	return envMapCache
+}
+
+func applyEnvironmentTemplate(body io.ReadCloser, logger Logger) (io.ReadCloser, error) {
+	bodyBytes, err := io.ReadAll(body)
+	if err != nil {
+		return nil, fmt.Errorf("reading body: %w", err)
+	}
+	defer func() {
+		err = body.Close()
+		if err != nil {
+			logger.Printf("applyEnvironmentTemplate: error closing body: %v", err)
+		}
+	}()
+
+	tmpl, err := template.New("migration").Parse(string(bodyBytes))
+	if err != nil {
+		return nil, fmt.Errorf("parsing template: %w", err)
+	}
+
+	r, w := io.Pipe()
+
+	go func() {
+		em := envMap()
+		err = tmpl.Execute(w, em)
+		if err != nil {
+			if logger != nil {
+				logger.Printf("applyEnvironmentTemplate: error executing template: %v", err)
+				if logger.Verbose() {
+					logger.Printf("applyEnvironmentTemplate: env map used for template execution: %v", em)
+				}
+			}
+		}
+		err = w.Close()
+		if err != nil {
+			if logger != nil {
+				logger.Printf("applyEnvironmentTemplate: error closing writer: %v", err)
+			}
+		}
+	}()
+
+	return r, nil
+}

template_test.go

@@ -0,0 +1,30 @@
+package migrate
+
+import (
+	"bytes"
+	"io"
+	"os"
+	"testing"
+)
+
+func Test_applyEnvironmentTemplate(t *testing.T) {
+	envMapCache = nil
+	_ = os.Setenv("WAREHOUSE_DB", "WH_STAGING")
+
+	migration := io.NopCloser(bytes.NewBuffer([]byte(`SELECT * FROM {{.WAREHOUSE_DB}}.STD.INVOICES`)))
+
+	gotReader, err := applyEnvironmentTemplate(migration, nil)
+	if err != nil {
+		t.Fatalf("expected no error applying template")
+	}
+
+	gotBytes, err := io.ReadAll(gotReader)
+	if err != nil {
+		t.Fatalf("expected no error reading")
+	}
+	got := string(gotBytes)
+	want := `SELECT * FROM WH_STAGING.STD.INVOICES`
+	if got != want {
+		t.Fatalf("expected [%s] but got [%s]", want, got)
+	}
+}