Fill additional db.* attributes on DataSource#getConnection() (open-telemetry#8966)

Author: Mateusz Rzeszutek

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

@@ -19,13 +19,14 @@
 import io.opentelemetry.instrumentation.jdbc.internal.JdbcNetAttributesGetter;
 import io.opentelemetry.javaagent.bootstrap.internal.CommonConfig;
 import io.opentelemetry.javaagent.bootstrap.internal.InstrumentationConfig;
+import io.opentelemetry.javaagent.bootstrap.jdbc.DbInfo;
 import javax.sql.DataSource;
 
 public final class JdbcSingletons {
   private static final String INSTRUMENTATION_NAME = "io.opentelemetry.jdbc";
 
   private static final Instrumenter<DbRequest, Void> STATEMENT_INSTRUMENTER;
-  public static final Instrumenter<DataSource, Void> DATASOURCE_INSTRUMENTER =
+  public static final Instrumenter<DataSource, DbInfo> DATASOURCE_INSTRUMENTER =
       createDataSourceInstrumenter(GlobalOpenTelemetry.get());
 
   static {
@@ -56,7 +57,7 @@ public static Instrumenter<DbRequest, Void> statementInstrumenter() {
     return STATEMENT_INSTRUMENTER;
   }
 
-  public static Instrumenter<DataSource, Void> dataSourceInstrumenter() {
+  public static Instrumenter<DataSource, DbInfo> dataSourceInstrumenter() {
     return DATASOURCE_INSTRUMENTER;
   }
 

instrumentation/jdbc/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jdbc/datasource/DataSourceInstrumentation.java

@@ -8,12 +8,16 @@
 import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.implementsInterface;
 import static io.opentelemetry.javaagent.instrumentation.jdbc.JdbcSingletons.dataSourceInstrumenter;
 import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.returns;
 
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
+import io.opentelemetry.instrumentation.jdbc.internal.JdbcUtils;
 import io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge;
+import io.opentelemetry.javaagent.bootstrap.jdbc.DbInfo;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import java.sql.Connection;
 import javax.sql.DataSource;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
@@ -29,7 +33,8 @@ public ElementMatcher<TypeDescription> typeMatcher() {
   @Override
   public void transform(TypeTransformer transformer) {
     transformer.applyAdviceToMethod(
-        named("getConnection"), DataSourceInstrumentation.class.getName() + "$GetConnectionAdvice");
+        named("getConnection").and(returns(named("java.sql.Connection"))),
+        DataSourceInstrumentation.class.getName() + "$GetConnectionAdvice");
   }
 
   @SuppressWarnings("unused")
@@ -47,21 +52,29 @@ public static void start(
         return;
       }
 
-      context = dataSourceInstrumenter().start(parentContext, ds);
-      scope = context.makeCurrent();
+      if (dataSourceInstrumenter().shouldStart(parentContext, ds)) {
+        context = dataSourceInstrumenter().start(parentContext, ds);
+        scope = context.makeCurrent();
+      }
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void stopSpan(
         @Advice.This DataSource ds,
+        @Advice.Return Connection connection,
+        @Advice.Thrown Throwable throwable,
         @Advice.Local("otelContext") Context context,
-        @Advice.Local("otelScope") Scope scope,
-        @Advice.Thrown Throwable throwable) {
+        @Advice.Local("otelScope") Scope scope) {
       if (scope == null) {
         return;
       }
       scope.close();
-      dataSourceInstrumenter().end(context, ds, null, throwable);
+      DbInfo dbInfo = null;
+      Connection realConnection = JdbcUtils.unwrapConnection(connection);
+      if (realConnection != null) {
+        dbInfo = JdbcUtils.extractDbInfo(realConnection);
+      }
+      dataSourceInstrumenter().end(context, ds, dbInfo, throwable);
     }
   }
 }

instrumentation/jdbc/javaagent/src/test/groovy/JdbcInstrumentationTest.groovy

@@ -580,6 +580,10 @@ class JdbcInstrumentationTest extends AgentInstrumentationSpecification {
           attributes {
             "$SemanticAttributes.CODE_NAMESPACE" datasource.class.name
             "$SemanticAttributes.CODE_FUNCTION" "getConnection"
+            "$SemanticAttributes.DB_SYSTEM" system
+            "$SemanticAttributes.DB_USER" { user == null | user == it }
+            "$SemanticAttributes.DB_NAME" "jdbcunittest"
+            "$SemanticAttributes.DB_CONNECTION_STRING" connectionString
           }
         }
         if (recursive) {
@@ -590,20 +594,24 @@ class JdbcInstrumentationTest extends AgentInstrumentationSpecification {
             attributes {
               "$SemanticAttributes.CODE_NAMESPACE" datasource.class.name
               "$SemanticAttributes.CODE_FUNCTION" "getConnection"
+              "$SemanticAttributes.DB_SYSTEM" system
+              "$SemanticAttributes.DB_USER" { user == null | user == it }
+              "$SemanticAttributes.DB_NAME" "jdbcunittest"
+              "$SemanticAttributes.DB_CONNECTION_STRING" connectionString
             }
           }
         }
       }
     }
 
     where:
