Add result set wrapping for jdbc library instrumentation

instrumentation/jdbc/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jdbc/PreparedStatementInstrumentation.java

@@ -12,6 +12,7 @@
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
 import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;
 import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.namedOneOf;
 import static net.bytebuddy.matcher.ElementMatchers.not;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments;
@@ -45,7 +46,7 @@ public ElementMatcher<TypeDescription> typeMatcher() {
   public void transform(TypeTransformer transformer) {
     transformer.applyAdviceToMethod(
         nameStartsWith("execute")
-            .and(not(named("executeBatch")))
+            .and(not(namedOneOf("executeBatch", "executeLargeBatch")))
             .and(takesArguments(0))
             .and(isPublic()),
         PreparedStatementInstrumentation.class.getName() + "$PreparedStatementAdvice");

instrumentation/jdbc/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jdbc/StatementInstrumentation.java

@@ -12,6 +12,7 @@
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
 import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;
 import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.namedOneOf;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments;
 
@@ -52,7 +53,7 @@ public void transform(TypeTransformer transformer) {
         named("clearBatch").and(isPublic()),
         StatementInstrumentation.class.getName() + "$ClearBatchAdvice");
     transformer.applyAdviceToMethod(
-        named("executeBatch").and(takesNoArguments()).and(isPublic()),
+        namedOneOf("executeBatch", "executeLargeBatch").and(takesNoArguments()).and(isPublic()),
         StatementInstrumentation.class.getName() + "$ExecuteBatchAdvice");
   }
 

instrumentation/jdbc/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/jdbc/test/JdbcInstrumentationTest.java

@@ -25,6 +25,7 @@
 import static io.opentelemetry.semconv.incubating.DbIncubatingAttributes.DB_SYSTEM_NAME;
 import static io.opentelemetry.semconv.incubating.DbIncubatingAttributes.DB_USER;
 import static java.util.Arrays.asList;
+import static org.junit.jupiter.api.Assumptions.assumeTrue;
 
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Maps;
@@ -61,6 +62,7 @@
 import javax.sql.DataSource;
 import org.apache.derby.jdbc.EmbeddedDataSource;
 import org.apache.derby.jdbc.EmbeddedDriver;
+import org.assertj.core.api.ThrowingConsumer;
 import org.h2.jdbcx.JdbcDataSource;
 import org.hsqldb.jdbc.JDBCDriver;
 import org.junit.jupiter.api.AfterAll;
@@ -803,10 +805,81 @@ void testPreparedStatementUpdate(
       String url,
       String table)
       throws SQLException {
+    testPreparedStatementUpdateImpl(
+        system,
+        connection,
+        username,
+        query,
+        spanName,
+        url,
+        table,
+        statement -> assertThat(statement.executeUpdate()).isEqualTo(0));
+  }
+
+  static Stream<Arguments> preparedStatementLargeUpdateStream() throws SQLException {
+    return Stream.of(
+        Arguments.of(
+            "h2",
+            new org.h2.Driver().connect(jdbcUrls.get("h2"), null),
+            null,
+            "CREATE TABLE PS_LARGE_H2 (id INTEGER not NULL, PRIMARY KEY ( id ))",
+            "CREATE TABLE jdbcunittest.PS_LARGE_H2",
+            "h2:mem:",
+            "PS_LARGE_H2"),
+        Arguments.of(
+            "h2",
+            cpDatasources.get("tomcat").get("h2").getConnection(),
+            null,
+            "CREATE TABLE PS_LARGE_H2_TOMCAT (id INTEGER not NULL, PRIMARY KEY ( id ))",
+            "CREATE TABLE jdbcunittest.PS_LARGE_H2_TOMCAT",
+            "h2:mem:",
+            "PS_LARGE_H2_TOMCAT"),
+        Arguments.of(
+            "h2",
+            cpDatasources.get("hikari").get("h2").getConnection(),
+            null,
+            "CREATE TABLE PS_LARGE_H2_HIKARI (id INTEGER not NULL, PRIMARY KEY ( id ))",
+            "CREATE TABLE jdbcunittest.PS_LARGE_H2_HIKARI",
+            "h2:mem:",
+            "PS_LARGE_H2_HIKARI"));
+  }
+
+  @ParameterizedTest
+  @MethodSource("preparedStatementLargeUpdateStream")
+  void testPreparedStatementLargeUpdate(
+      String system,
+      Connection connection,
+      String username,
+      String query,
+      String spanName,
+      String url,
+      String table)
+      throws SQLException {
+    testPreparedStatementUpdateImpl(
+        system,
+        connection,
+        username,
+        query,
+        spanName,
+        url,
+        table,
+        statement -> assertThat(statement.executeLargeUpdate()).isEqualTo(0));
+  }
+
+  void testPreparedStatementUpdateImpl(
+      String system,
+      Connection connection,
+      String username,
+      String query,
+      String spanName,
+      String url,
+      String table,
+      ThrowingConsumer<PreparedStatement> action)
+      throws SQLException {
     String sql = connection.nativeSQL(query);
     PreparedStatement statement = connection.prepareStatement(sql);
     cleanup.deferCleanup(statement);
-    testing.runWithSpan("parent", () -> assertThat(statement.executeUpdate()).isEqualTo(0));
+    testing.runWithSpan("parent", () -> action.accept(statement));
 
     testing.waitAndAssertTraces(
         trace ->
@@ -1333,7 +1406,39 @@ static Stream<Arguments> batchStream() throws SQLException {
   @MethodSource("batchStream")
   void testBatch(String system, Connection connection, String username, String url)
       throws SQLException {
-    String tableName = "simple_batch_test";
+    testBatchImpl(
+        system,
+        connection,
+        username,
+        url,
+        "simple_batch_test",
+        statement -> assertThat(statement.executeBatch()).isEqualTo(new int[] {1, 1}));
+  }
+
+  @ParameterizedTest
+  @MethodSource("batchStream")
+  void testLargeBatch(String system, Connection connection, String username, String url)
+      throws SQLException {
+    // derby and hsqldb used in this test don't support executeLargeBatch
+    assumeTrue("h2".equals(system));
+
+    testBatchImpl(
+        system,
+        connection,
+        username,
+        url,
+        "simple_batch_test_large",
+        statement -> assertThat(statement.executeLargeBatch()).isEqualTo(new long[] {1, 1}));
+  }
+
+  private static void testBatchImpl(
+      String system,
+      Connection connection,
+      String username,
+      String url,
+      String tableName,
+      ThrowingConsumer<Statement> action)
+      throws SQLException {
     Statement createTable = connection.createStatement();
     createTable.execute("CREATE TABLE " + tableName + " (id INTEGER not NULL, PRIMARY KEY ( id ))");
     cleanup.deferCleanup(createTable);
@@ -1347,8 +1452,7 @@ void testBatch(String system, Connection connection, String username, String url
     statement.clearBatch();
     statement.addBatch("INSERT INTO " + tableName + " VALUES(1)");
     statement.addBatch("INSERT INTO " + tableName + " VALUES(2)");
-    testing.runWithSpan(
-        "parent", () -> assertThat(statement.executeBatch()).isEqualTo(new int[] {1, 1}));
+    testing.runWithSpan("parent", () -> action.accept(statement));
 
     testing.waitAndAssertTraces(
         trace ->

instrumentation/jdbc/library/build.gradle.kts

@@ -16,6 +16,15 @@ dependencies {
 }
 
 tasks {
+  // We cannot use "--release" javac option here because that will forbid using apis that were added
+  // in later versions. In JDBC wrappers we wish to implement delegation for methods that are not
+  // present in jdk8.
+  compileJava {
+    sourceCompatibility = "1.8"
+    targetCompatibility = "1.8"
+    options.release.set(null as Int?)
+  }
+
   shadowJar {
     dependencies {
       // including only current module excludes its transitive dependencies

instrumentation/jdbc/library/src/main/java/io/opentelemetry/instrumentation/jdbc/OpenTelemetryDriver.java

@@ -244,7 +244,7 @@ public Connection connect(String url, Properties info) throws SQLException {
 
     Instrumenter<DbRequest, Void> statementInstrumenter =
         JdbcInstrumenterFactory.createStatementInstrumenter(openTelemetry);
-    return new OpenTelemetryConnection(connection, dbInfo, statementInstrumenter);
+    return OpenTelemetryConnection.create(connection, dbInfo, statementInstrumenter);
   }
 
   @Override

instrumentation/jdbc/library/src/main/java/io/opentelemetry/instrumentation/jdbc/datasource/OpenTelemetryDataSource.java

@@ -93,14 +93,14 @@ public OpenTelemetryDataSource(DataSource delegate, OpenTelemetry openTelemetry)
   public Connection getConnection() throws SQLException {
     Connection connection = wrapCall(delegate::getConnection);
     DbInfo dbInfo = getDbInfo(connection);
-    return new OpenTelemetryConnection(connection, dbInfo, statementInstrumenter);
+    return OpenTelemetryConnection.create(connection, dbInfo, statementInstrumenter);
   }
 
   @Override
   public Connection getConnection(String username, String password) throws SQLException {
     Connection connection = wrapCall(() -> delegate.getConnection(username, password));
     DbInfo dbInfo = getDbInfo(connection);
-    return new OpenTelemetryConnection(connection, dbInfo, statementInstrumenter);
+    return OpenTelemetryConnection.create(connection, dbInfo, statementInstrumenter);
   }
 
   @Override

instrumentation/jdbc/library/src/main/java/io/opentelemetry/instrumentation/jdbc/internal/OpenTelemetryCallableStatement.java

@@ -35,17 +35,15 @@
 import java.sql.Ref;
 import java.sql.RowId;
 import java.sql.SQLException;
+import java.sql.SQLType;
 import java.sql.SQLXML;
 import java.sql.Time;
 import java.sql.Timestamp;
 import java.util.Calendar;
 import java.util.Map;
 
-/**
- * This class is internal and is hence not for public use. Its APIs are unstable and can change at
- * any time.
- */
-public class OpenTelemetryCallableStatement<S extends CallableStatement>
+@SuppressWarnings("OverloadMethodsDeclarationOrder")
+class OpenTelemetryCallableStatement<S extends CallableStatement>
     extends OpenTelemetryPreparedStatement<S> implements CallableStatement {
 
   public OpenTelemetryCallableStatement(
@@ -489,6 +487,7 @@ public void setBinaryStream(String parameterName, InputStream x) throws SQLExcep
     delegate.setBinaryStream(parameterName, x);
   }
 
+  @SuppressWarnings("UngroupedOverloads")
   @Override
   public void setObject(String parameterName, Object x, int targetSqlType, int scale)
       throws SQLException {
@@ -710,4 +709,51 @@ public Reader getCharacterStream(int parameterIndex) throws SQLException {
   public Reader getCharacterStream(String parameterName) throws SQLException {
     return delegate.getCharacterStream(parameterName);
   }
+
+  // JDBC 4.2
+
+  @Override
+  public void setObject(String parameterName, Object x, SQLType targetSqlType, int scaleOrLength)
+      throws SQLException {
+    delegate.setObject(parameterName, x, targetSqlType, scaleOrLength);
+  }
+
+  @Override
+  public void setObject(String parameterName, Object x, SQLType targetSqlType) throws SQLException {
+    delegate.setObject(parameterName, x, targetSqlType);
+  }
+
+  @Override
+  public void registerOutParameter(int parameterIndex, SQLType sqlType) throws SQLException {
+    delegate.registerOutParameter(parameterIndex, sqlType);
+  }
+
+  @Override
+  public void registerOutParameter(int parameterIndex, SQLType sqlType, int scale)
+      throws SQLException {
+    delegate.registerOutParameter(parameterIndex, sqlType, scale);
+  }
+
+  @Override
+  public void registerOutParameter(int parameterIndex, SQLType sqlType, String typeName)
+      throws SQLException {
+    delegate.registerOutParameter(parameterIndex, sqlType, typeName);
+  }
+
+  @Override
+  public void registerOutParameter(String parameterName, SQLType sqlType) throws SQLException {
+    delegate.registerOutParameter(parameterName, sqlType);
+  }
+
+  @Override
+  public void registerOutParameter(String parameterName, SQLType sqlType, int scale)
+      throws SQLException {
+    delegate.registerOutParameter(parameterName, sqlType, scale);
+  }
+
+  @Override
+  public void registerOutParameter(String parameterName, SQLType sqlType, String typeName)
+      throws SQLException {
+    delegate.registerOutParameter(parameterName, sqlType, typeName);
+  }
 }