73 lines
1.6 KiB
Go
73 lines
1.6 KiB
Go
package transformers
|
|
|
|
import (
|
|
"context"
|
|
"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/etl"
|
|
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
|
|
)
|
|
|
|
type PostgresTransformer struct {
|
|
sourceTable config.SourceTableInfo
|
|
}
|
|
|
|
func NewPostgresTransformer(sourceTable config.SourceTableInfo) etl.Transformer {
|
|
return &PostgresTransformer{sourceTable: sourceTable}
|
|
}
|
|
|
|
func (pgTr *PostgresTransformer) 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 := computePostgresTransformationPlan(columns)
|
|
|
|
acc := &batchAccumulator{batchSize: batchSize}
|
|
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
|
|
case batch, ok := <-chBatchesIn:
|
|
if !ok {
|
|
acc.flush(ctx, chBatchesOut, wgActiveBatches)
|
|
return
|
|
}
|
|
|
|
if len(transformationPlan) > 0 {
|
|
if err := ProcessBatchWithRetries(ctx, &batch, transformationPlan, retryConfig); err != nil {
|
|
sendTransformError(ctx, err, chJobErrorsOut)
|
|
return
|
|
}
|
|
}
|
|
|
|
if batchSize <= 0 {
|
|
wgActiveBatches.Add(1)
|
|
select {
|
|
case chBatchesOut <- batch:
|
|
case <-ctx.Done():
|
|
wgActiveBatches.Done()
|
|
return
|
|
}
|
|
continue
|
|
}
|
|
|
|
acc.add(batch)
|
|
if acc.ready() {
|
|
if !acc.flush(ctx, chBatchesOut, wgActiveBatches) {
|
|
return
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|