[kbn-server-route-repository] Allow custom response (#182679)

This PR changes `registerRoutes` to:
1. Pass the `response` object into the `handler` inside of the
`wrappedHandler` to give access to the route handler to create Kibana
responses
2. Check if the result of the `handler` is already a Kibana response and
only wrap it into a Kibana response if it isn't.

---------

Co-authored-by: kibanamachine <42973632+kibanamachine@users.noreply.github.com>

packages/core/http/core-http-server/src/router/response.test.ts

@@ -82,4 +82,8 @@ describe('isKibanaResponse', () => {
       })
     ).toEqual(false);
   });
+
+  it('handles undefined inputs', () => {
+    expect(isKibanaResponse(undefined)).toEqual(false);
+  });
 });

packages/core/http/core-http-server/src/router/response.ts

@@ -55,7 +55,11 @@ export interface IKibanaResponse<T extends HttpResponsePayload | ResponseError =
   readonly options: HttpResponseOptions;
 }
 
-export function isKibanaResponse(response: Record<string, any>): response is IKibanaResponse {
+export function isKibanaResponse(response?: Record<string, any>): response is IKibanaResponse {
+  if (!response) {
+    return false;
+  }
+
   const { status, options, payload, ...rest } = response;
 
   if (Object.keys(rest).length !== 0) {

packages/kbn-server-route-repository/src/register_routes.test.ts

@@ -0,0 +1,209 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+import * as t from 'io-ts';
+import { CoreSetup, kibanaResponseFactory } from '@kbn/core/server';
+import { loggerMock } from '@kbn/logging-mocks';
+import { registerRoutes } from './register_routes';
+import { routeValidationObject } from './route_validation_object';
+import { NEVER } from 'rxjs';
+
+describe('registerRoutes', () => {
+  const get = jest.fn();
+
+  const getAddVersion = jest.fn();
+  const getWithVersion = jest.fn((_options) => {
+    return {
+      addVersion: getAddVersion,
+    };
+  });
+
+  const createRouter = jest.fn().mockReturnValue({
+    get,
+    versioned: {
+      get: getWithVersion,
+    },
+  });
+
+  const internalOptions = {
+    internal: true,
+  };
+  const publicOptions = {
+    public: true,
+  };
+
+  const internalHandler = jest.fn().mockResolvedValue('internal');
+  const publicHandler = jest
+    .fn()
+    .mockResolvedValue(
+      kibanaResponseFactory.custom({ statusCode: 201, body: { message: 'public' } })
+    );
+  const errorHandler = jest.fn().mockRejectedValue(new Error('error'));
+
+  const mockLogger = loggerMock.create();
+  const mockService = jest.fn();
+
+  const mockContext = {};
+  const mockRequest = {
+    body: {
+      bodyParam: 'body',
+    },
+    query: {
+      queryParam: 'query',
+    },
+    params: {
+      pathParam: 'path',
+    },
+    events: {
+      aborted$: NEVER,
+    },
+  };
+
+  beforeEach(() => {
+    jest.clearAllMocks();
+
+    const coreSetup = {
+      http: {
+        createRouter,
+      },
+    } as unknown as CoreSetup;
+
+    const paramsRt = t.type({
+      body: t.type({
+        bodyParam: t.string,
+      }),
+      query: t.type({
+        queryParam: t.string,
+      }),
+      path: t.type({
+        pathParam: t.string,
+      }),
+    });
+
+    registerRoutes({
+      core: coreSetup,
+      repository: {
+        internal: {
+          endpoint: 'GET /internal/app/feature',
+          handler: internalHandler,
+          params: paramsRt,
+          options: internalOptions,
+        },
+        public: {
+          endpoint: 'GET /api/app/feature version',
+          handler: publicHandler,
+          params: paramsRt,
+          options: publicOptions,
+        },
+        error: {
+          endpoint: 'GET /internal/app/feature/error',
+          handler: errorHandler,
+          params: paramsRt,
+          options: internalOptions,
+        },
+      },
+      dependencies: {
+        aService: mockService,
+      },
+      logger: mockLogger,
+    });
+  });
+
+  it('creates a router and defines the routes', () => {
+    expect(createRouter).toHaveBeenCalledTimes(1);
+
+    expect(get).toHaveBeenCalledTimes(2);
+
+    const [internalRoute] = get.mock.calls[0];
+    expect(internalRoute.path).toEqual('/internal/app/feature');
+    expect(internalRoute.options).toEqual(internalOptions);
+    expect(internalRoute.validate).toEqual(routeValidationObject);
+
+    const [errorRoute] = get.mock.calls[1];
+    expect(errorRoute.path).toEqual('/internal/app/feature/error');
+    expect(errorRoute.options).toEqual(internalOptions);
+    expect(errorRoute.validate).toEqual(routeValidationObject);
+
+    expect(getWithVersion).toHaveBeenCalledTimes(1);
+    const [publicRoute] = getWithVersion.mock.calls[0];
+    expect(publicRoute.path).toEqual('/api/app/feature');
+    expect(publicRoute.options).toEqual(publicOptions);
+    expect(publicRoute.access).toEqual('public');
+
+    expect(getAddVersion).toHaveBeenCalledTimes(1);
+    const [versionedRoute] = getAddVersion.mock.calls[0];
+    expect(versionedRoute.version).toEqual('version');
+    expect(versionedRoute.validate).toEqual({
+      request: routeValidationObject,
+    });
+  });
+
+  it('calls the route handler with all dependencies', async () => {
+    const [_, internalRouteHandler] = get.mock.calls[0];
+    await internalRouteHandler(mockContext, mockRequest, kibanaResponseFactory);
+
+    const [args] = internalHandler.mock.calls[0];
+    expect(Object.keys(args).sort()).toEqual(
+      ['aService', 'request', 'response', 'context', 'params', 'logger'].sort()
+    );
+
+    const { params, logger, aService, request, response, context } = args;
+    expect(params).toEqual({
+      body: {
+        bodyParam: 'body',
+      },
+      query: {
+        queryParam: 'query',
+      },
+      path: {
+        pathParam: 'path',
+      },
+    });
+    expect(request).toBe(mockRequest);
+    expect(response).toBe(kibanaResponseFactory);
+    expect(context).toBe(mockContext);
+    expect(aService).toBe(mockService);
+    expect(logger).toBe(mockLogger);
+  });
+
+  it('wraps a plain route handler result into a response', async () => {
+    const [_, internalRouteHandler] = get.mock.calls[0];
+    const internalResult = await internalRouteHandler(
+      mockContext,
+      mockRequest,
+      kibanaResponseFactory
+    );
+
+    expect(internalHandler).toHaveBeenCalledTimes(1);
+    expect(internalResult).toEqual({
+      status: 200,
+      payload: 'internal',
+      options: { body: 'internal' },
+    });
+  });
+
+  it('allows for route handlers to define a custom response', async () => {
+    const [_, publicRouteHandler] = getAddVersion.mock.calls[0];
+    const publicResult = await publicRouteHandler(mockContext, mockRequest, kibanaResponseFactory);
+
+    expect(publicHandler).toHaveBeenCalledTimes(1);
+    expect(publicResult).toEqual({ status: 201, payload: { message: 'public' }, options: {} });
+  });
+
+  it('translates errors thrown in a route handler to an error response', async () => {
+    const [_, errorRouteHandler] = get.mock.calls[1];
+    const errorResult = await errorRouteHandler(mockContext, mockRequest, kibanaResponseFactory);
+
+    expect(errorHandler).toHaveBeenCalledTimes(1);
+    expect(errorResult).toEqual({
+      status: 500,
+      payload: { message: 'error', attributes: { data: {} } },
+      options: {},
+    });
+  });
+});

packages/kbn-server-route-repository/src/register_routes.ts

@@ -9,6 +9,7 @@ import { errors } from '@elastic/elasticsearch';
 import { isBoom } from '@hapi/boom';
 import type { RequestHandlerContext } from '@kbn/core-http-request-handler-context-server';
 import type { KibanaRequest, KibanaResponseFactory } from '@kbn/core-http-server';
+import { isKibanaResponse } from '@kbn/core-http-server';
 import type { CoreSetup } from '@kbn/core-lifecycle-server';
 import type { Logger } from '@kbn/logging';
 import * as t from 'io-ts';
@@ -58,23 +59,24 @@ export function registerRoutes({
           runtimeType
         );
 
-        const { aborted, data } = await Promise.race([
+        const { aborted, result } = await Promise.race([
           handler({
             request,
+            response,
             context,
             params: validatedParams,
             logger,
             ...dependencies,
           }).then((value) => {
             return {
               aborted: false,
-              data: value,
+              result: value,
             };
           }),
           request.events.aborted$.toPromise().then(() => {
             return {
               aborted: true,
-              data: undefined,
+              result: undefined,
             };
           }),
         ]);
@@ -83,9 +85,12 @@ export function registerRoutes({
           return response.custom(CLIENT_CLOSED_REQUEST);
         }
 
-        const body = data || {};
-
-        return response.ok({ body });
+        if (isKibanaResponse(result)) {
+          return result;
+        } else {
+          const body = result || {};
+          return response.ok({ body });
+        }
       } catch (error) {
         logger.error(error);
 

packages/kbn-server-route-repository/src/test_types.ts

@@ -6,6 +6,7 @@
  * Side Public License, v 1.
  */
 import * as t from 'io-ts';
+import { kibanaResponseFactory } from '@kbn/core/server';
 import { createServerRouteFactory } from './create_server_route_factory';
 import { decodeRequestParams } from './decode_request_params';
 import { EndpointOf, ReturnOf, RouteRepositoryClient } from './typings';
@@ -124,6 +125,24 @@ const repository = {
       };
     },
   }),
+  ...createServerRoute({
+    endpoint: 'GET /internal/endpoint_returning_result',
+    handler: async () => {
+      return {
+        result: true,
+      };
+    },
+  }),
+  ...createServerRoute({
+    endpoint: 'GET /internal/endpoint_returning_kibana_response',
+    handler: async () => {
+      return kibanaResponseFactory.ok({
+        body: {
+          result: true,
+        },
+      });
+    },
+  }),
 };
 
 type TestRepository = typeof repository;
@@ -150,6 +169,14 @@ const noParamsInvalid: ReturnOf<TestRepository, 'GET /internal/endpoint_without_
   paramsForMe: true,
 };
 
+assertType<ReturnOf<TestRepository, 'GET /internal/endpoint_returning_result'>>({
+  result: true,
+});
+
+assertType<ReturnOf<TestRepository, 'GET /internal/endpoint_returning_kibana_response'>>({
+  result: true,
+});
+
 // RouteRepositoryClient
 
 type TestClient = RouteRepositoryClient<TestRepository, { timeout: number }>;
@@ -214,6 +241,22 @@ client('GET /internal/endpoint_with_params', {
   }>(res);
 });
 
+client('GET /internal/endpoint_returning_result', {
+  timeout: 1,
+}).then((res) => {
+  assertType<{
+    result: boolean;
+  }>(res);
+});
+
+client('GET /internal/endpoint_returning_kibana_response', {
+  timeout: 1,
+}).then((res) => {
+  assertType<{
+    result: boolean;
+  }>(res);
+});
+
 // decodeRequestParams should return the type of the codec that is passed
 assertType<{ path: { serviceName: string } }>(
   decodeRequestParams(

packages/kbn-server-route-repository/src/typings.ts

@@ -6,6 +6,7 @@
  * Side Public License, v 1.
  */
 
+import { IKibanaResponse } from '@kbn/core-http-server';
 import * as t from 'io-ts';
 import { RequiredKeys } from 'utility-types';
 
@@ -94,7 +95,9 @@ export type ReturnOf<
   infer TReturnType,
   ServerRouteCreateOptions
 >
-  ? TReturnType
+  ? TReturnType extends IKibanaResponse<infer TWrappedResponseType>
+    ? TWrappedResponseType
+    : TReturnType
   : never;
 
 export type DecodedRequestParamsOf<

packages/kbn-server-route-repository/tsconfig.json

@@ -18,6 +18,8 @@
     "@kbn/core-http-server",
     "@kbn/core-lifecycle-server",
     "@kbn/logging",
+    "@kbn/core",
+    "@kbn/logging-mocks",
   ],
   "exclude": [
     "target/**/*",