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) } }