ref(serverless): Convert AWSServices to function

Author: AbhiPrasad

packages/serverless/src/awsservices.ts

@@ -1,5 +1,6 @@
-import { startInactiveSpan } from '@sentry/node';
-import type { Integration, Span } from '@sentry/types';
+import { convertIntegrationFnToClass, defineIntegration } from '@sentry/core';
+import { getClient, startInactiveSpan } from '@sentry/node';
+import type { Client, Integration, IntegrationClass, IntegrationFn, Span } 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.
@@ -15,62 +16,68 @@ interface AWSService {
   serviceIdentifier: string;
 }
 
-/** AWS service requests tracking */
-export class AWSServices implements Integration {
-  /**
-   * @inheritDoc
-   */
-  public static id: string = 'AWSServices';
+const INTEGRATION_NAME = 'AWSServices';
 
-  /**
-   * @inheritDoc
-   */
-  public name: string;
+export const SETUP_CLIENTS = new WeakMap<Client, boolean>();
 
-  private readonly _optional: boolean;
-
-  public constructor(options: { optional?: boolean } = {}) {
-    this.name = AWSServices.id;
+const _awsServicesIntegration = ((options: { optional?: boolean } = {}) => {
+  const optional = options.optional || false;
+  return {
+    name: INTEGRATION_NAME,
+    setupOnce() {
+      try {
+        // eslint-disable-next-line @typescript-eslint/no-var-requires
+        const awsModule = require('aws-sdk/global') as typeof AWS;
+        fill(awsModule.Service.prototype, 'makeRequest', wrapMakeRequest);
+      } catch (e) {
+        if (!optional) {
+          throw e;
+        }
+      }
+    },
+    setup(client) {
+      SETUP_CLIENTS.set(client, true);
+    },
+  };
+}) satisfies IntegrationFn;
 
-    this._optional = options.optional || false;
-  }
+export const awsServicesIntegration = defineIntegration(_awsServicesIntegration);
 
-  /**
-   * @inheritDoc
-   */
-  public setupOnce(): void {
-    try {
-      // eslint-disable-next-line @typescript-eslint/no-var-requires
-      const awsModule = require('aws-sdk/global') as typeof AWS;
-      fill(awsModule.Service.prototype, 'makeRequest', wrapMakeRequest);
-    } catch (e) {
-      if (!this._optional) {
-        throw e;
-      }
-    }
-  }
-}
+/**
+ * AWS Service Request Tracking.
+ *
+ * @deprecated Use `awsServicesIntegration()` instead.
+ */
+// eslint-disable-next-line deprecation/deprecation
+export const AWSServices = convertIntegrationFnToClass(
+  INTEGRATION_NAME,
+  awsServicesIntegration,
+) as IntegrationClass<Integration>;
 
