From 10a3f79258032c55680beb0ed26e50bb59e19b52 Mon Sep 17 00:00:00 2001
From: rochdev <roch.devost@datadoghq.com>
Date: Mon, 18 Jun 2018 14:57:50 -0400
Subject: [PATCH] fix graphql integration when using parse/execute directly

---
 src/plugins/graphql.js       | 92 +++++++++++++++++++++---------------
 test/plugins/graphql.spec.js | 20 ++++++++
 2 files changed, 75 insertions(+), 37 deletions(-)

diff --git a/src/plugins/graphql.js b/src/plugins/graphql.js
index 09efc0bf186..089ce0b7320 100644
--- a/src/plugins/graphql.js
+++ b/src/plugins/graphql.js
@@ -3,46 +3,48 @@
 const shimmer = require('shimmer')
 const platform = require('../platform')
 
-function createWrapGraphql (tracer, config, defaultFieldResolver) {
-  return function wrapGraphql (graphql) {
-    return function graphqlWithTrace () {
-      const source = arguments[1] || arguments[0].source
-      const contextValue = arguments[3] || arguments[0].contextValue || {}
-
-      if (arguments.length === 1) {
-        arguments[0].contextValue = contextValue
-      } else {
-        arguments[3] = contextValue
-        arguments.length = Math.max(arguments.length, 4)
+function createWrapExecute (tracer, config, defaultFieldResolver) {
+  return function wrapExecute (execute) {
+    return function executeWithTrace () {
+      const args = normalizeArgs(arguments)
+      const schema = args.schema
+      const document = args.document
+      const contextValue = args.contextValue || {}
+      const fieldResolver = args.fieldResolver || defaultFieldResolver
+
+      if (!schema || !document || typeof fieldResolver !== 'function') {
+        return execute.apply(this, arguments)
       }
 
+      args.fieldResolver = wrapResolve(fieldResolver, tracer, config)
+      args.contextValue = contextValue
+
       Object.defineProperties(contextValue, {
         _datadog_operation: { value: {} },
         _datadog_fields: { value: {} },
-        _datadog_source: { value: source }
+        _datadog_source: { value: document._datadog_source }
       })
 
-      return graphql.apply(this, arguments)
+      if (!schema._datadog_patched) {
+        wrapFields(schema._queryType._fields, tracer, config, [])
+        schema._datadog_patched = true
+      }
+
+      return call(execute, this, [args], defer(tracer), () => finishOperation(contextValue))
     }
   }
 }
 
-function createWrapExecute (tracer, config, defaultFieldResolver) {
-  return function wrapExecute (execute) {
-    return function executeWithTrace () {
-      const schema = arguments[0]
-      const contextValue = arguments[3]
-      const fieldResolver = arguments[6] || defaultFieldResolver
-
-      arguments[6] = wrapResolve(fieldResolver, tracer, config)
-      arguments[3] = contextValue
+function createWrapParse () {
+  return function wrapParse (parse) {
+    return function parseWithTrace (source) {
+      const document = parse.apply(this, arguments)
 
-      if (!schema._datadog_patched) {
-        wrapFields(schema._queryType._fields, tracer, config, [])
-        schema._datadog_patched = true
-      }
+      Object.defineProperties(document, {
+        _datadog_source: { value: source }
+      })
 
-      return call(execute, this, arguments, defer(tracer), () => finishOperation(contextValue))
+      return document
     }
   }
 }
@@ -135,6 +137,22 @@ function getFieldParent (tracer, config, contextValue, info, path) {
   return contextValue._datadog_fields[path.slice(0, -1).join('.')].span
 }
 
+function normalizeArgs (args) {
+  if (args.length === 1) {
+    return args
+  }
+
+  return {
+    schema: args[0],
+    document: args[1],
+    rootValue: args[2],
+    contextValue: args[3],
+    variableValues: args[4],
+    operationName: args[5],
+    fieldResolver: args[6]
+  }
+}
+
 function createOperationSpan (tracer, config, contextValue, info) {
   const type = info.operation.operation
   const name = info.operation.name && info.operation.name.value
@@ -218,24 +236,24 @@ function addError (span, error) {
 module.exports = [
   {
     name: 'graphql',
-    file: 'graphql.js',
+    file: 'execution/execute.js',
     versions: ['0.13.x'],
-    patch (graphql, tracer, config) {
-      shimmer.wrap(graphql, 'graphql', createWrapGraphql(tracer, config))
+    patch (execute, tracer, config) {
+      shimmer.wrap(execute, 'execute', createWrapExecute(tracer, config, execute.defaultFieldResolver))
     },
-    unpatch (graphql) {
-      shimmer.unwrap(graphql, 'graphql')
+    unpatch (execute) {
+      shimmer.unwrap(execute, 'execute')
     }
   },
   {
     name: 'graphql',
-    file: 'execution/execute.js',
+    file: 'language/parser.js',
     versions: ['0.13.x'],
-    patch (execute, tracer, config) {
-      shimmer.wrap(execute, 'execute', createWrapExecute(tracer, config, execute.defaultFieldResolver))
+    patch (parser, tracer, config) {
+      shimmer.wrap(parser, 'parse', createWrapParse(tracer, config))
     },
-    unpatch (execute) {
-      shimmer.unwrap(execute, 'execute')
+    unpatch (parser) {
+      shimmer.unwrap(parser, 'parse')
     }
   }
 ]
diff --git a/test/plugins/graphql.spec.js b/test/plugins/graphql.spec.js
index b3c04edabb9..2ab3a67989a 100644
--- a/test/plugins/graphql.spec.js
+++ b/test/plugins/graphql.spec.js
@@ -289,6 +289,26 @@ describe('Plugin', () => {
           })
       })
 
+      it('should handle calling low level APIs directly', done => {
+        const source = `query MyQuery { hello(name: "world") }`
+        const document = graphql.parse(source)
+
+        agent
+          .use(traces => {
+            const spans = sort(traces[0])
+
+            expect(spans).to.have.length(3)
+            expect(spans[0]).to.have.property('service', 'test-graphql')
+            expect(spans[0]).to.have.property('name', 'graphql.query')
+            expect(spans[0]).to.have.property('resource', 'query MyQuery')
+            expect(spans[0].meta).to.have.property('graphql.document', source)
+          })
+          .then(done)
+          .catch(done)
+
+        graphql.execute(schema, document)
+      })
+
       it('should handle exceptions', done => {
         const error = new Error('test')