feat: added instrumentation for @aws-sdk/client-redshift-data (#2875)

lib/instrumentation/aws-sdk/v3/redshift.js

@@ -0,0 +1,103 @@
+/*
+ * Copyright 2021 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+const { OperationSpec, QuerySpec } = require('../../../shim/specs')
+const InstrumentationDescriptor = require('../../../instrumentation-descriptor')
+const {
+  params: { DatastoreParameters }
+} = require('../../../shim/specs')
+const UNKNOWN = 'Unknown'
+
+function getRedshiftQuerySpec(shim, original, name, args) {
+  const [{ input }] = args
+
+  const isBatch = Array.isArray(input.Sqls) === true
+  const SqlQuery = isBatch ? input.Sqls[0] : input.Sql
+
+  return new QuerySpec({
+    query: SqlQuery,
+    parameters: setRedshiftParameters(this.endpoint, input),
+    callback: shim.LAST,
+    opaque: true,
+    promise: true
+  })
+}
+
+function getRedshiftOperationSpec(shim, original, name, args) {
+  const [{ input }] = args
+  return new OperationSpec({
+    name: this.commandName,
+    parameters: setRedshiftParameters(this.endpoint, input),
+    callback: shim.LAST,
+    opaque: true,
+    promise: true
+  })
+}
+
+async function getEndpoint(config) {
+  if (typeof config.endpoint === 'function') {
+    return await config.endpoint()
+  }
+
+  const region = await config.region()
+  return new URL(`https://redshift-data.${region}.amazonaws.com`)
+}
+
+function redshiftMiddleware(shim, config, next, context) {
+  const { commandName } = context
+  return async function wrappedMiddleware(args) {
+    let endpoint = null
+    try {
+      endpoint = await getEndpoint(config)
+    } catch (err) {
+      shim.logger.debug(err, 'Failed to get the endpoint.')
+    }
+
+    let wrappedNext
+
+    if (commandName === 'ExecuteStatementCommand') {
+      const getSpec = getRedshiftQuerySpec.bind({ endpoint, commandName })
+      wrappedNext = shim.recordQuery(next, getSpec)
+    } else if (commandName === 'BatchExecuteStatementCommand') {
+      const getSpec = getRedshiftQuerySpec.bind({ endpoint, commandName })
+      wrappedNext = shim.recordBatchQuery(next, getSpec)
+    } else {
+      const getSpec = getRedshiftOperationSpec.bind({ endpoint, commandName })
+      wrappedNext = shim.recordOperation(next, getSpec)
+    }
+
+    return wrappedNext(args)
+  }
+}
+
+function setRedshiftParameters(endpoint, params) {
+  return new DatastoreParameters({
+    host: endpoint && (endpoint.host || endpoint.hostname),
+    port_path_or_id: (endpoint && endpoint.port) || 443,
+    database_name: (params && params.Database) || UNKNOWN
+  })
+}
+
+const redshiftMiddlewareConfig = [
+  {
+    middleware: redshiftMiddleware,
+    init(shim) {
+      shim.setDatastore(shim.REDSHIFT)
+      return true
+    },
+    type: InstrumentationDescriptor.TYPE_DATASTORE,
+    config: {
+      name: 'NewRelicRedshiftMiddleware',
+      step: 'initialize',
+      priority: 'high',
+      override: true
+    }
+  }
+]
+
+module.exports = {
+  redshiftMiddlewareConfig
+}

lib/instrumentation/aws-sdk/v3/smithy-client.js

@@ -11,6 +11,7 @@ const { sqsMiddlewareConfig } = require('./sqs')
 const { dynamoMiddlewareConfig } = require('./dynamodb')
 const { lambdaMiddlewareConfig } = require('./lambda')
 const { bedrockMiddlewareConfig } = require('./bedrock')
+const { redshiftMiddlewareConfig } = require('./redshift')
 const MIDDLEWARE = Symbol('nrMiddleware')
 
 const middlewareByClient = {
@@ -20,7 +21,8 @@ const middlewareByClient = {
   SNS: [...middlewareConfig, snsMiddlewareConfig],
   SQS: [...middlewareConfig, sqsMiddlewareConfig],
   DynamoDB: [...middlewareConfig, ...dynamoMiddlewareConfig],
-  DynamoDBDocument: [...middlewareConfig, ...dynamoMiddlewareConfig]
+  DynamoDBDocument: [...middlewareConfig, ...dynamoMiddlewareConfig],
+  RedshiftData: [...middlewareConfig, ...redshiftMiddlewareConfig]
 }
 
 module.exports = function instrumentSmithyClient(shim, smithyClientExport) {

lib/shim/datastore-shim.js

@@ -41,6 +41,7 @@ const DATASTORE_NAMES = {
   OPENSEARCH: 'OpenSearch',
   POSTGRES: 'Postgres',
   REDIS: 'Redis',
+  REDSHIFT: 'Redshift',
   PRISMA: 'Prisma'
 }
 

test/versioned/aws-sdk-v3/package.json

@@ -10,7 +10,8 @@
     {"name": "@aws-sdk/lib-dynamodb", "minAgentVersion": "8.7.1"},
     {"name": "@aws-sdk/smithy-client", "minSupported": "3.47.0", "minAgentVersion": "8.7.1"},
     {"name": "@smithy/smithy-client", "minSupported": "2.0.0", "minAgentVersion": "11.0.0"},
-    {"name": "@aws-sdk/client-bedrock-runtime", "minAgentVersion": "11.13.0"}
+    {"name": "@aws-sdk/client-bedrock-runtime", "minAgentVersion": "11.13.0"},
+    {"name": "@aws-sdk/client-redshift-data", "minAgentVersion": "12.12.0"}
   ],
   "version": "0.0.0",
   "private": true,
@@ -209,6 +210,20 @@
         "bedrock-embeddings.test.js",
         "bedrock-negative-tests.test.js"
       ]
+    },
+    {
+      "engines": {
+        "node": ">=18.0"
+      },
+      "dependencies": {
+        "@aws-sdk/client-redshift-data": {
+          "versions": ">=3.474.0",
+          "samples": 2
+        }
+      },
+      "files": [
+        "redshift-data.test.js"
+      ]
     }
   ],
   "dependencies": {}

test/versioned/aws-sdk-v3/redshift-data.test.js

@@ -0,0 +1,191 @@
+/*
+ * Copyright 2020 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+const assert = require('node:assert')
+const test = require('node:test')
+const { createEmptyResponseServer, FAKE_CREDENTIALS } = require('../../lib/aws-server-stubs')
+const common = require('./common')
+const helper = require('../../lib/agent_helper')
+const { match } = require('../../lib/custom-assertions')
+
+test('Redshift-data', async (t) => {
+  t.beforeEach(async (ctx) => {
+    ctx.nr = {}
+    const server = createEmptyResponseServer()
+
+    await new Promise((resolve) => {
+      server.listen(0, resolve)
+    })
+
+    ctx.nr.server = server
+    ctx.nr.agent = helper.instrumentMockedAgent()
+
+    const lib = require('@aws-sdk/client-redshift-data')
+
+    ctx.nr.redshiftCommands = {
+      ExecuteStatementCommand: lib.ExecuteStatementCommand,
+      BatchExecuteStatementCommand: lib.BatchExecuteStatementCommand,
+      DescribeStatementCommand: lib.DescribeStatementCommand,
+      GetStatementResultCommand: lib.GetStatementResultCommand,
+      ListDatabasesCommand: lib.ListDatabasesCommand
+    }
+
+    const endpoint = `http://localhost:${server.address().port}`
+    ctx.nr.client = new lib.RedshiftDataClient({
+      credentials: FAKE_CREDENTIALS,
+      endpoint,
+      region: 'us-east-1'
+    })
+
+    ctx.nr.tests = createTests()
+  })
+
+  t.afterEach(common.afterEach)
+
+  await t.test('client commands', (t, end) => {
+    const { redshiftCommands, client, agent, tests } = t.nr
+    helper.runInTransaction(agent, async function (tx) {
+      for (const test of tests) {
+        const CommandClass = redshiftCommands[test.command]
+        const command = new CommandClass(test.params)
+        await client.send(command)
+      }
+
+      tx.end()
+
+      const args = [end, tests, tx]
+      setImmediate(finish, ...args)
+    })
+  })
+})
+
+function finish(end, tests, tx) {
+  const root = tx.trace.root
+  const segments = common.checkAWSAttributes({ trace: tx.trace, segment: root, pattern: common.DATASTORE_PATTERN })
+  assert.equal(segments.length, tests.length, `should have ${tests.length} aws datastore segments`)
+
+  const externalSegments = common.checkAWSAttributes({ trace: tx.trace, segment: root, pattern: common.EXTERN_PATTERN })
+  assert.equal(externalSegments.length, 0, 'should not have any External segments')
+
+  segments.forEach((segment, i) => {
+    const command = tests[i].command
+
+    if (tests[i].command === 'ExecuteStatementCommand' || tests[i].command === 'BatchExecuteStatementCommand') {
+      assert.equal(
+        segment.name,
+        `Datastore/statement/Redshift/${tests[i].tableName}/${tests[i].queryType}`,
+        'should have table name and query type in segment name'
+      )
+    } else {
+      assert.equal(
+        segment.name,
+        `Datastore/operation/Redshift/${command}`,
+        'should have command in segment name'
+      )
+    }
+
+    const attrs = segment.attributes.get(common.SEGMENT_DESTINATION)
+    attrs.port_path_or_id = parseInt(attrs.port_path_or_id, 10)
+    match(attrs, {
+      host: String,
+      port_path_or_id: Number,
+      product: 'Redshift',
+      database_name: String,
+      'aws.operation': command,
+      'aws.requestId': String,
+      'aws.region': 'us-east-1',
+      'aws.service': 'Redshift Data',
+    })
+
+    assert(attrs.host, 'should have host')
+  })
+
+  end()
+}
+
+function createTests() {
+  const insertData = insertDataIntoTable()
+  const selectData = selectDataFromTable()
+  const updateData = updateDataInTable()
+  const deleteData = deleteDataFromTable()
+  const insertBatchData = insertBatchDataIntoTable()
+  const describeSqlStatement = describeStatement()
+  const getSqlStatement = getStatement()
+  const getDatabases = listDatabases()
+
+  return [
+    { params: insertData, tableName, queryType: 'insert', command: 'ExecuteStatementCommand' },
+    { params: selectData, tableName, queryType: 'select', command: 'ExecuteStatementCommand' },
+    { params: updateData, tableName, queryType: 'update', command: 'ExecuteStatementCommand' },
+    { params: deleteData, tableName, queryType: 'delete', command: 'ExecuteStatementCommand' },
+    { params: insertBatchData, tableName, queryType: 'insert', command: 'BatchExecuteStatementCommand' },
+    { params: describeSqlStatement, command: 'DescribeStatementCommand' },
+    { params: getSqlStatement, command: 'GetStatementResultCommand' },
+    { params: getDatabases, command: 'ListDatabasesCommand' }
+  ]
+}
+
+const commonParams = {
+  Database: 'dev',
+  DbUser: 'a_user',
+  ClusterIdentifier: 'a_cluster'
+}
+
+const tableName = 'test_table'
+
+function insertDataIntoTable() {
+  return {
+    ...commonParams,
+    Sql: `INSERT INTO ${tableName} (id, name) VALUES (1, 'test')`
+  }
+}
+
+function selectDataFromTable() {
+  return {
+    ...commonParams,
+    Sql: `SELECT id, name FROM ${tableName}`
+  }
+}
+
+function updateDataInTable() {
+  return {
+    ...commonParams,
+    Sql: `UPDATE ${tableName} SET name = 'updated' WHERE id = 1`
+  }
+}
+
+function deleteDataFromTable() {
+  return {
+    ...commonParams,
+    Sql: `DELETE FROM ${tableName} WHERE id = 1`
+  }
+}
+
+function insertBatchDataIntoTable() {
+  return {
+    ...commonParams,
+    Sqls: ['INSERT INTO test_table (id, name) VALUES (2, \'test2\')', 'INSERT INTO test_table (id, name) VALUES (3, \'test3\')']
+  }
+}
+
+function describeStatement() {
+  return {
+    Id: 'a_statement_id'
+  }
+}
+
+function getStatement() {
+  return {
+    Id: 'a_statement_id',
+    NextToken: 'a_token'
+  }
+}
+
+function listDatabases() {
+  return {
+    ...commonParams,
+  }
+}