Merge pull request #2004 from newrelic/vertx4-pg-client

Vertx 4.4.2 Postgresql SqlClient Instrumentation

agent-bridge-datastore/src/main/java/com/newrelic/agent/bridge/datastore/R2dbcOperation.java

@@ -13,7 +13,7 @@
 import java.util.regex.Pattern;
 
 public class R2dbcOperation {
-
+    public static final OperationAndTableName UNKNOWN_OPERATION_AND_TABLE_NAME = new OperationAndTableName("unknown", "unknown");
     static final Pattern VALID_METRIC_NAME_MATCHER = Pattern.compile("[a-zA-Z0-9.$_@]+");
     static final int PATTERN_SWITCHES = Pattern.CASE_INSENSITIVE | Pattern.DOTALL;
     static final Pattern COMMENT_PATTERN = Pattern.compile("/\\*.*?\\*/", Pattern.DOTALL);

instrumentation/vertx-postgres-sqlclient-4.4.2/build.gradle

@@ -0,0 +1,30 @@
+jar {
+    manifest {
+        attributes 'Implementation-Title': 'com.newrelic.instrumentation.vertx-sqlclient-4.4.2'
+    }
+}
+
+dependencies {
+    implementation(project(":agent-bridge"))
+    implementation(project(":agent-bridge-datastore"))
+    implementation("io.vertx:vertx-pg-client:4.4.2")
+
+    testImplementation("io.vertx:vertx-core:4.4.2")
+    testImplementation("io.vertx:vertx-web:4.4.2")
+    testImplementation('org.testcontainers:postgresql:1.20.1')
+    testImplementation('com.ongres.scram:client:2.1')
+}
+
+verifyInstrumentation {
+    passes 'io.vertx:vertx-pg-client:[4.4.2,)'
+    excludeRegex '.*SNAPSHOT'
+    excludeRegex '.*milestone.*'
+    excludeRegex '.*alpha.*'
+    excludeRegex '.*Beta.*'
+    excludeRegex '.*CR.*'
+}
+
+site {
+    title 'Vertx'
+    type 'Framework'
+}

instrumentation/vertx-postgres-sqlclient-4.4.2/src/main/java/com/nr/vertx/sqlclient/instrumentation/NRSqlClientWrapper.java

@@ -0,0 +1,47 @@
+/*
+ *
+ *  * Copyright 2024 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+package com.nr.vertx.sqlclient.instrumentation;
+
+import com.newrelic.api.agent.DatastoreParameters;
+import com.newrelic.api.agent.NewRelic;
+import com.newrelic.api.agent.Segment;
+import com.newrelic.api.agent.Token;
+import com.newrelic.api.agent.Trace;
+import io.vertx.core.Handler;
+
+import java.util.logging.Level;
+
+public class NRSqlClientWrapper<E> implements Handler<E> {
+    private final Handler<E> delegate;
+    private Token token;
+    private Segment segment;
+
+    public NRSqlClientWrapper(Handler<E> delegate, Segment segment) {
+        this.delegate = delegate;
+        token = NewRelic.getAgent().getTransaction().getToken();
+        this.segment = segment;
+    }
+
+    @Override
+    @Trace(async = true)
+    public void handle(E event) {
+        if (token != null) {
+            token.linkAndExpire();
+            token = null;
+        }
+
+        if (segment != null) {
+            segment.end();
+            segment = null;
+        }
+
+        if (delegate != null) {
+            NewRelic.getAgent().getTracedMethod().setMetricName("Java", delegate.getClass().getName(), "handle");
+            delegate.handle(event);
+        }
+    }
+}

instrumentation/vertx-postgres-sqlclient-4.4.2/src/main/java/com/nr/vertx/sqlclient/instrumentation/SqlClientUtils.java

@@ -0,0 +1,42 @@
+/*
+ *
+ *  * Copyright 2024 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+package com.nr.vertx.sqlclient.instrumentation;
+
+import com.newrelic.agent.bridge.datastore.OperationAndTableName;
+import com.newrelic.agent.bridge.datastore.R2dbcOperation;
+import com.newrelic.api.agent.NewRelic;
+import io.vertx.sqlclient.impl.command.CommandBase;
+import io.vertx.sqlclient.impl.command.PrepareStatementCommand;
+import io.vertx.sqlclient.impl.command.QueryCommandBase;
+
+import java.util.logging.Level;
+
+public class SqlClientUtils {
+    public static <R> OperationAndTableName extractSqlFromSqlClientCommand(CommandBase<R> cmd) {
+        String sql = null;
+
+        if (cmd != null) {
+            if (cmd instanceof QueryCommandBase) {
+                QueryCommandBase<?> qCmd = (QueryCommandBase<?>)cmd;
+                sql = qCmd.sql();
+            }
+            if (cmd instanceof PrepareStatementCommand) {
+                PrepareStatementCommand pCmd = (PrepareStatementCommand)cmd;
+                sql = pCmd.sql();
+            }
+
+            if (sql != null) {
+                OperationAndTableName operationAndTableName = R2dbcOperation.extractFrom(sql);
+                if (operationAndTableName != null) {
+                    return operationAndTableName;
+                }
+            }
+        }
+
+        return R2dbcOperation.UNKNOWN_OPERATION_AND_TABLE_NAME;
+    }
+}

instrumentation/vertx-postgres-sqlclient-4.4.2/src/main/java/io/vertx/pgclient/impl/PgSocketConnection_Instrumentation.java

@@ -0,0 +1,51 @@
+/*
+ *
+ *  * Copyright 2024 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+package io.vertx.pgclient.impl;
+
+import com.newrelic.agent.bridge.datastore.DatastoreVendor;
+import com.newrelic.agent.bridge.datastore.OperationAndTableName;
+import com.newrelic.api.agent.DatastoreParameters;
+import com.newrelic.api.agent.NewRelic;
+import com.newrelic.api.agent.Segment;
+import com.newrelic.api.agent.Trace;
+import com.newrelic.api.agent.weaver.MatchType;
+import com.newrelic.api.agent.weaver.Weave;
+import com.newrelic.api.agent.weaver.Weaver;
+import com.nr.vertx.sqlclient.instrumentation.NRSqlClientWrapper;
+import com.nr.vertx.sqlclient.instrumentation.SqlClientUtils;
+import io.vertx.core.AsyncResult;
+import io.vertx.core.Handler;
+import io.vertx.pgclient.PgConnectOptions;
+import io.vertx.sqlclient.impl.command.CommandBase;
+
+import java.util.logging.Level;
+
+@Weave(type = MatchType.ExactClass, originalName = "io.vertx.pgclient.impl.PgSocketConnection")
+public abstract class PgSocketConnection_Instrumentation {
+    @Trace
+    protected <R> void doSchedule(CommandBase<R> cmd, Handler<AsyncResult<R>> handler) {
+        if (!(handler instanceof NRSqlClientWrapper)) {
+            OperationAndTableName operationAndTableName = SqlClientUtils.extractSqlFromSqlClientCommand(cmd);
+            PgConnectOptions pgConnectOptions = connectOptions();
+            Segment segment = NewRelic.getAgent().getTransaction().startSegment("Query");
+
+            DatastoreParameters databaseParams = DatastoreParameters.product(DatastoreVendor.Postgres.name())
+                    .collection(operationAndTableName.getTableName())
+                    .operation(operationAndTableName.getOperation())
+                    .instance(pgConnectOptions.getHost(), pgConnectOptions.getPort())
+                    .databaseName(pgConnectOptions.getDatabase())
+                    .build();
+
+            segment.reportAsExternal(databaseParams);
+            handler = new NRSqlClientWrapper(handler, segment);
+        }
+
+        Weaver.callOriginal();
+    }
+
+    protected abstract PgConnectOptions connectOptions();
+}

instrumentation/vertx-postgres-sqlclient-4.4.2/src/main/java/io/vertx/sqlclient/impl/SqlClientBase_Instrumentation.java

@@ -0,0 +1,43 @@
+/*
+ *
+ *  * Copyright 2024 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+package io.vertx.sqlclient.impl;
+
+import com.newrelic.api.agent.Trace;
+import com.newrelic.api.agent.weaver.MatchType;
+import com.newrelic.api.agent.weaver.Weave;
+import com.newrelic.api.agent.weaver.Weaver;
+
+import io.vertx.core.AsyncResult;
+import io.vertx.core.Handler;
+import io.vertx.sqlclient.Query;
+import io.vertx.sqlclient.Row;
+import io.vertx.sqlclient.RowSet;
+import io.vertx.sqlclient.SqlResult;
+import io.vertx.sqlclient.Tuple;
+
+import java.util.List;
+
+@Weave(type = MatchType.ExactClass, originalName = "io.vertx.sqlclient.impl.SqlClientBase")
+public abstract class SqlClientBase_Instrumentation {
+    @Trace
+    public abstract Query<RowSet<Row>> query(String sql);
+
+    @Weave(type = MatchType.ExactClass, originalName = "io.vertx.sqlclient.impl.SqlClientBase$QueryImpl")
+    private static abstract class QueryImpl_Instrumentation<T, R extends SqlResult<T>>  {
+        @Trace
+        public abstract void execute(Handler<AsyncResult<R>> handler);
+    }
+
+    @Weave(type = MatchType.ExactClass, originalName = "io.vertx.sqlclient.impl.SqlClientBase$PreparedQueryImpl")
+    private static abstract class PreparedQueryImpl_Instrumentation<T, R extends SqlResult<T>>  {
+        @Trace
+        public abstract void execute(Tuple arguments, Handler<AsyncResult<R>> handler);
+
+        @Trace
+        public abstract void executeBatch(List<Tuple> batch, Handler<AsyncResult<R>> handler);
+    }
+}