apply configured service name mapping to DBM-injected dddbs (#7064)

* apply configured service name mapping to DBM-injected dddbs * fix nullref that broke some tests * add test * simplify test * clean conf after tests * make DBM test a forked test * apply suggestions
DataDog · Jun 5, 2024 · 23584d8 · 23584d8
1 parent 6860183
commit 23584d8
Show file tree

Hide file tree

Showing 5 changed files with 85 additions and 12 deletions.
diff --git a/.../main/java/datadog/trace/instrumentation/jdbc/DBMCompatibleConnectionInstrumentation.java b/.../main/java/datadog/trace/instrumentation/jdbc/DBMCompatibleConnectionInstrumentation.java
@@ -3,6 +3,8 @@
 import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.hasInterface;
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.nameStartsWith;
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.traceConfig;
 import static datadog.trace.instrumentation.jdbc.JDBCDecorator.DECORATE;
 import static datadog.trace.instrumentation.jdbc.JDBCDecorator.logQueryInfoInjection;
 import static net.bytebuddy.matcher.ElementMatchers.returns;
@@ -68,6 +70,8 @@ public DBMCompatibleConnectionInstrumentation() {
     "org.mariadb.jdbc.Connection",
     // aws-mysql-jdbc
     "software.aws.rds.jdbc.mysql.shading.com.mysql.cj.jdbc.ConnectionImpl",
+    // for testing purposes
+    "test.TestConnection"
   };
 
   @Override
@@ -116,22 +120,19 @@ public static String onEnter(
         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);
+        }
         if (dbInfo.getType().equals("sqlserver")) {
           sql =
               SQLCommenter.append(
-                  sql,
-                  DECORATE.getDbService(dbInfo),
-                  dbInfo.getType(),
-                  dbInfo.getHost(),
-                  dbInfo.getDb());
+                  sql, dbService, dbInfo.getType(), dbInfo.getHost(), dbInfo.getDb());
         } else {
           sql =
               SQLCommenter.prepend(
-                  sql,
-                  DECORATE.getDbService(dbInfo),
-                  dbInfo.getType(),
-                  dbInfo.getHost(),
-                  dbInfo.getDb());
+                  sql, dbService, dbInfo.getType(), dbInfo.getHost(), dbInfo.getDb());
         }
         return inputSql;
       }

