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 }