Hibernate: instrument querying

dd-java-agent/instrumentation/hibernate/hibernate.gradle

@@ -1,33 +1,32 @@
 // Set properties before any plugins get loaded
 project.ext {
   // Execute tests on all JVMs, even rare and outdated ones
-//  coreJavaInstrumentation = true
+  coreJavaInstrumentation = true
-//  minJavaVersionForTests = JavaVersion.VERSION_1_7
+  minJavaVersionForTests = JavaVersion.VERSION_1_7
 }
 
-//muzzle {
+muzzle {
-//  pass {
+  pass {
-//    group = "org.hibernate"
+    group = "org.hibernate"
-//    module = "hibernate-core"
+    module = "hibernate-core"
-//    versions = "[5.4.2.Final]"
+    versions = "[5.0.0.Final, 5.+]"
-//    assertInverse = true
+    assertInverse = true
-//  }
+  }
-//}
+}
 
 apply from: "${rootDir}/gradle/java.gradle"
 
 apply plugin: 'org.unbroken-dome.test-sets'
 
-//testSets {
+testSets {
-//  latestDepTest {
+  latestDepTest {
-//    dirName = 'test'
+    dirName = 'test'
-//  }
+  }
-//}
+}
 
 dependencies {
-  compileOnly group: 'org.hibernate', name: 'hibernate-core', version: '5.4.1.Final'
+  compileOnly group: 'org.hibernate', name: 'hibernate-core', version: '5.0.0.Final'
 
-//  compile project(':dd-trace-api')
   compile project(':dd-java-agent:agent-tooling')
 
   compile deps.bytebuddy
@@ -36,8 +35,12 @@ dependencies {
   implementation deps.autoservice
 
   testCompile project(':dd-java-agent:testing')
+  testCompile project(':dd-java-agent:instrumentation:jdbc')
 
 
-  testCompile group: 'org.hibernate', name: 'hibernate-core', version: '5.4.1.Final'
+  testCompile group: 'org.hibernate', name: 'hibernate-core', version: '5.0.0.Final'
   testCompile group: 'com.h2database', name: 'h2', version: '1.4.197'
+
+  latestDepTestCompile group: 'org.hibernate', name: 'hibernate-core', version: '5.+'
+  latestDepTestCompile group: 'com.h2database', name: 'h2', version: '1.4.197'
 }

dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/IteratorInstrumentation.java

@@ -0,0 +1,80 @@
+package datadog.trace.instrumentation.hibernate;
+
+import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.safeHasSuperType;
+import static net.bytebuddy.matcher.ElementMatchers.isInterface;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.not;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.bootstrap.ContextStore;
+import datadog.trace.bootstrap.InstrumentationContext;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import org.hibernate.engine.HibernateIterator;
+
+@AutoService(Instrumenter.class)
+public class IteratorInstrumentation extends Instrumenter.Default {
+
+  public IteratorInstrumentation() {
+    super("hibernate");
+  }
+
+  @Override
+  public Map<String, String> contextStore() {
+    final Map<String, String> map = new HashMap<>();
+    map.put("org.hibernate.engine.HibernateIterator", SessionState.class.getName());
+    return Collections.unmodifiableMap(map);
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      "datadog.trace.instrumentation.hibernate.SessionMethodUtils",
+      "datadog.trace.instrumentation.hibernate.SessionState",
+    };
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return not(isInterface())
+        .and(safeHasSuperType(named("org.hibernate.engine.HibernateIterator")));
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    final Map<ElementMatcher<? super MethodDescription>, String> transformers = new HashMap<>();
+    transformers.put(
+        isMethod().and(named("next").or(named("remove"))), IteratorAdvice.class.getName());
+    return transformers;
+  }
+
+  public static class IteratorAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static SessionState startMethod(
+        @Advice.This final HibernateIterator iterator, @Advice.Origin("#m") final String name) {
+
+      final ContextStore<HibernateIterator, SessionState> contextStore =
+          InstrumentationContext.get(HibernateIterator.class, SessionState.class);
+
+      return SessionMethodUtils.startScopeFrom(
+          contextStore, iterator, "hibernate.iterator." + name);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void endMethod(
+        @Advice.This final HibernateIterator iterator,
+        @Advice.Enter final SessionState state,
+        @Advice.Thrown final Throwable throwable) {
+
+      SessionMethodUtils.closeScope(state, throwable);
+    }
+  }
+}

dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/QueryInstrumentation.java

@@ -1,7 +1,6 @@
 package datadog.trace.instrumentation.hibernate;
 
 import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.safeHasSuperType;
-import static net.bytebuddy.matcher.ElementMatchers.isDeclaredBy;
 import static net.bytebuddy.matcher.ElementMatchers.isInterface;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.named;
@@ -19,7 +18,7 @@ import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
-import org.hibernate.query.Query;
+import org.hibernate.Query;
 
 @AutoService(Instrumenter.class)
 public class QueryInstrumentation extends Instrumenter.Default {
@@ -31,7 +30,7 @@ public class QueryInstrumentation extends Instrumenter.Default {
   @Override
   public Map<String, String> contextStore() {
     final Map<String, String> map = new HashMap<>();
-    map.put("org.hibernate.query.Query", SessionState.class.getName());
+    map.put("org.hibernate.Query", SessionState.class.getName());
     return Collections.unmodifiableMap(map);
   }
 
@@ -45,7 +44,7 @@ public class QueryInstrumentation extends Instrumenter.Default {
 
   @Override
   public ElementMatcher<TypeDescription> typeMatcher() {
-    return not(isInterface()).and(safeHasSuperType(named("org.hibernate.query.Query")));
+    return not(isInterface()).and(safeHasSuperType(named("org.hibernate.Query")));
   }
 
   @Override
@@ -53,27 +52,32 @@ public class QueryInstrumentation extends Instrumenter.Default {
     final Map<ElementMatcher<? super MethodDescription>, String> transformers = new HashMap<>();
     transformers.put(
         isMethod()
-            .and(named("list"))
+            .and(
-            .and(isDeclaredBy(safeHasSuperType(named("org.hibernate.Query"))))
+                named("list")
+                    .or(named("executeUpdate"))
+                    .or(named("uniqueResult"))
+                    .or(named("iterate"))
+                    .or(named("scroll"))) // TODO(will): Instrument the ScrollableWhatever returned
             .and(takesArguments(0)),
-        QueryListAdvice.class.getName());
+        QueryMethodAdvice.class.getName());
 
     return transformers;
   }
 
-  public static class QueryListAdvice {
+  public static class QueryMethodAdvice {
 
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static SessionState startList(@Advice.This final Query query) {
+    public static SessionState startMethod(
+        @Advice.This final Query query, @Advice.Origin("#m") final String name) {
 
       final ContextStore<Query, SessionState> contextStore =
           InstrumentationContext.get(Query.class, SessionState.class);
 
-      return SessionMethodUtils.startScopeFrom(contextStore, query, "hibernate.query.list");
+      return SessionMethodUtils.startScopeFrom(contextStore, query, "hibernate.query." + name);
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
-    public static void endList(
+    public static void endMethod(
         @Advice.This final Query query,
         @Advice.Enter final SessionState state,
         @Advice.Thrown final Throwable throwable) {

dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/SessionInstrumentation.java

@@ -23,10 +23,12 @@ import java.util.Map;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.implementation.bytecode.assign.Assigner;
 import net.bytebuddy.matcher.ElementMatcher;
+import org.hibernate.Query;
 import org.hibernate.SharedSessionContract;
 import org.hibernate.Transaction;
-import org.hibernate.query.Query;
+import org.hibernate.engine.HibernateIterator;
 
 @AutoService(Instrumenter.class)
 public class SessionInstrumentation extends Instrumenter.Default {
@@ -39,8 +41,9 @@ public class SessionInstrumentation extends Instrumenter.Default {
   public Map<String, String> contextStore() {
     final Map<String, String> map = new HashMap<>();
     map.put("org.hibernate.SharedSessionContract", SessionState.class.getName());
-    map.put("org.hibernate.query.Query", SessionState.class.getName());
+    map.put("org.hibernate.Query", SessionState.class.getName());
     map.put("org.hibernate.Transaction", SessionState.class.getName());
+    map.put("org.hibernate.engine.HibernateIterator", SessionState.class.getName());
     return Collections.unmodifiableMap(map);
   }
 
@@ -76,7 +79,12 @@ public class SessionInstrumentation extends Instrumenter.Default {
                     .or(named("lock"))
                     .or(named("refresh"))
                     .or(named("insert"))
-                    .or(named("delete"))),
+                    .or(named("delete"))
+                    // Iterator methods.
+                    .or(named("iterate"))
+                    // Lazy-load methods.
+                    .or(named("immediateLoad"))
+                    .or(named("internalLoad"))),
         SessionMethodAdvice.class.getName());
     // Handle the generic and non-generic 'get' separately.
     transformers.put(
@@ -96,7 +104,7 @@ public class SessionInstrumentation extends Instrumenter.Default {
         GetTransactionAdvice.class.getName());
 
     transformers.put(
-        isMethod().and(returns(safeHasSuperType(named("org.hibernate.query.Query")))),
+        isMethod().and(returns(safeHasSuperType(named("org.hibernate.Query")))),
         GetQueryAdvice.class.getName());
 
     return transformers;
@@ -131,7 +139,7 @@ public class SessionInstrumentation extends Instrumenter.Default {
   public static class SessionMethodAdvice {
 
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static SessionState startSave(
+    public static SessionState startMethod(
         @Advice.This final SharedSessionContract session,
         @Advice.Origin("#m") final String name,
         @Advice.Argument(0) final Object entity) {
@@ -143,12 +151,26 @@ public class SessionInstrumentation extends Instrumenter.Default {
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
-    public static void endSave(
+    public static void endMethod(
         @Advice.This final SharedSessionContract session,
         @Advice.Enter final SessionState sessionState,
-        @Advice.Thrown final Throwable throwable) {
+        @Advice.Thrown final Throwable throwable,
+        @Advice.Return(typing = Assigner.Typing.DYNAMIC) final Object returned) {
 
       SessionMethodUtils.closeScope(sessionState, throwable);
+
+      // Attach instrumentation to any returned object.
+      if (returned == null) {
+        return;
+      }
+      final ContextStore<SharedSessionContract, SessionState> sessionContextStore =
+          InstrumentationContext.get(SharedSessionContract.class, SessionState.class);
+      if (returned instanceof HibernateIterator) {
+        final ContextStore<HibernateIterator, SessionState> iteratorContextStore =
+            InstrumentationContext.get(HibernateIterator.class, SessionState.class);
+        SessionMethodUtils.attachSpanFromStore(
+            sessionContextStore, session, iteratorContextStore, (HibernateIterator) returned);
+      }
     }
   }
 
@@ -163,7 +185,7 @@ public class SessionInstrumentation extends Instrumenter.Default {
       final ContextStore<Query, SessionState> queryContextStore =
           InstrumentationContext.get(Query.class, SessionState.class);
 
-      SessionMethodUtils.attachSpanFromSession(
+      SessionMethodUtils.attachSpanFromStore(
           sessionContextStore, session, queryContextStore, query);
     }
   }
@@ -180,7 +202,7 @@ public class SessionInstrumentation extends Instrumenter.Default {
       final ContextStore<Transaction, SessionState> transactionContextStore =
           InstrumentationContext.get(Transaction.class, SessionState.class);
 
-      SessionMethodUtils.attachSpanFromSession(
+      SessionMethodUtils.attachSpanFromStore(
           sessionContextStore, session, transactionContextStore, transaction);
     }
   }

dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/SessionMethodUtils.java

@@ -11,7 +11,6 @@ import io.opentracing.tag.Tags;
 import io.opentracing.util.GlobalTracer;
 import java.util.Collections;
 import javax.persistence.Entity;
-import org.hibernate.SharedSessionContract;
 
 public class SessionMethodUtils {
 
@@ -82,31 +81,27 @@ public class SessionMethodUtils {
 
     final Scope scope = sessionState.getMethodScope();
     final Span span = scope.span();
-    if (throwable != null) {
+    if (span != null) {
-      Tags.ERROR.set(span, true);
+      if (throwable != null) {
-      span.log(Collections.singletonMap(ERROR_OBJECT, throwable));
+        Tags.ERROR.set(span, true);
+        span.log(Collections.singletonMap(ERROR_OBJECT, throwable));
+      }
+      span.finish();
     }
 
-    span.finish();
     scope.close();
     sessionState.setMethodScope(null);
   }
 
   // Copies a span from the given Session ContextStore into the targetContextStore. Used to
   // propagate a Span from a Session to transient Session objects such as Transaction and Query.
-  public static <T> void attachSpanFromSession(
+  public static <S, T> void attachSpanFromStore(
-      final ContextStore<SharedSessionContract, SessionState> sessionContextStore,
+      final ContextStore<S, SessionState> sourceContextStore,
-      final SharedSessionContract session,
+      final S source,
       final ContextStore<T, SessionState> targetContextStore,
       final T target) {
 
-    if (!(session instanceof SharedSessionContract)
+    final SessionState state = sourceContextStore.get(source);
-        || sessionContextStore == null
-        || targetContextStore == null) {
-      return;
-    }
-
-    final SessionState state = sessionContextStore.get((SharedSessionContract) session);
     if (state == null) {
       return;
     }

dd-java-agent/instrumentation/hibernate/src/test/groovy/QueryTest.groovy

@@ -2,6 +2,7 @@ import datadog.trace.agent.test.AgentTestRunner
 import datadog.trace.api.DDSpanTypes
 import datadog.trace.api.DDTags
 import io.opentracing.tag.Tags
+import org.hibernate.Query
 import org.hibernate.Session
 import org.hibernate.SessionFactory
 import org.hibernate.boot.MetadataSources
@@ -40,23 +41,18 @@ class QueryTest extends AgentTestRunner {
     }
   }
 
-  def "test hibernate query"() {
+  def "test hibernate query.#queryMethodName"() {
     setup:
 
     Session session = sessionFactory.openSession()
     session.beginTransaction()
-    List result = session.createQuery("from Value").list()
+    queryInteraction(session)
-    for (Value value : (List<Value>) result) {
-      System.out.println(value.getName())
-    }
     session.getTransaction().commit()
     session.close()
 
     expect:
-    result.size() == 2
-
     assertTraces(1) {
-      trace(0, 3) {
+      trace(0, 4) {
         span(0) {
           serviceName "hibernate"
           resourceName "hibernate.session"
@@ -84,8 +80,8 @@ class QueryTest extends AgentTestRunner {
         }
         span(2) {
           serviceName "hibernate"
-          resourceName "hibernate.query.list"
+          resourceName "hibernate.query.$queryMethodName"
-          operationName "hibernate.query.list"
+          operationName "hibernate.query.$queryMethodName"
           spanType DDSpanTypes.HIBERNATE
           childOf span(0)
           tags {
@@ -94,8 +90,35 @@ class QueryTest extends AgentTestRunner {
             defaultTags()
           }
         }
+        span(3) {
+          serviceName "h2"
+          childOf span(2)
+        }
       }
     }
+
+    where:
+    queryMethodName | isError | queryInteraction
+    "list"          | false   | { sess ->
+      Query q = sess.createQuery("from Value")
+      q.list()
+    }
+    "executeUpdate" | false   | { sess ->
+      Query q = sess.createQuery("update Value set name = 'alyx'")
+      q.executeUpdate()
+    }
+    "uniqueResult"  | false   | { sess ->
+      Query q = sess.createQuery("from Value where id = 1")
+      q.uniqueResult()
+    }
+    "iterate"       | false   | { sess ->
+      Query q = sess.createQuery("from Value")
+      q.iterate()
+    }
+    "scroll"        | false   | { sess ->
+      Query q = sess.createQuery("from Value")
+      q.scroll()
+    }
   }
 
 }

dd-java-agent/instrumentation/hibernate/src/test/groovy/SessionTest.groovy

@@ -2,7 +2,6 @@ import datadog.trace.agent.test.AgentTestRunner
 import datadog.trace.api.DDSpanTypes
 import datadog.trace.api.DDTags
 import io.opentracing.tag.Tags
-import lombok.NonNull
 import org.hibernate.*
 import org.hibernate.boot.MetadataSources
 import org.hibernate.boot.registry.StandardServiceRegistry
@@ -14,6 +13,12 @@ class SessionTest extends AgentTestRunner {
   @Shared
   private SessionFactory sessionFactory
 
+  @Shared
+  private Map<String, Closure> sessionBuilders
+
+  @Shared
+  private List<Value> prepopulated
+
   def setupSpec() {
     final StandardServiceRegistry registry =
       new StandardServiceRegistryBuilder()
@@ -24,6 +29,25 @@ class SessionTest extends AgentTestRunner {
     } catch (Exception e) {
       StandardServiceRegistryBuilder.destroy(registry)
     }
+
+    // Test two different types of Session. Groovy doesn't allow testing the cross-product/combinations of two data
+    // tables, so we get this hack instead.
+    sessionBuilders = new HashMap<>();
+    sessionBuilders.put("Session", { return sessionFactory.openSession() })
+    sessionBuilders.put("StatelessSession", { return sessionFactory.openStatelessSession() })
+
+    // Pre-populate the DB, so delete/update can be tested.
+    Session writer = sessionFactory.openSession()
+    writer.beginTransaction()
+    prepopulated = new ArrayList<>()
+    for (int i = 0; i < 2; i++) {
+      prepopulated.add(new Value("Hello :)"))
+      writer.save(prepopulated.get(i))
+    }
+    writer.getTransaction().commit()
+    writer.close()
+    TEST_WRITER.waitForTraces(1)
+    TEST_WRITER.clear()
   }
 
   def cleanupSpec() {
@@ -32,46 +56,17 @@ class SessionTest extends AgentTestRunner {
     }
   }
 
-  @lombok.Value
-  static class SessionTypeTest {
-    // Lombok can't get the types right for the constructor.
-    SessionTypeTest(Closure fn) {
-      sessionBuilder = fn
-    }
 
-    @NonNull
+  def "test hibernate action #testName"() {
-    Closure sessionBuilder
-    @NonNull
-    Value value = new Value("Hello :)")
-  }
-
-  def "test hibernate #testName"() {
     setup:
-    // Test two different types of Session. Groovy doesn't allow testing the cross-product/combinations of two data
-    // tables, so we get this hack instead.
-    Map<String, SessionTypeTest> sessionTests = new HashMap<>();
-    sessionTests.put("Session", new SessionTypeTest({ return sessionFactory.openSession() }))
-    sessionTests.put("StatelessSession", new SessionTypeTest({ return sessionFactory.openStatelessSession() }))
-
-    // Pre-populate the DB, so delete/update can be tested.
-    Session writer = sessionFactory.openSession()
-    writer.beginTransaction()
-    for (SessionTypeTest sessionTypeTest : sessionTests.values()) {
-      writer.save(sessionTypeTest.value)
-    }
-    writer.getTransaction().commit()
-    writer.close()
-    TEST_WRITER.waitForTraces(1)
-    TEST_WRITER.clear()
 
     // Test for each implementation of Session.
     for (String sessionImplementation : sessionImplementations) {
-      SessionTypeTest sessionTypeTest = sessionTests.get(sessionImplementation)
+      def session = sessionBuilders.get(sessionImplementation)()
-      def session = sessionTypeTest.getSessionBuilder().call()
       session.beginTransaction()
 
       try {
-        sessionMethodTest.call(session, sessionTypeTest.getValue())
+        sessionMethodTest.call(session, prepopulated.get(0))
       } catch (Exception e) {
         // We expected this, we should see the error field set on the span.
       }
@@ -83,7 +78,7 @@ class SessionTest extends AgentTestRunner {
     expect:
     assertTraces(sessionImplementations.size()) {
       for (int i = 0; i < sessionImplementations.size(); i++) {
-        trace(i, 3) {
+        trace(i, 4) {
           span(0) {
             serviceName "hibernate"
             resourceName "hibernate.session"
@@ -127,64 +122,295 @@ class SessionTest extends AgentTestRunner {
               defaultTags()
             }
           }
+          span(3) {
+            serviceName "h2"
+            childOf span(2)
+          }
         }
       }
     }
 
     where:
-    testName                  | methodName     | resource         | isError | sessionImplementations          | sessionMethodTest
+    testName                                  | methodName | resource | isError | sessionImplementations          | sessionMethodTest
-    "replicate"               | "replicate"    | "Value"          | false   | ["Session"]                     | { sesh, val ->
+    "lock"                                    | "lock"     | "Value"  | false   | ["Session"]                     | { sesh, val ->
+      sesh.lock(val, LockMode.READ)
+    }
+    "refresh"                                 | "refresh"  | "Value"  | false   | ["Session", "StatelessSession"] | { sesh, val ->
+      sesh.refresh(val)
+    }
+    "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"))
+    }
+    "update (StatelessSession)"               | "update"   | "Value"  | false   | ["StatelessSession"]            | { sesh, val ->
+      val.setName("New name")
+      sesh.update(val)
+    }
+    "update by entityName (StatelessSession)" | "update"   | "Value"  | false   | ["StatelessSession"]            | { sesh, val ->
+      val.setName("New name")
+      sesh.update("Value", val)
+    }
+    "delete (Session)"                        | "delete"   | "Value"  | false   | ["StatelessSession"]            | { sesh, val ->
+      sesh.delete(val)
+    }
+  }
+
+  def "test hibernate replicate: #testName"() {
+    setup:
+
+    // Test for each implementation of Session.
+    def session = sessionFactory.openSession()
+    session.beginTransaction()
+
+    try {
+      sessionMethodTest.call(session, prepopulated.get(0))
+    } catch (Exception e) {
+      // We expected this, we should see the error field set on the span.
+    }
+
+    session.getTransaction().commit()
+    session.close()
+
+    expect:
+    assertTraces(1) {
+      trace(0, 5) {
+        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 "h2"
+          childOf span(1)
+        }
+        span(3) {
+          serviceName "hibernate"
+          resourceName resource
+          operationName "hibernate.$methodName"
+          spanType DDSpanTypes.HIBERNATE
+          childOf span(0)
+          if (isError) {
+            errored true
+          }
+          tags {
+            if (isError) {
+              errorTags(MappingException, "Unknown entity: java.lang.Long")
+            }
+            "$Tags.COMPONENT.key" "hibernate-java"
+            "$DDTags.SPAN_TYPE" DDSpanTypes.HIBERNATE
+            defaultTags()
+          }
+        }
+        span(4) {
+          serviceName "h2"
+          childOf span(3)
+        }
+      }
+
+    }
+
+    where:
+    testName                  | methodName  | resource | isError | sessionMethodTest
+    "replicate"               | "replicate" | "Value"  | false   | { sesh, val ->
       Value replicated = new Value(val.getName() + " replicated")
       replicated.setId(val.getId())
       sesh.replicate(replicated, ReplicationMode.OVERWRITE)
     }
-    "replicate by entityName" | "replicate"    | "Value"          | false   | ["Session"]                     | { sesh, val ->
+    "replicate by entityName" | "replicate" | "Value"  | false   | { sesh, val ->
       Value replicated = new Value(val.getName() + " replicated")
       replicated.setId(val.getId())
       sesh.replicate("Value", replicated, ReplicationMode.OVERWRITE)
     }
-    "failed replicate"        | "replicate"    | "unknown object" | true    | ["Session"]                     | { sesh, val ->
+  }
-      sesh.replicate(new Long(123) /* Not a valid entity */, ReplicationMode.OVERWRITE)
+
+  def "test hibernate failed replicate"() {
+    setup:
+
+    // Test for each implementation of Session.
+    def session = sessionFactory.openSession()
+    session.beginTransaction()
+
+    try {
+      session.replicate(new Long(123) /* Not a valid entity */, ReplicationMode.OVERWRITE)
+    } catch (Exception e) {
+      // We expected this, we should see the error field set on the span.
     }
-    "save"                    | "save"         | "Value"          | false   | ["Session"]                     | { sesh, val ->
+
+    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 "unknown object"
+          operationName "hibernate.replicate"
+          spanType DDSpanTypes.HIBERNATE
+          childOf span(0)
+          tags {
+            errorTags(MappingException, "Unknown entity: java.lang.Long")
+
+            "$Tags.COMPONENT.key" "hibernate-java"
+            "$DDTags.SPAN_TYPE" DDSpanTypes.HIBERNATE
+            defaultTags()
+          }
+        }
+      }
+
+    }
+  }
+
+
+  def "test hibernate commit action #testName"() {
+    setup:
+
+    // Test for each implementation of Session.
+    for (String sessionImplementation : sessionImplementations) {
+      def session = sessionBuilders.get(sessionImplementation)()
+      session.beginTransaction()
+
+      try {
+        sessionMethodTest.call(session, prepopulated.get(0))
+      } catch (Exception e) {
+        // We expected this, we should see the error field set on the span.
+      }
+
+      session.getTransaction().commit()
+      session.close()
+    }
+
+    expect:
+    assertTraces(sessionImplementations.size()) {
+      for (int i = 0; i < sessionImplementations.size(); i++) {
+        trace(i, 4) {
+          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 "h2"
+            childOf span(1)
+          }
+          span(3) {
+            serviceName "hibernate"
+            resourceName resource
+            operationName "hibernate.$methodName"
+            spanType DDSpanTypes.HIBERNATE
+            childOf span(0)
+            if (isError) {
+              errored true
+            }
+            tags {
+              if (isError) {
+                errorTags(MappingException, "Unknown entity: java.lang.Long")
+              }
+              "$Tags.COMPONENT.key" "hibernate-java"
+              "$DDTags.SPAN_TYPE" DDSpanTypes.HIBERNATE
+              defaultTags()
+            }
+          }
+        }
+      }
+    }
+
+    where:
+    testName                         | methodName     | resource | isError | sessionImplementations | sessionMethodTest
+    "save"                           | "save"         | "Value"  | false   | ["Session"]            | { sesh, val ->
       sesh.save(new Value("Another value"))
     }
-    "saveOrUpdate save"       | "saveOrUpdate" | "Value"          | false   | ["Session"]                     | { sesh, val ->
+    "saveOrUpdate save"              | "saveOrUpdate" | "Value"  | false   | ["Session"]            | { sesh, val ->
       sesh.saveOrUpdate(new Value("Value"))
     }
-    "saveOrUpdate update"     | "saveOrUpdate" | "Value"          | false   | ["Session"]                     | { sesh, val ->
+    "saveOrUpdate update"            | "saveOrUpdate" | "Value"  | false   | ["Session"]            | { sesh, val ->
       val.setName("New name")
       sesh.saveOrUpdate(val)
     }
-    "update"                  | "update"       | "Value"          | false   | ["Session", "StatelessSession"] | { sesh, val ->
+    "merge"                          | "merge"        | "Value"  | false   | ["Session"]            | { sesh, val ->
+      sesh.merge(new Value("merge me in"))
+    }
+    "persist"                        | "persist"      | "Value"  | false   | ["Session"]            | { sesh, val ->
+      sesh.persist(new Value("merge me in"))
+    }
+    "update (Session)"               | "update"       | "Value"  | false   | ["Session"]            | { sesh, val ->
       val.setName("New name")
       sesh.update(val)
     }
-    "update by entityName"    | "update"       | "Value"          | false   | ["Session", "StatelessSession"] | { sesh, val ->
+    "update by entityName (Session)" | "update"       | "Value"  | false   | ["Session"]            | { sesh, val ->
       val.setName("New name")
       sesh.update("Value", val)
     }
-    "merge"                   | "merge"        | "Value"          | false   | ["Session"]                     | { sesh, val ->
+    "delete (Session)"               | "delete"       | "Value"  | false   | ["Session"]            | { sesh, val ->
-      sesh.merge(new Value("merge me in"))
-    }
-    "persist"                 | "persist"      | "Value"          | false   | ["Session"]                     | { sesh, val ->
-      sesh.persist(new Value("merge me in"))
-    }
-    "lock"                    | "lock"         | "Value"          | false   | ["Session"]                     | { sesh, val ->
-      sesh.lock(val, LockMode.READ)
-    }
-    "refresh"                 | "refresh"      | "Value"          | false   | ["Session", "StatelessSession"] | { sesh, val ->
-      sesh.refresh(val)
-    }
-    "delete"                  | "delete"       | "Value"          | false   | ["Session", "StatelessSession"] | { sesh, val ->
       sesh.delete(val)
     }
-    "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"))
-    }
   }
 
 
@@ -199,7 +425,7 @@ class SessionTest extends AgentTestRunner {
 
     expect:
     assertTraces(1) {
-      trace(0, 3) {
+      trace(0, 4) {
         span(0) {
           serviceName "hibernate"
           resourceName "hibernate.session"
@@ -237,18 +463,18 @@ class SessionTest extends AgentTestRunner {
             defaultTags()
           }
         }
+        span(3) {
+          serviceName "h2"
+          childOf span(2)
+        }
       }
     }
 
     where:
-    queryMethodName                     | queryBuildMethod
+    queryMethodName  | queryBuildMethod
-    "createQuery"                       | { sess -> sess.createQuery("from Value") }
+    "createQuery"    | { sess -> sess.createQuery("from Value") }
-    "createQuery with resultClass"      | { sess -> sess.createQuery("from Value", Value.class) }
+    "getNamedQuery"  | { sess -> sess.getNamedQuery("TestNamedQuery") }
-    "createNamedQuery"                  | { sess -> sess.createNamedQuery("TestNamedQuery") }
+    "createSQLQuery" | { sess -> sess.createSQLQuery("SELECT * FROM Value") }
-    "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") }
   }
 }