R2DBC Postgres Instrumentation

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

@@ -26,15 +26,21 @@ public class R2dbcObfuscator {
     private static final Pattern ALL_UNMATCHED_PATTERN;
     private static final Pattern MYSQL_DIALECT_PATTERN;
     private static final Pattern MYSQL_UNMATCHED_PATTERN;
+    private static final Pattern POSTGRES_DIALECT_PATTERN;
+    private static final Pattern POSTGRES_UNMATCHED_PATTERN;
 
     public static final QueryConverter<String> QUERY_CONVERTER;
     public static final QueryConverter<String> MYSQL_QUERY_CONVERTER;
+    public static final QueryConverter<String> POSTGRES_QUERY_CONVERTER;
 
     static {
         ALL_DIALECTS_PATTERN = Pattern.compile(String.join("|", SINGLE_QUOTE, DOUBLE_QUOTE, DOLLAR_QUOTE, COMMENT, MULTILINE_COMMENT, UUID, HEX, BOOLEAN, NUMBER), PATTERN_SWITCHES);
         ALL_UNMATCHED_PATTERN = Pattern.compile("'|\"|/\\*|\\*/|\\$", PATTERN_SWITCHES);
         MYSQL_DIALECT_PATTERN = Pattern.compile(String.join("|", SINGLE_QUOTE, DOUBLE_QUOTE, COMMENT, MULTILINE_COMMENT, HEX, BOOLEAN, NUMBER), PATTERN_SWITCHES);
         MYSQL_UNMATCHED_PATTERN = Pattern.compile("'|\"|/\\*|\\*/", PATTERN_SWITCHES);
+        POSTGRES_DIALECT_PATTERN = Pattern.compile(String.join(SINGLE_QUOTE, DOLLAR_QUOTE, COMMENT, MULTILINE_COMMENT, UUID, BOOLEAN, NUMBER), PATTERN_SWITCHES);
+        POSTGRES_UNMATCHED_PATTERN = Pattern.compile("'|/\\*|\\*/|\\$(?!\\?)", PATTERN_SWITCHES);
+
 
         QUERY_CONVERTER = new QueryConverter<String>() {
             @Override
@@ -59,6 +65,18 @@ public String toObfuscatedQueryString(String statement) {
                 return obfuscateSql(statement, MYSQL_DIALECT_PATTERN, MYSQL_UNMATCHED_PATTERN);
             }
         };
+
+        POSTGRES_QUERY_CONVERTER = new QueryConverter<String>() {
+            @Override
+            public String toRawQueryString(String statement) {
+                return statement;
+            }
+
+            @Override
+            public String toObfuscatedQueryString(String statement) {
+                return obfuscateSql(statement, POSTGRES_DIALECT_PATTERN, POSTGRES_UNMATCHED_PATTERN);
+            }
+        };
     }
 
     private static String obfuscateSql(String sql, Pattern dialect, Pattern unmatched) {

instrumentation/r2dbc-postgresql/build.gradle

@@ -0,0 +1,20 @@
+dependencies {
+    implementation(project(":agent-bridge"))
+    implementation(project(":agent-bridge-datastore"))
+    implementation("org.postgresql:r2dbc-postgresql:0.9.1.RELEASE")
+    testImplementation("ru.yandex.qatools.embed:postgresql-embedded:2.10")
+}
+
+jar {
+    manifest { attributes 'Implementation-Title': 'com.newrelic.instrumentation.r2dbc-postgresql' }
+}
+
+verifyInstrumentation {
+    passesOnly 'org.postgresql:r2dbc-postgresql:[0.9.0,)'
+    excludeRegex(".*(M1|M2|RC).*")
+}
+
+site {
+    title 'PostgreSQL R2DBC'
+    type 'Datastore'
+}

instrumentation/r2dbc-postgresql/src/main/java/io/r2dbc/postgresql/PostgresqlStatement_Instrumentation.java

@@ -0,0 +1,21 @@
+package io.r2dbc.postgresql;
+
+import com.newrelic.api.agent.weaver.MatchType;
+import com.newrelic.api.agent.weaver.Weave;
+import com.newrelic.api.agent.weaver.Weaver;
+import io.r2dbc.postgresql.api.PostgresqlResult;
+import reactor.core.publisher.Flux;
+
+@Weave(type = MatchType.ExactClass, originalName = "io.r2dbc.postgresql.PostgresqlStatement")
+final class PostgresqlStatement_Instrumentation {
+    private final TokenizedSql tokenizedSql = Weaver.callOriginal();
+    private final ConnectionResources resources = Weaver.callOriginal();
+
+    public Flux<PostgresqlResult> execute() {
+        Flux<PostgresqlResult> request = Weaver.callOriginal();
+        if(request != null && tokenizedSql != null && resources != null) {
+             return R2dbcUtils.wrapRequest(request, tokenizedSql.getSql(), resources.getConfiguration());
+        }
+        return request;
+    }
+}

instrumentation/r2dbc-postgresql/src/main/java/io/r2dbc/postgresql/R2dbcUtils.java

@@ -0,0 +1,47 @@
+package io.r2dbc.postgresql;
+
+import com.newrelic.agent.bridge.NoOpTransaction;
+import com.newrelic.agent.bridge.datastore.DatastoreVendor;
+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 io.r2dbc.postgresql.api.PostgresqlResult;
+import org.reactivestreams.Subscription;
+import reactor.core.publisher.Flux;
+
+import java.util.function.Consumer;
+
+public class R2dbcUtils {
+    public static Flux<PostgresqlResult> wrapRequest(Flux<PostgresqlResult> request, String sql, 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))
+                        .doFinally((type) -> segment.end());
+            }
+        }
+        return request;
+    }
+
+    private static Consumer<Subscription> reportExecution(String sql, PostgresqlConnectionConfiguration connectionConfiguration, Segment segment) {
+        return (subscription) -> {
+            OperationAndTableName sqlOperation = R2dbcOperation.extractFrom(sql);
+            if (sqlOperation != null) {
+                segment.reportAsExternal(DatastoreParameters
+                        .product(DatastoreVendor.Postgres.name())
+                        .collection(sqlOperation.getTableName())
+                        .operation(sqlOperation.getOperation())
+                        .instance(connectionConfiguration.getHost(), connectionConfiguration.getPort())
+                        .databaseName(connectionConfiguration.getDatabase())
+                        .slowQuery(sql, R2dbcObfuscator.POSTGRES_QUERY_CONVERTER)
+                        .build());
+            }
+        };
+    }
+}

