Implement support for Protocol Buffer references (imports)

bin/avdlToAVSC.sh

@@ -1,3 +1,5 @@
+#!/bin/sh
+
 avdl_path=$1
 avsc_name=$2
 

src/@types.ts

@@ -11,6 +11,7 @@ export enum SchemaType {
 export interface SchemaHelper {
   validate(schema: Schema): void
   getSubject(confluentSchema: ConfluentSchema, schema: Schema, separator: string): ConfluentSubject
+  referencedSchemas(schema: string): Promise<string[]>
 }
 
 export type AvroOptions = Partial<ForSchemaOptions>

src/AvroHelper.ts

@@ -44,4 +44,11 @@ export default class AvroHelper implements SchemaHelper {
     }
     return subject
   }
+
+  public async referencedSchemas(
+    // eslint-disable-next-line @typescript-eslint/no-unused-vars
+    _schema: string,
+  ): Promise<string[]> {
+    return []
+  }
 }

src/JsonHelper.ts

@@ -14,4 +14,11 @@ export default class JsonHelper implements SchemaHelper {
   ): ConfluentSubject {
     throw new ConfluentSchemaRegistryError('not implemented yet')
   }
+
+  public async referencedSchemas(
+    // eslint-disable-next-line @typescript-eslint/no-unused-vars
+    _schema: string,
+  ): Promise<string[]> {
+    return []
+  }
 }

src/ProtoHelper.ts

@@ -1,17 +1,34 @@
-// @ts-nocheck
 import { Schema, SchemaHelper, ConfluentSubject, ConfluentSchema } from './@types'
 import { ConfluentSchemaRegistryError } from './errors'
