From d6c47f2b94d7176c1cd591ff122c80bee44a7357 Mon Sep 17 00:00:00 2001 From: Christian Gati Date: Thu, 17 Jul 2025 15:47:35 -0700 Subject: [PATCH 1/5] feat: add SQLCommenter custom fields API with thread-local context Implements a thread-local context system for injecting custom key-value pairs into SQL comments, similar to SLF4J MDC pattern. This allows external users to add custom metadata to SQL traces. Features: - Thread-local context management for custom fields - Scoped execution with withSQLCommentFields() pattern - Public API in dd-trace-api with reflection-based access - Proper URL encoding for SQL comment safety - Comprehensive test coverage for all functionality - Maintains backward compatibility with existing SQLCommenter The implementation includes: - SQLCommenterContext: Thread-local context system - SQLCommenter API: Public interface for external users - Integration with existing SQL comment injection - Exception handling with context restoration - Thread isolation and nested context support --- .../instrumentation/jdbc/SQLCommenter.java | 38 ++- .../jdbc/SQLCommenterContext.java | 153 +++++++++ .../jdbc/StatementInstrumentation.java | 6 +- .../src/test/groovy/SQLCommenterTest.groovy | 182 ++++++++++ .../jdbc/SQLCommenterContextTest.java | 301 +++++++++++++++++ .../java/datadog/trace/api/SQLCommenter.java | 317 ++++++++++++++++++ 6 files changed, 991 insertions(+), 6 deletions(-) create mode 100644 dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/SQLCommenterContext.java create mode 100644 dd-java-agent/instrumentation/jdbc/src/test/java/datadog/trace/instrumentation/jdbc/SQLCommenterContextTest.java create mode 100644 dd-trace-api/src/main/java/datadog/trace/api/SQLCommenter.java diff --git a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/SQLCommenter.java b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/SQLCommenter.java index 075ccbbbcc7..b2d4574fda7 100644 --- a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/SQLCommenter.java +++ b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/SQLCommenter.java @@ -8,6 +8,7 @@ import java.io.UnsupportedEncodingException; import java.net.URLEncoder; import java.nio.charset.StandardCharsets; +import java.util.Map; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -117,7 +118,9 @@ public static String inject( final String parentService = config.getServiceName(); final String env = config.getEnv(); final String version = config.getVersion(); - final int commentSize = capacity(traceParent, parentService, dbService, env, version); + final Map customFields = SQLCommenterContext.getCopyOfContextMap(); + final int commentSize = + capacity(traceParent, parentService, dbService, env, version, customFields); StringBuilder sb = new StringBuilder(sql.length() + commentSize); boolean commentAdded = false; String peerService = peerServiceObj != null ? peerServiceObj.toString() : null; @@ -137,7 +140,8 @@ public static String inject( peerService, env, version, - traceParent); + traceParent, + customFields); sb.append(CLOSE_COMMENT); } else { sb.append(OPEN_COMMENT); @@ -152,7 +156,8 @@ public static String inject( peerService, env, version, - traceParent); + traceParent, + customFields); sb.append(CLOSE_COMMENT); sb.append(SPACE); @@ -226,7 +231,8 @@ protected static boolean toComment( final String peerService, final String env, final String version, - final String traceparent) { + final String traceparent, + final Map customFields) { int emptySize = sb.length(); append(sb, PARENT_SERVICE, parentService, false); @@ -241,6 +247,14 @@ protected static boolean toComment( if (injectTrace) { append(sb, TRACEPARENT, traceparent, sb.length() > emptySize); } + + // Add custom fields from SQLCommenterContext + if (customFields != null) { + for (Map.Entry entry : customFields.entrySet()) { + append(sb, encode(entry.getKey()), entry.getValue(), sb.length() > emptySize); + } + } + return sb.length() > emptySize; } @@ -268,7 +282,8 @@ private static int capacity( final String parentService, final String dbService, final String env, - final String version) { + final String version, + final Map customFields) { int len = INITIAL_CAPACITY; if (null != traceparent) { len += traceparent.length(); @@ -285,6 +300,19 @@ private static int capacity( if (null != version) { len += version.length(); } + + // Add capacity for custom fields + if (customFields != null) { + for (Map.Entry entry : customFields.entrySet()) { + if (entry.getKey() != null && entry.getValue() != null) { + // Account for URL encoding overhead (approximately 3x for worst case) + len += entry.getKey().length() * 3; + len += entry.getValue().length() * 3; + len += 4; // equals, comma, and two quotes + } + } + } + return len; } diff --git a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/SQLCommenterContext.java b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/SQLCommenterContext.java new file mode 100644 index 00000000000..c766880314b --- /dev/null +++ b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/SQLCommenterContext.java @@ -0,0 +1,153 @@ +package datadog.trace.instrumentation.jdbc; + +import java.util.HashMap; +import java.util.Map; + +/** + * Thread-local context system for custom SQL comment fields. This class provides functionality + * similar to SLF4J's MDC but specifically for SQL comment injection. External users can add custom + * key-value pairs that will be included in SQL comments. + */ +public class SQLCommenterContext { + + private static final ThreadLocal> contextHolder = + new ThreadLocal>() { + @Override + protected Map initialValue() { + return new HashMap<>(); + } + }; + + /** + * Gets a copy of the current context map. + * + * @return a copy of the current context map, or null if no context is set + */ + public static Map getCopyOfContextMap() { + Map contextMap = contextHolder.get(); + return contextMap.isEmpty() ? null : new HashMap<>(contextMap); + } + + /** + * Sets the context map for the current thread. + * + * @param contextMap the new context map to set + */ + public static void setContextMap(Map contextMap) { + if (contextMap == null) { + contextHolder.remove(); + } else { + contextHolder.set(new HashMap<>(contextMap)); + } + } + + /** + * Puts a key-value pair into the current context. + * + * @param key the key + * @param value the value (will be converted to string) + */ + public static void put(String key, Object value) { + if (key == null) { + throw new IllegalArgumentException("Key cannot be null"); + } + Map contextMap = contextHolder.get(); + contextMap.put(key, value != null ? value.toString() : null); + } + + /** + * Gets a value from the current context. + * + * @param key the key to look up + * @return the value associated with the key, or null if not found + */ + public static String get(String key) { + return contextHolder.get().get(key); + } + + /** + * Removes a key from the current context. + * + * @param key the key to remove + * @return the previous value associated with the key, or null if not found + */ + public static String remove(String key) { + return contextHolder.get().remove(key); + } + + /** Clears the current context. */ + public static void clear() { + contextHolder.remove(); + } + + /** + * Returns true if the current context is empty. + * + * @return true if the current context has no entries + */ + public static boolean isEmpty() { + return contextHolder.get().isEmpty(); + } + + /** + * Executes a block of code with the given context map, restoring the original context afterward. + * This method is similar to the withLoggingFields pattern from the provided Kotlin code. + * + * @param contextMap the context map to use during execution + * @param block the code block to execute + * @param the return type of the block + * @return the result of executing the block + */ + public static T withSQLCommentFields( + Map contextMap, SQLCommenterBlock block) throws Exception { + Map oldContextMap = getCopyOfContextMap(); + try { + Map newContextMap = + new HashMap<>(oldContextMap != null ? oldContextMap : new HashMap<>()); + if (contextMap != null) { + contextMap.forEach( + (key, value) -> newContextMap.put(key, value != null ? value.toString() : null)); + } + setContextMap(newContextMap); + return block.execute(); + } finally { + setContextMap(oldContextMap); + } + } + + /** + * Executes a block of code with the given key-value pairs, restoring the original context + * afterward. + * + * @param block the code block to execute + * @param keyValuePairs alternating key-value pairs to add to the context + * @param the return type of the block + * @return the result of executing the block + */ + @SuppressWarnings("unchecked") + public static T withSQLCommentFields(SQLCommenterBlock block, Object... keyValuePairs) + throws Exception { + if (keyValuePairs.length % 2 != 0) { + throw new IllegalArgumentException("Key-value pairs must be provided in pairs"); + } + + Map contextMap = new HashMap<>(); + for (int i = 0; i < keyValuePairs.length; i += 2) { + String key = keyValuePairs[i].toString(); + Object value = keyValuePairs[i + 1]; + contextMap.put(key, value); + } + + return withSQLCommentFields(contextMap, block); + } + + /** + * Functional interface for code blocks that can be executed with a SQL comment context. + * + * @param the return type of the block + */ + @FunctionalInterface + public interface SQLCommenterBlock { + T execute() throws Exception; + } +} 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 33b4ff4af31..cca546356da 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 @@ -58,7 +58,11 @@ public Map contextStore() { @Override public String[] helperClassNames() { - return new String[] {packageName + ".JDBCDecorator", packageName + ".SQLCommenter"}; + return new String[] { + packageName + ".JDBCDecorator", + packageName + ".SQLCommenter", + packageName + ".SQLCommenterContext" + }; } // prepend mode will prepend the SQL comment to the raw sql query diff --git a/dd-java-agent/instrumentation/jdbc/src/test/groovy/SQLCommenterTest.groovy b/dd-java-agent/instrumentation/jdbc/src/test/groovy/SQLCommenterTest.groovy index 79afb02602c..16f234c7647 100644 --- a/dd-java-agent/instrumentation/jdbc/src/test/groovy/SQLCommenterTest.groovy +++ b/dd-java-agent/instrumentation/jdbc/src/test/groovy/SQLCommenterTest.groovy @@ -4,6 +4,7 @@ import datadog.trace.bootstrap.instrumentation.api.AgentTracer import datadog.trace.bootstrap.instrumentation.api.Tags import datadog.trace.instrumentation.jdbc.SQLCommenter +import datadog.trace.instrumentation.jdbc.SQLCommenterContext import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace @@ -145,4 +146,185 @@ class SQLCommenterTest extends AgentTestRunner { "SELECT * FROM foo" | "SqlCommenter" | "Test" | "my-service" | "mysql" | "h" | "n" | "TestVersion" | true | true | "00-00000000000000007fffffffffffffff-000000024cb016ea-00" | "testPeer" | "SELECT * FROM foo /*ddps='SqlCommenter',dddbs='my-service',ddh='h',dddb='n',ddprs='testPeer',dde='Test',ddpv='TestVersion',traceparent='00-00000000000000007fffffffffffffff-000000024cb016ea-00'*/" "SELECT * FROM foo" | "SqlCommenter" | "Test" | "my-service" | "postgres" | "h" | "n" | "TestVersion" | true | true | "00-00000000000000007fffffffffffffff-000000024cb016ea-00" | "testPeer" | "SELECT * FROM foo /*ddps='SqlCommenter',dddbs='my-service',ddh='h',dddb='n',ddprs='testPeer',dde='Test',ddpv='TestVersion',traceparent='00-00000000000000007fffffffffffffff-000000024cb016ea-00'*/" } + + def "test SQL comment injection with custom fields"() { + setup: + injectSysConfig("dd.service", "SqlCommenter") + injectSysConfig("dd.env", "Test") + injectSysConfig("dd.version", "TestVersion") + + when: + String sqlWithComment = "" + + // Set up custom fields in context + SQLCommenterContext.put("tenant_id", "abc123") + SQLCommenterContext.put("request_id", "req-456") + SQLCommenterContext.put("user_id", 789) + + sqlWithComment = SQLCommenter.inject( + "SELECT * FROM foo", + "my-service", + "mysql", + "h", + "n", + "00-00000000000000007fffffffffffffff-000000024cb016ea-00", + true, + true + ) + + SQLCommenterContext.clear() + + then: + // Should contain all standard fields plus custom fields + sqlWithComment.contains("ddps='SqlCommenter'") + sqlWithComment.contains("dddbs='my-service'") + sqlWithComment.contains("ddh='h'") + sqlWithComment.contains("dddb='n'") + sqlWithComment.contains("dde='Test'") + sqlWithComment.contains("ddpv='TestVersion'") + sqlWithComment.contains("traceparent='00-00000000000000007fffffffffffffff-000000024cb016ea-00'") + sqlWithComment.contains("tenant_id='abc123'") + sqlWithComment.contains("request_id='req-456'") + sqlWithComment.contains("user_id='789'") + sqlWithComment.startsWith("SELECT * FROM foo /*") + sqlWithComment.endsWith("*/") + } + + def "test SQL comment injection with custom fields using withSQLCommentFields"() { + setup: + injectSysConfig("dd.service", "SqlCommenter") + injectSysConfig("dd.env", "Test") + injectSysConfig("dd.version", "TestVersion") + + when: + String sqlWithComment = SQLCommenterContext.withSQLCommentFields( + ["tenant_id": "abc123", "request_id": "req-456"], { + return SQLCommenter.inject( + "SELECT * FROM foo", + "my-service", + "mysql", + "h", + "n", + "00-00000000000000007fffffffffffffff-000000024cb016ea-00", + true, + true + ) + } + ) + + then: + // Should contain all standard fields plus custom fields + sqlWithComment.contains("ddps='SqlCommenter'") + sqlWithComment.contains("dddbs='my-service'") + sqlWithComment.contains("tenant_id='abc123'") + sqlWithComment.contains("request_id='req-456'") + sqlWithComment.startsWith("SELECT * FROM foo /*") + sqlWithComment.endsWith("*/") + + // Context should be cleared after the block + SQLCommenterContext.isEmpty() + } + + def "test SQL comment injection with empty custom fields"() { + setup: + injectSysConfig("dd.service", "SqlCommenter") + injectSysConfig("dd.env", "Test") + injectSysConfig("dd.version", "TestVersion") + + when: + String sqlWithComment = SQLCommenter.inject( + "SELECT * FROM foo", + "my-service", + "mysql", + "h", + "n", + "00-00000000000000007fffffffffffffff-000000024cb016ea-00", + true, + true + ) + + then: + // Should contain only standard fields + sqlWithComment.contains("ddps='SqlCommenter'") + sqlWithComment.contains("dddbs='my-service'") + sqlWithComment.contains("ddh='h'") + sqlWithComment.contains("dddb='n'") + sqlWithComment.contains("dde='Test'") + sqlWithComment.contains("ddpv='TestVersion'") + sqlWithComment.contains("traceparent='00-00000000000000007fffffffffffffff-000000024cb016ea-00'") + sqlWithComment.startsWith("SELECT * FROM foo /*") + sqlWithComment.endsWith("*/") + } + + def "test SQL comment injection with custom fields containing special characters"() { + setup: + injectSysConfig("dd.service", "SqlCommenter") + injectSysConfig("dd.env", "Test") + injectSysConfig("dd.version", "TestVersion") + + when: + String sqlWithComment = "" + + // Set up custom fields with special characters that need URL encoding + SQLCommenterContext.put("special_key", "value with spaces") + SQLCommenterContext.put("unicode_key", "value_with_émojis_🎉") + SQLCommenterContext.put("symbols_key", "value&with=symbols") + + sqlWithComment = SQLCommenter.inject( + "SELECT * FROM foo", + "my-service", + "mysql", + "h", + "n", + null, + false, + true + ) + + SQLCommenterContext.clear() + + then: + // Should contain URL-encoded custom fields (URLEncoder uses + for spaces) + sqlWithComment.contains("special_key='value+with+spaces'") + sqlWithComment.contains("unicode_key=") // Should be URL encoded + sqlWithComment.contains("symbols_key='value%26with%3Dsymbols'") + sqlWithComment.startsWith("SELECT * FROM foo /*") + sqlWithComment.endsWith("*/") + } + + def "test SQL comment injection with null and empty custom field values"() { + setup: + injectSysConfig("dd.service", "SqlCommenter") + injectSysConfig("dd.env", "Test") + injectSysConfig("dd.version", "TestVersion") + + when: + String sqlWithComment = "" + + // Set up custom fields with null and empty values + SQLCommenterContext.put("null_key", null) + SQLCommenterContext.put("empty_key", "") + SQLCommenterContext.put("valid_key", "valid_value") + + sqlWithComment = SQLCommenter.inject( + "SELECT * FROM foo", + "my-service", + "mysql", + "h", + "n", + null, + false, + true + ) + + SQLCommenterContext.clear() + + then: + // Should not contain null or empty fields, but should contain valid field + !sqlWithComment.contains("null_key") + !sqlWithComment.contains("empty_key") + sqlWithComment.contains("valid_key='valid_value'") + sqlWithComment.startsWith("SELECT * FROM foo /*") + sqlWithComment.endsWith("*/") + } } diff --git a/dd-java-agent/instrumentation/jdbc/src/test/java/datadog/trace/instrumentation/jdbc/SQLCommenterContextTest.java b/dd-java-agent/instrumentation/jdbc/src/test/java/datadog/trace/instrumentation/jdbc/SQLCommenterContextTest.java new file mode 100644 index 00000000000..06d03a8ab2d --- /dev/null +++ b/dd-java-agent/instrumentation/jdbc/src/test/java/datadog/trace/instrumentation/jdbc/SQLCommenterContextTest.java @@ -0,0 +1,301 @@ +package datadog.trace.instrumentation.jdbc; + +import static org.junit.Assert.*; + +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +public class SQLCommenterContextTest { + + @Before + public void setUp() { + SQLCommenterContext.clear(); + } + + @After + public void tearDown() { + SQLCommenterContext.clear(); + } + + @Test + public void testPutAndGet() { + assertNull(SQLCommenterContext.get("test_key")); + + SQLCommenterContext.put("test_key", "test_value"); + assertEquals("test_value", SQLCommenterContext.get("test_key")); + + SQLCommenterContext.put("test_key", 123); + assertEquals("123", SQLCommenterContext.get("test_key")); + + SQLCommenterContext.put("test_key", null); + assertNull(SQLCommenterContext.get("test_key")); + } + + @Test + public void testRemove() { + SQLCommenterContext.put("test_key", "test_value"); + assertEquals("test_value", SQLCommenterContext.get("test_key")); + + String removed = SQLCommenterContext.remove("test_key"); + assertEquals("test_value", removed); + assertNull(SQLCommenterContext.get("test_key")); + + assertNull(SQLCommenterContext.remove("nonexistent")); + } + + @Test + public void testClear() { + SQLCommenterContext.put("key1", "value1"); + SQLCommenterContext.put("key2", "value2"); + assertFalse(SQLCommenterContext.isEmpty()); + + SQLCommenterContext.clear(); + assertTrue(SQLCommenterContext.isEmpty()); + assertNull(SQLCommenterContext.get("key1")); + assertNull(SQLCommenterContext.get("key2")); + } + + @Test + public void testIsEmpty() { + assertTrue(SQLCommenterContext.isEmpty()); + + SQLCommenterContext.put("key", "value"); + assertFalse(SQLCommenterContext.isEmpty()); + + SQLCommenterContext.clear(); + assertTrue(SQLCommenterContext.isEmpty()); + } + + @Test + public void testGetCopyOfContextMap() { + assertNull(SQLCommenterContext.getCopyOfContextMap()); + + SQLCommenterContext.put("key1", "value1"); + SQLCommenterContext.put("key2", "value2"); + + Map contextMap = SQLCommenterContext.getCopyOfContextMap(); + assertNotNull(contextMap); + assertEquals(2, contextMap.size()); + assertEquals("value1", contextMap.get("key1")); + assertEquals("value2", contextMap.get("key2")); + + // Verify it's a copy (modifications don't affect original) + contextMap.put("key3", "value3"); + assertNull(SQLCommenterContext.get("key3")); + } + + @Test + public void testSetContextMap() { + Map contextMap = new HashMap<>(); + contextMap.put("key1", "value1"); + contextMap.put("key2", "value2"); + + SQLCommenterContext.setContextMap(contextMap); + assertEquals("value1", SQLCommenterContext.get("key1")); + assertEquals("value2", SQLCommenterContext.get("key2")); + + // Verify the internal map is a copy + contextMap.put("key3", "value3"); + assertNull(SQLCommenterContext.get("key3")); + + // Test setting null clears the context + SQLCommenterContext.setContextMap(null); + assertTrue(SQLCommenterContext.isEmpty()); + } + + @Test + public void testWithSQLCommentFieldsMap() throws Exception { + SQLCommenterContext.put("existing_key", "existing_value"); + + Map customFields = new HashMap<>(); + customFields.put("custom_key", "custom_value"); + customFields.put("number_key", 42); + customFields.put("null_key", null); + + String result = + SQLCommenterContext.withSQLCommentFields( + customFields, + () -> { + // Inside the block, we should have both existing and custom fields + assertEquals("existing_value", SQLCommenterContext.get("existing_key")); + assertEquals("custom_value", SQLCommenterContext.get("custom_key")); + assertEquals("42", SQLCommenterContext.get("number_key")); + assertNull(SQLCommenterContext.get("null_key")); + + return "block_executed"; + }); + + assertEquals("block_executed", result); + + // After the block, only the original context should remain + assertEquals("existing_value", SQLCommenterContext.get("existing_key")); + assertNull(SQLCommenterContext.get("custom_key")); + assertNull(SQLCommenterContext.get("number_key")); + } + + @Test + public void testWithSQLCommentFieldsVarArgs() throws Exception { + SQLCommenterContext.put("existing_key", "existing_value"); + + String result = + SQLCommenterContext.withSQLCommentFields( + () -> { + assertEquals("existing_value", SQLCommenterContext.get("existing_key")); + assertEquals("custom_value", SQLCommenterContext.get("custom_key")); + assertEquals("42", SQLCommenterContext.get("number_key")); + return "block_executed"; + }, + "custom_key", + "custom_value", + "number_key", + 42); + + assertEquals("block_executed", result); + + // After the block, only the original context should remain + assertEquals("existing_value", SQLCommenterContext.get("existing_key")); + assertNull(SQLCommenterContext.get("custom_key")); + assertNull(SQLCommenterContext.get("number_key")); + } + + @Test + public void testWithSQLCommentFieldsNested() throws Exception { + String result = + SQLCommenterContext.withSQLCommentFields( + () -> { + SQLCommenterContext.put("level1", "value1"); + + return SQLCommenterContext.withSQLCommentFields( + () -> { + assertEquals("value1", SQLCommenterContext.get("level1")); + assertEquals("value2", SQLCommenterContext.get("level2")); + return "nested_executed"; + }, + "level2", + "value2"); + }, + "outer_key", + "outer_value"); + + assertEquals("nested_executed", result); + + // After nested execution, context should be empty + assertTrue(SQLCommenterContext.isEmpty()); + } + + @Test + public void testWithSQLCommentFieldsExceptionHandling() { + SQLCommenterContext.put("existing_key", "existing_value"); + + RuntimeException expectedException = new RuntimeException("test exception"); + + try { + SQLCommenterContext.withSQLCommentFields( + () -> { + SQLCommenterContext.put("temp_key", "temp_value"); + throw expectedException; + }, + "custom_key", + "custom_value"); + fail("Should have thrown exception"); + } catch (Exception e) { + assertEquals(expectedException, e); + } + + // Context should be restored even after exception + assertEquals("existing_value", SQLCommenterContext.get("existing_key")); + assertNull(SQLCommenterContext.get("temp_key")); + assertNull(SQLCommenterContext.get("custom_key")); + } + + @Test + public void testThreadIsolation() throws Exception { + SQLCommenterContext.put("main_key", "main_value"); + + ExecutorService executor = Executors.newFixedThreadPool(2); + CountDownLatch latch = new CountDownLatch(2); + + // Start two threads with different contexts + Future future1 = + executor.submit( + () -> { + try { + SQLCommenterContext.put("thread1_key", "thread1_value"); + latch.countDown(); + latch.await(5, TimeUnit.SECONDS); + + // Each thread should only see its own context + assertEquals("thread1_value", SQLCommenterContext.get("thread1_key")); + assertNull(SQLCommenterContext.get("thread2_key")); + assertNull( + SQLCommenterContext.get( + "main_key")); // Thread local, so main context not visible + + return "thread1_done"; + } catch (Exception e) { + throw new RuntimeException(e); + } + }); + + Future future2 = + executor.submit( + () -> { + try { + SQLCommenterContext.put("thread2_key", "thread2_value"); + latch.countDown(); + latch.await(5, TimeUnit.SECONDS); + + // Each thread should only see its own context + assertEquals("thread2_value", SQLCommenterContext.get("thread2_key")); + assertNull(SQLCommenterContext.get("thread1_key")); + assertNull( + SQLCommenterContext.get( + "main_key")); // Thread local, so main context not visible + + return "thread2_done"; + } catch (Exception e) { + throw new RuntimeException(e); + } + }); + + assertEquals("thread1_done", future1.get(10, TimeUnit.SECONDS)); + assertEquals("thread2_done", future2.get(10, TimeUnit.SECONDS)); + + // Main thread should still have its context + assertEquals("main_value", SQLCommenterContext.get("main_key")); + assertNull(SQLCommenterContext.get("thread1_key")); + assertNull(SQLCommenterContext.get("thread2_key")); + + executor.shutdown(); + } + + @Test + public void testVarArgsValidation() { + try { + SQLCommenterContext.withSQLCommentFields(() -> "test", "odd_number_of_args"); + fail("Should have thrown IllegalArgumentException"); + } catch (IllegalArgumentException e) { + assertEquals("Key-value pairs must be provided in pairs", e.getMessage()); + } catch (Exception e) { + fail("Should have thrown IllegalArgumentException, but got: " + e.getClass().getSimpleName()); + } + } + + @Test + public void testPutWithNullKey() { + try { + SQLCommenterContext.put(null, "value"); + fail("Should have thrown IllegalArgumentException"); + } catch (IllegalArgumentException e) { + assertEquals("Key cannot be null", e.getMessage()); + } + } +} diff --git a/dd-trace-api/src/main/java/datadog/trace/api/SQLCommenter.java b/dd-trace-api/src/main/java/datadog/trace/api/SQLCommenter.java new file mode 100644 index 00000000000..1c8d60b448e --- /dev/null +++ b/dd-trace-api/src/main/java/datadog/trace/api/SQLCommenter.java @@ -0,0 +1,317 @@ +package datadog.trace.api; + +import java.util.Map; +import java.util.concurrent.Callable; + +/** + * Public API for adding custom key-value pairs to SQL comments generated by the Datadog Java + * tracer. This API allows external users to inject custom metadata into SQL comment traces. + * + *

This API only works when the dd-trace-java agent is loaded. If the agent is not present, these + * methods will have no effect. Applications should handle the optional dependency by checking for + * class availability before calling these methods. + * + *

Usage example: + * + * 

+ * // Add a custom field for a single operation
+ * SQLCommenter.withCustomFields("tenant_id", "abc123", "request_id", "xyz789", () -> {
+ *   // Execute database operations here
+ *   return myRepository.findUser(userId);
+ * });
+ *
+ * // Add custom fields using a map
+ * Map<String, Object> customFields = Map.of(
+ *   "tenant_id", "abc123",
+ *   "operation", "user_lookup"
+ * );
+ * SQLCommenter.withCustomFields(customFields, () -> {
+ *   return myRepository.findUser(userId);
+ * });
+ *
+ * // Add a field for the current thread
+ * SQLCommenter.put("correlation_id", requestId);
+ * // All subsequent SQL operations will include this field
+ * SQLCommenter.clear(); // Clear when done
+ *
+ */ +public class SQLCommenter { + + /** + * Adds a custom field to the current thread's SQL comment context. + * + * @param key the field key + * @param value the field value (will be converted to string) + * @throws IllegalArgumentException if key is null + */ + public static void put(String key, Object value) { + getContext().put(key, value); + } + + /** + * Gets a custom field from the current thread's SQL comment context. + * + * @param key the field key + * @return the field value, or null if not found + */ + public static String get(String key) { + return getContext().get(key); + } + + /** + * Removes a custom field from the current thread's SQL comment context. + * + * @param key the field key + * @return the previous value, or null if not found + */ + public static String remove(String key) { + return getContext().remove(key); + } + + /** Clears all custom fields from the current thread's SQL comment context. */ + public static void clear() { + getContext().clear(); + } + + /** + * Returns true if the current thread has no custom SQL comment fields. + * + * @return true if no custom fields are set + */ + public static boolean isEmpty() { + return getContext().isEmpty(); + } + + /** + * Gets a copy of the current thread's SQL comment context map. + * + * @return a copy of the current context map, or null if empty + */ + public static Map getCopyOfContextMap() { + return getContext().getCopyOfContextMap(); + } + + /** + * Executes a block of code with the given custom fields, restoring the original context + * afterward. + * + * @param customFields the custom fields to use during execution + * @param callable the code block to execute + * @param the return type of the callable + * @return the result of executing the callable + * @throws Exception if the callable throws an exception + */ + public static T withCustomFields(Map customFields, Callable callable) + throws Exception { + return getContext().withSQLCommentFields(customFields, callable::call); + } + + /** + * Executes a block of code with the given custom fields, restoring the original context + * afterward. + * + * @param callable the code block to execute + * @param keyValuePairs alternating key-value pairs to add to the context + * @param the return type of the callable + * @return the result of executing the callable + * @throws Exception if the callable throws an exception + */ + public static T withCustomFields(Callable callable, Object... keyValuePairs) + throws Exception { + return getContext().withSQLCommentFields(callable::call, keyValuePairs); + } + + /** + * Executes a runnable with the given custom fields, restoring the original context afterward. + * + * @param customFields the custom fields to use during execution + * @param runnable the code block to execute + */ + public static void withCustomFields(Map customFields, Runnable runnable) { + try { + withCustomFields( + customFields, + () -> { + runnable.run(); + return null; + }); + } catch (Exception e) { + if (e instanceof RuntimeException) { + throw (RuntimeException) e; + } + throw new RuntimeException(e); + } + } + + /** + * Executes a runnable with the given custom fields, restoring the original context afterward. + * + * @param runnable the code block to execute + * @param keyValuePairs alternating key-value pairs to add to the context + */ + public static void withCustomFields(Runnable runnable, Object... keyValuePairs) { + try { + withCustomFields( + () -> { + runnable.run(); + return null; + }, + keyValuePairs); + } catch (Exception e) { + if (e instanceof RuntimeException) { + throw (RuntimeException) e; + } + throw new RuntimeException(e); + } + } + + // Direct access to the context - simpler than the provider pattern + private static SQLCommenterContextAccess getContext() { + return SQLCommenterContextHolder.INSTANCE; + } + + private static class SQLCommenterContextHolder { + private static final SQLCommenterContextAccess INSTANCE = createInstance(); + + private static SQLCommenterContextAccess createInstance() { + try { + Class contextClass = + Class.forName("datadog.trace.instrumentation.jdbc.SQLCommenterContext"); + return new DirectSQLCommenterContextAccess(contextClass); + } catch (ClassNotFoundException | LinkageError e) { + // This should not happen if the dd-trace-java agent is properly loaded + throw new RuntimeException( + "SQLCommenterContext is not available. Ensure dd-trace-java agent is loaded.", e); + } + } + } + + // Direct access interface - much simpler than the provider pattern + private interface SQLCommenterContextAccess { + void put(String key, Object value); + + String get(String key); + + String remove(String key); + + void clear(); + + boolean isEmpty(); + + Map getCopyOfContextMap(); + + T withSQLCommentFields(Map contextMap, SQLCommenterBlock block); + + T withSQLCommentFields(SQLCommenterBlock block, Object... keyValuePairs); + } + + @FunctionalInterface + private interface SQLCommenterBlock { + T execute() throws Exception; + } + + private static class DirectSQLCommenterContextAccess implements SQLCommenterContextAccess { + private final Class contextClass; + + public DirectSQLCommenterContextAccess(Class contextClass) { + this.contextClass = contextClass; + } + + @Override + public void put(String key, Object value) { + try { + contextClass.getMethod("put", String.class, Object.class).invoke(null, key, value); + } catch (Exception e) { + throw new RuntimeException("Failed to put SQL comment field", e); + } + } + + @Override + public String get(String key) { + try { + return (String) contextClass.getMethod("get", String.class).invoke(null, key); + } catch (Exception e) { + throw new RuntimeException("Failed to get SQL comment field", e); + } + } + + @Override + public String remove(String key) { + try { + return (String) contextClass.getMethod("remove", String.class).invoke(null, key); + } catch (Exception e) { + throw new RuntimeException("Failed to remove SQL comment field", e); + } + } + + @Override + public void clear() { + try { + contextClass.getMethod("clear").invoke(null); + } catch (Exception e) { + throw new RuntimeException("Failed to clear SQL comment context", e); + } + } + + @Override + public boolean isEmpty() { + try { + return (Boolean) contextClass.getMethod("isEmpty").invoke(null); + } catch (Exception e) { + throw new RuntimeException("Failed to check if SQL comment context is empty", e); + } + } + + @Override + @SuppressWarnings("unchecked") + public Map getCopyOfContextMap() { + try { + return (Map) contextClass.getMethod("getCopyOfContextMap").invoke(null); + } catch (Exception e) { + throw new RuntimeException("Failed to get copy of SQL comment context map", e); + } + } + + @Override + @SuppressWarnings("unchecked") + public T withSQLCommentFields(Map contextMap, SQLCommenterBlock block) { + try { + Class blockClass = + Class.forName( + "datadog.trace.instrumentation.jdbc.SQLCommenterContext$SQLCommenterBlock"); + Object blockProxy = + java.lang.reflect.Proxy.newProxyInstance( + blockClass.getClassLoader(), + new Class[] {blockClass}, + (proxy, method, args) -> block.execute()); + return (T) + contextClass + .getMethod("withSQLCommentFields", Map.class, blockClass) + .invoke(null, contextMap, blockProxy); + } catch (Exception e) { + throw new RuntimeException("Failed to execute with SQL comment fields", e); + } + } + + @Override + @SuppressWarnings("unchecked") + public T withSQLCommentFields(SQLCommenterBlock block, Object... keyValuePairs) { + try { + Class blockClass = + Class.forName( + "datadog.trace.instrumentation.jdbc.SQLCommenterContext$SQLCommenterBlock"); + Object blockProxy = + java.lang.reflect.Proxy.newProxyInstance( + blockClass.getClassLoader(), + new Class[] {blockClass}, + (proxy, method, args) -> block.execute()); + return (T) + contextClass + .getMethod("withSQLCommentFields", blockClass, Object[].class) + .invoke(null, blockProxy, keyValuePairs); + } catch (Exception e) { + throw new RuntimeException("Failed to execute with SQL comment fields", e); + } + } + } +} From b721eeed26c72f7091b912cf9fe23aaed3203792 Mon Sep 17 00:00:00 2001 From: Christian Gati Date: Fri, 18 Jul 2025 11:51:32 -0700 Subject: [PATCH 2/5] Remove reference to Kotlin code example --- .../datadog/trace/instrumentation/jdbc/SQLCommenterContext.java | 1 - 1 file changed, 1 deletion(-) diff --git a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/SQLCommenterContext.java b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/SQLCommenterContext.java index c766880314b..ec34dd6d516 100644 --- a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/SQLCommenterContext.java +++ b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/SQLCommenterContext.java @@ -91,7 +91,6 @@ public static boolean isEmpty() { /** * Executes a block of code with the given context map, restoring the original context afterward. - * This method is similar to the withLoggingFields pattern from the provided Kotlin code. * * @param contextMap the context map to use during execution * @param block the code block to execute From d8a38372c0480fb123f062014472a67909746583 Mon Sep 17 00:00:00 2001 From: Christian Gati Date: Fri, 18 Jul 2025 11:54:01 -0700 Subject: [PATCH 3/5] Fix html-encoded comment --- dd-trace-api/src/main/java/datadog/trace/api/SQLCommenter.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dd-trace-api/src/main/java/datadog/trace/api/SQLCommenter.java b/dd-trace-api/src/main/java/datadog/trace/api/SQLCommenter.java index 1c8d60b448e..4d687e98374 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/SQLCommenter.java +++ b/dd-trace-api/src/main/java/datadog/trace/api/SQLCommenter.java @@ -21,7 +21,7 @@ * }); * * // Add custom fields using a map - * Map<String, Object> customFields = Map.of( + * Map customFields = Map.of( * "tenant_id", "abc123", * "operation", "user_lookup" * ); From 75d5761d6325293be3e2a94ac976701f824a20ab Mon Sep 17 00:00:00 2001 From: Christian Gati Date: Fri, 18 Jul 2025 15:10:59 -0700 Subject: [PATCH 4/5] simplify capacity calculation --- .../datadog/trace/instrumentation/jdbc/SQLCommenter.java | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/SQLCommenter.java b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/SQLCommenter.java index b2d4574fda7..57abf11b9c0 100644 --- a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/SQLCommenter.java +++ b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/SQLCommenter.java @@ -305,9 +305,8 @@ private static int capacity( if (customFields != null) { for (Map.Entry entry : customFields.entrySet()) { if (entry.getKey() != null && entry.getValue() != null) { - // Account for URL encoding overhead (approximately 3x for worst case) - len += entry.getKey().length() * 3; - len += entry.getValue().length() * 3; + len += entry.getKey().length(); + len += entry.getValue().length(); len += 4; // equals, comma, and two quotes } } From 5ca884f128d376d1462044ce1e2837b0c1010d5f Mon Sep 17 00:00:00 2001 From: Christian Gati Date: Fri, 18 Jul 2025 15:41:43 -0700 Subject: [PATCH 5/5] simplify the api to getting and setting the context map --- .../jdbc/SQLCommenterContext.java | 109 ------- .../jdbc/SQLCommenterContextTest.java | 279 +++++------------- .../java/datadog/trace/api/SQLCommenter.java | 258 ++-------------- 3 files changed, 99 insertions(+), 547 deletions(-) diff --git a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/SQLCommenterContext.java b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/SQLCommenterContext.java index ec34dd6d516..4f010c90d97 100644 --- a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/SQLCommenterContext.java +++ b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/SQLCommenterContext.java @@ -40,113 +40,4 @@ public static void setContextMap(Map contextMap) { contextHolder.set(new HashMap<>(contextMap)); } } - - /** - * Puts a key-value pair into the current context. - * - * @param key the key - * @param value the value (will be converted to string) - */ - public static void put(String key, Object value) { - if (key == null) { - throw new IllegalArgumentException("Key cannot be null"); - } - Map contextMap = contextHolder.get(); - contextMap.put(key, value != null ? value.toString() : null); - } - - /** - * Gets a value from the current context. - * - * @param key the key to look up - * @return the value associated with the key, or null if not found - */ - public static String get(String key) { - return contextHolder.get().get(key); - } - - /** - * Removes a key from the current context. - * - * @param key the key to remove - * @return the previous value associated with the key, or null if not found - */ - public static String remove(String key) { - return contextHolder.get().remove(key); - } - - /** Clears the current context. */ - public static void clear() { - contextHolder.remove(); - } - - /** - * Returns true if the current context is empty. - * - * @return true if the current context has no entries - */ - public static boolean isEmpty() { - return contextHolder.get().isEmpty(); - } - - /** - * Executes a block of code with the given context map, restoring the original context afterward. - * - * @param contextMap the context map to use during execution - * @param block the code block to execute - * @param the return type of the block - * @return the result of executing the block - */ - public static T withSQLCommentFields( - Map contextMap, SQLCommenterBlock block) throws Exception { - Map oldContextMap = getCopyOfContextMap(); - try { - Map newContextMap = - new HashMap<>(oldContextMap != null ? oldContextMap : new HashMap<>()); - if (contextMap != null) { - contextMap.forEach( - (key, value) -> newContextMap.put(key, value != null ? value.toString() : null)); - } - setContextMap(newContextMap); - return block.execute(); - } finally { - setContextMap(oldContextMap); - } - } - - /** - * Executes a block of code with the given key-value pairs, restoring the original context - * afterward. - * - * @param block the code block to execute - * @param keyValuePairs alternating key-value pairs to add to the context - * @param the return type of the block - * @return the result of executing the block - */ - @SuppressWarnings("unchecked") - public static T withSQLCommentFields(SQLCommenterBlock block, Object... keyValuePairs) - throws Exception { - if (keyValuePairs.length % 2 != 0) { - throw new IllegalArgumentException("Key-value pairs must be provided in pairs"); - } - - Map contextMap = new HashMap<>(); - for (int i = 0; i < keyValuePairs.length; i += 2) { - String key = keyValuePairs[i].toString(); - Object value = keyValuePairs[i + 1]; - contextMap.put(key, value); - } - - return withSQLCommentFields(contextMap, block); - } - - /** - * Functional interface for code blocks that can be executed with a SQL comment context. - * - * @param the return type of the block - */ - @FunctionalInterface - public interface SQLCommenterBlock { - T execute() throws Exception; - } } diff --git a/dd-java-agent/instrumentation/jdbc/src/test/java/datadog/trace/instrumentation/jdbc/SQLCommenterContextTest.java b/dd-java-agent/instrumentation/jdbc/src/test/java/datadog/trace/instrumentation/jdbc/SQLCommenterContextTest.java index 06d03a8ab2d..d54565f0b8d 100644 --- a/dd-java-agent/instrumentation/jdbc/src/test/java/datadog/trace/instrumentation/jdbc/SQLCommenterContextTest.java +++ b/dd-java-agent/instrumentation/jdbc/src/test/java/datadog/trace/instrumentation/jdbc/SQLCommenterContextTest.java @@ -17,208 +17,76 @@ public class SQLCommenterContextTest { @Before public void setUp() { - SQLCommenterContext.clear(); + SQLCommenterContext.setContextMap(null); } @After public void tearDown() { - SQLCommenterContext.clear(); - } - - @Test - public void testPutAndGet() { - assertNull(SQLCommenterContext.get("test_key")); - - SQLCommenterContext.put("test_key", "test_value"); - assertEquals("test_value", SQLCommenterContext.get("test_key")); - - SQLCommenterContext.put("test_key", 123); - assertEquals("123", SQLCommenterContext.get("test_key")); - - SQLCommenterContext.put("test_key", null); - assertNull(SQLCommenterContext.get("test_key")); - } - - @Test - public void testRemove() { - SQLCommenterContext.put("test_key", "test_value"); - assertEquals("test_value", SQLCommenterContext.get("test_key")); - - String removed = SQLCommenterContext.remove("test_key"); - assertEquals("test_value", removed); - assertNull(SQLCommenterContext.get("test_key")); - - assertNull(SQLCommenterContext.remove("nonexistent")); - } - - @Test - public void testClear() { - SQLCommenterContext.put("key1", "value1"); - SQLCommenterContext.put("key2", "value2"); - assertFalse(SQLCommenterContext.isEmpty()); - - SQLCommenterContext.clear(); - assertTrue(SQLCommenterContext.isEmpty()); - assertNull(SQLCommenterContext.get("key1")); - assertNull(SQLCommenterContext.get("key2")); - } - - @Test - public void testIsEmpty() { - assertTrue(SQLCommenterContext.isEmpty()); - - SQLCommenterContext.put("key", "value"); - assertFalse(SQLCommenterContext.isEmpty()); - - SQLCommenterContext.clear(); - assertTrue(SQLCommenterContext.isEmpty()); + SQLCommenterContext.setContextMap(null); } @Test - public void testGetCopyOfContextMap() { + public void testGetCopyOfContextMapEmpty() { assertNull(SQLCommenterContext.getCopyOfContextMap()); - - SQLCommenterContext.put("key1", "value1"); - SQLCommenterContext.put("key2", "value2"); - - Map contextMap = SQLCommenterContext.getCopyOfContextMap(); - assertNotNull(contextMap); - assertEquals(2, contextMap.size()); - assertEquals("value1", contextMap.get("key1")); - assertEquals("value2", contextMap.get("key2")); - - // Verify it's a copy (modifications don't affect original) - contextMap.put("key3", "value3"); - assertNull(SQLCommenterContext.get("key3")); } @Test - public void testSetContextMap() { + public void testSetAndGetContextMap() { Map contextMap = new HashMap<>(); contextMap.put("key1", "value1"); contextMap.put("key2", "value2"); SQLCommenterContext.setContextMap(contextMap); - assertEquals("value1", SQLCommenterContext.get("key1")); - assertEquals("value2", SQLCommenterContext.get("key2")); - // Verify the internal map is a copy - contextMap.put("key3", "value3"); - assertNull(SQLCommenterContext.get("key3")); + Map retrievedMap = SQLCommenterContext.getCopyOfContextMap(); + assertNotNull(retrievedMap); + assertEquals(2, retrievedMap.size()); + assertEquals("value1", retrievedMap.get("key1")); + assertEquals("value2", retrievedMap.get("key2")); - // Test setting null clears the context - SQLCommenterContext.setContextMap(null); - assertTrue(SQLCommenterContext.isEmpty()); + // Verify it's a copy (modifications don't affect original) + retrievedMap.put("key3", "value3"); + Map retrievedAgain = SQLCommenterContext.getCopyOfContextMap(); + assertEquals(2, retrievedAgain.size()); + assertNull(retrievedAgain.get("key3")); } @Test - public void testWithSQLCommentFieldsMap() throws Exception { - SQLCommenterContext.put("existing_key", "existing_value"); - - Map customFields = new HashMap<>(); - customFields.put("custom_key", "custom_value"); - customFields.put("number_key", 42); - customFields.put("null_key", null); - - String result = - SQLCommenterContext.withSQLCommentFields( - customFields, - () -> { - // Inside the block, we should have both existing and custom fields - assertEquals("existing_value", SQLCommenterContext.get("existing_key")); - assertEquals("custom_value", SQLCommenterContext.get("custom_key")); - assertEquals("42", SQLCommenterContext.get("number_key")); - assertNull(SQLCommenterContext.get("null_key")); + public void testSetContextMapMakesInternalCopy() { + Map contextMap = new HashMap<>(); + contextMap.put("key1", "value1"); - return "block_executed"; - }); + SQLCommenterContext.setContextMap(contextMap); - assertEquals("block_executed", result); + // Modify the original map + contextMap.put("key2", "value2"); - // After the block, only the original context should remain - assertEquals("existing_value", SQLCommenterContext.get("existing_key")); - assertNull(SQLCommenterContext.get("custom_key")); - assertNull(SQLCommenterContext.get("number_key")); + // Internal context should not be affected + Map retrievedMap = SQLCommenterContext.getCopyOfContextMap(); + assertEquals(1, retrievedMap.size()); + assertEquals("value1", retrievedMap.get("key1")); + assertNull(retrievedMap.get("key2")); } @Test - public void testWithSQLCommentFieldsVarArgs() throws Exception { - SQLCommenterContext.put("existing_key", "existing_value"); - - String result = - SQLCommenterContext.withSQLCommentFields( - () -> { - assertEquals("existing_value", SQLCommenterContext.get("existing_key")); - assertEquals("custom_value", SQLCommenterContext.get("custom_key")); - assertEquals("42", SQLCommenterContext.get("number_key")); - return "block_executed"; - }, - "custom_key", - "custom_value", - "number_key", - 42); - - assertEquals("block_executed", result); - - // After the block, only the original context should remain - assertEquals("existing_value", SQLCommenterContext.get("existing_key")); - assertNull(SQLCommenterContext.get("custom_key")); - assertNull(SQLCommenterContext.get("number_key")); - } + public void testSetContextMapNull() { + Map contextMap = new HashMap<>(); + contextMap.put("key1", "value1"); + SQLCommenterContext.setContextMap(contextMap); - @Test - public void testWithSQLCommentFieldsNested() throws Exception { - String result = - SQLCommenterContext.withSQLCommentFields( - () -> { - SQLCommenterContext.put("level1", "value1"); - - return SQLCommenterContext.withSQLCommentFields( - () -> { - assertEquals("value1", SQLCommenterContext.get("level1")); - assertEquals("value2", SQLCommenterContext.get("level2")); - return "nested_executed"; - }, - "level2", - "value2"); - }, - "outer_key", - "outer_value"); - - assertEquals("nested_executed", result); - - // After nested execution, context should be empty - assertTrue(SQLCommenterContext.isEmpty()); - } + // Verify context is set + assertNotNull(SQLCommenterContext.getCopyOfContextMap()); - @Test - public void testWithSQLCommentFieldsExceptionHandling() { - SQLCommenterContext.put("existing_key", "existing_value"); - - RuntimeException expectedException = new RuntimeException("test exception"); - - try { - SQLCommenterContext.withSQLCommentFields( - () -> { - SQLCommenterContext.put("temp_key", "temp_value"); - throw expectedException; - }, - "custom_key", - "custom_value"); - fail("Should have thrown exception"); - } catch (Exception e) { - assertEquals(expectedException, e); - } - - // Context should be restored even after exception - assertEquals("existing_value", SQLCommenterContext.get("existing_key")); - assertNull(SQLCommenterContext.get("temp_key")); - assertNull(SQLCommenterContext.get("custom_key")); + // Clear with null + SQLCommenterContext.setContextMap(null); + assertNull(SQLCommenterContext.getCopyOfContextMap()); } @Test public void testThreadIsolation() throws Exception { - SQLCommenterContext.put("main_key", "main_value"); + Map mainContext = new HashMap<>(); + mainContext.put("main_key", "main_value"); + SQLCommenterContext.setContextMap(mainContext); ExecutorService executor = Executors.newFixedThreadPool(2); CountDownLatch latch = new CountDownLatch(2); @@ -228,16 +96,19 @@ public void testThreadIsolation() throws Exception { executor.submit( () -> { try { - SQLCommenterContext.put("thread1_key", "thread1_value"); + Map thread1Context = new HashMap<>(); + thread1Context.put("thread1_key", "thread1_value"); + SQLCommenterContext.setContextMap(thread1Context); + latch.countDown(); latch.await(5, TimeUnit.SECONDS); // Each thread should only see its own context - assertEquals("thread1_value", SQLCommenterContext.get("thread1_key")); - assertNull(SQLCommenterContext.get("thread2_key")); - assertNull( - SQLCommenterContext.get( - "main_key")); // Thread local, so main context not visible + Map context = SQLCommenterContext.getCopyOfContextMap(); + assertEquals(1, context.size()); + assertEquals("thread1_value", context.get("thread1_key")); + assertNull(context.get("thread2_key")); + assertNull(context.get("main_key")); // Thread local, so main context not visible return "thread1_done"; } catch (Exception e) { @@ -249,16 +120,19 @@ public void testThreadIsolation() throws Exception { executor.submit( () -> { try { - SQLCommenterContext.put("thread2_key", "thread2_value"); + Map thread2Context = new HashMap<>(); + thread2Context.put("thread2_key", "thread2_value"); + SQLCommenterContext.setContextMap(thread2Context); + latch.countDown(); latch.await(5, TimeUnit.SECONDS); // Each thread should only see its own context - assertEquals("thread2_value", SQLCommenterContext.get("thread2_key")); - assertNull(SQLCommenterContext.get("thread1_key")); - assertNull( - SQLCommenterContext.get( - "main_key")); // Thread local, so main context not visible + Map context = SQLCommenterContext.getCopyOfContextMap(); + assertEquals(1, context.size()); + assertEquals("thread2_value", context.get("thread2_key")); + assertNull(context.get("thread1_key")); + assertNull(context.get("main_key")); // Thread local, so main context not visible return "thread2_done"; } catch (Exception e) { @@ -270,32 +144,31 @@ public void testThreadIsolation() throws Exception { assertEquals("thread2_done", future2.get(10, TimeUnit.SECONDS)); // Main thread should still have its context - assertEquals("main_value", SQLCommenterContext.get("main_key")); - assertNull(SQLCommenterContext.get("thread1_key")); - assertNull(SQLCommenterContext.get("thread2_key")); + Map mainContextAfter = SQLCommenterContext.getCopyOfContextMap(); + assertEquals(1, mainContextAfter.size()); + assertEquals("main_value", mainContextAfter.get("main_key")); + assertNull(mainContextAfter.get("thread1_key")); + assertNull(mainContextAfter.get("thread2_key")); executor.shutdown(); } @Test - public void testVarArgsValidation() { - try { - SQLCommenterContext.withSQLCommentFields(() -> "test", "odd_number_of_args"); - fail("Should have thrown IllegalArgumentException"); - } catch (IllegalArgumentException e) { - assertEquals("Key-value pairs must be provided in pairs", e.getMessage()); - } catch (Exception e) { - fail("Should have thrown IllegalArgumentException, but got: " + e.getClass().getSimpleName()); - } - } - - @Test - public void testPutWithNullKey() { - try { - SQLCommenterContext.put(null, "value"); - fail("Should have thrown IllegalArgumentException"); - } catch (IllegalArgumentException e) { - assertEquals("Key cannot be null", e.getMessage()); - } + public void testContextPersistence() { + // Set initial context + Map context1 = new HashMap<>(); + context1.put("key1", "value1"); + SQLCommenterContext.setContextMap(context1); + + // Update context + Map context2 = new HashMap<>(); + context2.put("key2", "value2"); + SQLCommenterContext.setContextMap(context2); + + // Should only have the latest context + Map retrieved = SQLCommenterContext.getCopyOfContextMap(); + assertEquals(1, retrieved.size()); + assertEquals("value2", retrieved.get("key2")); + assertNull(retrieved.get("key1")); } } diff --git a/dd-trace-api/src/main/java/datadog/trace/api/SQLCommenter.java b/dd-trace-api/src/main/java/datadog/trace/api/SQLCommenter.java index 4d687e98374..d1ef62298c2 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/SQLCommenter.java +++ b/dd-trace-api/src/main/java/datadog/trace/api/SQLCommenter.java @@ -1,11 +1,10 @@ package datadog.trace.api; import java.util.Map; -import java.util.concurrent.Callable; /** - * Public API for adding custom key-value pairs to SQL comments generated by the Datadog Java - * tracer. This API allows external users to inject custom metadata into SQL comment traces. + * Public API for accessing SQL comment context managed by the Datadog Java tracer. This API + * provides essential primitives for custom context management integration. * *

This API only works when the dd-trace-java agent is loaded. If the agent is not present, these * methods will have no effect. Applications should handle the optional dependency by checking for @@ -14,74 +13,23 @@ *

Usage example: * * 

- * // Add a custom field for a single operation
- * SQLCommenter.withCustomFields("tenant_id", "abc123", "request_id", "xyz789", () -> {
- *   // Execute database operations here
- *   return myRepository.findUser(userId);
- * });
+ * // Get current context
+ * Map currentContext = SQLCommenter.getCopyOfContextMap();
  *
- * // Add custom fields using a map
- * Map customFields = Map.of(
- *   "tenant_id", "abc123",
- *   "operation", "user_lookup"
- * );
- * SQLCommenter.withCustomFields(customFields, () -> {
- *   return myRepository.findUser(userId);
- * });
+ * // Merge with custom fields and set new context
+ * Map newContext = new HashMap<>(currentContext != null ? currentContext : Map.of());
+ * newContext.put("tenant_id", "abc123");
+ * newContext.put("request_id", "xyz789");
+ * SQLCommenter.setContextMap(newContext);
  *
- * // Add a field for the current thread
- * SQLCommenter.put("correlation_id", requestId);
- * // All subsequent SQL operations will include this field
- * SQLCommenter.clear(); // Clear when done
+ * // Execute database operations...
+ *
+ * // Restore original context
+ * SQLCommenter.setContextMap(currentContext);
  *
*/ public class SQLCommenter { - /** - * Adds a custom field to the current thread's SQL comment context. - * - * @param key the field key - * @param value the field value (will be converted to string) - * @throws IllegalArgumentException if key is null - */ - public static void put(String key, Object value) { - getContext().put(key, value); - } - - /** - * Gets a custom field from the current thread's SQL comment context. - * - * @param key the field key - * @return the field value, or null if not found - */ - public static String get(String key) { - return getContext().get(key); - } - - /** - * Removes a custom field from the current thread's SQL comment context. - * - * @param key the field key - * @return the previous value, or null if not found - */ - public static String remove(String key) { - return getContext().remove(key); - } - - /** Clears all custom fields from the current thread's SQL comment context. */ - public static void clear() { - getContext().clear(); - } - - /** - * Returns true if the current thread has no custom SQL comment fields. - * - * @return true if no custom fields are set - */ - public static boolean isEmpty() { - return getContext().isEmpty(); - } - /** * Gets a copy of the current thread's SQL comment context map. * @@ -92,80 +40,15 @@ public static Map getCopyOfContextMap() { } /** - * Executes a block of code with the given custom fields, restoring the original context - * afterward. - * - * @param customFields the custom fields to use during execution - * @param callable the code block to execute - * @param the return type of the callable - * @return the result of executing the callable - * @throws Exception if the callable throws an exception - */ - public static T withCustomFields(Map customFields, Callable callable) - throws Exception { - return getContext().withSQLCommentFields(customFields, callable::call); - } - - /** - * Executes a block of code with the given custom fields, restoring the original context - * afterward. - * - * @param callable the code block to execute - * @param keyValuePairs alternating key-value pairs to add to the context - * @param the return type of the callable - * @return the result of executing the callable - * @throws Exception if the callable throws an exception - */ - public static T withCustomFields(Callable callable, Object... keyValuePairs) - throws Exception { - return getContext().withSQLCommentFields(callable::call, keyValuePairs); - } - - /** - * Executes a runnable with the given custom fields, restoring the original context afterward. + * Sets the context map for the current thread. * - * @param customFields the custom fields to use during execution - * @param runnable the code block to execute + * @param contextMap the new context map to set, or null to clear */ - public static void withCustomFields(Map customFields, Runnable runnable) { - try { - withCustomFields( - customFields, - () -> { - runnable.run(); - return null; - }); - } catch (Exception e) { - if (e instanceof RuntimeException) { - throw (RuntimeException) e; - } - throw new RuntimeException(e); - } + public static void setContextMap(Map contextMap) { + getContext().setContextMap(contextMap); } - /** - * Executes a runnable with the given custom fields, restoring the original context afterward. - * - * @param runnable the code block to execute - * @param keyValuePairs alternating key-value pairs to add to the context - */ - public static void withCustomFields(Runnable runnable, Object... keyValuePairs) { - try { - withCustomFields( - () -> { - runnable.run(); - return null; - }, - keyValuePairs); - } catch (Exception e) { - if (e instanceof RuntimeException) { - throw (RuntimeException) e; - } - throw new RuntimeException(e); - } - } - - // Direct access to the context - simpler than the provider pattern + // Direct access to the context private static SQLCommenterContextAccess getContext() { return SQLCommenterContextHolder.INSTANCE; } @@ -186,28 +69,11 @@ private static SQLCommenterContextAccess createInstance() { } } - // Direct access interface - much simpler than the provider pattern + // Direct access interface private interface SQLCommenterContextAccess { - void put(String key, Object value); - - String get(String key); - - String remove(String key); - - void clear(); - - boolean isEmpty(); - Map getCopyOfContextMap(); - T withSQLCommentFields(Map contextMap, SQLCommenterBlock block); - - T withSQLCommentFields(SQLCommenterBlock block, Object... keyValuePairs); - } - - @FunctionalInterface - private interface SQLCommenterBlock { - T execute() throws Exception; + void setContextMap(Map contextMap); } private static class DirectSQLCommenterContextAccess implements SQLCommenterContextAccess { @@ -217,51 +83,6 @@ public DirectSQLCommenterContextAccess(Class contextClass) { this.contextClass = contextClass; } - @Override - public void put(String key, Object value) { - try { - contextClass.getMethod("put", String.class, Object.class).invoke(null, key, value); - } catch (Exception e) { - throw new RuntimeException("Failed to put SQL comment field", e); - } - } - - @Override - public String get(String key) { - try { - return (String) contextClass.getMethod("get", String.class).invoke(null, key); - } catch (Exception e) { - throw new RuntimeException("Failed to get SQL comment field", e); - } - } - - @Override - public String remove(String key) { - try { - return (String) contextClass.getMethod("remove", String.class).invoke(null, key); - } catch (Exception e) { - throw new RuntimeException("Failed to remove SQL comment field", e); - } - } - - @Override - public void clear() { - try { - contextClass.getMethod("clear").invoke(null); - } catch (Exception e) { - throw new RuntimeException("Failed to clear SQL comment context", e); - } - } - - @Override - public boolean isEmpty() { - try { - return (Boolean) contextClass.getMethod("isEmpty").invoke(null); - } catch (Exception e) { - throw new RuntimeException("Failed to check if SQL comment context is empty", e); - } - } - @Override @SuppressWarnings("unchecked") public Map getCopyOfContextMap() { @@ -273,44 +94,11 @@ public Map getCopyOfContextMap() { } @Override - @SuppressWarnings("unchecked") - public T withSQLCommentFields(Map contextMap, SQLCommenterBlock block) { - try { - Class blockClass = - Class.forName( - "datadog.trace.instrumentation.jdbc.SQLCommenterContext$SQLCommenterBlock"); - Object blockProxy = - java.lang.reflect.Proxy.newProxyInstance( - blockClass.getClassLoader(), - new Class[] {blockClass}, - (proxy, method, args) -> block.execute()); - return (T) - contextClass - .getMethod("withSQLCommentFields", Map.class, blockClass) - .invoke(null, contextMap, blockProxy); - } catch (Exception e) { - throw new RuntimeException("Failed to execute with SQL comment fields", e); - } - } - - @Override - @SuppressWarnings("unchecked") - public T withSQLCommentFields(SQLCommenterBlock block, Object... keyValuePairs) { + public void setContextMap(Map contextMap) { try { - Class blockClass = - Class.forName( - "datadog.trace.instrumentation.jdbc.SQLCommenterContext$SQLCommenterBlock"); - Object blockProxy = - java.lang.reflect.Proxy.newProxyInstance( - blockClass.getClassLoader(), - new Class[] {blockClass}, - (proxy, method, args) -> block.execute()); - return (T) - contextClass - .getMethod("withSQLCommentFields", blockClass, Object[].class) - .invoke(null, blockProxy, keyValuePairs); + contextClass.getMethod("setContextMap", Map.class).invoke(null, contextMap); } catch (Exception e) { - throw new RuntimeException("Failed to execute with SQL comment fields", e); + throw new RuntimeException("Failed to set SQL comment context map", e); } } }