Add AWSResources integration

This integration traces AWS service calls as spans.

Author: marshall-lee

packages/serverless/README.md

@@ -41,3 +41,14 @@ exports.handler = Sentry.AWSLambda.wrapHandler((event, context, callback) => {
   throw new Error('oh, hello there!');
 });
 ```
+
+If you also want to trace AWS service calls, add an `AWSResources` integration. They will be available to observe as nested spans for each transaction.
+
+```javascript
+Sentry.init({
+  dsn: '__DSN__',
+  integrations: [
+    new Sentry.AWSResources(),
+  ],
+});
+```

packages/serverless/package.json

@@ -27,8 +27,10 @@
     "@sentry-internal/eslint-config-sdk": "5.25.0",
     "@types/aws-lambda": "^8.10.62",
     "@types/node": "^14.6.4",
+    "aws-sdk": "^2.765.0",
     "eslint": "7.6.0",
     "jest": "^24.7.1",
+    "nock": "^13.0.4",
     "npm-run-all": "^4.1.2",
     "prettier": "1.19.0",
     "rimraf": "^2.6.3",

packages/serverless/src/awsresources.ts

@@ -0,0 +1,105 @@
+import { getCurrentHub } from '@sentry/node';
+import { Integration, Span, Transaction } from '@sentry/types';
+import { fill } from '@sentry/utils';
+// 'aws-sdk/global' import is expected to be type-only so it's erased in the final .js file.
+// When TypeScript compiler is upgraded, use `import type` syntax to explicitly assert that we don't want to load a module here.
+import * as AWS from 'aws-sdk/global';
+
+type GenericParams = { [key: string]: any }; // eslint-disable-line @typescript-eslint/no-explicit-any
+type MakeRequestCallback<TResult> = (err: AWS.AWSError, data: TResult) => void;
+// This interace could be replaced with just type alias once the `strictBindCallApply` mode is enabled.
+interface MakeRequestFunction<TParams, TResult> extends CallableFunction {
+  (operation: string, params?: TParams, callback?: MakeRequestCallback<TResult>): AWS.Request<TResult, AWS.AWSError>;
+}
+interface AWSService {
+  serviceIdentifier: string;
+}
+
+/** AWS resource requests tracking */
+export class AWSResources implements Integration {
+  /**
+   * @inheritDoc
+   */
+  public static id: string = 'AWSResources';
+
+  /**
+   * @inheritDoc
+   */
+  public name: string = AWSResources.id;
+
+  /**
+   * @inheritDoc
+   */
+  public setupOnce(): void {
+    const awsModule = require('aws-sdk/global') as typeof AWS;
+    fill(
+      awsModule.Service.prototype,
+      'makeRequest',
+      <TService extends AWSService, TResult>(
+        orig: MakeRequestFunction<GenericParams, TResult>,
+      ): MakeRequestFunction<GenericParams, TResult> =>
+        function(this: TService, operation: string, params?: GenericParams, callback?: MakeRequestCallback<TResult>) {
+          let transaction: Transaction | undefined;
+          let span: Span | undefined;
+          const scope = getCurrentHub().getScope();
+          if (scope) {
+            transaction = scope.getTransaction();
+          }
+          const req = orig.call(this, operation, params);
+          req.on('afterBuild', () => {
+            if (transaction) {
+              span = transaction.startChild({
+                description: describe(this, operation, params),
+                op: 'request',
+              });
+            }
+          });
+          req.on('complete', () => {
+            if (span) {
+              span.finish();
+            }
+          });
+
+          if (callback) {
+            req.send(callback);
+          }
+          return req;
+        },
+    );
+  }
+}
+
+/** Describes an operation on generic AWS service */
+function describe<TService extends AWSService>(service: TService, operation: string, params?: GenericParams): string {
+  let ret = `aws.${service.serviceIdentifier}.${operation}`;
+  if (params === undefined) {
+    return ret;
+  }
+  switch (service.serviceIdentifier) {
+    case 's3':
+      ret += describeS3Operation(operation, params);
+      break;
+    case 'lambda':
+      ret += describeLambdaOperation(operation, params);
+      break;
+  }
+  return ret;
+}
+
+/** Describes an operation on AWS Lambda service */
+function describeLambdaOperation(_operation: string, params: GenericParams): string {
+  let ret = '';
+  if ('FunctionName' in params) {
+    ret += ` ${params.FunctionName}`;
+  }
+  return ret;
+}
+
+/** Describes an operation on AWS S3 service */
+function describeS3Operation(_operation: string, params: GenericParams): string {
+  let ret = '';
+  if ('Bucket' in params) {
+    ret += ` ${params.Bucket}`;
+  }
+  return ret;
+}

packages/serverless/src/index.ts

@@ -3,4 +3,5 @@ import * as AWSLambda from './awslambda';
 import * as GCPFunction from './gcpfunction';
 export { AWSLambda, GCPFunction };
 
+export * from './awsresources';
 export * from '@sentry/node';

packages/serverless/test/__mocks__/@sentry/node.ts

@@ -6,19 +6,25 @@ export const Severity = {
 };
 export const fakeParentScope = {
   setSpan: jest.fn(),
+  getTransaction: jest.fn(() => fakeTransaction),
 };
 export const fakeHub = {
   configureScope: jest.fn((fn: (arg: any) => any) => fn(fakeParentScope)),
+  getScope: jest.fn(() => fakeParentScope),
 };
 export const fakeScope = {
   addEventProcessor: jest.fn(),
   setTransactionName: jest.fn(),
   setTag: jest.fn(),
   setContext: jest.fn(),
 };
+export const fakeSpan = {
+  finish: jest.fn(),
+};
 export const fakeTransaction = {
   finish: jest.fn(),
   setHttpStatus: jest.fn(),
+  startChild: jest.fn(() => fakeSpan),
 };
 export const getCurrentHub = jest.fn(() => fakeHub);
 export const startTransaction = jest.fn(_ => fakeTransaction);
@@ -30,8 +36,12 @@ export const flush = jest.fn(() => Promise.resolve());
 export const resetMocks = (): void => {
   fakeTransaction.setHttpStatus.mockClear();
   fakeTransaction.finish.mockClear();
+  fakeTransaction.startChild.mockClear();
+  fakeSpan.finish.mockClear();
   fakeParentScope.setSpan.mockClear();
+  fakeParentScope.getTransaction.mockClear();
   fakeHub.configureScope.mockClear();
+  fakeHub.getScope.mockClear();
 
   fakeScope.addEventProcessor.mockClear();
   fakeScope.setTransactionName.mockClear();

packages/serverless/test/awsresources.test.ts

@@ -0,0 +1,69 @@
+import * as AWS from 'aws-sdk';
+import * as nock from 'nock';
+
+import * as Sentry from '../src';
+import { AWSResources } from '../src/awsresources';
+
+/**
+ * Why @ts-ignore some Sentry.X calls
+ *
+ * A hack-ish way to contain everything related to mocks in the same __mocks__ file.
+ * Thanks to this, we don't have to do more magic than necessary. Just add and export desired method and assert on it.
+ */
+
+describe('AWSResources', () => {
+  beforeAll(() => {
+    new AWSResources().setupOnce();
+  });
+  afterEach(() => {
+    // @ts-ignore see "Why @ts-ignore" note
+    Sentry.resetMocks();
+  });
+  afterAll(() => {
+    nock.restore();
+  });
+
+  describe('S3', () => {
+    const s3 = new AWS.S3({ accessKeyId: '-', secretAccessKey: '-' });
+
+    test('getObject', async () => {
+      nock('https://foo.s3.amazonaws.com')
+        .get('/bar')
+        .reply(200, 'contents');
+      const data = await s3.getObject({ Bucket: 'foo', Key: 'bar' }).promise();
+      expect(data.Body?.toString('utf-8')).toEqual('contents');
+      // @ts-ignore see "Why @ts-ignore" note
+      expect(Sentry.fakeTransaction.startChild).toBeCalledWith({ op: 'request', description: 'aws.s3.getObject foo' });
+      // @ts-ignore see "Why @ts-ignore" note
+      expect(Sentry.fakeSpan.finish).toBeCalled();
+    });
+
+    test('getObject with callback', done => {
+      expect.assertions(3);
+      nock('https://foo.s3.amazonaws.com')
+        .get('/bar')
+        .reply(200, 'contents');
+      s3.getObject({ Bucket: 'foo', Key: 'bar' }, (err, data) => {
+        expect(err).toBeNull();
+        expect(data.Body?.toString('utf-8')).toEqual('contents');
+        done();
+      });
+      // @ts-ignore see "Why @ts-ignore" note
+      expect(Sentry.fakeTransaction.startChild).toBeCalledWith({ op: 'request', description: 'aws.s3.getObject foo' });
+    });
+  });
+
+  describe('Lambda', () => {
+    const lambda = new AWS.Lambda({ accessKeyId: '-', secretAccessKey: '-', region: 'eu-north-1' });
+
+    test('invoke', async () => {
+      nock('https://lambda.eu-north-1.amazonaws.com')
+        .post('/2015-03-31/functions/foo/invocations')
+        .reply(201, 'reply');
+      const data = await lambda.invoke({ FunctionName: 'foo' }).promise();
+      expect(data.Payload?.toString('utf-8')).toEqual('reply');
+      // @ts-ignore see "Why @ts-ignore" note
+      expect(Sentry.fakeTransaction.startChild).toBeCalledWith({ op: 'request', description: 'aws.lambda.invoke foo' });
+    });
+  });
+});