Use AOP proxy for DataSource telemetry to preserve bean type

instrumentation/spring/spring-boot-autoconfigure/src/main/java/io/opentelemetry/instrumentation/spring/autoconfigure/internal/instrumentation/jdbc/DataSourcePostProcessor.java

@@ -11,7 +11,10 @@
 import io.opentelemetry.instrumentation.spring.autoconfigure.internal.properties.InstrumentationConfigUtil;
 import io.opentelemetry.sdk.autoconfigure.spi.ConfigProperties;
 import javax.sql.DataSource;
+import org.aopalliance.intercept.MethodInterceptor;
+import org.springframework.aop.framework.ProxyFactory;
 import org.springframework.aop.scope.ScopedProxyUtils;
+import org.springframework.aop.support.AopUtils;
 import org.springframework.beans.factory.ObjectProvider;
 import org.springframework.beans.factory.config.BeanPostProcessor;
 import org.springframework.core.Ordered;
@@ -50,22 +53,32 @@ public Object postProcessAfterInitialization(Object bean, String beanName) {
         && !isRoutingDatasource(bean)
         && !ScopedProxyUtils.isScopedTarget(beanName)) {
       DataSource dataSource = (DataSource) bean;
-      return JdbcTelemetry.builder(openTelemetryProvider.getObject())
-          .setStatementSanitizationEnabled(
-              InstrumentationConfigUtil.isStatementSanitizationEnabled(
-                  configPropertiesProvider.getObject(),
-                  "otel.instrumentation.jdbc.statement-sanitizer.enabled"))
-          .setCaptureQueryParameters(
-              configPropertiesProvider
-                  .getObject()
-                  .getBoolean(
-                      "otel.instrumentation.jdbc.experimental.capture-query-parameters", false))
-          .setTransactionInstrumenterEnabled(
-              configPropertiesProvider
-                  .getObject()
-                  .getBoolean("otel.instrumentation.jdbc.experimental.transaction.enabled", false))
-          .build()
-          .wrap(dataSource);
+      DataSource wrapped =
+          JdbcTelemetry.builder(openTelemetryProvider.getObject())
+              .setStatementSanitizationEnabled(
+                  InstrumentationConfigUtil.isStatementSanitizationEnabled(
+                      configPropertiesProvider.getObject(),
+                      "otel.instrumentation.jdbc.statement-sanitizer.enabled"))
+              .setCaptureQueryParameters(
+                  configPropertiesProvider
+                      .getObject()
+                      .getBoolean(
+                          "otel.instrumentation.jdbc.experimental.capture-query-parameters", false))
+              .setTransactionInstrumenterEnabled(
+                  configPropertiesProvider
+                      .getObject()
+                      .getBoolean(
+                          "otel.instrumentation.jdbc.experimental.transaction.enabled", false))
+              .build()
+              .wrap(dataSource);
+      ProxyFactory proxyFactory = new ProxyFactory(new Class<?>[] {DataSource.class});
+      proxyFactory.setTarget(bean);
+      proxyFactory.addAdvice(
+          (MethodInterceptor)
+              invocation ->
+                  AopUtils.invokeJoinpointUsingReflection(
+                      wrapped, invocation.getMethod(), invocation.getArguments()));
+      return proxyFactory.getProxy();
     }
     return bean;
   }

instrumentation/spring/spring-boot-autoconfigure/src/test/java/io/opentelemetry/instrumentation/spring/autoconfigure/internal/instrumentation/jdbc/DataSourcePostProcessorTest.java

