kylesoda/go-migrate-db
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add exponential backoff retry strategy

Browse Source 
Introduce a retry policy with exponential backoff and jitter for extractor and loader errors, with configurable max attempts and delay caps.
This commit is contained in:
Kylesoda
2026-04-13 14:00:00 -05:00
parent 7f3d2b8cc4
commit 74abf12dcf
21 changed files with 847 additions and 684 deletions
Download Patch File Download Diff File Expand all files Collapse all files

108
cmd/go_migrate/batch-generator.go
View File

@@ -1,108 +0,0 @@
package main
import (
"context"
"database/sql"
"fmt"
"time"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
"github.com/google/uuid"
)
func estimateTotalRowsMssql(ctx context.Context, db *sql.DB, tableInfo config.SourceTableInfo) (int64, error) {
query := `
SELECT
SUM(p.rows) AS count
FROM sys.tables t
JOIN sys.schemas s ON t.schema_id = s.schema_id
JOIN sys.partitions p ON t.object_id = p.object_id
WHERE s.name = @schema AND t.name = @table AND p.index_id IN (0, 1)
GROUP BY t.name`
ctxTimeout, cancel := context.WithTimeout(ctx, time.Second*20)
defer cancel()
var rowsCount int64
err := db.QueryRowContext(ctxTimeout, query, sql.Named("schema", tableInfo.Schema), sql.Named("table", tableInfo.Table)).Scan(&rowsCount)
if err != nil {
return 0, err
}
return rowsCount, nil
}
func calculateBatchesMssql(ctx context.Context, db *sql.DB, tableInfo config.SourceTableInfo, batchCount int64) ([]models.Partition, error) {
query := fmt.Sprintf(`
SELECT
MIN([%s]) AS lower_limit,
MAX([%s]) AS upper_limit
FROM
(SELECT [%s], NTILE(@batchCount) OVER (ORDER BY [%s]) AS batch_id FROM [%s].[%s]) AS T
GROUP BY batch_id
ORDER BY batch_id`,
tableInfo.PrimaryKey,
tableInfo.PrimaryKey,
tableInfo.PrimaryKey,
tableInfo.PrimaryKey,
tableInfo.Schema,
tableInfo.Table)
ctxTimeout, cancel := context.WithTimeout(ctx, time.Second*20)
defer cancel()
rows, err := db.QueryContext(ctxTimeout, query, sql.Named("batchCount", batchCount))
if err != nil {
return nil, err
}
defer rows.Close()
batches := make([]models.Partition, 0, batchCount)
for rows.Next() {
batch := models.Partition{
Id: uuid.New(),
ShouldUseRange: true,
RetryCounter: 0,
IsLowerLimitInclusive: true,
}
if err := rows.Scan(&batch.LowerLimit, &batch.UpperLimit); err != nil {
return nil, err
}
batches = append(batches, batch)
}
if err := rows.Err(); err != nil {
return nil, err
}
return batches, nil
}
func batchGeneratorMssql(ctx context.Context, db *sql.DB, tableInfo config.SourceTableInfo, rowsPerBatch int64) ([]models.Partition, error) {
rowsCount, err := estimateTotalRowsMssql(ctx, db, tableInfo)
if err != nil {
return nil, err
}
var batchCount int64 = 1
if rowsCount > rowsPerBatch {
batchCount = rowsCount / rowsPerBatch
} else {
return []models.Partition{{
Id: uuid.New(),
ShouldUseRange: false,
RetryCounter: 0,
}}, nil
}
batches, err := calculateBatchesMssql(ctx, db, tableInfo, batchCount)
if err != nil {
return nil, err
}
return batches, nil
}

44
cmd/go_migrate/colum-type.go
View File

@@ -1,44 +0,0 @@
package main
type ColumnType struct {
name string
hasMaxLength bool
hasPrecisionScale bool
userType string
systemType string
unifiedType string
nullable bool
maxLength int64
precision int64
scale int64
}
func (c *ColumnType) Name() string {
return c.name
}
func (c *ColumnType) UserType() string {
return c.userType
}
func (c *ColumnType) SystemType() string {
return c.systemType
}
func (c *ColumnType) Length() (length int64, ok bool) {
return c.maxLength, c.hasMaxLength
}
func (c *ColumnType) DecimalSize() (precision, scale int64, ok bool) {
return c.precision, c.scale, c.hasPrecisionScale
}
func (c *ColumnType) Nullable() bool {
return c.nullable
}
func (c *ColumnType) Type() string {
return c.unifiedType
}

316
cmd/go_migrate/inspect-columns.go
View File

@@ -1,316 +0,0 @@
package main
import (
"context"
"database/sql"
"errors"
"fmt"
"strings"
"sync"
"time"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
"github.com/jackc/pgx/v5/pgxpool"
_ "github.com/microsoft/go-mssqldb"
log "github.com/sirupsen/logrus"
)
func GetUnifiedType(systemType string) string {
systemType = strings.ToLower(systemType)
if systemType == "varchar" || systemType == "char" || systemType == "nvarchar" || systemType == "nchar" || systemType == "text" || systemType == "ntext" {
return "STRING"
}
if systemType == "int" || systemType == "int4" || systemType == "integer" || systemType == "smallint" || systemType == "int2" || systemType == "bigint" || systemType == "int8" || systemType == "tinyint" {
return "INTEGER"
}
if systemType == "decimal" || systemType == "numeric" {
return "DECIMAL"
}
if systemType == "float" || systemType == "real" || systemType == "double precision" {
return "FLOAT"
}
if systemType == "bit" || systemType == "boolean" {
return "BOOLEAN"
}
if systemType == "date" {
return "DATE"
}
if systemType == "time" || systemType == "time without time zone" {
return "TIME"
}
if systemType == "datetime" || systemType == "datetime2" || systemType == "timestamp" || systemType == "timestamptz" || systemType == "timestamp with time zone" {
return "TIMESTAMP"
}
if systemType == "binary" || systemType == "varbinary" || systemType == "image" || systemType == "bytea" {
return "BINARY"
}
if systemType == "uniqueidentifier" || systemType == "uuid" {
return "UUID"
}
if systemType == "json" {
return "JSON"
}
if systemType == "geometry" || systemType == "geography" {
return "GEOMETRY"
}
return strings.ToUpper(systemType)
}
func MapPostgresColumn(column ColumnType, maxLength *int64, precision *int64, scale *int64) models.ColumnType {
stringTypes := map[string]bool{
"varchar": true, "char": true, "character": true, "text": true, "character varying": true,
}
decimalTypes := map[string]bool{
"decimal": true, "numeric": true,
}
if stringTypes[column.systemType] {
if maxLength != nil {
column.maxLength = *maxLength
column.hasMaxLength = true
} else {
column.maxLength = -1
column.hasMaxLength = false
}
column.hasPrecisionScale = false
column.precision = -1
column.scale = -1
} else if decimalTypes[column.systemType] {
column.hasMaxLength = false
column.maxLength = -1
if precision != nil && scale != nil {
column.precision = *precision
column.scale = *scale
column.hasPrecisionScale = true
} else {
column.precision = -1
column.scale = -1
column.hasPrecisionScale = false
}
} else {
column.hasMaxLength = false
column.maxLength = -1
column.hasPrecisionScale = false
column.precision = -1
column.scale = -1
}
column.unifiedType = GetUnifiedType(column.systemType)
colType := models.NewColumnType(
column.name,
column.hasMaxLength,
column.hasPrecisionScale,
column.userType,
column.systemType,
column.unifiedType,
column.nullable,
column.maxLength,
column.precision,
column.scale,
)
return colType
}
func GetColumnTypesPostgres(db *pgxpool.Pool, tableInfo config.TargetTableInfo) ([]models.ColumnType, error) {
query := `
SELECT
c.column_name AS name,
c.data_type AS user_type,
c.udt_name AS system_type,
(CASE WHEN c.is_nullable = 'YES' THEN TRUE ELSE FALSE END) AS nullable,
c.character_maximum_length AS max_length,
c.numeric_precision AS precision,
c.numeric_scale AS scale
FROM information_schema.columns c
WHERE c.table_schema = $1 AND c.table_name = $2
ORDER BY c.ordinal_position;
`
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
defer cancel()
rows, err := db.Query(ctx, query, tableInfo.Schema, tableInfo.Table)
if err != nil {
return nil, fmt.Errorf("Error querying column types: %w", err)
}
defer rows.Close()
var colTypes []models.ColumnType
for rows.Next() {
var column ColumnType
var scanMaxLength *int64
var scanPrecision *int64
var scanScale *int64
if err := rows.Scan(
&column.name,
&column.userType,
&column.systemType,
&column.nullable,
&scanMaxLength,
&scanPrecision,
&scanScale,
); err != nil {
return nil, fmt.Errorf("Error scanning column type results: %w", err)
}
colTypes = append(colTypes, MapPostgresColumn(column, scanMaxLength, scanPrecision, scanScale))
}
return colTypes, nil
}
func MapMssqlColumn(column ColumnType) models.ColumnType {
stringTypes := map[string]bool{
"varchar": true, "char": true, "nvarchar": true, "nchar": true, "text": true, "ntext": true,
}
decimalTypes := map[string]bool{
"decimal": true, "numeric": true,
}
if stringTypes[column.systemType] {
column.hasMaxLength = true
if column.systemType == "nvarchar" || column.systemType == "nchar" {
if column.maxLength > 0 {
column.maxLength = column.maxLength / 2
}
}
column.hasPrecisionScale = false
column.precision = -1
column.scale = -1
} else if decimalTypes[column.systemType] {
column.hasMaxLength = false
column.maxLength = -1
column.hasPrecisionScale = true
} else {
column.hasMaxLength = false
column.maxLength = -1
column.hasPrecisionScale = false
column.precision = -1
column.scale = -1
}
column.unifiedType = GetUnifiedType(column.systemType)
colType := models.NewColumnType(
column.name,
column.hasMaxLength,
column.hasPrecisionScale,
column.userType,
column.systemType,
column.unifiedType,
column.nullable,
column.maxLength,
column.precision,
column.scale,
)
return colType
}
func GetColumnTypesMssql(db *sql.DB, tableInfo config.SourceTableInfo) ([]models.ColumnType, error) {
query := `
SELECT
c.name AS name,
t.name AS user_type,
CASE WHEN t.is_user_defined = 0 THEN t.name ELSE bt.name END AS system_type,
c.is_nullable AS nullable,
c.max_length AS max_length,
c.precision AS precision,
c.scale AS scale
FROM sys.columns c
JOIN sys.types t ON c.user_type_id = t.user_type_id
LEFT JOIN sys.types bt ON t.is_user_defined = 1 AND bt.user_type_id = t.system_type_id
JOIN sys.tables st ON c.object_id = st.object_id
JOIN sys.schemas s ON st.schema_id = s.schema_id
WHERE s.name = @schema AND st.name = @table
ORDER BY c.column_id;
`
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
defer cancel()
rows, err := db.QueryContext(ctx, query, sql.Named("schema", tableInfo.Schema), sql.Named("table", tableInfo.Table))
if err != nil {
return nil, fmt.Errorf("Error querying column types: %w", err)
}
defer rows.Close()
var colTypes []models.ColumnType
for rows.Next() {
var column ColumnType
if err := rows.Scan(
&column.name,
&column.userType,
&column.systemType,
&column.nullable,
&column.maxLength,
&column.precision,
&column.scale,
); err != nil {
return nil, fmt.Errorf("Error scanning column type results: %W", err)
}
if strings.HasPrefix(column.name, "graph_id") && column.systemType == "bigint" {
continue
}
colTypes = append(colTypes, MapMssqlColumn(column))
}
return colTypes, nil
}
func GetColumnTypes(
sourceDb *sql.DB,
targetDb *pgxpool.Pool,
sourceTable config.SourceTableInfo,
targetTable config.TargetTableInfo,
) ([]models.ColumnType, []models.ColumnType, error) {
var sourceDbErr error
var targetDbErr error
var sourceColTypes []models.ColumnType
var targetColTypes []models.ColumnType
var wg sync.WaitGroup
wg.Go(func() {
sourceColTypes, sourceDbErr = GetColumnTypesMssql(sourceDb, sourceTable)
if sourceDbErr != nil {
log.Error("Error (sourceDb): ", sourceDbErr)
}
})
wg.Go(func() {
targetColTypes, targetDbErr = GetColumnTypesPostgres(targetDb, targetTable)
if targetDbErr != nil {
log.Error("Error (targetDb): ", targetDbErr)
}
})
wg.Wait()
if sourceDbErr != nil || targetDbErr != nil {
return nil, nil, errors.New("Error querying column types")
}
return sourceColTypes, targetColTypes, nil
}

