diff --git a/cmd/go_migrate/loader.go b/cmd/go_migrate/loader.go index 0a9f13e..abc4c8e 100644 --- a/cmd/go_migrate/loader.go +++ b/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 diff --git a/scripts/mssql-copy-in/main.go b/scripts/mssql-copy-in/main.go new file mode 100644 index 0000000..847cb4e --- /dev/null +++ b/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() +}