Skip to content

Commit 7dceae9

Browse files
feat: added instrumentation for @aws-sdk/client-redshift-data (#2875)
1 parent 46462d0 commit 7dceae9

File tree

5 files changed

+314
-2
lines changed

5 files changed

+314
-2
lines changed
+103
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
/*
2+
* Copyright 2021 New Relic Corporation. All rights reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
'use strict'
7+
const { OperationSpec, QuerySpec } = require('../../../shim/specs')
8+
const InstrumentationDescriptor = require('../../../instrumentation-descriptor')
9+
const {
10+
params: { DatastoreParameters }
11+
} = require('../../../shim/specs')
12+
const UNKNOWN = 'Unknown'
13+
14+
function getRedshiftQuerySpec(shim, original, name, args) {
15+
const [{ input }] = args
16+
17+
const isBatch = Array.isArray(input.Sqls) === true
18+
const SqlQuery = isBatch ? input.Sqls[0] : input.Sql
19+
20+
return new QuerySpec({
21+
query: SqlQuery,
22+
parameters: setRedshiftParameters(this.endpoint, input),
23+
callback: shim.LAST,
24+
opaque: true,
25+
promise: true
26+
})
27+
}
28+
29+
function getRedshiftOperationSpec(shim, original, name, args) {
30+
const [{ input }] = args
31+
return new OperationSpec({
32+
name: this.commandName,
33+
parameters: setRedshiftParameters(this.endpoint, input),
34+
callback: shim.LAST,
35+
opaque: true,
36+
promise: true
37+
})
38+
}
39+
40+
async function getEndpoint(config) {
41+
if (typeof config.endpoint === 'function') {
42+
return await config.endpoint()
43+
}
44+
45+
const region = await config.region()
46+
return new URL(`https://redshift-data.${region}.amazonaws.com`)
47+
}
48+
49+
function redshiftMiddleware(shim, config, next, context) {
50+
const { commandName } = context
51+
return async function wrappedMiddleware(args) {
52+
let endpoint = null
53+
try {
54+
endpoint = await getEndpoint(config)
55+
} catch (err) {
56+
shim.logger.debug(err, 'Failed to get the endpoint.')
57+
}
58+
59+
let wrappedNext
60+
61+
if (commandName === 'ExecuteStatementCommand') {
62+
const getSpec = getRedshiftQuerySpec.bind({ endpoint, commandName })
63+
wrappedNext = shim.recordQuery(next, getSpec)
64+
} else if (commandName === 'BatchExecuteStatementCommand') {
65+
const getSpec = getRedshiftQuerySpec.bind({ endpoint, commandName })
66+
wrappedNext = shim.recordBatchQuery(next, getSpec)
67+
} else {
68+
const getSpec = getRedshiftOperationSpec.bind({ endpoint, commandName })
69+
wrappedNext = shim.recordOperation(next, getSpec)
70+
}
71+
72+
return wrappedNext(args)
73+
}
74+
}
75+
76+
function setRedshiftParameters(endpoint, params) {
77+
return new DatastoreParameters({
78+
host: endpoint && (endpoint.host || endpoint.hostname),
79+
port_path_or_id: (endpoint && endpoint.port) || 443,
80+
database_name: (params && params.Database) || UNKNOWN
81+
})
82+
}
83+
84+
const redshiftMiddlewareConfig = [
85+
{
86+
middleware: redshiftMiddleware,
87+
init(shim) {
88+
shim.setDatastore(shim.REDSHIFT)
89+
return true
90+
},
91+
type: InstrumentationDescriptor.TYPE_DATASTORE,
92+
config: {
93+
name: 'NewRelicRedshiftMiddleware',
94+
step: 'initialize',
95+
priority: 'high',
96+
override: true
97+
}
98+
}
99+
]
100+
101+
module.exports = {
102+
redshiftMiddlewareConfig
103+
}

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

+3-1
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ const { sqsMiddlewareConfig } = require('./sqs')
1111
const { dynamoMiddlewareConfig } = require('./dynamodb')
1212
const { lambdaMiddlewareConfig } = require('./lambda')
1313
const { bedrockMiddlewareConfig } = require('./bedrock')
14+
const { redshiftMiddlewareConfig } = require('./redshift')
1415
const MIDDLEWARE = Symbol('nrMiddleware')
1516

1617
const middlewareByClient = {
@@ -20,7 +21,8 @@ const middlewareByClient = {
2021
SNS: [...middlewareConfig, snsMiddlewareConfig],
2122
SQS: [...middlewareConfig, sqsMiddlewareConfig],
2223
DynamoDB: [...middlewareConfig, ...dynamoMiddlewareConfig],
23-
DynamoDBDocument: [...middlewareConfig, ...dynamoMiddlewareConfig]
24+
DynamoDBDocument: [...middlewareConfig, ...dynamoMiddlewareConfig],
25+
RedshiftData: [...middlewareConfig, ...redshiftMiddlewareConfig]
2426
}
2527

2628
module.exports = function instrumentSmithyClient(shim, smithyClientExport) {

lib/shim/datastore-shim.js

+1
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ const DATASTORE_NAMES = {
4141
OPENSEARCH: 'OpenSearch',
4242
POSTGRES: 'Postgres',
4343
REDIS: 'Redis',
44+
REDSHIFT: 'Redshift',
4445
PRISMA: 'Prisma'
4546
}
4647

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

+16-1
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,8 @@
1010
{"name": "@aws-sdk/lib-dynamodb", "minAgentVersion": "8.7.1"},
1111
{"name": "@aws-sdk/smithy-client", "minSupported": "3.47.0", "minAgentVersion": "8.7.1"},
1212
{"name": "@smithy/smithy-client", "minSupported": "2.0.0", "minAgentVersion": "11.0.0"},
13-
{"name": "@aws-sdk/client-bedrock-runtime", "minAgentVersion": "11.13.0"}
13+
{"name": "@aws-sdk/client-bedrock-runtime", "minAgentVersion": "11.13.0"},
14+
{"name": "@aws-sdk/client-redshift-data", "minAgentVersion": "12.12.0"}
1415
],
1516
"version": "0.0.0",
1617
"private": true,
@@ -209,6 +210,20 @@
209210
"bedrock-embeddings.test.js",
210211
"bedrock-negative-tests.test.js"
211212
]
213+
},
214+
{
215+
"engines": {
216+
"node": ">=18.0"
217+
},
218+
"dependencies": {
219+
"@aws-sdk/client-redshift-data": {
220+
"versions": ">=3.474.0",
221+
"samples": 2
222+
}
223+
},
224+
"files": [
225+
"redshift-data.test.js"
226+
]
212227
}
213228
],
214229
"dependencies": {}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,191 @@
1+
/*
2+
* Copyright 2020 New Relic Corporation. All rights reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
'use strict'
7+
const assert = require('node:assert')
8+
const test = require('node:test')
9+
const { createEmptyResponseServer, FAKE_CREDENTIALS } = require('../../lib/aws-server-stubs')
10+
const common = require('./common')
11+
const helper = require('../../lib/agent_helper')
12+
const { match } = require('../../lib/custom-assertions')
13+
14+
test('Redshift-data', async (t) => {
15+
t.beforeEach(async (ctx) => {
16+
ctx.nr = {}
17+
const server = createEmptyResponseServer()
18+
19+
await new Promise((resolve) => {
20+
server.listen(0, resolve)
21+
})
22+
23+
ctx.nr.server = server
24+
ctx.nr.agent = helper.instrumentMockedAgent()
25+
26+
const lib = require('@aws-sdk/client-redshift-data')
27+
28+
ctx.nr.redshiftCommands = {
29+
ExecuteStatementCommand: lib.ExecuteStatementCommand,
30+
BatchExecuteStatementCommand: lib.BatchExecuteStatementCommand,
31+
DescribeStatementCommand: lib.DescribeStatementCommand,
32+
GetStatementResultCommand: lib.GetStatementResultCommand,
33+
ListDatabasesCommand: lib.ListDatabasesCommand
34+
}
35+
36+
const endpoint = `http://localhost:${server.address().port}`
37+
ctx.nr.client = new lib.RedshiftDataClient({
38+
credentials: FAKE_CREDENTIALS,
39+
endpoint,
40+
region: 'us-east-1'
41+
})
42+
43+
ctx.nr.tests = createTests()
44+
})
45+
46+
t.afterEach(common.afterEach)
47+
48+
await t.test('client commands', (t, end) => {
49+
const { redshiftCommands, client, agent, tests } = t.nr
50+
helper.runInTransaction(agent, async function (tx) {
51+
for (const test of tests) {
52+
const CommandClass = redshiftCommands[test.command]
53+
const command = new CommandClass(test.params)
54+
await client.send(command)
55+
}
56+
57+
tx.end()
58+
59+
const args = [end, tests, tx]
60+
setImmediate(finish, ...args)
61+
})
62+
})
63+
})
64+
65+
function finish(end, tests, tx) {
66+
const root = tx.trace.root
67+
const segments = common.checkAWSAttributes({ trace: tx.trace, segment: root, pattern: common.DATASTORE_PATTERN })
68+
assert.equal(segments.length, tests.length, `should have ${tests.length} aws datastore segments`)
69+
70+
const externalSegments = common.checkAWSAttributes({ trace: tx.trace, segment: root, pattern: common.EXTERN_PATTERN })
71+
assert.equal(externalSegments.length, 0, 'should not have any External segments')
72+
73+
segments.forEach((segment, i) => {
74+
const command = tests[i].command
75+
76+
if (tests[i].command === 'ExecuteStatementCommand' || tests[i].command === 'BatchExecuteStatementCommand') {
77+
assert.equal(
78+
segment.name,
79+
`Datastore/statement/Redshift/${tests[i].tableName}/${tests[i].queryType}`,
80+
'should have table name and query type in segment name'
81+
)
82+
} else {
83+
assert.equal(
84+
segment.name,
85+
`Datastore/operation/Redshift/${command}`,
86+
'should have command in segment name'
87+
)
88+
}
89+
90+
const attrs = segment.attributes.get(common.SEGMENT_DESTINATION)
91+
attrs.port_path_or_id = parseInt(attrs.port_path_or_id, 10)
92+
match(attrs, {
93+
host: String,
94+
port_path_or_id: Number,
95+
product: 'Redshift',
96+
database_name: String,
97+
'aws.operation': command,
98+
'aws.requestId': String,
99+
'aws.region': 'us-east-1',
100+
'aws.service': 'Redshift Data',
101+
})
102+
103+
assert(attrs.host, 'should have host')
104+
})
105+
106+
end()
107+
}
108+
109+
function createTests() {
110+
const insertData = insertDataIntoTable()
111+
const selectData = selectDataFromTable()
112+
const updateData = updateDataInTable()
113+
const deleteData = deleteDataFromTable()
114+
const insertBatchData = insertBatchDataIntoTable()
115+
const describeSqlStatement = describeStatement()
116+
const getSqlStatement = getStatement()
117+
const getDatabases = listDatabases()
118+
119+
return [
120+
{ params: insertData, tableName, queryType: 'insert', command: 'ExecuteStatementCommand' },
121+
{ params: selectData, tableName, queryType: 'select', command: 'ExecuteStatementCommand' },
122+
{ params: updateData, tableName, queryType: 'update', command: 'ExecuteStatementCommand' },
123+
{ params: deleteData, tableName, queryType: 'delete', command: 'ExecuteStatementCommand' },
124+
{ params: insertBatchData, tableName, queryType: 'insert', command: 'BatchExecuteStatementCommand' },
125+
{ params: describeSqlStatement, command: 'DescribeStatementCommand' },
126+
{ params: getSqlStatement, command: 'GetStatementResultCommand' },
127+
{ params: getDatabases, command: 'ListDatabasesCommand' }
128+
]
129+
}
130+
131+
const commonParams = {
132+
Database: 'dev',
133+
DbUser: 'a_user',
134+
ClusterIdentifier: 'a_cluster'
135+
}
136+
137+
const tableName = 'test_table'
138+
139+
function insertDataIntoTable() {
140+
return {
141+
...commonParams,
142+
Sql: `INSERT INTO ${tableName} (id, name) VALUES (1, 'test')`
143+
}
144+
}
145+
146+
function selectDataFromTable() {
147+
return {
148+
...commonParams,
149+
Sql: `SELECT id, name FROM ${tableName}`
150+
}
151+
}
152+
153+
function updateDataInTable() {
154+
return {
155+
...commonParams,
156+
Sql: `UPDATE ${tableName} SET name = 'updated' WHERE id = 1`
157+
}
158+
}
159+
160+
function deleteDataFromTable() {
161+
return {
162+
...commonParams,
163+
Sql: `DELETE FROM ${tableName} WHERE id = 1`
164+
}
165+
}
166+
167+
function insertBatchDataIntoTable() {
168+
return {
169+
...commonParams,
170+
Sqls: ['INSERT INTO test_table (id, name) VALUES (2, \'test2\')', 'INSERT INTO test_table (id, name) VALUES (3, \'test3\')']
171+
}
172+
}
173+
174+
function describeStatement() {
175+
return {
176+
Id: 'a_statement_id'
177+
}
178+
}
179+
180+
function getStatement() {
181+
return {
182+
Id: 'a_statement_id',
183+
NextToken: 'a_token'
184+
}
185+
}
186+
187+
function listDatabases() {
188+
return {
189+
...commonParams,
190+
}
191+
}

0 commit comments

Comments
 (0)