Allow configuration of backfill batch size (#406)

It can be controlled via the `--backfill-batch-size` command line
parameter or by setting the `PGROLL_BACKFILL_BATCH_SIZE` environment
variable.

It can also be set programatically via the `roll.WithBackfillBatchSize`
function.

If unset, it will default to 1000.

Part of #168

Author: ryanslade

cmd/flags/flags.go

@@ -22,6 +22,8 @@ func LockTimeout() int {
 	return viper.GetInt("LOCK_TIMEOUT")
 }
 
+func BackfillBatchSize() int { return viper.GetInt("BACKFILL_BATCH_SIZE") }
+
 func Role() string {
 	return viper.GetString("ROLE")
 }

cmd/root.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
+
 	"github.com/xataio/pgroll/cmd/flags"
 	"github.com/xataio/pgroll/pkg/roll"
 	"github.com/xataio/pgroll/pkg/state"
@@ -23,12 +24,14 @@ func init() {
 	rootCmd.PersistentFlags().String("schema", "public", "Postgres schema to use for the migration")
 	rootCmd.PersistentFlags().String("pgroll-schema", "pgroll", "Postgres schema to use for pgroll internal state")
 	rootCmd.PersistentFlags().Int("lock-timeout", 500, "Postgres lock timeout in milliseconds for pgroll DDL operations")
+	rootCmd.PersistentFlags().Int("backfill-batch-size", roll.DefaultBackfillBatchSize, "Number of rows backfilled in each batch")
 	rootCmd.PersistentFlags().String("role", "", "Optional postgres role to set when executing migrations")
 
 	viper.BindPFlag("PG_URL", rootCmd.PersistentFlags().Lookup("postgres-url"))
 	viper.BindPFlag("SCHEMA", rootCmd.PersistentFlags().Lookup("schema"))
 	viper.BindPFlag("STATE_SCHEMA", rootCmd.PersistentFlags().Lookup("pgroll-schema"))
 	viper.BindPFlag("LOCK_TIMEOUT", rootCmd.PersistentFlags().Lookup("lock-timeout"))
+	viper.BindPFlag("BACKFILL_BATCH_SIZE", rootCmd.PersistentFlags().Lookup("backfill-batch-size"))
 	viper.BindPFlag("ROLE", rootCmd.PersistentFlags().Lookup("role"))
 }
 
@@ -44,6 +47,7 @@ func NewRoll(ctx context.Context) (*roll.Roll, error) {
 	stateSchema := flags.StateSchema()
 	lockTimeout := flags.LockTimeout()
 	role := flags.Role()
+	backfillBatchSize := flags.BackfillBatchSize()
 
 	state, err := state.New(ctx, pgURL, stateSchema)
 	if err != nil {
@@ -53,6 +57,7 @@ func NewRoll(ctx context.Context) (*roll.Roll, error) {
 	return roll.New(ctx, pgURL, schema, state,
 		roll.WithLockTimeoutMs(lockTimeout),
 		roll.WithRole(role),
+		roll.WithBackfillBatchSize(backfillBatchSize),
 	)
 }
 

pkg/migrations/backfill.go

@@ -19,7 +19,7 @@ import (
 // 2. Get the first batch of rows from the table, ordered by the primary key.
 // 3. Update each row in the batch, setting the value of the primary key column to itself.
 // 4. Repeat steps 2 and 3 until no more rows are returned.
-func Backfill(ctx context.Context, conn db.DB, table *schema.Table, cbs ...CallbackFn) error {
+func Backfill(ctx context.Context, conn db.DB, table *schema.Table, batchSize int, cbs ...CallbackFn) error {
 	// get the backfill column
 	identityColumn := getIdentityColumn(table)
 	if identityColumn == nil {
@@ -31,7 +31,7 @@ func Backfill(ctx context.Context, conn db.DB, table *schema.Table, cbs ...Callb
 		table:          table,
 		identityColumn: identityColumn,
 		lastValue:      nil,
-		batchSize:      1000,
+		batchSize:      batchSize,
 	}
 
 	// Update each batch of rows, invoking callbacks for each one.

pkg/roll/execute.go

@@ -277,7 +277,7 @@ func (m *Roll) ensureView(ctx context.Context, version, name string, table schem
 
 func (m *Roll) performBackfills(ctx context.Context, tables []*schema.Table, cbs ...migrations.CallbackFn) error {
 	for _, table := range tables {
-		if err := migrations.Backfill(ctx, m.pgConn, table, cbs...); err != nil {
+		if err := migrations.Backfill(ctx, m.pgConn, table, m.backfillBatchSize, cbs...); err != nil {
 			errRollback := m.Rollback(ctx)
 
 			return errors.Join(

pkg/roll/options.go

@@ -25,6 +25,9 @@ type options struct {
 	// additional entries to add to the search_path during migration execution
 	searchPath []string
 
+	// the number of rows to backfill in each batch
+	backfillBatchSize int
+
 	migrationHooks MigrationHooks
 }
 
@@ -99,3 +102,10 @@ func WithSearchPath(schemas ...string) Option {
 		o.searchPath = schemas
 	}
 }
+
+// WithBackfillBatchSize sets the number of rows backfilled in each batch.
+func WithBackfillBatchSize(batchSize int) Option {
+	return func(o *options) {
+		o.backfillBatchSize = batchSize
+	}
+}

pkg/roll/roll.go

@@ -17,7 +17,10 @@ import (
 
 type PGVersion int
 
-const PGVersion15 PGVersion = 15
+const (
+	PGVersion15              PGVersion = 15
+	DefaultBackfillBatchSize int       = 1000
+)
 
 type Roll struct {
 	pgConn db.DB
@@ -31,10 +34,11 @@ type Roll struct {
 	// disable creation of version schema for raw SQL migrations
 	noVersionSchemaForRawSQL bool
 
-	migrationHooks MigrationHooks
-	state          *state.State
-	pgVersion      PGVersion
-	sqlTransformer migrations.SQLTransformer
+	migrationHooks    MigrationHooks
+	state             *state.State
+	pgVersion         PGVersion
+	sqlTransformer    migrations.SQLTransformer
+	backfillBatchSize int
 }
 
 // New creates a new Roll instance
@@ -43,6 +47,9 @@ func New(ctx context.Context, pgURL, schema string, state *state.State, opts ...
 	for _, o := range opts {
 		o(rollOpts)
 	}
+	if rollOpts.backfillBatchSize <= 0 {
+		rollOpts.backfillBatchSize = DefaultBackfillBatchSize
+	}
 
 	conn, err := setupConn(ctx, pgURL, schema, *rollOpts)
 	if err != nil {
@@ -71,6 +78,7 @@ func New(ctx context.Context, pgURL, schema string, state *state.State, opts ...
 		noVersionSchemaForRawSQL: rollOpts.noVersionSchemaForRawSQL,
 		migrationHooks:           rollOpts.migrationHooks,
 		sqlTransformer:           sqlTransformer,
+		backfillBatchSize:        rollOpts.backfillBatchSize,
 	}, nil
 }