Vertx 4.4.2 Postgresql SqlClient Instrumentation

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,53 @@
+/*
+ *
+ *  * 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;
+    private final DatastoreParameters params;
+
+    public NRSqlClientWrapper(Handler<E> delegate, Segment segment, DatastoreParameters dbParams) {
+        this.delegate = delegate;
+        token = NewRelic.getAgent().getTransaction().getToken();
+        this.segment = segment;
+        params = dbParams;
+    }
+
+    @Override
+    @Trace(async = true)
+    public void handle(E event) {
+        if (token != null) {
+            token.linkAndExpire();
+            token = null;
+        }
+
+        if (segment != null) {
+            if (params != null) {
+                segment.reportAsExternal(params);
+            }
+            segment.end();
+            segment = null;
+        } else if (params != null) {
+            NewRelic.getAgent().getTracedMethod().reportAsExternal(params);
+        }
+
+        if (delegate != null) {
+            delegate.handle(event);
+        }
+    }
+}

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

@@ -0,0 +1,40 @@
+/*
+ *
+ *  * 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) {
+        OperationAndTableName operationAndTableName = null;
+        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 = R2dbcOperation.extractFrom(sql);
+            }
+        }
+
+        return operationAndTableName;
+    }
+}

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

@@ -0,0 +1,52 @@
+/*
+ *
+ *  * 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);
+            if (operationAndTableName != null) {
+                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();
+
+                handler = new NRSqlClientWrapper(handler, segment, databaseParams);
+            }
+
+        }
+        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,51 @@
+/*
+ *
+ *  * 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 Query<RowSet<Row>> query(String sql) {
+        return Weaver.callOriginal();
+    }
+
+    @Weave(type = MatchType.ExactClass, originalName = "io.vertx.sqlclient.impl.SqlClientBase$QueryImpl")
+    private static abstract class QueryImpl_Instrumentation<T, R extends SqlResult<T>>  {
+        @Trace
+        public void execute(Handler<AsyncResult<R>> handler) {
+            Weaver.callOriginal();
+        }
+    }
+
+    @Weave(type = MatchType.ExactClass, originalName = "io.vertx.sqlclient.impl.SqlClientBase$PreparedQueryImpl")
+    private static abstract class PreparedQueryImpl_Instrumentation<T, R extends SqlResult<T>>  {
+        @Trace
+        public void execute(Tuple arguments, Handler<AsyncResult<R>> handler) {
+            Weaver.callOriginal();
+        }
+
+        @Trace
+        public void executeBatch(List<Tuple> batch, Handler<AsyncResult<R>> handler) {
+            Weaver.callOriginal();
+        }
+    }
+}