-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #25 from appoptics/NH-5774-trace-context-in-sql-query
NH-5774: Trace context in sql query
- Loading branch information
Showing
7 changed files
with
130 additions
and
4 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
53 changes: 53 additions & 0 deletions
53
...rc/main/java/com/appoptics/opentelemetry/instrumentation/AoConnectionInstrumentation.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,53 @@ | ||
| package com.appoptics.opentelemetry.instrumentation; | ||
|
|
||
| import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed; | ||
| import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.implementsInterface; | ||
| import static io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge.currentContext; | ||
| import static net.bytebuddy.matcher.ElementMatchers.*; | ||
|
|
||
| import com.tracelytics.joboe.config.ConfigManager; | ||
| import com.tracelytics.joboe.config.ConfigProperty; | ||
| import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation; | ||
| import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer; | ||
| import net.bytebuddy.asm.Advice; | ||
| import net.bytebuddy.description.type.TypeDescription; | ||
| import net.bytebuddy.matcher.ElementMatcher; | ||
|
|
||
| public class AoConnectionInstrumentation implements TypeInstrumentation { | ||
|
|
||
| @Override | ||
| public ElementMatcher<ClassLoader> classLoaderOptimization() { | ||
| return hasClassesNamed("java.sql.Connection"); | ||
| } | ||
|
|
||
| @Override | ||
| public ElementMatcher<TypeDescription> typeMatcher() { | ||
| Boolean sqlTag = ConfigManager.getConfigOptional(ConfigProperty.AGENT_SQL_TAG, false); | ||
| if (!sqlTag) { | ||
| return none(); | ||
| } | ||
| return named("com.mysql.cj.jdbc.ConnectionImpl") // only inject MySQL JDBC driver | ||
| .and(implementsInterface(named("java.sql.Connection"))); | ||
| } | ||
|
|
||
| @Override | ||
| public void transform(TypeTransformer transformer) { | ||
| transformer.applyAdviceToMethod( | ||
| nameStartsWith("prepare") | ||
| .and(takesArgument(0, String.class)) | ||
| // Also include CallableStatement, which is a sub type of PreparedStatement | ||
| .and(returns(implementsInterface(named("java.sql.PreparedStatement")))), | ||
| AoConnectionInstrumentation.class.getName() + "$PrepareAdvice"); | ||
| } | ||
|
|
||
| @SuppressWarnings("unused") | ||
| public static class PrepareAdvice { | ||
|
|
||
| @Advice.OnMethodEnter(suppress = Throwable.class) | ||
| public static void injectComment( | ||
| @Advice.Argument(value = 0, readOnly = false) String sql) { | ||
| sql = TraceContextInjector.inject(currentContext(), sql); | ||
| } | ||
| } | ||
| } | ||
|
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
24 changes: 24 additions & 0 deletions
24
.../jdbc/src/main/java/com/appoptics/opentelemetry/instrumentation/TraceContextInjector.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,24 @@ | ||
| package com.appoptics.opentelemetry.instrumentation; | ||
|
|
||
| import io.opentelemetry.api.trace.Span; | ||
| import io.opentelemetry.api.trace.SpanContext; | ||
| import io.opentelemetry.context.Context; | ||
|
|
||
| public class TraceContextInjector { | ||
| public static String inject(Context context, String sql) { | ||
| if (sql.contains("traceparent")) { | ||
| return sql; | ||
| } | ||
|
|
||
| Span span = Span.fromContext(context); | ||
| SpanContext spanContext = span.getSpanContext(); | ||
| if (!(spanContext.isValid() && spanContext.isSampled())) { | ||
| return sql; | ||
| } | ||
| String flags = "01"; // only inject into sampled requests | ||
| String traceContext = "00-" + spanContext.getTraceId() + "-" + spanContext.getSpanId() + "-" + flags; | ||
| String tag = String.format("/*traceparent:'%s'*/", traceContext); | ||
| span.setAttribute("QueryTag", tag); | ||
| return String.format("%s %s", tag, sql); | ||
| } | ||
| } |
29 changes: 29 additions & 0 deletions
29
instrumentation/jdbc/src/test/groovy/InjectTraceContextTest.groovy
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,29 @@ | ||
| import com.appoptics.opentelemetry.instrumentation.TraceContextInjector | ||
| import io.opentelemetry.api.GlobalOpenTelemetry | ||
| import io.opentelemetry.context.Context | ||
| import spock.lang.Specification | ||
| import spock.lang.Unroll | ||
|
|
||
| import java.util.regex.Pattern | ||
|
|
||
| @Unroll | ||
| class InjectTraceContextTest extends Specification { | ||
| def "inject trace context from span to sql"(String sql) { | ||
| setup: | ||
| def tracer = GlobalOpenTelemetry.getTracer("test") | ||
| def testScope = tracer.spanBuilder("test").startSpan().makeCurrent() | ||
| Pattern pattern = Pattern.compile("/\\*traceparent:'00-[a-f0-9]{32}-[a-f0-9]{16}-0[0|1]'\\*/ .+") | ||
|
|
||
| expect: | ||
| pattern.matcher(TraceContextInjector.inject(Context.current(), sql)).matches() | ||
|
|
||
| cleanup: | ||
| testScope.close() | ||
|
|
||
| where: | ||
| sql |_ | ||
| "select name from students" |_ | ||
| "insert into students values('tom', 1)" |_ | ||
| "/* comment */ select name from students" |_ | ||
| } | ||
| } |