-    datasource                               | init
-    new JdbcDataSource()                     | { ds -> ds.setURL(jdbcUrls.get("h2")) }
-    new EmbeddedDataSource()                 | { ds -> ds.jdbcurl = jdbcUrls.get("derby") }
-    cpDatasources.get("hikari").get("h2")    | null
-    cpDatasources.get("hikari").get("derby") | null
-    cpDatasources.get("c3p0").get("h2")      | null
-    cpDatasources.get("c3p0").get("derby")   | null
+    datasource                               | init                                         | system  | user  | connectionString
+    new JdbcDataSource()                     | { ds -> ds.setURL(jdbcUrls.get("h2")) }      | "h2"    | null  | "h2:mem:"
+    new EmbeddedDataSource()                 | { ds -> ds.jdbcurl = jdbcUrls.get("derby") } | "derby" | "APP" | "derby:memory:"
+    cpDatasources.get("hikari").get("h2")    | null                                         | "h2"    | null  | "h2:mem:"
+    cpDatasources.get("hikari").get("derby") | null                                         | "derby" | "APP" | "derby:memory:"
+    cpDatasources.get("c3p0").get("h2")      | null                                         | "h2"    | null  | "h2:mem:"
+    cpDatasources.get("c3p0").get("derby")   | null                                         | "derby" | "APP" | "derby:memory:"
 
     // Tomcat's pool doesn't work because the getConnection method is
     // implemented in a parent class that doesn't implement DataSource

instrumentation/jdbc/library/src/main/java/io/opentelemetry/instrumentation/jdbc/datasource/OpenTelemetryDataSource.java

@@ -45,7 +45,7 @@
 public class OpenTelemetryDataSource implements DataSource, AutoCloseable {
 
   private final DataSource delegate;
-  private final Instrumenter<DataSource, Void> dataSourceInstrumenter;
+  private final Instrumenter<DataSource, DbInfo> dataSourceInstrumenter;
   private final Instrumenter<DbRequest, Void> statementInstrumenter;
   private volatile DbInfo cachedDbInfo;
 
@@ -128,25 +128,29 @@ public void close() throws Exception {
     }
   }
 
