db/engine/engine.go

package engine

import (
	"database/sql"
	"fmt"
	"log"

	"git.fsdpf.net/go/db"
	sqlite3dialect "git.fsdpf.net/go/db/dialect/sqlite3"
	sqlite3vtab "git.fsdpf.net/go/db/dialect/sqlite3/vtab"
)

type Engine struct {
	configs map[string]DBConfig
	dbs     map[string]*sql.DB
}

var _engine *Engine

func init() {
	_engine = &Engine{
		configs: make(map[string]DBConfig),
		dbs:     make(map[string]*sql.DB),
	}
}

func (e Engine) Connection(name string) *db.Database {
	cfg, ok := e.configs[name]

	if !ok {
		panic(fmt.Errorf("database connection %s not configured", name))
	}

	_db, ok := e.dbs[name]

	if !ok {
		_db = e.MakeConnection(cfg)
		e.dbs[name] = _db

		// vtable 连接：同时创建一个 :memory: 连接供虚拟表查询执行，
		// 原文件连接仅用于 _vtab_cache 持久化，两者互不阻塞。
		if cfg.Driver == "vtable" {
			// 使用命名共享内存数据库，确保连接池中所有连接共享同一份内存数据，
			// 避免匿名 :memory: 各连接独立导致虚表在其他连接不可见的问题。
			sharedDSN := fmt.Sprintf("file:%s?mode=memory&cache=shared", name)
			memDB, err := sql.Open(sqlite3vtab.DriverName, sharedDSN)
			if err != nil {
				panic(fmt.Sprintf("vtable: open memory connection for %s: %v", name, err))
			}
			e.dbs["__"+name] = memDB
			e.configs["__"+name] = DBConfig{Driver: "vtable"}
		}
	}

	return db.New(cfg.Driver, _db)
}

func (e Engine) MakeConnection(cfg DBConfig) (db *sql.DB) {
	dsn := cfg.ToDSN()

	driverName := cfg.Driver
	switch cfg.Driver {
	case "mysql":
	case "sqlite3":
		driverName = sqlite3dialect.DriverWithIF
	case "vtable":
		driverName = sqlite3vtab.DriverName
	case "sqlserver":
	case "postgres":
	case "duckdb":
	default:
		panic(fmt.Sprintf("Unsupported driver: %s", cfg.Driver))
	}

	db, err := sql.Open(driverName, dsn)

	if err != nil {
		panic(err)
	}

	if err := db.Ping(); err != nil {
		panic(err)
	}

	if cfg.Driver == "duckdb" {
		for _, ext := range cfg.DuckDB.Extensions {
			if _, err := db.Exec("INSTALL " + ext); err != nil {
				panic(fmt.Sprintf("duckdb: install extension %q: %v", ext, err))
			}
			if _, err := db.Exec("LOAD " + ext); err != nil {
				panic(fmt.Sprintf("duckdb: load extension %q: %v", ext, err))
			}
		}
	}

	// 应用连接池配置（对所有驱动生效）
	if cfg.MaxOpenConns > 0 {
		db.SetMaxOpenConns(cfg.MaxOpenConns)
	}
	if cfg.MaxIdleConns > 0 {
		db.SetMaxIdleConns(cfg.MaxIdleConns)
	}
	if cfg.ConnMaxLifetime > 0 {
		db.SetConnMaxLifetime(cfg.ConnMaxLifetime)
	}
	if cfg.ConnMaxIdleTime > 0 {
		db.SetConnMaxIdleTime(cfg.ConnMaxIdleTime)
	}

	return db
}

// Shutdown 关闭所有已建立的数据库连接，使 DuckDB 等驱动得以完成 WAL checkpoint。
// 实现 do.Shutdownable 接口，由 DI 容器在关闭时调用。
func (e Engine) Shutdown() error {
	for name, sqlDB := range e.dbs {
		if err := sqlDB.Close(); err != nil {
			log.Printf("engine: closing connection %q: %v", name, err)
		}
	}
	return nil
}

func Open(cfgs map[string]DBConfig) Engine {
	for n, cfg := range cfgs {
		_engine.configs[n] = cfg
	}

	return *_engine
}

func Mock(cfgs map[string]MockDBConfig) Engine {
	for k, cfg := range cfgs {
		_engine.dbs[k] = cfg.Mock
		_engine.configs[k] = DBConfig{Driver: cfg.Driver}
	}
	return *_engine
}