Add result set wrapping for jdbc library instrumentation (#13646)

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

@@ -12,6 +12,7 @@ import static io.opentelemetry.javaagent.instrumentation.jdbc.JdbcSingletons.sta
 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 class PreparedStatementInstrumentation implements TypeInstrumentation {
   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 io.opentelemetry.javaagent.instrumentation.jdbc.JdbcSingletons.sta
 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 class StatementInstrumentation implements TypeInstrumentation {
         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_SYST
 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 java.util.stream.Stream;
 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 @@ class JdbcInstrumentationTest {
       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 @@ class JdbcInstrumentationTest {
   @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 @@ class JdbcInstrumentationTest {
     statement.clearBatch();
     statement.addBatch("INSERT INTO " + tableName + " VALUES(1)");
     statement.addBatch("INSERT INTO " + tableName + " VALUES(2)");
-    testing.runWithSpan(
+    testing.runWithSpan("parent", () -> action.accept(statement));
-        "parent", () -> assertThat(statement.executeBatch()).isEqualTo(new int[] {1, 1}));
 
     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 final class OpenTelemetryDriver implements Driver {
 
     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 class OpenTelemetryDataSource implements DataSource, AutoCloseable {
   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.NClob;
 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;
 
-/**
+@SuppressWarnings("OverloadMethodsDeclarationOrder")
- * This class is internal and is hence not for public use. Its APIs are unstable and can change at
+class OpenTelemetryCallableStatement<S extends CallableStatement>
- * any time.
- */
-public class OpenTelemetryCallableStatement<S extends CallableStatement>
     extends OpenTelemetryPreparedStatement<S> implements CallableStatement {
 
   public OpenTelemetryCallableStatement(
@@ -489,6 +487,7 @@ public class OpenTelemetryCallableStatement<S extends CallableStatement>
     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 class OpenTelemetryCallableStatement<S extends CallableStatement>
   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);
+  }
 }

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

@@ -35,6 +35,7 @@ import java.sql.SQLException;
 import java.sql.SQLWarning;
 import java.sql.SQLXML;
 import java.sql.Savepoint;
+import java.sql.ShardingKey;
 import java.sql.Statement;
 import java.sql.Struct;
 import java.util.Map;
@@ -47,17 +48,36 @@ import java.util.concurrent.Executor;
  */
 public class OpenTelemetryConnection implements Connection {
 
-  private final Connection delegate;
+  private static final boolean hasJdbc43 = hasJdbc43();
+  protected final Connection delegate;
   private final DbInfo dbInfo;
   protected final Instrumenter<DbRequest, Void> statementInstrumenter;
 
-  public OpenTelemetryConnection(
+  protected OpenTelemetryConnection(
       Connection delegate, DbInfo dbInfo, Instrumenter<DbRequest, Void> statementInstrumenter) {
     this.delegate = delegate;
     this.dbInfo = dbInfo;
     this.statementInstrumenter = statementInstrumenter;
   }
 
+  // visible for testing
+  static boolean hasJdbc43() {
+    try {
+      Class.forName("java.sql.ShardingKey");
+      return true;
+    } catch (ClassNotFoundException exception) {
+      return false;
+    }
+  }
+
+  public static Connection create(
+      Connection delegate, DbInfo dbInfo, Instrumenter<DbRequest, Void> statementInstrumenter) {
+    if (hasJdbc43) {
+      return new OpenTelemetryConnectionJdbc43(delegate, dbInfo, statementInstrumenter);
+    }
+    return new OpenTelemetryConnection(delegate, dbInfo, statementInstrumenter);
+  }
+
   @Override
   public Statement createStatement() throws SQLException {
     Statement statement = delegate.createStatement();
@@ -369,4 +389,50 @@ public class OpenTelemetryConnection implements Connection {
   public DbInfo getDbInfo() {
     return dbInfo;
   }
+
+  // JDBC 4.3
+  static class OpenTelemetryConnectionJdbc43 extends OpenTelemetryConnection {
+    OpenTelemetryConnectionJdbc43(
+        Connection delegate, DbInfo dbInfo, Instrumenter<DbRequest, Void> statementInstrumenter) {
+      super(delegate, dbInfo, statementInstrumenter);
+    }
+
+    @SuppressWarnings("Since15")
+    @Override
+    public void beginRequest() throws SQLException {
+      delegate.beginRequest();
+    }
+
+    @SuppressWarnings("Since15")
+    @Override
+    public void endRequest() throws SQLException {
+      delegate.endRequest();
+    }
+
+    @SuppressWarnings("Since15")
+    @Override
+    public boolean setShardingKeyIfValid(
+        ShardingKey shardingKey, ShardingKey superShardingKey, int timeout) throws SQLException {
+      return delegate.setShardingKeyIfValid(shardingKey, superShardingKey, timeout);
+    }
+
+    @SuppressWarnings("Since15")
+    @Override
+    public boolean setShardingKeyIfValid(ShardingKey shardingKey, int timeout) throws SQLException {
+      return delegate.setShardingKeyIfValid(shardingKey, timeout);
+    }
+
+    @SuppressWarnings("Since15")
+    @Override
+    public void setShardingKey(ShardingKey shardingKey, ShardingKey superShardingKey)
+        throws SQLException {
+      delegate.setShardingKey(shardingKey, superShardingKey);
+    }
+
+    @SuppressWarnings("Since15")
+    @Override
+    public void setShardingKey(ShardingKey shardingKey) throws SQLException {
+      delegate.setShardingKey(shardingKey);
+    }
+  }
 }

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

@@ -38,17 +38,15 @@ import java.sql.ResultSet;
 import java.sql.ResultSetMetaData;
 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;
 
-/**
+@SuppressWarnings("OverloadMethodsDeclarationOrder")
- * This class is internal and is hence not for public use. Its APIs are unstable and can change at
+class OpenTelemetryPreparedStatement<S extends PreparedStatement> extends OpenTelemetryStatement<S>
- * any time.
+    implements PreparedStatement {
- */
-public class OpenTelemetryPreparedStatement<S extends PreparedStatement>
-    extends OpenTelemetryStatement<S> implements PreparedStatement {
 
   public OpenTelemetryPreparedStatement(
       S delegate,
@@ -61,7 +59,7 @@ public class OpenTelemetryPreparedStatement<S extends PreparedStatement>
 
   @Override
   public ResultSet executeQuery() throws SQLException {
-    return wrapCall(query, delegate::executeQuery);
+    return new OpenTelemetryResultSet(wrapCall(query, delegate::executeQuery), this);
   }
 
   @Override
@@ -374,4 +372,22 @@ public class OpenTelemetryPreparedStatement<S extends PreparedStatement>
     DbRequest request = DbRequest.create(dbInfo, query, batchSize);
     return wrapCall(request, callable);
   }
+
+  // JDBC 4.2
+
+  @Override
+  public void setObject(int parameterIndex, Object x, SQLType targetSqlType, int scaleOrLength)
+      throws SQLException {
+    delegate.setObject(parameterIndex, x, targetSqlType, scaleOrLength);
+  }
+
+  @Override
+  public void setObject(int parameterIndex, Object x, SQLType targetSqlType) throws SQLException {
+    delegate.setObject(parameterIndex, x, targetSqlType);
+  }
+
+  @Override
+  public long executeLargeUpdate() throws SQLException {
+    return wrapCall(query, delegate::executeLargeUpdate);
+  }
 }

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

@@ -32,11 +32,7 @@ import java.sql.Statement;
 import java.util.ArrayList;
 import java.util.List;
 
-/**
+class OpenTelemetryStatement<S extends Statement> implements Statement {
- * This class is internal and is hence not for public use. Its APIs are unstable and can change at
- * any time.
- */
-public class OpenTelemetryStatement<S extends Statement> implements Statement {
 
   protected final S delegate;
   protected final OpenTelemetryConnection connection;
@@ -180,7 +176,7 @@ public class OpenTelemetryStatement<S extends Statement> implements Statement {
 
   @Override
   public ResultSet getResultSet() throws SQLException {
-    return delegate.getResultSet();
+    return new OpenTelemetryResultSet(delegate.getResultSet(), this);
   }
 
   @Override
@@ -250,7 +246,7 @@ public class OpenTelemetryStatement<S extends Statement> implements Statement {
 
   @Override
   public ResultSet getGeneratedKeys() throws SQLException {
-    return delegate.getGeneratedKeys();
+    return new OpenTelemetryResultSet(delegate.getGeneratedKeys(), this);
   }
 
   @Override
@@ -293,6 +289,74 @@ public class OpenTelemetryStatement<S extends Statement> implements Statement {
     return delegate.isWrapperFor(iface);
   }
 
+  // JDBC 4.2
+
+  @Override
+  public long getLargeUpdateCount() throws SQLException {
+    return delegate.getLargeUpdateCount();
+  }
+
+  @Override
+  public void setLargeMaxRows(long max) throws SQLException {
+    delegate.setLargeMaxRows(max);
+  }
+
+  @Override
+  public long getLargeMaxRows() throws SQLException {
+    return delegate.getLargeMaxRows();
+  }
+
+  @Override
+  public long[] executeLargeBatch() throws SQLException {
+    return wrapBatchCall(delegate::executeLargeBatch);
+  }
+
+  @Override
+  public long executeLargeUpdate(String sql) throws SQLException {
+    return wrapCall(sql, () -> delegate.executeLargeUpdate(sql));
+  }
+
+  @Override
+  public long executeLargeUpdate(String sql, int autoGeneratedKeys) throws SQLException {
+    return wrapCall(sql, () -> delegate.executeLargeUpdate(sql, autoGeneratedKeys));
+  }
+
+  @Override
+  public long executeLargeUpdate(String sql, int[] columnIndexes) throws SQLException {
+    return wrapCall(sql, () -> delegate.executeLargeUpdate(sql, columnIndexes));
+  }
+
+  @Override
+  public long executeLargeUpdate(String sql, String[] columnNames) throws SQLException {
+    return wrapCall(sql, () -> delegate.executeLargeUpdate(sql, columnNames));
+  }
+
+  // JDBC 4.3
+
+  @SuppressWarnings("Since15")
+  @Override
+  public String enquoteLiteral(String val) throws SQLException {
+    return delegate.enquoteLiteral(val);
+  }
+
+  @SuppressWarnings("Since15")
+  @Override
+  public String enquoteIdentifier(String identifier, boolean alwaysQuote) throws SQLException {
+    return delegate.enquoteIdentifier(identifier, alwaysQuote);
+  }
+
+  @SuppressWarnings("Since15")
+  @Override
+  public boolean isSimpleIdentifier(String identifier) throws SQLException {
+    return delegate.isSimpleIdentifier(identifier);
+  }
+
+  @SuppressWarnings("Since15")
+  @Override
+  public String enquoteNCharLiteral(String val) throws SQLException {
+    return delegate.enquoteNCharLiteral(val);
+  }
+
   protected <T, E extends Exception> T wrapCall(String sql, ThrowingSupplier<T, E> callable)
       throws E {
     DbRequest request = DbRequest.create(dbInfo, sql);

instrumentation/jdbc/library/src/test/java/io/opentelemetry/instrumentation/jdbc/OpenTelemetryDriverTest.java

@@ -128,7 +128,7 @@ public class OpenTelemetryDriverTest {
     Connection connection = OpenTelemetryDriver.INSTANCE.connect("jdbc:otel:test:", null);
     OpenTelemetryDriver.removeDriverCandidate(driver);
 
-    assertThat(connection).isExactlyInstanceOf(OpenTelemetryConnection.class);
+    assertThat(connection).isInstanceOf(OpenTelemetryConnection.class);
   }
 
   @DisplayName("verify remove driver candidate")
@@ -220,7 +220,7 @@ public class OpenTelemetryDriverTest {
     Connection connection2 = OpenTelemetryDriver.INSTANCE.connect("jdbc:otel:test:", null);
 
     assertThat(connection2).isNotNull();
-    assertThat(connection2).isExactlyInstanceOf(OpenTelemetryConnection.class);
+    assertThat(connection2).isInstanceOf(OpenTelemetryConnection.class);
   }
 
   @DisplayName("Verify get property info with test driver url")

instrumentation/jdbc/library/src/test/java/io/opentelemetry/instrumentation/jdbc/datasource/OpenTelemetryDataSourceTest.java

@@ -67,7 +67,7 @@ class OpenTelemetryDataSourceTest {
                                     ? null
                                     : "postgresql://127.0.0.1:5432"))));
 
-    assertThat(connection).isExactlyInstanceOf(OpenTelemetryConnection.class);
+    assertThat(connection).isInstanceOf(OpenTelemetryConnection.class);
     DbInfo dbInfo = ((OpenTelemetryConnection) connection).getDbInfo();
     assertDbInfo(dbInfo);
   }
@@ -83,7 +83,7 @@ class OpenTelemetryDataSourceTest {
 
     assertThat(testing.waitForTraces(0)).isEmpty();
 
-    assertThat(connection).isExactlyInstanceOf(OpenTelemetryConnection.class);
+    assertThat(connection).isInstanceOf(OpenTelemetryConnection.class);
     DbInfo dbInfo = ((OpenTelemetryConnection) connection).getDbInfo();
     assertDbInfo(dbInfo);
   }

instrumentation/jdbc/library/src/test/java/io/opentelemetry/instrumentation/jdbc/internal/OpenTelemetryConnectionTest.java

@@ -30,6 +30,7 @@ import io.opentelemetry.instrumentation.jdbc.internal.dbinfo.DbInfo;
 import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
 import io.opentelemetry.instrumentation.testing.junit.LibraryInstrumentationExtension;
 import java.sql.PreparedStatement;
+import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
 import org.junit.jupiter.api.Test;
@@ -93,6 +94,33 @@ class OpenTelemetryConnectionTest {
         "parent",
         () -> {
           assertThat(statement.execute()).isTrue();
+          ResultSet resultSet = statement.getResultSet();
+          assertThat(resultSet).isInstanceOf(OpenTelemetryResultSet.class);
+          assertThat(resultSet.getStatement()).isEqualTo(statement);
+        });
+
+    jdbcTraceAssertion(dbInfo, query);
+
+    statement.close();
+    connection.close();
+  }
+
+  @Test
+  void testVerifyPrepareStatementQuery() throws SQLException {
+    Instrumenter<DbRequest, Void> instrumenter =
+        createStatementInstrumenter(testing.getOpenTelemetry());
+    DbInfo dbInfo = getDbInfo();
+    OpenTelemetryConnection connection =
+        new OpenTelemetryConnection(new TestConnection(), dbInfo, instrumenter);
+    String query = "SELECT * FROM users";
+    PreparedStatement statement = connection.prepareStatement(query);
+
+    testing.runWithSpan(
+        "parent",
+        () -> {
+          ResultSet resultSet = statement.executeQuery();
+          assertThat(resultSet).isInstanceOf(OpenTelemetryResultSet.class);
+          assertThat(resultSet.getStatement()).isEqualTo(statement);
         });
 
     jdbcTraceAssertion(dbInfo, query);

instrumentation/jdbc/library/src/test/java/io/opentelemetry/instrumentation/jdbc/internal/WrapperTest.java

@@ -0,0 +1,41 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.jdbc.internal;
+
+import io.opentelemetry.instrumentation.jdbc.internal.OpenTelemetryConnection.OpenTelemetryConnectionJdbc43;
+import java.lang.reflect.Method;
+import java.sql.CallableStatement;
+import java.sql.Connection;
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
+import java.sql.Statement;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+class WrapperTest {
+
+  @Test
+  void wrapperImplementsAllMethods() throws Exception {
+    validate(Statement.class, OpenTelemetryStatement.class);
+    validate(PreparedStatement.class, OpenTelemetryPreparedStatement.class);
+    validate(CallableStatement.class, OpenTelemetryCallableStatement.class);
+    validate(
+        Connection.class,
+        OpenTelemetryConnection.hasJdbc43()
+            ? OpenTelemetryConnectionJdbc43.class
+            : OpenTelemetryConnection.class);
+    validate(ResultSet.class, OpenTelemetryResultSet.class);
+  }
+
+  void validate(Class<?> jdbcClass, Class<?> wrapperClass) throws Exception {
+    for (Method method : jdbcClass.getMethods()) {
+      Method result = wrapperClass.getMethod(method.getName(), method.getParameterTypes());
+      if (!result.getDeclaringClass().getName().startsWith("io.opentelemetry")) {
+        Assertions.fail(wrapperClass.getName() + " does not override " + method);
+      }
+    }
+  }
+}