refactor: consolidate error handling and retry configuration

Centralize the error channel size, retry policy and max-failed batches/partition limits in the job config.

internal/app/etl/extractors/consume.go

@@ -0,0 +1,94 @@
+package extractors
+
+import (
+	"context"
+	"errors"
+	"slices"
+	"strings"
+	"sync"
+	"sync/atomic"
+
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/custom_errors"
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
+	"github.com/sirupsen/logrus"
+)
+
+func (ex *GenericExtractor) Consume(
+	ctx context.Context,
+	tableInfo config.SourceTableInfo,
+	columns []models.ColumnType,
+	batchSize int,
+	retryConfig config.RetryConfig,
+	chPartitionsIn <-chan models.Partition,
+	chBatchesOut chan<- models.Batch,
+	chErrorsOut chan<- custom_errors.JobError,
+	wgActivePartitions *sync.WaitGroup,
+	rowsRead *int64,
+) {
+	indexPrimaryKey := slices.IndexFunc(columns, func(col models.ColumnType) bool {
+		return strings.EqualFold(col.Name(), tableInfo.PrimaryKey)
+	})
+
+	if indexPrimaryKey == -1 {
+		select {
+		case <-ctx.Done():
+			return
+		case chErrorsOut <- custom_errors.JobError{
+			ShouldCancelJob: true,
+			Msg:             "Primary key not found in provided columns",
+		}:
+		}
+
+		return
+	}
+
+	for {
+		if ctx.Err() != nil {
+			return
+		}
+
+		select {
+		case <-ctx.Done():
+			return
+		case partition, ok := <-chPartitionsIn:
+			if !ok {
+				return
+			}
+
+			rowsReadResult, err := ex.ProcessPartitionWithRetries(
+				ctx,
+				tableInfo,
+				columns,
+				batchSize,
+				partition,
+				indexPrimaryKey,
+				retryConfig,
+				chBatchesOut,
+			)
+			wgActivePartitions.Done()
+
+			if rowsReadResult > 0 {
+				current := atomic.LoadInt64(rowsRead)
+				logrus.Debugf("Rows read: +%v [current=%v] (%s.%s)", rowsReadResult, current, tableInfo.Schema, tableInfo.Table)
+				atomic.AddInt64(rowsRead, int64(rowsReadResult))
+			}
+
+			if err != nil {
+				if jobError, ok := errors.AsType[*custom_errors.JobError](err); ok {
+					select {
+					case <-ctx.Done():
+						return
+					case chErrorsOut <- *jobError:
+					}
+				} else {
+					select {
+					case <-ctx.Done():
+						return
+					case chErrorsOut <- custom_errors.JobError{ShouldCancelJob: false, Msg: err.Error(), Prev: err}:
+					}
+				}
+			}
+		}
+	}
+}

internal/app/etl/extractors/main.go

@@ -2,16 +2,7 @@ package extractors
 
 import (
 	"context"
-	"errors"
-	"fmt"
-	"slices"
-	"strings"
-	"sync"
-	"sync/atomic"
 
-	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
-	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/convert"
-	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/custom_errors"
 	dbwrapper "git.ksdemosapps.com/kylesoda/go-migrate/internal/app/db-wrapper"
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
 	"github.com/google/uuid"
@@ -25,210 +16,27 @@ func NewExtractor(db dbwrapper.DbWrapper) GenericExtractor {
 	return GenericExtractor{db: db}
 }
 
-func errorFromLastRow(
-	lastRow models.UnknownRowValues,
-	indexPrimaryKey int,
-	partition models.Partition,
-	previousError error,
+func sendBatch(ctx context.Context, chBatchesOut chan<- models.Batch, batch models.Batch) error {
+	select {
+	case chBatchesOut <- batch:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}
+
+func flush(
+	ctx context.Context,
+	partition *models.Partition,
+	batchSize int,
+	batchRows []models.UnknownRowValues,
+	chBatchesOut chan<- models.Batch,
 ) error {
-	lastIdRawValue := lastRow[indexPrimaryKey]
-
-	lastId, ok := convert.ToInt64(lastIdRawValue)
-	if !ok {
-		currentPartition := partition
-		currentPartition.RetryCounter = 3
-		return &custom_errors.ExtractorError{
-			Partition: currentPartition,
-			HasLastId: true,
-			Msg:       fmt.Sprintf("Couldn't cast last id value as int: %s", previousError.Error()),
-		}
+	if len(batchRows) == 0 {
+		return nil
 	}
 
-	return &custom_errors.ExtractorError{
-		Partition: partition,
-		HasLastId: true,
-		LastId:    lastId,
-		Msg:       previousError.Error(),
-	}
-}
-
-func (ex *GenericExtractor) ProcessPartition(
-	ctx context.Context,
-	tableInfo config.SourceTableInfo,
-	columns []models.ColumnType,
-	batchSize int,
-	partition models.Partition,
-	indexPrimaryKey int,
-	chBatchesOut chan<- models.Batch,
-) (int, error) {
-	rowsRead := 0
-	query := dbwrapper.ExtractionQuery{
-		Schema:     tableInfo.Schema,
-		Table:      tableInfo.Table,
-		PrimaryKey: tableInfo.PrimaryKey,
-		LowerLimit: dbwrapper.ExtractorQueryLimit{
-			IsValid:     partition.HasRange && partition.Range.Min > 0,
-			IsInclusive: partition.Range.IsMinInclusive,
-			Value:       partition.Range.Min,
-		},
-		UpperLimit: dbwrapper.ExtractorQueryLimit{
-			IsValid:     partition.HasRange && partition.Range.Max > 0,
-			IsInclusive: partition.Range.IsMaxInclusive,
-			Value:       partition.Range.Max,
-		},
-	}
-
-	rows, err := ex.db.QueryFromObject(ctx, query)
-
-	if err != nil {
-		return rowsRead, &custom_errors.ExtractorError{Partition: partition, HasLastId: false, Msg: err.Error()}
-	}
-	defer rows.Close()
-
-	batchRows := make([]models.UnknownRowValues, 0, batchSize)
-
-	for rows.Next() {
-		rowValues := make([]any, len(columns))
-		scanArgs := make([]any, len(columns))
-
-		for i := range rowValues {
-			scanArgs[i] = &rowValues[i]
-		}
-
-		if err := rows.Scan(scanArgs...); err != nil {
-			if len(batchRows) == 0 {
-				return rowsRead, &custom_errors.ExtractorError{Partition: partition, HasLastId: false, Msg: err.Error()}
-			}
-
-			lastRow := batchRows[len(batchRows)-1]
-
-			select {
-			case chBatchesOut <- models.Batch{Id: uuid.New(), PartitionId: partition.Id, Rows: batchRows, RetryCounter: 0}:
-			case <-ctx.Done():
-				return rowsRead, ctx.Err()
-			}
-
-			return rowsRead, errorFromLastRow(lastRow, indexPrimaryKey, partition, err)
-		}
-		rowsRead++
-
-		batchRows = append(batchRows, rowValues)
-		if len(batchRows) >= batchSize {
-			select {
-			case chBatchesOut <- models.Batch{Id: uuid.New(), PartitionId: partition.Id, Rows: batchRows, RetryCounter: 0}:
-			case <-ctx.Done():
-				return rowsRead, ctx.Err()
-			}
-
-			batchRows = make([]models.UnknownRowValues, 0, batchSize)
-		}
-	}
-
-	if err := rows.Err(); err != nil {
-		if errors.Is(err, ctx.Err()) {
-			return rowsRead, ctx.Err()
-		}
-
-		if len(batchRows) > 0 {
-			lastRow := batchRows[len(batchRows)-1]
-			return rowsRead, errorFromLastRow(lastRow, indexPrimaryKey, partition, err)
-		}
-
-		return rowsRead, &custom_errors.ExtractorError{Partition: partition, HasLastId: false, Msg: err.Error()}
-	}
-
-	if len(batchRows) > 0 {
-		select {
-		case chBatchesOut <- models.Batch{Id: uuid.New(), PartitionId: partition.Id, Rows: batchRows, RetryCounter: 0}:
-		case <-ctx.Done():
-			return rowsRead, ctx.Err()
-		}
-	}
-
-	return rowsRead, nil
-}
-
-func (ex *GenericExtractor) Consume(
-	ctx context.Context,
-	tableInfo config.SourceTableInfo,
-	columns []models.ColumnType,
-	batchSize int,
-	chPartitionsIn <-chan models.Partition,
-	chBatchesOut chan<- models.Batch,
-	chErrorsOut chan<- custom_errors.ExtractorError,
-	chJobErrorsOut chan<- custom_errors.JobError,
-	wgActivePartitions *sync.WaitGroup,
-	rowsRead *int64,
-) {
-	indexPrimaryKey := slices.IndexFunc(columns, func(col models.ColumnType) bool {
-		return strings.EqualFold(col.Name(), tableInfo.PrimaryKey)
-	})
-
-	if indexPrimaryKey == -1 {
-		select {
-		case <-ctx.Done():
-			return
-		case chJobErrorsOut <- custom_errors.JobError{
-			ShouldCancelJob: true,
-			Msg:             "Primary key not found in provided columns",
-		}:
-		}
-
-		return
-	}
-
-	for {
-		if ctx.Err() != nil {
-			return
-		}
-
-		select {
-		case <-ctx.Done():
-			return
-		case partition, ok := <-chPartitionsIn:
-			if !ok {
-				return
-			}
-
-			rowsReadResult, err := ex.ProcessPartition(
-				ctx,
-				tableInfo,
-				columns,
-				batchSize,
-				partition,
-				indexPrimaryKey,
-				chBatchesOut,
-			)
-
-			if rowsReadResult > 0 {
-				atomic.AddInt64(rowsRead, int64(rowsReadResult))
-			}
-
-			if err != nil {
-				if exError, ok := errors.AsType[*custom_errors.ExtractorError](err); ok {
-					select {
-					case <-ctx.Done():
-						return
-					case chErrorsOut <- *exError:
-					}
-				} else if jobError, ok := errors.AsType[*custom_errors.JobError](err); ok {
-					select {
-					case <-ctx.Done():
-						return
-					case chJobErrorsOut <- *jobError:
-					}
-				} else {
-					select {
-					case <-ctx.Done():
-						return
-					case chErrorsOut <- custom_errors.ExtractorError{Partition: partition, Msg: err.Error()}:
-					}
-				}
-
-				continue
-			}
-
-			wgActivePartitions.Done()
-		}
-	}
+	batch := models.Batch{Id: uuid.New(), PartitionId: partition.Id, Rows: batchRows}
+	batchRows = make([]models.UnknownRowValues, 0, batchSize)
+	return sendBatch(ctx, chBatchesOut, batch)
 }

internal/app/etl/extractors/process-with-retries.go

@@ -0,0 +1,75 @@
+package extractors
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"time"
+
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/custom_errors"
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
+	"github.com/google/uuid"
+	// "github.com/sirupsen/logrus"
+)
+
+func (ex *GenericExtractor) ProcessPartitionWithRetries(
+	ctx context.Context,
+	tableInfo config.SourceTableInfo,
+	columns []models.ColumnType,
+	batchSize int,
+	partition models.Partition,
+	indexPrimaryKey int,
+	retryConfig config.RetryConfig,
+	chBatchesOut chan<- models.Batch,
+) (int64, error) {
+	var totalRowsRead int64
+	currentParitition := partition
+
+	for {
+		rowsRead, err := ex.ProcessPartition(
+			ctx,
+			tableInfo,
+			columns,
+			batchSize,
+			currentParitition,
+			indexPrimaryKey,
+			chBatchesOut,
+		)
+		// logrus.Debugf("Partition %v finished processing (%s.%s)", partition.Id, tableInfo.Schema, tableInfo.Table)
+		totalRowsRead += rowsRead
+
+		if err == nil {
+			return totalRowsRead, nil
+		}
+
+		if exError, ok := errors.AsType[*custom_errors.ExtractorError](err); ok {
+			currentParitition.RetryCounter++
+
+			if currentParitition.RetryCounter >= retryConfig.Attempts {
+				return totalRowsRead, &custom_errors.JobError{
+					Msg:  fmt.Sprintf("Partition %v reached max retries", exError.Partition.Id),
+					Prev: err,
+				}
+			}
+
+			if exError.HasLastId {
+				currentParitition.ParentId = exError.Partition.Id
+				currentParitition.Id = uuid.New()
+				currentParitition.Range.Min = exError.LastId
+				currentParitition.Range.IsMinInclusive = false
+			}
+
+			delay := custom_errors.ComputeBackoffDelay(
+				currentParitition.RetryCounter,
+				retryConfig.BaseDelayMs,
+				retryConfig.MaxDelayMs,
+				retryConfig.MaxJitterMs,
+			)
+			time.Sleep(delay)
+			continue
+		}
+
+		return totalRowsRead, err
+	}
+}

internal/app/etl/extractors/process.go

@@ -0,0 +1,116 @@
+package extractors
+
+import (
+	"context"
+	"fmt"
+
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/convert"
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/custom_errors"
+	dbwrapper "git.ksdemosapps.com/kylesoda/go-migrate/internal/app/db-wrapper"
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
+	// "github.com/sirupsen/logrus"
+)
+
+func errorFromLastPartitionRow(
+	lastRow models.UnknownRowValues,
+	indexPrimaryKey int,
+	partition models.Partition,
+	previousError error,
+) error {
+	lastIdRawValue := lastRow[indexPrimaryKey]
+
+	lastId, ok := convert.ToInt64(lastIdRawValue)
+	if !ok {
+		currentPartition := partition
+		currentPartition.RetryCounter = 3
+		return &custom_errors.ExtractorError{
+			Partition: currentPartition,
+			HasLastId: true,
+			Msg:       fmt.Sprintf("Couldn't cast last id value as int: %s", previousError.Error()),
+		}
+	}
+
+	return &custom_errors.ExtractorError{
+		Partition: partition,
+		HasLastId: true,
+		LastId:    lastId,
+		Msg:       previousError.Error(),
+	}
+}
+
+func (ex *GenericExtractor) ProcessPartition(
+	ctx context.Context,
+	tableInfo config.SourceTableInfo,
+	columns []models.ColumnType,
+	batchSize int,
+	partition models.Partition,
+	indexPrimaryKey int,
+	chBatchesOut chan<- models.Batch,
+) (int64, error) {
+	query := dbwrapper.ExtractionQuery{
+		Schema:     tableInfo.Schema,
+		Table:      tableInfo.Table,
+		PrimaryKey: tableInfo.PrimaryKey,
+		Columns:    columns,
+		LowerLimit: dbwrapper.ExtractorQueryLimit{
+			IsValid:     partition.HasRange && partition.Range.Min > 0,
+			IsInclusive: partition.Range.IsMinInclusive,
+			Value:       partition.Range.Min,
+		},
+		UpperLimit: dbwrapper.ExtractorQueryLimit{
+			IsValid:     partition.HasRange && partition.Range.Max > 0,
+			IsInclusive: partition.Range.IsMaxInclusive,
+			Value:       partition.Range.Max,
+		},
+	}
+
+	// logrus.Debugf("Processing partition: %+v (%s.%s)", query, tableInfo.Schema, tableInfo.Table)
+	rows, err := ex.db.QueryFromObject(ctx, query)
+	if err != nil {
+		return 0, err
+	}
+	defer rows.Close()
+
+	batchRows := make([]models.UnknownRowValues, 0, batchSize)
+	var rowsRead int64 = 0
+
+	for rows.Next() {
+		rowValues := make([]any, len(columns))
+		scanArgs := make([]any, len(columns))
+
+		for i := range rowValues {
+			scanArgs[i] = &rowValues[i]
+		}
+
+		if err := rows.Scan(scanArgs...); err != nil {
+			if len(batchRows) == 0 {
+				return rowsRead, err
+			}
+
+			if err := flush(ctx, &partition, batchSize, batchRows, chBatchesOut); err != nil {
+				return rowsRead, err
+			}
+
+			lastRow := batchRows[len(batchRows)-1]
+			return rowsRead, errorFromLastPartitionRow(lastRow, indexPrimaryKey, partition, err)
+		}
+		rowsRead++
+
+		batchRows = append(batchRows, rowValues)
+		if len(batchRows) >= batchSize {
+			// logrus.Debugf("Batch size reached, flushing batch with %v rows (rowsRead=%v)", len(batchRows), rowsRead)
+			if err := flush(ctx, &partition, batchSize, batchRows, chBatchesOut); err != nil {
+				// logrus.Warnf("Error flushing rows: %v", err)
+				return rowsRead, err
+			}
+			batchRows = make([]models.UnknownRowValues, 0, batchSize)
+		}
+	}
+
+	if err := flush(ctx, &partition, batchSize, batchRows, chBatchesOut); err != nil {
+		return rowsRead, err
+	}
+
+	return rowsRead, rows.Err()
+}

internal/app/etl/table_analyzers/main.go

@@ -7,6 +7,7 @@ import (
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/etl"
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
 	"github.com/google/uuid"
+	"github.com/sirupsen/logrus"
 )
 
 func PartitionRangeGenerator(
@@ -32,6 +33,7 @@ func PartitionRangeGenerator(
 	}
 
 	rowsCount, err := tableAnalyzer.EstimateTotalRows(ctx, tableInfo)
+	logrus.Infof("Estimated rows in source: %v (%s.%s)", rowsCount, tableInfo.Schema, tableInfo.Table)
 	if err != nil {
 		return nil, err
 	}
@@ -51,5 +53,7 @@ func PartitionRangeGenerator(
 		return nil, err
 	}
 
+	// logrus.Debugf("Partitions: %+v (%s.%s)", partitions, tableInfo.Schema, tableInfo.Table)
+
 	return partitions, nil
 }

internal/app/etl/table_analyzers/mssql.go

@@ -234,6 +234,7 @@ ORDER BY batch_id`,
 			RetryCounter: 0,
 			Range: models.PartitionRange{
 				IsMinInclusive: true,
+				IsMaxInclusive: true,
 			},
 		}