Show completion estimate during backfill

cmd/start.go

@@ -55,8 +55,20 @@ func runMigrationFromFile(ctx context.Context, m *roll.Roll, fileName string, co
 
 func runMigration(ctx context.Context, m *roll.Roll, migration *migrations.Migration, complete bool) error {
 	sp, _ := pterm.DefaultSpinner.WithText("Starting migration...").Start()
-	cb := func(n int64) {
-		sp.UpdateText(fmt.Sprintf("%d records complete...", n))
+	cb := func(n int64, total int64) {
+		var percent float64
+		if total > 0 {
+			percent = float64(n) / float64(total) * 100
+		}
+		if percent > 100 {
+			// This can happen if we're on the last batch
+			percent = 100
+		}
+		if total > 0 {
+			sp.UpdateText(fmt.Sprintf("%d records complete... (%.2f%%)", n, percent))
+		} else {
+			sp.UpdateText(fmt.Sprintf("%d records complete...", n))
+		}
 	}
 
 	err := m.Start(ctx, migration, cb)

pkg/db/db.go

@@ -20,6 +20,7 @@ const (
 
 type DB interface {
 	ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error)
+	QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row
 	WithRetryableTransaction(ctx context.Context, f func(context.Context, *sql.Tx) error) error
 	Close() error
 }
@@ -52,6 +53,11 @@ func (db *RDB) ExecContext(ctx context.Context, query string, args ...interface{
 	}
 }
 
+// QueryRowContext wraps sql.DB.QueryRowContext.
+func (db *RDB) QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row {
+	return db.DB.QueryRowContext(ctx, query, args...)
+}
+
 // WithRetryableTransaction runs `f` in a transaction, retrying on lock_timeout errors.
 func (db *RDB) WithRetryableTransaction(ctx context.Context, f func(context.Context, *sql.Tx) error) error {
 	b := backoff.New(maxBackoffDuration, backoffInterval)

pkg/migrations/backfill.go

@@ -29,13 +29,18 @@ func Backfill(ctx context.Context, conn db.DB, table *schema.Table, batchSize in
 		return BackfillNotPossibleError{Table: table.Name}
 	}
 
+	total, err := getRowCount(ctx, conn, table.Name)
+	if err != nil {
+		return fmt.Errorf("get row count for %q: %w", table.Name, err)
+	}
+
 	// Create a batcher for the table.
 	b := newBatcher(table, batchSize)
 
 	// Update each batch of rows, invoking callbacks for each one.
 	for batch := 0; ; batch++ {
 		for _, cb := range cbs {
-			cb(int64(batch * batchSize))
+			cb(int64(batch*batchSize), total)
 		}
 
 		if err := b.updateBatch(ctx, conn); err != nil {
@@ -55,6 +60,35 @@ func Backfill(ctx context.Context, conn db.DB, table *schema.Table, batchSize in
 	return nil
 }
 
+// getRowCount will attempt to get the row count for the given table. It first attempts to get an
+// estimate and if that is zero, falls back to a full table scan.
+func getRowCount(ctx context.Context, conn db.DB, tableName string) (int64, error) {
+	var total int64
+	// Try and get estimated row count
+	var currentSchema string
+	row := conn.QueryRowContext(ctx, "select current_schema()")
+	if err := row.Scan(&currentSchema); err != nil {
+		return 0, fmt.Errorf("getting current schema: %w", err)
+	}
+	row = conn.QueryRowContext(ctx, `
+	  SELECT n_live_tup AS estimate
+	  FROM pg_stat_user_tables
+	  WHERE schemaname = $1 AND relname = $2`, currentSchema, tableName)
+	if err := row.Scan(&total); err != nil {
+		return 0, fmt.Errorf("scanning row count estimate for %q: %w", tableName, err)
+	}
+	if total > 0 {
+		return total, nil
+	}
+
+	// If the estimate is zero, fall back to full count
+	row = conn.QueryRowContext(ctx, fmt.Sprintf(`SELECT count(*) from %s`, tableName))
+	if err := row.Scan(&total); err != nil {
+		return 0, fmt.Errorf("scanning row count for %q: %w", tableName, err)
+	}
+	return total, nil
+}
+
 // checkBackfill will return an error if the backfill operation is not supported.
 func checkBackfill(table *schema.Table) error {
 	cols := getIdentityColumns(table)

pkg/migrations/migrations.go

@@ -12,7 +12,7 @@ import (
 	"github.com/xataio/pgroll/pkg/schema"
 )
 
-type CallbackFn func(int64)
+type CallbackFn func(done int64, total int64)
 
 // Operation is an operation that can be applied to a schema
 type Operation interface {

pkg/roll/execute_test.go

@@ -655,7 +655,7 @@ func TestCallbacksAreInvokedOnMigrationStart(t *testing.T) {
 
 		// Define a mock callback
 		invoked := false
-		cb := func(n int64) { invoked = true }
+		cb := func(n, total int64) { invoked = true }
 
 		// Start a migration that requires a backfill
 		err = mig.Start(ctx, &migrations.Migration{