Add sqlcommenter support for jdbc and r2dbc

instrumentation-api-incubator/src/main/java/io/opentelemetry/instrumentation/api/incubator/config/internal/CommonConfig.java

@@ -29,6 +29,7 @@ public final class CommonConfig {
   private final Set<String> knownHttpRequestMethods;
   private final EnduserConfig enduserConfig;
   private final boolean statementSanitizationEnabled;
+  private final boolean sqlCommenterEnabled;
   private final boolean emitExperimentalHttpClientTelemetry;
   private final boolean emitExperimentalHttpServerTelemetry;
   private final boolean redactQueryParameters;
@@ -56,6 +57,8 @@ public CommonConfig(InstrumentationConfig config) {
                 new ArrayList<>(HttpConstants.KNOWN_METHODS)));
     statementSanitizationEnabled =
         config.getBoolean("otel.instrumentation.common.db-statement-sanitizer.enabled", true);
+    sqlCommenterEnabled =
+        config.getBoolean("otel.instrumentation.common.db-sqlcommenter.enabled", false);
     emitExperimentalHttpClientTelemetry =
         config.getBoolean("otel.instrumentation.http.client.emit-experimental-telemetry", false);
     redactQueryParameters =
@@ -107,6 +110,10 @@ public boolean isStatementSanitizationEnabled() {
     return statementSanitizationEnabled;
   }
 
+  public boolean isSqlCommenterEnabled() {
+    return sqlCommenterEnabled;
+  }
+
   public boolean shouldEmitExperimentalHttpClientTelemetry() {
     return emitExperimentalHttpClientTelemetry;
   }

instrumentation-api-incubator/src/main/java/io/opentelemetry/instrumentation/api/incubator/semconv/db/internal/SqlCommenterUtil.java

@@ -0,0 +1,80 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.api.incubator.semconv.db.internal;
+
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
+import io.opentelemetry.context.Context;
+import java.io.UnsupportedEncodingException;
+import java.net.URLEncoder;
+
+/**
+ * This class is internal and experimental. Its APIs are unstable and can change at any time. Its
+ * APIs (or a version of them) may be promoted to the public stable API in the future, but no
+ * guarantees are made.
+ */
+public final class SqlCommenterUtil {
+
+  /**
+   * Append comment containing tracing information at the end of the query. See <a
+   * href="https://google.github.io/sqlcommenter/spec/">sqlcommenter</a> for the description of the
+   * algorithm.
+   */
+  public static String processQuery(String query) {
+    if (!Span.current().getSpanContext().isValid()) {
+      return query;
+    }
+    // skip queries that contain comments
+    if (containsSqlComment(query)) {
+      return query;
+    }
+
+    class State {
+      String traceparent;
+      String tracestate;
+    }
+
+    State state = new State();
+
+    W3CTraceContextPropagator.getInstance()
+        .inject(
+            Context.current(),
+            state,
+            (carrier, key, value) -> {
+              if ("traceparent".equals(key)) {
+                carrier.traceparent = value;
+              } else if ("tracestate".equals(key)) {
+                carrier.tracestate = value;
+              }
+            });
+    try {
+      // we know that the traceparent doesn't contain anything that needs to be encoded
+      query += " /*traceparent='" + state.traceparent + "'";
+      if (state.tracestate != null) {
+        query += ", tracestate=" + serialize(state.tracestate);
+      }
+      query += "*/";
+    } catch (UnsupportedEncodingException exception) {
+      // this exception should never happen as UTF-8 encoding is always available
+    }
+    return query;
+  }
+
+  private static boolean containsSqlComment(String query) {
+    return query.contains("--") || query.contains("/*");
+  }
+
+  private static String serialize(String value) throws UnsupportedEncodingException {
+    // specification requires percent encoding, here we use the java build in url encoder that
+    // encodes space as '+' instead of '%20' as required
+    String result = URLEncoder.encode(value, "UTF-8").replace("+", "%20");
+    // specification requires escaping ' with a backslash, we skip this because URLEncoder already
+    // encodes the '
+    return "'" + result + "'";
+  }
+
+  private SqlCommenterUtil() {}
+}

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

