refactor: enhance config handling and unify main entrypoint

Unify configuration parsing logic across entrypoints and clean up legacy script-side config paths.

internal/app/config/migration.go

@@ -8,10 +8,12 @@ import (
 )
 
 type RetryConfig struct {
-	Attempts    int `yaml:"attempts"`
-	BaseDelayMs int `yaml:"base_delay_ms"`
-	MaxDelayMs  int `yaml:"max_delay_ms"`
-	MaxJitterMs int `yaml:"max_jitter_ms"`
+	Attempts             int `yaml:"attempts"`
+	BaseDelayMs          int `yaml:"base_delay_ms"`
+	MaxDelayMs           int `yaml:"max_delay_ms"`
+	MaxJitterMs          int `yaml:"max_jitter_ms"`
+	MaxFailedPartitions  int `yaml:"max_failed_partitions"`
+	MaxFailedBatchesLoad int `yaml:"max_failed_batches_load"`
 }
 
 type ToStorageColumnConfig struct {
@@ -25,22 +27,25 @@ type ToStorageConfig struct {
 }
 
 type JobConfig struct {
-	BatchesPerPartition     int         `yaml:"batches_per_partition"`
-	MaxExtractors           int         `yaml:"max_extractors"`
-	ExtractorBatchSize      int         `yaml:"extractor_batch_size"`
-	ExtractorQueueSize      int         `yaml:"extractor_queue_size"`
-	MaxTransformers         int         `yaml:"max_transformers"`
-	TransformerBatchSize    int         `yaml:"transformer_batch_size"`
-	TransformerQueueSize    int         `yaml:"transformer_queue_size"`
-	MaxLoaders              int         `yaml:"max_loaders"`
-	LoaderBatchSize         int         `yaml:"loader_batch_size"`
-	TruncateTarget          bool        `yaml:"truncate_target"`
-	TruncateMethod          string      `yaml:"truncate_method"`
-	MaxPartitionErrrors     int         `yaml:"max_partition_errrors"`
-	MaxExtractorBatchErrors int         `yaml:"max_extractor_batch_errors"`
-	Retry                   RetryConfig `yaml:"retry"`
-	RowsPerPartition        int64
-	ToStorage               ToStorageConfig `yaml:"to_storage"`
+	BatchesPerPartition  int         `yaml:"batches_per_partition"`
+	MaxExtractors        int         `yaml:"max_extractors"`
+	ExtractorBatchSize   int         `yaml:"extractor_batch_size"`
+	ExtractorQueueSize   int         `yaml:"extractor_queue_size"`
+	MaxTransformers      int         `yaml:"max_transformers"`
+	TransformerBatchSize int         `yaml:"transformer_batch_size"`
+	TransformerQueueSize int         `yaml:"transformer_queue_size"`
+	MaxLoaders           int         `yaml:"max_loaders"`
+	LoaderBatchSize      int         `yaml:"loader_batch_size"`
+	TruncateTarget       bool        `yaml:"truncate_target"`
+	TruncateMethod       string      `yaml:"truncate_method"`
+	Retry                RetryConfig `yaml:"retry"`
+	RowsPerPartition     int64
+	ToStorage            ToStorageConfig `yaml:"to_storage"`
+}
+
+type FromJsonItem struct {
+	Column string `yaml:"column"`
+	Field  string `yaml:"field"`
 }
 
 type TableInfo struct {
@@ -49,8 +54,9 @@ type TableInfo struct {
 }
 
 type SourceTableInfo struct {
-	TableInfo  `yaml:",inline"`
-	PrimaryKey string `yaml:"primary_key"`
+	TableInfo       `yaml:",inline"`
+	PrimaryKey      string         `yaml:"primary_key"`
+	FromJsonColumns []FromJsonItem `yaml:"from_json"`
 }
 
 type TargetTableInfo struct {

internal/app/custom_errors/backoff.go

@@ -1,7 +1,6 @@
 package custom_errors
 
 import (
-	"context"
 	"math/rand"
 	"time"
 )
@@ -40,22 +39,3 @@ func ComputeBackoffDelay(retryCounter int, baseDelayMs int, maxDelayMs int, maxJ
 
 	return delay
 }
-
-func requeueWithBackoff(ctx context.Context, delay time.Duration, enqueue func()) {
-	if delay <= 0 {
-		enqueue()
-		return
-	}
-
-	go func() {
-		timer := time.NewTimer(delay)
-		defer timer.Stop()
-
-		select {
-		case <-ctx.Done():
-			return
-		case <-timer.C:
-			enqueue()
-		}
-	}()
-}

internal/app/custom_errors/loader.error.go

@@ -1,107 +0,0 @@
-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"
-)
-
-type LoaderError struct {
-	Batch models.Batch
-	Msg   string
-}
-
-func (e *LoaderError) Error() string {
-	return e.Msg
-}
-
-func LoaderErrorHandler(
-	ctx context.Context,
-	retryConfig config.RetryConfig,
-	maxChunkErrors int,
-	chErrorsIn <-chan LoaderError,
-	chBatchesOut chan<- models.Batch,
-	chJobErrorsOut chan<- JobError,
-	wgActiveBatches *sync.WaitGroup,
-) {
-	definitiveErrors := 0
-
-	for {
-		if ctx.Err() != nil {
-			return
-		}
-
-		select {
-		case <-ctx.Done():
-			return
-
-		case err, ok := <-chErrorsIn:
-			if !ok {
-				return
-			}
-
-			if err.Batch.RetryCounter >= retryConfig.Attempts {
-				wgActiveBatches.Done()
-				definitiveErrors++
-				jobError := JobError{
-					ShouldCancelJob: false,
-					Msg:             fmt.Sprintf("Batch %v reached max retries (%d)", err.Batch.Id, retryConfig.Attempts),
-					Prev:            &err,
-				}
-
-				select {
-				case chJobErrorsOut <- jobError:
-				case <-ctx.Done():
-					return
-				}
-
-				if maxChunkErrors > 0 && definitiveErrors >= maxChunkErrors {
-					fatalError := JobError{
-						ShouldCancelJob: true,
-						Msg:             fmt.Sprintf("Chunk error limit reached (%d)", maxChunkErrors),
-						Prev:            &err,
-					}
-
-					select {
-					case chJobErrorsOut <- fatalError:
-					case <-ctx.Done():
-						return
-					}
-				}
-
-				continue
-			} else {
-				jobError := JobError{
-					ShouldCancelJob: false,
-					Msg:             fmt.Sprintf("Temporal error in batch %v (retries: %d)", err.Batch.Id, err.Batch.RetryCounter),
-					Prev:            &err,
-				}
-
-				select {
-				case chJobErrorsOut <- jobError:
-				case <-ctx.Done():
-					return
-				}
-			}
-
-			err.Batch.RetryCounter++
-			delay := ComputeBackoffDelay(
-				err.Batch.RetryCounter,
-				retryConfig.BaseDelayMs,
-				retryConfig.MaxDelayMs,
-				retryConfig.MaxJitterMs,
-			)
-
-			requeueWithBackoff(ctx, delay, func() {
-				select {
-				case chBatchesOut <- err.Batch:
-				case <-ctx.Done():
-					return
-				}
-			})
-		}
-	}
-}

internal/app/custom_errors/types.go

@@ -14,3 +14,12 @@ type ExtractorError struct {
 func (e *ExtractorError) Error() string {
 	return e.Msg
 }
+
+type LoaderError struct {
+	Batch models.Batch
+	Msg   string
+}
+
+func (e *LoaderError) Error() string {
+	return e.Msg
+}

internal/app/db-wrapper/mssql.go

@@ -6,7 +6,9 @@ import (
 	"fmt"
 	"strings"
 
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
 	dbdialects "git.ksdemosapps.com/kylesoda/go-migrate/internal/app/db-wrapper/db_dialects"
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
 	mssql "github.com/microsoft/go-mssqldb"
 )
 
@@ -177,25 +179,53 @@ func (mw *mssqlDbWrapper) SaveMassive(ctx context.Context, schema string, table
 	return rowsAffected, nil
 }
 
-func (mw *mssqlDbWrapper) QueryFromObject(ctx context.Context, q ExtractionQuery) (RowsResult, error) {
+func buildExtractQueryMssql(q ExtractionQuery) (string, error) {
 	var sbQuery strings.Builder
 
 	sbQuery.WriteString("SELECT ")
 
-	if len(q.Columns) == 0 {
-		sbQuery.WriteString("*")
-	} else {
-		for i, col := range q.Columns {
-			fmt.Fprintf(&sbQuery, "[%s]", col.Name())
+	hasRegularColumns := len(q.Columns) > 0
+	hasJsonColumns := len(q.FromJsonColumns) > 0
 
+	resolvedJson := make(map[string][]config.FromJsonItem, len(q.FromJsonColumns))
+	if hasJsonColumns {
+		for _, jsonConfig := range q.FromJsonColumns {
+			actualColumnName, err := findColumnByPattern(q.Columns, jsonConfig.Column)
+			if err != nil {
+				return "", err
+			}
+			resolvedJson[actualColumnName] = append(resolvedJson[actualColumnName], jsonConfig)
+		}
+	}
+
+	selectParts := make([]string, 0, len(q.Columns)+len(q.FromJsonColumns))
+	if hasRegularColumns {
+		for _, col := range q.Columns {
+			jsonConfigs, isJsonColumn := resolvedJson[col.Name()]
+			if isJsonColumn {
+				for _, jsonConfig := range jsonConfigs {
+					jsonPath := buildJsonPathMssql(jsonConfig.Field)
+					jsonExpr := fmt.Sprintf("JSON_VALUE([%s], '%s') AS [%s]", col.Name(), jsonPath, col.Name())
+					selectParts = append(selectParts, jsonExpr)
+				}
+				continue
+			}
+
+			colExpr := fmt.Sprintf("[%s]", col.Name())
 			switch col.Type() {
 			case "GEOMETRY":
-				fmt.Fprintf(&sbQuery, ".STAsBinary() AS [%s]", col.Name())
+				colExpr = fmt.Sprintf("[%s].STAsBinary() AS [%s]", col.Name(), col.Name())
 			}
+			selectParts = append(selectParts, colExpr)
+		}
+	} else if !hasJsonColumns {
+		selectParts = append(selectParts, "*")
+	}
 
-			if i < len(q.Columns)-1 {
-				sbQuery.WriteString(", ")
-			}
+	for i, part := range selectParts {
+		sbQuery.WriteString(part)
+		if i < len(selectParts)-1 {
+			sbQuery.WriteString(", ")
 		}
 	}
 
@@ -231,7 +261,37 @@ func (mw *mssqlDbWrapper) QueryFromObject(ctx context.Context, q ExtractionQuery
 
 	fmt.Fprintf(&sbQuery, " ORDER BY [%s] ASC", q.PrimaryKey)
 
-	queryString := sbQuery.String()
+	return sbQuery.String(), nil
+}
+
+func findColumnByPattern(columns []models.ColumnType, pattern string) (string, error) {
+	if pattern == "" {
+		return "", fmt.Errorf("column pattern cannot be empty")
+	}
+
+	if before, ok := strings.CutSuffix(pattern, "*"); ok {
+		prefix := before
+		for _, col := range columns {
+			if strings.HasPrefix(col.Name(), prefix) {
+				return col.Name(), nil
+			}
+		}
+		return "", fmt.Errorf("no column found matching pattern '%s'", pattern)
+	}
+
+	for _, col := range columns {
+		if col.Name() == pattern {
+			return col.Name(), nil
+		}
+	}
+	return "", fmt.Errorf("column '%s' not found in table columns", pattern)
+}
+
+func (mw *mssqlDbWrapper) QueryFromObject(ctx context.Context, q ExtractionQuery) (RowsResult, error) {
+	queryString, err := buildExtractQueryMssql(q)
+	if err != nil {
+		return nil, err
+	}
 
 	// logrus.Debugf("Query: %s", queryString)
 
@@ -247,3 +307,11 @@ func (mw *mssqlDbWrapper) QueryFromObject(ctx context.Context, q ExtractionQuery
 
 	return mw.Query(ctx, queryString, queryArgs...)
 }
+
+func buildJsonPathMssql(field string) string {
+	if len(field) > 0 && field[0] == '.' {
+		field = field[1:]
+	}
+
+	return "$." + field
+}

internal/app/db-wrapper/mssql_test.go

@@ -0,0 +1,396 @@
+package dbwrapper
+
+import (
+	"strings"
+	"testing"
+
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
+)
+
+func TestBuildExtractQueryMssql_NoJsonColumns(t *testing.T) {
+	q := ExtractionQuery{
+		Schema:     "dbo",
+		Table:      "Users",
+		PrimaryKey: "ID",
+		Columns: []models.ColumnType{
+			models.NewColumnType("ID", false, false, "INT", "int", "INT", false, 0, 0, 0),
+			models.NewColumnType("Name", true, false, "VARCHAR", "varchar", "VARCHAR", true, 255, 0, 0),
+		},
+		FromJsonColumns: []config.FromJsonItem{},
+		LowerLimit:      ExtractorQueryLimit{IsValid: false},
+		UpperLimit:      ExtractorQueryLimit{IsValid: false},
+	}
+
+	query, err := buildExtractQueryMssql(q)
+	if err != nil {
+		t.Fatalf("Expected no error, got: %v", err)
+	}
+
+	if !strings.Contains(query, "SELECT [ID], [Name]") {
+		t.Errorf("Expected columns in query, got: %s", query)
+	}
+
+	if !strings.Contains(query, "FROM [dbo].[Users]") {
+		t.Errorf("Expected FROM clause, got: %s", query)
+	}
+
+	if !strings.Contains(query, "ORDER BY [ID] ASC") {
+		t.Errorf("Expected ORDER BY clause, got: %s", query)
+	}
+}
+
+func TestBuildExtractQueryMssql_WithJsonColumns_ExactColumnMatch(t *testing.T) {
+	// Test that the actual column name is used as alias, not a generated one
+	q := ExtractionQuery{
+		Schema:     "dbo",
+		Table:      "Events",
+		PrimaryKey: "EventID",
+		Columns: []models.ColumnType{
+			models.NewColumnType("EventID", false, false, "INT", "int", "INT", false, 0, 0, 0),
+			models.NewColumnType("EventData", true, false, "VARCHAR", "varchar", "VARCHAR", true, 500, 0, 0),
+		},
+		FromJsonColumns: []config.FromJsonItem{
+			{Column: "EventData", Field: ".userId"},
+			{Column: "EventData", Field: ".timestamp"},
+		},
+		LowerLimit: ExtractorQueryLimit{IsValid: false},
+		UpperLimit: ExtractorQueryLimit{IsValid: false},
+	}
+
+	query, err := buildExtractQueryMssql(q)
+	if err != nil {
+		t.Fatalf("Expected no error, got: %v", err)
+	}
+
+	if !strings.HasPrefix(query, "SELECT [EventID], JSON_VALUE([EventData], '$.userId') AS [EventData], JSON_VALUE([EventData], '$.timestamp') AS [EventData]") {
+		t.Errorf("Expected JSON columns to replace EventData in-order, got: %s", query)
+	}
+
+	if strings.Contains(query, "SELECT [EventID], [EventData]") {
+		t.Errorf("Expected EventData to be replaced by JSON extraction, got: %s", query)
+	}
+
+	// Alias should be exactly "EventData", not "EventData_userId"
+	if !strings.Contains(query, "JSON_VALUE([EventData], '$.userId') AS [EventData]") {
+		t.Errorf("Expected JSON alias to be [EventData], got: %s", query)
+	}
+
+	if !strings.Contains(query, "JSON_VALUE([EventData], '$.timestamp') AS [EventData]") {
+		t.Errorf("Expected JSON alias to be [EventData], got: %s", query)
+	}
+
+	// Should have comma separating them
+	if !strings.Contains(query, "JSON_VALUE([EventData], '$.userId') AS [EventData], JSON_VALUE([EventData], '$.timestamp') AS [EventData]") {
+		t.Errorf("Expected comma-separated JSON values, got: %s", query)
+	}
+}
+
+func TestBuildExtractQueryMssql_WithWildcardPattern(t *testing.T) {
+	// Test that wildcard pattern matching finds the correct column
+	q := ExtractionQuery{
+		Schema:     "dbo",
+		Table:      "Events",
+		PrimaryKey: "ID",
+		Columns: []models.ColumnType{
+			models.NewColumnType("ID", false, false, "INT", "int", "INT", false, 0, 0, 0),
+			models.NewColumnType("NodeMetadata", true, false, "VARCHAR", "varchar", "VARCHAR", true, 500, 0, 0),
+		},
+		FromJsonColumns: []config.FromJsonItem{
+			{Column: "NodeMeta*", Field: ".id"},
+		},
+		LowerLimit: ExtractorQueryLimit{IsValid: false},
+		UpperLimit: ExtractorQueryLimit{IsValid: false},
+	}
+
+	query, err := buildExtractQueryMssql(q)
+	if err != nil {
+		t.Fatalf("Expected no error, got: %v", err)
+	}
+
+	// Should find "NodeMetadata" from pattern "NodeMeta*" and use it as alias
+	if !strings.Contains(query, "JSON_VALUE([NodeMetadata], '$.id') AS [NodeMetadata]") {
+		t.Errorf("Expected to find and use NodeMetadata column by pattern, got: %s", query)
+	}
+
+	if strings.Contains(query, "SELECT [ID], [NodeMetadata]") {
+		t.Errorf("Expected NodeMetadata to be replaced by JSON extraction, got: %s", query)
+	}
+}
+
+func TestBuildExtractQueryMssql_ColumnNotFound_Error(t *testing.T) {
+	// Test that an error is returned when column is not found
+	q := ExtractionQuery{
+		Schema:     "dbo",
+		Table:      "Events",
+		PrimaryKey: "ID",
+		Columns: []models.ColumnType{
+			models.NewColumnType("ID", false, false, "INT", "int", "INT", false, 0, 0, 0),
+		},
+		FromJsonColumns: []config.FromJsonItem{
+			{Column: "NonExistentColumn", Field: ".id"},
+		},
+		LowerLimit: ExtractorQueryLimit{IsValid: false},
+		UpperLimit: ExtractorQueryLimit{IsValid: false},
+	}
+
+	query, err := buildExtractQueryMssql(q)
+	if err == nil {
+		t.Fatalf("Expected error for missing column, got no error. Query: %s", query)
+	}
+
+	if !strings.Contains(err.Error(), "NonExistentColumn") {
+		t.Errorf("Expected error message to contain column name, got: %v", err)
+	}
+}
+
+func TestBuildExtractQueryMssql_WildcardPatternNotMatched_Error(t *testing.T) {
+	// Test that an error is returned when wildcard pattern doesn't match any column
+	q := ExtractionQuery{
+		Schema:     "dbo",
+		Table:      "Events",
+		PrimaryKey: "ID",
+		Columns: []models.ColumnType{
+			models.NewColumnType("ID", false, false, "INT", "int", "INT", false, 0, 0, 0),
+			models.NewColumnType("EventData", true, false, "VARCHAR", "varchar", "VARCHAR", true, 500, 0, 0),
+		},
+		FromJsonColumns: []config.FromJsonItem{
+			{Column: "NonMatching*", Field: ".id"},
+		},
+		LowerLimit: ExtractorQueryLimit{IsValid: false},
+		UpperLimit: ExtractorQueryLimit{IsValid: false},
+	}
+
+	query, err := buildExtractQueryMssql(q)
+	if err == nil {
+		t.Fatalf("Expected error for non-matching wildcard pattern, got no error. Query: %s", query)
+	}
+
+	if !strings.Contains(err.Error(), "NonMatching*") {
+		t.Errorf("Expected error message to contain pattern, got: %v", err)
+	}
+}
+
+func TestBuildExtractQueryMssql_NestedJsonFields(t *testing.T) {
+	q := ExtractionQuery{
+		Schema:     "dbo",
+		Table:      "Data",
+		PrimaryKey: "ID",
+		Columns: []models.ColumnType{
+			models.NewColumnType("ID", false, false, "INT", "int", "INT", false, 0, 0, 0),
+			models.NewColumnType("NodeData", true, false, "VARCHAR", "varchar", "VARCHAR", true, 500, 0, 0),
+		},
+		FromJsonColumns: []config.FromJsonItem{
+			{Column: "NodeData", Field: ".user.name"},
+			{Column: "NodeData", Field: ".user.email"},
+		},
+		LowerLimit: ExtractorQueryLimit{IsValid: false},
+		UpperLimit: ExtractorQueryLimit{IsValid: false},
+	}
+
+	query, err := buildExtractQueryMssql(q)
+	if err != nil {
+		t.Fatalf("Expected no error, got: %v", err)
+	}
+
+	if !strings.Contains(query, "JSON_VALUE([NodeData], '$.user.name') AS [NodeData]") {
+		t.Errorf("Expected nested JSON path for user.name, got: %s", query)
+	}
+
+	if !strings.Contains(query, "JSON_VALUE([NodeData], '$.user.email') AS [NodeData]") {
+		t.Errorf("Expected nested JSON path for user.email, got: %s", query)
+	}
+
+	if strings.Contains(query, "SELECT [ID], [NodeData]") {
+		t.Errorf("Expected NodeData to be replaced by JSON extraction, got: %s", query)
+	}
+}
+
+func TestBuildExtractQueryMssql_WithRangeLimits(t *testing.T) {
+	q := ExtractionQuery{
+		Schema:     "dbo",
+		Table:      "Products",
+		PrimaryKey: "ProductID",
+		Columns: []models.ColumnType{
+			models.NewColumnType("ProductID", false, false, "INT", "int", "INT", false, 0, 0, 0),
+			models.NewColumnType("Details", true, false, "VARCHAR", "varchar", "VARCHAR", true, 500, 0, 0),
+		},
+		FromJsonColumns: []config.FromJsonItem{
+			{Column: "Details", Field: ".price"},
+		},
+		LowerLimit: ExtractorQueryLimit{IsValid: true, IsInclusive: true, Value: 100},
+		UpperLimit: ExtractorQueryLimit{IsValid: true, IsInclusive: false, Value: 500},
+	}
+
+	query, err := buildExtractQueryMssql(q)
+	if err != nil {
+		t.Fatalf("Expected no error, got: %v", err)
+	}
+
+	if !strings.Contains(query, "WHERE [ProductID] >= @min") {
+		t.Errorf("Expected WHERE clause with >=, got: %s", query)
+	}
+
+	if !strings.Contains(query, "[ProductID] < @max") {
+		t.Errorf("Expected upper limit with <, got: %s", query)
+	}
+
+	if !strings.Contains(query, "JSON_VALUE([Details], '$.price') AS [Details]") {
+		t.Errorf("Expected JSON_VALUE for Details, got: %s", query)
+	}
+
+	if strings.Contains(query, "SELECT [ProductID], [Details]") {
+		t.Errorf("Expected Details to be replaced by JSON extraction, got: %s", query)
+	}
+}
+
+func TestBuildJsonPathMssql(t *testing.T) {
+	tests := []struct {
+		input    string
+		expected string
+	}{
+		{".id", "$.id"},
+		{"id", "$.id"},
+		{".user.name", "$.user.name"},
+		{"user.name", "$.user.name"},
+		{".location.coordinates.lat", "$.location.coordinates.lat"},
+		{"", "$."},
+	}
+
+	for _, tt := range tests {
+		result := buildJsonPathMssql(tt.input)
+		if result != tt.expected {
+			t.Errorf("buildJsonPathMssql(%q) = %q, want %q", tt.input, result, tt.expected)
+		}
+	}
+}
+
+func TestFindColumnByPattern_ExactMatch(t *testing.T) {
+	columns := []models.ColumnType{
+		models.NewColumnType("ID", false, false, "INT", "int", "INT", false, 0, 0, 0),
+		models.NewColumnType("Metadata", true, false, "VARCHAR", "varchar", "VARCHAR", true, 500, 0, 0),
+		models.NewColumnType("EventData", true, false, "VARCHAR", "varchar", "VARCHAR", true, 500, 0, 0),
+	}
+
+	result, err := findColumnByPattern(columns, "Metadata")
+	if err != nil {
+		t.Fatalf("Expected no error, got: %v", err)
+	}
+
+	if result != "Metadata" {
+		t.Errorf("Expected 'Metadata', got '%s'", result)
+	}
+}
+
+func TestFindColumnByPattern_WildcardMatch(t *testing.T) {
+	columns := []models.ColumnType{
+		models.NewColumnType("ID", false, false, "INT", "int", "INT", false, 0, 0, 0),
+		models.NewColumnType("NodeMetadata", true, false, "VARCHAR", "varchar", "VARCHAR", true, 500, 0, 0),
+		models.NewColumnType("EventData", true, false, "VARCHAR", "varchar", "VARCHAR", true, 500, 0, 0),
+	}
+
+	result, err := findColumnByPattern(columns, "NodeMeta*")
+	if err != nil {
+		t.Fatalf("Expected no error, got: %v", err)
+	}
+
+	if result != "NodeMetadata" {
+		t.Errorf("Expected 'NodeMetadata', got '%s'", result)
+	}
+}
+
+func TestFindColumnByPattern_NotFound(t *testing.T) {
+	columns := []models.ColumnType{
+		models.NewColumnType("ID", false, false, "INT", "int", "INT", false, 0, 0, 0),
+		models.NewColumnType("Metadata", true, false, "VARCHAR", "varchar", "VARCHAR", true, 500, 0, 0),
+	}
+
+	result, err := findColumnByPattern(columns, "NonExistent")
+	if err == nil {
+		t.Fatalf("Expected error, got no error. Result: %s", result)
+	}
+
+	if !strings.Contains(err.Error(), "NonExistent") {
+		t.Errorf("Expected error to contain column name, got: %v", err)
+	}
+}
+
+func TestFindColumnByPattern_WildcardNotFound(t *testing.T) {
+	columns := []models.ColumnType{
+		models.NewColumnType("ID", false, false, "INT", "int", "INT", false, 0, 0, 0),
+		models.NewColumnType("Metadata", true, false, "VARCHAR", "varchar", "VARCHAR", true, 500, 0, 0),
+	}
+
+	result, err := findColumnByPattern(columns, "Event*")
+	if err == nil {
+		t.Fatalf("Expected error, got no error. Result: %s", result)
+	}
+
+	if !strings.Contains(err.Error(), "Event*") {
+		t.Errorf("Expected error to contain pattern, got: %v", err)
+	}
+}
+
+func TestBuildExtractQueryMssql_OnlyJsonColumns(t *testing.T) {
+	// Test when all columns are used via JSON extraction
+	q := ExtractionQuery{
+		Schema:     "dbo",
+		Table:      "Data",
+		PrimaryKey: "ID",
+		Columns: []models.ColumnType{
+			models.NewColumnType("ID", false, false, "INT", "int", "INT", false, 0, 0, 0),
+			models.NewColumnType("JsonData", true, false, "VARCHAR", "varchar", "VARCHAR", true, 500, 0, 0),
+		},
+		FromJsonColumns: []config.FromJsonItem{
+			{Column: "JsonData", Field: ".field1"},
+		},
+		LowerLimit: ExtractorQueryLimit{IsValid: false},
+		UpperLimit: ExtractorQueryLimit{IsValid: false},
+	}
+
+	query, err := buildExtractQueryMssql(q)
+	if err != nil {
+		t.Fatalf("Expected no error, got: %v", err)
+	}
+
+	if !strings.HasPrefix(query, "SELECT [ID], JSON_VALUE([JsonData], '$.field1') AS [JsonData]") {
+		t.Errorf("Expected JsonData to be replaced by JSON extraction, got: %s", query)
+	}
+
+	if strings.Contains(query, "SELECT [ID], [JsonData]") {
+		t.Errorf("Expected JsonData to be excluded from raw selection, got: %s", query)
+	}
+}
+
+func TestBuildExtractQueryMssql_JsonColumnsReplaceInOrder(t *testing.T) {
+	q := ExtractionQuery{
+		Schema:     "dbo",
+		Table:      "Users",
+		PrimaryKey: "UserID",
+		Columns: []models.ColumnType{
+			models.NewColumnType("UserID", false, false, "INT", "int", "INT", false, 0, 0, 0),
+			models.NewColumnType("Name", true, false, "VARCHAR", "varchar", "VARCHAR", false, 255, 0, 0),
+			models.NewColumnType("Email", true, false, "VARCHAR", "varchar", "VARCHAR", false, 255, 0, 0),
+			models.NewColumnType("Metadata", true, false, "NVARCHAR", "nvarchar", "NVARCHAR", true, 4000, 0, 0),
+			models.NewColumnType("Profile", true, false, "NVARCHAR", "nvarchar", "NVARCHAR", true, 4000, 0, 0),
+			models.NewColumnType("Settings", true, false, "NVARCHAR", "nvarchar", "NVARCHAR", true, 4000, 0, 0),
+		},
+		FromJsonColumns: []config.FromJsonItem{
+			{Column: "Metadata", Field: ".id"},
+			{Column: "Profile", Field: ".id"},
+			{Column: "Settings", Field: ".id"},
+		},
+		LowerLimit: ExtractorQueryLimit{IsValid: false},
+		UpperLimit: ExtractorQueryLimit{IsValid: false},
+	}
+
+	query, err := buildExtractQueryMssql(q)
+	if err != nil {
+		t.Fatalf("Expected no error, got: %v", err)
+	}
+
+	expected := "SELECT [UserID], [Name], [Email], JSON_VALUE([Metadata], '$.id') AS [Metadata], JSON_VALUE([Profile], '$.id') AS [Profile], JSON_VALUE([Settings], '$.id') AS [Settings] FROM [dbo].[Users] ORDER BY [UserID] ASC"
+	if query != expected {
+		t.Errorf("Unexpected query.\nExpected: %s\nGot: %s", expected, query)
+	}
+}

internal/app/db-wrapper/types.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
 )
 
@@ -33,12 +34,13 @@ type ExtractorQueryLimit struct {
 }
 
 type ExtractionQuery struct {
-	Schema     string
-	Table      string
-	PrimaryKey string
-	Columns    []models.ColumnType
-	LowerLimit ExtractorQueryLimit
-	UpperLimit ExtractorQueryLimit
+	Schema          string
+	Table           string
+	PrimaryKey      string
+	Columns         []models.ColumnType
+	LowerLimit      ExtractorQueryLimit
+	UpperLimit      ExtractorQueryLimit
+	FromJsonColumns []config.FromJsonItem
 }
 
 type DbWrapper interface {

internal/app/etl/extractors/consume.go

@@ -25,6 +25,8 @@ func (ex *GenericExtractor) Consume(
 	chErrorsOut chan<- custom_errors.JobError,
 	wgActivePartitions *sync.WaitGroup,
 	rowsRead *int64,
+	failedPartitionsCount *int32,
+	fromJsonColumns []config.FromJsonItem,
 ) {
 	indexPrimaryKey := slices.IndexFunc(columns, func(col models.ColumnType) bool {
 		return strings.EqualFold(col.Name(), tableInfo.PrimaryKey)
@@ -65,16 +67,18 @@ func (ex *GenericExtractor) Consume(
 				indexPrimaryKey,
 				retryConfig,
 				chBatchesOut,
+				fromJsonColumns,
 			)
 			wgActivePartitions.Done()
 
 			if rowsReadResult > 0 {
 				current := atomic.LoadInt64(rowsRead)
-				logrus.Debugf("Rows read: +%v [current=%v] (%s.%s)", rowsReadResult, current, tableInfo.Schema, tableInfo.Table)
+				logrus.Debugf("Rows read (partition extracted): +%v [current=%v] (%s.%s)", rowsReadResult, current, tableInfo.Schema, tableInfo.Table)
 				atomic.AddInt64(rowsRead, int64(rowsReadResult))
 			}
 
 			if err != nil {
+				atomic.AddInt32(failedPartitionsCount, 1)
 				if jobError, ok := errors.AsType[*custom_errors.JobError](err); ok {
 					select {
 					case <-ctx.Done():
@@ -88,6 +92,16 @@ func (ex *GenericExtractor) Consume(
 					case chErrorsOut <- custom_errors.JobError{ShouldCancelJob: false, Msg: err.Error(), Prev: err}:
 					}
 				}
+
+				currentFPCount := atomic.LoadInt32(failedPartitionsCount)
+				if currentFPCount > int32(retryConfig.MaxFailedPartitions) {
+					select {
+					case <-ctx.Done():
+						return
+					case chErrorsOut <- custom_errors.JobError{ShouldCancelJob: true, Msg: "Max failed partitions reached"}:
+						return
+					}
+				}
 			}
 		}
 	}

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

@@ -22,6 +22,7 @@ func (ex *GenericExtractor) ProcessPartitionWithRetries(
 	indexPrimaryKey int,
 	retryConfig config.RetryConfig,
 	chBatchesOut chan<- models.Batch,
+	fromJsonColumns []config.FromJsonItem,
 ) (int64, error) {
 	var totalRowsRead int64
 	currentParitition := partition
@@ -35,6 +36,7 @@ func (ex *GenericExtractor) ProcessPartitionWithRetries(
 			currentParitition,
 			indexPrimaryKey,
 			chBatchesOut,
+			fromJsonColumns,
 		)
 		// logrus.Debugf("Partition %v finished processing (%s.%s)", partition.Id, tableInfo.Schema, tableInfo.Table)
 		totalRowsRead += rowsRead
@@ -48,7 +50,7 @@ func (ex *GenericExtractor) ProcessPartitionWithRetries(
 
 			if currentParitition.RetryCounter >= retryConfig.Attempts {
 				return totalRowsRead, &custom_errors.JobError{
-					Msg:  fmt.Sprintf("Partition %v reached max retries", exError.Partition.Id),
+					Msg:  fmt.Sprintf("Partition %v reached max retries (%d)", currentParitition.Id, currentParitition.RetryCounter),
 					Prev: err,
 				}
 			}

internal/app/etl/extractors/process.go

@@ -47,6 +47,7 @@ func (ex *GenericExtractor) ProcessPartition(
 	partition models.Partition,
 	indexPrimaryKey int,
 	chBatchesOut chan<- models.Batch,
+	fromJsonColumns []config.FromJsonItem,
 ) (int64, error) {
 	query := dbwrapper.ExtractionQuery{
 		Schema:     tableInfo.Schema,
@@ -63,6 +64,7 @@ func (ex *GenericExtractor) ProcessPartition(
 			IsInclusive: partition.Range.IsMaxInclusive,
 			Value:       partition.Range.Max,
 		},
+		FromJsonColumns: fromJsonColumns,
 	}
 
 	// logrus.Debugf("Processing partition: %+v (%s.%s)", query, tableInfo.Schema, tableInfo.Table)

internal/app/etl/loaders/consume.go

@@ -0,0 +1,154 @@
+package loaders
+
+import (
+	"context"
+	"errors"
+	"sync"
+	"sync/atomic"
+
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/custom_errors"
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
+	"github.com/google/uuid"
+	"github.com/sirupsen/logrus"
+)
+
+func (gl *GenericLoader) Consume(
+	ctx context.Context,
+	tableInfo config.TargetTableInfo,
+	columns []models.ColumnType,
+	retryConfig config.RetryConfig,
+	batchSize int,
+	chBatchesIn <-chan models.Batch,
+	chErrorsOut chan<- custom_errors.JobError,
+	wgActiveBatches *sync.WaitGroup,
+	rowsLoaded *int64,
+	failedBatchesCount *int32,
+) {
+	colNames := mapSlice(columns, func(col models.ColumnType) string {
+		return col.Name()
+	})
+
+	var accRows []models.UnknownRowValues
+	var parentBatchesId []uuid.UUID
+	pendingDone := 0
+
+	defer func() {
+		for range pendingDone {
+			wgActiveBatches.Done()
+		}
+	}()
+
+	flush := func() bool {
+		if len(accRows) == 0 {
+			return true
+		}
+		count := len(parentBatchesId)
+		superBatch := models.Batch{
+			Id:              uuid.New(),
+			ParentBatchesId: parentBatchesId,
+			Rows:            accRows,
+		}
+		processedRows, err := gl.ProcessBatchWithRetries(ctx, tableInfo, colNames, retryConfig, superBatch)
+		for range count {
+			wgActiveBatches.Done()
+		}
+		pendingDone -= count
+		accRows = nil
+		parentBatchesId = nil
+
+		if err != nil {
+			atomic.AddInt32(failedBatchesCount, 1)
+			if jobError, ok := errors.AsType[*custom_errors.JobError](err); ok {
+				select {
+				case <-ctx.Done():
+					return false
+				case chErrorsOut <- *jobError:
+				}
+			} else {
+				select {
+				case <-ctx.Done():
+					return false
+				case chErrorsOut <- custom_errors.JobError{ShouldCancelJob: false, Msg: err.Error(), Prev: err}:
+				}
+			}
+
+			if atomic.LoadInt32(failedBatchesCount) > int32(retryConfig.MaxFailedBatchesLoad) {
+				select {
+				case <-ctx.Done():
+				case chErrorsOut <- custom_errors.JobError{ShouldCancelJob: true, Msg: "Max failed batches (load) reached"}:
+				}
+				return false
+			}
+			return true
+		}
+
+		current := atomic.LoadInt64(rowsLoaded)
+		logrus.Debugf("Rows loaded (batch loaded): +%v [current=%v] (%s.%s)", processedRows, current, tableInfo.Schema, tableInfo.Table)
+		atomic.AddInt64(rowsLoaded, int64(processedRows))
+		return true
+	}
+
+	for {
+		if ctx.Err() != nil {
+			return
+		}
+
+		select {
+		case <-ctx.Done():
+			return
+		case batch, ok := <-chBatchesIn:
+			if !ok {
+				flush()
+				return
+			}
+
+			if batchSize <= 0 {
+				processedRows, err := gl.ProcessBatchWithRetries(ctx, tableInfo, colNames, retryConfig, batch)
+				wgActiveBatches.Done()
+
+				if err != nil {
+					atomic.AddInt32(failedBatchesCount, 1)
+					if jobError, ok := errors.AsType[*custom_errors.JobError](err); ok {
+						select {
+						case <-ctx.Done():
+							return
+						case chErrorsOut <- *jobError:
+						}
+					} else {
+						select {
+						case <-ctx.Done():
+							return
+						case chErrorsOut <- custom_errors.JobError{ShouldCancelJob: false, Msg: err.Error(), Prev: err}:
+						}
+					}
+
+					if atomic.LoadInt32(failedBatchesCount) > int32(retryConfig.MaxFailedBatchesLoad) {
+						select {
+						case <-ctx.Done():
+							return
+						case chErrorsOut <- custom_errors.JobError{ShouldCancelJob: true, Msg: "Max failed batches (load) reached"}:
+							return
+						}
+					}
+					continue
+				}
+
+				current := atomic.LoadInt64(rowsLoaded)
+				logrus.Debugf("Rows loaded: +%v [current=%v] (%s.%s)", processedRows, current, tableInfo.Schema, tableInfo.Table)
+				atomic.AddInt64(rowsLoaded, int64(processedRows))
+				continue
+			}
+
+			pendingDone++
+			accRows = append(accRows, batch.Rows...)
+			parentBatchesId = append(parentBatchesId, batch.Id)
+
+			if len(accRows) >= batchSize {
+				if !flush() {
+					return
+				}
+			}
+		}
+	}
+}

internal/app/etl/loaders/main.go

@@ -1,117 +1,13 @@
 package loaders
 
 import (
-	"context"
-	"errors"
-	"fmt"
-	"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/jackc/pgx/v5/pgconn"
 )
 
 type GenericLoader struct {
 	db dbwrapper.DbWrapper
 }
 
-func NewGenericLoader(db dbwrapper.DbWrapper) etl.Loader {
-	return &GenericLoader{db: db}
-}
-
-func (gl *GenericLoader) ProcessBatch(
-	ctx context.Context,
-	tableInfo config.TargetTableInfo,
-	colNames []string,
-	batch models.Batch,
-) (int, error) {
-	_, err := gl.db.SaveMassive(
-		ctx,
-		tableInfo.Schema,
-		tableInfo.Table,
-		colNames,
-		batch.Rows,
-	)
-
-	if err != nil {
-		var pgErr *pgconn.PgError
-		if errors.As(err, &pgErr) {
-			if pgErr.Code == "23505" {
-				return 0, &custom_errors.JobError{
-					ShouldCancelJob: true,
-					Msg:             fmt.Sprintf("Fatal error in table %s.%s", tableInfo.Schema, tableInfo.Table),
-					Prev:            err,
-				}
-			}
-		}
-
-		return 0, &custom_errors.LoaderError{Batch: batch, Msg: err.Error()}
-	}
-
-	return len(batch.Rows), nil
-}
-
-func (gl *GenericLoader) Exec(
-	ctx context.Context,
-	tableInfo config.TargetTableInfo,
-	columns []models.ColumnType,
-	chBatchesIn <-chan models.Batch,
-	chErrorsOut chan<- custom_errors.LoaderError,
-	chJobErrorsOut chan<- custom_errors.JobError,
-	wgActiveBatches *sync.WaitGroup,
-	rowsLoaded *int64,
-) {
-	colNames := mapSlice(columns, func(col models.ColumnType) string {
-		return col.Name()
-	})
-
-	for {
-		if ctx.Err() != nil {
-			return
-		}
-
-		select {
-		case <-ctx.Done():
-			return
-		case batch, ok := <-chBatchesIn:
-			if !ok {
-				return
-			}
-
-			processedRows, err := gl.ProcessBatch(ctx, tableInfo, colNames, batch)
-
-			if err != nil {
-				var ldError *custom_errors.LoaderError
-				var jobError *custom_errors.JobError
-				if errors.As(err, &ldError) {
-					select {
-					case <-ctx.Done():
-						return
-					case chErrorsOut <- *ldError:
-					}
-				} else if errors.As(err, &jobError) {
-					select {
-					case <-ctx.Done():
-						return
-					case chJobErrorsOut <- *jobError:
-					}
-				} else {
-					select {
-					case <-ctx.Done():
-						return
-					case chErrorsOut <- custom_errors.LoaderError{Batch: batch, Msg: err.Error()}:
-					}
-				}
-
-				continue
-			}
-
-			wgActiveBatches.Done()
-			atomic.AddInt64(rowsLoaded, int64(processedRows))
-		}
-	}
+func NewGenericLoader(db dbwrapper.DbWrapper) GenericLoader {
+	return GenericLoader{db: db}
 }

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

@@ -0,0 +1,49 @@
+package loaders
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"time"
+
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/custom_errors"
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
+)
+
+func (gl *GenericLoader) ProcessBatchWithRetries(
+	ctx context.Context,
+	tableInfo config.TargetTableInfo,
+	colNames []string,
+	retryConfig config.RetryConfig,
+	batch models.Batch,
+) (int64, error) {
+	for {
+		rowsLoaded, err := gl.ProcessBatch(ctx, tableInfo, colNames, batch)
+		if err == nil {
+			return rowsLoaded, nil
+		}
+
+		if btError, ok := errors.AsType[*custom_errors.LoaderError](err); ok {
+			batch.RetryCounter++
+
+			if batch.RetryCounter >= retryConfig.Attempts {
+				return rowsLoaded, &custom_errors.JobError{
+					Msg:  fmt.Sprintf("Batch %v reached max retries (%d)", batch.Id, batch.RetryCounter),
+					Prev: btError,
+				}
+			}
+
+			delay := custom_errors.ComputeBackoffDelay(
+				batch.RetryCounter,
+				retryConfig.BaseDelayMs,
+				retryConfig.MaxDelayMs,
+				retryConfig.MaxJitterMs,
+			)
+			time.Sleep(delay)
+			continue
+		}
+
+		return rowsLoaded, err
+	}
+}

internal/app/etl/loaders/process.go

@@ -0,0 +1,43 @@
+package loaders
+
+import (
+	"context"
+	"errors"
+	"fmt"
+
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/custom_errors"
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
+	"github.com/jackc/pgx/v5/pgconn"
+)
+
+func (gl *GenericLoader) ProcessBatch(
+	ctx context.Context,
+	tableInfo config.TargetTableInfo,
+	colNames []string,
+	batch models.Batch,
+) (int64, error) {
+	_, err := gl.db.SaveMassive(
+		ctx,
+		tableInfo.Schema,
+		tableInfo.Table,
+		colNames,
+		batch.Rows,
+	)
+
+	if err != nil {
+		if pgErr, ok := errors.AsType[*pgconn.PgError](err); ok {
+			if pgErr.Code == "23505" {
+				return 0, &custom_errors.JobError{
+					ShouldCancelJob: true,
+					Msg:             fmt.Sprintf("Fatal error in table %s.%s", tableInfo.Schema, tableInfo.Table),
+					Prev:            err,
+				}
+			}
+		}
+
+		return 0, &custom_errors.LoaderError{Batch: batch, Msg: err.Error()}
+	}
+
+	return int64(len(batch.Rows)), nil
+}

internal/app/etl/transformers/consume.go

@@ -0,0 +1,117 @@
+package transformers
+
+import (
+	"context"
+	"errors"
+	"sync"
+
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/custom_errors"
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
+	"github.com/google/uuid"
+)
+
+func (mssqlTr *MssqlTransformer) Consume(
+	ctx context.Context,
+	columns []models.ColumnType,
+	retryConfig config.RetryConfig,
+	batchSize int,
+	chBatchesIn <-chan models.Batch,
+	chBatchesOut chan<- models.Batch,
+	chJobErrorsOut chan<- custom_errors.JobError,
+	wgActiveBatches *sync.WaitGroup,
+) {
+	transformationPlan := computeTransformationPlan(columns)
+	storagePlan := computeStorageTransformationPlan(ctx, mssqlTr.azureClient, mssqlTr.toStorage, columns, mssqlTr.sourceTable)
+	transformationPlan = append(transformationPlan, storagePlan...)
+
+	var accRows []models.UnknownRowValues
+	var parentBatchesId []uuid.UUID
+	var firstPartitionId uuid.UUID
+
+	flush := func() bool {
+		if len(accRows) == 0 {
+			return true
+		}
+		out := models.Batch{
+			Id:              uuid.New(),
+			PartitionId:     firstPartitionId,
+			ParentBatchesId: parentBatchesId,
+			Rows:            accRows,
+		}
+		select {
+		case chBatchesOut <- out:
+			wgActiveBatches.Add(1)
+		case <-ctx.Done():
+			return false
+		}
+		accRows = nil
+		parentBatchesId = nil
+		firstPartitionId = uuid.Nil
+		return true
+	}
+
+	for {
+		if ctx.Err() != nil {
+			return
+		}
+
+		select {
+		case <-ctx.Done():
+			return
+
+		case batch, ok := <-chBatchesIn:
+			if !ok {
+				flush()
+				return
+			}
+
+			if len(transformationPlan) > 0 {
+				err := ProcessBatchWithRetries(ctx, &batch, transformationPlan, retryConfig)
+				if err != nil {
+					if errors.Is(err, ctx.Err()) {
+						return
+					}
+
+					if jobError, ok := errors.AsType[*custom_errors.JobError](err); ok {
+						select {
+						case chJobErrorsOut <- *jobError:
+						case <-ctx.Done():
+							return
+						}
+					} else {
+						select {
+						case chJobErrorsOut <- custom_errors.JobError{ShouldCancelJob: true, Msg: "Transformation failed", Prev: err}:
+						case <-ctx.Done():
+							return
+						}
+					}
+
+					return
+				}
+			}
+
+			if batchSize <= 0 {
+				select {
+				case chBatchesOut <- batch:
+					wgActiveBatches.Add(1)
+				case <-ctx.Done():
+					return
+				}
+				continue
+			}
+
+			if len(parentBatchesId) == 0 {
+				firstPartitionId = batch.PartitionId
+			}
+			accRows = append(accRows, batch.Rows...)
+			parentBatchesId = append(parentBatchesId, batch.Id)
+
+			if len(accRows) >= batchSize {
+				if !flush() {
+					return
+				}
+			}
+		}
+	}
+}

internal/app/etl/transformers/mssql.go

@@ -1,20 +1,9 @@
 package transformers
 
 import (
-	"context"
-	"errors"
-	"fmt"
-	"strings"
-	"sync"
-	"time"
-
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/azure"
 	"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"
-	log "github.com/sirupsen/logrus"
 )
 
 type MssqlTransformer struct {
@@ -30,196 +19,3 @@ func NewMssqlTransformer(toStorage config.ToStorageConfig, sourceTable config.So
 		azureClient: azureClient,
 	}
 }
-
-func computeTransformationPlan(columns []models.ColumnType) []etl.ColumnTransformPlan {
-	var plan []etl.ColumnTransformPlan
-
-	for i, col := range columns {
-		switch col.SystemType() {
-		case "uniqueidentifier":
-			plan = append(plan, etl.ColumnTransformPlan{
-				Index: i,
-				Fn: func(v any) (any, error) {
-					if b, ok := v.([]byte); ok && b != nil {
-						return mssqlUuidToBigEndian(b)
-					}
-					return v, nil
-				},
-			})
-
-		case "geometry", "geography":
-			plan = append(plan, etl.ColumnTransformPlan{
-				Index: i,
-				Fn: func(v any) (any, error) {
-					if b, ok := v.([]byte); ok && b != nil {
-						return wkbToEwkbWithSrid(b, 4326)
-					}
-					return v, nil
-				},
-			})
-
-		case "datetime", "datetime2":
-			plan = append(plan, etl.ColumnTransformPlan{
-				Index: i,
-				Fn: func(v any) (any, error) {
-					if t, ok := v.(time.Time); ok {
-						return ensureUTC(t), nil
-					}
-					return v, nil
-				},
-			})
-		}
-	}
-
-	return plan
-}
-
-func computeStorageTransformationPlan(
-	ctx context.Context,
-	azureClient *azure.Client,
-	toStorage config.ToStorageConfig,
-	sourceColumns []models.ColumnType,
-	sourceTable config.SourceTableInfo,
-) []etl.ColumnTransformPlan {
-	if azureClient == nil || len(toStorage.Columns) == 0 {
-		return nil
-	}
-
-	colIndex := make(map[string]int, len(sourceColumns))
-	for i, col := range sourceColumns {
-		colIndex[strings.ToUpper(col.Name())] = i
-	}
-
-	var plan []etl.ColumnTransformPlan
-	for _, storageCol := range toStorage.Columns {
-		if storageCol.Mode != "REFERENCE_ONLY" {
-			log.Warnf("to_storage: unsupported mode %q for column %s — skipping", storageCol.Mode, storageCol.Source)
-			continue
-		}
-
-		idx, ok := colIndex[strings.ToUpper(storageCol.Source)]
-		if !ok {
-			log.Warnf("to_storage: source column %q not found in source schema — skipping", storageCol.Source)
-			continue
-		}
-
-		sourceColName := storageCol.Source
-		schema := sourceTable.Schema
-		table := sourceTable.Table
-
-		plan = append(plan, etl.ColumnTransformPlan{
-			Index: idx,
-			Fn: func(v any) (any, error) {
-				if v == nil {
-					return nil, nil
-				}
-				b, ok := v.([]byte)
-				if !ok {
-					log.Warnf("to_storage: expected []byte for %s.%s.%s, got %T — passing through",
-						schema, table, sourceColName, v)
-					return v, nil
-				}
-				start := time.Now()
-				blobPath := fmt.Sprintf("%s/%s/%s", schema, table, uuid.New().String())
-				blobURL, err := azureClient.UploadAndGetURL(ctx, blobPath, b)
-				if err != nil {
-					return nil, fmt.Errorf("uploading %s.%s.%s: %w", schema, table, sourceColName, err)
-				}
-				log.Debugf(`Succesfully uploaded "%s", (%vms)`, blobURL, time.Since(start).Milliseconds())
-				return blobURL, nil
-			},
-		})
-	}
-	return plan
-}
-
-const processBatchCtxCheck = 4096
-
-func (mssqlTr *MssqlTransformer) ProcessBatch(
-	ctx context.Context,
-	batch *models.Batch,
-	transformationPlan []etl.ColumnTransformPlan,
-) error {
-	for i, rowValues := range batch.Rows {
-		if i%processBatchCtxCheck == 0 {
-			if err := ctx.Err(); err != nil {
-				return err
-			}
-		}
-
-		for _, task := range transformationPlan {
-			val := rowValues[task.Index]
-			if val == nil {
-				continue
-			}
-
-			transformed, err := task.Fn(val)
-			if err != nil {
-				return err
-			}
-			rowValues[task.Index] = transformed
-		}
-	}
-
-	return nil
-}
-
-func (mssqlTr *MssqlTransformer) Exec(
-	ctx context.Context,
-	columns []models.ColumnType,
-	chBatchesIn <-chan models.Batch,
-	chBatchesOut chan<- models.Batch,
-	chJobErrorsOut chan<- custom_errors.JobError,
-	wgActiveBatches *sync.WaitGroup,
-) {
-	transformationPlan := computeTransformationPlan(columns)
-	storagePlan := computeStorageTransformationPlan(ctx, mssqlTr.azureClient, mssqlTr.toStorage, columns, mssqlTr.sourceTable)
-	transformationPlan = append(transformationPlan, storagePlan...)
-
-	for {
-		if ctx.Err() != nil {
-			return
-		}
-
-		select {
-		case <-ctx.Done():
-			return
-
-		case batch, ok := <-chBatchesIn:
-			if !ok {
-				return
-			}
-
-			if len(transformationPlan) == 0 {
-				select {
-				case chBatchesOut <- batch:
-					wgActiveBatches.Add(1)
-					continue
-				case <-ctx.Done():
-					return
-				}
-			}
-
-			err := mssqlTr.ProcessBatch(ctx, &batch, transformationPlan)
-			if err != nil {
-				if errors.Is(err, ctx.Err()) {
-					return
-				}
-
-				select {
-				case chJobErrorsOut <- custom_errors.JobError{ShouldCancelJob: true, Msg: "Transformation failed", Prev: err}:
-				case <-ctx.Done():
-				}
-				return
-			}
-
-			select {
-			case chBatchesOut <- batch:
-			case <-ctx.Done():
-				return
-			}
-
-			wgActiveBatches.Add(1)
-		}
-	}
-}

internal/app/etl/transformers/plan.go

@@ -0,0 +1,122 @@
+package transformers
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"time"
+
+	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/azure"
+	"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"
+	"github.com/sirupsen/logrus"
+)
+
+func computeTransformationPlan(columns []models.ColumnType) []etl.ColumnTransformPlan {
+	var plan []etl.ColumnTransformPlan
+
+	for i, col := range columns {
+		switch col.SystemType() {
+		case "uniqueidentifier":
+			plan = append(plan, etl.ColumnTransformPlan{
+				Index: i,
+				Fn: func(v any) (any, error) {
+					if b, ok := v.([]byte); ok && b != nil {
+						return mssqlUuidToBigEndian(b)
+					}
+					return v, nil
+				},
+			})
+
+		case "geometry", "geography":
+			plan = append(plan, etl.ColumnTransformPlan{
+				Index: i,
+				Fn: func(v any) (any, error) {
+					if b, ok := v.([]byte); ok && b != nil {
+						return wkbToEwkbWithSrid(b, 4326)
+					}
+					return v, nil
+				},
+			})
+
+		case "datetime", "datetime2":
+			plan = append(plan, etl.ColumnTransformPlan{
+				Index: i,
+				Fn: func(v any) (any, error) {
+					if t, ok := v.(time.Time); ok {
+						return ensureUTC(t), nil
+					}
+					return v, nil
+				},
+			})
+		}
+	}
+
+	return plan
+}
+
+func computeStorageTransformationPlan(
+	ctx context.Context,
+	azureClient *azure.Client,
+	toStorage config.ToStorageConfig,
+	sourceColumns []models.ColumnType,
+	sourceTable config.SourceTableInfo,
+) []etl.ColumnTransformPlan {
+	if azureClient == nil || len(toStorage.Columns) == 0 {
+		return nil
+	}
+
+	colIndex := make(map[string]int, len(sourceColumns))
+	for i, col := range sourceColumns {
+		colIndex[strings.ToUpper(col.Name())] = i
+	}
+
+	var plan []etl.ColumnTransformPlan
+	for _, storageCol := range toStorage.Columns {
+		if storageCol.Mode != "REFERENCE_ONLY" {
+			logrus.Warnf("to_storage: unsupported mode %q for column %s — skipping", storageCol.Mode, storageCol.Source)
+			continue
+		}
+
+		idx, ok := colIndex[strings.ToUpper(storageCol.Source)]
+		if !ok {
+			logrus.Warnf("to_storage: source column %q not found in source schema — skipping", storageCol.Source)
+			continue
+		}
+
+		sourceColName := storageCol.Source
+		schema := sourceTable.Schema
+		table := sourceTable.Table
+
+		plan = append(plan, etl.ColumnTransformPlan{
+			Index: idx,
+			Fn: func(v any) (any, error) {
+				if v == nil {
+					return nil, nil
+				}
+				b, ok := v.([]byte)
+				if !ok {
+					logrus.Warnf("to_storage: expected []byte for %s.%s.%s, got %T — passing through",
+						schema, table, sourceColName, v)
+					return v, nil
+				}
+				// start := time.Now()
+				blobPath := fmt.Sprintf("%s/%s/%s", schema, table, uuid.New().String())
+				blobURL, err := azureClient.UploadAndGetURL(ctx, blobPath, b)
+				if err != nil {
+					return nil, &custom_errors.JobError{
+						Msg:  fmt.Sprintf("Error uploading %s.%s.%s", schema, table, sourceColName),
+						Prev: err,
+					}
+				}
+
+				// logrus.Debugf(`Succesfully uploaded "%s", (%vms)`, blobURL, time.Since(start).Milliseconds())
+				return blobURL, nil
+			},
+		})
+	}
+	return plan
+}

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

@@ -0,0 +1,73 @@
+package transformers
+
+import (
+	"context"
+	"errors"
+	"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"
+)
+
+const processBatchCtxCheck = 4096
+
+func ProcessBatchWithRetries(
+	ctx context.Context,
+	batch *models.Batch,
+	transformationPlan []etl.ColumnTransformPlan,
+	retryConfig config.RetryConfig,
+) error {
+	for i, rowValues := range batch.Rows {
+		if i%processBatchCtxCheck == 0 {
+			if err := ctx.Err(); err != nil {
+				return err
+			}
+		}
+
+		for _, task := range transformationPlan {
+			val := rowValues[task.Index]
+			if val == nil {
+				continue
+			}
+
+			var lastErr error
+			success := false
+
+			for attempt := 0; attempt < retryConfig.Attempts; attempt++ {
+				transformed, err := task.Fn(val)
+				if err == nil {
+					rowValues[task.Index] = transformed
+					success = true
+					break
+				}
+
+				lastErr = err
+				if jobError, ok := errors.AsType[*custom_errors.JobError](err); ok {
+					if jobError.ShouldCancelJob {
+						return jobError
+					}
+				}
+
+				if attempt == retryConfig.Attempts-1 {
+					break
+				}
+
+				delay := custom_errors.ComputeBackoffDelay(
+					attempt,
+					retryConfig.BaseDelayMs,
+					retryConfig.MaxDelayMs,
+					retryConfig.MaxJitterMs,
+				)
+				time.Sleep(delay)
+			}
+
+			if !success {
+				return lastErr
+			}
+		}
+	}
+
+	return nil
+}

internal/app/etl/transformers/types.go

@@ -1 +0,0 @@
-package transformers

internal/app/etl/types.go

@@ -9,18 +9,6 @@ import (
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
 )
 
-type Extractor interface {
-	ProcessPartition(
-		ctx context.Context,
-		tableInfo config.SourceTableInfo,
-		columns []models.ColumnType,
-		batchSize int,
-		partition models.Partition,
-		indexPrimaryKey int,
-		chBatchesOut chan<- models.Batch,
-	) (int, error)
-}
-
 type TransformerFunc func(any) (any, error)
 
 type ColumnTransformPlan struct {
@@ -29,42 +17,18 @@ type ColumnTransformPlan struct {
 }
 
 type Transformer interface {
-	ProcessBatch(
-		ctx context.Context,
-		batch *models.Batch,
-		transformationPlan []ColumnTransformPlan,
-	) error
-
-	Exec(
+	Consume(
 		ctx context.Context,
 		columns []models.ColumnType,
+		retryConfig config.RetryConfig,
+		batchSize int,
 		chBatchesIn <-chan models.Batch,
-		chBactchesOut chan<- models.Batch,
+		chBatchesOut chan<- models.Batch,
 		chJobErrorsOut chan<- custom_errors.JobError,
 		wgActiveBatches *sync.WaitGroup,
 	)
 }
 
-type Loader interface {
-	ProcessBatch(
-		ctx context.Context,
-		tableInfo config.TargetTableInfo,
-		colNames []string,
-		batch models.Batch,
-	) (int, error)
-
-	Exec(
-		ctx context.Context,
-		tableInfo config.TargetTableInfo,
-		columns []models.ColumnType,
-		chBatchesIn <-chan models.Batch,
-		chErrorsOut chan<- custom_errors.LoaderError,
-		chJobErrorsOut chan<- custom_errors.JobError,
-		wgActiveBatches *sync.WaitGroup,
-		rowsLoaded *int64,
-	)
-}
-
 type TableAnalyzer interface {
 	QueryColumnTypes(
 		ctx context.Context,

internal/app/models/main.go

@@ -9,10 +9,11 @@ import (
 type UnknownRowValues = []any
 
 type Batch struct {
-	Id           uuid.UUID
-	PartitionId  uuid.UUID
-	Rows         []UnknownRowValues
-	RetryCounter int
+	Id              uuid.UUID
+	PartitionId     uuid.UUID
+	ParentBatchesId []uuid.UUID
+	Rows            []UnknownRowValues
+	RetryCounter    int
 }
 
 type PartitionRange struct {