From 7c2712660999c1d4201c9d07c089dbc0ba0c6012 Mon Sep 17 00:00:00 2001
From: Charlie Brown <carbonrobot@gmail.com>
Date: Thu, 9 Nov 2023 16:19:18 -0500
Subject: [PATCH] feat: support global context updates (#49)

---
 README.md                                      | 11 +++++++++++
 src/context.ts                                 | 14 +++++++++++---
 src/index.ts                                   |  2 +-
 src/tests/index.spec.ts                        | 17 +++++++++++++++++
 tap-snapshots/src/tests/index.spec.ts.test.cjs |  4 ++++
 5 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 0ccaf37..0072187 100644
--- a/README.md
+++ b/README.md
@@ -59,6 +59,17 @@ other Cloudwatch aware tools such as Datadog and Splunk.
 }
 ```
 
+## Context Updates
+
+The request logging context can be updated downstream by calling the `updateContext` function. Any duplicate values will overwrite previous values.
+
+_Note: If you provide a custom storage context, you will need to update that storage context directly (this is not typical)_
+
+```
+import { GlobalContextStorageProvider } from 'pino-lambda';
+GlobalContextStorageProvider.updateContext({ userId: '12' });
+```
+
 ## Lambda request tracing
 
 With context tracing enabled, all instances of `pino` that use one of the built in formatters will automatically log the request id and other details so you don't need to pass an instance of a logger to all of your functions.
diff --git a/src/context.ts b/src/context.ts
index 7115d30..56518af 100644
--- a/src/context.ts
+++ b/src/context.ts
@@ -2,16 +2,24 @@
 const CONTEXT_SYMBOL = Symbol.for('aws.lambda.runtime.context');
 
 export interface ContextMap {
-  awsRequestId: string,
-  [key: string]: string
- };
+  awsRequestId: string;
+  [key: string]: string;
+}
 
 export interface ContextStorageProvider {
   getContext: () => ContextMap;
   setContext: (map: ContextMap) => void;
+  updateContext: (values: Record<string, string>) => void;
 }
 
 export const GlobalContextStorageProvider: ContextStorageProvider = {
   getContext: () => (global as any)[CONTEXT_SYMBOL] as ContextMap,
   setContext: (map: ContextMap) => ((global as any)[CONTEXT_SYMBOL] = map),
+  updateContext: (values: Record<string, string>) => {
+    const ctx = GlobalContextStorageProvider.getContext();
+    (global as any)[CONTEXT_SYMBOL] = {
+      ...ctx,
+      ...values,
+    };
+  },
 };
diff --git a/src/index.ts b/src/index.ts
index 077d2ea..0007a2e 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -1,4 +1,4 @@
-// reexport all public types
+export * from './context';
 export * from './destination';
 export * from './formatters';
 export * from './request';
diff --git a/src/tests/index.spec.ts b/src/tests/index.spec.ts
index 9673011..5eafa38 100644
--- a/src/tests/index.spec.ts
+++ b/src/tests/index.spec.ts
@@ -13,6 +13,8 @@ import {
   LogData,
 } from '../';
 
+import { GlobalContextStorageProvider } from '../context';
+
 sinon.useFakeTimers(Date.UTC(2016, 11, 1, 6, 0, 0, 0));
 
 tap.test('should log a simple info message', (t) => {
@@ -85,6 +87,21 @@ tap.test('should log correlation headers', (t) => {
   t.end();
 });
 
+tap.test('should allow modifying the global context', (t) => {
+  const { log, output, withRequest } = createLogger();
+
+  withRequest(
+    { headers: { 'x-correlation-data': 'abbb' } },
+    { awsRequestId: '98875' },
+  );
+
+  GlobalContextStorageProvider.updateContext({ userId: '12' });
+
+  log.error('Context updates');
+  t.matchSnapshot(output.buffer);
+  t.end();
+});
+
 tap.test('should set correlation if to trace id if present', (t) => {
   const { log, output, withRequest } = createLogger();
 
diff --git a/tap-snapshots/src/tests/index.spec.ts.test.cjs b/tap-snapshots/src/tests/index.spec.ts.test.cjs
index 40389a8..b7c589f 100644
--- a/tap-snapshots/src/tests/index.spec.ts.test.cjs
+++ b/tap-snapshots/src/tests/index.spec.ts.test.cjs
@@ -17,6 +17,10 @@ exports[`src/tests/index.spec.ts TAP should allow default pino formatter > must
 {"awsRequestId":"431234","x-correlation-id":"431234","level":30,"time":1480572000000,"msg":"Message with pino formatter"}
 `
 
+exports[`src/tests/index.spec.ts TAP should allow modifying the global context > must match snapshot 1`] = `
+2016-12-01T06:00:00.000Z	98875	ERROR	Context updates	{"awsRequestId":"98875","x-correlation-data":"abbb","x-correlation-id":"98875","userId":"12","level":50,"time":1480572000000,"msg":"Context updates"}
+`
+
 exports[`src/tests/index.spec.ts TAP should allow removing default request data > must match snapshot 1`] = `
 2016-12-01T06:00:00.000Z	431234	INFO	Message with trace ID	{"awsRequestId":"431234","level":30,"time":1480572000000,"msg":"Message with trace ID"}
 `