@@ -0,0 +1,62 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.api.incubator.semconv.db.internal;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanContext;
+import io.opentelemetry.api.trace.TraceFlags;
+import io.opentelemetry.api.trace.TraceState;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+
+class SqlCommenterUtilTest {
+
+  @ParameterizedTest
+  @ValueSource(strings = {"SELECT /**/ 1", "SELECT 1 --", "SELECT '/*'"})
+  void skipQueriesWithComments(String query) {
+    Context parent =
+        Context.root()
+            .with(
+                Span.wrap(
+                    SpanContext.create(
+                        "ff01020304050600ff0a0b0c0d0e0f00",
+                        "090a0b0c0d0e0f00",
+                        TraceFlags.getSampled(),
+                        TraceState.getDefault())));
+
+    try (Scope ignore = parent.makeCurrent()) {
+      assertThat(SqlCommenterUtil.processQuery(query)).isEqualTo(query);
+    }
+  }
+
+  @ParameterizedTest
+  @ValueSource(booleans = {true, false})
+  void sqlCommenter(boolean hasTraceState) {
+    TraceState state =
+        hasTraceState ? TraceState.builder().put("test", "test'").build() : TraceState.getDefault();
+    Context parent =
+        Context.root()
+            .with(
+                Span.wrap(
+                    SpanContext.create(
+                        "ff01020304050600ff0a0b0c0d0e0f00",
+                        "090a0b0c0d0e0f00",
+                        TraceFlags.getSampled(),
+                        state)));
+
+    try (Scope ignore = parent.makeCurrent()) {
+      assertThat(SqlCommenterUtil.processQuery("SELECT 1"))
+          .isEqualTo(
+              hasTraceState
+                  ? "SELECT 1 /*traceparent='00-ff01020304050600ff0a0b0c0d0e0f00-090a0b0c0d0e0f00-01', tracestate='test%3Dtest%27'*/"
+                  : "SELECT 1 /*traceparent='00-ff01020304050600ff0a0b0c0d0e0f00-090a0b0c0d0e0f00-01'*/");
+    }
+  }
+}

instrumentation/jdbc/README.md

@@ -1,5 +1,6 @@
 # Settings for the JDBC instrumentation
 
