Skip to content

Commit 211b20e

Browse files
dhilpipredhilpipre
authored
Merge pull request #2 from newrelic-experimental/mods_from_es
used same logic as ElasticSearch instrumentation
2 parents 471c52c + 0aee37d commit 211b20e

File tree

26 files changed

+3747
-508
lines changed

26 files changed

+3747
-508
lines changed

‎opensearch-java-client-2.0/build.gradle

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
2+
// Build.gradle generated for instrumentation module opensearch-java-client
3+
4+
apply plugin: 'java'
5+
6+
dependencies {
7+
implementation 'org.opensearch.client:opensearch-java:2.0.0'
8+
implementation 'org.opensearch.client:opensearch-rest-client:2.4.1'
9+
implementation 'org.apache.httpcomponents.client5:httpclient5:5.1.4'
10+
11+
// New Relic Java Agent dependencies
12+
implementation 'com.newrelic.agent.java:newrelic-agent:8.4.0'
13+
implementation 'com.newrelic.agent.java:newrelic-api:8.4.0'
14+
implementation fileTree(include: ['*.jar'], dir: '../libs')
15+
implementation fileTree(include: ['*.jar'], dir: '../test-lib')
16+
}
17+
18+
jar {
19+
manifest {
20+
attributes 'Implementation-Title': 'com.newrelic.instrumentation.labs.opensearch-java-client-2.0'
21+
attributes 'Implementation-Vendor': 'New Relic Labs'
22+
attributes 'Implementation-Vendor-Id': 'com.newrelic.labs'
23+
attributes 'Implementation-Version': 1.0
24+
}
25+
}
26+
27+
verifyInstrumentation {
28+
passes 'org.opensearch.client:opensearch-java:[2.0.0,2.2.0)'
29+
// Verifier plugin documentation:
30+
// https://github.com/newrelic/newrelic-gradle-verify-instrumentation
31+
// Example:
32+
// passes 'javax.servlet:servlet-api:[2.2,2.5]'
33+
// exclude 'javax.servlet:servlet-api:2.4.public_draft'
34+
}

‎opensearch-java-client/src/main/java/com/newrelic/instrumentation/labs/opensearch/javaclient/AttributeUtils.java renamed to ‎opensearch-java-client-2.0/src/main/java/com/newrelic/instrumentation/labs/opensearch/javaclient/AttributeUtils.java

File renamed without changes.

‎opensearch-java-client/src/main/java/com/newrelic/instrumentation/labs/opensearch/javaclient/NRCompletion.java renamed to ‎opensearch-java-client-2.0/src/main/java/com/newrelic/instrumentation/labs/opensearch/javaclient/NRCompletion.java

Lines changed: 28 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package com.newrelic.instrumentation.labs.opensearch.javaclient;
22

3+
import java.util.Map;
34
import java.util.function.BiConsumer;
45

