go-migrate/cmd/go_migrate/extractor.go

package main

import (
	"context"
	"database/sql"
	"errors"
	"slices"
	"strings"
	"sync"
	"time"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5/pgxpool"
	_ "github.com/microsoft/go-mssqldb"
	log "github.com/sirupsen/logrus"
)

type UnknownRowValues = []any

type Chunk struct {
	Id           uuid.UUID
	BatchId      uuid.UUID
	Data         []UnknownRowValues
	RetryCounter int
}

func extractFromMssql(
	ctx context.Context,
	db *sql.DB,
	job MigrationJob,
	columns []ColumnType,
	chunkSize int,
	chBatchesIn <-chan Batch,
	chChunksOut chan<- Chunk,
	chErrorsOut chan<- ExtractorError,
	chJobErrorsOut chan<- JobError,
	wgActiveBatches *sync.WaitGroup,
) {
	indexPrimaryKey := slices.IndexFunc(columns, func(col ColumnType) bool {
		return strings.EqualFold(col.name, job.PrimaryKey)
	})

	if indexPrimaryKey == -1 {
		jobError := JobError{
			ShouldCancelJob: true,
			Msg:             "Primary key not found in provided columns",
		}

		select {
		case <-ctx.Done():
			return
		case chJobErrorsOut <- jobError:
		}

		return
	}

	for {
		if ctx.Err() != nil {
			return
		}

		select {
		case <-ctx.Done():
			return
		case batch, ok := <-chBatchesIn:
			if !ok {
				return
			}

			if abort := processBatch(ctx, db, job, columns, chunkSize, batch, indexPrimaryKey, chChunksOut, chErrorsOut, wgActiveBatches); abort {
				return
			}
		}
	}
}

func processBatch(
	ctx context.Context,
	db *sql.DB,
	job MigrationJob,
	columns []ColumnType,
	chunkSize int,
	batch Batch,
	indexPrimaryKey int,
	chChunksOut chan<- Chunk,
	chErrorsOut chan<- ExtractorError,
	wgActiveBatches *sync.WaitGroup,
) (abort bool) {
	query := buildExtractQueryMssql(job, columns, batch.ShouldUseRange, batch.IsLowerLimitInclusive)
	log.Debug("Query used to extract data from mssql: ", query)

	var queryArgs []any
	if batch.ShouldUseRange {
		queryArgs = append(queryArgs,
			sql.Named("min", batch.LowerLimit),
			sql.Named("max", batch.UpperLimit),
		)
	}

	queryStartTime := time.Now()
	rows, err := db.QueryContext(ctx, query, queryArgs...)
	if err != nil {
		select {
		case chErrorsOut <- ExtractorError{Batch: batch, HasLastId: false, Msg: err.Error()}:
		case <-ctx.Done():
			return true
		}
		return false
	}
	defer rows.Close()
	log.Debugf("Query executed in %v", time.Since(queryStartTime))

	rowsChunk := make([]UnknownRowValues, 0, chunkSize)
	totalRowsExtracted := 0
	chunkStartTime := time.Now()

	for rows.Next() {
		values := make([]any, len(columns))
		scanArgs := make([]any, len(columns))

		for i := range values {
			scanArgs[i] = &values[i]
		}

		if err := rows.Scan(scanArgs...); err != nil {
			if len(rowsChunk) == 0 {
				select {
				case chErrorsOut <- ExtractorError{Batch: batch, HasLastId: false, Msg: err.Error()}:
				case <-ctx.Done():
					return true
				}
				return false
			}

			lastRow := rowsChunk[len(rowsChunk)-1]
			select {
			case chErrorsOut <- ExtractorErrorFromLastRowMssql(lastRow, indexPrimaryKey, &batch, err):
			case <-ctx.Done():
				return true
			}

			select {
			case chChunksOut <- Chunk{Id: uuid.New(), BatchId: batch.Id, Data: rowsChunk, RetryCounter: 0}:
			case <-ctx.Done():
				return true
			}

			return false
		}

		rowsChunk = append(rowsChunk, values)
		totalRowsExtracted++

		if len(rowsChunk) >= chunkSize {
			chunkDuration := time.Since(chunkStartTime)
			rowsPerSec := float64(chunkSize) / chunkDuration.Seconds()
			log.Infof("Extracted chunk: %d rows in %v (%.0f rows/sec) - Total: %d rows", len(rowsChunk), chunkDuration, rowsPerSec, totalRowsExtracted)

			select {
			case chChunksOut <- Chunk{Id: uuid.New(), BatchId: batch.Id, Data: rowsChunk, RetryCounter: 0}:
			case <-ctx.Done():
				return true
			}

			rowsChunk = make([]UnknownRowValues, 0, chunkSize)
			chunkStartTime = time.Now()
		}
	}

	if err := rows.Err(); err != nil {
		if errors.Is(err, ctx.Err()) {
			return true
		}

		if len(rowsChunk) == 0 {
			select {
			case chErrorsOut <- ExtractorError{Batch: batch, HasLastId: false, Msg: err.Error()}:
			case <-ctx.Done():
				return true
			}
			return false
		}

		lastRow := rowsChunk[len(rowsChunk)-1]
		select {
		case chErrorsOut <- ExtractorErrorFromLastRowMssql(lastRow, indexPrimaryKey, &batch, err):
		case <-ctx.Done():
			return true
		}
		return false
	}

	if len(rowsChunk) > 0 {
		chunkDuration := time.Since(chunkStartTime)
		rowsPerSec := float64(len(rowsChunk)) / chunkDuration.Seconds()
		log.Infof("Extracted final chunk: %d rows in %v (%.0f rows/sec) - Total: %d rows", len(rowsChunk), chunkDuration, rowsPerSec, totalRowsExtracted)
		select {
		case chChunksOut <- Chunk{Id: uuid.New(), BatchId: batch.Id, Data: rowsChunk, RetryCounter: 0}:
		case <-ctx.Done():
			return true
		}
	}

	wgActiveBatches.Done()
	return false
}

func extractFromPostgres(ctx context.Context, job MigrationJob, columns []ColumnType, chunkSize int, db *pgxpool.Pool, out chan<- []UnknownRowValues) error {
	query := buildExtractQueryPostgres(job, columns)
	log.Debug("Query used to extract data from postgres: ", query)

	rows, err := db.Query(ctx, query)
	if err != nil {
		return err
	}
	defer rows.Close()

	rowsChunk := make([]UnknownRowValues, 0, chunkSize)

	for rows.Next() {
		values, err := rows.Values()
		if err != nil {
			return err
		}

		rowsChunk = append(rowsChunk, values)

		if len(rowsChunk) >= chunkSize {
			out <- rowsChunk
			rowsChunk = make([]UnknownRowValues, 0, chunkSize)
			log.Infof("Chunk send... %+v", job)
		}
	}

	if len(rowsChunk) > 0 {
		out <- rowsChunk
		log.Infof("Chunk send... %+v", job)
	}

	return nil
}