To enable trace when transaction commit or rollback in jdbc

instrumentation-api-incubator/src/main/jflex/SqlSanitizer.jflex

@@ -323,6 +323,12 @@ WHITESPACE           = [ \t\r\n]+
     }
   }
 
+  private class Commit extends Operation {
+  }
+
+  private class Rollback extends Operation {
+  }
+
   private class Create extends DdlOperation {
   }
 
@@ -461,6 +467,20 @@ WHITESPACE           = [ \t\r\n]+
           appendCurrentFragment();
           if (isOverLimit()) return YYEOF;
       }
+  "COMMIT" {
+          if (!insideComment) {
+            setOperation(new Commit());
+          }
+          appendCurrentFragment();
+          if (isOverLimit()) return YYEOF;
+      }
+  "ROLLBACK" {
+          if (!insideComment) {
+            setOperation(new Rollback());
+          }
+          appendCurrentFragment();
+          if (isOverLimit()) return YYEOF;
+      }
 
   {COMMA} {
           if (!insideComment && !extractionDone) {

instrumentation-api-incubator/src/test/java/io/opentelemetry/instrumentation/api/incubator/semconv/db/SqlStatementSanitizerTest.java

@@ -71,6 +71,17 @@ void checkDdlOperationStatementsAreOk(
     assertThat(result.getMainIdentifier()).isEqualTo(expected.getMainIdentifier());
   }
 
+  @ParameterizedTest
+  @ArgumentsSource(TransactionArgs.class)
+  void checkTransactionOperationStatementsAreOk(
+      String actual, Function<String, SqlStatementInfo> expectFunc) {
+    SqlStatementInfo result = SqlStatementSanitizer.create(true).sanitize(actual);
+    SqlStatementInfo expected = expectFunc.apply(actual);
+    assertThat(result.getFullStatement()).isEqualTo(expected.getFullStatement());
+    assertThat(result.getOperation()).isEqualTo(expected.getOperation());
+    assertThat(result.getMainIdentifier()).isEqualTo(expected.getMainIdentifier());
+  }
+
   @Test
   void lotsOfTicksDontCauseStackOverflowOrLongRuntimes() {
     String s = "'";
@@ -402,4 +413,21 @@ public Stream<? extends Arguments> provideArguments(ExtensionContext context) {
               "CREATE PROCEDURE p AS SELECT * FROM table GO", expect("CREATE PROCEDURE", null)));
     }
   }
+
+  static class TransactionArgs implements ArgumentsProvider {
+    static Function<String, SqlStatementInfo> expect(String operation, String identifier) {
+      return sql -> SqlStatementInfo.create(sql, operation, identifier);
+    }
+
+    @Override
+    public Stream<? extends Arguments> provideArguments(ExtensionContext context) {
+      return Stream.of(
+          Arguments.of("COMMIT", expect("COMMIT", null)),
+          Arguments.of("commit", expect("COMMIT", null)),
+          Arguments.of("/*COMMIT*/", expect(null, null)),
+          Arguments.of("ROLLBACK", expect("ROLLBACK", null)),
+          Arguments.of("rollback", expect("ROLLBACK", null)),
+          Arguments.of("/*ROLLBACK*/", expect(null, null)));
+    }
+  }
 }

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

@@ -20,6 +20,8 @@
 
 package io.opentelemetry.instrumentation.jdbc.internal;
 
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
 import io.opentelemetry.instrumentation.jdbc.internal.dbinfo.DbInfo;
 import java.sql.Array;
