refactor: replace Exec with Consume method in MssqlTransformer; enhance retry handling and streamline transformation logic

internal/app/etl/transformers/consume.go

@@ -0,0 +1,82 @@
+package transformers
+
+import (
+	"context"
+	"errors"
+	"sync"
+
+	"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"
+)
+
+func (mssqlTr *MssqlTransformer) Consume(
+	ctx context.Context,
+	columns []models.ColumnType,
+	retryConfig config.RetryConfig,
+	chBatchesIn <-chan models.Batch,
+	chBatchesOut chan<- models.Batch,
+	chJobErrorsOut chan<- custom_errors.JobError,
+	wgActiveBatches *sync.WaitGroup,
+) {
+	transformationPlan := computeTransformationPlan(columns)
+	storagePlan := computeStorageTransformationPlan(ctx, mssqlTr.azureClient, mssqlTr.toStorage, columns, mssqlTr.sourceTable)
+	transformationPlan = append(transformationPlan, storagePlan...)
+
+	for {
+		if ctx.Err() != nil {
+			return
+		}
+
+		select {
+		case <-ctx.Done():
+			return
+
+		case batch, ok := <-chBatchesIn:
+			if !ok {
+				return
+			}
+
+			if len(transformationPlan) == 0 {
+				select {
+				case chBatchesOut <- batch:
+					wgActiveBatches.Add(1)
+					continue
+				case <-ctx.Done():
+					return
+				}
+			}
+
+			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
+			}
+
+			select {
+			case chBatchesOut <- batch:
+			case <-ctx.Done():
+				return
+			}
+
+			wgActiveBatches.Add(1)
+		}
+	}
+}