diff --git a/cmd/go_migrate/main.go b/cmd/go_migrate/main.go index 6b1c61e..9a24884 100644 --- a/cmd/go_migrate/main.go +++ b/cmd/go_migrate/main.go @@ -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) } diff --git a/cmd/go_migrate/process.go b/cmd/go_migrate/process.go index 512e040..efd4a9c 100644 --- a/cmd/go_migrate/process.go +++ b/cmd/go_migrate/process.go @@ -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 } diff --git a/config.yaml b/config.yaml index 5e611e0..acd5fa1 100644 --- a/config.yaml +++ b/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 diff --git a/internal/app/config/migration.go b/internal/app/config/migration.go index 51baa7c..39309f7 100644 --- a/internal/app/config/migration.go +++ b/internal/app/config/migration.go @@ -8,10 +8,12 @@ import ( ) type RetryConfig struct { - Attempts int `yaml:"attempts"` - BaseDelayMs int `yaml:"base_delay_ms"` - MaxDelayMs int `yaml:"max_delay_ms"` - MaxJitterMs int `yaml:"max_jitter_ms"` + Attempts int `yaml:"attempts"` + 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 { @@ -25,22 +27,25 @@ type ToStorageConfig struct { } type JobConfig struct { - BatchesPerPartition int `yaml:"batches_per_partition"` - MaxExtractors int `yaml:"max_extractors"` - ExtractorBatchSize int `yaml:"extractor_batch_size"` - ExtractorQueueSize int `yaml:"extractor_queue_size"` - MaxTransformers int `yaml:"max_transformers"` - TransformerBatchSize int `yaml:"transformer_batch_size"` - TransformerQueueSize int `yaml:"transformer_queue_size"` - MaxLoaders int `yaml:"max_loaders"` - 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"` + BatchesPerPartition int `yaml:"batches_per_partition"` + MaxExtractors int `yaml:"max_extractors"` + ExtractorBatchSize int `yaml:"extractor_batch_size"` + ExtractorQueueSize int `yaml:"extractor_queue_size"` + MaxTransformers int `yaml:"max_transformers"` + TransformerBatchSize int `yaml:"transformer_batch_size"` + TransformerQueueSize int `yaml:"transformer_queue_size"` + MaxLoaders int `yaml:"max_loaders"` + LoaderBatchSize int `yaml:"loader_batch_size"` + TruncateTarget bool `yaml:"truncate_target"` + TruncateMethod string `yaml:"truncate_method"` + 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 { @@ -49,8 +54,9 @@ type TableInfo struct { } type SourceTableInfo struct { - TableInfo `yaml:",inline"` - PrimaryKey string `yaml:"primary_key"` + TableInfo `yaml:",inline"` + PrimaryKey string `yaml:"primary_key"` + FromJsonColumns []FromJsonItem `yaml:"from_json"` } type TargetTableInfo struct { diff --git a/internal/app/custom_errors/backoff.go b/internal/app/custom_errors/backoff.go index 64e928a..49052c6 100644 --- a/internal/app/custom_errors/backoff.go +++ b/internal/app/custom_errors/backoff.go @@ -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() - } - }() -} diff --git a/internal/app/custom_errors/loader.error.go b/internal/app/custom_errors/loader.error.go deleted file mode 100644 index 927e446..0000000 --- a/internal/app/custom_errors/loader.error.go +++ /dev/null @@ -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 - } - }) - } - } -} diff --git a/internal/app/custom_errors/extractor.error.go b/internal/app/custom_errors/types.go similarity index 69% rename from internal/app/custom_errors/extractor.error.go rename to internal/app/custom_errors/types.go index 1ab5888..10b8cfb 100644 --- a/internal/app/custom_errors/extractor.error.go +++ b/internal/app/custom_errors/types.go @@ -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 +} diff --git a/internal/app/db-wrapper/mssql.go b/internal/app/db-wrapper/mssql.go index 03c9284..0fcd972 100644 --- a/internal/app/db-wrapper/mssql.go +++ b/internal/app/db-wrapper/mssql.go @@ -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,25 +179,53 @@ 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 { - sbQuery.WriteString(", ") - } + for i, part := range selectParts { + sbQuery.WriteString(part) + if i < len(selectParts)-1 { + sbQuery.WriteString(", ") } } @@ -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 +} diff --git a/internal/app/db-wrapper/mssql_test.go b/internal/app/db-wrapper/mssql_test.go new file mode 100644 index 0000000..d7abf0e --- /dev/null +++ b/internal/app/db-wrapper/mssql_test.go @@ -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) + } +} diff --git a/internal/app/db-wrapper/types.go b/internal/app/db-wrapper/types.go index e194710..7bc99e0 100644 --- a/internal/app/db-wrapper/types.go +++ b/internal/app/db-wrapper/types.go @@ -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" ) @@ -33,12 +34,13 @@ type ExtractorQueryLimit struct { } type ExtractionQuery struct { - Schema string - Table string - PrimaryKey string - Columns []models.ColumnType - LowerLimit ExtractorQueryLimit - UpperLimit ExtractorQueryLimit + Schema string + Table string + PrimaryKey string + Columns []models.ColumnType + LowerLimit ExtractorQueryLimit + UpperLimit ExtractorQueryLimit + FromJsonColumns []config.FromJsonItem } type DbWrapper interface { diff --git a/internal/app/etl/extractors/consume.go b/internal/app/etl/extractors/consume.go index 9884288..2fd9984 100644 --- a/internal/app/etl/extractors/consume.go +++ b/internal/app/etl/extractors/consume.go @@ -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 + } + } } } } diff --git a/internal/app/etl/extractors/process-with-retries.go b/internal/app/etl/extractors/process-with-retries.go index 5837306..d316dde 100644 --- a/internal/app/etl/extractors/process-with-retries.go +++ b/internal/app/etl/extractors/process-with-retries.go @@ -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, } } diff --git a/internal/app/etl/extractors/process.go b/internal/app/etl/extractors/process.go index d624551..1a56912 100644 --- a/internal/app/etl/extractors/process.go +++ b/internal/app/etl/extractors/process.go @@ -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) diff --git a/internal/app/etl/loaders/consume.go b/internal/app/etl/loaders/consume.go new file mode 100644 index 0000000..0b93aa9 --- /dev/null +++ b/internal/app/etl/loaders/consume.go @@ -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 + } + } + } + } +} diff --git a/internal/app/etl/loaders/main.go b/internal/app/etl/loaders/main.go index 523d6c8..1b6d20e 100644 --- a/internal/app/etl/loaders/main.go +++ b/internal/app/etl/loaders/main.go @@ -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} } diff --git a/internal/app/etl/loaders/process-with-retries.go b/internal/app/etl/loaders/process-with-retries.go new file mode 100644 index 0000000..6869192 --- /dev/null +++ b/internal/app/etl/loaders/process-with-retries.go @@ -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 + } +} diff --git a/internal/app/etl/loaders/process.go b/internal/app/etl/loaders/process.go new file mode 100644 index 0000000..4e48988 --- /dev/null +++ b/internal/app/etl/loaders/process.go @@ -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 +} diff --git a/internal/app/etl/transformers/consume.go b/internal/app/etl/transformers/consume.go new file mode 100644 index 0000000..bd3a92d --- /dev/null +++ b/internal/app/etl/transformers/consume.go @@ -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 + } + } + } + } +} diff --git a/internal/app/etl/transformers/mssql.go b/internal/app/etl/transformers/mssql.go index 0c238e5..da0710c 100644 --- a/internal/app/etl/transformers/mssql.go +++ b/internal/app/etl/transformers/mssql.go @@ -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) - } - } -} diff --git a/internal/app/etl/transformers/plan.go b/internal/app/etl/transformers/plan.go new file mode 100644 index 0000000..5758dba --- /dev/null +++ b/internal/app/etl/transformers/plan.go @@ -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 +} diff --git a/internal/app/etl/transformers/process-with-retries.go b/internal/app/etl/transformers/process-with-retries.go new file mode 100644 index 0000000..dbf42ef --- /dev/null +++ b/internal/app/etl/transformers/process-with-retries.go @@ -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 +} diff --git a/internal/app/etl/transformers/types.go b/internal/app/etl/transformers/types.go deleted file mode 100644 index b3fa142..0000000 --- a/internal/app/etl/transformers/types.go +++ /dev/null @@ -1 +0,0 @@ -package transformers diff --git a/internal/app/etl/types.go b/internal/app/etl/types.go index b2eb87b..a1970ae 100644 --- a/internal/app/etl/types.go +++ b/internal/app/etl/types.go @@ -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, diff --git a/internal/app/models/main.go b/internal/app/models/main.go index 60eb73e..5becf6a 100644 --- a/internal/app/models/main.go +++ b/internal/app/models/main.go @@ -9,10 +9,11 @@ import ( type UnknownRowValues = []any type Batch struct { - Id uuid.UUID - PartitionId uuid.UUID - Rows []UnknownRowValues - RetryCounter int + Id uuid.UUID + PartitionId uuid.UUID + ParentBatchesId []uuid.UUID + Rows []UnknownRowValues + RetryCounter int } type PartitionRange struct {