83 lines
1.7 KiB
Go
83 lines
1.7 KiB
Go
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"
|
|
)
|
|
|
|
func (mssqlTr *MssqlTransformer) Consume(
|
|
ctx context.Context,
|
|
columns []models.ColumnType,
|
|
retryConfig config.RetryConfig,
|
|
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 := 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
|
|
}
|
|
|
|
select {
|
|
case chBatchesOut <- batch:
|
|
case <-ctx.Done():
|
|
return
|
|
}
|
|
|
|
wgActiveBatches.Add(1)
|
|
}
|
|
}
|
|
}
|