ref(serverless): Convert GoogleCloudGrpc to function

Author: AbhiPrasad

packages/serverless/src/awsservices.ts

@@ -18,7 +18,7 @@ interface AWSService {
 
 const INTEGRATION_NAME = 'AWSServices';
 
-export const SETUP_CLIENTS = new WeakMap<Client, boolean>();
+const SETUP_CLIENTS = new WeakMap<Client, boolean>();
 
 const _awsServicesIntegration = ((options: { optional?: boolean } = {}) => {
   const optional = options.optional || false;

packages/serverless/src/gcpfunction/index.ts

@@ -7,7 +7,7 @@ import {
 } from '@sentry/node';
 import type { Integration, Options, SdkMetadata } from '@sentry/types';
 
-import { GoogleCloudGrpc } from '../google-cloud-grpc';
+import { googleCloudGrpcIntegration } from '../google-cloud-grpc';
 import { GoogleCloudHttp } from '../google-cloud-http';
 
 export * from './http';
@@ -19,15 +19,15 @@ export const defaultIntegrations: Integration[] = [
   // eslint-disable-next-line deprecation/deprecation
   ...defaultNodeIntegrations,
   new GoogleCloudHttp({ optional: true }), // We mark this integration optional since '@google-cloud/common' module could be missing.
-  new GoogleCloudGrpc({ optional: true }), // We mark this integration optional since 'google-gax' module could be missing.
+  googleCloudGrpcIntegration({ optional: true }), // We mark this integration optional since 'google-gax' module could be missing.
 ];
 
 /** Get the default integrations for the GCP SDK. */
 export function getDefaultIntegrations(options: Options): Integration[] {
   return [
     ...getDefaultNodeIntegrations(options),
     new GoogleCloudHttp({ optional: true }), // We mark this integration optional since '@google-cloud/common' module could be missing.
-    new GoogleCloudGrpc({ optional: true }), // We mark this integration optional since 'google-gax' module could be missing.
+    googleCloudGrpcIntegration({ optional: true }), // We mark this integration optional since 'google-gax' module could be missing.
   ];
 }
 

packages/serverless/src/google-cloud-grpc.ts

@@ -1,6 +1,7 @@
 import type { EventEmitter } from 'events';
+import { convertIntegrationFnToClass, defineIntegration, getClient } from '@sentry/core';
 import { startInactiveSpan } from '@sentry/node';
-import type { Integration } from '@sentry/types';
+import type { Client, Integration, IntegrationClass, IntegrationFn } from '@sentry/types';
 import { fill } from '@sentry/utils';
 
 interface GrpcFunction extends CallableFunction {
@@ -25,45 +26,49 @@ interface Stub {
   [key: string]: GrpcFunctionObject;
 }
 
-/** Google Cloud Platform service requests tracking for GRPC APIs */
-export class GoogleCloudGrpc implements Integration {
-  /**
-   * @inheritDoc
-   */
-  public static id: string = 'GoogleCloudGrpc';
+const SERVICE_PATH_REGEX = /^(\w+)\.googleapis.com$/;
 
-  /**
-   * @inheritDoc
-   */
-  public name: string;
+const INTEGRATION_NAME = 'GoogleCloudGrpc';
 
-  private readonly _optional: boolean;
+const SETUP_CLIENTS = new WeakMap<Client, boolean>();
 
-  public constructor(options: { optional?: boolean } = {}) {
-    this.name = GoogleCloudGrpc.id;
+const _googleCloudGrpcIntegration = ((options: { optional?: boolean } = {}) => {
+  const optional = options.optional || false;
+  return {
+    name: INTEGRATION_NAME,
+    setupOnce() {
+      try {
+        // eslint-disable-next-line @typescript-eslint/no-var-requires
+        const gaxModule = require('google-gax');
+        fill(
+          gaxModule.GrpcClient.prototype, // eslint-disable-line @typescript-eslint/no-unsafe-member-access
+          'createStub',
+          wrapCreateStub,
+        );
+      } catch (e) {
+        if (!optional) {
+          throw e;
+        }
+      }
+    },
+    setup(client) {
+      SETUP_CLIENTS.set(client, true);
+    },
+  };
+}) satisfies IntegrationFn;
 
-    this._optional = options.optional || false;
-  }
+export const googleCloudGrpcIntegration = defineIntegration(_googleCloudGrpcIntegration);
 
-  /**
-   * @inheritDoc
-   */
-  public setupOnce(): void {
-    try {
-      // eslint-disable-next-line @typescript-eslint/no-var-requires
-      const gaxModule = require('google-gax');
-      fill(
-        gaxModule.GrpcClient.prototype, // eslint-disable-line @typescript-eslint/no-unsafe-member-access
-        'createStub',
-        wrapCreateStub,
-      );
-    } catch (e) {
-      if (!this._optional) {
-        throw e;
-      }
-    }
-  }
-}
+/**
+ * Google Cloud Platform service requests tracking for GRPC APIs.
+ *
+ * @deprecated Use `googleCloudGrpcIntegration()` instead.
+ */
+// eslint-disable-next-line deprecation/deprecation
+export const GoogleCloudGrpc = convertIntegrationFnToClass(
+  INTEGRATION_NAME,
+  googleCloudGrpcIntegration,
+) as IntegrationClass<Integration>;
 
 /** Returns a wrapped function that returns a stub with tracing enabled */
 function wrapCreateStub(origCreate: CreateStubFunc): CreateStubFunc {
@@ -107,23 +112,26 @@ function fillGrpcFunction(stub: Stub, serviceIdentifier: string, methodName: str
         if (typeof ret?.on !== 'function') {
           return ret;
         }
-        const span = startInactiveSpan({
-          name: `${callType} ${methodName}`,
-          op: `grpc.${serviceIdentifier}`,
-          origin: 'auto.grpc.serverless',
-        });
-        ret.on('status', () => {
-          if (span) {
-            span.end();
-          }
-        });
+        // only instrument if integration is active on client
+        if (SETUP_CLIENTS.has(getClient() as Client)) {
+          const span = startInactiveSpan({
+            name: `${callType} ${methodName}`,
+            op: `grpc.${serviceIdentifier}`,
+            origin: 'auto.grpc.serverless',
+          });
+          ret.on('status', () => {
+            if (span) {
+              span.end();
+            }
+          });
+        }
         return ret;
       },
   );
 }
 
 /** Identifies service by its address */
 function identifyService(servicePath: string): string {
-  const match = servicePath.match(/^(\w+)\.googleapis.com$/);
+  const match = servicePath.match(SERVICE_PATH_REGEX);
   return match ? match[1] : servicePath;
 }

packages/serverless/test/awsservices.test.ts

@@ -1,4 +1,4 @@
-import * as SentryNode from '@sentry/node';
+import { NodeClient, createTransport, setCurrentClient } from '@sentry/node';
 import * as AWS from 'aws-sdk';
 import * as nock from 'nock';
 
@@ -18,22 +18,22 @@ jest.mock('@sentry/node', () => {
 });
 
 describe('awsServicesIntegration', () => {
-  const mockClient = new SentryNode.NodeClient({
+  const mockClient = new NodeClient({
     tracesSampleRate: 1.0,
     integrations: [],
     dsn: 'https://withAWSServices@domain/123',
-    transport: () => SentryNode.createTransport({ recordDroppedEvent: () => undefined }, _ => Promise.resolve({})),
+    transport: () => createTransport({ recordDroppedEvent: () => undefined }, _ => Promise.resolve({})),
     stackParser: () => [],
   });
 
   const integration = awsServicesIntegration();
   mockClient.addIntegration(integration);
 
-  const mockClientWithoutIntegration = new SentryNode.NodeClient({
+  const mockClientWithoutIntegration = new NodeClient({
     tracesSampleRate: 1.0,
     integrations: [],
     dsn: 'https://withoutAWSServices@domain/123',
-    transport: () => SentryNode.createTransport({ recordDroppedEvent: () => undefined }, _ => Promise.resolve({})),
+    transport: () => createTransport({ recordDroppedEvent: () => undefined }, _ => Promise.resolve({})),
     stackParser: () => [],
   });
 
@@ -42,7 +42,7 @@ describe('awsServicesIntegration', () => {
   });
 
   beforeEach(() => {
-    SentryNode.setCurrentClient(mockClient);
+    setCurrentClient(mockClient);
     mockSpanEnd.mockClear();
     mockStartInactiveSpan.mockClear();
   });
@@ -64,7 +64,7 @@ describe('awsServicesIntegration', () => {
     });
 
     test('getObject with integration-less client', async () => {
-      SentryNode.setCurrentClient(mockClientWithoutIntegration);
+      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();
@@ -86,7 +86,7 @@ describe('awsServicesIntegration', () => {
     });
 
     test('getObject with callback with integration-less client', done => {
-      SentryNode.setCurrentClient(mockClientWithoutIntegration);
+      setCurrentClient(mockClientWithoutIntegration);
       expect.assertions(1);
       nock('https://foo.s3.amazonaws.com').get('/bar').reply(200, 'contents');
       s3.getObject({ Bucket: 'foo', Key: 'bar' }, () => {
@@ -112,7 +112,7 @@ describe('awsServicesIntegration', () => {
     });
 
     test('invoke with integration-less client', async () => {
-      SentryNode.setCurrentClient(mockClientWithoutIntegration);
+      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();

packages/serverless/test/google-cloud-grpc.test.ts

@@ -5,14 +5,27 @@ import { EventEmitter } from 'events';
 import * as fs from 'fs';
 import * as path from 'path';
 import { PubSub } from '@google-cloud/pubsub';
-import * as SentryNode from '@sentry/node';
 import * as http2 from 'http2';
 import * as nock from 'nock';
 
-import { GoogleCloudGrpc } from '../src/google-cloud-grpc';
+import { NodeClient, createTransport, setCurrentClient } from '@sentry/node';
+import { googleCloudGrpcIntegration } from '../src/google-cloud-grpc';
 
 const spyConnect = jest.spyOn(http2, 'connect');
 
+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 };
+    },
+  };
+});
+
 /** Fake HTTP2 stream */
 class FakeStream extends EventEmitter {
   public rstCode: number = 0;
@@ -70,18 +83,24 @@ function mockHttp2Session(): FakeSession {
 }
 
 describe('GoogleCloudGrpc tracing', () => {
-  beforeAll(() => {
-    new GoogleCloudGrpc().setupOnce();
+  const mockClient = new NodeClient({
+    tracesSampleRate: 1.0,
+    integrations: [],
+    dsn: 'https://withAWSServices@domain/123',
+    transport: () => createTransport({ recordDroppedEvent: () => undefined }, _ => Promise.resolve({})),
+    stackParser: () => [],
   });
 
+  const integration = googleCloudGrpcIntegration();
+  mockClient.addIntegration(integration);
+
   beforeEach(() => {
     nock('https://www.googleapis.com').post('/oauth2/v4/token').reply(200, []);
+    setCurrentClient(mockClient);
+    mockSpanEnd.mockClear();
+    mockStartInactiveSpan.mockClear();
   });
-  afterEach(() => {
-    // @ts-expect-error see "Why @ts-expect-error" note
-    SentryNode.resetMocks();
-    spyConnect.mockClear();
-  });
+
   afterAll(() => {
     nock.restore();
     spyConnect.mockRestore();
@@ -115,16 +134,19 @@ describe('GoogleCloudGrpc tracing', () => {
       resolveTxt.mockReset();
     });
 
+    afterAll(async () => {
+      await pubsub.close();
+    });
+
     test('publish', async () => {
       mockHttp2Session().mockUnaryRequest(Buffer.from('00000000120a1031363337303834313536363233383630', 'hex'));
       const resp = await pubsub.topic('nicetopic').publish(Buffer.from('data'));
       expect(resp).toEqual('1637084156623860');
-      expect(SentryNode.startInactiveSpan).toBeCalledWith({
+      expect(mockStartInactiveSpan).toBeCalledWith({
         op: 'grpc.pubsub',
         origin: 'auto.grpc.serverless',
         name: 'unary call publish',
       });
-      await pubsub.close();
     });
   });
 });