refactor: streamline error handling in migration process; consolidate failed partitions and batches tracking

internal/app/etl/extractors/consume.go

@@ -25,6 +25,7 @@ func (ex *GenericExtractor) Consume(
 	chErrorsOut chan<- custom_errors.JobError,
 	wgActivePartitions *sync.WaitGroup,
 	rowsRead *int64,
+	failedPartitionsCount *int32,
 	fromJsonColumns []config.FromJsonItem,
 ) {
 	indexPrimaryKey := slices.IndexFunc(columns, func(col models.ColumnType) bool {
@@ -77,6 +78,7 @@ func (ex *GenericExtractor) Consume(
 			}
 
 			if err != nil {
+				atomic.AddInt32(failedPartitionsCount, 1)
 				if jobError, ok := errors.AsType[*custom_errors.JobError](err); ok {
 					select {
 					case <-ctx.Done():
@@ -90,6 +92,16 @@ func (ex *GenericExtractor) Consume(
 					case chErrorsOut <- custom_errors.JobError{ShouldCancelJob: false, Msg: err.Error(), Prev: err}:
 					}
 				}
+
+				currentFPCount := atomic.LoadInt32(failedPartitionsCount)
+				if currentFPCount > int32(retryConfig.MaxFailedPartitions) {
+					select {
+					case <-ctx.Done():
+						return
+					case chErrorsOut <- custom_errors.JobError{ShouldCancelJob: true, Msg: "Max failed partitions reached"}:
+						return
+					}
+				}
 			}
 		}
 	}

internal/app/etl/loaders/consume.go

@@ -20,6 +20,7 @@ func (gl *GenericLoader) Consume(
 	chErrorsOut chan<- custom_errors.JobError,
 	wgActiveBatches *sync.WaitGroup,
 	rowsLoaded *int64,
+	failedBatchesCount *int32,
 ) {
 	colNames := mapSlice(columns, func(col models.ColumnType) string {
 		return col.Name()
@@ -42,6 +43,7 @@ func (gl *GenericLoader) Consume(
 			wgActiveBatches.Done()
 
 			if err != nil {
+				atomic.AddInt32(failedBatchesCount, 1)
 				if jobError, ok := errors.AsType[*custom_errors.JobError](err); ok {
 					select {
 					case <-ctx.Done():
@@ -55,9 +57,21 @@ func (gl *GenericLoader) Consume(
 					case chErrorsOut <- custom_errors.JobError{ShouldCancelJob: false, Msg: err.Error(), Prev: err}:
 					}
 				}
-			} else {
-				atomic.AddInt64(rowsLoaded, int64(processedRows))
+
+				currentFBCount := atomic.LoadInt32(failedBatchesCount)
+				if currentFBCount > int32(retryConfig.MaxFailedBatchesLoad) {
+					select {
+					case <-ctx.Done():
+						return
+					case chErrorsOut <- custom_errors.JobError{ShouldCancelJob: true, Msg: "Max failed batches (load) reached"}:
+						return
+					}
+				}
+
+				continue
 			}
+
+			atomic.AddInt64(rowsLoaded, int64(processedRows))
 		}
 	}
 }