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

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);

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") }
   }
 }
 

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;