feat: add backoff package for retry logic and implement example usage in new backoff-test script

cmd/go_migrate/connect.go

@@ -1,77 +0,0 @@
-package main
-
-import (
-	"context"
-	"database/sql"
-	"errors"
-	"fmt"
-	"sync"
-	"time"
-
-	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
-	"github.com/jackc/pgx/v5/pgxpool"
-	_ "github.com/microsoft/go-mssqldb"
-	log "github.com/sirupsen/logrus"
-)
-
-func connectToSqlServer() (*sql.DB, error) {
-	db, err := sql.Open("sqlserver", config.App.SourceDbUrl)
-	if err != nil {
-		return nil, fmt.Errorf("Unable to connect to sqlserver: %w", err)
-	}
-
-	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
-	defer cancel()
-
-	if err := db.PingContext(ctx); err != nil {
-		return nil, fmt.Errorf("Unable to ping sqlserver: %w", err)
-	}
-
-	return db, nil
-}
-
-func connectToPostgres() (*pgxpool.Pool, error) {
-	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
-	defer cancel()
-
-	pool, err := pgxpool.New(ctx, config.App.TargetDbUrl)
-	if err != nil {
-		return nil, fmt.Errorf("Unable to connect to postgres: %w", err)
-	}
-
-	if err := pool.Ping(ctx); err != nil {
-		pool.Close()
-		return nil, fmt.Errorf("Unable to ping postgres: %w", err)
-	}
-
-	return pool, nil
-}
-
-func connectToDatabases() (*sql.DB, *pgxpool.Pool, error) {
-	var sourceDbErr, targetDbErr error
-	var sourceDb *sql.DB
-	var targetDb *pgxpool.Pool
-	var wg sync.WaitGroup
-
-	wg.Go(func() {
-		sourceDb, sourceDbErr = connectToSqlServer()
-		if sourceDbErr != nil {
-			log.Error("Unable to connect to source db: ", sourceDbErr)
-		}
-	})
-
-	wg.Go(func() {
-		targetDb, targetDbErr = connectToPostgres()
-		if targetDbErr != nil {
-			log.Error("Unable to connect to target db: ", targetDbErr)
-		}
-	})
-
-	wg.Wait()
-
-	if sourceDbErr != nil || targetDbErr != nil {
-		return nil, nil, errors.New("Unable to connect to databases")
-	}
-
-	return sourceDb, targetDb, nil
-}

cmd/go_migrate/main.go

@@ -11,6 +11,7 @@ import (
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/etl/loaders"
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/etl/table_analyzers"
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/etl/transformers"
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
 	log "github.com/sirupsen/logrus"
 	"golang.org/x/sync/errgroup"
 )
