109 lines
2.6 KiB
Go
109 lines
2.6 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"fmt"
|
|
"time"
|
|
|
|
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
|
|
"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
|
|
"github.com/google/uuid"
|
|
)
|
|
|
|
func estimateTotalRowsMssql(ctx context.Context, db *sql.DB, tableInfo config.SourceTableInfo) (int64, 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`
|
|
|
|
ctxTimeout, cancel := context.WithTimeout(ctx, time.Second*20)
|
|
defer cancel()
|
|
|
|
var rowsCount int64
|
|
err := db.QueryRowContext(ctxTimeout, query, sql.Named("schema", tableInfo.Schema), sql.Named("table", tableInfo.Table)).Scan(&rowsCount)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
return rowsCount, nil
|
|
}
|
|
|
|
func calculateBatchesMssql(ctx context.Context, db *sql.DB, tableInfo config.SourceTableInfo, batchCount int64) ([]models.Partition, 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 batch_id FROM [%s].[%s]) AS T
|
|
GROUP BY batch_id
|
|
ORDER BY batch_id`,
|
|
tableInfo.PrimaryKey,
|
|
tableInfo.PrimaryKey,
|
|
tableInfo.PrimaryKey,
|
|
tableInfo.PrimaryKey,
|
|
tableInfo.Schema,
|
|
tableInfo.Table)
|
|
|
|
ctxTimeout, cancel := context.WithTimeout(ctx, time.Second*20)
|
|
defer cancel()
|
|
|
|
rows, err := db.QueryContext(ctxTimeout, query, sql.Named("batchCount", batchCount))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
|
|
batches := make([]models.Partition, 0, batchCount)
|
|
|
|
for rows.Next() {
|
|
batch := models.Partition{
|
|
Id: uuid.New(),
|
|
ShouldUseRange: true,
|
|
RetryCounter: 0,
|
|
IsLowerLimitInclusive: true,
|
|
}
|
|
|
|
if err := rows.Scan(&batch.LowerLimit, &batch.UpperLimit); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
batches = append(batches, batch)
|
|
}
|
|
|
|
if err := rows.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return batches, nil
|
|
}
|
|
|
|
func partitionGeneratorMssql(ctx context.Context, db *sql.DB, tableInfo config.SourceTableInfo, rowsPerBatch int64) ([]models.Partition, error) {
|
|
rowsCount, err := estimateTotalRowsMssql(ctx, db, tableInfo)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
var batchCount int64 = 1
|
|
if rowsCount > rowsPerBatch {
|
|
batchCount = rowsCount / rowsPerBatch
|
|
} else {
|
|
return []models.Partition{{
|
|
Id: uuid.New(),
|
|
ShouldUseRange: false,
|
|
RetryCounter: 0,
|
|
}}, nil
|
|
}
|
|
|
|
batches, err := calculateBatchesMssql(ctx, db, tableInfo, batchCount)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return batches, nil
|
|
}
|