Compare commits

...

16 Commits

Author SHA1 Message Date
nityanandagohain
f6a823c000 fix: remove unnecessary function 2026-08-06 15:46:50 +05:30
nityanandagohain
23f03973c5 fix: more cleanup 2026-08-06 12:23:06 +05:30
nityanandagohain
0c0e969cfc fix: more cleanup 2026-08-06 00:06:41 +05:30
nityanandagohain
c870efa12d Merge remote-tracking branch 'origin/main' into issue_5602 2026-08-05 11:30:36 +05:30
nityanandagohain
d250f190a7 Squashed commit of the following:
commit 6d0b607f41b80583b2f4dc015c728d57fd0a3f6a
Author: nityanandagohain <nityanandagohain@gmail.com>
Date:   Wed Jul 15 10:30:43 2026 +0530

    fix: update tests

commit 20f42474b0901e49cc745155a2c70bb829988a94
Merge: 36dca5dab0 31efe177a4
Author: nityanandagohain <nityanandagohain@gmail.com>
Date:   Tue Jul 14 23:09:00 2026 +0530

    Merge branch 'issue_5601': scoped-trace builder package + port aggregations

    Merges the telemetryscopedtraces refactor (generic scoped-trace topology moved
    out of telemetryai, feature-flagged gen_ai key enrichment, tracefield. context
    classification) and ports the scalar/time-series work onto it:

    - trace_aggregation.go / trace_having.go now live in telemetryscopedtraces
    - trace-level HAVING resolution composes the canonical pkg/variables
      replacement (dynamic __all__ drops a condition for any operator) with the
      standard filter pipeline, so values still bind as args
    - buildPerTraceScan adopts the merged embedExpr placeholder/arg checking
    - ColumnProvider gains ActivityGateAlias; gen_ai provider gates trace-level
      aggregations on llm_call_count

commit 36dca5dab0c0abb08aeb2645fee07f28cc764eb8
Author: nityanandagohain <nityanandagohain@gmail.com>
Date:   Tue Jul 14 23:09:00 2026 +0530

    feat: trace-level aggregations for scalar/time-series and trace-scoped span list

    - trace.-prefixed aggregations (avg(trace.output_tokens), count(trace.trace_id),
      arithmetic) run over window-clipped per-trace values via a native CTE pipeline
    - a trace-level filter condition qualifies traces in scalar/time-series/raw via
      __qualified / __trace_scope, resource-fingerprint pruned like the trace list
    - trace-level filter conditions resolve through the standard filter pipeline
      (bound args, operators, query variables) against the per-trace column aliases
    - per-trace rows with no LLM span in the window are dropped from trace-level
      aggregations (LLM-activity gate)
    - targeted rejections for group-by/order-by on trace-level columns
2026-07-15 12:38:57 +05:30
nityanandagohain
97c49c870b feat: support ai trace aggregate filtering in ai span list 2026-07-15 12:27:55 +05:30
nityanandagohain
31efe177a4 fix: address comments 2026-07-14 18:53:30 +05:30
nityanandagohain
d502d12ac3 fix: update openapi 2026-07-10 14:27:07 +05:30
nityanandagohain
bd9f15a716 fix: update integration test 2026-07-10 14:21:16 +05:30
nityanandagohain
813ef988c9 fix: edge cases and correct cost key 2026-07-10 12:06:34 +05:30
nityanandagohain
40e6799285 fix: add resource fingerprint cte 2026-07-10 00:36:25 +05:30
nityanandagohain
1caa60a3cd fix: cleanup and more tests 2026-07-09 23:54:05 +05:30
nityanandagohain
3f781f0083 fix: more cleanup 2026-07-09 12:39:01 +05:30
nityanandagohain
6aec05cf7a fix: more tests 2026-07-09 08:45:22 +05:30
nityanandagohain
683a52f35a fix: take perf into consideration 2026-07-09 08:45:22 +05:30
nityanandagohain
e924fa1e62 feat: support llm trace list and span list 2026-07-09 08:45:20 +05:30
14 changed files with 2054 additions and 110 deletions

View File

@@ -74,5 +74,9 @@ func Scope() scopedtraces.TraceScope {
FieldKeys: gateKeys,
Columns: columns,
DefaultOrderAlias: "last_activity_time",
// A trace whose only in-scope spans are tool/agent spans has NULL output_tokens,
// so avg(trace.output_tokens) skips it while count(trace.trace_id) counts it.
// Gating on llm_call_count > 0 keeps both looking at the same trace set.
ActivityGateAlias: "llm_call_count",
}
}

View File

@@ -993,8 +993,8 @@ func TestBuild_TraceList_MultiVariantGateKey(t *testing.T) {
assert.Contains(t, got, "mapContains(attributes_string, 'gen_ai.tool.name') OR mapContains(attributes_number, 'gen_ai.tool.name')")
}
// `trace.` marks a trace-level aggregate; `tracefield.` routes trace-level too but is
// not a rewritable alias, so the HAVING rewriter rejects it.
// A `trace.`-prefixed aggregate in the filter box and the same condition in the
// explicit Having box build the same query; output-only aggregates are rejected.
func TestBuild_TraceList_TraceContextPrefix(t *testing.T) {
b := newTestBuilder(t)
build := func(q qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]) (*qbtypes.Statement, error) {
@@ -1002,19 +1002,14 @@ func TestBuild_TraceList_TraceContextPrefix(t *testing.T) {
return b.Build(context.Background(), valuer.UUID{}, testStartMs, testEndMs, qbtypes.RequestTypeTrace, q, nil)
}
_, err := build(qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
viaTrace, err := build(qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Filter: &qbtypes.Filter{Expression: "trace.output_tokens > 1000"}})
require.NoError(t, err)
_, err = build(qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Filter: &qbtypes.Filter{Expression: "tracefield.output_tokens > 1000"}})
require.Error(t, err)
assert.Contains(t, err.Error(), "Invalid references in `Having` expression: [tracefield.output_tokens]")
_, err = build(qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Having: &qbtypes.Having{Expression: "tracefield.output_tokens > 1000"}})
require.Error(t, err)
assert.Contains(t, err.Error(), "Invalid references in `Having` expression: [tracefield.output_tokens]")
viaHaving, err := build(qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Having: &qbtypes.Having{Expression: "trace.output_tokens > 1000"}})
require.NoError(t, err)
assert.Equal(t, viaTrace.Query, viaHaving.Query)
_, err = build(qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Filter: &qbtypes.Filter{Expression: "trace.span_count > 3"}})
@@ -1022,7 +1017,8 @@ func TestBuild_TraceList_TraceContextPrefix(t *testing.T) {
assert.Contains(t, err.Error(), "cannot be used")
}
// Query variables in a trace-level condition are substituted into the HAVING.
// Query variables in a trace-level condition resolve like span filters: bound args,
// list/IN handling, dynamic __all__ dropping the condition.
func TestBuild_TraceList_VariableInAggregateFilter(t *testing.T) {
b := newTestBuilder(t)
build := func(expr string, vars map[string]qbtypes.VariableItem) (*qbtypes.Statement, error) {
@@ -1034,17 +1030,18 @@ func TestBuild_TraceList_VariableInAggregateFilter(t *testing.T) {
}, vars)
}
// scalar variable -> literal in HAVING
// scalar variable -> bound arg via the filter pipeline
stmt, err := build("trace.output_tokens > $threshold",
map[string]qbtypes.VariableItem{"threshold": {Value: 700}})
require.NoError(t, err)
assert.Contains(t, stmt.Query, "HAVING output_tokens > 700")
assert.Contains(t, stmt.Query, "HAVING output_tokens > ?")
assert.Contains(t, stmt.Args, float64(700))
// list variable with IN
stmt, err = build("trace.llm_call_count IN $counts",
map[string]qbtypes.VariableItem{"counts": {Value: []any{1, 2}}})
require.NoError(t, err)
assert.Contains(t, stmt.Query, "HAVING llm_call_count IN")
assert.Contains(t, stmt.Query, "HAVING llm_call_count IN (?, ?)")
// dynamic __all__ -> condition dropped, no HAVING at all
stmt, err = build("trace.output_tokens > $threshold",
@@ -1052,7 +1049,7 @@ func TestBuild_TraceList_VariableInAggregateFilter(t *testing.T) {
require.NoError(t, err)
assert.NotContains(t, stmt.Query, "HAVING")
// unresolved variable -> rejected, not compared as a literal
// unresolved variable -> rejected, though only as an unknown aggregate today
_, err = build("trace.output_tokens > $missing", map[string]qbtypes.VariableItem{"other": {Value: 1}})
require.Error(t, err)
}

View File

@@ -0,0 +1,227 @@
package aistatementbuilder
import (
"context"
"strings"
"testing"
"time"
qbtypes "github.com/SigNoz/signoz/pkg/types/querybuildertypes/querybuildertypesv5"
"github.com/SigNoz/signoz/pkg/types/telemetrytypes"
"github.com/SigNoz/signoz/pkg/valuer"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// Build tests for scalar / time-series through the gen_ai scope; the
// rewriteTraceAggregation unit tests live in scopedtracesstatementbuilder.
// Mixing span- and trace-level aggregations across one query is rejected.
func TestBuild_Aggregation_MixedDomainsRejected(t *testing.T) {
b := newTestBuilder(t)
_, err := b.Build(context.Background(), valuer.UUID{}, testStartMs, testEndMs, qbtypes.RequestTypeScalar,
qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Signal: telemetrytypes.SignalTraces,
Aggregations: []qbtypes.TraceAggregation{
{Expression: "avg(trace.output_tokens)"},
{Expression: "sum(gen_ai.usage.output_tokens)"},
},
}, nil)
require.ErrorContains(t, err, "cannot be mixed")
}
// Output-only aggregates are rejected in trace-level filters on the aggregation
// path too (the raw and trace-list paths are covered elsewhere).
func TestBuild_Aggregation_OutputOnlyFilterRejected(t *testing.T) {
b := newTestBuilder(t)
_, err := b.Build(context.Background(), valuer.UUID{}, testStartMs, testEndMs, qbtypes.RequestTypeScalar,
qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Signal: telemetrytypes.SignalTraces,
Aggregations: []qbtypes.TraceAggregation{{Expression: "count()"}},
Filter: &qbtypes.Filter{Expression: "trace.span_count > 3"},
}, nil)
require.ErrorContains(t, err, `aggregate "span_count" cannot be used`)
}
// Trace-level columns are rejected as group-by / order keys with a targeted error;
// ordering by the aggregation's own alias stays valid.
func TestBuild_Aggregation_GroupByOrderValidation(t *testing.T) {
b := newTestBuilder(t)
ctx := context.Background()
_, err := b.Build(ctx, valuer.UUID{}, testStartMs, testEndMs, qbtypes.RequestTypeScalar,
qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Signal: telemetrytypes.SignalTraces,
Aggregations: []qbtypes.TraceAggregation{{Expression: "avg(trace.output_tokens)"}},
GroupBy: []qbtypes.GroupByKey{{TelemetryFieldKey: telemetrytypes.TelemetryFieldKey{Name: "trace.llm_call_count"}}},
}, nil)
require.ErrorContains(t, err, `grouping by trace-level aggregate "trace.llm_call_count" is not supported`)
_, err = b.Build(ctx, valuer.UUID{}, testStartMs, testEndMs, qbtypes.RequestTypeScalar,
qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Signal: telemetrytypes.SignalTraces,
Aggregations: []qbtypes.TraceAggregation{{Expression: "avg(trace.output_tokens)"}},
Order: []qbtypes.OrderBy{{Key: qbtypes.OrderByKey{TelemetryFieldKey: telemetrytypes.TelemetryFieldKey{Name: "trace.total_tokens"}}}},
}, nil)
require.ErrorContains(t, err, `ordering by trace-level aggregate "trace.total_tokens" is not supported`)
_, err = b.Build(ctx, valuer.UUID{}, testStartMs, testEndMs, qbtypes.RequestTypeScalar,
qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Signal: telemetrytypes.SignalTraces,
Aggregations: []qbtypes.TraceAggregation{{Expression: "avg(trace.output_tokens)", Alias: "avg_out"}},
Order: []qbtypes.OrderBy{{Key: qbtypes.OrderByKey{TelemetryFieldKey: telemetrytypes.TelemetryFieldKey{Name: "avg_out"}}, Direction: qbtypes.OrderDirectionAsc}},
}, nil)
require.NoError(t, err)
}
// Variables in trace-level conditions bind as args on the aggregation path;
// unknown $vars fail with a variable error, __all__ drops the condition.
func TestBuild_Aggregation_VariablesInTraceFilter(t *testing.T) {
b := newTestBuilder(t)
ctx := context.Background()
q := qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Signal: telemetrytypes.SignalTraces,
Aggregations: []qbtypes.TraceAggregation{{Expression: "avg(trace.output_tokens)"}},
Filter: &qbtypes.Filter{Expression: "trace.output_tokens > $threshold"},
}
stmt, err := b.Build(ctx, valuer.UUID{}, testStartMs, testEndMs, qbtypes.RequestTypeScalar, q,
map[string]qbtypes.VariableItem{"threshold": {Type: qbtypes.TextBoxVariableType, Value: float64(1000)}})
require.NoError(t, err)
assert.Contains(t, stmt.Query, "HAVING output_tokens > ?")
assert.Contains(t, stmt.Args, float64(1000))
// an unresolved $var is only rejected as an unknown aggregate today; a targeted
// "unknown variable" error is a separate concern
_, err = b.Build(ctx, valuer.UUID{}, testStartMs, testEndMs, qbtypes.RequestTypeScalar, q, nil)
require.ErrorContains(t, err, `aggregate "$threshold" cannot be used`)
stmt, err = b.Build(ctx, valuer.UUID{}, testStartMs, testEndMs, qbtypes.RequestTypeScalar, q,
map[string]qbtypes.VariableItem{"threshold": {Type: qbtypes.DynamicVariableType, Value: "__all__"}})
require.NoError(t, err)
assert.NotContains(t, stmt.Query, "__qualified")
// list variables render as IN with bound args
stmt, err = b.Build(ctx, valuer.UUID{}, testStartMs, testEndMs, qbtypes.RequestTypeScalar,
qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Signal: telemetrytypes.SignalTraces,
Aggregations: []qbtypes.TraceAggregation{{Expression: "count(trace.trace_id)"}},
Filter: &qbtypes.Filter{Expression: "trace.llm_call_count IN $counts"},
}, map[string]qbtypes.VariableItem{
"counts": {Type: qbtypes.QueryVariableType, Value: []any{float64(1), float64(2)}},
})
require.NoError(t, err)
assert.Contains(t, stmt.Query, "HAVING llm_call_count IN (?, ?)")
}
// Resource conditions prune the qualification scan: __qualified references the
// __resource_filter CTE, the delegated __trace_scope inlines the fingerprint subquery.
func TestBuild_Aggregation_QualificationResourcePruned(t *testing.T) {
b := newTestBuilder(t)
ctx := context.Background()
filter := &qbtypes.Filter{Expression: "service.name = 'api' AND trace.output_tokens > 1000"}
stmt, err := b.Build(ctx, valuer.UUID{}, testStartMs, testEndMs, qbtypes.RequestTypeScalar,
qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Signal: telemetrytypes.SignalTraces,
Aggregations: []qbtypes.TraceAggregation{{Expression: "avg(trace.output_tokens)"}},
Filter: filter,
}, nil)
require.NoError(t, err)
qualified := stmt.Query[strings.Index(stmt.Query, "__qualified"):strings.Index(stmt.Query, "__scoped_traces")]
assert.Contains(t, qualified, "resource_fingerprint GLOBAL IN (SELECT fingerprint FROM __resource_filter)")
stmt, err = b.Build(ctx, valuer.UUID{}, testStartMs, testEndMs, qbtypes.RequestTypeRaw,
qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Signal: telemetrytypes.SignalTraces,
Filter: filter,
Limit: 10,
}, nil)
require.NoError(t, err)
scope := stmt.Query[strings.Index(stmt.Query, "__trace_scope"):]
assert.Contains(t, scope, "resource_fingerprint GLOBAL IN (SELECT fingerprint FROM (SELECT")
}
// ---------------------------------------------------------------------------
// Full-query goldens — native trace-domain pipeline
// ---------------------------------------------------------------------------
// Scalar over per-trace values: one window-clipped per-trace scan, outer avg.
func TestBuild_FullSQL_Scalar_TraceAgg(t *testing.T) {
b := newTestBuilder(t)
stmt, err := b.Build(context.Background(), valuer.UUID{}, testStartMs, testEndMs, qbtypes.RequestTypeScalar,
qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Signal: telemetrytypes.SignalTraces,
Aggregations: []qbtypes.TraceAggregation{{Expression: "avg(trace.output_tokens)"}},
}, nil)
require.NoError(t, err)
assertSQLEqual(t, `
WITH __scoped_traces AS (
SELECT trace_id,
sum(multiIf(mapContains(attributes_number, 'gen_ai.usage.output_tokens'), toFloat64(attributes_number['gen_ai.usage.output_tokens']), NULL)) AS output_tokens
FROM signoz_traces.distributed_signoz_index_v3
WHERE timestamp >= '1747947419000000000'
AND timestamp < '1747983448000000000'
AND ts_bucket_start >= 1747945619
AND ts_bucket_start <= 1747983448
AND (mapContains(attributes_string, 'gen_ai.request.model') OR mapContains(attributes_string, 'gen_ai.tool.name') OR mapContains(attributes_string, 'gen_ai.agent.name'))
GROUP BY trace_id
HAVING (countIf(mapContains(attributes_string, 'gen_ai.request.model'))) > 0
)
SELECT avg(output_tokens) AS __result_0
FROM __scoped_traces
ORDER BY __result_0 DESC
SETTINGS distributed_product_mode='allow', max_memory_usage=10000000000
`, stmt)
}
// Time series: the per-trace scan buckets by span time, the outer aggregation per bucket.
func TestBuild_FullSQL_TimeSeries_TraceAgg(t *testing.T) {
b := newTestBuilder(t)
stmt, err := b.Build(context.Background(), valuer.UUID{}, testStartMs, testEndMs, qbtypes.RequestTypeTimeSeries,
qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Signal: telemetrytypes.SignalTraces,
StepInterval: qbtypes.Step{Duration: 60 * time.Second},
Aggregations: []qbtypes.TraceAggregation{{Expression: "avg(trace.output_tokens)"}},
}, nil)
require.NoError(t, err)
assertSQLEqual(t, `
WITH __scoped_traces AS (
SELECT trace_id,
toStartOfInterval(timestamp, INTERVAL 60 SECOND) AS ts,
sum(multiIf(mapContains(attributes_number, 'gen_ai.usage.output_tokens'), toFloat64(attributes_number['gen_ai.usage.output_tokens']), NULL)) AS output_tokens
FROM signoz_traces.distributed_signoz_index_v3
WHERE timestamp >= '1747947419000000000'
AND timestamp < '1747983448000000000'
AND ts_bucket_start >= 1747945619
AND ts_bucket_start <= 1747983448
AND (mapContains(attributes_string, 'gen_ai.request.model') OR mapContains(attributes_string, 'gen_ai.tool.name') OR mapContains(attributes_string, 'gen_ai.agent.name'))
GROUP BY trace_id, ts
HAVING (countIf(mapContains(attributes_string, 'gen_ai.request.model'))) > 0
)
SELECT ts, avg(output_tokens) AS __result_0
FROM __scoped_traces
GROUP BY ts
SETTINGS distributed_product_mode='allow', max_memory_usage=10000000000
`, stmt)
}
// Span-level scalar with a trace-level filter: delegated, constrained by __trace_scope.
func TestBuild_FullSQL_Scalar_SpanAgg_TraceScoped(t *testing.T) {
b := newTestBuilder(t)
stmt, err := b.Build(context.Background(), valuer.UUID{}, testStartMs, testEndMs, qbtypes.RequestTypeScalar,
qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Signal: telemetrytypes.SignalTraces,
Aggregations: []qbtypes.TraceAggregation{{Expression: "sum(gen_ai.usage.output_tokens)"}},
Filter: &qbtypes.Filter{Expression: "trace.output_tokens > 1000"},
}, nil)
require.NoError(t, err)
got := renderSQL(t, stmt)
assert.Contains(t, got, "__trace_scope AS (")
assert.Contains(t, got, "HAVING output_tokens > 1000")
assert.Contains(t, got, "trace_id GLOBAL IN (SELECT trace_id FROM __trace_scope)")
assert.Contains(t, got, "AS __result_0")
}

