Repository: xataio/pgstream (go)

Issue:
Snapshot progress tracking
When a longer snapshot is under process, it would be cool if there is a way to provide progress status to the user. Currently, I am using `trace` log level BUT it is too noisy, specially when the data in rows in large (think `jsonb`). It would be cool if there is a mode that can do user readable progress, like

|Table|Progress|
|---|---|
|User | 4/20 |
|Post | 10/200 |

Pull request:
[Snapshot] Data progress tracking
This PR adds progress tracking to the data snapshots. The initial idea was to have a progress bar for each table being snapshotted. However, libraries supporting multiple progress bars ([multibar](https://github.com/sethgrid/multibar), [mpb](https://github.com/vbauerster/mpb) and [uiprogress](https://github.com/gosuri/uiprogress)) didn't show the progress properly, duplicating the result rather than updating the existing bar. 

Instead, this PR reports the progress for all tables in a schema in the same progress bar, using the [progressbar](https://github.com/schollz/progressbar) library.

The progress uses the total bytes for the snapshotted tables (using `pg_table_size` to avoid including indices which are snapshotted separately), and the average row bytes to calculate the progress made by each batch.

The logging of intermediate table snapshotting steps has been downgraded from `INFO` to `DEBUG` in favour of relying on the progress bar. 


Fixes #286 

Patch:
diff --git a/go.mod b/go.mod
index 356bc0b..b8e95cf 100644
--- a/go.mod
+++ b/go.mod
@@ -25,6 +25,7 @@ require (
 	github.com/pterm/pterm v0.12.81
 	github.com/rs/xid v1.6.0
 	github.com/rs/zerolog v1.34.0
+	github.com/schollz/progressbar/v3 v3.18.0
 	github.com/segmentio/kafka-go v0.4.48
 	github.com/spf13/cobra v1.9.1
 	github.com/spf13/pflag v1.0.6
@@ -111,6 +112,7 @@ require (
 	github.com/mattn/go-colorable v0.1.14 // indirect
 	github.com/mattn/go-isatty v0.0.20 // indirect
 	github.com/mattn/go-runewidth v0.0.16 // indirect
+	github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db // indirect
 	github.com/mitchellh/copystructure v1.2.0 // indirect
 	github.com/mitchellh/reflectwalk v1.0.2 // indirect
 	github.com/moby/docker-image-spec v1.3.1 // indirect
diff --git a/go.sum b/go.sum
index f80c24a..ae74a86 100644
--- a/go.sum
+++ b/go.sum
@@ -62,6 +62,8 @@ github.com/cenkalti/backoff/v5 v5.0.2/go.mod h1:rkhZdG3JZukswDf7f0cwqPNk4K0sa+F9
 github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=
 github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
 github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
+github.com/chengxilo/virtualterm v1.0.4 h1:Z6IpERbRVlfB8WkOmtbHiDbBANU7cimRIof7mk9/PwM=
+github.com/chengxilo/virtualterm v1.0.4/go.mod h1:DyxxBZz/x1iqJjFxTFcr6/x+jSpqN0iwWCOK1q10rlY=
 github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDkc90ppPyw=
 github.com/cloudwego/base64x v0.1.5 h1:XPciSp1xaq2VCSt6lF0phncD4koWyULpl5bUxbfCyP4=
 github.com/cloudwego/base64x v0.1.5/go.mod h1:0zlkT4Wn5C6NdauXdJRhSKRlJvmclQ1hhJgA0rcu/8w=
@@ -290,6 +292,8 @@ github.com/mattn/go-runewidth v0.0.16 h1:E5ScNMtiwvlvB5paMFdw9p4kSQzbXFikJ5SQO6T
 github.com/mattn/go-runewidth v0.0.16/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
 github.com/mdelapenya/tlscert v0.2.0 h1:7H81W6Z/4weDvZBNOfQte5GpIMo0lGYEeWbkGp5LJHI=
 github.com/mdelapenya/tlscert v0.2.0/go.mod h1:O4njj3ELLnJjGdkN7M/vIVCpZ+Cf0L6muqOG4tLSl8o=
+github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db h1:62I3jR2EmQ4l5rM/4FEfDWcRD+abF5XlKShorW5LRoQ=
+github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db/go.mod h1:l0dey0ia/Uv7NcFFVbCLtqEBQbrT4OCwCSKTEv6enCw=
 github.com/mitchellh/copystructure v1.2.0 h1:vpKXTN4ewci03Vljg/q9QvCGUDttBOGBIa15WveJJGw=
 github.com/mitchellh/copystructure v1.2.0/go.mod h1:qLl+cE2AmVv+CoeAwDPye/v+N2HKCj9FbZEVFJRxO9s=
 github.com/mitchellh/mapstructure v1.5.0 h1:jeMsZIYE/09sWLaz43PL7Gy6RuMjD2eJVyuac5Z2hdY=
@@ -378,6 +382,8 @@ github.com/russross/blackfriday/v2 v2.1.0 h1:JIOH55/0cWyOuilr9/qlrm0BSXldqnqwMsf
 github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
 github.com/sagikazarmark/locafero v0.7.0 h1:5MqpDsTGNDhY8sGp0Aowyf0qKsPrhewaLSsFaodPcyo=
 github.com/sagikazarmark/locafero v0.7.0/go.mod h1:2za3Cg5rMaTMoG/2Ulr9AwtFaIppKXTRYnozin4aB5k=
+github.com/schollz/progressbar/v3 v3.18.0 h1:uXdoHABRFmNIjUfte/Ex7WtuyVslrw2wVPQmCN62HpA=
+github.com/schollz/progressbar/v3 v3.18.0/go.mod h1:IsO3lpbaGuzh8zIMzgY3+J8l4C8GjO0Y9S69eFvNsec=
 github.com/segmentio/kafka-go v0.4.48 h1:9jyu9CWK4W5W+SroCe8EffbrRZVqAOkuaLd/ApID4Vs=
 github.com/segmentio/kafka-go v0.4.48/go.mod h1:HjF6XbOKh0Pjlkr5GVZxt6CsjjwnmhVOfURM5KMd8qg=
 github.com/segmentio/ksuid v1.0.4 h1:sBo2BdShXjmcugAMwjugoGUdUV0pcxY5mW4xKRn3v4c=
diff --git a/internal/postgres/pg_restore.go b/internal/postgres/pg_restore.go
index c62f325..ea0a3d4 100644
--- a/internal/postgres/pg_restore.go
+++ b/internal/postgres/pg_restore.go
@@ -93,7 +93,7 @@ func RunPGRestore(ctx context.Context, opts PGRestoreOptions, dump []byte) (stri
 	// TODO: add streaming support when large data output is required
 	out, err := cmd.CombinedOutput()
 	if err != nil || strings.Contains(string(out), "ERROR") {
-		return "", fmt.Errorf("error running pg_restore: %w", parsePgRestoreOutputErrs(out))
+		return "", fmt.Errorf("error restoring dump: %w", parsePgRestoreOutputErrs(out))
 	}
 
 	return string(out), nil
diff --git a/internal/progress/mocks/mock_progress_bar.go b/internal/progress/mocks/mock_progress_bar.go
new file mode 100644
index 0000000..125c9fd
--- /dev/null
+++ b/internal/progress/mocks/mock_progress_bar.go
@@ -0,0 +1,30 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package mocks
+
+type Bar struct {
+	AddFn   func(int) error
+	Add64Fn func(int64) error
+	CloseFn func() error
+}
+
+func (b *Bar) Add(n int) error {
+	if b.AddFn != nil {
+		return b.AddFn(n)
+	}
+	return nil
+}
+
+func (b *Bar) Add64(n int64) error {
+	if b.Add64Fn != nil {
+		return b.Add64Fn(n)
+	}
+	return nil
+}
+
+func (b *Bar) Close() error {
+	if b.CloseFn != nil {
+		return b.CloseFn()
+	}
+	return nil
+}
diff --git a/internal/progress/progress_bar.go b/internal/progress/progress_bar.go
new file mode 100644
index 0000000..16ce067
--- /dev/null
+++ b/internal/progress/progress_bar.go
@@ -0,0 +1,51 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package progress
+
+import (
+	"fmt"
+
+	"github.com/schollz/progressbar/v3"
+)
+
+type Bar interface {
+	Add(int) error
+	Add64(int64) error
+	Close() error
+}
+
+type ProgressBar struct {
+	*progressbar.ProgressBar
+}
+
+func NewBar(max int, opts ...progressbar.Option) *ProgressBar {
+	return &ProgressBar{
+		ProgressBar: progressbar.NewOptions(max, opts...),
+	}
+}
+
+func NewBytesBar(totalBytes int64, description string) *ProgressBar {
+	return &ProgressBar{
+		ProgressBar: progressbar.NewOptions64(totalBytes,
+			progressbar.OptionShowCount(),
+			progressbar.OptionSetRenderBlankState(true),
+			progressbar.OptionSetPredictTime(true),
+			progressbar.OptionSetWidth(20),
+			progressbar.OptionSetPredictTime(true),
+			progressbar.OptionEnableColorCodes(true),
+			progressbar.OptionShowBytes(true),
+			progressbar.OptionShowTotalBytes(true),
+			progressbar.OptionShowElapsedTimeOnFinish(),
+			progressbar.OptionSetDescription(description),
+			progressbar.OptionOnCompletion(func() {
+				fmt.Printf("\n") //nolint:forbidigo
+			}),
+			progressbar.OptionSetTheme(progressbar.Theme{
+				Saucer:        "[green]=[reset]",
+				SaucerHead:    "[green]>[reset]",
+				SaucerPadding: " ",
+				BarStart:      "[",
+				BarEnd:        "]",
+			})),
+	}
+}
diff --git a/pkg/snapshot/generator/postgres/data/instrumented_table_snapshot_generator.go b/pkg/snapshot/generator/postgres/data/instrumented_table_snapshot_generator.go
index 0ee6da2..6ed5e39 100644
--- a/pkg/snapshot/generator/postgres/data/instrumented_table_snapshot_generator.go
+++ b/pkg/snapshot/generator/postgres/data/instrumented_table_snapshot_generator.go
@@ -22,11 +22,11 @@ func newInstrumentedTableSnapshotGenerator(fn snapshotTableFn, i *otel.Instrumen
 	}
 }
 
-func (i *instrumentedTableSnapshotGenerator) snapshotTable(ctx context.Context, snapshotID string, schema, table string) (err error) {
+func (i *instrumentedTableSnapshotGenerator) snapshotTable(ctx context.Context, snapshotID string, table *table) (err error) {
 	ctx, span := otel.StartSpan(ctx, i.tracer, "tableSnapshotGenerator.SnapshotTable", trace.WithAttributes([]attribute.KeyValue{
-		{Key: "schema", Value: attribute.StringValue(schema)},
-		{Key: "table", Value: attribute.StringValue(table)},
+		{Key: "schema", Value: attribute.StringValue(table.schema)},
+		{Key: "table", Value: attribute.StringValue(table.name)},
 	}...))
 	defer otel.CloseSpan(span, err)
-	return i.snapshotTableFn(ctx, snapshotID, schema, table)
+	return i.snapshotTableFn(ctx, snapshotID, table)
 }
diff --git a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go
index 151cd5b..fa3e6f7 100644
--- a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go
+++ b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go
@@ -6,11 +6,13 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"strings"
 	"sync"
 
 	"github.com/jackc/pgx/v5/pgconn"
 	pglib "github.com/xataio/pgstream/internal/postgres"
 	pglibinstrumentation "github.com/xataio/pgstream/internal/postgres/instrumentation"
+	"github.com/xataio/pgstream/internal/progress"
 	loglib "github.com/xataio/pgstream/pkg/log"
 	"github.com/xataio/pgstream/pkg/otel"
 	"github.com/xataio/pgstream/pkg/snapshot"
@@ -33,15 +35,20 @@ type SnapshotGenerator struct {
 	// Function called for processing produced rows.
 	rowsProcessor          snapshot.RowsProcessor
 	tableSnapshotGenerator snapshotTableFn
+
+	progressTracking   bool
+	progressBars       map[string]progress.Bar
+	progressBarBuilder func(totalBytes int64, description string) progress.Bar
 }
 
 type mapper interface {
 	TypeForOID(context.Context, uint32) (string, error)
 }
 
-type tablePageInfo struct {
+type tableInfo struct {
 	pageCount    int
 	avgPageBytes int64
+	avgRowBytes  int64
 }
 
 type pageRange struct {
@@ -54,7 +61,13 @@ type schemaTables struct {
 	tables []string
 }
 
-type snapshotTableFn func(ctx context.Context, snapshotID string, schema, table string) error
+type table struct {
+	schema  string
+	name    string
+	rowSize int64
+}
+
+type snapshotTableFn func(ctx context.Context, snapshotID string, table *table) error
 
 type Option func(sg *SnapshotGenerator)
 
@@ -106,6 +119,16 @@ func WithInstrumentation(i *otel.Instrumentation) Option {
 	}
 }
 
+func WithProgressTracking() Option {
+	return func(sg *SnapshotGenerator) {
+		sg.progressTracking = true
+		sg.progressBars = make(map[string]progress.Bar)
+		sg.progressBarBuilder = func(totalBytes int64, description string) progress.Bar {
+			return progress.NewBytesBar(totalBytes, description)
+		}
+	}
+}
+
 func (sg *SnapshotGenerator) CreateSnapshot(ctx context.Context, ss *snapshot.Snapshot) (err error) {
 	defer func() {
 		// make sure we close the rows processor once the snapshot is completed.
@@ -120,9 +143,9 @@ func (sg *SnapshotGenerator) CreateSnapshot(ctx context.Context, ss *snapshot.Sn
 	for i := uint(0); i < sg.snapshotWorkers; i++ {
 		errGroup.Go(func() error {
 			for schemaTables := range schemaTablesChan {
-				sg.logger.Info("creating snapshot for schema", loglib.Fields{"schema": schemaTables.schema, "tables": schemaTables.tables})
+				sg.logger.Info("creating data snapshot", loglib.Fields{"schema": schemaTables.schema, "tables": schemaTables.tables})
 				if err := sg.createSchemaSnapshot(ctx, schemaTables); err != nil {
-					sg.logger.Error(err, "creating snapshot for schema", loglib.Fields{"schema": schemaTables.schema, "tables": schemaTables.tables, "error": err.Error()})
+					sg.logger.Error(err, "creating data snapshot", loglib.Fields{"schema": schemaTables.schema, "tables": schemaTables.tables, "error": err.Error()})
 					schemaErrs[schemaTables.schema] = err
 				}
 			}
@@ -158,13 +181,24 @@ func (sg *SnapshotGenerator) createSchemaSnapshot(ctx context.Context, schemaTab
 	// The transaction snapshot is available for use only until the end of the
 	// transaction that exported it.
 	// https://www.postgresql.org/docs/current/functions-admin.html#FUNCTIONS-SNAPSHOT-SYNCHRONIZATION
-	return sg.conn.ExecInTxWithOptions(ctx, func(tx pglib.Tx) error {
+	return sg.conn.ExecInTxWithOptions(ctx, func(tx pglib.Tx) (err error) {
 		snapshotID, err := sg.exportSnapshot(ctx, tx)
 		if err != nil {
 			return snapshot.NewSchemaErrors(schemaTables.schema, err)
 		}
 
-		tableChan := make(chan string, len(schemaTables.tables))
+		if sg.progressTracking {
+			if err := sg.addProgressBar(ctx, snapshotID, schemaTables); err != nil {
+				return err
+			}
+			defer func() {
+				if err == nil {
+					sg.markProgressBarCompleted(schemaTables.schema)
+				}
+			}()
+		}
+
+		tableChan := make(chan *table, len(schemaTables.tables))
 		// a map of table errors per worker to avoid race conditions
 		workerTableErrs := make([]map[string]error, sg.schemaWorkers)
 		wg := &sync.WaitGroup{}
@@ -172,11 +206,14 @@ func (sg *SnapshotGenerator) createSchemaSnapshot(ctx context.Context, schemaTab
 		for i := uint(0); i < sg.schemaWorkers; i++ {
 			wg.Add(1)
 			workerTableErrs[i] = make(map[string]error, len(schemaTables.tables))
-			go sg.createSnapshotWorker(ctx, wg, schemaTables.schema, snapshotID, tableChan, workerTableErrs[i])
+			go sg.createSnapshotWorker(ctx, wg, snapshotID, tableChan, workerTableErrs[i])
 		}
 
-		for _, table := range schemaTables.tables {
-			tableChan <- table
+		for _, tableName := range schemaTables.tables {
+			tableChan <- &table{
+				schema: schemaTables.schema,
+				name:   tableName,
+			}
 		}
 
 		close(tableChan)
@@ -186,20 +223,20 @@ func (sg *SnapshotGenerator) createSchemaSnapshot(ctx context.Context, schemaTab
 	}, snapshotTxOptions())
 }
 
-func (sg *SnapshotGenerator) createSnapshotWorker(ctx context.Context, wg *sync.WaitGroup, schema, snapshotID string, tableChan <-chan string, tableErrMap map[string]error) {
+func (sg *SnapshotGenerator) createSnapshotWorker(ctx context.Context, wg *sync.WaitGroup, snapshotID string, tableChan <-chan *table, tableErrMap map[string]error) {
 	defer wg.Done()
-	for table := range tableChan {
-		logFields := loglib.Fields{"schema": schema, "table": table, "snapshotID": snapshotID}
-		sg.logger.Info("snapshotting table", logFields)
+	for t := range tableChan {
+		logFields := loglib.Fields{"schema": t.schema, "table": t.name, "snapshotID": snapshotID}
+		sg.logger.Debug("snapshotting table", logFields)
 
-		if err := sg.tableSnapshotGenerator(ctx, snapshotID, schema, table); err != nil {
+		if err := sg.tableSnapshotGenerator(ctx, snapshotID, t); err != nil {
 			sg.logger.Error(err, "snapshotting table", logFields)
 			// errors will get notified unless the table doesn't exist
 			if !errors.Is(err, pglib.ErrNoRows) {
-				tableErrMap[table] = err
+				tableErrMap[t.name] = err
 			}
 		}
-		sg.logger.Info("table snapshot completed", logFields)
+		sg.logger.Debug("table snapshot completed", logFields)
 	}
 }
 
@@ -240,38 +277,38 @@ func (sg *SnapshotGenerator) collectSchemaErrors(workerSchemaErrs map[string]err
 	return nil
 }
 
-func (sg *SnapshotGenerator) snapshotTable(ctx context.Context, snapshotID string, schema, table string) error {
-	tablePageInfo, err := sg.getTablePageInfo(ctx, schema, table, snapshotID)
+func (sg *SnapshotGenerator) snapshotTable(ctx context.Context, snapshotID string, table *table) error {
+	tableInfo, err := sg.getTableInfo(ctx, table.schema, table.name, snapshotID)
 	if err != nil {
 		return err
 	}
-	// the table is empty
-	if tablePageInfo.pageCount < 0 {
+	if tableInfo.isEmpty() {
 		return nil
 	}
+	table.rowSize = tableInfo.avgRowBytes
 
 	// If one page range fails, we abort the entire table snapshot. The
 	// snapshot relies on the transaction snapshot id to ensure all workers
 	// have the same table view, which allows us to use the ctid to
 	// parallelise the work.
-	rangeChan := make(chan pageRange, tablePageInfo.pageCount)
+	rangeChan := make(chan pageRange, tableInfo.pageCount)
 	errGroup, ctx := errgroup.WithContext(ctx)
 	for i := uint(0); i < sg.tableWorkers; i++ {
 		errGroup.Go(func() error {
-			return sg.snapshotTableRangeWorker(ctx, snapshotID, schema, table, rangeChan)
+			return sg.snapshotTableRangeWorker(ctx, snapshotID, table, rangeChan)
 		})
 	}
 
 	// automatically determine the batch page size based on the average page
 	// size and the configured batch bytes limit.
-	batchPageSize := tablePageInfo.getBatchPageSize(sg.batchBytes)
+	batchPageSize := tableInfo.getBatchPageSize(sg.batchBytes)
 	sg.logger.Debug("batch page size for table", loglib.Fields{
-		"schema": schema, "table": table, "batch_page_size": batchPageSize,
-		"page_count": tablePageInfo.pageCount, "avg_page_bytes": tablePageInfo.avgPageBytes,
+		"schema": table.schema, "table": table.name, "batch_page_size": batchPageSize,
+		"page_count": tableInfo.pageCount, "avg_page_bytes": tableInfo.avgPageBytes,
 	})
 	// page count returned by postgres starts at 0, so we need to include it
 	// when creating the page ranges.
-	for start := uint(0); start <= uint(tablePageInfo.pageCount); start += batchPageSize {
+	for start := uint(0); start <= uint(tableInfo.pageCount); start += batchPageSize {
 		rangeChan <- pageRange{
 			start: start,
 			end:   start + batchPageSize,
@@ -283,9 +320,9 @@ func (sg *SnapshotGenerator) snapshotTable(ctx context.Context, snapshotID strin
 	return errGroup.Wait()
 }
 
-func (sg *SnapshotGenerator) snapshotTableRangeWorker(ctx context.Context, snapshotID, schema, table string, pageRangeChan <-chan (pageRange)) error {
+func (sg *SnapshotGenerator) snapshotTableRangeWorker(ctx context.Context, snapshotID string, table *table, pageRangeChan <-chan (pageRange)) error {
 	for pageRange := range pageRangeChan {
-		if err := sg.snapshotTableRange(ctx, snapshotID, schema, table, pageRange); err != nil {
+		if err := sg.snapshotTableRange(ctx, snapshotID, table, pageRange); err != nil {
 			return err
 		}
 	}
@@ -294,13 +331,13 @@ func (sg *SnapshotGenerator) snapshotTableRangeWorker(ctx context.Context, snaps
 
 var pageRangeQuery = "SELECT * FROM %s WHERE ctid BETWEEN '(%d,0)' AND '(%d,0)'"
 
-func (sg *SnapshotGenerator) snapshotTableRange(ctx context.Context, snapshotID, schema, table string, pageRange pageRange) error {
+func (sg *SnapshotGenerator) snapshotTableRange(ctx context.Context, snapshotID string, table *table, pageRange pageRange) error {
 	return sg.execInSnapshotTx(ctx, snapshotID, func(tx pglib.Tx) error {
 		sg.logger.Debug(fmt.Sprintf("querying table page range %d-%d", pageRange.start, pageRange.end), loglib.Fields{
-			"schema": schema, "table": table, "snapshotID": snapshotID,
+			"schema": table.schema, "table": table.name, "snapshotID": snapshotID,
 		})
 
-		query := fmt.Sprintf(pageRangeQuery, pglib.QuoteQualifiedIdentifier(schema, table), pageRange.start, pageRange.end)
+		query := fmt.Sprintf(pageRangeQuery, pglib.QuoteQualifiedIdentifier(table.schema, table.name), pageRange.start, pageRange.end)
 		rows, err := tx.Query(ctx, query)
 		if err != nil {
 			return fmt.Errorf("querying table rows: %w", err)
@@ -326,8 +363,8 @@ func (sg *SnapshotGenerator) snapshotTableRange(ctx context.Context, snapshotID,
 				}
 
 				if err := sg.rowsProcessor.ProcessRow(ctx, &snapshot.Row{
-					Schema:  schema,
-					Table:   table,
+					Schema:  table.schema,
+					Table:   table.name,
 					Columns: columns,
 				}); err != nil {
 					return fmt.Errorf("processing snapshot row: %w", err)
@@ -335,7 +372,14 @@ func (sg *SnapshotGenerator) snapshotTableRange(ctx context.Context, snapshotID,
 			}
 		}
 
-		sg.logger.Debug(fmt.Sprintf("%d rows processed", rowCount))
+		if sg.progressTracking {
+			sg.logger.Debug(fmt.Sprintf("adding %d rows to progress bar", rowCount))
+			sg.progressBars[table.schema].Add64(int64(rowCount) * table.rowSize)
+		}
+
+		sg.logger.Debug(fmt.Sprintf("%d rows processed", rowCount), loglib.Fields{
+			"schema": table.schema, "table": table.name, "snapshotID": snapshotID,
+		})
 
 		return rows.Err()
 	})
@@ -360,11 +404,37 @@ func (sg *SnapshotGenerator) toSnapshotColumns(ctx context.Context, fieldDescrip
 	return columns
 }
 
+func (sg *SnapshotGenerator) addProgressBar(ctx context.Context, snapshotID string, schemaTables *schemaTables) error {
+	totalBytes, err := sg.getSnapshotSchemaTotalBytes(ctx, snapshotID, schemaTables.schema, schemaTables.tables)
+	if err != nil {
+		return err
+	}
+
+	sg.progressBars[schemaTables.schema] = sg.progressBarBuilder(totalBytes, fmt.Sprintf("[cyan][%s][reset] Snapshotting data...", schemaTables.schema))
+	return nil
+}
+
+func (sg *SnapshotGenerator) markProgressBarCompleted(schema string) {
+	bar, found := sg.progressBars[schema]
+	if !found {
+		return
+	}
+
+	bar.Close()
+	delete(sg.progressBars, schema)
+}
+
 // use pg_table_size instead of pg_total_relation_size since we only care about the size of the table itself and toast tables, not indices.
 // pg_relation_size will return only the size of the table itself, without toast tables.
-const tablePageInfoQuery = `SELECT
+const tableInfoQuery = `SELECT
   c.relpages AS page_count,
-  (pg_table_size(c.oid) / COALESCE(NULLIF(c.relpages, 0),1)) AS avg_page_size_bytes
+  (pg_table_size(c.oid) / COALESCE(NULLIF(c.relpages, 0),1)) AS avg_page_size_bytes,
+  CASE
+	WHEN c.reltuples > 0 THEN
+		ROUND(pg_table_size(c.oid) / c.reltuples)
+	ELSE
+		0
+  END AS avg_row_size
 FROM
   pg_class c
   JOIN pg_namespace n ON n.oid = c.relnamespace
@@ -373,14 +443,14 @@ WHERE
   AND n.nspname = $2
   AND c.relkind = 'r';`
 
-func (sg *SnapshotGenerator) getTablePageInfo(ctx context.Context, schemaName, tableName, snapshotID string) (*tablePageInfo, error) {
-	tablePageInfo := &tablePageInfo{}
+func (sg *SnapshotGenerator) getTableInfo(ctx context.Context, schemaName, tableName, snapshotID string) (*tableInfo, error) {
+	tableInfo := &tableInfo{}
 	err := sg.execInSnapshotTx(ctx, snapshotID, func(tx pglib.Tx) error {
-		if err := tx.QueryRow(ctx, tablePageInfoQuery, tableName, schemaName).Scan(&tablePageInfo.pageCount, &tablePageInfo.avgPageBytes); err != nil {
+		if err := tx.QueryRow(ctx, tableInfoQuery, tableName, schemaName).Scan(&tableInfo.pageCount, &tableInfo.avgPageBytes, &tableInfo.avgRowBytes); err != nil {
 			return fmt.Errorf("getting page information for table %s.%s: %w", schemaName, tableName, err)
 		}
 
-		sg.logger.Debug(fmt.Sprintf("table page count: %d, page avg bytes: %d", tablePageInfo.pageCount, tablePageInfo.avgPageBytes), loglib.Fields{
+		sg.logger.Debug(fmt.Sprintf("table page count: %d, page avg bytes: %d", tableInfo.pageCount, tableInfo.avgPageBytes), loglib.Fields{
 			"schema": schemaName, "table": tableName, "snapshotID": snapshotID,
 		})
 		return nil
@@ -389,7 +459,33 @@ func (sg *SnapshotGenerator) getTablePageInfo(ctx context.Context, schemaName, t
 		return nil, err
 	}
 
-	return tablePageInfo, nil
+	return tableInfo, nil
+}
+
+const tablesBytesQuery = `SELECT SUM(pg_table_size(c.oid)) FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace WHERE n.nspname = '%s' AND c.relname IN (%s) AND c.relkind = 'r';`
+
+func (sg *SnapshotGenerator) getSnapshotSchemaTotalBytes(ctx context.Context, snapshotID, schema string, tables []string) (int64, error) {
+	paramRefs := make([]string, 0, len(tables))
+	tableParams := make([]any, 0, len(tables))
+	for i, table := range tables {
+		tableParams = append(tableParams, table)
+		paramRefs = append(paramRefs, fmt.Sprintf("$%d", i+1))
+	}
+
+	totalBytes := int64(0)
+	query := fmt.Sprintf(tablesBytesQuery, schema, strings.Join(paramRefs, ","))
+	sg.logger.Debug("querying total bytes for schema", loglib.Fields{
+		"schema": schema, "tables": tables, "query": query, "snapshotID": snapshotID,
+	})
+	err := sg.execInSnapshotTx(ctx, snapshotID, func(tx pglib.Tx) error {
+		err := tx.QueryRow(ctx, query, tableParams...).Scan(&totalBytes)
+		if err != nil {
+			return fmt.Errorf("retrieving total bytes for schema: %w", err)
+		}
+		return nil
+	})
+
+	return totalBytes, err
 }
 
 const exportSnapshotQuery = `SELECT pg_export_snapshot()`
@@ -427,7 +523,7 @@ func snapshotTxOptions() pglib.TxOptions {
 	}
 }
 
-func (t *tablePageInfo) getBatchPageSize(bytes uint64) uint {
+func (t *tableInfo) getBatchPageSize(bytes uint64) uint {
 	// at least one page is needed to process the table
 	if t.pageCount == 0 {
 		return 1
@@ -451,3 +547,7 @@ func (t *tablePageInfo) getBatchPageSize(bytes uint64) uint {
 
 	return uint(batchPageSize)
 }
+
+func (t *tableInfo) isEmpty() bool {
+	return t.pageCount < 0
+}
diff --git a/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator.go b/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator.go
index b64e2b5..7b4bc4c 100644
--- a/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator.go
+++ b/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator.go
@@ -237,7 +237,7 @@ func (s *SnapshotGenerator) restoreDump(ctx context.Context, schemaTables map[st
 				return err
 			}
 			ignoredErrors := pgrestoreErr.GetIgnoredErrors()
-			s.logger.Warn(nil, fmt.Sprintf("pg_restore: %d errors ignored", len(ignoredErrors)), loglib.Fields{"errors_ignored": ignoredErrors})
+			s.logger.Warn(nil, fmt.Sprintf("restore: %d errors ignored", len(ignoredErrors)), loglib.Fields{"errors_ignored": ignoredErrors})
 		default:
 			return err
 		}
diff --git a/pkg/wal/listener/snapshot/builder/wal_listener_snapshot_generator_builder.go b/pkg/wal/listener/snapshot/builder/wal_listener_snapshot_generator_builder.go
index 963a4c9..eb28763 100644
--- a/pkg/wal/listener/snapshot/builder/wal_listener_snapshot_generator_builder.go
+++ b/pkg/wal/listener/snapshot/builder/wal_listener_snapshot_generator_builder.go
@@ -57,6 +57,7 @@ func NewSnapshotGenerator(ctx context.Context, cfg *SnapshotListenerConfig, p li
 	// postgres data snapshot generator layer
 	opts := []pgsnapshotgenerator.Option{
 		pgsnapshotgenerator.WithLogger(logger),
+		pgsnapshotgenerator.WithProgressTracking(),
 	}
 	if instrumentation.IsEnabled() {
 		opts = append(opts, pgsnapshotgenerator.WithInstrumentation(instrumentation))


Test patch:
diff --git a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_test.go b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_test.go
index b30893f..4b5f71e 100644
--- a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_test.go
+++ b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_test.go
@@ -17,6 +17,8 @@ import (
 	"github.com/xataio/pgstream/internal/log/zerolog"
 	pglib "github.com/xataio/pgstream/internal/postgres"
 	pgmocks "github.com/xataio/pgstream/internal/postgres/mocks"
+	"github.com/xataio/pgstream/internal/progress"
+	progressmocks "github.com/xataio/pgstream/internal/progress/mocks"
 	"github.com/xataio/pgstream/pkg/snapshot"
 	"github.com/xataio/pgstream/pkg/snapshot/generator/mocks"
 	"github.com/xataio/pgstream/pkg/wal"
@@ -43,6 +45,8 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 	testSnapshotID := "test-snapshot-id"
 	testPageCount := 0 // 0 means 1 page
 	testPageAvgBytes := int64(1024)
+	testTotalBytes := int64(2048)
+	testRowBytes := int64(512)
 	testUUID := uuid.New().String()
 	testColumns := []snapshot.Column{
 		{Name: "id", Type: "uuid", Value: testUUID},
@@ -57,6 +61,20 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 		}
 	}
 
+	validTableInfoScanFn := func(args ...any) error {
+		require.Len(t, args, 3)
+		pageCount, ok := args[0].(*int)
+		require.True(t, ok, fmt.Sprintf("pageCount, expected *int, got %T", args[0]))
+		*pageCount = testPageCount
+		pageAvgBytes, ok := args[1].(*int64)
+		require.True(t, ok, fmt.Sprintf("pageAvgBytes, expected *int64, got %T", args[1]))
+		*pageAvgBytes = testPageAvgBytes
+		rowAvgBytes, ok := args[2].(*int64)
+		require.True(t, ok, fmt.Sprintf("rowAvgBytes, expected *int64, got %T", args[2]))
+		*rowAvgBytes = testRowBytes
+		return nil
+	}
+
 	errTest := errors.New("oh noes")
 
 	tests := []struct {
@@ -64,6 +82,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 		querier       pglib.Querier
 		snapshot      *snapshot.Snapshot
 		schemaWorkers uint
+		progressBar   *progressmocks.Bar
 
 		wantRows []*snapshot.Row
 		wantErr  error
@@ -98,17 +117,88 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 								return pglib.CommandTag{}, nil
 							},
 							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
-								require.Equal(t, tablePageInfoQuery, query)
+								require.Equal(t, tableInfoQuery, query)
 								require.Equal(t, []any{testTable1, testSchema}, args)
+								return &pgmocks.Row{
+									ScanFn: validTableInfoScanFn,
+								}
+							},
+						}
+						return f(&mockTx)
+					case 3:
+						mockTx := pgmocks.Tx{
+							ExecFn: func(ctx context.Context, _ uint, query string, args ...any) (pglib.CommandTag, error) {
+								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
+								require.Len(t, args, 0)
+								return pglib.CommandTag{}, nil
+							},
+							QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
+								require.Equal(t, fmt.Sprintf(pageRangeQuery, quotedSchemaTable1, 0, 1), query)
+								require.Len(t, args, 0)
+								return &pgmocks.Rows{
+									CloseFn: func() {},
+									NextFn:  func(i uint) bool { return i == 1 },
+									FieldDescriptionsFn: func() []pgconn.FieldDescription {
+										return []pgconn.FieldDescription{
+											{Name: "id", DataTypeOID: pgtype.UUIDOID},
+											{Name: "name", DataTypeOID: pgtype.TextOID},
+										}
+									},
+									ValuesFn: func() ([]any, error) {
+										return []any{testUUID, "alice"}, nil
+									},
+									ErrFn: func() error { return nil },
+								}, nil
+							},
+						}
+						return f(&mockTx)
+					default:
+						return fmt.Errorf("unexpected call to ExecInTxWithOptions: %d", i)
+					}
+				},
+			},
+
+			wantErr:  nil,
+			wantRows: []*snapshot.Row{testRow(testTable1, testColumns)},
+		},
+		{
+			name: "ok - with progress tracking",
+			querier: &pgmocks.Querier{
+				ExecInTxWithOptionsFn: func(_ context.Context, i uint, f func(tx pglib.Tx) error, to pglib.TxOptions) error {
+					require.Equal(t, txOptions, to)
+					switch i {
+					case 1:
+						mockTx := pgmocks.Tx{
+							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
+								require.Equal(t, exportSnapshotQuery, query)
+								return &pgmocks.Row{
+									ScanFn: func(args ...any) error {
+										require.Len(t, args, 1)
+										snapshotID, ok := args[0].(*string)
+										require.True(t, ok, fmt.Sprintf("snapshotID, expected *string, got %T", args[0]))
+										*snapshotID = testSnapshotID
+										return nil
+									},
+								}
+							},
+						}
+						return f(&mockTx)
+					case 2:
+						mockTx := pgmocks.Tx{
+							ExecFn: func(ctx context.Context, _ uint, query string, args ...any) (pglib.CommandTag, error) {
+								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
+								require.Len(t, args, 0)
+								return pglib.CommandTag{}, nil
+							},
+							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
+								require.Equal(t, fmt.Sprintf(tablesBytesQuery, testSchema, "$1"), query)
+								require.Equal(t, []any{testTable1}, args)
 								return &pgmocks.Row{
 									ScanFn: func(args ...any) error {
-										require.Len(t, args, 2)
-										pageCount, ok := args[0].(*int)
-										require.True(t, ok, fmt.Sprintf("pageCount, expected *int, got %T", args[0]))
-										*pageCount = testPageCount
-										pageAvgBytes, ok := args[1].(*int64)
-										require.True(t, ok, fmt.Sprintf("pageAvgBytes, expected *int, got %T", args[1]))
-										*pageAvgBytes = testPageAvgBytes
+										require.Len(t, args, 1)
+										totalBytes, ok := args[0].(*int64)
+										require.True(t, ok, fmt.Sprintf("totalBytes, expected *int64, got %T", args[0]))
+										*totalBytes = testTotalBytes
 										return nil
 									},
 								}
@@ -116,6 +206,22 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 						}
 						return f(&mockTx)
 					case 3:
+						mockTx := pgmocks.Tx{
+							ExecFn: func(ctx context.Context, _ uint, query string, args ...any) (pglib.CommandTag, error) {
+								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
+								require.Len(t, args, 0)
+								return pglib.CommandTag{}, nil
+							},
+							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
+								require.Equal(t, tableInfoQuery, query)
+								require.Equal(t, []any{testTable1, testSchema}, args)
+								return &pgmocks.Row{
+									ScanFn: validTableInfoScanFn,
+								}
+							},
+						}
+						return f(&mockTx)
+					case 4:
 						mockTx := pgmocks.Tx{
 							ExecFn: func(ctx context.Context, _ uint, query string, args ...any) (pglib.CommandTag, error) {
 								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
@@ -147,6 +253,12 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 					}
 				},
 			},
+			progressBar: &progressmocks.Bar{
+				Add64Fn: func(n int64) error {
+					require.Equal(t, testRowBytes, n) // only 1 row processed
+					return nil
+				},
+			},
 
 			wantErr:  nil,
 			wantRows: []*snapshot.Row{testRow(testTable1, testColumns)},
@@ -169,18 +281,9 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 									},
 								}
 							}
-							if query == tablePageInfoQuery {
+							if query == tableInfoQuery {
 								return &pgmocks.Row{
-									ScanFn: func(args ...any) error {
-										require.Len(t, args, 2)
-										pageCount, ok := args[0].(*int)
-										require.True(t, ok, fmt.Sprintf("pageCount, expected *int, got %T", args[0]))
-										*pageCount = testPageCount
-										pageAvgBytes, ok := args[1].(*int64)
-										require.True(t, ok, fmt.Sprintf("pageAvgBytes, expected *int, got %T", args[1]))
-										*pageAvgBytes = testPageAvgBytes
-										return nil
-									},
+									ScanFn: validTableInfoScanFn,
 								}
 							}
 							return &pgmocks.Row{
@@ -257,19 +360,10 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 								return pglib.CommandTag{}, nil
 							},
 							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
-								require.Equal(t, tablePageInfoQuery, query)
+								require.Equal(t, tableInfoQuery, query)
 								require.Equal(t, []any{testTable1, testSchema}, args)
 								return &pgmocks.Row{
-									ScanFn: func(args ...any) error {
-										require.Len(t, args, 2)
-										pageCount, ok := args[0].(*int)
-										require.True(t, ok, fmt.Sprintf("pageCount, expected *int, got %T", args[0]))
-										*pageCount = testPageCount
-										pageAvgBytes, ok := args[1].(*int64)
-										require.True(t, ok, fmt.Sprintf("pageAvgBytes, expected *int, got %T", args[1]))
-										*pageAvgBytes = testPageAvgBytes
-										return nil
-									},
+									ScanFn: validTableInfoScanFn,
 								}
 							},
 						}
@@ -341,19 +435,10 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 								return pglib.CommandTag{}, nil
 							},
 							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
-								require.Equal(t, tablePageInfoQuery, query)
+								require.Equal(t, tableInfoQuery, query)
 								require.Equal(t, []any{testTable1, testSchema}, args)
 								return &pgmocks.Row{
-									ScanFn: func(args ...any) error {
-										require.Len(t, args, 2)
-										pageCount, ok := args[0].(*int)
-										require.True(t, ok, fmt.Sprintf("pageCount, expected *int, got %T", args[0]))
-										*pageCount = testPageCount
-										pageAvgBytes, ok := args[1].(*int64)
-										require.True(t, ok, fmt.Sprintf("pageAvgBytes, expected *int, got %T", args[1]))
-										*pageAvgBytes = testPageAvgBytes
-										return nil
-									},
+									ScanFn: validTableInfoScanFn,
 								}
 							},
 						}
@@ -491,7 +576,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 								return pglib.CommandTag{}, nil
 							},
 							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
-								require.Equal(t, tablePageInfoQuery, query)
+								require.Equal(t, tableInfoQuery, query)
 								require.Equal(t, []any{testTable1, testSchema}, args)
 								return &pgmocks.Row{
 									ScanFn: func(args ...any) error {
@@ -547,19 +632,10 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 								return pglib.CommandTag{}, nil
 							},
 							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
-								require.Equal(t, tablePageInfoQuery, query)
+								require.Equal(t, tableInfoQuery, query)
 								require.Equal(t, []any{testTable1, testSchema}, args)
 								return &pgmocks.Row{
-									ScanFn: func(args ...any) error {
-										require.Len(t, args, 2)
-										pageCount, ok := args[0].(*int)
-										require.True(t, ok, fmt.Sprintf("pageCount, expected *int, got %T", args[0]))
-										*pageCount = testPageCount
-										pageAvgBytes, ok := args[1].(*int64)
-										require.True(t, ok, fmt.Sprintf("pageAvgBytes, expected *int, got %T", args[1]))
-										*pageAvgBytes = testPageAvgBytes
-										return nil
-									},
+									ScanFn: validTableInfoScanFn,
 								}
 							},
 						}
@@ -619,19 +695,10 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 								return pglib.CommandTag{}, nil
 							},
 							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
-								require.Equal(t, tablePageInfoQuery, query)
+								require.Equal(t, tableInfoQuery, query)
 								require.Equal(t, []any{testTable1, testSchema}, args)
 								return &pgmocks.Row{
-									ScanFn: func(args ...any) error {
-										require.Len(t, args, 2)
-										pageCount, ok := args[0].(*int)
-										require.True(t, ok, fmt.Sprintf("pageCount, expected *int, got %T", args[0]))
-										*pageCount = testPageCount
-										pageAvgBytes, ok := args[1].(*int64)
-										require.True(t, ok, fmt.Sprintf("pageAvgBytes, expected *int, got %T", args[1]))
-										*pageAvgBytes = testPageAvgBytes
-										return nil
-									},
+									ScanFn: validTableInfoScanFn,
 								}
 							},
 						}
@@ -696,19 +763,10 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 								return pglib.CommandTag{}, nil
 							},
 							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
-								require.Equal(t, tablePageInfoQuery, query)
+								require.Equal(t, tableInfoQuery, query)
 								require.Equal(t, []any{testTable1, testSchema}, args)
 								return &pgmocks.Row{
-									ScanFn: func(args ...any) error {
-										require.Len(t, args, 2)
-										pageCount, ok := args[0].(*int)
-										require.True(t, ok, fmt.Sprintf("pageCount, expected *int, got %T", args[0]))
-										*pageCount = testPageCount
-										pageAvgBytes, ok := args[1].(*int64)
-										require.True(t, ok, fmt.Sprintf("pageAvgBytes, expected *int, got %T", args[1]))
-										*pageAvgBytes = testPageAvgBytes
-										return nil
-									},
+									ScanFn: validTableInfoScanFn,
 								}
 							},
 						}
@@ -777,19 +835,10 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 								return pglib.CommandTag{}, nil
 							},
 							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
-								require.Equal(t, tablePageInfoQuery, query)
+								require.Equal(t, tableInfoQuery, query)
 								require.Equal(t, []any{testTable1, testSchema}, args)
 								return &pgmocks.Row{
-									ScanFn: func(args ...any) error {
-										require.Len(t, args, 2)
-										pageCount, ok := args[0].(*int)
-										require.True(t, ok, fmt.Sprintf("pageCount, expected *int, got %T", args[0]))
-										*pageCount = testPageCount
-										pageAvgBytes, ok := args[1].(*int64)
-										require.True(t, ok, fmt.Sprintf("pageAvgBytes, expected *int, got %T", args[1]))
-										*pageAvgBytes = testPageAvgBytes
-										return nil
-									},
+									ScanFn: validTableInfoScanFn,
 								}
 							},
 						}
@@ -846,7 +895,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 									},
 								}
 							}
-							if query == tablePageInfoQuery {
+							if query == tableInfoQuery {
 								return &pgmocks.Row{
 									ScanFn: func(args ...any) error {
 										return errTest
@@ -884,6 +933,65 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 			},
 			wantRows: []*snapshot.Row{},
 		},
+		{
+			name: "error - adding progress bar",
+			querier: &pgmocks.Querier{
+				ExecInTxWithOptionsFn: func(_ context.Context, i uint, f func(tx pglib.Tx) error, to pglib.TxOptions) error {
+					require.Equal(t, txOptions, to)
+					switch i {
+					case 1:
+						mockTx := pgmocks.Tx{
+							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
+								require.Equal(t, exportSnapshotQuery, query)
+								return &pgmocks.Row{
+									ScanFn: func(args ...any) error {
+										require.Len(t, args, 1)
+										snapshotID, ok := args[0].(*string)
+										require.True(t, ok, fmt.Sprintf("snapshotID, expected *string, got %T", args[0]))
+										*snapshotID = testSnapshotID
+										return nil
+									},
+								}
+							},
+						}
+						return f(&mockTx)
+					case 2:
+						mockTx := pgmocks.Tx{
+							ExecFn: func(ctx context.Context, _ uint, query string, args ...any) (pglib.CommandTag, error) {
+								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
+								require.Len(t, args, 0)
+								return pglib.CommandTag{}, nil
+							},
+							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
+								require.Equal(t, fmt.Sprintf(tablesBytesQuery, testSchema, "$1"), query)
+								require.Equal(t, []any{testTable1}, args)
+								return &pgmocks.Row{
+									ScanFn: func(args ...any) error {
+										return errTest
+									},
+								}
+							},
+						}
+						return f(&mockTx)
+					default:
+						return fmt.Errorf("unexpected call to ExecInTxWithOptions: %d", i)
+					}
+				},
+			},
+			progressBar: &progressmocks.Bar{
+				Add64Fn: func(n int64) error {
+					return errors.New("Add64Fn should not be called")
+				},
+			},
+
+			wantErr: snapshot.Errors{
+				testSchema: &snapshot.SchemaErrors{
+					Schema:       testSchema,
+					GlobalErrors: []string{fmt.Sprintf("retrieving total bytes for schema: %v", errTest)},
+				},
+			},
+			wantRows: []*snapshot.Row{},
+		},
 	}
 
 	for _, tc := range tests {
@@ -903,10 +1011,15 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 						return nil
 					},
 				},
-				schemaWorkers:   1,
-				tableWorkers:    1,
-				batchBytes:      1024 * 1024, // 1MB
-				snapshotWorkers: 1,
+				schemaWorkers:    1,
+				tableWorkers:     1,
+				batchBytes:       1024 * 1024, // 1MB
+				snapshotWorkers:  1,
+				progressTracking: tc.progressBar != nil,
+				progressBars:     make(map[string]progress.Bar),
+				progressBarBuilder: func(totalBytes int64, description string) progress.Bar {
+					return tc.progressBar
+				},
 			}
 			sg.tableSnapshotGenerator = sg.snapshotTable
 
@@ -939,15 +1052,15 @@ func TestTablePageInfo_getBatchPageSize(t *testing.T) {
 	t.Parallel()
 
 	tests := []struct {
-		name          string
-		tablePageInfo *tablePageInfo
-		bytes         uint64
+		name      string
+		tableInfo *tableInfo
+		bytes     uint64
 
 		wantPageSize uint
 	}{
 		{
 			name: "no avg page bytes - return total pages",
-			tablePageInfo: &tablePageInfo{
+			tableInfo: &tableInfo{
 				pageCount:    1,
 				avgPageBytes: 0,
 			},
@@ -955,7 +1068,7 @@ func TestTablePageInfo_getBatchPageSize(t *testing.T) {
 		},
 		{
 			name: "single page - return total pages",
-			tablePageInfo: &tablePageInfo{
+			tableInfo: &tableInfo{
 				pageCount:    0,
 				avgPageBytes: 0,
 			},
@@ -963,7 +1076,7 @@ func TestTablePageInfo_getBatchPageSize(t *testing.T) {
 		},
 		{
 			name: "multiple pages with average page bytes",
-			tablePageInfo: &tablePageInfo{
+			tableInfo: &tableInfo{
 				pageCount:    100,
 				avgPageBytes: 10,
 			},
@@ -972,7 +1085,7 @@ func TestTablePageInfo_getBatchPageSize(t *testing.T) {
 		},
 		{
 			name: "page size != page count",
-			tablePageInfo: &tablePageInfo{
+			tableInfo: &tableInfo{
 				pageCount:    100,
 				avgPageBytes: 10,
 			},
@@ -981,7 +1094,7 @@ func TestTablePageInfo_getBatchPageSize(t *testing.T) {
 		},
 		{
 			name: "bytes > avgPageBytes",
-			tablePageInfo: &tablePageInfo{
+			tableInfo: &tableInfo{
 				pageCount:    100,
 				avgPageBytes: 11,
 			},
@@ -990,7 +1103,7 @@ func TestTablePageInfo_getBatchPageSize(t *testing.T) {
 		},
 		{
 			name: "batch page size > total pages",
-			tablePageInfo: &tablePageInfo{
+			tableInfo: &tableInfo{
 				pageCount:    1,
 				avgPageBytes: 10,
 			},
@@ -1003,7 +1116,7 @@ func TestTablePageInfo_getBatchPageSize(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
-			pageSize := tc.tablePageInfo.getBatchPageSize(tc.bytes)
+			pageSize := tc.tableInfo.getBatchPageSize(tc.bytes)
 			require.Equal(t, tc.wantPageSize, pageSize, "wanted page size %d, got %d", tc.wantPageSize, pageSize)
 		})
 	}