Skip to content

Commit 4f98a95

Browse files
Model changes for hashcode and id (#191)
Signed-off-by: Siddhant Deshmukh <deshsid@amazon.com> (cherry picked from commit 1d991b8) Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
1 parent 2c98064 commit 4f98a95

File tree

8 files changed

+105
-29
lines changed

8 files changed

+105
-29
lines changed

src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -288,7 +288,7 @@ private void constructSearchQueryRecord(final SearchPhaseContext context, final
288288
// Add hashcode attribute when grouping is enabled
289289
if (queryInsightsService.isGroupingEnabled()) {
290290
String hashcode = queryShapeGenerator.getShapeHashCodeAsString(queryShape);
291-
attributes.put(Attribute.ID, hashcode);
291+
attributes.put(Attribute.QUERY_GROUP_HASHCODE, hashcode);
292292
}
293293
}
294294

src/main/java/org/opensearch/plugin/insights/core/service/grouper/MinMaxHeapQueryGrouper.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -271,7 +271,7 @@ int numberOfTopGroups() {
271271
private String getGroupingId(final SearchQueryRecord searchQueryRecord) {
272272
switch (groupingType) {
273273
case SIMILARITY:
274-
return searchQueryRecord.getAttributes().get(Attribute.ID).toString();
274+
return searchQueryRecord.getAttributes().get(Attribute.QUERY_GROUP_HASHCODE).toString();
275275
case NONE:
276276
throw new IllegalArgumentException("Should not try to group queries if grouping type is NONE");
277277
default:

src/main/java/org/opensearch/plugin/insights/rules/model/Attribute.java

+2-2
Original file line numberDiff line numberDiff line change
@@ -58,9 +58,9 @@ public enum Attribute {
5858
*/
5959
LABELS,
6060
/**
61-
* Query Group hashcode or query hashcode representing a unique identifier for the query/group
61+
* Query Group hashcode
6262
*/
63-
ID,
63+
QUERY_GROUP_HASHCODE,
6464
/**
6565
* Grouping type of the query record (none, similarity)
6666
*/

src/main/java/org/opensearch/plugin/insights/rules/model/SearchQueryRecord.java

+38-5
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
import java.util.List;
1616
import java.util.Map;
1717
import java.util.Objects;
18+
import java.util.UUID;
1819
import org.apache.logging.log4j.LogManager;
1920
import org.apache.logging.log4j.Logger;
2021
import org.opensearch.Version;
@@ -41,6 +42,8 @@ public class SearchQueryRecord implements ToXContentObject, Writeable {
4142
private final long timestamp;
4243
private final Map<MetricType, Measurement> measurements;
4344
private final Map<Attribute, Object> attributes;
45+
private final String id;
46+
4447
/**
4548
* Timestamp
4649
*/
@@ -93,12 +96,16 @@ public class SearchQueryRecord implements ToXContentObject, Writeable {
9396
* Grouping type of the query record (none, similarity)
9497
*/
9598
public static final String GROUP_BY = "group_by";
96-
9799
/**
98-
* Query Group hashcode or query hashcode representing a unique identifier for the query/group
100+
* UUID
99101
*/
100102
public static final String ID = "id";
101103

104+
/**
105+
* Query Group hashcode
106+
*/
107+
public static final String QUERY_GROUP_HASHCODE = "query_group_hashcode";
108+
102109
public static final String MEASUREMENTS = "measurements";
103110
private String groupingId;
104111

@@ -111,6 +118,7 @@ public class SearchQueryRecord implements ToXContentObject, Writeable {
111118
*/
112119
public SearchQueryRecord(final StreamInput in) throws IOException, ClassCastException {
113120
this.timestamp = in.readLong();
121+
this.id = in.readString();
114122
if (in.getVersion().onOrAfter(Version.V_2_17_0)) {
115123
measurements = new LinkedHashMap<>();
116124
in.readOrderedMap(MetricType::readFromStream, Measurement::readFromStream)
@@ -137,12 +145,30 @@ public SearchQueryRecord(final StreamInput in) throws IOException, ClassCastExce
137145
* @param attributes A list of Attributes associated with this query
138146
*/
139147
public SearchQueryRecord(final long timestamp, Map<MetricType, Measurement> measurements, final Map<Attribute, Object> attributes) {
148+
this(timestamp, measurements, attributes, UUID.randomUUID().toString());
149+
}
150+
151+
/**
152+
* Constructor of SearchQueryRecord
153+
*
154+
* @param timestamp The timestamp of the query.
155+
* @param measurements A list of Measurement associated with this query
156+
* @param attributes A list of Attributes associated with this query
157+
* @param id unique id for a search query record
158+
*/
159+
public SearchQueryRecord(
160+
final long timestamp,
161+
Map<MetricType, Measurement> measurements,
162+
final Map<Attribute, Object> attributes,
163+
String id
164+
) {
140165
if (measurements == null) {
141166
throw new IllegalArgumentException("Measurements cannot be null");
142167
}
143168
this.measurements = measurements;
144169
this.attributes = attributes;
145170
this.timestamp = timestamp;
171+
this.id = id;
146172
}
147173

148174
/**
@@ -156,6 +182,7 @@ public static SearchQueryRecord fromXContent(XContentParser parser) throws IOExc
156182
long timestamp = 0L;
157183
Map<MetricType, Measurement> measurements = new HashMap<>();
158184
Map<Attribute, Object> attributes = new HashMap<>();
185+
String id = null;
159186

160187
parser.nextToken();
161188
XContentParserUtils.ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
@@ -167,6 +194,9 @@ public static SearchQueryRecord fromXContent(XContentParser parser) throws IOExc
167194
case TIMESTAMP:
168195
timestamp = parser.longValue();
169196
break;
197+
case ID:
198+
id = parser.text();
199+
break;
170200
case LATENCY:
171201
case CPU:
172202
case MEMORY:
@@ -179,8 +209,8 @@ public static SearchQueryRecord fromXContent(XContentParser parser) throws IOExc
179209
case GROUP_BY:
180210
attributes.put(Attribute.GROUP_BY, parser.text());
181211
break;
182-
case ID:
183-
attributes.put(Attribute.ID, parser.text());
212+
case QUERY_GROUP_HASHCODE:
213+
attributes.put(Attribute.QUERY_GROUP_HASHCODE, parser.text());
184214
break;
185215
case SOURCE:
186216
XContentParserUtils.ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
@@ -264,7 +294,7 @@ public static SearchQueryRecord fromXContent(XContentParser parser) throws IOExc
264294
log.error("Error when parsing through search hit", e);
265295
}
266296
}
267-
return new SearchQueryRecord(timestamp, measurements, attributes);
297+
return new SearchQueryRecord(timestamp, measurements, attributes, id);
268298
}
269299

270300
/**
@@ -337,6 +367,8 @@ public void addAttribute(final Attribute attribute, final Object value) {
337367
public XContentBuilder toXContent(final XContentBuilder builder, final Params params) throws IOException {
338368
builder.startObject();
339369
builder.field("timestamp", timestamp);
370+
builder.field("id", id);
371+
340372
for (Map.Entry<Attribute, Object> entry : attributes.entrySet()) {
341373
builder.field(entry.getKey().toString(), entry.getValue());
342374
}
@@ -358,6 +390,7 @@ public XContentBuilder toXContent(final XContentBuilder builder, final Params pa
358390
@Override
359391
public void writeTo(final StreamOutput out) throws IOException {
360392
out.writeLong(timestamp);
393+
out.writeString(id);
361394
if (out.getVersion().onOrAfter(Version.V_2_17_0)) {
362395
out.writeMap(
363396
measurements,

src/test/java/org/opensearch/plugin/insights/QueryInsightsTestUtils.java

+7-7
Original file line numberDiff line numberDiff line change
@@ -139,7 +139,7 @@ public static List<SearchQueryRecord> generateQueryInsightRecords(
139139
attributes.put(Attribute.TOTAL_SHARDS, randomIntBetween(1, 100));
140140
attributes.put(Attribute.INDICES, randomArray(1, 3, Object[]::new, () -> randomAlphaOfLengthBetween(5, 10)));
141141
attributes.put(Attribute.PHASE_LATENCY_MAP, phaseLatencyMap);
142-
attributes.put(Attribute.ID, Objects.hashCode(i));
142+
attributes.put(Attribute.QUERY_GROUP_HASHCODE, Objects.hashCode(i));
143143
attributes.put(Attribute.GROUP_BY, GroupingType.NONE);
144144
attributes.put(
145145
Attribute.TASK_RESOURCE_USAGES,
@@ -200,13 +200,13 @@ public static List<List<SearchQueryRecord>> generateMultipleQueryInsightsRecords
200200

201201
public static void populateSameQueryHashcodes(List<SearchQueryRecord> searchQueryRecords) {
202202
for (SearchQueryRecord record : searchQueryRecords) {
203-
record.getAttributes().put(Attribute.ID, 1);
203+
record.getAttributes().put(Attribute.QUERY_GROUP_HASHCODE, 1);
204204
}
205205
}
206206

207207
public static void populateHashcode(List<SearchQueryRecord> searchQueryRecords, int hash) {
208208
for (SearchQueryRecord record : searchQueryRecords) {
209-
record.getAttributes().put(Attribute.ID, hash);
209+
record.getAttributes().put(Attribute.QUERY_GROUP_HASHCODE, hash);
210210
}
211211
}
212212

@@ -223,7 +223,7 @@ public static TopQueries createRandomTopQueries() {
223223
return new TopQueries(node, records);
224224
}
225225

226-
public static TopQueries createFixedTopQueries() {
226+
public static TopQueries createFixedTopQueries(String id) {
227227
DiscoveryNode node = new DiscoveryNode(
228228
"node_for_top_queries_test",
229229
buildNewFakeTransportAddress(),
@@ -232,12 +232,12 @@ public static TopQueries createFixedTopQueries() {
232232
VersionUtils.randomVersion(random())
233233
);
234234
List<SearchQueryRecord> records = new ArrayList<>();
235-
records.add(createFixedSearchQueryRecord());
235+
records.add(createFixedSearchQueryRecord(id));
236236

237237
return new TopQueries(node, records);
238238
}
239239

240-
public static SearchQueryRecord createFixedSearchQueryRecord() {
240+
public static SearchQueryRecord createFixedSearchQueryRecord(String id) {
241241
long timestamp = 1706574180000L;
242242
Map<MetricType, Measurement> measurements = Map.of(MetricType.LATENCY, new Measurement(1L));
243243

@@ -256,7 +256,7 @@ public static SearchQueryRecord createFixedSearchQueryRecord() {
256256
)
257257
);
258258

259-
return new SearchQueryRecord(timestamp, measurements, attributes);
259+
return new SearchQueryRecord(timestamp, measurements, attributes, id);
260260
}
261261

262262
public static void compareJson(ToXContent param1, ToXContent param2) throws IOException {

src/test/java/org/opensearch/plugin/insights/core/service/grouper/MinMaxHeapQueryGrouperTests.java

+2-2
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ public void testWithAllDifferentHashcodes() {
4444
Set<Integer> hashcodeSet = new HashSet<>();
4545
for (SearchQueryRecord record : records) {
4646
groupedRecord = minMaxHeapQueryGrouper.add(record);
47-
int hashcode = (int) groupedRecord.getAttributes().get(Attribute.ID);
47+
int hashcode = (int) groupedRecord.getAttributes().get(Attribute.QUERY_GROUP_HASHCODE);
4848
hashcodeSet.add(hashcode);
4949
}
5050
assertEquals(numOfRecords, hashcodeSet.size());
@@ -58,7 +58,7 @@ public void testWithAllSameHashcodes() {
5858
Set<Integer> hashcodeSet = new HashSet<>();
5959
for (SearchQueryRecord record : records) {
6060
groupedRecord = minMaxHeapQueryGrouper.add(record);
61-
int hashcode = (int) groupedRecord.getAttributes().get(Attribute.ID);
61+
int hashcode = (int) groupedRecord.getAttributes().get(Attribute.QUERY_GROUP_HASHCODE);
6262
hashcodeSet.add(hashcode);
6363
}
6464
assertEquals(1, hashcodeSet.size());

src/test/java/org/opensearch/plugin/insights/rules/action/top_queries/TopQueriesResponseTests.java

+49-6
Original file line numberDiff line numberDiff line change
@@ -40,18 +40,61 @@ public void testSerialize() throws Exception {
4040
}
4141

4242
public void testToXContent() throws IOException {
43-
char[] expectedXcontent =
44-
"{\"top_queries\":[{\"timestamp\":1706574180000,\"node_id\":\"node_for_top_queries_test\",\"phase_latency_map\":{\"expand\":1,\"query\":10,\"fetch\":1},\"task_resource_usages\":[{\"action\":\"action\",\"taskId\":2,\"parentTaskId\":1,\"nodeId\":\"id\",\"taskResourceUsage\":{\"cpu_time_in_nanos\":1000,\"memory_in_bytes\":2000}},{\"action\":\"action2\",\"taskId\":3,\"parentTaskId\":1,\"nodeId\":\"id2\",\"taskResourceUsage\":{\"cpu_time_in_nanos\":2000,\"memory_in_bytes\":1000}}],\"search_type\":\"query_then_fetch\",\"measurements\":{\"latency\":{\"number\":1,\"count\":1,\"aggregationType\":\"NONE\"}}}]}"
45-
.toCharArray();
46-
TopQueries topQueries = QueryInsightsTestUtils.createFixedTopQueries();
43+
String id = "sample_id";
44+
45+
char[] expectedXContent = ("{"
46+
+ "\"top_queries\":[{"
47+
+ "\"timestamp\":1706574180000,"
48+
+ "\"node_id\":\"node_for_top_queries_test\","
49+
+ "\"phase_latency_map\":{"
50+
+ "\"expand\":1,"
51+
+ "\"query\":10,"
52+
+ "\"fetch\":1"
53+
+ "},"
54+
+ "\"task_resource_usages\":[{"
55+
+ "\"action\":\"action\","
56+
+ "\"taskId\":2,"
57+
+ "\"parentTaskId\":1,"
58+
+ "\"nodeId\":\"id\","
59+
+ "\"taskResourceUsage\":{"
60+
+ "\"cpu_time_in_nanos\":1000,"
61+
+ "\"memory_in_bytes\":2000"
62+
+ "}"
63+
+ "},{"
64+
+ "\"action\":\"action2\","
65+
+ "\"taskId\":3,"
66+
+ "\"parentTaskId\":1,"
67+
+ "\"nodeId\":\"id2\","
68+
+ "\"taskResourceUsage\":{"
69+
+ "\"cpu_time_in_nanos\":2000,"
70+
+ "\"memory_in_bytes\":1000"
71+
+ "}"
72+
+ "}],"
73+
+ "\"search_type\":\"query_then_fetch\","
74+
+ "\"measurements\":{"
75+
+ "\"latency\":{"
76+
+ "\"number\":1,"
77+
+ "\"count\":1,"
78+
+ "\"aggregationType\":\"NONE\""
79+
+ "}"
80+
+ "},"
81+
+ "\"id\":\""
82+
+ id
83+
+ "\""
84+
+ "}]"
85+
+ "}").toCharArray();
86+
87+
TopQueries topQueries = QueryInsightsTestUtils.createFixedTopQueries(id);
4788
ClusterName clusterName = new ClusterName("test-cluster");
4889
TopQueriesResponse response = new TopQueriesResponse(clusterName, List.of(topQueries), new ArrayList<>(), 10, MetricType.LATENCY);
90+
4991
XContentBuilder builder = MediaTypeRegistry.contentBuilder(MediaTypeRegistry.JSON);
5092
char[] xContent = BytesReference.bytes(response.toXContent(builder, ToXContent.EMPTY_PARAMS)).utf8ToString().toCharArray();
51-
Arrays.sort(expectedXcontent);
93+
94+
Arrays.sort(expectedXContent);
5295
Arrays.sort(xContent);
5396

54-
assertEquals(Arrays.hashCode(expectedXcontent), Arrays.hashCode(xContent));
97+
assertEquals(Arrays.hashCode(expectedXContent), Arrays.hashCode(xContent));
5598
}
5699

57100
/**

src/test/java/org/opensearch/plugin/insights/rules/model/SearchQueryRecordTests.java

+5-5
Original file line numberDiff line numberDiff line change
@@ -45,19 +45,19 @@ public void testAllMetricTypes() {
4545
}
4646

4747
public void testCompare() {
48-
SearchQueryRecord record1 = QueryInsightsTestUtils.createFixedSearchQueryRecord();
49-
SearchQueryRecord record2 = QueryInsightsTestUtils.createFixedSearchQueryRecord();
48+
SearchQueryRecord record1 = QueryInsightsTestUtils.createFixedSearchQueryRecord("id");
49+
SearchQueryRecord record2 = QueryInsightsTestUtils.createFixedSearchQueryRecord("id");
5050
assertEquals(0, SearchQueryRecord.compare(record1, record2, MetricType.LATENCY));
5151
}
5252

5353
public void testEqual() {
54-
SearchQueryRecord record1 = QueryInsightsTestUtils.createFixedSearchQueryRecord();
55-
SearchQueryRecord record2 = QueryInsightsTestUtils.createFixedSearchQueryRecord();
54+
SearchQueryRecord record1 = QueryInsightsTestUtils.createFixedSearchQueryRecord("id");
55+
SearchQueryRecord record2 = QueryInsightsTestUtils.createFixedSearchQueryRecord("id");
5656
assertEquals(record1, record2);
5757
}
5858

5959
public void testFromXContent() {
60-
SearchQueryRecord record = QueryInsightsTestUtils.createFixedSearchQueryRecord();
60+
SearchQueryRecord record = QueryInsightsTestUtils.createFixedSearchQueryRecord("id");
6161
try (XContentParser recordParser = createParser(JsonXContent.jsonXContent, record.toString())) {
6262
SearchQueryRecord parsedRecord = SearchQueryRecord.fromXContent(recordParser);
6363
QueryInsightsTestUtils.checkRecordsEquals(List.of(record), List.of(parsedRecord));

0 commit comments

Comments
 (0)