Support stored procedures
MODIFY
dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/CriteriaInstrumentation.java
dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/CriteriaInstrumentation.java
@@ -39,6 +39,11 @@ public class CriteriaInstrumentation extends Instrumenter.Default {
return new String[] {
"datadog.trace.instrumentation.hibernate.SessionMethodUtils",
"datadog.trace.instrumentation.hibernate.SessionState",
+ "datadog.trace.agent.decorator.BaseDecorator",
+ "datadog.trace.agent.decorator.ClientDecorator",
+ "datadog.trace.agent.decorator.DatabaseClientDecorator",
+ "datadog.trace.agent.decorator.OrmClientDecorator",
+ packageName + ".HibernateDecorator",
};
}

MODIFY
dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/IteratorInstrumentation.java
dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/IteratorInstrumentation.java
@@ -39,6 +39,11 @@ public class IteratorInstrumentation extends Instrumenter.Default {
return new String[] {
"datadog.trace.instrumentation.hibernate.SessionMethodUtils",
"datadog.trace.instrumentation.hibernate.SessionState",
+ "datadog.trace.agent.decorator.BaseDecorator",
+ "datadog.trace.agent.decorator.ClientDecorator",
+ "datadog.trace.agent.decorator.DatabaseClientDecorator",
+ "datadog.trace.agent.decorator.OrmClientDecorator",
+ packageName + ".HibernateDecorator",
};
}

ADD
dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/ProcedureCallInstrumentation.java
+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.procedure.ProcedureCall;
+
+@AutoService(Instrumenter.class)
+public class ProcedureCallInstrumentation extends Instrumenter.Default {
+
+ public ProcedureCallInstrumentation() {
+ super("hibernate");
+ }
+
+ @Override
+ public Map<String, String> contextStore() {
+ final Map<String, String> map = new HashMap<>();
+ map.put("org.hibernate.procedure.ProcedureCall", SessionState.class.getName());
+ return Collections.unmodifiableMap(map);
+ }
+
+ @Override
+ public String[] helperClassNames() {
+ return new String[] {
+ "datadog.trace.instrumentation.hibernate.SessionMethodUtils",
+ "datadog.trace.instrumentation.hibernate.SessionState",
+ "datadog.trace.agent.decorator.BaseDecorator",
+ "datadog.trace.agent.decorator.ClientDecorator",
+ "datadog.trace.agent.decorator.DatabaseClientDecorator",
+ "datadog.trace.agent.decorator.OrmClientDecorator",
+ packageName + ".HibernateDecorator",
+ };
+ }
+
+ @Override
+ public ElementMatcher<TypeDescription> typeMatcher() {
+ return not(isInterface()).and(safeHasSuperType(named("org.hibernate.procedure.ProcedureCall")));
+ }
+
+ @Override
+ public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+ final Map<ElementMatcher<? super MethodDescription>, String> transformers = new HashMap<>();
+ transformers.put(
+ isMethod().and(named("getOutputs")), ProcedureCallMethodAdvice.class.getName());
+
+ return transformers;
+ }
+
+ public static class ProcedureCallMethodAdvice {
+
+ @Advice.OnMethodEnter(suppress = Throwable.class)
+ public static SessionState startMethod(
+ @Advice.This final ProcedureCall call, @Advice.Origin("#m") final String name) {
+
+ final ContextStore<ProcedureCall, SessionState> contextStore =
+ InstrumentationContext.get(ProcedureCall.class, SessionState.class);
+
+ final SessionState state =
+ SessionMethodUtils.startScopeFrom(
+ contextStore, call, "hibernate.procedure." + name, call.getProcedureName());
+ return state;
+ }
+
+ @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+ public static void endMethod(
+ @Advice.Enter final SessionState state, @Advice.Thrown final Throwable throwable) {
+ SessionMethodUtils.closeScope(state, throwable, null);
+ }
+ }
+}

