137 lines
3.3 KiB
Go
137 lines
3.3 KiB
Go
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 loadRowsPostgres(ctx context.Context, job MigrationJob, columns []ColumnType, db *pgxpool.Pool, chChunksIn <-chan Chunk) error {
|
|
chunkCount := 0
|
|
totalRowsLoaded := 0
|
|
|
|
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
|
|
}
|
|
|
|
chunkCount++
|
|
totalRowsLoaded += len(chunk.Data)
|
|
copyDuration := time.Since(copyStartTime)
|
|
chunkDuration := time.Since(chunkStartTime)
|
|
rowsPerSec := float64(len(chunk.Data)) / chunkDuration.Seconds()
|
|
|
|
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)
|
|
}
|
|
|
|
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
|
|
}
|
|
|
|
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))
|
|
}
|
|
}
|
|
}
|
|
}
|