extract serializer and allow passing extra descriptors

Author: gimmyxd

__tests__/directory/v3/serializeResponse.test.ts

@@ -59,7 +59,7 @@ import {
   nullReaderProxy,
   nullWriterProxy,
 } from "./null";
-import { serializeResponse } from "./serializer";
+import { DsRegistry } from "./serializer";
 import {
   CheckRequest,
   CheckResponse,
@@ -133,6 +133,8 @@ export class DirectoryV3 {
   ImporterClient: Client<typeof Importer>;
   ExporterClient: Client<typeof Exporter>;
   ModelClient: Client<typeof Model>;
+  registry: DsRegistry;
+
   CreateTransport: (
     config: ServiceConfig | undefined,
     fallback: ServiceConfig | undefined,
@@ -275,12 +277,12 @@ export class DirectoryV3 {
     this.ExporterClient = !!exporterGrpcTransport
       ? createClient(Exporter, exporterGrpcTransport)
       : nullExporterProxy();
-
     this.ModelClient = !!modelGrpcTransport
       ? createClient(Model, modelGrpcTransport)
       : nullModelProxy();
 
     this.CreateTransport = createTransport;
+    this.registry = new DsRegistry(...(config.additionalDescriptors || []));
   }
 
   async check(
@@ -293,7 +295,7 @@ export class DirectoryV3 {
         options,
       );
 
-      return serializeResponse(response);
+      return this.registry.serializeResponse(response);
     } catch (error) {
       throw handleError(error, "check");
     }
@@ -313,7 +315,7 @@ export class DirectoryV3 {
         options,
       );
 
-      return serializeResponse(response);
+      return this.registry.serializeResponse(response);
     } catch (error) {
       throw handleError(error, "object");
     }
@@ -330,7 +332,7 @@ export class DirectoryV3 {
         create(GetObjectsRequestSchema, params),
         options,
       );
-      return serializeResponse(response);
+      return this.registry.serializeResponse(response);
     } catch (error) {
       throw handleError(error, "objects");
     }
@@ -346,7 +348,7 @@ export class DirectoryV3 {
         options,
       );
 
-      return serializeResponse(response);
+      return this.registry.serializeResponse(response);
     } catch (error) {
       throw handleError(error, "objectMany");
     }
@@ -365,7 +367,7 @@ export class DirectoryV3 {
         options,
       );
 
-      return serializeResponse(response);
+      return this.registry.serializeResponse(response);
     } catch (error) {
       throw handleError(error, "setObject");
     }
@@ -381,7 +383,7 @@ export class DirectoryV3 {
         options,
       );
 
-      return serializeResponse(response);
+      return this.registry.serializeResponse(response);
     } catch (error) {
       throw handleError(error, "deleteObject");
     }
@@ -397,7 +399,7 @@ export class DirectoryV3 {
         options,
       );
 
-      return serializeResponse(response);
+      return this.registry.serializeResponse(response);
     } catch (error) {
       throw handleError(error, "relation");
     }
@@ -416,7 +418,7 @@ export class DirectoryV3 {
         options,
       );
 
-      return serializeResponse(response);
+      return this.registry.serializeResponse(response);
     } catch (error) {
       throw handleError(error, "setRelation");
     }
@@ -432,7 +434,7 @@ export class DirectoryV3 {
         options,
       );
 
-      return serializeResponse(response);
+      return this.registry.serializeResponse(response);
     } catch (error) {
       throw handleError(error, "deleteRelation");
     }
@@ -448,7 +450,7 @@ export class DirectoryV3 {
       }
 
       const response = await this.ReaderClient.getRelations(params, options);
-      return serializeResponse(response);
+      return this.registry.serializeResponse(response);
     } catch (error) {
       throw handleError(error, "relations");
     }
@@ -464,7 +466,7 @@ export class DirectoryV3 {
         options,
       );
 
-      return serializeResponse(response);
+      return this.registry.serializeResponse(response);
     } catch (error) {
       throw handleError(error, "graph");
     }
@@ -527,7 +529,7 @@ export class DirectoryV3 {
       return {
         body,
         updatedAt: metadata?.updatedAt
-          ? serializeResponse(metadata?.updatedAt)
+          ? this.registry.serializeResponse(metadata?.updatedAt)
           : undefined,
         etag: metadata?.etag,
       };
@@ -553,7 +555,7 @@ export class DirectoryV3 {
         options,
       );
 
-      return serializeResponse(response);
+      return this.registry.serializeResponse(response);
     } catch (error) {
       throw handleError(error, "setManifest");
     }