7
cmd/go_migrate/main.go
View File

@@ -9,6 +9,7 @@ import (
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config" "git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/etl/extractors" "git.ksdemosapps.com/kylesoda/go-migrate/internal/app/etl/extractors"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/etl/loaders" "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/etl/transformers"
"github.com/jackc/pgx/v5/pgxpool" "github.com/jackc/pgx/v5/pgxpool"
log "github.com/sirupsen/logrus" log "github.com/sirupsen/logrus"
@@ -90,6 +91,8 @@ func processMigrationJobs(
chJobs := make(chan config.Job, len(jobs)) chJobs := make(chan config.Job, len(jobs))
var wgJobs sync.WaitGroup var wgJobs sync.WaitGroup
sourceTableAnalyzer := table_analyzers.NewMssqlTableAnalyzer(sourceDb)
targetTableAnalyzer := table_analyzers.NewPostgresTableAnalyzer(targetDb)
extractor := extractors.NewMssqlExtractor(sourceDb) extractor := extractors.NewMssqlExtractor(sourceDb)
transformer := transformers.NewMssqlTransformer() transformer := transformers.NewMssqlTransformer()
loader := loaders.NewPostgresLoader(targetDb) loader := loaders.NewPostgresLoader(targetDb)
@@ -100,8 +103,8 @@ func processMigrationJobs(
log.Infof("[worker %d] >>> Processing job: %s.%s <<<", i, job.SourceTable.Schema, job.SourceTable.Table) log.Infof("[worker %d] >>> Processing job: %s.%s <<<", i, job.SourceTable.Schema, job.SourceTable.Table)
res := processMigrationJob( res := processMigrationJob(
ctx, ctx,
sourceDb, sourceTableAnalyzer,
targetDb, targetTableAnalyzer,
extractor, extractor,
transformer, transformer,
loader, loader,

136
cmd/go_migrate/process.go
View File

@@ -2,7 +2,6 @@ package main
import ( import (
"context" "context"
"database/sql"
"sync" "sync"
"sync/atomic" "sync/atomic"
"time" "time"
@@ -10,21 +9,24 @@ import (
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config" "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/custom_errors"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/etl" "git.ksdemosapps.com/kylesoda/go-migrate/internal/app/etl"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/etl/table_analyzers"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models" "git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
"github.com/jackc/pgx/v5/pgxpool"
_ "github.com/microsoft/go-mssqldb"
log "github.com/sirupsen/logrus" log "github.com/sirupsen/logrus"
"golang.org/x/sync/errgroup"
) )
func processMigrationJob( func processMigrationJob(
ctx context.Context, ctx context.Context,
sourceDb *sql.DB, sourceTableAnalyzer etl.TableAnalyzer,
targetDb *pgxpool.Pool, targetTableAnalyzer etl.TableAnalyzer,
extractor etl.Extractor, extractor etl.Extractor,
transformer etl.Transformer, transformer etl.Transformer,
loader etl.Loader, loader etl.Loader,
job config.Job, job config.Job,
) JobResult { ) JobResult {
jobCtx, cancel := context.WithCancel(ctx)
defer cancel()
result := JobResult{ result := JobResult{
JobName: job.Name, JobName: job.Name,
StartTime: time.Now(), StartTime: time.Now(),
@@ -32,47 +34,85 @@ func processMigrationJob(
var rowsRead, rowsLoaded, rowsFailed int64 var rowsRead, rowsLoaded, rowsFailed int64
sourceColTypes, targetColTypes, err := GetColumnTypes(sourceDb, targetDb, job.SourceTable, job.TargetTable) var wgQueryColumnTypes errgroup.Group
var sourceColTypes, targetColTypes []models.ColumnType
wgQueryColumnTypes.Go(func() error {
var err error
sourceColTypes, err = sourceTableAnalyzer.QueryColumnTypes(jobCtx, job.SourceTable.TableInfo)
if err != nil {
return err
}
return nil
})
wgQueryColumnTypes.Go(func() error {
var err error
targetColTypes, err = targetTableAnalyzer.QueryColumnTypes(jobCtx, job.TargetTable.TableInfo)
if err != nil {
return err
}
return nil
})
err := wgQueryColumnTypes.Wait()
if err != nil { if err != nil {
result.Error = err result.Error = err
return result return result
} }
logColumnTypes(sourceColTypes, "Source col types") partitions, err := table_analyzers.PartitionRangeGenerator(
logColumnTypes(targetColTypes, "Target col types") jobCtx,
sourceTableAnalyzer,
jobCtx, cancel := context.WithCancel(ctx) job.SourceTable.TableInfo,
defer cancel() job.SourceTable.PrimaryKey,
job.RowsPerPartition,
batches, err := batchGeneratorMssql(jobCtx, sourceDb, job.SourceTable, job.RowsPerBatch) )
if err != nil { if err != nil {
log.Error("Unexpected error calculating batch ranges: ", err) log.Error("Unexpected error calculating batch ranges: ", err)
} }
chJobErrors := make(chan custom_errors.JobError, job.QueueSize) chJobErrors := make(chan custom_errors.JobError, job.QueueSize)
chBatches := make(chan models.Partition, job.QueueSize)
chExtractorErrors := make(chan custom_errors.ExtractorError, job.QueueSize) chExtractorErrors := make(chan custom_errors.ExtractorError, job.QueueSize)
chChunksRaw := make(chan models.Batch, job.QueueSize)
chChunksTransformed := make(chan models.Batch, job.QueueSize)
chLoadersErrors := make(chan custom_errors.LoaderError, job.QueueSize) chLoadersErrors := make(chan custom_errors.LoaderError, job.QueueSize)
chPartitions := make(chan models.Partition, job.QueueSize)
chBatchesRaw := make(chan models.Batch, job.QueueSize)
chBatchesTransformed := make(chan models.Batch, job.QueueSize)
var wgActivePartitions sync.WaitGroup
var wgActiveBatches sync.WaitGroup var wgActiveBatches sync.WaitGroup
var wgActiveChunks sync.WaitGroup
var wgExtractors sync.WaitGroup var wgExtractors sync.WaitGroup
var wgTransformers sync.WaitGroup var wgTransformers sync.WaitGroup
var wgLoaders sync.WaitGroup var wgLoaders sync.WaitGroup
go func() { go func() {
if err := custom_errors.JobErrorHandler(jobCtx, chJobErrors); err != nil { if err := custom_errors.JobErrorHandler(jobCtx, chJobErrors); err != nil {
log.Error("Fatal error received from JobErrorHandler, canceling job... - ", err)
cancel() cancel()
result.Error = err result.Error = err
} }
}() }()
go custom_errors.ExtractorErrorHandler(jobCtx, job.Retry.Attempts, chExtractorErrors, chBatches, chJobErrors, &wgActiveBatches) go custom_errors.ExtractorErrorHandler(
go custom_errors.LoaderErrorHandler(jobCtx, job.Retry.Attempts, chLoadersErrors, chChunksTransformed, chJobErrors, &wgActiveChunks) jobCtx,
job.Retry,
chExtractorErrors,
chPartitions,
chJobErrors,
&wgActivePartitions,
)
go custom_errors.LoaderErrorHandler(
jobCtx,
job.Retry,
chLoadersErrors,
chBatchesTransformed,
chJobErrors,
&wgActiveBatches,
)
maxExtractors := min(job.MaxExtractors, len(batches)) maxExtractors := min(job.MaxExtractors, len(partitions))
log.Infof("Starting %d extractor(s)...", maxExtractors) log.Infof("Starting %d extractor(s)...", maxExtractors)
for range maxExtractors { for range maxExtractors {
@@ -81,21 +121,21 @@ func processMigrationJob(
jobCtx, jobCtx,
job.SourceTable, job.SourceTable,
sourceColTypes, sourceColTypes,
job.ChunkSize, job.BatchSize,
chBatches, chPartitions,
chChunksRaw, chBatchesRaw,
chExtractorErrors, chExtractorErrors,
chJobErrors, chJobErrors,
&wgActiveBatches, &wgActivePartitions,
&rowsRead, &rowsRead,
) )
}) })
} }
wgActiveBatches.Add(len(batches)) wgActivePartitions.Add(len(partitions))
go func() { go func() {
for _, batch := range batches { for _, batch := range partitions {
chBatches <- batch chPartitions <- batch
} }
}() }()
@@ -106,10 +146,10 @@ func processMigrationJob(
transformer.Exec( transformer.Exec(
jobCtx, jobCtx,
sourceColTypes, sourceColTypes,
chChunksRaw, chBatchesRaw,
chChunksTransformed, chBatchesTransformed,
chJobErrors, chJobErrors,
&wgActiveChunks, &wgActiveBatches,
) )
}) })
} }
@@ -122,35 +162,49 @@ func processMigrationJob(
jobCtx, jobCtx,
job.TargetTable, job.TargetTable,
targetColTypes, targetColTypes,
chChunksTransformed, chBatchesTransformed,
chLoadersErrors, chLoadersErrors,
chJobErrors, chJobErrors,
&wgActiveChunks, &wgActiveBatches,
&rowsLoaded, &rowsLoaded,
) )
}) })
} }
go func() { go func() {
wgActiveBatches.Wait() log.Debugf("Waiting for goroutines (%v)", job.Name)
close(chBatches)
wgActivePartitions.Wait()
log.Debugf("wgActivePartitions is empty (%v)", job.Name)
close(chPartitions)
log.Debugf("chPartitions is closed (%v)", job.Name)
close(chExtractorErrors) close(chExtractorErrors)
log.Debugf("chExtractorErrors is closed (%v)", job.Name)
wgExtractors.Wait() wgExtractors.Wait()
close(chChunksRaw) log.Debugf("wgExtractors is empty (%v)", job.Name)
close(chBatchesRaw)
log.Debugf("chBatchesRaw is closed (%v)", job.Name)
wgTransformers.Wait() wgTransformers.Wait()
log.Debugf("wgTransformers is empty (%v)", job.Name)
wgActiveChunks.Wait() wgActiveBatches.Wait()
close(chChunksTransformed) log.Debugf("wgActiveBatches is empty (%v)", job.Name)
close(chBatchesTransformed)
log.Debugf("chBatchesTransformed is empty (%v)", job.Name)
close(chLoadersErrors) close(chLoadersErrors)
log.Debugf("chLoadersErrors is empty (%v)", job.Name)
wgLoaders.Wait() wgLoaders.Wait()
log.Debugf("wgLoaders is empty (%v)", job.Name)
cancel() cancel()
}() }()
log.Debugf("waiting for local context to be done (%v)", job.Name)
<-jobCtx.Done() <-jobCtx.Done()
log.Debugf("local context done (%v)", job.Name)
if ctx.Err() != nil { if ctx.Err() != nil {
result.Error = ctx.Err() result.Error = ctx.Err()
@@ -163,11 +217,3 @@ func processMigrationJob(
return result return result
} }
func logColumnTypes(columnTypes []models.ColumnType, label string) {
log.Debug(label)
for _, col := range columnTypes {
log.Debugf("%+v", col)
}
}

7
config.yaml
View File

@@ -6,12 +6,15 @@ defaults:
max_extractors: 2 max_extractors: 2
max_loaders: 4 max_loaders: 4
queue_size: 8 queue_size: 8
chunk_size: 25000 batch_size: 25000
chunks_per_batch: 8 batches_per_partition: 8
truncate_target: true truncate_target: true
truncate_method: TRUNCATE # TRUNCATE | DELETE truncate_method: TRUNCATE # TRUNCATE | DELETE
retry: retry:
attempts: 3 attempts: 3
base_delay_ms: 500
max_delay_ms: 10000
max_jitter_ms: 500
jobs: jobs:
- name: demo_users - name: demo_users

2
go.mod
View File

@@ -10,6 +10,7 @@ require (
github.com/microsoft/go-mssqldb v1.9.8 github.com/microsoft/go-mssqldb v1.9.8
github.com/sirupsen/logrus v1.9.4 github.com/sirupsen/logrus v1.9.4
github.com/twpayne/go-geom v1.6.1 github.com/twpayne/go-geom v1.6.1
golang.org/x/sync v0.19.0
gopkg.in/yaml.v3 v3.0.1 gopkg.in/yaml.v3 v3.0.1
) )
@@ -23,7 +24,6 @@ require (
github.com/rogpeppe/go-internal v1.14.1 // indirect github.com/rogpeppe/go-internal v1.14.1 // indirect
github.com/shopspring/decimal v1.4.0 // indirect github.com/shopspring/decimal v1.4.0 // indirect
golang.org/x/crypto v0.48.0 // indirect golang.org/x/crypto v0.48.0 // indirect
golang.org/x/sync v0.19.0 // indirect
golang.org/x/sys v0.41.0 // indirect golang.org/x/sys v0.41.0 // indirect
golang.org/x/text v0.34.0 // indirect golang.org/x/text v0.34.0 // indirect
) )

