118 lines
2.5 KiB
Go
118 lines
2.5 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"
|
|
"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
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|