Unwrapped java.sql PreparedStatement and Connection

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/matcher/AdditionalLibraryIgnoresMatcher.java

@@ -270,6 +270,7 @@ public boolean matches(final T target) {
           || name.startsWith("org.h2.jdbcx.")
           // Some runnables that get instrumented
           || name.equals("org.h2.util.Task")
+          || name.equals("org.h2.util.MathUtils$1")
           || name.equals("org.h2.store.FileLock")
           || name.equals("org.h2.engine.DatabaseCloser")
           || name.equals("org.h2.engine.OnExitDatabaseCloser")) {

dd-java-agent/instrumentation/jdbc/jdbc.gradle

@@ -9,6 +9,10 @@ apply from: "$rootDir/gradle/java.gradle"
 apply plugin: 'org.unbroken-dome.test-sets'
 
 testSets {
+  oldH2Test {
+    dirName = 'test'
+  }
+
   latestDepTest {
     dirName = 'test'
   }
@@ -26,6 +30,11 @@ dependencies {
   testCompile group: 'com.zaxxer', name: 'HikariCP', version: '2.4.0'
   testCompile group: 'com.mchange', name: 'c3p0', version: '0.9.5'
 
+  // Test pre jdk 1.6 H2
+  oldH2TestCompile(group: 'com.h2database', name: 'h2', version: '1.3.168') {
+    force = true
+  }
+
   latestDepTestCompile group: 'com.h2database', name: 'h2', version: '+'
   latestDepTestCompile group: 'org.apache.derby', name: 'derby', version: '10.14.+'
   latestDepTestCompile group: 'org.hsqldb', name: 'hsqldb', version: '+'
@@ -35,3 +44,5 @@ dependencies {
   latestDepTestCompile group: 'com.zaxxer', name: 'HikariCP', version: '+'
   latestDepTestCompile group: 'com.mchange', name: 'c3p0', version: '+'
 }
+
+test.dependsOn oldH2Test

dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/JDBCUtils.java

@@ -2,12 +2,54 @@
 
 import datadog.trace.bootstrap.ExceptionLogger;
 import java.lang.reflect.Field;
+import java.lang.reflect.Modifier;
 import java.sql.Connection;
+import java.sql.PreparedStatement;
 import java.sql.Statement;
 
 public abstract class JDBCUtils {
   private static Field c3poField = null;
 
+  // Cache if the class's isWrapperFor() or unwrap() methods are usable
+  private static final ClassValue<Boolean> CAN_UNWRAP =
+      new ClassValue<Boolean>() {
+        @Override
+        protected Boolean computeValue(Class<?> type) {
+          // DB drivers compiled before JDK 1.6 do not implement java.sql.Wrapper methods and throw
+          // AbstractMethodError when called
+          // h2 before v1.3.168 and jdts.jdbc up to version 1.3 are in this category
+          // Additionally h2, instead of returning "false" in versions 1.3.169+, throws an exception
+          // so those classes are hardcoded below
+          try {
+            return !"org.h2.jdbc.JdbcConnection".equals(type.getName())
+                && !"org.h2.jdbcx.JdbcXAConnection.PooledJdbcConnection".equals(type.getName())
+                && !"org.h2.jdbc.JdbcPreparedStatement".equals(type.getName())
+                && !"org.h2.jdbc.JdbcCallableStatement".equals(type.getName())
+                && !Modifier.isAbstract(type.getMethod("isWrapperFor", Class.class).getModifiers())
+                && !Modifier.isAbstract(type.getMethod("unwrap", Class.class).getModifiers());
+          } catch (NoSuchMethodException ignored) {
+          }
+          return false;
+        }
+      };
+
+  public static PreparedStatement unwrappedStatement(PreparedStatement statement) {
+    if (CAN_UNWRAP.get(statement.getClass())) {
+      try {
+        // technically this could be recursive. In practice, one level is enough
+        // Recursive would require cycle checking
+        if (statement.isWrapperFor(PreparedStatement.class)) {
+          return statement.unwrap(PreparedStatement.class);
+        }
+      } catch (Exception ignored) {
+        // note that we will not even call this method unless it has been shown not to
+        // throw AbstractMethodError
+      }
+    }
+
+    return statement;
+  }
+
   /**
    * @param statement
    * @return the unwrapped connection or null if exception was thrown.
@@ -23,35 +65,41 @@ public static Connection connectionFromStatement(final Statement statement) {
         }
       }
 
-      try {
-        // unwrap the connection to cache the underlying actual connection and to not cache proxy
-        // objects
-        if (connection.isWrapperFor(Connection.class)) {
-          connection = connection.unwrap(Connection.class);
-        }
-      } catch (final Exception | AbstractMethodError e) {
-        if (connection != null) {
-          // Attempt to work around c3po delegating to an connection that doesn't support
-          // unwrapping.
-          final Class<? extends Connection> connectionClass = connection.getClass();
-          if (connectionClass.getName().equals("com.mchange.v2.c3p0.impl.NewProxyConnection")) {
-            final Field inner = connectionClass.getDeclaredField("inner");
-            inner.setAccessible(true);
-            c3poField = inner;
-            return (Connection) c3poField.get(connection);
-          }
-        }
-
-        // perhaps wrapping isn't supported?
-        // ex: org.h2.jdbc.JdbcConnection v1.3.175
-        // or: jdts.jdbc which always throws `AbstractMethodError` (at least up to version 1.3)
-        // Stick with original connection.
+      if (CAN_UNWRAP.get(connection.getClass())) {
+        return tryUnwrap(connection);
+      } else {
+        return connection;
       }
     } catch (final Throwable e) {
       // Had some problem getting the connection.
       ExceptionLogger.LOGGER.debug("Could not get connection for StatementAdvice", e);
       return null;
     }
+  }
+
+  private static Connection tryUnwrap(Connection connection) throws ReflectiveOperationException {
+    try {
+      // technically this could be recursive. In practice, one level is enough
+      if (connection.isWrapperFor(Connection.class)) {
+        return connection.unwrap(Connection.class);
+      } else {
+        return connection;
+      }
+    } catch (Exception ignored) {
+      // note that we will not even call this method unless it has been shown not to
+      // throw AbstractMethodError
+    }
+
+    // Attempt to work around c3po delegating to an connection that doesn't support
+    // unwrapping.
+    final Class<? extends Connection> connectionClass = connection.getClass();
+    if (connectionClass.getName().equals("com.mchange.v2.c3p0.impl.NewProxyConnection")) {
+      final Field inner = connectionClass.getDeclaredField("inner");
+      inner.setAccessible(true);
+      c3poField = inner;
+      return (Connection) c3poField.get(connection);
+    }
+
     return connection;
   }
 }

dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/PreparedStatementInstrumentation.java

@@ -7,6 +7,7 @@
 import static datadog.trace.instrumentation.jdbc.JDBCDecorator.DATABASE_QUERY;
 import static datadog.trace.instrumentation.jdbc.JDBCDecorator.DECORATE;
 import static datadog.trace.instrumentation.jdbc.JDBCUtils.connectionFromStatement;
+import static datadog.trace.instrumentation.jdbc.JDBCUtils.unwrappedStatement;
 import static java.util.Collections.singletonMap;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
 import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;
@@ -58,7 +59,7 @@ public ElementMatcher<TypeDescription> typeMatcher() {
   @Override
   public String[] helperClassNames() {
     return new String[] {
-      packageName + ".JDBCUtils", packageName + ".JDBCDecorator",
+      packageName + ".JDBCUtils", packageName + ".JDBCUtils$1", packageName + ".JDBCDecorator",
     };
   }
 
@@ -73,7 +74,9 @@ public static class PreparedStatementAdvice {
 
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static AgentScope onEnter(@Advice.This final PreparedStatement statement) {
-      final Connection connection = connectionFromStatement(statement);
+      PreparedStatement actualStatement = unwrappedStatement(statement);
+
+      final Connection connection = connectionFromStatement(actualStatement);
       if (connection == null) {
         return null;
       }
@@ -84,7 +87,8 @@ public static AgentScope onEnter(@Advice.This final PreparedStatement statement)
       }
 
       UTF8BytesString sql =
-          InstrumentationContext.get(PreparedStatement.class, UTF8BytesString.class).get(statement);
+          InstrumentationContext.get(PreparedStatement.class, UTF8BytesString.class)
+              .get(actualStatement);
 
       final AgentSpan span = startSpan(DATABASE_QUERY);
       DECORATE.afterStart(span);

dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java

@@ -46,7 +46,7 @@ public ElementMatcher<TypeDescription> typeMatcher() {
   @Override
   public String[] helperClassNames() {
     return new String[] {
-      packageName + ".JDBCUtils", packageName + ".JDBCDecorator",
+      packageName + ".JDBCUtils", packageName + ".JDBCUtils$1", packageName + ".JDBCDecorator",
     };
   }