instrumentation/r2dbc-postgresql/src/test/java/com/nr/agent/instrumentation/r2dbc/PostgresqlInstrumentedTest.java

@@ -0,0 +1,71 @@
+package com.nr.agent.instrumentation.r2dbc;
+
+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 ru.yandex.qatools.embed.postgresql.EmbeddedPostgres;
+
+import static org.junit.Assert.assertEquals;
+import static ru.yandex.qatools.embed.postgresql.distribution.Version.Main.V9_6;
+
+@RunWith(InstrumentationTestRunner.class)
+@InstrumentationTestConfig(includePrefixes = "io.r2dbc.postgresql")
+public class PostgresqlInstrumentedTest {
+
+    public static final EmbeddedPostgres postgres = new EmbeddedPostgres(V9_6);
+    public static Connection connection;
+
+    @Before
+    public void setup() throws Exception {
+        String databaseName = "Postgres" + System.currentTimeMillis();
+        final String url = postgres.start("localhost", 5432, databaseName, "user", "password");
+        final String updatedUrl = url.replace("jdbc", "r2dbc").replace("localhost", "user:password@localhost").replace("?user=user&password=password", "");
+        ConnectionFactory connectionFactory = ConnectionFactories.get(updatedUrl);
+        connection = Mono.from(connectionFactory.create()).block();
+        Mono.from(connection.createStatement("CREATE TABLE IF NOT EXISTS USERS(id int primary key, first_name varchar(255), last_name varchar(255), age int);").execute()).block();
+        Mono.from(connection.createStatement("TRUNCATE TABLE USERS;").execute()).block();
+    }
+
+    @AfterClass
+    public static void teardown() {
+        Mono.from(connection.close()).block();
+        postgres.stop();
+    }
+
+    @Test
+    public void testBasicRequests() {
+        //Given
+        Introspector introspector = InstrumentationTestRunner.getIntrospector();
+        DatastoreHelper helper = new DatastoreHelper("Postgres");
+
+        //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-postgresql/src/test/java/com/nr/agent/instrumentation/r2dbc/PostgresqlNoInstrumentationTest.java

@@ -0,0 +1,62 @@
+package com.nr.agent.instrumentation.r2dbc;
+
+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 ru.yandex.qatools.embed.postgresql.EmbeddedPostgres;
+
+import static org.junit.Assert.assertEquals;
+import static ru.yandex.qatools.embed.postgresql.distribution.Version.Main.V9_6;
+
+@RunWith(InstrumentationTestRunner.class)
+@InstrumentationTestConfig(includePrefixes = "none")
+public class PostgresqlNoInstrumentationTest {
+
+    public static final EmbeddedPostgres postgres = new EmbeddedPostgres(V9_6);
+    public static Connection connection;
+
+    @Before
+    public void setup() throws Exception {
+        String databaseName = "Postgres" + System.currentTimeMillis();
+        final String url = postgres.start("localhost", 5432, databaseName, "user", "password");
+        final String updatedUrl = url.replace("jdbc", "r2dbc").replace("localhost", "user:password@localhost").replace("?user=user&password=password", "");
+        ConnectionFactory connectionFactory = ConnectionFactories.get(updatedUrl);
+        connection = Mono.from(connectionFactory.create()).block();
+        Mono.from(connection.createStatement("CREATE TABLE IF NOT EXISTS USERS(id int primary key, first_name varchar(255), last_name varchar(255), age int);").execute()).block();
+        Mono.from(connection.createStatement("TRUNCATE TABLE USERS;").execute()).block();
+    }
+
+    @AfterClass
+    public static void teardown() {
+        Mono.from(connection.close()).block();
+        postgres.stop();
+    }
+
+    @Test
+    public void testBasicRequests() {
+        //Given
+        Introspector introspector = InstrumentationTestRunner.getIntrospector();
+        DatastoreHelper helper = new DatastoreHelper("Postgres");
+
+        //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-postgresql/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();
+    }
+}

settings.gradle

@@ -245,6 +245,7 @@ include 'instrumentation:ning-async-http-client-1.6.1'
 include 'instrumentation:r2dbc-h2'
 include 'instrumentation:r2dbc-mariadb'
 include 'instrumentation:r2dbc-mysql'
+include 'instrumentation:r2dbc-postgresql'
 include 'instrumentation:rabbit-amqp-2.7'
 include 'instrumentation:rabbit-amqp-3.5.0'
 include 'instrumentation:rabbit-amqp-5.0.0'