diff --git a/dd-java-agent/instrumentation-testing/src/main/java/datadog/trace/agent/test/assertions/TagsMatcher.java b/dd-java-agent/instrumentation-testing/src/main/java/datadog/trace/agent/test/assertions/TagsMatcher.java index 7fe7cd0645e..c3a5e338f56 100644 --- a/dd-java-agent/instrumentation-testing/src/main/java/datadog/trace/agent/test/assertions/TagsMatcher.java +++ b/dd-java-agent/instrumentation-testing/src/main/java/datadog/trace/agent/test/assertions/TagsMatcher.java @@ -2,6 +2,7 @@ import static datadog.trace.api.DDTags.BASE_SERVICE; import static datadog.trace.api.DDTags.DD_INTEGRATION; +import static datadog.trace.api.DDTags.DD_SVC_SRC; import static datadog.trace.api.DDTags.DJM_ENABLED; import static datadog.trace.api.DDTags.DSM_ENABLED; import static datadog.trace.api.DDTags.ERROR_MSG; @@ -57,6 +58,7 @@ public static TagsMatcher defaultTags() { tagMatchers.put(PARENT_ID, any()); tagMatchers.put(SPAN_LINKS, any()); // this is checked by LinksAsserter tagMatchers.put(DD_INTEGRATION, any()); + tagMatchers.put(DD_SVC_SRC, any()); tagMatchers.put(TRACER_HOST, any()); for (String tagName : REQUIRED_CODE_ORIGIN_TAGS) { diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/build.gradle b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/build.gradle new file mode 100644 index 00000000000..3f820b16c72 --- /dev/null +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/build.gradle @@ -0,0 +1,19 @@ +muzzle { + pass { + group = "io.r2dbc" + module = "r2dbc-spi" + versions = "[1.0.0.RELEASE,)" + } +} + +apply from: "$rootDir/gradle/java.gradle" + +addTestSuiteForDir('latestDepTest', 'test') + +dependencies { + compileOnly group: 'io.r2dbc', name: 'r2dbc-spi', version: '1.0.0.RELEASE' + + testImplementation group: 'io.r2dbc', name: 'r2dbc-spi', version: '1.0.0.RELEASE' + + latestDepTestImplementation group: 'io.r2dbc', name: 'r2dbc-spi', version: '1.+' +} diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/BatchInstrumentation.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/BatchInstrumentation.java new file mode 100644 index 00000000000..a21f6efb98e --- /dev/null +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/BatchInstrumentation.java @@ -0,0 +1,67 @@ +package datadog.trace.instrumentation.r2dbc; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; +import static datadog.trace.instrumentation.r2dbc.R2dbcDecorator.DECORATE; +import static datadog.trace.instrumentation.r2dbc.R2dbcDecorator.R2DBC_BATCH; +import static net.bytebuddy.matcher.ElementMatchers.isMethod; +import static net.bytebuddy.matcher.ElementMatchers.isPublic; +import static net.bytebuddy.matcher.ElementMatchers.takesArguments; + +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import net.bytebuddy.asm.Advice; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.matcher.ElementMatcher; +import org.reactivestreams.Publisher; + +public class BatchInstrumentation + implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice { + + @Override + public String hierarchyMarkerType() { + return "io.r2dbc.spi.Batch"; + } + + @Override + public ElementMatcher hierarchyMatcher() { + return implementsInterface(named("io.r2dbc.spi.Batch")); + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice( + isMethod().and(isPublic()).and(named("execute")).and(takesArguments(0)), + BatchInstrumentation.class.getName() + "$BatchExecuteAdvice"); + } + + public static class BatchExecuteAdvice { + + @Advice.OnMethodEnter(suppress = Throwable.class) + public static AgentScope onEnter() { + AgentSpan span = startSpan("r2dbc", R2DBC_BATCH); + DECORATE.afterStart(span); + return activateSpan(span); + } + + @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) + public static void onExit( + @Advice.Enter final AgentScope scope, + @Advice.Return(readOnly = false) Publisher publisher, + @Advice.Thrown final Throwable throwable) { + AgentSpan span = scope.span(); + if (throwable != null) { + DECORATE.onError(span, throwable); + DECORATE.beforeFinish(span); + scope.close(); + span.finish(); + } else { + publisher = new TracingPublisher<>(publisher, span); + scope.close(); + } + } + } +} diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/ConnectionInstrumentation.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/ConnectionInstrumentation.java new file mode 100644 index 00000000000..d479d3b6980 --- /dev/null +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/ConnectionInstrumentation.java @@ -0,0 +1,77 @@ +package datadog.trace.instrumentation.r2dbc; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static datadog.trace.instrumentation.r2dbc.R2dbcDecorator.INJECT_COMMENT; +import static net.bytebuddy.matcher.ElementMatchers.isMethod; +import static net.bytebuddy.matcher.ElementMatchers.isPublic; +import static net.bytebuddy.matcher.ElementMatchers.takesArgument; +import static net.bytebuddy.matcher.ElementMatchers.takesArguments; + +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.bootstrap.InstrumentationContext; +import io.r2dbc.spi.Connection; +import io.r2dbc.spi.ConnectionMetadata; +import io.r2dbc.spi.Statement; +import net.bytebuddy.asm.Advice; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.matcher.ElementMatcher; + +public class ConnectionInstrumentation + implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice { + + @Override + public String hierarchyMarkerType() { + return "io.r2dbc.spi.Connection"; + } + + @Override + public ElementMatcher hierarchyMatcher() { + return implementsInterface(named("io.r2dbc.spi.Connection")); + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice( + isMethod() + .and(isPublic()) + .and(named("createStatement")) + .and(takesArguments(1)) + .and(takesArgument(0, String.class)), + ConnectionInstrumentation.class.getName() + "$CreateStatementAdvice"); + } + + public static class CreateStatementAdvice { + + @Advice.OnMethodEnter(suppress = Throwable.class) + public static R2dbcConnectionInfo onEnter( + @Advice.This final Connection connection, + @Advice.Argument(value = 0, readOnly = false) String sql) { + final String originalSql = sql; + String dbType = null; + try { + ConnectionMetadata metadata = connection.getMetadata(); + if (metadata != null) { + String productName = metadata.getDatabaseProductName(); + if (productName != null) { + dbType = productName.toLowerCase(); + } + } + } catch (Throwable ignored) { + // Connection may be closed or metadata unavailable + } + if (INJECT_COMMENT) { + sql = R2dbcSQLCommenter.inject(sql, null, dbType, null, null, null); + } + return R2dbcConnectionInfo.of(originalSql, dbType, null, null, null); + } + + @Advice.OnMethodExit(suppress = Throwable.class) + public static void onExit( + @Advice.Enter final R2dbcConnectionInfo info, @Advice.Return final Statement statement) { + if (statement != null && info != null) { + InstrumentationContext.get(Statement.class, R2dbcConnectionInfo.class).put(statement, info); + } + } + } +} diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcConnectionInfo.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcConnectionInfo.java new file mode 100644 index 00000000000..02825ea483f --- /dev/null +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcConnectionInfo.java @@ -0,0 +1,53 @@ +package datadog.trace.instrumentation.r2dbc; + +/** + * Holds connection metadata and the original SQL string associated with an R2DBC Statement. Used + * for span tagging and DBM SQL comment injection. + */ +public final class R2dbcConnectionInfo { + private final String sql; + private final String dbType; + private final String dbInstance; + private final String dbUser; + private final String dbHostname; + + private R2dbcConnectionInfo( + String sql, String dbType, String dbInstance, String dbUser, String dbHostname) { + this.sql = sql; + this.dbType = dbType; + this.dbInstance = dbInstance; + this.dbUser = dbUser; + this.dbHostname = dbHostname; + } + + public String getSql() { + return sql; + } + + public String getDbType() { + return dbType; + } + + public String getDbInstance() { + return dbInstance; + } + + public String getDbUser() { + return dbUser; + } + + public String getDbHostname() { + return dbHostname; + } + + /** Creates an info object with only the SQL string (no connection metadata). */ + public static R2dbcConnectionInfo ofSql(String sql) { + return new R2dbcConnectionInfo(sql, null, null, null, null); + } + + /** Creates an info object with SQL string and connection metadata. */ + public static R2dbcConnectionInfo of( + String sql, String dbType, String dbInstance, String dbUser, String dbHostname) { + return new R2dbcConnectionInfo(sql, dbType, dbInstance, dbUser, dbHostname); + } +} diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcDecorator.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcDecorator.java new file mode 100644 index 00000000000..ac3c5bf1214 --- /dev/null +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcDecorator.java @@ -0,0 +1,129 @@ +package datadog.trace.instrumentation.r2dbc; + +import datadog.trace.api.Config; +import datadog.trace.api.naming.SpanNaming; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes; +import datadog.trace.bootstrap.instrumentation.api.Tags; +import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString; +import datadog.trace.bootstrap.instrumentation.decorator.DBTypeProcessingDatabaseClientDecorator; + +public class R2dbcDecorator extends DBTypeProcessingDatabaseClientDecorator { + public static final R2dbcDecorator DECORATE = new R2dbcDecorator(); + + private static final String R2DBC = "r2dbc"; + public static final CharSequence R2DBC_QUERY = + UTF8BytesString.create(SpanNaming.instance().namingSchema().database().operation(R2DBC)); + public static final CharSequence R2DBC_BATCH = UTF8BytesString.create("r2dbc.batch"); + private static final String SERVICE_NAME = + SpanNaming.instance().namingSchema().database().service(R2DBC); + private static final CharSequence COMPONENT_NAME = UTF8BytesString.create(R2DBC); + + public static final String DBM_PROPAGATION_MODE = Config.get().getDbmPropagationMode(); + + public static final boolean INJECT_COMMENT = + DBM_PROPAGATION_MODE.equals(Config.DBM_PROPAGATION_MODE_FULL) + || DBM_PROPAGATION_MODE.equals(Config.DBM_PROPAGATION_MODE_STATIC) + || DBM_PROPAGATION_MODE.equals(Config.DBM_PROPAGATION_MODE_DYNAMIC_SERVICE); + + public static final boolean INJECT_TRACE_CONTEXT = + DBM_PROPAGATION_MODE.equals(Config.DBM_PROPAGATION_MODE_FULL); + + @Override + protected String[] instrumentationNames() { + return new String[] {R2DBC}; + } + + @Override + protected String service() { + return SERVICE_NAME; + } + + @Override + protected CharSequence component() { + return COMPONENT_NAME; + } + + @Override + protected CharSequence spanType() { + return InternalSpanTypes.SQL; + } + + @Override + protected String dbType() { + return R2DBC; + } + + @Override + protected String dbUser(Void connection) { + return null; + } + + @Override + protected String dbInstance(Void connection) { + return null; + } + + @Override + protected CharSequence dbHostname(Void connection) { + return null; + } + + /** + * Apply connection metadata tags to the span for peer service computation. Overrides db.type with + * the actual database product name when available, and sets peer.hostname and db.instance which + * feed into PeerServiceCalculator. + */ + public void onConnection(AgentSpan span, R2dbcConnectionInfo info) { + if (info != null) { + if (info.getDbType() != null) { + processDatabaseType(span, info.getDbType()); + } + if (info.getDbInstance() != null) { + onInstance(span, info.getDbInstance()); + } + if (info.getDbUser() != null) { + span.setTag(Tags.DB_USER, info.getDbUser()); + } + if (info.getDbHostname() != null) { + span.setTag(Tags.PEER_HOSTNAME, info.getDbHostname()); + } + } + } + + /** Extracts the first SQL keyword (SELECT, INSERT, UPDATE, DELETE, etc.) from a SQL string. */ + public static String extractDbOperation(String sql) { + if (sql == null || sql.isEmpty()) { + return null; + } + // Skip leading whitespace and SQL comments (e.g., /* DBM comment */) + int start = 0; + int len = sql.length(); + while (start < len) { + // Skip whitespace + if (Character.isWhitespace(sql.charAt(start))) { + start++; + continue; + } + // Skip block comments /* ... */ + if (start + 1 < len && sql.charAt(start) == '/' && sql.charAt(start + 1) == '*') { + int endComment = sql.indexOf("*/", start + 2); + if (endComment == -1) { + return null; + } + start = endComment + 2; + continue; + } + break; + } + // Find the end of the first word + int end = start; + while (end < len && !Character.isWhitespace(sql.charAt(end))) { + end++; + } + if (start == end) { + return null; + } + return sql.substring(start, end).toUpperCase(); + } +} diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcInstrumenterModule.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcInstrumenterModule.java new file mode 100644 index 00000000000..e6620363456 --- /dev/null +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcInstrumenterModule.java @@ -0,0 +1,45 @@ +package datadog.trace.instrumentation.r2dbc; + +import com.google.auto.service.AutoService; +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.agent.tooling.InstrumenterModule; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +@AutoService(InstrumenterModule.class) +public class R2dbcInstrumenterModule extends InstrumenterModule.Tracing { + + public R2dbcInstrumenterModule() { + super("r2dbc"); + } + + @Override + public String[] helperClassNames() { + return new String[] { + packageName + ".R2dbcConnectionInfo", + packageName + ".R2dbcDecorator", + packageName + ".R2dbcSQLCommenter", + packageName + ".TracingPublisher", + packageName + ".TracingPublisher$TracingSubscriber", + packageName + ".TracingPublisher$TracingSubscriber$TracingSubscription", + }; + } + + @Override + public Map contextStore() { + Map contextStore = new HashMap<>(); + contextStore.put("io.r2dbc.spi.Statement", packageName + ".R2dbcConnectionInfo"); + return contextStore; + } + + @Override + public List typeInstrumentations() { + List instrumenters = new ArrayList<>(3); + instrumenters.add(new ConnectionInstrumentation()); + instrumenters.add(new StatementInstrumentation()); + instrumenters.add(new BatchInstrumentation()); + return instrumenters; + } +} diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcSQLCommenter.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcSQLCommenter.java new file mode 100644 index 00000000000..c6c5e2d3123 --- /dev/null +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcSQLCommenter.java @@ -0,0 +1,65 @@ +package datadog.trace.instrumentation.r2dbc; + +import datadog.trace.bootstrap.instrumentation.dbm.SharedDBCommenter; + +/** + * Injects DBM trace context as SQL comments for R2DBC queries. Delegates comment content building + * to {@link SharedDBCommenter} and wraps the result in SQL block comment delimiters. + */ +public final class R2dbcSQLCommenter { + private static final String OPEN_COMMENT = "/*"; + private static final String CLOSE_COMMENT = "*/"; + + private R2dbcSQLCommenter() {} + + /** + * Injects a DBM SQL comment into the given query. By default the comment is prepended. + * + * @param sql the original SQL string + * @param dbService the database service name + * @param dbType the database type (e.g. "postgresql", "mysql") + * @param hostname the database hostname + * @param dbName the database name + * @param traceParent the W3C traceparent string (null for service-only mode) + * @return the SQL with injected comment, or the original if injection is not applicable + */ + public static String inject( + String sql, + String dbService, + String dbType, + String hostname, + String dbName, + String traceParent) { + if (sql == null || sql.isEmpty()) { + return sql; + } + // Check if a DD comment is already present + if (hasDDComment(sql)) { + return sql; + } + String commentContent = + SharedDBCommenter.buildComment(dbService, dbType, hostname, dbName, traceParent); + if (commentContent == null) { + return sql; + } + // Prepend the comment to the SQL + StringBuilder sb = new StringBuilder(sql.length() + commentContent.length() + 5); + sb.append(OPEN_COMMENT); + sb.append(commentContent); + sb.append(CLOSE_COMMENT); + sb.append(' '); + sb.append(sql); + return sb.toString(); + } + + private static boolean hasDDComment(String sql) { + if (!sql.startsWith(OPEN_COMMENT)) { + return false; + } + int endIdx = sql.indexOf(CLOSE_COMMENT); + if (endIdx > 2) { + return SharedDBCommenter.containsTraceComment(sql, 2, endIdx); + } + return false; + } +} diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/StatementInstrumentation.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/StatementInstrumentation.java new file mode 100644 index 00000000000..b1d0359060e --- /dev/null +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/StatementInstrumentation.java @@ -0,0 +1,96 @@ +package datadog.trace.instrumentation.r2dbc; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; +import static datadog.trace.instrumentation.r2dbc.R2dbcDecorator.DECORATE; +import static datadog.trace.instrumentation.r2dbc.R2dbcDecorator.INJECT_COMMENT; +import static datadog.trace.instrumentation.r2dbc.R2dbcDecorator.INJECT_TRACE_CONTEXT; +import static datadog.trace.instrumentation.r2dbc.R2dbcDecorator.R2DBC_QUERY; +import static net.bytebuddy.matcher.ElementMatchers.isMethod; +import static net.bytebuddy.matcher.ElementMatchers.isPublic; +import static net.bytebuddy.matcher.ElementMatchers.takesArguments; + +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.bootstrap.InstrumentationContext; +import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.InstrumentationTags; +import datadog.trace.bootstrap.instrumentation.api.Tags; +import io.r2dbc.spi.Statement; +import net.bytebuddy.asm.Advice; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.matcher.ElementMatcher; +import org.reactivestreams.Publisher; + +public class StatementInstrumentation + implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice { + + @Override + public String hierarchyMarkerType() { + return "io.r2dbc.spi.Statement"; + } + + @Override + public ElementMatcher hierarchyMatcher() { + return implementsInterface(named("io.r2dbc.spi.Statement")); + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice( + isMethod().and(isPublic()).and(named("execute")).and(takesArguments(0)), + StatementInstrumentation.class.getName() + "$StatementExecuteAdvice"); + } + + public static class StatementExecuteAdvice { + + @Advice.OnMethodEnter(suppress = Throwable.class) + public static AgentScope onEnter(@Advice.This final Statement statement) { + R2dbcConnectionInfo info = + InstrumentationContext.get(Statement.class, R2dbcConnectionInfo.class).get(statement); + + String sql = info != null ? info.getSql() : null; + + AgentSpan span = startSpan("r2dbc", R2DBC_QUERY); + DECORATE.afterStart(span); + + if (sql != null) { + DECORATE.onStatement(span, sql); + String dbOperation = R2dbcDecorator.extractDbOperation(sql); + if (dbOperation != null) { + span.setTag(Tags.DB_OPERATION, dbOperation); + } + } + + if (info != null) { + DECORATE.onConnection(span, info); + } + + if (INJECT_COMMENT && INJECT_TRACE_CONTEXT && info != null) { + span.setTag(InstrumentationTags.DBM_TRACE_INJECTED, true); + } + + return activateSpan(span); + } + + @SuppressWarnings("unchecked") + @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) + public static void onExit( + @Advice.Enter final AgentScope scope, + @Advice.Return(readOnly = false) Publisher publisher, + @Advice.Thrown final Throwable throwable) { + AgentSpan span = scope.span(); + if (throwable != null) { + DECORATE.onError(span, throwable); + DECORATE.beforeFinish(span); + scope.close(); + span.finish(); + } else { + publisher = new TracingPublisher<>(publisher, span); + scope.close(); + } + } + } +} diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/TracingPublisher.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/TracingPublisher.java new file mode 100644 index 00000000000..99786becc39 --- /dev/null +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/TracingPublisher.java @@ -0,0 +1,106 @@ +package datadog.trace.instrumentation.r2dbc; + +import static datadog.trace.instrumentation.r2dbc.R2dbcDecorator.DECORATE; + +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import java.util.concurrent.atomic.AtomicBoolean; +import org.reactivestreams.Publisher; +import org.reactivestreams.Subscriber; +import org.reactivestreams.Subscription; + +/** + * A wrapping Publisher that finishes the associated span when the downstream subscriber receives + * onComplete, onError, or cancels its subscription. + */ +public final class TracingPublisher implements Publisher { + + private final Publisher delegate; + private final AgentSpan span; + + public TracingPublisher(Publisher delegate, AgentSpan span) { + this.delegate = delegate; + this.span = span; + } + + @Override + public void subscribe(Subscriber subscriber) { + delegate.subscribe(new TracingSubscriber<>(subscriber, span)); + } + + static final class TracingSubscriber implements Subscriber { + + private final Subscriber delegate; + private final AgentSpan span; + // Guards against finishing the span more than once when terminal signals (onComplete/onError) + // race with a downstream cancel(). + private final AtomicBoolean finished = new AtomicBoolean(false); + + TracingSubscriber(Subscriber delegate, AgentSpan span) { + this.delegate = delegate; + this.span = span; + } + + private void finishSpan() { + if (finished.compareAndSet(false, true)) { + DECORATE.beforeFinish(span); + span.finish(); + } + } + + @Override + public void onSubscribe(Subscription s) { + // Wrap the subscription so a downstream cancel() (e.g. take(1), timeout, disconnect) finishes + // the span instead of leaving it open indefinitely. + delegate.onSubscribe(new TracingSubscription(s)); + } + + @Override + public void onNext(T t) { + delegate.onNext(t); + } + + final class TracingSubscription implements Subscription { + private final Subscription delegate; + + TracingSubscription(Subscription delegate) { + this.delegate = delegate; + } + + @Override + public void request(long n) { + delegate.request(n); + } + + @Override + public void cancel() { + try { + delegate.cancel(); + } finally { + finishSpan(); + } + } + } + + @Override + public void onError(Throwable t) { + try { + if (finished.compareAndSet(false, true)) { + DECORATE.onError(span, t); + DECORATE.beforeFinish(span); + span.finish(); + } + } finally { + delegate.onError(t); + } + } + + @Override + public void onComplete() { + try { + finishSpan(); + } finally { + delegate.onComplete(); + } + } + } +} diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/datadog/trace/instrumentation/r2dbc/R2dbcDBMForkedTest.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/datadog/trace/instrumentation/r2dbc/R2dbcDBMForkedTest.java new file mode 100644 index 00000000000..56ecebd2791 --- /dev/null +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/datadog/trace/instrumentation/r2dbc/R2dbcDBMForkedTest.java @@ -0,0 +1,253 @@ +package datadog.trace.instrumentation.r2dbc; + +import static datadog.trace.agent.test.assertions.SpanMatcher.span; +import static datadog.trace.agent.test.assertions.TagsMatcher.defaultTags; +import static datadog.trace.agent.test.assertions.TagsMatcher.error; +import static datadog.trace.agent.test.assertions.TagsMatcher.tag; +import static datadog.trace.agent.test.assertions.TraceMatcher.trace; +import static datadog.trace.agent.test.utils.TraceUtils.runnableUnderTrace; +import static datadog.trace.bootstrap.instrumentation.api.Tags.COMPONENT; +import static datadog.trace.bootstrap.instrumentation.api.Tags.DB_OPERATION; +import static datadog.trace.bootstrap.instrumentation.api.Tags.DB_TYPE; +import static datadog.trace.bootstrap.instrumentation.api.Tags.SPAN_KIND; +import static datadog.trace.bootstrap.instrumentation.api.Tags.SPAN_KIND_CLIENT; +import static datadog.trace.test.junit.utils.assertions.Matchers.is; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import datadog.trace.agent.test.AbstractInstrumentationTest; +import datadog.trace.api.DDSpanTypes; +import datadog.trace.bootstrap.instrumentation.api.InstrumentationTags; +import io.r2dbc.spi.Result; +import io.r2dbc.spi.Statement; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import org.junit.jupiter.api.Test; +import org.reactivestreams.Publisher; +import org.reactivestreams.Subscriber; +import org.reactivestreams.Subscription; +import r2dbc.test.TestConnection; +import r2dbc.test.TestStatement; + +/** + * Forked test that runs with DBM propagation mode set to "full". This enables SQL comment injection + * and the _dd.dbm_trace_injected tag on database spans. + * + *

Uses ForkedTest suffix so the test runs in a separate JVM where the DBM config is applied + * before R2dbcDecorator's static fields are initialized. + */ +@datadog.trace.test.junit.utils.config.WithConfig(key = "dbm.propagation.mode", value = "full") +class R2dbcDBMForkedTest extends AbstractInstrumentationTest { + + @Test + void statementExecuteInjectsDBMCommentAndSetsTraceTags() { + TestConnection connection = new TestConnection(); + Statement statement = connection.createStatement("SELECT * FROM users"); + + // Verify SQL comment was injected into the statement received by the driver + String injectedSql = ((TestStatement) statement).getSql(); + assertNotNull(injectedSql, "SQL should not be null"); + assertTrue( + injectedSql.startsWith("/*"), "SQL should start with DBM comment, got: " + injectedSql); + assertTrue( + injectedSql.contains("ddps="), + "DBM comment should contain parent service tag, got: " + injectedSql); + assertTrue( + injectedSql.endsWith("SELECT * FROM users"), + "SQL should end with original query, got: " + injectedSql); + + runnableUnderTrace( + "parent", + () -> { + Publisher publisher = statement.execute(); + blockOnPublisher(publisher); + }); + + // Verify the span has the original SQL as resource (not the DBM-commented version) + // and includes the _dd.dbm_trace_injected tag + assertTraces( + trace( + span().root().operationName("parent"), + span() + .childOfPrevious() + .operationName("r2dbc.query") + .resourceName("SELECT * FROM users") + .type(DDSpanTypes.SQL) + .measured() + .tags( + defaultTags(), + tag(SPAN_KIND, is(SPAN_KIND_CLIENT)), + tag(COMPONENT, is("r2dbc")), + tag(DB_TYPE, is("testdb")), + tag(DB_OPERATION, is("SELECT")), + tag(InstrumentationTags.DBM_TRACE_INJECTED, is(true))))); + } + + @Test + void statementExecuteWithInsertInjectsDBMComment() { + TestConnection connection = new TestConnection(); + Statement statement = connection.createStatement("INSERT INTO users (name) VALUES ('Alice')"); + + // Verify SQL comment injection + String injectedSql = ((TestStatement) statement).getSql(); + assertTrue( + injectedSql.startsWith("/*"), "SQL should start with DBM comment, got: " + injectedSql); + assertTrue( + injectedSql.endsWith("INSERT INTO users (name) VALUES ('Alice')"), + "SQL should end with original query, got: " + injectedSql); + + runnableUnderTrace( + "parent", + () -> { + Publisher publisher = statement.execute(); + blockOnPublisher(publisher); + }); + + assertTraces( + trace( + span().root().operationName("parent"), + span() + .childOfPrevious() + .operationName("r2dbc.query") + .resourceName("INSERT INTO users (name) VALUES ('Alice')") + .type(DDSpanTypes.SQL) + .measured() + .tags( + defaultTags(), + tag(SPAN_KIND, is(SPAN_KIND_CLIENT)), + tag(COMPONENT, is("r2dbc")), + tag(DB_TYPE, is("testdb")), + tag(DB_OPERATION, is("INSERT")), + tag(InstrumentationTags.DBM_TRACE_INJECTED, is(true))))); + } + + @Test + void statementExecuteWithErrorStillHasDBMTags() { + TestConnection connection = new TestConnection(); + Statement statement = connection.createStatement("SELECT * FROM nonexistent"); + ((TestStatement) statement).setFailOnExecute(true); + + try { + runnableUnderTrace( + "parent", + () -> { + Publisher publisher = statement.execute(); + blockOnPublisher(publisher); + }); + } catch (RuntimeException ignored) { + // expected + } + + assertTraces( + trace( + span().root().operationName("parent").error(), + span() + .childOfPrevious() + .operationName("r2dbc.query") + .resourceName("SELECT * FROM nonexistent") + .type(DDSpanTypes.SQL) + .error() + .measured() + .tags( + defaultTags(), + tag(SPAN_KIND, is(SPAN_KIND_CLIENT)), + tag(COMPONENT, is("r2dbc")), + tag(DB_TYPE, is("testdb")), + tag(DB_OPERATION, is("SELECT")), + tag(InstrumentationTags.DBM_TRACE_INJECTED, is(true)), + error(RuntimeException.class, "query failed"), + tag("error.type", is(RuntimeException.class.getName())), + tag("error.message", is("query failed"))))); + } + + @Test + void dbmCommentIsNotInjectedTwice() { + TestConnection connection = new TestConnection(); + Statement statement = connection.createStatement("SELECT 1"); + + // Get the injected SQL + String injectedSql = ((TestStatement) statement).getSql(); + assertTrue(injectedSql.startsWith("/*"), "Should have DBM comment"); + + // Count how many comment blocks there are + int commentCount = 0; + int idx = 0; + while ((idx = injectedSql.indexOf("/*", idx)) != -1) { + commentCount++; + idx += 2; + } + assertTrue( + commentCount == 1, + "Should have exactly one comment block, found " + commentCount + " in: " + injectedSql); + + runnableUnderTrace( + "parent", + () -> { + Publisher publisher = statement.execute(); + blockOnPublisher(publisher); + }); + + // Span resource should use the original SQL, not the commented one + assertTraces( + trace( + span().root().operationName("parent"), + span() + .childOfPrevious() + .operationName("r2dbc.query") + .resourceName("SELECT 1") + .type(DDSpanTypes.SQL) + .measured() + .tags( + defaultTags(), + tag(SPAN_KIND, is(SPAN_KIND_CLIENT)), + tag(COMPONENT, is("r2dbc")), + tag(DB_TYPE, is("testdb")), + tag(DB_OPERATION, is("SELECT")), + tag(InstrumentationTags.DBM_TRACE_INJECTED, is(true))))); + } + + /** Subscribes to a Publisher and blocks until completion or error. */ + private static void blockOnPublisher(Publisher publisher) { + List results = new ArrayList<>(); + RuntimeException[] error = new RuntimeException[1]; + CountDownLatch latch = new CountDownLatch(1); + publisher.subscribe( + new Subscriber() { + @Override + public void onSubscribe(Subscription s) { + s.request(Long.MAX_VALUE); + } + + @Override + public void onNext(T t) { + results.add(t); + } + + @Override + public void onError(Throwable t) { + if (t instanceof RuntimeException) { + error[0] = (RuntimeException) t; + } else { + error[0] = new RuntimeException(t); + } + latch.countDown(); + } + + @Override + public void onComplete() { + latch.countDown(); + } + }); + try { + latch.await(5, TimeUnit.SECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException("Interrupted waiting for publisher", e); + } + if (error[0] != null) { + throw error[0]; + } + } +} diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/datadog/trace/instrumentation/r2dbc/R2dbcInstrumentationTest.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/datadog/trace/instrumentation/r2dbc/R2dbcInstrumentationTest.java new file mode 100644 index 00000000000..186e199ce48 --- /dev/null +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/datadog/trace/instrumentation/r2dbc/R2dbcInstrumentationTest.java @@ -0,0 +1,314 @@ +package datadog.trace.instrumentation.r2dbc; + +import static datadog.trace.agent.test.assertions.SpanMatcher.span; +import static datadog.trace.agent.test.assertions.TagsMatcher.defaultTags; +import static datadog.trace.agent.test.assertions.TagsMatcher.error; +import static datadog.trace.agent.test.assertions.TagsMatcher.tag; +import static datadog.trace.agent.test.assertions.TraceMatcher.trace; +import static datadog.trace.agent.test.utils.TraceUtils.runnableUnderTrace; +import static datadog.trace.bootstrap.instrumentation.api.Tags.COMPONENT; +import static datadog.trace.bootstrap.instrumentation.api.Tags.DB_OPERATION; +import static datadog.trace.bootstrap.instrumentation.api.Tags.DB_TYPE; +import static datadog.trace.bootstrap.instrumentation.api.Tags.SPAN_KIND; +import static datadog.trace.bootstrap.instrumentation.api.Tags.SPAN_KIND_CLIENT; +import static datadog.trace.test.junit.utils.assertions.Matchers.is; +import static org.junit.jupiter.api.Assertions.assertFalse; + +import datadog.trace.agent.test.AbstractInstrumentationTest; +import datadog.trace.agent.test.assertions.TraceMatcher; +import datadog.trace.api.DDSpanTypes; +import io.r2dbc.spi.Batch; +import io.r2dbc.spi.Result; +import io.r2dbc.spi.Statement; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import org.junit.jupiter.api.Test; +import org.reactivestreams.Publisher; +import org.reactivestreams.Subscriber; +import org.reactivestreams.Subscription; +import r2dbc.test.TestBatch; +import r2dbc.test.TestConnection; +import r2dbc.test.TestStatement; + +class R2dbcInstrumentationTest extends AbstractInstrumentationTest { + + @Test + void statementExecuteCreatesSpan() { + TestConnection connection = new TestConnection(); + Statement statement = connection.createStatement("SELECT * FROM users"); + + List[] results = new List[1]; + runnableUnderTrace( + "parent", + () -> { + Publisher publisher = statement.execute(); + results[0] = blockOnPublisher(publisher); + }); + assertFalse(results[0].isEmpty(), "Statement execute should return results"); + + assertTraces( + trace( + span().root().operationName("parent"), + span() + .childOfPrevious() + .operationName("r2dbc.query") + .resourceName("SELECT * FROM users") + .type(DDSpanTypes.SQL) + .measured() + .tags( + defaultTags(), + tag(SPAN_KIND, is(SPAN_KIND_CLIENT)), + tag(COMPONENT, is("r2dbc")), + tag(DB_TYPE, is("testdb")), + tag(DB_OPERATION, is("SELECT"))))); + } + + @Test + void statementExecuteWithInsert() { + TestConnection connection = new TestConnection(); + Statement statement = connection.createStatement("INSERT INTO users (name) VALUES ($1)"); + + List[] results = new List[1]; + runnableUnderTrace( + "parent", + () -> { + Publisher publisher = statement.execute(); + results[0] = blockOnPublisher(publisher); + }); + assertFalse(results[0].isEmpty(), "Statement execute should return results"); + + assertTraces( + trace( + span().root().operationName("parent"), + span() + .childOfPrevious() + .operationName("r2dbc.query") + .resourceName("INSERT INTO users (name) VALUES ($1)") + .type(DDSpanTypes.SQL) + .measured() + .tags( + defaultTags(), + tag(SPAN_KIND, is(SPAN_KIND_CLIENT)), + tag(COMPONENT, is("r2dbc")), + tag(DB_TYPE, is("testdb")), + tag(DB_OPERATION, is("INSERT"))))); + } + + @Test + void statementExecuteWithError() { + TestConnection connection = new TestConnection(); + Statement statement = connection.createStatement("SELECT * FROM nonexistent"); + ((TestStatement) statement).setFailOnExecute(true); + + try { + runnableUnderTrace( + "parent", + () -> { + Publisher publisher = statement.execute(); + blockOnPublisher(publisher); + }); + } catch (RuntimeException ignored) { + // expected + } + + assertTraces( + trace( + span().root().operationName("parent").error(), + span() + .childOfPrevious() + .operationName("r2dbc.query") + .resourceName("SELECT * FROM nonexistent") + .type(DDSpanTypes.SQL) + .error() + .measured() + .tags( + defaultTags(), + tag(SPAN_KIND, is(SPAN_KIND_CLIENT)), + tag(COMPONENT, is("r2dbc")), + tag(DB_TYPE, is("testdb")), + tag(DB_OPERATION, is("SELECT")), + error(RuntimeException.class, "query failed")))); + } + + @Test + void batchExecuteCreatesSpan() { + TestConnection connection = new TestConnection(); + Batch batch = connection.createBatch(); + batch.add("INSERT INTO users (name) VALUES ('Alice')"); + batch.add("INSERT INTO users (name) VALUES ('Bob')"); + + List[] results = new List[1]; + runnableUnderTrace( + "parent", + () -> { + Publisher publisher = batch.execute(); + results[0] = blockOnPublisher(publisher); + }); + assertFalse(results[0].isEmpty(), "Batch execute should return results"); + + assertTraces( + trace( + span().root().operationName("parent"), + span() + .childOfPrevious() + .operationName("r2dbc.batch") + .resourceName("r2dbc.batch") + .type(DDSpanTypes.SQL) + .measured() + .tags( + defaultTags(), + tag(SPAN_KIND, is(SPAN_KIND_CLIENT)), + tag(COMPONENT, is("r2dbc")), + tag(DB_TYPE, is("r2dbc"))))); + } + + @Test + void batchExecuteWithError() { + TestConnection connection = new TestConnection(); + TestBatch batch = (TestBatch) connection.createBatch(); + batch.add("INSERT INTO users (name) VALUES ('Alice')"); + batch.setFailOnExecute(true); + + try { + runnableUnderTrace( + "parent", + () -> { + Publisher publisher = batch.execute(); + blockOnPublisher(publisher); + }); + } catch (RuntimeException ignored) { + // expected + } + + assertTraces( + trace( + span().root().operationName("parent").error(), + span() + .childOfPrevious() + .operationName("r2dbc.batch") + .resourceName("r2dbc.batch") + .type(DDSpanTypes.SQL) + .error() + .measured() + .tags( + defaultTags(), + tag(SPAN_KIND, is(SPAN_KIND_CLIENT)), + tag(COMPONENT, is("r2dbc")), + tag(DB_TYPE, is("r2dbc")), + error(RuntimeException.class, "batch failed")))); + } + + @Test + void multipleStatementExecutesCreateSeparateSpans() { + TestConnection connection = new TestConnection(); + Statement statement1 = connection.createStatement("SELECT * FROM users"); + Statement statement2 = connection.createStatement("UPDATE users SET name = 'test'"); + + runnableUnderTrace( + "parent", + () -> { + blockOnPublisher(statement1.execute()); + blockOnPublisher(statement2.execute()); + }); + + assertTraces( + trace( + TraceMatcher.SORT_BY_ANCESTRY, + span().root().operationName("parent"), + span() + .childOfPrevious() + .operationName("r2dbc.query") + .resourceName("SELECT * FROM users") + .type(DDSpanTypes.SQL) + .measured() + .tags( + defaultTags(), + tag(SPAN_KIND, is(SPAN_KIND_CLIENT)), + tag(COMPONENT, is("r2dbc")), + tag(DB_TYPE, is("testdb")), + tag(DB_OPERATION, is("SELECT"))), + span() + .childOfIndex(0) + .operationName("r2dbc.query") + .resourceName("UPDATE users SET name = 'test'") + .type(DDSpanTypes.SQL) + .measured() + .tags( + defaultTags(), + tag(SPAN_KIND, is(SPAN_KIND_CLIENT)), + tag(COMPONENT, is("r2dbc")), + tag(DB_TYPE, is("testdb")), + tag(DB_OPERATION, is("UPDATE"))))); + } + + @Test + void statementExecuteWithNoParentSpan() { + TestConnection connection = new TestConnection(); + Statement statement = connection.createStatement("SELECT 1"); + + Publisher publisher = statement.execute(); + List results = blockOnPublisher(publisher); + assertFalse(results.isEmpty(), "Statement execute should return results"); + + assertTraces( + trace( + span() + .root() + .operationName("r2dbc.query") + .resourceName("SELECT 1") + .type(DDSpanTypes.SQL) + .measured() + .tags( + defaultTags(), + tag(SPAN_KIND, is(SPAN_KIND_CLIENT)), + tag(COMPONENT, is("r2dbc")), + tag(DB_TYPE, is("testdb")), + tag(DB_OPERATION, is("SELECT"))))); + } + + /** Subscribes to a Publisher and blocks until completion or error. Returns collected results. */ + private static List blockOnPublisher(Publisher publisher) { + List results = new ArrayList<>(); + RuntimeException[] error = new RuntimeException[1]; + CountDownLatch latch = new CountDownLatch(1); + publisher.subscribe( + new Subscriber() { + @Override + public void onSubscribe(Subscription s) { + s.request(Long.MAX_VALUE); + } + + @Override + public void onNext(T t) { + results.add(t); + } + + @Override + public void onError(Throwable t) { + if (t instanceof RuntimeException) { + error[0] = (RuntimeException) t; + } else { + error[0] = new RuntimeException(t); + } + latch.countDown(); + } + + @Override + public void onComplete() { + latch.countDown(); + } + }); + try { + latch.await(5, TimeUnit.SECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException("Interrupted waiting for publisher", e); + } + if (error[0] != null) { + throw error[0]; + } + return results; + } +} diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestBatch.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestBatch.java new file mode 100644 index 00000000000..f62e0bf3f1c --- /dev/null +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestBatch.java @@ -0,0 +1,54 @@ +package r2dbc.test; + +import io.r2dbc.spi.Batch; +import io.r2dbc.spi.Result; +import java.util.ArrayList; +import java.util.List; +import org.reactivestreams.Publisher; +import org.reactivestreams.Subscription; + +public class TestBatch implements Batch { + private final List statements = new ArrayList<>(); + private boolean failOnExecute; + + public void setFailOnExecute(boolean fail) { + this.failOnExecute = fail; + } + + @Override + public Batch add(String sql) { + statements.add(sql); + return this; + } + + @Override + public Publisher execute() { + if (failOnExecute) { + return subscriber -> { + subscriber.onSubscribe( + new Subscription() { + @Override + public void request(long n) { + subscriber.onError(new RuntimeException("batch failed")); + } + + @Override + public void cancel() {} + }); + }; + } + return subscriber -> { + subscriber.onSubscribe( + new Subscription() { + @Override + public void request(long n) { + subscriber.onNext(new TestResult()); + subscriber.onComplete(); + } + + @Override + public void cancel() {} + }); + }; + } +} diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestConnection.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestConnection.java new file mode 100644 index 00000000000..ae9ba2d510b --- /dev/null +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestConnection.java @@ -0,0 +1,141 @@ +package r2dbc.test; + +import io.r2dbc.spi.Batch; +import io.r2dbc.spi.Connection; +import io.r2dbc.spi.ConnectionMetadata; +import io.r2dbc.spi.IsolationLevel; +import io.r2dbc.spi.Statement; +import io.r2dbc.spi.TransactionDefinition; +import io.r2dbc.spi.ValidationDepth; +import java.time.Duration; +import org.reactivestreams.Publisher; +import org.reactivestreams.Subscription; + +public class TestConnection implements Connection { + @Override + public Publisher beginTransaction() { + return emptyPublisher(); + } + + @Override + public Publisher beginTransaction(TransactionDefinition definition) { + return emptyPublisher(); + } + + @Override + public Publisher close() { + return emptyPublisher(); + } + + @Override + public Publisher commitTransaction() { + return emptyPublisher(); + } + + @Override + public Batch createBatch() { + return new TestBatch(); + } + + @Override + public Publisher createSavepoint(String name) { + return emptyPublisher(); + } + + @Override + public Statement createStatement(String sql) { + return new TestStatement(sql); + } + + @Override + public boolean isAutoCommit() { + return true; + } + + @Override + public ConnectionMetadata getMetadata() { + return new ConnectionMetadata() { + @Override + public String getDatabaseProductName() { + return "TestDB"; + } + + @Override + public String getDatabaseVersion() { + return "1.0"; + } + }; + } + + @Override + public IsolationLevel getTransactionIsolationLevel() { + return IsolationLevel.READ_COMMITTED; + } + + @Override + public Publisher releaseSavepoint(String name) { + return emptyPublisher(); + } + + @Override + public Publisher rollbackTransaction() { + return emptyPublisher(); + } + + @Override + public Publisher rollbackTransactionToSavepoint(String name) { + return emptyPublisher(); + } + + @Override + public Publisher setAutoCommit(boolean autoCommit) { + return emptyPublisher(); + } + + @Override + public Publisher setLockWaitTimeout(Duration timeout) { + return emptyPublisher(); + } + + @Override + public Publisher setStatementTimeout(Duration timeout) { + return emptyPublisher(); + } + + @Override + public Publisher setTransactionIsolationLevel(IsolationLevel isolationLevel) { + return emptyPublisher(); + } + + @Override + public Publisher validate(ValidationDepth depth) { + return subscriber -> { + subscriber.onSubscribe( + new Subscription() { + @Override + public void request(long n) { + subscriber.onNext(true); + subscriber.onComplete(); + } + + @Override + public void cancel() {} + }); + }; + } + + private static Publisher emptyPublisher() { + return subscriber -> { + subscriber.onSubscribe( + new Subscription() { + @Override + public void request(long n) { + subscriber.onComplete(); + } + + @Override + public void cancel() {} + }); + }; + } +} diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestConnectionFactory.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestConnectionFactory.java new file mode 100644 index 00000000000..0ff3815d8ee --- /dev/null +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestConnectionFactory.java @@ -0,0 +1,31 @@ +package r2dbc.test; + +import io.r2dbc.spi.Connection; +import io.r2dbc.spi.ConnectionFactory; +import io.r2dbc.spi.ConnectionFactoryMetadata; +import org.reactivestreams.Publisher; +import org.reactivestreams.Subscription; + +public class TestConnectionFactory implements ConnectionFactory { + @Override + public Publisher create() { + return subscriber -> { + subscriber.onSubscribe( + new Subscription() { + @Override + public void request(long n) { + subscriber.onNext(new TestConnection()); + subscriber.onComplete(); + } + + @Override + public void cancel() {} + }); + }; + } + + @Override + public ConnectionFactoryMetadata getMetadata() { + return () -> "TestDB"; + } +} diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestResult.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestResult.java new file mode 100644 index 00000000000..531b5489368 --- /dev/null +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestResult.java @@ -0,0 +1,67 @@ +package r2dbc.test; + +import io.r2dbc.spi.Result; +import io.r2dbc.spi.Row; +import io.r2dbc.spi.RowMetadata; +import java.util.function.BiFunction; +import java.util.function.Function; +import java.util.function.Predicate; +import org.reactivestreams.Publisher; +import org.reactivestreams.Subscription; + +public class TestResult implements Result { + @Override + public Publisher getRowsUpdated() { + return subscriber -> { + subscriber.onSubscribe( + new Subscription() { + @Override + public void request(long n) { + subscriber.onNext(0L); + subscriber.onComplete(); + } + + @Override + public void cancel() {} + }); + }; + } + + @Override + public Publisher map(BiFunction mappingFunction) { + return subscriber -> { + subscriber.onSubscribe( + new Subscription() { + @Override + public void request(long n) { + subscriber.onComplete(); + } + + @Override + public void cancel() {} + }); + }; + } + + @Override + public Result filter(Predicate filter) { + return this; + } + + @Override + public Publisher flatMap( + Function> mappingFunction) { + return subscriber -> { + subscriber.onSubscribe( + new Subscription() { + @Override + public void request(long n) { + subscriber.onComplete(); + } + + @Override + public void cancel() {} + }); + }; + } +} diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestStatement.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestStatement.java new file mode 100644 index 00000000000..1ba5136d330 --- /dev/null +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestStatement.java @@ -0,0 +1,79 @@ +package r2dbc.test; + +import io.r2dbc.spi.Result; +import io.r2dbc.spi.Statement; +import org.reactivestreams.Publisher; +import org.reactivestreams.Subscription; + +public class TestStatement implements Statement { + private final String sql; + private boolean failOnExecute; + + public TestStatement(String sql) { + this.sql = sql; + } + + public String getSql() { + return sql; + } + + public void setFailOnExecute(boolean fail) { + this.failOnExecute = fail; + } + + @Override + public Statement add() { + return this; + } + + @Override + public Statement bind(int index, Object value) { + return this; + } + + @Override + public Statement bind(String name, Object value) { + return this; + } + + @Override + public Statement bindNull(int index, Class type) { + return this; + } + + @Override + public Statement bindNull(String name, Class type) { + return this; + } + + @Override + public Publisher execute() { + if (failOnExecute) { + return subscriber -> { + subscriber.onSubscribe( + new Subscription() { + @Override + public void request(long n) { + subscriber.onError(new RuntimeException("query failed")); + } + + @Override + public void cancel() {} + }); + }; + } + return subscriber -> { + subscriber.onSubscribe( + new Subscription() { + @Override + public void request(long n) { + subscriber.onNext(new TestResult()); + subscriber.onComplete(); + } + + @Override + public void cancel() {} + }); + }; + } +} diff --git a/metadata/agent-jar-checks.properties b/metadata/agent-jar-checks.properties index fd04d637e35..1b03884d0bb 100644 --- a/metadata/agent-jar-checks.properties +++ b/metadata/agent-jar-checks.properties @@ -144,6 +144,7 @@ expected.integrations = IastInstrumentation,\ play-ws,\ protobuf,\ quartz,\ + r2dbc,\ ratpack,\ ratpack-request-body,\ reactive-streams,\ diff --git a/metadata/supported-configurations.json b/metadata/supported-configurations.json index 850c6055d58..821faad6f93 100644 --- a/metadata/supported-configurations.json +++ b/metadata/supported-configurations.json @@ -9065,6 +9065,30 @@ "aliases": ["DD_TRACE_INTEGRATION_QUARTZ_ENABLED", "DD_INTEGRATION_QUARTZ_ENABLED"] } ], + "DD_TRACE_R2DBC_ANALYTICS_ENABLED": [ + { + "version": "A", + "type": "boolean", + "default": "false", + "aliases": ["DD_R2DBC_ANALYTICS_ENABLED"] + } + ], + "DD_TRACE_R2DBC_ANALYTICS_SAMPLE_RATE": [ + { + "version": "A", + "type": "decimal", + "default": "1.0", + "aliases": ["DD_R2DBC_ANALYTICS_SAMPLE_RATE"] + } + ], + "DD_TRACE_R2DBC_ENABLED": [ + { + "version": "A", + "type": "boolean", + "default": "true", + "aliases": ["DD_TRACE_INTEGRATION_R2DBC_ENABLED", "DD_INTEGRATION_R2DBC_ENABLED"] + } + ], "DD_TRACE_RABBITMQ_ANALYTICS_ENABLED": [ { "version": "A", diff --git a/settings.gradle.kts b/settings.gradle.kts index 86aa23cc9d4..4df106170de 100644 --- a/settings.gradle.kts +++ b/settings.gradle.kts @@ -537,6 +537,7 @@ include( ":dd-java-agent:instrumentation:play:play-appsec-common", ":dd-java-agent:instrumentation:protobuf-3.0", ":dd-java-agent:instrumentation:quartz-2.0", + ":dd-java-agent:instrumentation:r2dbc:r2dbc-1.0", ":dd-java-agent:instrumentation:rabbitmq-amqp-2.7", ":dd-java-agent:instrumentation:ratpack-1.5", ":dd-java-agent:instrumentation:reactive-streams-1.0", diff --git a/utils/test-junit-utils/src/main/java/datadog/trace/test/junit/utils/assertions/Is.java b/utils/test-junit-utils/src/main/java/datadog/trace/test/junit/utils/assertions/Is.java index f7a3346e64e..2deeb06f697 100644 --- a/utils/test-junit-utils/src/main/java/datadog/trace/test/junit/utils/assertions/Is.java +++ b/utils/test-junit-utils/src/main/java/datadog/trace/test/junit/utils/assertions/Is.java @@ -27,6 +27,9 @@ public String failureReason() { @Override public boolean test(T t) { + if (this.expected instanceof CharSequence && t instanceof CharSequence) { + return this.expected.toString().equals(t.toString()); + } return this.expected.equals(t); } }