From ddd44bb0b236707a671685e34b7762a28ffa5fe1 Mon Sep 17 00:00:00 2001
From: Ryan Slade <ryanslade@gmail.com>
Date: Mon, 28 Oct 2024 13:15:11 +0100
Subject: [PATCH] Respect context deadline during batch backfill (#438)

---
 pkg/migrations/backfill.go | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/pkg/migrations/backfill.go b/pkg/migrations/backfill.go
index 1e0e5120..dc346405 100644
--- 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 {