modify: Add statement timeout query parmeter for MySQL (golang-migrat…

…e#778) * modify: Add statement timeout query parmeter for MySQL * fix: parsing and assignment of statementTimeout * fix: bad merge
nicheinc · Jul 27, 2022 · 3bba76f · 3bba76f
1 parent 0d585e2
commit 3bba76f
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 7 deletions.
diff --git a/database/mysql/README.md b/database/mysql/README.md
@@ -6,6 +6,7 @@
 |------------|---------------------|-------------|
 | `x-migrations-table` | `MigrationsTable` | Name of the migrations table |
 | `x-no-lock` | `NoLock` | Set to `true` to skip `GET_LOCK`/`RELEASE_LOCK` statements. Useful for [multi-master MySQL flavors](https://www.percona.com/doc/percona-xtradb-cluster/LATEST/features/pxc-strict-mode.html#explicit-table-locking). Only run migrations from one host when this is enabled. |
+| `x-statement-timeout` | `StatementTimeout` | Abort any statement that takes more than the specified number of milliseconds, functionally similar to [Server-side SELECT statement timeouts](https://dev.mysql.com/blog-archive/server-side-select-statement-timeouts/) but enforced by the client. Available for all versions of MySQL, not just >=5.7. | 
 | `dbname` | `DatabaseName` | The name of the database to connect to |
 | `user` | | The user to sign in as |
 | `password` | | The user's password | 

diff --git a/database/mysql/mysql.go b/database/mysql/mysql.go
@@ -14,6 +14,7 @@ import (
 	"os"
 	"strconv"
 	"strings"
+	"time"
 
 	"go.uber.org/atomic"
 
@@ -39,9 +40,10 @@ var (
 )
 
 type Config struct {
-	MigrationsTable string
-	DatabaseName    string
-	NoLock          bool
+	MigrationsTable  string
+	DatabaseName     string
+	NoLock           bool
+	StatementTimeout time.Duration
 }
 
 type Mysql struct {
@@ -242,15 +244,25 @@ func (m *Mysql) Open(url string) (database.Driver, error) {
 		}
 	}
 
+	statementTimeoutParam := customParams["x-statement-timeout"]
+	statementTimeout := 0
+	if statementTimeoutParam != "" {
+		statementTimeout, err = strconv.Atoi(statementTimeoutParam)
+		if err != nil {
+			return nil, fmt.Errorf("could not parse x-statement-timeout as float: %w", err)
+		}
+	}
+
 	db, err := sql.Open("mysql", config.FormatDSN())
 	if err != nil {
 		return nil, err
 	}
 
 	mx, err := WithInstance(db, &Config{
-		DatabaseName:    config.DBName,
-		MigrationsTable: customParams["x-migrations-table"],
-		NoLock:          noLock,
+		DatabaseName:     config.DBName,
+		MigrationsTable:  customParams["x-migrations-table"],
+		NoLock:           noLock,
+		StatementTimeout: time.Duration(statementTimeout) * time.Millisecond,
 	})
 	if err != nil {
 		return nil, err
@@ -328,8 +340,15 @@ func (m *Mysql) Run(migration io.Reader) error {
 		return err
 	}
 
+	ctx := context.Background()
+	if m.config.StatementTimeout != 0 {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(ctx, m.config.StatementTimeout)
+		defer cancel()
+	}
+
 	query := string(migr[:])
-	if _, err := m.conn.ExecContext(context.Background(), query); err != nil {
+	if _, err := m.conn.ExecContext(ctx, query); err != nil {
 		return database.Error{OrigErr: err, Err: "migration failed", Query: migr}
 	}