refactor: enhance config handling and unify main entrypoint

Unify configuration parsing logic across entrypoints and clean up legacy script-side config paths.
This commit is contained in:
2026-05-09 12:00:00 -05:00
parent bf15979c00
commit 6f8a332719
24 changed files with 1177 additions and 575 deletions

View File

@@ -2,6 +2,7 @@ package main
import (
"context"
"flag"
"sync"
"time"
@@ -19,7 +20,18 @@ import (
func main() {
configureLog()
migrationConfig, err := config.ReadMigrationConfig()
configPath := flag.String("config", "", "path to migration config file")
flag.Parse()
if flag.NArg() > 1 {
log.Fatalf("only one config file path is allowed")
}
if *configPath == "" && flag.NArg() == 1 {
*configPath = flag.Arg(0)
}
migrationConfig, err := config.ReadMigrationConfig(*configPath)
if err != nil {
log.Fatalf("error leyendo configuracion: %v", err)
}
@@ -86,7 +98,7 @@ func main() {
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("=== Migration completed successfully! ===")
log.Infof("Total migration time: %v", totalDuration)
}

View File

@@ -13,6 +13,7 @@ import (
dbwrapper "git.ksdemosapps.com/kylesoda/go-migrate/internal/app/db-wrapper"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/etl"
"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/table_analyzers"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/etl/transformers"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
@@ -20,6 +21,8 @@ import (
"golang.org/x/sync/errgroup"
)
const jobErrorsChannelSize int = 100
func buildTruncateQuery(targetDbType, schema, table, truncateMethod string) string {
if truncateMethod == "DELETE" {
if targetDbType == "postgres" {
@@ -41,7 +44,7 @@ func processMigrationJob(
targetTableAnalyzer etl.TableAnalyzer,
extractor extractors.GenericExtractor,
azureClient *azure.Client,
loader etl.Loader,
loader loaders.GenericLoader,
job config.Job,
targetDbType string,
) models.JobResult {
@@ -54,8 +57,6 @@ func processMigrationJob(
StartTime: time.Now(),
}
var rowsRead, rowsLoaded, rowsFailed int64
var wgQueryColumnTypes errgroup.Group
var sourceColTypes, targetColTypes []models.ColumnType
@@ -110,17 +111,14 @@ func processMigrationJob(
log.Error("Unexpected error calculating batch ranges: ", err)
}
chJobErrors := make(chan custom_errors.JobError, job.ExtractorQueueSize)
chLoadersErrors := make(chan custom_errors.LoaderError, job.ExtractorQueueSize)
chPartitions := make(chan models.Partition, job.ExtractorQueueSize)
chJobErrors := make(chan custom_errors.JobError, jobErrorsChannelSize)
chPartitions := make(chan models.Partition)
chBatchesRaw := make(chan models.Batch, job.ExtractorQueueSize)
chBatchesTransformed := make(chan models.Batch, job.ExtractorQueueSize)
chBatchesTransformed := make(chan models.Batch, job.TransformerQueueSize)
var wgActivePartitions sync.WaitGroup
var wgActiveBatches sync.WaitGroup
var wgExtractors sync.WaitGroup
var wgTransformers sync.WaitGroup
var wgLoaders sync.WaitGroup
var wgActivePartitions, wgActiveBatches, wgExtractors, wgTransformers, wgLoaders sync.WaitGroup
var rowsRead, rowsLoaded, rowsFailed int64
var failedPartitionsCount, failedBatchesLoadCount int32
go func() {
if err := custom_errors.JobErrorHandler(localCtx, chJobErrors); err != nil {
@@ -130,18 +128,8 @@ func processMigrationJob(
}
}()
go custom_errors.LoaderErrorHandler(
localCtx,
job.Retry,
job.MaxExtractorBatchErrors,
chLoadersErrors,
chBatchesTransformed,
chJobErrors,
&wgActiveBatches,
)
maxExtractors := min(job.MaxExtractors, len(partitions))
log.Infof("Starting %d extractor(s)...", maxExtractors)
log.Infof("Starting %d extractor(s)... (%v)", maxExtractors, job.Name)
for range maxExtractors {
wgExtractors.Go(func() {
@@ -156,6 +144,8 @@ func processMigrationJob(
chJobErrors,
&wgActivePartitions,
&rowsRead,
&failedPartitionsCount,
job.SourceTable.FromJsonColumns,
)
})
}
@@ -167,13 +157,15 @@ func processMigrationJob(
}
}()
log.Infof("Starting %d transformer(s)...", maxExtractors)
log.Infof("Starting %d transformer(s)... (%v)", maxExtractors, job.Name)
for range maxExtractors {
wgTransformers.Go(func() {
transformer.Exec(
transformer.Consume(
localCtx,
sourceColTypes,
job.Retry,
job.TransformerBatchSize,
chBatchesRaw,
chBatchesTransformed,
chJobErrors,
@@ -182,48 +174,48 @@ func processMigrationJob(
})
}
log.Infof("Starting %d loader(s)...", job.MaxLoaders)
log.Infof("Starting %d loader(s)... (%v)", job.MaxLoaders, job.Name)
for range job.MaxLoaders {
wgLoaders.Go(func() {
loader.Exec(
loader.Consume(
localCtx,
job.TargetTable,
targetColTypes,
job.Retry,
job.LoaderBatchSize,
chBatchesTransformed,
chLoadersErrors,
chJobErrors,
&wgActiveBatches,
&rowsLoaded,
&failedBatchesLoadCount,
)
})
}
go func() {
log.Debugf("Waiting for goroutines (%v)", job.Name)
// log.Debugf("Waiting for goroutines (%v)", job.Name)
wgActivePartitions.Wait()
log.Debugf("wgActivePartitions is empty (%v)", job.Name)
// log.Debugf("wgActivePartitions is empty (%v)", job.Name)
close(chPartitions)
log.Debugf("chPartitions is closed (%v)", job.Name)
// log.Debugf("chPartitions is closed (%v)", job.Name)
wgExtractors.Wait()
log.Debugf("wgExtractors is empty (%v)", job.Name)
// log.Debugf("wgExtractors is empty (%v)", job.Name)
close(chBatchesRaw)
log.Debugf("chBatchesRaw is closed (%v)", job.Name)
// log.Debugf("chBatchesRaw is closed (%v)", job.Name)
wgTransformers.Wait()
log.Debugf("wgTransformers is empty (%v)", job.Name)
// log.Debugf("wgTransformers is empty (%v)", job.Name)
close(chBatchesTransformed)
// log.Debugf("chBatchesTransformed is closed (%v)", job.Name)
wgActiveBatches.Wait()
log.Debugf("wgActiveBatches is empty (%v)", job.Name)
close(chBatchesTransformed)
log.Debugf("chBatchesTransformed is empty (%v)", job.Name)
close(chLoadersErrors)
log.Debugf("chLoadersErrors is empty (%v)", job.Name)
// log.Debugf("wgActiveBatches is empty (%v)", job.Name)
wgLoaders.Wait()
log.Debugf("wgLoaders is empty (%v)", job.Name)
// log.Debugf("wgLoaders is empty (%v)", job.Name)
cancel()
}()
@@ -235,9 +227,9 @@ func processMigrationJob(
}
}
log.Debugf("waiting for local context to be done (%v)", job.Name)
// log.Debugf("waiting for local context to be done (%v)", job.Name)
<-localCtx.Done()
log.Debugf("local context done (%v)", job.Name)
// log.Debugf("local context done (%v)", job.Name)
if ctx.Err() != nil {
result.Error = ctx.Err()
@@ -252,5 +244,9 @@ func processMigrationJob(
result.Error = fmt.Errorf("Row count mismatch: extracted %d rows but loaded %d rows (failed: %d)", result.RowsRead, result.RowsLoaded, result.RowsFailed)
}
if result.RowsRead == 0 {
log.Warnf("No rows extracted from (%v)", job.Name)
}
return result
}

View File

@@ -3,24 +3,24 @@ source_db_type: sqlserver
target_db_type: postgres
defaults:
batches_per_partition: 8
batches_per_partition: 4
max_extractors: 2
extractor_batch_size: 25000
extractor_batch_size: 5000
extractor_queue_size: 8
max_transformers: 2
transformer_batch_size: 25000
transformer_batch_size: 12500
transformer_queue_size: 8
max_loaders: 4
loader_batch_size: 25000
truncate_target: true
truncate_method: TRUNCATE # TRUNCATE | DELETE
max_partition_errrors: 5
max_extractor_batch_errors: 5
retry:
attempts: 3
base_delay_ms: 500
max_delay_ms: 10000
max_jitter_ms: 500
max_failed_partitions: 5
max_failed_batches_load: 5
jobs:
- name: demo_users
@@ -46,6 +46,9 @@ jobs:
schema: analytics
table: events
primary_key: ID_events
from_json:
- column: $node_id*
field: id
target:
schema: analytics
table: events
@@ -67,10 +70,15 @@ jobs:
- source: DATA
target: FILE_URL
mode: REFERENCE_ONLY
max_extractors: 8
batches_per_partition: 20
max_extractors: 32
extractor_batch_size: 1
extractor_queue_size: 100
max_transformers: 48
transformer_batch_size: 500
transformer_queue_size: 8
max_loaders: 4
queue_size: 32
batch_size: 1
loader_batch_size: 500
retry:
attempts: 5
base_delay_ms: 1000

View File

@@ -12,6 +12,8 @@ type RetryConfig struct {
BaseDelayMs int `yaml:"base_delay_ms"`
MaxDelayMs int `yaml:"max_delay_ms"`
MaxJitterMs int `yaml:"max_jitter_ms"`
MaxFailedPartitions int `yaml:"max_failed_partitions"`
MaxFailedBatchesLoad int `yaml:"max_failed_batches_load"`
}
type ToStorageColumnConfig struct {
@@ -36,13 +38,16 @@ type JobConfig struct {
LoaderBatchSize int `yaml:"loader_batch_size"`
TruncateTarget bool `yaml:"truncate_target"`
TruncateMethod string `yaml:"truncate_method"`
MaxPartitionErrrors int `yaml:"max_partition_errrors"`
MaxExtractorBatchErrors int `yaml:"max_extractor_batch_errors"`
Retry RetryConfig `yaml:"retry"`
RowsPerPartition int64
ToStorage ToStorageConfig `yaml:"to_storage"`
}
type FromJsonItem struct {
Column string `yaml:"column"`
Field string `yaml:"field"`
}
type TableInfo struct {
Schema string `yaml:"schema"`
Table string `yaml:"table"`
@@ -51,6 +56,7 @@ type TableInfo struct {
type SourceTableInfo struct {
TableInfo `yaml:",inline"`
PrimaryKey string `yaml:"primary_key"`
FromJsonColumns []FromJsonItem `yaml:"from_json"`
}
type TargetTableInfo struct {

View File

@@ -1,7 +1,6 @@
package custom_errors
import (
"context"
"math/rand"
"time"
)
@@ -40,22 +39,3 @@ func ComputeBackoffDelay(retryCounter int, baseDelayMs int, maxDelayMs int, maxJ
return delay
}
func requeueWithBackoff(ctx context.Context, delay time.Duration, enqueue func()) {
if delay <= 0 {
enqueue()
return
}
go func() {
timer := time.NewTimer(delay)
defer timer.Stop()
select {
case <-ctx.Done():
return
case <-timer.C:
enqueue()
}
}()
}

View File

@@ -1,107 +0,0 @@
package custom_errors
import (
"context"
"fmt"
"sync"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
)
type LoaderError struct {
Batch models.Batch
Msg string
}
func (e *LoaderError) Error() string {
return e.Msg
}
func LoaderErrorHandler(
ctx context.Context,
retryConfig config.RetryConfig,
maxChunkErrors int,
chErrorsIn <-chan LoaderError,
chBatchesOut chan<- models.Batch,
chJobErrorsOut chan<- JobError,
wgActiveBatches *sync.WaitGroup,
) {
definitiveErrors := 0
for {
if ctx.Err() != nil {
return
}
select {
case <-ctx.Done():
return
case err, ok := <-chErrorsIn:
if !ok {
return
}
if err.Batch.RetryCounter >= retryConfig.Attempts {
wgActiveBatches.Done()
definitiveErrors++
jobError := JobError{
ShouldCancelJob: false,
Msg: fmt.Sprintf("Batch %v reached max retries (%d)", err.Batch.Id, retryConfig.Attempts),
Prev: &err,
}
select {
case chJobErrorsOut <- jobError:
case <-ctx.Done():
return
}
if maxChunkErrors > 0 && definitiveErrors >= maxChunkErrors {
fatalError := JobError{
ShouldCancelJob: true,
Msg: fmt.Sprintf("Chunk error limit reached (%d)", maxChunkErrors),
Prev: &err,
}
select {
case chJobErrorsOut <- fatalError:
case <-ctx.Done():
return
}
}
continue
} else {
jobError := JobError{
ShouldCancelJob: false,
Msg: fmt.Sprintf("Temporal error in batch %v (retries: %d)", err.Batch.Id, err.Batch.RetryCounter),
Prev: &err,
}
select {
case chJobErrorsOut <- jobError:
case <-ctx.Done():
return
}
}
err.Batch.RetryCounter++
delay := ComputeBackoffDelay(
err.Batch.RetryCounter,
retryConfig.BaseDelayMs,
retryConfig.MaxDelayMs,
retryConfig.MaxJitterMs,
)
requeueWithBackoff(ctx, delay, func() {
select {
case chBatchesOut <- err.Batch:
case <-ctx.Done():
return
}
})
}
}
}

View File

@@ -14,3 +14,12 @@ type ExtractorError struct {
func (e *ExtractorError) Error() string {
return e.Msg
}
type LoaderError struct {
Batch models.Batch
Msg string
}
func (e *LoaderError) Error() string {
return e.Msg
}

View File

@@ -6,7 +6,9 @@ import (
"fmt"
"strings"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
dbdialects "git.ksdemosapps.com/kylesoda/go-migrate/internal/app/db-wrapper/db_dialects"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
mssql "github.com/microsoft/go-mssqldb"
)
@@ -177,27 +179,55 @@ func (mw *mssqlDbWrapper) SaveMassive(ctx context.Context, schema string, table
return rowsAffected, nil
}
func (mw *mssqlDbWrapper) QueryFromObject(ctx context.Context, q ExtractionQuery) (RowsResult, error) {
func buildExtractQueryMssql(q ExtractionQuery) (string, error) {
var sbQuery strings.Builder
sbQuery.WriteString("SELECT ")
if len(q.Columns) == 0 {
sbQuery.WriteString("*")
} else {
for i, col := range q.Columns {
fmt.Fprintf(&sbQuery, "[%s]", col.Name())
hasRegularColumns := len(q.Columns) > 0
hasJsonColumns := len(q.FromJsonColumns) > 0
resolvedJson := make(map[string][]config.FromJsonItem, len(q.FromJsonColumns))
if hasJsonColumns {
for _, jsonConfig := range q.FromJsonColumns {
actualColumnName, err := findColumnByPattern(q.Columns, jsonConfig.Column)
if err != nil {
return "", err
}
resolvedJson[actualColumnName] = append(resolvedJson[actualColumnName], jsonConfig)
}
}
selectParts := make([]string, 0, len(q.Columns)+len(q.FromJsonColumns))
if hasRegularColumns {
for _, col := range q.Columns {
jsonConfigs, isJsonColumn := resolvedJson[col.Name()]
if isJsonColumn {
for _, jsonConfig := range jsonConfigs {
jsonPath := buildJsonPathMssql(jsonConfig.Field)
jsonExpr := fmt.Sprintf("JSON_VALUE([%s], '%s') AS [%s]", col.Name(), jsonPath, col.Name())
selectParts = append(selectParts, jsonExpr)
}
continue
}
colExpr := fmt.Sprintf("[%s]", col.Name())
switch col.Type() {
case "GEOMETRY":
fmt.Fprintf(&sbQuery, ".STAsBinary() AS [%s]", col.Name())
colExpr = fmt.Sprintf("[%s].STAsBinary() AS [%s]", col.Name(), col.Name())
}
selectParts = append(selectParts, colExpr)
}
} else if !hasJsonColumns {
selectParts = append(selectParts, "*")
}
if i < len(q.Columns)-1 {
for i, part := range selectParts {
sbQuery.WriteString(part)
if i < len(selectParts)-1 {
sbQuery.WriteString(", ")
}
}
}
fmt.Fprintf(&sbQuery, " FROM [%s].[%s]", q.Schema, q.Table)
@@ -231,7 +261,37 @@ func (mw *mssqlDbWrapper) QueryFromObject(ctx context.Context, q ExtractionQuery
fmt.Fprintf(&sbQuery, " ORDER BY [%s] ASC", q.PrimaryKey)
queryString := sbQuery.String()
return sbQuery.String(), nil
}
func findColumnByPattern(columns []models.ColumnType, pattern string) (string, error) {
if pattern == "" {
return "", fmt.Errorf("column pattern cannot be empty")
}
if before, ok := strings.CutSuffix(pattern, "*"); ok {
prefix := before
for _, col := range columns {
if strings.HasPrefix(col.Name(), prefix) {
return col.Name(), nil
}
}
return "", fmt.Errorf("no column found matching pattern '%s'", pattern)
}
for _, col := range columns {
if col.Name() == pattern {
return col.Name(), nil
}
}
return "", fmt.Errorf("column '%s' not found in table columns", pattern)
}
func (mw *mssqlDbWrapper) QueryFromObject(ctx context.Context, q ExtractionQuery) (RowsResult, error) {
queryString, err := buildExtractQueryMssql(q)
if err != nil {
return nil, err
}
// logrus.Debugf("Query: %s", queryString)
@@ -247,3 +307,11 @@ func (mw *mssqlDbWrapper) QueryFromObject(ctx context.Context, q ExtractionQuery
return mw.Query(ctx, queryString, queryArgs...)
}
func buildJsonPathMssql(field string) string {
if len(field) > 0 && field[0] == '.' {
field = field[1:]
}
return "$." + field
}

View File

@@ -0,0 +1,396 @@
package dbwrapper
import (
"strings"
"testing"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
)
func TestBuildExtractQueryMssql_NoJsonColumns(t *testing.T) {
q := ExtractionQuery{
Schema: "dbo",
Table: "Users",
PrimaryKey: "ID",
Columns: []models.ColumnType{
models.NewColumnType("ID", false, false, "INT", "int", "INT", false, 0, 0, 0),
models.NewColumnType("Name", true, false, "VARCHAR", "varchar", "VARCHAR", true, 255, 0, 0),
},
FromJsonColumns: []config.FromJsonItem{},
LowerLimit: ExtractorQueryLimit{IsValid: false},
UpperLimit: ExtractorQueryLimit{IsValid: false},
}
query, err := buildExtractQueryMssql(q)
if err != nil {
t.Fatalf("Expected no error, got: %v", err)
}
if !strings.Contains(query, "SELECT [ID], [Name]") {
t.Errorf("Expected columns in query, got: %s", query)
}
if !strings.Contains(query, "FROM [dbo].[Users]") {
t.Errorf("Expected FROM clause, got: %s", query)
}
if !strings.Contains(query, "ORDER BY [ID] ASC") {
t.Errorf("Expected ORDER BY clause, got: %s", query)
}
}
func TestBuildExtractQueryMssql_WithJsonColumns_ExactColumnMatch(t *testing.T) {
// Test that the actual column name is used as alias, not a generated one
q := ExtractionQuery{
Schema: "dbo",
Table: "Events",
PrimaryKey: "EventID",
Columns: []models.ColumnType{
models.NewColumnType("EventID", false, false, "INT", "int", "INT", false, 0, 0, 0),
models.NewColumnType("EventData", true, false, "VARCHAR", "varchar", "VARCHAR", true, 500, 0, 0),
},
FromJsonColumns: []config.FromJsonItem{
{Column: "EventData", Field: ".userId"},
{Column: "EventData", Field: ".timestamp"},
},
LowerLimit: ExtractorQueryLimit{IsValid: false},
UpperLimit: ExtractorQueryLimit{IsValid: false},
}
query, err := buildExtractQueryMssql(q)
if err != nil {
t.Fatalf("Expected no error, got: %v", err)
}
if !strings.HasPrefix(query, "SELECT [EventID], JSON_VALUE([EventData], '$.userId') AS [EventData], JSON_VALUE([EventData], '$.timestamp') AS [EventData]") {
t.Errorf("Expected JSON columns to replace EventData in-order, got: %s", query)
}
if strings.Contains(query, "SELECT [EventID], [EventData]") {
t.Errorf("Expected EventData to be replaced by JSON extraction, got: %s", query)
}
// Alias should be exactly "EventData", not "EventData_userId"
if !strings.Contains(query, "JSON_VALUE([EventData], '$.userId') AS [EventData]") {
t.Errorf("Expected JSON alias to be [EventData], got: %s", query)
}
if !strings.Contains(query, "JSON_VALUE([EventData], '$.timestamp') AS [EventData]") {
t.Errorf("Expected JSON alias to be [EventData], got: %s", query)
}
// Should have comma separating them
if !strings.Contains(query, "JSON_VALUE([EventData], '$.userId') AS [EventData], JSON_VALUE([EventData], '$.timestamp') AS [EventData]") {
t.Errorf("Expected comma-separated JSON values, got: %s", query)
}
}
func TestBuildExtractQueryMssql_WithWildcardPattern(t *testing.T) {
// Test that wildcard pattern matching finds the correct column
q := ExtractionQuery{
Schema: "dbo",
Table: "Events",
PrimaryKey: "ID",
Columns: []models.ColumnType{
models.NewColumnType("ID", false, false, "INT", "int", "INT", false, 0, 0, 0),
models.NewColumnType("NodeMetadata", true, false, "VARCHAR", "varchar", "VARCHAR", true, 500, 0, 0),
},
FromJsonColumns: []config.FromJsonItem{
{Column: "NodeMeta*", Field: ".id"},
},
LowerLimit: ExtractorQueryLimit{IsValid: false},
UpperLimit: ExtractorQueryLimit{IsValid: false},
}
query, err := buildExtractQueryMssql(q)
if err != nil {
t.Fatalf("Expected no error, got: %v", err)
}
// Should find "NodeMetadata" from pattern "NodeMeta*" and use it as alias
if !strings.Contains(query, "JSON_VALUE([NodeMetadata], '$.id') AS [NodeMetadata]") {
t.Errorf("Expected to find and use NodeMetadata column by pattern, got: %s", query)
}
if strings.Contains(query, "SELECT [ID], [NodeMetadata]") {
t.Errorf("Expected NodeMetadata to be replaced by JSON extraction, got: %s", query)
}
}
func TestBuildExtractQueryMssql_ColumnNotFound_Error(t *testing.T) {
// Test that an error is returned when column is not found
q := ExtractionQuery{
Schema: "dbo",
Table: "Events",
PrimaryKey: "ID",
Columns: []models.ColumnType{
models.NewColumnType("ID", false, false, "INT", "int", "INT", false, 0, 0, 0),
},
FromJsonColumns: []config.FromJsonItem{
{Column: "NonExistentColumn", Field: ".id"},
},
LowerLimit: ExtractorQueryLimit{IsValid: false},
UpperLimit: ExtractorQueryLimit{IsValid: false},
}
query, err := buildExtractQueryMssql(q)
if err == nil {
t.Fatalf("Expected error for missing column, got no error. Query: %s", query)
}
if !strings.Contains(err.Error(), "NonExistentColumn") {
t.Errorf("Expected error message to contain column name, got: %v", err)
}
}
func TestBuildExtractQueryMssql_WildcardPatternNotMatched_Error(t *testing.T) {
// Test that an error is returned when wildcard pattern doesn't match any column
q := ExtractionQuery{
Schema: "dbo",
Table: "Events",
PrimaryKey: "ID",
Columns: []models.ColumnType{
models.NewColumnType("ID", false, false, "INT", "int", "INT", false, 0, 0, 0),
models.NewColumnType("EventData", true, false, "VARCHAR", "varchar", "VARCHAR", true, 500, 0, 0),
},
FromJsonColumns: []config.FromJsonItem{
{Column: "NonMatching*", Field: ".id"},
},
LowerLimit: ExtractorQueryLimit{IsValid: false},
UpperLimit: ExtractorQueryLimit{IsValid: false},
}
query, err := buildExtractQueryMssql(q)
if err == nil {
t.Fatalf("Expected error for non-matching wildcard pattern, got no error. Query: %s", query)
}
if !strings.Contains(err.Error(), "NonMatching*") {
t.Errorf("Expected error message to contain pattern, got: %v", err)
}
}
func TestBuildExtractQueryMssql_NestedJsonFields(t *testing.T) {
q := ExtractionQuery{
Schema: "dbo",
Table: "Data",
PrimaryKey: "ID",
Columns: []models.ColumnType{
models.NewColumnType("ID", false, false, "INT", "int", "INT", false, 0, 0, 0),
models.NewColumnType("NodeData", true, false, "VARCHAR", "varchar", "VARCHAR", true, 500, 0, 0),
},
FromJsonColumns: []config.FromJsonItem{
{Column: "NodeData", Field: ".user.name"},
{Column: "NodeData", Field: ".user.email"},
},
LowerLimit: ExtractorQueryLimit{IsValid: false},
UpperLimit: ExtractorQueryLimit{IsValid: false},
}
query, err := buildExtractQueryMssql(q)
if err != nil {
t.Fatalf("Expected no error, got: %v", err)
}
if !strings.Contains(query, "JSON_VALUE([NodeData], '$.user.name') AS [NodeData]") {
t.Errorf("Expected nested JSON path for user.name, got: %s", query)
}
if !strings.Contains(query, "JSON_VALUE([NodeData], '$.user.email') AS [NodeData]") {
t.Errorf("Expected nested JSON path for user.email, got: %s", query)
}
if strings.Contains(query, "SELECT [ID], [NodeData]") {
t.Errorf("Expected NodeData to be replaced by JSON extraction, got: %s", query)
}
}
func TestBuildExtractQueryMssql_WithRangeLimits(t *testing.T) {
q := ExtractionQuery{
Schema: "dbo",
Table: "Products",
PrimaryKey: "ProductID",
Columns: []models.ColumnType{
models.NewColumnType("ProductID", false, false, "INT", "int", "INT", false, 0, 0, 0),
models.NewColumnType("Details", true, false, "VARCHAR", "varchar", "VARCHAR", true, 500, 0, 0),
},
FromJsonColumns: []config.FromJsonItem{
{Column: "Details", Field: ".price"},
},
LowerLimit: ExtractorQueryLimit{IsValid: true, IsInclusive: true, Value: 100},
UpperLimit: ExtractorQueryLimit{IsValid: true, IsInclusive: false, Value: 500},
}
query, err := buildExtractQueryMssql(q)
if err != nil {
t.Fatalf("Expected no error, got: %v", err)
}
if !strings.Contains(query, "WHERE [ProductID] >= @min") {
t.Errorf("Expected WHERE clause with >=, got: %s", query)
}
if !strings.Contains(query, "[ProductID] < @max") {
t.Errorf("Expected upper limit with <, got: %s", query)
}
if !strings.Contains(query, "JSON_VALUE([Details], '$.price') AS [Details]") {
t.Errorf("Expected JSON_VALUE for Details, got: %s", query)
}
if strings.Contains(query, "SELECT [ProductID], [Details]") {
t.Errorf("Expected Details to be replaced by JSON extraction, got: %s", query)
}
}
func TestBuildJsonPathMssql(t *testing.T) {
tests := []struct {
input string
expected string
}{
{".id", "$.id"},
{"id", "$.id"},
{".user.name", "$.user.name"},
{"user.name", "$.user.name"},
{".location.coordinates.lat", "$.location.coordinates.lat"},
{"", "$."},
}
for _, tt := range tests {
result := buildJsonPathMssql(tt.input)
if result != tt.expected {
t.Errorf("buildJsonPathMssql(%q) = %q, want %q", tt.input, result, tt.expected)
}
}
}
func TestFindColumnByPattern_ExactMatch(t *testing.T) {
columns := []models.ColumnType{
models.NewColumnType("ID", false, false, "INT", "int", "INT", false, 0, 0, 0),
models.NewColumnType("Metadata", true, false, "VARCHAR", "varchar", "VARCHAR", true, 500, 0, 0),
models.NewColumnType("EventData", true, false, "VARCHAR", "varchar", "VARCHAR", true, 500, 0, 0),
}
result, err := findColumnByPattern(columns, "Metadata")
if err != nil {
t.Fatalf("Expected no error, got: %v", err)
}
if result != "Metadata" {
t.Errorf("Expected 'Metadata', got '%s'", result)
}
}
func TestFindColumnByPattern_WildcardMatch(t *testing.T) {
columns := []models.ColumnType{
models.NewColumnType("ID", false, false, "INT", "int", "INT", false, 0, 0, 0),
models.NewColumnType("NodeMetadata", true, false, "VARCHAR", "varchar", "VARCHAR", true, 500, 0, 0),
models.NewColumnType("EventData", true, false, "VARCHAR", "varchar", "VARCHAR", true, 500, 0, 0),
}
result, err := findColumnByPattern(columns, "NodeMeta*")
if err != nil {
t.Fatalf("Expected no error, got: %v", err)
}
if result != "NodeMetadata" {
t.Errorf("Expected 'NodeMetadata', got '%s'", result)
}
}
func TestFindColumnByPattern_NotFound(t *testing.T) {
columns := []models.ColumnType{
models.NewColumnType("ID", false, false, "INT", "int", "INT", false, 0, 0, 0),
models.NewColumnType("Metadata", true, false, "VARCHAR", "varchar", "VARCHAR", true, 500, 0, 0),
}
result, err := findColumnByPattern(columns, "NonExistent")
if err == nil {
t.Fatalf("Expected error, got no error. Result: %s", result)
}
if !strings.Contains(err.Error(), "NonExistent") {
t.Errorf("Expected error to contain column name, got: %v", err)
}
}
func TestFindColumnByPattern_WildcardNotFound(t *testing.T) {
columns := []models.ColumnType{
models.NewColumnType("ID", false, false, "INT", "int", "INT", false, 0, 0, 0),
models.NewColumnType("Metadata", true, false, "VARCHAR", "varchar", "VARCHAR", true, 500, 0, 0),
}
result, err := findColumnByPattern(columns, "Event*")
if err == nil {
t.Fatalf("Expected error, got no error. Result: %s", result)
}
if !strings.Contains(err.Error(), "Event*") {
t.Errorf("Expected error to contain pattern, got: %v", err)
}
}
func TestBuildExtractQueryMssql_OnlyJsonColumns(t *testing.T) {
// Test when all columns are used via JSON extraction
q := ExtractionQuery{
Schema: "dbo",
Table: "Data",
PrimaryKey: "ID",
Columns: []models.ColumnType{
models.NewColumnType("ID", false, false, "INT", "int", "INT", false, 0, 0, 0),
models.NewColumnType("JsonData", true, false, "VARCHAR", "varchar", "VARCHAR", true, 500, 0, 0),
},
FromJsonColumns: []config.FromJsonItem{
{Column: "JsonData", Field: ".field1"},
},
LowerLimit: ExtractorQueryLimit{IsValid: false},
UpperLimit: ExtractorQueryLimit{IsValid: false},
}
query, err := buildExtractQueryMssql(q)
if err != nil {
t.Fatalf("Expected no error, got: %v", err)
}
if !strings.HasPrefix(query, "SELECT [ID], JSON_VALUE([JsonData], '$.field1') AS [JsonData]") {
t.Errorf("Expected JsonData to be replaced by JSON extraction, got: %s", query)
}
if strings.Contains(query, "SELECT [ID], [JsonData]") {
t.Errorf("Expected JsonData to be excluded from raw selection, got: %s", query)
}
}
func TestBuildExtractQueryMssql_JsonColumnsReplaceInOrder(t *testing.T) {
q := ExtractionQuery{
Schema: "dbo",
Table: "Users",
PrimaryKey: "UserID",
Columns: []models.ColumnType{
models.NewColumnType("UserID", false, false, "INT", "int", "INT", false, 0, 0, 0),
models.NewColumnType("Name", true, false, "VARCHAR", "varchar", "VARCHAR", false, 255, 0, 0),
models.NewColumnType("Email", true, false, "VARCHAR", "varchar", "VARCHAR", false, 255, 0, 0),
models.NewColumnType("Metadata", true, false, "NVARCHAR", "nvarchar", "NVARCHAR", true, 4000, 0, 0),
models.NewColumnType("Profile", true, false, "NVARCHAR", "nvarchar", "NVARCHAR", true, 4000, 0, 0),
models.NewColumnType("Settings", true, false, "NVARCHAR", "nvarchar", "NVARCHAR", true, 4000, 0, 0),
},
FromJsonColumns: []config.FromJsonItem{
{Column: "Metadata", Field: ".id"},
{Column: "Profile", Field: ".id"},
{Column: "Settings", Field: ".id"},
},
LowerLimit: ExtractorQueryLimit{IsValid: false},
UpperLimit: ExtractorQueryLimit{IsValid: false},
}
query, err := buildExtractQueryMssql(q)
if err != nil {
t.Fatalf("Expected no error, got: %v", err)
}
expected := "SELECT [UserID], [Name], [Email], JSON_VALUE([Metadata], '$.id') AS [Metadata], JSON_VALUE([Profile], '$.id') AS [Profile], JSON_VALUE([Settings], '$.id') AS [Settings] FROM [dbo].[Users] ORDER BY [UserID] ASC"
if query != expected {
t.Errorf("Unexpected query.\nExpected: %s\nGot: %s", expected, query)
}
}

View File

@@ -4,6 +4,7 @@ import (
"context"
"errors"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
)
@@ -39,6 +40,7 @@ type ExtractionQuery struct {
Columns []models.ColumnType
LowerLimit ExtractorQueryLimit
UpperLimit ExtractorQueryLimit
FromJsonColumns []config.FromJsonItem
}
type DbWrapper interface {

View File

@@ -25,6 +25,8 @@ func (ex *GenericExtractor) Consume(
chErrorsOut chan<- custom_errors.JobError,
wgActivePartitions *sync.WaitGroup,
rowsRead *int64,
failedPartitionsCount *int32,
fromJsonColumns []config.FromJsonItem,
) {
indexPrimaryKey := slices.IndexFunc(columns, func(col models.ColumnType) bool {
return strings.EqualFold(col.Name(), tableInfo.PrimaryKey)
@@ -65,16 +67,18 @@ func (ex *GenericExtractor) Consume(
indexPrimaryKey,
retryConfig,
chBatchesOut,
fromJsonColumns,
)
wgActivePartitions.Done()
if rowsReadResult > 0 {
current := atomic.LoadInt64(rowsRead)
logrus.Debugf("Rows read: +%v [current=%v] (%s.%s)", rowsReadResult, current, tableInfo.Schema, tableInfo.Table)
logrus.Debugf("Rows read (partition extracted): +%v [current=%v] (%s.%s)", rowsReadResult, current, tableInfo.Schema, tableInfo.Table)
atomic.AddInt64(rowsRead, int64(rowsReadResult))
}
if err != nil {
atomic.AddInt32(failedPartitionsCount, 1)
if jobError, ok := errors.AsType[*custom_errors.JobError](err); ok {
select {
case <-ctx.Done():
@@ -88,6 +92,16 @@ func (ex *GenericExtractor) Consume(
case chErrorsOut <- custom_errors.JobError{ShouldCancelJob: false, Msg: err.Error(), Prev: err}:
}
}
currentFPCount := atomic.LoadInt32(failedPartitionsCount)
if currentFPCount > int32(retryConfig.MaxFailedPartitions) {
select {
case <-ctx.Done():
return
case chErrorsOut <- custom_errors.JobError{ShouldCancelJob: true, Msg: "Max failed partitions reached"}:
return
}
}
}
}
}

View File

@@ -22,6 +22,7 @@ func (ex *GenericExtractor) ProcessPartitionWithRetries(
indexPrimaryKey int,
retryConfig config.RetryConfig,
chBatchesOut chan<- models.Batch,
fromJsonColumns []config.FromJsonItem,
) (int64, error) {
var totalRowsRead int64
currentParitition := partition
@@ -35,6 +36,7 @@ func (ex *GenericExtractor) ProcessPartitionWithRetries(
currentParitition,
indexPrimaryKey,
chBatchesOut,
fromJsonColumns,
)
// logrus.Debugf("Partition %v finished processing (%s.%s)", partition.Id, tableInfo.Schema, tableInfo.Table)
totalRowsRead += rowsRead
@@ -48,7 +50,7 @@ func (ex *GenericExtractor) ProcessPartitionWithRetries(
if currentParitition.RetryCounter >= retryConfig.Attempts {
return totalRowsRead, &custom_errors.JobError{
Msg: fmt.Sprintf("Partition %v reached max retries", exError.Partition.Id),
Msg: fmt.Sprintf("Partition %v reached max retries (%d)", currentParitition.Id, currentParitition.RetryCounter),
Prev: err,
}
}

View File

@@ -47,6 +47,7 @@ func (ex *GenericExtractor) ProcessPartition(
partition models.Partition,
indexPrimaryKey int,
chBatchesOut chan<- models.Batch,
fromJsonColumns []config.FromJsonItem,
) (int64, error) {
query := dbwrapper.ExtractionQuery{
Schema: tableInfo.Schema,
@@ -63,6 +64,7 @@ func (ex *GenericExtractor) ProcessPartition(
IsInclusive: partition.Range.IsMaxInclusive,
Value: partition.Range.Max,
},
FromJsonColumns: fromJsonColumns,
}
// logrus.Debugf("Processing partition: %+v (%s.%s)", query, tableInfo.Schema, tableInfo.Table)

View File

@@ -0,0 +1,154 @@
package loaders
import (
"context"
"errors"
"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/models"
"github.com/google/uuid"
"github.com/sirupsen/logrus"
)
func (gl *GenericLoader) Consume(
ctx context.Context,
tableInfo config.TargetTableInfo,
columns []models.ColumnType,
retryConfig config.RetryConfig,
batchSize int,
chBatchesIn <-chan models.Batch,
chErrorsOut chan<- custom_errors.JobError,
wgActiveBatches *sync.WaitGroup,
rowsLoaded *int64,
failedBatchesCount *int32,
) {
colNames := mapSlice(columns, func(col models.ColumnType) string {
return col.Name()
})
var accRows []models.UnknownRowValues
var parentBatchesId []uuid.UUID
pendingDone := 0
defer func() {
for range pendingDone {
wgActiveBatches.Done()
}
}()
flush := func() bool {
if len(accRows) == 0 {
return true
}
count := len(parentBatchesId)
superBatch := models.Batch{
Id: uuid.New(),
ParentBatchesId: parentBatchesId,
Rows: accRows,
}
processedRows, err := gl.ProcessBatchWithRetries(ctx, tableInfo, colNames, retryConfig, superBatch)
for range count {
wgActiveBatches.Done()
}
pendingDone -= count
accRows = nil
parentBatchesId = nil
if err != nil {
atomic.AddInt32(failedBatchesCount, 1)
if jobError, ok := errors.AsType[*custom_errors.JobError](err); ok {
select {
case <-ctx.Done():
return false
case chErrorsOut <- *jobError:
}
} else {
select {
case <-ctx.Done():
return false
case chErrorsOut <- custom_errors.JobError{ShouldCancelJob: false, Msg: err.Error(), Prev: err}:
}
}
if atomic.LoadInt32(failedBatchesCount) > int32(retryConfig.MaxFailedBatchesLoad) {
select {
case <-ctx.Done():
case chErrorsOut <- custom_errors.JobError{ShouldCancelJob: true, Msg: "Max failed batches (load) reached"}:
}
return false
}
return true
}
current := atomic.LoadInt64(rowsLoaded)
logrus.Debugf("Rows loaded (batch loaded): +%v [current=%v] (%s.%s)", processedRows, current, tableInfo.Schema, tableInfo.Table)
atomic.AddInt64(rowsLoaded, int64(processedRows))
return true
}
for {
if ctx.Err() != nil {
return
}
select {
case <-ctx.Done():
return
case batch, ok := <-chBatchesIn:
if !ok {
flush()
return
}
if batchSize <= 0 {
processedRows, err := gl.ProcessBatchWithRetries(ctx, tableInfo, colNames, retryConfig, batch)
wgActiveBatches.Done()
if err != nil {
atomic.AddInt32(failedBatchesCount, 1)
if jobError, ok := errors.AsType[*custom_errors.JobError](err); ok {
select {
case <-ctx.Done():
return
case chErrorsOut <- *jobError:
}
} else {
select {
case <-ctx.Done():
return
case chErrorsOut <- custom_errors.JobError{ShouldCancelJob: false, Msg: err.Error(), Prev: err}:
}
}
if atomic.LoadInt32(failedBatchesCount) > int32(retryConfig.MaxFailedBatchesLoad) {
select {
case <-ctx.Done():
return
case chErrorsOut <- custom_errors.JobError{ShouldCancelJob: true, Msg: "Max failed batches (load) reached"}:
return
}
}
continue
}
current := atomic.LoadInt64(rowsLoaded)
logrus.Debugf("Rows loaded: +%v [current=%v] (%s.%s)", processedRows, current, tableInfo.Schema, tableInfo.Table)
atomic.AddInt64(rowsLoaded, int64(processedRows))
continue
}
pendingDone++
accRows = append(accRows, batch.Rows...)
parentBatchesId = append(parentBatchesId, batch.Id)
if len(accRows) >= batchSize {
if !flush() {
return
}
}
}
}
}

View File

@@ -1,117 +1,13 @@
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"
dbwrapper "git.ksdemosapps.com/kylesoda/go-migrate/internal/app/db-wrapper"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/etl"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
"github.com/jackc/pgx/v5/pgconn"
)
type GenericLoader struct {
db dbwrapper.DbWrapper
}
func NewGenericLoader(db dbwrapper.DbWrapper) etl.Loader {
return &GenericLoader{db: db}
}
func (gl *GenericLoader) ProcessBatch(
ctx context.Context,
tableInfo config.TargetTableInfo,
colNames []string,
batch models.Batch,
) (int, error) {
_, err := gl.db.SaveMassive(
ctx,
tableInfo.Schema,
tableInfo.Table,
colNames,
batch.Rows,
)
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.%s", tableInfo.Schema, tableInfo.Table),
Prev: err,
}
}
}
return 0, &custom_errors.LoaderError{Batch: batch, Msg: err.Error()}
}
return len(batch.Rows), nil
}
func (gl *GenericLoader) Exec(
ctx context.Context,
tableInfo config.TargetTableInfo,
columns []models.ColumnType,
chBatchesIn <-chan models.Batch,
chErrorsOut chan<- custom_errors.LoaderError,
chJobErrorsOut chan<- custom_errors.JobError,
wgActiveBatches *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 batch, ok := <-chBatchesIn:
if !ok {
return
}
processedRows, err := gl.ProcessBatch(ctx, tableInfo, colNames, batch)
if err != nil {
var ldError *custom_errors.LoaderError
var jobError *custom_errors.JobError
if errors.As(err, &ldError) {
select {
case <-ctx.Done():
return
case chErrorsOut <- *ldError:
}
} else if errors.As(err, &jobError) {
select {
case <-ctx.Done():
return
case chJobErrorsOut <- *jobError:
}
} else {
select {
case <-ctx.Done():
return
case chErrorsOut <- custom_errors.LoaderError{Batch: batch, Msg: err.Error()}:
}
}
continue
}
wgActiveBatches.Done()
atomic.AddInt64(rowsLoaded, int64(processedRows))
}
}
func NewGenericLoader(db dbwrapper.DbWrapper) GenericLoader {
return GenericLoader{db: db}
}

View File

@@ -0,0 +1,49 @@
package loaders
import (
"context"
"errors"
"fmt"
"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/models"
)
func (gl *GenericLoader) ProcessBatchWithRetries(
ctx context.Context,
tableInfo config.TargetTableInfo,
colNames []string,
retryConfig config.RetryConfig,
batch models.Batch,
) (int64, error) {
for {
rowsLoaded, err := gl.ProcessBatch(ctx, tableInfo, colNames, batch)
if err == nil {
return rowsLoaded, nil
}
if btError, ok := errors.AsType[*custom_errors.LoaderError](err); ok {
batch.RetryCounter++
if batch.RetryCounter >= retryConfig.Attempts {
return rowsLoaded, &custom_errors.JobError{
Msg: fmt.Sprintf("Batch %v reached max retries (%d)", batch.Id, batch.RetryCounter),
Prev: btError,
}
}
delay := custom_errors.ComputeBackoffDelay(
batch.RetryCounter,
retryConfig.BaseDelayMs,
retryConfig.MaxDelayMs,
retryConfig.MaxJitterMs,
)
time.Sleep(delay)
continue
}
return rowsLoaded, err
}
}

View File

@@ -0,0 +1,43 @@
package loaders
import (
"context"
"errors"
"fmt"
"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"
"github.com/jackc/pgx/v5/pgconn"
)
func (gl *GenericLoader) ProcessBatch(
ctx context.Context,
tableInfo config.TargetTableInfo,
colNames []string,
batch models.Batch,
) (int64, error) {
_, err := gl.db.SaveMassive(
ctx,
tableInfo.Schema,
tableInfo.Table,
colNames,
batch.Rows,
)
if err != nil {
if pgErr, ok := errors.AsType[*pgconn.PgError](err); ok {
if pgErr.Code == "23505" {
return 0, &custom_errors.JobError{
ShouldCancelJob: true,
Msg: fmt.Sprintf("Fatal error in table %s.%s", tableInfo.Schema, tableInfo.Table),
Prev: err,
}
}
}
return 0, &custom_errors.LoaderError{Batch: batch, Msg: err.Error()}
}
return int64(len(batch.Rows)), nil
}

View File

@@ -0,0 +1,117 @@
package transformers
import (
"context"
"errors"
"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"
"github.com/google/uuid"
)
func (mssqlTr *MssqlTransformer) Consume(
ctx context.Context,
columns []models.ColumnType,
retryConfig config.RetryConfig,
batchSize int,
chBatchesIn <-chan models.Batch,
chBatchesOut chan<- models.Batch,
chJobErrorsOut chan<- custom_errors.JobError,
wgActiveBatches *sync.WaitGroup,
) {
transformationPlan := computeTransformationPlan(columns)
storagePlan := computeStorageTransformationPlan(ctx, mssqlTr.azureClient, mssqlTr.toStorage, columns, mssqlTr.sourceTable)
transformationPlan = append(transformationPlan, storagePlan...)
var accRows []models.UnknownRowValues
var parentBatchesId []uuid.UUID
var firstPartitionId uuid.UUID
flush := func() bool {
if len(accRows) == 0 {
return true
}
out := models.Batch{
Id: uuid.New(),
PartitionId: firstPartitionId,
ParentBatchesId: parentBatchesId,
Rows: accRows,
}
select {
case chBatchesOut <- out:
wgActiveBatches.Add(1)
case <-ctx.Done():
return false
}
accRows = nil
parentBatchesId = nil
firstPartitionId = uuid.Nil
return true
}
for {
if ctx.Err() != nil {
return
}
select {
case <-ctx.Done():
return
case batch, ok := <-chBatchesIn:
if !ok {
flush()
return
}
if len(transformationPlan) > 0 {
err := ProcessBatchWithRetries(ctx, &batch, transformationPlan, retryConfig)
if err != nil {
if errors.Is(err, ctx.Err()) {
return
}
if jobError, ok := errors.AsType[*custom_errors.JobError](err); ok {
select {
case chJobErrorsOut <- *jobError:
case <-ctx.Done():
return
}
} else {
select {
case chJobErrorsOut <- custom_errors.JobError{ShouldCancelJob: true, Msg: "Transformation failed", Prev: err}:
case <-ctx.Done():
return
}
}
return
}
}
if batchSize <= 0 {
select {
case chBatchesOut <- batch:
wgActiveBatches.Add(1)
case <-ctx.Done():
return
}
continue
}
if len(parentBatchesId) == 0 {
firstPartitionId = batch.PartitionId
}
accRows = append(accRows, batch.Rows...)
parentBatchesId = append(parentBatchesId, batch.Id)
if len(accRows) >= batchSize {
if !flush() {
return
}
}
}
}
}

View File

@@ -1,20 +1,9 @@
package transformers
import (
"context"
"errors"
"fmt"
"strings"
"sync"
"time"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/azure"
"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"
log "github.com/sirupsen/logrus"
)
type MssqlTransformer struct {
@@ -30,196 +19,3 @@ func NewMssqlTransformer(toStorage config.ToStorageConfig, sourceTable config.So
azureClient: azureClient,
}
}
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
}
func computeStorageTransformationPlan(
ctx context.Context,
azureClient *azure.Client,
toStorage config.ToStorageConfig,
sourceColumns []models.ColumnType,
sourceTable config.SourceTableInfo,
) []etl.ColumnTransformPlan {
if azureClient == nil || len(toStorage.Columns) == 0 {
return nil
}
colIndex := make(map[string]int, len(sourceColumns))
for i, col := range sourceColumns {
colIndex[strings.ToUpper(col.Name())] = i
}
var plan []etl.ColumnTransformPlan
for _, storageCol := range toStorage.Columns {
if storageCol.Mode != "REFERENCE_ONLY" {
log.Warnf("to_storage: unsupported mode %q for column %s — skipping", storageCol.Mode, storageCol.Source)
continue
}
idx, ok := colIndex[strings.ToUpper(storageCol.Source)]
if !ok {
log.Warnf("to_storage: source column %q not found in source schema — skipping", storageCol.Source)
continue
}
sourceColName := storageCol.Source
schema := sourceTable.Schema
table := sourceTable.Table
plan = append(plan, etl.ColumnTransformPlan{
Index: idx,
Fn: func(v any) (any, error) {
if v == nil {
return nil, nil
}
b, ok := v.([]byte)
if !ok {
log.Warnf("to_storage: expected []byte for %s.%s.%s, got %T — passing through",
schema, table, sourceColName, v)
return v, nil
}
start := time.Now()
blobPath := fmt.Sprintf("%s/%s/%s", schema, table, uuid.New().String())
blobURL, err := azureClient.UploadAndGetURL(ctx, blobPath, b)
if err != nil {
return nil, fmt.Errorf("uploading %s.%s.%s: %w", schema, table, sourceColName, err)
}
log.Debugf(`Succesfully uploaded "%s", (%vms)`, blobURL, time.Since(start).Milliseconds())
return blobURL, nil
},
})
}
return plan
}
const processBatchCtxCheck = 4096
func (mssqlTr *MssqlTransformer) ProcessBatch(
ctx context.Context,
batch *models.Batch,
transformationPlan []etl.ColumnTransformPlan,
) error {
for i, rowValues := range batch.Rows {
if i%processBatchCtxCheck == 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,
chBatchesIn <-chan models.Batch,
chBatchesOut chan<- models.Batch,
chJobErrorsOut chan<- custom_errors.JobError,
wgActiveBatches *sync.WaitGroup,
) {
transformationPlan := computeTransformationPlan(columns)
storagePlan := computeStorageTransformationPlan(ctx, mssqlTr.azureClient, mssqlTr.toStorage, columns, mssqlTr.sourceTable)
transformationPlan = append(transformationPlan, storagePlan...)
for {
if ctx.Err() != nil {
return
}
select {
case <-ctx.Done():
return
case batch, ok := <-chBatchesIn:
if !ok {
return
}
if len(transformationPlan) == 0 {
select {
case chBatchesOut <- batch:
wgActiveBatches.Add(1)
continue
case <-ctx.Done():
return
}
}
err := mssqlTr.ProcessBatch(ctx, &batch, transformationPlan)
if err != nil {
if errors.Is(err, ctx.Err()) {
return
}
select {
case chJobErrorsOut <- custom_errors.JobError{ShouldCancelJob: true, Msg: "Transformation failed", Prev: err}:
case <-ctx.Done():
}
return
}
select {
case chBatchesOut <- batch:
case <-ctx.Done():
return
}
wgActiveBatches.Add(1)
}
}
}

View File

@@ -0,0 +1,122 @@
package transformers
import (
"context"
"fmt"
"strings"
"time"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/azure"
"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/sirupsen/logrus"
)
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
}
func computeStorageTransformationPlan(
ctx context.Context,
azureClient *azure.Client,
toStorage config.ToStorageConfig,
sourceColumns []models.ColumnType,
sourceTable config.SourceTableInfo,
) []etl.ColumnTransformPlan {
if azureClient == nil || len(toStorage.Columns) == 0 {
return nil
}
colIndex := make(map[string]int, len(sourceColumns))
for i, col := range sourceColumns {
colIndex[strings.ToUpper(col.Name())] = i
}
var plan []etl.ColumnTransformPlan
for _, storageCol := range toStorage.Columns {
if storageCol.Mode != "REFERENCE_ONLY" {
logrus.Warnf("to_storage: unsupported mode %q for column %s — skipping", storageCol.Mode, storageCol.Source)
continue
}
idx, ok := colIndex[strings.ToUpper(storageCol.Source)]
if !ok {
logrus.Warnf("to_storage: source column %q not found in source schema — skipping", storageCol.Source)
continue
}
sourceColName := storageCol.Source
schema := sourceTable.Schema
table := sourceTable.Table
plan = append(plan, etl.ColumnTransformPlan{
Index: idx,
Fn: func(v any) (any, error) {
if v == nil {
return nil, nil
}
b, ok := v.([]byte)
if !ok {
logrus.Warnf("to_storage: expected []byte for %s.%s.%s, got %T — passing through",
schema, table, sourceColName, v)
return v, nil
}
// start := time.Now()
blobPath := fmt.Sprintf("%s/%s/%s", schema, table, uuid.New().String())
blobURL, err := azureClient.UploadAndGetURL(ctx, blobPath, b)
if err != nil {
return nil, &custom_errors.JobError{
Msg: fmt.Sprintf("Error uploading %s.%s.%s", schema, table, sourceColName),
Prev: err,
}
}
// logrus.Debugf(`Succesfully uploaded "%s", (%vms)`, blobURL, time.Since(start).Milliseconds())
return blobURL, nil
},
})
}
return plan
}

View File

@@ -0,0 +1,73 @@
package transformers
import (
"context"
"errors"
"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"
)
const processBatchCtxCheck = 4096
func ProcessBatchWithRetries(
ctx context.Context,
batch *models.Batch,
transformationPlan []etl.ColumnTransformPlan,
retryConfig config.RetryConfig,
) error {
for i, rowValues := range batch.Rows {
if i%processBatchCtxCheck == 0 {
if err := ctx.Err(); err != nil {
return err
}
}
for _, task := range transformationPlan {
val := rowValues[task.Index]
if val == nil {
continue
}
var lastErr error
success := false
for attempt := 0; attempt < retryConfig.Attempts; attempt++ {
transformed, err := task.Fn(val)
if err == nil {
rowValues[task.Index] = transformed
success = true
break
}
lastErr = err
if jobError, ok := errors.AsType[*custom_errors.JobError](err); ok {
if jobError.ShouldCancelJob {
return jobError
}
}
if attempt == retryConfig.Attempts-1 {
break
}
delay := custom_errors.ComputeBackoffDelay(
attempt,
retryConfig.BaseDelayMs,
retryConfig.MaxDelayMs,
retryConfig.MaxJitterMs,
)
time.Sleep(delay)
}
if !success {
return lastErr
}
}
}
return nil
}

View File

@@ -1 +0,0 @@
package transformers

View File

@@ -9,18 +9,6 @@ import (
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
)
type Extractor interface {
ProcessPartition(
ctx context.Context,
tableInfo config.SourceTableInfo,
columns []models.ColumnType,
batchSize int,
partition models.Partition,
indexPrimaryKey int,
chBatchesOut chan<- models.Batch,
) (int, error)
}
type TransformerFunc func(any) (any, error)
type ColumnTransformPlan struct {
@@ -29,42 +17,18 @@ type ColumnTransformPlan struct {
}
type Transformer interface {
ProcessBatch(
ctx context.Context,
batch *models.Batch,
transformationPlan []ColumnTransformPlan,
) error
Exec(
Consume(
ctx context.Context,
columns []models.ColumnType,
retryConfig config.RetryConfig,
batchSize int,
chBatchesIn <-chan models.Batch,
chBactchesOut chan<- models.Batch,
chBatchesOut chan<- models.Batch,
chJobErrorsOut chan<- custom_errors.JobError,
wgActiveBatches *sync.WaitGroup,
)
}
type Loader interface {
ProcessBatch(
ctx context.Context,
tableInfo config.TargetTableInfo,
colNames []string,
batch models.Batch,
) (int, error)
Exec(
ctx context.Context,
tableInfo config.TargetTableInfo,
columns []models.ColumnType,
chBatchesIn <-chan models.Batch,
chErrorsOut chan<- custom_errors.LoaderError,
chJobErrorsOut chan<- custom_errors.JobError,
wgActiveBatches *sync.WaitGroup,
rowsLoaded *int64,
)
}
type TableAnalyzer interface {
QueryColumnTypes(
ctx context.Context,

View File

@@ -11,6 +11,7 @@ type UnknownRowValues = []any
type Batch struct {
Id uuid.UUID
PartitionId uuid.UUID
ParentBatchesId []uuid.UUID
Rows []UnknownRowValues
RetryCounter int
}