refactor: streamline error handling and remove redundant code in mssql extractor

internal/app/etl/extractors/mssql.go

@@ -12,7 +12,6 @@ import (
 	"time"
 
 	"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"
@@ -72,34 +71,6 @@ func buildExtractQueryMssql(
 	return sbQuery.String()
 }
 
-func errorFromLastRow(
-	lastRow models.UnknownRowValues,
-	indexPrimaryKey int,
-	partition models.Partition,
-	previousError error,
-) *custom_errors.ExtractorError {
-	lastIdRawValue := lastRow[indexPrimaryKey]
-
-	lastId, ok := convert.ToInt64(lastIdRawValue)
-	if !ok {
-		currentPartition := partition
-		currentPartition.RetryCounter = 3
-		return &custom_errors.ExtractorError{
-			Partition: currentPartition,
-			HasLastId: true,
-			Msg:       fmt.Sprintf("Couldn't cast last id value as int: %s", previousError.Error()),
-		}
-
-	}
-
-	return &custom_errors.ExtractorError{
-		Partition: partition,
-		HasLastId: true,
-		LastId:    lastId,
-		Msg:       previousError.Error(),
-	}
-}
-
 func (mssqlEx *MssqlExtractor) Extract(
 	ctx context.Context,
 	tableInfo config.SourceTableInfo,
@@ -113,81 +84,52 @@ func (mssqlEx *MssqlExtractor) Extract(
 
 	var queryArgs []any
 	if partition.HasRange {
-		queryArgs = append(queryArgs,
-			sql.Named("min", partition.Range.Min),
-			sql.Named("max", partition.Range.Max),
-		)
+		queryArgs = append(queryArgs, sql.Named("min", partition.Range.Min), sql.Named("max", partition.Range.Max))
 	}
 
-	var rowsRead int64 = 0
 	rows, err := mssqlEx.db.Query(ctx, query, queryArgs...)
 	if err != nil {
-		return rowsRead, &custom_errors.ExtractorError{Partition: partition, HasLastId: false, Msg: err.Error()}
+		return 0, err
 	}
 	defer rows.Close()
 
 	batchRows := make([]models.UnknownRowValues, 0, batchSize)
+	var rowsRead int64 = 0
+
+	rowValues := make([]any, len(columns))
+	scanArgs := make([]any, len(columns))
+	for i := range rowValues {
+		scanArgs[i] = &rowValues[i]
+	}
 
 	for rows.Next() {
-		rowValues := make([]any, len(columns))
-		scanArgs := make([]any, len(columns))
-
-		for i := range rowValues {
-			scanArgs[i] = &rowValues[i]
-		}
-
 		if err := rows.Scan(scanArgs...); err != nil {
 			if len(batchRows) == 0 {
-				return rowsRead, &custom_errors.ExtractorError{Partition: partition, HasLastId: false, Msg: err.Error()}
+				return rowsRead, err
+			}
+
+			if err := flush(ctx, &partition, batchSize, batchRows, chBatchesOut); err != nil {
+				return rowsRead, err
 			}
 
 			lastRow := batchRows[len(batchRows)-1]
-
-			select {
-			case chBatchesOut <- models.Batch{Id: uuid.New(), PartitionId: partition.Id, Rows: batchRows, RetryCounter: 0}:
-			case <-ctx.Done():
-				return rowsRead, ctx.Err()
-			}
-
-			return rowsRead, errorFromLastRow(lastRow, indexPrimaryKey, partition, err)
+			return rowsRead, errorFromLastPartitionRow(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 rowsRead, ctx.Err()
+			if err := flush(ctx, &partition, batchSize, batchRows, chBatchesOut); err != nil {
+				return rowsRead, err
 			}
-
-			batchRows = make([]models.UnknownRowValues, 0, batchSize)
-		}
-
-	}
-
-	if err := rows.Err(); err != nil {
-		if errors.Is(err, ctx.Err()) {
-			return rowsRead, ctx.Err()
-		}
-
-		if len(batchRows) > 0 {
-			lastRow := batchRows[len(batchRows)-1]
-			return rowsRead, errorFromLastRow(lastRow, indexPrimaryKey, partition, err)
-		}
-
-		return rowsRead, &custom_errors.ExtractorError{Partition: partition, HasLastId: false, Msg: err.Error()}
-	}
-
-	if len(batchRows) > 0 {
-		select {
-		case chBatchesOut <- models.Batch{Id: uuid.New(), PartitionId: partition.Id, Rows: batchRows, RetryCounter: 0}:
-		case <-ctx.Done():
-			return rowsRead, ctx.Err()
 		}
 	}
 
-	return rowsRead, nil
+	if err := flush(ctx, &partition, batchSize, batchRows, chBatchesOut); err != nil {
+		return rowsRead, err
+	}
+
+	return rowsRead, rows.Err()
 }
 
 func (mssqlEx *MssqlExtractor) ExtractWithRetries(
@@ -252,7 +194,7 @@ func (mssqlEx *MssqlExtractor) Consume(
 	batchSize int,
 	chPartitionsIn <-chan models.Partition,
 	chBatchesOut chan<- models.Batch,
-	chJobErrorsOut chan<- custom_errors.JobError,
+	chErrorsOut chan<- custom_errors.JobError,
 	wgActivePartitions *sync.WaitGroup,
 	rowsRead *int64,
 ) {
@@ -264,12 +206,11 @@ func (mssqlEx *MssqlExtractor) Consume(
 		select {
 		case <-ctx.Done():
 			return
-		case chJobErrorsOut <- custom_errors.JobError{
+		case chErrorsOut <- custom_errors.JobError{
 			ShouldCancelJob: true,
 			Msg:             "Primary key not found in provided columns",
 		}:
 		}
-
 		return
 	}
 
@@ -307,13 +248,13 @@ func (mssqlEx *MssqlExtractor) Consume(
 					select {
 					case <-ctx.Done():
 						return
-					case chJobErrorsOut <- *jobError:
+					case chErrorsOut <- *jobError:
 					}
 				} else {
 					select {
 					case <-ctx.Done():
 						return
-					case chJobErrorsOut <- custom_errors.JobError{ShouldCancelJob: false, Msg: err.Error(), Prev: err}:
+					case chErrorsOut <- custom_errors.JobError{ShouldCancelJob: false, Msg: err.Error(), Prev: err}:
 					}
 				}