From 1c8d7ac5cc00c5acee6f512934c000d21a8b084f Mon Sep 17 00:00:00 2001 From: Saraj Munjal Date: Wed, 31 May 2023 12:23:55 -0700 Subject: [PATCH] Add context support to Exec methods --- bindata_test.go | 12 ++++++----- doc.go | 31 +++++++++++++-------------- migrate.go | 57 ++++++++++++++++++++++++++++++++++++++++++------- migrate_test.go | 38 +++++++++++++++++++++++++++++++++ 4 files changed, 109 insertions(+), 29 deletions(-) diff --git a/bindata_test.go b/bindata_test.go index e8afed6f..d120d3fa 100644 --- a/bindata_test.go +++ b/bindata_test.go @@ -94,11 +94,13 @@ var _bindata = map[string]func() ([]byte, error){ // directory embedded in the file by go-bindata. // For example if you run go-bindata on data/... and data contains the // following hierarchy: -// data/ -// foo.txt -// img/ -// a.png -// b.png +// +// data/ +// foo.txt +// img/ +// a.png +// b.png +// // then AssetDir("data") would return []string{"foo.txt", "img"} // AssetDir("data/img") would return []string{"a.png", "b.png"} // AssetDir("foo.txt") and AssetDir("notexist") would return an error diff --git a/doc.go b/doc.go index eb4ed857..8ff186d0 100644 --- a/doc.go +++ b/doc.go @@ -1,24 +1,23 @@ /* - SQL Schema migration tool for Go. Key features: - * Usable as a CLI tool or as a library - * Supports SQLite, PostgreSQL, MySQL, MSSQL and Oracle databases (through gorp) - * Can embed migrations into your application - * Migrations are defined with SQL for full flexibility - * Atomic migrations - * Up/down migrations to allow rollback - * Supports multiple database types in one project + - Usable as a CLI tool or as a library + - Supports SQLite, PostgreSQL, MySQL, MSSQL and Oracle databases (through gorp) + - Can embed migrations into your application + - Migrations are defined with SQL for full flexibility + - Atomic migrations + - Up/down migrations to allow rollback + - Supports multiple database types in one project -Installation +# Installation To install the library and command line program, use the following: go get -v github.com/rubenv/sql-migrate/... -Command-line tool +# Command-line tool The main command is called sql-migrate. @@ -77,7 +76,7 @@ Use the status command to see the state of the applied migrations: | 2_record.sql | no | +---------------+-----------------------------------------+ -MySQL Caveat +# MySQL Caveat If you are using MySQL, you must append ?parseTime=true to the datasource configuration. For example: @@ -89,7 +88,7 @@ If you are using MySQL, you must append ?parseTime=true to the datasource config See https://github.com/go-sql-driver/mysql#parsetime for more information. -Library +# Library Import sql-migrate into your application: @@ -137,7 +136,7 @@ Note that n can be greater than 0 even if there is an error: any migration that The full set of capabilities can be found in the API docs below. -Writing migrations +# Writing migrations Migrations are defined in SQL files, which contain a set of SQL statements. Special comments are used to distinguish up and down migrations. @@ -183,7 +182,7 @@ Normally each migration is run within a transaction in order to guarantee that i -- +migrate Down DROP INDEX people_unique_id_idx; -Embedding migrations with packr +# Embedding migrations with packr If you like your Go applications self-contained (that is: a single binary): use packr (https://github.com/gobuffalo/packr) to embed the migration files. @@ -202,7 +201,7 @@ If you already have a box and would like to use a subdirectory: Dir: "./migrations", } -Embedding migrations with bindata +# Embedding migrations with bindata As an alternative, but slightly less maintained, you can use bindata (https://github.com/shuLhan/go-bindata) to embed the migration files. @@ -226,7 +225,7 @@ Both Asset and AssetDir are functions provided by bindata. Then proceed as usual. -Extending +# Extending Adding a new migration source means implementing MigrationSource. diff --git a/migrate.go b/migrate.go index 5ff70325..acb5a412 100644 --- a/migrate.go +++ b/migrate.go @@ -2,6 +2,7 @@ package migrate import ( "bytes" + "context" "database/sql" "errors" "fmt" @@ -429,12 +430,24 @@ type SqlExecutor interface { // // Returns the number of applied migrations. func Exec(db *sql.DB, dialect string, m MigrationSource, dir MigrationDirection) (int, error) { - return ExecMax(db, dialect, m, dir, 0) + return ExecMaxContext(context.Background(), db, dialect, m, dir, 0) } // Returns the number of applied migrations. func (ms MigrationSet) Exec(db *sql.DB, dialect string, m MigrationSource, dir MigrationDirection) (int, error) { - return ms.ExecMax(db, dialect, m, dir, 0) + return ms.ExecMaxContext(context.Background(), db, dialect, m, dir, 0) +} + +// Execute a set of migrations with an input context. +// +// Returns the number of applied migrations. +func ExecContext(ctx context.Context, db *sql.DB, dialect string, m MigrationSource, dir MigrationDirection) (int, error) { + return ExecMaxContext(ctx, db, dialect, m, dir, 0) +} + +// Returns the number of applied migrations. +func (ms MigrationSet) ExecContext(ctx context.Context, db *sql.DB, dialect string, m MigrationSource, dir MigrationDirection) (int, error) { + return ms.ExecMaxContext(ctx, db, dialect, m, dir, 0) } // Execute a set of migrations @@ -446,50 +459,78 @@ func ExecMax(db *sql.DB, dialect string, m MigrationSource, dir MigrationDirecti return migSet.ExecMax(db, dialect, m, dir, max) } +// Execute a set of migrations with an input context. +// +// Will apply at most `max` migrations. Pass 0 for no limit (or use Exec). +// +// Returns the number of applied migrations. +func ExecMaxContext(ctx context.Context, db *sql.DB, dialect string, m MigrationSource, dir MigrationDirection, max int) (int, error) { + return migSet.ExecMaxContext(ctx, db, dialect, m, dir, max) +} + // Execute a set of migrations // // Will apply at the target `version` of migration. Cannot be a negative value. // // Returns the number of applied migrations. func ExecVersion(db *sql.DB, dialect string, m MigrationSource, dir MigrationDirection, version int64) (int, error) { + return ExecVersionContext(context.Background(), db, dialect, m, dir, version) +} + +// Execute a set of migrations with an input context. +// +// Will apply at the target `version` of migration. Cannot be a negative value. +// +// Returns the number of applied migrations. +func ExecVersionContext(ctx context.Context, db *sql.DB, dialect string, m MigrationSource, dir MigrationDirection, version int64) (int, error) { if version < 0 { return 0, fmt.Errorf("target version %d should not be negative", version) } - return migSet.ExecVersion(db, dialect, m, dir, version) + return migSet.ExecVersionContext(ctx, db, dialect, m, dir, version) } // Returns the number of applied migrations. func (ms MigrationSet) ExecMax(db *sql.DB, dialect string, m MigrationSource, dir MigrationDirection, max int) (int, error) { + return ms.ExecMaxContext(context.Background(), db, dialect, m, dir, max) +} + +// Returns the number of applied migrations, but applies with an input context. +func (ms MigrationSet) ExecMaxContext(ctx context.Context, db *sql.DB, dialect string, m MigrationSource, dir MigrationDirection, max int) (int, error) { migrations, dbMap, err := ms.PlanMigration(db, dialect, m, dir, max) if err != nil { return 0, err } - return ms.applyMigrations(dir, migrations, dbMap) + return ms.applyMigrations(ctx, dir, migrations, dbMap) } // Returns the number of applied migrations. func (ms MigrationSet) ExecVersion(db *sql.DB, dialect string, m MigrationSource, dir MigrationDirection, version int64) (int, error) { + return ms.ExecVersionContext(context.Background(), db, dialect, m, dir, version) +} + +func (ms MigrationSet) ExecVersionContext(ctx context.Context, db *sql.DB, dialect string, m MigrationSource, dir MigrationDirection, version int64) (int, error) { migrations, dbMap, err := ms.PlanMigrationToVersion(db, dialect, m, dir, version) if err != nil { return 0, err } - return ms.applyMigrations(dir, migrations, dbMap) + return ms.applyMigrations(ctx, dir, migrations, dbMap) } // Applies the planned migrations and returns the number of applied migrations. -func (ms MigrationSet) applyMigrations(dir MigrationDirection, migrations []*PlannedMigration, dbMap *gorp.DbMap) (int, error) { +func (ms MigrationSet) applyMigrations(ctx context.Context, dir MigrationDirection, migrations []*PlannedMigration, dbMap *gorp.DbMap) (int, error) { applied := 0 for _, migration := range migrations { var executor SqlExecutor var err error if migration.DisableTransaction { - executor = dbMap + executor = dbMap.WithContext(ctx) } else { - executor, err = dbMap.Begin() + e, err := dbMap.Begin() if err != nil { return applied, newTxError(migration, err) } + executor = e.WithContext(ctx) } for _, stmt := range migration.Queries { diff --git a/migrate_test.go b/migrate_test.go index 89073a35..0730c61b 100644 --- a/migrate_test.go +++ b/migrate_test.go @@ -1,8 +1,10 @@ package migrate import ( + "context" "database/sql" "net/http" + "time" "github.com/go-gorp/gorp/v3" "github.com/gobuffalo/packr/v2" @@ -758,3 +760,39 @@ func (s *SqliteMigrateSuite) TestGetMigrationDbMapWithDisableCreateTable(c *C) { _, err := migSet.getMigrationDbMap(s.Db, "postgres") c.Assert(err, IsNil) } + +func (s *SqliteMigrateSuite) TestContextTimeout(c *C) { + // This statement will run for a long time: 1,000,000 iterations of the fibonacci sequence + fibonacciLoopStmt := `WITH RECURSIVE + fibo (curr, next) + AS + ( SELECT 1,1 + UNION ALL + SELECT next, curr+next FROM fibo + LIMIT 1000000 ) + SELECT group_concat(curr) FROM fibo; + ` + migrations := &MemoryMigrationSource{ + Migrations: []*Migration{ + sqliteMigrations[0], + sqliteMigrations[1], + { + Id: "125", + Up: []string{fibonacciLoopStmt}, + Down: []string{}, // Not important here + }, + { + Id: "125", + Up: []string{"INSERT INTO people (id, first_name) VALUES (1, 'Test')", "SELECT fail"}, + Down: []string{}, // Not important here + }, + }, + } + + // Should never run the insert + ctx, cancelFunc := context.WithTimeout(context.Background(), 10*time.Millisecond) + defer cancelFunc() + n, err := ExecContext(ctx, s.Db, "sqlite3", migrations, Up) + c.Assert(err, Not(IsNil)) + c.Assert(n, Equals, 2) +}