Add R2DBC MariaDB Instrumentation (#799)

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

@@ -0,0 +1,19 @@
+package com.newrelic.agent.bridge.datastore;
+
+public class OperationAndTableName {
+    private final String operation;
+    private final String tableName;
+
+    public OperationAndTableName(String operation, String tableName) {
+        this.operation = operation;
+        this.tableName = tableName;
+    }
+
+    public String getOperation() {
+        return operation;
+    }
+
+    public String getTableName() {
+        return tableName;
+    }
+}

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

@@ -29,15 +29,15 @@ public class R2dbcOperation {
         OPERATION_PATTERNS.put("EXEC", new Pattern[]{Pattern.compile(".*(?:exec|execute)\\s+(?!as\\s+)([^\\s(,=;]*+);?\\s*+(?:[^=]|$).*", PATTERN_SWITCHES), Pattern.compile(".*(?:exec|execute)\\s+[^\\s(,]*.*?=(?:\\s|)([^\\s]*)", PATTERN_SWITCHES)});
     }
 
-    public static String[] extractFrom(String sql) {
+    public static OperationAndTableName extractFrom(String sql) {
         try {
             String strippedSql = COMMENT_PATTERN.matcher(sql).replaceAll("");
             for (Map.Entry<String, Pattern[]> operation : OPERATION_PATTERNS.entrySet()) {
                 for (Pattern pattern : operation.getValue()) {
                     Matcher matcher = pattern.matcher(strippedSql);
                     if(matcher.find()) {
                         String model = matcher.groupCount() > 0 ? removeBrackets(unquoteDatabaseName(matcher.group(1).trim())) : "unknown";
-                        return new String[] { operation.getKey(), VALID_METRIC_NAME_MATCHER.matcher(model).matches() ? model : "ParseError" };
+                        return new OperationAndTableName(operation.getKey(), VALID_METRIC_NAME_MATCHER.matcher(model).matches() ? model : "ParseError");
                     }
                 }
             }

instrumentation/r2dbc-h2/src/main/java/com/nr/agent/instrumentation/r2dbc/R2dbcUtils.java

@@ -3,6 +3,7 @@
 import com.newrelic.agent.bridge.NoOpTransaction;
 import com.newrelic.agent.bridge.datastore.DatastoreVendor;
 import com.newrelic.agent.bridge.datastore.JdbcHelper;
+import com.newrelic.agent.bridge.datastore.OperationAndTableName;
 import com.newrelic.agent.bridge.datastore.R2dbcObfuscator;
 import com.newrelic.agent.bridge.datastore.R2dbcOperation;
 import com.newrelic.api.agent.DatastoreParameters;
@@ -31,12 +32,12 @@ public static Flux<H2Result> wrapRequest(Flux<H2Result> request, String sql, Str
 
     private static Consumer<Subscription> reportExecution(String sql, String databaseName, String url, Segment segment) {
         return (subscription) -> {
-            String[] sqlOperationCollection = R2dbcOperation.extractFrom(sql);
-            if (sqlOperationCollection != null) {
+            OperationAndTableName sqlOperation = R2dbcOperation.extractFrom(sql);
+            if (sqlOperation != null) {
                 segment.reportAsExternal(DatastoreParameters
                         .product(DatastoreVendor.H2.name())
-                        .collection(sqlOperationCollection[1])
-                        .operation(sqlOperationCollection[0])
+                        .collection(sqlOperation.getTableName())
+                        .operation(sqlOperation.getOperation())
                         .instance("localhost", JdbcHelper.parseInMemoryIdentifier(url))
                         .databaseName(databaseName)
                         .slowQuery(sql, R2dbcObfuscator.QUERY_CONVERTER)

instrumentation/r2dbc-mariadb/build.gradle

@@ -0,0 +1,20 @@
+dependencies {
+    implementation(project(":agent-bridge"))
+    implementation(project(":agent-bridge-datastore"))
+    implementation("org.mariadb:r2dbc-mariadb:1.0.2")
+    testImplementation("ch.vorburger.mariaDB4j:mariaDB4j:2.2.1")
+}
+
+jar {
+    manifest { attributes 'Implementation-Title': 'com.newrelic.instrumentation.r2dbc-mariadb' }
+}
+
+verifyInstrumentation {
+    passesOnly 'org.mariadb:r2dbc-mariadb:[1.0.2,)'
+    excludeRegex(".*(alpha|beta|rc).*")
+}
+
+site {
+    title 'MariaDB R2DBC'
+    type 'Datastore'
+}

instrumentation/r2dbc-mariadb/src/main/java/com/nr/agent/instrumentation/r2dbc/R2dbcUtils.java

@@ -0,0 +1,47 @@
+package com.nr.agent.instrumentation.r2dbc;
+
+import com.newrelic.agent.bridge.NoOpTransaction;
+import com.newrelic.agent.bridge.datastore.OperationAndTableName;
+import com.newrelic.agent.bridge.datastore.R2dbcObfuscator;
+import com.newrelic.agent.bridge.datastore.R2dbcOperation;
+import com.newrelic.api.agent.DatastoreParameters;
+import com.newrelic.api.agent.NewRelic;
+import com.newrelic.api.agent.Segment;
+import com.newrelic.api.agent.Transaction;
+import org.mariadb.r2dbc.api.MariadbResult;
+import org.mariadb.r2dbc.client.Client;
+import org.reactivestreams.Subscription;
+import reactor.core.publisher.Flux;
+
+import java.util.function.Consumer;
+
+public class R2dbcUtils {
+    public static Flux<MariadbResult> wrapRequest(Flux<MariadbResult> request, String sql, Client client) {
+        if(request != null) {
+            Transaction transaction = NewRelic.getAgent().getTransaction();
+            if(transaction != null && !(transaction instanceof NoOpTransaction)) {
+                Segment segment = transaction.startSegment("execute");
+                return request
+                        .doOnSubscribe(reportExecution(sql, client, segment))
+                        .doFinally((type) -> segment.end());
+            }
+        }
+        return request;
+    }
+
+    private static Consumer<Subscription> reportExecution(String sql, Client client, Segment segment) {
+        return (subscription) -> {
+            OperationAndTableName sqlOperation = R2dbcOperation.extractFrom(sql);
+            if (sqlOperation != null) {
+                segment.reportAsExternal(DatastoreParameters
+                        .product("MariaDB")
+                        .collection(sqlOperation.getTableName())
+                        .operation(sqlOperation.getOperation())
+                        .instance(client.getConf().getHost(), client.getConf().getPort())
+                        .databaseName(client.getConf().getDatabase())
+                        .slowQuery(sql, R2dbcObfuscator.MYSQL_QUERY_CONVERTER)
+                        .build());
+            }
+        };
+    }
+}

instrumentation/r2dbc-mariadb/src/main/java/org/mariadb/r2dbc/MariadbClientParameterizedQueryStatement_Instrumentation.java

@@ -0,0 +1,20 @@
+package org.mariadb.r2dbc;
+
+import com.newrelic.api.agent.weaver.MatchType;
+import com.newrelic.api.agent.weaver.Weave;
+import com.newrelic.api.agent.weaver.Weaver;
+import com.nr.agent.instrumentation.r2dbc.R2dbcUtils;
+import org.mariadb.r2dbc.api.MariadbResult;
+import org.mariadb.r2dbc.client.Client;
+import reactor.core.publisher.Flux;
+
+@Weave(type = MatchType.ExactClass, originalName = "org.mariadb.r2dbc.MariadbClientParameterizedQueryStatement")
+final class MariadbClientParameterizedQueryStatement_Instrumentation {
+    private final Client client = Weaver.callOriginal();
+    private final String sql = Weaver.callOriginal();
+
+    public Flux<MariadbResult> execute() {
+        Flux<MariadbResult> request = Weaver.callOriginal();
+        return R2dbcUtils.wrapRequest(request, sql, client);
+    }
+}

instrumentation/r2dbc-mariadb/src/main/java/org/mariadb/r2dbc/MariadbServerParameterizedQueryStatement_Instrumentation.java

@@ -0,0 +1,20 @@
+package org.mariadb.r2dbc;
+
+import com.newrelic.api.agent.weaver.MatchType;
+import com.newrelic.api.agent.weaver.Weave;
+import com.newrelic.api.agent.weaver.Weaver;
+import com.nr.agent.instrumentation.r2dbc.R2dbcUtils;
+import org.mariadb.r2dbc.api.MariadbResult;
+import org.mariadb.r2dbc.client.Client;
+import reactor.core.publisher.Flux;
+
+@Weave(type = MatchType.ExactClass, originalName = "org.mariadb.r2dbc.MariadbServerParameterizedQueryStatement")
+final class MariadbServerParameterizedQueryStatement_Instrumentation {
+    private final Client client = Weaver.callOriginal();
+    private final String initialSql = Weaver.callOriginal();
+
+    public Flux<MariadbResult> execute() {
+        Flux<MariadbResult> request = Weaver.callOriginal();
+        return R2dbcUtils.wrapRequest(request, initialSql, client);
+    }
+}

instrumentation/r2dbc-mariadb/src/main/java/org/mariadb/r2dbc/MariadbSimpleQueryStatement_Instrumentation.java

@@ -0,0 +1,20 @@
+package org.mariadb.r2dbc;
+
+import com.newrelic.api.agent.weaver.MatchType;
+import com.newrelic.api.agent.weaver.Weave;
+import com.newrelic.api.agent.weaver.Weaver;
+import com.nr.agent.instrumentation.r2dbc.R2dbcUtils;
+import org.mariadb.r2dbc.api.MariadbResult;
+import org.mariadb.r2dbc.client.Client;
+import reactor.core.publisher.Flux;
+
+@Weave(type = MatchType.ExactClass, originalName = "org.mariadb.r2dbc.MariadbSimpleQueryStatement")
+final class MariadbSimpleQueryStatement_Instrumentation {
+    private final Client client = Weaver.callOriginal();
+    private final String sql = Weaver.callOriginal();
+
+    public Flux<MariadbResult> execute() {
+        Flux<MariadbResult> request = Weaver.callOriginal();
+        return R2dbcUtils.wrapRequest(request, sql, client);
+    }
+}

instrumentation/r2dbc-mariadb/src/test/java/com/nr/agent/instrumentation/r2dbc/MariadbInstrumentedTest.java

@@ -0,0 +1,71 @@
+package com.nr.agent.instrumentation.r2dbc;
+
+import ch.vorburger.mariadb4j.DB;
+import ch.vorburger.mariadb4j.DBConfigurationBuilder;
+import com.newrelic.agent.introspec.DatastoreHelper;
+import com.newrelic.agent.introspec.InstrumentationTestConfig;
+import com.newrelic.agent.introspec.InstrumentationTestRunner;
+import com.newrelic.agent.introspec.Introspector;
+import io.r2dbc.spi.Connection;
+import io.r2dbc.spi.ConnectionFactories;
+import io.r2dbc.spi.ConnectionFactory;
+import org.junit.AfterClass;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import reactor.core.publisher.Mono;
+
+import static org.junit.Assert.assertEquals;
+
+@RunWith(InstrumentationTestRunner.class)
+@InstrumentationTestConfig(includePrefixes = "org.mariadb.r2dbc")
+public class MariadbInstrumentedTest {
+
+    public static DB mariaDb;
+    public Connection connection;
+
+    @Before
+    public void setup() throws Exception {
+        String databaseName = "MariaDB" + System.currentTimeMillis();
+        DBConfigurationBuilder builder = DBConfigurationBuilder.newBuilder().setPort(0);
+        mariaDb = DB.newEmbeddedDB(builder.build());
+        mariaDb.start();
+        mariaDb.createDB(databaseName);
+        mariaDb.source("users.sql", "user", "password", databaseName);
+        ConnectionFactory connectionFactory = ConnectionFactories.get(builder.getURL(databaseName).replace("mysql", "mariadb").replace("jdbc", "r2dbc").replace("localhost", "user:password@localhost"));
+        connection = Mono.from(connectionFactory.create()).block();
+    }
+
+    @AfterClass
+    public static void teardown() throws Exception {
+        mariaDb.stop();
+    }
+
+    @Test
+    public void testSelect() {
+        //Given
+        Introspector introspector = InstrumentationTestRunner.getIntrospector();
+        DatastoreHelper helper = new DatastoreHelper("MariaDB");
+
+        //When
+        R2dbcTestUtils.basicRequests(connection);
+
+        //Then
+        assertEquals(1, introspector.getFinishedTransactionCount(1000));
+        assertEquals(1, introspector.getTransactionNames().size());
+        String transactionName = introspector.getTransactionNames().stream().findFirst().orElse("");
+        helper.assertScopedStatementMetricCount(transactionName, "INSERT", "USERS", 1);
+        helper.assertScopedStatementMetricCount(transactionName, "SELECT", "USERS", 3);
+        helper.assertScopedStatementMetricCount(transactionName, "UPDATE", "USERS", 1);
+        helper.assertScopedStatementMetricCount(transactionName, "DELETE", "USERS", 1);
+        helper.assertAggregateMetrics();
+        helper.assertUnscopedOperationMetricCount("INSERT", 1);
+        helper.assertUnscopedOperationMetricCount("SELECT", 3);
+        helper.assertUnscopedOperationMetricCount("UPDATE", 1);
+        helper.assertUnscopedOperationMetricCount("DELETE", 1);
+        helper.assertUnscopedStatementMetricCount("INSERT", "USERS", 1);
+        helper.assertUnscopedStatementMetricCount("SELECT", "USERS", 3);
+        helper.assertUnscopedStatementMetricCount("UPDATE", "USERS", 1);
+        helper.assertUnscopedStatementMetricCount("DELETE", "USERS", 1);
+    }
+}

instrumentation/r2dbc-mariadb/src/test/java/com/nr/agent/instrumentation/r2dbc/MariadbNoInstrumentationTest.java

@@ -0,0 +1,62 @@
+package com.nr.agent.instrumentation.r2dbc;
+
+import ch.vorburger.mariadb4j.DB;
+import ch.vorburger.mariadb4j.DBConfigurationBuilder;
+import com.newrelic.agent.introspec.DatastoreHelper;
+import com.newrelic.agent.introspec.InstrumentationTestConfig;
+import com.newrelic.agent.introspec.InstrumentationTestRunner;
+import com.newrelic.agent.introspec.Introspector;
+import io.r2dbc.spi.Connection;
+import io.r2dbc.spi.ConnectionFactories;
+import io.r2dbc.spi.ConnectionFactory;
+import org.junit.AfterClass;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import reactor.core.publisher.Mono;
+
+import static org.junit.Assert.assertEquals;
+
+@RunWith(InstrumentationTestRunner.class)
+@InstrumentationTestConfig(includePrefixes = "none")
+public class MariadbNoInstrumentationTest {
+
+    public static DB mariaDb;
+    public Connection connection;
+
+    @Before
+    public void setup() throws Exception {
+        String databaseName = "MariaDB" + System.currentTimeMillis();
+        DBConfigurationBuilder builder = DBConfigurationBuilder.newBuilder().setPort(0);
+        mariaDb = DB.newEmbeddedDB(builder.build());
+        mariaDb.start();
+        mariaDb.createDB(databaseName);
+        mariaDb.source("users.sql", "user", "password", databaseName);
+        ConnectionFactory connectionFactory = ConnectionFactories.get(builder.getURL(databaseName).replace("mysql", "mariadb").replace("jdbc", "r2dbc").replace("localhost", "user:password@localhost"));
+        connection = Mono.from(connectionFactory.create()).block();
+    }
+
+    @AfterClass
+    public static void teardown() throws Exception {
+        mariaDb.stop();
+    }
+
+    @Test
+    public void testSelect() {
+        //Given
+        Introspector introspector = InstrumentationTestRunner.getIntrospector();
+        DatastoreHelper helper = new DatastoreHelper("MariaDB");
+
+        //When
+        R2dbcTestUtils.basicRequests(connection);
+
+        //Then
+        assertEquals(1, introspector.getFinishedTransactionCount(1000));
+        assertEquals(1, introspector.getTransactionNames().size());
+        String transactionName = introspector.getTransactionNames().stream().findFirst().orElse("");
+        helper.assertScopedStatementMetricCount(transactionName, "INSERT", "USERS", 0);
+        helper.assertScopedStatementMetricCount(transactionName, "SELECT", "USERS", 0);
+        helper.assertScopedStatementMetricCount(transactionName, "UPDATE", "USERS", 0);
+        helper.assertScopedStatementMetricCount(transactionName, "DELETE", "USERS", 0);
+    }
+}

instrumentation/r2dbc-mariadb/src/test/java/com/nr/agent/instrumentation/r2dbc/R2dbcTestUtils.java

@@ -0,0 +1,17 @@
+package com.nr.agent.instrumentation.r2dbc;
+
+import com.newrelic.api.agent.Trace;
+import io.r2dbc.spi.Connection;
+import reactor.core.publisher.Mono;
+
+public class R2dbcTestUtils {
+    @Trace(dispatcher = true)
+    public static void basicRequests(Connection connection) {
+        Mono.from(connection.createStatement("INSERT INTO USERS(id, first_name, last_name, age) VALUES(1, 'Max', 'Power', 30)").execute()).block();
+        Mono.from(connection.createStatement("SELECT * FROM USERS WHERE last_name='Power'").execute()).block();
+        Mono.from(connection.createStatement("UPDATE USERS SET age = 36 WHERE last_name = 'Power'").execute()).block();
+        Mono.from(connection.createStatement("SELECT * FROM USERS WHERE last_name='Power'").execute()).block();
+        Mono.from(connection.createStatement("DELETE FROM USERS WHERE last_name = 'Power'").execute()).block();
+        Mono.from(connection.createStatement("SELECT * FROM USERS").execute()).block();
+    }
+}

instrumentation/r2dbc-mariadb/src/test/resources/users.sql

@@ -0,0 +1,2 @@
+CREATE TABLE IF NOT EXISTS USERS(id int primary key, first_name varchar(255), last_name varchar(255), age int);
+TRUNCATE TABLE USERS

settings.gradle

@@ -239,6 +239,7 @@ include 'instrumentation:ning-async-http-client-1.0'
 include 'instrumentation:ning-async-http-client-1.1'
 include 'instrumentation:ning-async-http-client-1.6.1'
 include 'instrumentation:r2dbc-h2'
+include 'instrumentation:r2dbc-mariadb'
 include 'instrumentation:rabbit-amqp-2.7'
 include 'instrumentation:rabbit-amqp-3.5.0'
 include 'instrumentation:rabbit-amqp-5.0.0'