Merge pull request newrelic#209 from bizob2828/allow-inheritance

fix: updated smithy-client instrumentation to properly handle all types of clients for dynamo, sqs, sns

lib/v3/smithy-client.js

@@ -13,10 +13,10 @@ const MIDDLEWARE = Symbol('nrMiddleware')
 
 const middlewareByClient = {
   Client: middlewareConfig,
-  SNSClient: [...middlewareConfig, snsMiddlewareConfig],
-  SQSClient: [...middlewareConfig, sqsMiddlewareConfig],
-  DynamoDBClient: [...middlewareConfig, dynamoMiddlewareConfig],
-  DynamoDBDocumentClient: [...middlewareConfig, dynamoMiddlewareConfig]
+  SNS: [...middlewareConfig, snsMiddlewareConfig],
+  SQS: [...middlewareConfig, sqsMiddlewareConfig],
+  DynamoDB: [...middlewareConfig, dynamoMiddlewareConfig],
+  DynamoDBDocument: [...middlewareConfig, dynamoMiddlewareConfig]
 }
 
 module.exports = function instrumentSmithyClient(shim, smithyClientExport) {
@@ -29,7 +29,11 @@ module.exports = function instrumentSmithyClient(shim, smithyClientExport) {
 
 function wrapSend(shim, send) {
   return function wrappedSend() {
-    const client = this?.constructor?.name
+    // most clients we want to instrument aside from smithy-client
+    // extend themselves to provide different names(i.e. - SNS and SNSClient)
+    // we want to handle these the same by registering the sns middleware
+    const client = this.constructor.name.replace(/Client$/, '')
+    shim.logger.trace('Sending with client %s', client)
     const config = this.config
     const middlewares = middlewareByClient[client] || middlewareByClient.Client
 
@@ -41,6 +45,7 @@ function wrapSend(shim, send) {
     if (!this[MIDDLEWARE]) {
       this[MIDDLEWARE] = true
       for (const mw of middlewares) {
+        shim.logger.trace('Registering middleware %s for %s', mw.config.name, client)
         const localShim = shim.makeSpecializedShim(mw.type, client)
         // copy the shim id from parent so if you check if something is wrapped
         // it will be across all instrumentation

tests/versioned/v3/lib-dynamodb.tap.js

@@ -16,6 +16,7 @@ tap.test('DynamoDB', (t) => {
   let helper = null
   let ddbCommands = null
   let DynamoDBDocumentClient = null
+  let DynamoDBDocument = null
 
   let tableName = null
   let tests = null
@@ -33,6 +34,7 @@ tap.test('DynamoDB', (t) => {
     helper = utils.TestAgent.makeInstrumented()
     common.registerInstrumentation(helper)
     const lib = require('@aws-sdk/lib-dynamodb')
+    DynamoDBDocument = lib.DynamoDBDocument
     DynamoDBDocumentClient = lib.DynamoDBDocumentClient
     const { DynamoDBClient } = require('@aws-sdk/client-dynamodb')
     ddbCommands = {
@@ -141,6 +143,27 @@ tap.test('DynamoDB', (t) => {
       setImmediate(finish, ...args)
     })
   })
+
+  t.test('DynamoDBDocument client from commands', (t) => {
+    const docClientFrom = DynamoDBDocument.from(client)
+    helper.runInTransaction(async function (tx) {
+      for (let i = 0; i < tests.length; i++) {
+        const cfg = tests[i]
+        t.comment(`Testing ${cfg.operation}`)
+
+        try {
+          await docClientFrom.send(new ddbCommands[cfg.command](cfg.params))
+        } catch (err) {
+          t.error(err)
+        }
+      }
+
+      tx.end()
+
+      const args = [t, tests, tx]
+      setImmediate(finish, ...args)
+    })
+  })
 })
 
 function finish(t, tests, tx) {