Make sqlcommenter more configurable (#15169)

Author: laurit

instrumentation-api-incubator/src/main/java/io/opentelemetry/instrumentation/api/incubator/semconv/db/internal/SqlCommenter.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 io.opentelemetry.context.propagation.TextMapPropagator;
+import java.util.function.BiFunction;
+import java.util.function.Predicate;
+
+/**
+ * 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 SqlCommenter {
+  private final boolean enabled;
+  private final BiFunction<Object, Boolean, TextMapPropagator> propagator;
+  private final Predicate<Object> prepend;
+
+  SqlCommenter(
+      boolean enabled,
+      BiFunction<Object, Boolean, TextMapPropagator> propagator,
+      Predicate<Object> prepend) {
+    this.enabled = enabled;
+    this.propagator = propagator;
+    this.prepend = prepend;
+  }
+
+  public static SqlCommenterBuilder builder() {
+    return new SqlCommenterBuilder();
+  }
+
+  public static SqlCommenter noop() {
+    return builder().build();
+  }
+
+  /**
+   * Augments the given SQL query with comment containing tracing context.
+   *
+   * @param connection connection object, e.g. JDBC connection or R2DBC connection, that is used to
+   *     execute the query
+   * @param sql original query
+   * @param executed whether the query is immediately executed after being processed, e.g. {@link
+   *     java.sql.Statement#execute(String)}, or may be executed later, e.g. {@link
+   *     java.sql.Connection#prepareStatement(String)}
+   * @return modified query
+   */
+  public String processQuery(Object connection, String sql, boolean executed) {
+    if (!enabled) {
+      return sql;
+    }
+
+    return SqlCommenterUtil.processQuery(
+        sql, propagator.apply(connection, executed), prepend.test(connection));
+  }
+
+  public boolean isEnabled() {
+    return enabled;
+  }
+}

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

@@ -0,0 +1,89 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.api.incubator.semconv.db.internal;
+
+import com.google.errorprone.annotations.CanIgnoreReturnValue;
+import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
+import io.opentelemetry.context.propagation.TextMapPropagator;
+import java.sql.Connection;
+import java.sql.Statement;
+import java.util.function.BiFunction;
+import java.util.function.Predicate;
+
+/**
+ * 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 SqlCommenterBuilder {
+  private boolean enabled;
+  private BiFunction<Object, Boolean, TextMapPropagator> propagator =
+      (unused1, unused2) -> W3CTraceContextPropagator.getInstance();
+  private Predicate<Object> prepend = unused -> false;
+
+  SqlCommenterBuilder() {}
+
+  /** Enable adding sqlcommenter comments to sql queries. Default is disabled. */
+  @CanIgnoreReturnValue
+  public SqlCommenterBuilder setEnabled(boolean enabled) {
+    this.enabled = enabled;
+    return this;
+  }
+
+  /**
+   * Prepend the sqlcommenter comment to the query instead of appending it. Default is to append.
+   */
+  @CanIgnoreReturnValue
+  public SqlCommenterBuilder setPrepend(boolean prepend) {
+    this.prepend = unused -> prepend;
+    return this;
+  }
+
+  /**
+   * Prepend the sqlcommenter comment to the query instead of appending it. Default is to append.
+   *
+   * @param prepend a predicate that receives the database connection. Connection may be a jdbc
+   *     Connection, R2DBC Connection, or any other connection type used by the data access
+   *     framework performing the operation.
+   */
+  @CanIgnoreReturnValue
+  public SqlCommenterBuilder setPrepend(Predicate<Object> prepend) {
+    this.prepend = prepend;
+    return this;
+  }
+
+  /**
+   * Set the propagator used to inject tracing context into sql comments. Default is W3C Trace
+   * Context propagator.
+   */
+  @CanIgnoreReturnValue
+  public SqlCommenterBuilder setPropagator(TextMapPropagator propagator) {
+    this.propagator = (unused1, unused2) -> propagator;
+    return this;
+  }
+
+  /**
+   * Set the propagator used to inject tracing context into sql comments. Default is W3C Trace
+   * Context propagator.
+   *
+   * @param propagator a function that receives the database connection and whether the query is
+   *     executed only once or could be reused. Connection may be a JDBC connection, R2DBC
+   *     connection, or any other connection type used by the data access framework performing the
+   *     operation. If the second argument to the function is true, the query is executed only once
+   *     (e.g. JDBC {@link Statement#execute(String)}) immediately after processing. If false, the
+   *     query could be reused (e.g. JDBC {@link Connection#prepareStatement(String)}).
+   */
+  @CanIgnoreReturnValue
+  public SqlCommenterBuilder setPropagator(
+      BiFunction<Object, Boolean, TextMapPropagator> propagator) {
+    this.propagator = propagator;
+    return this;
+  }
+
+  public SqlCommenter build() {
+    return new SqlCommenter(enabled, propagator, prepend);
+  }
+}

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

@@ -6,25 +6,22 @@
 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 io.opentelemetry.context.propagation.TextMapPropagator;
 import java.io.UnsupportedEncodingException;
 import java.net.URLEncoder;
-import javax.annotation.Nullable;
+import java.util.Iterator;
+import java.util.LinkedHashMap;
+import java.util.Map;
 