13
internal/app/config/migration.go
View File

@@ -9,18 +9,21 @@ import (
type RetryConfig struct { type RetryConfig struct {
Attempts int `yaml:"attempts"` Attempts int `yaml:"attempts"`
BaseDelayMs int `yaml:"base_delay_ms"`
MaxDelayMs int `yaml:"max_delay_ms"`
MaxJitterMs int `yaml:"max_jitter_ms"`
} }
type JobConfig struct { type JobConfig struct {
MaxExtractors int `yaml:"max_extractors"` MaxExtractors int `yaml:"max_extractors"`
MaxLoaders int `yaml:"max_loaders"` MaxLoaders int `yaml:"max_loaders"`
QueueSize int `yaml:"queue_size"` QueueSize int `yaml:"queue_size"`
ChunkSize int `yaml:"chunk_size"` BatchSize int `yaml:"batch_size"`
ChunksPerBatch int `yaml:"chunks_per_batch"` BatchesPerPartition int `yaml:"batches_per_partition"`
RowsPerBatch int64
TruncateTarget bool `yaml:"truncate_target"` TruncateTarget bool `yaml:"truncate_target"`
TruncateMethod string `yaml:"truncate_method"` TruncateMethod string `yaml:"truncate_method"`
Retry RetryConfig `yaml:"retry"` Retry RetryConfig `yaml:"retry"`
RowsPerPartition int64
} }
type TableInfo struct { type TableInfo struct {
@@ -71,7 +74,7 @@ func (c *MigrationConfig) UnmarshalYAML(value *yaml.Node) error {
c.MaxParallelWorkers = raw.MaxParallelWorkers c.MaxParallelWorkers = raw.MaxParallelWorkers
c.Defaults = raw.Defaults c.Defaults = raw.Defaults
c.Defaults.RowsPerBatch = int64(raw.Defaults.ChunkSize * raw.Defaults.ChunksPerBatch) c.Defaults.RowsPerPartition = int64(raw.Defaults.BatchSize * raw.Defaults.BatchesPerPartition)
for _, node := range raw.Jobs { for _, node := range raw.Jobs {
job := Job{ job := Job{
@@ -82,7 +85,7 @@ func (c *MigrationConfig) UnmarshalYAML(value *yaml.Node) error {
return err return err
} }
job.RowsPerBatch = int64(job.ChunkSize * job.ChunksPerBatch) job.RowsPerPartition = int64(job.BatchSize * job.BatchesPerPartition)
c.Jobs = append(c.Jobs, job) c.Jobs = append(c.Jobs, job)
} }

61
internal/app/custom_errors/backoff.go Normal file
View File

@@ -0,0 +1,61 @@
package custom_errors
import (
"context"
"math/rand"
"time"
)
func computeBackoffDelay(retryCounter int, baseDelayMs int, maxDelayMs int, maxJitterMs int) time.Duration {
if retryCounter < 0 {
retryCounter = 0
}
delay := max(time.Duration(baseDelayMs)*time.Millisecond, 0)
maxDelay := time.Duration(maxDelayMs) * time.Millisecond
for i := 0; i < retryCounter; i++ {
if maxDelayMs > 0 && delay >= maxDelay {
delay = maxDelay
break
}
if delay == 0 {
break
}
delay *= 2
}
if maxDelayMs > 0 && delay > maxDelay {
delay = maxDelay
}
if maxJitterMs > 0 {
jitter := time.Duration(rand.Intn(maxJitterMs+1)) * time.Millisecond
delay += jitter
}
if delay < 0 {
delay = 0
}
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()
}
}()
}

58
internal/app/custom_errors/extractor.error.go
View File

@@ -5,12 +5,13 @@ import (
"fmt" "fmt"
"sync" "sync"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models" "git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
"github.com/google/uuid" "github.com/google/uuid"
) )
type ExtractorError struct { type ExtractorError struct {
Batch models.Partition Partition models.Partition
LastId int64 LastId int64
HasLastId bool HasLastId bool
Msg string Msg string
@@ -22,11 +23,11 @@ func (e *ExtractorError) Error() string {
func ExtractorErrorHandler( func ExtractorErrorHandler(
ctx context.Context, ctx context.Context,
maxRetryAttempts int, retryConfig config.RetryConfig,
chErrorsIn <-chan ExtractorError, chErrorsIn <-chan ExtractorError,
chBatchesOut chan<- models.Partition, chPartitionsOut chan<- models.Partition,
chJobErrorsOut chan<- JobError, chJobErrorsOut chan<- JobError,
wgActiveBatches *sync.WaitGroup, wgActivePartitions *sync.WaitGroup,
) { ) {
for { for {
if ctx.Err() != nil { if ctx.Err() != nil {
@@ -42,10 +43,11 @@ func ExtractorErrorHandler(
return return
} }
if err.Batch.RetryCounter >= maxRetryAttempts { if err.Partition.RetryCounter >= retryConfig.Attempts {
wgActivePartitions.Done()
jobError := JobError{ jobError := JobError{
ShouldCancelJob: false, ShouldCancelJob: false,
Msg: fmt.Sprintf("batch %v reached max retries (%d)", err.Batch.Id, maxRetryAttempts), Msg: fmt.Sprintf("Partition %v reached max retries (%d)", err.Partition.Id, retryConfig.Attempts),
Prev: &err, Prev: &err,
} }
@@ -55,25 +57,45 @@ func ExtractorErrorHandler(
return return
} }
wgActiveBatches.Done()
continue continue
} } else {
jobError := JobError{
newBatch := err.Batch ShouldCancelJob: false,
newBatch.RetryCounter++ Msg: fmt.Sprintf("Temporal error in partition %v (retries: %d)", err.Partition.Id, err.Partition.RetryCounter),
Prev: &err,
if err.HasLastId {
newBatch.ParentId = err.Batch.Id
newBatch.Id = uuid.New()
newBatch.LowerLimit = err.LastId
newBatch.IsLowerLimitInclusive = false
} }
select { select {
case chBatchesOut <- newBatch: case chJobErrorsOut <- jobError:
case <-ctx.Done(): case <-ctx.Done():
return return
} }
} }
newPartition := err.Partition
newPartition.RetryCounter++
delay := computeBackoffDelay(
newPartition.RetryCounter,
retryConfig.BaseDelayMs,
retryConfig.MaxDelayMs,
retryConfig.MaxJitterMs,
)
if err.HasLastId {
newPartition.ParentId = err.Partition.Id
newPartition.Id = uuid.New()
newPartition.LowerLimit = err.LastId
newPartition.IsLowerLimitInclusive = false
}
requeueWithBackoff(ctx, delay, func() {
select {
case chPartitionsOut <- newPartition:
case <-ctx.Done():
return
}
})
}
} }
} }

4
internal/app/custom_errors/job.error.go
View File

@@ -37,11 +37,11 @@ func JobErrorHandler(ctx context.Context, chErrorsIn <-chan JobError) error {
} }
if err.ShouldCancelJob { if err.ShouldCancelJob {
log.Error(err.Msg, " - ", err.Prev) log.Errorf("(Fatal job error) - %v - %v", err.Msg, err.Prev)
return &err return &err
} }
log.Error(err.Msg, " - ", err.Prev) log.Errorf("%v - %v", err.Msg, err.Prev)
} }
} }
} }

