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:
@@ -2,34 +2,54 @@ package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"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"
|
||||
dbwrapper "git.ksdemosapps.com/kylesoda/go-migrate/internal/app/db-wrapper"
|
||||
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/etl"
|
||||
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/etl/extractors"
|
||||
"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/models"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"golang.org/x/sync/errgroup"
|
||||
)
|
||||
|
||||
func buildTruncateQuery(targetDbType, schema, table, truncateMethod string) string {
|
||||
if truncateMethod == "DELETE" {
|
||||
if targetDbType == "postgres" {
|
||||
return fmt.Sprintf(`DELETE FROM "%s"."%s"`, schema, table)
|
||||
}
|
||||
return fmt.Sprintf(`DELETE FROM [%s].[%s]`, schema, table)
|
||||
}
|
||||
|
||||
if targetDbType == "postgres" {
|
||||
return fmt.Sprintf(`TRUNCATE TABLE "%s"."%s"`, schema, table)
|
||||
}
|
||||
return fmt.Sprintf(`TRUNCATE TABLE [%s].[%s]`, schema, table)
|
||||
}
|
||||
|
||||
func processMigrationJob(
|
||||
ctx context.Context,
|
||||
targetDbWrapper dbwrapper.DbWrapper,
|
||||
sourceTableAnalyzer etl.TableAnalyzer,
|
||||
targetTableAnalyzer etl.TableAnalyzer,
|
||||
extractor etl.Extractor,
|
||||
transformer etl.Transformer,
|
||||
azureClient *azure.Client,
|
||||
loader etl.Loader,
|
||||
job config.Job,
|
||||
) JobResult {
|
||||
targetDbType string,
|
||||
) models.JobResult {
|
||||
transformer := transformers.NewMssqlTransformer(job.ToStorage, job.SourceTable, azureClient)
|
||||
localCtx, cancel := context.WithCancel(ctx)
|
||||
defer cancel()
|
||||
|
||||
result := JobResult{
|
||||
result := models.JobResult{
|
||||
JobName: job.Name,
|
||||
StartTime: time.Now(),
|
||||
}
|
||||
@@ -65,7 +85,13 @@ func processMigrationJob(
|
||||
return result
|
||||
}
|
||||
|
||||
for _, query := range job.PreSQL {
|
||||
preSqlQueries := job.TargetTable.PreSQL
|
||||
if job.TruncateTarget {
|
||||
truncateQuery := buildTruncateQuery(targetDbType, job.TargetTable.Schema, job.TargetTable.Table, job.TruncateMethod)
|
||||
preSqlQueries = append([]string{truncateQuery}, job.TargetTable.PreSQL...)
|
||||
}
|
||||
|
||||
for _, query := range preSqlQueries {
|
||||
if _, err := targetDbWrapper.Exec(localCtx, query); err != nil {
|
||||
result.Error = err
|
||||
return result
|
||||
@@ -78,6 +104,7 @@ func processMigrationJob(
|
||||
job.SourceTable.TableInfo,
|
||||
job.SourceTable.PrimaryKey,
|
||||
job.RowsPerPartition,
|
||||
job.Range,
|
||||
)
|
||||
if err != nil {
|
||||
log.Error("Unexpected error calculating batch ranges: ", err)
|
||||
@@ -128,8 +155,9 @@ func processMigrationJob(
|
||||
|
||||
for range maxExtractors {
|
||||
wgExtractors.Go(func() {
|
||||
extractor.Exec(
|
||||
extractors.Consume(
|
||||
localCtx,
|
||||
extractor,
|
||||
job.SourceTable,
|
||||
sourceColTypes,
|
||||
job.BatchSize,
|
||||
@@ -213,7 +241,7 @@ func processMigrationJob(
|
||||
cancel()
|
||||
}()
|
||||
|
||||
for _, query := range job.PostSQL {
|
||||
for _, query := range job.TargetTable.PostSQL {
|
||||
if _, err := targetDbWrapper.Exec(localCtx, query); err != nil {
|
||||
result.Error = err
|
||||
return result
|
||||
@@ -233,5 +261,9 @@ func processMigrationJob(
|
||||
result.RowsLoaded = atomic.LoadInt64(&rowsLoaded)
|
||||
result.RowsFailed = atomic.LoadInt64(&rowsFailed)
|
||||
|
||||
if result.RowsRead != result.RowsLoaded {
|
||||
result.Error = fmt.Errorf("Row count mismatch: extracted %d rows but loaded %d rows (failed: %d)", result.RowsRead, result.RowsLoaded, result.RowsFailed)
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user