-  private <T, E extends SQLException> T wrapCall(ThrowingSupplier<T, E> callable) throws E {
+  private Connection wrapCall(ThrowingSupplier<Connection, SQLException> getConnection)
+      throws SQLException {
     Context parentContext = Context.current();
 
-    if (!Span.fromContext(parentContext).getSpanContext().isValid()) {
+    if (!Span.fromContext(parentContext).getSpanContext().isValid()
+        || !dataSourceInstrumenter.shouldStart(parentContext, delegate)) {
       // this instrumentation is already very noisy, and calls to getConnection outside of an
       // existing trace do not tend to be very interesting
-      return callable.call();
+      return getConnection.call();
     }
 
-    Context context = this.dataSourceInstrumenter.start(parentContext, delegate);
-    T result;
+    Context context = dataSourceInstrumenter.start(parentContext, delegate);
+    Connection connection = null;
+    Throwable error = null;
     try (Scope ignored = context.makeCurrent()) {
-      result = callable.call();
+      connection = getConnection.call();
+      return connection;
     } catch (Throwable t) {
-      this.dataSourceInstrumenter.end(context, delegate, null, t);
+      error = t;
       throw t;
+    } finally {
+      dataSourceInstrumenter.end(context, delegate, getDbInfo(connection), error);
     }
-    this.dataSourceInstrumenter.end(context, delegate, null, null);
-    return result;
   }
 
   private DbInfo getDbInfo(Connection connection) {

instrumentation/jdbc/library/src/main/java/io/opentelemetry/instrumentation/jdbc/internal/DataSourceCodeAttributesGetter.java

@@ -8,7 +8,8 @@
 import io.opentelemetry.instrumentation.api.instrumenter.code.CodeAttributesGetter;
 import javax.sql.DataSource;
 
-final class DataSourceCodeAttributesGetter implements CodeAttributesGetter<DataSource> {
+enum DataSourceCodeAttributesGetter implements CodeAttributesGetter<DataSource> {
+  INSTANCE;
 
   @Override
   public Class<?> getCodeClass(DataSource dataSource) {

instrumentation/jdbc/library/src/main/java/io/opentelemetry/instrumentation/jdbc/internal/DataSourceDbAttributesExtractor.java

@@ -0,0 +1,44 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.jdbc.internal;
+
+import static io.opentelemetry.instrumentation.api.internal.AttributesExtractorUtil.internalSet;
+
+import io.opentelemetry.api.common.AttributesBuilder;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.api.instrumenter.AttributesExtractor;
+import io.opentelemetry.instrumentation.jdbc.internal.dbinfo.DbInfo;
+import io.opentelemetry.semconv.trace.attributes.SemanticAttributes;
+import javax.annotation.Nullable;
+import javax.sql.DataSource;
+
+enum DataSourceDbAttributesExtractor implements AttributesExtractor<DataSource, DbInfo> {
+  INSTANCE;
+
+  @Override
+  public void onStart(AttributesBuilder attributes, Context parentContext, DataSource dataSource) {}
+
+  @Override
+  public void onEnd(
+      AttributesBuilder attributes,
+      Context context,
+      DataSource dataSource,
+      @Nullable DbInfo dbInfo,
+      @Nullable Throwable error) {
+    if (dbInfo == null) {
+      return;
+    }
+    internalSet(attributes, SemanticAttributes.DB_SYSTEM, dbInfo.getSystem());
+    internalSet(attributes, SemanticAttributes.DB_USER, dbInfo.getUser());
+    internalSet(attributes, SemanticAttributes.DB_NAME, getName(dbInfo));
+    internalSet(attributes, SemanticAttributes.DB_CONNECTION_STRING, dbInfo.getShortUrl());
+  }
+
+  private static String getName(DbInfo dbInfo) {
+    String name = dbInfo.getName();
+    return name == null ? dbInfo.getDb() : name;
+  }
+}

instrumentation/jdbc/library/src/main/java/io/opentelemetry/instrumentation/jdbc/internal/DataSourceInstrumenterFactory.java

@@ -9,22 +9,24 @@
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
 import io.opentelemetry.instrumentation.api.instrumenter.code.CodeAttributesExtractor;
 import io.opentelemetry.instrumentation.api.instrumenter.code.CodeSpanNameExtractor;
+import io.opentelemetry.instrumentation.jdbc.internal.dbinfo.DbInfo;
 import javax.sql.DataSource;
 
 /**
  * This class is internal and is hence not for public use. Its APIs are unstable and can change at
  * any time.
  */
 public final class DataSourceInstrumenterFactory {
+
   private static final String INSTRUMENTATION_NAME = "io.opentelemetry.jdbc";
-  private static final DataSourceCodeAttributesGetter codeAttributesGetter =
-      new DataSourceCodeAttributesGetter();
 
-  public static Instrumenter<DataSource, Void> createDataSourceInstrumenter(
+  public static Instrumenter<DataSource, DbInfo> createDataSourceInstrumenter(
       OpenTelemetry openTelemetry) {
-    return Instrumenter.<DataSource, Void>builder(
-            openTelemetry, INSTRUMENTATION_NAME, CodeSpanNameExtractor.create(codeAttributesGetter))
-        .addAttributesExtractor(CodeAttributesExtractor.create(codeAttributesGetter))
+    DataSourceCodeAttributesGetter getter = DataSourceCodeAttributesGetter.INSTANCE;
+    return Instrumenter.<DataSource, DbInfo>builder(
+            openTelemetry, INSTRUMENTATION_NAME, CodeSpanNameExtractor.create(getter))
+        .addAttributesExtractor(CodeAttributesExtractor.create(getter))
+        .addAttributesExtractor(DataSourceDbAttributesExtractor.INSTANCE)
         .buildInstrumenter();
   }
 

instrumentation/jdbc/library/src/main/java/io/opentelemetry/instrumentation/jdbc/internal/JdbcUtils.java

@@ -26,12 +26,19 @@ public final class JdbcUtils {
 
   @Nullable private static Field c3poField = null;
 
-  /** Returns the unwrapped connection or null if exception was thrown. */
   public static Connection connectionFromStatement(Statement statement) {
-    Connection connection;
     try {
-      connection = statement.getConnection();
+      return unwrapConnection(statement.getConnection());
+    } catch (Throwable e) {
+      // Had some problem getting the connection.
+      logger.log(FINE, "Could not get connection from a statement", e);
+      return null;
+    }
+  }
 
+  /** Returns the unwrapped connection or null if exception was thrown. */
+  public static Connection unwrapConnection(Connection connection) {
+    try {
       if (c3poField != null) {
         if (connection.getClass().getName().equals("com.mchange.v2.c3p0.impl.NewProxyConnection")) {
           return (Connection) c3poField.get(connection);
@@ -64,7 +71,7 @@ public static Connection connectionFromStatement(Statement statement) {
       }
     } catch (Throwable e) {
       // Had some problem getting the connection.
-      logger.log(FINE, "Could not get connection for StatementAdvice", e);
+      logger.log(FINE, "Could not unwrap connection", e);
       return null;
     }
     return connection;