feat: refactor chunk handling in extractor and transformer for improved data processing

cmd/go_migrate/extractor.go

@@ -7,6 +7,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/google/uuid"
 	"github.com/jackc/pgx/v5/pgxpool"
 	_ "github.com/microsoft/go-mssqldb"
 	log "github.com/sirupsen/logrus"
@@ -14,6 +15,13 @@ import (
 
 type UnknownRowValues = []any
 
+type Chunk struct {
+	Id           uuid.UUID
+	BatchId      uuid.UUID
+	Data         []UnknownRowValues
+	RetryCounter int
+}
+
 func extractFromMssql(
 	ctx context.Context,
 	db *sql.DB,
@@ -21,7 +29,7 @@ func extractFromMssql(
 	columns []ColumnType,
 	chunkSize int,
 	chBatchesIn <-chan Batch,
-	chChunksOut chan<- []UnknownRowValues,
+	chChunksOut chan<- Chunk,
 	chErrorsOut chan<- ExtractorError,
 	chJobErrorsOut chan<- JobError,
 ) {
@@ -34,7 +42,11 @@ func extractFromMssql(
 			ShouldCancelJob: true,
 			Msg:             "Primary key not found in provided columns",
 		}
-		chJobErrorsOut <- jobError
+		select {
+		case chJobErrorsOut <- jobError:
+		case <-ctx.Done():
+			return
+		}
 		return
 	}
 
@@ -89,8 +101,13 @@ func extractFromMssql(
 					}
 
 					lastRow := rowsChunk[len(rowsChunk)-1]
-					chChunksOut <- rowsChunk
 					chErrorsOut <- ExtractorErrorFromLastRowMssql(lastRow, indexPrimaryKey, &batch, err)
+					chChunksOut <- Chunk{
+						Id:           uuid.New(),
+						BatchId:      batch.Id,
+						Data:         rowsChunk,
+						RetryCounter: 0,
+					}
 					return
 				}
 
@@ -102,7 +119,12 @@ func extractFromMssql(
 					rowsPerSec := float64(chunkSize) / chunkDuration.Seconds()
 					log.Infof("Extracted chunk: %d rows in %v (%.0f rows/sec) - Total: %d rows",
 						len(rowsChunk), chunkDuration, rowsPerSec, totalRowsExtracted)
-					chChunksOut <- rowsChunk
+					chChunksOut <- Chunk{
+						Id:           uuid.New(),
+						BatchId:      batch.Id,
+						Data:         rowsChunk,
+						RetryCounter: 0,
+					}
 					rowsChunk = make([]UnknownRowValues, 0, chunkSize)
 					chunkStartTime = time.Now()
 				}
@@ -113,7 +135,12 @@ func extractFromMssql(
 				rowsPerSec := float64(len(rowsChunk)) / chunkDuration.Seconds()
 				log.Infof("Extracted final chunk: %d rows in %v (%.0f rows/sec) - Total: %d rows",
 					len(rowsChunk), chunkDuration, rowsPerSec, totalRowsExtracted)
-				chChunksOut <- rowsChunk
+				chChunksOut <- Chunk{
+					Id:           uuid.New(),
+					BatchId:      batch.Id,
+					Data:         rowsChunk,
+					RetryCounter: 0,
+				}
 			}
 
 			if err := rows.Err(); err != nil {