View File

@@ -0,0 +1,103 @@
package aistatementbuilder
import (
"context"
"testing"
qbtypes "github.com/SigNoz/signoz/pkg/types/querybuildertypes/querybuildertypesv5"
"github.com/SigNoz/signoz/pkg/types/telemetrytypes"
"github.com/SigNoz/signoz/pkg/valuer"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// Span list with a mixed filter: spans matching the span-level part, in traces whose
// window-clipped aggregates satisfy the trace-level part (__trace_scope).
func TestBuild_SpanList_TraceScoped(t *testing.T) {
b := newTestBuilder(t)
stmt, err := b.Build(context.Background(), valuer.UUID{}, testStartMs, testEndMs, qbtypes.RequestTypeRaw,
qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Signal: telemetrytypes.SignalTraces,
Filter: &qbtypes.Filter{Expression: "gen_ai.request.model = 'gpt-4o-mini' AND trace.output_tokens > 1000"},
Limit: 10,
}, nil)
require.NoError(t, err)
got := renderSQL(t, stmt)
assert.Contains(t, got, "__trace_scope AS (")
assert.Contains(t, got, "HAVING output_tokens > 1000")
assert.Contains(t, got, "trace_id GLOBAL IN (SELECT trace_id FROM __trace_scope)")
assert.Contains(t, got, "gpt-4o-mini")
}
// Without a trace-level condition nothing changes: the span list stays a single
// gated span scan (no __trace_scope CTE).
func TestBuild_SpanList_NoTraceFilter_NoScope(t *testing.T) {
b := newTestBuilder(t)
stmt, err := b.Build(context.Background(), valuer.UUID{}, testStartMs, testEndMs, qbtypes.RequestTypeRaw,
qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Signal: telemetrytypes.SignalTraces,
Filter: &qbtypes.Filter{Expression: "gen_ai.request.model = 'gpt-4o-mini'"},
Limit: 10,
}, nil)
require.NoError(t, err)
assert.NotContains(t, stmt.Query, "__trace_scope")
}
// The span-list trace-level filter shares the trace list's rules: output-only
// aggregates, OR-mixing, and trace-level order keys are rejected — while bare span
// columns sharing an aggregate alias name (duration_nano) stay orderable.
func TestBuild_SpanList_TraceFilter_Validation(t *testing.T) {
b := newTestBuilder(t)
build := func(q qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]) error {
q.Signal = telemetrytypes.SignalTraces
_, err := b.Build(context.Background(), valuer.UUID{}, testStartMs, testEndMs, qbtypes.RequestTypeRaw, q, nil)
return err
}
err := build(qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Filter: &qbtypes.Filter{Expression: "trace.span_count > 3"},
})
require.ErrorContains(t, err, `aggregate "span_count" cannot be used`)
err = build(qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Filter: &qbtypes.Filter{Expression: "trace.output_tokens > 1000 OR kind_string = 'Client'"},
})
require.ErrorContains(t, err, "cannot be combined")
err = build(qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Order: []qbtypes.OrderBy{{Key: qbtypes.OrderByKey{TelemetryFieldKey: telemetrytypes.TelemetryFieldKey{Name: "trace.output_tokens"}}}},
})
require.ErrorContains(t, err, `ordering the span list by trace-level aggregate "trace.output_tokens" is not supported`)
err = build(qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Order: []qbtypes.OrderBy{{Key: qbtypes.OrderByKey{TelemetryFieldKey: telemetrytypes.TelemetryFieldKey{Name: "duration_nano"}}, Direction: qbtypes.OrderDirectionDesc}},
Limit: 10,
})
require.NoError(t, err, "bare duration_nano is a span column, not a trace-level key")
}
// Variables in a trace-level condition on the span list get the trace list's
// treatment: resolved and bound as args, __all__ drops the condition (no scope CTE).
func TestBuild_SpanList_TraceFilter_Variables(t *testing.T) {
b := newTestBuilder(t)
build := func(expr string, vars map[string]qbtypes.VariableItem) (*qbtypes.Statement, error) {
return b.Build(context.Background(), valuer.UUID{}, testStartMs, testEndMs, qbtypes.RequestTypeRaw,
qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Signal: telemetrytypes.SignalTraces,
Filter: &qbtypes.Filter{Expression: expr},
Limit: 10,
}, vars)
}
stmt, err := build("trace.output_tokens > $threshold",
map[string]qbtypes.VariableItem{"threshold": {Value: 700}})
require.NoError(t, err)
assert.Contains(t, stmt.Query, "HAVING output_tokens > ?")
assert.Contains(t, stmt.Args, float64(700))
stmt, err = build("trace.output_tokens > $threshold",
map[string]qbtypes.VariableItem{"threshold": {Type: qbtypes.DynamicVariableType, Value: "__all__"}})
require.NoError(t, err)
assert.NotContains(t, stmt.Query, "__trace_scope")
}

View File

@@ -16,6 +16,9 @@ type TraceScope struct {
Columns []TraceColumn
// DefaultOrderAlias is sorted by (desc) when the query gives no order.
DefaultOrderAlias string
// ActivityGateAlias names the column that must be > 0 for a per-trace row to feed
// trace-level (trace.) aggregations; empty disables the gate.
ActivityGateAlias string
}
// TraceColumn is one per-trace output column.

View File

