add test

Author: laurit

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

@@ -1330,7 +1330,7 @@ void testProxyStatement() throws Exception {
     cleanup.deferCleanup(statement);
     cleanup.deferCleanup(connection);
 
-    Statement proxyStatement = ProxyStatementFactory.proxyStatement(statement);
+    Statement proxyStatement = ProxyStatementFactory.proxyStatementWithCustomClassLoader(statement);
     ResultSet resultSet =
         testing.runWithSpan("parent", () -> proxyStatement.executeQuery("SELECT 3"));
 
@@ -1737,4 +1737,100 @@ static Stream<Arguments> transactionOperationsStream() throws SQLException {
         Arguments.of(
             "hsqldb", new JDBCDriver().connect(jdbcUrls.get("hsqldb"), null), "SA", "hsqldb:mem:"));
   }
+
+  private static PreparedStatement wrapPreparedStatement(PreparedStatement statement) {
+    return ProxyStatementFactory.proxyPreparedStatement(
+        (proxy, method, args) -> {
+          if ("isWrapperFor".equals(method.getName())
+              && args.length == 1
+              && args[0] == PreparedStatement.class) {
+            return true;
+          }
+          if ("unwrap".equals(method.getName())
+              && args.length == 1
+              && args[0] == PreparedStatement.class) {
+            return statement;
+          }
+          return testing.runWithSpan("wrapper", () -> method.invoke(statement, args));
+        });
+  }
+
+  // test that tracing does not start from a wrapper
+  // regression test for
+  // https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/14733
+  @Test
+  void testPreparedStatementWrapper() throws SQLException {
+    Connection connection = new org.h2.Driver().connect(jdbcUrls.get("h2"), null);
+    Connection proxyConnection =
+        ProxyStatementFactory.proxy(
+            Connection.class,
+            (proxy, method, args) -> {
+              // we don't implement unwrapping here as that would also cause the executeQuery
+              // instrumentation to get skipped for the prepared statement wrapper
+              if ("prepareStatement".equals(method.getName())) {
+                return wrapPreparedStatement((PreparedStatement) method.invoke(connection, args));
+              }
+              return method.invoke(connection, args);
+            });
+    PreparedStatement statement = proxyConnection.prepareStatement("SELECT 3");
+    cleanup.deferCleanup(statement);
+    cleanup.deferCleanup(connection);
+
+    ResultSet resultSet = testing.runWithSpan("parent", () -> statement.executeQuery());
+
+    resultSet.next();
+    assertThat(resultSet.getInt(1)).isEqualTo(3);
+    testing.waitAndAssertTraces(
+        trace ->
+            trace.hasSpansSatisfyingExactly(
+                span -> span.hasName("parent").hasKind(SpanKind.INTERNAL).hasNoParent(),
+                span ->
+                    span.hasName("wrapper").hasKind(SpanKind.INTERNAL).hasParent(trace.getSpan(0)),
+                span ->
+                    span.hasName("SELECT " + dbNameLower)
+                        .hasKind(SpanKind.CLIENT)
+                        .hasParent(trace.getSpan(1))));
+  }
+
+  // test that tracing does not start from a wrapper
+  // regression test for
+  // https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/14733
+  @Test
+  void testStatementWrapper() throws SQLException {
+    Connection connection = new org.h2.Driver().connect(jdbcUrls.get("h2"), null);
+    Statement statement = connection.createStatement();
+    cleanup.deferCleanup(statement);
+    cleanup.deferCleanup(connection);
+
+    Statement proxyStatement =
+        ProxyStatementFactory.proxyStatement(
+            (proxy, method, args) -> {
+              if ("isWrapperFor".equals(method.getName())
+                  && args.length == 1
+                  && args[0] == Statement.class) {
+                return true;
+              }
+              if ("unwrap".equals(method.getName())
+                  && args.length == 1
+                  && args[0] == Statement.class) {
+                return statement;
+              }
+              return testing.runWithSpan("wrapper", () -> method.invoke(statement, args));
+            });
+    ResultSet resultSet =
+        testing.runWithSpan("parent", () -> proxyStatement.executeQuery("SELECT 3"));
+
+    resultSet.next();
+    assertThat(resultSet.getInt(1)).isEqualTo(3);
+    testing.waitAndAssertTraces(
+        trace ->
+            trace.hasSpansSatisfyingExactly(
+                span -> span.hasName("parent").hasKind(SpanKind.INTERNAL).hasNoParent(),
+                span ->
+                    span.hasName("wrapper").hasKind(SpanKind.INTERNAL).hasParent(trace.getSpan(0)),
+                span ->
+                    span.hasName("SELECT " + dbNameLower)
+                        .hasKind(SpanKind.CLIENT)
+                        .hasParent(trace.getSpan(1))));
+  }
 }

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

