fix: redis多数据源

2023-09-16 17:54:19 +08:00
parent 3bf1ce62ca
commit 03c73db8c2
3 changed files with 202 additions and 60 deletions
--- a/config/etc/default/restconf.yaml
+++ b/config/etc/default/restconf.yaml
@@ -43,12 +43,23 @@ database:
  name: omc_db
  backup: /usr/local/omc/database

-# Redis 缓存数据
+# Redis 缓存数据，数据源声明全小写
 redis:
-  port: 6379
-  host: 127.0.0.1
-  password: ""
-  db: 0
+  dataSource:
+    # OMC系统使用库
+    default:
+      port: 6379 # Redis port
+      host: "127.0.0.1" # Redis host
+      password: ""
+      db: 10 # Redis db_num
+    # UDM网元用户库
+    udmuser:
+      port: 6379 # Redis port
+      host: "127.0.0.1"
+      password: ""
+      db: 0 # Redis db_num
+  # 多个数据源时可以用这个指定默认的数据源
+  defaultDataSourceName: "default"

 mml:
  port: 4100
--- a/lib/core/redis/redis.go
+++ b/lib/core/redis/redis.go
@@ -12,7 +12,7 @@ import (
 )

 // Redis连接实例
-var rdb *redis.Client
+var rdbMap = make(map[string]*redis.Client)

 // 声明定义限流脚本命令
 var rateLimitCommand = redis.NewScript(`
@@ -32,31 +32,55 @@ return tonumber(current);`)
 // 连接Redis实例
 func Connect() {
 	ctx := context.Background()
-	client := conf.Get("redis").(map[string]any)
-	address := fmt.Sprintf("%s:%d", client["host"], client["port"])
-	// 创建连接
-	rdb = redis.NewClient(&redis.Options{
-		Addr:     address,
-		Password: client["password"].(string),
-		DB:       client["db"].(int),
-	})
-	// 测试数据库连接
-	pong, err := rdb.Ping(ctx).Result()
-	if err != nil {
-		log.Fatalf("failed error ping redis: %v", err)
+	// 读取数据源配置
+	datasource := conf.Get("redis.dataSource").(map[string]any)
+	for k, v := range datasource {
+		client := v.(map[string]any)
+		// 创建连接
+		address := fmt.Sprintf("%s:%d", client["host"], client["port"])
+		rdb := redis.NewClient(&redis.Options{
+			Addr:     address,
+			Password: client["password"].(string),
+			DB:       client["db"].(int),
+		})
+		// 测试数据库连接
+		pong, err := rdb.Ping(ctx).Result()
+		if err != nil {
+			log.Fatalf("failed error ping redis %s is %v", client["host"], err)
+		}
+		log.Infof("redis %s %s connection is successful.", client["host"], pong)
+		rdbMap[k] = rdb
 	}
-	log.Infof("redis %s connection is successful.", pong)
 }

 // 关闭Redis实例
 func Close() {
-	if err := rdb.Close(); err != nil {
-		log.Fatalf("fatal error db close: %s", err)
+	for _, rdb := range rdbMap {
+		if err := rdb.Close(); err != nil {
+			log.Errorf("fatal error db close: %s", err)
+		}
 	}
 }

+// 获取默认实例
+func DefaultRDB() *redis.Client {
+	source := conf.Get("redis.defaultDataSourceName").(string)
+	return rdbMap[source]
+}
+
+// 获取实例
+func RDB(source string) *redis.Client {
+	return rdbMap[source]
+}
+
 // Info 获取redis服务信息
-func Info() map[string]map[string]string {
+func Info(source string) map[string]map[string]string {
+	// 数据源
+	rdb := DefaultRDB()
+	if source != "" {
+		rdb = RDB(source)
+	}
+
 	ctx := context.Background()
 	info, err := rdb.Info(ctx).Result()
 	if err != nil {
@@ -83,7 +107,13 @@ func Info() map[string]map[string]string {
 }

 // KeySize 获取redis当前连接可用键Key总数信息
-func KeySize() int64 {
+func KeySize(source string) int64 {
+	// 数据源
+	rdb := DefaultRDB()
+	if source != "" {
+		rdb = RDB(source)
+	}
+
 	ctx := context.Background()
 	size, err := rdb.DBSize(ctx).Result()
 	if err != nil {
@@ -93,7 +123,13 @@ func KeySize() int64 {
 }

 // CommandStats 获取redis命令状态信息
-func CommandStats() []map[string]string {
+func CommandStats(source string) []map[string]string {
+	// 数据源
+	rdb := DefaultRDB()
+	if source != "" {
+		rdb = RDB(source)
+	}
+
 	ctx := context.Background()
 	commandstats, err := rdb.Info(ctx, "commandstats").Result()
 	if err != nil {
@@ -117,17 +153,29 @@ func CommandStats() []map[string]string {
 }

 // 获取键的剩余有效时间（秒）
-func GetExpire(key string) float64 {
+func GetExpire(source string, key string) (float64, error) {
+	// 数据源
+	rdb := DefaultRDB()
+	if source != "" {
+		rdb = RDB(source)
+	}
+
 	ctx := context.Background()
 	ttl, err := rdb.TTL(ctx, key).Result()
 	if err != nil {
-		return 0
+		return 0, err
 	}
-	return ttl.Seconds()
+	return ttl.Seconds(), nil
 }

 // 获得缓存数据的key列表
-func GetKeys(pattern string) []string {
+func GetKeys(source string, pattern string) ([]string, error) {
+	// 数据源
+	rdb := DefaultRDB()
+	if source != "" {
+		rdb = RDB(source)
+	}
+
 	// 初始化变量
 	var keys []string
 	var cursor uint64 = 0
@@ -138,7 +186,7 @@ func GetKeys(pattern string) []string {
 		batchKeys, nextCursor, err := rdb.Scan(ctx, cursor, pattern, 100).Result()
 		if err != nil {
 			log.Errorf("Failed to scan keys: %v", err)
-			break
+			return keys, err
 		}
 		cursor = nextCursor
 		keys = append(keys, batchKeys...)
@@ -147,91 +195,163 @@ func GetKeys(pattern string) []string {
 			break
 		}
 	}
-	return keys
+	return keys, nil
 }

 // 批量获得缓存数据
-func GetBatch(keys []string) []any {
+func GetBatch(source string, keys []string) ([]any, error) {
 	if len(keys) == 0 {
-		return []any{}
+		return []any{}, fmt.Errorf("not keys")
 	}
+
+	// 数据源
+	rdb := DefaultRDB()
+	if source != "" {
+		rdb = RDB(source)
+	}
+
 	// 获取缓存数据
 	result, err := rdb.MGet(context.Background(), keys...).Result()
 	if err != nil {
 		log.Errorf("Failed to get batch data: %v", err)
-		return []any{}
+		return []any{}, err
 	}
-	return result
+	return result, nil
 }

 // 获得缓存数据
-func Get(key string) string {
+func Get(source, key string) (string, error) {
+	// 数据源
+	rdb := DefaultRDB()
+	if source != "" {
+		rdb = RDB(source)
+	}
+
 	ctx := context.Background()
 	value, err := rdb.Get(ctx, key).Result()
 	if err == redis.Nil || err != nil {
-		return ""
+		return "", err
 	}
-	return value
+	return value, nil
 }

 // 获得缓存数据Hash
-func GetHash(key string) map[string]string {
+func GetHash(source, key string) (map[string]string, error) {
+	// 数据源
+	rdb := DefaultRDB()
+	if source != "" {
+		rdb = RDB(source)
+	}
+
 	ctx := context.Background()
 	value, err := rdb.HGetAll(ctx, key).Result()
 	if err == redis.Nil || err != nil {
-		return map[string]string{}
+		return map[string]string{}, err
 	}
-	return value
+	return value, nil
 }

 // 判断是否存在
-func Has(keys ...string) bool {
+func Has(source string, keys ...string) (bool, error) {
+	// 数据源
+	rdb := DefaultRDB()
+	if source != "" {
+		rdb = RDB(source)
+	}
+
 	ctx := context.Background()
 	exists, err := rdb.Exists(ctx, keys...).Result()
 	if err != nil {
-		return false
+		return false, err
 	}
-	return exists >= 1
+	return exists >= 1, nil
 }

 // 设置缓存数据
-func Set(key string, value any) bool {
+func Set(source, key string, value any) (bool, error) {
+	// 数据源
+	rdb := DefaultRDB()
+	if source != "" {
+		rdb = RDB(source)
+	}
+
 	ctx := context.Background()
 	err := rdb.Set(ctx, key, value, 0).Err()
-	return err == nil
+	if err != nil {
+		log.Errorf("redis lua script err %v", err)
+		return false, err
+	}
+	return true, nil
 }

 // 设置缓存数据与过期时间
-func SetByExpire(key string, value any, expiration time.Duration) bool {
+func SetByExpire(source, key string, value any, expiration time.Duration) (bool, error) {
+	// 数据源
+	rdb := DefaultRDB()
+	if source != "" {
+		rdb = RDB(source)
+	}
+
 	ctx := context.Background()
 	err := rdb.Set(ctx, key, value, expiration).Err()
-	return err == nil
+	if err != nil {
+		log.Errorf("redis lua script err %v", err)
+		return false, err
+	}
+	return true, nil
 }

 // 删除单个
-func Del(key string) bool {
+func Del(source string, key string) (bool, error) {
+	// 数据源
+	rdb := DefaultRDB()
+	if source != "" {
+		rdb = RDB(source)
+	}
+
 	ctx := context.Background()
 	err := rdb.Del(ctx, key).Err()
-	return err == nil
+	if err != nil {
+		log.Errorf("redis lua script err %v", err)
+		return false, err
+	}
+	return true, nil
 }

 // 删除多个
-func DelKeys(keys []string) bool {
+func DelKeys(source string, keys []string) (bool, error) {
 	if len(keys) == 0 {
-		return false
+		return false, fmt.Errorf("no keys")
 	}
+
+	// 数据源
+	rdb := DefaultRDB()
+	if source != "" {
+		rdb = RDB(source)
+	}
+
 	ctx := context.Background()
 	err := rdb.Del(ctx, keys...).Err()
-	return err == nil
+	if err != nil {
+		log.Errorf("redis lua script err %v", err)
+		return false, err
+	}
+	return true, nil
 }

 // 限流查询并记录
-func RateLimit(limitKey string, time, count int64) int64 {
+func RateLimit(source, limitKey string, time, count int64) (int64, error) {
+	// 数据源
+	rdb := DefaultRDB()
+	if source != "" {
+		rdb = RDB(source)
+	}
+
 	ctx := context.Background()
 	result, err := rateLimitCommand.Run(ctx, rdb, []string{limitKey}, time, count).Result()
 	if err != nil {
 		log.Errorf("redis lua script err %v", err)
-		return 0
+		return 0, err
 	}
-	return result.(int64)
+	return result.(int64), err
 }
--- a/restagent/etc/restconf.yaml
+++ b/restagent/etc/restconf.yaml
@@ -43,12 +43,23 @@ database:
  name: omc_db
  backup: d:/local.git/ems.agt/restagent/database

-# Redis 缓存数据
+# Redis 缓存数据，数据源声明全小写
 redis:
-  port: 6379
-  host: "192.168.4.228"
-  password: ""
-  db: 0
+  dataSource:
+    # OMC系统使用库
+    default:
+      port: 6379 # Redis port
+      host: "192.168.0.229" # Redis host
+      password: ""
+      db: 10 # Redis db_num
+    # UDM网元用户库
+    udmuser:
+      port: 6379 # Redis port
+      host: "192.168.0.229"
+      password: ""
+      db: 0 # Redis db_num
+  # 多个数据源时可以用这个指定默认的数据源
+  defaultDataSourceName: "default"

 mml:
  port: 4100