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:
@@ -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)
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
24
config.yaml
24
config.yaml
@@ -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
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
|
||||
396
internal/app/db-wrapper/mssql_test.go
Normal file
396
internal/app/db-wrapper/mssql_test.go
Normal 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)
|
||||
}
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
154
internal/app/etl/loaders/consume.go
Normal file
154
internal/app/etl/loaders/consume.go
Normal 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
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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}
|
||||
}
|
||||
|
||||
49
internal/app/etl/loaders/process-with-retries.go
Normal file
49
internal/app/etl/loaders/process-with-retries.go
Normal 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
|
||||
}
|
||||
}
|
||||
43
internal/app/etl/loaders/process.go
Normal file
43
internal/app/etl/loaders/process.go
Normal 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
|
||||
}
|
||||
117
internal/app/etl/transformers/consume.go
Normal file
117
internal/app/etl/transformers/consume.go
Normal 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
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
122
internal/app/etl/transformers/plan.go
Normal file
122
internal/app/etl/transformers/plan.go
Normal 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
|
||||
}
|
||||
73
internal/app/etl/transformers/process-with-retries.go
Normal file
73
internal/app/etl/transformers/process-with-retries.go
Normal 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
|
||||
}
|
||||
@@ -1 +0,0 @@
|
||||
package transformers
|
||||
@@ -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,
|
||||
|
||||
@@ -11,6 +11,7 @@ type UnknownRowValues = []any
|
||||
type Batch struct {
|
||||
Id uuid.UUID
|
||||
PartitionId uuid.UUID
|
||||
ParentBatchesId []uuid.UUID
|
||||
Rows []UnknownRowValues
|
||||
RetryCounter int
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user