feat: implement loader error handling and refactor chunk processing in migration job

cmd/go_migrate/loader.go

@@ -12,39 +12,70 @@ import (
 	log "github.com/sirupsen/logrus"
 )
 
-func loadRowsPostgres(ctx context.Context, job MigrationJob, columns []ColumnType, db *pgxpool.Pool, chChunksIn <-chan Chunk) error {
-	chunkCount := 0
-	totalRowsLoaded := 0
+func loadRowsPostgres(
+	ctx context.Context,
+	db *pgxpool.Pool,
+	job MigrationJob,
+	columns []ColumnType,
+	chChunksIn <-chan Chunk,
+	chErrorsOut chan<- LoaderError,
+) {
+	tableId := pgx.Identifier{job.Schema, job.Table}
+	colNames := Map(columns, func(col ColumnType) string {
+		return col.name
+	})
 
-	for chunk := range chChunksIn {
-		chunkStartTime := time.Now()
-		identifier := pgx.Identifier{job.Schema, job.Table}
-		colNames := Map(columns, func(col ColumnType) string {
-			return col.name
-		})
-
-		copyStartTime := time.Now()
-		_, err := db.CopyFrom(
-			ctx,
-			identifier,
-			colNames,
-			pgx.CopyFromRows(chunk.Data),
-		)
-
-		if err != nil {
-			return err
+	for {
+		if ctx.Err() != nil {
+			return
 		}
 
-		chunkCount++
-		totalRowsLoaded += len(chunk.Data)
-		copyDuration := time.Since(copyStartTime)
-		chunkDuration := time.Since(chunkStartTime)
-		rowsPerSec := float64(len(chunk.Data)) / chunkDuration.Seconds()
+		select {
+		case <-ctx.Done():
+			return
+		case chunk, ok := <-chChunksIn:
+			if !ok {
+				return
+			}
 
-		log.Infof("Loaded chunk #%d: %d rows in %v (copy: %v, %.0f rows/sec) - Total: %d rows", chunkCount, len(chunk.Data), chunkDuration, copyDuration, rowsPerSec, totalRowsLoaded)
+			if abort := loadChunkPostgres(ctx, db, tableId, colNames, chunk, chErrorsOut); abort {
+				return
+			}
+		}
+	}
+}
+
+func loadChunkPostgres(
+	ctx context.Context,
+	db *pgxpool.Pool,
+	identifier pgx.Identifier,
+	colNames []string,
+	chunk Chunk,
+	chErrorsOut chan<- LoaderError,
+) (abort bool) {
+	chunkStartTime := time.Now()
+	_, err := db.CopyFrom(
+		ctx,
+		identifier,
+		colNames,
+		pgx.CopyFromRows(chunk.Data),
+	)
+
+	if err != nil {
+		select {
+		case chErrorsOut <- LoaderError{Chunk: chunk, Msg: err.Error()}:
+		case <-ctx.Done():
+			return true
+		}
+		return false
 	}
 
-	return nil
+	chunkDuration := time.Since(chunkStartTime)
+	rowsPerSec := float64(len(chunk.Data)) / chunkDuration.Seconds()
+
+	log.Infof("Loaded chunk: %d rows in %v (%.0f rows/sec)", len(chunk.Data), chunkDuration, rowsPerSec)
+
+	return false
 }
 
 func loadRowsMssql(ctx context.Context, job MigrationJob, columns []ColumnType, db *sql.DB, in <-chan []UnknownRowValues) error {