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

.gitignore

@@ -4,7 +4,6 @@
 *.dll
 *.so
 *.dylib
-bin/
 
 # Test binary, built with `go test -c`
 *.test

Makefile

@@ -1,80 +0,0 @@
-.PHONY: build build-linux build-windows build-all clean help
-
-# Variables
-BINARY_NAME=go-migrate
-CMD_PATH=./cmd/go_migrate
-OUTPUT_DIR=bin
-VERSION?=$(shell git describe --tags --always --dirty 2>/dev/null || echo "dev")
-BUILD_TIME=$(shell date -u '+%Y-%m-%d_%H:%M:%S')
-GIT_COMMIT=$(shell git rev-parse --short HEAD 2>/dev/null || echo "unknown")
-
-# Flags de compilación
-LD_FLAGS=-ldflags="-s -w -X main.Version=$(VERSION) -X main.BuildTime=$(BUILD_TIME) -X main.GitCommit=$(GIT_COMMIT)"
-
-# Default: compilar para el SO actual
-build: build-$(OS)
-
-ifeq ($(OS),Windows_NT)
-build-native: build-windows
-else
-build-native: build-linux
-endif
-
-# Compilar para Linux (sin CGO para máxima compatibilidad)
-build-linux:
-	@echo "Compilando para Linux..."
-	@mkdir -p $(OUTPUT_DIR)
-	CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build \
-		$(LD_FLAGS) \
-		-o $(OUTPUT_DIR)/$(BINARY_NAME)-linux-amd64 \
-		$(CMD_PATH)
-	@echo "Binario creado: $(OUTPUT_DIR)/$(BINARY_NAME)-linux-amd64"
-
-# Compilar para Windows
-build-windows:
-	@echo "Compilando para Windows..."
-	@mkdir -p $(OUTPUT_DIR)
-	CGO_ENABLED=0 GOOS=windows GOARCH=amd64 go build \
-		$(LD_FLAGS) \
-		-o $(OUTPUT_DIR)/$(BINARY_NAME)-windows-amd64.exe \
-		$(CMD_PATH)
-	@echo "Binario creado: $(OUTPUT_DIR)/$(BINARY_NAME)-windows-amd64.exe"
-
-# Compilar para ambas plataformas
-build-all: build-linux build-windows
-	@echo ""
-	@echo "Binarios compilados:"
-	@ls -lh $(OUTPUT_DIR)/$(BINARY_NAME)*
-
-# Compilar para Linux arm64 (opcional, para Raspberry Pi, etc.)
-build-linux-arm64:
-	@echo "Compilando para Linux ARM64..."
-	@mkdir -p $(OUTPUT_DIR)
-	CGO_ENABLED=0 GOOS=linux GOARCH=arm64 go build \
-		$(LD_FLAGS) \
-		-o $(OUTPUT_DIR)/$(BINARY_NAME)-linux-arm64 \
-		$(CMD_PATH)
-	@echo "Binario creado: $(OUTPUT_DIR)/$(BINARY_NAME)-linux-arm64"
-
-# Limpiar binarios
-clean:
-	@echo "Limpiando binarios..."
-	@rm -rf $(OUTPUT_DIR)
-	@echo "Limpieza completada"
-
-# Ayuda
-help:
-	@echo "Comandos disponibles:"
-	@echo ""
-	@echo "  make build              - Compilar para el SO actual (Linux/Windows)"
-	@echo "  make build-linux        - Compilar para Linux x86_64"
-	@echo "  make build-windows      - Compilar para Windows x86_64"
-	@echo "  make build-linux-arm64  - Compilar para Linux ARM64 (opcional)"
-	@echo "  make build-all          - Compilar para Linux y Windows"
-	@echo "  make clean              - Eliminar binarios compilados"
-	@echo "  make help               - Mostrar esta ayuda"
-	@echo ""
-	@echo "Ejemplos de uso:"
-	@echo "  make build-all          # Crear binarios para ambas plataformas"
-	@echo "  make build-linux OS=    # Crear solo para Linux"
-	@echo ""

cmd/go_migrate/main.go

