db/engine/engine_config.go

package engine

import (
	"database/sql"
	"fmt"
	"net/url"
	"time"

	"github.com/mattn/go-sqlite3"
)

type MockDBConfig struct {
	Driver string
	Mock   *sql.DB
}

// DBConfig 数据库配置结构体
type DBConfig struct {
	Driver          string
	Host            string
	Port            string
	Database        string
	Username        string
	Password        string
	Charset         string
	Prefix          string
	ConnMaxLifetime time.Duration
	ConnMaxIdleTime time.Duration
	MaxIdleConns    int
	MaxOpenConns    int
	ParseTime       bool
	EnableLog       bool
	ReadHosts       []string
	WriteHosts      []string

	MySQL struct {
		Dsn             string
		Collation       string
		UnixSocket      string
		MultiStatements bool
	}

	PostgreSQL struct {
		Sslmode         string
		TLS             string
		SearchPath      string
		ConnectTimeout  int
		ApplicationName string
	}

	SQLite struct {
		File        string
		Journal     string
		Locking     string
		Mode        string
		Synchronous int
		Cache       string
		BusyTimeout int
		Raw         func(*sqlite3.SQLiteConn) error
	}

	SQLServer struct {
		Instance          string
		Encrypt           string
		TrustServerCert   bool
		AppName           string
		FailoverPartner   string
		PacketSize        int
		WorkstationID     string
		ConnectionTimeout int
		KeepAlive         int
		Dsn               string
	}
}

// Option 配置选项类型
type Option func(*DBConfig)

// ToDSN 生成对应驱动的 DSN
func (c *DBConfig) ToDSN() string {
	switch c.Driver {
	case "mysql":
		return c.toMySQLDSN()
	case "pgsql":
		return c.toPostgreSQLDSN()
	case "sqlite3":
		return c.toSQLiteDSN()
	case "sqlserver":
		return c.toSQLServerDSN()
	default:
		panic(fmt.Sprintf("Unsupported driver for DSN generation: %s", c.Driver))
	}
}

// toMySQLDSN 生成 MySQL DSN
func (c *DBConfig) toMySQLDSN() string {
	if c.MySQL.Dsn != "" {
		return c.MySQL.Dsn
	}
	var dsn string
	if c.Username != "" {
		if c.Password != "" {
			dsn = fmt.Sprintf("%s:%s@", c.Username, c.Password)
		} else {
			dsn = fmt.Sprintf("%s@", c.Username)
		}
	}
	if c.MySQL.UnixSocket != "" {
		dsn += fmt.Sprintf("unix(%s)", c.MySQL.UnixSocket)
	} else {
		host := c.Host
		if host == "" && len(c.WriteHosts) > 0 {
			host = c.WriteHosts[0]
		}
		if host == "" && len(c.ReadHosts) > 0 {
			host = c.ReadHosts[0]
		}
		if c.Port == "" {
			c.Port = "3306"
		}
		dsn += fmt.Sprintf("tcp(%s:%s)", host, c.Port)
	}
	dsn += fmt.Sprintf("/%s", c.Database)
	params := url.Values{}
	if c.Charset != "" {
		params.Add("charset", c.Charset)
	}
	if c.ParseTime {
		params.Add("parseTime", "true")
	}
	if c.MySQL.Collation != "" {
		params.Add("collation", c.MySQL.Collation)
	}
	if c.MySQL.MultiStatements {
		params.Add("multiStatements", "true")
	}
	if len(params) > 0 {
		dsn += "?" + params.Encode()
	}
	return dsn
}

