feat: support for new enrichment logic in traces (#6438)

* feat: support for new enrichment logic in traces

* fix: default test added

* fix: update func name in links

* Update pkg/query-service/utils/logs_test.go

Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>

---------

Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>

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) {

pkg/query-service/constants/constants.go

@@ -589,6 +589,60 @@ var StaticFieldsTraces = map[string]v3.AttributeKey{
 		DataType: v3.AttributeKeyDataTypeString,
 		IsColumn: true,
 	},
+
+	// new support
+	"response_status_code": {
+		Key:      "response_status_code",
+		DataType: v3.AttributeKeyDataTypeString,
+		IsColumn: true,
+	},
+	"external_http_url": {
+		Key:      "external_http_url",
+		DataType: v3.AttributeKeyDataTypeString,
+		IsColumn: true,
+	},
+	"http_url": {
+		Key:      "http_url",
+		DataType: v3.AttributeKeyDataTypeString,
+		IsColumn: true,
+	},
+	"external_http_method": {
+		Key:      "external_http_method",
+		DataType: v3.AttributeKeyDataTypeString,
+		IsColumn: true,
+	},
+	"http_method": {
+		Key:      "http_method",
+		DataType: v3.AttributeKeyDataTypeString,
+		IsColumn: true,
+	},
+	"http_host": {
+		Key:      "http_host",
+		DataType: v3.AttributeKeyDataTypeString,
+		IsColumn: true,
+	},
+	"db_name": {
+		Key:      "db_name",
+		DataType: v3.AttributeKeyDataTypeString,
+		IsColumn: true,
+	},
+	"db_operation": {
+		Key:      "db_operation",
+		DataType: v3.AttributeKeyDataTypeString,
+		IsColumn: true,
+	},
+	"has_error": {
+		Key:      "has_error",
+		DataType: v3.AttributeKeyDataTypeBool,
+		IsColumn: true,
+	},
+	"is_remote": {
+		Key:      "is_remote",
+		DataType: v3.AttributeKeyDataTypeString,
+		IsColumn: true,
+	},
+	// the simple attributes are not present here as
+	// they are taken care by new format <attribute_type>_<attribute_datatype>_'<attribute_key>'
 }
 
 const TRACE_V4_MAX_PAGINATION_LIMIT = 10000

pkg/query-service/contextlinks/links.go

@@ -183,7 +183,7 @@ func PrepareFilters(labels map[string]string, whereClauseItems []v3.FilterItem,
 			var attrFound bool
 
 			// as of now this logic will only apply for logs
-			for _, tKey := range utils.GenerateLogEnrichmentKeys(v3.AttributeKey{Key: key}) {
+			for _, tKey := range utils.GenerateEnrichmentKeys(v3.AttributeKey{Key: key}) {
 				if val, ok := keys[tKey]; ok {
 					attributeKey = val
 					attrFound = true

pkg/query-service/utils/logs.go

@@ -40,8 +40,8 @@ func GetListTsRanges(start, end int64) []LogsListTsRange {
 }
 
 // This tries to see all possible fields that it can fall back to if some meta is missing
-// check Test_GenerateLogEnrichmentKeys for example
-func GenerateLogEnrichmentKeys(field v3.AttributeKey) []string {
+// check Test_GenerateEnrichmentKeys for example
+func GenerateEnrichmentKeys(field v3.AttributeKey) []string {
 	names := []string{}
 	if field.Type != v3.AttributeKeyTypeUnspecified && field.DataType != v3.AttributeKeyDataTypeUnspecified {
 		names = append(names, field.Key+"##"+field.Type.String()+"##"+field.DataType.String())

pkg/query-service/utils/logs_test.go

@@ -53,7 +53,7 @@ func TestListTsRange(t *testing.T) {
 	}
 }
 
-func Test_GenerateLogEnrichmentKeys(t *testing.T) {
+func Test_GenerateEnrichmentKeys(t *testing.T) {
 	type args struct {
 		field v3.AttributeKey
 	}
@@ -96,8 +96,8 @@ func Test_GenerateLogEnrichmentKeys(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			if got := GenerateLogEnrichmentKeys(tt.args.field); !reflect.DeepEqual(got, tt.want) {
-				t.Errorf("generateLogEnrichmentKeys() = %v, want %v", got, tt.want)
+			if got := GenerateEnrichmentKeys(tt.args.field); !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("GenerateEnrichmentKeys() = %v, want %v", got, tt.want)
 			}
 		})
 	}