Merge pull request #10299 from nestjs/feat/durable-providers-payload

feat(core): allow setting payload for durable trees
nestjs · Sep 19, 2022 · a439cd9 · a439cd9
2 parents b9ea9a4 + 4a43505
commit a439cd9
Show file tree

Hide file tree

Showing 9 changed files with 79 additions and 18 deletions.
diff --git a/integration/scopes/e2e/durable-providers.spec.ts b/integration/scopes/e2e/durable-providers.spec.ts
@@ -23,9 +23,13 @@ describe('Durable providers', () => {
   });
 
   describe('when service is durable', () => {
-    const performHttpCall = (tenantId: number, end: (err?: any) => void) =>
+    const performHttpCall = (
+      tenantId: number,
+      end: (err?: any) => void,
+      endpoint = '/durable',
+    ) =>
       request(server)
-        .get('/durable')
+        .get(endpoint)
         .set({ ['x-tenant-id']: tenantId })
         .end((err, res) => {
           if (err) return end(err);
@@ -67,6 +71,19 @@ describe('Durable providers', () => {
       );
       expect(result.text).equal('Hello world! Counter: 1');
     });
+
+    it(`should register a custom per-tenant request payload`, async () => {
+      let result: request.Response;
+      result = await new Promise<request.Response>(resolve =>
+        performHttpCall(1, resolve, '/durable/echo'),
+      );
+      expect(result.body).deep.equal({ tenantId: '1' });
+
+      result = await new Promise<request.Response>(resolve =>
+        performHttpCall(3, resolve, '/durable/echo'),
+      );
+      expect(result.body).deep.equal({ tenantId: '3' });
+    });
   });
 
   after(async () => {

diff --git a/integration/scopes/src/durable/durable-context-id.strategy.ts b/integration/scopes/src/durable/durable-context-id.strategy.ts
@@ -14,8 +14,10 @@ export class DurableContextIdStrategy implements ContextIdStrategy {
       tenantSubTreeId = { id: +tenantId } as ContextId;
       tenants.set(tenantId, tenantSubTreeId);
     }
-
-    return (info: HostComponentInfo) =>
-      info.isTreeDurable ? tenantSubTreeId : contextId;
+    return {
+      resolve: (info: HostComponentInfo) =>
+        info.isTreeDurable ? tenantSubTreeId : contextId,
+      payload: { tenantId },
+    };
   }
 }
diff --git a/integration/scopes/src/durable/durable.controller.ts b/integration/scopes/src/durable/durable.controller.ts
@@ -9,4 +9,9 @@ export class DurableController {
   greeting(): string {
     return this.durableService.greeting();
   }
+
+  @Get('echo')
+  echo() {
+    return this.durableService.requestPayload;
+  }
 }
diff --git a/integration/scopes/src/durable/durable.service.ts b/integration/scopes/src/durable/durable.service.ts
@@ -1,9 +1,12 @@
-import { Injectable, Scope } from '@nestjs/common';
+import { Inject, Injectable, Scope } from '@nestjs/common';
+import { REQUEST } from '@nestjs/core';
 
 @Injectable({ scope: Scope.REQUEST, durable: true })
 export class DurableService {
   public instanceCounter = 0;
 
+  constructor(@Inject(REQUEST) public readonly requestPayload: unknown) {}
+
   greeting() {
     ++this.instanceCounter;
     return `Hello world! Counter: ${this.instanceCounter}`;

diff --git a/packages/core/helpers/context-id-factory.ts b/packages/core/helpers/context-id-factory.ts
@@ -1,3 +1,4 @@
+import { isObject } from '@nestjs/common/utils/shared.utils';
 import { ContextId, HostComponentInfo } from '../injector/instance-wrapper';
 import { REQUEST_CONTEXT_ID } from '../router/request/request-constants';
 
@@ -13,18 +14,30 @@ export function createContextId(): ContextId {
   return { id: Math.random() };
 }
 
+export type ContextIdResolverFn = (info: HostComponentInfo) => ContextId;
+
+export interface ContextIdResolver {
+  /**
+   * Payload associated with the custom context id
+   */
+  payload: unknown;
+  /**
+   * A context id resolver function
+   */
+  resolve: ContextIdResolverFn;
+}
+
 export interface ContextIdStrategy<T = any> {
   /**
    * Allows to attach a parent context id to the existing child context id.
    * This lets you construct durable DI sub-trees that can be shared between contexts.
    * @param contextId auto-generated child context id
    * @param request request object
-   * @returns a context id resolver function
    */
   attach(
     contextId: ContextId,
     request: T,
-  ): ((info: HostComponentInfo) => ContextId) | undefined;
+  ): ContextIdResolverFn | ContextIdResolver | undefined;
 }
 
 export class ContextIdFactory {
@@ -60,7 +73,13 @@ export class ContextIdFactory {
       return ContextIdFactory.create();
     }
     const contextId = createContextId();
-    contextId.getParent = this.strategy.attach(contextId, request);
+    const resolverObjectOrFunction = this.strategy.attach(contextId, request);
+    if (this.isContextIdResolverWithPayload(resolverObjectOrFunction)) {
+      contextId.getParent = resolverObjectOrFunction.resolve;
+      contextId.payload = resolverObjectOrFunction.payload;
+    } else {
+      contextId.getParent = resolverObjectOrFunction;
+    }
     return contextId;
   }
 
@@ -72,4 +91,10 @@ export class ContextIdFactory {
   public static apply(strategy: ContextIdStrategy) {
     this.strategy = strategy;
   }
+
+  private static isContextIdResolverWithPayload(
+    resolverOrResolverFn: ContextIdResolver | ContextIdResolverFn,
+  ): resolverOrResolverFn is ContextIdResolver {
+    return isObject(resolverOrResolverFn);
+  }
 }
diff --git a/packages/core/injector/instance-wrapper.ts b/packages/core/injector/instance-wrapper.ts
@@ -1,9 +1,5 @@
 import { Logger, LoggerService, Provider, Scope, Type } from '@nestjs/common';
-import {
-  ClassProvider,
-  FactoryProvider,
-  ValueProvider,
-} from '@nestjs/common/interfaces';
+import { FactoryProvider } from '@nestjs/common/interfaces';
 import { clc } from '@nestjs/common/utils/cli-colors.util';
 import { randomStringGenerator } from '@nestjs/common/utils/random-string-generator.util';
 import {
@@ -36,6 +32,7 @@ export interface HostComponentInfo {
 
 export interface ContextId {
   readonly id: number;
+  payload?: unknown;
   getParent?(info: HostComponentInfo): ContextId;
 }
 

diff --git a/packages/core/middleware/middleware-module.ts b/packages/core/middleware/middleware-module.ts
@@ -217,7 +217,10 @@ export class MiddlewareModule {
               writable: false,
               configurable: false,
             });
-            this.container.registerRequestProvider(req, contextId);
+            this.container.registerRequestProvider(
+              contextId.getParent ? contextId.payload : req,
+              contextId,
+            );
           }
           const contextInstance = await this.injector.loadPerContext(
             instance,

diff --git a/packages/core/router/router-explorer.ts b/packages/core/router/router-explorer.ts
@@ -425,7 +425,10 @@ export class RouterExplorer {
         writable: false,
         configurable: false,
       });
-      this.container.registerRequestProvider(request, contextId);
+      this.container.registerRequestProvider(
+        contextId.getParent ? contextId.payload : request,
+        contextId,
+      );
     }
     return contextId;
   }

diff --git a/packages/microservices/listeners-controller.ts b/packages/microservices/listeners-controller.ts
@@ -182,7 +182,10 @@ export class ListenersController {
             reqCtx as BaseRpcContext,
           );
           contextId = this.getContextId(request);
-          this.container.registerRequestProvider(request, contextId);
+          this.container.registerRequestProvider(
+            contextId.getParent ? contextId.payload : request,
+            contextId,
+          );
           dataOrContextHost = request;
         }
         const contextInstance = await this.injector.loadPerContext(
@@ -237,7 +240,10 @@ export class ListenersController {
         writable: false,
         configurable: false,
       });
-      this.container.registerRequestProvider(request, contextId);
+      this.container.registerRequestProvider(
+        contextId.getParent ? contextId.payload : request,
+        contextId,
+      );
     }
     return contextId;
   }