// toPostgreSQLDSN 生成 PostgreSQL DSN
func (c *DBConfig) toPostgreSQLDSN() string {
	params := url.Values{}
	if c.Host != "" {
		params.Add("host", c.Host)
	} else if len(c.WriteHosts) > 0 {
		params.Add("host", c.WriteHosts[0])
	} else if len(c.ReadHosts) > 0 {
		params.Add("host", c.ReadHosts[0])
	}
	if c.Port == "" {
		c.Port = "5432" // 默认 PostgreSQL 端口
	}
	params.Add("port", c.Port)
	params.Add("dbname", c.Database)
	params.Add("user", c.Username)
	if c.Password != "" {
		params.Add("password", c.Password)
	}
	if c.PostgreSQL.Sslmode != "" {
		params.Add("sslmode", c.PostgreSQL.Sslmode)
	}
	if c.PostgreSQL.ConnectTimeout > 0 {
		params.Add("connect_timeout", fmt.Sprintf("%d", c.PostgreSQL.ConnectTimeout))
	}
	if c.PostgreSQL.ApplicationName != "" {
		params.Add("application_name", c.PostgreSQL.ApplicationName)
	}
	return "postgres://" + c.Username + ":" + url.QueryEscape(c.Password) + "@" + params.Get("host") + ":" + c.Port + "/" + c.Database + "?" + params.Encode()
}

// toSQLiteDSN 生成 SQLite DSN
func (c *DBConfig) toSQLiteDSN() string {
	if c.SQLite.File == "" {
		return ""
	}
	dsn := c.SQLite.File
	params := url.Values{}
	if c.SQLite.Journal != "" {
		params.Add("_journal", c.SQLite.Journal)
	}
	if c.SQLite.Locking != "" {
		params.Add("_locking", c.SQLite.Locking)
	}
	if c.SQLite.Mode != "" {
		params.Add("_mode", c.SQLite.Mode)
	}
	if c.SQLite.Synchronous > 0 {
		params.Add("_synchronous", fmt.Sprintf("%d", c.SQLite.Synchronous))
	}
	if c.SQLite.Cache != "" {
		params.Add("_cache", c.SQLite.Cache)
	}
	if c.SQLite.BusyTimeout > 0 {
		params.Add("_busy_timeout", fmt.Sprintf("%d", c.SQLite.BusyTimeout))
	}
	if len(params) > 0 {
		dsn += "?" + params.Encode()
	}
	return dsn
}

// toSQLServerDSN 生成 SQL Server DSN
func (c *DBConfig) toSQLServerDSN() string {
	if c.SQLServer.Dsn != "" {
		return c.SQLServer.Dsn
	}
	host := c.Host
	if host == "" && len(c.WriteHosts) > 0 {
		host = c.WriteHosts[0]
	}
	if host == "" && len(c.ReadHosts) > 0 {
		host = c.ReadHosts[0]
	}
	if c.Port == "" {
		c.Port = "1433" // 默认 SQL Server 端口
	}
	if c.SQLServer.Instance != "" {
		host += "\\" + c.SQLServer.Instance
	}
	params := url.Values{}
	params.Add("server", host)
	params.Add("database", c.Database)
	if c.SQLServer.Encrypt != "" {
		params.Add("encrypt", c.SQLServer.Encrypt)
	}
	if c.SQLServer.TrustServerCert {
		params.Add("TrustServerCertificate", "true")
	}
	if c.SQLServer.AppName != "" {
		params.Add("app name", c.SQLServer.AppName)
	}
	if c.SQLServer.ConnectionTimeout > 0 {
		params.Add("connection timeout", fmt.Sprintf("%d", c.SQLServer.ConnectionTimeout))
	}
	return "sqlserver://" + c.Username + ":" + url.QueryEscape(c.Password) + "@" + host + ":" + c.Port + "?" + params.Encode()
}

