Preserve JDBC instrumentation when unwrapping standard interfaces (#20947)

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: trask <218610+trask@users.noreply.github.com>

Author: Copilot

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

@@ -454,12 +454,15 @@ public int getNetworkTimeout() throws SQLException {
 
   @Override
   public <T> T unwrap(Class<T> iface) throws SQLException {
+    if (iface.isInstance(this)) {
+      return iface.cast(this);
+    }
     return delegate.unwrap(iface);
   }
 
   @Override
   public boolean isWrapperFor(Class<?> iface) throws SQLException {
-    return delegate.isWrapperFor(iface);
+    return iface.isInstance(this) || delegate.isWrapperFor(iface);
   }
 
   // visible for testing

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

@@ -1030,11 +1030,14 @@ public void updateObject(String columnLabel, Object x, SQLType targetSqlType)
 
   @Override
   public <T> T unwrap(Class<T> iface) throws SQLException {
+    if (iface.isInstance(this)) {
+      return iface.cast(this);
+    }
     return delegate.unwrap(iface);
   }
 
   @Override
   public boolean isWrapperFor(Class<?> iface) throws SQLException {
-    return delegate.isWrapperFor(iface);
+    return iface.isInstance(this) || delegate.isWrapperFor(iface);
   }
 }

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

@@ -307,12 +307,15 @@ public boolean isCloseOnCompletion() throws SQLException {
 
   @Override
   public <T> T unwrap(Class<T> iface) throws SQLException {
+    if (iface.isInstance(this)) {
+      return iface.cast(this);
+    }
     return delegate.unwrap(iface);
   }
 
   @Override
   public boolean isWrapperFor(Class<?> iface) throws SQLException {
-    return delegate.isWrapperFor(iface);
+    return iface.isInstance(this) || delegate.isWrapperFor(iface);
   }
 
   // JDBC 4.2

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

@@ -7,10 +7,15 @@
 
 import static io.opentelemetry.instrumentation.jdbc.internal.JdbcInstrumenterFactory.createStatementInstrumenter;
 import static io.opentelemetry.instrumentation.jdbc.internal.JdbcInstrumenterFactory.createTransactionInstrumenter;
+import static java.util.stream.Collectors.toList;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.assertj.core.api.Assertions.fail;
 import static org.junit.jupiter.api.Assumptions.assumeTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoInteractions;
+import static org.mockito.Mockito.when;
 
 import io.opentelemetry.api.OpenTelemetry;
 import io.opentelemetry.instrumentation.api.incubator.semconv.db.internal.SqlCommenter;
@@ -26,9 +31,14 @@
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
+import java.sql.Wrapper;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.stream.Stream;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
 
 class WrapperTest {
 
@@ -112,6 +122,102 @@ void jdbc45MethodsPropagateSqlException() throws Exception {
         .isSameAs(expected);
   }
 
+  @ParameterizedTest
+  @MethodSource("wrappers")
+  void supportedInterfacesReturnWrapperWithoutConsultingDelegate(Wrapper wrapper, Wrapper delegate)
+      throws SQLException {
+    for (Class<?> iface : supportedInterfaces(wrapper)) {
+      assertThat(wrapper.unwrap(iface)).isSameAs(wrapper);
+      assertThat(wrapper.isWrapperFor(iface)).isTrue();
+    }
+    verifyNoInteractions(delegate);
+  }
+
+  @ParameterizedTest
+  @MethodSource("wrappers")
+  void delegatesVendorInterfaces(Wrapper wrapper, Wrapper delegate) throws SQLException {
+    VendorJdbcObject vendorObject = mock(VendorJdbcObject.class);
+    when(delegate.unwrap(VendorJdbcObject.class)).thenReturn(vendorObject);
+    when(delegate.isWrapperFor(VendorJdbcObject.class)).thenReturn(true);
+
+    assertThat(wrapper.unwrap(VendorJdbcObject.class)).isSameAs(vendorObject);
+    assertThat(wrapper.isWrapperFor(VendorJdbcObject.class)).isTrue();
+    verify(delegate).unwrap(VendorJdbcObject.class);
+    verify(delegate).isWrapperFor(VendorJdbcObject.class);
+  }
+
+  @ParameterizedTest
+  @MethodSource("wrappers")
+  void preservesUnsupportedInterfaceBehavior(Wrapper wrapper, Wrapper delegate)
+      throws SQLException {
+    SQLException exception = new SQLException("Unsupported interface");
+    when(delegate.unwrap(Runnable.class)).thenThrow(exception);
+    when(delegate.isWrapperFor(Runnable.class)).thenReturn(false);
+
+    assertThatThrownBy(() -> wrapper.unwrap(Runnable.class)).isSameAs(exception);
+    assertThat(wrapper.isWrapperFor(Runnable.class)).isFalse();
+    verify(delegate).unwrap(Runnable.class);
+    verify(delegate).isWrapperFor(Runnable.class);
+
+    when(delegate.isWrapperFor(Runnable.class)).thenThrow(exception);
+    assertThatThrownBy(() -> wrapper.isWrapperFor(Runnable.class)).isSameAs(exception);
+  }
+
+  private static List<Class<?>> supportedInterfaces(Wrapper wrapper) {
+    return Stream.of(
+            Connection.class,
+            Statement.class,
+            PreparedStatement.class,
+            CallableStatement.class,
+            ResultSet.class,
+            Wrapper.class,
+            AutoCloseable.class)
+        .filter(iface -> iface.isInstance(wrapper))
+        .collect(toList());
+  }
+
+  private static Stream<Arguments> wrappers() {
+    Connection connectionDelegate = mock(Connection.class);
+    OpenTelemetryConnection connection = (OpenTelemetryConnection) wrap(connectionDelegate);
+    Statement statementDelegate = mock(Statement.class);
+    PreparedStatement preparedDelegate = mock(PreparedStatement.class);
+    CallableStatement callableDelegate = mock(CallableStatement.class);
+    ResultSet resultSetDelegate = mock(ResultSet.class);
+    DbInfo dbInfo = DbInfo.builder().build();
+    Statement statement =
+        new OpenTelemetryStatement<>(
+            statementDelegate,
+            connection,
+            dbInfo,
+            connection.statementInstrumenter,
+            SqlCommenter.noop());
+
+    return Stream.of(
+        Arguments.of(connection, connectionDelegate),
+        Arguments.of(statement, statementDelegate),
+        Arguments.of(
+            new OpenTelemetryPreparedStatement<>(
+                preparedDelegate,
+                connection,
+                dbInfo,
+                "SELECT 1",
+                connection.statementInstrumenter,
+                false,
+                SqlCommenter.noop()),
+            preparedDelegate),
+        Arguments.of(
+            new OpenTelemetryCallableStatement<>(
+                callableDelegate,
+                connection,
+                dbInfo,
+                "CALL test()",
+                connection.statementInstrumenter,
+                false,
+                SqlCommenter.noop()),
+            callableDelegate),
+        Arguments.of(OpenTelemetryResultSet.wrap(resultSetDelegate, statement), resultSetDelegate));
+  }
+
   private static Connection wrap(Connection delegate) {
     OpenTelemetry noop = OpenTelemetry.noop();
     return OpenTelemetryConnection.create(
@@ -132,6 +238,8 @@ private static Object invokeConnectionMethod(Connection wrapper, String name, St
     return Connection.class.getMethod(name, String.class).invoke(wrapper, arg);
   }
 
+  private interface VendorJdbcObject {}
+
   void validate(Class<?> jdbcClass, Class<?> wrapperClass) throws NoSuchMethodException {
     for (Method method : jdbcClass.getMethods()) {
       Method result = wrapperClass.getMethod(method.getName(), method.getParameterTypes());