package main import ( "context" "database/sql" "fmt" "time" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgxpool" mssql "github.com/microsoft/go-mssqldb" log "github.com/sirupsen/logrus" ) func fakeLoader(job MigrationJob, columns []ColumnType, in <-chan [][]any) { for rows := range in { log.Debugf("Chunk received, loading data into...") for i, rowValues := range rows { if i%100 == 0 { logSampleRow(job, columns, rowValues, fmt.Sprintf("row %d", i)) } } } } func loadRowsPostgres(ctx context.Context, job MigrationJob, columns []ColumnType, db *pgxpool.Pool, in <-chan []UnknownRowValues) error { chunkCount := 0 totalRowsLoaded := 0 for rows := range in { 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(rows), ) if err != nil { return err } chunkCount++ totalRowsLoaded += len(rows) copyDuration := time.Since(copyStartTime) chunkDuration := time.Since(chunkStartTime) rowsPerSec := float64(len(rows)) / chunkDuration.Seconds() log.Infof("Loaded chunk #%d: %d rows in %v (copy: %v, %.0f rows/sec) - Total: %d rows", chunkCount, len(rows), chunkDuration, copyDuration, rowsPerSec, totalRowsLoaded) } return nil } func loadRowsMssql(ctx context.Context, job MigrationJob, columns []ColumnType, db *sql.DB, in <-chan []UnknownRowValues) error { chunkCount := 0 totalRowsLoaded := 0 for rows := range in { chunkStartTime := time.Now() tx, err := db.BeginTx(ctx, nil) if err != nil { return fmt.Errorf("error starting transaction: %w", err) } fullTableName := fmt.Sprintf("[%s].[%s]", job.Schema, job.Table) colNames := Map(columns, func(col ColumnType) string { return col.name }) stmt, err := tx.PrepareContext(ctx, mssql.CopyIn(fullTableName, mssql.BulkOptions{}, colNames...)) if err != nil { tx.Rollback() return fmt.Errorf("error preparing bulk copy statement: %w", err) } copyStartTime := time.Now() for _, row := range rows { _, err = stmt.ExecContext(ctx, row...) if err != nil { stmt.Close() tx.Rollback() return fmt.Errorf("error executing row insert: %w", err) } } result, err := stmt.ExecContext(ctx) if err != nil { stmt.Close() tx.Rollback() return fmt.Errorf("error flushing bulk data: %w", err) } err = stmt.Close() if err != nil { tx.Rollback() return fmt.Errorf("error closing statement: %w", err) } if err := tx.Commit(); err != nil { return fmt.Errorf("error committing transaction: %w", err) } rowsAffected, _ := result.RowsAffected() chunkCount++ totalRowsLoaded += int(rowsAffected) copyDuration := time.Since(copyStartTime) chunkDuration := time.Since(chunkStartTime) rowsPerSec := float64(len(rows)) / chunkDuration.Seconds() log.Infof("Loaded chunk #%d (MSSQL): %d rows in %v (copy: %v, %.0f rows/sec) - Total: %d rows", chunkCount, len(rows), chunkDuration, copyDuration, rowsPerSec, totalRowsLoaded) } return nil } func Map[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 }