// NewDBConfig 创建新的 DBConfig 实例，检查关键参数
func NewDBConfig(driver string, options ...Option) DBConfig {
	if driver == "" {
		panic("Driver is required and cannot be empty")
	}

	// 支持的驱动类型
	validDrivers := map[string]bool{
		"mysql":     true,
		"pgsql":     true,
		"sqlite3":   true,
		"sqlserver": true,
	}
	if !validDrivers[driver] {
		panic(fmt.Sprintf("Unsupported driver: %s", driver))
	}

	// 初始化配置
	config := &DBConfig{
		Driver:          driver,
		Charset:         "utf8mb4",
		ConnMaxLifetime: 2 * time.Hour,
		ConnMaxIdleTime: 30 * time.Minute,
		MaxIdleConns:    10,
		MaxOpenConns:    100,
	}

	// 应用所有选项
	for _, opt := range options {
		opt(config)
	}

	// 检查关键参数
	switch config.Driver {
	case "mysql", "pgsql", "sqlserver":
		if config.Host == "" && len(config.ReadHosts) == 0 && len(config.WriteHosts) == 0 {
			panic(fmt.Sprintf("Host, ReadHosts, or WriteHosts is required for %s driver", config.Driver))
		}
		if config.Database == "" {
			panic(fmt.Sprintf("Database is required for %s driver", config.Driver))
		}
		if config.Username == "" {
			panic(fmt.Sprintf("Username is required for %s driver", config.Driver))
		}
		if config.Password == "" {
			panic(fmt.Sprintf("Password is required for %s driver", config.Driver))
		}
	case "sqlite3":
		if config.SQLite.File == "" {
			panic("File is required for sqlite3 driver")
		}
	}

	return *config
}

// 通用的 WithOption 配置函数
func WithHost(host string) Option {
	return func(c *DBConfig) {
		c.Host = host
	}
}

func WithPort(port string) Option {
	return func(c *DBConfig) {
		c.Port = port
	}
}

func WithDatabase(database string) Option {
	return func(c *DBConfig) {
		c.Database = database
	}
}

func WithUsername(username string) Option {
	return func(c *DBConfig) {
		c.Username = username
	}
}

func WithPassword(password string) Option {
	return func(c *DBConfig) {
		c.Password = password
	}
}

func WithParseTime(parseTime bool) Option {
	return func(c *DBConfig) {
		c.ParseTime = parseTime
	}
}

func WithCharset(charset string) Option {
	return func(c *DBConfig) {
		c.Charset = charset
	}
}

func WithReadHosts(hosts []string) Option {
	return func(c *DBConfig) {
		c.ReadHosts = hosts
	}
}

func WithWriteHosts(hosts []string) Option {
	return func(c *DBConfig) {
		c.WriteHosts = hosts
	}
}

// MySQL 专用选项
func WithMySQLCollation(collation string) Option {
	return func(c *DBConfig) {
		if c.Driver != "mysql" {
			panic("WithMySQLCollation is only valid for mysql driver")
		}
		c.MySQL.Collation = collation
	}
}

func WithMySQLUnixSocket(socket string) Option {
	return func(c *DBConfig) {
		if c.Driver != "mysql" {
			panic("WithMySQLUnixSocket is only valid for mysql driver")
		}
		c.MySQL.UnixSocket = socket
	}
}

// PostgreSQL 专用选项
func WithPgSslmode(sslmode string) Option {
	return func(c *DBConfig) {
		if c.Driver != "pgsql" {
			panic("WithPgSslmode is only valid for pgsql driver")
		}
		c.PostgreSQL.Sslmode = sslmode
	}
}

// SQLite 专用选项
func WithSQLiteFile(file string) Option {
	return func(c *DBConfig) {
		if c.Driver != "sqlite3" {
			panic("WithSQLiteFile is only valid for sqlite3 driver")
		}
		c.SQLite.File = file
	}
}

// SQLite 注册函数
func WithSQLiteRaw(raw func(*sqlite3.SQLiteConn) error) Option {
	return func(c *DBConfig) {
		if c.Driver != "sqlite3" {
			panic("WithSQLiteRegFn is only valid for sqlite3 driver")
		}
		c.SQLite.Raw = raw
	}
}

func WithSQLiteJournal(journal string) Option {
	return func(c *DBConfig) {
		if c.Driver != "sqlite3" {
			panic("WithSQLiteJournal is only valid for sqlite3 driver")
		}
		c.SQLite.Journal = journal
	}
}

// SQL Server 专用选项
func WithSQLServerInstance(instance string) Option {
	return func(c *DBConfig) {
		if c.Driver != "sqlserver" {
			panic("WithSQLServerInstance is only valid for sqlserver driver")
		}
		c.SQLServer.Instance = instance
	}
}