From 7daad26f5209deba1814f68aa41cc3b8bd4460a0 Mon Sep 17 00:00:00 2001
From: rochdev <roch.devost@datadoghq.com>
Date: Mon, 15 Jul 2019 16:04:18 -0400
Subject: [PATCH] fix missing spans for express built-in middleware

---
 packages/datadog-plugin-express/src/index.js  | 23 ++++++-------------
 .../datadog-plugin-express/test/index.spec.js | 22 ++++++++++--------
 2 files changed, 20 insertions(+), 25 deletions(-)

diff --git a/packages/datadog-plugin-express/src/index.js b/packages/datadog-plugin-express/src/index.js
index 1462f6a3f08..1b530545cf8 100644
--- a/packages/datadog-plugin-express/src/index.js
+++ b/packages/datadog-plugin-express/src/index.js
@@ -1,36 +1,27 @@
 'use strict'
 
-const METHODS = require('methods').concat('use', 'route', 'param', 'all')
 const web = require('../../dd-trace/src/plugins/util/web')
 const routerPlugin = require('../../datadog-plugin-router/src')
 
-function createWrapMethod (tracer, config) {
+function createWrapHandle (tracer, config) {
   config = web.normalizeConfig(config)
 
-  function ddTrace (req, res, next) {
-    web.instrument(tracer, config, req, res, 'express.request')
+  return function wrapHandle (handle) {
+    return function handleWithTrace (req, res) {
+      web.instrument(tracer, config, req, res, 'express.request')
 
-    next()
-  }
-
-  return function wrapMethod (original) {
-    return function methodWithTrace () {
-      if (!this._datadog_trace_patched && !this._router) {
-        this._datadog_trace_patched = true
-        this.use(ddTrace)
-      }
-      return original.apply(this, arguments)
+      return handle.apply(this, arguments)
     }
   }
 }
 
 function patch (express, tracer, config) {
-  this.wrap(express.application, METHODS, createWrapMethod(tracer, config))
+  this.wrap(express.application, 'handle', createWrapHandle(tracer, config))
   routerPlugin.patch.call(this, { prototype: express.Router }, tracer, config)
 }
 
 function unpatch (express) {
-  this.unwrap(express.application, METHODS)
+  this.unwrap(express.application, 'handle')
   routerPlugin.unpatch.call(this, { prototype: express.Router })
 }
 
diff --git a/packages/datadog-plugin-express/test/index.spec.js b/packages/datadog-plugin-express/test/index.spec.js
index 17f4ba69f69..45d9084af88 100644
--- a/packages/datadog-plugin-express/test/index.spec.js
+++ b/packages/datadog-plugin-express/test/index.spec.js
@@ -154,19 +154,23 @@ describe('Plugin', () => {
               .use(traces => {
                 const spans = sort(traces[0])
 
-                expect(spans).to.have.length(5)
+                expect(spans).to.have.length(7)
 
                 expect(spans[0]).to.have.property('resource', 'GET /app/user/:id')
                 expect(spans[0]).to.have.property('name', 'express.request')
-                expect(spans[1]).to.have.property('resource', 'named')
+                expect(spans[1]).to.have.property('resource', 'query')
                 expect(spans[1]).to.have.property('name', 'express.middleware')
                 expect(spans[1].parent_id.toString()).to.equal(spans[0].trace_id.toString())
-                expect(spans[2]).to.have.property('resource', 'router')
+                expect(spans[2]).to.have.property('resource', 'expressInit')
                 expect(spans[2]).to.have.property('name', 'express.middleware')
-                expect(spans[3].resource).to.match(/^bound\s.*$/)
+                expect(spans[3]).to.have.property('resource', 'named')
                 expect(spans[3]).to.have.property('name', 'express.middleware')
-                expect(spans[4]).to.have.property('resource', '<anonymous>')
+                expect(spans[4]).to.have.property('resource', 'router')
                 expect(spans[4]).to.have.property('name', 'express.middleware')
+                expect(spans[5].resource).to.match(/^bound\s.*$/)
+                expect(spans[5]).to.have.property('name', 'express.middleware')
+                expect(spans[6]).to.have.property('resource', '<anonymous>')
+                expect(spans[6]).to.have.property('name', 'express.middleware')
               })
               .then(done)
               .catch(done)
@@ -773,10 +777,10 @@ describe('Plugin', () => {
               .use(traces => {
                 const spans = sort(traces[0])
 
-                expect(spans[1]).to.have.property('error', 1)
-                expect(spans[1].meta).to.have.property('error.type', error.name)
-                expect(spans[1].meta).to.have.property('error.msg', error.message)
-                expect(spans[1].meta).to.have.property('error.stack', error.stack)
+                expect(spans[3]).to.have.property('error', 1)
+                expect(spans[3].meta).to.have.property('error.type', error.name)
+                expect(spans[3].meta).to.have.property('error.msg', error.message)
+                expect(spans[3].meta).to.have.property('error.stack', error.stack)
               })
               .then(done)
               .catch(done)