feat(observe): add response headers and status code (#141)

Signed-off-by: Milan Gallas <Milan.Gallas@ibm.com>
Co-authored-by: Milan Gallas <Milan.Gallas@ibm.com>

src/observe/observe.module.ts

@@ -17,11 +17,21 @@
 import { FastifyPluginAsyncJsonSchemaToTs } from '@fastify/type-provider-json-schema-to-ts';
 
 import { listSpans, getTrace } from './observe.service.js';
-import { traceReadParamsSchema, traceReadQuerySchema } from './dtos/trace-read.js';
-import { spanReadParamsSchema, spanReadQuerySchema } from './dtos/span-read.js';
+import {
+  TraceReadParams,
+  traceReadParamsSchema,
+  TraceReadQuery,
+  traceReadQuerySchema
+} from './dtos/trace-read.js';
+import {
+  SpanReadParams,
+  spanReadParamsSchema,
+  SpanReadQuery,
+  spanReadQuerySchema
+} from './dtos/span-read.js';
 
 export const observeModule: FastifyPluginAsyncJsonSchemaToTs = async (app) => {
-  app.get(
+  app.get<{ Params: TraceReadParams; Querystring: TraceReadQuery }>(
     '/traces/:id',
     {
       preHandler: app.auth(),
@@ -31,12 +41,12 @@ export const observeModule: FastifyPluginAsyncJsonSchemaToTs = async (app) => {
         hide: true
       }
     },
-    async (req) => {
-      return getTrace({ ...req.params, ...req.query });
+    async (req, reply) => {
+      return getTrace(req, reply);
     }
   );
 
-  app.get(
+  app.get<{ Params: SpanReadParams; Querystring: SpanReadQuery }>(
     '/traces/:trace_id/spans',
     {
       preHandler: app.auth(),
@@ -46,8 +56,8 @@ export const observeModule: FastifyPluginAsyncJsonSchemaToTs = async (app) => {
         hide: true
       }
     },
-    async (req) => {
-      return listSpans({ ...req.query, ...req.params });
+    async (req, reply) => {
+      return listSpans(req, reply);
     }
   );
 };

src/observe/observe.service.ts

@@ -14,50 +14,49 @@
  * limitations under the License.
  */
 
+import { FastifyReply, FastifyRequest } from 'fastify';
+
 import { client } from './api/client.js';
 import { SpanReadParams, SpanReadQuery } from './dtos/span-read.js';
 import { assertTracePermission, assertClient, processApiProxyResponse } from './utils.js';
 import { TraceReadParams, TraceReadQuery } from './dtos/trace-read.js';
 
-export async function getTrace({
-  id,
-  include_mlflow,
-  include_mlflow_tree,
-  include_tree
-}: TraceReadParams & TraceReadQuery) {
-  await assertTracePermission({ traceId: id });
+export async function getTrace(
+  req: FastifyRequest<{ Params: TraceReadParams; Querystring: TraceReadQuery }>,
+  reply: FastifyReply
+) {
+  await assertTracePermission({ traceId: req.params.id });
   assertClient(client);
 
   return processApiProxyResponse(
     client.GET('/v1/traces/{id}', {
       params: {
         path: {
-          id
+          id: req.params.id
         },
-        query: {
-          include_mlflow,
-          include_mlflow_tree,
-          include_tree
-        }
+        query: req.query
       }
-    })
+    }),
+    reply
   );
 }
 
-export async function listSpans(props: SpanReadQuery & SpanReadParams) {
-  await assertTracePermission({ traceId: props.trace_id });
+export async function listSpans(
+  req: FastifyRequest<{ Params: SpanReadParams; Querystring: SpanReadQuery }>,
+  reply: FastifyReply
+) {
+  await assertTracePermission({ traceId: req.params.trace_id });
   assertClient(client);
 
-  const { trace_id, ...restQueryObject } = props;
-
   return processApiProxyResponse(
     client.GET('/v1/traces/{trace_id}/spans', {
       params: {
         path: {
-          trace_id: trace_id
+          trace_id: req.params.trace_id
         },
-        query: restQueryObject
+        query: req.query
       }
-    })
+    }),
+    reply
   );
 }

src/observe/utils.ts

@@ -15,6 +15,7 @@
  */
 
 import { FetchResponse } from 'openapi-fetch';
+import { FastifyReply } from 'fastify';
 
 import { Client } from './api/client.js';
 
@@ -42,18 +43,44 @@ export function assertClient(client: Client | undefined): asserts client is Clie
   }
 }
 
+function pickHeaders(headers: Headers, keys: string[]): Record<string, string> {
+  const result: Record<string, string> = {};
+
+  keys.forEach((key) => {
+    if (headers.has(key)) {
+      result[key] = headers.get(key) as string;
+    }
+  });
+
+  return result;
+}
+
+function isAPIErrorCode(value: any): value is APIErrorCode {
+  return Object.values(APIErrorCode).includes(value);
+}
+
 type MediaType = `${string}/${string}`;
 export async function processApiProxyResponse<T, O, M extends MediaType>(
-  response: Promise<FetchResponse<T, O, M>>
+  response: Promise<FetchResponse<T, O, M>>,
+  reply: FastifyReply
 ): Promise<FetchResponse<T, O, M>['data']> {
-  const { data, error } = await response;
+  const {
+    data,
+    error,
+    response: { headers }
+  } = await response;
+
+  // apply only allowed headers from Observe API response
+  reply.headers(pickHeaders(headers, ['Retry-After']));
 
   if (error) {
     getTraceLogger().error({ err: error }, 'Observe API: Invalid response');
+    const errorCode = error.code.toLocaleLowerCase();
+
     throw new APIError(
       {
-        message: 'Observe API: Invalid response',
-        code: APIErrorCode.SERVICE_ERROR
+        message: error.message ?? 'Observe API: Invalid response',
+        code: isAPIErrorCode(errorCode) ? errorCode : APIErrorCode.SERVICE_ERROR
       },
       {
         cause: error