Scert/backend/internal/database/database.go

package database

import (
	"context"
	"database/sql"
	"fmt"
	"net/url"
	"path/filepath"
	"strings"
	"time"

	"github.com/go-sql-driver/mysql"
	"github.com/jackc/pgx/v5"
	_ "github.com/jackc/pgx/v5/stdlib"
	_ "modernc.org/sqlite"

	"git.pcux.cn/xuans/Scert/backend/internal/config"
)

type Database struct {
	DB *sql.DB
}

func Open(ctx context.Context, cfg config.DatabaseConfig) (*Database, error) {
	driver := strings.ToLower(strings.TrimSpace(cfg.Driver))
	if driver != "sqlite" && driver != "mysql" && driver != "postgres" {
		return nil, fmt.Errorf("数据库驱动无效：%s", cfg.Driver)
	}
	dsn := strings.TrimSpace(cfg.DSN)
	if driver == "mysql" {
		var err error
		dsn, err = normalizeMySQLDSN(dsn)
		if err != nil {
			return nil, err
		}
	}
	if driver == "postgres" {
		var err error
		dsn, err = normalizePostgresDSN(dsn)
		if err != nil {
			return nil, err
		}
	}
	if driver == "sqlite" {
		if cfg.DSN != ":memory:" && len(cfg.DSN) > 0 && cfg.DSN != "file::memory:?cache=shared" {
			if err := ensureSQLiteDir(cfg.DSN); err != nil {
				return nil, err
			}
		}
	}
	sqlDriver := driver
	if driver == "postgres" {
		sqlDriver = "pgx"
	}
	db, err := sql.Open(sqlDriver, dsn)
	if err != nil {
		return nil, fmt.Errorf("打开数据库失败：%w", err)
	}
	db.SetMaxOpenConns(cfg.MaxOpenConns)
	db.SetMaxIdleConns(cfg.MaxIdleConns)
	db.SetConnMaxLifetime(time.Duration(cfg.ConnMaxLifetimeMins) * time.Minute)
	pingCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
	defer cancel()
	if err := db.PingContext(pingCtx); err != nil {
		_ = db.Close()
		return nil, fmt.Errorf("连接数据库失败：%w", err)
	}
	if driver == "sqlite" {
		if err := configureSQLite(pingCtx, db); err != nil {
			_ = db.Close()
			return nil, err
		}
	}
	return &Database{DB: db}, nil
}

func configureSQLite(ctx context.Context, db *sql.DB) error {
	if _, err := db.ExecContext(ctx, "PRAGMA busy_timeout = 5000"); err != nil {
		return fmt.Errorf("配置 SQLite 忙等待失败：%w", err)
	}
	var mode string
	if _, err := db.ExecContext(ctx, "PRAGMA journal_mode = WAL"); err != nil {
		return fmt.Errorf("配置 SQLite WAL 失败：%w", err)
	}
	if err := db.QueryRowContext(ctx, "PRAGMA journal_mode").Scan(&mode); err != nil {
		return fmt.Errorf("读取 SQLite 日志模式失败：%w", err)
	}
	if strings.EqualFold(strings.TrimSpace(mode), "memory") || strings.EqualFold(strings.TrimSpace(mode), "wal") {
		return nil
	}
	return fmt.Errorf("SQLite 日志模式无效：%s", mode)
}

func normalizeMySQLDSN(dsn string) (string, error) {
	cfg, err := mysql.ParseDSN(dsn)
	if err != nil {
		return "", fmt.Errorf("解析 MySQL 连接信息失败：%w", err)
	}
	cfg.ParseTime = true
	cfg.Loc = time.UTC
	return cfg.FormatDSN(), nil
}

func normalizePostgresDSN(dsn string) (string, error) {
	parsed, err := pgx.ParseConfig(dsn)
	if err != nil {
		return "", fmt.Errorf("解析 PostgreSQL 连接信息失败：%w", err)
	}
	if strings.TrimSpace(parsed.Host) == "" || strings.TrimSpace(parsed.Database) == "" || strings.TrimSpace(parsed.User) == "" {
		return "", fmt.Errorf("PostgreSQL 连接信息必须包含主机、数据库和用户名")
	}
	sslMode := postgresSSLMode(dsn)
	if sslMode != "verify-full" {
		return "", fmt.Errorf("PostgreSQL 连接必须使用 sslmode=verify-full")
	}
	if parsed.TLSConfig == nil || parsed.TLSConfig.InsecureSkipVerify {
		return "", fmt.Errorf("PostgreSQL TLS 校验配置无效")
	}
	return dsn, nil
}

func postgresSSLMode(dsn string) string {
	dsn = strings.TrimSpace(dsn)
	if parsed, err := url.Parse(dsn); err == nil && (parsed.Scheme == "postgres" || parsed.Scheme == "postgresql") {
		return strings.ToLower(strings.TrimSpace(parsed.Query().Get("sslmode")))
	}
	for _, field := range strings.Fields(dsn) {
		key, value, found := strings.Cut(field, "=")
		if found && strings.EqualFold(strings.TrimSpace(key), "sslmode") {
			return strings.ToLower(strings.Trim(strings.TrimSpace(value), "'\""))
		}
	}
	return ""
}

// SanitizeDSN 返回可写入日志的数据库连接信息，不包含 MySQL 密码。
func SanitizeDSN(driver, dsn string) string {
	switch strings.ToLower(strings.TrimSpace(driver)) {
	case "mysql":
		cfg, err := mysql.ParseDSN(dsn)
		if err != nil {
			return "mysql://[连接信息无效]"
		}
		if cfg.Passwd != "" {
			cfg.Passwd = "********"
		}
		return cfg.FormatDSN()
	case "postgres":
		parsed, err := url.Parse(dsn)
		if err == nil && (parsed.Scheme == "postgres" || parsed.Scheme == "postgresql") {
			if parsed.User != nil {
				username := parsed.User.Username()
				if _, hasPassword := parsed.User.Password(); hasPassword {
					parsed.User = url.UserPassword(username, "********")
				}
			}
			return parsed.String()
		}
		return sanitizePostgresKeywordDSN(dsn)
	}
	return dsn
}

func sanitizePostgresKeywordDSN(dsn string) string {
	fields := strings.Fields(dsn)
	for index, field := range fields {
		key, _, found := strings.Cut(field, "=")
		if found && strings.EqualFold(strings.TrimSpace(key), "password") {
			fields[index] = key + "=********"
		}
	}
	return strings.Join(fields, " ")
}

func ensureSQLiteDir(dsn string) error {
	path := dsn
	if strings.HasPrefix(path, "file:") {
		path = strings.TrimPrefix(strings.SplitN(path, "?", 2)[0], "file:")
	}
	return osMkdirAll(filepath.Dir(path))
}

var osMkdirAll = func(path string) error {
	if path == "." || path == "" {
		return nil
	}
	return mkdirAll(path)
}

var mkdirAll = func(path string) error {
	return makeDirectory(path)
}

func (d *Database) Close() error {
	if d == nil || d.DB == nil {
		return nil
	}
	return d.DB.Close()
}

// Ping 检查数据库连接是否可用。
func (d *Database) Ping(ctx context.Context) error {
	if d == nil || d.DB == nil {
		return fmt.Errorf("数据库连接不能为空")
	}
	if err := d.DB.PingContext(ctx); err != nil {
		return fmt.Errorf("数据库健康检查失败：%w", err)
	}
	return nil
}