MODIFY
dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/QueryInstrumentation.java
dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/QueryInstrumentation.java
@@ -41,6 +41,11 @@ public class QueryInstrumentation extends Instrumenter.Default {
return new String[] {
"datadog.trace.instrumentation.hibernate.SessionMethodUtils",
"datadog.trace.instrumentation.hibernate.SessionState",
+ "datadog.trace.agent.decorator.BaseDecorator",
+ "datadog.trace.agent.decorator.ClientDecorator",
+ "datadog.trace.agent.decorator.DatabaseClientDecorator",
+ "datadog.trace.agent.decorator.OrmClientDecorator",
+ packageName + ".HibernateDecorator",
};
}

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
@@ -28,6 +28,7 @@ import org.hibernate.Query;
import org.hibernate.SharedSessionContract;
import org.hibernate.Transaction;
import org.hibernate.engine.HibernateIterator;
+import org.hibernate.procedure.ProcedureCall;
@AutoService(Instrumenter.class)
public class SessionInstrumentation extends Instrumenter.Default {
@@ -43,6 +44,7 @@ public class SessionInstrumentation extends Instrumenter.Default {
map.put("org.hibernate.Query", SessionState.class.getName());
map.put("org.hibernate.Transaction", SessionState.class.getName());
map.put("org.hibernate.Criteria", SessionState.class.getName());
+ map.put("org.hibernate.procedure.ProcedureCall", SessionState.class.getName());
map.put("org.hibernate.engine.HibernateIterator", SessionState.class.getName());
return Collections.unmodifiableMap(map);
}
@@ -116,6 +118,10 @@ public class SessionInstrumentation extends Instrumenter.Default {
isMethod().and(returns(safeHasSuperType(named("org.hibernate.Criteria")))),
GetCriteriaAdvice.class.getName());
+ transformers.put(
+ isMethod().and(returns(safeHasSuperType(named("org.hibernate.procedure.ProcedureCall")))),
+ GetProcedureCallAdvice.class.getName());
+
return transformers;
}
@@ -228,4 +234,21 @@ public class SessionInstrumentation extends Instrumenter.Default {
sessionContextStore, session, criteriaContextStore, criteria);
}
}
+
+ public static class GetProcedureCallAdvice {
+
+ @Advice.OnMethodExit(suppress = Throwable.class)
+ public static void getProcedureCall(
+ @Advice.This final SharedSessionContract session,
+ @Advice.Return final ProcedureCall returned) {
+
+ final ContextStore<SharedSessionContract, SessionState> sessionContextStore =
+ InstrumentationContext.get(SharedSessionContract.class, SessionState.class);
+ final ContextStore<ProcedureCall, SessionState> returnedContextStore =
+ InstrumentationContext.get(ProcedureCall.class, SessionState.class);
+
+ SessionMethodUtils.attachSpanFromStore(
+ sessionContextStore, session, returnedContextStore, returned);
+ }
+ }
}

MODIFY
dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/TransactionInstrumentation.java
dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/TransactionInstrumentation.java
@@ -40,6 +40,11 @@ public class TransactionInstrumentation extends Instrumenter.Default {
return new String[] {
"datadog.trace.instrumentation.hibernate.SessionMethodUtils",
"datadog.trace.instrumentation.hibernate.SessionState",
+ "datadog.trace.agent.decorator.BaseDecorator",
+ "datadog.trace.agent.decorator.ClientDecorator",
+ "datadog.trace.agent.decorator.DatabaseClientDecorator",
+ "datadog.trace.agent.decorator.OrmClientDecorator",
+ packageName + ".HibernateDecorator",
};
}

MODIFY
dd-java-agent/instrumentation/hibernate/src/test/groovy/CriteriaTest.groovy
dd-java-agent/instrumentation/hibernate/src/test/groovy/CriteriaTest.groovy
-import datadog.trace.agent.test.AgentTestRunner
import datadog.trace.api.DDSpanTypes
import datadog.trace.api.DDTags
import io.opentracing.tag.Tags
import org.hibernate.Criteria
import org.hibernate.Session
-import org.hibernate.SessionFactory
-import org.hibernate.boot.MetadataSources
-import org.hibernate.boot.registry.StandardServiceRegistry
-import org.hibernate.boot.registry.StandardServiceRegistryBuilder
import org.hibernate.criterion.Order
import org.hibernate.criterion.Restrictions
-import spock.lang.Shared
-class CriteriaTest extends AgentTestRunner {
-
- @Shared
- private SessionFactory sessionFactory
-
- def setupSpec() {
- final StandardServiceRegistry registry =
- new StandardServiceRegistryBuilder()
- .configure()
- .build()
- try {
- sessionFactory = new MetadataSources(registry).buildMetadata().buildSessionFactory()
- } catch (Exception e) {
- StandardServiceRegistryBuilder.destroy(registry)
- return
- }
-
- Session session = sessionFactory.openSession()
- session.beginTransaction()
- session.save(new Value("A Hibernate value to be serialized"))
- session.save(new Value("Another value"))
- session.getTransaction().commit()
- session.close()
- }
-
- def cleanupSpec() {
- if (sessionFactory != null) {
- sessionFactory.close()
- }
- }
+class CriteriaTest extends AbstractHibernateTest {
def "test criteria.#methodName"() {
setup:

Support stored procedures