diff --git a/packages/commons/src/middleware/cleanupMiddlewares.ts b/packages/commons/src/middleware/cleanupMiddlewares.ts new file mode 100644 index 0000000000..e35719b9d7 --- /dev/null +++ b/packages/commons/src/middleware/cleanupMiddlewares.ts @@ -0,0 +1,77 @@ +import { + TRACER_KEY, + METRICS_KEY, + LOGGER_KEY, + IDEMPOTENCY_KEY, +} from './constants'; +import type { MiddyLikeRequest, CleanupFunction } from '../types/middy'; + +// Typeguard to assert that an object is of Function type +const isFunction = (obj: unknown): obj is CleanupFunction => { + return typeof obj === 'function'; +}; + +/** + * Function used to cleanup Powertools for AWS resources when a Middy + * middleware [returns early](https://middy.js.org/docs/intro/early-interrupt) + * and terminates the middleware chain. + * + * When a middleware returns early, all the middleware lifecycle functions + * that come after it are not executed. This means that if a middleware + * was relying on certain logic to be run during the `after` or `onError` + * lifecycle functions, that logic will not be executed. + * + * This is the case for the middlewares that are part of Powertools for AWS + * which rely on these lifecycle functions to perform cleanup operations + * like closing the current segment in the tracer or flushing any stored + * metrics. + * + * When authoring a middleware that might return early, you can use this + * function to cleanup Powertools resources. This function will check if + * any cleanup function is present in the `request.internal` object and + * execute it. + * + * @example + * ```typescript + * import middy from '@middy/core'; + * import { cleanupMiddlewares } from '@aws-lambda-powertools/commons/lib/middleware'; + * + * // Example middleware that returns early + * const myCustomMiddleware = (): middy.MiddlewareObj => { + * const before = async (request: middy.Request): Promise => { + * // If the request is a GET, return early (as an example) + * if (request.event.httpMethod === 'GET') { + * // Cleanup Powertools resources + * await cleanupMiddlewares(request); + * // Then return early + * return 'GET method not supported'; + * } + * }; + * + * return { + * before, + * }; + * }; + * ``` + * + * @param request - The Middy request object + * @param options - An optional object that can be used to pass options to the function + */ +const cleanupMiddlewares = async (request: MiddyLikeRequest): Promise => { + const cleanupFunctionNames = [ + TRACER_KEY, + METRICS_KEY, + LOGGER_KEY, + IDEMPOTENCY_KEY, + ]; + for (const functionName of cleanupFunctionNames) { + if (Object(request.internal).hasOwnProperty(functionName)) { + const functionReference = request.internal[functionName]; + if (isFunction(functionReference)) { + await functionReference(request); + } + } + } +}; + +export { cleanupMiddlewares }; diff --git a/packages/commons/src/middleware/constants.ts b/packages/commons/src/middleware/constants.ts new file mode 100644 index 0000000000..9f9b30bcb6 --- /dev/null +++ b/packages/commons/src/middleware/constants.ts @@ -0,0 +1,12 @@ +/** + * These constants are used to store cleanup functions in Middy's `request.internal` object. + * They are used by the `cleanupPowertools` function to check if any cleanup function + * is present and execute it. + */ +const PREFIX = 'powertools-for-aws'; +const TRACER_KEY = `${PREFIX}.tracer`; +const METRICS_KEY = `${PREFIX}.metrics`; +const LOGGER_KEY = `${PREFIX}.logger`; +const IDEMPOTENCY_KEY = `${PREFIX}.idempotency`; + +export { TRACER_KEY, METRICS_KEY, LOGGER_KEY, IDEMPOTENCY_KEY }; diff --git a/packages/commons/src/middleware/index.ts b/packages/commons/src/middleware/index.ts new file mode 100644 index 0000000000..85f7388af3 --- /dev/null +++ b/packages/commons/src/middleware/index.ts @@ -0,0 +1,2 @@ +export * from './cleanupMiddlewares'; +export * from './constants'; diff --git a/packages/commons/src/types/middy.ts b/packages/commons/src/types/middy.ts index fe060ece0a..79a8cb5c3d 100644 --- a/packages/commons/src/types/middy.ts +++ b/packages/commons/src/types/middy.ts @@ -1,4 +1,4 @@ -import { Context } from 'aws-lambda'; +import type { Context } from 'aws-lambda'; /** * We need to define these types and interfaces here because we can't import them from @middy/core. @@ -22,14 +22,14 @@ type Request< }; }; -declare type MiddlewareFn< +type MiddlewareFn< TEvent = unknown, TResult = unknown, TErr = Error, TContext extends Context = Context > = (request: Request) => unknown; -export type MiddlewareLikeObj< +type MiddlewareLikeObj< TEvent = unknown, TResult = unknown, TErr = Error, @@ -40,9 +40,21 @@ export type MiddlewareLikeObj< onError?: MiddlewareFn; }; -export type MiddyLikeRequest = { +type MiddyLikeRequest = { event: unknown; context: Context; response: unknown | null; error: Error | null; + internal: { + [key: string]: unknown; + }; }; + +/** + * Cleanup function that is used to cleanup resources when a middleware returns early. + * Each Powertools for AWS middleware that needs to perform cleanup operations will + * store a cleanup function with this signature in the `request.internal` object. + */ +type CleanupFunction = (request: MiddyLikeRequest) => Promise; + +export { MiddlewareLikeObj, MiddyLikeRequest, CleanupFunction }; diff --git a/packages/commons/tests/unit/cleanupMiddlewares.test.ts b/packages/commons/tests/unit/cleanupMiddlewares.test.ts new file mode 100644 index 0000000000..12583d9107 --- /dev/null +++ b/packages/commons/tests/unit/cleanupMiddlewares.test.ts @@ -0,0 +1,66 @@ +/** + * Test Middy cleanupMiddlewares function + * + * @group unit/commons/cleanupMiddlewares + */ +import { + cleanupMiddlewares, + TRACER_KEY, + METRICS_KEY, +} from '../../src/middleware'; +import { helloworldContext as context } from '../../src/samples/resources/contexts/hello-world'; + +describe('Function: cleanupMiddlewares', () => { + it('calls the cleanup function that are present', async () => { + // Prepare + const mockCleanupFunction1 = jest.fn(); + const mockCleanupFunction2 = jest.fn(); + const mockRequest = { + event: {}, + context: context, + response: null, + error: null, + internal: { + [TRACER_KEY]: mockCleanupFunction1, + [METRICS_KEY]: mockCleanupFunction2, + }, + }; + + // Act + await cleanupMiddlewares(mockRequest); + + // Assess + expect(mockCleanupFunction1).toHaveBeenCalledTimes(1); + expect(mockCleanupFunction1).toHaveBeenCalledWith(mockRequest); + expect(mockCleanupFunction2).toHaveBeenCalledTimes(1); + expect(mockCleanupFunction2).toHaveBeenCalledWith(mockRequest); + }); + it('resolves successfully if no cleanup function is present', async () => { + // Prepare + const mockRequest = { + event: {}, + context: context, + response: null, + error: null, + internal: {}, + }; + + // Act & Assess + await expect(cleanupMiddlewares(mockRequest)).resolves.toBeUndefined(); + }); + it('resolves successfully if cleanup function is not a function', async () => { + // Prepare + const mockRequest = { + event: {}, + context: context, + response: null, + error: null, + internal: { + [TRACER_KEY]: 'not a function', + }, + }; + + // Act & Assess + await expect(cleanupMiddlewares(mockRequest)).resolves.toBeUndefined(); + }); +});