feat: enhance concurrency management by adding WaitGroup support in extractors and loaders

cmd/go_migrate/loader.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
+	"sync"
 	"time"
 
 	"github.com/jackc/pgx/v5"
@@ -19,6 +20,7 @@ func loadRowsPostgres(
 	columns []ColumnType,
 	chChunksIn <-chan Chunk,
 	chErrorsOut chan<- LoaderError,
+	wgActiveChunks *sync.WaitGroup,
 ) {
 	tableId := pgx.Identifier{job.Schema, job.Table}
 	colNames := Map(columns, func(col ColumnType) string {
@@ -38,7 +40,7 @@ func loadRowsPostgres(
 				return
 			}
 
-			if abort := loadChunkPostgres(ctx, db, tableId, colNames, chunk, chErrorsOut); abort {
+			if abort := loadChunkPostgres(ctx, db, tableId, colNames, chunk, chErrorsOut, wgActiveChunks); abort {
 				return
 			}
 		}
@@ -52,6 +54,7 @@ func loadChunkPostgres(
 	colNames []string,
 	chunk Chunk,
 	chErrorsOut chan<- LoaderError,
+	wgActiveChunks *sync.WaitGroup,
 ) (abort bool) {
 	chunkStartTime := time.Now()
 	_, err := db.CopyFrom(
@@ -75,6 +78,7 @@ func loadChunkPostgres(
 
 	log.Infof("Loaded chunk: %d rows in %v (%.0f rows/sec)", len(chunk.Data), chunkDuration, rowsPerSec)
 
+	wgActiveChunks.Done()
 	return false
 }