Skip to content

Commit

Permalink
fix(core): capture command data for SDK v3 clients (#611)
Browse files Browse the repository at this point in the history 
* fix(core): capture additional command data with SDK v3

* adds unit tests for abandoned PR #527

* adds output mapping to test

* removes leftover util import

* requested changes on #611

---------

Co-authored-by: SergeiPoluektov <sergei@hellohakuna.com>
  • Loading branch information
@danno-s @SergeyPoluektov
danno-s and SergeyPoluektov authored Sep 18, 2023
1 parent e32dd95 commit 5eb97ef
Show file tree
Hide file tree
Showing 2 changed files with 93 additions and 28 deletions.
11 changes: 9 additions & 2 deletions packages/core/lib/patchers/aws3_p.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@ const buildAttributesFromMetadata = async (
service: string,
operation: string,
region: string,
commandInput: any,
res: any | null,
error: SdkError | null,
): Promise<[ServiceSegment, HttpResponse]> => {
Expand All @@ -49,8 +50,10 @@ const buildAttributesFromMetadata = async (
extendedRequestId,
requestId,
retryCount: attempts,
data: res?.output,
request: {
operation,
params: commandInput,
httpRequest: {
region,
statusCode,
Expand Down Expand Up @@ -93,11 +96,13 @@ function addFlags(http: HttpResponse, subsegment: Subsegment, err?: SdkError): v
const getXRayMiddleware = (config: RegionResolvedConfig, manualSegment?: SegmentLike): BuildMiddleware<any, any> => (next: any, context: any) => async (args: any) => {
const segment = contextUtils.isAutomaticMode() ? contextUtils.resolveSegment() : manualSegment;
const {clientName, commandName} = context;
const operation: string = commandName.slice(0, -7); // Strip trailing "Command" string
const commandInput = args?.input ?? {};
const commandOperation: string = commandName.slice(0, -7); // Strip trailing "Command" string
const operation: string = commandOperation.charAt(0).toLowerCase() + commandOperation.slice(1);
const service: string = clientName.slice(0, -6); // Strip trailing "Client" string

if (!segment) {
const output = service + '.' + operation.charAt(0).toLowerCase() + operation.slice(1);
const output = service + '.' + operation;

if (!contextUtils.isAutomaticMode()) {
logger.getLogger().info('Call ' + output + ' requires a segment object' +
Expand Down Expand Up @@ -148,6 +153,7 @@ const getXRayMiddleware = (config: RegionResolvedConfig, manualSegment?: Segment
service,
operation,
await config.region(),
commandInput,
res,
null,
);
Expand All @@ -164,6 +170,7 @@ const getXRayMiddleware = (config: RegionResolvedConfig, manualSegment?: Segment
service,
operation,
await config.region(),
commandInput,
null,
err,
);
Expand Down
110 changes: 84 additions & 26 deletions packages/core/test/unit/patchers/aws3_p.test.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,38 +18,38 @@ chai.use(sinonChai);

var traceId = '1-57fbe041-2c7ad569f5d6ff149137be86';

describe('AWS v3 patcher', function() {
describe('#captureAWSClient', function() {
describe('AWS v3 patcher', function () {
describe('#captureAWSClient', function () {
var sandbox, useMiddleware;

var awsClient = {
send: function() {},
send: function () { },
config: {
serviceId: 's3',
},
middlewareStack: constructStack(),
};

beforeEach(function() {
beforeEach(function () {
sandbox = sinon.createSandbox();
useMiddleware = sandbox.stub(awsClient.middlewareStack, 'use');
});

afterEach(function() {
afterEach(function () {
sandbox.restore();
});

it('should call middlewareStack.use and return the service', function() {
it('should call middlewareStack.use and return the service', function () {
const patched = awsPatcher.captureAWSClient(awsClient);
useMiddleware.should.have.been.calledOnce;
assert.equal(patched, awsClient);
});
});

describe('#captureAWSRequest', function() {
var awsClient, awsRequest, sandbox, segment, stubResolve, addNewSubsegmentStub, sub;
describe('#captureAWSRequest', function () {
let awsClient, awsRequest, ddbClient, ddbGetRequest, sandbox, segment, stubResolve, addNewSubsegmentStub, sub;

before(function() {
before(function () {
awsClient = {
send: async (req) => {
const context = {
Expand All @@ -74,9 +74,34 @@ describe('AWS v3 patcher', function() {
},
middlewareStack: constructStack(),
};

ddbClient = {
send: async (req) => {
const middlewareContext = {
clientName: 'DynamoDBClient',
commandName: 'GetItemCommand',
};
const handler = awsClient.middlewareStack.resolve((args) => {
const error = req.response.error;
if (error) {
const err = new Error(error.message);
err.name = error.code;
err.$metadata = req.response.$metadata;
throw err;
}
return args;
}, middlewareContext);
await handler(req);
return req.response;
},
config: {
region: async () => 'us-east-1',
},
middlewareStack: constructStack(),
};
});

beforeEach(function() {
beforeEach(function () {
sandbox = sinon.createSandbox();

awsRequest = new (class ListBucketsCommand {
Expand All @@ -99,14 +124,39 @@ describe('AWS v3 patcher', function() {
}
})();

ddbGetRequest = new (class GetItemCommand {
constructor() {
this.request = {
method: 'GET',
url: '/',
connection: {
remoteAddress: 'localhost'
},
headers: {},
};
this.input = {
TableName: 'TestTableName',
ConsistentRead: true,
};
this.response = {};
this.output = {
ConsumedCapacity: 10,
$metadata: {
requestId: '123',
extendedRequestId: '456',
},
};
}
})();

segment = new Segment('testSegment', traceId);
segment.additionalTraceData = {'Foo': 'bar'};
segment.additionalTraceData = { 'Foo': 'bar' };
sub = segment.addNewSubsegment('subseg');
stubResolve = sandbox.stub(contextUtils, 'resolveSegment').returns(segment);
addNewSubsegmentStub = sandbox.stub(segment, 'addNewSubsegment').returns(sub);
});

afterEach(function() {
afterEach(function () {
sandbox.restore();
});

Expand All @@ -125,13 +175,13 @@ describe('AWS v3 patcher', function() {

awsClient.send(awsRequest);

setTimeout(function() {
setTimeout(function () {
logStub.should.have.been.calledOnce;
done();
}, 50);
});

it('should inject the tracing headers', async function() {
it('should inject the tracing headers', async function () {
await awsClient.send(awsRequest);
assert.isTrue(addNewSubsegmentStub.calledWith('S3'));

Expand All @@ -140,7 +190,7 @@ describe('AWS v3 patcher', function() {
assert.match(awsRequest.request.headers['X-Amzn-Trace-Id'], expected);
});

it('should close on complete with no errors when code 200 is seen', async function() {
it('should close on complete with no errors when code 200 is seen', async function () {
const closeStub = sandbox.stub(sub, 'close').returns();
sandbox.stub(sub, 'addAttribute').returns();
sandbox.stub(Aws.prototype, 'init').returns();
Expand All @@ -154,7 +204,7 @@ describe('AWS v3 patcher', function() {
closeStub.should.have.been.calledWithExactly();
});

it('should mark the subsegment as throttled and error if code 429 is seen', async function() {
it('should mark the subsegment as throttled and error if code 429 is seen', async function () {
const throttleStub = sandbox.stub(sub, 'addThrottleFlag').returns();

sandbox.stub(sub, 'addAttribute').returns();
Expand All @@ -171,7 +221,7 @@ describe('AWS v3 patcher', function() {
assert.isTrue(sub.error);
});

it('should mark the subsegment as throttled and error if code service.throttledError returns true, regardless of status code', async function() {
it('should mark the subsegment as throttled and error if code service.throttledError returns true, regardless of status code', async function () {
const throttleStub = sandbox.stub(sub, 'addThrottleFlag').returns();

sandbox.stub(sub, 'addAttribute').returns();
Expand All @@ -188,7 +238,7 @@ describe('AWS v3 patcher', function() {
assert.isTrue(sub.error);
});

it('should capture an error on the response and mark exception as remote', async function() {
it('should capture an error on the response and mark exception as remote', async function () {
const closeStub = sandbox.stub(sub, 'close').returns();
const getCauseStub = sandbox.stub(Utils, 'getCauseTypeFromHttpStatus').returns();

Expand All @@ -205,7 +255,15 @@ describe('AWS v3 patcher', function() {
await awsClient.send(awsRequest).catch(() => null);

getCauseStub.should.have.been.calledWithExactly(500);
closeStub.should.have.been.calledWithExactly(sinon.match({ message: error.message, name: error.code}), true);
closeStub.should.have.been.calledWithExactly(sinon.match({ message: error.message, name: error.code }), true);
});

it('should pass params into aws subsegment', async function () {
await ddbClient.send(ddbGetRequest);

assert.isTrue(addNewSubsegmentStub.calledWith('DynamoDB'));
addNewSubsegmentStub.returnValues[0].should.have.property('aws');
addNewSubsegmentStub.returnValues[0].aws.should.include({ consistent_read: true, table_name: 'TestTableName', consumed_capacity: 10 });
});
});

Expand All @@ -220,7 +278,7 @@ describe('AWS v3 patcher', function() {

awsClient.send(awsRequest);

setTimeout(function() {
setTimeout(function () {
logStub.should.have.been.calledOnce;
done();
}, 50);
Expand Down Expand Up @@ -250,10 +308,10 @@ describe('AWS v3 patcher', function() {
});


describe('#captureAWSRequest-Unsampled', function() {
describe('#captureAWSRequest-Unsampled', function () {
var awsClient, awsRequest, sandbox, segment, stubResolve, addNewSubsegmentStub, sub, service, addNewServiceSubsegmentStub;

before(function() {
before(function () {
awsClient = {
send: async (req) => {
const context = {
Expand All @@ -280,7 +338,7 @@ describe('AWS v3 patcher', function() {
};
});

beforeEach(function() {
beforeEach(function () {
sandbox = sinon.createSandbox();

awsRequest = new (class ListBucketsCommand {
Expand All @@ -304,7 +362,7 @@ describe('AWS v3 patcher', function() {
})();

segment = new Segment('testSegment', traceId);
segment.additionalTraceData = {'Foo': 'bar'};
segment.additionalTraceData = { 'Foo': 'bar' };
sub = segment.addNewSubsegmentWithoutSampling('subseg');
service = sub.addNewSubsegmentWithoutSampling('service');

Expand All @@ -314,7 +372,7 @@ describe('AWS v3 patcher', function() {

});

afterEach(function() {
afterEach(function () {
sandbox.restore();
});

Expand All @@ -328,7 +386,7 @@ describe('AWS v3 patcher', function() {
});


it('should inject the tracing headers', async function() {
it('should inject the tracing headers', async function () {
await awsClient.send(awsRequest);
assert.isTrue(addNewServiceSubsegmentStub.calledWith('S3'));

Expand Down

0 comments on commit 5eb97ef

Please sign in to comment.