ai-agent/internal/bootstrap/db.go

package bootstrap

import (
	"fmt"
	"log"
	"os"
	"strings"
	"time"

	"code.tczkiot.com/wlw/ai-agent/internal/pkg/config"

	"github.com/mlogclub/simple/sqls"
	"gorm.io/driver/postgres"
	"gorm.io/gorm"

	// "gorm.io/gorm/logger"
	"gorm.io/gorm/logger"
	"gorm.io/gorm/schema"
)

func InitDB(cfg config.DBConfig) (*gorm.DB, error) {
	dialector, err := newDialector(cfg)
	if err != nil {
		return nil, err
	}

	db, err := gorm.Open(dialector, &gorm.Config{
		Logger: logger.New(
			log.New(os.Stdout, "\r\n", log.LstdFlags),
			logger.Config{
				SlowThreshold:             time.Second,
				LogLevel:                  logger.Warn,
				IgnoreRecordNotFoundError: true,
				Colorful:                  true,
			},
		),
		NamingStrategy: schema.NamingStrategy{
			TablePrefix:   "t_",
			SingularTable: true,
		},
	})
	if err != nil {
		return nil, err
	}

	sqlDB, err := db.DB()
	if err != nil {
		return nil, err
	}
	if cfg.MaxIdleConns > 0 {
		sqlDB.SetMaxIdleConns(cfg.MaxIdleConns)
	}
	if cfg.MaxOpenConns > 0 {
		sqlDB.SetMaxOpenConns(cfg.MaxOpenConns)
	}
	if cfg.ConnMaxIdleTimeSeconds > 0 {
		sqlDB.SetConnMaxIdleTime(time.Duration(cfg.ConnMaxIdleTimeSeconds) * time.Second)
	}
	if cfg.ConnMaxLifetimeSeconds > 0 {
		sqlDB.SetConnMaxLifetime(time.Duration(cfg.ConnMaxLifetimeSeconds) * time.Second)
	}

	sqls.SetDB(db)
	return db, nil
}

// UseDatabase scopes the host application's connection to AI Agent tables.
// The connection pool is shared, while the naming strategy is copied so the
// host database naming rules remain untouched.
func UseDatabase(database *gorm.DB, tablePrefix string) error {
	moduleDB, err := ScopedDatabase(database, tablePrefix)
	if err != nil {
		return err
	}
	sqls.SetDB(moduleDB)
	return nil
}

// ScopedDatabase shares the host connection pool while applying the AI Agent
// table prefix to an isolated GORM session.
func ScopedDatabase(database *gorm.DB, tablePrefix string) (*gorm.DB, error) {
	if database == nil {
		return nil, fmt.Errorf("database is required")
	}
	moduleDB := database.Session(&gorm.Session{NewDB: true})
	moduleConfig := *moduleDB.Config
	moduleConfig.NamingStrategy = schema.NamingStrategy{
		TablePrefix:   tablePrefix,
		SingularTable: true,
	}
	moduleDB.Config = &moduleConfig
	return moduleDB, nil
}

func newDialector(cfg config.DBConfig) (gorm.Dialector, error) {
	switch strings.ToLower(strings.TrimSpace(cfg.Type)) {
	case "postgres", "postgresql":
		return postgres.Open(cfg.DSN), nil
	default:
		return nil, fmt.Errorf("unsupported db type %q: only postgres is supported", cfg.Type)
	}
}