@@ -12,50 +12,61 @@
 
 public final class ProxyStatementFactory {
 
-  public static Statement proxyStatement(Statement statement) throws Exception {
+  public static Statement proxyStatementWithCustomClassLoader(Statement statement)
+      throws Exception {
     TestClassLoader classLoader = new TestClassLoader(ProxyStatementFactory.class.getClassLoader());
     Class<?> testInterface = classLoader.loadClass(TestInterface.class.getName());
     if (testInterface.getClassLoader() != classLoader) {
       throw new IllegalStateException("wrong class loader");
     }
     InvocationHandler invocationHandler = (proxy, method, args) -> method.invoke(statement, args);
-    Statement proxyStatement =
-        (Statement)
-            Proxy.newProxyInstance(
-                classLoader, new Class<?>[] {Statement.class, testInterface}, invocationHandler);
-    // adding package private interface TestInterface to jdk proxy forces defining the proxy class
-    // in the same package as the package private interface
-    if (!proxyStatement
-        .getClass()
-        .getName()
-        .startsWith("io.opentelemetry.javaagent.instrumentation.jdbc.test")) {
-      throw new IllegalStateException("proxy statement is in wrong package");
-    }
+    return proxy(
+        Statement.class,
+        classLoader,
+        new Class<?>[] {Statement.class, testInterface},
+        invocationHandler);
+  }
 
-    return proxyStatement;
+  public static Statement proxyStatement(InvocationHandler invocationHandler) {
+    return proxy(Statement.class, invocationHandler);
   }
 
   public static PreparedStatement proxyPreparedStatement(PreparedStatement statement) {
     InvocationHandler invocationHandler = (proxy, method, args) -> method.invoke(statement, args);
-    PreparedStatement proxyStatement =
-        (PreparedStatement)
-            Proxy.newProxyInstance(
-                ProxyStatementFactory.class.getClassLoader(),
-                new Class<?>[] {PreparedStatement.class, TestInterface.class},
-                invocationHandler);
+    return proxyPreparedStatement(invocationHandler);
+  }
+
+  public static PreparedStatement proxyPreparedStatement(InvocationHandler invocationHandler) {
+    return proxy(PreparedStatement.class, invocationHandler);
+  }
+
+  public static <T> T proxy(Class<T> clazz, InvocationHandler invocationHandler) {
+    return proxy(
+        clazz,
+        ProxyStatementFactory.class.getClassLoader(),
+        new Class<?>[] {clazz, TestInterface.class},
+        invocationHandler);
+  }
+
+  public static <T> T proxy(
+      Class<T> clazz,
+      ClassLoader classLoader,
+      Class<?>[] interfaces,
+      InvocationHandler invocationHandler) {
+    T proxy = clazz.cast(Proxy.newProxyInstance(classLoader, interfaces, invocationHandler));
 
     // adding package private interface TestInterface to jdk proxy forces defining the proxy class
     // in the same package as the package private interface
     // by default we ignore jdk proxies, having the proxy in a different package ensures it gets
     // instrumented
-    if (!proxyStatement
+    if (!proxy
         .getClass()
         .getName()
         .startsWith("io.opentelemetry.javaagent.instrumentation.jdbc.test")) {
-      throw new IllegalStateException("proxy statement is in wrong package");
+      throw new IllegalStateException("proxy is in wrong package");
     }
 
-    return proxyStatement;
+    return proxy;
   }
 
   private ProxyStatementFactory() {}