refactor: streamline error handling and processing in GenericExtractor; implement partition processing with retries

internal/app/etl/extractors/process.go

@@ -2,7 +2,6 @@ package extractors
 
 import (
 	"context"
-	"errors"
 	"fmt"
 
 	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
@@ -10,10 +9,10 @@ import (
 	"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/models"
-	"github.com/google/uuid"
+	// "github.com/sirupsen/logrus"
 )
 
-func errorFromLastRow(
+func errorFromLastPartitionRow(
 	lastRow models.UnknownRowValues,
 	indexPrimaryKey int,
 	partition models.Partition,
@@ -48,8 +47,7 @@ func (ex *GenericExtractor) ProcessPartition(
 	partition models.Partition,
 	indexPrimaryKey int,
 	chBatchesOut chan<- models.Batch,
-) (int, error) {
-	rowsRead := 0
+) (int64, error) {
 	query := dbwrapper.ExtractionQuery{
 		Schema:     tableInfo.Schema,
 		Table:      tableInfo.Table,
@@ -67,15 +65,15 @@ func (ex *GenericExtractor) ProcessPartition(
 		},
 	}
 
-	// logrus.Debugf("Querying with: %+v (%s.%s)", query, tableInfo.Schema, tableInfo.Table)
+	// logrus.Debugf("Processing partition: %+v (%s.%s)", query, tableInfo.Schema, tableInfo.Table)
 	rows, err := ex.db.QueryFromObject(ctx, query)
-
 	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
 
 	for rows.Next() {
 		rowValues := make([]any, len(columns))
@@ -87,53 +85,32 @@ func (ex *GenericExtractor) ProcessPartition(
 
 		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()
+			// logrus.Debugf("Batch size reached, flushing batch with %v rows (rowsRead=%v)", len(batchRows), rowsRead)
+			if err := flush(ctx, &partition, batchSize, batchRows, chBatchesOut); err != nil {
+				// logrus.Warnf("Error flushing rows: %v", err)
+				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 err := flush(ctx, &partition, batchSize, batchRows, chBatchesOut); err != nil {
+		return rowsRead, err
 	}
 
-	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
+	return rowsRead, rows.Err()
 }