yx-dataset-server/boot/boot.go

package boot

import (
	"context"
	"filesdk"
	"os"

	"github.com/dgrijalva/jwt-go"
	redisLib "github.com/go-redis/redis/v8"
	"github.com/gogf/gf/v2/frame/g"
	"github.com/gogf/gf/v2/net/ghttp"
	"go.uber.org/dig"
	"gorm.io/gorm"

	"yx-dataset-server/app/bll"
	"yx-dataset-server/app/bll/impl"
	"yx-dataset-server/app/errors"
	"yx-dataset-server/library/auth"
	"yx-dataset-server/library/dingtalk"
	"yx-dataset-server/library/gplus"
	"yx-dataset-server/library/logger"
	"yx-dataset-server/library/middleware"
	"yx-dataset-server/library/ragflow"
	"yx-dataset-server/library/redis"
	"yx-dataset-server/library/utils"
	"yx-dataset-server/router"
)

// VERSION 定义应用版本号
const VERSION = "1.0.0"

func init() {
	// 初始化logger
	logger.Init(utils.GetConfig("common.run_mode").String())
	logger.SetVersion(VERSION)
	logger.SetTraceIdFunc(utils.NewTraceID)
	ctx := logger.NewTraceIDContext(context.Background(), utils.NewTraceID())
	Init(ctx)
}

// 初始化App,
// TODO: 返回释放回调，暂时没调用
func Init(ctx context.Context) func() {
	logger.Printf(ctx, "服务启动，运行模式：%s，版本号：%s，进程号：%d", utils.GetConfig("common.run_mode").String(), VERSION, os.Getpid())
	// 初始化依赖注入容器
	container, call := buildContainer(ctx)
	// 初始化路由注册
	s := g.Server()
	// 每个请求生成新的追踪Id,如果上下文件中没有trace-id
	s.Use(middleware.TraceIdMiddleware())
	if utils.GetConfig("prometheus.enable").Bool() {
		p := middleware.NewPrometheus("ghttp")
		p.Use(s)
	}
	// 统一处理内部错误
	s.Use(func(r *ghttp.Request) {
		r.Middleware.Next()
		if err := r.GetError(); err != nil {
			gplus.ResError(r, err)
		}
	})
	_ = auth.InitAuth(ctx, container)
	//_ = initJwtAuth(ctx, container)
	initFileSDK(ctx)

	router.InitRouters(s, container)

	// 启动时批量拉起所有钉钉卡片回复模式的机器人
	if err := container.Invoke(func(b bll.IRobotConfig) error {
		return b.RestartAllRobots(ctx)
	}); err != nil {
		logger.Printf(ctx, "启动钉钉机器人（卡片流式模式）失败：%v", err)
	}

	return func() {
		if call != nil {
			call()
		}
	}
}

// 初始化jwt认证,可以把相关配置放到config.toml中
func initJwtAuth(ctx context.Context, container *dig.Container) error {
	var opts []auth.Option
	opts = append(opts, auth.SetExpired(utils.GetConfig("jwt.expired").Int()))
	opts = append(opts, auth.SetSigningKey([]byte(utils.GetConfig("jwt.signing_key").String())))
	opts = append(opts, auth.SetKeyfunc(func(t *jwt.Token) (interface{}, error) {
		if _, ok := t.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, errors.ErrInvalidToken
		}
		return []byte(utils.GetConfig("jwt.signing_key").String()), nil
	}))

	switch utils.GetConfig("jwt.signing_method").String() {
	case "HS256":
		opts = append(opts, auth.SetSigningMethod(jwt.SigningMethodHS256))
	case "HS384":
		opts = append(opts, auth.SetSigningMethod(jwt.SigningMethodHS384))
	case "HS512":
		opts = append(opts, auth.SetSigningMethod(jwt.SigningMethodHS512))
	}
	var store auth.Storer
	switch utils.GetConfig("jwt.store").String() {
	case "file":

	case "redis":
		//store = redis.NewStore(&redis.Config{
		//	Addr:      rcfg.Addr,
		//	Password:  rcfg.Password,
		//	DB:        cfg.RedisDB,
		//	KeyPrefix: cfg.RedisPrefix,
		//})
	}
	return container.Provide(func() auth.Auther { return auth.New(store, opts...) })
}

// 初始化redis
func initRedis(ctx context.Context, container *dig.Container) func() {
	addr := utils.GetConfig("redis.addr").String()
	password := utils.GetConfig("redis.password").String()
	db := utils.GetConfig("redis.db").Int()
	redisCli := redis.Init(ctx, addr, password, db)
	logger.Printf(ctx, "REDIS初始化成功，当前服务器[%s]", addr)
	// 注入redis client
	_ = container.Provide(func() *redisLib.Client {
		return redisCli.GetRedisClient()
	})
	// 将 Redis 作为钉钉 access_token 的 L2 持久化存储注入
	dingtalk.SetTokenStorage(dingtalk.NewRedisTokenStorage(redisCli.GetRedisClient(), ""))
	logger.Printf(ctx, "钉钉 access_token Redis 缓存已启用")
	return func() {
		_ = redisCli.Close
	}
}

// 初始化redis
func initRagHttpClient(ctx context.Context, container *dig.Container) func() {
	addr := utils.GetConfig("ragFlow.base_url").String()
	apiKey := utils.GetConfig("ragFlow.api_key").String()
	ragClient := ragflow.Init(ctx, addr, apiKey)
	// 注入redis client
	_ = container.Provide(func() *ragflow.Client {
		return ragClient
	})
	return nil
}

// 初始化存储，目前只初始化gorm
func initStore(ctx context.Context, container *dig.Container) (func(), error) {
	var storeCall func()
	db, err := initGorm()
	if err != nil {
		return storeCall, err
	}
	// 如果自动映射数据表
	if utils.GetConfig("gorm.enable_auto_migrate").Bool() {
		err = autoMigrate(db)
		if err != nil {
			return storeCall, err
		}
	}
	// 注入DB
	_ = container.Provide(func() *gorm.DB { return db })
	// 注入model接口
	_ = InjectModel(container)
	storeCall = func() {
		sqlDb, _ := db.DB()
		_ = sqlDb.Close()
	}
	logger.Printf(ctx, "MYSQL初始化成功, 服务器[%s], 数据库[%s]",
		utils.GetConfig("mysql.host"),
		utils.GetConfig("mysql.db_name"))
	return storeCall, nil
}

// 构建依赖注入容器
func buildContainer(ctx context.Context) (*dig.Container, func()) {
	container := dig.New()
	// 初始化存储模块
	storeCall, err := initStore(ctx, container)
	if err != nil {
		panic(err)
	}
	// 初始化redis模块
	var redisCall func()
	if utils.GetConfig("redis.enable").Bool() {
		redisCall = initRedis(ctx, container)
	}
	initRagHttpClient(ctx, container)
	// 注入bll
	impl.Inject(container)

	return container, func() {
		if storeCall != nil {
			storeCall()
		}
		if redisCall != nil {
			redisCall()
		}
	}
}

func initFileSDK(ctx context.Context) {
	fileConfig := filesdk.Config{
		Host: utils.GetConfig("file_server.host").String(),
	}
	fileConfig.Init()
}