feat: add extract query builders for both dialects
Add QueryFromObject helpers to build select queries from a table schema, with limits and geometry support.
This commit is contained in:
@@ -17,7 +17,7 @@ func buildExtractQueryMssql(job MigrationJob, columns []ColumnType) string {
|
||||
sbColumns.WriteString("]")
|
||||
|
||||
if col.unifiedType == "GEOMETRY" {
|
||||
sbColumns.WriteString(".STAsBinary() AS [")
|
||||
sbColumns.WriteString(".STAsWKB() AS [")
|
||||
sbColumns.WriteString(col.name)
|
||||
sbColumns.WriteString("]")
|
||||
}
|
||||
@@ -28,7 +28,7 @@ func buildExtractQueryMssql(job MigrationJob, columns []ColumnType) string {
|
||||
}
|
||||
}
|
||||
|
||||
return fmt.Sprintf(`SELECT %s FROM [%s].[%s] ORDER BY [%s] ASC`, sbColumns.String(), job.Schema, job.Table, job.PrimaryKey)
|
||||
return fmt.Sprintf(`SELECT %s FROM [%s].[%s] WITH (NOLOCK)`, sbColumns.String(), job.Schema, job.Table)
|
||||
}
|
||||
|
||||
func buildExtractQueryPostgres(job MigrationJob, columns []ColumnType) string {
|
||||
@@ -56,5 +56,5 @@ func buildExtractQueryPostgres(job MigrationJob, columns []ColumnType) string {
|
||||
}
|
||||
}
|
||||
|
||||
return fmt.Sprintf(`SELECT %s FROM "%s"."%s" ORDER BY "%s" ASC`, sbColumns.String(), job.Schema, job.Table, job.PrimaryKey)
|
||||
return fmt.Sprintf(`SELECT %s FROM "%s"."%s"`, sbColumns.String(), job.Schema, job.Table)
|
||||
}
|
||||
|
||||
@@ -1,88 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
_ "github.com/microsoft/go-mssqldb"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
type UnknownRowValues []any
|
||||
|
||||
func extractFromMssql(ctx context.Context, job MigrationJob, columns []ColumnType, chunkSize int, db *sql.DB, out chan<- []UnknownRowValues) error {
|
||||
query := buildExtractQueryMssql(job, columns)
|
||||
log.Debug("Query used to extract data from mssql: ", query)
|
||||
|
||||
rows, err := db.QueryContext(ctx, query)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
rowsChunk := make([]UnknownRowValues, 0, chunkSize)
|
||||
|
||||
for rows.Next() {
|
||||
values := make([]any, len(columns))
|
||||
scanArgs := make([]any, len(columns))
|
||||
|
||||
for i := range values {
|
||||
scanArgs[i] = &values[i]
|
||||
}
|
||||
|
||||
if err := rows.Scan(scanArgs...); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
rowsChunk = append(rowsChunk, values)
|
||||
|
||||
if len(rowsChunk) >= chunkSize {
|
||||
out <- rowsChunk
|
||||
rowsChunk = make([]UnknownRowValues, 0, chunkSize)
|
||||
log.Infof("Chunk send... %+v", job)
|
||||
}
|
||||
}
|
||||
|
||||
if len(rowsChunk) > 0 {
|
||||
out <- rowsChunk
|
||||
log.Infof("Chunk send... %+v", job)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func extractFromPostgres(ctx context.Context, job MigrationJob, columns []ColumnType, chunkSize int, db *pgxpool.Pool, out chan<- []UnknownRowValues) error {
|
||||
query := buildExtractQueryPostgres(job, columns)
|
||||
log.Debug("Query used to extract data from postgres: ", query)
|
||||
|
||||
rows, err := db.Query(ctx, query)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
rowsChunk := make([]UnknownRowValues, 0, chunkSize)
|
||||
|
||||
for rows.Next() {
|
||||
values, err := rows.Values()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
rowsChunk = append(rowsChunk, values)
|
||||
|
||||
if len(rowsChunk) >= chunkSize {
|
||||
out <- rowsChunk
|
||||
rowsChunk = make([]UnknownRowValues, 0, chunkSize)
|
||||
log.Infof("Chunk send... %+v", job)
|
||||
}
|
||||
}
|
||||
|
||||
if len(rowsChunk) > 0 {
|
||||
out <- rowsChunk
|
||||
log.Infof("Chunk send... %+v", job)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -7,23 +7,15 @@ import (
|
||||
type MigrationJob struct {
|
||||
Schema string
|
||||
Table string
|
||||
PrimaryKey string
|
||||
}
|
||||
|
||||
var migrationJobs []MigrationJob = []MigrationJob{
|
||||
{
|
||||
Schema: "demo",
|
||||
Table: "users",
|
||||
PrimaryKey: "id",
|
||||
},
|
||||
}
|
||||
|
||||
const (
|
||||
NumExtractors int = 2
|
||||
ChunkSize int = 20
|
||||
QueueSize int = 10
|
||||
)
|
||||
|
||||
func main() {
|
||||
configureLog()
|
||||
log.Info("Starting migration...")
|
||||
|
||||
@@ -1,52 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
)
|
||||
|
||||
func mssqlUuidToBigEndian(mssqlUuid []byte) []byte {
|
||||
if len(mssqlUuid) != 16 {
|
||||
return mssqlUuid
|
||||
}
|
||||
pgUuid := make([]byte, 16)
|
||||
pgUuid[0], pgUuid[1], pgUuid[2], pgUuid[3] = mssqlUuid[3], mssqlUuid[2], mssqlUuid[1], mssqlUuid[0]
|
||||
pgUuid[4], pgUuid[5] = mssqlUuid[5], mssqlUuid[4]
|
||||
pgUuid[6], pgUuid[7] = mssqlUuid[7], mssqlUuid[6]
|
||||
copy(pgUuid[8:], mssqlUuid[8:])
|
||||
|
||||
return pgUuid
|
||||
}
|
||||
|
||||
const sridFlag = 0x20000000
|
||||
|
||||
func wkbToEwkbWithSrid(geometry []byte, srid int) []byte {
|
||||
if len(geometry) < 5 {
|
||||
return geometry
|
||||
}
|
||||
|
||||
var byteOrder binary.ByteOrder
|
||||
if geometry[0] == 0 {
|
||||
byteOrder = binary.BigEndian
|
||||
} else {
|
||||
byteOrder = binary.LittleEndian
|
||||
}
|
||||
|
||||
wkbType := byteOrder.Uint32(geometry[1:5])
|
||||
if wkbType&sridFlag != 0 {
|
||||
return geometry
|
||||
}
|
||||
|
||||
ewkbType := wkbType | sridFlag
|
||||
|
||||
result := make([]byte, len(geometry)+4)
|
||||
|
||||
result[0] = geometry[0]
|
||||
|
||||
byteOrder.PutUint32(result[1:5], ewkbType)
|
||||
|
||||
byteOrder.PutUint32(result[5:9], uint32(srid))
|
||||
|
||||
copy(result[9:], geometry[5:])
|
||||
|
||||
return result
|
||||
}
|
||||
@@ -1,11 +1,7 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"sync"
|
||||
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
|
||||
_ "github.com/microsoft/go-mssqldb"
|
||||
@@ -21,43 +17,12 @@ func processMigrationJob(sourceDb *sql.DB, targetDb *pgxpool.Pool, job Migration
|
||||
logColumnTypes(sourceColTypes, "Source col types")
|
||||
logColumnTypes(targetColTypes, "Target col types")
|
||||
|
||||
chRowsExtract := make(chan []UnknownRowValues, QueueSize)
|
||||
chRowsTransform := make(chan []UnknownRowValues)
|
||||
mssqlContext := context.Background()
|
||||
sourceQuery := buildExtractQueryMssql(job, sourceColTypes)
|
||||
|
||||
go func() {
|
||||
if err := extractFromMssql(mssqlContext, job, sourceColTypes, ChunkSize, sourceDb, chRowsExtract); err != nil {
|
||||
log.Error("Unexpected error extrating data from mssql: ", err)
|
||||
}
|
||||
close(chRowsExtract)
|
||||
}()
|
||||
log.Debug(sourceQuery)
|
||||
|
||||
go func() {
|
||||
transformRowsMssql(sourceColTypes, chRowsExtract, chRowsTransform)
|
||||
close(chRowsTransform)
|
||||
}()
|
||||
|
||||
var wgFakeLoaders sync.WaitGroup
|
||||
|
||||
wgFakeLoaders.Go(func() {
|
||||
fakeLoader(job, sourceColTypes, chRowsTransform)
|
||||
})
|
||||
|
||||
chRowsExtractPostgres := make(chan []UnknownRowValues, QueueSize)
|
||||
postgresContext := context.Background()
|
||||
|
||||
go func() {
|
||||
if err := extractFromPostgres(postgresContext, job, sourceColTypes, ChunkSize, targetDb, chRowsExtractPostgres); err != nil {
|
||||
log.Error("Unexpected error extrating data from postgres: ", err)
|
||||
}
|
||||
close(chRowsExtractPostgres)
|
||||
}()
|
||||
|
||||
wgFakeLoaders.Go(func() {
|
||||
fakeLoader(job, targetColTypes, chRowsExtractPostgres)
|
||||
})
|
||||
|
||||
wgFakeLoaders.Wait()
|
||||
targetQuery := buildExtractQueryPostgres(job, targetColTypes)
|
||||
log.Debug(targetQuery)
|
||||
}
|
||||
|
||||
func logColumnTypes(columnTypes []ColumnType, label string) {
|
||||
@@ -67,45 +32,3 @@ func logColumnTypes(columnTypes []ColumnType, label string) {
|
||||
log.Infof("%+v", col)
|
||||
}
|
||||
}
|
||||
|
||||
func transformRowsMssql(columns []ColumnType, in <-chan []UnknownRowValues, out chan<- []UnknownRowValues) {
|
||||
for rows := range in {
|
||||
log.Debugf("Chunk received, transforming...")
|
||||
|
||||
for _, rowValues := range rows {
|
||||
for i, col := range columns {
|
||||
value := rowValues[i]
|
||||
if col.SystemType() == "uniqueidentifier" {
|
||||
if b, ok := value.([]byte); ok {
|
||||
rowValues[i] = mssqlUuidToBigEndian(b)
|
||||
}
|
||||
} else if col.SystemType() == "geometry" || col.SystemType() == "geography" {
|
||||
if b, ok := value.([]byte); ok {
|
||||
rowValues[i] = wkbToEwkbWithSrid(b, 4326)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
out <- rows
|
||||
}
|
||||
}
|
||||
|
||||
func logSampleRow(job MigrationJob, columns []ColumnType, rowValues UnknownRowValues, tag string) {
|
||||
log.Infof("[%s.%s] Sample row: (%s)", job.Schema, job.Table, tag)
|
||||
for i, col := range columns {
|
||||
log.Infof("%s (%T): %v", col.Name(), rowValues[i], rowValues[i])
|
||||
}
|
||||
}
|
||||
|
||||
func fakeLoader(job MigrationJob, columns []ColumnType, in <-chan []UnknownRowValues) {
|
||||
for rows := range in {
|
||||
log.Debugf("Chunk received, loading data into...")
|
||||
|
||||
for i, rowValues := range rows {
|
||||
if i%100 == 0 {
|
||||
logSampleRow(job, columns, rowValues, fmt.Sprintf("row %d", i))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user