-/**
- * 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 {
+final class SqlCommenterUtil {
 
   /**
-   * Append comment containing tracing information at the end of the query. See <a
+   * Append comment containing tracing information to 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) {
+  public static String processQuery(String query, TextMapPropagator propagator, boolean prepend) {
     if (!Span.current().getSpanContext().isValid()) {
       return query;
     }
@@ -33,51 +30,53 @@ public static String processQuery(String query) {
       return query;
     }
 
-    class State {
-      @Nullable String traceparent;
-      @Nullable String tracestate;
-    }
+    Map<String, String> state = new LinkedHashMap<>();
+    propagator.inject(
+        Context.current(),
+        state,
+        (carrier, key, value) -> {
+          if (carrier == null) {
+            return;
+          }
+          carrier.put(key, value);
+        });
 
-    State state = new State();
+    if (state.isEmpty()) {
+      return query;
+    }
 
-    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;
-              }
-            });
+    StringBuilder stringBuilder = new StringBuilder("/*");
     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);
+      for (Iterator<Map.Entry<String, String>> iterator = state.entrySet().iterator();
+          iterator.hasNext(); ) {
+        Map.Entry<String, String> entry = iterator.next();
+        stringBuilder
+            .append(serialize(entry.getKey()))
+            .append("='")
+            .append(serialize(entry.getValue()))
+            .append("'");
+        if (iterator.hasNext()) {
+          stringBuilder.append(", ");
+        }
       }
-      query += "*/";
     } catch (UnsupportedEncodingException exception) {
       // this exception should never happen as UTF-8 encoding is always available
     }
-    return query;
+    stringBuilder.append("*/");
+
+    return prepend ? stringBuilder + " " + query : query + " " + stringBuilder;
   }
 
   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
+    // specification requires percent encoding, here we use the java built 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 + "'";
+    return URLEncoder.encode(value, "UTF-8").replace("+", "%20");
   }
 
   private SqlCommenterUtil() {}

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

@@ -11,12 +11,16 @@
 import io.opentelemetry.api.trace.SpanContext;
 import io.opentelemetry.api.trace.TraceFlags;
 import io.opentelemetry.api.trace.TraceState;
+import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
+import io.opentelemetry.context.propagation.TextMapPropagator;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
+import org.junitpioneer.jupiter.cartesian.CartesianTest;
 
 class SqlCommenterUtilTest {
+  private static final TextMapPropagator propagator = W3CTraceContextPropagator.getInstance();
 
   @ParameterizedTest
   @ValueSource(strings = {"SELECT /**/ 1", "SELECT 1 --", "SELECT '/*'"})
@@ -32,13 +36,14 @@ void skipQueriesWithComments(String query) {
                         TraceState.getDefault())));
 
     try (Scope ignore = parent.makeCurrent()) {
-      assertThat(SqlCommenterUtil.processQuery(query)).isEqualTo(query);
+      assertThat(SqlCommenterUtil.processQuery(query, propagator, false)).isEqualTo(query);
     }
   }
 
-  @ParameterizedTest
-  @ValueSource(booleans = {true, false})
-  void sqlCommenter(boolean hasTraceState) {
+  @CartesianTest
+  void sqlCommenter(
+      @CartesianTest.Values(booleans = {true, false}) boolean hasTraceState,
+      @CartesianTest.Values(booleans = {true, false}) boolean prepend) {
     TraceState state =
         hasTraceState ? TraceState.builder().put("test", "test'").build() : TraceState.getDefault();
     Context parent =
@@ -52,11 +57,12 @@ void sqlCommenter(boolean hasTraceState) {
                         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'*/");
+      String fragment =
+          hasTraceState
+              ? "/*traceparent='00-ff01020304050600ff0a0b0c0d0e0f00-090a0b0c0d0e0f00-01', tracestate='test%3Dtest%27'*/"
+              : "/*traceparent='00-ff01020304050600ff0a0b0c0d0e0f00-090a0b0c0d0e0f00-01'*/";
+      assertThat(SqlCommenterUtil.processQuery("SELECT 1", propagator, prepend))
+          .isEqualTo(prepend ? fragment + " SELECT 1" : "SELECT 1 " + fragment);
     }
   }
 }

instrumentation/jdbc/javaagent/build.gradle.kts

@@ -63,11 +63,17 @@ tasks {
   }
 
   val testSqlCommenter by registering(Test::class) {
+    testClassesDirs = sourceSets.test.get().output.classesDirs
+    classpath = sourceSets.test.get().runtimeClasspath
+
     filter {
       includeTestsMatching("SqlCommenterTest")
     }
     include("**/SqlCommenterTest.*")
-    jvmArgs("-Dotel.instrumentation.jdbc.experimental.sqlcommenter.enabled=true")
+    // This property is read in TestAgentSqlCommenterCustomizer, we use it instead of the
+    // otel.instrumentation.jdbc.experimental.sqlcommenter.enabled to test that the
+    // SqlCommenterCustomizer is run.
+    jvmArgs("-Dotel.testing.sqlcommenter.enabled=true")
   }
 
   val testStableSemconv by registering(Test::class) {

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

@@ -99,11 +99,12 @@ public Context storeInContext(Context context) {
 
     @AssignReturned.ToArguments(@ToArgument(value = 0, index = 0))
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static Object[] processSql(@Advice.Argument(0) String sql) {
+    public static Object[] processSql(
+        @Advice.This Connection connection, @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);
+        String processSql = JdbcSingletons.processSql(connection, sql, false);
         Scope scope = PrepareContext.init(context, sql).makeCurrent();
         return new Object[] {processSql, scope};
       } else {