From 7f3d2b8cc4516f3e6754a2ac77428c8c8bb73472 Mon Sep 17 00:00:00 2001 From: Kylesoda <249518290+kylesoda@users.noreply.github.com> Date: Sat, 11 Apr 2026 18:00:00 -0500 Subject: [PATCH] refactor: rename Batch to Partition across pipeline Rename internal types and channels from Batch to Partition for consistency with the data model. --- cmd/go_migrate/batch-generator.go | 42 ++- cmd/go_migrate/build-extract-query.go | 73 ----- cmd/go_migrate/extractor.go | 242 ---------------- cmd/go_migrate/inspect-columns.go | 63 +++- cmd/go_migrate/loader.go | 191 ------------ cmd/go_migrate/main.go | 132 +++++++-- cmd/go_migrate/metrics.go | 13 + cmd/go_migrate/process.go | 117 +++++--- config.yaml | 25 +- go.mod | 4 +- go.sum | 11 +- internal/app/config/main.go | 22 +- internal/app/config/migration.go | 125 ++++++++ internal/app/convert/main.go | 18 ++ .../app/custom_errors/extractor.error.go | 41 +-- .../app/custom_errors/job.error.go | 4 +- .../app/custom_errors/loader.error.go | 13 +- internal/app/etl/extractors/mssql.go | 273 ++++++++++++++++++ internal/app/etl/extractors/postgres.go | 128 ++++++++ internal/app/etl/extractors/types.go | 1 + internal/app/etl/loaders/postgres.go | 124 ++++++++ internal/app/etl/loaders/types.go | 1 + .../app/etl/transformers/mssql.go | 175 +++++------ internal/app/etl/transformers/types.go | 1 + .../app/etl/transformers/utils.go | 2 +- internal/app/etl/types.go | 98 +++++++ internal/app/models/colum-type.go | 70 +++++ internal/app/models/main.go | 22 ++ 28 files changed, 1250 insertions(+), 781 deletions(-) delete mode 100644 cmd/go_migrate/build-extract-query.go delete mode 100644 cmd/go_migrate/extractor.go delete mode 100644 cmd/go_migrate/loader.go create mode 100644 cmd/go_migrate/metrics.go create mode 100644 internal/app/config/migration.go create mode 100644 internal/app/convert/main.go rename cmd/go_migrate/extractor-error-handler.go => internal/app/custom_errors/extractor.error.go (56%) rename cmd/go_migrate/job-error-handler.go => internal/app/custom_errors/job.error.go (87%) rename cmd/go_migrate/loader-error-handler.go => internal/app/custom_errors/loader.error.go (79%) create mode 100644 internal/app/etl/extractors/mssql.go create mode 100644 internal/app/etl/extractors/postgres.go create mode 100644 internal/app/etl/extractors/types.go create mode 100644 internal/app/etl/loaders/postgres.go create mode 100644 internal/app/etl/loaders/types.go rename cmd/go_migrate/transformer.go => internal/app/etl/transformers/mssql.go (54%) create mode 100644 internal/app/etl/transformers/types.go rename cmd/go_migrate/mssql-transform.go => internal/app/etl/transformers/utils.go (98%) create mode 100644 internal/app/etl/types.go create mode 100644 internal/app/models/colum-type.go create mode 100644 internal/app/models/main.go diff --git a/cmd/go_migrate/batch-generator.go b/cmd/go_migrate/batch-generator.go index 78311a9..b0488b9 100644 --- a/cmd/go_migrate/batch-generator.go +++ b/cmd/go_migrate/batch-generator.go @@ -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 } diff --git a/cmd/go_migrate/build-extract-query.go b/cmd/go_migrate/build-extract-query.go deleted file mode 100644 index a43f09e..0000000 --- a/cmd/go_migrate/build-extract-query.go +++ /dev/null @@ -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) -} diff --git a/cmd/go_migrate/extractor.go b/cmd/go_migrate/extractor.go deleted file mode 100644 index 3756d48..0000000 --- a/cmd/go_migrate/extractor.go +++ /dev/null @@ -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 -} diff --git a/cmd/go_migrate/inspect-columns.go b/cmd/go_migrate/inspect-columns.go index 4c07f46..6da878c 100644 --- a/cmd/go_migrate/inspect-columns.go +++ b/cmd/go_migrate/inspect-columns.go @@ -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) } diff --git a/cmd/go_migrate/loader.go b/cmd/go_migrate/loader.go deleted file mode 100644 index ac096cc..0000000 --- a/cmd/go_migrate/loader.go +++ /dev/null @@ -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)) - } - } - } -} diff --git a/cmd/go_migrate/main.go b/cmd/go_migrate/main.go index 0f9cee7..6afdf01 100644 --- a/cmd/go_migrate/main.go +++ b/cmd/go_migrate/main.go @@ -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 +} diff --git a/cmd/go_migrate/metrics.go b/cmd/go_migrate/metrics.go new file mode 100644 index 0000000..b540c8c --- /dev/null +++ b/cmd/go_migrate/metrics.go @@ -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 +} diff --git a/cmd/go_migrate/process.go b/cmd/go_migrate/process.go index aeb267f..21b3fc7 100644 --- a/cmd/go_migrate/process.go +++ b/cmd/go_migrate/process.go @@ -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]) - } -} diff --git a/config.yaml b/config.yaml index 8796758..ceebcca 100644 --- a/config.yaml +++ b/config.yaml @@ -1,11 +1,13 @@ -max_parallel_workers: 2 +max_parallel_workers: 4 +source_db_type: sqlserver +target_db_type: postgres defaults: - max_extractors: 4 - max_loaders: 8 + max_extractors: 2 + max_loaders: 4 queue_size: 8 - chunk_size: 50000 - chunks_per_batch: 10 + chunk_size: 25000 + chunks_per_batch: 8 truncate_target: true truncate_method: TRUNCATE # TRUNCATE | DELETE retry: @@ -21,19 +23,6 @@ jobs: 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 diff --git a/go.mod b/go.mod index 8d3a13e..af6b560 100644 --- a/go.mod +++ b/go.mod @@ -4,13 +4,13 @@ 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 + gopkg.in/yaml.v3 v3.0.1 ) require ( @@ -19,6 +19,8 @@ require ( 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 + github.com/kr/text v0.2.0 // indirect + github.com/rogpeppe/go-internal v1.14.1 // indirect github.com/shopspring/decimal v1.4.0 // indirect golang.org/x/crypto v0.48.0 // indirect golang.org/x/sync v0.19.0 // indirect diff --git a/go.sum b/go.sum index 33febce..7fecac7 100644 --- a/go.sum +++ b/go.sum @@ -16,13 +16,12 @@ github.com/alecthomas/assert/v2 v2.10.0 h1:jjRCHsj6hBJhkmhznrCzoNpbA3zqy0fYiUcYZ 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/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= 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= @@ -43,6 +42,10 @@ github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= github.com/joho/godotenv v1.5.1 h1:7eLL/+HRGLY0ldzfGMeQkb7vMd0as4CfYvUVzLqw0N0= github.com/joho/godotenv v1.5.1/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwAbqwq4= +github.com/kr/pretty v0.3.0 h1:WgNl7dwNpEZ6jJ9k1snq4pZsg7DOEN8hP9Xw0Tsjwk0= +github.com/kr/pretty v0.3.0/go.mod h1:640gp4NfQd8pI5XOwp5fnNeVWj67G7CFk/SaSQn7NBk= +github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= +github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc= github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw= github.com/microsoft/go-mssqldb v1.9.8 h1:d4IFMvF/o+HdpXUqbBfzHvn/NlFA75YGcfHUUvDFJEM= @@ -51,6 +54,8 @@ github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c h1:+mdjkGKdHQG3305AYmd github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c/go.mod h1:7rwL4CYBLnjLxUqIJNnCWiEdr3bn6IUYi15bNlnbCCU= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ= +github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc= github.com/shopspring/decimal v1.4.0 h1:bxl37RwXBklmTi0C79JfXCEBD1cqqHt0bbgBAGFp81k= github.com/shopspring/decimal v1.4.0/go.mod h1:gawqmDU56v4yIKSwfBSFip1HdCCXN8/+DMd9qYNcwME= github.com/sirupsen/logrus v1.9.4 h1:TsZE7l11zFCLZnZ+teH4Umoq5BhEIfIzfRDZ1Uzql2w= @@ -73,6 +78,8 @@ golang.org/x/sys v0.41.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks= golang.org/x/text v0.34.0 h1:oL/Qq0Kdaqxa1KbNeMKwQq0reLCCaFtqu2eNuSeNHbk= golang.org/x/text v0.34.0/go.mod h1:homfLqTYRFyVYemLBFl5GgL/DWEiH5wcsQ5gSh1yziA= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/internal/app/config/main.go b/internal/app/config/main.go index 43c503c..94e751d 100644 --- a/internal/app/config/main.go +++ b/internal/app/config/main.go @@ -8,10 +8,8 @@ import ( ) type appConfig struct { - SourceDbUrl string - SourceDbType string - TargetDbUrl string - TargetDbType string + SourceDbUrl string + TargetDbUrl string } func loadEnv() { @@ -29,26 +27,14 @@ func getAppConfig() appConfig { log.Fatal("SOURCE_DB_URL environment variable not set") } - sourceDbType := os.Getenv("SOURCE_DB_TYPE") - if sourceDbType == "" { - log.Fatal("SOURCE_DB_TYPE environment variable not set") - } - targetDbUrl := os.Getenv("TARGET_DB_URL") if targetDbUrl == "" { log.Fatal("TARGET_DB_URL environment variable not set") } - targetDbType := os.Getenv("TARGET_DB_TYPE") - if targetDbType == "" { - log.Fatal("TARGET_DB_TYPE environment variable not set") - } - return appConfig{ - SourceDbUrl: sourceDbUrl, - SourceDbType: sourceDbType, - TargetDbUrl: targetDbUrl, - TargetDbType: targetDbType, + SourceDbUrl: sourceDbUrl, + TargetDbUrl: targetDbUrl, } } diff --git a/internal/app/config/migration.go b/internal/app/config/migration.go new file mode 100644 index 0000000..9bb6f4c --- /dev/null +++ b/internal/app/config/migration.go @@ -0,0 +1,125 @@ +package config + +import ( + "fmt" + "os" + + "gopkg.in/yaml.v3" +) + +type RetryConfig struct { + Attempts int `yaml:"attempts"` +} + +type JobConfig struct { + MaxExtractors int `yaml:"max_extractors"` + MaxLoaders int `yaml:"max_loaders"` + QueueSize int `yaml:"queue_size"` + ChunkSize int `yaml:"chunk_size"` + ChunksPerBatch int `yaml:"chunks_per_batch"` + RowsPerBatch int64 + TruncateTarget bool `yaml:"truncate_target"` + TruncateMethod string `yaml:"truncate_method"` + Retry RetryConfig `yaml:"retry"` +} + +type TableInfo struct { + Schema string `yaml:"schema"` + Table string `yaml:"table"` +} + +type TargetTableInfo struct { + TableInfo `yaml:",inline"` +} + +type SourceTableInfo struct { + TableInfo `yaml:",inline"` + PrimaryKey string `yaml:"primary_key"` +} + +type Job struct { + Name string `yaml:"name"` + Enabled bool `yaml:"enabled"` + SourceTable SourceTableInfo `yaml:"source"` + TargetTable TargetTableInfo `yaml:"target"` + PreSQL []string `yaml:"pre_sql"` + PostSQL []string `yaml:"post_sql"` + JobConfig `yaml:",inline"` +} + +type MigrationConfig struct { + MaxParallelWorkers int `yaml:"max_parallel_workers"` + SourceDbType string `yaml:"source_db_type"` + TargetDbType string `yaml:"target_db_type"` + Defaults JobConfig `yaml:"defaults"` + Jobs []Job `yaml:"jobs"` +} + +type rawConfig struct { + MaxParallelWorkers int `yaml:"max_parallel_workers"` + SourceDbType string `yaml:"source_db_type"` + TargetDbType string `yaml:"target_db_type"` + Defaults JobConfig `yaml:"defaults"` + Jobs []yaml.Node `yaml:"jobs"` +} + +func (c *MigrationConfig) UnmarshalYAML(value *yaml.Node) error { + var raw rawConfig + if err := value.Decode(&raw); err != nil { + return err + } + + c.MaxParallelWorkers = raw.MaxParallelWorkers + c.Defaults = raw.Defaults + c.Defaults.RowsPerBatch = int64(raw.Defaults.ChunkSize * raw.Defaults.ChunksPerBatch) + + for _, node := range raw.Jobs { + job := Job{ + JobConfig: raw.Defaults, + } + + if err := node.Decode(&job); err != nil { + return err + } + + job.RowsPerBatch = int64(job.ChunkSize * job.ChunksPerBatch) + + c.Jobs = append(c.Jobs, job) + } + + return nil +} + +const defaultConfigFileName string = "config.yaml" + +func filenamesOrDefault(filenames []string) []string { + if len(filenames) == 0 { + return []string{defaultConfigFileName} + } + return filenames +} + +func ReadMigrationConfig(filenames ...string) (MigrationConfig, error) { + filenames = filenamesOrDefault(filenames) + var data []byte + var err error + + for _, filename := range filenames { + data, err = os.ReadFile(filename) + if err != nil { + continue + } + break + } + + if err != nil { + return MigrationConfig{}, fmt.Errorf("Error reading config file: %v", err) + } + + var config MigrationConfig + if err := yaml.Unmarshal(data, &config); err != nil { + return MigrationConfig{}, fmt.Errorf("Error parsing config file: %v", err) + } + + return config, nil +} diff --git a/internal/app/convert/main.go b/internal/app/convert/main.go new file mode 100644 index 0000000..294d0dc --- /dev/null +++ b/internal/app/convert/main.go @@ -0,0 +1,18 @@ +package convert + +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/extractor-error-handler.go b/internal/app/custom_errors/extractor.error.go similarity index 56% rename from cmd/go_migrate/extractor-error-handler.go rename to internal/app/custom_errors/extractor.error.go index 7f8e7fd..ad4ccf0 100644 --- a/cmd/go_migrate/extractor-error-handler.go +++ b/internal/app/custom_errors/extractor.error.go @@ -1,15 +1,16 @@ -package main +package custom_errors import ( "context" "fmt" "sync" + "git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models" "github.com/google/uuid" ) type ExtractorError struct { - Batch + Batch models.Partition LastId int64 HasLastId bool Msg string @@ -19,12 +20,11 @@ func (e *ExtractorError) Error() string { return e.Msg } -const maxRetryAttempts = 3 - -func extractorErrorHandler( +func ExtractorErrorHandler( ctx context.Context, + maxRetryAttempts int, chErrorsIn <-chan ExtractorError, - chBatchesOut chan<- Batch, + chBatchesOut chan<- models.Partition, chJobErrorsOut chan<- JobError, wgActiveBatches *sync.WaitGroup, ) { @@ -42,10 +42,10 @@ func extractorErrorHandler( return } - if err.RetryCounter >= maxRetryAttempts { + if err.Batch.RetryCounter >= maxRetryAttempts { jobError := JobError{ ShouldCancelJob: false, - Msg: fmt.Sprintf("batch %v reached max retries (%d)", err.Id, maxRetryAttempts), + Msg: fmt.Sprintf("batch %v reached max retries (%d)", err.Batch.Id, maxRetryAttempts), Prev: &err, } @@ -63,7 +63,7 @@ func extractorErrorHandler( newBatch.RetryCounter++ if err.HasLastId { - newBatch.ParentId = err.Id + newBatch.ParentId = err.Batch.Id newBatch.Id = uuid.New() newBatch.LowerLimit = err.LastId newBatch.IsLowerLimitInclusive = false @@ -77,26 +77,3 @@ func extractorErrorHandler( } } } - -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/job-error-handler.go b/internal/app/custom_errors/job.error.go similarity index 87% rename from cmd/go_migrate/job-error-handler.go rename to internal/app/custom_errors/job.error.go index 4fc89ff..ca359af 100644 --- a/cmd/go_migrate/job-error-handler.go +++ b/internal/app/custom_errors/job.error.go @@ -1,4 +1,4 @@ -package main +package custom_errors import ( "context" @@ -21,7 +21,7 @@ func (e *JobError) Error() string { return e.Msg } -func jobErrorHandler(ctx context.Context, chErrorsIn <-chan JobError) error { +func JobErrorHandler(ctx context.Context, chErrorsIn <-chan JobError) error { for { if ctx.Err() != nil { return nil diff --git a/cmd/go_migrate/loader-error-handler.go b/internal/app/custom_errors/loader.error.go similarity index 79% rename from cmd/go_migrate/loader-error-handler.go rename to internal/app/custom_errors/loader.error.go index 205fa73..4f03e9d 100644 --- a/cmd/go_migrate/loader-error-handler.go +++ b/internal/app/custom_errors/loader.error.go @@ -1,13 +1,15 @@ -package main +package custom_errors import ( "context" "fmt" "sync" + + "git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models" ) type LoaderError struct { - Chunk + models.Batch Msg string } @@ -15,10 +17,11 @@ func (e *LoaderError) Error() string { return e.Msg } -func loaderErrorHandler( +func LoaderErrorHandler( ctx context.Context, + maxRetryAttempts int, chErrorsIn <-chan LoaderError, - chChunksOut chan<- Chunk, + chChunksOut chan<- models.Batch, chJobErrorsOut chan<- JobError, wgActiveChunks *sync.WaitGroup, ) { @@ -56,7 +59,7 @@ func loaderErrorHandler( err.RetryCounter++ select { - case chChunksOut <- err.Chunk: + case chChunksOut <- err.Batch: case <-ctx.Done(): return } diff --git a/internal/app/etl/extractors/mssql.go b/internal/app/etl/extractors/mssql.go new file mode 100644 index 0000000..629377d --- /dev/null +++ b/internal/app/etl/extractors/mssql.go @@ -0,0 +1,273 @@ +package extractors + +import ( + "context" + "database/sql" + "errors" + "fmt" + "slices" + "strings" + "sync" + "sync/atomic" + + "git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config" + "git.ksdemosapps.com/kylesoda/go-migrate/internal/app/convert" + "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/google/uuid" +) + +type MssqlExtractor struct { + db *sql.DB +} + +func NewMssqlExtractor(db *sql.DB) etl.Extractor { + return &MssqlExtractor{db: db} +} + +func buildExtractQueryMssql( + tableInfo config.SourceTableInfo, + columns []models.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.Type() == "GEOMETRY" { + fmt.Fprintf(&sbQuery, ".STAsBinary() AS [%s]", col.Name()) + } + + if i < len(columns)-1 { + sbQuery.WriteString(", ") + } + } + } + + fmt.Fprintf(&sbQuery, " FROM [%s].[%s]", tableInfo.Schema, tableInfo.Table) + + if includeRange { + fmt.Fprintf(&sbQuery, " WHERE [%s]", tableInfo.PrimaryKey) + if isMinInclusive { + sbQuery.WriteString(" >=") + } else { + sbQuery.WriteString(" >") + } + + fmt.Fprintf(&sbQuery, " @min AND [%s] <= @max", tableInfo.PrimaryKey) + } + + fmt.Fprintf(&sbQuery, " ORDER BY [%s] ASC", tableInfo.PrimaryKey) + + return sbQuery.String() +} + +func extractorErrorFromLastRowMssql( + lastRow models.UnknownRowValues, + indexPrimaryKey int, + batch *models.Partition, + previousError error, +) *custom_errors.ExtractorError { + lastIdRawValue := lastRow[indexPrimaryKey] + + lastId, ok := convert.ToInt64(lastIdRawValue) + if !ok { + currentBatch := *batch + currentBatch.RetryCounter = 3 + return &custom_errors.ExtractorError{ + Batch: currentBatch, + HasLastId: true, + Msg: fmt.Sprintf("Couldn't cast last id value as int: %s", previousError.Error()), + } + + } + + return &custom_errors.ExtractorError{ + Batch: *batch, + HasLastId: true, + LastId: lastId, + Msg: previousError.Error(), + } +} + +func (mssqlEx *MssqlExtractor) ProcessBatch( + ctx context.Context, + tableInfo config.SourceTableInfo, + columns []models.ColumnType, + chunkSize int, + batch models.Partition, + indexPrimaryKey int, + chChunksOut chan<- models.Batch, + rowsRead *int64, +) error { + query := buildExtractQueryMssql(tableInfo, columns, batch.ShouldUseRange, batch.IsLowerLimitInclusive) + + var queryArgs []any + if batch.ShouldUseRange { + queryArgs = append(queryArgs, + sql.Named("min", batch.LowerLimit), + sql.Named("max", batch.UpperLimit), + ) + } + + rows, err := mssqlEx.db.QueryContext(ctx, query, queryArgs...) + if err != nil { + return &custom_errors.ExtractorError{Batch: batch, HasLastId: false, Msg: err.Error()} + } + defer rows.Close() + + rowsChunk := make([]models.UnknownRowValues, 0, chunkSize) + + 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 { + return &custom_errors.ExtractorError{Batch: batch, HasLastId: false, Msg: err.Error()} + } + + lastRow := rowsChunk[len(rowsChunk)-1] + + select { + case chChunksOut <- models.Batch{Id: uuid.New(), PartitionId: batch.Id, Data: rowsChunk, RetryCounter: 0}: + case <-ctx.Done(): + return nil + } + + atomic.AddInt64(rowsRead, int64(len(rowsChunk))) + + return extractorErrorFromLastRowMssql(lastRow, indexPrimaryKey, &batch, err) + } + + rowsChunk = append(rowsChunk, values) + + if len(rowsChunk) >= chunkSize { + select { + case chChunksOut <- models.Batch{Id: uuid.New(), PartitionId: batch.Id, Data: rowsChunk, RetryCounter: 0}: + case <-ctx.Done(): + return nil + } + + atomic.AddInt64(rowsRead, int64(len(rowsChunk))) + rowsChunk = make([]models.UnknownRowValues, 0, chunkSize) + } + } + + if err := rows.Err(); err != nil { + if errors.Is(err, ctx.Err()) { + return ctx.Err() + } + + if len(rowsChunk) == 0 { + return &custom_errors.ExtractorError{Batch: batch, HasLastId: false, Msg: err.Error()} + } + + lastRow := rowsChunk[len(rowsChunk)-1] + return extractorErrorFromLastRowMssql(lastRow, indexPrimaryKey, &batch, err) + } + + if len(rowsChunk) > 0 { + select { + case chChunksOut <- models.Batch{Id: uuid.New(), PartitionId: batch.Id, Data: rowsChunk, RetryCounter: 0}: + case <-ctx.Done(): + return nil + } + + atomic.AddInt64(rowsRead, int64(len(rowsChunk))) + } + + return nil +} + +func (mssqlEx *MssqlExtractor) Exec( + ctx context.Context, + tableInfo config.SourceTableInfo, + columns []models.ColumnType, + chunkSize int, + chBatchesIn <-chan models.Partition, + chChunksOut chan<- models.Batch, + chErrorsOut chan<- custom_errors.ExtractorError, + chJobErrorsOut chan<- custom_errors.JobError, + wgActiveBatches *sync.WaitGroup, + rowsRead *int64, +) { + indexPrimaryKey := slices.IndexFunc(columns, func(col models.ColumnType) bool { + return strings.EqualFold(col.Name(), tableInfo.PrimaryKey) + }) + + if indexPrimaryKey == -1 { + select { + case <-ctx.Done(): + return + case chJobErrorsOut <- custom_errors.JobError{ + ShouldCancelJob: true, + Msg: "Primary key not found in provided columns", + }: + } + + return + } + + for { + if ctx.Err() != nil { + return + } + + select { + case <-ctx.Done(): + return + case batch, ok := <-chBatchesIn: + if !ok { + return + } + + err := mssqlEx.ProcessBatch( + ctx, + tableInfo, + columns, + chunkSize, + batch, + indexPrimaryKey, + chChunksOut, + rowsRead, + ) + + if err != nil { + var exError *custom_errors.ExtractorError + if errors.As(err, &exError) { + select { + case <-ctx.Done(): + return + case chErrorsOut <- *exError: + } + } + + var jobError *custom_errors.JobError + if errors.As(err, &jobError) { + select { + case <-ctx.Done(): + return + case chJobErrorsOut <- *jobError: + } + } + + return + } + + wgActiveBatches.Done() + } + } +} diff --git a/internal/app/etl/extractors/postgres.go b/internal/app/etl/extractors/postgres.go new file mode 100644 index 0000000..374b5da --- /dev/null +++ b/internal/app/etl/extractors/postgres.go @@ -0,0 +1,128 @@ +package extractors + +import ( + "context" + "errors" + "fmt" + "strings" + "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/google/uuid" + "github.com/jackc/pgx/v5/pgxpool" +) + +type PostgresExtractor struct { + db *pgxpool.Pool +} + +func NewPostgresExtractor(pool *pgxpool.Pool) etl.Extractor { + return &PostgresExtractor{db: pool} +} + +func buildExtractQueryPostgres(sourceDbInfo config.SourceTableInfo, columns []models.ColumnType) string { + var sbColumns strings.Builder + + if len(columns) == 0 { + sbColumns.WriteString("*") + } else { + for i, col := range columns { + if col.Type() == "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(), sourceDbInfo.Schema, sourceDbInfo.Table, sourceDbInfo.PrimaryKey) +} + +func (postgresEx *PostgresExtractor) ProcessBatch( + ctx context.Context, + tableInfo config.SourceTableInfo, + columns []models.ColumnType, + chunkSize int, + batch models.Partition, + indexPrimaryKey int, + chChunksOut chan<- models.Batch, + rowsRead *int64, +) error { + query := buildExtractQueryPostgres(tableInfo, columns) + + if batch.ShouldUseRange { + return errors.New("Batch config not yet supported") + } + + rows, err := postgresEx.db.Query(ctx, query) + if err != nil { + return &custom_errors.ExtractorError{Batch: batch, HasLastId: false, Msg: err.Error()} + } + defer rows.Close() + + rowsChunk := make([]models.UnknownRowValues, 0, chunkSize) + + for rows.Next() { + values, err := rows.Values() + if err != nil { + return errors.New("Unexpected error reading rows from source") + } + + rowsChunk = append(rowsChunk, values) + + if len(rowsChunk) >= chunkSize { + select { + case chChunksOut <- models.Batch{Id: uuid.New(), PartitionId: batch.Id, Data: rowsChunk, RetryCounter: 0}: + case <-ctx.Done(): + return nil + } + + atomic.AddInt64(rowsRead, int64(len(rowsChunk))) + rowsChunk = make([]models.UnknownRowValues, 0, chunkSize) + } + } + + if err := rows.Err(); err != nil { + return errors.New("Unexpected error reading rows from source") + } + + if len(rowsChunk) > 0 { + select { + case chChunksOut <- models.Batch{Id: uuid.New(), PartitionId: batch.Id, Data: rowsChunk, RetryCounter: 0}: + case <-ctx.Done(): + return nil + } + + atomic.AddInt64(rowsRead, int64(len(rowsChunk))) + } + + return nil +} + +func (postgresEx *PostgresExtractor) Exec( + ctx context.Context, + tableInfo config.SourceTableInfo, + columns []models.ColumnType, + chunkSize int, + chBatchesIn <-chan models.Partition, + chChunksOut chan<- models.Batch, + chErrorsOut chan<- custom_errors.ExtractorError, + chJobErrorsOut chan<- custom_errors.JobError, + wgActiveBatches *sync.WaitGroup, + rowsRead *int64, +) { +} diff --git a/internal/app/etl/extractors/types.go b/internal/app/etl/extractors/types.go new file mode 100644 index 0000000..85defa0 --- /dev/null +++ b/internal/app/etl/extractors/types.go @@ -0,0 +1 @@ +package extractors diff --git a/internal/app/etl/loaders/postgres.go b/internal/app/etl/loaders/postgres.go new file mode 100644 index 0000000..c9d2276 --- /dev/null +++ b/internal/app/etl/loaders/postgres.go @@ -0,0 +1,124 @@ +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) ProcessChunk( + ctx context.Context, + tableInfo config.TargetTableInfo, + colNames []string, + chunk models.Batch, +) (int, error) { + tableId := pgx.Identifier{tableInfo.Schema, tableInfo.Table} + _, err := postgresLd.db.CopyFrom( + ctx, + tableId, + colNames, + pgx.CopyFromRows(chunk.Data), + ) + + 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: chunk, Msg: err.Error()} + } + + return len(chunk.Data), nil +} + +func (postgresLd *PostgresLoader) Exec( + ctx context.Context, + tableInfo config.TargetTableInfo, + columns []models.ColumnType, + chChunksIn <-chan models.Batch, + chErrorsOut chan<- custom_errors.LoaderError, + chJobErrorsOut chan<- custom_errors.JobError, + wgActiveChunks *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 chunk, ok := <-chChunksIn: + if !ok { + return + } + + processedRows, err := postgresLd.ProcessChunk(ctx, tableInfo, colNames, chunk) + + if err != nil { + var ldError *custom_errors.LoaderError + if errors.As(err, &ldError) { + select { + case <-ctx.Done(): + return + case chErrorsOut <- *ldError: + } + } + + var jobError *custom_errors.JobError + if errors.As(err, &jobError) { + select { + case <-ctx.Done(): + return + case chJobErrorsOut <- *jobError: + } + } + + return + } + + wgActiveChunks.Done() + atomic.AddInt64(rowsLoaded, int64(processedRows)) + } + } +} diff --git a/internal/app/etl/loaders/types.go b/internal/app/etl/loaders/types.go new file mode 100644 index 0000000..c88d5fb --- /dev/null +++ b/internal/app/etl/loaders/types.go @@ -0,0 +1 @@ +package loaders diff --git a/cmd/go_migrate/transformer.go b/internal/app/etl/transformers/mssql.go similarity index 54% rename from cmd/go_migrate/transformer.go rename to internal/app/etl/transformers/mssql.go index 4107a0c..23dd68e 100644 --- a/cmd/go_migrate/transformer.go +++ b/internal/app/etl/transformers/mssql.go @@ -1,4 +1,4 @@ -package main +package transformers import ( "context" @@ -6,22 +6,97 @@ import ( "sync" "time" - log "github.com/sirupsen/logrus" + "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" ) -type transformerFunc func(any) (any, error) +type MssqlTransformer struct{} -type columnTransformPlan struct { - index int - fn transformerFunc +func NewMssqlTransformer() etl.Transformer { + return &MssqlTransformer{} } -func transformRowsMssql( +func computeTransformationPlan(columns []models.ColumnType) []etl.ColumnTransformPlan { + var plan []etl.ColumnTransformPlan + + for i, col := range columns { + switch col.SystemType() { + case "uniqueidentifier": + plan = append(plan, etl.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, etl.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, etl.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 (mssqlTr *MssqlTransformer) ProcessChunk( ctx context.Context, - columns []ColumnType, - chChunksIn <-chan Chunk, - chChunksOut chan<- Chunk, - chJobErrorsOut chan<- JobError, + chunk *models.Batch, + transformationPlan []etl.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 +} + +func (mssqlTr *MssqlTransformer) Exec( + ctx context.Context, + columns []models.ColumnType, + chChunksIn <-chan models.Batch, + chChunksOut chan<- models.Batch, + chJobErrorsOut chan<- custom_errors.JobError, wgActiveChunks *sync.WaitGroup, ) { transformationPlan := computeTransformationPlan(columns) @@ -50,23 +125,19 @@ func transformRowsMssql( } } - chunkStartTime := time.Now() - - err := processChunk(ctx, &chunk, transformationPlan) + err := mssqlTr.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 chJobErrorsOut <- custom_errors.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(): @@ -77,73 +148,3 @@ func transformRowsMssql( } } } - -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/internal/app/etl/transformers/types.go b/internal/app/etl/transformers/types.go new file mode 100644 index 0000000..b3fa142 --- /dev/null +++ b/internal/app/etl/transformers/types.go @@ -0,0 +1 @@ +package transformers diff --git a/cmd/go_migrate/mssql-transform.go b/internal/app/etl/transformers/utils.go similarity index 98% rename from cmd/go_migrate/mssql-transform.go rename to internal/app/etl/transformers/utils.go index a5ae098..00b3939 100644 --- a/cmd/go_migrate/mssql-transform.go +++ b/internal/app/etl/transformers/utils.go @@ -1,4 +1,4 @@ -package main +package transformers import ( "encoding/binary" diff --git a/internal/app/etl/types.go b/internal/app/etl/types.go new file mode 100644 index 0000000..05fd79d --- /dev/null +++ b/internal/app/etl/types.go @@ -0,0 +1,98 @@ +package etl + +import ( + "context" + "sync" + + "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/models" +) + +type Extractor interface { + ProcessBatch( + ctx context.Context, + tableInfo config.SourceTableInfo, + columns []models.ColumnType, + chunkSize int, + batch models.Partition, + indexPrimaryKey int, + chChunksOut chan<- models.Batch, + rowsRead *int64, + ) error + + Exec( + ctx context.Context, + tableInfo config.SourceTableInfo, + columns []models.ColumnType, + chunkSize int, + chBatchesIn <-chan models.Partition, + chChunksOut chan<- models.Batch, + chErrorsOut chan<- custom_errors.ExtractorError, + chJobErrorsOut chan<- custom_errors.JobError, + wgActiveBatches *sync.WaitGroup, + rowsRead *int64, + ) +} + +type TransformerFunc func(any) (any, error) + +type ColumnTransformPlan struct { + Index int + Fn TransformerFunc +} + +type Transformer interface { + ProcessChunk( + ctx context.Context, + chunk *models.Batch, + transformationPlan []ColumnTransformPlan, + ) error + + Exec( + ctx context.Context, + columns []models.ColumnType, + chChunksIn <-chan models.Batch, + chChunksOut chan<- models.Batch, + chJobErrorsOut chan<- custom_errors.JobError, + wgActiveChunks *sync.WaitGroup, + ) +} + +type Loader interface { + ProcessChunk( + ctx context.Context, + tableInfo config.TargetTableInfo, + colNames []string, + chunk models.Batch, + ) (int, error) + + Exec( + ctx context.Context, + tableInfo config.TargetTableInfo, + columns []models.ColumnType, + chChunksIn <-chan models.Batch, + chErrorsOut chan<- custom_errors.LoaderError, + chJobErrorsOut chan<- custom_errors.JobError, + wgActiveChunks *sync.WaitGroup, + rowsLoaded *int64, + ) +} + +type TableAnalizer interface { + QueryColumnTypes( + ctx context.Context, + tableInfo config.TableInfo, + ) ([]models.ColumnType, error) + + EstimateTotalRows( + ctx context.Context, + tableInfo config.TableInfo, + ) (int64, error) + + CalculatePartitionRanges( + ctx context.Context, + tableInfo config.TableInfo, + totalPartitions int, + ) (models.Partition, error) +} diff --git a/internal/app/models/colum-type.go b/internal/app/models/colum-type.go new file mode 100644 index 0000000..81e6a62 --- /dev/null +++ b/internal/app/models/colum-type.go @@ -0,0 +1,70 @@ +package models + +type ColumnType struct { + name string + + hasMaxLength bool + hasPrecisionScale bool + + userType string + systemType string + unifiedType string + nullable bool + maxLength int64 + precision int64 + scale int64 +} + +func (c *ColumnType) Name() string { + return c.name +} + +func (c *ColumnType) UserType() string { + return c.userType +} + +func (c *ColumnType) SystemType() string { + return c.systemType +} + +func (c *ColumnType) Length() (length int64, ok bool) { + return c.maxLength, c.hasMaxLength +} + +func (c *ColumnType) DecimalSize() (precision, scale int64, ok bool) { + return c.precision, c.scale, c.hasPrecisionScale +} + +func (c *ColumnType) Nullable() bool { + return c.nullable +} + +func (c *ColumnType) Type() string { + return c.unifiedType +} + +func NewColumnType( + name string, + hasMaxLength bool, + hasPrecisionScale bool, + userType string, + systemType string, + unifiedType string, + nullable bool, + maxLength int64, + precision int64, + scale int64, +) ColumnType { + return ColumnType{ + name, + hasMaxLength, + hasPrecisionScale, + userType, + systemType, + unifiedType, + nullable, + maxLength, + precision, + scale, + } +} diff --git a/internal/app/models/main.go b/internal/app/models/main.go new file mode 100644 index 0000000..6be796e --- /dev/null +++ b/internal/app/models/main.go @@ -0,0 +1,22 @@ +package models + +import "github.com/google/uuid" + +type UnknownRowValues = []any + +type Batch struct { + Id uuid.UUID + PartitionId uuid.UUID + Data []UnknownRowValues + RetryCounter int +} + +type Partition struct { + Id uuid.UUID + ParentId uuid.UUID + LowerLimit int64 + UpperLimit int64 + IsLowerLimitInclusive bool + ShouldUseRange bool + RetryCounter int +}