Add sqlcommenter support for jdbc and r2dbc

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

@@ -34,6 +34,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;
@@ -87,6 +88,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 =
@@ -138,6 +141,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,84 @@
+/*
+ * 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;
+import javax.annotation.Nullable;
+
+/**
+ * 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 {
+      @Nullable String traceparent;
+      @Nullable String tracestate;
+    }
+
+    State state = new State();
+
+    W3CTraceContextPropagator.getInstance()
+        .inject(
+            Context.current(),
+            state,
+            (carrier, key, value) -> {
+              if (carrier == null) {
+                return;
+              }
+              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,7 +1,8 @@
 # Settings for the JDBC instrumentation
 
-| System property                                                   | Type    | Default | Description                                                                                                                                                                                                                                                                   |
-|-------------------------------------------------------------------|---------|---------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
-| `otel.instrumentation.jdbc.statement-sanitizer.enabled`           | Boolean | `true`  | Enables the DB statement sanitization.                                                                                                                                                                                                                                        |
-| `otel.instrumentation.jdbc.experimental.capture-query-parameters` | Boolean | `false` | Enable the capture of query parameters as span attributes. Enabling this option disables the statement sanitization. <p>WARNING: captured query parameters may contain sensitive information such as passwords, personally identifiable information or protected health info. |
-| `otel.instrumentation.jdbc.experimental.transaction.enabled`      | Boolean | `false` | Enables experimental instrumentation to create spans for COMMIT and ROLLBACK operations.                                                                                                                                                                                      |
+| System property                                                   | Type    | Default | Description                                                                                                                                                                                                                                                                                                                                  |
+|-------------------------------------------------------------------|---------|---------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| `otel.instrumentation.jdbc.statement-sanitizer.enabled`           | Boolean | `true`  | Enables the DB statement sanitization.                                                                                                                                                                                                                                                                                                       |
+| `otel.instrumentation.jdbc.experimental.capture-query-parameters` | Boolean | `false` | Enable the capture of query parameters as span attributes. Enabling this option disables the statement sanitization. <p>WARNING: captured query parameters may contain sensitive information such as passwords, personally identifiable information or protected health info.                                                                |
+| `otel.instrumentation.jdbc.experimental.transaction.enabled`      | Boolean | `false` | Enables experimental instrumentation to create spans for COMMIT and ROLLBACK operations.                                                                                                                                                                                                                                                     |
+| `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

@@ -66,12 +66,21 @@ tasks {
     include("**/SlickTest.*")
   }
 
