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,8 @@
 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.InstrumentationTags.INSTRUMENTATION_TIME_MS;
 import static datadog.trace.bootstrap.instrumentation.api.Tags.*;
 
 import datadog.trace.api.Config;
@@ -53,6 +55,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 +252,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 +320,44 @@ public long setContextInfo(Connection connection, DBInfo dbInfo) {
     return spanID;
   }
 
+  /**
+   * Executes `SET application_name` statement on the Postgres DB to set the trace parent in
+   * `pg_stat_activity.application_name`. This is used for prepared statements where it isn't
+   * possible to propagate trace parent with the comment. Downside: makes an additional round trip
+   * to the database.
+   *
+   * @param span The span of the instrumented statement
+   * @param connection The same connection as the one that will be used for the actual statement
+   */
+  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;
+
+      connection.setClientInfo("ApplicationName", traceContext);
+    } catch (Throwable e) {
+      if (log.isDebugEnabled()) {
+        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(INSTRUMENTATION_TIME_MS, elapsed);
+    }
+  }
+
   @Override
   protected void postProcessServiceAndOperationName(
       AgentSpan span, DatabaseClientDecorator.NamingEntry namingEntry) {