Feature/cls-interceptor

package.json

@@ -16,6 +16,7 @@
         "cls",
         "continuation-local-storage",
         "als",
+        "AsyncLocalStorage",
         "async_hooks",
         "request context"
     ],

src/index.ts

@@ -1,6 +1,7 @@
 export * from './lib/cls-service-manager';
 export * from './lib/cls.constants';
 export * from './lib/cls.middleware';
+export * from './lib/cls.interceptor';
 export * from './lib/cls.module';
 export * from './lib/cls.service';
 export * from './lib/cls.decorators';

src/lib/cls.constants.ts

@@ -4,3 +4,4 @@ export const CLS_ID = 'CLS_ID';
 export const CLS_DEFAULT_NAMESPACE = 'CLS_DEFAULT_NAMESPACE';
 export const CLS_MIDDLEWARE_OPTIONS = 'ClsMiddlewareOptions';
 export const CLS_GUARD_OPTIONS = 'ClsGuardOptions';
+export const CLS_INTERCEPTOR_OPTIONS = 'ClsInterceptorOptions';

src/lib/cls.interceptor.ts

@@ -0,0 +1,44 @@
+import {
+    CallHandler,
+    ExecutionContext,
+    Inject,
+    Injectable,
+    NestInterceptor,
+} from '@nestjs/common';
+import { Observable } from 'rxjs';
+import { CLS_ID } from '..';
+import { ClsServiceManager } from './cls-service-manager';
+import { CLS_INTERCEPTOR_OPTIONS } from './cls.constants';
+import { ClsInterceptorOptions } from './cls.interfaces';
+
+@Injectable()
+export class ClsInterceptor implements NestInterceptor {
+    constructor(
+        @Inject(CLS_INTERCEPTOR_OPTIONS)
+        private readonly options?: Omit<ClsInterceptorOptions, 'mount'>,
+    ) {
+        this.options = { ...new ClsInterceptorOptions(), ...options };
+    }
+
+    intercept(context: ExecutionContext, next: CallHandler): Observable<any> {
+        const cls = ClsServiceManager.getClsService(this.options.namespaceName);
+        return new Observable((subscriber) => {
+            cls.run(async () => {
+                if (this.options.generateId) {
+                    const id = await this.options.idGenerator(context);
+                    cls.set(CLS_ID, id);
+                }
+                if (this.options.setup) {
+                    await this.options.setup(cls, context);
+                }
+                next.handle()
+                    .pipe()
+                    .subscribe({
+                        next: (res) => subscriber.next(res),
+                        error: (err) => subscriber.error(err),
+                        complete: () => subscriber.complete(),
+                    });
+            });
+        });
+    }
+}

src/lib/cls.interfaces.ts

