diff --git a/package-lock.json b/package-lock.json index 433ad38564..49a40b6e28 100644 --- a/package-lock.json +++ b/package-lock.json @@ -10687,10 +10687,22 @@ "name": "@aws-lambda-powertools/commons", "version": "2.27.0", "license": "MIT-0", + "dependencies": { + "@aws/lambda-invoke-store": "0.1.0" + }, "devDependencies": { "@aws-lambda-powertools/testing-utils": "file:../testing" } }, + "packages/commons/node_modules/@aws/lambda-invoke-store": { + "version": "0.1.0", + "resolved": "https://registry.npmjs.org/@aws/lambda-invoke-store/-/lambda-invoke-store-0.1.0.tgz", + "integrity": "sha512-I1y5yahbSFTfKldV4qoKv2IEZ20QOhn5rPvWwGnswZ8hssN7tsLANLg9tL8dp2klz2MZDGL5jZrvBwplIWtM8A==", + "license": "Apache-2.0", + "engines": { + "node": ">=18.0.0" + } + }, "packages/event-handler": { "name": "@aws-lambda-powertools/event-handler", "version": "2.27.0", diff --git a/packages/commons/package.json b/packages/commons/package.json index 5d4bb706b8..6ad22ac5bf 100644 --- a/packages/commons/package.json +++ b/packages/commons/package.json @@ -110,6 +110,9 @@ "serverless", "nodejs" ], + "dependencies": { + "@aws/lambda-invoke-store": "0.1.0" + }, "devDependencies": { "@aws-lambda-powertools/testing-utils": "file:../testing" } diff --git a/packages/commons/src/envUtils.ts b/packages/commons/src/envUtils.ts index cfc4b2adfa..751d109bfe 100644 --- a/packages/commons/src/envUtils.ts +++ b/packages/commons/src/envUtils.ts @@ -1,3 +1,4 @@ +import { InvokeStore } from '@aws/lambda-invoke-store'; import { POWERTOOLS_DEV_ENV_VAR, POWERTOOLS_SERVICE_NAME_ENV_VAR, @@ -249,15 +250,19 @@ const getServiceName = (): string => { }; /** - * Get the AWS X-Ray Trace data from the environment variable. + * Get the AWS X-Ray Trace data from the lambda RIC async context or the `_X_AMZN_TRACE_ID` environment variable. * - * The method parses the environment variable `_X_AMZN_TRACE_ID` and returns an object with the key-value pairs. + * Checks the async context first and if that returns undefined, falls back to the environment variable + * + * The method parses the value and returns an object with the key-value pairs. */ const getXrayTraceDataFromEnv = (): Record | undefined => { - const xRayTraceEnv = getStringFromEnv({ - key: XRAY_TRACE_ID_ENV_VAR, - defaultValue: '', - }); + const xRayTraceEnv = + InvokeStore.getXRayTraceId() ?? + getStringFromEnv({ + key: XRAY_TRACE_ID_ENV_VAR, + defaultValue: '', + }); if (xRayTraceEnv === '') { return undefined; } @@ -280,8 +285,10 @@ const getXrayTraceDataFromEnv = (): Record | undefined => { /** * Determine if the current invocation is part of a sampled X-Ray trace. * - * The AWS X-Ray Trace data available in the environment variable has this format: + * The AWS X-Ray Trace data is available in either the RIC async context or the `_X_AMZN_TRACE_ID` environment variable has this format: * `Root=1-5759e988-bd862e3fe1be46a994272793;Parent=557abcec3ee5a047;Sampled=1`, + * + * Checks the async context first and if that returns undefined, falls back to the environment variable */ const isRequestXRaySampled = (): boolean => { const xRayTraceData = getXrayTraceDataFromEnv(); @@ -289,9 +296,11 @@ const isRequestXRaySampled = (): boolean => { }; /** - * Get the value of the `_X_AMZN_TRACE_ID` environment variable. + * AWS X-Ray Trace id from the lambda RIC async context or the `_X_AMZN_TRACE_ID` environment variable. + * + * Checks the async context first and if that returns undefined, falls back to the environment variable * - * The AWS X-Ray Trace data available in the environment variable has this format: + * The AWS X-Ray Trace data has this format: * `Root=1-5759e988-bd862e3fe1be46a994272793;Parent=557abcec3ee5a047;Sampled=1`, * * The actual Trace ID is: `1-5759e988-bd862e3fe1be46a994272793`. diff --git a/packages/commons/tests/unit/envUtils.test.ts b/packages/commons/tests/unit/envUtils.test.ts index 1c32bceb3d..74f7df741d 100644 --- a/packages/commons/tests/unit/envUtils.test.ts +++ b/packages/commons/tests/unit/envUtils.test.ts @@ -1,4 +1,5 @@ -import { beforeEach, describe, expect, it } from 'vitest'; +import { InvokeStore } from '@aws/lambda-invoke-store'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; import { getBooleanFromEnv, getNumberFromEnv, @@ -10,16 +11,14 @@ import { } from '../../src/envUtils.js'; describe('Functions: envUtils', () => { - const env = process.env; - beforeEach(() => { - process.env = { ...env }; + vi.unstubAllEnvs(); }); describe('Function: getStringFromEnv', () => { it('returns the value of the environment variable', () => { // Prepare - process.env.TEST_ENV = 'testValue'; + vi.stubEnv('TEST_ENV', 'testValue'); // Act const result = getStringFromEnv({ key: 'TEST_ENV' }); @@ -29,9 +28,6 @@ describe('Functions: envUtils', () => { }); it('returns the default value if the environment variable is not set', () => { - // Prepare - process.env.TEST_ENV = undefined; - // Act const result = getStringFromEnv({ key: 'TEST_ENV', @@ -43,9 +39,6 @@ describe('Functions: envUtils', () => { }); it('throws an error if the environment variable is not set', () => { - // Prepare - process.env.TEST_ENV = undefined; - // Act & Assess expect(() => getStringFromEnv({ key: 'TEST_ENV' })).toThrowError( 'Environment variable TEST_ENV is required' @@ -54,7 +47,7 @@ describe('Functions: envUtils', () => { it('returns the trimmed value of the environment variable', () => { // Prepare - process.env.TEST_ENV = ' testValue '; + vi.stubEnv('TEST_ENV', ' testValue '); // Act const result = getStringFromEnv({ key: 'TEST_ENV' }); @@ -64,9 +57,6 @@ describe('Functions: envUtils', () => { }); it('uses the provided error message if the environment variable is not set', () => { - // Prepare - process.env.TEST_ENV = undefined; - // Act & Assess expect(() => getStringFromEnv({ @@ -80,7 +70,7 @@ describe('Functions: envUtils', () => { describe('Function: getNumberFromEnv', () => { it('returns the value of the environment variable as a number', () => { // Prepare - process.env.TEST_ENV = '123'; + vi.stubEnv('TEST_ENV', '123'); // Act const result = getNumberFromEnv({ key: 'TEST_ENV' }); @@ -90,9 +80,6 @@ describe('Functions: envUtils', () => { }); it('returns the default value if the environment variable is not set', () => { - // Prepare - process.env.TEST_ENV = undefined; - // Act const result = getNumberFromEnv({ key: 'TEST_ENV', @@ -105,7 +92,7 @@ describe('Functions: envUtils', () => { it('throws an error if the environment variable is not a number', () => { // Prepare - process.env.TEST_ENV = 'notANumber'; + vi.stubEnv('TEST_ENV', 'notANumber'); // Act & Assess expect(() => getNumberFromEnv({ key: 'TEST_ENV' })).toThrowError( @@ -117,7 +104,7 @@ describe('Functions: envUtils', () => { describe('Function: getBooleanFromEnv', () => { it('returns true if the environment variable is set to a truthy value', () => { // Prepare - process.env.TEST_ENV = 'true'; + vi.stubEnv('TEST_ENV', 'true'); // Act const result = getBooleanFromEnv({ key: 'TEST_ENV' }); @@ -128,7 +115,7 @@ describe('Functions: envUtils', () => { it('returns false if the environment variable is set to a falsy value', () => { // Prepare - process.env.TEST_ENV = 'false'; + vi.stubEnv('TEST_ENV', 'false'); // Act const result = getBooleanFromEnv({ key: 'TEST_ENV' }); @@ -138,9 +125,6 @@ describe('Functions: envUtils', () => { }); it('returns the default value if the environment variable is not set', () => { - // Prepare - process.env.TEST_ENV = undefined; - // Act const result = getBooleanFromEnv({ key: 'TEST_ENV', @@ -153,7 +137,7 @@ describe('Functions: envUtils', () => { it('throws an error if the environment variable value is not a boolean', () => { // Prepare - process.env.TEST_ENV = 'notABoolean'; + vi.stubEnv('TEST_ENV', 'notABoolean'); // Act & Assess expect(() => getBooleanFromEnv({ key: 'TEST_ENV' })).toThrowError( @@ -172,7 +156,7 @@ describe('Functions: envUtils', () => { 'returns true if the environment variable is set to a truthy value: %s', (value, expected) => { // Prepare - process.env.TEST_ENV = value; + vi.stubEnv('TEST_ENV', value); // Act const result = getBooleanFromEnv({ @@ -196,7 +180,7 @@ describe('Functions: envUtils', () => { 'returns false if the environment variable is set to a falsy value: %s', (value, expected) => { // Prepare - process.env.TEST_ENV = value; + vi.stubEnv('TEST_ENV', value); // Act const result = getBooleanFromEnv({ @@ -213,7 +197,7 @@ describe('Functions: envUtils', () => { describe('Function: isDevMode', () => { it('returns true if the environment variable is set to a truthy value', () => { // Prepare - process.env.POWERTOOLS_DEV = 'true'; + vi.stubEnv('POWERTOOLS_DEV', 'true'); // Act const result = isDevMode(); @@ -224,7 +208,7 @@ describe('Functions: envUtils', () => { it('returns false if the environment variable is set to a falsy value', () => { // Prepare - process.env.POWERTOOLS_DEV = 'false'; + vi.stubEnv('POWERTOOLS_DEV', 'false'); // Act const result = isDevMode(); @@ -234,9 +218,6 @@ describe('Functions: envUtils', () => { }); it('returns false if the environment variable is not set', () => { - // Prepare - process.env.POWERTOOLS_DEV = undefined; - // Act const result = isDevMode(); @@ -248,7 +229,7 @@ describe('Functions: envUtils', () => { describe('Function: getServiceName', () => { it('returns the service name from the environment variable', () => { // Prepare - process.env.POWERTOOLS_SERVICE_NAME = 'testService'; + vi.stubEnv('POWERTOOLS_SERVICE_NAME', 'testService'); // Act const result = getServiceName(); @@ -259,7 +240,7 @@ describe('Functions: envUtils', () => { it('returns an empty string if the environment variable is not set', () => { // Prepare - process.env.POWERTOOLS_SERVICE_NAME = undefined; + vi.stubEnv('POWERTOOLS_SERVICE_NAME', undefined); // Act const result = getServiceName(); @@ -270,64 +251,135 @@ describe('Functions: envUtils', () => { }); describe('Function: getXrayTraceIdFromEnv', () => { - it('returns the value of the environment variable _X_AMZN_TRACE_ID', () => { + it.each<{ description: string; traceData: string; expected: string }>([ + { + description: + 'returns the value of the environment variable _X_AMZN_TRACE_ID', + traceData: 'abcd123456789', + expected: 'abcd123456789', + }, + { + description: + 'returns the value of the Root X-Ray segment ID properly formatted', + traceData: + 'Root=1-5759e988-bd862e3fe1be46a994272793;Parent=557abcec3ee5a047;Sampled=1', + expected: '1-5759e988-bd862e3fe1be46a994272793', + }, + ])('$description', ({ traceData, expected }) => { // Prepare - process.env._X_AMZN_TRACE_ID = 'abcd123456789'; + vi.stubEnv('_X_AMZN_TRACE_ID', traceData); // Act const value = getXRayTraceIdFromEnv(); // Assess - expect(value).toEqual('abcd123456789'); + expect(value).toEqual(expected); }); - it('returns the value of the Root X-Ray segment ID properly formatted', () => { - // Prepare - process.env._X_AMZN_TRACE_ID = - 'Root=1-5759e988-bd862e3fe1be46a994272793;Parent=557abcec3ee5a047;Sampled=1'; - - // Act - const value = getXRayTraceIdFromEnv(); - - // Assess - expect(value).toEqual('1-5759e988-bd862e3fe1be46a994272793'); + it.each<{ description: string; traceData: string; expected: string }>([ + { + description: 'returns trace id from async context', + traceData: 'xyz987654321', + expected: 'xyz987654321', + }, + { + description: + 'returns the Root X-Ray segment ID properly formatted from async context', + traceData: + 'Root=1-6849f099-ce973f4ea2c57e4f9a382904;Parent=668bfc7d9aa5b120;Sampled=0', + expected: '1-6849f099-ce973f4ea2c57e4f9a382904', + }, + ])('$description', ({ traceData, expected }) => { + InvokeStore.run( + { + [InvokeStore.PROTECTED_KEYS.X_RAY_TRACE_ID]: traceData, + }, + () => { + // Act + const value = getXRayTraceIdFromEnv(); + + // Assess + expect(value).toEqual(expected); + } + ); }); }); describe('Function: isRequestXRaySampled', () => { - it('returns true if the Sampled flag is set in the _X_AMZN_TRACE_ID environment variable', () => { + it.each<{ + description: string; + traceData: string | undefined; + expected: boolean; + }>([ + { + description: + 'returns true if the Sampled flag is set in the _X_AMZN_TRACE_ID environment variable', + traceData: + 'Root=1-5759e988-bd862e3fe1be46a994272793;Parent=557abcec3ee5a047;Sampled=1', + expected: true, + }, + { + description: + 'returns false if the Sampled flag is not set in the _X_AMZN_TRACE_ID environment variable', + traceData: + 'Root=1-5759e988-bd862e3fe1be46a994272793;Parent=557abcec3ee5a047', + expected: false, + }, + { + description: + 'returns false when no _X_AMZN_TRACE_ID environment variable is present', + traceData: undefined, + expected: false, + }, + ])('$description', ({ traceData, expected }) => { // Prepare - process.env._X_AMZN_TRACE_ID = - 'Root=1-5759e988-bd862e3fe1be46a994272793;Parent=557abcec3ee5a047;Sampled=1'; + vi.stubEnv('_X_AMZN_TRACE_ID', traceData); // Act const value = isRequestXRaySampled(); // Assess - expect(value).toEqual(true); + expect(value).toEqual(expected); }); - it('returns false if the Sampled flag is not set in the _X_AMZN_TRACE_ID environment variable', () => { - // Prepare - process.env._X_AMZN_TRACE_ID = - 'Root=1-5759e988-bd862e3fe1be46a994272793;Parent=557abcec3ee5a047'; - - // Act - const value = isRequestXRaySampled(); - - // Assess - expect(value).toEqual(false); - }); - - it('returns false when no _X_AMZN_TRACE_ID environment variable is present', () => { - // Prepare - process.env._X_AMZN_TRACE_ID = undefined; - - // Act - const value = isRequestXRaySampled(); - - // Assess - expect(value).toEqual(false); + it.each<{ + description: string; + traceData: string | undefined; + expected: boolean; + }>([ + { + description: + 'returns true if the Sampled flag is set from async context', + traceData: + 'Root=1-7a5bc3d2-ef456789abcdef012345678;Parent=9f8e7d6c5b4a3210;Sampled=1', + expected: true, + }, + { + description: + 'returns false if the Sampled flag is not set from async context', + traceData: + 'Root=1-8b6cd4e3-fg567890bcdefg123456789;Parent=0g9f8e7d6c5b4321', + expected: false, + }, + { + description: + 'returns false when no trace ID is present in async context', + traceData: undefined, + expected: false, + }, + ])('$description', ({ traceData, expected }) => { + InvokeStore.run( + { + [InvokeStore.PROTECTED_KEYS.X_RAY_TRACE_ID]: traceData, + }, + () => { + // Act + const value = isRequestXRaySampled(); + + // Assess + expect(value).toEqual(expected); + } + ); }); }); });