From 1aaafa46388c5357c6bf1625e91aa87dd5ad56ed Mon Sep 17 00:00:00 2001
From: Nityananda Gohain <nityanandagohain@gmail.com>
Date: Thu, 29 Feb 2024 15:32:37 +0530
Subject: [PATCH] feat: use attribute if present for json query (#4458)

* feat: use attribute if present for json query

* fix: refractor test cases
---
 pkg/query-service/app/logs/v3/enrich_query.go |  14 ++
 .../app/logs/v3/enrich_query_test.go          | 122 ++++++++++++++++++
 2 files changed, 136 insertions(+)

diff --git a/pkg/query-service/app/logs/v3/enrich_query.go b/pkg/query-service/app/logs/v3/enrich_query.go
index fe8ed8b9ed..c8a5a797b2 100644
--- a/pkg/query-service/app/logs/v3/enrich_query.go
+++ b/pkg/query-service/app/logs/v3/enrich_query.go
@@ -103,6 +103,7 @@ func enrichLogsQuery(query *v3.BuilderQuery, fields map[string]v3.AttributeKey)
 		for i := 0; i < len(query.Filters.Items); i++ {
 			query.Filters.Items[i] = jsonFilterEnrich(query.Filters.Items[i])
 			if query.Filters.Items[i].Key.IsJSON {
+				query.Filters.Items[i] = jsonReplaceField(query.Filters.Items[i], fields)
 				continue
 			}
 			query.Filters.Items[i].Key = enrichFieldWithMetadata(query.Filters.Items[i].Key, fields)
@@ -181,6 +182,19 @@ func jsonFilterEnrich(filter v3.FilterItem) v3.FilterItem {
 	return filter
 }
 
+func jsonReplaceField(filter v3.FilterItem, fields map[string]v3.AttributeKey) v3.FilterItem {
+	key, found := strings.CutPrefix(filter.Key.Key, "body.")
+	if !found {
+		return filter
+	}
+
+	if field, ok := fields[key]; ok && field.DataType == filter.Key.DataType {
+		filter.Key = field
+	}
+
+	return filter
+}
+
 func parseStrValue(valueStr string, operator v3.FilterOperator) (string, interface{}) {
 
 	valueType := "string"
diff --git a/pkg/query-service/app/logs/v3/enrich_query_test.go b/pkg/query-service/app/logs/v3/enrich_query_test.go
index 8b831f56ef..4903139610 100644
--- a/pkg/query-service/app/logs/v3/enrich_query_test.go
+++ b/pkg/query-service/app/logs/v3/enrich_query_test.go
@@ -456,6 +456,128 @@ func TestJsonEnrich(t *testing.T) {
 	}
 }
 
+func TestJsonReplaceField(t *testing.T) {
+	fields := map[string]v3.AttributeKey{
+		"method.name": {
+			Key:      "method.name",
+			DataType: v3.AttributeKeyDataTypeString,
+			Type:     v3.AttributeKeyTypeTag,
+		},
+		"status": {
+			Key:      "status",
+			DataType: v3.AttributeKeyDataTypeInt64,
+			Type:     v3.AttributeKeyTypeTag,
+		},
+		"data.error": {
+			Key:      "data.error",
+			DataType: v3.AttributeKeyDataTypeString,
+			Type:     v3.AttributeKeyTypeTag,
+			IsColumn: true,
+		},
+	}
+	var TestJsonReplaceFieldData = []struct {
+		Name   string
+		Filter v3.FilterItem
+		Result v3.FilterItem
+	}{
+		{
+			Name: "key in nested json",
+			Filter: v3.FilterItem{
+				Key: v3.AttributeKey{
+					Key:      "body.method.name",
+					DataType: v3.AttributeKeyDataTypeString,
+					Type:     v3.AttributeKeyTypeUnspecified,
+				},
+				Operator: "has",
+				Value:    "index_service",
+			},
+			Result: v3.FilterItem{
+				Key: v3.AttributeKey{
+					Key:      "method.name",
+					DataType: v3.AttributeKeyDataTypeString,
+					Type:     v3.AttributeKeyTypeTag,
+					IsJSON:   false,
+				},
+				Operator: "has",
+				Value:    "index_service",
+			},
+		},
+		{
+			Name: "key at top level",
+			Filter: v3.FilterItem{
+				Key: v3.AttributeKey{
+					Key:      "body.status",
+					DataType: v3.AttributeKeyDataTypeInt64,
+					Type:     v3.AttributeKeyTypeUnspecified,
+				},
+				Operator: "=",
+				Value:    10,
+			},
+			Result: v3.FilterItem{
+				Key: v3.AttributeKey{
+					Key:      "status",
+					DataType: v3.AttributeKeyDataTypeInt64,
+					Type:     v3.AttributeKeyTypeTag,
+					IsJSON:   false,
+				},
+				Operator: "=",
+				Value:    10,
+			},
+		},
+		{
+			Name: "key not present",
+			Filter: v3.FilterItem{
+				Key: v3.AttributeKey{
+					Key:      "body.status.code",
+					DataType: v3.AttributeKeyDataTypeInt64,
+					Type:     v3.AttributeKeyTypeUnspecified,
+				},
+				Operator: "=",
+				Value:    10,
+			},
+			Result: v3.FilterItem{
+				Key: v3.AttributeKey{
+					Key:      "body.status.code",
+					DataType: v3.AttributeKeyDataTypeInt64,
+					Type:     v3.AttributeKeyTypeUnspecified,
+					IsJSON:   false,
+				},
+				Operator: "=",
+				Value:    10,
+			},
+		},
+		{
+			Name: "key materialized",
+			Filter: v3.FilterItem{
+				Key: v3.AttributeKey{
+					Key:      "body.data.error",
+					DataType: v3.AttributeKeyDataTypeString,
+					Type:     v3.AttributeKeyTypeUnspecified,
+				},
+				Operator: "=",
+				Value:    10,
+			},
+			Result: v3.FilterItem{
+				Key: v3.AttributeKey{
+					Key:      "data.error",
+					DataType: v3.AttributeKeyDataTypeString,
+					Type:     v3.AttributeKeyTypeTag,
+					IsJSON:   false,
+					IsColumn: true,
+				},
+				Operator: "=",
+				Value:    10,
+			},
+		},
+	}
+	for _, tt := range TestJsonReplaceFieldData {
+		Convey(tt.Name, t, func() {
+			res := jsonReplaceField(tt.Filter, fields)
+			So(res, ShouldResemble, tt.Result)
+		})
+	}
+}
+
 var testParseStrValueData = []struct {
 	Name       string
 	Operator   v3.FilterOperator