diff --git a/lib/instrumentation/aws-sdk/v3/redshift.js b/lib/instrumentation/aws-sdk/v3/redshift.js new file mode 100644 index 0000000000..29752559f6 --- /dev/null +++ b/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 +} diff --git a/lib/instrumentation/aws-sdk/v3/smithy-client.js b/lib/instrumentation/aws-sdk/v3/smithy-client.js index b6fe5404b3..a57aa87811 100644 --- a/lib/instrumentation/aws-sdk/v3/smithy-client.js +++ b/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) { diff --git a/lib/shim/datastore-shim.js b/lib/shim/datastore-shim.js index 3e2c7836a5..3e3a142cad 100644 --- a/lib/shim/datastore-shim.js +++ b/lib/shim/datastore-shim.js @@ -41,6 +41,7 @@ const DATASTORE_NAMES = { OPENSEARCH: 'OpenSearch', POSTGRES: 'Postgres', REDIS: 'Redis', + REDSHIFT: 'Redshift', PRISMA: 'Prisma' } diff --git a/test/versioned/aws-sdk-v3/package.json b/test/versioned/aws-sdk-v3/package.json index 4a8470d7b8..f3d272ca55 100644 --- a/test/versioned/aws-sdk-v3/package.json +++ b/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": {} diff --git a/test/versioned/aws-sdk-v3/redshift-data.test.js b/test/versioned/aws-sdk-v3/redshift-data.test.js new file mode 100644 index 0000000000..90b25bdfdf --- /dev/null +++ b/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, + } +}