@@ -566,7 +568,7 @@ export class DirectoryV3 {
     try {
       const response = await this.ModelClient.deleteManifest(params!, options);
 
-      return serializeResponse(response);
+      return this.registry.serializeResponse(response);
     } catch (error) {
       throw handleError(error, "deleteManifest");
     }
@@ -600,9 +602,10 @@ export async function readAsyncIterable<T>(
 export async function serializeAsyncIterable<T extends Message>(
   gen: AsyncIterable<T>,
 ): Promise<T[]> {
+  const registry = new DsRegistry();
   const out: T[] = [];
   for await (const x of gen) {
-    out.push(serializeResponse<T>(x));
+    out.push(registry.serializeResponse<T>(x));
   }
   return out;
 }

lib/directory/v3/index.ts

@@ -4,7 +4,14 @@ import { file_aserto_directory_importer_v3_importer } from "@aserto/node-directo
 import { file_aserto_directory_model_v3_model } from "@aserto/node-directory/src/gen/cjs/aserto/directory/model/v3/model_pb";
 import { file_aserto_directory_reader_v3_reader } from "@aserto/node-directory/src/gen/cjs/aserto/directory/reader/v3/reader_pb";
 import { file_aserto_directory_writer_v3_writer } from "@aserto/node-directory/src/gen/cjs/aserto/directory/writer/v3/writer_pb";
-import type { Registry } from "@bufbuild/protobuf";
+import type {
+  DescEnum,
+  DescExtension,
+  DescFile,
+  DescMessage,
+  DescService,
+  Registry,
+} from "@bufbuild/protobuf";
 import {
   createRegistry,
   Message,
@@ -16,31 +23,45 @@ import { file_google_protobuf_timestamp } from "@bufbuild/protobuf/wkt";
 
 import { InvalidSchemaError } from "../../errors";
 
-const registry: Registry = createRegistry(
-  file_aserto_directory_common_v3_common,
-  file_aserto_directory_reader_v3_reader,
-  file_aserto_directory_writer_v3_writer,
-  file_aserto_directory_exporter_v3_exporter,
-  file_aserto_directory_importer_v3_importer,
-  file_aserto_directory_model_v3_model,
-  file_google_protobuf_timestamp,
-);
+class DsRegistry {
+  registry: Registry;
 
-/**
- * Serializes a response message into JSON format based on its schema.
- *
- * @param response - The message object to be serialized, adhering to MessageShape.
- * @returns The serialized message in JSON format.
- * @throws InvalidSchemaError if the schema for the given message type is not registered.
- */
-export function serializeResponse<T extends Message>(
-  response: MessageShape<GenMessage<T>>,
-): T {
-  const schema = registry.getMessage(response.$typeName);
-  if (!schema) {
-    throw new InvalidSchemaError(
-      `schema not registered for type: [${response.$typeName}]`,
+  constructor(
+    ...input: (
+      | Registry
+      | DescFile
+      | DescMessage
+      | DescEnum
+      | DescExtension
+      | DescService
+    )[]
+  ) {
+    this.registry = createRegistry(
+      file_aserto_directory_common_v3_common,
+      file_aserto_directory_reader_v3_reader,
+      file_aserto_directory_writer_v3_writer,
+      file_aserto_directory_exporter_v3_exporter,
+      file_aserto_directory_importer_v3_importer,
+      file_aserto_directory_model_v3_model,
+      file_google_protobuf_timestamp,
+      ...input,
     );
   }
-  return toJson(schema, response, { alwaysEmitImplicit: true }) as unknown as T;
+
+  serializeResponse<T extends Message>(
+    response: MessageShape<GenMessage<T>>,
+  ): T {
+    const schema = this.registry.getMessage(response.$typeName);
+    if (!schema) {
+      throw new InvalidSchemaError(
+        `schema not registered for type: [${response.$typeName}]`,
+      );
+    }
+    return toJson(schema, response, {
+      alwaysEmitImplicit: true,
+      registry: this.registry,
+    }) as unknown as T;
+  }
 }
+
+export { DsRegistry };

lib/directory/v3/serializer.ts

@@ -43,6 +43,14 @@ import {
   SetRelationRequest as SetRelationRequest$,
   SetRelationResponse as SetRelationResponse$,
 } from "@aserto/node-directory/src/gen/cjs/aserto/directory/writer/v3/writer_pb";
+import {
+  DescEnum,
+  DescExtension,
+  DescFile,
+  DescMessage,
+  DescService,
+  Registry,
+} from "@bufbuild/protobuf";
 import { Timestamp } from "@bufbuild/protobuf/wkt";
 
 import { NestedOmit, NestedOptional, Optional } from "../../util/types";
@@ -88,7 +96,17 @@ export type DirectoryV3Config = ServiceConfig & {
   importer?: ServiceConfig;
   exporter?: ServiceConfig;
   model?: ServiceConfig;
+} & {
+  additionalDescriptors?: (
+    | Registry
+    | DescFile
+    | DescMessage
+    | DescEnum
+    | DescExtension
+    | DescService
+  )[];
 };
+
 export type Object$ = Optional<
   Omit<Object$$, "$typeName">,
   "etag" | "displayName"

lib/directory/v3/types.ts

@@ -53,8 +53,8 @@ export {
   policyInstance,
   PolicyPathMapper,
   queryOptions,
-  serializeAsyncIterable,
   readAsyncIterable,
+  serializeAsyncIterable,
   SubIdentityMapper,
 };