diff --git a/cmd/go_migrate/batch-generator.go b/cmd/go_migrate/batch-generator.go new file mode 100644 index 0000000..78311a9 --- /dev/null +++ b/cmd/go_migrate/batch-generator.go @@ -0,0 +1,110 @@ +package main + +import ( + "context" + "database/sql" + "fmt" + "time" + + "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) { + query := ` +SELECT + SUM(p.rows) AS count +FROM sys.tables t +JOIN sys.schemas s ON t.schema_id = s.schema_id +JOIN sys.partitions p ON t.object_id = p.object_id +WHERE s.name = @schema AND t.name = @table AND p.index_id IN (0, 1) +GROUP BY t.name` + + ctxTimeout, cancel := context.WithTimeout(ctx, time.Second*20) + defer cancel() + + var rowsCount int64 + err := db.QueryRowContext(ctxTimeout, query, sql.Named("schema", job.Schema), sql.Named("table", job.Table)).Scan(&rowsCount) + if err != nil { + return 0, err + } + + return rowsCount, nil +} + +func calculateBatchesMssql(ctx context.Context, db *sql.DB, job MigrationJob, batchCount int64) ([]Batch, error) { + query := fmt.Sprintf(` +SELECT + MIN([%s]) AS lower_limit, + MAX([%s]) AS upper_limit +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) + + ctxTimeout, cancel := context.WithTimeout(ctx, time.Second*20) + defer cancel() + + rows, err := db.QueryContext(ctxTimeout, query, sql.Named("batchCount", batchCount)) + if err != nil { + return nil, err + } + defer rows.Close() + + batches := make([]Batch, 0, batchCount) + + for rows.Next() { + batch := Batch{ + Id: uuid.New(), + ShouldUseRange: true, + RetryCounter: 0, + IsLowerLimitInclusive: true, + } + + if err := rows.Scan(&batch.LowerLimit, &batch.UpperLimit); err != nil { + return nil, err + } + + batches = append(batches, batch) + } + + if err := rows.Err(); err != nil { + return nil, err + } + + return batches, nil +} + +func batchGeneratorMssql(ctx context.Context, db *sql.DB, job MigrationJob) ([]Batch, error) { + rowsCount, err := estimateTotalRowsMssql(ctx, db, job) + if err != nil { + return nil, err + } + + var batchCount int64 = 1 + if rowsCount > RowsPerBatch { + batchCount = rowsCount / RowsPerBatch + } else { + return []Batch{{ + Id: uuid.New(), + ShouldUseRange: false, + RetryCounter: 0, + }}, nil + } + + batches, err := calculateBatchesMssql(ctx, db, job, batchCount) + if err != nil { + return nil, err + } + + return batches, nil +} diff --git a/cmd/go_migrate/build-extract-query.go b/cmd/go_migrate/build-extract-query.go index 3a00640..a43f09e 100644 --- a/cmd/go_migrate/build-extract-query.go +++ b/cmd/go_migrate/build-extract-query.go @@ -5,30 +5,43 @@ import ( "strings" ) -func buildExtractQueryMssql(job MigrationJob, columns []ColumnType) string { - var sbColumns strings.Builder +func buildExtractQueryMssql(job MigrationJob, columns []ColumnType, includeRange bool, isMinInclusive bool) string { + var sbQuery strings.Builder + + sbQuery.WriteString("SELECT ") if len(columns) == 0 { - sbColumns.WriteString("*") + sbQuery.WriteString("*") } else { for i, col := range columns { - sbColumns.WriteString("[") - sbColumns.WriteString(col.name) - sbColumns.WriteString("]") + fmt.Fprintf(&sbQuery, "[%s]", col.name) if col.unifiedType == "GEOMETRY" { - sbColumns.WriteString(".STAsWKB() AS [") - sbColumns.WriteString(col.name) - sbColumns.WriteString("]") + fmt.Fprintf(&sbQuery, ".STAsBinary() AS [%s]", col.name) } if i < len(columns)-1 { - sbColumns.WriteString(", ") + sbQuery.WriteString(", ") } } } - return fmt.Sprintf(`SELECT %s FROM [%s].[%s] WITH (NOLOCK)`, sbColumns.String(), job.Schema, job.Table) + 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 { @@ -56,5 +69,5 @@ func buildExtractQueryPostgres(job MigrationJob, columns []ColumnType) string { } } - return fmt.Sprintf(`SELECT %s FROM "%s"."%s"`, sbColumns.String(), job.Schema, job.Table) + return fmt.Sprintf(`SELECT %s FROM "%s"."%s" ORDER BY "%s" ASC`, sbColumns.String(), job.Schema, job.Table, job.PrimaryKey) } diff --git a/cmd/go_migrate/extractor-error-handler.go b/cmd/go_migrate/extractor-error-handler.go new file mode 100644 index 0000000..7f8e7fd --- /dev/null +++ b/cmd/go_migrate/extractor-error-handler.go @@ -0,0 +1,102 @@ +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(), + } +} diff --git a/cmd/go_migrate/extractor.go b/cmd/go_migrate/extractor.go new file mode 100644 index 0000000..3756d48 --- /dev/null +++ b/cmd/go_migrate/extractor.go @@ -0,0 +1,242 @@ +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 +} diff --git a/cmd/go_migrate/inspect-columns.go b/cmd/go_migrate/inspect-columns.go index b97043c..4c07f46 100644 --- a/cmd/go_migrate/inspect-columns.go +++ b/cmd/go_migrate/inspect-columns.go @@ -242,6 +242,10 @@ ORDER BY c.column_id; return nil, fmt.Errorf("Error scanning column type results: %W", err) } + if strings.HasPrefix(column.name, "graph_id") && column.systemType == "bigint" { + continue + } + colTypes = append(colTypes, MapMssqlColumn(column)) } diff --git a/cmd/go_migrate/job-error-handler.go b/cmd/go_migrate/job-error-handler.go new file mode 100644 index 0000000..4fc89ff --- /dev/null +++ b/cmd/go_migrate/job-error-handler.go @@ -0,0 +1,47 @@ +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) + } + } +} diff --git a/cmd/go_migrate/loader-error-handler.go b/cmd/go_migrate/loader-error-handler.go new file mode 100644 index 0000000..205fa73 --- /dev/null +++ b/cmd/go_migrate/loader-error-handler.go @@ -0,0 +1,65 @@ +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 + } + } + } +} diff --git a/cmd/go_migrate/loader.go b/cmd/go_migrate/loader.go new file mode 100644 index 0000000..ac096cc --- /dev/null +++ b/cmd/go_migrate/loader.go @@ -0,0 +1,191 @@ +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)) + } + } + } +} diff --git a/cmd/go_migrate/log.go b/cmd/go_migrate/log.go index 0aff8e9..a7bad26 100644 --- a/cmd/go_migrate/log.go +++ b/cmd/go_migrate/log.go @@ -10,6 +10,8 @@ func configureLog() { log.SetFormatter(&log.TextFormatter{ FullTimestamp: true, TimestampFormat: time.StampMilli, + DisableSorting: false, + PadLevelText: true, }) - log.SetLevel(log.DebugLevel) + log.SetLevel(log.InfoLevel) } diff --git a/cmd/go_migrate/main.go b/cmd/go_migrate/main.go index f9c18b9..0f9cee7 100644 --- a/cmd/go_migrate/main.go +++ b/cmd/go_migrate/main.go @@ -1,25 +1,49 @@ package main import ( + "context" + "time" + log "github.com/sirupsen/logrus" ) type MigrationJob struct { - Schema string - Table string + Schema string + Table string + PrimaryKey string } var migrationJobs []MigrationJob = []MigrationJob{ { - Schema: "demo", - Table: "users", + 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() - log.Info("Starting migration...") - // log.Debugf("Migration jobs: %+v", migrationJobs) + 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 { @@ -30,9 +54,11 @@ func main() { defer targetDb.Close() for _, job := range migrationJobs { - log.Infof("Processing job: %+v", job) - processMigrationJob(sourceDb, targetDb, job) + log.Infof(">>> Processing job: %s.%s <<<", job.Schema, job.Table) + processMigrationJob(ctx, sourceDb, targetDb, job) } - log.Info("Migration completed successfully!") + totalDuration := time.Since(startTime) + log.Infof("=== Migration completed successfully! ===") + log.Infof("Total migration time: %v", totalDuration) } diff --git a/cmd/go_migrate/mssql-transform.go b/cmd/go_migrate/mssql-transform.go new file mode 100644 index 0000000..a5ae098 --- /dev/null +++ b/cmd/go_migrate/mssql-transform.go @@ -0,0 +1,80 @@ +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 + } +} diff --git a/cmd/go_migrate/process.go b/cmd/go_migrate/process.go index 6f1c099..aeb267f 100644 --- a/cmd/go_migrate/process.go +++ b/cmd/go_migrate/process.go @@ -1,14 +1,26 @@ package main import ( + "context" "database/sql" + "sync" + "time" + "github.com/jackc/pgx/v5/pgxpool" _ "github.com/microsoft/go-mssqldb" log "github.com/sirupsen/logrus" ) -func processMigrationJob(sourceDb *sql.DB, targetDb *pgxpool.Pool, job MigrationJob) { +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) + sourceColTypes, targetColTypes, err := GetColumnTypes(sourceDb, targetDb, job) if err != nil { log.Fatal("Unexpected error: ", err) @@ -17,18 +29,108 @@ func processMigrationJob(sourceDb *sql.DB, targetDb *pgxpool.Pool, job Migration logColumnTypes(sourceColTypes, "Source col types") logColumnTypes(targetColTypes, "Target col types") - sourceQuery := buildExtractQueryMssql(job, sourceColTypes) + jobCtx, cancel := context.WithCancel(ctx) + defer cancel() - log.Debug(sourceQuery) + batches, err := batchGeneratorMssql(jobCtx, sourceDb, job) + if err != nil { + log.Error("Unexpected error calculating batch ranges: ", err) + } - targetQuery := buildExtractQueryPostgres(job, targetColTypes) - log.Debug(targetQuery) + 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) + + var wgActiveBatches sync.WaitGroup + var wgActiveChunks sync.WaitGroup + var wgExtractors sync.WaitGroup + var wgTransformers sync.WaitGroup + var wgLoaders sync.WaitGroup + + go func() { + if err := jobErrorHandler(jobCtx, chJobErrors); err != nil { + cancel() + } + }() + + go extractorErrorHandler(jobCtx, chExtractorErrors, chBatches, chJobErrors, &wgActiveBatches) + go loaderErrorHandler(jobCtx, chLoadersErrors, chChunksTransformed, chJobErrors, &wgActiveChunks) + + maxExtractors := min(NumExtractors, len(batches)) + log.Infof("Starting %d extractors...", maxExtractors) + extractStartTime := time.Now() + + for range maxExtractors { + wgExtractors.Go(func() { + extractFromMssql(jobCtx, sourceDb, job, sourceColTypes, ChunkSize, chBatches, chChunksRaw, chExtractorErrors, chJobErrors, &wgActiveBatches) + }) + } + + wgActiveBatches.Add(len(batches)) + go func() { + for _, batch := range batches { + chBatches <- batch + } + }() + + log.Infof("Starting %d transformers...", maxExtractors) + transformStartTime := time.Now() + + for range maxExtractors { + wgTransformers.Go(func() { + transformRowsMssql(jobCtx, sourceColTypes, chChunksRaw, chChunksTransformed, chJobErrors, &wgActiveChunks) + }) + } + + log.Infof("Starting %d PostgreSQL loader(s)...", NumLoaders) + loadStartTime := time.Now() + + for range NumLoaders { + wgLoaders.Go(func() { + loadRowsPostgres(jobCtx, targetDb, job, targetColTypes, chChunksTransformed, chLoadersErrors, chJobErrors, &wgActiveChunks) + }) + } + + go func() { + wgActiveBatches.Wait() + close(chBatches) + 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)) } func logColumnTypes(columnTypes []ColumnType, label string) { - log.Info(label) + log.Debug(label) for _, col := range columnTypes { - log.Infof("%+v", col) + 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]) } } diff --git a/cmd/go_migrate/transformer.go b/cmd/go_migrate/transformer.go new file mode 100644 index 0000000..4107a0c --- /dev/null +++ b/cmd/go_migrate/transformer.go @@ -0,0 +1,149 @@ +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 +} diff --git a/config.yaml b/config.yaml new file mode 100644 index 0000000..8796758 --- /dev/null +++ b/config.yaml @@ -0,0 +1,46 @@ +max_parallel_workers: 2 + +defaults: + max_extractors: 4 + max_loaders: 8 + queue_size: 8 + chunk_size: 50000 + chunks_per_batch: 10 + truncate_target: true + truncate_method: TRUNCATE # TRUNCATE | DELETE + retry: + attempts: 3 + +jobs: + - name: demo_users + enabled: true + source: + schema: demo + table: users + primary_key: id + target: + schema: demo + table: users + max_extractors: 2 # overrides default config + max_loaders: 4 # overrides default config + queue_size: 4 # overrides default config + chunk_size: 25000 # overrides default config + chunks_per_batch: 8 # overrides default config + truncate_target: false # overrides default config + truncate_method: DELETE # overrides default config + retry: + attempts: 5 # overrides default config + pre_sql: + - "SELECT 1" + post_sql: + - "SELECT 2" + + - name: analytics_events + enabled: true + source: + schema: analytics + table: events + primary_key: ID_events + target: + schema: analytics + table: events diff --git a/go.mod b/go.mod index b6526e4..8d3a13e 100644 --- a/go.mod +++ b/go.mod @@ -3,16 +3,19 @@ module git.ksdemosapps.com/kylesoda/go-migrate go 1.25.7 require ( + github.com/gaspardle/go-mssqlclrgeo v0.0.0-20160129143314-97ceabf987a4 + github.com/goccy/go-yaml v1.19.2 + github.com/google/uuid v1.6.0 github.com/jackc/pgx/v5 v5.9.1 github.com/joho/godotenv v1.5.1 github.com/microsoft/go-mssqldb v1.9.8 github.com/sirupsen/logrus v1.9.4 + github.com/twpayne/go-geom v1.6.1 ) require ( github.com/golang-sql/civil v0.0.0-20220223132316-b832511892a9 // indirect github.com/golang-sql/sqlexp v0.1.0 // indirect - github.com/google/uuid v1.6.0 // indirect github.com/jackc/pgpassfile v1.0.0 // indirect github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect github.com/jackc/puddle/v2 v2.2.2 // indirect diff --git a/go.sum b/go.sum index 14a1c24..33febce 100644 --- a/go.sum +++ b/go.sum @@ -10,9 +10,19 @@ github.com/Azure/azure-sdk-for-go/sdk/security/keyvault/internal v1.2.0 h1:nCYfg github.com/Azure/azure-sdk-for-go/sdk/security/keyvault/internal v1.2.0/go.mod h1:ucUjca2JtSZboY8IoUqyQyuuXvwbMBVwFOm0vdQPNhA= github.com/AzureAD/microsoft-authentication-library-for-go v1.6.0 h1:XRzhVemXdgvJqCH0sFfrBUTnUJSBrBf7++ypk+twtRs= github.com/AzureAD/microsoft-authentication-library-for-go v1.6.0/go.mod h1:HKpQxkWaGLJ+D/5H8QRpyQXA1eKjxkFlOMwck5+33Jk= +github.com/DATA-DOG/go-sqlmock v1.5.2 h1:OcvFkGmslmlZibjAjaHm3L//6LiuBgolP7OputlJIzU= +github.com/DATA-DOG/go-sqlmock v1.5.2/go.mod h1:88MAG/4G7SMwSE3CeA0ZKzrT5CiOU3OJ+JlNzwDqpNU= +github.com/alecthomas/assert/v2 v2.10.0 h1:jjRCHsj6hBJhkmhznrCzoNpbA3zqy0fYiUcYZP/GkPY= +github.com/alecthomas/assert/v2 v2.10.0/go.mod h1:Bze95FyfUr7x34QZrjL+XP+0qgp/zg8yS+TtBj1WA3k= +github.com/alecthomas/repr v0.4.0 h1:GhI2A8MACjfegCPVq9f1FLvIBS+DrQ2KQBFZP1iFzXc= +github.com/alecthomas/repr v0.4.0/go.mod h1:Fr0507jx4eOXV7AlPV6AVZLYrLIuIeSOWtW57eE/O/4= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/gaspardle/go-mssqlclrgeo v0.0.0-20160129143314-97ceabf987a4 h1:4vH4+3zfwZTqoJEFw7DsTaH1V8jgVwnyeDvNi2TxzAc= +github.com/gaspardle/go-mssqlclrgeo v0.0.0-20160129143314-97ceabf987a4/go.mod h1:jlB0I5BIfcJBGdV6rRGPthSBfeY86RGkSAwcsldbHJc= +github.com/goccy/go-yaml v1.19.2 h1:PmFC1S6h8ljIz6gMRBopkjP1TVT7xuwrButHID66PoM= +github.com/goccy/go-yaml v1.19.2/go.mod h1:XBurs7gK8ATbW4ZPGKgcbrY1Br56PdM69F7LkFRi1kA= github.com/golang-jwt/jwt/v5 v5.3.1 h1:kYf81DTWFe7t+1VvL7eS+jKFVWaUnK9cB1qbwn63YCY= github.com/golang-jwt/jwt/v5 v5.3.1/go.mod h1:fxCRLWMO43lRc8nhHWY6LGqRcf+1gQWArsqaEUEa5bE= github.com/golang-sql/civil v0.0.0-20220223132316-b832511892a9 h1:au07oEsX2xN0ktxqI+Sida1w446QrXBRJ0nee3SNZlA= @@ -21,6 +31,8 @@ github.com/golang-sql/sqlexp v0.1.0 h1:ZCD6MBpcuOVfGVqsEmY5/4FtYiKz6tSyUv9LPEDei github.com/golang-sql/sqlexp v0.1.0/go.mod h1:J4ad9Vo8ZCWQ2GMrC4UCQy1JpCbwU9m3EOqtpKwwwHI= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/hexops/gotextdiff v1.0.3 h1:gitA9+qJrrTCsiCl7+kh75nPqQt1cx4ZkudSTLoUqJM= +github.com/hexops/gotextdiff v1.0.3/go.mod h1:pSWU5MAI3yDq+fZBTazCSJysOMbxWL1BSow5/V2vxeg= github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM= github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg= github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo= @@ -48,6 +60,8 @@ github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UV github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= +github.com/twpayne/go-geom v1.6.1 h1:iLE+Opv0Ihm/ABIcvQFGIiFBXd76oBIar9drAwHFhR4= +github.com/twpayne/go-geom v1.6.1/go.mod h1:Kr+Nly6BswFsKM5sd31YaoWS5PeDDH2NftJTK7Gd028= golang.org/x/crypto v0.48.0 h1:/VRzVqiRSggnhY7gNRxPauEQ5Drw9haKdM0jqfcCFts= golang.org/x/crypto v0.48.0/go.mod h1:r0kV5h3qnFPlQnBSrULhlsRfryS2pmewsg+XfMgkVos= golang.org/x/net v0.51.0 h1:94R/GTO7mt3/4wIKpcR5gkGmRLOuE/2hNGeWq/GBIFo=