feat: integrate azure storage handling in migration pipeline

Stream transformer output to an Azure blob container when a job declares ToStorage, with per-job prefix and credential configuration.

internal/app/etl/extractors/main.go

@@ -0,0 +1,101 @@
+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.ExtractorError,
+	chJobErrorsOut chan<- custom_errors.JobError,
+	wgActivePartitions *sync.WaitGroup,
+	rowsRead *int64,
+) {
+	indexPrimaryKey := slices.IndexFunc(columns, func(col models.ColumnType) bool {
+		return strings.EqualFold(col.Name(), tableInfo.PrimaryKey)
+	})
+
+	if indexPrimaryKey == -1 {
+		select {
+		case <-ctx.Done():
+			return
+		case chJobErrorsOut <- custom_errors.JobError{
+			ShouldCancelJob: true,
+			Msg:             "Primary key not found in provided columns",
+		}:
+		}
+
+		return
+	}
+
+	for {
+		if ctx.Err() != nil {
+			return
+		}
+
+		select {
+		case <-ctx.Done():
+			return
+		case partition, ok := <-chPartitionsIn:
+			if !ok {
+				return
+			}
+
+			rowsReadResult, err := extractor.Exec(
+				ctx,
+				tableInfo,
+				columns,
+				batchSize,
+				partition,
+				indexPrimaryKey,
+				chBatchesOut,
+			)
+
+			if rowsReadResult > 0 {
+				atomic.AddInt64(rowsRead, int64(rowsReadResult))
+			}
+
+			if err != nil {
+				if exError, ok := errors.AsType[*custom_errors.ExtractorError](err); ok {
+					select {
+					case <-ctx.Done():
+						return
+					case chErrorsOut <- *exError:
+					}
+				} else if jobError, ok := errors.AsType[*custom_errors.JobError](err); ok {
+					select {
+					case <-ctx.Done():
+						return
+					case chJobErrorsOut <- *jobError:
+					}
+				} else {
+					select {
+					case <-ctx.Done():
+						return
+					case chErrorsOut <- custom_errors.ExtractorError{Partition: partition, Msg: err.Error()}:
+					}
+				}
+
+				continue
+			}
+
+			wgActivePartitions.Done()
+		}
+	}
+}

internal/app/etl/extractors/mssql.go

