refactor: enhance logging and batch processing in migration; adjust configuration parameters for improved performance

internal/app/etl/transformers/consume.go

@@ -8,12 +8,14 @@ import (
 	"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,
@@ -23,6 +25,32 @@ func (mssqlTr *MssqlTransformer) Consume(
 	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
@@ -34,49 +62,56 @@ func (mssqlTr *MssqlTransformer) Consume(
 
 		case batch, ok := <-chBatchesIn:
 			if !ok {
+				flush()
 				return
 			}
 
-			if len(transformationPlan) == 0 {
+			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)
-					continue
 				case <-ctx.Done():
 					return
 				}
+				continue
 			}
 
-			err := ProcessBatchWithRetries(ctx, &batch, transformationPlan, retryConfig)
-			if err != nil {
-				if errors.Is(err, ctx.Err()) {
+			if len(parentBatchesId) == 0 {
+				firstPartitionId = batch.PartitionId
+			}
+			accRows = append(accRows, batch.Rows...)
+			parentBatchesId = append(parentBatchesId, batch.Id)
+
+			if len(accRows) >= batchSize {
+				if !flush() {
 					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
 			}
-
-			select {
-			case chBatchesOut <- batch:
-			case <-ctx.Done():
-				return
-			}
-
-			wgActiveBatches.Add(1)
 		}
 	}
 }