refactor: Moved relevant aws instrumentation and rely on agent commons

merged/aws-sdk/nr-hooks.js → lib/instrumentation/aws-sdk/nr-hooks.js

@@ -9,18 +9,18 @@ const instrumentations = [
   {
     type: 'conglomerate',
     moduleName: 'aws-sdk',
-    onRequire: require('./lib/v2/instrumentation')
+    onRequire: require('./v2/instrumentation')
   },
   {
     type: 'conglomerate',
     moduleName: '@aws-sdk/smithy-client',
-    onRequire: require('./lib/v3/smithy-client'),
+    onRequire: require('./v3/smithy-client'),
     shimName: 'aws-sdk'
   },
   {
     type: 'conglomerate',
     moduleName: '@smithy/smithy-client',
-    onRequire: require('./lib/v3/smithy-client'),
+    onRequire: require('./v3/smithy-client'),
     shimName: 'aws-sdk'
   }
 ]

merged/aws-sdk/lib/util.js → lib/instrumentation/aws-sdk/util.js

@@ -5,16 +5,9 @@
 
 'use strict'
 const UNKNOWN = 'Unknown'
-const DESTINATIONS = {
-  TRANS_EVENT: 0x01,
-
-  // This magic number is brought to you by:
-  // https://github.com/newrelic/node-newrelic/blob/10762a7/lib/config/attribute-filter.js#L10-L23
-  // We hard code it here because we'd have a cyclic dependency if we tried
-  // to import it from `newrelic` (`newrelic` uses this module to provide
-  // the AWS instrumentation).
-  TRANS_SCOPE: 0x01 | 0x02 | 0x04 | 0x08
-}
+const {
+  params: { DatastoreParameters }
+} = require('../../shim/specs')
 
 function grabLastUrlSegment(url = '/') {
   // cast URL as string, and an empty
@@ -27,12 +20,11 @@ function grabLastUrlSegment(url = '/') {
 /**
  * Retrieves the db segment params from endpoint and command parameters
  *
- * @param {function} DatastoreParameters constructor of `shim.spec.DatastoreParameters`
  * @param {Object} endpoint instance of ddb endpoint
  * @param {Object} params parameters passed to a ddb command
  * @returns {Object}
  */
-function setDynamoParameters(DatastoreParameters, endpoint, params) {
+function setDynamoParameters(endpoint, params) {
   return new DatastoreParameters({
     host: endpoint && (endpoint.host || endpoint.hostname),
     port_path_or_id: (endpoint && endpoint.port) || 443,
@@ -42,6 +34,5 @@ function setDynamoParameters(DatastoreParameters, endpoint, params) {
 
 module.exports = {
   grabLastUrlSegment,
-  setDynamoParameters,
-  DESTINATIONS
+  setDynamoParameters
 }

merged/aws-sdk/lib/v2/core.js → lib/instrumentation/aws-sdk/v2/core.js

@@ -6,6 +6,7 @@
 'use strict'
 
 const UNKNOWN = 'Unknown'
+const symbols = require('../../../symbols')
 
 function validate(shim, AWS) {
   if (!shim.isFunction(AWS.NodeHttpClient)) {
@@ -30,7 +31,7 @@ function wrapHandleRequest(shim, handleRequest) {
       if (!httpRequest.headers) {
         httpRequest.headers = Object.create(null)
       }
-      httpRequest.headers[shim.DISABLE_DT] = true
+      httpRequest.headers[symbols.disableDT] = true
     } else {
       shim.logger.debug('Unknown arguments to AWS.NodeHttpClient#handleRequest!')
     }

merged/aws-sdk/lib/v2/dynamodb.js → lib/instrumentation/aws-sdk/v2/dynamodb.js

@@ -28,6 +28,10 @@ const DOC_CLIENT_OPERATIONS = [
   'query',
   'scan'
 ]
+const {
+  OperationSpec,
+  params: { DatastoreParameters }
+} = require('../../../shim/specs')
 
 const { setDynamoParameters } = require('../util')
 
@@ -43,13 +47,9 @@ function instrument(shim, AWS) {
       function wrapMethod(shim, original, operationName, args) {
         const params = args[0]
 
-        return new shim.specs.OperationSpec({
+        return new OperationSpec({
           name: operationName,
-          parameters: setDynamoParameters(
-            shim.specs.params.DatastoreParameters,
-            this.endpoint,
-            params
-          ),
+          parameters: setDynamoParameters(this.endpoint, params),
           callback: shim.LAST,
           opaque: true
         })
@@ -74,9 +74,9 @@ function instrument(shim, AWS) {
       // the eventual cached endpoint to be hit is not known at this point.
       const endpoint = this.service && this.service.endpoint
 
-      return new shim.specs.OperationSpec({
+      return new OperationSpec({
         name: dynamoOperation,
-        parameters: new shim.specs.params.DatastoreParameters({
+        parameters: new DatastoreParameters({
           host: endpoint?.host,
           port_path_or_id: endpoint?.port,
           collection: params?.TableName || 'Unknown'

merged/aws-sdk/lib/v2/instrumentation-helper.js → lib/instrumentation/aws-sdk/v2/instrumentation-helper.js

@@ -8,6 +8,7 @@
 /**
  * Series of tests to determine if the library
  * has the features needed to provide instrumentation
+ * @param AWS
  */
 const instrumentationSupported = function instrumentationSupported(AWS) {
   // instrumentation requires the serviceClientOperationsMap property

merged/aws-sdk/lib/v2/sns.js → lib/instrumentation/aws-sdk/v2/sns.js

@@ -4,6 +4,7 @@
  */
 
 'use strict'
+const { MessageSpec } = require('../../../shim/specs')
 
 module.exports = {
   name: 'sns',
@@ -27,7 +28,7 @@ function instrument(shim, AWS) {
 }
 
 function wrapPublish(shim, original, name, args) {
-  return new shim.specs.MessageSpec({
+  return new MessageSpec({
     callback: shim.LAST,
     destinationName: getDestinationName(args[0]),
     destinationType: shim.TOPIC,

merged/aws-sdk/lib/v2/sqs.js → lib/instrumentation/aws-sdk/v2/sqs.js

@@ -6,6 +6,7 @@
 'use strict'
 
 const { grabLastUrlSegment } = require('../util')
+const { MessageSpec } = require('../../../shim/specs')
 module.exports = {
   name: 'sqs',
   type: 'message',
@@ -38,7 +39,7 @@ function recordMessageApi(shim, original, name, args) {
   const params = args[0]
   const queueName = grabLastUrlSegment(params.QueueUrl)
 
-  return new shim.specs.MessageSpec({
+  return new MessageSpec({
     callback: shim.LAST,
     destinationName: queueName,
     destinationType: shim.QUEUE,

merged/aws-sdk/lib/v3/bedrock.js → lib/instrumentation/aws-sdk/v3/bedrock.js

@@ -12,10 +12,11 @@ const {
   BedrockCommand,
   BedrockResponse,
   StreamHandler
-} = require('../llm')
+} = require('../../../llm-events/aws-bedrock')
 
-const { DESTINATIONS } = require('../util')
-const AI_PREFIX = 'Supportability/Nodejs/ML'
+const { DESTINATIONS } = require('../../../config/attribute-filter')
+const { AI } = require('../../../metrics/names')
+const { RecorderSpec } = require('../../../shim/specs')
 
 let TRACKING_METRIC
 
@@ -82,6 +83,7 @@ function addLlmMeta({ agent, segment }) {
  * @param {object} params.segment active segment
  * @param {BedrockCommand} params.bedrockCommand parsed input
  * @param {Error|null} params.err error from request if exists
+ * @param params.bedrockResponse
  */
 function recordChatCompletionMessages({ agent, segment, bedrockCommand, bedrockResponse, err }) {
   const summary = new LlmChatCompletionSummary({
@@ -133,6 +135,7 @@ function recordChatCompletionMessages({ agent, segment, bedrockCommand, bedrockR
  * @param {object} params.segment active segment
  * @param {BedrockCommand} params.bedrockCommand parsed input
  * @param {Error|null} params.err error from request if exists
+ * @param params.bedrockResponse
  */
 function recordEmbeddingMessage({ agent, segment, bedrockCommand, bedrockResponse, err }) {
   const embedding = new LlmEmbedding({
@@ -176,6 +179,8 @@ function createBedrockResponse({ bedrockCommand, response, err }) {
  *
  * @param {object} params { config, commandName } aws config and command name
  * @param {Shim} _shim instance of shim
+ * @param params.commandName
+ * @param shim
  * @param {function} _original original middleware function
  * @param {string} _name function name
  * @param {array} args argument passed to middleware
@@ -188,7 +193,7 @@ function getBedrockSpec({ commandName }, shim, _original, _name, args) {
   const bedrockCommand = new BedrockCommand(input)
   const { modelType } = bedrockCommand
 
-  return new shim.specs.RecorderSpec({
+  return new RecorderSpec({
     promise: true,
     name: `Llm/${modelType}/Bedrock/${commandName}`,
     // eslint-disable-next-line max-params
@@ -219,7 +224,7 @@ function getBedrockSpec({ commandName }, shim, _original, _name, args) {
         shim.logger.warn(
           'ai_monitoring.streaming.enabled is set to `false`, stream will not be instrumented.'
         )
-        agent.metrics.getOrCreateMetric(`${AI_PREFIX}/Streaming/Disabled`).incrementCallCount()
+        agent.metrics.getOrCreateMetric(AI.STREAMING_DISABLED).incrementCallCount()
         addLlmMeta({ agent, segment })
       }
     }
@@ -283,7 +288,7 @@ module.exports.bedrockMiddlewareConfig = {
       return false
     }
 
-    TRACKING_METRIC = `${AI_PREFIX}/Bedrock/${shim.pkgVersion}`
+    TRACKING_METRIC = `${AI.BEDROCK.TRACKING_PREFIX}/${shim.pkgVersion}`
     return true
   },
   type: 'generic',

merged/aws-sdk/lib/v3/common.js → lib/instrumentation/aws-sdk/v3/common.js

@@ -12,8 +12,9 @@ const UNKNOWN = 'Unknown'
  * header to all outgoing requests
  *
  * @param {Shim} shim
+ * @param config
  * @param {function} next next function in middleware chain
- * @return {function}
+ * @returns {function}
  *
  */
 function headerMiddleware(shim, config, next) {
@@ -34,7 +35,8 @@ function headerMiddleware(shim, config, next) {
  * @param {function} next next function in middleware chain
  * @param {Object} contxt AWS command context
  * cons
- * @return {function}
+ * @param context
+ * @returns {function}
  */
 function attrMiddleware(shim, config, next, context) {
   return async function wrappedMiddleware(args) {

merged/aws-sdk/lib/v3/dynamodb.js → lib/instrumentation/aws-sdk/v3/dynamodb.js

@@ -5,6 +5,7 @@
 
 'use strict'
 const { setDynamoParameters } = require('../util')
+const { OperationSpec } = require('../../../shim/specs')
 
 /**
  * Returns the spec for Dynamo commands
@@ -17,9 +18,9 @@ const { setDynamoParameters } = require('../util')
  */
 function getDynamoSpec(shim, original, name, args) {
   const [{ input }] = args
-  return new shim.specs.OperationSpec({
+  return new OperationSpec({
     name: this.commandName,
-    parameters: setDynamoParameters(shim.specs.params.DatastoreParameters, this.endpoint, input),
+    parameters: setDynamoParameters(this.endpoint, input),
     callback: shim.LAST,
     opaque: true,
     promise: true
@@ -33,6 +34,7 @@ function getDynamoSpec(shim, original, name, args) {
  * unless you provide an endpoint override to your library constructor
  * This function will derive the endpoint in that scenario by grabbing the region
  * and building the URL
+ * @param config
  */
 async function getEndpoint(config) {
   if (typeof config.endpoint === 'function') {

merged/aws-sdk/lib/v3/sns.js → lib/instrumentation/aws-sdk/v3/sns.js

@@ -4,12 +4,14 @@
  */
 
 'use strict'
+const { MessageSpec } = require('../../../shim/specs')
 
 /**
  * Middleware hook that records the middleware chain
  * when command is `PublishCommand`
  *
  * @param {Shim} shim
+ * @param config
  * @param {function} next middleware function
  * @param {Object} context
  * @returns {function}
@@ -28,12 +30,13 @@ function snsMiddleware(shim, config, next, context) {
  *
  * @param {Shim} shim
  * @param {original} original original middleware function
+ * @param name
  * @param {Array} args to the middleware function
  * @returns {Object}
  */
 function getSnsSpec(shim, original, name, args) {
   const [command] = args
-  return new shim.specs.MessageSpec({
+  return new MessageSpec({
     promise: true,
     callback: shim.LAST,
     destinationName: getDestinationName(command.input),
@@ -46,7 +49,9 @@ function getSnsSpec(shim, original, name, args) {
  * Helper to set the appropriate destinationName based on
  * the command input
  *
- * @param {Object}
+ * @param root0
+ * @param root0.TopicArn
+ * @param root0.TargetArn
  */
 function getDestinationName({ TopicArn, TargetArn }) {
   return TopicArn || TargetArn || 'PhoneNumber' // We don't want the value of PhoneNumber

merged/aws-sdk/lib/v3/sqs.js → lib/instrumentation/aws-sdk/v3/sqs.js

@@ -10,12 +10,14 @@ const { grabLastUrlSegment } = require('../util')
 const SEND_COMMANDS = ['SendMessageCommand', 'SendMessageBatchCommand']
 
 const RECEIVE_COMMANDS = ['ReceiveMessageCommand']
+const { MessageSpec } = require('../../../shim/specs')
 
 /**
  * Middleware hook that records the middleware chain
  * when command is `PublishCommand`
  *
  * @param {Shim} shim
+ * @param config
  * @param {function} next middleware function
  * @param {Object} context
  * @returns {function}
@@ -36,13 +38,14 @@ function sqsMiddleware(shim, config, next, context) {
  *
  * @param {Shim} shim
  * @param {original} original original middleware function
+ * @param name
  * @param {Array} args to the middleware function
  * @returns {Object}
  */
 function getSqsSpec(shim, original, name, args) {
   const [command] = args
   const { QueueUrl } = command.input
-  return new shim.specs.MessageSpec({
+  return new MessageSpec({
     callback: shim.LAST,
     destinationName: grabLastUrlSegment(QueueUrl),
     destinationType: shim.QUEUE,

lib/instrumentations.js

@@ -19,7 +19,7 @@ module.exports = function instrumentations() {
     '@prisma/client': { type: InstrumentationDescriptor.TYPE_DATASTORE },
     '@redis/client': { type: InstrumentationDescriptor.TYPE_DATASTORE },
     'amqplib': { module: './instrumentation/amqplib' },
-    'aws-sdk': { module: '@newrelic/aws-sdk' },
+    'aws-sdk': { module: './instrumentation/aws-sdk' },
     'bluebird': { type: InstrumentationDescriptor.TYPE_PROMISE },
     'bunyan': { type: InstrumentationDescriptor.TYPE_GENERIC },
     'cassandra-driver': { type: InstrumentationDescriptor.TYPE_DATASTORE },

merged/aws-sdk/lib/llm/event.js → lib/llm-events/aws-bedrock/event.js

@@ -6,7 +6,7 @@
 'use strict'
 
 const { randomUUID } = require('crypto')
-const { DESTINATIONS } = require('../util')
+const { DESTINATIONS } = require('../../config/attribute-filter')
 
 /**
  * @typedef {object} LlmEventParams

lib/metrics/names.js

@@ -180,6 +180,10 @@ AI.OPENAI = {
   COMPLETION: `${AI.COMPLETION}/OpenAI/create`
 }
 
+AI.BEDROCK = {
+  TRACKING_PREFIX: `${AI.TRACKING_PREFIX}/Bedrock`
+}
+
 AI.LANGCHAIN = {
   TRACKING_PREFIX: `${AI.TRACKING_PREFIX}/Langchain`,
   EMBEDDING: `${AI.EMBEDDING}/Langchain`,

lib/shim/shim.js

@@ -55,10 +55,6 @@ module.exports = Shim
 Shim.defineProperty = defineProperty
 Shim.defineProperties = defineProperties
 
-// This is for backwards compat for external libraries like aws-sdk that expect the symbol to be defined here
-defineProperty(Shim, 'DISABLE_DT', symbols.disableDT)
-defineProperty(Shim.prototype, 'DISABLE_DT', symbols.disableDT)
-
 // Copy the argument index enumeration onto the shim.
 Shim.prototype.ARG_INDEXES = specs.ARG_INDEXES
 defineProperties(Shim.prototype, specs.ARG_INDEXES)