mirror of
https://github.com/truewhile/MeBox.git
synced 2026-09-29 03:26:37 +08:00
优化
This commit is contained in:
truewhile
@@ -0,0 +1,223 @@
|
||||
package database
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"go.uber.org/zap"
|
||||
"gorm.io/driver/postgres"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/logger"
|
||||
|
||||
"github.com/ShukeBta/MMTL/internal/config"
|
||||
"github.com/ShukeBta/MMTL/internal/model"
|
||||
)
|
||||
|
||||
// DatabaseStatus describes the currently active database engine and runtime metrics.
|
||||
type DatabaseStatus struct {
|
||||
Type string `json:"type"`
|
||||
DSN string `json:"dsn,omitempty"`
|
||||
DBPath string `json:"db_path,omitempty"`
|
||||
OpenConns int `json:"open_conns"`
|
||||
InUse int `json:"in_use"`
|
||||
Idle int `json:"idle"`
|
||||
MaxOpenConns int `json:"max_open_conns"`
|
||||
TableCounts map[string]int64 `json:"table_counts"`
|
||||
}
|
||||
|
||||
// PostgresTestResult returns latency and version info after testing connection.
|
||||
type PostgresTestResult struct {
|
||||
Success bool `json:"success"`
|
||||
LatencyMS int64 `json:"latency_ms"`
|
||||
Version string `json:"version,omitempty"`
|
||||
Message string `json:"message,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
// DatabaseMigrationResult returns row counts and execution duration of migration.
|
||||
type DatabaseMigrationResult struct {
|
||||
Success bool `json:"success"`
|
||||
TotalRows int64 `json:"total_rows"`
|
||||
TableRows map[string]int64 `json:"table_rows"`
|
||||
DurationMS int64 `json:"duration_ms"`
|
||||
Message string `json:"message,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
// InspectDatabaseStatus queries the currently active database for metrics and table rows.
|
||||
func InspectDatabaseStatus(db *gorm.DB, cfg *config.Config) *DatabaseStatus {
|
||||
st := &DatabaseStatus{
|
||||
Type: "sqlite",
|
||||
TableCounts: make(map[string]int64),
|
||||
}
|
||||
if cfg != nil {
|
||||
st.DBPath = cfg.Database.DBPath
|
||||
if cfg.Database.Type == "postgres" || (cfg.Database.Type == "auto" && strings.TrimSpace(cfg.Database.DSN) != "") {
|
||||
st.Type = "postgres"
|
||||
st.DSN = MaskDSN(cfg.Database.DSN)
|
||||
}
|
||||
}
|
||||
if isPostgres(db) {
|
||||
st.Type = "postgres"
|
||||
}
|
||||
|
||||
if db != nil {
|
||||
if sqlDB, err := db.DB(); err == nil {
|
||||
stats := sqlDB.Stats()
|
||||
st.OpenConns = stats.OpenConnections
|
||||
st.InUse = stats.InUse
|
||||
st.Idle = stats.Idle
|
||||
st.MaxOpenConns = stats.MaxOpenConnections
|
||||
}
|
||||
|
||||
// Count rows for major model tables
|
||||
for _, m := range model.AllModels() {
|
||||
if tbl, err := modelTableName(db, m); err == nil {
|
||||
if db.Migrator().HasTable(tbl) {
|
||||
var count int64
|
||||
if err := db.Raw("SELECT COUNT(1) FROM " + quoteIdent(tbl)).Scan(&count).Error; err == nil {
|
||||
st.TableCounts[tbl] = count
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return st
|
||||
}
|
||||
|
||||
// TestPostgres establishes a temporary connection to verify reachability and permissions.
|
||||
func TestPostgres(dsn string) (*PostgresTestResult, error) {
|
||||
dsn = strings.TrimSpace(dsn)
|
||||
if dsn == "" {
|
||||
return &PostgresTestResult{
|
||||
Success: false,
|
||||
Error: "PostgreSQL DSN 不能为空",
|
||||
}, nil
|
||||
}
|
||||
|
||||
start := time.Now()
|
||||
testDB, err := gorm.Open(postgres.Open(dsn), &gorm.Config{
|
||||
Logger: logger.Default.LogMode(logger.Silent),
|
||||
})
|
||||
if err != nil {
|
||||
return &PostgresTestResult{
|
||||
Success: false,
|
||||
Error: fmt.Sprintf("连接失败: %v", err),
|
||||
}, nil
|
||||
}
|
||||
|
||||
sqlDB, err := testDB.DB()
|
||||
if err != nil {
|
||||
return &PostgresTestResult{
|
||||
Success: false,
|
||||
Error: fmt.Sprintf("获取底层连接失败: %v", err),
|
||||
}, nil
|
||||
}
|
||||
defer sqlDB.Close()
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
if err := sqlDB.PingContext(ctx); err != nil {
|
||||
return &PostgresTestResult{
|
||||
Success: false,
|
||||
Error: fmt.Sprintf("Ping 超时或失败: %v", err),
|
||||
}, nil
|
||||
}
|
||||
|
||||
var version string
|
||||
if err := testDB.WithContext(ctx).Raw("SELECT version()").Scan(&version).Error; err != nil {
|
||||
version = "PostgreSQL (unknown version)"
|
||||
}
|
||||
|
||||
latency := time.Since(start).Milliseconds()
|
||||
return &PostgresTestResult{
|
||||
Success: true,
|
||||
LatencyMS: latency,
|
||||
Version: version,
|
||||
Message: "连接成功",
|
||||
}, nil
|
||||
}
|
||||
|
||||
// MigrateCurrentToPostgres performs schema initialization and full table data copy into target PostgreSQL.
|
||||
func MigrateCurrentToPostgres(src *gorm.DB, targetDSN string, batchSize int, log *zap.Logger) (*DatabaseMigrationResult, error) {
|
||||
targetDSN = strings.TrimSpace(targetDSN)
|
||||
if targetDSN == "" {
|
||||
return nil, fmt.Errorf("target PostgreSQL DSN cannot be empty")
|
||||
}
|
||||
if src == nil {
|
||||
return nil, fmt.Errorf("current database is not available")
|
||||
}
|
||||
|
||||
started := time.Now()
|
||||
targetDB, err := gorm.Open(postgres.Open(targetDSN), &gorm.Config{
|
||||
Logger: newGormLogger(log),
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("open target PostgreSQL: %w", err)
|
||||
}
|
||||
targetSQLDB, err := targetDB.DB()
|
||||
if err == nil {
|
||||
defer targetSQLDB.Close()
|
||||
}
|
||||
|
||||
// 1. 初始化目标库 Schema、类型与索引
|
||||
if err := AutoMigrate(targetDB); err != nil {
|
||||
return nil, fmt.Errorf("auto migrate target PostgreSQL: %w", err)
|
||||
}
|
||||
|
||||
// 2. 安全重置目标数据库的初始默认数据
|
||||
if err := resetBootstrapTargetBeforeSQLiteMigrationIfSafe(src, targetDB, log); err != nil {
|
||||
return nil, fmt.Errorf("reset target bootstrap data: %w", err)
|
||||
}
|
||||
|
||||
// 3. 执行数据批量复制
|
||||
tableRows, totalRows, err := copyModelTables(src, targetDB, batchSize)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("copy tables: %w", err)
|
||||
}
|
||||
|
||||
// 4. 标记迁移完成
|
||||
if err := markSQLiteMigrationComplete(targetDB); err != nil {
|
||||
return nil, fmt.Errorf("mark migration complete: %w", err)
|
||||
}
|
||||
|
||||
duration := time.Since(started).Milliseconds()
|
||||
return &DatabaseMigrationResult{
|
||||
Success: true,
|
||||
TotalRows: totalRows,
|
||||
TableRows: tableRows,
|
||||
DurationMS: duration,
|
||||
Message: fmt.Sprintf("成功迁移 %d 条记录至 PostgreSQL", totalRows),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// MaskDSN masks the password in a connection string for safe API responses.
|
||||
func MaskDSN(rawDSN string) string {
|
||||
rawDSN = strings.TrimSpace(rawDSN)
|
||||
if rawDSN == "" {
|
||||
return ""
|
||||
}
|
||||
if u, err := url.Parse(rawDSN); err == nil && u.User != nil {
|
||||
if pass, hasPassword := u.User.Password(); hasPassword && pass != "" {
|
||||
rawUserPass := u.User.String()
|
||||
user := u.User.Username()
|
||||
maskedUserPass := user + ":******"
|
||||
return strings.Replace(rawDSN, rawUserPass+"@", maskedUserPass+"@", 1)
|
||||
}
|
||||
}
|
||||
// Fallback for keyword-style DSN (e.g. host=... password=...)
|
||||
if strings.Contains(rawDSN, "password=") {
|
||||
parts := strings.Fields(rawDSN)
|
||||
for i, p := range parts {
|
||||
if strings.HasPrefix(p, "password=") {
|
||||
parts[i] = "password=******"
|
||||
}
|
||||
}
|
||||
return strings.Join(parts, " ")
|
||||
}
|
||||
return rawDSN
|
||||
}
|
||||
Reference in New Issue
Block a user