41
internal/app/custom_errors/loader.error.go
View File

@@ -5,11 +5,12 @@ import (
"fmt" "fmt"
"sync" "sync"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models" "git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
) )
type LoaderError struct { type LoaderError struct {
models.Batch Batch models.Batch
Msg string Msg string
} }
@@ -19,11 +20,11 @@ func (e *LoaderError) Error() string {
func LoaderErrorHandler( func LoaderErrorHandler(
ctx context.Context, ctx context.Context,
maxRetryAttempts int, retryConfig config.RetryConfig,
chErrorsIn <-chan LoaderError, chErrorsIn <-chan LoaderError,
chChunksOut chan<- models.Batch, chBatchesOut chan<- models.Batch,
chJobErrorsOut chan<- JobError, chJobErrorsOut chan<- JobError,
wgActiveChunks *sync.WaitGroup, wgActiveBatches *sync.WaitGroup,
) { ) {
for { for {
if ctx.Err() != nil { if ctx.Err() != nil {
@@ -39,10 +40,11 @@ func LoaderErrorHandler(
return return
} }
if err.RetryCounter >= maxRetryAttempts { if err.Batch.RetryCounter >= retryConfig.Attempts {
wgActiveBatches.Done()
jobError := JobError{ jobError := JobError{
ShouldCancelJob: false, ShouldCancelJob: false,
Msg: fmt.Sprintf("chunk %v reached max retries (%d)", err.Id, maxRetryAttempts), Msg: fmt.Sprintf("Batch %v reached max retries (%d)", err.Batch.Id, retryConfig.Attempts),
Prev: &err, Prev: &err,
} }
@@ -52,17 +54,36 @@ func LoaderErrorHandler(
return return
} }
wgActiveChunks.Done()
continue continue
} else {
jobError := JobError{
ShouldCancelJob: false,
Msg: fmt.Sprintf("Temporal error in batch %v (retries: %d)", err.Batch.Id, err.Batch.RetryCounter),
Prev: &err,
} }
err.RetryCounter++
select { select {
case chChunksOut <- err.Batch: case chJobErrorsOut <- jobError:
case <-ctx.Done(): case <-ctx.Done():
return 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
}
})
}
} }
} }

110
internal/app/etl/extractors/mssql.go
View File