@@ -4,7 +4,6 @@ import (
"context"
"fmt"
"log/slog"
"sort"
"strings"
"github.com/SigNoz/signoz/pkg/errors"
@@ -19,7 +18,6 @@ import (
qbtypes "github.com/SigNoz/signoz/pkg/types/querybuildertypes/querybuildertypesv5"
"github.com/SigNoz/signoz/pkg/types/telemetrytypes"
"github.com/SigNoz/signoz/pkg/valuer"
qbvariables "github.com/SigNoz/signoz/pkg/variables"
"github.com/huandu/go-sqlbuilder"
)
@@ -114,14 +112,26 @@ func (b *scopedTraceStatementBuilder) Build(
case qbtypes.RequestTypeTrace:
return b.buildTraceListQuery(ctx, orgID, querybuilder.ToNanoSecs(start), querybuilder.ToNanoSecs(end), query, variables)
case qbtypes.RequestTypeRaw:
if err := b.validateGroupByAndOrder(requestType, query); err != nil {
return nil, err
}
return b.buildDelegated(ctx, orgID, start, end, requestType, query, variables)
case qbtypes.RequestTypeScalar, qbtypes.RequestTypeTimeSeries:
return b.buildAggregation(ctx, orgID, start, end, requestType, query, variables)
default:
return nil, ErrUnsupportedRequestType
}
}
// buildDelegated ANDs the base gate into the user filter and delegates to the
// standard trace builder (the span-list / raw path).
// traceScopedStatementBuilder is the delegate's optional capability of constraining a
// query to a set of trace ids (implemented by the traces statement builder).
type traceScopedStatementBuilder interface {
BuildTraceScoped(ctx context.Context, orgID valuer.UUID, start, end uint64, requestType qbtypes.RequestType, query qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation], variables map[string]qbtypes.VariableItem, traceScope *qbtypes.Statement) (*qbtypes.Statement, error)
}
// buildDelegated (span list + span-level scalar/time-series) ANDs the gate into the
// filter's span-level part and delegates to the standard trace builder; a trace-level
// part becomes a qualification the delegate constrains trace_id by.
func (b *scopedTraceStatementBuilder) buildDelegated(
ctx context.Context,
orgID valuer.UUID,
@@ -130,17 +140,42 @@ func (b *scopedTraceStatementBuilder) buildDelegated(
query qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation],
variables map[string]qbtypes.VariableItem,
) (*qbtypes.Statement, error) {
var spanExpr, traceExpr string
var err error
if query.Filter != nil && strings.TrimSpace(query.Filter.Expression) != "" {
spanExpr, traceExpr, err = querybuilder.SplitFilterForAggregates(query.Filter.Expression, b.aggregateAliasSet())
if err != nil {
return nil, err
}
}
gate := b.scope.FilterExpression
expr := gate
if query.Filter != nil && strings.TrimSpace(query.Filter.Expression) != "" {
expr = fmt.Sprintf("(%s) AND (%s)", gate, query.Filter.Expression)
if strings.TrimSpace(spanExpr) != "" {
expr = fmt.Sprintf("(%s) AND (%s)", gate, spanExpr)
}
// shallow copy; only Filter is replaced, caller's query untouched
gated := query
gated.Filter = &qbtypes.Filter{Expression: expr}
return b.traceStmtBuilder.Build(ctx, orgID, start, end, requestType, gated, variables)
if strings.TrimSpace(traceExpr) == "" {
return b.traceStmtBuilder.Build(ctx, orgID, start, end, requestType, gated, variables)
}
scoped, ok := b.traceStmtBuilder.(traceScopedStatementBuilder)
if !ok {
return nil, errors.NewInternalf(errors.CodeInternal, "trace statement builder does not support trace-scoped queries")
}
scope, err := b.buildQualifiedStatement(ctx, orgID, querybuilder.ToNanoSecs(start), querybuilder.ToNanoSecs(end), traceExpr, query, variables)
if err != nil {
return nil, err
}
if scope == nil {
// every trace-level condition was dropped by variable resolution
return b.traceStmtBuilder.Build(ctx, orgID, start, end, requestType, gated, variables)
}
return scoped.BuildTraceScoped(ctx, orgID, start, end, requestType, gated, variables, scope)
}
// buildTraceListQuery wires the CTE pipeline (start/end are nanoseconds):
@@ -184,19 +219,17 @@ func (b *scopedTraceStatementBuilder) buildTraceListQuery(
if err != nil {
return nil, err
}
orderableSet := orderableAliasSet(resolved)
resourceFrag, resourceArgs, resourcePred, err := b.maybeAttachResourceFilter(ctx, orgID, query, start, end, variables)
if err != nil {
return nil, err
}
fp, err := b.splitFilter(ctx, orgID, query, b.aggregateAliasSet(), orderableSet, start, end, variables, matchedSB)
fp, err := b.splitFilter(ctx, orgID, query, b.aggregateAliasSet(), start, end, variables, matchedSB)
if err != nil {
return nil, err
}
matchedFrag, matchedArgs, err := b.buildMatchedCTE(matchedSB, start, end, startBucket, endBucket, resolved, orders, orderableSet, maskExpr, fp, resourcePred, limit, query.Offset)
matchedFrag, matchedArgs, err := b.buildMatchedCTE(matchedSB, start, end, startBucket, endBucket, resolved, orders, maskExpr, fp, resourcePred, limit, query.Offset)
if err != nil {
return nil, err
}
@@ -380,27 +413,27 @@ func (b *scopedTraceStatementBuilder) resolveListOrders(order []qbtypes.OrderBy,
return orders, nil
}
// filterParts is the user filter split into a span-level predicate and a trace-level
// HAVING expression.
// filterParts is the user filter split into a span-level predicate and the resolved
// trace-level HAVING (nil when there is none).
type filterParts struct {
spanPred string
hasSpanFilter bool
havingExpr string
having *traceHaving
warnings []string
warningsURL string
}
// splitFilter splits query.Filter into a span-level predicate (args bound into sb)
// and a trace-level HAVING (explicit query.Having ANDed on), then validates the
// trace-level part against the matched-pass aggregates.
func (b *scopedTraceStatementBuilder) splitFilter(ctx context.Context, orgID valuer.UUID, query qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation], classifySet, orderableSet map[string]struct{}, start, end uint64, variables map[string]qbtypes.VariableItem, sb *sqlbuilder.SelectBuilder) (filterParts, error) {
// splitFilter splits query.Filter into a span-level predicate and a trace-level
// HAVING (explicit query.Having ANDed on before resolution); args bind into sb.
func (b *scopedTraceStatementBuilder) splitFilter(ctx context.Context, orgID valuer.UUID, query qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation], classifySet map[string]struct{}, start, end uint64, variables map[string]qbtypes.VariableItem, sb *sqlbuilder.SelectBuilder) (filterParts, error) {
var fp filterParts
havingExpr := ""
if query.Filter != nil && strings.TrimSpace(query.Filter.Expression) != "" {
spanExpr, traceExpr, err := querybuilder.SplitFilterForAggregates(query.Filter.Expression, classifySet)
if err != nil {
return fp, err
}
fp.havingExpr = traceExpr
havingExpr = traceExpr
if strings.TrimSpace(spanExpr) != "" {
pred, warnings, url, err := b.resolveSpanPredicate(ctx, orgID, start, end, spanExpr, variables, sb)
if err != nil {
@@ -415,23 +448,17 @@ func (b *scopedTraceStatementBuilder) splitFilter(ctx context.Context, orgID val
}
}
if query.Having != nil && strings.TrimSpace(query.Having.Expression) != "" {
if fp.havingExpr != "" {
fp.havingExpr = fmt.Sprintf("(%s) AND (%s)", fp.havingExpr, query.Having.Expression)
if havingExpr != "" {
havingExpr = fmt.Sprintf("(%s) AND (%s)", havingExpr, query.Having.Expression)
} else {
fp.havingExpr = query.Having.Expression
havingExpr = query.Having.Expression
}
}
// the HAVING is a plain text rewrite, so substitute variables here
if strings.TrimSpace(fp.havingExpr) != "" && len(variables) > 0 {
replaced, err := qbvariables.ReplaceVariablesInExpression(fp.havingExpr, variables)
if err != nil {
return fp, err
}
fp.havingExpr = replaced
}
if err := validateAggregateFilter(fp.havingExpr, orderableSet); err != nil {
having, err := b.resolveTraceHaving(ctx, havingExpr, variables, sb)
if err != nil {
return fp, err
}
fp.having = having
return fp, nil
}
@@ -473,8 +500,8 @@ func (b *scopedTraceStatementBuilder) resolveSpanPredicate(ctx context.Context,
// span filter + HAVING + ORDER BY + LIMIT/OFFSET, selecting only the aliases ORDER BY
// / HAVING reference. Expressions carry $n markers bound to sb, so each can appear
// several times and every occurrence resolves to the same arg.
func (b *scopedTraceStatementBuilder) buildMatchedCTE(sb *sqlbuilder.SelectBuilder, start, end, startBucket, endBucket uint64, resolved []resolvedColumn, orders []listOrder, orderableSet map[string]struct{}, maskExpr string, fp filterParts, resourcePred string, limit, offset int) (string, []any, error) {
needed := neededMatchedAliases(orders, fp.havingExpr, orderableSet)
func (b *scopedTraceStatementBuilder) buildMatchedCTE(sb *sqlbuilder.SelectBuilder, start, end, startBucket, endBucket uint64, resolved []resolvedColumn, orders []listOrder, maskExpr string, fp filterParts, resourcePred string, limit, offset int) (string, []any, error) {
needed := neededMatchedAliases(orders, fp.having)
selects := []string{"trace_id"}
for _, rc := range resolved {
if _, ok := needed[rc.alias]; !ok {
@@ -511,22 +538,8 @@ func (b *scopedTraceStatementBuilder) buildMatchedCTE(sb *sqlbuilder.SelectBuild
having = append(having, "countIf("+maskExpr+") > 0")
having = append(having, "countIf("+fp.spanPred+") > 0")
}
if strings.TrimSpace(fp.havingExpr) != "" {
// the rewriter matches raw key text, so map the trace. form alongside the bare name
columnMap := make(map[string]string, len(orderableSet)*2)
for a := range orderableSet {
columnMap[a] = quoteAlias(a)
columnMap[telemetrytypes.FieldContextTrace.StringValue()+"."+a] = quoteAlias(a)
}
hv, err := querybuilder.NewHavingExpressionRewriter().Rewrite(fp.havingExpr, columnMap)
if err != nil {
return "", nil, err
}
if hv != "" {
// escape user text so a literal $ isn't read as an arg marker; the countIf
// entries hold live $n markers and must stay unescaped
having = append(having, sqlbuilder.Escape(hv))
}
if fp.having != nil {
having = append(having, fp.having.pred)
}
if len(having) > 0 {
sb.Having(strings.Join(having, " AND "))
@@ -592,59 +605,36 @@ func (b *scopedTraceStatementBuilder) aggregateAliasSet() map[string]struct{} {
return set
}
// orderableAliasSet is the subset of aliases computable in the matched pass.
func orderableAliasSet(resolved []resolvedColumn) map[string]struct{} {
set := make(map[string]struct{})
for _, rc := range resolved {
if rc.orderable {
set[rc.alias] = struct{}{}
}
}
return set
}
// neededMatchedAliases is the minimal alias set the matched pass must select: those
// in ORDER BY plus those in the aggregate HAVING.
func neededMatchedAliases(orders []listOrder, havingExpr string, orderableSet map[string]struct{}) map[string]struct{} {
// in ORDER BY plus those the resolved trace-level HAVING touches.
func neededMatchedAliases(orders []listOrder, having *traceHaving) map[string]struct{} {
needed := make(map[string]struct{})
for _, o := range orders {
needed[o.alias] = struct{}{}
}
for _, name := range traceAggregateNames(havingExpr) {
if _, ok := orderableSet[name]; ok {
if having != nil {
for name := range having.used {
needed[name] = struct{}{}
}
}
return needed
}
// traceAggregateNames extracts the aggregate names a trace-level HAVING references;
// only unspecified- and trace-context selectors name aggregates.
func traceAggregateNames(havingExpr string) []string {
var names []string
for _, sel := range querybuilder.QueryStringToKeysSelectors(havingExpr) {
if sel.FieldContext == telemetrytypes.FieldContextUnspecified || sel.FieldContext == telemetrytypes.FieldContextTrace {
names = append(names, sel.Name)
}
}
return names
}
// validateAggregateFilter rejects a trace-level filter referencing an aggregate not
// computable in the matched pass.
// validateAggregateFilter rejects trace-level filters on aggregates not computable in
// the matched pass (e.g. span_count) with a targeted top-level error; inside the
// where-clause visitor it would surface only as a detail of a combined error. Only
// unspecified- and trace-context selectors name aggregates.
func validateAggregateFilter(havingExpr string, orderableSet map[string]struct{}) error {
if strings.TrimSpace(havingExpr) == "" {
return nil
}
allowed := make([]string, 0, len(orderableSet))
for a := range orderableSet {
allowed = append(allowed, a)
}
sort.Strings(allowed)
for _, name := range traceAggregateNames(havingExpr) {
if _, ok := orderableSet[name]; !ok {
for _, sel := range querybuilder.QueryStringToKeysSelectors(havingExpr) {
if sel.FieldContext != telemetrytypes.FieldContextUnspecified && sel.FieldContext != telemetrytypes.FieldContextTrace {
continue
}
if _, ok := orderableSet[sel.Name]; !ok {
return errors.NewInvalidInputf(errors.CodeInvalidInput,
"aggregate %q cannot be used in the trace-list filter; filterable aggregates: %s", name, strings.Join(allowed, ", "))
"aggregate %q cannot be used in a trace-level filter; filterable aggregates: %s", sel.Name, strings.Join(sortedAliases(orderableSet), ", "))
}
}
return nil

View File

@@ -0,0 +1,789 @@
package scopedtracesstatementbuilder
import (
"context"
"fmt"
"sort"
"strings"
chparser "github.com/AfterShip/clickhouse-sql-parser/parser"
"github.com/SigNoz/signoz/pkg/errors"
"github.com/SigNoz/signoz/pkg/querybuilder"
"github.com/SigNoz/signoz/pkg/telemetryschema/tracestelemetryschema"
qbtypes "github.com/SigNoz/signoz/pkg/types/querybuildertypes/querybuildertypesv5"
"github.com/SigNoz/signoz/pkg/types/telemetrytypes"
"github.com/SigNoz/signoz/pkg/valuer"
"github.com/huandu/go-sqlbuilder"
)
// Scalar / time-series for scoped-trace queries. The `trace.` prefix picks the
// domain per expression: span-level (bare keys) delegates to the standard trace
// builder with the gate ANDed in; trace-level aggregates window-clipped per-trace
// values through the native pipeline (buildTraceAggregationQuery):
//
// __qualified traces whose whole-window aggregates satisfy the trace-level
// │ filter part; present only when the filter has one.
// ▼
// __scoped_traces per-trace values: windowed, mask-pruned GROUP BY trace_id
// │ (+ ts bucket for time series, + group-by columns), activity-
// ▼ gated rows only.
// main outer aggregation over the per-trace rows → __result_i.
// traceAggregation is one aggregation rewritten to run over the per-trace scan.
type traceAggregation struct {
expr string // rewritten SQL over the per-trace column aliases
used map[string]struct{} // per-trace aliases referenced
isRate bool
}
// buildAggregation routes scalar/time-series requests by aggregation domain.
func (b *scopedTraceStatementBuilder) buildAggregation(
ctx context.Context,
orgID valuer.UUID,
start, end uint64,
requestType qbtypes.RequestType,
query qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation],
variables map[string]qbtypes.VariableItem,
) (*qbtypes.Statement, error) {
traceAggs, err := b.classifyAggregations(query.Aggregations)
if err != nil {
return nil, err
}
if err := b.validateGroupByAndOrder(requestType, query); err != nil {
return nil, err
}
if len(traceAggs) == 0 {
return b.buildDelegated(ctx, orgID, start, end, requestType, query, variables)
}
return b.buildTraceAggregationQuery(ctx, orgID, querybuilder.ToNanoSecs(start), querybuilder.ToNanoSecs(end), requestType, query, variables, traceAggs)
}
// classifyAggregations returns the rewritten trace-domain aggregations, nil when all
// are span-domain; mixing the two domains is rejected.
func (b *scopedTraceStatementBuilder) classifyAggregations(aggs []qbtypes.TraceAggregation) ([]traceAggregation, error) {
traceCols := b.orderableColumnSet()
var out []traceAggregation
spanCount := 0
for _, agg := range aggs {
ta, isTrace, err := rewriteTraceAggregation(agg.Expression, traceCols)
if err != nil {
return nil, err
}
if isTrace {
out = append(out, *ta)
} else {
spanCount++
}
}
if len(out) > 0 && spanCount > 0 {
return nil, errors.NewInvalidInputf(errors.CodeInvalidInput,
"span-level and trace-level (trace.) aggregations cannot be mixed in one query")
}
return out, nil
}
// orderableColumnSet is the static per-trace column set usable in trace-level
// aggregations and filters.
func (b *scopedTraceStatementBuilder) orderableColumnSet() map[string]struct{} {
set := make(map[string]struct{})
for _, c := range b.scope.Columns {
if c.Orderable {
set[c.Alias] = struct{}{}
}
}
return set
}
// validateGroupByAndOrder rejects trace-level columns as group-by / order keys with a
// targeted error (not the field mapper's generic "field not found"); order keys naming
// an aggregation (alias / expression / index) are exempt.
func (b *scopedTraceStatementBuilder) validateGroupByAndOrder(requestType qbtypes.RequestType, query qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]) error {
aliases := b.aggregateAliasSet()
for _, gb := range query.GroupBy {
if isTraceLevelKey(gb.Name, gb.FieldContext, aliases) {
return errors.NewInvalidInputf(errors.CodeInvalidInput,
"grouping by trace-level aggregate %q is not supported; group by span attributes instead (e.g. service.name)", gb.Name)
}
}
for _, o := range query.Order {
if _, isAgg := traceAggOrderIndex(o, query); isAgg {
continue
}
if !isTraceLevelKey(o.Key.Name, o.Key.FieldContext, aliases) {
continue
}
if requestType == qbtypes.RequestTypeRaw {
return errors.NewInvalidInputf(errors.CodeInvalidInput,
"ordering the span list by trace-level aggregate %q is not supported; order by span columns instead (e.g. timestamp, duration_nano)", o.Key.Name)
}
return errors.NewInvalidInputf(errors.CodeInvalidInput,
"ordering by trace-level aggregate %q is not supported; order by the aggregation itself (its alias or expression) or a group-by key", o.Key.Name)
}
return nil
}
// isTraceLevelKey reports whether a group-by / order key explicitly (trace. prefix or
// trace context) names a per-trace aggregate; bare names pass through since they may
// be span columns sharing an alias name (duration_nano, timestamp).
func isTraceLevelKey(name string, fieldContext telemetrytypes.FieldContext, aliases map[string]struct{}) bool {
stripped := strings.TrimPrefix(name, "trace.")
if _, ok := aliases[stripped]; !ok {
return false
}
return stripped != name || fieldContext == telemetrytypes.FieldContextTrace
}
// rewriteTraceAggregation rewrites an aggregation over trace.-prefixed columns to run
// on the per-trace scan (trace.output_tokens → output_tokens, functions mapped via
// AggreFuncMap); a pure span-level expression returns isTrace=false for the delegate.
func rewriteTraceAggregation(expr string, traceCols map[string]struct{}) (*traceAggregation, bool, error) {
p := chparser.NewParser("SELECT " + expr)
stmts, err := p.ParseStmts()
if err != nil {
return nil, false, errors.WrapInvalidInputf(err, errors.CodeInvalidInput, "failed to parse aggregation expression %q", expr)
}
if len(stmts) == 0 {
return nil, false, errors.NewInvalidInputf(errors.CodeInvalidInput, "invalid aggregation expression %q", expr)
}
sel, ok := stmts[0].(*chparser.SelectQuery)
if !ok || len(sel.SelectItems) == 0 {
return nil, false, errors.NewInvalidInputf(errors.CodeInvalidInput, "invalid aggregation expression %q", expr)
}
v := &traceAggVisitor{traceCols: traceCols, used: make(map[string]struct{})}
if err := sel.SelectItems[0].Accept(v); err != nil {
return nil, false, err
}
if !v.hasTrace {
return nil, false, nil
}
if v.hasSpan {
return nil, false, errors.NewInvalidInputf(errors.CodeInvalidInput,
"aggregation %q mixes trace-level (trace.) and span-level columns; use one domain per aggregation", expr)
}
return &traceAggregation{expr: chparser.Format(sel.SelectItems[0]), used: v.used, isRate: v.isRate}, true, nil
}
// traceAggVisitor classifies column references and rewrites trace.-prefixed ones in
// place; the ancestor stack tells a column identifier from a path segment, function
// name, or alias, and rejects trace. columns inside *If combinators.
type traceAggVisitor struct {
chparser.DefaultASTVisitor
traceCols map[string]struct{}
used map[string]struct{}
stack []chparser.Expr
hasTrace bool
hasSpan bool
isRate bool
}
func (v *traceAggVisitor) Enter(expr chparser.Expr) { v.stack = append(v.stack, expr) }
func (v *traceAggVisitor) Leave(expr chparser.Expr) { v.stack = v.stack[:len(v.stack)-1] }
// parent is the node enclosing the one currently being visited (the visited node
// itself is the stack top).
func (v *traceAggVisitor) parent() chparser.Expr {
if len(v.stack) < 2 {
return nil
}
return v.stack[len(v.stack)-2]
}
// enclosingCombinator returns the name of a surrounding *If-combinator function, if any.
func (v *traceAggVisitor) enclosingCombinator() (string, bool) {
for _, e := range v.stack {
fn, ok := e.(*chparser.FunctionExpr)
if !ok {
continue
}
if agg, known := querybuilder.AggreFuncMap[valuer.NewString(strings.ToLower(fn.Name.Name))]; known && agg.FuncCombinator {
return fn.Name.Name, true
}
}
return "", false
}
// VisitPath classifies a dotted reference (trace.output_tokens); trace-level ones are
// rewritten in place to the bare per-trace alias.
func (v *traceAggVisitor) VisitPath(p *chparser.Path) error {
ref := chparser.Format(p)
col, isTrace := traceColumnRef(ref)
if !isTrace {
v.hasSpan = true
return nil
}
if err := v.acceptTraceColumn(ref, col); err != nil {
return err
}
p.Fields = p.Fields[len(p.Fields)-1:]
p.Fields[0].Name = col
return nil
}
// VisitIdent classifies a plain identifier (a backquoted `trace.output_tokens` is
// trace-level); path segments, function names, and aliases are structural, not columns.
func (v *traceAggVisitor) VisitIdent(i *chparser.Ident) error {
switch parent := v.parent().(type) {
case *chparser.Path:
return nil // segments are classified whole by VisitPath
case *chparser.FunctionExpr:
if parent.Name == i {
return nil
}
case *chparser.ColumnExpr:
if parent.Alias == i {
return nil
}
}
col, isTrace := traceColumnRef(i.Name)
if !isTrace {
v.hasSpan = true
return nil
}
if err := v.acceptTraceColumn(i.Name, col); err != nil {
return err
}
i.Name = col
return nil
}
// acceptTraceColumn validates one trace-level column reference and records it.
func (v *traceAggVisitor) acceptTraceColumn(ref, col string) error {
if name, in := v.enclosingCombinator(); in {
return errors.NewInvalidInputf(errors.CodeInvalidInput,
"%q over trace-level (trace.) columns is not supported; put the trace-level condition in the filter expression instead", name)
}
// trace_id is always selected by the per-trace scan (count(trace.trace_id)
// counts traces); everything else must be a scope column.
if col != "trace_id" {
if _, known := v.traceCols[col]; !known {
return errors.NewInvalidInputf(errors.CodeInvalidInput,
"unknown trace-level aggregation column %q; usable columns: %s", ref, strings.Join(sortedAliases(v.traceCols), ", "))
}
v.used[col] = struct{}{}
}
v.hasTrace = true
return nil
}
// VisitFunctionExpr validates and maps the function name. Children were already
// visited (post-order), so classification is complete for this subtree.
func (v *traceAggVisitor) VisitFunctionExpr(fn *chparser.FunctionExpr) error {
name := strings.ToLower(fn.Name.Name)
aggFunc, ok := querybuilder.AggreFuncMap[valuer.NewString(name)]
if !ok {
return errors.NewInvalidInputf(errors.CodeInvalidInput, "unrecognized function: %s", name)
}
if fn.Params != nil && fn.Params.Items != nil && len(fn.Params.Items.Items) > 0 && aggFunc.FuncCombinator {
// combinator predicates over span columns stay span-level (countIf(has_error=true))
v.hasSpan = true
return nil
}
fn.Name.Name = aggFunc.FuncName
if aggFunc.Rate {
v.isRate = true
}
return nil
}
// traceColumnRef reports whether text is a pure trace.-prefixed column reference
// (trace.output_tokens) and returns the bare column name.
func traceColumnRef(text string) (string, bool) {
text = strings.TrimSpace(text)
rest, ok := strings.CutPrefix(text, "trace.")
if !ok {
return "", false
}
if rest == "" || strings.ContainsAny(rest, " ()'\"`,+-*/<>=!") {
return "", false
}
return rest, true
}
func sortedAliases(set map[string]struct{}) []string {
out := make([]string, 0, len(set))
for a := range set {
out = append(out, a)
}
sort.Strings(out)
return out
}
// ---------------------------------------------------------------------------
// Qualification + per-trace scan
// ---------------------------------------------------------------------------
// buildQualifiedStatement builds the delegate's __trace_scope: trace ids whose
// window-clipped aggregates satisfy the trace-level filter, resource-pruned inline
// (the caller embeds it standalone). start/end are ns; nil when every condition was
// dropped by variable resolution.
func (b *scopedTraceStatementBuilder) buildQualifiedStatement(
ctx context.Context,
orgID valuer.UUID,
start, end uint64,
traceExpr string,
query qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation],
variables map[string]qbtypes.VariableItem,
) (*qbtypes.Statement, error) {
keys, err := b.fetchKeys(ctx, orgID)
if err != nil {
return nil, err
}
sb := sqlbuilder.NewSelectBuilder()
maskExpr, resolved, err := b.resolveFor(ctx, orgID, start, end, keys, sb)
if err != nil {
return nil, err
}
having, err := b.resolveTraceHaving(ctx, traceExpr, variables, sb)
if err != nil {
return nil, err
}
if having == nil {
return nil, nil
}
var resourcePred string
// nil when the filter has no resource-attribute conditions
if stmt, err := b.resourceFilterStmtBuilder.Build(ctx, orgID, start, end, qbtypes.RequestTypeRaw, query, variables); err != nil {
return nil, err
} else if stmt != nil {
inlined, err := embedExpr(sb, stmt.Query, stmt.Args)
if err != nil {
return nil, err
}
resourcePred = fmt.Sprintf("resource_fingerprint GLOBAL IN (SELECT fingerprint FROM (%s))", inlined)
}
sql, args := b.buildPerTraceScan(sb, start, end, resolved, maskExpr, perTraceScanOpts{
needed: having.used,
havingPred: having.pred,
resourcePred: resourcePred,
})
return &qbtypes.Statement{Query: sql, Args: args}, nil
}
// embedExpr inlines a pre-built statement into sb, replacing each `?` with a builder
// Var; a count mismatch would silently shift args into the wrong slots, so error out.
func embedExpr(sb *sqlbuilder.SelectBuilder, expr string, args []any) (string, error) {
if n := strings.Count(expr, "?"); n != len(args) {
return "", errors.NewInternalf(errors.CodeInternal,
"scoped trace builder: %d placeholders != %d args embedding %q", n, len(args), expr)
}
var out strings.Builder
ai := 0
for i := 0; i < len(expr); i++ {
if expr[i] == '?' {
out.WriteString(sb.Var(args[ai]))
ai++
continue
}
out.WriteByte(expr[i])
}
return out.String(), nil
}
// groupColumn is a resolved span-attribute group-by column (arg-free expression).
type groupColumn struct {
name string
expr string
}
// perTraceScanOpts parametrize one windowed, mask-pruned GROUP BY trace_id scan.
// All expressions are already resolved against the scan's builder.
type perTraceScanOpts struct {
stepSeconds int64 // >0 → bucket per-trace values by time (ts column)
groupCols []groupColumn
needed map[string]struct{} // per-trace aliases to select
spanPred string // resolved span-level filter, ANDed per span
resourcePred string // resource-fingerprint prune (CTE reference or inline subquery)
qualified bool // constrain to __qualified
havingPred string // resolved HAVING predicate over the selected aliases
activityExpr string // aggregate expr that must be > 0 for a row to survive
}
// buildPerTraceScan renders the scan: window + gate mask (+ span filter, resource
// prune, qualification), grouped by trace_id (+ ts bucket, group-by columns).
func (b *scopedTraceStatementBuilder) buildPerTraceScan(sb *sqlbuilder.SelectBuilder, start, end uint64, resolved []resolvedColumn, maskExpr string, o perTraceScanOpts) (string, []any) {
startBucket := start/querybuilder.NsToSeconds - querybuilder.BucketAdjustment
endBucket := end / querybuilder.NsToSeconds
selects := []string{"trace_id"}
if o.stepSeconds > 0 {
selects = append(selects, fmt.Sprintf("toStartOfInterval(timestamp, INTERVAL %d SECOND) AS ts", o.stepSeconds))
}
for _, gc := range o.groupCols {
selects = append(selects, fmt.Sprintf("toString(%s) AS `%s`", gc.expr, gc.name))
}
for _, rc := range resolved {
if _, ok := o.needed[rc.alias]; !ok {
continue
}
selects = append(selects, rc.expr+" AS "+quoteAlias(rc.alias))
}
sb.Select(selects...)
sb.From(fmt.Sprintf("%s.%s", tracestelemetryschema.DBName, tracestelemetryschema.SpanIndexV3TableName))
where := []string{
sb.GE("timestamp", fmt.Sprintf("%d", start)),
sb.L("timestamp", fmt.Sprintf("%d", end)),
sb.GE("ts_bucket_start", startBucket),
sb.LE("ts_bucket_start", endBucket),
maskExpr,
}
if strings.TrimSpace(o.spanPred) != "" {
where = append(where, o.spanPred)
}
if o.resourcePred != "" {
where = append(where, o.resourcePred)
}
if o.qualified {
where = append(where, "trace_id GLOBAL IN (SELECT trace_id FROM __qualified)")
}
sb.Where(where...)
groupBy := []string{"trace_id"}
if o.stepSeconds > 0 {
groupBy = append(groupBy, "ts")
}
for _, gc := range o.groupCols {
groupBy = append(groupBy, "`"+gc.name+"`")
}
sb.GroupBy(groupBy...)
var having []string
if strings.TrimSpace(o.activityExpr) != "" {
having = append(having, "("+o.activityExpr+") > 0")
}
if strings.TrimSpace(o.havingPred) != "" {
having = append(having, o.havingPred)
}
if len(having) > 0 {
sb.Having(strings.Join(having, " AND "))
}
return sb.BuildWithFlavor(sqlbuilder.ClickHouse)
}
// resolveGroupColumns resolves span-attribute group-by keys through the field mapper
// (metadata-aware), for selection inside the per-trace scan.
func (b *scopedTraceStatementBuilder) resolveGroupColumns(ctx context.Context, orgID valuer.UUID, start, end uint64, groupBy []qbtypes.GroupByKey) ([]groupColumn, error) {
if len(groupBy) == 0 {
return nil, nil
}
selectors := make([]*telemetrytypes.FieldKeySelector, 0, len(groupBy))
for i := range groupBy {
selectors = append(selectors, &telemetrytypes.FieldKeySelector{
Name: groupBy[i].Name,
Signal: telemetrytypes.SignalTraces,
FieldContext: groupBy[i].FieldContext,
FieldDataType: groupBy[i].FieldDataType,
SelectorMatchType: telemetrytypes.FieldSelectorMatchTypeExact,
})
}
keys, _, err := b.metadataStore.GetKeysMulti(ctx, orgID, selectors)
if err != nil {
return nil, err
}
out := make([]groupColumn, 0, len(groupBy))
for i := range groupBy {
expr, err := b.fm.ColumnExpressionFor(ctx, orgID, start, end, &groupBy[i].TelemetryFieldKey, telemetrytypes.FieldDataTypeString, keys)
if err != nil {
return nil, err
}
out = append(out, groupColumn{name: groupBy[i].Name, expr: sqlbuilder.Escape(expr)})
}
return out, nil
}
// ---------------------------------------------------------------------------
// Native trace-domain aggregation query
// ---------------------------------------------------------------------------
// scanContext is one per-scan resolution: a fresh builder with the mask, columns,
// span predicate, and optionally the trace-level HAVING resolved against it.
type scanContext struct {
sb *sqlbuilder.SelectBuilder
maskExpr string
resolved []resolvedColumn
spanPred string
having *traceHaving
warnings []string
warnURL string
}
// newScanContext resolves everything a per-trace scan embeds against a fresh builder.
func (b *scopedTraceStatementBuilder) newScanContext(
ctx context.Context,
orgID valuer.UUID,
start, end uint64,
keys map[string][]*telemetrytypes.TelemetryFieldKey,
spanExpr, traceExpr string,
variables map[string]qbtypes.VariableItem,
) (*scanContext, error) {
sc := &scanContext{sb: sqlbuilder.NewSelectBuilder()}
var err error
sc.maskExpr, sc.resolved, err = b.resolveFor(ctx, orgID, start, end, keys, sc.sb)
if err != nil {
return nil, err
}
if strings.TrimSpace(spanExpr) != "" {
pred, warns, url, err := b.resolveSpanPredicate(ctx, orgID, start, end, spanExpr, variables, sc.sb)
if err != nil {
return nil, err
}
sc.spanPred, sc.warnings, sc.warnURL = pred, warns, url
}
if strings.TrimSpace(traceExpr) != "" {
sc.having, err = b.resolveTraceHaving(ctx, traceExpr, variables, sc.sb)
if err != nil {
return nil, err
}
}
return sc, nil
}
// buildTraceAggregationQuery builds the native pipeline (see the file comment).
// start/end are ns.
func (b *scopedTraceStatementBuilder) buildTraceAggregationQuery(
ctx context.Context,
orgID valuer.UUID,
start, end uint64,
requestType qbtypes.RequestType,
query qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation],
variables map[string]qbtypes.VariableItem,
traceAggs []traceAggregation,
) (*qbtypes.Statement, error) {
keys, err := b.fetchKeys(ctx, orgID)
if err != nil {
return nil, err
}
var spanExpr, traceExpr string
if query.Filter != nil && strings.TrimSpace(query.Filter.Expression) != "" {
spanExpr, traceExpr, err = querybuilder.SplitFilterForAggregates(query.Filter.Expression, b.aggregateAliasSet())
if err != nil {
return nil, err
}
}
resourceFrag, resourceArgs, resourcePred, err := b.maybeAttachResourceFilter(ctx, orgID, query, start, end, variables)
if err != nil {
return nil, err
}
var cteFragments []string
var cteArgs [][]any
if resourceFrag != "" {
cteFragments = append(cteFragments, resourceFrag)
cteArgs = append(cteArgs, resourceArgs)
}
// __qualified: its own scan resolution, HAVING = the trace-level filter part
qualified := false
if strings.TrimSpace(traceExpr) != "" {
qsc, err := b.newScanContext(ctx, orgID, start, end, keys, "", traceExpr, variables)
if err != nil {
return nil, err
}
if qsc.having != nil {
qsql, qargs := b.buildPerTraceScan(qsc.sb, start, end, qsc.resolved, qsc.maskExpr, perTraceScanOpts{
needed: qsc.having.used,
havingPred: qsc.having.pred,
resourcePred: resourcePred,
})
cteFragments = append(cteFragments, fmt.Sprintf("__qualified AS (%s)", qsql))
cteArgs = append(cteArgs, qargs)
qualified = true
}
}
groupCols, err := b.resolveGroupColumns(ctx, orgID, start, end, query.GroupBy)
if err != nil {
return nil, err
}
groupNames := make([]string, 0, len(groupCols))
for _, gc := range groupCols {
groupNames = append(groupNames, "`"+gc.name+"`")
}
needed := make(map[string]struct{})
for _, ta := range traceAggs {
for a := range ta.used {
needed[a] = struct{}{}
}
}
stepSeconds := int64(0)
rateInterval := (end - start) / querybuilder.NsToSeconds
if requestType == qbtypes.RequestTypeTimeSeries {
stepSeconds = int64(query.StepInterval.Seconds())
rateInterval = uint64(stepSeconds)
}
// outer aggregation over the per-trace rows
sb := sqlbuilder.NewSelectBuilder()
selects := []string{}
if stepSeconds > 0 {
selects = append(selects, "ts")
}
selects = append(selects, groupNames...)
for i, ta := range traceAggs {
selects = append(selects, fmt.Sprintf("%s AS __result_%d", ta.rendered(rateInterval), i))
}
sb.Select(selects...)
sb.From("__scoped_traces")
// grouped, limited time series → rank groups on whole-window per-trace values
// (exact for non-composable aggregates) and constrain the main query to the top-N.
if requestType == qbtypes.RequestTypeTimeSeries && query.Limit > 0 && len(groupCols) > 0 {
tsc, err := b.newScanContext(ctx, orgID, start, end, keys, spanExpr, "", variables)
if err != nil {
return nil, err
}
totalSQL, totalArgs := b.buildPerTraceScan(tsc.sb, start, end, tsc.resolved, tsc.maskExpr, perTraceScanOpts{
groupCols: groupCols,
needed: needed,
spanPred: tsc.spanPred,
resourcePred: resourcePred,
qualified: qualified,
activityExpr: activityGate(b.scope, tsc.resolved),
})
cteFragments = append(cteFragments, fmt.Sprintf("__scoped_traces_total AS (%s)", totalSQL))
cteArgs = append(cteArgs, totalArgs)
limitSQL, limitArgs := outerLimitSQL(query, traceAggs, groupNames, (end-start)/querybuilder.NsToSeconds)
cteFragments = append(cteFragments, fmt.Sprintf("__limit_cte AS (%s)", limitSQL))
cteArgs = append(cteArgs, limitArgs)
tuple := "(" + strings.Join(groupNames, ", ") + ")"
sb.Where(fmt.Sprintf("%s IN (SELECT %s FROM __limit_cte)", tuple, strings.Join(groupNames, ", ")))
}
msc, err := b.newScanContext(ctx, orgID, start, end, keys, spanExpr, "", variables)
if err != nil {
return nil, err
}
perTraceSQL, perTraceArgs := b.buildPerTraceScan(msc.sb, start, end, msc.resolved, msc.maskExpr, perTraceScanOpts{
stepSeconds: stepSeconds,
groupCols: groupCols,
needed: needed,
spanPred: msc.spanPred,
resourcePred: resourcePred,
qualified: qualified,
// drop rows with no gated activity in their window/bucket slice, so all
// aggregations see the same trace set (see TraceScope.ActivityGateAlias)
activityExpr: activityGate(b.scope, msc.resolved),
})
cteFragments = append(cteFragments, fmt.Sprintf("__scoped_traces AS (%s)", perTraceSQL))
cteArgs = append(cteArgs, perTraceArgs)
groupBys := []string{}
if stepSeconds > 0 {
groupBys = append(groupBys, "ts")
}
groupBys = append(groupBys, groupNames...)
if len(groupBys) > 0 {
sb.GroupBy(groupBys...)
}
if query.Having != nil && strings.TrimSpace(query.Having.Expression) != "" {
rewritten, err := querybuilder.NewHavingExpressionRewriter().RewriteForTraces(query.Having.Expression, query.Aggregations)
if err != nil {
return nil, err
}
sb.Having(sqlbuilder.Escape(rewritten))
}
if requestType == qbtypes.RequestTypeTimeSeries {
if len(query.Order) != 0 {
for _, orderBy := range query.Order {
if _, ok := traceAggOrderIndex(orderBy, query); !ok {
sb.OrderBy(fmt.Sprintf("`%s` %s", orderBy.Key.Name, orderBy.Direction.StringValue()))
}
}
sb.OrderBy("ts desc")
}
} else {
for _, orderBy := range query.Order {
if idx, ok := traceAggOrderIndex(orderBy, query); ok {
sb.OrderBy(fmt.Sprintf("__result_%d %s", idx, orderBy.Direction.StringValue()))
} else {
sb.OrderBy(fmt.Sprintf("`%s` %s", orderBy.Key.Name, orderBy.Direction.StringValue()))
}
}
if len(query.Order) == 0 {
sb.OrderBy("__result_0 DESC")
}
if query.Limit > 0 {
sb.Limit(query.Limit)
}
}
mainSQL, mainArgs := sb.BuildWithFlavor(sqlbuilder.ClickHouse)
finalSQL := querybuilder.CombineCTEs(cteFragments) + mainSQL + " SETTINGS distributed_product_mode='allow', max_memory_usage=10000000000"
finalArgs := querybuilder.PrependArgs(cteArgs, mainArgs)
return &qbtypes.Statement{
Query: finalSQL,
Args: finalArgs,
Warnings: msc.warnings,
WarningsDocURL: msc.warnURL,
}, nil
}
// activityGate resolves the scope's activity-gate column to its aggregate expression;
// empty when the scope declares none.
func activityGate(scope TraceScope, resolved []resolvedColumn) string {
if scope.ActivityGateAlias == "" {
return ""
}
for _, rc := range resolved {
if rc.alias == scope.ActivityGateAlias {
return rc.expr
}
}
return ""
}
// rendered returns the outer aggregation SQL, dividing rate aggregations by the
// interval (step for time series, window length for scalar).
func (ta traceAggregation) rendered(rateInterval uint64) string {
if ta.isRate {
return fmt.Sprintf("%s/%d", ta.expr, rateInterval)
}
return ta.expr
}
// outerLimitSQL renders the top-N group selection for a grouped, limited time series:
// outer aggregations over whole-window per-trace values, ranked and limited.
func outerLimitSQL(query qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation], traceAggs []traceAggregation, groupNames []string, windowSeconds uint64) (string, []any) {
sb := sqlbuilder.NewSelectBuilder()
selects := append([]string{}, groupNames...)
for i, ta := range traceAggs {
selects = append(selects, fmt.Sprintf("%s AS __result_%d", ta.rendered(windowSeconds), i))
}
sb.Select(selects...)
sb.From("__scoped_traces_total")
sb.GroupBy(groupNames...)
for _, orderBy := range query.Order {
if idx, ok := traceAggOrderIndex(orderBy, query); ok {
sb.OrderBy(fmt.Sprintf("__result_%d %s", idx, orderBy.Direction.StringValue()))
} else {
sb.OrderBy(fmt.Sprintf("`%s` %s", orderBy.Key.Name, orderBy.Direction.StringValue()))
}
}
if len(query.Order) == 0 {
sb.OrderBy("__result_0 DESC")
}
sb.Limit(query.Limit)
return sb.BuildWithFlavor(sqlbuilder.ClickHouse)
}
// traceAggOrderIndex reports whether an order key refers to the i-th aggregation
// (by alias, expression, or index), mirroring the trace builder.
func traceAggOrderIndex(k qbtypes.OrderBy, q qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]) (int, bool) {
for i, agg := range q.Aggregations {
if k.Key.Name == agg.Alias ||
k.Key.Name == agg.Expression ||
k.Key.Name == fmt.Sprintf("%d", i) {
return i, true
}
}
return 0, false
}

View File

@@ -0,0 +1,60 @@
package scopedtracesstatementbuilder
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestRewriteTraceAggregation(t *testing.T) {
cols := map[string]struct{}{
"input_tokens": {}, "output_tokens": {}, "total_tokens": {}, "llm_call_count": {}, "max_llm_latency_ns": {},
}
cases := []struct {
name string
expr string
isTrace bool
want string // rewritten expr, only checked when isTrace
used []string
wantErr string
}{
{name: "avg trace col", expr: "avg(trace.output_tokens)", isTrace: true, want: "avg(output_tokens)", used: []string{"output_tokens"}},
{name: "sum trace col", expr: "sum(trace.total_tokens)", isTrace: true, want: "sum(total_tokens)", used: []string{"total_tokens"}},
{name: "count traces", expr: "count(trace.trace_id)", isTrace: true, want: "count(trace_id)"},
{name: "p90 trace col", expr: "p90(trace.max_llm_latency_ns)", isTrace: true, want: "quantile(0.90)(max_llm_latency_ns)", used: []string{"max_llm_latency_ns"}},
{name: "arithmetic between trace cols", expr: "avg(trace.output_tokens + trace.input_tokens)", isTrace: true, want: "avg(output_tokens + input_tokens)", used: []string{"output_tokens", "input_tokens"}},
{name: "arithmetic with constant", expr: "sum(trace.output_tokens * 1.5)", isTrace: true, want: "sum(output_tokens * 1.5)", used: []string{"output_tokens"}},
{name: "ratio of two aggregations", expr: "sum(trace.output_tokens)/count(trace.trace_id)", isTrace: true, want: "sum(output_tokens) / count(trace_id)", used: []string{"output_tokens"}},
{name: "backquoted trace col", expr: "avg(`trace.output_tokens`)", isTrace: true, want: "avg(`output_tokens`)", used: []string{"output_tokens"}},
{name: "bare count is span-level", expr: "count()", isTrace: false},
{name: "span attribute is span-level", expr: "sum(gen_ai.usage.output_tokens)", isTrace: false},
{name: "countIf span predicate is span-level", expr: "countIf(has_error = true)", isTrace: false},
{name: "mixed domains in one expression", expr: "sum(trace.output_tokens) + sum(gen_ai.usage.input_tokens)", wantErr: "mixes trace-level"},
{name: "mixed domains in one function", expr: "sum(trace.output_tokens + gen_ai.usage.input_tokens)", wantErr: "mixes trace-level"},
{name: "output-only column rejected", expr: "avg(trace.span_count)", wantErr: "unknown trace-level aggregation column"},
{name: "unknown column rejected", expr: "avg(trace.bogus)", wantErr: "unknown trace-level aggregation column"},
{name: "countIf over trace col rejected", expr: "countIf(trace.output_tokens > 1000)", wantErr: "not supported"},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
ta, isTrace, err := rewriteTraceAggregation(tc.expr, cols)
if tc.wantErr != "" {
require.ErrorContains(t, err, tc.wantErr)
return
}
require.NoError(t, err)
assert.Equal(t, tc.isTrace, isTrace)
if !tc.isTrace {
return
}
assert.Equal(t, tc.want, ta.expr)
for _, u := range tc.used {
assert.Contains(t, ta.used, u)
}
assert.Len(t, ta.used, len(tc.used))
})
}
}

View File

@@ -0,0 +1,146 @@
package scopedtracesstatementbuilder
import (
"context"
"strings"
"github.com/SigNoz/signoz/pkg/errors"
"github.com/SigNoz/signoz/pkg/querybuilder"
qbtypes "github.com/SigNoz/signoz/pkg/types/querybuildertypes/querybuildertypesv5"
"github.com/SigNoz/signoz/pkg/types/telemetrytypes"
"github.com/SigNoz/signoz/pkg/valuer"
qbvariables "github.com/SigNoz/signoz/pkg/variables"
"github.com/huandu/go-sqlbuilder"
)
// traceHaving is the resolved trace-level filter part: a HAVING predicate over the
// per-trace aliases plus the aliases it references (so scans select only those).
type traceHaving struct {
pred string
used map[string]struct{}
}
// resolveTraceHaving resolves a trace-level filter through the standard pipeline
// (variable replacement, then PrepareWhereClause against the per-trace aliases), so
// operators, bound args, and __all__ behave exactly as in span filters. Returns nil
// when the expression is empty or every condition was dropped; args bind into sb.
func (b *scopedTraceStatementBuilder) resolveTraceHaving(ctx context.Context, expr string, variables map[string]qbtypes.VariableItem, sb *sqlbuilder.SelectBuilder) (*traceHaving, error) {
if strings.TrimSpace(expr) == "" {
return nil, nil
}
// variables are replaced before validation so their literals are not mistaken for
// aggregate names; an unresolved $var is left in place and fails validation below
// (as an unknown aggregate — targeted variable errors are a separate concern)
if len(variables) > 0 {
replaced, err := qbvariables.ReplaceVariablesInExpression(expr, variables)
if err != nil {
return nil, err
}
expr = replaced
if strings.TrimSpace(expr) == "" {
return nil, nil
}
}
allowed := b.orderableColumnSet()
// upfront targeted errors; the visitor folds them into a combined "Found N errors"
if err := validateAggregateFilter(expr, allowed); err != nil {
return nil, err
}
// both spellings resolve here: the key parser strips the trace. prefix into
// FieldContextTrace, which matches this entry's context
fieldKeys := make(map[string][]*telemetrytypes.TelemetryFieldKey, len(allowed))
for alias := range allowed {
key := &telemetrytypes.TelemetryFieldKey{Name: alias, FieldContext: telemetrytypes.FieldContextTrace}
fieldKeys[alias] = []*telemetrytypes.TelemetryFieldKey{key}
}
cb := &aliasConditionBuilder{allowed: allowed, used: make(map[string]struct{})}
prepared, err := querybuilder.PrepareWhereClause(expr, querybuilder.FilterExprVisitorOpts{
Context: ctx,
Logger: b.logger,
ConditionBuilder: cb,
FieldKeys: fieldKeys,
Variables: variables,
Builder: sb,
})
if err != nil {
return nil, err
}
if prepared.IsEmpty() {
return nil, nil
}
return &traceHaving{pred: prepared.Expr, used: cb.used}, nil
}
// aliasConditionBuilder renders filter conditions directly against the per-trace
// aliases, recording the ones it touches; a key resolving to no alias is an error.
type aliasConditionBuilder struct {
allowed map[string]struct{}
used map[string]struct{}
}
var _ qbtypes.ConditionBuilder = (*aliasConditionBuilder)(nil)
func (c *aliasConditionBuilder) ConditionFor(
_ context.Context,
_ valuer.UUID,
_, _ uint64,
key *telemetrytypes.TelemetryFieldKey,
keys map[string][]*telemetrytypes.TelemetryFieldKey,
_ qbtypes.ConditionBuilderOptions,
op qbtypes.FilterOperator,
value any,
sb *sqlbuilder.SelectBuilder,
) ([]string, []string, error) {
matching := keys[key.Name]
if len(matching) == 0 {
return nil, nil, errors.NewInvalidInputf(errors.CodeInvalidInput,
"aggregate %q cannot be used in a trace-level filter; filterable aggregates: %s",
key.Name, strings.Join(sortedAliases(c.allowed), ", "))
}
alias := matching[0].Name
c.used[alias] = struct{}{}
col := quoteAlias(alias)
var cond string
switch op {
case qbtypes.FilterOperatorEqual:
cond = sb.E(col, value)
case qbtypes.FilterOperatorNotEqual:
cond = sb.NE(col, value)
case qbtypes.FilterOperatorGreaterThan:
cond = sb.G(col, value)
case qbtypes.FilterOperatorGreaterThanOrEq:
cond = sb.GE(col, value)
case qbtypes.FilterOperatorLessThan:
cond = sb.L(col, value)
case qbtypes.FilterOperatorLessThanOrEq:
cond = sb.LE(col, value)
case qbtypes.FilterOperatorIn, qbtypes.FilterOperatorNotIn:
values, ok := value.([]any)
if !ok {
values = []any{value}
}
if op == qbtypes.FilterOperatorIn {
cond = sb.In(col, values...)
} else {
cond = sb.NotIn(col, values...)
}
case qbtypes.FilterOperatorBetween, qbtypes.FilterOperatorNotBetween:
values, ok := value.([]any)
if !ok || len(values) != 2 {
return nil, nil, errors.NewInvalidInputf(errors.CodeInvalidInput,
"between on trace-level aggregate %q requires exactly two values", alias)
}
if op == qbtypes.FilterOperatorBetween {
cond = sb.Between(col, values[0], values[1])
} else {
cond = sb.NotBetween(col, values[0], values[1])
}
default:
return nil, nil, errors.NewInvalidInputf(errors.CodeInvalidInput,
"trace-level aggregate %q supports only comparison operators (=, !=, <, <=, >, >=, in, between)", alias)
}
return []string{cond}, nil, nil
}

View File

@@ -32,6 +32,9 @@ type traceQueryStatementBuilder struct {
resourceFilterResolver *resourcefilter.ResourceFingerprintResolver[qbtypes.TraceAggregation]
aggExprRewriter qbtypes.AggExprRewriter
skipResourceFingerprintEnabled bool
// traceScope, set only on the per-call copy made by BuildTraceScoped, constrains
// queries to spans whose trace_id is in the __trace_scope CTE.
traceScope *qbtypes.Statement
}
var _ qbtypes.StatementBuilder[qbtypes.TraceAggregation] = (*traceQueryStatementBuilder)(nil)
@@ -95,6 +98,33 @@ func NewTraceQueryStatementBuilder(
}
}
// BuildTraceScoped is Build constrained to trace_ids selected by traceScope; the
// receiver is copied so the shared builder stays stateless.
func (b *traceQueryStatementBuilder) BuildTraceScoped(
ctx context.Context,
orgID valuer.UUID,
start uint64,
end uint64,
requestType qbtypes.RequestType,
query qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation],
variables map[string]qbtypes.VariableItem,
traceScope *qbtypes.Statement,
) (*qbtypes.Statement, error) {
scoped := *b
scoped.traceScope = traceScope
return scoped.Build(ctx, orgID, start, end, requestType, query, variables)
}
// attachTraceScope adds the trace-scope condition to sb and returns the CTE fragment
// + args to prepend; both empty when no scope is set.
func (b *traceQueryStatementBuilder) attachTraceScope(sb *sqlbuilder.SelectBuilder) (string, []any) {
if b.traceScope == nil {
return "", nil
}
sb.Where("trace_id GLOBAL IN (SELECT trace_id FROM __trace_scope)")
return fmt.Sprintf("__trace_scope AS (%s)", b.traceScope.Query), b.traceScope.Args
}
// Build builds a SQL query for traces based on the given parameters.
func (b *traceQueryStatementBuilder) Build(
ctx context.Context,
@@ -318,6 +348,11 @@ func (b *traceQueryStatementBuilder) buildListQuery(
cteArgs = append(cteArgs, args)
}
if scopeFrag, scopeArgs := b.attachTraceScope(sb); scopeFrag != "" {
cteFragments = append(cteFragments, scopeFrag)
cteArgs = append(cteArgs, scopeArgs)
}
for i, field := range query.SelectFields {
expr, err := b.fm.ColumnExpressionFor(ctx, orgID, start, end, &field, telemetrytypes.FieldDataTypeUnspecified, keys)
if err != nil {
@@ -519,6 +554,11 @@ func (b *traceQueryStatementBuilder) buildTimeSeriesQuery(
cteArgs = append(cteArgs, args)
}
if scopeFrag, scopeArgs := b.attachTraceScope(sb); scopeFrag != "" {
cteFragments = append(cteFragments, scopeFrag)
cteArgs = append(cteArgs, scopeArgs)
}
sb.SelectMore(fmt.Sprintf(
"toStartOfInterval(timestamp, INTERVAL %d SECOND) AS ts",
int64(query.StepInterval.Seconds()),
@@ -679,6 +719,13 @@ func (b *traceQueryStatementBuilder) buildScalarQuery(
cteArgs = append(cteArgs, args)
}
// skipResourceCTE means this scalar is embedded as a CTE of a time-series query,
// which has already emitted the __trace_scope fragment — add only the condition.
if scopeFrag, scopeArgs := b.attachTraceScope(sb); scopeFrag != "" && !skipResourceCTE {
cteFragments = append(cteFragments, scopeFrag)
cteArgs = append(cteArgs, scopeArgs)
}
allAggChArgs := []any{}
fieldNames := make([]string, 0, len(query.GroupBy))

View File

@@ -37,10 +37,10 @@ def ai_trace(
*,
now: datetime,
service: str,
user: str,
in_tokens: int | None,
out_tokens: int,
cost: float,
user: str = "user",
cost: float = 0.1,
model: str = "gpt-4o-mini",
environment: str = "production",
) -> list[Traces]:
@@ -79,6 +79,28 @@ def ai_trace(
]
def tool_only_trace(*, now: datetime, service: str) -> list[Traces]:
"""Root + one tool span: passes the gen_ai gate but has NO LLM span."""
trace_id = TraceIdGenerator.trace_id()
root_id = TraceIdGenerator.span_id()
resources = {"service.name": service}
return [
root_span(now=now, trace_id=trace_id, span_id=root_id, resources=resources, duration_s=2),
Traces(
timestamp=now - timedelta(seconds=4),
duration=timedelta(seconds=0.5),
trace_id=trace_id,
span_id=TraceIdGenerator.span_id(),
parent_span_id=root_id,
name="execute_tool",
kind=TracesKind.SPAN_KIND_INTERNAL,
status_code=TracesStatusCode.STATUS_CODE_OK,
resources=resources,
attributes={"gen_ai.tool.name": "get_weather", "gen_ai.tool.type": "function"},
),
]
def ai_trace_mixed_spans(*, now: datetime, service: str, user: str) -> list[Traces]:
"""Root + LLM + tool + agent spans; only the LLM span carries gen_ai.request.model."""
trace_id = TraceIdGenerator.trace_id()

View File

@@ -77,8 +77,8 @@ def test_ai_list_having_aggregate_filter(
get_token: Callable[[str, str], str],
insert_traces: Callable[[list[Traces]], None],
) -> None:
"""Span + aggregate condition in one filter box splits into WHERE + HAVING; bare
and `trace.` spellings behave identically; an output-only aggregate is rejected."""
"""One filter box splits into WHERE + HAVING; bare and `trace.` spellings behave
identically; an output-only aggregate is rejected."""
now = datetime.now(tz=UTC).replace(second=0, microsecond=0)
service = "ai-it-having"
@@ -214,6 +214,40 @@ def test_ai_span_list_excludes_non_gen_ai_spans(
assert "POST /api/chat" not in names # root span excluded
def test_ai_span_list_trace_level_filter(
signoz: types.SigNoz,
create_user_admin: None, # pylint: disable=unused-argument
get_token: Callable[[str, str], str],
insert_traces: Callable[[list[Traces]], None],
) -> None:
"""Span list (raw) with a trace-level condition qualifies whole traces first:
out-tokens 100 / 300 and `trace.output_tokens > 100` keeps only the large trace's span."""
now = datetime.now(tz=UTC).replace(second=0, microsecond=0)
service = "ai-it-spanlist-tracefilter"
small = ai_trace(now=now, service=service, user="a", in_tokens=10, out_tokens=100, cost=0.1)
large = ai_trace(now=now, service=service, user="b", in_tokens=30, out_tokens=300, cost=0.2)
insert_traces(small + large)
token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
start_ms, end_ms = query_window(now)
query = BuilderQuery(
signal="traces",
query_type="builder_ai_query",
name="A",
filter_expression=f"service.name = '{service}' AND trace.output_tokens > 100",
limit=10,
)
resp = make_query_request(signoz, token, start_ms, end_ms, [query.to_dict()], request_type=RequestType.RAW)
assert resp.status_code == HTTPStatus.OK, resp.text
rows = resp.json()["data"]["data"]["results"][0]["rows"]
assert len(rows) == 1, f"expected only the large trace's LLM span, got {len(rows)} rows"
body = json.dumps(rows)
assert large[0].trace_id in body
assert small[0].trace_id not in body
def test_ai_list_having_or_aggregates(
signoz: types.SigNoz,
create_user_admin: None, # pylint: disable=unused-argument
@@ -327,9 +361,8 @@ def test_ai_list_nested_group_span_or_and_aggregate(
get_token: Callable[[str, str], str],
insert_traces: Callable[[list[Traces]], None],
) -> None:
"""service.name = X AND (has_error = true OR gen_ai.request.model = 'gpt-4o') AND
total_tokens > 100: the nested OR group must not flatten, span predicates go to
WHERE, the aggregate to HAVING."""
"""A nested (span OR span) group ANDed with an aggregate must not flatten: span
predicates go to WHERE, the aggregate to HAVING."""
now = datetime.now(tz=UTC).replace(second=0, microsecond=0)
service = "ai-it-nested"

View File

@@ -0,0 +1,489 @@
"""
builder_ai_query scalar / time-series aggregations. The `trace.` prefix picks the
domain per expression: bare keys aggregate over gen_ai spans, trace.* over
window-clipped per-trace values; a trace-level filter condition qualifies whole
traces in every request type. Tests isolate their data via unique service.name.
"""
from collections.abc import Callable
from datetime import UTC, datetime, timedelta
from http import HTTPStatus
import pytest
from fixtures import types
from fixtures.auth import USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD
from fixtures.querier import (
Aggregation,
BuilderQuery,
OrderBy,
RequestType,
TelemetryFieldKey,
get_scalar_table_data,
make_query_request,
)
from fixtures.querierai import ai_trace, query_window, tool_only_trace
from fixtures.traces import TraceIdGenerator, Traces, TracesKind, TracesStatusCode
def scalar_query(
service: str,
expression: str,
*,
filter_extra: str = "",
group_by: list[TelemetryFieldKey] | None = None,
alias: str | None = None,
having: str | None = None,
limit: int | None = None,
) -> dict:
filter_expression = f"service.name = '{service}'"
if filter_extra:
filter_expression += f" AND {filter_extra}"
return BuilderQuery(
signal="traces",
query_type="builder_ai_query",
name="A",
filter_expression=filter_expression,
aggregations=[Aggregation(expression=expression, alias=alias)],
group_by=group_by,
having_expression=having,
limit=limit,
).to_dict()
def scalar_value(signoz: types.SigNoz, token: str, start_ms: int, end_ms: int, service: str, expression: str) -> float:
"""Run one single-aggregation scalar query and return its value."""
resp = make_query_request(
signoz,
token,
start_ms,
end_ms,
[scalar_query(service, expression)],
request_type=RequestType.SCALAR,
)
assert resp.status_code == HTTPStatus.OK, f"{expression}: {resp.text}"
data = get_scalar_table_data(resp.json())
assert len(data) == 1, f"{expression}: expected one row, got {data}"
return float(data[0][-1])
def series_values(response_json: dict) -> list[list[float]]:
"""Per-series lists of bucket values (bucket order as returned)."""
series = response_json["data"]["data"]["results"][0]["aggregations"][0]["series"]
return [[v["value"] for v in ser["values"]] for ser in series]
def test_ai_scalar_trace_level_aggregations(
signoz: types.SigNoz,
create_user_admin: None, # pylint: disable=unused-argument
get_token: Callable[[str, str], str],
insert_traces: Callable[[list[Traces]], None],
) -> None:
"""Trace-level scalars over per-trace values: out-tokens 100/300 give avg=200 and
count=2, while the span-level count() sees the two LLM spans (roots gated out)."""
now = datetime.now(tz=UTC).replace(second=0, microsecond=0)
service = "ai-it-agg-scalar"
insert_traces(ai_trace(now=now, service=service, in_tokens=10, out_tokens=100) + ai_trace(now=now, service=service, in_tokens=30, out_tokens=300))
token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
start_ms, end_ms = query_window(now)
def value(expression: str) -> float:
return scalar_value(signoz, token, start_ms, end_ms, service, expression)
assert value("avg(trace.output_tokens)") == pytest.approx(200)
assert value("count(trace.trace_id)") == 2
assert value("max(trace.total_tokens)") == pytest.approx(330)
assert value("p50(trace.output_tokens)") == pytest.approx(200) # AggreFuncMap -> quantile(0.50)
# arithmetic inside one function and between functions
assert value("avg(trace.output_tokens + trace.input_tokens)") == pytest.approx(220)
assert value("sum(trace.output_tokens)/count(trace.trace_id)") == pytest.approx(200)
# span-level domain still works through the same request type
assert value("count()") == 2 # the two LLM spans; roots are not gen_ai
assert value("sum(gen_ai.usage.output_tokens)") == pytest.approx(400)
# multiple trace-level aggregations in one query -> one column per aggregation
multi = BuilderQuery(
signal="traces",
query_type="builder_ai_query",
name="A",
filter_expression=f"service.name = '{service}'",
aggregations=[Aggregation(expression="avg(trace.output_tokens)"), Aggregation(expression="count(trace.trace_id)")],
)
resp = make_query_request(signoz, token, start_ms, end_ms, [multi.to_dict()], request_type=RequestType.SCALAR)
assert resp.status_code == HTTPStatus.OK, resp.text
data = get_scalar_table_data(resp.json())
assert len(data) == 1 and [float(v) for v in data[0]] == [pytest.approx(200), 2], data
def test_ai_scalar_trace_level_filter_qualifies_traces(
signoz: types.SigNoz,
create_user_admin: None, # pylint: disable=unused-argument
get_token: Callable[[str, str], str],
insert_traces: Callable[[list[Traces]], None],
) -> None:
"""`trace.output_tokens > 100` qualifies whole traces before aggregation: with
out-tokens 100/300 only the 300 trace survives, on both aggregation domains."""
now = datetime.now(tz=UTC).replace(second=0, microsecond=0)
service = "ai-it-agg-qualify"
insert_traces(ai_trace(now=now, service=service, in_tokens=10, out_tokens=100) + ai_trace(now=now, service=service, in_tokens=30, out_tokens=300))
token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
start_ms, end_ms = query_window(now)
for expression, expected in (
("sum(trace.output_tokens)", 300), # native trace-domain path
("sum(gen_ai.usage.output_tokens)", 300), # delegated span-domain path (__trace_scope)
):
resp = make_query_request(
signoz,
token,
start_ms,
end_ms,
[scalar_query(service, expression, filter_extra="trace.output_tokens > 100")],
request_type=RequestType.SCALAR,
)
assert resp.status_code == HTTPStatus.OK, resp.text
data = get_scalar_table_data(resp.json())
assert len(data) == 1 and float(data[0][-1]) == pytest.approx(expected), f"{expression}: {data}"
# the qualification also constrains delegated (span-domain) time series
ts = BuilderQuery(
signal="traces",
query_type="builder_ai_query",
name="A",
filter_expression=f"service.name = '{service}' AND trace.output_tokens > 100",
aggregations=[Aggregation(expression="sum(gen_ai.usage.output_tokens)")],
step_interval=60,
)
resp = make_query_request(signoz, token, start_ms, end_ms, [ts.to_dict()], request_type=RequestType.TIME_SERIES)
assert resp.status_code == HTTPStatus.OK, resp.text
values = series_values(resp.json())
assert values == [[pytest.approx(300)]], values
def test_ai_scalar_group_by_model(
signoz: types.SigNoz,
create_user_admin: None, # pylint: disable=unused-argument
get_token: Callable[[str, str], str],
insert_traces: Callable[[list[Traces]], None],
) -> None:
"""Trace-level aggregation grouped by a span attribute: per-model avg of per-trace tokens."""
now = datetime.now(tz=UTC).replace(second=0, microsecond=0)
service = "ai-it-agg-groupby"
insert_traces(ai_trace(now=now, service=service, in_tokens=10, out_tokens=100, model="gpt-4o") + ai_trace(now=now, service=service, in_tokens=10, out_tokens=300, model="gpt-4o") + ai_trace(now=now, service=service, in_tokens=10, out_tokens=50, model="gpt-4o-mini"))
token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
start_ms, end_ms = query_window(now)
resp = make_query_request(
signoz,
token,
start_ms,
end_ms,
[scalar_query(service, "avg(trace.output_tokens)", group_by=[TelemetryFieldKey(name="gen_ai.request.model")])],
request_type=RequestType.SCALAR,
)
assert resp.status_code == HTTPStatus.OK, resp.text
data = get_scalar_table_data(resp.json())
by_model = {row[0]: float(row[-1]) for row in data}
assert by_model == {"gpt-4o": pytest.approx(200), "gpt-4o-mini": pytest.approx(50)}, data
def test_ai_timeseries_trace_level_aggregation(
signoz: types.SigNoz,
create_user_admin: None, # pylint: disable=unused-argument
get_token: Callable[[str, str], str],
insert_traces: Callable[[list[Traces]], None],
) -> None:
"""Time-series over per-trace values: all spans fall in one step bucket, avg=200."""
now = datetime.now(tz=UTC).replace(second=0, microsecond=0)
service = "ai-it-agg-ts"
insert_traces(ai_trace(now=now, service=service, in_tokens=10, out_tokens=100) + ai_trace(now=now, service=service, in_tokens=30, out_tokens=300))
token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
start_ms, end_ms = query_window(now)
query = BuilderQuery(
signal="traces",
query_type="builder_ai_query",
name="A",
filter_expression=f"service.name = '{service}'",
aggregations=[Aggregation(expression="avg(trace.output_tokens)")],
step_interval=60,
)
resp = make_query_request(
signoz,
token,
start_ms,
end_ms,
[query.to_dict()],
request_type=RequestType.TIME_SERIES,
)
assert resp.status_code == HTTPStatus.OK, resp.text
values = series_values(resp.json())
assert values == [[pytest.approx(200)]], values
def test_ai_timeseries_top_n_groups(
signoz: types.SigNoz,
create_user_admin: None, # pylint: disable=unused-argument
get_token: Callable[[str, str], str],
insert_traces: Callable[[list[Traces]], None],
) -> None:
"""Grouped, limited time series ranks groups on whole-window per-trace values:
gpt-4o sums to 400 vs gpt-4o-mini's 50, so limit=1 keeps only gpt-4o."""
now = datetime.now(tz=UTC).replace(second=0, microsecond=0)
service = "ai-it-agg-topn"
insert_traces(ai_trace(now=now, service=service, in_tokens=10, out_tokens=300, model="gpt-4o") + ai_trace(now=now, service=service, in_tokens=10, out_tokens=100, model="gpt-4o") + ai_trace(now=now, service=service, in_tokens=10, out_tokens=50, model="gpt-4o-mini"))
token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
start_ms, end_ms = query_window(now)
query = BuilderQuery(
signal="traces",
query_type="builder_ai_query",
name="A",
filter_expression=f"service.name = '{service}'",
aggregations=[Aggregation(expression="sum(trace.output_tokens)")],
group_by=[TelemetryFieldKey(name="gen_ai.request.model")],
step_interval=60,
limit=1,
)
resp = make_query_request(signoz, token, start_ms, end_ms, [query.to_dict()], request_type=RequestType.TIME_SERIES)
assert resp.status_code == HTTPStatus.OK, resp.text
series = resp.json()["data"]["data"]["results"][0]["aggregations"][0]["series"]
assert len(series) == 1, f"limit=1 must keep only the top group, got {len(series)} series"
assert series[0]["labels"][0]["value"] == "gpt-4o", series[0]["labels"]
assert [v["value"] for v in series[0]["values"]] == [pytest.approx(400)]
def test_ai_timeseries_span_time_bucketing(
signoz: types.SigNoz,
create_user_admin: None, # pylint: disable=unused-argument
get_token: Callable[[str, str], str],
insert_traces: Callable[[list[Traces]], None],
) -> None:
"""Per-trace values are clipped per (bucket, trace): two LLM calls two minutes
apart contribute each call's tokens to its own bucket, not the total to both."""
now = datetime.now(tz=UTC).replace(second=0, microsecond=0)
service = "ai-it-agg-buckets"
trace_id = TraceIdGenerator.trace_id()
root_id = TraceIdGenerator.span_id()
resources = {"service.name": service}
def llm(offset_s: float, out_tokens: int) -> Traces:
return Traces(
timestamp=now - timedelta(seconds=offset_s),
duration=timedelta(seconds=1),
trace_id=trace_id,
span_id=TraceIdGenerator.span_id(),
parent_span_id=root_id,
name="chat",
kind=TracesKind.SPAN_KIND_CLIENT,
status_code=TracesStatusCode.STATUS_CODE_OK,
resources=resources,
attributes={"gen_ai.request.model": "gpt-4o-mini", "gen_ai.usage.output_tokens": out_tokens},
)
root = Traces(
timestamp=now - timedelta(seconds=130),
duration=timedelta(seconds=130),
trace_id=trace_id,
span_id=root_id,
parent_span_id="",
name="POST /api/chat",
kind=TracesKind.SPAN_KIND_SERVER,
status_code=TracesStatusCode.STATUS_CODE_OK,
resources=resources,
attributes={"http.request.method": "POST"},
)
insert_traces([root, llm(124, 100), llm(4, 300)])
token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
start_ms, end_ms = query_window(now)
query = BuilderQuery(
signal="traces",
query_type="builder_ai_query",
name="A",
filter_expression=f"service.name = '{service}'",
aggregations=[Aggregation(expression="avg(trace.output_tokens)")],
step_interval=60,
)
resp = make_query_request(signoz, token, start_ms, end_ms, [query.to_dict()], request_type=RequestType.TIME_SERIES)
assert resp.status_code == HTTPStatus.OK, resp.text
values = series_values(resp.json())
assert len(values) == 1, values
assert sorted(values[0]) == [pytest.approx(100), pytest.approx(300)], f"each call's tokens in its own bucket: {values}"
def test_ai_scalar_variables_in_trace_level_filter(
signoz: types.SigNoz,
create_user_admin: None, # pylint: disable=unused-argument
get_token: Callable[[str, str], str],
insert_traces: Callable[[list[Traces]], None],
) -> None:
"""Variables resolve inside trace-level conditions with span-filter semantics;
an unresolvable $var is a 400 (today via aggregate validation — a targeted
unknown-variable error is a separate concern)."""
now = datetime.now(tz=UTC).replace(second=0, microsecond=0)
service = "ai-it-agg-vars"
insert_traces(ai_trace(now=now, service=service, in_tokens=10, out_tokens=100) + ai_trace(now=now, service=service, in_tokens=30, out_tokens=300))
token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
start_ms, end_ms = query_window(now)
query = scalar_query(service, "sum(trace.output_tokens)", filter_extra="trace.output_tokens > $threshold")
resp = make_query_request(
signoz,
token,
start_ms,
end_ms,
[query],
request_type=RequestType.SCALAR,
variables={"threshold": {"type": "text", "value": 100}},
)
assert resp.status_code == HTTPStatus.OK, resp.text
data = get_scalar_table_data(resp.json())
assert len(data) == 1 and float(data[0][-1]) == pytest.approx(300), data
resp = make_query_request(
signoz,
token,
start_ms,
end_ms,
[query],
request_type=RequestType.SCALAR,
)
assert resp.status_code == HTTPStatus.BAD_REQUEST, resp.text
assert '"$threshold" cannot be used' in resp.text
# a dynamic variable resolved to __all__ drops the condition (both traces count)
resp = make_query_request(
signoz,
token,
start_ms,
end_ms,
[query],
request_type=RequestType.SCALAR,
variables={"threshold": {"type": "dynamic", "value": "__all__"}},
)
assert resp.status_code == HTTPStatus.OK, resp.text
data = get_scalar_table_data(resp.json())
assert len(data) == 1 and float(data[0][-1]) == pytest.approx(400), data
def test_ai_scalar_activity_gate_excludes_tool_only_traces(
signoz: types.SigNoz,
create_user_admin: None, # pylint: disable=unused-argument
get_token: Callable[[str, str], str],
insert_traces: Callable[[list[Traces]], None],
) -> None:
"""A tool-only trace (in the gate, no LLM span) must not feed trace-level
aggregations: count(trace.trace_id) skips it, span-level count() still sees it."""
now = datetime.now(tz=UTC).replace(second=0, microsecond=0)
service = "ai-it-agg-gate"
insert_traces(ai_trace(now=now, service=service, in_tokens=10, out_tokens=100) + tool_only_trace(now=now, service=service))
token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
start_ms, end_ms = query_window(now)
def value(expression: str) -> float:
return scalar_value(signoz, token, start_ms, end_ms, service, expression)
# count and avg agree on the trace set — the gate's purpose
assert value("count(trace.trace_id)") == 1, "tool-only trace must be dropped by the LLM-activity gate"
assert value("avg(trace.output_tokens)") == pytest.approx(100), "avg over the same gated trace set"
assert value("count()") == 2, "span-level count still sees the tool span"
def test_ai_scalar_having_on_aggregation(
signoz: types.SigNoz,
create_user_admin: None, # pylint: disable=unused-argument
get_token: Callable[[str, str], str],
insert_traces: Callable[[list[Traces]], None],
) -> None:
"""The outer having filters aggregation results per group (by alias)."""
now = datetime.now(tz=UTC).replace(second=0, microsecond=0)
service = "ai-it-agg-having"
insert_traces(ai_trace(now=now, service=service, in_tokens=10, out_tokens=300, model="gpt-4o") + ai_trace(now=now, service=service, in_tokens=10, out_tokens=50, model="gpt-4o-mini"))
token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
start_ms, end_ms = query_window(now)
resp = make_query_request(
signoz,
token,
start_ms,
end_ms,
[
scalar_query(
service,
"avg(trace.output_tokens)",
group_by=[TelemetryFieldKey(name="gen_ai.request.model")],
alias="avg_out",
having="avg_out > 100",
)
],
request_type=RequestType.SCALAR,
)
assert resp.status_code == HTTPStatus.OK, resp.text
data = get_scalar_table_data(resp.json())
assert len(data) == 1 and data[0][0] == "gpt-4o", data
def test_ai_aggregation_rejections(
signoz: types.SigNoz,
create_user_admin: None, # pylint: disable=unused-argument
get_token: Callable[[str, str], str],
insert_traces: Callable[[list[Traces]], None],
) -> None:
"""Targeted 400s: mixed domains, group-by on a trace column, raw order by a trace column."""
now = datetime.now(tz=UTC).replace(second=0, microsecond=0)
service = "ai-it-agg-reject"
insert_traces(ai_trace(now=now, service=service, in_tokens=10, out_tokens=100))
token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
start_ms, end_ms = query_window(now)
# span-level and trace-level aggregations cannot be mixed in one query
mixed = BuilderQuery(
signal="traces",
query_type="builder_ai_query",
name="A",
filter_expression=f"service.name = '{service}'",
aggregations=[Aggregation(expression="avg(trace.output_tokens)"), Aggregation(expression="count()")],
)
resp = make_query_request(signoz, token, start_ms, end_ms, [mixed.to_dict()], request_type=RequestType.SCALAR)
assert resp.status_code == HTTPStatus.BAD_REQUEST, resp.text
assert "cannot be mixed" in resp.text
# grouping by a trace-level per-trace column is rejected with a targeted error
bad_group = BuilderQuery(
signal="traces",
query_type="builder_ai_query",
name="A",
filter_expression=f"service.name = '{service}'",
aggregations=[Aggregation(expression="avg(trace.output_tokens)")],
group_by=[TelemetryFieldKey(name="trace.llm_call_count")],
)
resp = make_query_request(signoz, token, start_ms, end_ms, [bad_group.to_dict()], request_type=RequestType.SCALAR)
assert resp.status_code == HTTPStatus.BAD_REQUEST, resp.text
assert "grouping by trace-level aggregate" in resp.text
# ordering the span list by a trace-level column is rejected with a targeted error
bad_order = BuilderQuery(
signal="traces",
query_type="builder_ai_query",
name="A",
filter_expression=f"service.name = '{service}'",
order=[OrderBy(key=TelemetryFieldKey(name="trace.output_tokens"), direction="desc")],
limit=10,
)
resp = make_query_request(signoz, token, start_ms, end_ms, [bad_order.to_dict()], request_type=RequestType.RAW)
assert resp.status_code == HTTPStatus.BAD_REQUEST, resp.text
assert "ordering the span list by trace-level aggregate" in resp.text

View File

@@ -0,0 +1,34 @@
import pytest
from testcontainers.core.container import Network
from fixtures import types
from fixtures.signoz import create_signoz
@pytest.fixture(name="signoz", scope="package")
def signoz_ai_observability(
network: Network,
migrator: types.Operation, # pylint: disable=unused-argument
zeus: types.TestContainerDocker,
gateway: types.TestContainerDocker,
sqlstore: types.TestContainerSQL,
clickhouse: types.TestContainerClickhouse,
request: pytest.FixtureRequest,
pytestconfig: pytest.Config,
) -> types.SigNoz:
"""Package-scoped SigNoz with AI observability enabled: the flag gates the static
gen_ai key definitions (enrichWithGenAIKeys) — without it the gate keys only
resolve once a span carrying them has been ingested."""
return create_signoz(
network=network,
zeus=zeus,
gateway=gateway,
sqlstore=sqlstore,
clickhouse=clickhouse,
request=request,
pytestconfig=pytestconfig,
cache_key="signoz-ai-observability",
env_overrides={
"SIGNOZ_FLAGGER_CONFIG_BOOLEAN_ENABLE__AI__OBSERVABILITY": True,
},
)