后端重构
This commit is contained in:
Yvan
@@ -0,0 +1,189 @@
|
||||
package initialize
|
||||
|
||||
import (
|
||||
"kra/internal/conf"
|
||||
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
const SystemDB = "system"
|
||||
|
||||
// DBList 数据库列表管理器
|
||||
type DBList struct {
|
||||
dbs map[string]*gorm.DB
|
||||
}
|
||||
|
||||
// NewDBList 创建数据库列表
|
||||
func NewDBList() *DBList {
|
||||
return &DBList{
|
||||
dbs: make(map[string]*gorm.DB),
|
||||
}
|
||||
}
|
||||
|
||||
// InitDBList 初始化多数据库连接
|
||||
func InitDBList(dbList []*conf.SpecializedDB) (*DBList, error) {
|
||||
list := NewDBList()
|
||||
|
||||
for _, info := range dbList {
|
||||
if info.Disable {
|
||||
continue
|
||||
}
|
||||
|
||||
var db *gorm.DB
|
||||
var err error
|
||||
|
||||
switch info.Type {
|
||||
case "mysql":
|
||||
db, err = GormMysqlByConfig(specializedDBToMysql(info))
|
||||
case "pgsql":
|
||||
db, err = GormPgsqlByConfig(specializedDBToPgsql(info))
|
||||
case "sqlite":
|
||||
db, err = GormSqliteByConfig(specializedDBToSqlite(info))
|
||||
case "mssql":
|
||||
db, err = GormMssqlByConfig(specializedDBToMssql(info))
|
||||
case "oracle":
|
||||
db, err = GormOracleByConfig(specializedDBToOracle(info))
|
||||
default:
|
||||
continue
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if db != nil {
|
||||
list.dbs[info.AliasName] = db
|
||||
}
|
||||
}
|
||||
|
||||
return list, nil
|
||||
}
|
||||
|
||||
// Get 获取指定名称的数据库连接
|
||||
func (l *DBList) Get(name string) *gorm.DB {
|
||||
if l.dbs == nil {
|
||||
return nil
|
||||
}
|
||||
return l.dbs[name]
|
||||
}
|
||||
|
||||
// GetSystem 获取系统数据库连接
|
||||
func (l *DBList) GetSystem() *gorm.DB {
|
||||
return l.Get(SystemDB)
|
||||
}
|
||||
|
||||
// Set 设置数据库连接
|
||||
func (l *DBList) Set(name string, db *gorm.DB) {
|
||||
if l.dbs == nil {
|
||||
l.dbs = make(map[string]*gorm.DB)
|
||||
}
|
||||
l.dbs[name] = db
|
||||
}
|
||||
|
||||
// All 获取所有数据库连接
|
||||
func (l *DBList) All() map[string]*gorm.DB {
|
||||
return l.dbs
|
||||
}
|
||||
|
||||
// Close 关闭所有数据库连接
|
||||
func (l *DBList) Close() error {
|
||||
for _, db := range l.dbs {
|
||||
sqlDB, err := db.DB()
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
sqlDB.Close()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// specializedDBToMysql 转换为MySQL配置
|
||||
func specializedDBToMysql(s *conf.SpecializedDB) *conf.Mysql {
|
||||
return &conf.Mysql{
|
||||
Prefix: s.Prefix,
|
||||
Path: s.Path,
|
||||
Port: s.Port,
|
||||
Config: s.Config,
|
||||
DbName: s.DbName,
|
||||
Username: s.Username,
|
||||
Password: s.Password,
|
||||
Engine: s.Engine,
|
||||
MaxIdleConns: s.MaxIdleConns,
|
||||
MaxOpenConns: s.MaxOpenConns,
|
||||
LogMode: s.LogMode,
|
||||
LogZap: s.LogZap,
|
||||
Singular: s.Singular,
|
||||
}
|
||||
}
|
||||
|
||||
// specializedDBToPgsql 转换为PostgreSQL配置
|
||||
func specializedDBToPgsql(s *conf.SpecializedDB) *conf.Pgsql {
|
||||
return &conf.Pgsql{
|
||||
Prefix: s.Prefix,
|
||||
Path: s.Path,
|
||||
Port: s.Port,
|
||||
Config: s.Config,
|
||||
DbName: s.DbName,
|
||||
Username: s.Username,
|
||||
Password: s.Password,
|
||||
Engine: s.Engine,
|
||||
MaxIdleConns: s.MaxIdleConns,
|
||||
MaxOpenConns: s.MaxOpenConns,
|
||||
LogMode: s.LogMode,
|
||||
LogZap: s.LogZap,
|
||||
Singular: s.Singular,
|
||||
}
|
||||
}
|
||||
|
||||
// specializedDBToSqlite 转换为SQLite配置
|
||||
func specializedDBToSqlite(s *conf.SpecializedDB) *conf.Sqlite {
|
||||
return &conf.Sqlite{
|
||||
Prefix: s.Prefix,
|
||||
Path: s.Path,
|
||||
DbName: s.DbName,
|
||||
Engine: s.Engine,
|
||||
MaxIdleConns: s.MaxIdleConns,
|
||||
MaxOpenConns: s.MaxOpenConns,
|
||||
LogMode: s.LogMode,
|
||||
LogZap: s.LogZap,
|
||||
Singular: s.Singular,
|
||||
}
|
||||
}
|
||||
|
||||
// specializedDBToMssql 转换为MSSQL配置
|
||||
func specializedDBToMssql(s *conf.SpecializedDB) *conf.Mssql {
|
||||
return &conf.Mssql{
|
||||
Prefix: s.Prefix,
|
||||
Path: s.Path,
|
||||
Port: s.Port,
|
||||
Config: s.Config,
|
||||
DbName: s.DbName,
|
||||
Username: s.Username,
|
||||
Password: s.Password,
|
||||
Engine: s.Engine,
|
||||
MaxIdleConns: s.MaxIdleConns,
|
||||
MaxOpenConns: s.MaxOpenConns,
|
||||
LogMode: s.LogMode,
|
||||
LogZap: s.LogZap,
|
||||
Singular: s.Singular,
|
||||
}
|
||||
}
|
||||
|
||||
// specializedDBToOracle 转换为Oracle配置
|
||||
func specializedDBToOracle(s *conf.SpecializedDB) *conf.Oracle {
|
||||
return &conf.Oracle{
|
||||
Prefix: s.Prefix,
|
||||
Path: s.Path,
|
||||
Port: s.Port,
|
||||
Config: s.Config,
|
||||
DbName: s.DbName,
|
||||
Username: s.Username,
|
||||
Password: s.Password,
|
||||
Engine: s.Engine,
|
||||
MaxIdleConns: s.MaxIdleConns,
|
||||
MaxOpenConns: s.MaxOpenConns,
|
||||
LogMode: s.LogMode,
|
||||
LogZap: s.LogZap,
|
||||
Singular: s.Singular,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,69 @@
|
||||
package initialize
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"kra/internal/conf"
|
||||
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/logger"
|
||||
"gorm.io/gorm/schema"
|
||||
)
|
||||
|
||||
// GormConfig 创建GORM配置
|
||||
func GormConfig(prefix string, singular bool, logMode string) *gorm.Config {
|
||||
return &gorm.Config{
|
||||
Logger: logger.Default.LogMode(parseLogLevel(logMode)),
|
||||
NamingStrategy: schema.NamingStrategy{
|
||||
TablePrefix: prefix,
|
||||
SingularTable: singular,
|
||||
},
|
||||
DisableForeignKeyConstraintWhenMigrating: true,
|
||||
}
|
||||
}
|
||||
|
||||
// GormConfigWithSlowThreshold 创建带慢查询阈值的GORM配置
|
||||
func GormConfigWithSlowThreshold(prefix string, singular bool, logMode string, slowThreshold time.Duration) *gorm.Config {
|
||||
return &gorm.Config{
|
||||
Logger: logger.Default.LogMode(parseLogLevel(logMode)),
|
||||
NamingStrategy: schema.NamingStrategy{
|
||||
TablePrefix: prefix,
|
||||
SingularTable: singular,
|
||||
},
|
||||
DisableForeignKeyConstraintWhenMigrating: true,
|
||||
}
|
||||
}
|
||||
|
||||
// parseLogLevel 解析日志级别
|
||||
func parseLogLevel(level string) logger.LogLevel {
|
||||
switch level {
|
||||
case "silent":
|
||||
return logger.Silent
|
||||
case "error":
|
||||
return logger.Error
|
||||
case "warn":
|
||||
return logger.Warn
|
||||
case "info":
|
||||
return logger.Info
|
||||
default:
|
||||
return logger.Info
|
||||
}
|
||||
}
|
||||
|
||||
// NewDBByType 根据数据库类型创建数据库连接
|
||||
func NewDBByType(dbType string, bootstrap *conf.Bootstrap) (*gorm.DB, error) {
|
||||
switch dbType {
|
||||
case "mysql":
|
||||
return GormMysql(bootstrap.Mysql)
|
||||
case "pgsql":
|
||||
return GormPgsql(bootstrap.Pgsql)
|
||||
case "sqlite":
|
||||
return GormSqlite(bootstrap.Sqlite)
|
||||
case "mssql":
|
||||
return GormMssql(bootstrap.Mssql)
|
||||
case "oracle":
|
||||
return GormOracle(bootstrap.Oracle)
|
||||
default:
|
||||
return GormMysql(bootstrap.Mysql)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
package initialize
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"kra/internal/conf"
|
||||
|
||||
"gorm.io/driver/sqlserver"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// GormMssql 初始化MSSQL数据库
|
||||
func GormMssql(m *conf.Mssql) (*gorm.DB, error) {
|
||||
if m == nil || m.DbName == "" {
|
||||
return nil, fmt.Errorf("mssql config is empty")
|
||||
}
|
||||
return initMssqlDatabase(m)
|
||||
}
|
||||
|
||||
// GormMssqlByConfig 通过传入配置初始化MSSQL数据库
|
||||
func GormMssqlByConfig(m *conf.Mssql) (*gorm.DB, error) {
|
||||
return initMssqlDatabase(m)
|
||||
}
|
||||
|
||||
// initMssqlDatabase 初始化MSSQL数据库的辅助函数
|
||||
func initMssqlDatabase(m *conf.Mssql) (*gorm.DB, error) {
|
||||
if m.DbName == "" {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
dsn := MssqlDsn(m)
|
||||
mssqlConfig := sqlserver.Config{
|
||||
DSN: dsn,
|
||||
DefaultStringSize: 191,
|
||||
}
|
||||
|
||||
gormConfig := GormConfig(m.Prefix, m.Singular, m.LogMode)
|
||||
|
||||
db, err := gorm.Open(sqlserver.New(mssqlConfig), gormConfig)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if m.Engine != "" {
|
||||
db.InstanceSet("gorm:table_options", "ENGINE="+m.Engine)
|
||||
}
|
||||
|
||||
sqlDB, err := db.DB()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
sqlDB.SetMaxIdleConns(int(m.MaxIdleConns))
|
||||
sqlDB.SetMaxOpenConns(int(m.MaxOpenConns))
|
||||
|
||||
return db, nil
|
||||
}
|
||||
|
||||
// MssqlDsn 生成MSSQL DSN
|
||||
func MssqlDsn(m *conf.Mssql) string {
|
||||
return fmt.Sprintf("sqlserver://%s:%s@%s:%s?database=%s&%s",
|
||||
m.Username, m.Password, m.Path, m.Port, m.DbName, m.Config)
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
package initialize
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"kra/internal/conf"
|
||||
|
||||
_ "github.com/go-sql-driver/mysql"
|
||||
"gorm.io/driver/mysql"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// GormMysql 初始化MySQL数据库
|
||||
func GormMysql(m *conf.Mysql) (*gorm.DB, error) {
|
||||
if m == nil || m.DbName == "" {
|
||||
return nil, fmt.Errorf("mysql config is empty")
|
||||
}
|
||||
return initMysqlDatabase(m)
|
||||
}
|
||||
|
||||
// GormMysqlByConfig 通过传入配置初始化MySQL数据库
|
||||
func GormMysqlByConfig(m *conf.Mysql) (*gorm.DB, error) {
|
||||
return initMysqlDatabase(m)
|
||||
}
|
||||
|
||||
// initMysqlDatabase 初始化MySQL数据库的辅助函数
|
||||
func initMysqlDatabase(m *conf.Mysql) (*gorm.DB, error) {
|
||||
if m.DbName == "" {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
dsn := MysqlDsn(m)
|
||||
mysqlConfig := mysql.Config{
|
||||
DSN: dsn,
|
||||
DefaultStringSize: 191,
|
||||
SkipInitializeWithVersion: false,
|
||||
}
|
||||
|
||||
gormConfig := GormConfig(m.Prefix, m.Singular, m.LogMode)
|
||||
|
||||
db, err := gorm.Open(mysql.New(mysqlConfig), gormConfig)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if m.Engine != "" {
|
||||
db.InstanceSet("gorm:table_options", "ENGINE="+m.Engine)
|
||||
}
|
||||
|
||||
sqlDB, err := db.DB()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
sqlDB.SetMaxIdleConns(int(m.MaxIdleConns))
|
||||
sqlDB.SetMaxOpenConns(int(m.MaxOpenConns))
|
||||
|
||||
return db, nil
|
||||
}
|
||||
|
||||
// MysqlDsn 生成MySQL DSN
|
||||
func MysqlDsn(m *conf.Mysql) string {
|
||||
return fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?%s",
|
||||
m.Username, m.Password, m.Path, m.Port, m.DbName, m.Config)
|
||||
}
|
||||
@@ -0,0 +1,54 @@
|
||||
package initialize
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"kra/internal/conf"
|
||||
|
||||
oracle "github.com/dzwvip/gorm-oracle"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// GormOracle 初始化Oracle数据库
|
||||
func GormOracle(o *conf.Oracle) (*gorm.DB, error) {
|
||||
if o == nil || o.DbName == "" {
|
||||
return nil, fmt.Errorf("oracle config is empty")
|
||||
}
|
||||
return initOracleDatabase(o)
|
||||
}
|
||||
|
||||
// GormOracleByConfig 通过传入配置初始化Oracle数据库
|
||||
func GormOracleByConfig(o *conf.Oracle) (*gorm.DB, error) {
|
||||
return initOracleDatabase(o)
|
||||
}
|
||||
|
||||
// initOracleDatabase 初始化Oracle数据库的辅助函数
|
||||
func initOracleDatabase(o *conf.Oracle) (*gorm.DB, error) {
|
||||
if o.DbName == "" {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
dsn := OracleDsn(o)
|
||||
gormConfig := GormConfig(o.Prefix, o.Singular, o.LogMode)
|
||||
|
||||
db, err := gorm.Open(oracle.Open(dsn), gormConfig)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
sqlDB, err := db.DB()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
sqlDB.SetMaxIdleConns(int(o.MaxIdleConns))
|
||||
sqlDB.SetMaxOpenConns(int(o.MaxOpenConns))
|
||||
|
||||
return db, nil
|
||||
}
|
||||
|
||||
// OracleDsn 生成Oracle DSN
|
||||
func OracleDsn(o *conf.Oracle) string {
|
||||
return fmt.Sprintf("oracle://%s:%s@%s:%s/%s?%s",
|
||||
o.Username, o.Password, o.Path, o.Port, o.DbName, o.Config)
|
||||
}
|
||||
@@ -0,0 +1,59 @@
|
||||
package initialize
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"kra/internal/conf"
|
||||
|
||||
"gorm.io/driver/postgres"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// GormPgsql 初始化PostgreSQL数据库
|
||||
func GormPgsql(p *conf.Pgsql) (*gorm.DB, error) {
|
||||
if p == nil || p.DbName == "" {
|
||||
return nil, fmt.Errorf("pgsql config is empty")
|
||||
}
|
||||
return initPgsqlDatabase(p)
|
||||
}
|
||||
|
||||
// GormPgsqlByConfig 通过传入配置初始化PostgreSQL数据库
|
||||
func GormPgsqlByConfig(p *conf.Pgsql) (*gorm.DB, error) {
|
||||
return initPgsqlDatabase(p)
|
||||
}
|
||||
|
||||
// initPgsqlDatabase 初始化PostgreSQL数据库的辅助函数
|
||||
func initPgsqlDatabase(p *conf.Pgsql) (*gorm.DB, error) {
|
||||
if p.DbName == "" {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
dsn := PgsqlDsn(p)
|
||||
pgsqlConfig := postgres.Config{
|
||||
DSN: dsn,
|
||||
PreferSimpleProtocol: false,
|
||||
}
|
||||
|
||||
gormConfig := GormConfig(p.Prefix, p.Singular, p.LogMode)
|
||||
|
||||
db, err := gorm.Open(postgres.New(pgsqlConfig), gormConfig)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
sqlDB, err := db.DB()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
sqlDB.SetMaxIdleConns(int(p.MaxIdleConns))
|
||||
sqlDB.SetMaxOpenConns(int(p.MaxOpenConns))
|
||||
|
||||
return db, nil
|
||||
}
|
||||
|
||||
// PgsqlDsn 生成PostgreSQL DSN
|
||||
func PgsqlDsn(p *conf.Pgsql) string {
|
||||
return fmt.Sprintf("host=%s user=%s password=%s dbname=%s port=%s %s",
|
||||
p.Path, p.Username, p.Password, p.DbName, p.Port, p.Config)
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
package initialize
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"kra/internal/conf"
|
||||
|
||||
"github.com/glebarez/sqlite"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// GormSqlite 初始化SQLite数据库
|
||||
func GormSqlite(s *conf.Sqlite) (*gorm.DB, error) {
|
||||
if s == nil || s.DbName == "" {
|
||||
return nil, fmt.Errorf("sqlite config is empty")
|
||||
}
|
||||
return initSqliteDatabase(s)
|
||||
}
|
||||
|
||||
// GormSqliteByConfig 通过传入配置初始化SQLite数据库
|
||||
func GormSqliteByConfig(s *conf.Sqlite) (*gorm.DB, error) {
|
||||
return initSqliteDatabase(s)
|
||||
}
|
||||
|
||||
// initSqliteDatabase 初始化SQLite数据库的辅助函数
|
||||
func initSqliteDatabase(s *conf.Sqlite) (*gorm.DB, error) {
|
||||
if s.DbName == "" {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
dsn := SqliteDsn(s)
|
||||
gormConfig := GormConfig(s.Prefix, s.Singular, s.LogMode)
|
||||
|
||||
db, err := gorm.Open(sqlite.Open(dsn), gormConfig)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
sqlDB, err := db.DB()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
sqlDB.SetMaxIdleConns(int(s.MaxIdleConns))
|
||||
sqlDB.SetMaxOpenConns(int(s.MaxOpenConns))
|
||||
|
||||
return db, nil
|
||||
}
|
||||
|
||||
// SqliteDsn 生成SQLite DSN
|
||||
func SqliteDsn(s *conf.Sqlite) string {
|
||||
if s.Path != "" {
|
||||
return fmt.Sprintf("%s/%s", s.Path, s.DbName)
|
||||
}
|
||||
return s.DbName
|
||||
}
|
||||
@@ -0,0 +1,193 @@
|
||||
package initialize
|
||||
|
||||
import (
|
||||
"kra/internal/conf"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/redis/go-redis/v9"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// Initializer 初始化器
|
||||
type Initializer struct {
|
||||
config *conf.Bootstrap
|
||||
db *gorm.DB
|
||||
dbList *DBList
|
||||
redis *RedisClient
|
||||
redisList *RedisList
|
||||
timerManager *TimerManager
|
||||
pluginManager *PluginManager
|
||||
mcpServer *MCPServer
|
||||
}
|
||||
|
||||
// NewInitializer 创建初始化器
|
||||
func NewInitializer(config *conf.Bootstrap) *Initializer {
|
||||
return &Initializer{
|
||||
config: config,
|
||||
}
|
||||
}
|
||||
|
||||
// InitDB 初始化主数据库
|
||||
func (i *Initializer) InitDB() (*gorm.DB, error) {
|
||||
dbType := "mysql"
|
||||
if i.config.System != nil && i.config.System.DbType != "" {
|
||||
dbType = i.config.System.DbType
|
||||
}
|
||||
|
||||
db, err := NewDBByType(dbType, i.config)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
i.db = db
|
||||
return db, nil
|
||||
}
|
||||
|
||||
// InitDBList 初始化多数据库
|
||||
func (i *Initializer) InitDBList() (*DBList, error) {
|
||||
dbList, err := InitDBList(i.config.DbList)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// 如果有system数据库，设置为主数据库
|
||||
if sysDB := dbList.GetSystem(); sysDB != nil {
|
||||
i.db = sysDB
|
||||
}
|
||||
|
||||
i.dbList = dbList
|
||||
return dbList, nil
|
||||
}
|
||||
|
||||
// InitRedis 初始化主Redis
|
||||
func (i *Initializer) InitRedis() (*RedisClient, error) {
|
||||
if i.config.Redis == nil {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
client, err := NewRedisClient(i.config.Redis)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
i.redis = client
|
||||
return client, nil
|
||||
}
|
||||
|
||||
// InitRedisList 初始化多Redis
|
||||
func (i *Initializer) InitRedisList() (*RedisList, error) {
|
||||
redisList, err := InitRedisList(i.config.RedisList)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
i.redisList = redisList
|
||||
return redisList, nil
|
||||
}
|
||||
|
||||
// InitTimer 初始化定时任务
|
||||
func (i *Initializer) InitTimer() *TimerManager {
|
||||
manager := NewTimerManager(i.db)
|
||||
manager.Start()
|
||||
i.timerManager = manager
|
||||
return manager
|
||||
}
|
||||
|
||||
// InitPlugin 初始化插件管理器
|
||||
func (i *Initializer) InitPlugin() *PluginManager {
|
||||
manager := NewPluginManager(i.db)
|
||||
i.pluginManager = manager
|
||||
return manager
|
||||
}
|
||||
|
||||
// InitMCP 初始化MCP服务器
|
||||
func (i *Initializer) InitMCP() *MCPServer {
|
||||
server := NewMCPServer(i.config.Mcp)
|
||||
i.mcpServer = server
|
||||
return server
|
||||
}
|
||||
|
||||
// DB 获取主数据库
|
||||
func (i *Initializer) DB() *gorm.DB {
|
||||
return i.db
|
||||
}
|
||||
|
||||
// DBList 获取数据库列表
|
||||
func (i *Initializer) DBList() *DBList {
|
||||
return i.dbList
|
||||
}
|
||||
|
||||
// Redis 获取主Redis客户端
|
||||
func (i *Initializer) Redis() redis.UniversalClient {
|
||||
if i.redis == nil {
|
||||
return nil
|
||||
}
|
||||
return i.redis.Client()
|
||||
}
|
||||
|
||||
// RedisList 获取Redis列表
|
||||
func (i *Initializer) RedisList() *RedisList {
|
||||
return i.redisList
|
||||
}
|
||||
|
||||
// TimerManager 获取定时任务管理器
|
||||
func (i *Initializer) TimerManager() *TimerManager {
|
||||
return i.timerManager
|
||||
}
|
||||
|
||||
// PluginManager 获取插件管理器
|
||||
func (i *Initializer) PluginManager() *PluginManager {
|
||||
return i.pluginManager
|
||||
}
|
||||
|
||||
// MCPServer 获取MCP服务器
|
||||
func (i *Initializer) MCPServer() *MCPServer {
|
||||
return i.mcpServer
|
||||
}
|
||||
|
||||
// InstallPlugins 安装插件
|
||||
func (i *Initializer) InstallPlugins(privateGroup *gin.RouterGroup, publicGroup *gin.RouterGroup, engine *gin.Engine) {
|
||||
if i.pluginManager != nil {
|
||||
i.pluginManager.InstallPlugins(privateGroup, publicGroup, engine)
|
||||
}
|
||||
}
|
||||
|
||||
// RegisterMCPRoutes 注册MCP路由
|
||||
func (i *Initializer) RegisterMCPRoutes(engine *gin.Engine) {
|
||||
if i.mcpServer != nil {
|
||||
i.mcpServer.RegisterRoutes(engine)
|
||||
}
|
||||
}
|
||||
|
||||
// Close 关闭所有资源
|
||||
func (i *Initializer) Close() error {
|
||||
// 关闭定时任务
|
||||
if i.timerManager != nil {
|
||||
i.timerManager.Stop()
|
||||
}
|
||||
|
||||
// 关闭数据库连接
|
||||
if i.db != nil {
|
||||
sqlDB, err := i.db.DB()
|
||||
if err == nil {
|
||||
sqlDB.Close()
|
||||
}
|
||||
}
|
||||
|
||||
// 关闭数据库列表
|
||||
if i.dbList != nil {
|
||||
i.dbList.Close()
|
||||
}
|
||||
|
||||
// 关闭Redis连接
|
||||
if i.redis != nil {
|
||||
i.redis.Close()
|
||||
}
|
||||
|
||||
// 关闭Redis列表
|
||||
if i.redisList != nil {
|
||||
i.redisList.Close()
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,117 @@
|
||||
package initialize
|
||||
|
||||
import (
|
||||
"kra/internal/conf"
|
||||
mcpTool "kra/pkg/mcp/tool"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
"github.com/mark3labs/mcp-go/server"
|
||||
"go.uber.org/zap"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// MCPServer MCP服务器包装器
|
||||
type MCPServer struct {
|
||||
server *server.MCPServer
|
||||
sseServer *server.SSEServer
|
||||
config *conf.MCP
|
||||
}
|
||||
|
||||
// NewMCPServer 创建MCP服务器
|
||||
func NewMCPServer(config *conf.MCP) *MCPServer {
|
||||
if config == nil || !config.Enabled {
|
||||
return nil
|
||||
}
|
||||
|
||||
s := server.NewMCPServer(
|
||||
config.Name,
|
||||
config.Version,
|
||||
)
|
||||
|
||||
sseServer := server.NewSSEServer(s,
|
||||
server.WithSSEEndpoint(config.SsePath),
|
||||
server.WithMessageEndpoint(config.MessagePath),
|
||||
server.WithBaseURL(config.UrlPrefix))
|
||||
|
||||
return &MCPServer{
|
||||
server: s,
|
||||
sseServer: sseServer,
|
||||
config: config,
|
||||
}
|
||||
}
|
||||
|
||||
// Server 获取MCP服务器实例
|
||||
func (m *MCPServer) Server() *server.MCPServer {
|
||||
if m == nil {
|
||||
return nil
|
||||
}
|
||||
return m.server
|
||||
}
|
||||
|
||||
// SSEServer 获取SSE服务器实例
|
||||
func (m *MCPServer) SSEServer() *server.SSEServer {
|
||||
if m == nil {
|
||||
return nil
|
||||
}
|
||||
return m.sseServer
|
||||
}
|
||||
|
||||
// Config 获取MCP配置
|
||||
func (m *MCPServer) Config() *conf.MCP {
|
||||
if m == nil {
|
||||
return nil
|
||||
}
|
||||
return m.config
|
||||
}
|
||||
|
||||
// RegisterRoutes 注册MCP路由到Gin引擎
|
||||
func (m *MCPServer) RegisterRoutes(engine *gin.Engine) {
|
||||
if m == nil || m.sseServer == nil || m.config == nil {
|
||||
return
|
||||
}
|
||||
|
||||
// 如果不是独立模式，注册到主服务
|
||||
if !m.config.Separate {
|
||||
// 注册SSE端点
|
||||
engine.GET(m.config.SsePath, func(c *gin.Context) {
|
||||
m.sseServer.SSEHandler().ServeHTTP(c.Writer, c.Request)
|
||||
})
|
||||
|
||||
// 注册消息端点
|
||||
engine.POST(m.config.MessagePath, func(c *gin.Context) {
|
||||
m.sseServer.MessageHandler().ServeHTTP(c.Writer, c.Request)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// RegisterTool 注册MCP工具
|
||||
func (m *MCPServer) RegisterTool(tool mcp.Tool, handler server.ToolHandlerFunc) {
|
||||
if m == nil || m.server == nil {
|
||||
return
|
||||
}
|
||||
m.server.AddTool(tool, handler)
|
||||
}
|
||||
|
||||
// MCPToolRegistrar MCP工具注册器接口
|
||||
type MCPToolRegistrar interface {
|
||||
RegisterTools(s *server.MCPServer)
|
||||
}
|
||||
|
||||
// RegisterAllTools 注册所有MCP工具
|
||||
func (m *MCPServer) RegisterAllTools(registrars ...MCPToolRegistrar) {
|
||||
if m == nil || m.server == nil {
|
||||
return
|
||||
}
|
||||
for _, registrar := range registrars {
|
||||
registrar.RegisterTools(m.server)
|
||||
}
|
||||
}
|
||||
|
||||
// RegisterMCPTools 注册所有MCP工具（带依赖注入）
|
||||
func (m *MCPServer) RegisterMCPTools(db *gorm.DB, logger *zap.Logger, autoCodeConfig *mcpTool.AutoCodeConfig) {
|
||||
if m == nil || m.server == nil {
|
||||
return
|
||||
}
|
||||
mcpTool.RegisterAllToolsWithDependencies(m.server, db, logger, autoCodeConfig)
|
||||
}
|
||||
@@ -0,0 +1,99 @@
|
||||
package initialize
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// Plugin 插件接口 (v1版本)
|
||||
type Plugin interface {
|
||||
// Register 注册路由
|
||||
Register(group *gin.RouterGroup)
|
||||
// RouterPath 返回路由路径
|
||||
RouterPath() string
|
||||
}
|
||||
|
||||
// PluginV2 插件接口 (v2版本)
|
||||
type PluginV2 interface {
|
||||
// Register 注册到引擎
|
||||
Register(engine *gin.Engine)
|
||||
}
|
||||
|
||||
// PluginManager 插件管理器
|
||||
type PluginManager struct {
|
||||
db *gorm.DB
|
||||
pluginsV1 []Plugin
|
||||
pluginsV2 []PluginV2
|
||||
initialized bool
|
||||
}
|
||||
|
||||
// NewPluginManager 创建插件管理器
|
||||
func NewPluginManager(db *gorm.DB) *PluginManager {
|
||||
return &PluginManager{
|
||||
db: db,
|
||||
pluginsV1: make([]Plugin, 0),
|
||||
pluginsV2: make([]PluginV2, 0),
|
||||
}
|
||||
}
|
||||
|
||||
// RegisterPluginV1 注册V1版本插件
|
||||
func (m *PluginManager) RegisterPluginV1(plugin Plugin) {
|
||||
m.pluginsV1 = append(m.pluginsV1, plugin)
|
||||
}
|
||||
|
||||
// RegisterPluginV2 注册V2版本插件
|
||||
func (m *PluginManager) RegisterPluginV2(plugin PluginV2) {
|
||||
m.pluginsV2 = append(m.pluginsV2, plugin)
|
||||
}
|
||||
|
||||
// InstallPlugins 安装所有插件
|
||||
func (m *PluginManager) InstallPlugins(privateGroup *gin.RouterGroup, publicGroup *gin.RouterGroup, engine *gin.Engine) {
|
||||
if m.db == nil {
|
||||
fmt.Println("项目暂未初始化，无法安装插件，初始化后重启项目即可完成插件安装")
|
||||
return
|
||||
}
|
||||
|
||||
// 安装V1版本插件
|
||||
m.installPluginsV1(privateGroup, publicGroup)
|
||||
|
||||
// 安装V2版本插件
|
||||
m.installPluginsV2(engine)
|
||||
|
||||
m.initialized = true
|
||||
}
|
||||
|
||||
// installPluginsV1 安装V1版本插件
|
||||
func (m *PluginManager) installPluginsV1(privateGroup *gin.RouterGroup, publicGroup *gin.RouterGroup) {
|
||||
for _, plugin := range m.pluginsV1 {
|
||||
fmt.Printf("%s 注册开始!\n", plugin.RouterPath())
|
||||
pluginGroup := privateGroup.Group(plugin.RouterPath())
|
||||
plugin.Register(pluginGroup)
|
||||
fmt.Printf("%s 注册成功!\n", plugin.RouterPath())
|
||||
}
|
||||
}
|
||||
|
||||
// installPluginsV2 安装V2版本插件
|
||||
func (m *PluginManager) installPluginsV2(engine *gin.Engine) {
|
||||
for _, plugin := range m.pluginsV2 {
|
||||
plugin.Register(engine)
|
||||
}
|
||||
}
|
||||
|
||||
// PluginInit 初始化V1版本插件（兼容旧版方式）
|
||||
func PluginInit(group *gin.RouterGroup, plugins ...Plugin) {
|
||||
for _, plugin := range plugins {
|
||||
fmt.Printf("%s 注册开始!\n", plugin.RouterPath())
|
||||
pluginGroup := group.Group(plugin.RouterPath())
|
||||
plugin.Register(pluginGroup)
|
||||
fmt.Printf("%s 注册成功!\n", plugin.RouterPath())
|
||||
}
|
||||
}
|
||||
|
||||
// PluginInitV2 初始化V2版本插件（兼容旧版方式）
|
||||
func PluginInitV2(engine *gin.Engine, plugins ...PluginV2) {
|
||||
for _, plugin := range plugins {
|
||||
plugin.Register(engine)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,8 @@
|
||||
package initialize
|
||||
|
||||
import (
|
||||
"github.com/google/wire"
|
||||
)
|
||||
|
||||
// ProviderSet is initialize providers.
|
||||
var ProviderSet = wire.NewSet(NewTimerManager)
|
||||
@@ -0,0 +1,136 @@
|
||||
package initialize
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"kra/internal/conf"
|
||||
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
// RedisClient Redis客户端包装器
|
||||
type RedisClient struct {
|
||||
client redis.UniversalClient
|
||||
name string
|
||||
}
|
||||
|
||||
// NewRedisClient 创建Redis客户端
|
||||
func NewRedisClient(cfg *conf.Redis) (*RedisClient, error) {
|
||||
if cfg == nil {
|
||||
return nil, fmt.Errorf("redis config is nil")
|
||||
}
|
||||
|
||||
client, err := initRedisClient(cfg)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &RedisClient{
|
||||
client: client,
|
||||
name: cfg.Name,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// initRedisClient 初始化Redis客户端
|
||||
func initRedisClient(cfg *conf.Redis) (redis.UniversalClient, error) {
|
||||
var client redis.UniversalClient
|
||||
|
||||
// 使用集群模式
|
||||
if cfg.UseCluster {
|
||||
client = redis.NewClusterClient(&redis.ClusterOptions{
|
||||
Addrs: cfg.ClusterAddrs,
|
||||
Password: cfg.Password,
|
||||
})
|
||||
} else {
|
||||
// 使用单例模式
|
||||
client = redis.NewClient(&redis.Options{
|
||||
Addr: cfg.Addr,
|
||||
Password: cfg.Password,
|
||||
DB: int(cfg.Db),
|
||||
})
|
||||
}
|
||||
|
||||
// 测试连接
|
||||
pong, err := client.Ping(context.Background()).Result()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("redis connect ping failed, name: %s, err: %v", cfg.Name, err)
|
||||
}
|
||||
|
||||
_ = pong // 连接成功
|
||||
return client, nil
|
||||
}
|
||||
|
||||
// Client 获取Redis客户端
|
||||
func (r *RedisClient) Client() redis.UniversalClient {
|
||||
return r.client
|
||||
}
|
||||
|
||||
// Name 获取Redis名称
|
||||
func (r *RedisClient) Name() string {
|
||||
return r.name
|
||||
}
|
||||
|
||||
// Close 关闭Redis连接
|
||||
func (r *RedisClient) Close() error {
|
||||
if r.client != nil {
|
||||
return r.client.Close()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// RedisList Redis列表管理器
|
||||
type RedisList struct {
|
||||
clients map[string]*RedisClient
|
||||
}
|
||||
|
||||
// NewRedisList 创建Redis列表
|
||||
func NewRedisList() *RedisList {
|
||||
return &RedisList{
|
||||
clients: make(map[string]*RedisClient),
|
||||
}
|
||||
}
|
||||
|
||||
// InitRedisList 初始化多Redis连接
|
||||
func InitRedisList(redisList []*conf.Redis) (*RedisList, error) {
|
||||
list := NewRedisList()
|
||||
|
||||
for _, cfg := range redisList {
|
||||
client, err := NewRedisClient(cfg)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
list.clients[cfg.Name] = client
|
||||
}
|
||||
|
||||
return list, nil
|
||||
}
|
||||
|
||||
// Get 获取指定名称的Redis客户端
|
||||
func (l *RedisList) Get(name string) *RedisClient {
|
||||
if l.clients == nil {
|
||||
return nil
|
||||
}
|
||||
return l.clients[name]
|
||||
}
|
||||
|
||||
// Set 设置Redis客户端
|
||||
func (l *RedisList) Set(name string, client *RedisClient) {
|
||||
if l.clients == nil {
|
||||
l.clients = make(map[string]*RedisClient)
|
||||
}
|
||||
l.clients[name] = client
|
||||
}
|
||||
|
||||
// All 获取所有Redis客户端
|
||||
func (l *RedisList) All() map[string]*RedisClient {
|
||||
return l.clients
|
||||
}
|
||||
|
||||
// Close 关闭所有Redis连接
|
||||
func (l *RedisList) Close() error {
|
||||
for _, client := range l.clients {
|
||||
client.Close()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,105 @@
|
||||
package initialize
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"os"
|
||||
|
||||
"kra/internal/conf"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// justFilesFilesystem 只允许访问文件的文件系统
|
||||
type justFilesFilesystem struct {
|
||||
fs http.FileSystem
|
||||
}
|
||||
|
||||
func (fs justFilesFilesystem) Open(name string) (http.File, error) {
|
||||
f, err := fs.fs.Open(name)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
stat, err := f.Stat()
|
||||
if err != nil {
|
||||
f.Close()
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if stat.IsDir() {
|
||||
f.Close()
|
||||
return nil, os.ErrPermission
|
||||
}
|
||||
|
||||
return f, nil
|
||||
}
|
||||
|
||||
// RouterConfig 路由配置
|
||||
type RouterConfig struct {
|
||||
RouterPrefix string
|
||||
StorePath string
|
||||
SwaggerPath string
|
||||
}
|
||||
|
||||
// NewRouterConfig 创建路由配置
|
||||
func NewRouterConfig(system *conf.System, local *conf.Local) *RouterConfig {
|
||||
cfg := &RouterConfig{
|
||||
RouterPrefix: "",
|
||||
StorePath: "uploads/file",
|
||||
SwaggerPath: "/swagger/*any",
|
||||
}
|
||||
|
||||
if system != nil {
|
||||
cfg.RouterPrefix = system.RouterPrefix
|
||||
}
|
||||
|
||||
if local != nil && local.StorePath != "" {
|
||||
cfg.StorePath = local.StorePath
|
||||
}
|
||||
|
||||
return cfg
|
||||
}
|
||||
|
||||
// SetupStaticFiles 设置静态文件服务
|
||||
func SetupStaticFiles(engine *gin.Engine, storePath string) {
|
||||
if storePath != "" {
|
||||
engine.StaticFS(storePath, justFilesFilesystem{http.Dir(storePath)})
|
||||
}
|
||||
}
|
||||
|
||||
// SetupHealthCheck 设置健康检查端点
|
||||
func SetupHealthCheck(group *gin.RouterGroup) {
|
||||
group.GET("/health", func(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, "ok")
|
||||
})
|
||||
}
|
||||
|
||||
// RouterGroups 路由组集合
|
||||
type RouterGroups struct {
|
||||
Public *gin.RouterGroup
|
||||
Private *gin.RouterGroup
|
||||
}
|
||||
|
||||
// NewRouterGroups 创建路由组
|
||||
func NewRouterGroups(engine *gin.Engine, prefix string) *RouterGroups {
|
||||
return &RouterGroups{
|
||||
Public: engine.Group(prefix),
|
||||
Private: engine.Group(prefix),
|
||||
}
|
||||
}
|
||||
|
||||
// RouteInfo 路由信息
|
||||
type RouteInfo struct {
|
||||
Method string
|
||||
Path string
|
||||
Handler string
|
||||
HandlerFunc gin.HandlerFunc
|
||||
}
|
||||
|
||||
// Routes 路由列表
|
||||
type Routes []RouteInfo
|
||||
|
||||
// GetRoutes 获取所有路由信息
|
||||
func GetRoutes(engine *gin.Engine) gin.RoutesInfo {
|
||||
return engine.Routes()
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
package initialize
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"kra/internal/task"
|
||||
"kra/pkg/timer"
|
||||
|
||||
"github.com/robfig/cron/v3"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// TimerManager 定时任务管理器
|
||||
type TimerManager struct {
|
||||
timer timer.Timer
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
// NewTimerManager 创建定时任务管理器
|
||||
func NewTimerManager(db *gorm.DB) *TimerManager {
|
||||
return &TimerManager{
|
||||
timer: timer.NewTimerTask(),
|
||||
db: db,
|
||||
}
|
||||
}
|
||||
|
||||
// Timer 获取定时器实例
|
||||
func (m *TimerManager) Timer() timer.Timer {
|
||||
return m.timer
|
||||
}
|
||||
|
||||
// Start 启动定时任务
|
||||
func (m *TimerManager) Start() {
|
||||
go func() {
|
||||
var option []cron.Option
|
||||
option = append(option, cron.WithSeconds())
|
||||
|
||||
// 清理DB定时任务
|
||||
_, err := m.timer.AddTaskByFunc("ClearDB", "@daily", func() {
|
||||
err := task.ClearTable(m.db)
|
||||
if err != nil {
|
||||
fmt.Println("timer error:", err)
|
||||
}
|
||||
}, "定时清理数据库【日志，黑名单】内容", option...)
|
||||
if err != nil {
|
||||
fmt.Println("add timer error:", err)
|
||||
}
|
||||
|
||||
// 其他定时任务可以在这里添加
|
||||
}()
|
||||
}
|
||||
|
||||
// Stop 停止定时任务
|
||||
func (m *TimerManager) Stop() {
|
||||
if m.timer != nil {
|
||||
m.timer.Close()
|
||||
}
|
||||
}
|
||||
|
||||
// AddTask 添加定时任务
|
||||
func (m *TimerManager) AddTask(name, spec string, fn func(), desc string, options ...cron.Option) (cron.EntryID, error) {
|
||||
return m.timer.AddTaskByFunc(name, spec, fn, desc, options...)
|
||||
}
|
||||
|
||||
// RemoveTask 移除定时任务
|
||||
func (m *TimerManager) RemoveTask(name string) {
|
||||
m.timer.Clear(name)
|
||||
}
|
||||
@@ -0,0 +1,35 @@
|
||||
package initialize
|
||||
|
||||
import (
|
||||
"kra/pkg/utils"
|
||||
)
|
||||
|
||||
// InitValidator 初始化验证器规则
|
||||
func InitValidator() {
|
||||
// 分页验证规则
|
||||
_ = utils.RegisterRule("PageVerify",
|
||||
utils.Rules{
|
||||
"Page": {utils.NotEmpty()},
|
||||
"PageSize": {utils.NotEmpty()},
|
||||
},
|
||||
)
|
||||
|
||||
// ID验证规则
|
||||
_ = utils.RegisterRule("IdVerify",
|
||||
utils.Rules{
|
||||
"Id": {utils.NotEmpty()},
|
||||
},
|
||||
)
|
||||
|
||||
// 权限ID验证规则
|
||||
_ = utils.RegisterRule("AuthorityIdVerify",
|
||||
utils.Rules{
|
||||
"AuthorityId": {utils.NotEmpty()},
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
// init 包初始化时自动注册验证规则
|
||||
func init() {
|
||||
InitValidator()
|
||||
}
|
||||
Reference in New Issue
Block a user