signoz/pkg/telemetrytraces/trace_operator_statement_builder_test.go

426 lines
12 KiB
Go

package telemetrytraces
import (
"context"
"strings"
"testing"
"time"
"github.com/SigNoz/signoz/pkg/instrumentation/instrumentationtest"
qbtypes "github.com/SigNoz/signoz/pkg/types/querybuildertypes/querybuildertypesv5"
"github.com/SigNoz/signoz/pkg/types/telemetrytypes"
"github.com/SigNoz/signoz/pkg/types/telemetrytypes/telemetrytypestest"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestNewTraceOperatorStatementBuilder(t *testing.T) {
fm := NewFieldMapper()
cb := NewConditionBuilder(fm)
mockMetadataStore := telemetrytypestest.NewMockMetadataStore()
mockMetadataStore.KeysMap = buildCompleteFieldKeyMap()
// Create a mock trace statement builder (you'd need to implement this)
traceStmtBuilder := NewTraceQueryStatementBuilder(
instrumentationtest.New().ToProviderSettings(),
mockMetadataStore,
fm,
cb,
nil, // resource filter builder
nil, // agg expr rewriter
)
aggExprRewriter := &mockAggExprRewriter{}
settings := instrumentationtest.New().ToProviderSettings()
builder := NewTraceOperatorStatementBuilder(
settings,
mockMetadataStore,
fm,
cb,
traceStmtBuilder,
aggExprRewriter,
)
assert.NotNil(t, builder)
assert.NotNil(t, builder.logger)
assert.Equal(t, mockMetadataStore, builder.metadataStore)
assert.Equal(t, fm, builder.fm)
assert.Equal(t, cb, builder.cb)
assert.Equal(t, traceStmtBuilder, builder.traceStmtBuilder)
assert.Equal(t, aggExprRewriter, builder.aggExprRewriter)
}
func TestTraceOperatorStatementBuilder_Build(t *testing.T) {
testCases := []struct {
name string
requestType qbtypes.RequestType
query qbtypes.QueryBuilderTraceOperator
compositeQuery *qbtypes.CompositeQuery
expectedError string
validateSQL func(string) bool
}{
{
name: "simple direct descendant - raw",
requestType: qbtypes.RequestTypeRaw,
query: qbtypes.QueryBuilderTraceOperator{
Name: "test_trace_op",
Expression: "A => B",
StepInterval: qbtypes.Step{
Duration: time.Minute,
},
},
compositeQuery: &qbtypes.CompositeQuery{
Queries: []qbtypes.QueryEnvelope{
{
Type: qbtypes.QueryTypeBuilder,
Spec: qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Name: "A",
Signal: telemetrytypes.SignalTraces,
Filter: &qbtypes.Filter{
Expression: "service.name = 'cartservice'",
},
},
},
{
Type: qbtypes.QueryTypeBuilder,
Spec: qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Name: "B",
Signal: telemetrytypes.SignalTraces,
Filter: &qbtypes.Filter{
Expression: "service.name = 'shippingservice'",
},
},
},
},
},
expectedError: "",
validateSQL: func(sql string) bool {
return strings.Contains(sql, "WITH") &&
strings.Contains(sql, "base_spans AS") &&
strings.Contains(sql, "A AS") &&
strings.Contains(sql, "B AS") &&
strings.Contains(sql, "A_=>_B AS") && // The actual CTE name format
strings.Contains(sql, "INNER JOIN") &&
strings.Contains(sql, "timestamp") &&
strings.Contains(sql, "trace_id")
},
},
{
name: "OR operation - scalar",
requestType: qbtypes.RequestTypeScalar,
query: qbtypes.QueryBuilderTraceOperator{
Name: "or_trace_op",
Expression: "A || B",
Aggregations: []qbtypes.TraceAggregation{
{Expression: "count()"},
},
},
compositeQuery: &qbtypes.CompositeQuery{
Queries: []qbtypes.QueryEnvelope{
{
Type: qbtypes.QueryTypeBuilder,
Spec: qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Name: "A",
Signal: telemetrytypes.SignalTraces,
},
},
{
Type: qbtypes.QueryTypeBuilder,
Spec: qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Name: "B",
Signal: telemetrytypes.SignalTraces,
},
},
},
},
expectedError: "",
validateSQL: func(sql string) bool {
return strings.Contains(sql, "UNION DISTINCT") &&
strings.Contains(sql, "subQuery.serviceName") &&
strings.Contains(sql, "span_count")
},
},
{
name: "time series request",
requestType: qbtypes.RequestTypeTimeSeries,
query: qbtypes.QueryBuilderTraceOperator{
Name: "ts_trace_op",
Expression: "A && B",
StepInterval: qbtypes.Step{
Duration: time.Minute,
},
Aggregations: []qbtypes.TraceAggregation{
{Expression: "count()"},
},
},
compositeQuery: &qbtypes.CompositeQuery{
Queries: []qbtypes.QueryEnvelope{
{
Type: qbtypes.QueryTypeBuilder,
Spec: qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Name: "A",
Signal: telemetrytypes.SignalTraces,
},
},
{
Type: qbtypes.QueryTypeBuilder,
Spec: qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Name: "B",
Signal: telemetrytypes.SignalTraces,
},
},
},
},
expectedError: "",
validateSQL: func(sql string) bool {
return strings.Contains(sql, "toStartOfInterval") &&
strings.Contains(sql, "INTERVAL 60 SECOND")
},
},
{
name: "invalid expression",
requestType: qbtypes.RequestTypeRaw,
query: qbtypes.QueryBuilderTraceOperator{
Name: "invalid_trace_op",
Expression: "invalid expression syntax",
},
compositeQuery: &qbtypes.CompositeQuery{
Queries: []qbtypes.QueryEnvelope{},
},
expectedError: "invalid query reference",
},
{
name: "missing composite query context",
requestType: qbtypes.RequestTypeRaw,
query: qbtypes.QueryBuilderTraceOperator{
Name: "missing_context_op",
Expression: "A => B",
},
compositeQuery: nil, // This will result in missing context
expectedError: "composite query not found in context",
},
{
name: "missing referenced query",
requestType: qbtypes.RequestTypeRaw,
query: qbtypes.QueryBuilderTraceOperator{
Name: "missing_ref_op",
Expression: "A => C", // C is not defined in composite query
},
compositeQuery: &qbtypes.CompositeQuery{
Queries: []qbtypes.QueryEnvelope{
{
Type: qbtypes.QueryTypeBuilder,
Spec: qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Name: "A",
Signal: telemetrytypes.SignalTraces,
},
},
},
},
expectedError: "referenced query 'C' not found",
},
{
name: "complex expression with filters",
requestType: qbtypes.RequestTypeScalar,
query: qbtypes.QueryBuilderTraceOperator{
Name: "complex_trace_op",
Expression: "(A => B) && (C || D)",
Filter: &qbtypes.Filter{
Expression: "duration > 1000",
},
Order: []qbtypes.OrderBy{
{
Key: qbtypes.OrderByKey{
TelemetryFieldKey: telemetrytypes.TelemetryFieldKey{
Name: qbtypes.OrderByTraceDuration.StringValue(),
},
},
Direction: qbtypes.OrderDirectionDesc,
},
},
},
compositeQuery: &qbtypes.CompositeQuery{
Queries: []qbtypes.QueryEnvelope{
{
Type: qbtypes.QueryTypeBuilder,
Spec: qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Name: "A",
Signal: telemetrytypes.SignalTraces,
},
},
{
Type: qbtypes.QueryTypeBuilder,
Spec: qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Name: "B",
Signal: telemetrytypes.SignalTraces,
},
},
{
Type: qbtypes.QueryTypeBuilder,
Spec: qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Name: "C",
Signal: telemetrytypes.SignalTraces,
},
},
{
Type: qbtypes.QueryTypeBuilder,
Spec: qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Name: "D",
Signal: telemetrytypes.SignalTraces,
},
},
},
},
expectedError: "",
validateSQL: func(sql string) bool {
return strings.Contains(sql, "ORDER BY") &&
strings.Contains(sql, "subQuery.durationNano")
},
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
fm := NewFieldMapper()
cb := NewConditionBuilder(fm)
mockMetadataStore := telemetrytypestest.NewMockMetadataStore()
mockMetadataStore.KeysMap = buildCompleteFieldKeyMap()
traceStmtBuilder := NewTraceQueryStatementBuilder(
instrumentationtest.New().ToProviderSettings(),
mockMetadataStore,
fm,
cb,
nil, // resource filter builder
nil, // agg expr rewriter
)
aggExprRewriter := &mockAggExprRewriter{}
builder := &traceOperatorStatementBuilder{
metadataStore: mockMetadataStore,
fm: fm,
cb: cb,
traceStmtBuilder: traceStmtBuilder,
aggExprRewriter: aggExprRewriter,
}
var ctx context.Context
if tc.compositeQuery != nil {
ctx = context.WithValue(context.Background(), compositeQueryKey, tc.compositeQuery)
} else {
ctx = context.Background()
}
stmt, err := builder.Build(ctx, 1000, 2000, tc.requestType, tc.query)
if tc.expectedError != "" {
require.Error(t, err)
assert.Contains(t, err.Error(), tc.expectedError)
assert.Nil(t, stmt)
} else {
require.NoError(t, err)
require.NotNil(t, stmt)
assert.NotEmpty(t, stmt.Query)
if tc.validateSQL != nil {
assert.True(t, tc.validateSQL(stmt.Query), "SQL validation failed for: %s", stmt.Query)
}
}
})
}
}
func TestTraceOperatorStatementBuilder_Build_WithSelectFields(t *testing.T) {
fm := NewFieldMapper()
cb := NewConditionBuilder(fm)
mockMetadataStore := telemetrytypestest.NewMockMetadataStore()
mockMetadataStore.KeysMap = buildCompleteFieldKeyMap()
traceStmtBuilder := NewTraceQueryStatementBuilder(
instrumentationtest.New().ToProviderSettings(),
mockMetadataStore,
fm,
cb,
nil,
nil,
)
aggExprRewriter := &mockAggExprRewriter{}
builder := &traceOperatorStatementBuilder{
metadataStore: mockMetadataStore,
fm: fm,
cb: cb,
traceStmtBuilder: traceStmtBuilder,
aggExprRewriter: aggExprRewriter,
}
query := qbtypes.QueryBuilderTraceOperator{
Name: "trace_op_with_fields",
Expression: "A => B",
SelectFields: []telemetrytypes.TelemetryFieldKey{
{
Name: "http.method",
FieldContext: telemetrytypes.FieldContextAttribute,
FieldDataType: telemetrytypes.FieldDataTypeString,
Signal: telemetrytypes.SignalTraces,
},
{
Name: "http.status_code",
FieldContext: telemetrytypes.FieldContextAttribute,
FieldDataType: telemetrytypes.FieldDataTypeNumber,
Signal: telemetrytypes.SignalTraces,
},
},
}
compositeQuery := &qbtypes.CompositeQuery{
Queries: []qbtypes.QueryEnvelope{
{
Type: qbtypes.QueryTypeBuilder,
Spec: qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Name: "A",
Signal: telemetrytypes.SignalTraces,
},
},
{
Type: qbtypes.QueryTypeBuilder,
Spec: qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
Name: "B",
Signal: telemetrytypes.SignalTraces,
},
},
},
}
ctx := context.WithValue(context.Background(), compositeQueryKey, compositeQuery)
stmt, err := builder.Build(ctx, 1000, 2000, qbtypes.RequestTypeRaw, query)
require.NoError(t, err)
require.NotNil(t, stmt)
assert.NotEmpty(t, stmt.Query)
// Verify that the base_spans CTE includes the additional select fields
assert.Contains(t, stmt.Query, "base_spans AS")
}
// Mock implementation for testing
type mockAggExprRewriter struct{}
func (m *mockAggExprRewriter) Rewrite(ctx context.Context, expr string, rateInterval uint64, keys map[string][]*telemetrytypes.TelemetryFieldKey) (string, []any, error) {
// Simple mock implementation
return expr, []any{}, nil
}
func (m *mockAggExprRewriter) RewriteMulti(ctx context.Context, exprs []string, rateInterval uint64, keys map[string][]*telemetrytypes.TelemetryFieldKey) ([]string, [][]any, error) {
args := make([][]any, len(exprs))
for i := range args {
args[i] = []any{}
}
return exprs, args, nil
}