-/** */
+/**
+ * Patches AWS SDK request to create `http.client` spans.
+ */
 function wrapMakeRequest<TService extends AWSService, TResult>(
   orig: MakeRequestFunction<GenericParams, TResult>,
 ): MakeRequestFunction<GenericParams, TResult> {
   return function (this: TService, operation: string, params?: GenericParams, callback?: MakeRequestCallback<TResult>) {
     let span: Span | undefined;
     const req = orig.call(this, operation, params);
-    req.on('afterBuild', () => {
-      span = startInactiveSpan({
-        name: describe(this, operation, params),
-        op: 'http.client',
-        origin: 'auto.http.serverless',
+    // Only instrument if integration is active on client
+    if (SETUP_CLIENTS.has(getClient() as Client)) {
+      req.on('afterBuild', () => {
+        span = startInactiveSpan({
+          name: describe(this, operation, params),
+          op: 'http.client',
+          origin: 'auto.http.serverless',
+        });
       });
-    });
-    req.on('complete', () => {
-      if (span) {
-        span.end();
-      }
-    });
-
+      req.on('complete', () => {
+        if (span) {
+          span.end();
+        }
+      });
+    }
     if (callback) {
       req.send(callback);
     }

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

@@ -2,34 +2,72 @@ import * as SentryNode from '@sentry/node';
 import * as AWS from 'aws-sdk';
 import * as nock from 'nock';
 
-import { AWSServices } from '../src/awsservices';
+import { awsServicesIntegration } from '../src/awsservices';
 
-describe('AWSServices', () => {
-  beforeAll(() => {
-    new AWSServices().setupOnce();
+const mockSpanEnd = jest.fn();
+const mockStartInactiveSpan = jest.fn(spanArgs => ({ ...spanArgs }));
+
+jest.mock('@sentry/node', () => {
+  return {
+    ...jest.requireActual('@sentry/node'),
+    startInactiveSpan: (ctx: unknown) => {
+      mockStartInactiveSpan(ctx);
+      return { end: mockSpanEnd };
+    },
+  };
+});
+
+describe('awsServicesIntegration', () => {
+  const mockClient = new SentryNode.NodeClient({
+    tracesSampleRate: 1.0,
+    integrations: [],
+    dsn: 'https://withAWSServices@domain/123',
+    transport: () => SentryNode.createTransport({ recordDroppedEvent: () => undefined }, _ => Promise.resolve({})),
+    stackParser: () => [],
   });
-  afterEach(() => {
-    // @ts-expect-error see "Why @ts-expect-error" note
-    SentryNode.resetMocks();
+
+  const integration = awsServicesIntegration();
+  mockClient.addIntegration(integration);
+
+  const mockClientWithoutIntegration = new SentryNode.NodeClient({
+    tracesSampleRate: 1.0,
+    integrations: [],
+    dsn: 'https://withoutAWSServices@domain/123',
+    transport: () => SentryNode.createTransport({ recordDroppedEvent: () => undefined }, _ => Promise.resolve({})),
+    stackParser: () => [],
   });
+
   afterAll(() => {
     nock.restore();
   });
 
+  beforeEach(() => {
+    SentryNode.setCurrentClient(mockClient);
+    mockSpanEnd.mockClear();
+    mockStartInactiveSpan.mockClear();
+  });
+
   describe('S3 tracing', () => {
     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');
-      expect(SentryNode.startInactiveSpan).toBeCalledWith({
+      expect(mockStartInactiveSpan).toBeCalledWith({
         op: 'http.client',
         origin: 'auto.http.serverless',
         name: 'aws.s3.getObject foo',
       });
-      // @ts-expect-error see "Why @ts-expect-error" note
-      expect(SentryNode.fakeSpan.end).toBeCalled();
+
+      expect(mockSpanEnd).toHaveBeenCalledTimes(1);
+    });
+
+    test('getObject with integration-less client', async () => {
+      SentryNode.setCurrentClient(mockClientWithoutIntegration);
+      nock('https://foo.s3.amazonaws.com').get('/bar').reply(200, 'contents');
+      await s3.getObject({ Bucket: 'foo', Key: 'bar' }).promise();
+      expect(mockStartInactiveSpan).not.toBeCalled();
     });
 
     test('getObject with callback', done => {
@@ -40,12 +78,22 @@ describe('AWSServices', () => {
         expect(data.Body?.toString('utf-8')).toEqual('contents');
         done();
       });
-      expect(SentryNode.startInactiveSpan).toBeCalledWith({
+      expect(mockStartInactiveSpan).toBeCalledWith({
         op: 'http.client',
         origin: 'auto.http.serverless',
         name: 'aws.s3.getObject foo',
       });
     });
+
+    test('getObject with callback with integration-less client', done => {
+      SentryNode.setCurrentClient(mockClientWithoutIntegration);
+      expect.assertions(1);
+      nock('https://foo.s3.amazonaws.com').get('/bar').reply(200, 'contents');
+      s3.getObject({ Bucket: 'foo', Key: 'bar' }, () => {
+        done();
+      });
+      expect(mockStartInactiveSpan).not.toBeCalled();
+    });
   });
 
   describe('Lambda', () => {
@@ -55,11 +103,19 @@ describe('AWSServices', () => {
       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');
-      expect(SentryNode.startInactiveSpan).toBeCalledWith({
+      expect(mockStartInactiveSpan).toBeCalledWith({
         op: 'http.client',
         origin: 'auto.http.serverless',
         name: 'aws.lambda.invoke foo',
       });
+      expect(mockSpanEnd).toHaveBeenCalledTimes(1);
+    });
+
+    test('invoke with integration-less client', async () => {
+      SentryNode.setCurrentClient(mockClientWithoutIntegration);
+      nock('https://lambda.eu-north-1.amazonaws.com').post('/2015-03-31/functions/foo/invocations').reply(201, 'reply');
+      await lambda.invoke({ FunctionName: 'foo' }).promise();
+      expect(mockStartInactiveSpan).not.toBeCalled();
     });
   });
 });