migrating tedious to use runStores

Author: crysmags

packages/datadog-instrumentations/src/tedious.js

@@ -2,8 +2,7 @@
 
 const {
   channel,
-  addHook,
-  AsyncResource
+  addHook
 } = require('./helpers/instrument')
 const shimmer = require('../../datadog-shimmer')
 
@@ -22,30 +21,26 @@ addHook({ name: 'tedious', versions: ['>=1.0.0'] }, tedious => {
       return makeRequest.apply(this, arguments)
     }
 
-    const callbackResource = new AsyncResource('bound-anonymous-fn')
-    const asyncResource = new AsyncResource('bound-anonymous-fn')
-
     const connectionConfig = this.config
+    const ctx = { queryOrProcedure, connectionConfig }
 
-    return asyncResource.runInAsyncScope(() => {
-      const payload = { queryOrProcedure, connectionConfig }
-      startCh.publish(payload)
-      queryParent[queryField] = payload.sql
+    return startCh.runStores(ctx, () => {
+      queryParent[queryField] = ctx.sql
 
-      const cb = callbackResource.bind(request.callback, request)
-      request.callback = asyncResource.bind(function (error) {
+      const cb = request.callback
+      request.callback = function (error, ...args) {
         if (error) {
-          errorCh.publish(error)
+          ctx.error = error
+          errorCh.publish(ctx)
         }
-        finishCh.publish(undefined)
-
-        return cb.apply(this, arguments)
-      }, null, request)
+        return finishCh.runStores(ctx, cb, this, error, ...args)
+      }
 
       try {
         return makeRequest.apply(this, arguments)
       } catch (error) {
-        errorCh.publish(error)
+        ctx.error = error
+        errorCh.publish(ctx)
 
         throw error
       }

packages/datadog-plugin-tedious/src/index.js

@@ -8,7 +8,7 @@ class TediousPlugin extends DatabasePlugin {
   static get operation () { return 'request' } // TODO: change to match other database plugins
   static get system () { return 'mssql' }
 
-  start (payload) {
+  bindStart (payload) {
     const service = this.serviceName({ pluginConfig: this.config, system: this.system })
     const span = this.startSpan(this.operationName(), {
       service,
@@ -24,11 +24,12 @@ class TediousPlugin extends DatabasePlugin {
         'db.name': payload.connectionConfig.options.database,
         'db.instance': payload.connectionConfig.options.instanceName
       }
-    })
+    }, payload)
 
     // SQL Server includes comments when caching queries
     // For that reason we allow service mode but not full mode
     payload.sql = this.injectDbmQuery(span, payload.queryOrProcedure, service, true)
+    return payload.currentStore
   }
 }
 

packages/dd-trace/src/plugins/storage.js

@@ -11,12 +11,12 @@ class StoragePlugin extends ClientPlugin {
     this.system = this.constructor.system || this.component
   }
 
-  startSpan (name, options) {
+  startSpan (name, options, ctx) {
     if (!options.service && this.system) {
       options.service = `${this.tracer._service}-${this.system}`
     }
 
-    return super.startSpan(name, options)
+    return super.startSpan(name, options, ctx)
   }
 }