feat: add source and target database type fields

Add source_db_type and target_db_type to the migration config so the same binary can drive any supported direction.

cmd/go_migrate/main.go

@@ -2,18 +2,17 @@ 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"
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/db-wrapper"
 	"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() {
@@ -33,11 +32,33 @@ func main() {
 
 	log.Info("=== Starting migration ===")
 
-	sourceDb, targetDb, connError := connectToDatabases()
+	var wgConnect errgroup.Group
-	if connError != nil {
+	var sourceDb, targetDb dbwrapper.DbWrapper
-		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()
 
@@ -70,8 +91,8 @@ func main() {
 
 func processMigrationJobs(
 	ctx context.Context,
-	sourceDb *sql.DB,
+	sourceDb dbwrapper.DbWrapper,
-	targetDb *pgxpool.Pool,
+	targetDb dbwrapper.DbWrapper,
 	jobs []config.Job,
 	maxParallelWorkers int,
 ) []JobResult {
@@ -94,7 +115,6 @@ 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)
@@ -107,7 +127,7 @@ func processMigrationJobs(
 				log.Infof("[worker %d] >>> Processing job: %s.%s <<<", i, job.SourceTable.Schema, job.SourceTable.Table)
 				res := processMigrationJob(
 					ctx,
-					targetDbWrapper,
+					targetDb,
 					sourceTableAnalyzer,
 					targetTableAnalyzer,
 					extractor,
@@ -138,3 +158,19 @@ 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

@@ -8,7 +8,7 @@ import (
 
 	"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/db"
+	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/etl/table_analyzers"
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
@@ -18,8 +18,7 @@ import (
 
 func processMigrationJob(
 	ctx context.Context,
-	// sourceDbWrapper db.DbWrapper,
+	targetDbWrapper dbwrapper.DbWrapper,
-	targetDbWrapper db.DbWrapper,
 	sourceTableAnalyzer etl.TableAnalyzer,
 	targetTableAnalyzer etl.TableAnalyzer,
 	extractor etl.Extractor,

config.yaml

@@ -30,6 +30,11 @@ jobs:
       table: users
     pre_sql: 
       - 'SELECT 1'
+    range:
+      min: 1000000
+      max: 2000000
+      is_min_inclusive: false
+      is_max_inclusive: true
 
   - name: analytics_events
     enabled: true

internal/app/config/migration.go

@@ -50,6 +50,12 @@ 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 {
@@ -76,6 +82,8 @@ 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

@@ -103,8 +103,8 @@ func ExtractorErrorHandler(
 			if err.HasLastId {
 				newPartition.ParentId = err.Partition.Id
 				newPartition.Id = uuid.New()
-				newPartition.LowerLimit = err.LastId
+				newPartition.Range.Min = err.LastId
-				newPartition.IsLowerLimitInclusive = false
+				newPartition.Range.IsMinInclusive = false
 			}
 
 			requeueWithBackoff(ctx, delay, func() {

internal/app/db-wrapper/main.go

@@ -0,0 +1,19 @@
+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

@@ -0,0 +1,176 @@
+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

@@ -0,0 +1,128 @@
+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

@@ -0,0 +1,35 @@
+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

@@ -1,30 +0,0 @@
-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

@@ -1,47 +0,0 @@
-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

@@ -1,11 +0,0 @@
-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/mssql.go

@@ -13,16 +13,17 @@ import (
 	"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/etl"
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
 	"github.com/google/uuid"
 )
 
 type MssqlExtractor struct {
-	db *sql.DB
+	db dbwrapper.DbWrapper
 }
 
-func NewMssqlExtractor(db *sql.DB) etl.Extractor {
+func NewMssqlExtractor(db dbwrapper.DbWrapper) etl.Extractor {
 	return &MssqlExtractor{db: db}
 }
 
@@ -73,14 +74,14 @@ func buildExtractQueryMssql(
 func errorFromLastRow(
 	lastRow models.UnknownRowValues,
 	indexPrimaryKey int,
-	partition *models.Partition,
+	partition models.Partition,
 	previousError error,
 ) *custom_errors.ExtractorError {
 	lastIdRawValue := lastRow[indexPrimaryKey]
 
 	lastId, ok := convert.ToInt64(lastIdRawValue)
 	if !ok {
-		currentPartition := *partition
+		currentPartition := partition
 		currentPartition.RetryCounter = 3
 		return &custom_errors.ExtractorError{
 			Partition: currentPartition,
@@ -91,7 +92,7 @@ func errorFromLastRow(
 	}
 
 	return &custom_errors.ExtractorError{
-		Partition: *partition,
+		Partition: partition,
 		HasLastId: true,
 		LastId:    lastId,
 		Msg:       previousError.Error(),
@@ -106,21 +107,21 @@ func (mssqlEx *MssqlExtractor) ProcessPartition(
 	partition models.Partition,
 	indexPrimaryKey int,
 	chBatchesOut chan<- models.Batch,
-	rowsRead *int64,
+) (int, error) {
-) error {
+	query := buildExtractQueryMssql(tableInfo, columns, partition.HasRange, partition.Range.IsMinInclusive)
-	query := buildExtractQueryMssql(tableInfo, columns, partition.ShouldUseRange, partition.IsLowerLimitInclusive)
 
 	var queryArgs []any
-	if partition.ShouldUseRange {
+	if partition.HasRange {
 		queryArgs = append(queryArgs,
-			sql.Named("min", partition.LowerLimit),
+			sql.Named("min", partition.Range.Min),
-			sql.Named("max", partition.UpperLimit),
+			sql.Named("max", partition.Range.Max),
 		)
 	}
 
-	rows, err := mssqlEx.db.QueryContext(ctx, query, queryArgs...)
+	rowsRead := 0
+	rows, err := mssqlEx.db.Query(ctx, query, queryArgs...)
 	if err != nil {
-		return &custom_errors.ExtractorError{Partition: partition, HasLastId: false, Msg: err.Error()}
+		return rowsRead, &custom_errors.ExtractorError{Partition: partition, HasLastId: false, Msg: err.Error()}
 	}
 	defer rows.Close()
 
@@ -136,7 +137,7 @@ func (mssqlEx *MssqlExtractor) ProcessPartition(
 
 		if err := rows.Scan(scanArgs...); err != nil {
 			if len(batchRows) == 0 {
-				return &custom_errors.ExtractorError{Partition: partition, HasLastId: false, Msg: err.Error()}
+				return rowsRead, &custom_errors.ExtractorError{Partition: partition, HasLastId: false, Msg: err.Error()}
 			}
 
 			lastRow := batchRows[len(batchRows)-1]
@@ -144,52 +145,48 @@ func (mssqlEx *MssqlExtractor) ProcessPartition(
 			select {
 			case chBatchesOut <- models.Batch{Id: uuid.New(), PartitionId: partition.Id, Rows: batchRows, RetryCounter: 0}:
 			case <-ctx.Done():
-				return nil
+				return rowsRead, ctx.Err()
 			}
 
-			atomic.AddInt64(rowsRead, int64(len(batchRows)))
+			return rowsRead, errorFromLastRow(lastRow, indexPrimaryKey, partition, err)
-
-			return 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 nil
+				return rowsRead, ctx.Err()
 			}
 
-			atomic.AddInt64(rowsRead, int64(len(batchRows)))
 			batchRows = make([]models.UnknownRowValues, 0, batchSize)
 		}
+
 	}
 
 	if err := rows.Err(); err != nil {
 		if errors.Is(err, ctx.Err()) {
-			return ctx.Err()
+			return rowsRead, ctx.Err()
 		}
 
-		if len(batchRows) == 0 {
+		if len(batchRows) > 0 {
-			return &custom_errors.ExtractorError{Partition: partition, HasLastId: false, Msg: err.Error()}
+			lastRow := batchRows[len(batchRows)-1]
+			return rowsRead, errorFromLastRow(lastRow, indexPrimaryKey, partition, err)
 		}
 
-		lastRow := batchRows[len(batchRows)-1]
+		return rowsRead, &custom_errors.ExtractorError{Partition: partition, HasLastId: false, Msg: err.Error()}
-		return errorFromLastRow(lastRow, indexPrimaryKey, &partition, err)
 	}
 
 	if len(batchRows) > 0 {
 		select {
 		case chBatchesOut <- models.Batch{Id: uuid.New(), PartitionId: partition.Id, Rows: batchRows, RetryCounter: 0}:
 		case <-ctx.Done():
-			return nil
+			return rowsRead, ctx.Err()
 		}
-
-		atomic.AddInt64(rowsRead, int64(len(batchRows)))
 	}
 
-	return nil
+	return rowsRead, nil
 }
 
 func (mssqlEx *MssqlExtractor) Exec(
@@ -234,7 +231,7 @@ func (mssqlEx *MssqlExtractor) Exec(
 				return
 			}
 
-			err := mssqlEx.ProcessPartition(
+			rowsReadResult, err := mssqlEx.ProcessPartition(
 				ctx,
 				tableInfo,
 				columns,
@@ -242,9 +239,12 @@ func (mssqlEx *MssqlExtractor) Exec(
 				partition,
 				indexPrimaryKey,
 				chBatchesOut,
-				rowsRead,
 			)
 
+			if rowsReadResult > 0 {
+				atomic.AddInt64(rowsRead, int64(rowsReadResult))
+			}
+
 			if err != nil {
 				var exError *custom_errors.ExtractorError
 				var jobError *custom_errors.JobError

internal/app/etl/extractors/postgres.go

@@ -6,22 +6,21 @@ import (
 	"fmt"
 	"strings"
 	"sync"
-	"sync/atomic"
 
 	"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 *pgxpool.Pool
+	db dbwrapper.DbWrapper
 }
 
-func NewPostgresExtractor(pool *pgxpool.Pool) etl.Extractor {
+func NewPostgresExtractor(db dbwrapper.DbWrapper) etl.Extractor {
-	return &PostgresExtractor{db: pool}
+	return &PostgresExtractor{db: db}
 }
 
 func buildExtractQueryPostgres(sourceDbInfo config.SourceTableInfo, columns []models.ColumnType) string {
@@ -60,17 +59,17 @@ func (postgresEx *PostgresExtractor) ProcessPartition(
 	partition models.Partition,
 	indexPrimaryKey int,
 	chBatchesOut chan<- models.Batch,
-	rowsRead *int64,
+) (int, error) {
-) error {
 	query := buildExtractQueryPostgres(tableInfo, columns)
 
-	if partition.ShouldUseRange {
+	if partition.HasRange {
-		return errors.New("Batch config not yet supported")
+		return 0, errors.New("Batch config not yet supported")
 	}
 
+	rowsRead := 0
 	rows, err := postgresEx.db.Query(ctx, query)
 	if err != nil {
-		return &custom_errors.ExtractorError{Partition: partition, HasLastId: false, Msg: err.Error()}
+		return rowsRead, &custom_errors.ExtractorError{Partition: partition, HasLastId: false, Msg: err.Error()}
 	}
 	defer rows.Close()
 
@@ -79,8 +78,9 @@ func (postgresEx *PostgresExtractor) ProcessPartition(
 	for rows.Next() {
 		values, err := rows.Values()
 		if err != nil {
-			return errors.New("Unexpected error reading rows from source")
+			return rowsRead, errors.New("Unexpected error reading rows from source")
 		}
+		rowsRead++
 
 		batchRows = append(batchRows, values)
 
@@ -88,29 +88,26 @@ func (postgresEx *PostgresExtractor) ProcessPartition(
 			select {
 			case chBatchesOut <- models.Batch{Id: uuid.New(), PartitionId: partition.Id, Rows: batchRows, RetryCounter: 0}:
 			case <-ctx.Done():
-				return nil
+				return rowsRead, ctx.Err()
 			}
 
-			atomic.AddInt64(rowsRead, int64(len(batchRows)))
 			batchRows = make([]models.UnknownRowValues, 0, batchSize)
 		}
 	}
 
 	if err := rows.Err(); err != nil {
-		return errors.New("Unexpected error reading rows from source")
+		return rowsRead, errors.New("Unexpected error reading rows from source")
 	}
 
 	if len(batchRows) > 0 {
 		select {
 		case chBatchesOut <- models.Batch{Id: uuid.New(), PartitionId: partition.Id, Rows: batchRows, RetryCounter: 0}:
 		case <-ctx.Done():
-			return nil
+			return rowsRead, nil
 		}
-
-		atomic.AddInt64(rowsRead, int64(len(batchRows)))
 	}
 
-	return nil
+	return rowsRead, nil
 }
 
 func (postgresEx *PostgresExtractor) Exec(

internal/app/etl/extractors/types.go

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

internal/app/etl/loaders/postgres.go

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

internal/app/etl/table_analyzers/mssql.go

@@ -8,16 +8,17 @@ 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 *sql.DB
+	db dbwrapper.DbWrapper
 }
 
-func NewMssqlTableAnalyzer(db *sql.DB) etl.TableAnalyzer {
+func NewMssqlTableAnalyzer(db dbwrapper.DbWrapper) etl.TableAnalyzer {
 	return &MssqlTableAnalyzer{db: db}
 }
 
@@ -142,7 +143,7 @@ func (ta *MssqlTableAnalyzer) QueryColumnTypes(
 	localCtx, cancel := context.WithTimeout(ctx, 20*time.Second)
 	defer cancel()
 
-	rows, err := ta.db.QueryContext(localCtx, mssqlColumnMetadataQuery, sql.Named("schema", tableInfo.Schema), sql.Named("table", tableInfo.Table))
+	rows, err := ta.db.Query(localCtx, mssqlColumnMetadataQuery, sql.Named("schema", tableInfo.Schema), sql.Named("table", tableInfo.Table))
 	if err != nil {
 		return nil, err
 	}
@@ -187,7 +188,7 @@ GROUP BY t.name`
 	defer cancel()
 
 	var rowsCount int64
-	err := ta.db.QueryRowContext(ctxTimeout, query, sql.Named("schema", tableInfo.Schema), sql.Named("table", tableInfo.Table)).Scan(&rowsCount)
+	err := ta.db.QueryRow(ctxTimeout, query, sql.Named("schema", tableInfo.Schema), sql.Named("table", tableInfo.Table)).Scan(&rowsCount)
 	if err != nil {
 		return 0, err
 	}
@@ -218,7 +219,7 @@ ORDER BY batch_id`,
 	ctxTimeout, cancel := context.WithTimeout(ctx, time.Second*20)
 	defer cancel()
 
-	rows, err := ta.db.QueryContext(ctxTimeout, query, sql.Named("maxPartitions", maxPartitions))
+	rows, err := ta.db.Query(ctxTimeout, query, sql.Named("maxPartitions", maxPartitions))
 	if err != nil {
 		return nil, err
 	}
@@ -228,13 +229,15 @@ ORDER BY batch_id`,
 
 	for rows.Next() {
 		partition := models.Partition{
-			Id:                    uuid.New(),
+			Id:           uuid.New(),
-			ShouldUseRange:        true,
+			HasRange:     true,
-			RetryCounter:          0,
+			RetryCounter: 0,
-			IsLowerLimitInclusive: true,
+			Range: models.PartitionRange{
+				IsMinInclusive: true,
+			},
 		}
 
-		if err := rows.Scan(&partition.LowerLimit, &partition.UpperLimit); err != nil {
+		if err := rows.Scan(&partition.Range.Min, &partition.Range.Max); 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 *pgxpool.Pool
+	db dbwrapper.DbWrapper
 }
 
-func NewPostgresTableAnalyzer(db *pgxpool.Pool) etl.TableAnalyzer {
+func NewPostgresTableAnalyzer(db dbwrapper.DbWrapper) etl.TableAnalyzer {
 	return &PostgresTableAnalyzer{db: db}
 }
 

internal/app/etl/types.go

@@ -18,8 +18,7 @@ type Extractor interface {
 		partition models.Partition,
 		indexPrimaryKey int,
 		chBatchesOut chan<- models.Batch,
-		rowsRead *int64,
+	) (int, error)
-	) error
 
 	Exec(
 		ctx context.Context,

internal/app/models/main.go

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