Respect context deadline during batch backfill (#438)

xataio · Oct 28, 2024 · ddd44bb · ddd44bb
1 parent 88cc975
commit ddd44bb
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/pkg/migrations/backfill.go b/pkg/migrations/backfill.go
@@ -10,6 +10,7 @@ import (
 	"time"
 
 	"github.com/lib/pq"
+
 	"github.com/xataio/pgroll/pkg/db"
 	"github.com/xataio/pgroll/pkg/schema"
 )
@@ -33,7 +34,6 @@ func Backfill(ctx context.Context, conn db.DB, table *schema.Table, batchSize in
 		identityColumn: identityColumn,
 		lastValue:      nil,
 		batchSize:      batchSize,
-		batchDelay:     batchDelay,
 	}
 
 	// Update each batch of rows, invoking callbacks for each one.
@@ -49,7 +49,11 @@ func Backfill(ctx context.Context, conn db.DB, table *schema.Table, batchSize in
 			return err
 		}
 
-		time.Sleep(b.batchDelay)
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-time.After(batchDelay):
+		}
 	}
 
 	return nil
@@ -88,7 +92,6 @@ type batcher struct {
 	identityColumn *schema.Column
 	lastValue      *string
 	batchSize      int
-	batchDelay     time.Duration
 }
 
 func (b *batcher) updateBatch(ctx context.Context, conn db.DB) error {