Merge pull request #169 from DataDog/ark/c3p0

Instrument Connection Constructor to get connection info

dd-java-agent-ittests/src/test/groovy/com/datadoghq/agent/integration/jdbc/JDBCInstrumentationTest.groovy

@@ -2,7 +2,7 @@ package com.datadoghq.agent.integration.jdbc
 
 import com.datadoghq.trace.DDTracer
 import com.datadoghq.trace.writer.ListWriter
-import io.opentracing.util.GlobalTracer
+import dd.test.TestUtils
 import org.apache.derby.jdbc.EmbeddedDriver
 import org.h2.Driver
 import org.hsqldb.jdbc.JDBCDriver
@@ -10,7 +10,6 @@ import spock.lang.Shared
 import spock.lang.Specification
 import spock.lang.Unroll
 
-import java.lang.reflect.Field
 import java.sql.Connection
 import java.sql.PreparedStatement
 import java.sql.ResultSet
@@ -18,8 +17,8 @@ import java.sql.Statement
 
 class JDBCInstrumentationTest extends Specification {
 
-  ListWriter writer = new ListWriter()
+  final ListWriter writer = new ListWriter()
-  DDTracer tracer = new DDTracer(writer)
+  final DDTracer tracer = new DDTracer(writer)
 
   @Shared
   private Map<String, Connection> connections
@@ -43,16 +42,8 @@ class JDBCInstrumentationTest extends Specification {
   }
 
   def setup() {
-    try {
+    TestUtils.registerOrReplaceGlobalTracer(tracer)
-      GlobalTracer.register(tracer)
-    } catch (final Exception e) {
-      // Force it anyway using reflection
-      final Field field = GlobalTracer.getDeclaredField("tracer")
-      field.setAccessible(true)
-      field.set(null, tracer)
-    }
     writer.start()
-    assert GlobalTracer.isRegistered()
   }
 
   @Unroll
@@ -80,6 +71,7 @@ class JDBCInstrumentationTest extends Specification {
 
     def tags = span.context().tags
     tags["db.type"] == driver
+    tags["db.user"] == username
     tags["span.kind"] == "client"
     tags["component"] == "java-jdbc-statement"
 
@@ -88,16 +80,16 @@ class JDBCInstrumentationTest extends Specification {
 
     tags["thread.name"] != null
     tags["thread.id"] != null
-    tags.size() == 7
+    tags.size() == username == null ? 7 : 8
 
     cleanup:
     statement.close()
 
     where:
-    driver   | connection                | query
+    driver   | connection                | username | query
-    "h2"     | connections.get("h2")     | "SELECT 3"
+    "h2"     | connections.get("h2")     | null     | "SELECT 3"
-    "derby"  | connections.get("derby")  | "SELECT 3 FROM SYSIBM.SYSDUMMY1"
+    "derby"  | connections.get("derby")  | "APP"    | "SELECT 3 FROM SYSIBM.SYSDUMMY1"
-    "hsqldb" | connections.get("hsqldb") | "SELECT 3 FROM INFORMATION_SCHEMA.SYSTEM_USERS"
+    "hsqldb" | connections.get("hsqldb") | "SA"     | "SELECT 3 FROM INFORMATION_SCHEMA.SYSTEM_USERS"
   }
 
   @Unroll
@@ -126,6 +118,7 @@ class JDBCInstrumentationTest extends Specification {
 
     def tags = span.context().tags
     tags["db.type"] == driver
+    tags["db.user"] == username
     tags["span.kind"] == "client"
     tags["component"] == "java-jdbc-prepared_statement"
 
@@ -134,16 +127,16 @@ class JDBCInstrumentationTest extends Specification {
 
     tags["thread.name"] != null
     tags["thread.id"] != null
-    tags.size() == 7
+    tags.size() == username == null ? 7 : 8
 
     cleanup:
     statement.close()
 
     where:
-    driver   | connection                | query
+    driver   | connection                | username | query
-    "h2"     | connections.get("h2")     | "SELECT 3"
+    "h2"     | connections.get("h2")     | null     | "SELECT 3"
-    "derby"  | connections.get("derby")  | "SELECT 3 FROM SYSIBM.SYSDUMMY1"
+    "derby"  | connections.get("derby")  | "APP"    | "SELECT 3 FROM SYSIBM.SYSDUMMY1"
-    "hsqldb" | connections.get("hsqldb") | "SELECT 3 FROM INFORMATION_SCHEMA.SYSTEM_USERS"
+    "hsqldb" | connections.get("hsqldb") | "SA"     | "SELECT 3 FROM INFORMATION_SCHEMA.SYSTEM_USERS"
   }
 
   @Unroll
@@ -171,6 +164,7 @@ class JDBCInstrumentationTest extends Specification {
 
     def tags = span.context().tags
     tags["db.type"] == driver
+    tags["db.user"] == username
     tags["span.kind"] == "client"
     tags["component"] == "java-jdbc-prepared_statement"
 
@@ -179,16 +173,16 @@ class JDBCInstrumentationTest extends Specification {
 
     tags["thread.name"] != null
     tags["thread.id"] != null
-    tags.size() == 7
+    tags.size() == username == null ? 7 : 8
 
     cleanup:
     statement.close()
 
     where:
-    driver   | connection                | query
+    driver   | connection                | username | query
-    "h2"     | connections.get("h2")     | "SELECT 3"
+    "h2"     | connections.get("h2")     | null     | "SELECT 3"
-    "derby"  | connections.get("derby")  | "SELECT 3 FROM SYSIBM.SYSDUMMY1"
+    "derby"  | connections.get("derby")  | "APP"    | "SELECT 3 FROM SYSIBM.SYSDUMMY1"
-    "hsqldb" | connections.get("hsqldb") | "SELECT 3 FROM INFORMATION_SCHEMA.SYSTEM_USERS"
+    "hsqldb" | connections.get("hsqldb") | "SA"     | "SELECT 3 FROM INFORMATION_SCHEMA.SYSTEM_USERS"
   }
 
   @Unroll
@@ -217,6 +211,7 @@ class JDBCInstrumentationTest extends Specification {
 
     def tags = span.context().tags
     tags["db.type"] == driver
+    tags["db.user"] == username
     tags["span.kind"] == "client"
     tags["component"] == "java-jdbc-statement"
 
@@ -225,16 +220,16 @@ class JDBCInstrumentationTest extends Specification {
 
     tags["thread.name"] != null
     tags["thread.id"] != null
-    tags.size() == 7
+    tags.size() == username == null ? 7 : 8
 
     cleanup:
     statement.close()
 
     where:
-    driver   | connection                | query
+    driver   | connection                | username | query
-    "h2"     | connections.get("h2")     | "CREATE TABLE S_H2 (id INTEGER not NULL, PRIMARY KEY ( id ))"
+    "h2"     | connections.get("h2")     | null     | "CREATE TABLE S_H2 (id INTEGER not NULL, PRIMARY KEY ( id ))"
-    "derby"  | connections.get("derby")  | "CREATE TABLE S_DERBY (id INTEGER not NULL, PRIMARY KEY ( id ))"
+    "derby"  | connections.get("derby")  | "APP"    | "CREATE TABLE S_DERBY (id INTEGER not NULL, PRIMARY KEY ( id ))"
-    "hsqldb" | connections.get("hsqldb") | "CREATE TABLE PUBLIC.S_HSQLDB (id INTEGER not NULL, PRIMARY KEY ( id ))"
+    "hsqldb" | connections.get("hsqldb") | "SA"     | "CREATE TABLE PUBLIC.S_HSQLDB (id INTEGER not NULL, PRIMARY KEY ( id ))"
   }
 
   @Unroll
@@ -261,6 +256,7 @@ class JDBCInstrumentationTest extends Specification {
 
     def tags = span.context().tags
     tags["db.type"] == driver
+    tags["db.user"] == username
     tags["span.kind"] == "client"
     tags["component"] == "java-jdbc-prepared_statement"
 
@@ -269,15 +265,15 @@ class JDBCInstrumentationTest extends Specification {
 
     tags["thread.name"] != null
     tags["thread.id"] != null
-    tags.size() == 7
+    tags.size() == username == null ? 7 : 8
 
     cleanup:
     statement.close()
 
     where:
-    driver   | connection                | query
+    driver   | connection                | username | query
-    "h2"     | connections.get("h2")     | "CREATE TABLE PS_H2 (id INTEGER not NULL, PRIMARY KEY ( id ))"
+    "h2"     | connections.get("h2")     | null     | "CREATE TABLE PS_H2 (id INTEGER not NULL, PRIMARY KEY ( id ))"
     // Derby calls executeLargeUpdate from executeUpdate thus generating a nested span breaking this test.
-    "hsqldb" | connections.get("hsqldb") | "CREATE TABLE PUBLIC.PS_HSQLDB (id INTEGER not NULL, PRIMARY KEY ( id ))"
+    "hsqldb" | connections.get("hsqldb") | "SA"     | "CREATE TABLE PUBLIC.PS_HSQLDB (id INTEGER not NULL, PRIMARY KEY ( id ))"
   }
 }

dd-java-agent/src/main/java/com/datadoghq/agent/TracingAgent.java

@@ -21,6 +21,7 @@ import static dd.trace.ClassLoaderMatcher.isReflectionClassLoader;
 import static net.bytebuddy.matcher.ElementMatchers.any;
 import static net.bytebuddy.matcher.ElementMatchers.isBootstrapClassLoader;
 import static net.bytebuddy.matcher.ElementMatchers.nameContains;
+import static net.bytebuddy.matcher.ElementMatchers.nameMatches;
 import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;
 
 import dd.trace.Instrumenter;
@@ -77,6 +78,7 @@ public class TracingAgent {
             .or(nameStartsWith("org.slf4j."))
             .or(nameContains("javassist"))
             .or(nameContains(".asm."))
+            .or(nameMatches("com\\.mchange\\.v2\\.c3p0\\..*Proxy"))
             .ignore(
                 any(),
                 isBootstrapClassLoader()

dd-java-agent/src/main/java/com/datadoghq/agent/instrumentation/jdbc/ConnectionInstrumentation.java

@@ -1,6 +1,7 @@
 package com.datadoghq.agent.instrumentation.jdbc;
 
 import static net.bytebuddy.matcher.ElementMatchers.hasSuperType;
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
 import static net.bytebuddy.matcher.ElementMatchers.isInterface;
 import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;
 import static net.bytebuddy.matcher.ElementMatchers.named;
@@ -15,11 +16,13 @@ import java.sql.Connection;
 import java.sql.PreparedStatement;
 import java.util.Map;
 import java.util.WeakHashMap;
+import lombok.Data;
 import net.bytebuddy.agent.builder.AgentBuilder;
 import net.bytebuddy.asm.Advice;
 
 @AutoService(Instrumenter.class)
 public final class ConnectionInstrumentation implements Instrumenter {
+  public static final Map<Connection, DBInfo> connectionInfo = new WeakHashMap<>();
   public static final Map<PreparedStatement, String> preparedStatements = new WeakHashMap<>();
 
   @Override
@@ -32,15 +35,47 @@ public final class ConnectionInstrumentation implements Instrumenter {
                     nameStartsWith("prepare")
                         .and(takesArgument(0, String.class))
                         .and(returns(PreparedStatement.class)),
-                    ConnectionAdvice.class.getName()))
+                    ConnectionPrepareAdvice.class.getName()))
+        .transform(
+            DDAdvice.create().advice(isConstructor(), ConnectionConstructorAdvice.class.getName()))
         .asDecorator();
   }
 
-  public static class ConnectionAdvice {
+  public static class ConnectionPrepareAdvice {
     @Advice.OnMethodExit(suppress = Throwable.class)
     public static void addDBInfo(
         @Advice.Argument(0) final String sql, @Advice.Return final PreparedStatement statement) {
       preparedStatements.put(statement, sql);
     }
   }
+
+  public static class ConnectionConstructorAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void addDBInfo(@Advice.This final Connection connection) {
+      try {
+        final String url = connection.getMetaData().getURL();
+        if (url != null) {
+          // Remove end of url to prevent passwords from leaking:
+          final String sanitizedURL = url.replaceAll("[?;].*", "");
+          final String type = url.split(":")[1];
+          String user = connection.getMetaData().getUserName();
+          if (user != null && user.trim().equals("")) {
+            user = null;
+          }
+          connectionInfo.put(connection, new DBInfo(sanitizedURL, type, user));
+        }
+      } catch (Throwable t) {
+        // object may not be fully initialized.
+        // calling constructor will populate map
+      }
+    }
+  }
+
+  @Data
+  public static class DBInfo {
+    public static DBInfo UNKNOWN = new DBInfo("null", "unknown", null);
+    private final String url;
+    private final String type;
+    private final String user;
+  }
 }

dd-java-agent/src/main/java/com/datadoghq/agent/instrumentation/jdbc/DriverInstrumentation.java

@@ -1,60 +0,0 @@
-package com.datadoghq.agent.instrumentation.jdbc;
-
-import static net.bytebuddy.matcher.ElementMatchers.hasSuperType;
-import static net.bytebuddy.matcher.ElementMatchers.isInterface;
-import static net.bytebuddy.matcher.ElementMatchers.named;
-import static net.bytebuddy.matcher.ElementMatchers.not;
-import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
-
-import com.google.auto.service.AutoService;
-import dd.trace.DDAdvice;
-import dd.trace.Instrumenter;
-import java.sql.Connection;
-import java.sql.Driver;
-import java.util.Map;
-import java.util.Properties;
-import java.util.WeakHashMap;
-import lombok.Data;
-import net.bytebuddy.agent.builder.AgentBuilder;
-import net.bytebuddy.asm.Advice;
-
-@AutoService(Instrumenter.class)
-public final class DriverInstrumentation implements Instrumenter {
-  public static final Map<Connection, DBInfo> connectionInfo = new WeakHashMap<>();
-
-  @Override
-  public AgentBuilder instrument(final AgentBuilder agentBuilder) {
-    return agentBuilder
-        .type(not(isInterface()).and(hasSuperType(named(Driver.class.getName()))))
-        .transform(
-            DDAdvice.create()
-                .advice(
-                    named("connect").and(takesArguments(String.class, Properties.class)),
-                    DriverAdvice.class.getName()))
-        .asDecorator();
-  }
-
-  public static class DriverAdvice {
-    @Advice.OnMethodExit(suppress = Throwable.class)
-    public static void addDBInfo(
-        @Advice.Argument(0) final String url,
-        @Advice.Argument(1) final Properties info,
-        @Advice.Return final Connection connection) {
-      if (url != null) {
-        // Remove end of url to prevent passwords from leaking:
-        final String sanitizedURL = url.replaceAll("[?;].*", "");
-        final String type = url.split(":")[1];
-        final String dbUser = info == null ? null : info.getProperty("user");
-        connectionInfo.put(connection, new DBInfo(sanitizedURL, type, dbUser));
-      }
-    }
-  }
-
-  @Data
-  public static class DBInfo {
-    public static DBInfo UNKNOWN = new DBInfo("null", "unknown", null);
-    private final String url;
-    private final String type;
-    private final String user;
-  }
-}

dd-java-agent/src/main/java/com/datadoghq/agent/instrumentation/jdbc/PreparedStatementInstrumentation.java

@@ -51,9 +51,10 @@ public final class PreparedStatementInstrumentation implements Instrumenter {
         return NoopActiveSpanSource.NoopActiveSpan.INSTANCE;
       }
 
-      DriverInstrumentation.DBInfo dbInfo = DriverInstrumentation.connectionInfo.get(connection);
+      ConnectionInstrumentation.DBInfo dbInfo =
+          ConnectionInstrumentation.connectionInfo.get(connection);
       if (dbInfo == null) {
-        dbInfo = DriverInstrumentation.DBInfo.UNKNOWN;
+        dbInfo = ConnectionInstrumentation.DBInfo.UNKNOWN;
       }
 
       final ActiveSpan span =

dd-java-agent/src/main/java/com/datadoghq/agent/instrumentation/jdbc/StatementInstrumentation.java

@@ -50,9 +50,10 @@ public final class StatementInstrumentation implements Instrumenter {
         return NoopActiveSpanSource.NoopActiveSpan.INSTANCE;
       }
 
-      DriverInstrumentation.DBInfo dbInfo = DriverInstrumentation.connectionInfo.get(connection);
+      ConnectionInstrumentation.DBInfo dbInfo =
+          ConnectionInstrumentation.connectionInfo.get(connection);
       if (dbInfo == null) {
-        dbInfo = DriverInstrumentation.DBInfo.UNKNOWN;
+        dbInfo = ConnectionInstrumentation.DBInfo.UNKNOWN;
       }
 
       final ActiveSpan span =