@@ -95,10 +96,10 @@ func processMigrationJobs(
 	targetDb dbwrapper.DbWrapper,
 	jobs []config.Job,
 	maxParallelWorkers int,
-) []JobResult {
+) []models.JobResult {
 	if len(jobs) == 0 {
 		log.Info("No migration jobs configured")
-		return []JobResult{}
+		return []models.JobResult{}
 	}
 
 	if maxParallelWorkers <= 0 {
@@ -111,7 +112,7 @@ func processMigrationJobs(
 
 	log.Infof("Starting migration with %d parallel worker(s)", maxParallelWorkers)
 
-	chJobResults := make(chan JobResult, len(jobs))
+	chJobResults := make(chan models.JobResult, len(jobs))
 	chJobs := make(chan config.Job, len(jobs))
 	var wgJobs sync.WaitGroup
 
@@ -151,7 +152,7 @@ func processMigrationJobs(
 		close(chJobResults)
 	}()
 
-	var finalResults []JobResult
+	var finalResults []models.JobResult
 	for res := range chJobResults {
 		finalResults = append(finalResults, res)
 	}

cmd/go_migrate/metrics.go

@@ -1,13 +0,0 @@
-package main
-
-import "time"
-
-type JobResult struct {
-	JobName    string
-	StartTime  time.Time
-	Duration   time.Duration
-	RowsRead   int64
-	RowsLoaded int64
-	RowsFailed int64
-	Error      error
-}

cmd/go_migrate/process.go

@@ -26,11 +26,11 @@ func processMigrationJob(
 	transformer etl.Transformer,
 	loader etl.Loader,
 	job config.Job,
-) JobResult {
+) models.JobResult {
 	localCtx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
-	result := JobResult{
+	result := models.JobResult{
 		JobName:   job.Name,
 		StartTime: time.Now(),
 	}
@@ -129,7 +129,7 @@ func processMigrationJob(
 
 	for range maxExtractors {
 		wgExtractors.Go(func() {
-			extractor.Consume(
+			extractor.Exec(
 				localCtx,
 				job.SourceTable,
 				sourceColTypes,

go.mod

@@ -15,6 +15,8 @@ require (
 )
 
 require (
+	github.com/cenkalti/backoff v2.2.1+incompatible // indirect
+	github.com/cenkalti/backoff/v5 v5.0.3 // indirect
 	github.com/golang-sql/civil v0.0.0-20220223132316-b832511892a9 // indirect
 	github.com/golang-sql/sqlexp v0.1.0 // indirect
 	github.com/jackc/pgpassfile v1.0.0 // indirect

go.sum

@@ -16,6 +16,10 @@ github.com/alecthomas/assert/v2 v2.10.0 h1:jjRCHsj6hBJhkmhznrCzoNpbA3zqy0fYiUcYZ
 github.com/alecthomas/assert/v2 v2.10.0/go.mod h1:Bze95FyfUr7x34QZrjL+XP+0qgp/zg8yS+TtBj1WA3k=
 github.com/alecthomas/repr v0.4.0 h1:GhI2A8MACjfegCPVq9f1FLvIBS+DrQ2KQBFZP1iFzXc=
 github.com/alecthomas/repr v0.4.0/go.mod h1:Fr0507jx4eOXV7AlPV6AVZLYrLIuIeSOWtW57eE/O/4=
+github.com/cenkalti/backoff v2.2.1+incompatible h1:tNowT99t7UNflLxfYYSlKYsBpXdEet03Pg2g16Swow4=
+github.com/cenkalti/backoff v2.2.1+incompatible/go.mod h1:90ReRw6GdpyfrHakVjL/QHaoyV4aDUVVkXQJJJ3NXXM=
+github.com/cenkalti/backoff/v5 v5.0.3 h1:ZN+IMa753KfX5hd8vVaMixjnqRZ3y8CuJKRKj1xcsSM=
+github.com/cenkalti/backoff/v5 v5.0.3/go.mod h1:rkhZdG3JZukswDf7f0cwqPNk4K0sa+F97BxZthm/crw=
 github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=

internal/app/etl/extractors/mssql.go

@@ -9,7 +9,6 @@ import (
 	"strings"
 	"sync"
 	"sync/atomic"
-	"time"
 
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/convert"
@@ -100,7 +99,7 @@ func errorFromLastRow(
 	}
 }
 
-func (mssqlEx *MssqlExtractor) Extract(
+func (mssqlEx *MssqlExtractor) ProcessPartition(
 	ctx context.Context,
 	tableInfo config.SourceTableInfo,
 	columns []models.ColumnType,
@@ -108,7 +107,7 @@ func (mssqlEx *MssqlExtractor) Extract(
 	partition models.Partition,
 	indexPrimaryKey int,
 	chBatchesOut chan<- models.Batch,
-) (int64, error) {
+) (int, error) {
 	query := buildExtractQueryMssql(tableInfo, columns, partition.HasRange, partition.Range.IsMinInclusive)
 
 	var queryArgs []any
@@ -119,7 +118,7 @@ func (mssqlEx *MssqlExtractor) Extract(
 		)
 	}
 
-	var rowsRead int64 = 0
+	rowsRead := 0
 	rows, err := mssqlEx.db.Query(ctx, query, queryArgs...)
 	if err != nil {
 		return rowsRead, &custom_errors.ExtractorError{Partition: partition, HasLastId: false, Msg: err.Error()}
@@ -190,67 +189,14 @@ func (mssqlEx *MssqlExtractor) Extract(
 	return rowsRead, nil
 }
 
-func (mssqlEx *MssqlExtractor) ExtractWithRetries(
-	ctx context.Context,
-	tableInfo config.SourceTableInfo,
-	columns []models.ColumnType,
-	batchSize int,
-	partition models.Partition,
-	indexPrimaryKey int,
-	chBatchesOut chan<- models.Batch,
-) (int64, error) {
-	var totalRowsRead int64
-	var fatalErr error
-	delay := time.Duration(time.Second * 1)
-	currentParitition := partition
-
-	for fatalErr != nil || currentParitition.RetryCounter < 3 {
-		currentParitition.RetryCounter++
-		rowsRead, err := mssqlEx.Extract(
-			ctx,
-			tableInfo,
-			columns,
-			batchSize,
-			currentParitition,
-			indexPrimaryKey,
-			chBatchesOut,
-		)
-
-		if rowsRead > 0 {
-			totalRowsRead += int64(rowsRead)
-		}
-
-		if err != nil {
-			var exError *custom_errors.ExtractorError
-			if errors.As(err, &exError) {
-				if exError.HasLastId {
-					currentParitition.ParentId = exError.Partition.Id
-					currentParitition.Id = uuid.New()
-					currentParitition.Range.Min = exError.LastId
-					currentParitition.Range.IsMinInclusive = false
-				}
-
-				time.Sleep(delay)
-			} else {
-				fatalErr = err
-			}
-
-			continue
-		}
-
-		break
-	}
-
-	return totalRowsRead, fatalErr
-}
-
-func (mssqlEx *MssqlExtractor) Consume(
+func (mssqlEx *MssqlExtractor) Exec(
 	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,
@@ -285,7 +231,7 @@ func (mssqlEx *MssqlExtractor) Consume(
 				return
 			}
 
-			rowsReadResult, err := mssqlEx.ExtractWithRetries(
+			rowsReadResult, err := mssqlEx.ProcessPartition(
 				ctx,
 				tableInfo,
 				columns,
@@ -300,8 +246,15 @@ func (mssqlEx *MssqlExtractor) Consume(
 			}
 
 			if err != nil {
+				var exError *custom_errors.ExtractorError
 				var jobError *custom_errors.JobError
-				if errors.As(err, &jobError) {
+				if errors.As(err, &exError) {
+					select {
+					case <-ctx.Done():
+						return
+					case chErrorsOut <- *exError:
+					}
+				} else if errors.As(err, &jobError) {
 					select {
 					case <-ctx.Done():
 						return
@@ -311,7 +264,7 @@ func (mssqlEx *MssqlExtractor) Consume(
 					select {
 					case <-ctx.Done():
 						return
-					case chJobErrorsOut <- custom_errors.JobError{ShouldCancelJob: false, Msg: err.Error(), Prev: err}:
+					case chErrorsOut <- custom_errors.ExtractorError{Partition: partition, Msg: err.Error()}:
 					}
 				}
 

internal/app/etl/extractors/postgres.go

@@ -51,7 +51,7 @@ func buildExtractQueryPostgres(sourceDbInfo config.SourceTableInfo, columns []mo
 	return fmt.Sprintf(`SELECT %s FROM "%s"."%s"  ORDER BY "%s" ASC`, sbColumns.String(), sourceDbInfo.Schema, sourceDbInfo.Table, sourceDbInfo.PrimaryKey)
 }
 
-func (postgresEx *PostgresExtractor) Extract(
+func (postgresEx *PostgresExtractor) ProcessPartition(
 	ctx context.Context,
 	tableInfo config.SourceTableInfo,
 	columns []models.ColumnType,
@@ -110,7 +110,7 @@ func (postgresEx *PostgresExtractor) Extract(
 	return rowsRead, nil
 }
 
-func (postgresEx *PostgresExtractor) Consume(
+func (postgresEx *PostgresExtractor) Exec(
 	ctx context.Context,
 	tableInfo config.SourceTableInfo,
 	columns []models.ColumnType,

internal/app/etl/types.go

@@ -10,7 +10,7 @@ import (
 )
 
 type Extractor interface {
-	Extract(
+	ProcessPartition(
 		ctx context.Context,
 		tableInfo config.SourceTableInfo,
 		columns []models.ColumnType,
@@ -18,25 +18,16 @@ type Extractor interface {
 		partition models.Partition,
 		indexPrimaryKey int,
 		chBatchesOut chan<- models.Batch,
-	) (int64, error)
+	) (int, error)
 
-	ExtractWithRetries(
-		ctx context.Context,
-		tableInfo config.SourceTableInfo,
-		columns []models.ColumnType,
-		batchSize int,
-		partition models.Partition,
-		indexPrimaryKey int,
-		chBatchesOut chan<- models.Batch,
-	) (int64, error)
-
-	Consume(
+	Exec(
 		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,

internal/app/models/main.go

@@ -1,6 +1,10 @@
 package models
 
-import "github.com/google/uuid"
+import (
+	"time"
+
+	"github.com/google/uuid"
+)
 
 type UnknownRowValues = []any
 
@@ -25,3 +29,13 @@ type Partition struct {
 	HasRange     bool
 	RetryCounter int
 }
+
+type JobResult struct {
+	JobName    string
+	StartTime  time.Time
+	Duration   time.Duration
+	RowsRead   int64
+	RowsLoaded int64
+	RowsFailed int64
+	Error      error
+}

scripts/backoff-test/main.go

@@ -0,0 +1,58 @@
+package main
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net/http"
+	"strconv"
+
+	"github.com/cenkalti/backoff/v5"
+)
+
+func ExampleRetry() {
+	// Define an operation function that returns a value and an error.
+	// The value can be any type.
+	// We'll pass this operation to Retry function.
+	operation := func() (string, error) {
+		// An example request that may fail.
+		resp, err := http.Get("http://httpbin.org/get")
+		if err != nil {
+			return "", err
+		}
+		defer resp.Body.Close()
+
+		// If we are being rate limited, return a RetryAfter to specify how long to wait.
+		// This will also reset the backoff policy.
+		if resp.StatusCode == 429 {
+			seconds, err := strconv.ParseInt(resp.Header.Get("Retry-After"), 10, 64)
+			if err == nil {
+				return "", backoff.RetryAfter(int(seconds))
+			}
+		}
+
+		// In case of non-retriable error, return Permanent error to stop retrying.
+		// For this HTTP example, client errors are non-retriable.
+		if resp.StatusCode >= 400 && resp.StatusCode < 500 {
+			return "", backoff.Permanent(errors.New("bad request"))
+		}
+
+		// Return successful response.
+		return "hello", nil
+	}
+
+	result, err := backoff.Retry(context.TODO(), operation, backoff.WithBackOff(backoff.NewExponentialBackOff()))
+	if err != nil {
+		fmt.Println("Error:", err)
+		return
+	}
+
+	// Operation is successful after retries.
+
+	fmt.Println(result)
+	// Output: hello
+}
+
+func main() {
+	ExampleRetry()
+}