Add trace injection for prepared statements in Postgres

dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/AbstractPreparedStatementInstrumentation.java

@@ -5,6 +5,7 @@
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
 import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.DBM_TRACE_INJECTED;
 import static datadog.trace.instrumentation.jdbc.JDBCDecorator.DATABASE_QUERY;
+import static datadog.trace.instrumentation.jdbc.JDBCDecorator.DBM_TRACE_PREPARED_STATEMENTS;
 import static datadog.trace.instrumentation.jdbc.JDBCDecorator.DECORATE;
 import static datadog.trace.instrumentation.jdbc.JDBCDecorator.INJECT_COMMENT;
 import static datadog.trace.instrumentation.jdbc.JDBCDecorator.logMissingQueryInfo;
@@ -80,12 +81,19 @@ public static AgentScope onEnter(@Advice.This final Statement statement) {
                 connection, InstrumentationContext.get(Connection.class, DBInfo.class));
         final boolean injectTraceContext = DECORATE.shouldInjectTraceContext(dbInfo);
 
-        if (INJECT_COMMENT && injectTraceContext && DECORATE.isSqlServer(dbInfo)) {
-          // The span ID is pre-determined so that we can reference it when setting the context
-          final long spanID = DECORATE.setContextInfo(connection, dbInfo);
-          // we then force that pre-determined span ID for the span covering the actual query
-          span = AgentTracer.get().buildSpan(DATABASE_QUERY).withSpanId(spanID).start();
-          span.setTag(DBM_TRACE_INJECTED, true);
+        if (INJECT_COMMENT && injectTraceContext) {
+          if (DECORATE.isSqlServer(dbInfo)) {
+            // The span ID is pre-determined so that we can reference it when setting the context
+            final long spanID = DECORATE.setContextInfo(connection, dbInfo);
+            // we then force that pre-determined span ID for the span covering the actual query
+            span = AgentTracer.get().buildSpan(DATABASE_QUERY).withSpanId(spanID).start();
+            span.setTag(DBM_TRACE_INJECTED, true);
+          } else if (DECORATE.isPostgres(dbInfo) && DBM_TRACE_PREPARED_STATEMENTS) {
+            span = startSpan(DATABASE_QUERY);
+            DECORATE.setApplicationName(span, connection);
+          } else {
+            span = startSpan(DATABASE_QUERY);
+          }
         } else {
           span = startSpan(DATABASE_QUERY);
         }

dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/JDBCDecorator.java

@@ -1,6 +1,7 @@
 package datadog.trace.instrumentation.jdbc;
 
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
+import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.DBM_TRACE_INJECTED;
 import static datadog.trace.bootstrap.instrumentation.api.Tags.*;
 
 import datadog.trace.api.Config;
@@ -53,6 +54,8 @@ public class JDBCDecorator extends DatabaseClientDecorator<DBInfo> {
           || DBM_PROPAGATION_MODE.equals(DBM_PROPAGATION_MODE_STATIC);
   private static final boolean INJECT_TRACE_CONTEXT =
       DBM_PROPAGATION_MODE.equals(DBM_PROPAGATION_MODE_FULL);
+  public static final boolean DBM_TRACE_PREPARED_STATEMENTS =
+      Config.get().isDBMTracePreparedStatements();
 
   private volatile boolean warnedAboutDBMPropagationMode = false; // to log a warning only once
 
@@ -248,6 +251,10 @@ public String traceParent(AgentSpan span, int samplingPriority) {
     return sb.toString();
   }
 
+  public boolean isPostgres(final DBInfo dbInfo) {
+    return dbInfo.getType().startsWith("postgres");
+  }
+
   public boolean isSqlServer(final DBInfo dbInfo) {
     return "sqlserver".equals(dbInfo.getType());
   }
@@ -312,6 +319,44 @@ public long setContextInfo(Connection connection, DBInfo dbInfo) {
     return spanID;
   }
 
+  // TODO: add description
+  public void setApplicationName(AgentSpan span, Connection connection) {
+    final long startTime = System.currentTimeMillis();
+    try {
+
+      Integer priority = span.forceSamplingDecision();
+      if (priority == null) {
+        return;
+      }
+      final String traceParent = DECORATE.traceParent(span, priority);
+      final String traceContext = "_DD_" + traceParent;
+
+      // SET doesn't work with parameters
+      StringBuilder sql = new StringBuilder();
+      sql.append("SET application_name = '");
+      sql.append(traceContext);
+      sql.append("';");
+
+      try (Statement statement = connection.createStatement()) {
+        statement.execute(sql.toString());
+      } catch (SQLException e) {
+        throw e;
+      }
+    } catch (Exception e) {
+      log.debug(
+          "Failed to set extra DBM data in application_name for trace {}. "
+              + "To disable this behavior, set trace_prepared_statements to 'false'. "
+              + "See https://docs.datadoghq.com/database_monitoring/connect_dbm_and_apm/ for more info.{}",
+          span.getTraceId().toHexString(),
+          e);
+      DECORATE.onError(span, e);
+    } finally {
+      span.setTag(DBM_TRACE_INJECTED, true);
+      final long elapsed = System.currentTimeMillis() - startTime;
+      span.setTag("dd.instrumentation.time_ms", elapsed);
+    }
+  }
+
   @Override
   protected void postProcessServiceAndOperationName(
       AgentSpan span, DatabaseClientDecorator.NamingEntry namingEntry) {

dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java

@@ -103,6 +103,7 @@ public static AgentScope onEnter(
         if (span != null && INJECT_COMMENT) {
           String traceParent = null;
 
+          // TODO: don't propagate trace info in the comments for postgres if full+ mode is defined
           if (injectTraceContext) {
             Integer priority = span.forceSamplingDecision();
             if (priority != null) {

dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java

@@ -69,6 +69,7 @@ public final class ConfigDefaults {
   static final boolean DEFAULT_DB_CLIENT_HOST_SPLIT_BY_INSTANCE_TYPE_SUFFIX = false;
   static final boolean DEFAULT_DB_CLIENT_HOST_SPLIT_BY_HOST = false;
   static final String DEFAULT_DB_DBM_PROPAGATION_MODE_MODE = "disabled";
+  static final boolean DEFAULT_DB_DBM_TRACE_PREPARED_STATEMENTS = false;
   static final int DEFAULT_SCOPE_DEPTH_LIMIT = 100;
   static final int DEFAULT_SCOPE_ITERATION_KEEP_ALIVE = 30; // in seconds
   static final int DEFAULT_PARTIAL_FLUSH_MIN_SPANS = 1000;

dd-trace-api/src/main/java/datadog/trace/api/config/TraceInstrumentationConfig.java

@@ -68,6 +68,8 @@ public final class TraceInstrumentationConfig {
 
   public static final String DB_DBM_PROPAGATION_MODE_MODE = "dbm.propagation.mode";
 
+  public static final String DB_DBM_TRACE_PREPARED_STATEMENTS = "dbm.trace_prepared_statements";
+
   public static final String JDBC_CONNECTION_CLASS_NAME = "trace.jdbc.connection.class.name";
 
   public static final String HTTP_URL_CONNECTION_CLASS_NAME =

internal-api/src/main/java/datadog/trace/api/Config.java

@@ -373,6 +373,7 @@ public static String getHostName() {
   private final int remoteConfigMaxExtraServices;
 
   private final String DBMPropagationMode;
+  private final boolean DBMTracePreparedStatements;
 
   private final boolean debuggerEnabled;
   private final int debuggerUploadTimeout;
@@ -845,6 +846,10 @@ private Config(final ConfigProvider configProvider, final InstrumenterConfig ins
         configProvider.getString(
             DB_DBM_PROPAGATION_MODE_MODE, DEFAULT_DB_DBM_PROPAGATION_MODE_MODE);
 
+    DBMTracePreparedStatements =
+        configProvider.getBoolean(
+            DB_DBM_TRACE_PREPARED_STATEMENTS, DEFAULT_DB_DBM_TRACE_PREPARED_STATEMENTS);
+
     splitByTags = tryMakeImmutableSet(configProvider.getList(SPLIT_BY_TAGS));
 
     springDataRepositoryInterfaceResourceName =
@@ -3672,6 +3677,10 @@ public boolean isEnabled(
         Collections.singletonList(settingName), "", settingSuffix, defaultEnabled);
   }
 
+  public boolean isDBMTracePreparedStatements() {
+    return DBMTracePreparedStatements;
+  }
+
   public String getDBMPropagationMode() {
     return DBMPropagationMode;
   }