@@ -153,7 +155,7 @@ public CallableStatement prepareCall(
 
   @Override
   public void commit() throws SQLException {
-    delegate.commit();
+    wrapCall(delegate::commit, "COMMIT");
   }
 
   @Override
@@ -259,13 +261,13 @@ public Savepoint setSavepoint(String name) throws SQLException {
   @SuppressWarnings("UngroupedOverloads")
   @Override
   public void rollback() throws SQLException {
-    delegate.rollback();
+    wrapCall(delegate::rollback, "ROLLBACK");
   }
 
   @SuppressWarnings("UngroupedOverloads")
   @Override
   public void rollback(Savepoint savepoint) throws SQLException {
-    delegate.rollback(savepoint);
+    wrapCall(() -> delegate.rollback(savepoint), "ROLLBACK");
   }
 
   @Override
@@ -369,4 +371,29 @@ public boolean isWrapperFor(Class<?> iface) throws SQLException {
   public DbInfo getDbInfo() {
     return dbInfo;
   }
+
+  protected <E extends Exception> void wrapCall(ThrowingSupplier<E> callable, String statement)
+      throws E {
+    Context parentContext = Context.current();
+    DbRequest request = DbRequest.create(dbInfo, statement);
+    if (!this.statementInstrumenter.shouldStart(parentContext, request)) {
+      callable.call();
+      return;
+    }
+
+    Context context = this.statementInstrumenter.start(parentContext, request);
+    try (Scope ignored = context.makeCurrent()) {
+      callable.call();
+    } catch (Throwable t) {
+      this.statementInstrumenter.end(context, request, null, t);
+      throw t;
+    }
+    this.statementInstrumenter.end(context, request, null, null);
+  }
+
+  protected interface ThrowingSupplier<E extends Exception> {
+    String statement = null;
+
+    void call() throws E;
+  }
 }

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

@@ -163,6 +163,48 @@ void testVerifyPrepareCallReturnsOtelWrapper() throws Exception {
     connection.close();
   }
 
+  @Test
+  void testVerifyCommit() throws Exception {
+    Instrumenter<DbRequest, Void> instrumenter =
+        createStatementInstrumenter(testing.getOpenTelemetry());
+    DbInfo dbInfo = getDbInfo();
+    OpenTelemetryConnection connection =
+        new OpenTelemetryConnection(new TestConnection(), dbInfo, instrumenter);
+    String query = "COMMIT";
+    Statement statement = connection.createStatement();
+
+    testing.runWithSpan(
+        "parent",
+        () -> {
+          assertThat(statement.execute(query)).isTrue();
+        });
+    transactionTraceAssertion(dbInfo, query);
+
+    statement.close();
+    connection.close();
+  }
+
+  @Test
+  void testVerifyRollback() throws Exception {
+    Instrumenter<DbRequest, Void> instrumenter =
+        createStatementInstrumenter(testing.getOpenTelemetry());
+    DbInfo dbInfo = getDbInfo();
+    OpenTelemetryConnection connection =
+        new OpenTelemetryConnection(new TestConnection(), dbInfo, instrumenter);
+    String query = "ROLLBACK";
+    Statement statement = connection.createStatement();
+
+    testing.runWithSpan(
+        "parent",
+        () -> {
+          assertThat(statement.execute(query)).isTrue();
+        });
+    transactionTraceAssertion(dbInfo, query);
+
+    statement.close();
+    connection.close();
+  }
+
   private static DbInfo getDbInfo() {
     return DbInfo.builder()
         .system("my_system")
@@ -198,4 +240,26 @@ private static void jdbcTraceAssertion(DbInfo dbInfo, String query) {
                             equalTo(ServerAttributes.SERVER_ADDRESS, dbInfo.getHost()),
                             equalTo(ServerAttributes.SERVER_PORT, dbInfo.getPort()))));
   }
+
+  @SuppressWarnings("deprecation") // old semconv
+  private static void transactionTraceAssertion(DbInfo dbInfo, String query) {
+    testing.waitAndAssertTraces(
+        trace ->
+            trace.hasSpansSatisfyingExactly(
+                span -> span.hasName("parent").hasKind(SpanKind.INTERNAL).hasNoParent(),
+                span ->
+                    span.hasName(query + " " + dbInfo.getName())
+                        .hasKind(SpanKind.CLIENT)
+                        .hasParent(trace.getSpan(0))
+                        .hasAttributesSatisfyingExactly(
+                            equalTo(DbIncubatingAttributes.DB_SYSTEM, dbInfo.getSystem()),
+                            equalTo(DbIncubatingAttributes.DB_NAME, dbInfo.getName()),
+                            equalTo(DbIncubatingAttributes.DB_USER, dbInfo.getUser()),
+                            equalTo(
+                                DbIncubatingAttributes.DB_CONNECTION_STRING, dbInfo.getShortUrl()),
+                            equalTo(DbIncubatingAttributes.DB_STATEMENT, query),
+                            equalTo(DbIncubatingAttributes.DB_OPERATION, query),
+                            equalTo(ServerAttributes.SERVER_ADDRESS, dbInfo.getHost()),
+                            equalTo(ServerAttributes.SERVER_PORT, dbInfo.getPort()))));
+  }
 }