diff --git a/dd-java-agent/instrumentation/jdbc/src/test/groovy/DBMInjectionForkedTest.groovy b/dd-java-agent/instrumentation/jdbc/src/test/groovy/DBMInjectionForkedTest.groovy
@@ -0,0 +1,50 @@
+import datadog.trace.agent.test.AgentTestRunner
+import datadog.trace.api.config.TraceInstrumentationConfig
+import datadog.trace.api.config.TracerConfig
+import test.TestConnection
+import test.TestPreparedStatement
+import test.TestStatement
+
+class DBMInjectionForkedTest extends AgentTestRunner {
+
+  @Override
+  void configurePreAgent() {
+    super.configurePreAgent()
+
+    injectSysConfig(TraceInstrumentationConfig.DB_DBM_PROPAGATION_MODE_MODE, "full")
+    // to check that we use the remapped service name in dddbs
+    injectSysConfig("service.name", "my_service_name")
+    injectSysConfig(TracerConfig.SERVICE_MAPPING, "testdb:remapped_testdb")
+    injectSysConfig("dd.trace.jdbc.prepared.statement.class.name", "test.TestPreparedStatement")
+    injectSysConfig("dd.trace.jdbc.connection.class.name", "test.TestConnection")
+  }
+
+  static query = "SELECT 1"
+  static serviceInjection = "ddps='my_service_name',dddbs='remapped_testdb'"
+  static fullInjection = serviceInjection + ",traceparent='00-00000000000000000000000000000004-0000000000000003-01'"
+
+  def "prepared stmt"() {
+    setup:
+    def connection = new TestConnection(false)
+
+    when:
+    def statement = connection.prepareStatement(query) as TestPreparedStatement
+    statement.execute()
+
+    then:
+    // even in full propagation mode, we cannot inject trace info in prepared statements
+    assert statement.sql == "/*${serviceInjection}*/ ${query}"
+  }
+
+  def "single query"() {
+    setup:
+    def connection = new TestConnection(false)
+
+    when:
+    def statement = connection.createStatement() as TestStatement
+    statement.executeQuery(query)
+
+    then:
+    assert statement.sql == "/*${fullInjection}*/ ${query}"
+  }
+}
diff --git a/dd-java-agent/instrumentation/jdbc/src/test/groovy/test/TestConnection.groovy b/dd-java-agent/instrumentation/jdbc/src/test/groovy/test/TestConnection.groovy
@@ -36,7 +36,7 @@ class TestConnection implements Connection {
 
   @Override
   PreparedStatement prepareStatement(String sql) throws SQLException {
-    return new TestPreparedStatement(this)
+    return new TestPreparedStatement(this, sql)
   }
 
   @Override

diff --git a/dd-java-agent/instrumentation/jdbc/src/test/groovy/test/TestPreparedStatement.groovy b/dd-java-agent/instrumentation/jdbc/src/test/groovy/test/TestPreparedStatement.groovy
@@ -20,18 +20,22 @@ import java.sql.Timestamp
 
 class TestPreparedStatement implements PreparedStatement {
   private Connection connection
+  public String sql
 
-  TestPreparedStatement(Connection connection) {
+  TestPreparedStatement(Connection connection, String sql = null) {
     this.connection = connection
+    this.sql = sql
   }
 
   @Override
   ResultSet executeQuery(String sql) throws SQLException {
+    this.sql = sql
     return null
   }
 
   @Override
   int executeUpdate(String sql) throws SQLException {
+    this.sql = sql
     return 0
   }
 
@@ -87,6 +91,7 @@ class TestPreparedStatement implements PreparedStatement {
 
   @Override
   boolean execute(String sql) throws SQLException {
+    this.sql = sql
     return false
   }
 
@@ -163,31 +168,37 @@ class TestPreparedStatement implements PreparedStatement {
 
   @Override
   int executeUpdate(String sql, int autoGeneratedKeys) throws SQLException {
+    this.sql = sql
     return 0
   }
 
   @Override
   int executeUpdate(String sql, int[] columnIndexes) throws SQLException {
+    this.sql = sql
     return 0
   }
 
   @Override
   int executeUpdate(String sql, String[] columnNames) throws SQLException {
+    this.sql = sql
     return 0
   }
 
   @Override
   boolean execute(String sql, int autoGeneratedKeys) throws SQLException {
+    this.sql = sql
     return false
   }
 
   @Override
   boolean execute(String sql, int[] columnIndexes) throws SQLException {
+    this.sql = sql
     return false
   }
 
   @Override
   boolean execute(String sql, String[] columnNames) throws SQLException {
+    this.sql = sql
     return false
   }
 

diff --git a/dd-java-agent/instrumentation/jdbc/src/test/groovy/test/TestStatement.groovy b/dd-java-agent/instrumentation/jdbc/src/test/groovy/test/TestStatement.groovy
@@ -8,18 +8,21 @@ import java.sql.Statement
 
 class TestStatement implements Statement {
   final Connection connection
+  public String sql
 
   TestStatement(Connection connection) {
     this.connection = connection
   }
 
   @Override
   ResultSet executeQuery(String sql) throws SQLException {
+    this.sql = sql
     return null
   }
 
   @Override
   int executeUpdate(String sql) throws SQLException {
+    this.sql = sql
     return 0
   }
 
@@ -77,6 +80,7 @@ class TestStatement implements Statement {
 
   @Override
   boolean execute(String sql) throws SQLException {
+    this.sql = sql
     return false
   }
 
@@ -125,6 +129,7 @@ class TestStatement implements Statement {
 
   @Override
   void addBatch(String sql) throws SQLException {
+    this.sql = sql
   }
 
   @Override
@@ -153,31 +158,37 @@ class TestStatement implements Statement {
 
   @Override
   int executeUpdate(String sql, int autoGeneratedKeys) throws SQLException {
+    this.sql = sql
     return 0
   }
 
   @Override
   int executeUpdate(String sql, int[] columnIndexes) throws SQLException {
+    this.sql = sql
     return 0
   }
 
   @Override
   int executeUpdate(String sql, String[] columnNames) throws SQLException {
+    this.sql = sql
     return 0
   }
 
   @Override
   boolean execute(String sql, int autoGeneratedKeys) throws SQLException {
+    this.sql = sql
     return false
   }
 
   @Override
   boolean execute(String sql, int[] columnIndexes) throws SQLException {
+    this.sql = sql
     return false
   }
 
   @Override
   boolean execute(String sql, String[] columnNames) throws SQLException {
+    this.sql = sql
     return false
   }