@@ -0,0 +1,148 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.spring.autoconfigure.internal.instrumentation.jdbc;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import io.opentelemetry.api.OpenTelemetry;
+import io.opentelemetry.sdk.autoconfigure.spi.ConfigProperties;
+import io.opentelemetry.sdk.autoconfigure.spi.internal.DefaultConfigProperties;
+import java.io.PrintWriter;
+import java.sql.Connection;
+import java.sql.SQLException;
+import java.sql.SQLFeatureNotSupportedException;
+import java.util.Collections;
+import java.util.logging.Logger;
+import javax.sql.DataSource;
+import org.junit.jupiter.api.DisplayName;
+import org.junit.jupiter.api.Test;
+import org.springframework.aop.framework.AopProxyUtils;
+import org.springframework.beans.factory.config.BeanPostProcessor;
+import org.springframework.beans.factory.support.DefaultListableBeanFactory;
+
+class DataSourcePostProcessorTest {
+
+  private static final DefaultListableBeanFactory beanFactory = new DefaultListableBeanFactory();
+
+  static {
+    beanFactory.registerSingleton("openTelemetry", OpenTelemetry.noop());
+    beanFactory.registerSingleton(
+        "configProperties", DefaultConfigProperties.createFromMap(Collections.emptyMap()));
+  }
+
+  @Test
+  @DisplayName("when processed bean is NOT of type DataSource should return Object unchanged")
+  void returnsObject() {
+    BeanPostProcessor underTest =
+        new DataSourcePostProcessor(
+            beanFactory.getBeanProvider(OpenTelemetry.class),
+            beanFactory.getBeanProvider(ConfigProperties.class));
+
+    Object nonDataSource = new Object();
+    assertThat(underTest.postProcessAfterInitialization(nonDataSource, "testObject"))
+        .isSameAs(nonDataSource);
+  }
+
+  @Test
+  @DisplayName("when processed bean is of type DataSource should return DataSource proxy")
+  void returnsDataSourceProxy() {
+    BeanPostProcessor underTest =
+        new DataSourcePostProcessor(
+            beanFactory.getBeanProvider(OpenTelemetry.class),
+            beanFactory.getBeanProvider(ConfigProperties.class));
+
+    DataSource originalDataSource = new TestDataSource();
+
+    Object result = underTest.postProcessAfterInitialization(originalDataSource, "testDataSource");
+
+    assertThat(result).isInstanceOf(DataSource.class);
+    assertThat(result).isNotSameAs(originalDataSource);
+
+    Object target = AopProxyUtils.getSingletonTarget(result);
+    assertThat(target).isSameAs(originalDataSource);
+  }
+
+  @Test
+  @DisplayName("when processed bean is scoped proxy target should return unchanged")
+  void returnsScopedProxyTargetUnchanged() {
+    BeanPostProcessor underTest =
+        new DataSourcePostProcessor(
+            beanFactory.getBeanProvider(OpenTelemetry.class),
+            beanFactory.getBeanProvider(ConfigProperties.class));
+
+    DataSource dataSource = new TestDataSource();
+    String scopedTargetBeanName = "scopedTarget.testDataSource";
+
+    Object result = underTest.postProcessAfterInitialization(dataSource, scopedTargetBeanName);
+
+    assertThat(result).isSameAs(dataSource);
+  }
+
+  @Test
+  @DisplayName("proxy should delegate method calls to wrapped telemetry DataSource")
+  void proxyDelegatesMethodCalls() throws SQLException {
+    BeanPostProcessor underTest =
+        new DataSourcePostProcessor(
+            beanFactory.getBeanProvider(OpenTelemetry.class),
+            beanFactory.getBeanProvider(ConfigProperties.class));
+
+    DataSource originalDataSource = new TestDataSource();
+
+    Object result = underTest.postProcessAfterInitialization(originalDataSource, "testDataSource");
+
+    DataSource proxiedDataSource = (DataSource) result;
+    Connection connection = proxiedDataSource.getConnection();
+
+    assertThat(connection).isNotNull();
+  }
+
+  private static class TestDataSource implements DataSource {
+    @Override
+    public Connection getConnection() throws SQLException {
+      Connection mockConnection = mock(Connection.class);
+      when(mockConnection.getMetaData()).thenReturn(mock(java.sql.DatabaseMetaData.class));
+      return mockConnection;
+    }
+
+    @Override
+    public Connection getConnection(String username, String password) throws SQLException {
+      return getConnection();
+    }
+
+    @Override
+    public PrintWriter getLogWriter() throws SQLException {
+      return null;
+    }
+
+    @Override
+    public void setLogWriter(PrintWriter out) throws SQLException {}
+
+    @Override
+    public void setLoginTimeout(int seconds) throws SQLException {}
+
+    @Override
+    public int getLoginTimeout() throws SQLException {
+      return 0;
+    }
+
+    @Override
+    public Logger getParentLogger() throws SQLFeatureNotSupportedException {
+      throw new SQLFeatureNotSupportedException();
+    }
+
+    @Override
+    public <T> T unwrap(Class<T> iface) throws SQLException {
+      throw new SQLException("Not supported");
+    }
+
+    @Override
+    public boolean isWrapperFor(Class<?> iface) throws SQLException {
+      return false;
+    }
+  }
+}

smoke-tests-otel-starter/spring-boot-3.2/src/main/java/io/opentelemetry/spring/smoketest/RuntimeHints.java

@@ -27,5 +27,8 @@ public void registerHints(
             hint -> {
               hint.withMembers(MemberCategory.INVOKE_DECLARED_CONSTRUCTORS);
             });
+
+    // Register proxy hints for DataSource AOP proxy used by DataSourcePostProcessor
+    hints.proxies().registerJdkProxy(TypeReference.of("javax.sql.DataSource"));
   }
 }

smoke-tests-otel-starter/spring-boot-3/src/main/java/io/opentelemetry/spring/smoketest/RuntimeHints.java

@@ -5,9 +5,13 @@
 
 package io.opentelemetry.spring.smoketest;
 
+import javax.sql.DataSource;
+import org.springframework.aop.SpringProxy;
+import org.springframework.aop.framework.Advised;
 import org.springframework.aot.hint.MemberCategory;
 import org.springframework.aot.hint.RuntimeHintsRegistrar;
 import org.springframework.aot.hint.TypeReference;
+import org.springframework.core.DecoratingProxy;
 
 // Necessary for GraalVM native test
 public class RuntimeHints implements RuntimeHintsRegistrar {
@@ -27,5 +31,11 @@ public void registerHints(
             hint -> {
               hint.withMembers(MemberCategory.INVOKE_DECLARED_CONSTRUCTORS);
             });
+
+    // Register proxy hints for DataSource AOP proxy used by DataSourcePostProcessor
+    hints
+        .proxies()
+        .registerJdkProxy(
+            DataSource.class, SpringProxy.class, Advised.class, DecoratingProxy.class);
   }
 }