@@ -25,6 +25,11 @@ export class ClsModuleOptions {
      * Cls guard options
      */
     guard?: ClsGuardOptions = null;
+
+    /**
+     * Cls interceptor options
+     */
+    interceptor?: ClsInterceptorOptions = null;
 }
 
 export class ClsMiddlewareOptions {
@@ -103,3 +108,32 @@ export class ClsGuardOptions {
 
     readonly namespaceName?: string;
 }
+
+export class ClsInterceptorOptions {
+    /**
+     * whether to mount the interceptor globally
+     */
+    mount?: boolean; // default false
+
+    /**
+     * whether to automatically generate request ids
+     */
+    generateId?: boolean; // default false
+
+    /**
+     * the function to generate request ids inside the interceptor
+     */
+    idGenerator?: (context: ExecutionContext) => string | Promise<string> =
+        () => Math.random().toString(36).slice(-8);
+
+    /**
+     * Function that executes after the CLS context has been initialised.
+     * It can be used to put additional variables in the CLS context.
+     */
+    setup?: (
+        cls: ClsService,
+        context: ExecutionContext,
+    ) => void | Promise<void>;
+
+    readonly namespaceName?: string;
+}

src/lib/cls.module.ts

@@ -6,12 +6,23 @@ import {
     NestModule,
     Provider,
 } from '@nestjs/common';
-import { APP_GUARD, HttpAdapterHost, ModuleRef } from '@nestjs/core';
+import {
+    APP_GUARD,
+    APP_INTERCEPTOR,
+    HttpAdapterHost,
+    ModuleRef,
+} from '@nestjs/core';
+import { ClsInterceptor } from '..';
 import { ClsServiceManager, getClsServiceToken } from './cls-service-manager';
-import { CLS_GUARD_OPTIONS, CLS_MIDDLEWARE_OPTIONS } from './cls.constants';
+import {
+    CLS_GUARD_OPTIONS,
+    CLS_INTERCEPTOR_OPTIONS,
+    CLS_MIDDLEWARE_OPTIONS,
+} from './cls.constants';
 import { ClsGuard } from './cls.guard';
 import {
     ClsGuardOptions,
+    ClsInterceptorOptions,
     ClsMiddlewareOptions,
     ClsModuleOptions,
 } from './cls.interfaces';
@@ -80,6 +91,11 @@ export class ClsModule implements NestModule {
             ...options.guard,
             namespaceName: options.namespaceName,
         };
+        const clsInterceptorOptions = {
+            ...new ClsInterceptorOptions(),
+            ...options.interceptor,
+            namespaceName: options.namespaceName,
+        };
         const providers: Provider[] = [
             ...ClsServiceManager.getClsServicesAsProviders(),
             {
@@ -90,18 +106,28 @@ export class ClsModule implements NestModule {
                 provide: CLS_GUARD_OPTIONS,
                 useValue: clsGuardOptions,
             },
+            {
+                provide: CLS_INTERCEPTOR_OPTIONS,
+                useValue: clsInterceptorOptions,
+            },
         ];
-        const guardArr = [];
+        const enhancerArr = [];
         if (clsGuardOptions.mount) {
-            guardArr.push({
+            enhancerArr.push({
                 provide: APP_GUARD,
                 useClass: ClsGuard,
             });
         }
+        if (clsInterceptorOptions.mount) {
+            enhancerArr.push({
+                provide: APP_INTERCEPTOR,
+                useClass: ClsInterceptor,
+            });
+        }
 
         return {
             module: ClsModule,
-            providers: providers.concat(...guardArr),
+            providers: providers.concat(...enhancerArr),
             exports: providers,
             global: options.global,
         };

src/lib/cls.service.spec.ts

@@ -1,5 +1,6 @@
 import { Test, TestingModule } from '@nestjs/testing';
-import { ClsServiceManager, CLS_DEFAULT_NAMESPACE } from '..';
+import { ClsServiceManager } from './cls-service-manager';
+import { CLS_DEFAULT_NAMESPACE } from './cls.constants';
 import { ClsService } from './cls.service';
 
 describe('ClsService', () => {

src/lib/cls.service.ts

@@ -50,7 +50,7 @@ export class ClsService<S = Record<string, any>> {
     }
 
     /**
-     * Run the callback in a shared CLS context.
+     * Run the callback with a shared CLS context.
      * @param callback function to run
      * @returns whatever the callback returns
      */
@@ -59,12 +59,30 @@ export class ClsService<S = Record<string, any>> {
     }
 
     /**
-     * Run any following code in a shared CLS context.
+     * Run the callbacks with a shared CLS context.
+     * @param store the default context contents
+     * @param callback function to run
+     * @returns whatever the callback returns
+     */
+    runWith<T = any>(store: any, callback: () => T) {
+        return this.namespace.run(store ?? {}, callback);
+    }
+
+    /**
+     * Run any following code with a shared CLS context.
      */
     enter() {
         return this.namespace.enterWith({});
     }
 
+    /**
+     * Run any following code with a shared ClS context
+     * @param store the default context contents
+     */
+    enterWith(store: any = {}) {
+        return this.namespace.enterWith(store);
+    }
+
     /**
      * Run the callback outside of a shared CLS context
      * @param callback function to run

test/common/test.exception.ts

@@ -0,0 +1,8 @@
+export class TestException extends Error {
+    public response: any;
+    public extensions: { a: 1 };
+    constructor(response: any) {
+        super('TestException');
+        this.response = response;
+    }
+}

test/common/test.guard.ts

@@ -9,7 +9,7 @@ export class TestGuard implements CanActivate {
     canActivate(
         context: ExecutionContext,
     ): boolean | Promise<boolean> | Observable<boolean> {
-        this.cls.set('FROM_GUARD', this.cls.getId());
-        return this.cls.isActive();
+        if (this.cls.isActive()) this.cls.set('FROM_GUARD', this.cls.getId());
+        return true;
     }
 }

test/gql/expect-ids-gql.ts

@@ -1,7 +1,10 @@
 import { INestApplication } from '@nestjs/common';
 import request from 'supertest';
 
-export const expectIdsGql = (app: INestApplication) =>
+export const expectOkIdsGql = (
+    app: INestApplication,
+    options = { skipGuard: false },
+) =>
     request(app.getHttpServer())
         .post('/graphql')
         .send({
@@ -19,8 +22,36 @@ export const expectIdsGql = (app: INestApplication) =>
         .expect(200)
         .then((r) => {
             const body = r.body.data?.items[0];
-            const id = body.id;
-            expect(body.fromGuard).toEqual(id);
+            const id = body.id ?? 'no-id';
+            if (!options.skipGuard) expect(body.fromGuard).toEqual(id);
+            expect(body.fromInterceptor).toEqual(id);
+            expect(body.fromInterceptorAfter).toEqual(id);
+            expect(body.fromResolver).toEqual(id);
+            expect(body.fromService).toEqual(id);
+        });
+
+export const expectErrorIdsGql = (
+    app: INestApplication,
+    options = { skipGuard: false },
+) =>
+    request(app.getHttpServer())
+        .post('/graphql')
+        .send({
+            query: `query {
+                error {
+                    id
+                    fromGuard
+                    fromInterceptor
+                    fromInterceptorAfter
+                    fromResolver
+                    fromService
+                }
+            }`,
+        })
+        .then((r) => {
+            const body = r.body.errors?.[0].extensions.exception?.response;
+            const id = body.id ?? 'no-id';
+            if (!options.skipGuard) expect(body.fromGuard).toEqual(id);
             expect(body.fromInterceptor).toEqual(id);
             expect(body.fromInterceptorAfter).toEqual(id);
             expect(body.fromResolver).toEqual(id);