@@ -70,20 +70,20 @@ func buildExtractQueryMssql(
return sbQuery.String() return sbQuery.String()
} }
func extractorErrorFromLastRowMssql( func errorFromLastRow(
lastRow models.UnknownRowValues, lastRow models.UnknownRowValues,
indexPrimaryKey int, indexPrimaryKey int,
batch *models.Partition, partition *models.Partition,
previousError error, previousError error,
) *custom_errors.ExtractorError { ) *custom_errors.ExtractorError {
lastIdRawValue := lastRow[indexPrimaryKey] lastIdRawValue := lastRow[indexPrimaryKey]
lastId, ok := convert.ToInt64(lastIdRawValue) lastId, ok := convert.ToInt64(lastIdRawValue)
if !ok { if !ok {
currentBatch := *batch currentPartition := *partition
currentBatch.RetryCounter = 3 currentPartition.RetryCounter = 3
return &custom_errors.ExtractorError{ return &custom_errors.ExtractorError{
Batch: currentBatch, Partition: currentPartition,
HasLastId: true, HasLastId: true,
Msg: fmt.Sprintf("Couldn't cast last id value as int: %s", previousError.Error()), Msg: fmt.Sprintf("Couldn't cast last id value as int: %s", previousError.Error()),
} }
@@ -91,78 +91,78 @@ func extractorErrorFromLastRowMssql(
} }
return &custom_errors.ExtractorError{ return &custom_errors.ExtractorError{
Batch: *batch, Partition: *partition,
HasLastId: true, HasLastId: true,
LastId: lastId, LastId: lastId,
Msg: previousError.Error(), Msg: previousError.Error(),
} }
} }
func (mssqlEx *MssqlExtractor) ProcessBatch( func (mssqlEx *MssqlExtractor) ProcessPartition(
ctx context.Context, ctx context.Context,
tableInfo config.SourceTableInfo, tableInfo config.SourceTableInfo,
columns []models.ColumnType, columns []models.ColumnType,
chunkSize int, batchSize int,
batch models.Partition, partition models.Partition,
indexPrimaryKey int, indexPrimaryKey int,
chChunksOut chan<- models.Batch, chBatchesOut chan<- models.Batch,
rowsRead *int64, rowsRead *int64,
) error { ) error {
query := buildExtractQueryMssql(tableInfo, columns, batch.ShouldUseRange, batch.IsLowerLimitInclusive) query := buildExtractQueryMssql(tableInfo, columns, partition.ShouldUseRange, partition.IsLowerLimitInclusive)
var queryArgs []any var queryArgs []any
if batch.ShouldUseRange { if partition.ShouldUseRange {
queryArgs = append(queryArgs, queryArgs = append(queryArgs,
sql.Named("min", batch.LowerLimit), sql.Named("min", partition.LowerLimit),
sql.Named("max", batch.UpperLimit), sql.Named("max", partition.UpperLimit),
) )
} }
rows, err := mssqlEx.db.QueryContext(ctx, query, queryArgs...) rows, err := mssqlEx.db.QueryContext(ctx, query, queryArgs...)
if err != nil { if err != nil {
return &custom_errors.ExtractorError{Batch: batch, HasLastId: false, Msg: err.Error()} return &custom_errors.ExtractorError{Partition: partition, HasLastId: false, Msg: err.Error()}
} }
defer rows.Close() defer rows.Close()
rowsChunk := make([]models.UnknownRowValues, 0, chunkSize) batchRows := make([]models.UnknownRowValues, 0, batchSize)
for rows.Next() { for rows.Next() {
values := make([]any, len(columns)) rowValues := make([]any, len(columns))
scanArgs := make([]any, len(columns)) scanArgs := make([]any, len(columns))
for i := range values { for i := range rowValues {
scanArgs[i] = &values[i] scanArgs[i] = &rowValues[i]
} }
if err := rows.Scan(scanArgs...); err != nil { if err := rows.Scan(scanArgs...); err != nil {
if len(rowsChunk) == 0 { if len(batchRows) == 0 {
return &custom_errors.ExtractorError{Batch: batch, HasLastId: false, Msg: err.Error()} return &custom_errors.ExtractorError{Partition: partition, HasLastId: false, Msg: err.Error()}
} }
lastRow := rowsChunk[len(rowsChunk)-1] lastRow := batchRows[len(batchRows)-1]
select { select {
case chChunksOut <- models.Batch{Id: uuid.New(), PartitionId: batch.Id, Data: rowsChunk, RetryCounter: 0}: case chBatchesOut <- models.Batch{Id: uuid.New(), PartitionId: partition.Id, Rows: batchRows, RetryCounter: 0}:
case <-ctx.Done(): case <-ctx.Done():
return nil return nil
} }
atomic.AddInt64(rowsRead, int64(len(rowsChunk))) atomic.AddInt64(rowsRead, int64(len(batchRows)))
return extractorErrorFromLastRowMssql(lastRow, indexPrimaryKey, &batch, err) return errorFromLastRow(lastRow, indexPrimaryKey, &partition, err)
} }
rowsChunk = append(rowsChunk, values) batchRows = append(batchRows, rowValues)
if len(rowsChunk) >= chunkSize { if len(batchRows) >= batchSize {
select { select {
case chChunksOut <- models.Batch{Id: uuid.New(), PartitionId: batch.Id, Data: rowsChunk, RetryCounter: 0}: case chBatchesOut <- models.Batch{Id: uuid.New(), PartitionId: partition.Id, Rows: batchRows, RetryCounter: 0}:
case <-ctx.Done(): case <-ctx.Done():
return nil return nil
} }
atomic.AddInt64(rowsRead, int64(len(rowsChunk))) atomic.AddInt64(rowsRead, int64(len(batchRows)))
rowsChunk = make([]models.UnknownRowValues, 0, chunkSize) batchRows = make([]models.UnknownRowValues, 0, batchSize)
} }
} }
@@ -171,22 +171,22 @@ func (mssqlEx *MssqlExtractor) ProcessBatch(
return ctx.Err() return ctx.Err()
} }
if len(rowsChunk) == 0 { if len(batchRows) == 0 {
return &custom_errors.ExtractorError{Batch: batch, HasLastId: false, Msg: err.Error()} return &custom_errors.ExtractorError{Partition: partition, HasLastId: false, Msg: err.Error()}
} }
lastRow := rowsChunk[len(rowsChunk)-1] lastRow := batchRows[len(batchRows)-1]
return extractorErrorFromLastRowMssql(lastRow, indexPrimaryKey, &batch, err) return errorFromLastRow(lastRow, indexPrimaryKey, &partition, err)
} }
if len(rowsChunk) > 0 { if len(batchRows) > 0 {
select { select {
case chChunksOut <- models.Batch{Id: uuid.New(), PartitionId: batch.Id, Data: rowsChunk, RetryCounter: 0}: case chBatchesOut <- models.Batch{Id: uuid.New(), PartitionId: partition.Id, Rows: batchRows, RetryCounter: 0}:
case <-ctx.Done(): case <-ctx.Done():
return nil return nil
} }
atomic.AddInt64(rowsRead, int64(len(rowsChunk))) atomic.AddInt64(rowsRead, int64(len(batchRows)))
} }
return nil return nil
@@ -196,12 +196,12 @@ func (mssqlEx *MssqlExtractor) Exec(
ctx context.Context, ctx context.Context,
tableInfo config.SourceTableInfo, tableInfo config.SourceTableInfo,
columns []models.ColumnType, columns []models.ColumnType,
chunkSize int, batchSize int,
chBatchesIn <-chan models.Partition, chPartitionsIn <-chan models.Partition,
chChunksOut chan<- models.Batch, chBatchesOut chan<- models.Batch,
chErrorsOut chan<- custom_errors.ExtractorError, chErrorsOut chan<- custom_errors.ExtractorError,
chJobErrorsOut chan<- custom_errors.JobError, chJobErrorsOut chan<- custom_errors.JobError,
wgActiveBatches *sync.WaitGroup, wgActivePartitions *sync.WaitGroup,
rowsRead *int64, rowsRead *int64,
) { ) {
indexPrimaryKey := slices.IndexFunc(columns, func(col models.ColumnType) bool { indexPrimaryKey := slices.IndexFunc(columns, func(col models.ColumnType) bool {
@@ -229,45 +229,49 @@ func (mssqlEx *MssqlExtractor) Exec(
select { select {
case <-ctx.Done(): case <-ctx.Done():
return return
case batch, ok := <-chBatchesIn: case partition, ok := <-chPartitionsIn:
if !ok { if !ok {
return return
} }
err := mssqlEx.ProcessBatch( err := mssqlEx.ProcessPartition(
ctx, ctx,
tableInfo, tableInfo,
columns, columns,
chunkSize, batchSize,
batch, partition,
indexPrimaryKey, indexPrimaryKey,
chChunksOut, chBatchesOut,
rowsRead, rowsRead,
) )
if err != nil { if err != nil {
var exError *custom_errors.ExtractorError var exError *custom_errors.ExtractorError
var jobError *custom_errors.JobError
if errors.As(err, &exError) { if errors.As(err, &exError) {
select { select {
case <-ctx.Done(): case <-ctx.Done():
return return
case chErrorsOut <- *exError: case chErrorsOut <- *exError:
} }
} } else if errors.As(err, &jobError) {
var jobError *custom_errors.JobError
if errors.As(err, &jobError) {
select { select {
case <-ctx.Done(): case <-ctx.Done():
return return
case chJobErrorsOut <- *jobError: case chJobErrorsOut <- *jobError:
} }
} } else {
select {
case <-ctx.Done():
return return
case chErrorsOut <- custom_errors.ExtractorError{Partition: partition, Msg: err.Error()}:
}
} }
wgActiveBatches.Done() continue
}
wgActivePartitions.Done()
} }
} }
} }

38
internal/app/etl/extractors/postgres.go
View File

@@ -52,29 +52,29 @@ func buildExtractQueryPostgres(sourceDbInfo config.SourceTableInfo, columns []mo
return fmt.Sprintf(`SELECT %s FROM "%s"."%s" ORDER BY "%s" ASC`, sbColumns.String(), sourceDbInfo.Schema, sourceDbInfo.Table, sourceDbInfo.PrimaryKey) return fmt.Sprintf(`SELECT %s FROM "%s"."%s" ORDER BY "%s" ASC`, sbColumns.String(), sourceDbInfo.Schema, sourceDbInfo.Table, sourceDbInfo.PrimaryKey)
} }
func (postgresEx *PostgresExtractor) ProcessBatch( func (postgresEx *PostgresExtractor) ProcessPartition(
ctx context.Context, ctx context.Context,
tableInfo config.SourceTableInfo, tableInfo config.SourceTableInfo,
columns []models.ColumnType, columns []models.ColumnType,
chunkSize int, batchSize int,
batch models.Partition, partition models.Partition,
indexPrimaryKey int, indexPrimaryKey int,
chChunksOut chan<- models.Batch, chBatchesOut chan<- models.Batch,
rowsRead *int64, rowsRead *int64,
) error { ) error {
query := buildExtractQueryPostgres(tableInfo, columns) query := buildExtractQueryPostgres(tableInfo, columns)
if batch.ShouldUseRange { if partition.ShouldUseRange {
return errors.New("Batch config not yet supported") return errors.New("Batch config not yet supported")
} }
rows, err := postgresEx.db.Query(ctx, query) rows, err := postgresEx.db.Query(ctx, query)
if err != nil { if err != nil {
return &custom_errors.ExtractorError{Batch: batch, HasLastId: false, Msg: err.Error()} return &custom_errors.ExtractorError{Partition: partition, HasLastId: false, Msg: err.Error()}
} }
defer rows.Close() defer rows.Close()
rowsChunk := make([]models.UnknownRowValues, 0, chunkSize) batchRows := make([]models.UnknownRowValues, 0, batchSize)
for rows.Next() { for rows.Next() {
values, err := rows.Values() values, err := rows.Values()
@@ -82,17 +82,17 @@ func (postgresEx *PostgresExtractor) ProcessBatch(
return errors.New("Unexpected error reading rows from source") return errors.New("Unexpected error reading rows from source")
} }
rowsChunk = append(rowsChunk, values) batchRows = append(batchRows, values)
if len(rowsChunk) >= chunkSize { if len(batchRows) >= batchSize {
select { select {
case chChunksOut <- models.Batch{Id: uuid.New(), PartitionId: batch.Id, Data: rowsChunk, RetryCounter: 0}: case chBatchesOut <- models.Batch{Id: uuid.New(), PartitionId: partition.Id, Rows: batchRows, RetryCounter: 0}:
case <-ctx.Done(): case <-ctx.Done():
return nil return nil
} }
atomic.AddInt64(rowsRead, int64(len(rowsChunk))) atomic.AddInt64(rowsRead, int64(len(batchRows)))
rowsChunk = make([]models.UnknownRowValues, 0, chunkSize) batchRows = make([]models.UnknownRowValues, 0, batchSize)
} }
} }
@@ -100,14 +100,14 @@ func (postgresEx *PostgresExtractor) ProcessBatch(
return errors.New("Unexpected error reading rows from source") return errors.New("Unexpected error reading rows from source")
} }
if len(rowsChunk) > 0 { if len(batchRows) > 0 {
select { select {
case chChunksOut <- models.Batch{Id: uuid.New(), PartitionId: batch.Id, Data: rowsChunk, RetryCounter: 0}: case chBatchesOut <- models.Batch{Id: uuid.New(), PartitionId: partition.Id, Rows: batchRows, RetryCounter: 0}:
case <-ctx.Done(): case <-ctx.Done():
return nil return nil
} }
atomic.AddInt64(rowsRead, int64(len(rowsChunk))) atomic.AddInt64(rowsRead, int64(len(batchRows)))
} }
return nil return nil
@@ -117,12 +117,12 @@ func (postgresEx *PostgresExtractor) Exec(
ctx context.Context, ctx context.Context,
tableInfo config.SourceTableInfo, tableInfo config.SourceTableInfo,
columns []models.ColumnType, columns []models.ColumnType,
chunkSize int, batchSize int,
chBatchesIn <-chan models.Partition, chPartitionsIn <-chan models.Partition,
chChunksOut chan<- models.Batch, chBatchesOut chan<- models.Batch,
chErrorsOut chan<- custom_errors.ExtractorError, chErrorsOut chan<- custom_errors.ExtractorError,
chJobErrorsOut chan<- custom_errors.JobError, chJobErrorsOut chan<- custom_errors.JobError,
wgActiveBatches *sync.WaitGroup, wgActivePartitions *sync.WaitGroup,
rowsRead *int64, rowsRead *int64,
) { ) {
} }

