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:
2026-04-09 16:00:00 -05:00
parent c5987fa7f8
commit 4f181269a4
5 changed files with 11 additions and 236 deletions

View File

@@ -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)
}

View File

@@ -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
}

View File

@@ -5,25 +5,17 @@ import (
)
type MigrationJob struct {
Schema string
Table string
PrimaryKey string
Schema string
Table string
}
var migrationJobs []MigrationJob = []MigrationJob{
{
Schema: "demo",
Table: "users",
PrimaryKey: "id",
Schema: "demo",
Table: "users",
},
}
const (
NumExtractors int = 2
ChunkSize int = 20
QueueSize int = 10
)
func main() {
configureLog()
log.Info("Starting migration...")

View File

@@ -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
}

View File

@@ -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))
}
}
}
}