feat: support for new enrichment logic in traces

pkg/query-service/app/logs/v3/enrich_query.go

@@ -142,7 +142,7 @@ func enrichFieldWithMetadata(field v3.AttributeKey, fields map[string]v3.Attribu
 	}
 
 	// check if the field is present in the fields map
-	for _, key := range utils.GenerateLogEnrichmentKeys(field) {
+	for _, key := range utils.GenerateEnrichmentKeys(field) {
 		if val, ok := fields[key]; ok {
 			return val
 		}

pkg/query-service/app/traces/v4/enrich.go

@@ -0,0 +1,118 @@
+package v4
+
+import (
+	"go.signoz.io/signoz/pkg/query-service/constants"
+	v3 "go.signoz.io/signoz/pkg/query-service/model/v3"
+	"go.signoz.io/signoz/pkg/query-service/utils"
+)
+
+// if the field is timestamp/id/value we don't need to enrich
+// if the field is static we don't need to enrich
+// for all others we need to enrich
+// an attribute/resource can be materialized/dematerialized
+// but the query should work regardless and shouldn't fail
+func isEnriched(field v3.AttributeKey) bool {
+	// if it is timestamp/id dont check
+	if field.Key == "timestamp" || field.Key == constants.SigNozOrderByValue {
+		return true
+	}
+
+	// we need to check if the field is static and return false if isColumn is not set
+	if _, ok := constants.StaticFieldsTraces[field.Key]; ok && field.IsColumn {
+		return true
+	}
+
+	return false
+}
+
+func enrichKeyWithMetadata(key v3.AttributeKey, keys map[string]v3.AttributeKey) v3.AttributeKey {
+	if isEnriched(key) {
+		return key
+	}
+
+	if v, ok := constants.StaticFieldsTraces[key.Key]; ok {
+		return v
+	}
+
+	for _, key := range utils.GenerateEnrichmentKeys(key) {
+		if val, ok := keys[key]; ok {
+			return val
+		}
+	}
+
+	// enrich with default values if metadata is not found
+	if key.Type == "" {
+		key.Type = v3.AttributeKeyTypeTag
+	}
+	if key.DataType == "" {
+		key.DataType = v3.AttributeKeyDataTypeString
+	}
+	return key
+}
+
+func Enrich(params *v3.QueryRangeParamsV3, keys map[string]v3.AttributeKey) {
+	if params.CompositeQuery.QueryType != v3.QueryTypeBuilder {
+		return
+	}
+
+	for _, query := range params.CompositeQuery.BuilderQueries {
+		if query.DataSource == v3.DataSourceTraces {
+			EnrichTracesQuery(query, keys)
+		}
+	}
+}
+
+func EnrichTracesQuery(query *v3.BuilderQuery, keys map[string]v3.AttributeKey) {
+	// enrich aggregate attribute
+	query.AggregateAttribute = enrichKeyWithMetadata(query.AggregateAttribute, keys)
+
+	// enrich filter items
+	if query.Filters != nil && len(query.Filters.Items) > 0 {
+		for idx, filter := range query.Filters.Items {
+			query.Filters.Items[idx].Key = enrichKeyWithMetadata(filter.Key, keys)
+			// if the serviceName column is used, use the corresponding resource attribute as well during filtering
+			// since there is only one of these resource attributes we are adding it here directly.
+			// move it somewhere else if this list is big
+			if filter.Key.Key == "serviceName" {
+				query.Filters.Items[idx].Key = v3.AttributeKey{
+					Key:      "service.name",
+					DataType: v3.AttributeKeyDataTypeString,
+					Type:     v3.AttributeKeyTypeResource,
+					IsColumn: false,
+				}
+			}
+		}
+	}
+
+	// enrich group by
+	for idx, groupBy := range query.GroupBy {
+		query.GroupBy[idx] = enrichKeyWithMetadata(groupBy, keys)
+	}
+
+	// enrich order by
+	query.OrderBy = enrichOrderBy(query.OrderBy, keys)
+
+	// enrich select columns
+	for idx, selectColumn := range query.SelectColumns {
+		query.SelectColumns[idx] = enrichKeyWithMetadata(selectColumn, keys)
+	}
+
+}
+
+func enrichOrderBy(items []v3.OrderBy, keys map[string]v3.AttributeKey) []v3.OrderBy {
+	enrichedItems := []v3.OrderBy{}
+	for i := 0; i < len(items); i++ {
+		attributeKey := enrichKeyWithMetadata(v3.AttributeKey{
+			Key: items[i].ColumnName,
+		}, keys)
+		enrichedItems = append(enrichedItems, v3.OrderBy{
+			ColumnName: items[i].ColumnName,
+			Order:      items[i].Order,
+			Key:        attributeKey.Key,
+			DataType:   attributeKey.DataType,
+			Type:       attributeKey.Type,
+			IsColumn:   attributeKey.IsColumn,
+		})
+	}
+	return enrichedItems
+}

pkg/query-service/app/traces/v4/enrich_test.go

@@ -0,0 +1,196 @@
+package v4
+
+import (
+	"reflect"
+	"testing"
+
+	v3 "go.signoz.io/signoz/pkg/query-service/model/v3"
+)
+
+func TestEnrichTracesQuery(t *testing.T) {
+	type args struct {
+		query *v3.BuilderQuery
+		keys  map[string]v3.AttributeKey
+		want  *v3.BuilderQuery
+	}
+	tests := []struct {
+		name string
+		args args
+	}{
+		{
+			name: "test 1",
+			args: args{
+				query: &v3.BuilderQuery{
+					Filters: &v3.FilterSet{
+						Operator: "AND",
+						Items: []v3.FilterItem{
+							{Key: v3.AttributeKey{Key: "bytes", Type: v3.AttributeKeyTypeTag}, Value: 100, Operator: ">"},
+						},
+					},
+					OrderBy: []v3.OrderBy{},
+				},
+				keys: map[string]v3.AttributeKey{
+					"bytes##tag##int64": {Key: "bytes", DataType: v3.AttributeKeyDataTypeInt64, Type: v3.AttributeKeyTypeTag},
+				},
+				want: &v3.BuilderQuery{
+					Filters: &v3.FilterSet{
+						Operator: "AND",
+						Items: []v3.FilterItem{
+							{Key: v3.AttributeKey{Key: "bytes", Type: v3.AttributeKeyTypeTag, DataType: v3.AttributeKeyDataTypeInt64}, Value: 100, Operator: ">"},
+						},
+					},
+					OrderBy: []v3.OrderBy{},
+				},
+			},
+		},
+		{
+			name: "test service name",
+			args: args{
+				query: &v3.BuilderQuery{
+					Filters: &v3.FilterSet{
+						Operator: "AND",
+						Items: []v3.FilterItem{
+							{Key: v3.AttributeKey{Key: "serviceName", DataType: v3.AttributeKeyDataTypeString, IsColumn: true}, Value: "myservice", Operator: "="},
+							{Key: v3.AttributeKey{Key: "serviceName"}, Value: "myservice", Operator: "="},
+						},
+					},
+					OrderBy: []v3.OrderBy{},
+				},
+				keys: map[string]v3.AttributeKey{},
+				want: &v3.BuilderQuery{
+					Filters: &v3.FilterSet{
+						Operator: "AND",
+						Items: []v3.FilterItem{
+							{Key: v3.AttributeKey{Key: "service.name", Type: v3.AttributeKeyTypeResource, DataType: v3.AttributeKeyDataTypeString}, Value: "myservice", Operator: "="},
+							{Key: v3.AttributeKey{Key: "service.name", Type: v3.AttributeKeyTypeResource, DataType: v3.AttributeKeyDataTypeString}, Value: "myservice", Operator: "="},
+						},
+					},
+					OrderBy: []v3.OrderBy{},
+				},
+			},
+		},
+		{
+			name: "test mat attrs",
+			args: args{
+				query: &v3.BuilderQuery{
+					Filters: &v3.FilterSet{
+						Operator: "AND",
+						Items: []v3.FilterItem{
+							{Key: v3.AttributeKey{Key: "http.route", DataType: v3.AttributeKeyDataTypeString, IsColumn: true}, Value: "/api", Operator: "="},
+							{Key: v3.AttributeKey{Key: "msgSystem"}, Value: "name", Operator: "="},
+							{Key: v3.AttributeKey{Key: "external_http_url"}, Value: "name", Operator: "="},
+						},
+					},
+					OrderBy: []v3.OrderBy{},
+				},
+				keys: map[string]v3.AttributeKey{},
+				want: &v3.BuilderQuery{
+					Filters: &v3.FilterSet{
+						Operator: "AND",
+						Items: []v3.FilterItem{
+							{Key: v3.AttributeKey{Key: "http.route", DataType: v3.AttributeKeyDataTypeString, Type: v3.AttributeKeyTypeTag, IsColumn: true}, Value: "/api", Operator: "="},
+							{Key: v3.AttributeKey{Key: "msgSystem", DataType: v3.AttributeKeyDataTypeString, IsColumn: true}, Value: "name", Operator: "="},
+							{Key: v3.AttributeKey{Key: "external_http_url", DataType: v3.AttributeKeyDataTypeString, IsColumn: true}, Value: "name", Operator: "="},
+						},
+					},
+					OrderBy: []v3.OrderBy{},
+				},
+			},
+		},
+		{
+			name: "test aggregateattr, filter, groupby, order by",
+			args: args{
+				query: &v3.BuilderQuery{
+					AggregateOperator: v3.AggregateOperatorCount,
+					AggregateAttribute: v3.AttributeKey{
+						Key:      "http.route",
+						DataType: v3.AttributeKeyDataTypeString,
+						Type:     v3.AttributeKeyTypeTag,
+					},
+					Filters: &v3.FilterSet{
+						Operator: "AND",
+						Items: []v3.FilterItem{
+							{Key: v3.AttributeKey{Key: "http.route", DataType: v3.AttributeKeyDataTypeString}, Value: "/api", Operator: "="},
+						},
+					},
+					GroupBy: []v3.AttributeKey{
+						{Key: "http.route", DataType: v3.AttributeKeyDataTypeString},
+						{Key: "msgSystem", DataType: v3.AttributeKeyDataTypeString},
+					},
+					OrderBy: []v3.OrderBy{
+						{ColumnName: "httpRoute", Order: v3.DirectionAsc},
+					},
+				},
+				keys: map[string]v3.AttributeKey{
+					"http.route##tag##string": {Key: "http.route", DataType: v3.AttributeKeyDataTypeString, Type: v3.AttributeKeyTypeTag, IsColumn: true},
+				},
+				want: &v3.BuilderQuery{
+					AggregateAttribute: v3.AttributeKey{
+						Key:      "http.route",
+						DataType: v3.AttributeKeyDataTypeString,
+						Type:     v3.AttributeKeyTypeTag,
+						IsColumn: true,
+					},
+					Filters: &v3.FilterSet{
+						Operator: "AND",
+						Items: []v3.FilterItem{
+							{Key: v3.AttributeKey{Key: "http.route", DataType: v3.AttributeKeyDataTypeString, Type: v3.AttributeKeyTypeTag, IsColumn: true}, Value: "/api", Operator: "="},
+						},
+					},
+					GroupBy: []v3.AttributeKey{
+						{Key: "http.route", DataType: v3.AttributeKeyDataTypeString, Type: v3.AttributeKeyTypeTag, IsColumn: true},
+						{Key: "msgSystem", DataType: v3.AttributeKeyDataTypeString, IsJSON: false, IsColumn: true},
+					},
+					OrderBy: []v3.OrderBy{
+						{Key: "httpRoute", Order: v3.DirectionAsc, ColumnName: "httpRoute", DataType: v3.AttributeKeyDataTypeString, IsColumn: true},
+					},
+				},
+			},
+		},
+		{
+			name: "enrich default values",
+			args: args{
+				query: &v3.BuilderQuery{
+					Filters: &v3.FilterSet{
+						Items: []v3.FilterItem{
+							{Key: v3.AttributeKey{Key: "testattr"}},
+						},
+					},
+					OrderBy: []v3.OrderBy{{ColumnName: "timestamp", Order: v3.DirectionAsc}},
+				},
+				keys: map[string]v3.AttributeKey{},
+				want: &v3.BuilderQuery{
+					Filters: &v3.FilterSet{
+						Items: []v3.FilterItem{{Key: v3.AttributeKey{Key: "testattr", Type: v3.AttributeKeyTypeTag, DataType: v3.AttributeKeyDataTypeString}}},
+					},
+					// isColumn won't matter in timestamp as it will always be a column
+					OrderBy: []v3.OrderBy{{Key: "timestamp", Order: v3.DirectionAsc, ColumnName: "timestamp"}},
+				},
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			EnrichTracesQuery(tt.args.query, tt.args.keys)
+			// Check AggregateAttribute
+			if tt.args.query.AggregateAttribute.Key != "" && !reflect.DeepEqual(tt.args.query.AggregateAttribute, tt.args.want.AggregateAttribute) {
+				t.Errorf("EnrichTracesQuery() AggregateAttribute = %v, want %v", tt.args.query.AggregateAttribute, tt.args.want.AggregateAttribute)
+			}
+
+			// Check Filters
+			if tt.args.query.Filters != nil && !reflect.DeepEqual(tt.args.query.Filters, tt.args.want.Filters) {
+				t.Errorf("EnrichTracesQuery() Filters = %v, want %v", tt.args.query.Filters, tt.args.want.Filters)
+			}
+
+			// Check GroupBy
+			if tt.args.query.GroupBy != nil && !reflect.DeepEqual(tt.args.query.GroupBy, tt.args.want.GroupBy) {
+				t.Errorf("EnrichTracesQuery() GroupBy = %v, want %v", tt.args.query.GroupBy, tt.args.want.GroupBy)
+			}
+
+			// Check OrderBy
+			if tt.args.query.OrderBy != nil && !reflect.DeepEqual(tt.args.query.OrderBy, tt.args.want.OrderBy) {
+				t.Errorf("EnrichTracesQuery() OrderBy = %v, want %v", tt.args.query.OrderBy, tt.args.want.OrderBy)
+			}
+		})
+	}
+}

pkg/query-service/app/traces/v4/query_builder_test.go

@@ -124,6 +124,27 @@ func Test_getColumnName(t *testing.T) {
 			},
 			want: "attributes_string['xyz']",
 		},
+		{
+			name: "new composite column",
+			args: args{
+				key: v3.AttributeKey{Key: "response_status_code"},
+			},
+			want: "response_status_code",
+		},
+		{
+			name: "new composite column with metadata",
+			args: args{
+				key: v3.AttributeKey{Key: "response_status_code", DataType: v3.AttributeKeyDataTypeString, IsColumn: true},
+			},
+			want: "response_status_code",
+		},
+		{
+			name: "new normal column with metadata",
+			args: args{
+				key: v3.AttributeKey{Key: "http.route", DataType: v3.AttributeKeyDataTypeString, Type: v3.AttributeKeyTypeTag, IsColumn: true},
+			},
+			want: "`attribute_string_http$$route`",
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {