Improve type inference and type safety for ClsService (#16)

* feat: add recursive key types

* feat: typing for ClsService#get

* feat: typing for ClsService#set

* chore: lint - remove unused variables

* docs: add section on type safety

README.md

@@ -1,18 +1,20 @@
 # NestJS CLS
 
+> **New**: Release `2.0` brings advanced [type safety and type inference](#type-safety-and-type-inference), check below for more info.
+
 A continuation-local storage module compatible with [NestJS](https://nestjs.com/)'s dependency injection.
 
 _Continuous-local storage allows to store state and propagate it throughout callbacks and promise chains. It allows storing data throughout the lifetime of a web request or any other asynchronous duration. It is similar to thread-local storage in other languages._
 
 Some common use cases for CLS include:
 
--   Request ID tracing for logging purposes
+-   Tracing the Request ID and other metadata for logging purposes
 -   Making the Tenant ID available everywhere in multi-tenant apps
--   Globally setting the authentication level for the request
+-   Globally setting an authentication level for the request
 
-Most of these are theoretically solvable using _request-scoped_ providers or passing the context as a parameter, but these solutions are often clunky and come with a whole lot of other issues. Thus this package was born.
+Most of these are to some extent solvable using _request-scoped_ providers or passing the context as a parameter, but these solutions are often clunky and come with a whole lot of other issues.
 
-> **Note**: For versions < 1.2, this package used [cls-hooked](https://www.npmjs.com/package/cls-hooked) as a peer dependency, now it uses [AsyncLocalStorage](https://nodejs.org/api/async_context.html#async_context_class_asynclocalstorage) from Node's `async_hooks` directly. The API stays the same for now but I'll consider making it more friendly for version 2.
+> **Note**: This package uses [AsyncLocalStorage](https://nodejs.org/api/async_context.html#async_context_class_asynclocalstorage) from Node's `async_hooks` API. Most parts of it are marked as _stable_ now, see [Security considerations](#security-considerations) for more details.
 
 # Outline
 
@@ -28,6 +30,7 @@ Most of these are theoretically solvable using _request-scoped_ providers or pas
 -   [Request ID](#request-id)
 -   [Additional CLS Setup](#additional-cls-setup)
 -   [Breaking out of DI](#breaking-out-of-di)
+-   [Type safety and type inference](#type-safety-and-type-inference)
 -   [Security considerations](#security-considerations)
 -   [Compatibility considerations](#compatibility-considerations)
     -   [REST](#rest)
@@ -58,23 +61,22 @@ Below is an example of storing the client's IP address in an interceptor and ret
         // Register the ClsModule and automatically mount the ClsMiddleware
         ClsModule.register({
             global: true,
-            middleware: { mount: true }
+            middleware: { mount: true },
         }),
     ],
     providers: [AppService],
     controllers: [AppController],
 })
 export class TestHttpApp {}
 
-
 /* user-ip.interceptor.ts */
 @Injectable()
 export class UserIpInterceptor implements NestInterceptor {
     constructor(
         // Inject the ClsService into the interceptor to get
         // access to the current shared cls context.
-        private readonly cls: ClsService
-    )
+        private readonly cls: ClsService,
+    );
 
     intercept(context: ExecutionContext, next: CallHandler): Observable<any> {
         // Extract the client's ip address from the request...
@@ -86,7 +88,6 @@ export class UserIpInterceptor implements NestInterceptor {
     }
 }
 
-
 /* app.controller.ts */
 
 // By mounting the interceptor on the controller, it gets access
@@ -102,13 +103,12 @@ export class AppController {
     }
 }
 
-
 /* app.service.ts */
 @Injectable()
 export class AppService {
     constructor(
         // Inject ClsService to be able to retrieve data from the cls context.
-        private readonly cls: ClsService
+        private readonly cls: ClsService,
     ) {}
 
     sayHello() {
@@ -194,7 +194,9 @@ If you need any other guards to use the `ClsService`, it's preferable to mount `
 })
 export class AppModule {}
 ```
+
 or mount it directly on the Controller/Resolver with
+
 ```ts
 @UseGuards(ClsGuard);
 ```
@@ -214,26 +216,23 @@ ClsModule.register({
 ```
 
 Or mount it manually as `APP_INTERCEPTOR`, or directly on the Controller/Resolver with:
+
 ```ts
 @UseInterceptors(ClsInterceptor);
 ```
 
 > **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**).
 
-
-
 # API
 
 The injectable `ClsService` provides the following API to manipulate the cls context:
 
--   **_`set`_**`<T>(key: string, value: T): T`  
+-   **_`set`_**`(key: string, value: any): void`  
     Set a value on the CLS context.
--   **_`get`_**`<T>(key: string): T`  
-    Retrieve a value from the CLS context by key.
+-   **_`get`_**`(key?: string): any`  
+    Retrieve a value from the CLS context by key. Get the whole store if key is omitted.
 -   **_`getId`_**`(): string;`  
     Retrieve the request ID (a shorthand for `cls.get(CLS_ID)`)
--   **_`getStore`_**`(): any`  
-    Retrieve the object containing all properties of the current CLS context.
 -   **_`enter`_**`(): void;`  
     Run any following code in a shared CLS context.
 -   **_`enterWith`_**`(store: any): void;`  
@@ -284,7 +283,6 @@ The `ClsMiddlewareOptions` additionally takes the following parameters:
 -   **_`useEnterWith`_: `boolean`** (default _`false`_)  
     Set to `true` to set up the context using a call to [`AsyncLocalStorage#enterWith`](https://nodejs.org/api/async_context.html#async_context_asynclocalstorage_enterwith_store) instead of wrapping the `next()` call with the safer [`AsyncLocalStorage#run`](https://nodejs.org/api/async_context.html#async_context_asynclocalstorage_run_store_callback_args). Most of the time this should not be necessary, but [some frameworks](#graphql) are known to lose the context with `run`.
 
-
 # Request ID
 
 Because of a shared storage, CLS is an ideal tool for tracking request (correlation) IDs for the purpose of logging. This package provides an option to automatically generate request IDs in the middleware/guard/interceptor, if you pass `{ generateId: true }` to its options. By default, the generated ID is a string based on `Math.random()`, but you can provide a custom function in the `idGenerator` option.
@@ -341,7 +339,7 @@ The function receives the `ClsService` instance and the `Request` (or `Execution
 ClsModule.register({
     middleware: {
         mount: true,
-        setup: (cls, req) => {
+        setup: (cls, req: Request) => {
             // put some additional default info in the CLS
             cls.set('TENANT_ID', req.params('tenant_id'));
             cls.set('AUTH', { authenticated: false });
@@ -364,6 +362,71 @@ function helper() {
 
 > **Please note**: Only use this feature where absolutely necessary. Using this technique instead of dependency injection will make it difficult to mock the ClsService and your code will become harder to test.
 
+# Type safety and type inference
+
+By default the CLS context is untyped and allows setting and retrieving any `string` or `symbol` key from the context. Some safety can be enforced by using `CONSTANTS` instead of magic strings, but that might not be enough.
+
+Therefore, it is possible to specify a custom interface for the `ClsService` and get proper typing and automatic type inference when retrieving or setting values. This works even for _nested objects_ using a dot notation.
+
+To create a typed CLS Store, start by creating an interface that extends `ClsStore`.
+
+```ts
+export interface MyClsStore extends ClsStore {
+    tenantId: string;
+    user: {
+        id: number;
+        authorized: boolean;
+    };
+}
+```
+
+Then you can inject the `ClsService` with a type parameter `ClsService<MyClsStore>` and
+
+```ts
+export class MyService {
+    constructor(private readonly cls: ClsService<ClsStore>) {}
+
+    doTheThing() {
+        // a boolean type will be enforced here
+        this.cls.set('user.authorized', true);
+
+        // tenantId will be inferred as a stirng
+        const tenantId = this.cls.get('tenantId');
+
+        // userId will be inferred as a number
+        const userId = this.cls.get('user.id');
+
+        // user will be inferred as { id: number, authorized: boolean }
+        const user = this.cls.get('user');
+
+        // you'll even get intellisense for the keys, because the type
+        // will be inferred as:
+        // symbol | 'tenantId˙ | 'user' | 'user.id' | 'user.authorized'
+
+        // alternatively, since the `get` method returns the whole store
+        // when called without arguments, you can use object destructuring
+        const { tenantId, user } = this.cls.get();
+
+        // accessing a nonexistent property will result in a type error
+        const notExist = this.cls.get('user.name');
+    }
+}
+```
+
+Alternatively, if you feel like using `ClsService<MyClsStore>` everywhere is tedious, you can instead globally [augment the `ClsStore interface`](https://www.typescriptlang.org/docs/handbook/declaration-merging.html) and have strict typing of `ClsService` anywhere without the type parameter:
+
+```ts
+declare module 'nestjs-cls' {
+    interface ClsStore {
+        tenantId: string;
+        user: {
+            id: number;
+            authorized: boolean;
+        };
+    }
+}
+```
+
 # Security considerations
 
 It is often discussed whether [`AsyncLocalStorage`](https://nodejs.org/api/async_context.html) is safe to use for _concurrent requests_ (because of a possible context leak) and whether the context could be _lost_ throughout the duration of a request.
@@ -384,11 +447,11 @@ The `ClsInterceptor` only uses the safe `run()` method.
 
 The table below outlines the compatibility with some platforms:
 
-|                                                              |                        REST                         |                            GQL                             | WS |Others |
-| :----------------------------------------------------------: | :-------------------------------------------------: | :--------------------------------------------------------: |:--:|:----: |
-|                      **ClsMiddleware**                       |                          ✔                          | ✔<br>must be _mounted manually_<br>and use `useEnterWith: true` | ✖ | ✖    |
-|             **ClsGuard** <br>(uses `enterWith`)              |                          ✔                          |                             ✔                              | ✔[*](#websockets)  | ?    |
-| **ClsInterceptor** <br>(context inaccessible<br>in _Guards_) | ✔<br>context also inaccessible<br>in _Exception Filters_ |                             ✔                              |  ✔[*](#websockets) |?    |
+|                                                              |                           REST                           |                               GQL                               |         WS         | Others |
+| :----------------------------------------------------------: | :------------------------------------------------------: | :-------------------------------------------------------------: | :----------------: | :----: |
+|                      **ClsMiddleware**                       |                            ✔                             | ✔<br>must be _mounted manually_<br>and use `useEnterWith: true` |         ✖          |   ✖    |
+|             **ClsGuard** <br>(uses `enterWith`)              |                            ✔                             |                                ✔                                | ✔[\*](#websockets) |   ?    |
+| **ClsInterceptor** <br>(context inaccessible<br>in _Guards_) | ✔<br>context also inaccessible<br>in _Exception Filters_ |                                ✔                                | ✔[\*](#websockets) |   ?    |
 
 ## REST
 
@@ -419,7 +482,8 @@ Use the `ClsGuard` or `ClsInterceptor` to set up context with any other platform
 Below are listed platforms with which it is confirmed to work.
 
 ### Websockets
-*Websocket Gateways* don't respect globally bound enhancers, therefore it is required to bind the `ClsGuard` or `ClsIntercetor` manually on the `WebscocketGateway`. (See [#8](https://github.com/Papooch/nestjs-cls/issues/8))
+
+_Websocket Gateways_ don't respect globally bound enhancers, therefore it is required to bind the `ClsGuard` or `ClsIntercetor` manually on the `WebscocketGateway`. (See [#8](https://github.com/Papooch/nestjs-cls/issues/8))
 
 # Namespaces (experimental)
 

package.json

@@ -93,6 +93,12 @@
             "**/*.(t|j)s"
         ],
         "coverageDirectory": "../coverage",
-        "testEnvironment": "node"
+        "testEnvironment": "node",
+        "globals": {
+            "ts-jest": {
+                "isolatedModules": true,
+                "maxWorkers": 1
+            }
+        }
     }
 }

src/lib/cls.constants.ts

@@ -1,6 +1,6 @@
-export const CLS_REQ = 'CLS_REQUEST';
-export const CLS_RES = 'CLS_RESPONSE';
-export const CLS_ID = 'CLS_ID';
+export const CLS_REQ = Symbol('CLS_REQUEST');
+export const CLS_RES = Symbol('CLS_RESPONSE');
+export const CLS_ID = Symbol('CLS_ID');
 export const CLS_DEFAULT_NAMESPACE = 'CLS_DEFAULT_NAMESPACE';
 export const CLS_MODULE_OPTIONS = 'ClsModuleOptions';
 export const CLS_MIDDLEWARE_OPTIONS = 'ClsMiddlewareOptions';

src/lib/cls.guard.ts

@@ -23,7 +23,7 @@ export class ClsGuard implements CanActivate {
             cls.enter();
             if (this.options.generateId) {
                 const id = await this.options.idGenerator(context);
-                cls.set(CLS_ID, id);
+                cls.set<any>(CLS_ID, id);
             }
             if (this.options.setup) {
                 await this.options.setup(cls, context);

src/lib/cls.interceptor.ts

@@ -25,7 +25,7 @@ export class ClsInterceptor implements NestInterceptor {
             cls.run(async () => {
                 if (this.options.generateId) {
                     const id = await this.options.idGenerator(context);
-                    cls.set(CLS_ID, id);
+                    cls.set<any>(CLS_ID, id);
                 }
                 if (this.options.setup) {
                     await this.options.setup(cls, context);

src/lib/cls.interfaces.ts

@@ -68,14 +68,14 @@ export class ClsMiddlewareOptions {
     /**
      * the function to generate request ids inside the middleware
      */
-    idGenerator?: (req: Request) => string | Promise<string> = () =>
+    idGenerator?: (req: any) => 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, req: Request) => void | Promise<void>;
+    setup?: (cls: ClsService, req: any) => void | Promise<void>;
 
     /**
      * Whether to store the Request object to the CLS
@@ -159,3 +159,7 @@ export class ClsInterceptorOptions {
 
     readonly namespaceName?: string;
 }
+
+export interface ClsStore {
+    [key: symbol]: any;
+}

src/lib/cls.middleware.ts

@@ -21,10 +21,10 @@ export class ClsMiddleware implements NestMiddleware {
             this.options.useEnterWith && cls.enter();
             if (this.options.generateId) {
                 const id = await this.options.idGenerator(req);
-                cls.set(CLS_ID, id);
+                cls.set<any>(CLS_ID, id);
             }
-            if (this.options.saveReq) cls.set(CLS_REQ, req);
-            if (this.options.saveRes) cls.set(CLS_RES, res);
+            if (this.options.saveReq) cls.set<any>(CLS_REQ, req);
+            if (this.options.saveRes) cls.set<any>(CLS_RES, res);
             if (this.options.setup) {
                 await this.options.setup(cls, req);
             }