Instrument r2dbc-postgresql 1.0.0 (#1556)

Co-authored-by: Koen Punt <[email protected]>
newrelic · Nov 13, 2023 · 4bc6b2e · 4bc6b2e
1 parent 88be534
commit 4bc6b2e
Show file tree

Hide file tree

Showing 8 changed files with 81 additions and 134 deletions.
diff --git a/...9.0/src/test/java/com/nr/agent/instrumentation/r2dbc/PostgresqlNoInstrumentationTest.java b/...9.0/src/test/java/com/nr/agent/instrumentation/r2dbc/PostgresqlNoInstrumentationTest.java
diff --git a/instrumentation/r2dbc-postgresql-0.9.2/build.gradle b/instrumentation/r2dbc-postgresql-0.9.2/build.gradle
@@ -10,8 +10,8 @@ jar {
 }
 
 verifyInstrumentation {
-    passesOnly 'org.postgresql:r2dbc-postgresql:[0.9.2,1.0.0.RELEASE)'
-    excludeRegex(".*(M1|M2|RC).*")
+    passesOnly 'org.postgresql:r2dbc-postgresql:[0.9.2,)'
+    exclude 'org.postgresql:r2dbc-postgresql:1.0.0.RC1'
 }
 
 site {

diff --git a/...stgresql-0.9.2/src/main/java/io/r2dbc/postgresql/PostgresqlStatement_Instrumentation.java b/...stgresql-0.9.2/src/main/java/io/r2dbc/postgresql/PostgresqlStatement_Instrumentation.java
@@ -1,3 +1,9 @@
+/*
+ *
+ *  * Copyright 2023 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
 package io.r2dbc.postgresql;
 
 import com.newrelic.api.agent.weaver.MatchType;
@@ -14,7 +20,7 @@ final class PostgresqlStatement_Instrumentation {
     public Flux<PostgresqlResult> execute() {
         Flux<PostgresqlResult> request = Weaver.callOriginal();
         if(request != null && parsedSql != null && resources != null) {
-            return R2dbcUtils.wrapRequest(request, parsedSql.getSql(), resources.getConfiguration());
+            return R2dbcUtils.wrapRequest(request, parsedSql.getSql(), resources.getClient(), resources.getConfiguration());
         }
         return request;
     }

diff --git a/instrumentation/r2dbc-postgresql-0.9.2/src/main/java/io/r2dbc/postgresql/R2dbcUtils.java b/instrumentation/r2dbc-postgresql-0.9.2/src/main/java/io/r2dbc/postgresql/R2dbcUtils.java
@@ -1,3 +1,9 @@
+/*
+ *
+ *  * Copyright 2023 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
 package io.r2dbc.postgresql;
 
 import com.newrelic.agent.bridge.NoOpTransaction;
@@ -9,39 +15,61 @@
 import com.newrelic.api.agent.NewRelic;
 import com.newrelic.api.agent.Segment;
 import com.newrelic.api.agent.Transaction;
+import io.r2dbc.postgresql.PostgresqlConnectionConfiguration;
 import io.r2dbc.postgresql.api.PostgresqlResult;
+import io.r2dbc.postgresql.client.Client;
+import io.r2dbc.postgresql.client.ReactorNettyClient_Instrumentation;
 import org.reactivestreams.Subscription;
 import reactor.core.publisher.Flux;
+import reactor.netty.Connection;
+import reactor.util.annotation.Nullable;
 
+import java.net.InetSocketAddress;
 import java.util.function.Consumer;
 
 public class R2dbcUtils {
-    public static Flux<PostgresqlResult> wrapRequest(Flux<PostgresqlResult> request, String sql, PostgresqlConnectionConfiguration connectionConfiguration) {
+    public static Flux<PostgresqlResult> wrapRequest(Flux<PostgresqlResult> request, String sql, Client client, PostgresqlConnectionConfiguration connectionConfiguration) {
         if(request != null) {
             Transaction transaction = NewRelic.getAgent().getTransaction();
             if(transaction != null && !(transaction instanceof NoOpTransaction)) {
                 Segment segment = transaction.startSegment("execute");
                 return request
-                        .doOnSubscribe(reportExecution(sql, connectionConfiguration, segment))
+                        .doOnSubscribe(reportExecution(sql, client, connectionConfiguration, segment))
                         .doFinally((type) -> segment.end());
             }
         }
         return request;
     }
 
-    private static Consumer<Subscription> reportExecution(String sql, PostgresqlConnectionConfiguration connectionConfiguration, Segment segment) {
+    private static Consumer<Subscription> reportExecution(String sql, Client client, PostgresqlConnectionConfiguration connectionConfiguration, Segment segment) {
         return (subscription) -> {
             OperationAndTableName sqlOperation = R2dbcOperation.extractFrom(sql);
-            if (sqlOperation != null) {
+            InetSocketAddress socketAddress = extractSocketAddress(client);
+            if (sqlOperation != null && socketAddress != null) {
                 segment.reportAsExternal(DatastoreParameters
                         .product(DatastoreVendor.Postgres.name())
                         .collection(sqlOperation.getTableName())
                         .operation(sqlOperation.getOperation())
-                        .instance(connectionConfiguration.getHost(), connectionConfiguration.getPort())
+                        .instance(socketAddress.getHostName(), socketAddress.getPort())
                         .databaseName(connectionConfiguration.getDatabase())
                         .slowQuery(sql, R2dbcObfuscator.POSTGRES_QUERY_CONVERTER)
                         .build());
             }
         };
     }
+
+    public static @Nullable InetSocketAddress extractSocketAddress(Client client) {
+        try {
+            if(client instanceof ReactorNettyClient_Instrumentation) {
+                ReactorNettyClient_Instrumentation instrumentedClient = (ReactorNettyClient_Instrumentation) client;
+                Connection clientConnection = instrumentedClient.clientConnection;
+                if(clientConnection.channel().remoteAddress() != null && clientConnection.channel().remoteAddress() instanceof InetSocketAddress) {
+                    return (InetSocketAddress) clientConnection.channel().remoteAddress();
+                }
+            }
+            return null;
+        } catch(Exception exception) {
+            return null;
+        }
+    }
 }
diff --git a/...ql-0.9.2/src/main/java/io/r2dbc/postgresql/client/ReactorNettyClient_Instrumentation.java b/...ql-0.9.2/src/main/java/io/r2dbc/postgresql/client/ReactorNettyClient_Instrumentation.java
@@ -0,0 +1,27 @@
+/*
+ *
+ *  * Copyright 2023 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+package io.r2dbc.postgresql.client;
+
+import com.newrelic.api.agent.weaver.MatchType;
+import com.newrelic.api.agent.weaver.NewField;
+import com.newrelic.api.agent.weaver.Weave;
+import io.r2dbc.postgresql.client.ConnectionSettings;
+import reactor.netty.Connection;
+
+import java.util.Optional;
+import java.util.TimeZone;
+
+@Weave(type = MatchType.ExactClass, originalName = "io.r2dbc.postgresql.client.ReactorNettyClient")
+public abstract class ReactorNettyClient_Instrumentation {
+    @NewField
+    public final Connection clientConnection;
+
+    private ReactorNettyClient_Instrumentation(Connection connection, ConnectionSettings settings) {
+        this.clientConnection = connection;
+    }
+
+}
diff --git a/...ql-0.9.2/src/test/java/com/nr/agent/instrumentation/r2dbc/PostgresqlInstrumentedTest.java b/...ql-0.9.2/src/test/java/com/nr/agent/instrumentation/r2dbc/PostgresqlInstrumentedTest.java
@@ -1,3 +1,9 @@
+/*
+ *
+ *  * Copyright 2023 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
 package com.nr.agent.instrumentation.r2dbc;
 
 import com.newrelic.agent.introspec.DatastoreHelper;

diff --git a/...9.2/src/test/java/com/nr/agent/instrumentation/r2dbc/PostgresqlNoInstrumentationTest.java b/...9.2/src/test/java/com/nr/agent/instrumentation/r2dbc/PostgresqlNoInstrumentationTest.java
diff --git a/...dbc-postgresql-0.9.2/src/test/java/com/nr/agent/instrumentation/r2dbc/R2dbcTestUtils.java b/...dbc-postgresql-0.9.2/src/test/java/com/nr/agent/instrumentation/r2dbc/R2dbcTestUtils.java
@@ -1,3 +1,9 @@
+/*
+ *
+ *  * Copyright 2023 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
 package com.nr.agent.instrumentation.r2dbc;
 
 import com.newrelic.api.agent.Trace;