Add context support to Exec methods

Author: sarajmunjal

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

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.
 

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 {

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)
+}