feat(tracing): Add Prisma ORM integration. (#4931)

Co-authored-by: Abhijeet Prasad <[email protected]>

Author: onurtemizkan

packages/tracing/src/integrations/index.ts

@@ -2,6 +2,7 @@ export { Express } from './node/express';
 export { Postgres } from './node/postgres';
 export { Mysql } from './node/mysql';
 export { Mongo } from './node/mongo';
+export { Prisma } from './node/prisma';
 
 // TODO(v7): Remove this export
 // Please see `src/index.ts` for more details.

packages/tracing/src/integrations/node/prisma.ts

@@ -0,0 +1,99 @@
+import { Hub } from '@sentry/hub';
+import { EventProcessor, Integration } from '@sentry/types';
+import { isThenable, logger } from '@sentry/utils';
+
+import { IS_DEBUG_BUILD } from '../../flags';
+
+type PrismaAction =
+  | 'findUnique'
+  | 'findMany'
+  | 'findFirst'
+  | 'create'
+  | 'createMany'
+  | 'update'
+  | 'updateMany'
+  | 'upsert'
+  | 'delete'
+  | 'deleteMany'
+  | 'executeRaw'
+  | 'queryRaw'
+  | 'aggregate'
+  | 'count'
+  | 'runCommandRaw';
+
+interface PrismaMiddlewareParams {
+  model?: unknown;
+  action: PrismaAction;
+  args: unknown;
+  dataPath: string[];
+  runInTransaction: boolean;
+}
+
+type PrismaMiddleware<T = unknown> = (
+  params: PrismaMiddlewareParams,
+  next: (params: PrismaMiddlewareParams) => Promise<T>,
+) => Promise<T>;
+
+interface PrismaClient {
+  $use: (cb: PrismaMiddleware) => void;
+}
+
+/** Tracing integration for @prisma/client package */
+export class Prisma implements Integration {
+  /**
+   * @inheritDoc
+   */
+  public static id: string = 'Prisma';
+
+  /**
+   * @inheritDoc
+   */
+  public name: string = Prisma.id;
+
+  /**
+   * Prisma ORM Client Instance
+   */
+  private readonly _client?: PrismaClient;
+
+  /**
+   * @inheritDoc
+   */
+  public constructor(options: { client?: PrismaClient } = {}) {
+    this._client = options.client;
+  }
+
+  /**
+   * @inheritDoc
+   */
+  public setupOnce(_: (callback: EventProcessor) => void, getCurrentHub: () => Hub): void {
+    if (!this._client) {
+      IS_DEBUG_BUILD && logger.error('PrismaIntegration is missing a Prisma Client Instance');
+      return;
+    }
+
+    this._client.$use((params: PrismaMiddlewareParams, next: (params: PrismaMiddlewareParams) => Promise<unknown>) => {
+      const scope = getCurrentHub().getScope();
+      const parentSpan = scope?.getSpan();
+
+      const action = params.action;
+      const model = params.model;
+
+      const span = parentSpan?.startChild({
+        description: model ? `${model} ${action}` : action,
+        op: 'db.prisma',
+      });
+
+      const rv = next(params);
+
+      if (isThenable(rv)) {
+        return rv.then((res: unknown) => {
+          span?.finish();
+          return res;
+        });
+      }
+
+      span?.finish();
+      return rv;
+    });
+  }
+}

packages/tracing/test/integrations/node/prisma.test.ts

@@ -0,0 +1,61 @@
+/* eslint-disable @typescript-eslint/unbound-method */
+import { Hub, Scope } from '@sentry/hub';
+
+import { Prisma } from '../../../src/integrations/node/prisma';
+import { Span } from '../../../src/span';
+
+type PrismaMiddleware = (params: unknown, next: (params?: unknown) => Promise<unknown>) => Promise<unknown>;
+
+class PrismaClient {
+  public user: { create: () => Promise<unknown> | undefined } = {
+    create: () => this._middleware?.({ action: 'create', model: 'user' }, () => Promise.resolve('result')),
+  };
+
+  private _middleware?: PrismaMiddleware;
+
+  constructor() {
+    this._middleware = undefined;
+  }
+
+  public $use(cb: PrismaMiddleware) {
+    this._middleware = cb;
+  }
+}
+
+describe('setupOnce', function () {
+  const Client: PrismaClient = new PrismaClient();
+
+  let scope = new Scope();
+  let parentSpan: Span;
+  let childSpan: Span;
+
+  beforeAll(() => {
+    // @ts-ignore, not to export PrismaClient types from integration source
+    new Prisma({ client: Client }).setupOnce(
+      () => undefined,
+      () => new Hub(undefined, scope),
+    );
+  });
+
+  beforeEach(() => {
+    scope = new Scope();
+    parentSpan = new Span();
+    childSpan = parentSpan.startChild();
+    jest.spyOn(scope, 'getSpan').mockReturnValueOnce(parentSpan);
+    jest.spyOn(parentSpan, 'startChild').mockReturnValueOnce(childSpan);
+    jest.spyOn(childSpan, 'finish');
+  });
+
+  it('should add middleware with $use method correctly', done => {
+    void Client.user.create()?.then(res => {
+      expect(res).toBe('result');
+      expect(scope.getSpan).toBeCalled();
+      expect(parentSpan.startChild).toBeCalledWith({
+        description: 'user create',
+        op: 'db.prisma',
+      });
+      expect(childSpan.finish).toBeCalled();
+      done();
+    });
+  });
+});