chore: add tag type filter support in attribute keys (#7522)

pkg/query-service/app/clickhouseReader/reader.go

@@ -3928,11 +3928,16 @@ func (r *ClickHouseReader) GetLogAttributeKeys(ctx context.Context, req *v3.Filt
 	var rows driver.Rows
 	var response v3.FilterAttributeKeyResponse
 
+	tagTypeFilter := `tag_type != 'logfield'`
+	if req.TagType != "" {
+		tagTypeFilter = fmt.Sprintf(`tag_type != 'logfield' and tag_type = '%s'`, req.TagType)
+	}
+
 	if len(req.SearchText) != 0 {
-		query = fmt.Sprintf("select distinct tag_key, tag_type, tag_data_type from  %s.%s where tag_type != 'logfield' and tag_key ILIKE $1 limit $2", r.logsDB, r.logsTagAttributeTableV2)
+		query = fmt.Sprintf("select distinct tag_key, tag_type, tag_data_type from  %s.%s where %s and tag_key ILIKE $1 limit $2", r.logsDB, r.logsTagAttributeTableV2, tagTypeFilter)
 		rows, err = r.db.Query(ctx, query, fmt.Sprintf("%%%s%%", req.SearchText), req.Limit)
 	} else {
-		query = fmt.Sprintf("select distinct tag_key, tag_type, tag_data_type from %s.%s where tag_type != 'logfield' limit $1", r.logsDB, r.logsTagAttributeTableV2)
+		query = fmt.Sprintf("select distinct tag_key, tag_type, tag_data_type from %s.%s where %s limit $1", r.logsDB, r.logsTagAttributeTableV2, tagTypeFilter)
 		rows, err = r.db.Query(ctx, query, req.Limit)
 	}
 
@@ -3967,7 +3972,9 @@ func (r *ClickHouseReader) GetLogAttributeKeys(ctx context.Context, req *v3.Filt
 		response.AttributeKeys = append(response.AttributeKeys, key)
 	}
 
-	// add other attributes
+	// add other attributes only when the tagType is not specified
+	// i.e retrieve all attributes
+	if req.TagType == "" {
 		for _, f := range constants.StaticFieldsLogsV3 {
 			if (v3.AttributeKey{} == f) {
 				continue
@@ -3976,6 +3983,7 @@ func (r *ClickHouseReader) GetLogAttributeKeys(ctx context.Context, req *v3.Filt
 				response.AttributeKeys = append(response.AttributeKeys, f)
 			}
 		}
+	}
 
 	return &response, nil
 }
@@ -4715,7 +4723,12 @@ func (r *ClickHouseReader) GetTraceAttributeKeys(ctx context.Context, req *v3.Fi
 	var rows driver.Rows
 	var response v3.FilterAttributeKeyResponse
 
-	query = fmt.Sprintf("SELECT DISTINCT(tag_key), tag_type, tag_data_type FROM %s.%s WHERE tag_key ILIKE $1 and tag_type != 'spanfield' LIMIT $2", r.TraceDB, r.spanAttributeTableV2)
+	tagTypeFilter := `tag_type != 'spanfield'`
+	if req.TagType != "" {
+		tagTypeFilter = fmt.Sprintf(`tag_type != 'spanfield' and tag_type = '%s'`, req.TagType)
+	}
+
+	query = fmt.Sprintf("SELECT DISTINCT(tag_key), tag_type, tag_data_type FROM %s.%s WHERE tag_key ILIKE $1 and %s LIMIT $2", r.TraceDB, r.spanAttributeTableV2, tagTypeFilter)
 
 	rows, err = r.db.Query(ctx, query, fmt.Sprintf("%%%s%%", req.SearchText), req.Limit)
 
@@ -4760,7 +4773,9 @@ func (r *ClickHouseReader) GetTraceAttributeKeys(ctx context.Context, req *v3.Fi
 		fields = constants.DeprecatedStaticFieldsTraces
 	}
 
-	// add the new static fields
+	// add the new static fields only when the tagType is not specified
+	// i.e retrieve all attributes
+	if req.TagType == "" {
 		for _, f := range fields {
 			if (v3.AttributeKey{} == f) {
 				continue
@@ -4769,6 +4784,7 @@ func (r *ClickHouseReader) GetTraceAttributeKeys(ctx context.Context, req *v3.Fi
 				response.AttributeKeys = append(response.AttributeKeys, f)
 			}
 		}
+	}
 
 	return &response, nil
 }

pkg/query-service/app/parser.go

@@ -719,6 +719,21 @@ func parseFilterAttributeKeyRequest(r *http.Request) (*v3.FilterAttributeKeyRequ
 	aggregateOperator := v3.AggregateOperator(r.URL.Query().Get("aggregateOperator"))
 	aggregateAttribute := r.URL.Query().Get("aggregateAttribute")
 	limit, err := strconv.Atoi(r.URL.Query().Get("limit"))
+	tagType := v3.TagType(r.URL.Query().Get("tagType"))
+
+	// empty string is a valid tagType
+	// i.e retrieve all attributes
+	if tagType != "" {
+		// what is happening here?
+		// if tagType is undefined(uh oh javascript) or any invalid value, set it to empty string
+		// instead of failing the request. Ideally, we should fail the request.
+		// but we are not doing that to maintain backward compatibility.
+		if err := tagType.Validate(); err != nil {
+			// if the tagType is invalid, set it to empty string
+			tagType = ""
+		}
+	}
+
 	if err != nil {
 		limit = 50
 	}
@@ -739,6 +754,7 @@ func parseFilterAttributeKeyRequest(r *http.Request) (*v3.FilterAttributeKeyRequ
 		AggregateAttribute: aggregateAttribute,
 		Limit:              limit,
 		SearchText:         r.URL.Query().Get("searchText"),
+		TagType:            tagType,
 	}
 	return &req, nil
 }

pkg/query-service/app/parser_test.go

@@ -112,6 +112,7 @@ func TestParseFilterAttributeKeyRequest(t *testing.T) {
 		expectedSearchText string
 		expectErr          bool
 		errMsg             string
+		expectedTagType    v3.TagType
 	}{
 		{
 			desc:               "valid operator and data source",
@@ -168,6 +169,38 @@ func TestParseFilterAttributeKeyRequest(t *testing.T) {
 			expectedDataSource: v3.DataSourceTraces,
 			expectedLimit:      50,
 		},
+		{
+			desc:               "invalid tag type",
+			queryString:        "aggregateOperator=avg&dataSource=traces&tagType=invalid",
+			expectedOperator:   v3.AggregateOperatorAvg,
+			expectedDataSource: v3.DataSourceTraces,
+			expectedTagType:    "",
+			expectedLimit:      50,
+		},
+		{
+			desc:               "valid tag type",
+			queryString:        "aggregateOperator=avg&dataSource=traces&tagType=resource",
+			expectedOperator:   v3.AggregateOperatorAvg,
+			expectedDataSource: v3.DataSourceTraces,
+			expectedTagType:    v3.TagTypeResource,
+			expectedLimit:      50,
+		},
+		{
+			desc:               "valid tag type",
+			queryString:        "aggregateOperator=avg&dataSource=traces&tagType=scope",
+			expectedOperator:   v3.AggregateOperatorAvg,
+			expectedDataSource: v3.DataSourceTraces,
+			expectedTagType:    v3.TagTypeInstrumentationScope,
+			expectedLimit:      50,
+		},
+		{
+			desc:               "valid tag type",
+			queryString:        "aggregateOperator=avg&dataSource=traces&tagType=tag",
+			expectedOperator:   v3.AggregateOperatorAvg,
+			expectedDataSource: v3.DataSourceTraces,
+			expectedTagType:    v3.TagTypeTag,
+			expectedLimit:      50,
+		},
 	}
 
 	for _, reqCase := range reqCases {

pkg/query-service/model/v3/v3.go

@@ -248,6 +248,7 @@ func (q TagType) Validate() error {
 type FilterAttributeKeyRequest struct {
 	DataSource         DataSource        `json:"dataSource"`
 	AggregateOperator  AggregateOperator `json:"aggregateOperator"`
+	TagType            TagType           `json:"tagType"`
 	AggregateAttribute string            `json:"aggregateAttribute"`
 	SearchText         string            `json:"searchText"`
 	Limit              int               `json:"limit"`