Add oracle context propagation

instrumentation/jdbc/build.gradle.kts

@@ -16,10 +16,16 @@ dependencies {
 
   testInstrumentation("io.opentelemetry.javaagent.instrumentation:opentelemetry-javaagent-jdbc")
 
-  testLibrary("com.microsoft.sqlserver:mssql-jdbc:6.1.0.jre8")
   testCompileOnly("io.opentelemetry.javaagent:opentelemetry-javaagent-extension-api")
   testImplementation("org.testcontainers:testcontainers")
+
+  // SQL Server
+  testLibrary("com.microsoft.sqlserver:mssql-jdbc:6.1.0.jre8")
   testImplementation("org.testcontainers:mssqlserver")
+
+  // Oracle
+  testLibrary("com.oracle.database.jdbc:ojdbc8:23.9.0.25.07")
+  testImplementation("org.testcontainers:oracle-free")
 }
 
 tasks.withType<Test>().configureEach {

instrumentation/jdbc/src/main/java/com/splunk/opentelemetry/instrumentation/jdbc/sqlserver/SqlServerUtil.java → instrumentation/jdbc/src/main/java/com/splunk/opentelemetry/instrumentation/jdbc/AbstractDbContextPropagator.java

@@ -14,23 +14,22 @@
  * limitations under the License.
  */
 
-package com.splunk.opentelemetry.instrumentation.jdbc.sqlserver;
+package com.splunk.opentelemetry.instrumentation.jdbc;
 
 import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.instrumentation.api.util.VirtualField;
-import java.nio.charset.StandardCharsets;
 import java.sql.Connection;
-import java.sql.PreparedStatement;
 import java.sql.SQLException;
 import java.util.Objects;
 import java.util.concurrent.atomic.AtomicReference;
+import javax.annotation.Nullable;
 
-public final class SqlServerUtil {
+public abstract class AbstractDbContextPropagator {
   private static final VirtualField<Connection, String> connectionState =
       VirtualField.find(Connection.class, String.class);
 
-  public static void propagateContext(Connection connection) throws SQLException {
+  public void propagateContext(Connection connection) throws SQLException {
     AtomicReference<String> state = new AtomicReference<>();
     W3CTraceContextPropagator.getInstance()
         .inject(
@@ -46,17 +45,13 @@ public static void propagateContext(Connection connection) throws SQLException {
     if (traceparent == null && existingTraceparent != null) {
       // we need to clear existing tracing state from the connection
       connectionState.set(connection, null);
-      setContextInfo(connection, new byte[0]);
+      setContext(connection, null);
     } else if (!Objects.equals(traceparent, existingTraceparent)) {
       connectionState.set(connection, traceparent);
-      setContextInfo(connection, traceparent.getBytes(StandardCharsets.UTF_8));
+      setContext(connection, traceparent);
     }
   }
 
-  private static void setContextInfo(Connection connection, byte[] contextInfo)
-      throws SQLException {
-    PreparedStatement statement = connection.prepareStatement("set context_info ?");
-    statement.setBytes(1, contextInfo);
-    statement.executeUpdate();
-  }
+  protected abstract void setContext(Connection connection, @Nullable String contextInfo)
+      throws SQLException;
 }

instrumentation/jdbc/src/main/java/com/splunk/opentelemetry/instrumentation/jdbc/oracle/OracleContextPropagator.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.oracle;
+
+import com.splunk.opentelemetry.instrumentation.jdbc.AbstractDbContextPropagator;
+import java.sql.Connection;
+import java.sql.SQLException;
+import javax.annotation.Nullable;
+
+public final class OracleContextPropagator extends AbstractDbContextPropagator {
+  public static final OracleContextPropagator INSTANCE = new OracleContextPropagator();
+
+  private OracleContextPropagator() {}
+
+  @Override
+  protected void setContext(Connection connection, @Nullable String contextInfo)
+      throws SQLException {
+    connection.setClientInfo("OCSID.ACTION", contextInfo);
+  }
+}

instrumentation/jdbc/src/main/java/com/splunk/opentelemetry/instrumentation/jdbc/oracle/OracleInstrumentationModule.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.oracle;
+
+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 OracleInstrumentationModule extends InstrumentationModule {
+
+  public OracleInstrumentationModule() {
+    super("splunk-jdbc", "splunk-jdbc-oracle");
+  }
+
+  @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 OracleStatementInstrumentation());
+  }
+
+  @Override
+  public boolean isHelperClass(String className) {
+    return className.startsWith("com.splunk.opentelemetry.instrumentation");
+  }
+}

instrumentation/jdbc/src/main/java/com/splunk/opentelemetry/instrumentation/jdbc/oracle/OracleStatementInstrumentation.java

@@ -0,0 +1,71 @@
+/*
+ * 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.oracle;
+
+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 OracleStatementInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return namedOneOf(
+        "oracle.jdbc.driver.OraclePreparedStatementWrapper",
+        "oracle.jdbc.driver.OracleStatementWrapper");
+  }
+
+  @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(OracleContextPropagator.class);
+      if (callDepth.getAndIncrement() > 0) {
+        return;
+      }
+
+      OracleContextPropagator.INSTANCE.propagateContext(statement.getConnection());
+    }
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void onExit(@Advice.Local("splunkCallDepth") CallDepth callDepth) {
+      callDepth.decrementAndGet();
+    }
+  }
+}

instrumentation/jdbc/src/main/java/com/splunk/opentelemetry/instrumentation/jdbc/sqlserver/SqlServerContextPropagator.java

@@ -0,0 +1,40 @@
+/*
+ * 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.sqlserver;
+
+import com.splunk.opentelemetry.instrumentation.jdbc.AbstractDbContextPropagator;
+import java.nio.charset.StandardCharsets;
+import java.sql.Connection;
+import java.sql.PreparedStatement;
+import java.sql.SQLException;
+import javax.annotation.Nullable;
+
+public final class SqlServerContextPropagator extends AbstractDbContextPropagator {
+  public static final SqlServerContextPropagator INSTANCE = new SqlServerContextPropagator();
+
+  private SqlServerContextPropagator() {}
+
+  @Override
+  protected void setContext(Connection connection, @Nullable String contextInfo)
+      throws SQLException {
+    byte[] contextBytes =
+        contextInfo == null ? new byte[0] : contextInfo.getBytes(StandardCharsets.UTF_8);
+    PreparedStatement statement = connection.prepareStatement("set context_info ?");
+    statement.setBytes(1, contextBytes);
+    statement.executeUpdate();
+  }
+}

instrumentation/jdbc/src/main/java/com/splunk/opentelemetry/instrumentation/jdbc/sqlserver/SqlServerStatementInstrumentation.java

@@ -55,12 +55,12 @@ public static class SetContextAdvice {
     public static void onEnter(
         @Advice.This Statement statement, @Advice.Local("splunkCallDepth") CallDepth callDepth)
         throws Exception {
-      callDepth = CallDepth.forClass(SqlServerUtil.class);
+      callDepth = CallDepth.forClass(SqlServerContextPropagator.class);
       if (callDepth.getAndIncrement() > 0) {
         return;
       }
 
-      SqlServerUtil.propagateContext(statement.getConnection());
+      SqlServerContextPropagator.INSTANCE.propagateContext(statement.getConnection());
     }
 
     @Advice.OnMethodExit(suppress = Throwable.class)