refactor: rename Batch to Partition across pipeline
Rename internal types and channels from Batch to Partition for consistency with the data model.
This commit is contained in:
@@ -6,20 +6,12 @@ import (
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
|
||||
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
type Batch struct {
|
||||
Id uuid.UUID
|
||||
ParentId uuid.UUID
|
||||
LowerLimit int64
|
||||
UpperLimit int64
|
||||
IsLowerLimitInclusive bool
|
||||
ShouldUseRange bool
|
||||
RetryCounter int
|
||||
}
|
||||
|
||||
func estimateTotalRowsMssql(ctx context.Context, db *sql.DB, job MigrationJob) (int64, error) {
|
||||
func estimateTotalRowsMssql(ctx context.Context, db *sql.DB, tableInfo config.SourceTableInfo) (int64, error) {
|
||||
query := `
|
||||
SELECT
|
||||
SUM(p.rows) AS count
|
||||
@@ -33,7 +25,7 @@ GROUP BY t.name`
|
||||
defer cancel()
|
||||
|
||||
var rowsCount int64
|
||||
err := db.QueryRowContext(ctxTimeout, query, sql.Named("schema", job.Schema), sql.Named("table", job.Table)).Scan(&rowsCount)
|
||||
err := db.QueryRowContext(ctxTimeout, query, sql.Named("schema", tableInfo.Schema), sql.Named("table", tableInfo.Table)).Scan(&rowsCount)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
@@ -41,7 +33,7 @@ GROUP BY t.name`
|
||||
return rowsCount, nil
|
||||
}
|
||||
|
||||
func calculateBatchesMssql(ctx context.Context, db *sql.DB, job MigrationJob, batchCount int64) ([]Batch, error) {
|
||||
func calculateBatchesMssql(ctx context.Context, db *sql.DB, tableInfo config.SourceTableInfo, batchCount int64) ([]models.Partition, error) {
|
||||
query := fmt.Sprintf(`
|
||||
SELECT
|
||||
MIN([%s]) AS lower_limit,
|
||||
@@ -49,7 +41,13 @@ SELECT
|
||||
FROM
|
||||
(SELECT [%s], NTILE(@batchCount) OVER (ORDER BY [%s]) AS batch_id FROM [%s].[%s]) AS T
|
||||
GROUP BY batch_id
|
||||
ORDER BY batch_id`, job.PrimaryKey, job.PrimaryKey, job.PrimaryKey, job.PrimaryKey, job.Schema, job.Table)
|
||||
ORDER BY batch_id`,
|
||||
tableInfo.PrimaryKey,
|
||||
tableInfo.PrimaryKey,
|
||||
tableInfo.PrimaryKey,
|
||||
tableInfo.PrimaryKey,
|
||||
tableInfo.Schema,
|
||||
tableInfo.Table)
|
||||
|
||||
ctxTimeout, cancel := context.WithTimeout(ctx, time.Second*20)
|
||||
defer cancel()
|
||||
@@ -60,10 +58,10 @@ ORDER BY batch_id`, job.PrimaryKey, job.PrimaryKey, job.PrimaryKey, job.PrimaryK
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
batches := make([]Batch, 0, batchCount)
|
||||
batches := make([]models.Partition, 0, batchCount)
|
||||
|
||||
for rows.Next() {
|
||||
batch := Batch{
|
||||
batch := models.Partition{
|
||||
Id: uuid.New(),
|
||||
ShouldUseRange: true,
|
||||
RetryCounter: 0,
|
||||
@@ -84,24 +82,24 @@ ORDER BY batch_id`, job.PrimaryKey, job.PrimaryKey, job.PrimaryKey, job.PrimaryK
|
||||
return batches, nil
|
||||
}
|
||||
|
||||
func batchGeneratorMssql(ctx context.Context, db *sql.DB, job MigrationJob) ([]Batch, error) {
|
||||
rowsCount, err := estimateTotalRowsMssql(ctx, db, job)
|
||||
func batchGeneratorMssql(ctx context.Context, db *sql.DB, tableInfo config.SourceTableInfo, rowsPerBatch int64) ([]models.Partition, error) {
|
||||
rowsCount, err := estimateTotalRowsMssql(ctx, db, tableInfo)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var batchCount int64 = 1
|
||||
if rowsCount > RowsPerBatch {
|
||||
batchCount = rowsCount / RowsPerBatch
|
||||
if rowsCount > rowsPerBatch {
|
||||
batchCount = rowsCount / rowsPerBatch
|
||||
} else {
|
||||
return []Batch{{
|
||||
return []models.Partition{{
|
||||
Id: uuid.New(),
|
||||
ShouldUseRange: false,
|
||||
RetryCounter: 0,
|
||||
}}, nil
|
||||
}
|
||||
|
||||
batches, err := calculateBatchesMssql(ctx, db, job, batchCount)
|
||||
batches, err := calculateBatchesMssql(ctx, db, tableInfo, batchCount)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
@@ -1,73 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
func buildExtractQueryMssql(job MigrationJob, columns []ColumnType, includeRange bool, isMinInclusive bool) string {
|
||||
var sbQuery strings.Builder
|
||||
|
||||
sbQuery.WriteString("SELECT ")
|
||||
|
||||
if len(columns) == 0 {
|
||||
sbQuery.WriteString("*")
|
||||
} else {
|
||||
for i, col := range columns {
|
||||
fmt.Fprintf(&sbQuery, "[%s]", col.name)
|
||||
|
||||
if col.unifiedType == "GEOMETRY" {
|
||||
fmt.Fprintf(&sbQuery, ".STAsBinary() AS [%s]", col.name)
|
||||
}
|
||||
|
||||
if i < len(columns)-1 {
|
||||
sbQuery.WriteString(", ")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fmt.Fprintf(&sbQuery, " FROM [%s].[%s]", job.Schema, job.Table)
|
||||
|
||||
if includeRange {
|
||||
fmt.Fprintf(&sbQuery, " WHERE [%s]", job.PrimaryKey)
|
||||
if isMinInclusive {
|
||||
sbQuery.WriteString(" >=")
|
||||
} else {
|
||||
sbQuery.WriteString(" >")
|
||||
}
|
||||
|
||||
fmt.Fprintf(&sbQuery, " @min AND [%s] <= @max", job.PrimaryKey)
|
||||
}
|
||||
|
||||
fmt.Fprintf(&sbQuery, " ORDER BY [%s] ASC", job.PrimaryKey)
|
||||
|
||||
return sbQuery.String()
|
||||
}
|
||||
|
||||
func buildExtractQueryPostgres(job MigrationJob, columns []ColumnType) string {
|
||||
var sbColumns strings.Builder
|
||||
|
||||
if len(columns) == 0 {
|
||||
sbColumns.WriteString("*")
|
||||
} else {
|
||||
for i, col := range columns {
|
||||
if col.unifiedType == "GEOMETRY" {
|
||||
sbColumns.WriteString(`ST_AsEWKB("`)
|
||||
sbColumns.WriteString(col.name)
|
||||
sbColumns.WriteString(`") AS "`)
|
||||
sbColumns.WriteString(col.name)
|
||||
sbColumns.WriteString(`"`)
|
||||
} else {
|
||||
sbColumns.WriteString(`"`)
|
||||
sbColumns.WriteString(col.name)
|
||||
sbColumns.WriteString(`"`)
|
||||
}
|
||||
|
||||
if i < len(columns)-1 {
|
||||
sbColumns.WriteString(", ")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return fmt.Sprintf(`SELECT %s FROM "%s"."%s" ORDER BY "%s" ASC`, sbColumns.String(), job.Schema, job.Table, job.PrimaryKey)
|
||||
}
|
||||
@@ -1,102 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sync"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
type ExtractorError struct {
|
||||
Batch
|
||||
LastId int64
|
||||
HasLastId bool
|
||||
Msg string
|
||||
}
|
||||
|
||||
func (e *ExtractorError) Error() string {
|
||||
return e.Msg
|
||||
}
|
||||
|
||||
const maxRetryAttempts = 3
|
||||
|
||||
func extractorErrorHandler(
|
||||
ctx context.Context,
|
||||
chErrorsIn <-chan ExtractorError,
|
||||
chBatchesOut chan<- Batch,
|
||||
chJobErrorsOut chan<- JobError,
|
||||
wgActiveBatches *sync.WaitGroup,
|
||||
) {
|
||||
for {
|
||||
if ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
|
||||
case err, ok := <-chErrorsIn:
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
if err.RetryCounter >= maxRetryAttempts {
|
||||
jobError := JobError{
|
||||
ShouldCancelJob: false,
|
||||
Msg: fmt.Sprintf("batch %v reached max retries (%d)", err.Id, maxRetryAttempts),
|
||||
Prev: &err,
|
||||
}
|
||||
|
||||
select {
|
||||
case chJobErrorsOut <- jobError:
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
|
||||
wgActiveBatches.Done()
|
||||
continue
|
||||
}
|
||||
|
||||
newBatch := err.Batch
|
||||
newBatch.RetryCounter++
|
||||
|
||||
if err.HasLastId {
|
||||
newBatch.ParentId = err.Id
|
||||
newBatch.Id = uuid.New()
|
||||
newBatch.LowerLimit = err.LastId
|
||||
newBatch.IsLowerLimitInclusive = false
|
||||
}
|
||||
|
||||
select {
|
||||
case chBatchesOut <- newBatch:
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func ExtractorErrorFromLastRowMssql(lastRow UnknownRowValues, indexPrimaryKey int, batch *Batch, previousError error) ExtractorError {
|
||||
lastIdRawValue := lastRow[indexPrimaryKey]
|
||||
|
||||
lastId, ok := ToInt64(lastIdRawValue)
|
||||
if !ok {
|
||||
currentBatch := *batch
|
||||
currentBatch.RetryCounter = maxRetryAttempts
|
||||
return ExtractorError{
|
||||
Batch: currentBatch,
|
||||
HasLastId: true,
|
||||
Msg: fmt.Sprintf("Couldn't cast last id value as int: %s", previousError.Error()),
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
return ExtractorError{
|
||||
Batch: *batch,
|
||||
HasLastId: true,
|
||||
LastId: lastId,
|
||||
Msg: previousError.Error(),
|
||||
}
|
||||
}
|
||||
@@ -1,242 +0,0 @@
|
||||
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
|
||||
}
|
||||
@@ -9,6 +9,8 @@ import (
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
|
||||
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
_ "github.com/microsoft/go-mssqldb"
|
||||
log "github.com/sirupsen/logrus"
|
||||
@@ -66,7 +68,7 @@ func GetUnifiedType(systemType string) string {
|
||||
return strings.ToUpper(systemType)
|
||||
}
|
||||
|
||||
func MapPostgresColumn(column ColumnType, maxLength *int64, precision *int64, scale *int64) ColumnType {
|
||||
func MapPostgresColumn(column ColumnType, maxLength *int64, precision *int64, scale *int64) models.ColumnType {
|
||||
stringTypes := map[string]bool{
|
||||
"varchar": true, "char": true, "character": true, "text": true, "character varying": true,
|
||||
}
|
||||
@@ -108,10 +110,23 @@ func MapPostgresColumn(column ColumnType, maxLength *int64, precision *int64, sc
|
||||
|
||||
column.unifiedType = GetUnifiedType(column.systemType)
|
||||
|
||||
return column
|
||||
colType := models.NewColumnType(
|
||||
column.name,
|
||||
column.hasMaxLength,
|
||||
column.hasPrecisionScale,
|
||||
column.userType,
|
||||
column.systemType,
|
||||
column.unifiedType,
|
||||
column.nullable,
|
||||
column.maxLength,
|
||||
column.precision,
|
||||
column.scale,
|
||||
)
|
||||
|
||||
return colType
|
||||
}
|
||||
|
||||
func GetColumnTypesPostgres(db *pgxpool.Pool, migrationJob MigrationJob) ([]ColumnType, error) {
|
||||
func GetColumnTypesPostgres(db *pgxpool.Pool, tableInfo config.TargetTableInfo) ([]models.ColumnType, error) {
|
||||
query := `
|
||||
SELECT
|
||||
c.column_name AS name,
|
||||
@@ -129,13 +144,13 @@ ORDER BY c.ordinal_position;
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
|
||||
defer cancel()
|
||||
|
||||
rows, err := db.Query(ctx, query, migrationJob.Schema, migrationJob.Table)
|
||||
rows, err := db.Query(ctx, query, tableInfo.Schema, tableInfo.Table)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("Error querying column types: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var colTypes []ColumnType
|
||||
var colTypes []models.ColumnType
|
||||
|
||||
for rows.Next() {
|
||||
var column ColumnType
|
||||
@@ -161,7 +176,7 @@ ORDER BY c.ordinal_position;
|
||||
return colTypes, nil
|
||||
}
|
||||
|
||||
func MapMssqlColumn(column ColumnType) ColumnType {
|
||||
func MapMssqlColumn(column ColumnType) models.ColumnType {
|
||||
stringTypes := map[string]bool{
|
||||
"varchar": true, "char": true, "nvarchar": true, "nchar": true, "text": true, "ntext": true,
|
||||
}
|
||||
@@ -194,10 +209,23 @@ func MapMssqlColumn(column ColumnType) ColumnType {
|
||||
|
||||
column.unifiedType = GetUnifiedType(column.systemType)
|
||||
|
||||
return column
|
||||
colType := models.NewColumnType(
|
||||
column.name,
|
||||
column.hasMaxLength,
|
||||
column.hasPrecisionScale,
|
||||
column.userType,
|
||||
column.systemType,
|
||||
column.unifiedType,
|
||||
column.nullable,
|
||||
column.maxLength,
|
||||
column.precision,
|
||||
column.scale,
|
||||
)
|
||||
|
||||
return colType
|
||||
}
|
||||
|
||||
func GetColumnTypesMssql(db *sql.DB, migrationJob MigrationJob) ([]ColumnType, error) {
|
||||
func GetColumnTypesMssql(db *sql.DB, tableInfo config.SourceTableInfo) ([]models.ColumnType, error) {
|
||||
query := `
|
||||
SELECT
|
||||
c.name AS name,
|
||||
@@ -219,13 +247,13 @@ ORDER BY c.column_id;
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
|
||||
defer cancel()
|
||||
|
||||
rows, err := db.QueryContext(ctx, query, sql.Named("schema", migrationJob.Schema), sql.Named("table", migrationJob.Table))
|
||||
rows, err := db.QueryContext(ctx, query, sql.Named("schema", tableInfo.Schema), sql.Named("table", tableInfo.Table))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("Error querying column types: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var colTypes []ColumnType
|
||||
var colTypes []models.ColumnType
|
||||
|
||||
for rows.Next() {
|
||||
var column ColumnType
|
||||
@@ -252,22 +280,27 @@ ORDER BY c.column_id;
|
||||
return colTypes, nil
|
||||
}
|
||||
|
||||
func GetColumnTypes(sourceDb *sql.DB, targetDb *pgxpool.Pool, migrationJob MigrationJob) ([]ColumnType, []ColumnType, error) {
|
||||
func GetColumnTypes(
|
||||
sourceDb *sql.DB,
|
||||
targetDb *pgxpool.Pool,
|
||||
sourceTable config.SourceTableInfo,
|
||||
targetTable config.TargetTableInfo,
|
||||
) ([]models.ColumnType, []models.ColumnType, error) {
|
||||
var sourceDbErr error
|
||||
var targetDbErr error
|
||||
var sourceColTypes []ColumnType
|
||||
var targetColTypes []ColumnType
|
||||
var sourceColTypes []models.ColumnType
|
||||
var targetColTypes []models.ColumnType
|
||||
var wg sync.WaitGroup
|
||||
|
||||
wg.Go(func() {
|
||||
sourceColTypes, sourceDbErr = GetColumnTypesMssql(sourceDb, migrationJob)
|
||||
sourceColTypes, sourceDbErr = GetColumnTypesMssql(sourceDb, sourceTable)
|
||||
if sourceDbErr != nil {
|
||||
log.Error("Error (sourceDb): ", sourceDbErr)
|
||||
}
|
||||
})
|
||||
|
||||
wg.Go(func() {
|
||||
targetColTypes, targetDbErr = GetColumnTypesPostgres(targetDb, migrationJob)
|
||||
targetColTypes, targetDbErr = GetColumnTypesPostgres(targetDb, targetTable)
|
||||
if targetDbErr != nil {
|
||||
log.Error("Error (targetDb): ", targetDbErr)
|
||||
}
|
||||
|
||||
@@ -1,47 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
type JobError struct {
|
||||
ShouldCancelJob bool
|
||||
Msg string
|
||||
Prev error
|
||||
}
|
||||
|
||||
func (e *JobError) Error() string {
|
||||
if e.Prev != nil {
|
||||
return fmt.Sprintf("%s: %v", e.Msg, e.Prev)
|
||||
}
|
||||
|
||||
return e.Msg
|
||||
}
|
||||
|
||||
func jobErrorHandler(ctx context.Context, chErrorsIn <-chan JobError) error {
|
||||
for {
|
||||
if ctx.Err() != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil
|
||||
|
||||
case err, ok := <-chErrorsIn:
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
|
||||
if err.ShouldCancelJob {
|
||||
log.Error(err.Msg, " - ", err.Prev)
|
||||
return &err
|
||||
}
|
||||
|
||||
log.Error(err.Msg, " - ", err.Prev)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,65 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sync"
|
||||
)
|
||||
|
||||
type LoaderError struct {
|
||||
Chunk
|
||||
Msg string
|
||||
}
|
||||
|
||||
func (e *LoaderError) Error() string {
|
||||
return e.Msg
|
||||
}
|
||||
|
||||
func loaderErrorHandler(
|
||||
ctx context.Context,
|
||||
chErrorsIn <-chan LoaderError,
|
||||
chChunksOut chan<- Chunk,
|
||||
chJobErrorsOut chan<- JobError,
|
||||
wgActiveChunks *sync.WaitGroup,
|
||||
) {
|
||||
for {
|
||||
if ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
|
||||
case err, ok := <-chErrorsIn:
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
if err.RetryCounter >= maxRetryAttempts {
|
||||
jobError := JobError{
|
||||
ShouldCancelJob: false,
|
||||
Msg: fmt.Sprintf("chunk %v reached max retries (%d)", err.Id, maxRetryAttempts),
|
||||
Prev: &err,
|
||||
}
|
||||
|
||||
select {
|
||||
case chJobErrorsOut <- jobError:
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
|
||||
wgActiveChunks.Done()
|
||||
continue
|
||||
}
|
||||
|
||||
err.RetryCounter++
|
||||
|
||||
select {
|
||||
case chChunksOut <- err.Chunk:
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,191 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgconn"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
mssql "github.com/microsoft/go-mssqldb"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
func loadRowsPostgres(
|
||||
ctx context.Context,
|
||||
db *pgxpool.Pool,
|
||||
job MigrationJob,
|
||||
columns []ColumnType,
|
||||
chChunksIn <-chan Chunk,
|
||||
chErrorsOut chan<- LoaderError,
|
||||
chJobErrorsOut chan<- JobError,
|
||||
wgActiveChunks *sync.WaitGroup,
|
||||
) {
|
||||
tableId := pgx.Identifier{job.Schema, job.Table}
|
||||
colNames := Map(columns, func(col ColumnType) string {
|
||||
return col.name
|
||||
})
|
||||
|
||||
for {
|
||||
if ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case chunk, ok := <-chChunksIn:
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
if abort := loadChunkPostgres(ctx, db, tableId, colNames, chunk, chErrorsOut, chJobErrorsOut, wgActiveChunks); abort {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func loadChunkPostgres(
|
||||
ctx context.Context,
|
||||
db *pgxpool.Pool,
|
||||
identifier pgx.Identifier,
|
||||
colNames []string,
|
||||
chunk Chunk,
|
||||
chErrorsOut chan<- LoaderError,
|
||||
chJobErrorsOut chan<- JobError,
|
||||
wgActiveChunks *sync.WaitGroup,
|
||||
) (abort bool) {
|
||||
chunkStartTime := time.Now()
|
||||
_, err := db.CopyFrom(
|
||||
ctx,
|
||||
identifier,
|
||||
colNames,
|
||||
pgx.CopyFromRows(chunk.Data),
|
||||
)
|
||||
|
||||
if err != nil {
|
||||
var pgErr *pgconn.PgError
|
||||
if errors.As(err, &pgErr) {
|
||||
if pgErr.Code == "23505" {
|
||||
select {
|
||||
case chJobErrorsOut <- JobError{
|
||||
ShouldCancelJob: true,
|
||||
Msg: fmt.Sprintf("Fatal data integrity error in table %s", identifier.Sanitize()),
|
||||
Prev: err,
|
||||
}:
|
||||
case <-ctx.Done():
|
||||
}
|
||||
wgActiveChunks.Done()
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
select {
|
||||
case chErrorsOut <- LoaderError{Chunk: chunk, Msg: err.Error()}:
|
||||
case <-ctx.Done():
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
chunkDuration := time.Since(chunkStartTime)
|
||||
rowsPerSec := float64(len(chunk.Data)) / chunkDuration.Seconds()
|
||||
|
||||
log.Infof("Loaded chunk: %d rows in %v (%.0f rows/sec)", len(chunk.Data), chunkDuration, rowsPerSec)
|
||||
|
||||
wgActiveChunks.Done()
|
||||
return false
|
||||
}
|
||||
|
||||
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))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -2,48 +2,34 @@ package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
|
||||
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/etl/extractors"
|
||||
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/etl/loaders"
|
||||
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/etl/transformers"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
type MigrationJob struct {
|
||||
Schema string
|
||||
Table string
|
||||
PrimaryKey string
|
||||
}
|
||||
|
||||
var migrationJobs []MigrationJob = []MigrationJob{
|
||||
{
|
||||
Schema: "demo",
|
||||
Table: "users",
|
||||
PrimaryKey: "id",
|
||||
},
|
||||
{
|
||||
Schema: "analytics",
|
||||
Table: "events",
|
||||
PrimaryKey: "ID_events",
|
||||
},
|
||||
}
|
||||
|
||||
const (
|
||||
NumExtractors int = 4
|
||||
NumLoaders int = 8
|
||||
ChunkSize int = 25000
|
||||
QueueSize int = 8
|
||||
ChunksPerBatch int = 16
|
||||
RowsPerBatch int64 = int64(ChunkSize * ChunksPerBatch)
|
||||
)
|
||||
|
||||
func main() {
|
||||
configureLog()
|
||||
|
||||
migrationConfig, err := config.ReadMigrationConfig()
|
||||
if err != nil {
|
||||
log.Fatalf("error leyendo configuracion: %v", err)
|
||||
}
|
||||
|
||||
log.Debugf("Config: %+v", migrationConfig)
|
||||
|
||||
startTime := time.Now()
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
log.Info("=== Starting migration ===")
|
||||
log.Infof("Number of loaders: %d, Chunk size: %d", NumLoaders, ChunkSize)
|
||||
|
||||
sourceDb, targetDb, connError := connectToDatabases()
|
||||
if connError != nil {
|
||||
@@ -53,12 +39,94 @@ func main() {
|
||||
defer sourceDb.Close()
|
||||
defer targetDb.Close()
|
||||
|
||||
for _, job := range migrationJobs {
|
||||
log.Infof(">>> Processing job: %s.%s <<<", job.Schema, job.Table)
|
||||
processMigrationJob(ctx, sourceDb, targetDb, job)
|
||||
results := processMigrationJobs(ctx, sourceDb, targetDb, migrationConfig.Jobs, migrationConfig.MaxParallelWorkers)
|
||||
|
||||
log.Info("=== RESUMEN DE MIGRACIÓN ===")
|
||||
var totalProcessed, totalErrors int64
|
||||
|
||||
for _, res := range results {
|
||||
status := "OK"
|
||||
if res.Error != nil {
|
||||
status = "FAILED"
|
||||
}
|
||||
log.Infof("[%s] Status: %s | Read: %d | Loaded: %d | Errors: %d | Time: %v", res.JobName, status, res.RowsRead, res.RowsLoaded, res.RowsFailed, res.Duration)
|
||||
|
||||
totalProcessed += res.RowsLoaded
|
||||
if res.Error != nil {
|
||||
totalErrors++
|
||||
}
|
||||
}
|
||||
|
||||
log.Infof("Migración terminada. Tablas: %d, Errores: %d, Filas totales: %d", len(results), totalErrors, totalProcessed)
|
||||
|
||||
totalDuration := time.Since(startTime)
|
||||
log.Infof("=== Migration completed successfully! ===")
|
||||
log.Infof("Total migration time: %v", totalDuration)
|
||||
}
|
||||
|
||||
func processMigrationJobs(
|
||||
ctx context.Context,
|
||||
sourceDb *sql.DB,
|
||||
targetDb *pgxpool.Pool,
|
||||
jobs []config.Job,
|
||||
maxParallelWorkers int,
|
||||
) []JobResult {
|
||||
if len(jobs) == 0 {
|
||||
log.Info("No migration jobs configured")
|
||||
return []JobResult{}
|
||||
}
|
||||
|
||||
if maxParallelWorkers <= 0 {
|
||||
maxParallelWorkers = 1
|
||||
}
|
||||
|
||||
if maxParallelWorkers > len(jobs) {
|
||||
maxParallelWorkers = len(jobs)
|
||||
}
|
||||
|
||||
log.Infof("Starting migration with %d parallel worker(s)", maxParallelWorkers)
|
||||
|
||||
chJobResults := make(chan JobResult, len(jobs))
|
||||
chJobs := make(chan config.Job, len(jobs))
|
||||
var wgJobs sync.WaitGroup
|
||||
|
||||
extractor := extractors.NewMssqlExtractor(sourceDb)
|
||||
transformer := transformers.NewMssqlTransformer()
|
||||
loader := loaders.NewPostgresLoader(targetDb)
|
||||
|
||||
for i := range maxParallelWorkers {
|
||||
wgJobs.Go(func() {
|
||||
for job := range chJobs {
|
||||
log.Infof("[worker %d] >>> Processing job: %s.%s <<<", i, job.SourceTable.Schema, job.SourceTable.Table)
|
||||
res := processMigrationJob(
|
||||
ctx,
|
||||
sourceDb,
|
||||
targetDb,
|
||||
extractor,
|
||||
transformer,
|
||||
loader,
|
||||
job,
|
||||
)
|
||||
|
||||
chJobResults <- res
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
for _, job := range jobs {
|
||||
chJobs <- job
|
||||
}
|
||||
close(chJobs)
|
||||
|
||||
go func() {
|
||||
wgJobs.Wait()
|
||||
close(chJobResults)
|
||||
}()
|
||||
|
||||
var finalResults []JobResult
|
||||
for res := range chJobResults {
|
||||
finalResults = append(finalResults, res)
|
||||
}
|
||||
|
||||
return finalResults
|
||||
}
|
||||
|
||||
13
cmd/go_migrate/metrics.go
Normal file
13
cmd/go_migrate/metrics.go
Normal file
@@ -0,0 +1,13 @@
|
||||
package main
|
||||
|
||||
import "time"
|
||||
|
||||
type JobResult struct {
|
||||
JobName string
|
||||
StartTime time.Time
|
||||
Duration time.Duration
|
||||
RowsRead int64
|
||||
RowsLoaded int64
|
||||
RowsFailed int64
|
||||
Error error
|
||||
}
|
||||
@@ -1,80 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"errors"
|
||||
"time"
|
||||
)
|
||||
|
||||
func mssqlUuidToBigEndian(mssqlUuid []byte) ([]byte, error) {
|
||||
if len(mssqlUuid) != 16 {
|
||||
return nil, errors.New("Invalid uuid")
|
||||
}
|
||||
|
||||
pgUuid := make([]byte, 16)
|
||||
pgUuid[0], pgUuid[1], pgUuid[2], pgUuid[3] = mssqlUuid[3], mssqlUuid[2], mssqlUuid[1], mssqlUuid[0]
|
||||
pgUuid[4], pgUuid[5] = mssqlUuid[5], mssqlUuid[4]
|
||||
pgUuid[6], pgUuid[7] = mssqlUuid[7], mssqlUuid[6]
|
||||
copy(pgUuid[8:], mssqlUuid[8:])
|
||||
|
||||
return pgUuid, nil
|
||||
}
|
||||
|
||||
const sridFlag = 0x20000000
|
||||
|
||||
func wkbToEwkbWithSrid(geometry []byte, srid int) ([]byte, error) {
|
||||
if len(geometry) < 5 {
|
||||
return nil, errors.New("Invalid wkb")
|
||||
}
|
||||
|
||||
var byteOrder binary.ByteOrder
|
||||
if geometry[0] == 0 {
|
||||
byteOrder = binary.BigEndian
|
||||
} else {
|
||||
byteOrder = binary.LittleEndian
|
||||
}
|
||||
|
||||
wkbType := byteOrder.Uint32(geometry[1:5])
|
||||
if wkbType&sridFlag != 0 {
|
||||
return geometry, nil
|
||||
}
|
||||
|
||||
ewkbType := wkbType | sridFlag
|
||||
|
||||
result := make([]byte, len(geometry)+4)
|
||||
|
||||
result[0] = geometry[0]
|
||||
|
||||
byteOrder.PutUint32(result[1:5], ewkbType)
|
||||
|
||||
byteOrder.PutUint32(result[5:9], uint32(srid))
|
||||
|
||||
copy(result[9:], geometry[5:])
|
||||
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func ensureUTC(t time.Time) time.Time {
|
||||
if t.Location() == time.UTC {
|
||||
return t
|
||||
}
|
||||
|
||||
return time.Date(t.Year(), t.Month(), t.Day(), t.Hour(), t.Minute(), t.Second(), t.Nanosecond(), time.UTC)
|
||||
}
|
||||
|
||||
func ToInt64(v any) (int64, bool) {
|
||||
switch t := v.(type) {
|
||||
case int:
|
||||
return int64(t), true
|
||||
case int8:
|
||||
return int64(t), true
|
||||
case int16:
|
||||
return int64(t), true
|
||||
case int32:
|
||||
return int64(t), true
|
||||
case int64:
|
||||
return int64(t), true
|
||||
default:
|
||||
return 0, false
|
||||
}
|
||||
}
|
||||
@@ -4,10 +4,14 @@ import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"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/pgxpool"
|
||||
|
||||
_ "github.com/microsoft/go-mssqldb"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
@@ -16,14 +20,22 @@ func processMigrationJob(
|
||||
ctx context.Context,
|
||||
sourceDb *sql.DB,
|
||||
targetDb *pgxpool.Pool,
|
||||
job MigrationJob,
|
||||
) {
|
||||
jobStartTime := time.Now()
|
||||
log.Infof("Starting migration job: %s.%s [PK: %s]", job.Schema, job.Table, job.PrimaryKey)
|
||||
extractor etl.Extractor,
|
||||
transformer etl.Transformer,
|
||||
loader etl.Loader,
|
||||
job config.Job,
|
||||
) JobResult {
|
||||
result := JobResult{
|
||||
JobName: job.Name,
|
||||
StartTime: time.Now(),
|
||||
}
|
||||
|
||||
sourceColTypes, targetColTypes, err := GetColumnTypes(sourceDb, targetDb, job)
|
||||
var rowsRead, rowsLoaded, rowsFailed int64
|
||||
|
||||
sourceColTypes, targetColTypes, err := GetColumnTypes(sourceDb, targetDb, job.SourceTable, job.TargetTable)
|
||||
if err != nil {
|
||||
log.Fatal("Unexpected error: ", err)
|
||||
result.Error = err
|
||||
return result
|
||||
}
|
||||
|
||||
logColumnTypes(sourceColTypes, "Source col types")
|
||||
@@ -32,17 +44,17 @@ func processMigrationJob(
|
||||
jobCtx, cancel := context.WithCancel(ctx)
|
||||
defer cancel()
|
||||
|
||||
batches, err := batchGeneratorMssql(jobCtx, sourceDb, job)
|
||||
batches, err := batchGeneratorMssql(jobCtx, sourceDb, job.SourceTable, job.RowsPerBatch)
|
||||
if err != nil {
|
||||
log.Error("Unexpected error calculating batch ranges: ", err)
|
||||
}
|
||||
|
||||
chJobErrors := make(chan JobError, 50)
|
||||
chBatches := make(chan Batch, QueueSize)
|
||||
chExtractorErrors := make(chan ExtractorError, QueueSize)
|
||||
chChunksRaw := make(chan Chunk, QueueSize)
|
||||
chChunksTransformed := make(chan Chunk, QueueSize)
|
||||
chLoadersErrors := make(chan LoaderError, QueueSize)
|
||||
chJobErrors := make(chan custom_errors.JobError, job.QueueSize)
|
||||
chBatches := make(chan models.Partition, job.QueueSize)
|
||||
chExtractorErrors := make(chan custom_errors.ExtractorError, job.QueueSize)
|
||||
chChunksRaw := make(chan models.Batch, job.QueueSize)
|
||||
chChunksTransformed := make(chan models.Batch, job.QueueSize)
|
||||
chLoadersErrors := make(chan custom_errors.LoaderError, job.QueueSize)
|
||||
|
||||
var wgActiveBatches sync.WaitGroup
|
||||
var wgActiveChunks sync.WaitGroup
|
||||
@@ -51,21 +63,32 @@ func processMigrationJob(
|
||||
var wgLoaders sync.WaitGroup
|
||||
|
||||
go func() {
|
||||
if err := jobErrorHandler(jobCtx, chJobErrors); err != nil {
|
||||
if err := custom_errors.JobErrorHandler(jobCtx, chJobErrors); err != nil {
|
||||
cancel()
|
||||
result.Error = err
|
||||
}
|
||||
}()
|
||||
|
||||
go extractorErrorHandler(jobCtx, chExtractorErrors, chBatches, chJobErrors, &wgActiveBatches)
|
||||
go loaderErrorHandler(jobCtx, chLoadersErrors, chChunksTransformed, chJobErrors, &wgActiveChunks)
|
||||
go custom_errors.ExtractorErrorHandler(jobCtx, job.Retry.Attempts, chExtractorErrors, chBatches, chJobErrors, &wgActiveBatches)
|
||||
go custom_errors.LoaderErrorHandler(jobCtx, job.Retry.Attempts, chLoadersErrors, chChunksTransformed, chJobErrors, &wgActiveChunks)
|
||||
|
||||
maxExtractors := min(NumExtractors, len(batches))
|
||||
log.Infof("Starting %d extractors...", maxExtractors)
|
||||
extractStartTime := time.Now()
|
||||
maxExtractors := min(job.MaxExtractors, len(batches))
|
||||
log.Infof("Starting %d extractor(s)...", maxExtractors)
|
||||
|
||||
for range maxExtractors {
|
||||
wgExtractors.Go(func() {
|
||||
extractFromMssql(jobCtx, sourceDb, job, sourceColTypes, ChunkSize, chBatches, chChunksRaw, chExtractorErrors, chJobErrors, &wgActiveBatches)
|
||||
extractor.Exec(
|
||||
jobCtx,
|
||||
job.SourceTable,
|
||||
sourceColTypes,
|
||||
job.ChunkSize,
|
||||
chBatches,
|
||||
chChunksRaw,
|
||||
chExtractorErrors,
|
||||
chJobErrors,
|
||||
&wgActiveBatches,
|
||||
&rowsRead,
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
@@ -76,21 +99,35 @@ func processMigrationJob(
|
||||
}
|
||||
}()
|
||||
|
||||
log.Infof("Starting %d transformers...", maxExtractors)
|
||||
transformStartTime := time.Now()
|
||||
log.Infof("Starting %d transformer(s)...", maxExtractors)
|
||||
|
||||
for range maxExtractors {
|
||||
wgTransformers.Go(func() {
|
||||
transformRowsMssql(jobCtx, sourceColTypes, chChunksRaw, chChunksTransformed, chJobErrors, &wgActiveChunks)
|
||||
transformer.Exec(
|
||||
jobCtx,
|
||||
sourceColTypes,
|
||||
chChunksRaw,
|
||||
chChunksTransformed,
|
||||
chJobErrors,
|
||||
&wgActiveChunks,
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
log.Infof("Starting %d PostgreSQL loader(s)...", NumLoaders)
|
||||
loadStartTime := time.Now()
|
||||
log.Infof("Starting %d loader(s)...", job.MaxLoaders)
|
||||
|
||||
for range NumLoaders {
|
||||
for range job.MaxLoaders {
|
||||
wgLoaders.Go(func() {
|
||||
loadRowsPostgres(jobCtx, targetDb, job, targetColTypes, chChunksTransformed, chLoadersErrors, chJobErrors, &wgActiveChunks)
|
||||
loader.Exec(
|
||||
jobCtx,
|
||||
job.TargetTable,
|
||||
targetColTypes,
|
||||
chChunksTransformed,
|
||||
chLoadersErrors,
|
||||
chJobErrors,
|
||||
&wgActiveChunks,
|
||||
&rowsLoaded,
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
@@ -100,37 +137,37 @@ func processMigrationJob(
|
||||
close(chExtractorErrors)
|
||||
|
||||
wgExtractors.Wait()
|
||||
log.Infof("Extraction completed in %v", time.Since(extractStartTime))
|
||||
close(chChunksRaw)
|
||||
|
||||
wgTransformers.Wait()
|
||||
log.Infof("Transformation completed in %v", time.Since(transformStartTime))
|
||||
|
||||
wgActiveChunks.Wait()
|
||||
close(chChunksTransformed)
|
||||
close(chLoadersErrors)
|
||||
|
||||
wgLoaders.Wait()
|
||||
log.Infof("Loading completed in %v", time.Since(loadStartTime))
|
||||
|
||||
cancel()
|
||||
}()
|
||||
|
||||
<-jobCtx.Done()
|
||||
log.Infof("Migration job completed. Total time: %v", time.Since(jobStartTime))
|
||||
|
||||
if ctx.Err() != nil {
|
||||
result.Error = ctx.Err()
|
||||
}
|
||||
|
||||
result.Duration = time.Since(result.StartTime)
|
||||
result.RowsRead = atomic.LoadInt64(&rowsRead)
|
||||
result.RowsLoaded = atomic.LoadInt64(&rowsLoaded)
|
||||
result.RowsFailed = atomic.LoadInt64(&rowsFailed)
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
func logColumnTypes(columnTypes []ColumnType, label string) {
|
||||
func logColumnTypes(columnTypes []models.ColumnType, label string) {
|
||||
log.Debug(label)
|
||||
|
||||
for _, col := range columnTypes {
|
||||
log.Debugf("%+v", col)
|
||||
}
|
||||
}
|
||||
|
||||
func logSampleRow(job MigrationJob, columns []ColumnType, rowValues UnknownRowValues, tag string) {
|
||||
log.Infof("[%s.%s] Sample row: (%s)", job.Schema, job.Table, tag)
|
||||
for i, col := range columns {
|
||||
log.Infof("%s (%T): %v", col.Name(), rowValues[i], rowValues[i])
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,149 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
type transformerFunc func(any) (any, error)
|
||||
|
||||
type columnTransformPlan struct {
|
||||
index int
|
||||
fn transformerFunc
|
||||
}
|
||||
|
||||
func transformRowsMssql(
|
||||
ctx context.Context,
|
||||
columns []ColumnType,
|
||||
chChunksIn <-chan Chunk,
|
||||
chChunksOut chan<- Chunk,
|
||||
chJobErrorsOut chan<- JobError,
|
||||
wgActiveChunks *sync.WaitGroup,
|
||||
) {
|
||||
transformationPlan := computeTransformationPlan(columns)
|
||||
|
||||
for {
|
||||
if ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
|
||||
case chunk, ok := <-chChunksIn:
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
if len(transformationPlan) == 0 {
|
||||
select {
|
||||
case chChunksOut <- chunk:
|
||||
wgActiveChunks.Add(1)
|
||||
continue
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
chunkStartTime := time.Now()
|
||||
|
||||
err := processChunk(ctx, &chunk, transformationPlan)
|
||||
if err != nil {
|
||||
if errors.Is(err, ctx.Err()) {
|
||||
return
|
||||
}
|
||||
|
||||
select {
|
||||
case chJobErrorsOut <- JobError{ShouldCancelJob: true, Msg: "Transformation failed", Prev: err}:
|
||||
case <-ctx.Done():
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
log.Infof("Transformed chunk %s: %d rows in %v", chunk.Id, len(chunk.Data), time.Since(chunkStartTime))
|
||||
|
||||
select {
|
||||
case chChunksOut <- chunk:
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
|
||||
wgActiveChunks.Add(1)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func computeTransformationPlan(columns []ColumnType) []columnTransformPlan {
|
||||
var plan []columnTransformPlan
|
||||
|
||||
for i, col := range columns {
|
||||
switch col.SystemType() {
|
||||
case "uniqueidentifier":
|
||||
plan = append(plan, columnTransformPlan{
|
||||
index: i,
|
||||
fn: func(v any) (any, error) {
|
||||
if b, ok := v.([]byte); ok && b != nil {
|
||||
return mssqlUuidToBigEndian(b)
|
||||
}
|
||||
return v, nil
|
||||
},
|
||||
})
|
||||
|
||||
case "geometry", "geography":
|
||||
plan = append(plan, columnTransformPlan{
|
||||
index: i,
|
||||
fn: func(v any) (any, error) {
|
||||
if b, ok := v.([]byte); ok && b != nil {
|
||||
return wkbToEwkbWithSrid(b, 4326)
|
||||
}
|
||||
return v, nil
|
||||
},
|
||||
})
|
||||
|
||||
case "datetime", "datetime2":
|
||||
plan = append(plan, columnTransformPlan{
|
||||
index: i,
|
||||
fn: func(v any) (any, error) {
|
||||
if t, ok := v.(time.Time); ok {
|
||||
return ensureUTC(t), nil
|
||||
}
|
||||
return v, nil
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
return plan
|
||||
}
|
||||
|
||||
const processChunkCtxCheck = 4096
|
||||
|
||||
func processChunk(ctx context.Context, chunk *Chunk, transformationPlan []columnTransformPlan) error {
|
||||
for i, rowValues := range chunk.Data {
|
||||
if i%processChunkCtxCheck == 0 {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
for _, task := range transformationPlan {
|
||||
val := rowValues[task.index]
|
||||
if val == nil {
|
||||
continue
|
||||
}
|
||||
|
||||
transformed, err := task.fn(val)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
rowValues[task.index] = transformed
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
Reference in New Issue
Block a user