From bbcbcb621bb3acf1f36824c1c55b94695f0c28d6 Mon Sep 17 00:00:00 2001 From: Diogo Monteiro Date: Thu, 21 Nov 2019 15:16:48 -0500 Subject: [PATCH 1/2] Allowing a context to be set when running a goose command. --- down.go | 9 +++++---- goose.go | 20 +++++++++++++------- migration.go | 13 +++++++------ migration_sql.go | 5 +++-- redo.go | 7 ++++--- reset.go | 5 +++-- up.go | 13 +++++++------ 7 files changed, 42 insertions(+), 30 deletions(-) diff --git a/down.go b/down.go index 1db984736..ae76fc6d1 100644 --- a/down.go +++ b/down.go @@ -1,12 +1,13 @@ package goose import ( + "context" "database/sql" "fmt" ) // Down rolls back a single migration from the current version. -func Down(db *sql.DB, dir string) error { +func Down(ctx context.Context, db *sql.DB, dir string) error { currentVersion, err := GetDBVersion(db) if err != nil { return err @@ -22,11 +23,11 @@ func Down(db *sql.DB, dir string) error { return fmt.Errorf("no migration %v", currentVersion) } - return current.Down(db) + return current.Down(ctx, db) } // DownTo rolls back migrations to a specific version. -func DownTo(db *sql.DB, dir string, version int64) error { +func DownTo(ctx context.Context, db *sql.DB, dir string, version int64) error { migrations, err := CollectMigrations(dir, minVersion, maxVersion) if err != nil { return err @@ -49,7 +50,7 @@ func DownTo(db *sql.DB, dir string, version int64) error { return nil } - if err = current.Down(db); err != nil { + if err = current.Down(ctx, db); err != nil { return err } } diff --git a/goose.go b/goose.go index 8ce57d2a3..b7e49b6a4 100644 --- a/goose.go +++ b/goose.go @@ -1,6 +1,7 @@ package goose import ( + "context" "database/sql" "fmt" "strconv" @@ -24,13 +25,18 @@ func SetVerbose(v bool) { // Run runs a goose command. func Run(command string, db *sql.DB, dir string, args ...string) error { + return RunWithCtx(context.Background(), command, db, dir, args...) +} + +// Run runs a goose command propagating a context down the call stack. +func RunWithCtx(ctx context.Context, command string, db *sql.DB, dir string, args ...string) error { switch command { case "up": - if err := Up(db, dir); err != nil { + if err := Up(ctx, db, dir); err != nil { return err } case "up-by-one": - if err := UpByOne(db, dir); err != nil { + if err := UpByOne(ctx, db, dir); err != nil { return err } case "up-to": @@ -42,7 +48,7 @@ func Run(command string, db *sql.DB, dir string, args ...string) error { if err != nil { return fmt.Errorf("version must be a number (got '%s')", args[0]) } - if err := UpTo(db, dir, version); err != nil { + if err := UpTo(ctx, db, dir, version); err != nil { return err } case "create": @@ -58,7 +64,7 @@ func Run(command string, db *sql.DB, dir string, args ...string) error { return err } case "down": - if err := Down(db, dir); err != nil { + if err := Down(ctx, db, dir); err != nil { return err } case "down-to": @@ -70,7 +76,7 @@ func Run(command string, db *sql.DB, dir string, args ...string) error { if err != nil { return fmt.Errorf("version must be a number (got '%s')", args[0]) } - if err := DownTo(db, dir, version); err != nil { + if err := DownTo(ctx, db, dir, version); err != nil { return err } case "fix": @@ -78,11 +84,11 @@ func Run(command string, db *sql.DB, dir string, args ...string) error { return err } case "redo": - if err := Redo(db, dir); err != nil { + if err := Redo(ctx, db, dir); err != nil { return err } case "reset": - if err := Reset(db, dir); err != nil { + if err := Reset(ctx, db, dir); err != nil { return err } case "status": diff --git a/migration.go b/migration.go index 68e579f90..0c69907e7 100644 --- a/migration.go +++ b/migration.go @@ -1,6 +1,7 @@ package goose import ( + "context" "database/sql" "fmt" "os" @@ -35,22 +36,22 @@ func (m *Migration) String() string { } // Up runs an up migration. -func (m *Migration) Up(db *sql.DB) error { - if err := m.run(db, true); err != nil { +func (m *Migration) Up(ctx context.Context, db *sql.DB) error { + if err := m.run(ctx, db, true); err != nil { return err } return nil } // Down runs a down migration. -func (m *Migration) Down(db *sql.DB) error { - if err := m.run(db, false); err != nil { +func (m *Migration) Down(ctx context.Context, db *sql.DB) error { + if err := m.run(ctx, db, false); err != nil { return err } return nil } -func (m *Migration) run(db *sql.DB, direction bool) error { +func (m *Migration) run(ctx context.Context, db *sql.DB, direction bool) error { switch filepath.Ext(m.Source) { case ".sql": f, err := os.Open(m.Source) @@ -64,7 +65,7 @@ func (m *Migration) run(db *sql.DB, direction bool) error { return errors.Wrapf(err, "ERROR %v: failed to parse SQL migration file", filepath.Base(m.Source)) } - if err := runSQLMigration(db, statements, useTx, m.Version, direction); err != nil { + if err := runSQLMigration(ctx, db, statements, useTx, m.Version, direction); err != nil { return errors.Wrapf(err, "ERROR %v: failed to run SQL migration", filepath.Base(m.Source)) } diff --git a/migration_sql.go b/migration_sql.go index a6a35994d..6324cea8b 100644 --- a/migration_sql.go +++ b/migration_sql.go @@ -1,6 +1,7 @@ package goose import ( + "context" "database/sql" "regexp" @@ -15,13 +16,13 @@ import ( // // All statements following an Up or Down directive are grouped together // until another direction directive is found. -func runSQLMigration(db *sql.DB, statements []string, useTx bool, v int64, direction bool) error { +func runSQLMigration(ctx context.Context, db *sql.DB, statements []string, useTx bool, v int64, direction bool) error { if useTx { // TRANSACTION. verboseInfo("Begin transaction") - tx, err := db.Begin() + tx, err := db.BeginTx(ctx, nil) if err != nil { return errors.Wrap(err, "failed to begin transaction") } diff --git a/redo.go b/redo.go index 6f9049f67..8c1f1663a 100644 --- a/redo.go +++ b/redo.go @@ -1,11 +1,12 @@ package goose import ( + "context" "database/sql" ) // Redo rolls back the most recently applied migration, then runs it again. -func Redo(db *sql.DB, dir string) error { +func Redo(ctx context.Context, db *sql.DB, dir string) error { currentVersion, err := GetDBVersion(db) if err != nil { return err @@ -21,11 +22,11 @@ func Redo(db *sql.DB, dir string) error { return err } - if err := current.Down(db); err != nil { + if err := current.Down(ctx, db); err != nil { return err } - if err := current.Up(db); err != nil { + if err := current.Up(ctx, db); err != nil { return err } diff --git a/reset.go b/reset.go index 26b9dcb7f..5bed959fa 100644 --- a/reset.go +++ b/reset.go @@ -1,6 +1,7 @@ package goose import ( + "context" "database/sql" "sort" @@ -8,7 +9,7 @@ import ( ) // Reset rolls back all migrations -func Reset(db *sql.DB, dir string) error { +func Reset(ctx context.Context, db *sql.DB, dir string) error { migrations, err := CollectMigrations(dir, minVersion, maxVersion) if err != nil { return errors.Wrap(err, "failed to collect migrations") @@ -23,7 +24,7 @@ func Reset(db *sql.DB, dir string) error { if !statuses[migration.Version] { continue } - if err = migration.Down(db); err != nil { + if err = migration.Down(ctx, db); err != nil { return errors.Wrap(err, "failed to db-down") } } diff --git a/up.go b/up.go index c9ed1110f..4b4908c8a 100644 --- a/up.go +++ b/up.go @@ -1,11 +1,12 @@ package goose import ( + "context" "database/sql" ) // UpTo migrates up to a specific version. -func UpTo(db *sql.DB, dir string, version int64) error { +func UpTo(ctx context.Context, db *sql.DB, dir string, version int64) error { migrations, err := CollectMigrations(dir, minVersion, version) if err != nil { return err @@ -26,19 +27,19 @@ func UpTo(db *sql.DB, dir string, version int64) error { return err } - if err = next.Up(db); err != nil { + if err = next.Up(ctx, db); err != nil { return err } } } // Up applies all available migrations. -func Up(db *sql.DB, dir string) error { - return UpTo(db, dir, maxVersion) +func Up(ctx context.Context, db *sql.DB, dir string) error { + return UpTo(ctx, db, dir, maxVersion) } // UpByOne migrates up by a single version. -func UpByOne(db *sql.DB, dir string) error { +func UpByOne(ctx context.Context, db *sql.DB, dir string) error { migrations, err := CollectMigrations(dir, minVersion, maxVersion) if err != nil { return err @@ -57,7 +58,7 @@ func UpByOne(db *sql.DB, dir string) error { return err } - if err = next.Up(db); err != nil { + if err = next.Up(ctx, db); err != nil { return err } From bdab3e86e23bfcaf07225cb371d883a9cdbae9ad Mon Sep 17 00:00:00 2001 From: Diogo Monteiro Date: Thu, 21 Nov 2019 16:27:02 -0500 Subject: [PATCH 2/2] Setting context on go migrations. --- migration.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/migration.go b/migration.go index 0c69907e7..bc01e97be 100644 --- a/migration.go +++ b/migration.go @@ -79,7 +79,7 @@ func (m *Migration) run(ctx context.Context, db *sql.DB, direction bool) error { if !m.Registered { return errors.Errorf("ERROR %v: failed to run Go migration: Go functions must be registered and built into a custom binary (see https://github.com/pressly/goose/tree/master/examples/go-migrations)", m.Source) } - tx, err := db.Begin() + tx, err := db.BeginTx(ctx, nil) if err != nil { return errors.Wrap(err, "ERROR failed to begin transaction") }