+  val testSqlCommenter by registering(Test::class) {
+    filter {
+      includeTestsMatching("SqlCommenterTest")
+    }
+    include("**/SqlCommenterTest.*")
+    jvmArgs("-Dotel.instrumentation.jdbc.sqlcommenter.enabled=true")
+  }
+
   val testStableSemconv by registering(Test::class) {
     testClassesDirs = sourceSets.test.get().output.classesDirs
     classpath = sourceSets.test.get().runtimeClasspath
 
     filter {
       excludeTestsMatching("SlickTest")
+      excludeTestsMatching("SqlCommenterTest")
       excludeTestsMatching("PreparedStatementParametersTest")
     }
     jvmArgs("-Dotel.instrumentation.jdbc-datasource.enabled=true")
@@ -102,13 +111,18 @@ tasks {
   test {
     filter {
       excludeTestsMatching("SlickTest")
+      excludeTestsMatching("SqlCommenterTest")
       excludeTestsMatching("PreparedStatementParametersTest")
     }
     jvmArgs("-Dotel.instrumentation.jdbc-datasource.enabled=true")
   }
 
   check {
-    dependsOn(testSlick, testStableSemconv, testSlickStableSemconv, testCaptureParameters)
+    dependsOn(testSlick)
+    dependsOn(testSqlCommenter)
+    dependsOn(testStableSemconv)
+    dependsOn(testSlickStableSemconv)
+    dependsOn(testCaptureParameters)
   }
 }
 

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

@@ -18,16 +18,21 @@
 import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments;
 
 import io.opentelemetry.context.Context;
+import io.opentelemetry.context.ContextKey;
+import io.opentelemetry.context.ImplicitContextKeyed;
 import io.opentelemetry.context.Scope;
 import io.opentelemetry.instrumentation.jdbc.internal.DbRequest;
 import io.opentelemetry.instrumentation.jdbc.internal.JdbcData;
+import io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
 import java.sql.Connection;
 import java.sql.PreparedStatement;
 import java.util.Locale;
 import javax.annotation.Nullable;
 import net.bytebuddy.asm.Advice;
+import net.bytebuddy.asm.Advice.AssignReturned;
+import net.bytebuddy.asm.Advice.AssignReturned.ToArguments.ToArgument;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 
@@ -48,7 +53,7 @@ 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");
     transformer.applyAdviceToMethod(
@@ -59,14 +64,72 @@ public void transform(TypeTransformer transformer) {
   @SuppressWarnings("unused")
   public static class PrepareAdvice {
 
-    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static final class PrepareContext implements ImplicitContextKeyed {
+
+      private static final ContextKey<PrepareContext> KEY =
+          ContextKey.named("jdbc-prepare-context");
+
+      private final String originalSql;
+
+      private PrepareContext(String originalSql) {
+        this.originalSql = originalSql;
+      }
+
+      public String get() {
+        return originalSql;
+      }
+
+      @Nullable
+      public static PrepareContext get(Context context) {
+        return context.get(KEY);
+      }
+
+      public static Context init(Context context, String originalSql) {
+        if (context.get(KEY) != null) {
+          return context;
+        }
+        return context.with(new PrepareContext(originalSql));
+      }
+
+      @Override
+      public Context storeInContext(Context context) {
+        return context.with(KEY, this);
+      }
+    }
+
+    @AssignReturned.ToArguments(@ToArgument(value = 0, index = 0))
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static Object[] processSql(@Advice.Argument(0) String sql) {
+      Context context = Java8BytecodeBridge.currentContext();
+      if (PrepareContext.get(context) == null) {
+        // process sql only in the outermost prepare call and save the original sql in context
+        String processSql = JdbcSingletons.processSql(sql);
+        Scope scope = PrepareContext.init(context, sql).makeCurrent();
+        return new Object[] {processSql, scope};
+      } else {
+        return new Object[] {sql, null};
+      }
+    }
+
+    @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
     public static void addDbInfo(
-        @Advice.Argument(0) String sql, @Advice.Return PreparedStatement statement) {
-      if (JdbcSingletons.isWrapper(statement, PreparedStatement.class)) {
+        @Advice.Return PreparedStatement statement,
+        @Advice.Enter Object[] enterResult,
+        @Advice.Thrown Throwable error) {
+      Context context = Java8BytecodeBridge.currentContext();
+      PrepareContext prepareContext = PrepareContext.get(context);
+      Scope scope = (Scope) enterResult[1];
+      if (scope != null) {
+        scope.close();
+      }
+      if (error != null
+          || prepareContext == null
+          || JdbcSingletons.isWrapper(statement, PreparedStatement.class)) {
         return;
       }
 
-      JdbcData.preparedStatement.set(statement, sql);
+      String originalSql = prepareContext.get();
+      JdbcData.preparedStatement.set(statement, originalSql);
     }
   }
 

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

@@ -8,6 +8,7 @@
 import static io.opentelemetry.instrumentation.jdbc.internal.JdbcInstrumenterFactory.createDataSourceInstrumenter;
 
 import io.opentelemetry.api.GlobalOpenTelemetry;
+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.AttributesExtractor;
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
@@ -28,6 +29,11 @@ public final class JdbcSingletons {
   private static final Instrumenter<DbRequest, Void> TRANSACTION_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());
   public static final boolean CAPTURE_QUERY_PARAMETERS;
 
   static {
@@ -97,5 +103,9 @@ private static <T extends Wrapper> boolean isWrapperInternal(T object, Class<T>
     return false;
   }
 
+  public static String processSql(String sql) {
+    return SQLCOMMENTER_ENABLED ? SqlCommenterUtil.processQuery(sql) : sql;
+  }
+
   private JdbcSingletons() {}
 }