feat: implement MSSQL bulk loading functionality with transaction handling

cmd/go_migrate/loader.go

@@ -2,15 +2,18 @@ package main
 
 import (
 	"context"
+	"database/sql"
 	"fmt"
 	"time"
 
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgxpool"
+	mssql "github.com/microsoft/go-mssqldb"
 	log "github.com/sirupsen/logrus"
 )
 
 func fakeLoader(job MigrationJob, columns []ColumnType, in <-chan [][]any) {
+
 	for rows := range in {
 		log.Debugf("Chunk received, loading data into...")
 
@@ -51,8 +54,72 @@ func loadRowsPostgres(ctx context.Context, job MigrationJob, columns []ColumnTyp
 		chunkDuration := time.Since(chunkStartTime)
 		rowsPerSec := float64(len(rows)) / chunkDuration.Seconds()
 
-		log.Infof("Loaded chunk #%d: %d rows in %v (copy: %v, %.0f rows/sec) - Total: %d rows",
-			chunkCount, len(rows), chunkDuration, copyDuration, rowsPerSec, totalRowsLoaded)
+		log.Infof("Loaded chunk #%d: %d rows in %v (copy: %v, %.0f rows/sec) - Total: %d rows", chunkCount, len(rows), chunkDuration, copyDuration, rowsPerSec, totalRowsLoaded)
+	}
+
+	return nil
+}
+
+func loadRowsMssql(ctx context.Context, job MigrationJob, columns []ColumnType, db *sql.DB, in <-chan []UnknownRowValues) error {
+	chunkCount := 0
+	totalRowsLoaded := 0
+
+	for rows := range in {
+		chunkStartTime := time.Now()
+
+		tx, err := db.BeginTx(ctx, nil)
+		if err != nil {
+			return fmt.Errorf("error starting transaction: %w", err)
+		}
+
+		fullTableName := fmt.Sprintf("[%s].[%s]", job.Schema, job.Table)
+		colNames := Map(columns, func(col ColumnType) string {
+			return col.name
+		})
+
+		stmt, err := tx.PrepareContext(ctx, mssql.CopyIn(fullTableName, mssql.BulkOptions{}, colNames...))
+		if err != nil {
+			tx.Rollback()
+			return fmt.Errorf("error preparing bulk copy statement: %w", err)
+		}
+
+		copyStartTime := time.Now()
+
+		for _, row := range rows {
+			_, err = stmt.ExecContext(ctx, row...)
+			if err != nil {
+				stmt.Close()
+				tx.Rollback()
+				return fmt.Errorf("error executing row insert: %w", err)
+			}
+		}
+
+		result, err := stmt.ExecContext(ctx)
+		if err != nil {
+			stmt.Close()
+			tx.Rollback()
+			return fmt.Errorf("error flushing bulk data: %w", err)
+		}
+
+		err = stmt.Close()
+		if err != nil {
+			tx.Rollback()
+			return fmt.Errorf("error closing statement: %w", err)
+		}
+
+		if err := tx.Commit(); err != nil {
+			return fmt.Errorf("error committing transaction: %w", err)
+		}
+
+		rowsAffected, _ := result.RowsAffected()
+		chunkCount++
+		totalRowsLoaded += int(rowsAffected)
+
+		copyDuration := time.Since(copyStartTime)
+		chunkDuration := time.Since(chunkStartTime)
+		rowsPerSec := float64(len(rows)) / chunkDuration.Seconds()
+
+		log.Infof("Loaded chunk #%d (MSSQL): %d rows in %v (copy: %v, %.0f rows/sec) - Total: %d rows", chunkCount, len(rows), chunkDuration, copyDuration, rowsPerSec, totalRowsLoaded)
 	}
 
 	return nil

scripts/mssql-copy-in/main.go

@@ -0,0 +1,101 @@
+//go:build go1.10
+// +build go1.10
+
+package main
+
+import (
+	"database/sql"
+	"log"
+	"strings"
+	"unicode/utf8"
+
+	mssql "github.com/microsoft/go-mssqldb"
+)
+
+const (
+	createTestTable = `CREATE TABLE test_table(
+		id int IDENTITY(1,1) NOT NULL,
+		test_nvarchar nvarchar(50) NULL,
+		test_varchar varchar(50) NULL,
+		test_float float NULL,
+		test_datetime2_3 datetime2(3) NULL,
+		test_bitn bit NULL,
+		test_bigint bigint NOT NULL,
+		test_geom geometry NULL,
+	CONSTRAINT PK_table_test_id PRIMARY KEY CLUSTERED
+	(
+		id ASC
+	) ON [PRIMARY]);`
+	dropTestTable = "IF OBJECT_ID('test_table', 'U') IS NOT NULL DROP TABLE test_table;"
+)
+
+// This example shows how to perform bulk imports
+func main() {
+	db, err := sql.Open("sqlserver", "")
+	if err != nil {
+		log.Fatal("Open connection failed:", err.Error())
+	}
+	defer db.Close()
+
+	txn, err := db.Begin()
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	// Create table
+	_, err = db.Exec(createTestTable)
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer db.Exec(dropTestTable)
+
+	// mssqldb.CopyIn creates string to be consumed by Prepare
+	stmt, err := txn.Prepare(mssql.CopyIn("test_table", mssql.BulkOptions{}, "test_varchar", "test_nvarchar", "test_float", "test_bigint"))
+	if err != nil {
+		log.Fatal(err.Error())
+	}
+
+	for i := 0; i < 10; i++ {
+		_, err = stmt.Exec(generateString(0, 30), generateStringUnicode(0, 30), i, i)
+		if err != nil {
+			log.Fatal(err.Error())
+		}
+	}
+
+	result, err := stmt.Exec()
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	err = stmt.Close()
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	err = txn.Commit()
+	if err != nil {
+		log.Fatal(err)
+	}
+	rowCount, _ := result.RowsAffected()
+	log.Printf("%d row copied\n", rowCount)
+	log.Printf("bye\n")
+}
+
+func generateString(x int, n int) string {
+	letters := "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
+	b := make([]byte, n)
+	for i := range b {
+		b[i] = letters[(x+i)%len(letters)]
+	}
+	return string(b)
+}
+func generateStringUnicode(x int, n int) string {
+	letters := []byte("ab©💾é?ghïjklmnopqЯ☀tuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")
+	b := &strings.Builder{}
+	for i := 0; i < n; i++ {
+		r, sz := utf8.DecodeRune(letters[x%len(letters):])
+		x += sz
+		b.WriteRune(r)
+	}
+	return b.String()
+}