From 021cff9c21d6319550896c181fe61c8d7a517826 Mon Sep 17 00:00:00 2001
From: Oleg Kovalov <oleg@hey.com>
Date: Fri, 1 Jul 2022 22:22:26 +0200
Subject: [PATCH] Add Timeout

---
 dbump.go      | 16 +++++++++++++++-
 dbump_test.go | 29 +++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/dbump.go b/dbump.go
index 7065959..762c0a6 100644
--- a/dbump.go
+++ b/dbump.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"sort"
+	"time"
 )
 
 // ErrMigrationAlreadyLocked is returned only when migration lock is already hold.
@@ -29,6 +30,10 @@ type Config struct {
 	// Set mode explicitly to show how migration should be done.
 	Mode MigratorMode
 
+	// Timeout per migration step. Default is 0 which means no timeout.
+	// Only Migrator.DoStep method will be bounded with this timeout.
+	Timeout time.Duration
+
 	// DisableTx will run every migration not in a transaction.
 	// This completely depends on a specific Migrator implementation
 	// because not every database supports transaction, so this option can be no-op all the time.
@@ -219,7 +224,7 @@ func (m *mig) runMigrationsLocked(ctx context.Context, ms []*Migration) error {
 	for _, step := range m.prepareSteps(curr, target, ms) {
 		m.BeforeStep(ctx, step)
 
-		if err := m.DoStep(ctx, step); err != nil {
+		if err := m.step(ctx, step); err != nil {
 			return err
 		}
 
@@ -228,6 +233,15 @@ func (m *mig) runMigrationsLocked(ctx context.Context, ms []*Migration) error {
 	return nil
 }
 
+func (m *mig) step(ctx context.Context, step Step) error {
+	if m.Timeout != 0 {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(ctx, m.Timeout)
+		defer cancel()
+	}
+	return m.DoStep(ctx, step)
+}
+
 func (m *mig) getCurrAndTargetVersions(ctx context.Context, migrations int) (curr, target int, err error) {
 	curr, err = m.Version(ctx)
 	if err != nil {
diff --git a/dbump_test.go b/dbump_test.go
index 70f89aa..ff3a90c 100644
--- a/dbump_test.go
+++ b/dbump_test.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"reflect"
 	"testing"
+	"time"
 )
 
 func TestRunCheck(t *testing.T) {
@@ -265,6 +266,34 @@ func TestBeforeAfterStep(t *testing.T) {
 	mustEqual(t, mm.log, wantLog)
 }
 
+func TestTimeout(t *testing.T) {
+	wantLog := []string{
+		"lockdb", "init", "getversion",
+		"dostep", "{v:1 q:'SELECT 1;' notx:false}",
+		"unlockdb",
+	}
+
+	mm := &MockMigrator{
+		DoStepFn: func(ctx context.Context, step Step) error {
+			select {
+			case <-time.After(30 * time.Second):
+				return nil
+			case <-ctx.Done():
+				return ctx.Err()
+			}
+		},
+	}
+	cfg := Config{
+		Migrator: mm,
+		Loader:   NewSliceLoader(testdataMigrations),
+		Mode:     ModeUp,
+		Timeout:  20 * time.Millisecond,
+	}
+
+	failIfOk(t, Run(context.Background(), cfg))
+	mustEqual(t, mm.log, wantLog)
+}
+
 func TestLockless(t *testing.T) {
 	wantLog := []string{
 		"init",