From c27a0a7a9edadc3bbdaa02b977a7e6d70c48674d Mon Sep 17 00:00:00 2001
From: Nenad Noveljic <18366081+nenadnoveljic@users.noreply.github.com>
Date: Tue, 12 Nov 2024 17:44:03 +0100
Subject: [PATCH] trace injection for prepared statements

---
 ...tractPreparedStatementInstrumentation.java |  5 +++
 .../instrumentation/jdbc/JDBCDecorator.java   | 39 +++++++++++++++++++
 .../jdbc/StatementInstrumentation.java        |  1 +
 3 files changed, 45 insertions(+)

diff --git a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/AbstractPreparedStatementInstrumentation.java b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/AbstractPreparedStatementInstrumentation.java
index 733d40c87939..e430b712364d 100644
--- a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/AbstractPreparedStatementInstrumentation.java
+++ b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/AbstractPreparedStatementInstrumentation.java
@@ -88,6 +88,11 @@ public static AgentScope onEnter(@Advice.This final Statement statement) {
           span.setTag(DBM_TRACE_INJECTED, true);
         } else {
           span = startSpan(DATABASE_QUERY);
+          // TODO: call setApplicationName only if postgres && inject comment && inject trace
+          // context && trace_prepared_statements
+          // if (DECORATE.isPostgres(dbInfo)) {
+          DECORATE.setApplicationName(span, connection);
+          // }
         }
         DECORATE.afterStart(span);
         DECORATE.onConnection(span, dbInfo);
diff --git a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/JDBCDecorator.java b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/JDBCDecorator.java
index d3dea36abe87..8108f6508309 100644
--- a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/JDBCDecorator.java
+++ b/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;
@@ -248,6 +249,10 @@ public String traceParent(AgentSpan span, int samplingPriority) {
     return sb.toString();
   }
 
+  public boolean isPostgres(final DBInfo dbInfo) {
+    return "postgres".equals(dbInfo.getType());
+  }
+
   public boolean isSqlServer(final DBInfo dbInfo) {
     return "sqlserver".equals(dbInfo.getType());
   }
@@ -312,6 +317,40 @@ public long setContextInfo(Connection connection, DBInfo dbInfo) {
     return spanID;
   }
 
+  // TODO: add description
+  public void setApplicationName(AgentSpan span, Connection connection) {
+    // TODO: measure time spent in instrumentation and add it as a tag in span
+    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);
+    }
+    span.setTag(DBM_TRACE_INJECTED, true);
+  }
+
   @Override
   protected void postProcessServiceAndOperationName(
       AgentSpan span, DatabaseClientDecorator.NamingEntry namingEntry) {
diff --git a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java
index 75f433f255c2..832913af9dfb 100644
--- a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java
+++ b/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) {