@@ -2,17 +2,18 @@ package main
 
 import (
 	"context"
+	"database/sql"
 	"sync"
 	"time"
 
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
-	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/db-wrapper"
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/db"
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/etl/extractors"
 	"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"
+	"github.com/jackc/pgx/v5/pgxpool"
 	log "github.com/sirupsen/logrus"
-	"golang.org/x/sync/errgroup"
 )
 
 func main() {
@@ -32,33 +33,11 @@ func main() {
 
 	log.Info("=== Starting migration ===")
 
-	var wgConnect errgroup.Group
+	sourceDb, targetDb, connError := connectToDatabases()
-	var sourceDb, targetDb dbwrapper.DbWrapper
+	if connError != nil {
-
+		log.Fatal("Connection error: ", connError)
-	wgConnect.Go(func() error {
-		var err error
-		sourceDb, err = connectWithTimeout(ctx, migrationConfig.SourceDbType, config.App.SourceDbUrl, 20*time.Second)
-		if err != nil {
-			return err
-		}
-
-		return nil
-	})
-
-	wgConnect.Go(func() error {
-		var err error
-		targetDb, err = connectWithTimeout(ctx, migrationConfig.TargetDbType, config.App.TargetDbUrl, 20*time.Second)
-		if err != nil {
-			return err
-		}
-
-		return nil
-	})
-
-	if err := wgConnect.Wait(); err != nil {
-		log.Error("Connection error: ", err)
-		return
 	}
+
 	defer sourceDb.Close()
 	defer targetDb.Close()
 
@@ -91,8 +70,8 @@ func main() {
 
 func processMigrationJobs(
 	ctx context.Context,
-	sourceDb dbwrapper.DbWrapper,
+	sourceDb *sql.DB,
-	targetDb dbwrapper.DbWrapper,
+	targetDb *pgxpool.Pool,
 	jobs []config.Job,
 	maxParallelWorkers int,
 ) []JobResult {
@@ -115,6 +94,7 @@ func processMigrationJobs(
 	chJobs := make(chan config.Job, len(jobs))
 	var wgJobs sync.WaitGroup
 
+	targetDbWrapper := db.NewPostgresDbWrapper(targetDb)
 	sourceTableAnalyzer := table_analyzers.NewMssqlTableAnalyzer(sourceDb)
 	targetTableAnalyzer := table_analyzers.NewPostgresTableAnalyzer(targetDb)
 	extractor := extractors.NewMssqlExtractor(sourceDb)
@@ -127,7 +107,7 @@ func processMigrationJobs(
 				log.Infof("[worker %d] >>> Processing job: %s.%s <<<", i, job.SourceTable.Schema, job.SourceTable.Table)
 				res := processMigrationJob(
 					ctx,
-					targetDb,
+					targetDbWrapper,
 					sourceTableAnalyzer,
 					targetTableAnalyzer,
 					extractor,
@@ -158,19 +138,3 @@ func processMigrationJobs(
 
 	return finalResults
 }
-
-func connectWithTimeout(ctx context.Context, dbType string, dbUrl string, timeout time.Duration) (dbwrapper.DbWrapper, error) {
-	localCtx, cancel := context.WithTimeout(ctx, timeout)
-	defer cancel()
-
-	sourceDb, err := dbwrapper.New(dbType)
-	if err != nil {
-		return nil, err
-	}
-
-	if err = sourceDb.Connect(localCtx, dbUrl); err != nil {
-		return nil, err
-	}
-
-	return sourceDb, nil
-}

cmd/go_migrate/process.go

@@ -2,16 +2,14 @@ package main
 
 import (
 	"context"
-	"fmt"
 	"sync"
 	"sync/atomic"
 	"time"
 
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
 	"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/db"
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/etl"
-	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/etl/extractors"
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/etl/table_analyzers"
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
 	log "github.com/sirupsen/logrus"
@@ -20,7 +18,8 @@ import (
 
 func processMigrationJob(
 	ctx context.Context,
-	targetDbWrapper dbwrapper.DbWrapper,
+	// sourceDbWrapper db.DbWrapper,
+	targetDbWrapper db.DbWrapper,
 	sourceTableAnalyzer etl.TableAnalyzer,
 	targetTableAnalyzer etl.TableAnalyzer,
 	extractor etl.Extractor,
@@ -86,6 +85,7 @@ func processMigrationJob(
 	}
 
 	chJobErrors := make(chan custom_errors.JobError, job.QueueSize)
+	chExtractorErrors := make(chan custom_errors.ExtractorError, job.QueueSize)
 	chLoadersErrors := make(chan custom_errors.LoaderError, job.QueueSize)
 	chPartitions := make(chan models.Partition, job.QueueSize)
 	chBatchesRaw := make(chan models.Batch, job.QueueSize)
@@ -105,6 +105,15 @@ func processMigrationJob(
 		}
 	}()
 
+	go custom_errors.ExtractorErrorHandler(
+		localCtx,
+		job.Retry,
+		job.MaxPartitionErrrors,
+		chExtractorErrors,
+		chPartitions,
+		chJobErrors,
+		&wgActivePartitions,
+	)
 	go custom_errors.LoaderErrorHandler(
 		localCtx,
 		job.Retry,
@@ -120,14 +129,14 @@ func processMigrationJob(
 
 	for range maxExtractors {
 		wgExtractors.Go(func() {
-			extractors.Consume(
+			extractor.Exec(
 				localCtx,
-				extractor,
 				job.SourceTable,
 				sourceColTypes,
 				job.BatchSize,
 				chPartitions,
 				chBatchesRaw,
+				chExtractorErrors,
 				chJobErrors,
 				&wgActivePartitions,
 				&rowsRead,
@@ -181,6 +190,8 @@ func processMigrationJob(
 		log.Debugf("wgActivePartitions is empty (%v)", job.Name)
 		close(chPartitions)
 		log.Debugf("chPartitions is closed (%v)", job.Name)
+		close(chExtractorErrors)
+		log.Debugf("chExtractorErrors is closed (%v)", job.Name)
 
 		wgExtractors.Wait()
 		log.Debugf("wgExtractors is empty (%v)", job.Name)
@@ -223,9 +234,5 @@ func processMigrationJob(
 	result.RowsLoaded = atomic.LoadInt64(&rowsLoaded)
 	result.RowsFailed = atomic.LoadInt64(&rowsFailed)
 
-	if result.RowsRead != result.RowsLoaded {
-		result.Error = fmt.Errorf("Row count mismatch: extracted %d rows but loaded %d rows (failed: %d)", result.RowsRead, result.RowsLoaded, result.RowsFailed)
-	}
-
 	return result
 }

config.yaml

@@ -30,12 +30,6 @@ jobs:
       table: MANZANA
     pre_sql: 
       - 'SELECT 1'
-      # - 'TRUNCATE TABLE "Cartografia"."MANZANA"'
-    range:
-      min: 1000000
-      max: 2000000
-      is_min_inclusive: false
-      is_max_inclusive: true
 
   - name: red_puerto
     enabled: true
@@ -48,6 +42,5 @@ jobs:
       table: PUERTO
     pre_sql: 
       - 'SELECT 1'
-      # - 'TRUNCATE TABLE "Red"."PUERTO"'
     post_sql: 
       - "SELECT 1"

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/config/migration.go

@@ -50,12 +50,6 @@ type Job struct {
 	PreSQL      []string        `yaml:"pre_sql"`
 	PostSQL     []string        `yaml:"post_sql"`
 	JobConfig   `yaml:",inline"`
-	Range       struct {
-		Min            int64 `yaml:"min"`
-		Max            int64 `yaml:"max"`
-		IsMinInclusive bool  `yaml:"is_min_inclusive"`
-		IsMaxInclusive bool  `yaml:"is_max_inclusive"`
-	}
 }
 
 type MigrationConfig struct {
@@ -82,8 +76,6 @@ func (c *MigrationConfig) UnmarshalYAML(value *yaml.Node) error {
 
 	c.MaxParallelWorkers = raw.MaxParallelWorkers
 	c.Defaults = raw.Defaults
-	c.SourceDbType = raw.SourceDbType
-	c.TargetDbType = raw.TargetDbType
 	c.Defaults.RowsPerPartition = int64(raw.Defaults.BatchSize * raw.Defaults.BatchesPerPartition)
 
 	for _, node := range raw.Jobs {

internal/app/custom_errors/extractor.error.go

@@ -1,7 +1,13 @@
 package custom_errors
 
 import (
+	"context"
+	"fmt"
+	"sync"
+
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
+	"github.com/google/uuid"
 )
 
 type ExtractorError struct {
@@ -14,3 +20,100 @@ type ExtractorError struct {
 func (e *ExtractorError) Error() string {
 	return e.Msg
 }
+
+func ExtractorErrorHandler(
+	ctx context.Context,
+	retryConfig config.RetryConfig,
+	maxPartitionErrors int,
+	chErrorsIn <-chan ExtractorError,
+	chPartitionsOut chan<- models.Partition,
+	chJobErrorsOut chan<- JobError,
+	wgActivePartitions *sync.WaitGroup,
+) {
+	definitiveErrors := 0
+
+	for {
+		if ctx.Err() != nil {
+			return
+		}
+
+		select {
+		case <-ctx.Done():
+			return
+
+		case err, ok := <-chErrorsIn:
+			if !ok {
+				return
+			}
+
+			if err.Partition.RetryCounter >= retryConfig.Attempts {
+				wgActivePartitions.Done()
+				definitiveErrors++
+				jobError := JobError{
+					ShouldCancelJob: false,
+					Msg:             fmt.Sprintf("Partition %v reached max retries (%d)", err.Partition.Id, retryConfig.Attempts),
+					Prev:            &err,
+				}
+
+				select {
+				case chJobErrorsOut <- jobError:
+				case <-ctx.Done():
+					return
+				}
+
+				if maxPartitionErrors > 0 && definitiveErrors >= maxPartitionErrors {
+					fatalError := JobError{
+						ShouldCancelJob: true,
+						Msg:             fmt.Sprintf("Partition error limit reached (%d)", maxPartitionErrors),
+						Prev:            &err,
+					}
+
+					select {
+					case chJobErrorsOut <- fatalError:
+					case <-ctx.Done():
+						return
+					}
+				}
+
+				continue
+			} else {
+				jobError := JobError{
+					ShouldCancelJob: false,
+					Msg:             fmt.Sprintf("Temporal error in partition %v (retries: %d)", err.Partition.Id, err.Partition.RetryCounter),
+					Prev:            &err,
+				}
+
+				select {
+				case chJobErrorsOut <- jobError:
+				case <-ctx.Done():
+					return
+				}
+			}
+
+			newPartition := err.Partition
+			newPartition.RetryCounter++
+
+			delay := computeBackoffDelay(
+				newPartition.RetryCounter,
+				retryConfig.BaseDelayMs,
+				retryConfig.MaxDelayMs,
+				retryConfig.MaxJitterMs,
+			)
+
+			if err.HasLastId {
+				newPartition.ParentId = err.Partition.Id
+				newPartition.Id = uuid.New()
+				newPartition.LowerLimit = err.LastId
+				newPartition.IsLowerLimitInclusive = false
+			}
+
+			requeueWithBackoff(ctx, delay, func() {
+				select {
+				case chPartitionsOut <- newPartition:
+				case <-ctx.Done():
+					return
+				}
+			})
+		}
+	}
+}

internal/app/db-wrapper/main.go

@@ -1,19 +0,0 @@
-package dbwrapper
-
-import "fmt"
-
-type Factory func() DbWrapper
-
-var drivers = make(map[string]Factory)
-
-func Register(name string, factory Factory) {
-	drivers[name] = factory
-}
-
-func New(driverType string) (DbWrapper, error) {
-	factory, ok := drivers[driverType]
-	if !ok {
-		return nil, fmt.Errorf("driver not yet supported: %s", driverType)
-	}
-	return factory(), nil
-}

internal/app/db-wrapper/mssql.go

@@ -1,176 +0,0 @@
-package dbwrapper
-
-import (
-	"context"
-	"database/sql"
-	"fmt"
-
-	mssql "github.com/microsoft/go-mssqldb"
-)
-
-func init() {
-	Register("sqlserver", func() DbWrapper {
-		return &mssqlDbWrapper{dialect: "sqlserver"}
-	})
-}
-
-type mssqlRowResult struct {
-	row *sql.Row
-}
-
-func (mr *mssqlRowResult) Scan(dest ...any) error {
-	return mr.row.Scan(dest...)
-}
-
-type mssqlRowsResult struct {
-	columns []string
-	rows    *sql.Rows
-}
-
-func (mr *mssqlRowsResult) Close() error {
-	return mr.rows.Close()
-}
-
-func (mr *mssqlRowsResult) Columns() ([]string, error) {
-	if mr.columns != nil {
-		return mr.columns, nil
-	}
-
-	return mr.rows.Columns()
-}
-
-func (mr *mssqlRowsResult) Err() error {
-	return mr.rows.Err()
-}
-
-func (mr *mssqlRowsResult) Next() bool {
-	return mr.rows.Next()
-}
-
-func (mr *mssqlRowsResult) Scan(dest ...any) error {
-	return mr.rows.Scan(dest...)
-}
-
-func (mr *mssqlRowsResult) Values() ([]any, error) {
-	columns, err := mr.Columns()
-	if err != nil {
-		return nil, err
-	}
-
-	rowValues := make([]any, len(columns))
-	scanArgs := make([]any, len(columns))
-	for i := range rowValues {
-		scanArgs[i] = &rowValues[i]
-	}
-
-	if err := mr.rows.Scan(scanArgs...); err != nil {
-		return nil, err
-	}
-
-	return rowValues, nil
-}
-
-type mssqlDbWrapper struct {
-	db      *sql.DB
-	dialect string
-}
-
-func (mw *mssqlDbWrapper) Connect(ctx context.Context, dbUrl string) error {
-	db, err := sql.Open("sqlserver", dbUrl)
-	if err != nil {
-		return err
-	}
-
-	if err := db.PingContext(ctx); err != nil {
-		if err := db.Close(); err != nil {
-			return err
-		}
-		return err
-	}
-
-	mw.db = db
-	return nil
-}
-
-func (mw *mssqlDbWrapper) Close() error {
-	return mw.db.Close()
-}
-
-func (mw *mssqlDbWrapper) Exec(ctx context.Context, query string, args ...any) (ExecResult, error) {
-	result, execErr := mw.db.ExecContext(ctx, query, args...)
-	if execErr != nil {
-		return ExecResult{}, execErr
-	}
-
-	affectedRows, err := result.RowsAffected()
-	if err != nil {
-		return ExecResult{}, err
-	}
-
-	return ExecResult{AffectedRows: affectedRows}, nil
-}
-
-func (mw *mssqlDbWrapper) GetDialect() string {
-	return mw.dialect
-}
-
-func (mw *mssqlDbWrapper) Query(ctx context.Context, query string, args ...any) (RowsResult, error) {
-	rows, err := mw.db.QueryContext(ctx, query, args...)
-	if err != nil {
-		return nil, err
-	}
-
-	return &mssqlRowsResult{columns: nil, rows: rows}, nil
-}
-
-func (mw *mssqlDbWrapper) QueryRow(ctx context.Context, query string, args ...any) RowResult {
-	row := mw.db.QueryRowContext(ctx, query, args...)
-	return &mssqlRowResult{row: row}
-}
-
-func (mw *mssqlDbWrapper) SaveMassive(ctx context.Context, schema string, table string, columnNames []string, rows [][]any) (int64, error) {
-	tx, err := mw.db.BeginTx(ctx, nil)
-	if err != nil {
-		return 0, err
-	}
-
-	fullTableName := fmt.Sprintf("[%s].[%s]", schema, table)
-
-	stmt, err := tx.PrepareContext(ctx, mssql.CopyIn(fullTableName, mssql.BulkOptions{}, columnNames...))
-	if err != nil {
-		tx.Rollback()
-		return 0, err
-	}
-
-	for _, row := range rows {
-		_, err = stmt.ExecContext(ctx, row...)
-		if err != nil {
-			stmt.Close()
-			tx.Rollback()
-			return 0, err
-		}
-	}
-
-	result, err := stmt.ExecContext(ctx)
-	if err != nil {
-		stmt.Close()
-		tx.Rollback()
-		return 0, err
-	}
-
-	if err := stmt.Close(); err != nil {
-		tx.Rollback()
-		return 0, err
-	}
-
-	if err := tx.Commit(); err != nil {
-		return 0, err
-	}
-
-	rowsAffected, raErr := result.RowsAffected()
-	if raErr != nil {
-		return 0, nil
-	}
-
-	return rowsAffected, nil
-}

internal/app/db-wrapper/postgres.go

@@ -1,128 +0,0 @@
-package dbwrapper
-
-import (
-	"context"
-	"errors"
-
-	"github.com/jackc/pgx/v5"
-	"github.com/jackc/pgx/v5/pgxpool"
-)
-
-func init() {
-	Register("postgres", func() DbWrapper {
-		return &postgresDbWrapper{dialect: "postgres"}
-	})
-}
-
-type postgresRowResult struct {
-	row pgx.Row
-}
-
-func (pr *postgresRowResult) Scan(dest ...any) error {
-	return pr.row.Scan(dest...)
-}
-
-type postgresRowsResult struct {
-	columns []string
-	rows    pgx.Rows
-}
-
-func (pr *postgresRowsResult) Close() error {
-	pr.rows.Close()
-	return nil
-}
-
-func (pr *postgresRowsResult) Columns() ([]string, error) {
-	if pr.columns != nil {
-		return pr.columns, nil
-	}
-
-	rawColumns := pr.rows.FieldDescriptions()
-	if rawColumns == nil {
-		return nil, errors.New("error retrieving columns")
-	}
-
-	columns := make([]string, 0, len(rawColumns))
-	for _, rc := range rawColumns {
-		columns = append(columns, rc.Name)
-	}
-
-	return columns, nil
-}
-
-func (pr *postgresRowsResult) Err() error {
-	return pr.rows.Err()
-}
-
-func (pr *postgresRowsResult) Next() bool {
-	return pr.rows.Next()
-}
-
-func (pr *postgresRowsResult) Scan(dest ...any) error {
-	return pr.rows.Scan(dest...)
-}
-
-func (pr *postgresRowsResult) Values() ([]any, error) {
-	return pr.rows.Values()
-}
-
-type postgresDbWrapper struct {
-	db      *pgxpool.Pool
-	dialect string
-}
-
-func (pw *postgresDbWrapper) Connect(ctx context.Context, dbUrl string) error {
-	pool, err := pgxpool.New(ctx, dbUrl)
-	if err != nil {
-		return err
-	}
-
-	if err := pool.Ping(ctx); err != nil {
-		pool.Close()
-		return err
-	}
-
-	pw.db = pool
-	return nil
-}
-
-func (pw *postgresDbWrapper) Close() error {
-	pw.db.Close()
-	return nil
-}
-
-func (pw *postgresDbWrapper) Exec(ctx context.Context, query string, args ...any) (ExecResult, error) {
-	result, err := pw.db.Exec(ctx, query, args...)
-	if err != nil {
-		return ExecResult{}, err
-	}
-
-	return ExecResult{AffectedRows: result.RowsAffected()}, nil
-}
-
-func (pw *postgresDbWrapper) GetDialect() string {
-	return pw.dialect
-}
-
-func (pw *postgresDbWrapper) Query(ctx context.Context, query string, args ...any) (RowsResult, error) {
-	rows, err := pw.db.Query(ctx, query, args...)
-	if err != nil {
-		return nil, err
-	}
-
-	return &postgresRowsResult{columns: nil, rows: rows}, nil
-}
-
-func (pw *postgresDbWrapper) QueryRow(ctx context.Context, query string, args ...any) RowResult {
-	row := pw.db.QueryRow(ctx, query, args...)
-	return &postgresRowResult{row: row}
-}
-
-func (pw *postgresDbWrapper) SaveMassive(ctx context.Context, schema string, table string, columnNames []string, rows [][]any) (int64, error) {
-	affectedRows, err := pw.db.CopyFrom(ctx, pgx.Identifier{schema, table}, columnNames, pgx.CopyFromRows(rows))
-	if err != nil {
-		return 0, err
-	}
-
-	return affectedRows, nil
-}

internal/app/db-wrapper/types.go

@@ -1,35 +0,0 @@
-package dbwrapper
-
-import (
-	"context"
-	"errors"
-)
-
-var MethodNotSupported error = errors.New("Method not supported by driver... yet :P")
-
-type ExecResult struct {
-	AffectedRows int64
-}
-
-type RowsResult interface {
-	Close() error
-	Columns() ([]string, error)
-	Err() error
-	Next() bool
-	Scan(dest ...any) error
-	Values() ([]any, error)
-}
-
-type RowResult interface {
-	Scan(dest ...any) error
-}
-
-type DbWrapper interface {
-	Close() error
-	Connect(ctx context.Context, dbUrl string) error
-	Exec(ctx context.Context, query string, args ...any) (ExecResult, error)
-	GetDialect() string
-	Query(ctx context.Context, query string, args ...any) (RowsResult, error)
-	QueryRow(ctx context.Context, query string, args ...any) RowResult
-	SaveMassive(ctx context.Context, schema string, table string, columnNames []string, rows [][]any) (int64, error)
-}

internal/app/db/mssql.go

@@ -0,0 +1,30 @@
+package db
+
+import (
+	"context"
+	"database/sql"
+)
+
+type MssqlDbWrapper struct {
+	db *sql.DB
+}
+
+func NewMssqlDbWrapper(db *sql.DB) DbWrapper {
+	return &MssqlDbWrapper{db: db}
+}
+
+func (wrapper *MssqlDbWrapper) Exec(ctx context.Context, query string, args ...any) (DbWrapperResult, error) {
+	result, execErr := wrapper.db.ExecContext(ctx, query, args...)
+	if execErr != nil {
+		return DbWrapperResult{}, execErr
+	}
+
+	affectedRows, err := result.RowsAffected()
+	if err != nil {
+		return DbWrapperResult{}, err
+	}
+
+	return DbWrapperResult{
+		AffectedRows: affectedRows,
+	}, nil
+}

internal/app/db/postgres.go

@@ -0,0 +1,47 @@
+package db
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/jackc/pgx/v5/pgxpool"
+)
+
+func Connect(ctx context.Context, dbURL string) (*pgxpool.Pool, error) {
+	pool, err := pgxpool.New(ctx, dbURL)
+	if err != nil {
+		return nil, fmt.Errorf("unable to connect to database: %w", err)
+	}
+
+	if err := pool.Ping(ctx); err != nil {
+		pool.Close()
+		return nil, fmt.Errorf("unable to ping database: %w", err)
+	}
+
+	return pool, nil
+}
+
+func Close(pool *pgxpool.Pool) {
+	if pool != nil {
+		pool.Close()
+	}
+}
+
+type PostgresDbWrapper struct {
+	db *pgxpool.Pool
+}
+
+func NewPostgresDbWrapper(db *pgxpool.Pool) DbWrapper {
+	return &PostgresDbWrapper{db: db}
+}
+
+func (wrapper *PostgresDbWrapper) Exec(ctx context.Context, query string, args ...any) (DbWrapperResult, error) {
+	result, err := wrapper.db.Exec(ctx, query, args...)
+	if err != nil {
+		return DbWrapperResult{}, err
+	}
+
+	return DbWrapperResult{
+		AffectedRows: result.RowsAffected(),
+	}, nil
+}

internal/app/db/types.go

@@ -0,0 +1,11 @@
+package db
+
+import "context"
+
+type DbWrapperResult struct {
+	AffectedRows int64
+}
+
+type DbWrapper interface {
+	Exec(ctx context.Context, query string, args ...any) (DbWrapperResult, error)
+}

internal/app/etl/extractors/consumer.go

@@ -1,92 +0,0 @@
-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/etl"
-	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
-)
-
-func Consume(
-	ctx context.Context,
-	extractor etl.Extractor,
-	tableInfo config.SourceTableInfo,
-	columns []models.ColumnType,
-	batchSize int,
-	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 := extractWithRetries(
-				ctx,
-				extractor,
-				tableInfo,
-				columns,
-				batchSize,
-				partition,
-				indexPrimaryKey,
-				chBatchesOut,
-			)
-			wgActivePartitions.Done()
-
-			if rowsReadResult > 0 {
-				atomic.AddInt64(rowsRead, rowsReadResult)
-			}
-
-			if err != nil {
-				var jobError *custom_errors.JobError
-				if errors.As(err, &jobError) {
-					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/extract-with-retries.go

@@ -1,70 +0,0 @@
-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/etl"
-	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
-	"github.com/google/uuid"
-)
-
-func extractWithRetries(
-	ctx context.Context,
-	extractor etl.Extractor,
-	tableInfo config.SourceTableInfo,
-	columns []models.ColumnType,
-	batchSize int,
-	partition models.Partition,
-	indexPrimaryKey int,
-	chBatchesOut chan<- models.Batch,
-) (int64, error) {
-	var totalRowsRead int64
-	delay := time.Duration(time.Second * 1)
-	currentParitition := partition
-
-	for {
-		rowsRead, err := extractor.Exec(
-			ctx,
-			tableInfo,
-			columns,
-			batchSize,
-			currentParitition,
-			indexPrimaryKey,
-			chBatchesOut,
-		)
-		totalRowsRead += rowsRead
-
-		if err == nil {
-			return totalRowsRead, nil
-		}
-
-		var exError *custom_errors.ExtractorError
-		if errors.As(err, &exError) {
-			currentParitition.RetryCounter++
-
-			if currentParitition.RetryCounter > 3 {
-				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
-			}
-
-			time.Sleep(delay)
-			continue
-		}
-
-		return totalRowsRead, err
-	}
-}

internal/app/etl/extractors/main.go

@@ -1,64 +0,0 @@
-package extractors
-
-import (
-	"context"
-	"fmt"
-
-	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/convert"
-	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/custom_errors"
-	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
-	"github.com/google/uuid"
-)
-
-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 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 {
-	if len(batchRows) == 0 {
-		return nil
-	}
-
-	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/mssql.go

@@ -3,20 +3,26 @@ package extractors
 import (
 	"context"
 	"database/sql"
+	"errors"
 	"fmt"
+	"slices"
 	"strings"
+	"sync"
+	"sync/atomic"
 
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
-	dbwrapper "git.ksdemosapps.com/kylesoda/go-migrate/internal/app/db-wrapper"
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/convert"
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/custom_errors"
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/etl"
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
+	"github.com/google/uuid"
 )
 
 type MssqlExtractor struct {
-	db dbwrapper.DbWrapper
+	db *sql.DB
 }
 
-func NewMssqlExtractor(db dbwrapper.DbWrapper) etl.Extractor {
+func NewMssqlExtractor(db *sql.DB) etl.Extractor {
 	return &MssqlExtractor{db: db}
 }
 
@@ -64,7 +70,35 @@ func buildExtractQueryMssql(
 	return sbQuery.String()
 }
 
-func (mssqlEx *MssqlExtractor) Exec(
+func errorFromLastRow(
+	lastRow models.UnknownRowValues,
+	indexPrimaryKey int,
+	partition models.Partition,
+	previousError error,
+) *custom_errors.ExtractorError {
+	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 (mssqlEx *MssqlExtractor) ProcessPartition(
 	ctx context.Context,
 	tableInfo config.SourceTableInfo,
 	columns []models.ColumnType,
@@ -72,50 +106,171 @@ func (mssqlEx *MssqlExtractor) Exec(
 	partition models.Partition,
 	indexPrimaryKey int,
 	chBatchesOut chan<- models.Batch,
-) (int64, error) {
+) (int, error) {
-	query := buildExtractQueryMssql(tableInfo, columns, partition.HasRange, partition.Range.IsMinInclusive)
+	query := buildExtractQueryMssql(tableInfo, columns, partition.ShouldUseRange, partition.IsLowerLimitInclusive)
 
 	var queryArgs []any
-	if partition.HasRange {
+	if partition.ShouldUseRange {
-		queryArgs = append(queryArgs, sql.Named("min", partition.Range.Min), sql.Named("max", partition.Range.Max))
+		queryArgs = append(queryArgs,
+			sql.Named("min", partition.LowerLimit),
+			sql.Named("max", partition.UpperLimit),
+		)
 	}
 
-	rows, err := mssqlEx.db.Query(ctx, query, queryArgs...)
+	rowsRead := 0
+	rows, err := mssqlEx.db.QueryContext(ctx, query, queryArgs...)
 	if err != nil {
-		return 0, err
+		return rowsRead, &custom_errors.ExtractorError{Partition: partition, HasLastId: false, Msg: err.Error()}
 	}
 	defer rows.Close()
 
 	batchRows := make([]models.UnknownRowValues, 0, batchSize)
-	var rowsRead int64 = 0
 
 	for rows.Next() {
-		values, err := rows.Values()
+		rowValues := make([]any, len(columns))
-		if err != nil {
+		scanArgs := make([]any, len(columns))
-			if len(batchRows) == 0 {
-				return rowsRead, err
-			}
 
-			if err := flush(ctx, &partition, batchSize, batchRows, chBatchesOut); err != nil {
+		for i := range rowValues {
-				return rowsRead, err
+			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]
-			return rowsRead, errorFromLastPartitionRow(lastRow, indexPrimaryKey, partition, err)
+
+			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, values)
+		batchRows = append(batchRows, rowValues)
 		if len(batchRows) >= batchSize {
-			if err := flush(ctx, &partition, batchSize, batchRows, chBatchesOut); err != nil {
+			select {
-				return rowsRead, err
+			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()
 		}
 	}
 
-	if err := flush(ctx, &partition, batchSize, batchRows, chBatchesOut); err != nil {
+	return rowsRead, nil
-		return rowsRead, err
+}
+
+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,
+) {
+	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
 	}
 
-	return rowsRead, rows.Err()
+	for {
+		if ctx.Err() != nil {
+			return
+		}
+
+		select {
+		case <-ctx.Done():
+			return
+		case partition, ok := <-chPartitionsIn:
+			if !ok {
+				return
+			}
+
+			rowsReadResult, err := mssqlEx.ProcessPartition(
+				ctx,
+				tableInfo,
+				columns,
+				batchSize,
+				partition,
+				indexPrimaryKey,
+				chBatchesOut,
+			)
+
+			if rowsReadResult > 0 {
+				atomic.AddInt64(rowsRead, int64(rowsReadResult))
+			}
+
+			if err != nil {
+				var exError *custom_errors.ExtractorError
+				var jobError *custom_errors.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
+					case chJobErrorsOut <- *jobError:
+					}
+				} else {
+					select {
+					case <-ctx.Done():
+						return
+					case chErrorsOut <- custom_errors.ExtractorError{Partition: partition, Msg: err.Error()}:
+					}
+				}
+
+				continue
+			}
+
+			wgActivePartitions.Done()
+		}
+	}
 }

internal/app/etl/extractors/postgres.go

@@ -5,21 +5,22 @@ import (
 	"errors"
 	"fmt"
 	"strings"
+	"sync"
 
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
 	"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/etl"
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
 	"github.com/google/uuid"
+	"github.com/jackc/pgx/v5/pgxpool"
 )
 
 type PostgresExtractor struct {
-	db dbwrapper.DbWrapper
+	db *pgxpool.Pool
 }
 
-func NewPostgresExtractor(db dbwrapper.DbWrapper) etl.Extractor {
+func NewPostgresExtractor(pool *pgxpool.Pool) etl.Extractor {
-	return &PostgresExtractor{db: db}
+	return &PostgresExtractor{db: pool}
 }
 
 func buildExtractQueryPostgres(sourceDbInfo config.SourceTableInfo, columns []models.ColumnType) string {
@@ -50,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) Exec(
+func (postgresEx *PostgresExtractor) ProcessPartition(
 	ctx context.Context,
 	tableInfo config.SourceTableInfo,
 	columns []models.ColumnType,
@@ -58,14 +59,14 @@ func (postgresEx *PostgresExtractor) Exec(
 	partition models.Partition,
 	indexPrimaryKey int,
 	chBatchesOut chan<- models.Batch,
-) (int64, error) {
+) (int, error) {
 	query := buildExtractQueryPostgres(tableInfo, columns)
 
-	if partition.HasRange {
+	if partition.ShouldUseRange {
 		return 0, errors.New("Batch config not yet supported")
 	}
 
-	var rowsRead int64 = 0
+	rowsRead := 0
 	rows, err := postgresEx.db.Query(ctx, query)
 	if err != nil {
 		return rowsRead, &custom_errors.ExtractorError{Partition: partition, HasLastId: false, Msg: err.Error()}
@@ -108,3 +109,17 @@ func (postgresEx *PostgresExtractor) Exec(
 
 	return rowsRead, nil
 }
+
+func (postgresEx *PostgresExtractor) 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/etl/extractors/types.go

@@ -0,0 +1 @@
+package extractors

internal/app/etl/loaders/postgres.go

@@ -9,18 +9,19 @@ import (
 
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
 	"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/etl"
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
+	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgconn"
+	"github.com/jackc/pgx/v5/pgxpool"
 )
 
 type PostgresLoader struct {
-	db dbwrapper.DbWrapper
+	db *pgxpool.Pool
 }
 
-func NewPostgresLoader(db dbwrapper.DbWrapper) etl.Loader {
+func NewPostgresLoader(pool *pgxpool.Pool) etl.Loader {
-	return &PostgresLoader{db: db}
+	return &PostgresLoader{db: pool}
 }
 
 func mapSlice[T any, V any](input []T, mapper func(T) V) []V {
@@ -39,12 +40,12 @@ func (postgresLd *PostgresLoader) ProcessBatch(
 	colNames []string,
 	batch models.Batch,
 ) (int, error) {
-	_, err := postgresLd.db.SaveMassive(
+	tableId := pgx.Identifier{tableInfo.Schema, tableInfo.Table}
+	_, err := postgresLd.db.CopyFrom(
 		ctx,
-		tableInfo.Schema,
+		tableId,
-		tableInfo.Table,
 		colNames,
-		batch.Rows,
+		pgx.CopyFromRows(batch.Rows),
 	)
 
 	if err != nil {
@@ -53,7 +54,7 @@ func (postgresLd *PostgresLoader) ProcessBatch(
 			if pgErr.Code == "23505" {
 				return 0, &custom_errors.JobError{
 					ShouldCancelJob: true,
-					Msg:             fmt.Sprintf("Fatal error in table %s.%s", tableInfo.Schema, tableInfo.Table),
+					Msg:             fmt.Sprintf("Fatal error in table %s", tableId.Sanitize()),
 					Prev:            err,
 				}
 			}

internal/app/etl/table_analyzers/main.go

@@ -23,9 +23,9 @@ func PartitionRangeGenerator(
 
 	if rowsCount <= rowsPerPartition {
 		return []models.Partition{{
-			Id:           uuid.New(),
+			Id:             uuid.New(),
-			HasRange:     false,
+			ShouldUseRange: false,
-			RetryCounter: 0,
+			RetryCounter:   0,
 		}}, nil
 
 	}

internal/app/etl/table_analyzers/mssql.go

@@ -8,17 +8,16 @@ import (
 	"time"
 
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
-	dbwrapper "git.ksdemosapps.com/kylesoda/go-migrate/internal/app/db-wrapper"
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/etl"
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
 	"github.com/google/uuid"
 )
 
 type MssqlTableAnalyzer struct {
-	db dbwrapper.DbWrapper
+	db *sql.DB
 }
 
-func NewMssqlTableAnalyzer(db dbwrapper.DbWrapper) etl.TableAnalyzer {
+func NewMssqlTableAnalyzer(db *sql.DB) etl.TableAnalyzer {
 	return &MssqlTableAnalyzer{db: db}
 }
 
@@ -36,7 +35,7 @@ JOIN sys.types t ON c.user_type_id = t.user_type_id
 LEFT JOIN sys.types bt ON t.is_user_defined = 1 AND bt.user_type_id = t.system_type_id
 JOIN sys.tables st ON c.object_id = st.object_id
 JOIN sys.schemas s ON st.schema_id = s.schema_id
-WHERE s.name = @schema AND st.name = @table AND (c.is_hidden = 0 OR (c.graph_type IS NOT NULL AND c.name LIKE '$%'))
+WHERE s.name = @schema AND st.name = @table AND c.name NOT LIKE 'graph_id%'
 ORDER BY c.column_id;`
 
 type rawColumnMssql struct {
@@ -143,7 +142,7 @@ func (ta *MssqlTableAnalyzer) QueryColumnTypes(
 	localCtx, cancel := context.WithTimeout(ctx, 20*time.Second)
 	defer cancel()
 
-	rows, err := ta.db.Query(localCtx, mssqlColumnMetadataQuery, sql.Named("schema", tableInfo.Schema), sql.Named("table", tableInfo.Table))
+	rows, err := ta.db.QueryContext(localCtx, mssqlColumnMetadataQuery, sql.Named("schema", tableInfo.Schema), sql.Named("table", tableInfo.Table))
 	if err != nil {
 		return nil, err
 	}
@@ -184,11 +183,11 @@ JOIN sys.partitions p ON t.object_id = p.object_id
 WHERE s.name = @schema AND t.name = @table AND p.index_id IN (0, 1)
 GROUP BY t.name`
 
-	ctxTimeout, cancel := context.WithTimeout(ctx, 1*time.Minute)
+	ctxTimeout, cancel := context.WithTimeout(ctx, time.Second*20)
 	defer cancel()
 
 	var rowsCount int64
-	err := ta.db.QueryRow(ctxTimeout, query, sql.Named("schema", tableInfo.Schema), sql.Named("table", tableInfo.Table)).Scan(&rowsCount)
+	err := ta.db.QueryRowContext(ctxTimeout, query, sql.Named("schema", tableInfo.Schema), sql.Named("table", tableInfo.Table)).Scan(&rowsCount)
 	if err != nil {
 		return 0, err
 	}
@@ -216,10 +215,10 @@ ORDER BY batch_id`,
 		tableInfo.Schema,
 		tableInfo.Table)
 
-	ctxTimeout, cancel := context.WithTimeout(ctx, 1*time.Minute)
+	ctxTimeout, cancel := context.WithTimeout(ctx, time.Second*20)
 	defer cancel()
 
-	rows, err := ta.db.Query(ctxTimeout, query, sql.Named("maxPartitions", maxPartitions))
+	rows, err := ta.db.QueryContext(ctxTimeout, query, sql.Named("maxPartitions", maxPartitions))
 	if err != nil {
 		return nil, err
 	}
@@ -229,15 +228,13 @@ ORDER BY batch_id`,
 
 	for rows.Next() {
 		partition := models.Partition{
-			Id:           uuid.New(),
+			Id:                    uuid.New(),
-			HasRange:     true,
+			ShouldUseRange:        true,
-			RetryCounter: 0,
+			RetryCounter:          0,
-			Range: models.PartitionRange{
+			IsLowerLimitInclusive: true,
-				IsMinInclusive: true,
-			},
 		}
 
-		if err := rows.Scan(&partition.Range.Min, &partition.Range.Max); err != nil {
+		if err := rows.Scan(&partition.LowerLimit, &partition.UpperLimit); err != nil {
 			return nil, err
 		}
 

internal/app/etl/table_analyzers/postgres.go

@@ -6,16 +6,16 @@ import (
 	"time"
 
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
-	dbwrapper "git.ksdemosapps.com/kylesoda/go-migrate/internal/app/db-wrapper"
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/etl"
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
+	"github.com/jackc/pgx/v5/pgxpool"
 )
 
 type PostgresTableAnalyzer struct {
-	db dbwrapper.DbWrapper
+	db *pgxpool.Pool
 }
 
-func NewPostgresTableAnalyzer(db dbwrapper.DbWrapper) etl.TableAnalyzer {
+func NewPostgresTableAnalyzer(db *pgxpool.Pool) etl.TableAnalyzer {
 	return &PostgresTableAnalyzer{db: db}
 }
 
@@ -125,7 +125,7 @@ func (ta *PostgresTableAnalyzer) QueryColumnTypes(
 	ctx context.Context,
 	tableInfo config.TableInfo,
 ) ([]models.ColumnType, error) {
-	localCtx, cancel := context.WithTimeout(ctx, 1*time.Minute)
+	localCtx, cancel := context.WithTimeout(ctx, 20*time.Second)
 	defer cancel()
 
 	rows, err := ta.db.Query(localCtx, postgresColumnMetadataQuery, tableInfo.Schema, tableInfo.Table)

internal/app/etl/transformers/mssql.go

@@ -74,10 +74,6 @@ func (mssqlTr *MssqlTransformer) ProcessBatch(
 			}
 		}
 
-		if rowValues == nil {
-			continue
-		}
-
 		for _, task := range transformationPlan {
 			val := rowValues[task.Index]
 			if val == nil {

internal/app/etl/types.go

@@ -10,7 +10,7 @@ import (
 )
 
 type Extractor interface {
-	Exec(
+	ProcessPartition(
 		ctx context.Context,
 		tableInfo config.SourceTableInfo,
 		columns []models.ColumnType,
@@ -18,7 +18,20 @@ type Extractor interface {
 		partition models.Partition,
 		indexPrimaryKey int,
 		chBatchesOut chan<- models.Batch,
-	) (int64, error)
+	) (int, error)
+
+	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,
+	)
 }
 
 type TransformerFunc func(any) (any, error)

internal/app/models/main.go

@@ -11,17 +11,12 @@ type Batch struct {
 	RetryCounter int
 }
 
-type PartitionRange struct {
-	Min            int64
-	Max            int64
-	IsMinInclusive bool
-	IsMaxInclusive bool
-}
-
 type Partition struct {
-	Id           uuid.UUID
+	Id                    uuid.UUID
-	ParentId     uuid.UUID
+	ParentId              uuid.UUID
-	Range        PartitionRange
+	LowerLimit            int64
-	HasRange     bool
+	UpperLimit            int64
-	RetryCounter int
+	IsLowerLimitInclusive bool
+	ShouldUseRange        bool
+	RetryCounter          int
 }

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()
+}

scripts/pg-info-test/main.go

@@ -8,32 +8,13 @@ import (
 	"time"
 
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/db"
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgtype"
 	"github.com/jackc/pgx/v5/pgxpool"
 	log "github.com/sirupsen/logrus"
 )
 
-func Connect(ctx context.Context, dbURL string) (*pgxpool.Pool, error) {
-	pool, err := pgxpool.New(ctx, dbURL)
-	if err != nil {
-		return nil, fmt.Errorf("unable to connect to database: %w", err)
-	}
-
-	if err := pool.Ping(ctx); err != nil {
-		pool.Close()
-		return nil, fmt.Errorf("unable to ping database: %w", err)
-	}
-
-	return pool, nil
-}
-
-func Close(pool *pgxpool.Pool) {
-	if pool != nil {
-		pool.Close()
-	}
-}
-
 func main() {
 	log.SetFormatter(&log.TextFormatter{
 		FullTimestamp:   true,
@@ -46,8 +27,8 @@ func main() {
 	ctxSource, cancel := context.WithTimeout(context.Background(), 20*time.Second)
 	defer cancel()
 
-	sourcePool, err := Connect(ctxSource, config.App.SourceDbUrl)
+	sourcePool, err := db.Connect(ctxSource, config.App.SourceDbUrl)
-	defer Close(sourcePool)
+	defer db.Close(sourcePool)
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -56,8 +37,8 @@ func main() {
 	ctxTarget, cancel := context.WithTimeout(context.Background(), 20*time.Second)
 	defer cancel()
 
-	targetPool, err := Connect(ctxTarget, config.App.TargetDbUrl)
+	targetPool, err := db.Connect(ctxTarget, config.App.TargetDbUrl)
-	defer Close(targetPool)
+	defer db.Close(targetPool)
 	if err != nil {
 		log.Fatal(err)
 	}