From c0dea3307d29cff6692b4485cef824b0018a53b1 Mon Sep 17 00:00:00 2001
From: Ben Gotow <ben@foundry376.com>
Date: Fri, 26 Jul 2024 17:05:28 -0500
Subject: [PATCH] Add an escape hatch you can use to make arbitrary
 customizations to OpenAPI specs (#3)

---
 .vscode/settings.json |  5 +++--
 src/openAPI.ts        | 25 +++++++++++++++++++++----
 src/openAPIRoute.ts   |  6 ++++++
 3 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/.vscode/settings.json b/.vscode/settings.json
index 1589eb1..1183802 100644
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@@ -1,6 +1,7 @@
 {
   "editor.formatOnSave": true,
   "editor.codeActionsOnSave": {
-    "source.organizeImports": true
-  }
+    "source.organizeImports": "explicit"
+  },
+  "typescript.tsdk": "node_modules/typescript/lib"
 }
diff --git a/src/openAPI.ts b/src/openAPI.ts
index c91b438..a7c5684 100644
--- a/src/openAPI.ts
+++ b/src/openAPI.ts
@@ -3,6 +3,7 @@ import {
   OpenAPIGenerator,
   OpenAPIRegistry,
   ResponseConfig,
+  RouteConfig,
 } from "@asteasolutions/zod-to-openapi";
 import { RequestHandler, Router } from "express";
 import { z, ZodArray, ZodEffects, ZodObject } from "zod";
@@ -63,8 +64,19 @@ export function buildOpenAPIDocument(args: {
   // Attach all the API routes, referencing the named components where
   // possible, and falling back to inlining the Zod shapes.
   getRoutes(routers).forEach(({ path, method, handler }) => {
-    const { tag, body, params, query, response, description, summary, security, deprecated, responseContentType } =
-      getSchemaOfOpenAPIRoute(handler) || {};
+    const {
+      tag,
+      body,
+      params,
+      query,
+      response,
+      description,
+      summary,
+      security,
+      deprecated,
+      finalizeRouteConfig,
+      responseContentType,
+    } = getSchemaOfOpenAPIRoute(handler) || {};
 
     //Express: /path/to/:variable/something -> OpenAPI /path/to/{variable}/something
     const pathOpenAPIFormat = path
@@ -124,7 +136,7 @@ export function buildOpenAPIDocument(args: {
     } else {
       responses[204] = z.void().openapi({ description: "No content - successful operation" });
     }
-    registry.registerPath({
+    let openapiRouteConfig: RouteConfig = {
       tags: [tag || "default"],
       method: method,
       summary: summary,
@@ -138,7 +150,12 @@ export function buildOpenAPIDocument(args: {
         body: referencingNamedSchemas(body),
       },
       responses: responses,
-    });
+    };
+    if (finalizeRouteConfig) {
+      openapiRouteConfig = finalizeRouteConfig(openapiRouteConfig);
+    }
+
+    registry.registerPath(openapiRouteConfig);
   });
 
   const generator = new OpenAPIGenerator(registry.definitions);
diff --git a/src/openAPIRoute.ts b/src/openAPIRoute.ts
index 58c9b99..06401d5 100644
--- a/src/openAPIRoute.ts
+++ b/src/openAPIRoute.ts
@@ -1,4 +1,5 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
+import { RouteConfig } from "@asteasolutions/zod-to-openapi";
 import { NextFunction, Request, RequestHandler, Response } from "express";
 import { ZodError, ZodSchema, ZodTypeAny, z } from "zod";
 import { ErrorResponse } from "./schemas";
@@ -46,6 +47,11 @@ type SchemaDefinition<
   responseContentType?: string;
   /** Mark the route as deprecated in generated OpenAPI docs. Does not have any impact on routing. */
   deprecated?: boolean;
+
+  /** Provide a function to apply additional post-processing to the OpenAPI route configuration generated
+   * based on your API handler. This is the last function to run before the OpenAPI route is added to the registry.
+   */
+  finalizeRouteConfig?: (config: RouteConfig) => RouteConfig;
 };
 
 const check = <TType>(obj?: any, schema?: ZodSchema<TType>): z.SafeParseReturnType<TType, TType> => {