feat: integrate azure storage handling in migration pipeline
Stream transformer output to an Azure blob container when a job declares ToStorage, with per-job prefix and credential configuration.
This commit is contained in:
101
internal/app/etl/extractors/main.go
Normal file
101
internal/app/etl/extractors/main.go
Normal file
@@ -0,0 +1,101 @@
|
||||
package extractors
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"slices"
|
||||
"strings"
|
||||
"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/etl"
|
||||
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
|
||||
)
|
||||
|
||||
func Consume(
|
||||
ctx context.Context,
|
||||
extractor etl.Extractor,
|
||||
tableInfo config.SourceTableInfo,
|
||||
columns []models.ColumnType,
|
||||
batchSize int,
|
||||
chPartitionsIn <-chan models.Partition,
|
||||
chBatchesOut chan<- models.Batch,
|
||||
chErrorsOut chan<- custom_errors.ExtractorError,
|
||||
chJobErrorsOut chan<- custom_errors.JobError,
|
||||
wgActivePartitions *sync.WaitGroup,
|
||||
rowsRead *int64,
|
||||
) {
|
||||
indexPrimaryKey := slices.IndexFunc(columns, func(col models.ColumnType) bool {
|
||||
return strings.EqualFold(col.Name(), tableInfo.PrimaryKey)
|
||||
})
|
||||
|
||||
if indexPrimaryKey == -1 {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case chJobErrorsOut <- custom_errors.JobError{
|
||||
ShouldCancelJob: true,
|
||||
Msg: "Primary key not found in provided columns",
|
||||
}:
|
||||
}
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
for {
|
||||
if ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case partition, ok := <-chPartitionsIn:
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
rowsReadResult, err := extractor.Exec(
|
||||
ctx,
|
||||
tableInfo,
|
||||
columns,
|
||||
batchSize,
|
||||
partition,
|
||||
indexPrimaryKey,
|
||||
chBatchesOut,
|
||||
)
|
||||
|
||||
if rowsReadResult > 0 {
|
||||
atomic.AddInt64(rowsRead, int64(rowsReadResult))
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
if exError, ok := errors.AsType[*custom_errors.ExtractorError](err); ok {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case chErrorsOut <- *exError:
|
||||
}
|
||||
} else if jobError, ok := errors.AsType[*custom_errors.JobError](err); ok {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case chJobErrorsOut <- *jobError:
|
||||
}
|
||||
} else {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case chErrorsOut <- custom_errors.ExtractorError{Partition: partition, Msg: err.Error()}:
|
||||
}
|
||||
}
|
||||
|
||||
continue
|
||||
}
|
||||
|
||||
wgActivePartitions.Done()
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -5,10 +5,7 @@ import (
|
||||
"database/sql"
|
||||
"errors"
|
||||
"fmt"
|
||||
"slices"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
|
||||
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
|
||||
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/convert"
|
||||
@@ -32,6 +29,9 @@ func buildExtractQueryMssql(
|
||||
columns []models.ColumnType,
|
||||
includeRange bool,
|
||||
isMinInclusive bool,
|
||||
isMaxInclusive bool,
|
||||
hasMin bool,
|
||||
hasMax bool,
|
||||
) string {
|
||||
var sbQuery strings.Builder
|
||||
|
||||
@@ -55,15 +55,32 @@ func buildExtractQueryMssql(
|
||||
|
||||
fmt.Fprintf(&sbQuery, " FROM [%s].[%s]", tableInfo.Schema, tableInfo.Table)
|
||||
|
||||
if includeRange {
|
||||
fmt.Fprintf(&sbQuery, " WHERE [%s]", tableInfo.PrimaryKey)
|
||||
if isMinInclusive {
|
||||
sbQuery.WriteString(" >=")
|
||||
} else {
|
||||
sbQuery.WriteString(" >")
|
||||
if includeRange && (hasMin || hasMax) {
|
||||
sbQuery.WriteString(" WHERE ")
|
||||
|
||||
if hasMin {
|
||||
fmt.Fprintf(&sbQuery, "[%s]", tableInfo.PrimaryKey)
|
||||
if isMinInclusive {
|
||||
sbQuery.WriteString(" >=")
|
||||
} else {
|
||||
sbQuery.WriteString(" >")
|
||||
}
|
||||
sbQuery.WriteString(" @min")
|
||||
}
|
||||
|
||||
fmt.Fprintf(&sbQuery, " @min AND [%s] <= @max", tableInfo.PrimaryKey)
|
||||
if hasMin && hasMax {
|
||||
sbQuery.WriteString(" AND ")
|
||||
}
|
||||
|
||||
if hasMax {
|
||||
fmt.Fprintf(&sbQuery, "[%s]", tableInfo.PrimaryKey)
|
||||
if isMaxInclusive {
|
||||
sbQuery.WriteString(" <=")
|
||||
} else {
|
||||
sbQuery.WriteString(" <")
|
||||
}
|
||||
sbQuery.WriteString(" @max")
|
||||
}
|
||||
}
|
||||
|
||||
fmt.Fprintf(&sbQuery, " ORDER BY [%s] ASC", tableInfo.PrimaryKey)
|
||||
@@ -99,7 +116,7 @@ func errorFromLastRow(
|
||||
}
|
||||
}
|
||||
|
||||
func (mssqlEx *MssqlExtractor) ProcessPartition(
|
||||
func (mssqlEx *MssqlExtractor) Exec(
|
||||
ctx context.Context,
|
||||
tableInfo config.SourceTableInfo,
|
||||
columns []models.ColumnType,
|
||||
@@ -108,14 +125,16 @@ func (mssqlEx *MssqlExtractor) ProcessPartition(
|
||||
indexPrimaryKey int,
|
||||
chBatchesOut chan<- models.Batch,
|
||||
) (int, error) {
|
||||
query := buildExtractQueryMssql(tableInfo, columns, partition.HasRange, partition.Range.IsMinInclusive)
|
||||
hasMin := partition.HasRange && partition.Range.Min > 0
|
||||
hasMax := partition.HasRange && partition.Range.Max > 0
|
||||
query := buildExtractQueryMssql(tableInfo, columns, partition.HasRange, partition.Range.IsMinInclusive, partition.Range.IsMaxInclusive, hasMin, hasMax)
|
||||
|
||||
var queryArgs []any
|
||||
if partition.HasRange {
|
||||
queryArgs = append(queryArgs,
|
||||
sql.Named("min", partition.Range.Min),
|
||||
sql.Named("max", partition.Range.Max),
|
||||
)
|
||||
if hasMin {
|
||||
queryArgs = append(queryArgs, sql.Named("min", partition.Range.Min))
|
||||
}
|
||||
if hasMax {
|
||||
queryArgs = append(queryArgs, sql.Named("max", partition.Range.Max))
|
||||
}
|
||||
|
||||
rowsRead := 0
|
||||
@@ -188,90 +207,3 @@ func (mssqlEx *MssqlExtractor) ProcessPartition(
|
||||
|
||||
return rowsRead, nil
|
||||
}
|
||||
|
||||
func (mssqlEx *MssqlExtractor) Exec(
|
||||
ctx context.Context,
|
||||
tableInfo config.SourceTableInfo,
|
||||
columns []models.ColumnType,
|
||||
batchSize int,
|
||||
chPartitionsIn <-chan models.Partition,
|
||||
chBatchesOut chan<- models.Batch,
|
||||
chErrorsOut chan<- custom_errors.ExtractorError,
|
||||
chJobErrorsOut chan<- custom_errors.JobError,
|
||||
wgActivePartitions *sync.WaitGroup,
|
||||
rowsRead *int64,
|
||||
) {
|
||||
indexPrimaryKey := slices.IndexFunc(columns, func(col models.ColumnType) bool {
|
||||
return strings.EqualFold(col.Name(), tableInfo.PrimaryKey)
|
||||
})
|
||||
|
||||
if indexPrimaryKey == -1 {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case chJobErrorsOut <- custom_errors.JobError{
|
||||
ShouldCancelJob: true,
|
||||
Msg: "Primary key not found in provided columns",
|
||||
}:
|
||||
}
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
for {
|
||||
if ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case partition, ok := <-chPartitionsIn:
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
rowsReadResult, err := mssqlEx.ProcessPartition(
|
||||
ctx,
|
||||
tableInfo,
|
||||
columns,
|
||||
batchSize,
|
||||
partition,
|
||||
indexPrimaryKey,
|
||||
chBatchesOut,
|
||||
)
|
||||
|
||||
if rowsReadResult > 0 {
|
||||
atomic.AddInt64(rowsRead, int64(rowsReadResult))
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
var exError *custom_errors.ExtractorError
|
||||
var jobError *custom_errors.JobError
|
||||
if errors.As(err, &exError) {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case chErrorsOut <- *exError:
|
||||
}
|
||||
} else if errors.As(err, &jobError) {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case chJobErrorsOut <- *jobError:
|
||||
}
|
||||
} else {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case chErrorsOut <- custom_errors.ExtractorError{Partition: partition, Msg: err.Error()}:
|
||||
}
|
||||
}
|
||||
|
||||
continue
|
||||
}
|
||||
|
||||
wgActivePartitions.Done()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,10 +2,8 @@ package extractors
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
|
||||
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/custom_errors"
|
||||
@@ -23,7 +21,15 @@ func NewPostgresExtractor(db dbwrapper.DbWrapper) etl.Extractor {
|
||||
return &PostgresExtractor{db: db}
|
||||
}
|
||||
|
||||
func buildExtractQueryPostgres(sourceDbInfo config.SourceTableInfo, columns []models.ColumnType) string {
|
||||
func buildExtractQueryPostgres(
|
||||
sourceDbInfo config.SourceTableInfo,
|
||||
columns []models.ColumnType,
|
||||
includeRange bool,
|
||||
isMinInclusive bool,
|
||||
isMaxInclusive bool,
|
||||
hasMin bool,
|
||||
hasMax bool,
|
||||
) string {
|
||||
var sbColumns strings.Builder
|
||||
|
||||
if len(columns) == 0 {
|
||||
@@ -48,10 +54,44 @@ 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)
|
||||
query := fmt.Sprintf(`SELECT %s FROM "%s"."%s"`, sbColumns.String(), sourceDbInfo.Schema, sourceDbInfo.Table)
|
||||
|
||||
if includeRange && (hasMin || hasMax) {
|
||||
query += " WHERE "
|
||||
paramIdx := 1
|
||||
|
||||
if hasMin {
|
||||
query += fmt.Sprintf(`"%s"`, sourceDbInfo.PrimaryKey)
|
||||
if isMinInclusive {
|
||||
query += " >="
|
||||
} else {
|
||||
query += " >"
|
||||
}
|
||||
query += fmt.Sprintf(" $%d", paramIdx)
|
||||
paramIdx++
|
||||
}
|
||||
|
||||
if hasMin && hasMax {
|
||||
query += " AND "
|
||||
}
|
||||
|
||||
if hasMax {
|
||||
query += fmt.Sprintf(`"%s"`, sourceDbInfo.PrimaryKey)
|
||||
if isMaxInclusive {
|
||||
query += " <="
|
||||
} else {
|
||||
query += " <"
|
||||
}
|
||||
query += fmt.Sprintf(" $%d", paramIdx)
|
||||
}
|
||||
}
|
||||
|
||||
query += fmt.Sprintf(` ORDER BY "%s" ASC`, sourceDbInfo.PrimaryKey)
|
||||
|
||||
return query
|
||||
}
|
||||
|
||||
func (postgresEx *PostgresExtractor) ProcessPartition(
|
||||
func (postgresEx *PostgresExtractor) Exec(
|
||||
ctx context.Context,
|
||||
tableInfo config.SourceTableInfo,
|
||||
columns []models.ColumnType,
|
||||
@@ -60,14 +100,20 @@ func (postgresEx *PostgresExtractor) ProcessPartition(
|
||||
indexPrimaryKey int,
|
||||
chBatchesOut chan<- models.Batch,
|
||||
) (int, error) {
|
||||
query := buildExtractQueryPostgres(tableInfo, columns)
|
||||
hasMin := partition.HasRange && partition.Range.Min > 0
|
||||
hasMax := partition.HasRange && partition.Range.Max > 0
|
||||
query := buildExtractQueryPostgres(tableInfo, columns, partition.HasRange, partition.Range.IsMinInclusive, partition.Range.IsMaxInclusive, hasMin, hasMax)
|
||||
|
||||
if partition.HasRange {
|
||||
return 0, errors.New("Batch config not yet supported")
|
||||
var queryArgs []any
|
||||
if hasMin {
|
||||
queryArgs = append(queryArgs, partition.Range.Min)
|
||||
}
|
||||
if hasMax {
|
||||
queryArgs = append(queryArgs, partition.Range.Max)
|
||||
}
|
||||
|
||||
rowsRead := 0
|
||||
rows, err := postgresEx.db.Query(ctx, query)
|
||||
rows, err := postgresEx.db.Query(ctx, query, queryArgs...)
|
||||
if err != nil {
|
||||
return rowsRead, &custom_errors.ExtractorError{Partition: partition, HasLastId: false, Msg: err.Error()}
|
||||
}
|
||||
@@ -78,7 +124,7 @@ func (postgresEx *PostgresExtractor) ProcessPartition(
|
||||
for rows.Next() {
|
||||
values, err := rows.Values()
|
||||
if err != nil {
|
||||
return rowsRead, errors.New("Unexpected error reading rows from source")
|
||||
return rowsRead, &custom_errors.ExtractorError{Partition: partition, HasLastId: false, Msg: err.Error()}
|
||||
}
|
||||
rowsRead++
|
||||
|
||||
@@ -96,7 +142,7 @@ func (postgresEx *PostgresExtractor) ProcessPartition(
|
||||
}
|
||||
|
||||
if err := rows.Err(); err != nil {
|
||||
return rowsRead, errors.New("Unexpected error reading rows from source")
|
||||
return rowsRead, &custom_errors.ExtractorError{Partition: partition, HasLastId: false, Msg: err.Error()}
|
||||
}
|
||||
|
||||
if len(batchRows) > 0 {
|
||||
@@ -109,17 +155,3 @@ func (postgresEx *PostgresExtractor) ProcessPartition(
|
||||
|
||||
return rowsRead, nil
|
||||
}
|
||||
|
||||
func (postgresEx *PostgresExtractor) Exec(
|
||||
ctx context.Context,
|
||||
tableInfo config.SourceTableInfo,
|
||||
columns []models.ColumnType,
|
||||
batchSize int,
|
||||
chPartitionsIn <-chan models.Partition,
|
||||
chBatchesOut chan<- models.Batch,
|
||||
chErrorsOut chan<- custom_errors.ExtractorError,
|
||||
chJobErrorsOut chan<- custom_errors.JobError,
|
||||
wgActivePartitions *sync.WaitGroup,
|
||||
rowsRead *int64,
|
||||
) {
|
||||
}
|
||||
|
||||
@@ -15,31 +15,21 @@ import (
|
||||
"github.com/jackc/pgx/v5/pgconn"
|
||||
)
|
||||
|
||||
type PostgresLoader struct {
|
||||
type GenericLoader struct {
|
||||
db dbwrapper.DbWrapper
|
||||
}
|
||||
|
||||
func NewPostgresLoader(db dbwrapper.DbWrapper) etl.Loader {
|
||||
return &PostgresLoader{db: db}
|
||||
func NewGenericLoader(db dbwrapper.DbWrapper) etl.Loader {
|
||||
return &GenericLoader{db: db}
|
||||
}
|
||||
|
||||
func mapSlice[T any, V any](input []T, mapper func(T) V) []V {
|
||||
result := make([]V, len(input))
|
||||
|
||||
for i, v := range input {
|
||||
result[i] = mapper(v)
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
func (postgresLd *PostgresLoader) ProcessBatch(
|
||||
func (gl *GenericLoader) ProcessBatch(
|
||||
ctx context.Context,
|
||||
tableInfo config.TargetTableInfo,
|
||||
colNames []string,
|
||||
batch models.Batch,
|
||||
) (int, error) {
|
||||
_, err := postgresLd.db.SaveMassive(
|
||||
_, err := gl.db.SaveMassive(
|
||||
ctx,
|
||||
tableInfo.Schema,
|
||||
tableInfo.Table,
|
||||
@@ -65,7 +55,7 @@ func (postgresLd *PostgresLoader) ProcessBatch(
|
||||
return len(batch.Rows), nil
|
||||
}
|
||||
|
||||
func (postgresLd *PostgresLoader) Exec(
|
||||
func (gl *GenericLoader) Exec(
|
||||
ctx context.Context,
|
||||
tableInfo config.TargetTableInfo,
|
||||
columns []models.ColumnType,
|
||||
@@ -92,7 +82,7 @@ func (postgresLd *PostgresLoader) Exec(
|
||||
return
|
||||
}
|
||||
|
||||
processedRows, err := postgresLd.ProcessBatch(ctx, tableInfo, colNames, batch)
|
||||
processedRows, err := gl.ProcessBatch(ctx, tableInfo, colNames, batch)
|
||||
|
||||
if err != nil {
|
||||
var ldError *custom_errors.LoaderError
|
||||
@@ -1 +0,0 @@
|
||||
package loaders
|
||||
11
internal/app/etl/loaders/utils.go
Normal file
11
internal/app/etl/loaders/utils.go
Normal file
@@ -0,0 +1,11 @@
|
||||
package loaders
|
||||
|
||||
func mapSlice[T any, V any](input []T, mapper func(T) V) []V {
|
||||
result := make([]V, len(input))
|
||||
|
||||
for i, v := range input {
|
||||
result[i] = mapper(v)
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
@@ -15,7 +15,22 @@ func PartitionRangeGenerator(
|
||||
tableInfo config.TableInfo,
|
||||
partitionColumn string,
|
||||
rowsPerPartition int64,
|
||||
jobRange config.RangeConfig,
|
||||
) ([]models.Partition, error) {
|
||||
if jobRange.Min > 0 {
|
||||
return []models.Partition{{
|
||||
Id: uuid.New(),
|
||||
HasRange: true,
|
||||
RetryCounter: 0,
|
||||
Range: models.PartitionRange{
|
||||
Min: jobRange.Min,
|
||||
Max: jobRange.Max,
|
||||
IsMinInclusive: jobRange.IsMinInclusive,
|
||||
IsMaxInclusive: jobRange.IsMaxInclusive,
|
||||
},
|
||||
}}, nil
|
||||
}
|
||||
|
||||
rowsCount, err := tableAnalyzer.EstimateTotalRows(ctx, tableInfo)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
|
||||
@@ -36,7 +36,7 @@ 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%'
|
||||
WHERE s.name = @schema AND st.name = @table AND (c.is_hidden = 0 OR (c.graph_type IS NOT NULL AND c.name LIKE '$%'))
|
||||
ORDER BY c.column_id;`
|
||||
|
||||
type rawColumnMssql struct {
|
||||
@@ -184,7 +184,7 @@ 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)
|
||||
ctxTimeout, cancel := context.WithTimeout(ctx, 1*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
var rowsCount int64
|
||||
@@ -216,7 +216,7 @@ ORDER BY batch_id`,
|
||||
tableInfo.Schema,
|
||||
tableInfo.Table)
|
||||
|
||||
ctxTimeout, cancel := context.WithTimeout(ctx, time.Second*20)
|
||||
ctxTimeout, cancel := context.WithTimeout(ctx, 1*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
rows, err := ta.db.Query(ctxTimeout, query, sql.Named("maxPartitions", maxPartitions))
|
||||
|
||||
@@ -125,7 +125,7 @@ func (ta *PostgresTableAnalyzer) QueryColumnTypes(
|
||||
ctx context.Context,
|
||||
tableInfo config.TableInfo,
|
||||
) ([]models.ColumnType, error) {
|
||||
localCtx, cancel := context.WithTimeout(ctx, 20*time.Second)
|
||||
localCtx, cancel := context.WithTimeout(ctx, 1*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
rows, err := ta.db.Query(localCtx, postgresColumnMetadataQuery, tableInfo.Schema, tableInfo.Table)
|
||||
|
||||
@@ -3,18 +3,32 @@ 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{}
|
||||
type MssqlTransformer struct {
|
||||
toStorage config.ToStorageConfig
|
||||
sourceTable config.SourceTableInfo
|
||||
azureClient *azure.Client
|
||||
}
|
||||
|
||||
func NewMssqlTransformer() etl.Transformer {
|
||||
return &MssqlTransformer{}
|
||||
func NewMssqlTransformer(toStorage config.ToStorageConfig, sourceTable config.SourceTableInfo, azureClient *azure.Client) etl.Transformer {
|
||||
return &MssqlTransformer{
|
||||
toStorage: toStorage,
|
||||
sourceTable: sourceTable,
|
||||
azureClient: azureClient,
|
||||
}
|
||||
}
|
||||
|
||||
func computeTransformationPlan(columns []models.ColumnType) []etl.ColumnTransformPlan {
|
||||
@@ -60,6 +74,63 @@ func computeTransformationPlan(columns []models.ColumnType) []etl.ColumnTransfor
|
||||
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
|
||||
}
|
||||
blobPath := fmt.Sprintf("%s/%s/%s/%s.bin", schema, table, sourceColName, 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)
|
||||
}
|
||||
return blobURL, nil
|
||||
},
|
||||
})
|
||||
}
|
||||
return plan
|
||||
}
|
||||
|
||||
const processBatchCtxCheck = 4096
|
||||
|
||||
func (mssqlTr *MssqlTransformer) ProcessBatch(
|
||||
@@ -100,6 +171,8 @@ func (mssqlTr *MssqlTransformer) Exec(
|
||||
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 {
|
||||
|
||||
@@ -10,7 +10,7 @@ import (
|
||||
)
|
||||
|
||||
type Extractor interface {
|
||||
ProcessPartition(
|
||||
Exec(
|
||||
ctx context.Context,
|
||||
tableInfo config.SourceTableInfo,
|
||||
columns []models.ColumnType,
|
||||
@@ -19,19 +19,6 @@ type Extractor interface {
|
||||
indexPrimaryKey int,
|
||||
chBatchesOut chan<- models.Batch,
|
||||
) (int, error)
|
||||
|
||||
Exec(
|
||||
ctx context.Context,
|
||||
tableInfo config.SourceTableInfo,
|
||||
columns []models.ColumnType,
|
||||
batchSize int,
|
||||
chPartitionsIn <-chan models.Partition,
|
||||
chBatchesOut chan<- models.Batch,
|
||||
chErrorsOut chan<- custom_errors.ExtractorError,
|
||||
chJobErrorsOut chan<- custom_errors.JobError,
|
||||
wgActivePartitions *sync.WaitGroup,
|
||||
rowsRead *int64,
|
||||
)
|
||||
}
|
||||
|
||||
type TransformerFunc func(any) (any, error)
|
||||
|
||||
Reference in New Issue
Block a user