+import { parse } from 'protobufjs'
 
 export default class ProtoHelper implements SchemaHelper {
-  public validate(schema: Schema): void {
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars
+  public validate(_schema: Schema): void {
     return
   }
 
   public getSubject(
-    confluentSchema: ConfluentSchema,
-    schema: Schema,
-    separator: string,
+    // eslint-disable-next-line @typescript-eslint/no-unused-vars
+    _confluentSchema: ConfluentSchema,
+    // eslint-disable-next-line @typescript-eslint/no-unused-vars
+    _schema: Schema,
+    // eslint-disable-next-line @typescript-eslint/no-unused-vars
+    _separator: string,
   ): ConfluentSubject {
     throw new ConfluentSchemaRegistryError('not implemented yet')
   }
+
+  /**
+   * Get the schemas referenced by the provided schema.
+   *
+   * @param schema The schema to find references for
+   *
+   * @returns A list of imported/referenced schemas that is used by the given schema.
+   */
+  public async referencedSchemas(schema: string): Promise<string[]> {
+    const parsed = parse(schema)
+    const out: string[] = []
+    return out.concat(parsed.imports || []).concat(parsed.weakImports || [])
+  }
 }

src/ProtoSchema.ts

@@ -8,11 +8,20 @@ import {
 
 export default class ProtoSchema implements Schema {
   private message: Type
+  private root: protobuf.Root
 
-  constructor(schema: ConfluentSchema, opts?: ProtoOptions) {
+  constructor(schema: ConfluentSchema, opts?: ProtoOptions, references?: Schema[]) {
     const parsedMessage = protobuf.parse(schema.schema)
-    const root = parsedMessage.root
-    this.message = root.lookupType(this.getTypeName(parsedMessage, opts))
+    this.root = parsedMessage.root
+
+    if (references) {
+      const schemas = references as ProtoSchema[]
+      schemas.forEach(reference => {
+        this.root.add(reference.root)
+      })
+    }
+
+    this.message = this.root.lookupType(this.getTypeName(parsedMessage, opts))
   }
 
   private getNestedTypeName(parent: { [k: string]: ReflectionObject } | undefined): string {

src/SchemaRegistry.newApi.spec.ts

@@ -15,6 +15,7 @@ import encodedAnotherPersonV2Json from '../fixtures/json/encodedAnotherPersonV2'
 import encodedAnotherPersonV2Proto from '../fixtures/proto/encodedAnotherPersonV2'
 import encodedNestedV2Proto from '../fixtures/proto/encodedNestedV2'
 import wrongMagicByte from '../fixtures/wrongMagicByte'
+import ProtoSchema from './ProtoSchema'
 
 const REGISTRY_HOST = 'http://localhost:8982'
 const schemaRegistryAPIClientArgs = { host: REGISTRY_HOST }
@@ -456,6 +457,41 @@ describe('SchemaRegistry - new Api', () => {
       v3Opts = { [SchemaType.PROTOBUF]: { messageName: 'AnotherPerson' } },
       type = SchemaType.PROTOBUF
 
+    const protoV3 = `
+      syntax = "proto3";
+      package com.org.domain.fixtures;
+      message AnotherPerson {
+        string city = 2 [default = "Stockholm"];
+      }
+      `
+
+    const protoContact = `
+      syntax = "proto3";
+      package com.org.domain.fixtures;
+      message ContactMessage {
+        string email = 1;
+      }
+      `
+
+    const protoEmployee = `
+      syntax = "proto3";
+      package com.org.domain.fixtures;
+      import "Contact.proto";
+      message Employee {
+        string city = 1 [default = "Stockholm"];
+        ContactMessage contact = 2;
+      }
+      `
+
+    const protoCompany = `
+      syntax = "proto3";
+      package com.org.domain.fixtures;
+      import "Employee.proto";
+      message Company {
+        Employee employee = 1;
+      }
+      `
+
     beforeAll(() => {
       schemaRegistry = new SchemaRegistry(schemaRegistryArgs, v3Opts)
     })
@@ -544,6 +580,61 @@ describe('SchemaRegistry - new Api', () => {
         })
       })
 
+      it('register and encode protocol buffer v3 schema', async () => {
+        const confluentSchemaV3: ConfluentSchema = {
+          type,
+          schema: protoV3,
+        }
+
+        const schema1 = await schemaRegistry.register(confluentSchemaV3, {
+          subject: `${type}_test_protoV3-value`,
+        })
+
+        await schemaRegistry.encode(schema1.id, payload)
+      })
+
+      it('register and encode protocol buffer v3 schema with import', async () => {
+        const confluentSchemaV3: ConfluentSchema = {
+          type,
+          schema: protoCompany,
+        }
+
+        async function fetchSchema(referenceName: string): Promise<ConfluentSchema> {
+          if (referenceName === 'Contact.proto') {
+            return { type: SchemaType.PROTOBUF, schema: protoContact }
+          }
+          if (referenceName === 'Employee.proto') {
+            return { type: SchemaType.PROTOBUF, schema: protoEmployee }
+          }
+          throw `unknown reference ${referenceName}`
+        }
+
+        const schema1 = await schemaRegistry.register(confluentSchemaV3, {
+          subject: `${type}_test_protoImportsV3-value`,
+          fetchSchema,
+        })
+
+        // Check that we can encode with the cached version from the register() call
+        const payload = { employee: { contact: { email: '[email protected]' } } }
+        const encoded1 = await schemaRegistry.encode(schema1.id, payload)
+        const decoded1 = await schemaRegistry.decode(encoded1)
+
+        // Clear the cache and try again to exercise getSchema()
+        schemaRegistry.cache.clear()
+        const encoded2 = await schemaRegistry.encode(schema1.id, payload)
+        schemaRegistry.cache.clear()
+        const decoded2 = await schemaRegistry.decode(encoded2)
+
+        expect(encoded1).toEqual(encoded2)
+        expect(decoded1).toEqual(decoded2)
+
+        // Check the default value
+        expect(decoded1.employee.city).toEqual('Stockholm')
+
+        // Check the value in the field defined in imported schema
+        expect(decoded1.employee.contact.email).toEqual('[email protected]')
+      })
+
       it('decodes', async () => {
         const confluentSchemaV4: ConfluentSchema = {
           type,

src/SchemaRegistry.ts

@@ -20,6 +20,7 @@ import {
   ConfluentSubject,
   SchemaRegistryAPIClientOptions,
   AvroConfluentSchema,
+  SchemaHelper,
 } from './@types'
 import {
   helperTypeFromSchemaType,
@@ -34,6 +35,7 @@ interface RegisteredSchema {
 interface Opts {
   compatibility?: COMPATIBILITY
   separator?: string
+  fetchSchema?: (referenceName: string) => Promise<ConfluentSchema>
   subject: string
 }
 
@@ -81,6 +83,35 @@ export default class SchemaRegistry {
     return confluentSchema
   }
 
+  async registerReferences(
+    opts: Omit<Opts, 'subject'> & { subject?: string },
+    helper: SchemaHelper,
+    confluentSchema: ConfluentSchema,
+  ): Promise<{ subject: string; name: string; version: number }[]> {
+    const fetchSchema = opts.fetchSchema
+    if (!fetchSchema) {
+      return []
+    }
+
+    const subjects = await helper.referencedSchemas(confluentSchema.schema)
+
+    return await Promise.all(
+      subjects.map(async subject => {
+        const schema = await fetchSchema(subject)
+        const registeredSchema = await this.register(schema, { ...opts, subject })
+        const response = await this.api.Schema.versions({ id: registeredSchema.id })
+        const allVersions: [{ subject: string; version: number }] = response.data()
+
+        const subjectVersion = allVersions
+          .filter(it => it.subject === subject)
+          .map(it => it.version)
+          .reduce((a, b) => Math.max(a, b))
+
+        return { subject, name: subject, version: subjectVersion }
+      }),
+    )
+  }
+
   public async register(
     schema: Exclude<ConfluentSchema, AvroConfluentSchema>,
     userOpts: Opts,
@@ -97,12 +128,26 @@ export default class SchemaRegistry {
     schema: RawAvroSchema | ConfluentSchema,
     userOpts?: Opts,
   ): Promise<RegisteredSchema> {
-    const { compatibility, separator } = { ...DEFAULT_OPTS, ...userOpts }
+    const opts = { ...DEFAULT_OPTS, ...userOpts }
+    const { compatibility, separator } = opts
 
     const confluentSchema: ConfluentSchema = this.getConfluentSchema(schema)
-
     const helper = helperTypeFromSchemaType(confluentSchema.type)
-    const schemaInstance = schemaFromConfluentSchema(confluentSchema, this.options)
+
+    // If this schema depends on other schemas, we need to make sure those are registered first.
+    // Later, we pass their IDs to Schema Registry as our 'references'.
+    const references = await this.registerReferences(opts, helper, confluentSchema)
+    const referenceSchemas = await Promise.all(
+      references.map(async reference => {
+        return this.getSchemaForSubject(reference.subject, reference.version)
+      }),
+    )
+
+    const schemaInstance = schemaFromConfluentSchema(
+      confluentSchema,
+      this.options,
+      referenceSchemas,
+    )
     helper.validate(schemaInstance)
 
     let subject: ConfluentSubject
@@ -138,6 +183,7 @@ export default class SchemaRegistry {
       body: {
         schemaType: confluentSchema.type,
         schema: confluentSchema.schema,
+        references,
       },
     })
 
@@ -148,6 +194,15 @@ export default class SchemaRegistry {
     return registeredSchema
   }
 
+  public async getSchemaForSubject(subject: string, version: number): Promise<Schema | AvroSchema> {
+    const response = await this.api.Subject.version({
+      subject,
+      version,
+    })
+    const schema: { id: number } = response.data()
+    return await this.getSchema(schema.id)
+  }
+
   public async getSchema(registryId: number): Promise<Schema | AvroSchema> {
     const schema = this.cache.getSchema(registryId)
 
@@ -156,19 +211,35 @@ export default class SchemaRegistry {
     }
 
     const response = await this.getSchemaOriginRequest(registryId)
-    const foundSchema: { schema: string; schemaType: string } = response.data()
+    const foundSchema: {
+      schema: string
+      schemaType: string
+      references?: { name: string; subject: string; version: number }[]
+    } = response.data()
     const rawSchema = foundSchema.schema
     const schemaType = schemaTypeFromString(foundSchema.schemaType)
 
     if (schemaType === SchemaType.UNKNOWN) {
       throw new ConfluentSchemaRegistryError(`Unknown schema type ${foundSchema.schemaType}`)
     }
 
+    const referenceSchemas = await Promise.all(
+      (foundSchema.references || []).map(async reference => {
+        return this.getSchemaForSubject(reference.subject, reference.version)
+      }),
+    )
+
     const confluentSchema: ConfluentSchema = {
       type: schemaType,
       schema: rawSchema,
     }
-    const schemaInstance = schemaFromConfluentSchema(confluentSchema, this.options)
+
+    const schemaInstance = schemaFromConfluentSchema(
+      confluentSchema,
+      this.options,
+      referenceSchemas,
+    )
+
     return this.cache.setSchema(registryId, schemaInstance)
   }
 

src/api/index.ts

@@ -28,6 +28,7 @@ export interface SchemaRegistryAPIClientArgs {
 export type SchemaRegistryAPIClient = Client<{
   Schema: {
     find: (_: any) => any
+    versions: (_: any) => any
   }
   Subject: {
     all: (_: any) => any
@@ -65,6 +66,10 @@ export default ({
           method: 'get',
           path: '/schemas/ids/{id}',
         },
+        versions: {
+          method: 'get',
+          path: '/schemas/ids/{id}/versions',
+        },
       },
       Subject: {
         all: {

src/schemaTypeResolver.ts

@@ -65,6 +65,7 @@ export const helperTypeFromSchemaType = (
 export const schemaFromConfluentSchema = (
   confluentSchema: ConfluentSchema,
   options?: SchemaRegistryAPIClientOptions,
+  referenceSchemas?: Schema[],
 ): Schema | AvroSchema => {
   try {
     let schema: Schema
@@ -87,7 +88,7 @@ export const schemaFromConfluentSchema = (
       }
       case SchemaType.PROTOBUF: {
         const opts: ProtoOptions | undefined = (options as ProtocolOptions)?.[SchemaType.PROTOBUF]
-        schema = new ProtoSchema(confluentSchema, opts)
+        schema = new ProtoSchema(confluentSchema, opts, referenceSchemas)
         break
       }
       default: