Repository: iamolegga/nestjs-pino
Pull request #2156 (2024-12-23 09:05:41)

Option to `assign` response logs
There's been a few issues around not being able to assign response logs, mainly [this one](https://github.com/iamolegga/nestjs-pino/issues/608). The follow up has been to [document this feature as not supported](https://github.com/iamolegga/nestjs-pino?tab=readme-ov-file#assign-extra-fields-for-future-calls).

However I've found a way to assign response logs as well, rather painlessly. This is a proposal to add an option to support that. I've implemented this as an option such that the default behaviour wouldn't change. Also, because this uses `setBinding`, this effectively changes the existing response logger rather than create a child logger (which wouldn't work).

I couldn't find a reason why this solution hadn't been considered, but maybe there are downsides I haven't though of. Feel free to comment on that.

Cheers,

iamolegga: @flovouin thank you for fantastic work. Could you squash commits in one? And I'll do the rest
flovouin: Thanks for the quick response @iamolegga! I've merged the two commits into one.

Patch:
diff --git a/README.md b/README.md
index 850b32e..3bded5e 100644
--- a/README.md
+++ b/README.md
@@ -437,7 +437,7 @@ class MyService {
 }
 ```
 
-Due to the [limitation](https://github.com/pinojs/pino-http/issues/30) of the underlying `pino-http` `PinoLogger.assign` cannot extend `Request completed` logs.
+By default, this does not extend `Request completed` logs. Set the `assignResponse` parameter to `true` to also enrich response logs automatically emitted by `pino-http`.
 
 ## Change pino params at runtime
 
diff --git a/src/LoggerModule.ts b/src/LoggerModule.ts
index 5eb3e8f..6af9c73 100644
--- a/src/LoggerModule.ts
+++ b/src/LoggerModule.ts
@@ -75,9 +75,14 @@ export class LoggerModule implements NestModule {
       forRoutes = DEFAULT_ROUTES,
       pinoHttp,
       useExisting,
+      assignResponse,
     } = this.params;
 
-    const middlewares = createLoggerMiddlewares(pinoHttp || {}, useExisting);
+    const middlewares = createLoggerMiddlewares(
+      pinoHttp || {},
+      useExisting,
+      assignResponse,
+    );
 
     if (exclude) {
       consumer
@@ -93,9 +98,10 @@ export class LoggerModule implements NestModule {
 function createLoggerMiddlewares(
   params: NonNullable<Params['pinoHttp']>,
   useExisting = false,
+  assignResponse = false,
 ) {
   if (useExisting) {
-    return [bindLoggerMiddlewareFactory(useExisting)];
+    return [bindLoggerMiddlewareFactory(useExisting, assignResponse)];
   }
 
   const middleware = pinoHttp(
@@ -108,24 +114,31 @@ function createLoggerMiddlewares(
 
   // FIXME: params type here is pinoHttp.Options | pino.DestinationStream
   // pinoHttp has two overloads, each of them takes those types
-  return [middleware, bindLoggerMiddlewareFactory(useExisting)];
+  return [middleware, bindLoggerMiddlewareFactory(useExisting, assignResponse)];
 }
 
-function bindLoggerMiddlewareFactory(useExisting: boolean) {
+function bindLoggerMiddlewareFactory(
+  useExisting: boolean,
+  assignResponse: boolean,
+) {
   return function bindLoggerMiddleware(
     req: express.Request,
-    _res: express.Response,
+    res: express.Response,
     next: express.NextFunction,
   ) {
     let log = req.log;
+    let resLog = assignResponse ? res.log : undefined;
 
     if (!useExisting && req.allLogs) {
       log = req.allLogs[req.allLogs.length - 1]!;
     }
+    if (assignResponse && !useExisting && res.allLogs) {
+      resLog = res.allLogs[res.allLogs.length - 1]!;
+    }
 
     // eslint-disable-next-line @typescript-eslint/ban-ts-comment
     // @ts-ignore: run requires arguments for next but should not because it can
     // be called without arguments
-    storage.run(new Store(log), next);
+    storage.run(new Store(log, resLog), next);
   };
 }
diff --git a/src/PinoLogger.ts b/src/PinoLogger.ts
index d6461c0..fae2899 100644
--- a/src/PinoLogger.ts
+++ b/src/PinoLogger.ts
@@ -134,6 +134,7 @@ export class PinoLogger implements PinoMethods {
       );
     }
     store.logger = store.logger.child(fields);
+    store.responseLogger?.setBindings(fields);
   }
 
   protected call(method: pino.Level, ...args: Parameters<LoggerFn>) {
diff --git a/src/params.ts b/src/params.ts
index 02251c7..049605c 100644
--- a/src/params.ts
+++ b/src/params.ts
@@ -54,6 +54,13 @@ export interface Params {
    * {"level":30, ... "RENAME_CONTEXT_VALUE_HERE":"AppController" }
    */
   renameContext?: string;
+
+  /**
+   * Optional parameter to also assign the response logger during calls to
+   * `PinoLogger.assign`. By default, `assign` does not impact response logs
+   * (e.g.`Request completed`).
+   */
+  assignResponse?: boolean;
 }
 
 // for support of nestjs@8 we don't use
diff --git a/src/storage.ts b/src/storage.ts
index 7d7d76c..6cebafd 100644
--- a/src/storage.ts
+++ b/src/storage.ts
@@ -3,7 +3,10 @@ import { AsyncLocalStorage } from 'async_hooks';
 import { Logger } from 'pino';
 
 export class Store {
-  constructor(public logger: Logger) {}
+  constructor(
+    public logger: Logger,
+    public responseLogger?: Logger,
+  ) {}
 }
 
 export const storage = new AsyncLocalStorage<Store>();


Test patch:
diff --git a/__tests__/assign.spec.ts b/__tests__/assign.spec.ts
index 5f27b9e..c94d987 100644
--- a/__tests__/assign.spec.ts
+++ b/__tests__/assign.spec.ts
@@ -45,11 +45,14 @@ describe('assign', () => {
           controllers: [TestController],
           providers: [TestService],
         })
-          .forRoot()
+          .forRoot({ pinoHttp: { customSuccessMessage: () => 'success' } })
           .run();
 
         const wanted = logs.some((l) => l.msg === msg && l.foo === 'bar');
         expect(wanted).toBeTruthy();
+        const responseLog = logs.find((l) => l.msg === 'success');
+        expect(responseLog).toBeDefined();
+        expect(responseLog).not.toHaveProperty('foo');
       });
 
       it('out of request context', async () => {
@@ -92,6 +95,54 @@ describe('assign', () => {
         expect(log).toBeTruthy();
         expect(log).not.toHaveProperty('foo');
       });
+
+      it('response log', async () => {
+        const msg = Math.random().toString();
+
+        @Injectable()
+        class TestService {
+          private readonly logger = new Logger(TestService.name);
+
+          test() {
+            this.logger.log(msg);
+            return {};
+          }
+        }
+
+        @Controller('/')
+        class TestController {
+          constructor(
+            private readonly logger: PinoLogger,
+            private readonly service: TestService,
+          ) {}
+
+          @Get()
+          get() {
+            this.logger.assign({ foo: 'bar' });
+            this.logger.assign({ other: 'value' });
+            return this.service.test();
+          }
+        }
+
+        const logs = await new TestCase(new PlatformAdapter(), {
+          controllers: [TestController],
+          providers: [TestService],
+        })
+          .forRoot({
+            assignResponse: true,
+            pinoHttp: { customSuccessMessage: () => 'success' },
+          })
+          .run();
+
+        const hasServiceLog = logs.some(
+          (l) => l.msg === msg && l.foo === 'bar' && l.other === 'value',
+        );
+        expect(hasServiceLog).toBeTruthy();
+        const hasResponseLog = logs.some(
+          (l) => l.msg === 'success' && l.foo === 'bar' && l.other === 'value',
+        );
+        expect(hasResponseLog).toBeTruthy();
+      });
     });
   }
 });