Write base hash to current span when injected in DBM comment

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

@@ -103,6 +103,7 @@ public static AgentScope onEnter(@Advice.This final Statement statement) {
         DECORATE.afterStart(span);
         DECORATE.onConnection(span, dbInfo);
         DECORATE.onPreparedStatement(span, queryInfo);
+        DECORATE.withBaseHash(span);
 
         return activateSpan(span);
       } catch (SQLException e) {

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

@@ -17,6 +17,7 @@
 import datadog.trace.bootstrap.CallDepthThreadLocalMap;
 import datadog.trace.bootstrap.ContextStore;
 import datadog.trace.bootstrap.InstrumentationContext;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.jdbc.DBInfo;
 import datadog.trace.bootstrap.instrumentation.jdbc.DBQueryInfo;
 import java.sql.Connection;
@@ -121,13 +122,13 @@ public static String onEnter(
         return null;
       }
       final String inputSql = sql;
+      final AgentSpan activeSpan = activeSpan();
       final DBInfo dbInfo =
           JDBCDecorator.parseDBInfo(
               connection, InstrumentationContext.get(Connection.class, DBInfo.class));
       String dbService = DECORATE.getDbService(dbInfo);
       if (dbService != null) {
-        dbService =
-            traceConfig(activeSpan()).getServiceMapping().getOrDefault(dbService, dbService);
+        dbService = traceConfig(activeSpan).getServiceMapping().getOrDefault(dbService, dbService);
       }
 
       boolean append =

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

@@ -7,6 +7,7 @@
 import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.INSTRUMENTATION_TIME_MS;
 import static datadog.trace.bootstrap.instrumentation.api.Tags.*;
 
+import datadog.trace.api.BaseHash;
 import datadog.trace.api.Config;
 import datadog.trace.api.DDTraceId;
 import datadog.trace.api.naming.SpanNaming;
@@ -54,6 +55,9 @@ public class JDBCDecorator extends DatabaseClientDecorator<DBInfo> {
   public static final String DD_INSTRUMENTATION_PREFIX = "_DD_";
 
   public static final String DBM_PROPAGATION_MODE = Config.get().getDbmPropagationMode();
+  private static final boolean DBM_INJECT_SQL_BASE_HASH = Config.get().isDbmInjectSqlBaseHash();
+  private static final boolean PROPAGATE_PROCESS_TAGS =
+      Config.get().isExperimentalPropagateProcessTagsEnabled();
   public static final boolean INJECT_COMMENT =
       DBM_PROPAGATION_MODE.equals(DBM_PROPAGATION_MODE_FULL)
           || DBM_PROPAGATION_MODE.equals(DBM_PROPAGATION_MODE_STATIC);
@@ -246,6 +250,16 @@ public AgentSpan onPreparedStatement(AgentSpan span, DBQueryInfo dbQueryInfo) {
     return withQueryInfo(span, dbQueryInfo, JDBC_PREPARED_STATEMENT);
   }
 
+  /**
+   * Sets the base hash tag on the span if DBM hash injection is enabled. This is necessary so that
+   * the span (tags) and the query can be matched in the backend.
+   */
+  public void withBaseHash(AgentSpan span) {
+    if (INJECT_COMMENT && DBM_INJECT_SQL_BASE_HASH && PROPAGATE_PROCESS_TAGS) {
+      span.setTag(Tags.BASE_HASH, BaseHash.getBaseHashStr());
+    }
+  }
+
   private AgentSpan withQueryInfo(AgentSpan span, DBQueryInfo info, CharSequence component) {
     if (null != info) {
       span.setResourceName(info.getSql());

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

@@ -159,6 +159,7 @@ public static AgentScope onEnter(
                   appendComment);
         }
         DECORATE.onStatement(span, copy);
+        DECORATE.withBaseHash(span);
         return activateSpan(span);
       } catch (SQLException e) {
         // if we can't get the connection for any reason

dd-java-agent/instrumentation/jdbc/src/test/groovy/DBMInjectionForkedTest.groovy

@@ -1,6 +1,11 @@
 import datadog.trace.agent.test.InstrumentationSpecification
+import datadog.trace.api.BaseHash
+import datadog.trace.api.DDSpanTypes
+import datadog.trace.api.ProcessTags
+import datadog.trace.api.config.GeneralConfig
 import datadog.trace.api.config.TraceInstrumentationConfig
 import datadog.trace.api.config.TracerConfig
+import datadog.trace.bootstrap.instrumentation.api.Tags
 import test.TestConnection
 import test.TestPreparedStatement
 import test.TestStatement
@@ -77,3 +82,39 @@ class DBMAppendInjectionForkedTest extends InjectionTest {
     assert statement.sql == "${query} /*${serviceInjection},traceparent='00-00000000000000000000000000000004-0000000000000003-01'*/"
   }
 }
+
+class DBMBaseHashInjectionForkedTest extends InjectionTest {
+
+  @Override
+  void configurePreAgent() {
+    super.configurePreAgent()
+    injectSysConfig(TraceInstrumentationConfig.DB_DBM_INJECT_SQL_BASEHASH, "true")
+    injectSysConfig(GeneralConfig.EXPERIMENTAL_PROPAGATE_PROCESS_TAGS_ENABLED, "true")
+  }
+
+  def "base hash tag is set on span and matches the one in the SQL comment"() {
+    setup:
+    ProcessTags.reset()
+    BaseHash.updateBaseHash(123456789L)
+    def connection = new TestConnection(false)
+
+    when:
+    def statement = connection.prepareStatement(query) as TestPreparedStatement
+    statement.execute()
+
+    then:
+    // the same hash should be in the SQL comment
+    assert statement.sql.contains("ddsh='123456789'")
+    // and the base hash should be set on the jdbc span
+    assertTraces(1) {
+      trace(1) {
+        span {
+          spanType DDSpanTypes.SQL
+          tags(false) {
+            "$Tags.BASE_HASH" "123456789"
+          }
+        }
+      }
+    }
+  }
+}

internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/Tags.java

@@ -49,6 +49,7 @@ public class Tags {
   public static final String DB_SCHEMA = "db.schema";
   public static final String MESSAGE_BUS_DESTINATION = "message_bus.destination";
   public static final String DB_POOL_NAME = "db.pool.name";
+  public static final String BASE_HASH = "_dd.propagated_hash";
 
   public static final String TEST_SESSION_NAME = "test_session.name";
   public static final String TEST_MODULE = "test.module";