diff --git a/backend/modules/observability/domain/component/config/config.go b/backend/modules/observability/domain/component/config/config.go index c85000e4d..701c8c5bf 100644 --- a/backend/modules/observability/domain/component/config/config.go +++ b/backend/modules/observability/domain/component/config/config.go @@ -5,6 +5,7 @@ package config import ( "context" + "slices" "github.com/coze-dev/coze-loop/backend/modules/observability/domain/trace/entity/loop_span" "github.com/coze-dev/coze-loop/backend/pkg/conf" @@ -249,6 +250,8 @@ type TrajectoryMetadataConfig struct { // Spaces 按 workspace_id 配置允许写入的 metadata key 规则列表 Spaces map[int64][]loop_span.MetaKeyRule `mapstructure:"spaces" json:"spaces"` EnableSingleQuery bool `mapstructure:"enable_single_query" json:"enable_single_query"` + // EnableDoubleQuerySpaces 白名单:命中的 workspace_id 强制走双查询,即使全局开启了单查询 + EnableDoubleQuerySpaces []int64 `mapstructure:"enable_double_query_spaces" json:"enable_double_query_spaces"` } func (c *TrajectoryMetadataConfig) IsSingleQueryEnabled() bool { @@ -258,6 +261,13 @@ func (c *TrajectoryMetadataConfig) IsSingleQueryEnabled() bool { return c.EnableSingleQuery } +func (c *TrajectoryMetadataConfig) IsDoubleQueryEnabled(workspaceID int64) bool { + if c == nil { + return false + } + return slices.Contains(c.EnableDoubleQuerySpaces, workspaceID) +} + //go:generate mockgen -destination=mocks/config.go -package=mocks . ITraceConfig type ITraceConfig interface { GetSystemViews(ctx context.Context) ([]*SystemView, error) diff --git a/backend/modules/observability/domain/trace/service/trace_service.go b/backend/modules/observability/domain/trace/service/trace_service.go index 1862379e6..a1d69f4a4 100644 --- a/backend/modules/observability/domain/trace/service/trace_service.go +++ b/backend/modules/observability/domain/trace/service/trace_service.go @@ -2594,7 +2594,7 @@ func (r *TraceServiceImpl) GetTrajectories(ctx context.Context, workspaceID int6 maxBytes = backfillCfg.GetTrajectoryMaxBytes(workspaceID) } - if metaCfg.IsSingleQueryEnabled() { + if metaCfg.IsSingleQueryEnabled() && !metaCfg.IsDoubleQueryEnabled(workspaceID) { return r.getTrajectoriesSingleQuery(ctx, workspaceID, tenant, traceIDs, startTime, endTime, platformType, trajectoryConfig, metaRules, maxBytes) } diff --git a/backend/modules/observability/domain/trace/service/trace_trajectory_service_test.go b/backend/modules/observability/domain/trace/service/trace_trajectory_service_test.go index e1f20bb20..8eae775bd 100644 --- a/backend/modules/observability/domain/trace/service/trace_trajectory_service_test.go +++ b/backend/modules/observability/domain/trace/service/trace_trajectory_service_test.go @@ -206,6 +206,18 @@ func TestTrajectoryMetadataConfig_IsSingleQueryEnabled(t *testing.T) { assert.True(t, (&config.TrajectoryMetadataConfig{EnableSingleQuery: true}).IsSingleQueryEnabled()) } +func TestTrajectoryMetadataConfig_IsDoubleQueryEnabled(t *testing.T) { + var nilCfg *config.TrajectoryMetadataConfig + assert.False(t, nilCfg.IsDoubleQueryEnabled(123)) + + assert.False(t, (&config.TrajectoryMetadataConfig{}).IsDoubleQueryEnabled(123)) + + cfg := &config.TrajectoryMetadataConfig{EnableDoubleQuerySpaces: []int64{123, 456}} + assert.True(t, cfg.IsDoubleQueryEnabled(123)) + assert.True(t, cfg.IsDoubleQueryEnabled(456)) + assert.False(t, cfg.IsDoubleQueryEnabled(789)) +} + func TestTraceServiceImpl_GetTrajectories_SingleQuery_MaxBytesExceeded(t *testing.T) { ctrl := gomock.NewController(t) defer ctrl.Finish()