Bläddra i källkod

refactor(sharding): 重构阈值配置管理

- 移除 logRepository 和 wafLogRepository 中的 getMaxRowsForTable 方法
- 在 ShardingManager 中添加 GetMaxRowsForTable 方法,统一管理表的阈值配置
- 更新 ThresholdConfig 结构,增加 TableThresholds 字段用于配置不同表的阈值
- 调整 NewShardingManager 方法,设置默认阈值和特定表的阈值
fusu 17 timmar sedan
förälder
incheckning
303f6d24a4

+ 2 - 13
internal/repository/admin/waflog.go

@@ -273,7 +273,7 @@ func (r *wafLogRepository) AddWafLog(ctx context.Context, log *model.WafLog) err
 
 	
 	// 获取最优的写入表(考虑数据量阈值)
-	tableName, err := r.Manager.GetOptimalWriteTable(ctx, r.DBWithName(ctx, "admin"), log, r.getMaxRowsForTable("waf_log"))
+	tableName, err := r.Manager.GetOptimalWriteTable(ctx, r.DBWithName(ctx, "admin"), log, r.Manager.GetMaxRowsForTable("waf_log"))
 	if err != nil {
 		return fmt.Errorf("获取写入表失败: %v", err)
 	}
@@ -295,7 +295,7 @@ func (r *wafLogRepository) BatchAddWafLog(ctx context.Context, logs []*model.Waf
 		return nil
 	}
 
-	maxRows := r.getMaxRowsForTable("waf_log")
+	maxRows := r.Manager.GetMaxRowsForTable("waf_log")
 	
 	// 按表名分组
 	tableGroups := make(map[string][]*model.WafLog)
@@ -391,17 +391,6 @@ func (r *wafLogRepository) GetWafLogExportCount(ctx context.Context, req adminAp
 }
 
 
-// getMaxRowsForTable 获取指定表的最大行数配置
-func (r *wafLogRepository) getMaxRowsForTable(tableName string) int64 {
-	switch tableName {
-	case "log":
-		return 3000000 // 300万条
-	case "waf_log":
-		return 5000000 // 500万条
-	default:
-		return 3000000 // 默认300万条
-	}
-}
 
 // applyWafLogFilters 应用WafLog查询过滤条件
 func (r *wafLogRepository) applyWafLogFilters(query *gorm.DB, req adminApi.SearchWafLogParams) *gorm.DB {

+ 1 - 14
internal/repository/log.go

@@ -95,7 +95,7 @@ func (r *logRepository) AddLog(ctx context.Context, log *model.Log) error {
 
 	
 	// 获取最优的写入表(考虑数据量阈值)
-	tableName, err := r.Manager.GetOptimalWriteTable(ctx, r.DBWithName(ctx, "admin"), log, r.getMaxRowsForTable("log"))
+	tableName, err := r.Manager.GetOptimalWriteTable(ctx, r.DBWithName(ctx, "admin"), log, r.Manager.GetMaxRowsForTable("log"))
 	if err != nil {
 		return fmt.Errorf("获取写入表失败: %v", err)
 	}
@@ -131,16 +131,3 @@ func (r *logRepository) EditLog(ctx context.Context, log *model.Log) error {
 	
 	return r.DBWithName(ctx, "admin").Table(tableName).Updates(log).Error
 }
-
-
-// getMaxRowsForTable 获取指定表的最大行数配置
-func (r *logRepository) getMaxRowsForTable(tableName string) int64 {
-	switch tableName {
-	case "log":
-		return 3000000 // 300万条
-	case "waf_log":
-		return 5000000 // 500万条
-	default:
-		return 3000000 // 默认300万条
-	}
-}

+ 6 - 2
internal/repository/repository.go

@@ -417,10 +417,14 @@ m = g(r.sub, p.sub) && r.obj == p.obj && r.act == p.act
 func NewShardingManager(logger *log.Logger) *sharding.ShardingManager {
 	strategy := sharding.NewMonthlyShardingStrategy()
 	
-	// 配置阈值参数
+	// 配置阈值参数 - 统一管理所有表的阈值
 	thresholdConfig := &sharding.ThresholdConfig{
 		Enabled: true,
-		MaxRows: 5000000, // waf_log表500万条阈值
+		MaxRows: 3000000, // 默认300万条
+		TableThresholds: map[string]int64{
+			"log":     3000000, // log表300万条
+			"waf_log": 5000000, // waf_log表500万条
+		},
 	}
 	
 	return sharding.NewShardingManagerWithThreshold(strategy, logger, thresholdConfig)

+ 20 - 0
pkg/sharding/manager.go

@@ -23,6 +23,8 @@ type ThresholdConfig struct {
 	Enabled       bool
 	MaxRows       int64
 	CheckInterval time.Duration
+	// 不同表的阈值配置
+	TableThresholds map[string]int64
 }
 
 // ShardingManager 分表管理器
@@ -264,4 +266,22 @@ func (sm *ShardingManager) CheckAndCreateNewTable(ctx context.Context, db *gorm.
 	}
 	
 	return nil
+}
+
+// GetMaxRowsForTable 获取指定表的最大行数配置
+func (sm *ShardingManager) GetMaxRowsForTable(tableName string) int64 {
+	// 优先使用表级配置
+	if sm.thresholdConfig != nil && sm.thresholdConfig.TableThresholds != nil {
+		if maxRows, exists := sm.thresholdConfig.TableThresholds[tableName]; exists {
+			return maxRows
+		}
+	}
+	
+	// 使用默认配置
+	if sm.thresholdConfig != nil {
+		return sm.thresholdConfig.MaxRows
+	}
+	
+	// 最终默认值
+	return 3000000
 }