go-migrate-db/internal/app/etl/loaders/postgres.go

package loaders

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"sync/atomic"

	"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/etl"
	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgconn"
	"github.com/jackc/pgx/v5/pgxpool"
)

type PostgresLoader struct {
	db *pgxpool.Pool
}

func NewPostgresLoader(pool *pgxpool.Pool) etl.Loader {
	return &PostgresLoader{db: pool}
}

func mapSlice[T any, V any](input []T, mapper func(T) V) []V {
	result := make([]V, len(input))

	for i, v := range input {
		result[i] = mapper(v)
	}

	return result
}

func (postgresLd *PostgresLoader) ProcessBatch(
	ctx context.Context,
	tableInfo config.TargetTableInfo,
	colNames []string,
	batch models.Batch,
) (int, error) {
	tableId := pgx.Identifier{tableInfo.Schema, tableInfo.Table}
	_, err := postgresLd.db.CopyFrom(
		ctx,
		tableId,
		colNames,
		pgx.CopyFromRows(batch.Rows),
	)

	if err != nil {
		var pgErr *pgconn.PgError
		if errors.As(err, &pgErr) {
			if pgErr.Code == "23505" {
				return 0, &custom_errors.JobError{
					ShouldCancelJob: true,
					Msg:             fmt.Sprintf("Fatal error in table %s", tableId.Sanitize()),
					Prev:            err,
				}
			}
		}

		return 0, &custom_errors.LoaderError{Batch: batch, Msg: err.Error()}
	}

	return len(batch.Rows), nil
}

func (postgresLd *PostgresLoader) Exec(
	ctx context.Context,
	tableInfo config.TargetTableInfo,
	columns []models.ColumnType,
	chBatchesIn <-chan models.Batch,
	chErrorsOut chan<- custom_errors.LoaderError,
	chJobErrorsOut chan<- custom_errors.JobError,
	wgActiveBatches *sync.WaitGroup,
	rowsLoaded *int64,
) {
	colNames := mapSlice(columns, func(col models.ColumnType) string {
		return col.Name()
	})

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

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

			processedRows, err := postgresLd.ProcessBatch(ctx, tableInfo, colNames, batch)

			if err != nil {
				var ldError *custom_errors.LoaderError
				var jobError *custom_errors.JobError
				if errors.As(err, &ldError) {
					select {
					case <-ctx.Done():
						return
					case chErrorsOut <- *ldError:
					}
				} else if errors.As(err, &jobError) {
					select {
					case <-ctx.Done():
						return
					case chJobErrorsOut <- *jobError:
					}
				} else {
					select {
					case <-ctx.Done():
						return
					case chErrorsOut <- custom_errors.LoaderError{Batch: batch, Msg: err.Error()}:
					}
				}

				continue
			}

			wgActiveBatches.Done()
			atomic.AddInt64(rowsLoaded, int64(processedRows))
		}
	}
}