Support JDBC spans for Druid-wrapped Oracle statements

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

@@ -34,6 +34,7 @@ dependencies {
   testImplementation group: 'org.apache.tomcat', name: 'tomcat-juli', version: '7.0.19'
   testImplementation group: 'com.zaxxer', name: 'HikariCP', version: '2.4.0'
   testImplementation group: 'com.mchange', name: 'c3p0', version: '0.9.5'
+  testImplementation group: 'com.alibaba', name: 'druid', version: '1.2.23'
 
   testImplementation group: 'mysql', name: 'mysql-connector-java', version: '8.0.23'
   testImplementation group: 'org.postgresql', name: 'postgresql', version: '[9.4,42.2.18]'

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

@@ -112,6 +112,10 @@ public static class PreparedStatementAdvice {
 
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static AgentScope onEnter(@Advice.This final Statement statement) {
+      final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(Statement.class);
+      if (callDepth > 0) {
+        return null;
+      }
       try {
         final Connection connection = statement.getConnection();
         final DBQueryInfo queryInfo =

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

@@ -78,6 +78,8 @@ public DBMCompatibleConnectionInstrumentation() {
     "software.aws.rds.jdbc.mysql.shading.com.mysql.cj.jdbc.ConnectionImpl",
     // IBM Informix
     "com.informix.jdbc.IfmxConnection",
+    // Druid connection wrapper
+    "com.alibaba.druid.pool.DruidPooledConnection",
     // 达梦 DB
     "dm.jdbc.driver.DmdbConnection",
     // kingbase

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

@@ -48,6 +48,8 @@ public class DefaultConnectionInstrumentation extends AbstractConnectionInstrume
     // Sybase
     "com.sybase.jdbc2.jdbc.SybConnection",
     "com.sybase.jdbc4.jdbc.SybConnection",
+    // Druid connection wrapper
+    "com.alibaba.druid.pool.DruidPooledConnection",
     // for testing purposes
     "test.TestConnection",
     // 达梦db

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

@@ -127,6 +127,9 @@ public final class PreparedStatementInstrumentation extends AbstractPreparedStat
     "com.sybase.jdbc2.jdbc.SybCallableStatement",
     "com.sybase.jdbc4.jdbc.SybPreparedStatement",
     "com.sybase.jdbc4.jdbc.SybCallableStatement",
+    // Druid statement wrappers
+    "com.alibaba.druid.pool.DruidPooledPreparedStatement",
+    "com.alibaba.druid.pool.DruidPooledCallableStatement",
     // for testing purposes
     "test.TestPreparedStatement",
     // 达梦DB

dd-java-agent/instrumentation/jdbc/src/test/groovy/DruidPreparedStatementTest.groovy

@@ -0,0 +1,69 @@
+import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
+
+import com.alibaba.druid.pool.DruidDataSource
+import datadog.trace.agent.test.InstrumentationSpecification
+import datadog.trace.api.DDSpanTypes
+import datadog.trace.bootstrap.instrumentation.api.Tags
+
+class DruidPreparedStatementTest extends InstrumentationSpecification {
+
+  def "prepared statement on druid wrapper generates span"() {
+    setup:
+    def ds = new DruidDataSource()
+    ds.setUrl("jdbc:h2:mem:druid_wrapper_test;DB_CLOSE_DELAY=-1")
+    ds.setDriverClassName("org.h2.Driver")
+    ds.setInitialSize(1)
+    ds.setMaxActive(1)
+
+    def initConnection = ds.getConnection()
+    def initStatement = initConnection.createStatement()
+    initStatement.execute("CREATE TABLE IF NOT EXISTS T_TEST (ID INT PRIMARY KEY, NAME VARCHAR(32))")
+    initStatement.execute("MERGE INTO T_TEST KEY(ID) VALUES (1, 'alice')")
+    initStatement.close()
+    initConnection.close()
+
+    TEST_WRITER.waitForTraces(1)
+    TEST_WRITER.clear()
+
+    when:
+    runUnderTrace("parent") {
+      def connection = ds.getConnection()
+      def statement = connection.prepareStatement("SELECT NAME FROM T_TEST WHERE ID = ?")
+      statement.setInt(1, 1)
+      def resultSet = statement.executeQuery()
+      assert resultSet.next()
+      assert resultSet.getString(1) == "alice"
+      resultSet.close()
+      statement.close()
+      connection.close()
+    }
+
+    then:
+    assertTraces(1) {
+      trace(2) {
+        span {
+          operationName "parent"
+        }
+        span {
+          operationName "h2.query"
+          serviceName "h2"
+          resourceName "SELECT NAME FROM T_TEST WHERE ID = ?"
+          spanType DDSpanTypes.SQL
+          childOfPrevious()
+          errored false
+          measured true
+          tags(false) {
+            "$Tags.COMPONENT" "java-jdbc-prepared_statement"
+            "$Tags.SPAN_KIND" Tags.SPAN_KIND_CLIENT
+            "$Tags.DB_TYPE" "h2"
+            "$Tags.DB_INSTANCE" "druid_wrapper_test"
+            "$Tags.DB_OPERATION" "SELECT"
+          }
+        }
+      }
+    }
+
+    cleanup:
+    ds?.close()
+  }
+}