aws / aws-xray-sdk-node

The official AWS X-Ray SDK for Node.js.
Apache License 2.0
267 stars 156 forks source link

Adding custom attributes to AWS SDK v3 calls #601

Open paddymccarroll opened 1 year ago

paddymccarroll commented 1 year ago

Is it possible to add our own custom attributes to segments created using captureAWSv3Client? For example, I would like to add query information from DynamoDB calls.

wangzlei commented 1 year ago

yes, AWS SDK v2 instrumentation is by middleware stack, user can customize XRay SDK's AWS sdk v3 instrumentation here https://github.com/aws/aws-xray-sdk-node/blob/e21281fc98d7a8a4eb2d96b64de60c295f474c1b/packages/core/lib/patchers/aws3_p.ts

One idea I haven't tried. User can add a new middleware after xray sdk's, use API AWSXRay.getSegment(); to add more attributes on AWS SDK v2 subsegment.

paddymccarroll commented 1 year ago

Thanks for pointing me in the right direction, I'll give that a go.

petermyers commented 2 months ago

Thanks for pointing me in the right direction, I'll give that a go.

Were you able to get this to work? I'd also like to add query information to the captured calls.

petermyers commented 2 months ago

In case anyone ends up here in the future, I got this to work for sdk v3 by adding a middleware to add the input as metadata to the segment. AWSXray.getSegment() was returning the aws lambda facade segment, so I just did a find for the DynamoDB segment that gets attached as a subsegment.

Here's the full code:

export const captureSdkClientTraces = (sdkClient, useDdbMiddleware = false) => {
  const capturedClient = AWSXray.captureAWSv3Client(sdkClient);
  if (useDdbMiddleware) applyDdbMiddleware(capturedClient);
  return capturedClient;
};

const ddbMiddleware = (next, context) => async (args) => {
  const segment = AWSXray.getSegment()?.subsegments.find((subsegment) =>
    subsegment.id === extractParentIdFromTraceHeader(args.request?.headers['X-Amzn-Trace-Id']));
  const { input } = args;

  segment?.addMetadata('Input', input);

  return next(args);
};

const extractParentIdFromTraceHeader = (traceHeader) => {
  // Trace header of the form: Root=ROOT_ID;Parent=PARENT_ID;Sampled=0;Lineage=OTHER_ID:0
  const match = traceHeader?.match(/^.+(Parent=[a-zA-Z0-9]+);.*$/);
  return match?.[1].split('=')[1];
};

const applyDdbMiddleware = (capturedClient) => {
  capturedClient.middlewareStack.remove(DDB_MIDDLEWARE_NAME);
  capturedClient.middlewareStack.use({
    applyToStack: (stack) => stack.addRelativeTo(ddbMiddleware, {
      name: DDB_MIDDLEWARE_NAME,
      relation: 'after',
      toMiddleware: 'XRaySDKInstrumentation',
    }),
  });
};