feat 初次提交

2026-01-03 14:26:09 +08:00
parent 18eb946934
commit 3ea3a3ac6d
48 changed files with 5420 additions and 1 deletions
--- a/db/db.go
+++ b/db/db.go
@@ -0,0 +1,48 @@
+package db
+
+import (
+	"common/config"
+	"common/db/etcd"
+	"common/db/mysql"
+	"common/db/redis"
+	"common/log"
+)
+
+// ModuleDB 数据库模块
+type ModuleDB struct {
+}
+
+func (p *ModuleDB) Init(cfg *config.DBConfig) error {
+	// ETCD
+	if cfg.Etcd != nil {
+		if err := etcd.Init(cfg.Etcd); err != nil {
+			return err
+		}
+	}
+	// MYSQL
+	if cfg.MySQL != nil {
+		if err := mysql.Init(cfg.MySQL); err != nil {
+			return err
+		}
+	}
+	// REDIS
+	if cfg.Redis != nil {
+		if err := redis.Init(cfg.Redis); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (p *ModuleDB) Stop() error {
+	if err := etcd.Close(); err != nil {
+		log.Errorf("close etcd failed: %v", err)
+	}
+	if err := mysql.Close(); err != nil {
+		log.Errorf("close mysql failed: %v", err)
+	}
+	if err := redis.Close(); err != nil {
+		log.Errorf("close redis failed: %v", err)
+	}
+	return nil
+}
--- a/db/etcd/client.go
+++ b/db/etcd/client.go
@@ -0,0 +1,76 @@
+package etcd
+
+import (
+	"common/config"
+	"context"
+	"go.etcd.io/etcd/client/v3"
+	"time"
+)
+
+var instance *Client
+
+type Client struct {
+	cli *clientv3.Client
+}
+
+// Init 初始化
+func Init(cfg *config.EtcdConfig) error {
+	client, err := clientv3.New(clientv3.Config{
+		Endpoints:   cfg.Endpoints,
+		DialTimeout: 5 * time.Second,
+	})
+	instance = &Client{
+		cli: client,
+	}
+	return err
+}
+
+// Close 关闭
+func Close() error {
+	if instance != nil && instance.cli != nil {
+		return instance.cli.Close()
+	}
+	return nil
+}
+
+func GetClient() *Client {
+	return instance
+}
+
+// Get 获取数据
+func (c *Client) Get(key string, opts ...clientv3.OpOption) (*clientv3.GetResponse, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+	defer cancel()
+	return c.cli.Get(ctx, key, opts...)
+}
+
+// Put 创建数据
+func (c *Client) Put(key, val string, opts ...clientv3.OpOption) (*clientv3.PutResponse, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+	defer cancel()
+	return c.cli.Put(ctx, key, val, opts...)
+}
+
+// Grant 创建租约
+func (c *Client) Grant(ttl int64) (*clientv3.LeaseGrantResponse, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+	defer cancel()
+	return c.cli.Grant(ctx, ttl)
+}
+
+// KeepAlive 保活租约
+func (c *Client) KeepAlive(id clientv3.LeaseID) (<-chan *clientv3.LeaseKeepAliveResponse, error) {
+	return c.cli.KeepAlive(context.Background(), id)
+}
+
+// Revoke 撤销租约
+func (c *Client) Revoke(id clientv3.LeaseID) (*clientv3.LeaseRevokeResponse, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+	defer cancel()
+	return c.cli.Revoke(ctx, id)
+}
+
+// Watch 监听数据
+func (c *Client) Watch(key string, opts ...clientv3.OpOption) clientv3.WatchChan {
+	return c.cli.Watch(context.Background(), key, opts...)
+}
--- a/db/mongo/client.go
+++ b/db/mongo/client.go
@@ -0,0 +1,70 @@
+package mongo
+
+import (
+	"common/config"
+	"common/log"
+	"context"
+	"go.mongodb.org/mongo-driver/mongo"
+	"go.mongodb.org/mongo-driver/mongo/options"
+)
+
+type DBName string
+
+var clients = make(map[DBName]*mongo.Client)
+
+func Init(cfg map[string]*config.MongoConfig) error {
+	for name, oneConfig := range cfg {
+		if client, err := initOneClient(oneConfig); err != nil {
+			return err
+		} else {
+			clients[DBName(name)] = client
+		}
+	}
+	return nil
+}
+
+func initOneClient(cfg *config.MongoConfig) (*mongo.Client, error) {
+	opts := options.Client().
+		ApplyURI(cfg.URI)
+
+	client, err := mongo.Connect(context.Background(), opts)
+	if err != nil {
+		return nil, err
+	}
+
+	return client, nil
+}
+
+// GetClient 返回 mongo.Client，你可以通过 .Database("xxx") 获取具体数据库
+func GetClient(dbName DBName) *mongo.Client {
+	if c, ok := clients[dbName]; ok {
+		return c
+	}
+	log.Errorf("mongo client %s not found", dbName)
+	return nil
+}
+
+// GetDB 是便捷方法，直接返回 *mongo.Database（假设 dbName 对应数据库名）
+// 如果你的配置中 dbName 和实际数据库名一致，可这样用
+func GetDB(dbName DBName) *mongo.Database {
+	client := GetClient(dbName)
+	if client == nil {
+		return nil
+	}
+	// 假设配置中的 key 就是数据库名；若需分离，可在 config 中加字段
+	return client.Database(string(dbName))
+}
+
+func Close() error {
+	if clients == nil {
+		return nil
+	}
+	for name, client := range clients {
+		if client != nil {
+			if err := client.Disconnect(context.Background()); err != nil {
+				log.Errorf("close mongo client %s error: %v", name, err)
+			}
+		}
+	}
+	return nil
+}
--- a/db/mysql/client.go
+++ b/db/mysql/client.go
@@ -0,0 +1,85 @@
+package mysql
+
+import (
+	"common/config"
+	"common/log"
+	"gorm.io/driver/mysql"
+	"gorm.io/gorm"
+	"gorm.io/gorm/logger"
+	"time"
+)
+
+type DBName string
+
+var dbs = make(map[DBName]*gorm.DB)
+
+func Init(cfg map[string]*config.MySQLConfig) error {
+	for name, oneConfig := range cfg {
+		if db, err := initOneDB(oneConfig); err != nil {
+			return err
+		} else {
+			dbs[DBName(name)] = db
+		}
+	}
+	return nil
+}
+
+func initOneDB(cfg *config.MySQLConfig) (*gorm.DB, error) {
+	db, err := gorm.Open(
+		mysql.Open(cfg.Dsn),
+		&gorm.Config{
+			TranslateError: true,
+			Logger:         logger.Default.LogMode(parseLogLevel(cfg.LogLevel)),
+		},
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	sqlDB, err := db.DB()
+	if err != nil {
+		return nil, err
+	}
+	sqlDB.SetMaxOpenConns(int(cfg.MaxOpenConn))
+	sqlDB.SetMaxIdleConns(int(cfg.MaxIdleConn))
+	sqlDB.SetConnMaxLifetime(time.Duration(cfg.ConnMaxLifetimeSec) * time.Second)
+	sqlDB.SetConnMaxIdleTime(time.Duration(cfg.ConnMaxIdleTimeSec) * time.Second)
+
+	return db, nil
+}
+
+func parseLogLevel(logLevel string) logger.LogLevel {
+	switch logLevel {
+	case "silent":
+		return logger.Silent
+	case "error":
+		return logger.Error
+	case "warn":
+		return logger.Warn
+	case "info":
+		return logger.Info
+	}
+	return logger.Info
+}
+
+func GetDB(dbName DBName) *gorm.DB {
+	if q, ok := dbs[dbName]; ok {
+		return q
+	}
+	log.Errorf("db %s not found", dbName)
+	return nil
+}
+
+func Close() error {
+	if dbs == nil {
+		return nil
+	}
+	for name, db := range dbs {
+		if sqlDB, _ := db.DB(); sqlDB != nil {
+			if err := sqlDB.Close(); err != nil {
+				log.Errorf("close db %s error: %v", name, err)
+			}
+		}
+	}
+	return nil
+}
--- a/db/redis/cache_client.go
+++ b/db/redis/cache_client.go
@@ -0,0 +1,50 @@
+package redis
+
+import (
+	"context"
+	"encoding/json"
+	"github.com/panjf2000/gnet/v2/pkg/logging"
+	"github.com/redis/go-redis/v9"
+	"time"
+)
+
+var cacheInstance *CacheClient
+
+type CacheClient struct {
+	cli    *redis.Client
+	logger logging.Logger
+}
+
+func GetCacheClient() *CacheClient {
+	return cacheInstance
+}
+
+func (c *CacheClient) Set(ctx context.Context, key string, value interface{}, expiration time.Duration) {
+	bytes, err := json.Marshal(value)
+	if err != nil {
+		c.logger.Errorf("Set json.Marshal error: %v, key: %v, value: %v", err, key, value)
+		return
+	}
+	if err = c.cli.Set(ctx, key, bytes, expiration).Err(); err != nil {
+		c.logger.Errorf("Set redis.Set error: %v, key: %v, value: %v", err, key, value)
+	}
+}
+
+func (c *CacheClient) Del(ctx context.Context, keys ...string) {
+	if err := c.cli.Del(ctx, keys...).Err(); err != nil {
+		c.logger.Errorf("Set redis.Del error: %v, keys: %v", err, keys)
+	}
+}
+
+// Get 获取数据
+func (c *CacheClient) Get(ctx context.Context, key string, dst interface{}) bool {
+	data, err := c.cli.Get(ctx, key).Bytes()
+	if err != nil {
+		return false
+	}
+	if err = json.Unmarshal(data, dst); err != nil {
+		c.logger.Errorf("Get json.Unmarshal error: %v, key: %v", err, key)
+		return false
+	}
+	return true
+}
--- a/db/redis/client.go
+++ b/db/redis/client.go
@@ -0,0 +1,53 @@
+package redis
+
+import (
+	"common/config"
+	"common/log"
+	"context"
+	"github.com/redis/go-redis/v9"
+	"time"
+)
+
+var instance *Client
+
+type Client struct {
+	cli *redis.Client
+}
+
+func Init(cfg *config.RedisConfig) error {
+	client := redis.NewClient(&redis.Options{
+		Addr:     cfg.Addr,
+		Password: cfg.Password,
+		DB:       cfg.DB,
+	})
+	instance = &Client{
+		cli: client,
+	}
+	cacheInstance = &CacheClient{
+		cli:    client,
+		logger: log.GetLogger().Named("CACHE"),
+	}
+
+	_, err := client.Ping(context.Background()).Result()
+	return err
+}
+
+func Close() error {
+	if instance != nil && instance.cli != nil {
+		return instance.cli.Close()
+	}
+	return nil
+}
+
+func GetClient() *Client {
+	return instance
+}
+
+func (c *Client) Set(ctx context.Context, key string, value interface{}, expiration time.Duration) *redis.StatusCmd {
+	return c.cli.Set(ctx, key, value, expiration)
+}
+
+// Get 获取数据
+func (c *Client) Get(ctx context.Context, key string) *redis.StringCmd {
+	return c.cli.Get(ctx, key)
+}