Added instrumentation to attach State to all Query objects, so they can be instrumented
MODIFY
dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/SessionInstrumentation.java
dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/SessionInstrumentation.java
@@ -27,7 +27,6 @@ import net.bytebuddy.matcher.ElementMatcher;
import org.hibernate.SharedSessionContract;
import org.hibernate.Transaction;
import org.hibernate.query.Query;
-import org.hibernate.query.spi.QueryImplementor;
@AutoService(Instrumenter.class)
public class SessionInstrumentation extends Instrumenter.Default {
@@ -76,6 +75,7 @@ public class SessionInstrumentation extends Instrumenter.Default {
.or(named("persist"))
.or(named("lock"))
.or(named("refresh"))
+ .or(named("insert"))
.or(named("delete"))),
SessionMethodAdvice.class.getName());
// Handle the generic and non-generic 'get' separately.
@@ -95,7 +95,9 @@ public class SessionInstrumentation extends Instrumenter.Default {
.and(returns(named("org.hibernate.Transaction"))),
GetTransactionAdvice.class.getName());
- transformers.put(isMethod().and(named("createQuery")), GetQueryAdvice.class.getName());
+ transformers.put(
+ isMethod().and(returns(safeHasSuperType(named("org.hibernate.query.Query")))),
+ GetQueryAdvice.class.getName());
return transformers;
}
@@ -154,12 +156,7 @@ public class SessionInstrumentation extends Instrumenter.Default {
@Advice.OnMethodExit(suppress = Throwable.class)
public static void getQuery(
- @Advice.This final SharedSessionContract session,
- @Advice.Return final QueryImplementor query) {
-
- if (!(query instanceof Query)) {
- return;
- }
+ @Advice.This final SharedSessionContract session, @Advice.Return final Query query) {
final ContextStore<SharedSessionContract, SessionState> sessionContextStore =
InstrumentationContext.get(SharedSessionContract.class, SessionState.class);

MODIFY
dd-java-agent/instrumentation/hibernate/src/test/groovy/SessionTest.groovy
dd-java-agent/instrumentation/hibernate/src/test/groovy/SessionTest.groovy
@@ -182,6 +182,73 @@ class SessionTest extends AgentTestRunner {
"get" | "get" | "Value" | false | ["Session", "StatelessSession"] | { sesh, val ->
sesh.get("Value", val.getId())
}
+ "insert" | "insert" | "Value" | false | ["StatelessSession"] | { sesh, val ->
+ sesh.insert("Value", new Value("insert me"))
+ }
+ }
+
+
+ def "test attaches State to query created via #queryMethodName"() {
+ setup:
+ Session session = sessionFactory.openSession()
+ session.beginTransaction()
+ Query query = queryBuildMethod(session)
+ List result = query.list()
+ session.getTransaction().commit()
+ session.close()
+
+ expect:
+ assertTraces(1) {
+ trace(0, 3) {
+ span(0) {
+ serviceName "hibernate"
+ resourceName "hibernate.session"
+ operationName "hibernate.session"
+ spanType DDSpanTypes.HIBERNATE
+ parent()
+ tags {
+ "$Tags.COMPONENT.key" "hibernate-java"
+ "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_CLIENT
+ "$DDTags.SPAN_TYPE" DDSpanTypes.HIBERNATE
+ defaultTags()
+ }
+ }
+ span(1) {
+ serviceName "hibernate"
+ resourceName "hibernate.transaction.commit"
+ operationName "hibernate.transaction.commit"
+ spanType DDSpanTypes.HIBERNATE
+ childOf span(0)
+ tags {
+ "$Tags.COMPONENT.key" "hibernate-java"
+ "$DDTags.SPAN_TYPE" DDSpanTypes.HIBERNATE
+ defaultTags()
+ }
+ }
+ span(2) {
+ serviceName "hibernate"
+ resourceName "hibernate.query.list"
+ operationName "hibernate.query.list"
+ spanType DDSpanTypes.HIBERNATE
+ childOf span(0)
+ tags {
+ "$Tags.COMPONENT.key" "hibernate-java"
+ "$DDTags.SPAN_TYPE" DDSpanTypes.HIBERNATE
+ defaultTags()
+ }
+ }
+ }
+ }
+
+ where:
+ queryMethodName | queryBuildMethod
+ "createQuery" | { sess -> sess.createQuery("from Value") }
+ "createQuery with resultClass" | { sess -> sess.createQuery("from Value", Value.class) }
+ "createNamedQuery" | { sess -> sess.createNamedQuery("TestNamedQuery") }
+ "createNamedQuery with resultClass" | { sess -> sess.createNamedQuery("TestNamedQuery", Value.class) }
+ "getNamedQuery" | { sess -> sess.getNamedQuery("TestNamedQuery") }
+ "createSQLQuery" | { sess -> sess.createSQLQuery("SELECT * FROM Value") }
+ "createNativeQuery" | { sess -> sess.createNativeQuery("SELECT * FROM Value") }
}
}

MODIFY
dd-java-agent/instrumentation/hibernate/src/test/java/Value.java
dd-java-agent/instrumentation/hibernate/src/test/java/Value.java
@@ -4,9 +4,11 @@ import javax.persistence.GeneratedValue;
import javax.persistence.Id;
import javax.persistence.Table;
import org.hibernate.annotations.GenericGenerator;
+import org.hibernate.annotations.NamedQuery;
@Entity
@Table
+@NamedQuery(name = "TestNamedQuery", query = "from Value")
public class Value {
private Long id;

Added instrumentation to attach State to all Query objects, so they can be instrumented