92 lines
2.1 KiB
Go
92 lines
2.1 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"fmt"
|
|
)
|
|
|
|
type BatchRange struct {
|
|
LowerLimit int
|
|
UpperLimit int
|
|
validRange bool
|
|
}
|
|
|
|
func estimateTotalRowsMssql(ctx context.Context, db *sql.DB, job MigrationJob) (int, error) {
|
|
query := `
|
|
SELECT
|
|
SUM(p.rows) AS count
|
|
FROM sys.tables t
|
|
JOIN sys.schemas s ON t.schema_id = s.schema_id
|
|
JOIN sys.partitions p ON t.object_id = p.object_id
|
|
WHERE s.name = @schema AND t.name = @table AND p.index_id IN (0, 1)
|
|
GROUP BY t.name`
|
|
|
|
var rowsCount int
|
|
err := db.QueryRowContext(ctx, query, sql.Named("schema", job.Schema), sql.Named("table", job.Table)).Scan(&rowsCount)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
return rowsCount, nil
|
|
}
|
|
|
|
func calculateChunkRangesMssql(ctx context.Context, db *sql.DB, job MigrationJob, batchCount int) ([]BatchRange, error) {
|
|
query := fmt.Sprintf(`
|
|
SELECT
|
|
MIN([%s]) AS lower_limit,
|
|
MAX([%s]) AS upper_limit
|
|
FROM
|
|
(SELECT [%s], NTILE(@batchCount) OVER (ORDER BY [%s]) AS chunk_id FROM [%s].[%s]) AS T
|
|
GROUP BY chunk_id
|
|
ORDER BY chunk_id`, job.PrimaryKey, job.PrimaryKey, job.PrimaryKey, job.PrimaryKey, job.Schema, job.Table)
|
|
|
|
rows, err := db.QueryContext(ctx, query, sql.Named("batchCount", batchCount))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
|
|
batchRanges := make([]BatchRange, 0, batchCount)
|
|
|
|
for rows.Next() {
|
|
var br BatchRange
|
|
br.validRange = true
|
|
|
|
if err := rows.Scan(&br.LowerLimit, &br.UpperLimit); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
batchRanges = append(batchRanges, br)
|
|
}
|
|
|
|
if err := rows.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return batchRanges, nil
|
|
}
|
|
|
|
const estimatedRowsPerBatch = 100_000
|
|
|
|
func calculateBatchMetrics(ctx context.Context, db *sql.DB, job MigrationJob) ([]BatchRange, error) {
|
|
rowsCount, err := estimateTotalRowsMssql(ctx, db, job)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
batchCount := 1
|
|
if rowsCount > estimatedRowsPerBatch {
|
|
batchCount = rowsCount / estimatedRowsPerBatch
|
|
} else {
|
|
return []BatchRange{{validRange: false}}, nil
|
|
}
|
|
|
|
chunksRange, err := calculateChunkRangesMssql(ctx, db, job, batchCount)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return chunksRange, nil
|
|
}
|