feat: elasticrestclient710-fixQuery

newrelic-experimental · Feb 5, 2025 · 5191771 · 5191771
1 parent 019bf3b
commit 5191771
Show file tree

Hide file tree

Showing 2 changed files with 89 additions and 70 deletions.
diff --git a/.../main/java/com/newrelic/instrumentation/labs/esrestclient/RestClient_Instrumentation.java b/.../main/java/com/newrelic/instrumentation/labs/esrestclient/RestClient_Instrumentation.java
@@ -1,5 +1,9 @@
 package com.newrelic.instrumentation.labs.esrestclient;
 
+import org.elasticsearch.client.Cancellable;
+import org.elasticsearch.client.Request;
+import org.elasticsearch.client.Response;
+import org.elasticsearch.client.ResponseListener;
 import com.newrelic.api.agent.DatastoreParameters;
 import com.newrelic.api.agent.NewRelic;
 import com.newrelic.api.agent.Segment;
@@ -8,79 +12,81 @@
 import com.newrelic.api.agent.weaver.Weaver;
 
 
-import org.elasticsearch.client.Cancellable;
-import org.elasticsearch.client.Request;
-import org.elasticsearch.client.Response;
-import org.elasticsearch.client.ResponseListener;
-
-
 @Weave(originalName = "org.elasticsearch.client.RestClient")
 public abstract class RestClient_Instrumentation {
 
-    @Trace(dispatcher = true)
-    public Response performRequest(Request request)  {
-        // Start a segment for the synchronous request
-        Segment segment = NewRelic.getAgent().getTransaction().startSegment("ElasticsearchRequest", request.getMethod());
-
-        try {
-        	Utils.logRequestAttributes(request);
-
-            String endPoint= request.getEndpoint();
-
-
-            if(endPoint != null && !endPoint.isEmpty()) {
-    			NewRelic.getAgent().getTracedMethod().setMetricName("Custom","ES","RestClient",getClass().getSimpleName(),"performRequest",endPoint);
-    		} else {
-    			NewRelic.getAgent().getTracedMethod().setMetricName("Custom","ES","RestClient",getClass().getSimpleName(),"performRequest");
-    		}
-            // Log the request as a database call
-
-            segment.reportAsExternal(DatastoreParameters.product("Elasticsearch")
-                    .collection(request.getEndpoint())
-                    .operation(request.getMethod())
-                    .build());
-
-            // Call the original method
-            return Weaver.callOriginal();
-        } finally {
-            // End the segment
-            segment.end();
-        }
+  @Trace(dispatcher = true)
+  public Response performRequest(Request request) {
+    // Start a segment for the synchronous request
+    Segment segment = NewRelic.getAgent().getTransaction().startSegment("ElasticsearchRequest",
+        request.getMethod());
+
+    try {
+      Utils.logRequestAttributes(request);
+
+      String endPoint = request.getEndpoint();
+
+
+      if (endPoint != null && !endPoint.isEmpty()) {
+        NewRelic.getAgent().getTracedMethod().setMetricName("Custom", "ES", "RestClient",
+            getClass().getSimpleName(), "performRequest", endPoint);
+      } else {
+        NewRelic.getAgent().getTracedMethod().setMetricName("Custom", "ES", "RestClient",
+            getClass().getSimpleName(), "performRequest");
+      }
+      // Log the request as a database call
+
+      segment.reportAsExternal(DatastoreParameters.product("Elasticsearch")
+          .collection(Utils.extractIndex(request.getEndpoint())).operation(request.getMethod())
+          .noInstance().noDatabaseName().slowQuery(endPoint, new ESQueryConverter()).build());
+
+      // Call the original method
+      return Weaver.callOriginal();
+    } finally {
+      // End the segment
+      segment.end();
     }
+  }
+
+  @Trace(dispatcher = true)
+  public Cancellable performRequestAsync(Request request, ResponseListener responseListener) {
+    // Start a segment for the asynchronous request
+    Segment segment = NewRelic.getAgent().getTransaction().startSegment("ElasticsearchRequestAsync",
+        request.getMethod());
+
+    try {
+      // Add custom attributes for better traceability
+      Utils.logRequestAttributes(request);
+
+      String endPoint = request.getEndpoint();
+
+
+      if (endPoint != null && !endPoint.isEmpty()) {
+        NewRelic.getAgent().getTracedMethod().setMetricName("Custom", "ES", "RestClient",
+            getClass().getSimpleName(), "performRequestAsync", endPoint);
+      } else {
+        NewRelic.getAgent().getTracedMethod().setMetricName("Custom", "ES", "RestClient",
+            getClass().getSimpleName(), "performRequestAsync");
+      }
+
+      DatastoreParameters params = DatastoreParameters.product("Elasticsearch")
+          .collection(Utils.extractIndex(request.getEndpoint())).operation(request.getMethod())
+          .noInstance().noDatabaseName().slowQuery(endPoint, new ESQueryConverter()).build();
+
+      // Log the request as a database call
+      int hash = responseListener.hashCode();
+
+      NRHolder.putParams(hash, params);
+      NRHolder.putSegment(hash, segment);
+      NRHolder.putToken(hash, NewRelic.getAgent().getTransaction().getToken());
+
 
-    @Trace(dispatcher = true)
-    public Cancellable performRequestAsync(Request request, ResponseListener responseListener) {
-        // Start a segment for the asynchronous request
-        Segment segment = NewRelic.getAgent().getTransaction().startSegment("ElasticsearchRequestAsync", request.getMethod());
-
-        try {
-            // Add custom attributes for better traceability
-        	Utils.logRequestAttributes(request);
-
-            String endPoint= request.getEndpoint();
-
-
-            if(endPoint != null && !endPoint.isEmpty()) {
-    			NewRelic.getAgent().getTracedMethod().setMetricName("Custom","ES","RestClient",getClass().getSimpleName(),"performRequestAsync",endPoint);
-    		} else {
-    			NewRelic.getAgent().getTracedMethod().setMetricName("Custom","ES","RestClient",getClass().getSimpleName(),"performRequestAsync");
-    		}
-
-            DatastoreParameters params = DatastoreParameters.product("Elasticsearch").collection(request.getEndpoint()).operation(request.getMethod()).build();
-
-            // Log the request as a database call
-    		int hash = responseListener.hashCode();
-    		NRHolder.putParams(hash,params);
-    		NRHolder.putSegment(hash, segment);
-    		NRHolder.putToken(hash,NewRelic.getAgent().getTransaction().getToken());
-
-
-            // Call the original method with the wrapped listener
-           return  Weaver.callOriginal();
-        } catch (Exception e) {
-            // Ensure segment is ended in case of an exception
-            segment.end();
-            throw e;
-        }
+      // Call the original method with the wrapped listener
+      return Weaver.callOriginal();
+    } catch (Exception e) {
+      // Ensure segment is ended in case of an exception
+      segment.end();
+      throw e;
     }
-}
+  }
+}
diff --git a/esrestclient-7.10/src/main/java/com/newrelic/instrumentation/labs/esrestclient/Utils.java b/esrestclient-7.10/src/main/java/com/newrelic/instrumentation/labs/esrestclient/Utils.java
@@ -55,4 +55,17 @@ public static void logRequestAttributes(Object request) {
     addRequestAttributes(attributes, request);
     NewRelic.getAgent().getTracedMethod().addCustomAttributes(attributes);
   }
+
+  public static String extractIndex(String input) {
+    if (input == null || input.isEmpty()) {
+      return "noindex"; // or throw an exception based on your use case
+    }
+
+    int underscoreIndex = input.indexOf('_');
+    if (underscoreIndex == -1) {
+      return "noindex"; // or handle the case where there's no underscore
+    }
+
+    return input.substring(0, underscoreIndex);
+  }
 }