36
internal/app/etl/loaders/postgres.go
View File

@@ -34,18 +34,18 @@ func mapSlice[T any, V any](input []T, mapper func(T) V) []V {
return result return result
} }
func (postgresLd *PostgresLoader) ProcessChunk( func (postgresLd *PostgresLoader) ProcessBatch(
ctx context.Context, ctx context.Context,
tableInfo config.TargetTableInfo, tableInfo config.TargetTableInfo,
colNames []string, colNames []string,
chunk models.Batch, batch models.Batch,
) (int, error) { ) (int, error) {
tableId := pgx.Identifier{tableInfo.Schema, tableInfo.Table} tableId := pgx.Identifier{tableInfo.Schema, tableInfo.Table}
_, err := postgresLd.db.CopyFrom( _, err := postgresLd.db.CopyFrom(
ctx, ctx,
tableId, tableId,
colNames, colNames,
pgx.CopyFromRows(chunk.Data), pgx.CopyFromRows(batch.Rows),
) )
if err != nil { if err != nil {
@@ -60,20 +60,20 @@ func (postgresLd *PostgresLoader) ProcessChunk(
} }
} }
return 0, &custom_errors.LoaderError{Batch: chunk, Msg: err.Error()} return 0, &custom_errors.LoaderError{Batch: batch, Msg: err.Error()}
} }
return len(chunk.Data), nil return len(batch.Rows), nil
} }
func (postgresLd *PostgresLoader) Exec( func (postgresLd *PostgresLoader) Exec(
ctx context.Context, ctx context.Context,
tableInfo config.TargetTableInfo, tableInfo config.TargetTableInfo,
columns []models.ColumnType, columns []models.ColumnType,
chChunksIn <-chan models.Batch, chBatchesIn <-chan models.Batch,
chErrorsOut chan<- custom_errors.LoaderError, chErrorsOut chan<- custom_errors.LoaderError,
chJobErrorsOut chan<- custom_errors.JobError, chJobErrorsOut chan<- custom_errors.JobError,
wgActiveChunks *sync.WaitGroup, wgActiveBatches *sync.WaitGroup,
rowsLoaded *int64, rowsLoaded *int64,
) { ) {
colNames := mapSlice(columns, func(col models.ColumnType) string { colNames := mapSlice(columns, func(col models.ColumnType) string {
@@ -88,36 +88,40 @@ func (postgresLd *PostgresLoader) Exec(
select { select {
case <-ctx.Done(): case <-ctx.Done():
return return
case chunk, ok := <-chChunksIn: case batch, ok := <-chBatchesIn:
if !ok { if !ok {
return return
} }
processedRows, err := postgresLd.ProcessChunk(ctx, tableInfo, colNames, chunk) processedRows, err := postgresLd.ProcessBatch(ctx, tableInfo, colNames, batch)
if err != nil { if err != nil {
var ldError *custom_errors.LoaderError var ldError *custom_errors.LoaderError
var jobError *custom_errors.JobError
if errors.As(err, &ldError) { if errors.As(err, &ldError) {
select { select {
case <-ctx.Done(): case <-ctx.Done():
return return
case chErrorsOut <- *ldError: case chErrorsOut <- *ldError:
} }
} } else if errors.As(err, &jobError) {
var jobError *custom_errors.JobError
if errors.As(err, &jobError) {
select { select {
case <-ctx.Done(): case <-ctx.Done():
return return
case chJobErrorsOut <- *jobError: case chJobErrorsOut <- *jobError:
} }
} } else {
select {
case <-ctx.Done():
return return
case chErrorsOut <- custom_errors.LoaderError{Batch: batch, Msg: err.Error()}:
}
} }
wgActiveChunks.Done() continue
}
wgActiveBatches.Done()
atomic.AddInt64(rowsLoaded, int64(processedRows)) atomic.AddInt64(rowsLoaded, int64(processedRows))
} }
} }

40
internal/app/etl/table_analyzers/main.go Normal file
View File

@@ -0,0 +1,40 @@
package table_analyzers
import (
"context"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/etl"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
"github.com/google/uuid"
)
func PartitionRangeGenerator(
ctx context.Context,
tableAnalyzer etl.TableAnalyzer,
tableInfo config.TableInfo,
partitionColumn string,
rowsPerPartition int64,
) ([]models.Partition, error) {
rowsCount, err := tableAnalyzer.EstimateTotalRows(ctx, tableInfo)
if err != nil {
return nil, err
}
if rowsCount <= rowsPerPartition {
return []models.Partition{{
Id: uuid.New(),
ShouldUseRange: false,
RetryCounter: 0,
}}, nil
}
partitionsCount := rowsCount / rowsPerPartition
partitions, err := tableAnalyzer.CalculatePartitionRanges(ctx, tableInfo, partitionColumn, partitionsCount)
if err != nil {
return nil, err
}
return partitions, nil
}

249
internal/app/etl/table_analyzers/mssql.go Normal file
View File

@@ -0,0 +1,249 @@
package table_analyzers
import (
"context"
"database/sql"
"fmt"
"strings"
"time"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/etl"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
"github.com/google/uuid"
)
type MssqlTableAnalyzer struct {
db *sql.DB
}
func NewMssqlTableAnalyzer(db *sql.DB) etl.TableAnalyzer {
return &MssqlTableAnalyzer{db: db}
}
const mssqlColumnMetadataQuery string = `
SELECT
c.name AS name,
t.name AS user_type,
CASE WHEN t.is_user_defined = 0 THEN t.name ELSE bt.name END AS system_type,
c.is_nullable AS nullable,
c.max_length AS max_length,
c.precision AS precision,
c.scale AS scale
FROM sys.columns c
JOIN sys.types t ON c.user_type_id = t.user_type_id
LEFT JOIN sys.types bt ON t.is_user_defined = 1 AND bt.user_type_id = t.system_type_id
JOIN sys.tables st ON c.object_id = st.object_id
JOIN sys.schemas s ON st.schema_id = s.schema_id
WHERE s.name = @schema AND st.name = @table AND c.name NOT LIKE 'graph_id%'
ORDER BY c.column_id;`
type rawColumnMssql struct {
name string
userType string
systemType string
nullable bool
maxLength int64
precision int64
scale int64
}
func (ta *MssqlTableAnalyzer) systemTypeToUnifiedType(systemType string) string {
systemType = strings.ToLower(systemType)
if systemType == "varchar" || systemType == "char" || systemType == "nvarchar" || systemType == "nchar" || systemType == "text" || systemType == "ntext" {
return "STRING"
}
if systemType == "int" || systemType == "int4" || systemType == "integer" || systemType == "smallint" || systemType == "int2" || systemType == "bigint" || systemType == "int8" || systemType == "tinyint" {
return "INTEGER"
}
if systemType == "decimal" || systemType == "numeric" {
return "DECIMAL"
}
if systemType == "float" || systemType == "real" || systemType == "double precision" {
return "FLOAT"
}
if systemType == "bit" || systemType == "boolean" {
return "BOOLEAN"
}
if systemType == "date" {
return "DATE"
}
if systemType == "time" || systemType == "time without time zone" {
return "TIME"
}
if systemType == "datetime" || systemType == "datetime2" || systemType == "timestamp" || systemType == "timestamptz" || systemType == "timestamp with time zone" {
return "TIMESTAMP"
}
if systemType == "binary" || systemType == "varbinary" || systemType == "image" || systemType == "bytea" {
return "BINARY"
}
if systemType == "uniqueidentifier" || systemType == "uuid" {
return "UUID"
}
if systemType == "json" {
return "JSON"
}
if systemType == "geometry" || systemType == "geography" {
return "GEOMETRY"
}
return strings.ToUpper(systemType)
}
func (ta *MssqlTableAnalyzer) rawColumnToColumnType(rawColumn rawColumnMssql) models.ColumnType {
const nullValue int64 = -1
stringTypes := map[string]bool{"varchar": true, "char": true, "nvarchar": true, "nchar": true, "text": true, "ntext": true}
decimalTypes := map[string]bool{"decimal": true, "numeric": true}
if stringTypes[rawColumn.systemType] {
if rawColumn.systemType == "nvarchar" || rawColumn.systemType == "nchar" {
if rawColumn.maxLength > 0 {
rawColumn.maxLength = rawColumn.maxLength / 2
}
}
rawColumn.precision, rawColumn.scale = nullValue, nullValue
} else if decimalTypes[rawColumn.systemType] {
rawColumn.maxLength = nullValue
} else {
rawColumn.maxLength, rawColumn.precision, rawColumn.scale = nullValue, nullValue, nullValue
}
columnType := models.NewColumnType(
rawColumn.name,
rawColumn.maxLength != nullValue,
rawColumn.precision != nullValue || rawColumn.scale != nullValue,
rawColumn.userType,
rawColumn.systemType,
ta.systemTypeToUnifiedType(rawColumn.systemType),
rawColumn.nullable,
rawColumn.maxLength,
rawColumn.precision,
rawColumn.scale,
)
return columnType
}
func (ta *MssqlTableAnalyzer) QueryColumnTypes(
ctx context.Context,
tableInfo config.TableInfo,
) ([]models.ColumnType, error) {
localCtx, cancel := context.WithTimeout(ctx, 20*time.Second)
defer cancel()
rows, err := ta.db.QueryContext(localCtx, mssqlColumnMetadataQuery, sql.Named("schema", tableInfo.Schema), sql.Named("table", tableInfo.Table))
if err != nil {
return nil, err
}
defer rows.Close()
var columnTypes []models.ColumnType
for rows.Next() {
var rawColumn rawColumnMssql
if err := rows.Scan(
&rawColumn.name,
&rawColumn.userType,
&rawColumn.systemType,
&rawColumn.nullable,
&rawColumn.maxLength,
&rawColumn.precision,
&rawColumn.scale,
); err != nil {
return nil, err
}
columnTypes = append(columnTypes, ta.rawColumnToColumnType(rawColumn))
}
return columnTypes, nil
}
func (ta *MssqlTableAnalyzer) EstimateTotalRows(
ctx context.Context,
tableInfo config.TableInfo,
) (int64, error) {
query := `
SELECT SUM(p.rows) AS count
FROM sys.tables t
JOIN sys.schemas s ON t.schema_id = s.schema_id
JOIN sys.partitions p ON t.object_id = p.object_id
WHERE s.name = @schema AND t.name = @table AND p.index_id IN (0, 1)
GROUP BY t.name`
ctxTimeout, cancel := context.WithTimeout(ctx, time.Second*20)
defer cancel()
var rowsCount int64
err := ta.db.QueryRowContext(ctxTimeout, query, sql.Named("schema", tableInfo.Schema), sql.Named("table", tableInfo.Table)).Scan(&rowsCount)
if err != nil {
return 0, err
}
return rowsCount, nil
}
func (ta *MssqlTableAnalyzer) CalculatePartitionRanges(
ctx context.Context,
tableInfo config.TableInfo,
partitionColumn string,
maxPartitions int64,
) ([]models.Partition, error) {
query := fmt.Sprintf(`
SELECT
MIN([%s]) AS lower_limit,
MAX([%s]) AS upper_limit
FROM (SELECT [%s], NTILE(@maxPartitions) OVER (ORDER BY [%s]) AS batch_id FROM [%s].[%s]) AS T
GROUP BY batch_id
ORDER BY batch_id`,
partitionColumn,
partitionColumn,
partitionColumn,
partitionColumn,
tableInfo.Schema,
tableInfo.Table)
ctxTimeout, cancel := context.WithTimeout(ctx, time.Second*20)
defer cancel()
rows, err := ta.db.QueryContext(ctxTimeout, query, sql.Named("maxPartitions", maxPartitions))
if err != nil {
return nil, err
}
defer rows.Close()
partitions := make([]models.Partition, 0, maxPartitions)
for rows.Next() {
partition := models.Partition{
Id: uuid.New(),
ShouldUseRange: true,
RetryCounter: 0,
IsLowerLimitInclusive: true,
}
if err := rows.Scan(&partition.LowerLimit, &partition.UpperLimit); err != nil {
return nil, err
}
partitions = append(partitions, partition)
}
if err := rows.Err(); err != nil {
return nil, err
}
return partitions, nil
}

174
internal/app/etl/table_analyzers/postgres.go Normal file
View File

@@ -0,0 +1,174 @@
package table_analyzers
import (
"context"
"strings"
"time"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/etl"
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
"github.com/jackc/pgx/v5/pgxpool"
)
type PostgresTableAnalyzer struct {
db *pgxpool.Pool
}
func NewPostgresTableAnalyzer(db *pgxpool.Pool) etl.TableAnalyzer {
return &PostgresTableAnalyzer{db: db}
}
const postgresColumnMetadataQuery string = `
SELECT
c.column_name AS name,
c.data_type AS user_type,
c.udt_name AS system_type,
(CASE WHEN c.is_nullable = 'YES' THEN TRUE ELSE FALSE END) AS nullable,
COALESCE(c.character_maximum_length, -1) AS max_length,
COALESCE(c.numeric_precision, -1) AS precision,
COALESCE(c.numeric_scale, -1) AS scale
FROM information_schema.columns c
WHERE c.table_schema = $1 AND c.table_name = $2
ORDER BY c.ordinal_position;`
type rawColumnPostgres struct {
name string
userType string
systemType string
nullable bool
maxLength int64
precision int64
scale int64
}
func (ta *PostgresTableAnalyzer) systemTypeToUnifiedType(systemType string) string {
systemType = strings.ToLower(systemType)
if systemType == "varchar" || systemType == "char" || systemType == "nvarchar" || systemType == "nchar" || systemType == "text" || systemType == "ntext" {
return "STRING"
}
if systemType == "int" || systemType == "int4" || systemType == "integer" || systemType == "smallint" || systemType == "int2" || systemType == "bigint" || systemType == "int8" || systemType == "tinyint" {
return "INTEGER"
}
if systemType == "decimal" || systemType == "numeric" {
return "DECIMAL"
}
if systemType == "float" || systemType == "real" || systemType == "double precision" {
return "FLOAT"
}
if systemType == "bit" || systemType == "boolean" {
return "BOOLEAN"
}
if systemType == "date" {
return "DATE"
}
if systemType == "time" || systemType == "time without time zone" {
return "TIME"
}
if systemType == "datetime" || systemType == "datetime2" || systemType == "timestamp" || systemType == "timestamptz" || systemType == "timestamp with time zone" {
return "TIMESTAMP"
}
if systemType == "binary" || systemType == "varbinary" || systemType == "image" || systemType == "bytea" {
return "BINARY"
}
if systemType == "uniqueidentifier" || systemType == "uuid" {
return "UUID"
}
if systemType == "json" {
return "JSON"
}
if systemType == "geometry" || systemType == "geography" {
return "GEOMETRY"
}
return strings.ToUpper(systemType)
}
func (ta *PostgresTableAnalyzer) rawColumnToColumnType(rawColumn rawColumnPostgres) models.ColumnType {
const nullValue int64 = -1
stringTypes := map[string]bool{"varchar": true, "char": true, "text": true}
decimalTypes := map[string]bool{"decimal": true, "numeric": true}
if stringTypes[rawColumn.systemType] {
rawColumn.precision, rawColumn.scale = nullValue, nullValue
} else if decimalTypes[rawColumn.systemType] {
rawColumn.maxLength = nullValue
} else {
rawColumn.maxLength, rawColumn.precision, rawColumn.scale = nullValue, nullValue, nullValue
}
return models.NewColumnType(
rawColumn.name,
rawColumn.maxLength != nullValue,
rawColumn.precision != nullValue || rawColumn.scale != nullValue,
rawColumn.userType,
rawColumn.systemType,
ta.systemTypeToUnifiedType(rawColumn.systemType),
rawColumn.nullable,
rawColumn.maxLength,
rawColumn.precision,
rawColumn.scale,
)
}
func (ta *PostgresTableAnalyzer) QueryColumnTypes(
ctx context.Context,
tableInfo config.TableInfo,
) ([]models.ColumnType, error) {
localCtx, cancel := context.WithTimeout(ctx, 20*time.Second)
defer cancel()
rows, err := ta.db.Query(localCtx, postgresColumnMetadataQuery, tableInfo.Schema, tableInfo.Table)
if err != nil {
return nil, err
}
defer rows.Close()
var colTypes []models.ColumnType
for rows.Next() {
var column rawColumnPostgres
if err := rows.Scan(
&column.name,
&column.userType,
&column.systemType,
&column.nullable,
&column.maxLength,
&column.precision,
&column.scale,
); err != nil {
return nil, err
}
colTypes = append(colTypes, ta.rawColumnToColumnType(column))
}
return colTypes, nil
}
func (ta *PostgresTableAnalyzer) EstimateTotalRows(
ctx context.Context,
tableInfo config.TableInfo,
) (int64, error) {
return 0, nil
}
func (ta *PostgresTableAnalyzer) CalculatePartitionRanges(
ctx context.Context,
tableInfo config.TableInfo,
partitionColumn string,
maxPartitions int64,
) ([]models.Partition, error) {
return []models.Partition{}, nil
}

28
internal/app/etl/transformers/mssql.go
View File

@@ -60,15 +60,15 @@ func computeTransformationPlan(columns []models.ColumnType) []etl.ColumnTransfor
return plan return plan
} }
const processChunkCtxCheck = 4096 const processBatchCtxCheck = 4096
func (mssqlTr *MssqlTransformer) ProcessChunk( func (mssqlTr *MssqlTransformer) ProcessBatch(
ctx context.Context, ctx context.Context,
chunk *models.Batch, batch *models.Batch,
transformationPlan []etl.ColumnTransformPlan, transformationPlan []etl.ColumnTransformPlan,
) error { ) error {
for i, rowValues := range chunk.Data { for i, rowValues := range batch.Rows {
if i%processChunkCtxCheck == 0 { if i%processBatchCtxCheck == 0 {
if err := ctx.Err(); err != nil { if err := ctx.Err(); err != nil {
return err return err
} }
@@ -94,10 +94,10 @@ func (mssqlTr *MssqlTransformer) ProcessChunk(
func (mssqlTr *MssqlTransformer) Exec( func (mssqlTr *MssqlTransformer) Exec(
ctx context.Context, ctx context.Context,
columns []models.ColumnType, columns []models.ColumnType,
chChunksIn <-chan models.Batch, chBatchesIn <-chan models.Batch,
chChunksOut chan<- models.Batch, chBatchesOut chan<- models.Batch,
chJobErrorsOut chan<- custom_errors.JobError, chJobErrorsOut chan<- custom_errors.JobError,
wgActiveChunks *sync.WaitGroup, wgActiveBatches *sync.WaitGroup,
) { ) {
transformationPlan := computeTransformationPlan(columns) transformationPlan := computeTransformationPlan(columns)
@@ -110,22 +110,22 @@ func (mssqlTr *MssqlTransformer) Exec(
case <-ctx.Done(): case <-ctx.Done():
return return
case chunk, ok := <-chChunksIn: case batch, ok := <-chBatchesIn:
if !ok { if !ok {
return return
} }
if len(transformationPlan) == 0 { if len(transformationPlan) == 0 {
select { select {
case chChunksOut <- chunk: case chBatchesOut <- batch:
wgActiveChunks.Add(1) wgActiveBatches.Add(1)
continue continue
case <-ctx.Done(): case <-ctx.Done():
return return
} }
} }
err := mssqlTr.ProcessChunk(ctx, &chunk, transformationPlan) err := mssqlTr.ProcessBatch(ctx, &batch, transformationPlan)
if err != nil { if err != nil {
if errors.Is(err, ctx.Err()) { if errors.Is(err, ctx.Err()) {
return return
@@ -139,12 +139,12 @@ func (mssqlTr *MssqlTransformer) Exec(
} }
select { select {
case chChunksOut <- chunk: case chBatchesOut <- batch:
case <-ctx.Done(): case <-ctx.Done():
return return
} }
wgActiveChunks.Add(1) wgActiveBatches.Add(1)
} }
} }
} }

41
internal/app/etl/types.go
View File

@@ -10,14 +10,14 @@ import (
) )
type Extractor interface { type Extractor interface {
ProcessBatch( ProcessPartition(
ctx context.Context, ctx context.Context,
tableInfo config.SourceTableInfo, tableInfo config.SourceTableInfo,
columns []models.ColumnType, columns []models.ColumnType,
chunkSize int, batchSize int,
batch models.Partition, partition models.Partition,
indexPrimaryKey int, indexPrimaryKey int,
chChunksOut chan<- models.Batch, chBatchesOut chan<- models.Batch,
rowsRead *int64, rowsRead *int64,
) error ) error
@@ -25,12 +25,12 @@ type Extractor interface {
ctx context.Context, ctx context.Context,
tableInfo config.SourceTableInfo, tableInfo config.SourceTableInfo,
columns []models.ColumnType, columns []models.ColumnType,
chunkSize int, batchSize int,
chBatchesIn <-chan models.Partition, chPartitionsIn <-chan models.Partition,
chChunksOut chan<- models.Batch, chBatchesOut chan<- models.Batch,
chErrorsOut chan<- custom_errors.ExtractorError, chErrorsOut chan<- custom_errors.ExtractorError,
chJobErrorsOut chan<- custom_errors.JobError, chJobErrorsOut chan<- custom_errors.JobError,
wgActiveBatches *sync.WaitGroup, wgActivePartitions *sync.WaitGroup,
rowsRead *int64, rowsRead *int64,
) )
} }
@@ -43,43 +43,43 @@ type ColumnTransformPlan struct {
} }
type Transformer interface { type Transformer interface {
ProcessChunk( ProcessBatch(
ctx context.Context, ctx context.Context,
chunk *models.Batch, batch *models.Batch,
transformationPlan []ColumnTransformPlan, transformationPlan []ColumnTransformPlan,
) error ) error
Exec( Exec(
ctx context.Context, ctx context.Context,
columns []models.ColumnType, columns []models.ColumnType,
chChunksIn <-chan models.Batch, chBatchesIn <-chan models.Batch,
chChunksOut chan<- models.Batch, chBactchesOut chan<- models.Batch,
chJobErrorsOut chan<- custom_errors.JobError, chJobErrorsOut chan<- custom_errors.JobError,
wgActiveChunks *sync.WaitGroup, wgActiveBatches *sync.WaitGroup,
) )
} }
type Loader interface { type Loader interface {
ProcessChunk( ProcessBatch(
ctx context.Context, ctx context.Context,
tableInfo config.TargetTableInfo, tableInfo config.TargetTableInfo,
colNames []string, colNames []string,
chunk models.Batch, batch models.Batch,
) (int, error) ) (int, error)
Exec( Exec(
ctx context.Context, ctx context.Context,
tableInfo config.TargetTableInfo, tableInfo config.TargetTableInfo,
columns []models.ColumnType, columns []models.ColumnType,
chChunksIn <-chan models.Batch, chBatchesIn <-chan models.Batch,
chErrorsOut chan<- custom_errors.LoaderError, chErrorsOut chan<- custom_errors.LoaderError,
chJobErrorsOut chan<- custom_errors.JobError, chJobErrorsOut chan<- custom_errors.JobError,
wgActiveChunks *sync.WaitGroup, wgActiveBatches *sync.WaitGroup,
rowsLoaded *int64, rowsLoaded *int64,
) )
} }
type TableAnalizer interface { type TableAnalyzer interface {
QueryColumnTypes( QueryColumnTypes(
ctx context.Context, ctx context.Context,
tableInfo config.TableInfo, tableInfo config.TableInfo,
@@ -93,6 +93,7 @@ type TableAnalizer interface {
CalculatePartitionRanges( CalculatePartitionRanges(
ctx context.Context, ctx context.Context,
tableInfo config.TableInfo, tableInfo config.TableInfo,
totalPartitions int, partitionColumn string,
) (models.Partition, error) maxPartitions int64,
) ([]models.Partition, error)
} }

2
internal/app/models/main.go
View File

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