-| System property                                         | Type    | Default | Description                            |
-|---------------------------------------------------------|---------|---------|----------------------------------------|
-| `otel.instrumentation.jdbc.statement-sanitizer.enabled` | Boolean | `true`  | Enables the DB statement sanitization. |
+| System property                                         | Type    | Default | Description                                                                                                                                                                                                                                                                                                                                  |
+|---------------------------------------------------------|---------|---------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| `otel.instrumentation.jdbc.statement-sanitizer.enabled` | Boolean | `true`  | Enables the DB statement sanitization.                                                                                                                                                                                                                                                                                                       |
+| `otel.instrumentation.jdbc.sqlcommenter.enabled`        | Boolean | `false` | Enables augmenting queries with a comment containing the tracing information. See [sqlcommenter](https://google.github.io/sqlcommenter/) for more info. WARNING: augmenting queries with tracing context will make query texts unique, which may have adverse impact on database performance. Consult with database experts before enabling. |

instrumentation/jdbc/javaagent/build.gradle.kts

@@ -62,16 +62,26 @@ tasks {
     include("**/SlickTest.*")
   }
 
+  val testSqlCommenter by registering(Test::class) {
+    filter {
+      includeTestsMatching("SqlCommenterTest")
+    }
+    include("**/SqlCommenterTest.*")
+    jvmArgs("-Dotel.instrumentation.jdbc.sqlcommenter.enabled=true")
+  }
+
   test {
     filter {
       excludeTestsMatching("SlickTest")
+      excludeTestsMatching("SqlCommenterTest")
     }
     jvmArgs("-Dotel.instrumentation.jdbc-datasource.enabled=true")
   }
 
   val testStableSemconv by registering(Test::class) {
     filter {
       excludeTestsMatching("SlickTest")
+      excludeTestsMatching("SqlCommenterTest")
     }
     jvmArgs("-Dotel.instrumentation.jdbc-datasource.enabled=true")
     jvmArgs("-Dotel.semconv-stability.opt-in=database")
@@ -87,6 +97,7 @@ tasks {
 
   check {
     dependsOn(testSlick)
+    dependsOn(testSqlCommenter)
     dependsOn(testStableSemconv)
     dependsOn(testSlickStableSemconv)
   }

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

@@ -37,18 +37,26 @@ public void transform(TypeTransformer transformer) {
     transformer.applyAdviceToMethod(
         nameStartsWith("prepare")
             .and(takesArgument(0, String.class))
-            // Also include CallableStatement, which is a sub type of PreparedStatement
+            // Also include CallableStatement, which is a subtype of PreparedStatement
             .and(returns(implementsInterface(named("java.sql.PreparedStatement")))),
         ConnectionInstrumentation.class.getName() + "$PrepareAdvice");
   }
 
   @SuppressWarnings("unused")
   public static class PrepareAdvice {
 
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void processSql(
+        @Advice.Argument(value = 0, readOnly = false) String sql,
+        @Advice.Local("otelSql") String originalSql) {
+      originalSql = sql;
+      sql = JdbcSingletons.processSql(sql);
+    }
+
     @Advice.OnMethodExit(suppress = Throwable.class)
     public static void addDbInfo(
-        @Advice.Argument(0) String sql, @Advice.Return PreparedStatement statement) {
-      JdbcData.preparedStatement.set(statement, sql);
+        @Advice.Local("otelSql") String originalSql, @Advice.Return PreparedStatement statement) {
+      JdbcData.preparedStatement.set(statement, originalSql);
     }
   }
 }

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

@@ -11,6 +11,7 @@
 import io.opentelemetry.instrumentation.api.incubator.semconv.db.DbClientMetrics;
 import io.opentelemetry.instrumentation.api.incubator.semconv.db.DbClientSpanNameExtractor;
 import io.opentelemetry.instrumentation.api.incubator.semconv.db.SqlClientAttributesExtractor;
+import io.opentelemetry.instrumentation.api.incubator.semconv.db.internal.SqlCommenterUtil;
 import io.opentelemetry.instrumentation.api.incubator.semconv.net.PeerServiceAttributesExtractor;
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
 import io.opentelemetry.instrumentation.api.instrumenter.SpanKindExtractor;
@@ -29,6 +30,11 @@ public final class JdbcSingletons {
   private static final Instrumenter<DbRequest, Void> STATEMENT_INSTRUMENTER;
   public static final Instrumenter<DataSource, DbInfo> DATASOURCE_INSTRUMENTER =
       createDataSourceInstrumenter(GlobalOpenTelemetry.get(), true);
+  private static final boolean SQLCOMMENTER_ENABLED =
+      AgentInstrumentationConfig.get()
+          .getBoolean(
+              "otel.instrumentation.jdbc.sqlcommenter.enabled",
+              AgentCommonConfig.get().isSqlCommenterEnabled());
 
   static {
     JdbcAttributesGetter dbAttributesGetter = new JdbcAttributesGetter();
@@ -63,5 +69,9 @@ public static Instrumenter<DataSource, DbInfo> dataSourceInstrumenter() {
     return DATASOURCE_INSTRUMENTER;
   }
 
+  public static String processSql(String sql) {
+    return SQLCOMMENTER_ENABLED ? SqlCommenterUtil.processQuery(sql) : sql;
+  }
+
   private JdbcSingletons() {}
 }

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

@@ -65,7 +65,7 @@ public static class StatementAdvice {
 
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static void onEnter(
-        @Advice.Argument(0) String sql,
+        @Advice.Argument(value = 0, readOnly = false) String sql,
         @Advice.This Statement statement,
         @Advice.Local("otelCallDepth") CallDepth callDepth,
         @Advice.Local("otelRequest") DbRequest request,
@@ -86,6 +86,8 @@ public static void onEnter(
       Context parentContext = currentContext();
       request = DbRequest.create(statement, sql);
 
+      sql = JdbcSingletons.processSql(sql);
+
       if (request == null || !statementInstrumenter().shouldStart(parentContext, request)) {
         return;
       }
@@ -115,12 +117,15 @@ public static void stopSpan(
   @SuppressWarnings("unused")
   public static class AddBatchAdvice {
 
-    @Advice.OnMethodExit(suppress = Throwable.class)
-    public static void addBatch(@Advice.This Statement statement, @Advice.Argument(0) String sql) {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void addBatch(
+        @Advice.This Statement statement,
+        @Advice.Argument(value = 0, readOnly = false) String sql) {
       if (statement instanceof PreparedStatement) {
         return;
       }
       JdbcData.addStatementBatch(statement, sql);
+      sql = JdbcSingletons.processSql(sql);
     }
   }
 

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

@@ -1632,4 +1632,21 @@ void testPreparedBatch(String system, Connection connection, String username, St
                                 DB_OPERATION_BATCH_SIZE,
                                 emitStableDatabaseSemconv() ? 2L : null))));
   }
+
+  // test that sqlcommenter is not enabled by default
+  @Test
+  void testSqlCommenterNotEnabled() throws SQLException {
+    List<String> executedSql = new ArrayList<>();
+    Connection connection = new TestConnection(executedSql::add);
+    Statement statement = connection.createStatement();
+
+    cleanup.deferCleanup(statement);
+    cleanup.deferCleanup(connection);
+
+    String query = "SELECT 1";
+    testing.runWithSpan("parent", () -> statement.execute(query));
+
+    assertThat(executedSql).hasSize(1);
+    assertThat(executedSql.get(0)).isEqualTo(query);
+  }
 }