Add context propagation for postgres

instrumentation/jdbc/build.gradle.kts

@@ -13,6 +13,7 @@ muzzle {
 
 dependencies {
   compileOnly("io.opentelemetry:opentelemetry-api")
+  compileOnly(project(":custom"))
 
   testInstrumentation("io.opentelemetry.javaagent.instrumentation:opentelemetry-javaagent-jdbc")
 
@@ -26,8 +27,14 @@ dependencies {
   // Oracle
   testLibrary("com.oracle.database.jdbc:ojdbc8:23.9.0.25.07")
   testImplementation("org.testcontainers:testcontainers-oracle-free")
+
+  // PostgreSQL
+  testLibrary("org.postgresql:postgresql:42.1.1")
+  testImplementation("org.testcontainers:testcontainers-postgresql")
 }
 
 tasks.withType<Test>().configureEach {
+  systemProperty("testLatestDeps", findProperty("testLatestDeps") as Boolean)
   jvmArgs("-Dotel.instrumentation.splunk-jdbc.enabled=true")
+  jvmArgs("-Dotel.service.name=test-service")
 }

instrumentation/jdbc/src/main/java/com/splunk/opentelemetry/instrumentation/jdbc/PropagatorInitializer.java

@@ -0,0 +1,74 @@
+/*
+ * Copyright Splunk Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.splunk.opentelemetry.instrumentation.jdbc;
+
+import static io.opentelemetry.sdk.autoconfigure.AutoConfigureUtil.getResource;
+import static io.opentelemetry.semconv.ServiceAttributes.SERVICE_NAME;
+
+import com.google.auto.service.AutoService;
+import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.propagation.TextMapGetter;
+import io.opentelemetry.context.propagation.TextMapPropagator;
+import io.opentelemetry.context.propagation.TextMapSetter;
+import io.opentelemetry.javaagent.extension.AgentListener;
+import io.opentelemetry.sdk.autoconfigure.AutoConfiguredOpenTelemetrySdk;
+import io.opentelemetry.sdk.resources.Resource;
+import java.util.Collection;
+import java.util.Collections;
+
+@AutoService(AgentListener.class)
+public class PropagatorInitializer implements AgentListener {
+  // propagates service.name attribute
+  static TextMapPropagator defaultPropagator = TextMapPropagator.noop();
+  // propagates service.name attribute and traceparent
+  static TextMapPropagator traceContextPropagator = W3CTraceContextPropagator.getInstance();
+
+  @Override
+  public void afterAgent(AutoConfiguredOpenTelemetrySdk sdk) {
+    Resource resource = getResource(sdk);
+    String serviceName = resource.getAttribute(SERVICE_NAME);
+    if (!"unknown_service:java".equals(serviceName)) {
+      defaultPropagator = new ServiceAttributePropagator(serviceName);
+      traceContextPropagator =
+          TextMapPropagator.composite(defaultPropagator, traceContextPropagator);
+    }
+  }
+
+  private static class ServiceAttributePropagator implements TextMapPropagator {
+    private final String serviceName;
+
+    ServiceAttributePropagator(String serviceName) {
+      this.serviceName = serviceName;
+    }
+
+    @Override
+    public <C> void inject(Context context, C carrier, TextMapSetter<C> setter) {
+      setter.set(carrier, SERVICE_NAME.getKey(), serviceName);
+    }
+
+    @Override
+    public <C> Context extract(Context context, C carrier, TextMapGetter<C> getter) {
+      return context;
+    }
+
+    @Override
+    public Collection<String> fields() {
+      return Collections.emptyList();
+    }
+  }
+}

instrumentation/jdbc/src/main/java/com/splunk/opentelemetry/instrumentation/jdbc/SqlCommenterInitializer.java

@@ -0,0 +1,45 @@
+/*
+ * Copyright Splunk Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.splunk.opentelemetry.instrumentation.jdbc;
+
+import com.google.auto.service.AutoService;
+import io.opentelemetry.instrumentation.api.incubator.semconv.db.internal.SqlCommenterBuilder;
+import io.opentelemetry.javaagent.bootstrap.internal.AgentInstrumentationConfig;
+import io.opentelemetry.javaagent.bootstrap.internal.sqlcommenter.SqlCommenterCustomizer;
+
+@AutoService(SqlCommenterCustomizer.class)
+public class SqlCommenterInitializer implements SqlCommenterCustomizer {
+
+  @Override
+  public void customize(SqlCommenterBuilder sqlCommenterBuilder) {
+    sqlCommenterBuilder.setEnabled(
+        AgentInstrumentationConfig.get()
+            .getBoolean("otel.instrumentation.splunk-jdbc.enabled", false));
+    sqlCommenterBuilder.setPropagator(
+        (connection, executed) -> {
+          // for postgres we add traceparent to comments, oracle and mssql use other means to
+          // propagate context and other databases are currently unsupported
+          if (connection.getClass().getName().startsWith("org.postgresql.jdbc")) {
+            return PropagatorInitializer.traceContextPropagator;
+          }
+
+          // note that besides jdbc this applies to r2dbc and other data access apis that upstream
+          // has sqlcommenter support for
+          return PropagatorInitializer.defaultPropagator;
+        });
+  }
+}

instrumentation/jdbc/src/test/java/com/splunk/opentelemetry/instrumentation/jdbc/AbstractConnectionUsingDbContextPropagationTest.java

@@ -0,0 +1,94 @@
+/*
+ * Copyright Splunk Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.splunk.opentelemetry.instrumentation.jdbc;
+
+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.propagation.W3CTraceContextPropagator;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.propagation.TextMapGetter;
+import io.opentelemetry.javaagent.bootstrap.CallDepth;
+import java.sql.Connection;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.Collections;
+
+public abstract class AbstractConnectionUsingDbContextPropagationTest
+    extends AbstractDbContextPropagationTest {
+
+  @Override
+  protected void assertBeforeQuery(Connection connection) throws Exception {
+    assertNoContext(connection);
+  }
+
+  @Override
+  protected void assertAfterQuery(Connection connection, SpanContext parent, SpanContext jdbc)
+      throws Exception {
+    assertSameSpan(jdbc, getContext(connection));
+    assertNoContext(connection);
+  }
+
+  private static void assertSameSpan(SpanContext expected, Context context) {
+    SpanContext actual = Span.fromContext(context).getSpanContext();
+    assertThat(expected.getTraceId()).isEqualTo(actual.getTraceId());
+    assertThat(expected.getSpanId()).isEqualTo(actual.getSpanId());
+  }
+
+  private static Context toContext(String traceparent) {
+    if (traceparent == null) {
+      return Context.root();
+    }
+
+    return W3CTraceContextPropagator.getInstance()
+        .extract(
+            Context.root(),
+            traceparent,
+            new TextMapGetter<>() {
+              public String get(String carrier, String key) {
+                if ("traceparent".equals(key)) {
+                  return carrier;
+                }
+                return null;
+              }
+
+              @Override
+              public Iterable<String> keys(String carrier) {
+                return Collections.singleton("traceparent");
+              }
+            });
+  }
+
+  protected abstract String getTraceparent(Connection connection) throws SQLException;
+
+  private Context getContext(Connection connection) throws SQLException {
+    return toContext(getTraceparent(connection));
+  }
+
+  private void assertNoContext(Connection connection) throws SQLException {
+    CallDepth callDepthJdbc = CallDepth.forClass(Statement.class);
+    // disable jdbc instrumentation, so it wouldn't create a span for the statement execution
+    callDepthJdbc.getAndIncrement();
+    try (Statement statement = connection.createStatement()) {
+      statement.execute("SELECT 1");
+      assertSameSpan(SpanContext.getInvalid(), getContext(connection));
+    } finally {
+      callDepthJdbc.decrementAndGet();
+    }
+  }
+}