From 23584d8411a90f0db10b7dac2809eeba04748418 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Vandon?= Date: Wed, 5 Jun 2024 19:39:07 +0200 Subject: [PATCH] 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 --- ...BMCompatibleConnectionInstrumentation.java | 21 ++++---- .../test/groovy/DBMInjectionForkedTest.groovy | 50 +++++++++++++++++++ .../test/groovy/test/TestConnection.groovy | 2 +- .../groovy/test/TestPreparedStatement.groovy | 13 ++++- .../src/test/groovy/test/TestStatement.groovy | 11 ++++ 5 files changed, 85 insertions(+), 12 deletions(-) create mode 100644 dd-java-agent/instrumentation/jdbc/src/test/groovy/DBMInjectionForkedTest.groovy diff --git a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/DBMCompatibleConnectionInstrumentation.java b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/DBMCompatibleConnectionInstrumentation.java index 4a1bb04594d..8041f2037ff 100644 --- a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/DBMCompatibleConnectionInstrumentation.java +++ b/dd-java-agent/instrumentation/jdbc/src/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 new file mode 100644 index 00000000000..12ecd5e4b85 --- /dev/null +++ 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 index bd71d220d30..defb43d9792 100644 --- 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 index da319ee9280..6f97d81f87f 100644 --- 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 index 97853f66a27..5a6f4977d52 100644 --- 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,6 +8,7 @@ import java.sql.Statement class TestStatement implements Statement { final Connection connection + public String sql TestStatement(Connection connection) { this.connection = connection @@ -15,11 +16,13 @@ class TestStatement implements Statement { @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 }