Add context propagation for postgres

instrumentation/jdbc/build.gradle.kts

@@ -13,6 +13,8 @@ muzzle {
 
 dependencies {
   compileOnly("io.opentelemetry:opentelemetry-api")
+  compileOnly("io.opentelemetry.semconv:opentelemetry-semconv-incubating")
+  compileOnly(project(":custom"))
 
   testInstrumentation("io.opentelemetry.javaagent.instrumentation:opentelemetry-javaagent-jdbc")
 
@@ -26,8 +28,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,42 @@
+/*
+ * 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 static io.opentelemetry.semconv.incubating.DeploymentIncubatingAttributes.DEPLOYMENT_ENVIRONMENT_NAME;
+import static io.opentelemetry.semconv.incubating.ServiceIncubatingAttributes.SERVICE_NAMESPACE;
+
+import com.google.auto.service.AutoService;
+import io.opentelemetry.javaagent.extension.AgentListener;
+import io.opentelemetry.sdk.autoconfigure.AutoConfiguredOpenTelemetrySdk;
+import io.opentelemetry.sdk.resources.Resource;
+
+@AutoService(AgentListener.class)
+public class PropagatorInitializer implements AgentListener {
+
+  @Override
+  public void afterAgent(AutoConfiguredOpenTelemetrySdk sdk) {
+    Resource resource = getResource(sdk);
+    String serviceName = resource.getAttribute(SERVICE_NAME);
+    String serviceNamespace = resource.getAttribute(SERVICE_NAMESPACE);
+    String deploymentEnvironment = resource.getAttribute(DEPLOYMENT_ENVIRONMENT_NAME);
+
+    SqlCommenterInitializer.propagator =
+        new ServiceAttributePropagator(serviceName, serviceNamespace, deploymentEnvironment);
+  }
+}

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

@@ -0,0 +1,65 @@
+/*
+ * 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.semconv.ServiceAttributes.SERVICE_NAME;
+import static io.opentelemetry.semconv.incubating.DeploymentIncubatingAttributes.DEPLOYMENT_ENVIRONMENT_NAME;
+import static io.opentelemetry.semconv.incubating.ServiceIncubatingAttributes.SERVICE_NAMESPACE;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.propagation.TextMapGetter;
+import io.opentelemetry.context.propagation.TextMapPropagator;
+import io.opentelemetry.context.propagation.TextMapSetter;
+import java.util.Collection;
+import java.util.Collections;
+
+class ServiceAttributePropagator implements TextMapPropagator {
+  private final String serviceName;
+  private final String serviceNamespace;
+  private final String deploymentEnvironment;
+
+  ServiceAttributePropagator(
+      String serviceName, String serviceNamespace, String deploymentEnvironment) {
+    this.serviceName = serviceName;
+    this.serviceNamespace = serviceNamespace;
+    this.deploymentEnvironment = deploymentEnvironment;
+  }
+
+  @Override
+  public <C> void inject(Context context, C carrier, TextMapSetter<C> setter) {
+    setIfNotNull(setter, carrier, SERVICE_NAME.getKey(), serviceName);
+    setIfNotNull(setter, carrier, SERVICE_NAMESPACE.getKey(), serviceNamespace);
+    setIfNotNull(setter, carrier, DEPLOYMENT_ENVIRONMENT_NAME.getKey(), deploymentEnvironment);
+  }
+
+  private static <C> void setIfNotNull(
+      TextMapSetter<C> setter, C carrier, String key, String value) {
+    if (value != null) {
+      setter.set(carrier, key, value);
+    }
+  }
+
+  @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,42 @@
+/*
+ * 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.context.propagation.TextMapPropagator;
+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 {
+  // propagates service.name and other static attributes
+  static TextMapPropagator propagator = TextMapPropagator.noop();
+
+  @Override
+  public void customize(SqlCommenterBuilder sqlCommenterBuilder) {
+    sqlCommenterBuilder.setEnabled(
+        AgentInstrumentationConfig.get()
+            .getBoolean("otel.instrumentation.splunk-jdbc.enabled", false));
+    sqlCommenterBuilder.setPropagator(
+        (connection, executed) -> {
+          // note that besides jdbc this applies to r2dbc and other data access apis that upstream
+          // has sqlcommenter support for
+          return propagator;
+        });
+  }
+}

instrumentation/jdbc/src/main/java/com/splunk/opentelemetry/instrumentation/jdbc/postgresql/PostgreSqlContextPropagator.java

@@ -0,0 +1,34 @@
+/*
+ * 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.postgresql;
+
+import com.splunk.opentelemetry.instrumentation.jdbc.AbstractDbContextPropagator;
+import java.sql.Connection;
+import java.sql.SQLException;
+import javax.annotation.Nullable;
+
+public final class PostgreSqlContextPropagator extends AbstractDbContextPropagator {
+  public static final PostgreSqlContextPropagator INSTANCE = new PostgreSqlContextPropagator();
+
+  private PostgreSqlContextPropagator() {}
+
+  @Override
+  protected void setContext(Connection connection, @Nullable String contextInfo)
+      throws SQLException {
+    connection.setClientInfo("ApplicationName", contextInfo);
+  }
+}

instrumentation/jdbc/src/main/java/com/splunk/opentelemetry/instrumentation/jdbc/postgresql/PostgreSqlInstrumentationModule.java

@@ -0,0 +1,54 @@
+/*
+ * 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.postgresql;
+
+import static java.util.Collections.singletonList;
+
+import com.google.auto.service.AutoService;
+import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.sdk.autoconfigure.spi.ConfigProperties;
+import java.util.List;
+
+@AutoService(InstrumentationModule.class)
+public class PostgreSqlInstrumentationModule extends InstrumentationModule {
+
+  public PostgreSqlInstrumentationModule() {
+    super("splunk-jdbc", "splunk-jdbc-postgresql");
+  }
+
+  @Override
+  public int order() {
+    // run after jdbc instrumentation
+    return 2;
+  }
+
+  @Override
+  public boolean defaultEnabled(ConfigProperties config) {
+    return false;
+  }
+
+  @Override
+  public List<TypeInstrumentation> typeInstrumentations() {
+    return singletonList(new PostgreSqlStatementInstrumentation());
+  }
+
+  @Override
+  public boolean isHelperClass(String className) {
+    return className.startsWith("com.splunk.opentelemetry.instrumentation");
+  }
+}

instrumentation/jdbc/src/main/java/com/splunk/opentelemetry/instrumentation/jdbc/postgresql/PostgreSqlStatementInstrumentation.java

@@ -0,0 +1,69 @@
+/*
+ * 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.postgresql;
+
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;
+import static net.bytebuddy.matcher.ElementMatchers.namedOneOf;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments;
+
+import io.opentelemetry.javaagent.bootstrap.CallDepth;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import java.sql.Statement;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+public class PostgreSqlStatementInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return namedOneOf("org.postgresql.jdbc.PgPreparedStatement", "org.postgresql.jdbc.PgStatement");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        isPublic()
+            .and(
+                nameStartsWith("execute")
+                    .and(takesNoArguments().or(takesArgument(0, String.class)))),
+        this.getClass().getName() + "$SetActionAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class SetActionAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter(
+        @Advice.This Statement statement, @Advice.Local("splunkCallDepth") CallDepth callDepth)
+        throws Exception {
+      callDepth = CallDepth.forClass(PostgreSqlContextPropagator.class);
+      if (callDepth.getAndIncrement() > 0) {
+        return;
+      }
+
+      PostgreSqlContextPropagator.INSTANCE.propagateContext(statement.getConnection());
+    }
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void onExit(@Advice.Local("splunkCallDepth") CallDepth callDepth) {
+      callDepth.decrementAndGet();
+    }
+  }
+}