parmeter

DataDog · nenadnoveljic · Dec 12, 2024 · Nov 12, 2024 · Nov 14, 2024 · Nov 14, 2024
commit 11adebc69580dabe0b99addcd7b322c7c7f257fd
@@ -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,19 +81,21 @@ 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);
-          // 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);

@@ -54,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
 
@@ -250,7 +252,7 @@ public String traceParent(AgentSpan span, int samplingPriority) {
   }
 
   public boolean isPostgres(final DBInfo dbInfo) {
-    return "postgres".equals(dbInfo.getType());
+    return dbInfo.getType().startsWith("postgres");
   }
 
   public boolean isSqlServer(final DBInfo dbInfo) {

@@ -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;

@@ -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 =

@@ -372,6 +372,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;
@@ -844,6 +845,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 =
@@ -3665,6 +3670,10 @@ public boolean isEnabled(
         Collections.singletonList(settingName), "", settingSuffix, defaultEnabled);
   }
 
+  public boolean isDBMTracePreparedStatements() {
+    return DBMTracePreparedStatements;
+  }
+
   public String getDBMPropagationMode() {
     return DBMPropagationMode;
   }