feat: add UseCls decorator (#56)

README.md

@@ -32,6 +32,7 @@ You might also be interested in [The Author's Take](#the-authors-take) on the to
     -   [Using a Middleware](#using-a-middleware-http-only)
     -   [Using a Guard](#using-a-guard)
     -   [Using an Interceptor](#using-an-interceptor)
+    -   [Using an Decorator](#using-a-decorator)
 -   [Features and use cases](#features-and-use-cases)
     -   [Request ID](#request-id)
     -   [Additional CLS Setup](#additional-cls-setup)
@@ -244,6 +245,12 @@ Or mount it manually as `APP_INTERCEPTOR`, or directly on the Controller/Resolve
 
 > **Please note**: Since Nest's _Interceptors_ run after _Guards_, that means using this method makes CLS **unavailable in Guards** (and in case of REST Controllers, also in **Exception Filters**).
 
+## Using a Decorator
+
+The `@UseCls()` decorator can be used at a method level to declaratively wrap the method with a `cls.run()` call. This method should only be used [outside of the context of a web request](#usage-outside-of-web-request).
+
+> Note: Please keep in mind, that since the CLS context initialization _can_ be async, the `@UseCls()` decorator can _only_ be used on _async_ function (or those that return a `Promise`).
+
 # Features and use cases
 
 In addition to the basic functionality described in the [Quick start](#quick-start) chapter, this module provides several other features.
@@ -321,7 +328,7 @@ ClsModule.forRoot({
 
 Sometimes, a part of the app that relies on the CLS storage might need to be called outside of the context of a web request - for example, in a Cron job, while consuming a Queue or during the application bootstrap. In such cases, there are no enhancers that can be bound to the handler to set up the context.
 
-Therefore, you as the the developer are responsible for wrapping the execution with `ClsService#run` and set up the appropriate context variables.
+Therefore, you as the the developer are responsible for wrapping the execution with `ClsService#run`, or using the `@UseCls` decorator. In any case, if any following code depends on some context variables, these need to be set up manually.
 
 ```ts
 @Injectable()
@@ -332,26 +339,25 @@ export class CronController {
     );
 
     @Cron('45 * * * * *')
-    handleCronExample1() {
-        this.clsService.run(() => {
-            // either set up all context variables inside the wrapped `run` call
-            this.cls.set(CLS_ID, uuid());
+    async handleCronExample1() {
+        // either explicitly wrap the function body with
+        // a call to `ClsService#run` ...
+        await this.cls.run(async () => {
             this.cls.set('mode', 'cron');
-            this.someService.doTheThing();
+            await this.someService.doTheThing();
         });
     }
 
     @Cron('90 * * * * *')
-    handleCronExample2() {
-        // or create the context object beforehand...
-        const context = {
-            [CLS_ID]: uuid(),
-            mode: 'cron',
-        };
-        // ...and pass it to the `runWith` call
-        this.clsService.runWith(context, () => {
-            this.someService.doTheThing();
-        });
+    // ... or use the convenience decorator which
+    // does the wrapping for you seamlessly.
+    @UseCls({
+        setup: (cls) => {
+            cls.set('mode', 'cron');
+        },
+    })
+    async handleCronExample2() {
+        await this.someService.doTheThing();
     }
 }
 ```

src/index.ts

@@ -7,4 +7,5 @@ export * from './lib/cls.module';
 export * from './lib/cls.service';
 export * from './lib/cls.decorators';
 export * from './lib/cls.options';
+export * from './lib/use-cls.decorator';
 export { Terminal } from './types/terminal.type';

src/lib/cls.decorators.ts

@@ -10,7 +10,7 @@ export function InjectCls() {
 }
 
 /**
- * Mark a Proxy provider with this decorator to distinguis it from regular NestJS singleton providers
+ * Mark a Proxy provider with this decorator to distinguish it from regular NestJS singleton providers
  */
 export function InjectableProxy() {
     return (target: any) =>

src/lib/cls.options.ts

@@ -176,8 +176,28 @@ export class ClsInterceptorOptions {
 }
 
 export class ClsDecoratorOptions<T extends any[]> {
+    /**
+     * Whether to automatically generate request ids
+     */
     generateId?: boolean; // default false
+
+    /**
+     * The function to generate request ids inside the interceptor.
+     *
+     * Takes the same parameters in the same order as the decorated function.
+     *
+     * Note: To avoid type errors, you must list all parameters, even if they're not used,
+     * or type the decorator as `@UseCls<[arg1: Type1, arg2: Type2]>()`
+     */
     idGenerator?: (...args: T) => string | Promise<string> = getRandomString;
+
+    /**
+     * Function that executes after the CLS context has been initialised.
+     * Takes ClsService as the first parameter and then the same parameters in the same order as the decorated function.
+     *
+     * Note: To avoid type errors, you must list all parameters, even if they're not used,
+     * or type the decorator as `@UseCls<[arg1: Type1, arg2: Type2]>()`
+     */
     setup?: (cls: ClsService, ...args: T) => void | Promise<void>;
 
     /**

src/lib/use-cls.decorator.spec.ts

@@ -0,0 +1,73 @@
+import { Injectable } from '@nestjs/common';
+import { Test } from '@nestjs/testing';
+import { CLS_ID } from './cls.constants';
+import { ClsModule } from './cls.module';
+import { ClsService } from './cls.service';
+import { UseCls } from './use-cls.decorator';
+
+@Injectable()
+class TestClass {
+    constructor(private readonly cls: ClsService) {}
+
+    @UseCls()
+    async startContext(value: string) {
+        this.cls.set(CLS_ID, this.generateId());
+        this.cls.set('value', value);
+        return this.useContextVariables();
+    }
+
+    @UseCls<[string]>({
+        generateId: true,
+        idGenerator: function (this: TestClass) {
+            return this.generateId();
+        },
+        setup: (cls, value: string) => {
+            cls.set('value', value);
+        },
+    })
+    // eslint-disable-next-line @typescript-eslint/no-unused-vars
+    async startContextWithIdAndSetup(value: string) {
+        return this.useContextVariables();
+    }
+
+    private generateId() {
+        return 'the-id';
+    }
+
+    private useContextVariables() {
+        return {
+            id: this.cls.getId(),
+            value: this.cls.get('value'),
+        };
+    }
+}
+
+describe('@UseCls', () => {
+    let testClass: TestClass;
+
+    beforeEach(async () => {
+        const module = await Test.createTestingModule({
+            imports: [ClsModule],
+            providers: [TestClass],
+        }).compile();
+
+        testClass = module.get(TestClass);
+    });
+
+    it('wraps function with context', async () => {
+        const result = await testClass.startContext('something');
+        expect(result).toEqual({
+            id: 'the-id',
+            value: 'something',
+        });
+    });
+    it('calls id generator and setup and uses correct this', async () => {
+        const result = await testClass.startContextWithIdAndSetup(
+            'something else',
+        );
+        expect(result).toEqual({
+            id: 'the-id',
+            value: 'something else',
+        });
+    });
+});

src/lib/use-cls.decorator.ts

@@ -0,0 +1,79 @@
+import 'reflect-metadata';
+import { ClsServiceManager } from './cls-service-manager';
+import { CLS_ID } from './cls.constants';
+import { ClsDecoratorOptions } from './cls.options';
+
+/**
+ * Wraps the decorated method in a CLS context.
+ */
+export function UseCls(): (
+    target: any,
+    propertyKey: string | symbol,
+    descriptor: TypedPropertyDescriptor<(...args: any) => Promise<any>>,
+) => void;
+
+/**
+ * Wraps the decorated method in a CLS context.
+ *
+ * @param options takes similar options to the enhancers.
+ */
+export function UseCls<TArgs extends any[]>(
+    options: ClsDecoratorOptions<TArgs>,
+): (
+    target: any,
+    propertyKey: string | symbol,
+    descriptor: TypedPropertyDescriptor<(...args: TArgs) => Promise<any>>,
+) => void;
+
+export function UseCls<TArgs extends any[]>(
+    maybeOptions?: ClsDecoratorOptions<TArgs>,
+) {
+    return (
+        target: any,
+        propertyKey: string | symbol,
+        descriptor: TypedPropertyDescriptor<(...args: TArgs) => Promise<any>>,
+    ) => {
+        const options = { ...new ClsDecoratorOptions(), ...maybeOptions };
+        const cls = ClsServiceManager.getClsService();
+        const original = descriptor.value;
+        // console.log('original is', original?.toString());
+        if (typeof original !== 'function') {
+            throw new Error(
+                `The @UseCls decorator can be only used on functions, but ${propertyKey.toString()} is not a function.`,
+            );
+        }
+        descriptor.value = function (...args: TArgs) {
+            return cls.run(async () => {
+                if (options.generateId) {
+                    const id = await options.idGenerator?.apply(this, args);
+                    cls.set<string>(CLS_ID, id);
+                }
+                if (options.setup) {
+                    await options.setup.apply(this, [cls, ...args]);
+                }
+                if (options.resolveProxyProviders) {
+                    await cls.resolveProxyProviders();
+                }
+                return original.apply(this, args);
+            });
+        };
+        copyMetadata(original, descriptor.value);
+    };
+}
+
+/**
+ * Copies all metadata from one object to another.
+ * Useful for overwriting function definition in
+ * decorators while keeping all previously
+ * attached metadata
+ *
+ * @param from object to copy metadata from
+ * @param to object to copy metadata to
+ */
+function copyMetadata(from: any, to: any) {
+    const metadataKeys = Reflect.getMetadataKeys(from);
+    metadataKeys.map((key) => {
+        const value = Reflect.getMetadata(key, from);
+        Reflect.defineMetadata(key, value, to);
+    });
+}