@@ -5,10 +5,7 @@ import (
 	"database/sql"
 	"errors"
 	"fmt"
-	"slices"
 	"strings"
-	"sync"
-	"sync/atomic"
 
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/convert"
@@ -32,6 +29,9 @@ func buildExtractQueryMssql(
 	columns []models.ColumnType,
 	includeRange bool,
 	isMinInclusive bool,
+	isMaxInclusive bool,
+	hasMin bool,
+	hasMax bool,
 ) string {
 	var sbQuery strings.Builder
 
@@ -55,15 +55,32 @@ func buildExtractQueryMssql(
 
 	fmt.Fprintf(&sbQuery, " FROM [%s].[%s]", tableInfo.Schema, tableInfo.Table)
 
-	if includeRange {
-		fmt.Fprintf(&sbQuery, " WHERE [%s]", tableInfo.PrimaryKey)
-		if isMinInclusive {
-			sbQuery.WriteString(" >=")
-		} else {
-			sbQuery.WriteString(" >")
+	if includeRange && (hasMin || hasMax) {
+		sbQuery.WriteString(" WHERE ")
+
+		if hasMin {
+			fmt.Fprintf(&sbQuery, "[%s]", tableInfo.PrimaryKey)
+			if isMinInclusive {
+				sbQuery.WriteString(" >=")
+			} else {
+				sbQuery.WriteString(" >")
+			}
+			sbQuery.WriteString(" @min")
 		}
 
-		fmt.Fprintf(&sbQuery, " @min AND [%s] <= @max", tableInfo.PrimaryKey)
+		if hasMin && hasMax {
+			sbQuery.WriteString(" AND ")
+		}
+
+		if hasMax {
+			fmt.Fprintf(&sbQuery, "[%s]", tableInfo.PrimaryKey)
+			if isMaxInclusive {
+				sbQuery.WriteString(" <=")
+			} else {
+				sbQuery.WriteString(" <")
+			}
+			sbQuery.WriteString(" @max")
+		}
 	}
 
 	fmt.Fprintf(&sbQuery, " ORDER BY [%s] ASC", tableInfo.PrimaryKey)
@@ -99,7 +116,7 @@ func errorFromLastRow(
 	}
 }
 
-func (mssqlEx *MssqlExtractor) ProcessPartition(
+func (mssqlEx *MssqlExtractor) Exec(
 	ctx context.Context,
 	tableInfo config.SourceTableInfo,
 	columns []models.ColumnType,
@@ -108,14 +125,16 @@ func (mssqlEx *MssqlExtractor) ProcessPartition(
 	indexPrimaryKey int,
 	chBatchesOut chan<- models.Batch,
 ) (int, error) {
-	query := buildExtractQueryMssql(tableInfo, columns, partition.HasRange, partition.Range.IsMinInclusive)
+	hasMin := partition.HasRange && partition.Range.Min > 0
+	hasMax := partition.HasRange && partition.Range.Max > 0
+	query := buildExtractQueryMssql(tableInfo, columns, partition.HasRange, partition.Range.IsMinInclusive, partition.Range.IsMaxInclusive, hasMin, hasMax)
 
 	var queryArgs []any
-	if partition.HasRange {
-		queryArgs = append(queryArgs,
-			sql.Named("min", partition.Range.Min),
-			sql.Named("max", partition.Range.Max),
-		)
+	if hasMin {
+		queryArgs = append(queryArgs, sql.Named("min", partition.Range.Min))
+	}
+	if hasMax {
+		queryArgs = append(queryArgs, sql.Named("max", partition.Range.Max))
 	}
 
 	rowsRead := 0
@@ -188,90 +207,3 @@ func (mssqlEx *MssqlExtractor) ProcessPartition(
 
 	return rowsRead, nil
 }
-
-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
-	}
-
-	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

@@ -2,10 +2,8 @@ package extractors
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"strings"
-	"sync"
 
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/custom_errors"
@@ -23,7 +21,15 @@ func NewPostgresExtractor(db dbwrapper.DbWrapper) etl.Extractor {
 	return &PostgresExtractor{db: db}
 }
 
-func buildExtractQueryPostgres(sourceDbInfo config.SourceTableInfo, columns []models.ColumnType) string {
+func buildExtractQueryPostgres(
+	sourceDbInfo config.SourceTableInfo,
+	columns []models.ColumnType,
+	includeRange bool,
+	isMinInclusive bool,
+	isMaxInclusive bool,
+	hasMin bool,
+	hasMax bool,
+) string {
 	var sbColumns strings.Builder
 
 	if len(columns) == 0 {
@@ -48,10 +54,44 @@ 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)
+	query := fmt.Sprintf(`SELECT %s FROM "%s"."%s"`, sbColumns.String(), sourceDbInfo.Schema, sourceDbInfo.Table)
+
+	if includeRange && (hasMin || hasMax) {
+		query += " WHERE "
+		paramIdx := 1
+
+		if hasMin {
+			query += fmt.Sprintf(`"%s"`, sourceDbInfo.PrimaryKey)
+			if isMinInclusive {
+				query += " >="
+			} else {
+				query += " >"
+			}
+			query += fmt.Sprintf(" $%d", paramIdx)
+			paramIdx++
+		}
+
+		if hasMin && hasMax {
+			query += " AND "
+		}
+
+		if hasMax {
+			query += fmt.Sprintf(`"%s"`, sourceDbInfo.PrimaryKey)
+			if isMaxInclusive {
+				query += " <="
+			} else {
+				query += " <"
+			}
+			query += fmt.Sprintf(" $%d", paramIdx)
+		}
+	}
+
+	query += fmt.Sprintf(` ORDER BY "%s" ASC`, sourceDbInfo.PrimaryKey)
+
+	return query
 }
 
-func (postgresEx *PostgresExtractor) ProcessPartition(
+func (postgresEx *PostgresExtractor) Exec(
 	ctx context.Context,
 	tableInfo config.SourceTableInfo,
 	columns []models.ColumnType,
@@ -60,14 +100,20 @@ func (postgresEx *PostgresExtractor) ProcessPartition(
 	indexPrimaryKey int,
 	chBatchesOut chan<- models.Batch,
 ) (int, error) {
-	query := buildExtractQueryPostgres(tableInfo, columns)
+	hasMin := partition.HasRange && partition.Range.Min > 0
+	hasMax := partition.HasRange && partition.Range.Max > 0
+	query := buildExtractQueryPostgres(tableInfo, columns, partition.HasRange, partition.Range.IsMinInclusive, partition.Range.IsMaxInclusive, hasMin, hasMax)
 
-	if partition.HasRange {
-		return 0, errors.New("Batch config not yet supported")
+	var queryArgs []any
+	if hasMin {
+		queryArgs = append(queryArgs, partition.Range.Min)
+	}
+	if hasMax {
+		queryArgs = append(queryArgs, partition.Range.Max)
 	}
 
 	rowsRead := 0
-	rows, err := postgresEx.db.Query(ctx, query)
+	rows, err := postgresEx.db.Query(ctx, query, queryArgs...)
 	if err != nil {
 		return rowsRead, &custom_errors.ExtractorError{Partition: partition, HasLastId: false, Msg: err.Error()}
 	}
@@ -78,7 +124,7 @@ func (postgresEx *PostgresExtractor) ProcessPartition(
 	for rows.Next() {
 		values, err := rows.Values()
 		if err != nil {
-			return rowsRead, errors.New("Unexpected error reading rows from source")
+			return rowsRead, &custom_errors.ExtractorError{Partition: partition, HasLastId: false, Msg: err.Error()}
 		}
 		rowsRead++
 
@@ -96,7 +142,7 @@ func (postgresEx *PostgresExtractor) ProcessPartition(
 	}
 
 	if err := rows.Err(); err != nil {
-		return rowsRead, errors.New("Unexpected error reading rows from source")
+		return rowsRead, &custom_errors.ExtractorError{Partition: partition, HasLastId: false, Msg: err.Error()}
 	}
 
 	if len(batchRows) > 0 {
@@ -109,17 +155,3 @@ func (postgresEx *PostgresExtractor) ProcessPartition(
 
 	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,
-) {
-}