56
import com.newrelic.api.agent.DatastoreParameters;
@@ -10,29 +11,47 @@ public class NRCompletion<T> implements BiConsumer<T, Throwable> {
1011

1112
private Segment segment = null;
1213
private DatastoreParameters params = null;
14+
private Map<String,Object> attributes = null;
1315

14-
public NRCompletion(Segment s, DatastoreParameters p) {
15-
segment = s;
16+
public NRCompletion(String segmentName, DatastoreParameters p, Map<String, Object> attrs) {
17+
segment = NewRelic.getAgent().getTransaction().startSegment(segmentName);
1618
params = p;
19+
attributes = attrs;
1720
}
1821

1922
@Override
2023
public void accept(T t, Throwable u) {
21-
if(u != null) {
22-
NewRelic.noticeError(u);
23-
if(segment != null) {
24-
segment.ignore();
25-
segment = null;
26-
}
27-
} else if(t != null) {
24+
if(t != null) {
2825
if(segment != null) {
2926
if(params != null) {
3027
segment.reportAsExternal(params);
3128
}
29+
if(attributes != null && !attributes.isEmpty()) {
30+
segment.addCustomAttributes(attributes);
31+
}
3232
segment.end();
3333
segment = null;
3434
}
35+
} else if(u != null) {
36+
NewRelic.noticeError(u);
37+
if(segment != null) {
38+
segment.ignore();
39+
segment = null;
40+
}
3541
}
3642
}
3743

44+
public Segment getSegment() {
45+
return segment;
46+
}
47+
48+
public DatastoreParameters getParams() {
49+
return params;
50+
}
51+
52+
public Map<String, Object> getAttributes() {
53+
return attributes;
54+
}
55+
56+
3857
}

‎opensearch-java-client-2.0/src/main/java/com/newrelic/instrumentation/labs/opensearch/javaclient/OpenSearchQueryConverter.java

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
package com.newrelic.instrumentation.labs.opensearch.javaclient;
2+
3+
import java.util.Map;
4+
5+
import com.newrelic.api.agent.QueryConverter;
6+
7+
public class OpenSearchQueryConverter implements QueryConverter<QueryHolder > {
8+
9+
@Override
10+
public String toRawQueryString(QueryHolder rawQuery) {
11+
String collString = rawQuery.getCollection();
12+
String operation = rawQuery.getOperation();
13+
Map<String, String> params = rawQuery.getParams();
14+
15+
StringBuffer sb = new StringBuffer();
16+
sb.append(operation);
17+
sb.append(' ');
18+
sb.append(collString);
19+
sb.append(" Parameters: ");
20+
for(String key : params.keySet()) {
21+
sb.append(key + "=" + params.get(key) + ",");
22+
}
23+
String queryStr = sb.toString();
24+
return queryStr;
25+
}
26+
27+
@Override
28+
public String toObfuscatedQueryString(QueryHolder rawQuery) {
29+
String collString = rawQuery.getCollection();
30+
String operation = rawQuery.getOperation();
31+
Map<String, String> params = rawQuery.getParams();
32+
33+
StringBuffer sb = new StringBuffer();
34+
sb.append(operation);
35+
sb.append(' ');
36+
sb.append(collString);
37+
sb.append(" Parameters: ");
38+
for(String key : params.keySet()) {
39+
sb.append(key + "=?,");
40+
}
41+
String queryStr = sb.toString();
42+
return queryStr;
43+
}
44+
45+
46+
}

‎opensearch-java-client-2.0/src/main/java/com/newrelic/instrumentation/labs/opensearch/javaclient/QueryHolder.java

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
package com.newrelic.instrumentation.labs.opensearch.javaclient;
2+
3+
import java.util.Map;
4+
5+
public class QueryHolder {
6+
7+
private String operation = null;
8+
private String collection = null;
9+
private Map<String, String> params = null;
10+
11+
public QueryHolder(String op, String coll, Map<String,String> p) {
12+
operation = op;
13+
collection = coll;
14+
params = p;
15+
}
16+
17+
public String getOperation() {
18+
return operation;
19+
}
20+
21+
public String getCollection() {
22+
return collection;
23+
}
24+
25+
@Override
26+
public String toString() {
27+
return "Operation: " + operation + ", Collection: " + collection +", Parameters: " + params.toString();
28+
}
29+
30+
public Map<String, String> getParams() {
31+
return params;
32+
}
33+
34+
35+
}

‎opensearch-java-client-2.0/src/main/java/com/newrelic/instrumentation/labs/opensearch/javaclient/Utils.java

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
1+
package com.newrelic.instrumentation.labs.opensearch.javaclient;
2+
3+
import java.io.ByteArrayOutputStream;
4+
import java.util.HashMap;
5+
import java.util.Iterator;
6+
import java.util.List;
7+
import java.util.Map;
8+
9+
import org.opensearch.client.json.JsonpMapper;
10+
import org.opensearch.client.json.NdJsonpSerializable;
11+
import org.opensearch.client.opensearch._types.query_dsl.Query;
12+
13+
import com.newrelic.api.agent.DatastoreParameters;
14+
15+
import jakarta.json.stream.JsonGenerator;
16+
17+
public class Utils {
18+
19+
public static <RequestT> DatastoreParameters getParams(Object request, Map<String, String> collectionAttributes, Map<String,String> params, Query query) {
20+
String operation = getOperationFromRequest(request);
21+
StringBuffer sb = new StringBuffer();
22+
if(params == null) {
23+
params = new HashMap<String, String>();
24+
}
25+
for(String key : collectionAttributes.keySet()) {
26+
if(key.equals("id")) {
27+
params.put("id", collectionAttributes.get("id"));
28+
} else {
29+
sb.append(key + ":" + collectionAttributes.get(key) + " ");
30+
}
31+
}
32+
if(query != null) {
33+
Object value = query._get();
34+
if(value != null) {
35+
params.put("Query-JSON", value.toString());
36+
}
37+
}
38+
String collection = sb.toString();
39+
QueryHolder holder = new QueryHolder(operation, collection, params);
40+
OpenSearchQueryConverter converter = new OpenSearchQueryConverter();
41+
42+
return DatastoreParameters.product("ElasticSearch").collection(collection).operation(operation).noInstance().noDatabaseName().slowQuery(holder, converter).build();
43+
}
44+
45+
public static String getObjectString(Object obj) {
46+
47+
if(obj instanceof String) {
48+
return (String)obj;
49+
}
50+
if(obj instanceof List) {
51+
List<?> list = (List<?>)obj;
52+
int size = list.size();
53+
StringBuffer sb = new StringBuffer();
54+
for(int i=0;i<size;i++) {
55+
sb.append(list.get(i).toString());
56+
if(i < size-1) {
57+
sb.append(',');
58+
}
59+
}
60+
return sb.toString();
61+
}
62+
if(obj instanceof Query) {
63+
return getCollectionFromQuery((Query)obj, null);
64+
}
65+
66+
return null;
67+
}
68+
69+
public static String getCollectionFromList(List<String> list) {
70+
int size = list.size();
71+
StringBuffer sb = new StringBuffer();
72+
for(int i=0;i<size;i++) {
73+
sb.append(list.get(i));
74+
if(i < size-1) {
75+
sb.append(',');
76+
}
77+
}
78+
return sb.toString();
79+
}
80+
81+
public static <TDocument> String getCollectionFromQuery(Query query, Class<TDocument> tDocumentClass) {
82+
if(query != null) {
83+
return tDocumentClass != null ? "QueryType-" + query._kind().name() + "-" + tDocumentClass.getSimpleName() : "QueryType-" + query._kind().name();
84+
}
85+
return null;
86+
}
87+
88+
public static void recordRequest(Map<String, Object> attributes, Object payload, String requestURL, String method) {
89+
if(attributes != null) {
90+
if(payload != null) {
91+
attributes.put("Request-Payload", payload.toString());
92+
}
93+
if(requestURL != null) {
94+
attributes.put("RequestUrl", requestURL);
95+
}
96+
if(method != null) {
97+
attributes.put("Request-Method", method);
98+
}
99+
100+
}
101+
}
102+
103+
public static String getOperationFromRequest(Object obj) {
104+
if(obj instanceof String) {
105+
return ((String)obj).replace("Request", "");
106+
}
107+
return convertFirstCharToLowerCase(obj.getClass().getSimpleName().replace("Request", ""));
108+
}
109+
110+
public static String convertFirstCharToLowerCase(String s) {
111+
if(s == null || s.isEmpty()) {
112+
return s;
113+
}
114+
return s.substring(0, 1).toLowerCase() + s.substring(1);
115+
}
116+
117+
public static Object getRequestBody(Object request, JsonpMapper mapper) {
118+
ByteArrayOutputStream baos = new ByteArrayOutputStream();
119+
if (request instanceof NdJsonpSerializable) {
120+
writeNdJson((NdJsonpSerializable) request, baos, mapper);
121+
} else {
122+
JsonGenerator generator = mapper.jsonProvider().createGenerator(baos);
123+
mapper.serialize(request, generator);
124+
generator.close();
125+
}
126+
127+
return baos;
128+
}
129+
130+
private static void writeNdJson(NdJsonpSerializable value, ByteArrayOutputStream baos, JsonpMapper mapper) {
131+
Iterator<?> values = value._serializables();
132+
while(values.hasNext()) {
133+
Object item = values.next();
134+
if (item instanceof NdJsonpSerializable && item != value) { // do not recurse on the item itself
135+
writeNdJson((NdJsonpSerializable) item, baos, mapper);
136+
} else {
137+
JsonGenerator generator = mapper.jsonProvider().createGenerator(baos);
138+
mapper.serialize(item, generator);
139+
generator.close();
140+
baos.write('\n');
141+
}
142+
}
143+
}
144+
145+
}

0 commit comments

Comments
 (0)