Fix Protobuf field transform to use schema from SR (#138)

* Fix Protobuf field transform to use schema from SR

* Minor refactoring

* Minor formatting

Author: rayokota

schemaregistry/serde/avro.ts

@@ -285,7 +285,7 @@ async function transform(ctx: RuleContext, schema: Type, msg: any, fieldTransfor
       const recordSchema = schema as RecordType
       const record = msg as Record<string, any>
       for (const field of recordSchema.fields) {
-        await transformField(ctx, recordSchema, field, record, record[field.name], fieldTransform)
+        await transformField(ctx, recordSchema, field, record, fieldTransform)
       }
       return record
     default:
@@ -304,13 +304,12 @@ async function transformField(
   recordSchema: RecordType,
   field: Field,
   record: Record<string, any>,
-  val: any,
   fieldTransform: FieldTransform,
 ): Promise<void> {
   const fullName = recordSchema.name + '.' + field.name
   try {
     ctx.enterField(
-      val,
+      record,
       fullName,
       field.name,
       getType(field.type),

schemaregistry/serde/json.ts

@@ -350,8 +350,7 @@ async function transform(ctx: RuleContext, schema: DereferencedJSONSchema, path:
     case FieldType.RECORD:
       if (schema.properties != null) {
         for (let [propName, propSchema] of Object.entries(schema.properties)) {
-          let value = msg[propName]
-          await transformField(ctx, path, propName, msg, value, propSchema, fieldTransform)
+          await transformField(ctx, path, propName, msg, propSchema, fieldTransform)
         }
       }
       return msg
@@ -371,12 +370,13 @@ async function transform(ctx: RuleContext, schema: DereferencedJSONSchema, path:
   return msg
 }
 
-async function transformField(ctx: RuleContext, path: string, propName: string, msg: any, value: any,
-                        propSchema: DereferencedJSONSchema,
-                        fieldTransform: FieldTransform): Promise<void> {
+async function transformField(ctx: RuleContext, path: string, propName: string, msg: any,
+                              propSchema: DereferencedJSONSchema,
+                              fieldTransform: FieldTransform): Promise<void> {
   const fullName = path + '.' + propName
   try {
     ctx.enterField(msg, fullName, propName, getType(propSchema), getInlineTags(propSchema))
+    let value = msg[propName]
     const newVal = await transform(ctx, propSchema, fullName, value, fieldTransform)
     if (ctx.rule.kind === 'CONDITION') {
       if (newVal === false) {

schemaregistry/serde/protobuf.ts

@@ -32,7 +32,10 @@ import {
   file_google_protobuf_field_mask,
   file_google_protobuf_source_context,
   file_google_protobuf_struct,
-  file_google_protobuf_timestamp, file_google_protobuf_type, file_google_protobuf_wrappers,
+  file_google_protobuf_timestamp,
+  file_google_protobuf_type,
+  file_google_protobuf_wrappers,
+  FileDescriptorProto,
   FileDescriptorProtoSchema
 } from "@bufbuild/protobuf/wkt";
 import { BufferWrapper, MAX_VARINT_LEN_64 } from "./buffer-wrapper";
@@ -100,6 +103,7 @@ export type ProtobufSerializerConfig = SerializerConfig & {
  */
 export class ProtobufSerializer extends Serializer implements ProtobufSerde {
   registry: MutableRegistry
+  fileRegistry: FileRegistry
   schemaToDescCache: LRUCache<string, DescFile>
   descToSchemaCache: LRUCache<string, SchemaInfo>
 
@@ -113,6 +117,7 @@ export class ProtobufSerializer extends Serializer implements ProtobufSerde {
   constructor(client: Client, serdeType: SerdeType, conf: ProtobufSerializerConfig, ruleRegistry?: RuleRegistry) {
     super(client, serdeType, conf, ruleRegistry)
     this.registry = conf.registry ?? createMutableRegistry()
+    this.fileRegistry = createFileRegistry()
     this.schemaToDescCache = new LRUCache<string, DescFile>({ max: this.config().cacheCapacity ?? 1000 } )
     this.descToSchemaCache = new LRUCache<string, SchemaInfo>({ max: this.config().cacheCapacity ?? 1000 } )
     this.fieldTransformer = async (ctx: RuleContext, fieldTransform: FieldTransform, msg: any) => {
@@ -281,15 +286,41 @@ export class ProtobufSerializer extends Serializer implements ProtobufSerde {
   }
 
   async fieldTransform(ctx: RuleContext, fieldTransform: FieldTransform, msg: any): Promise<any> {
+    const fileDesc = await this.toFileDesc(this.client, ctx.target)
     const typeName = msg.$typeName
-    if (typeName == null) {
-      throw new SerializationError('message type name is empty')
+    const messageDesc = this.toMessageDescFromName(fileDesc, typeName)
+    return await transform(ctx, messageDesc, msg, fieldTransform)
+  }
+
+  async toFileDesc(client: Client, info: SchemaInfo): Promise<DescFile> {
+    const value = this.schemaToDescCache.get(stringify(info.schema))
+    if (value != null) {
+      return value
     }
-    const messageDesc = this.registry.getMessage(typeName)
-    if (messageDesc == null) {
-      throw new SerializationError('message descriptor not in registry')
+    const fileDesc = await this.parseFileDesc(client, info)
+    if (fileDesc == null) {
+      throw new SerializationError('file descriptor not found')
     }
-    return await transform(ctx, messageDesc, msg, fieldTransform)
+    this.schemaToDescCache.set(stringify(info.schema), fileDesc)
+    return fileDesc
+  }
+
+  async parseFileDesc(client: Client, info: SchemaInfo): Promise<DescFile | undefined> {
+    const deps = new Map<string, string>()
+    await this.resolveReferences(client, info, deps, 'serialized')
+    const fileDesc = fromBinary(FileDescriptorProtoSchema, Buffer.from(info.schema, 'base64'))
+    const fileRegistry = newFileRegistry(fileDesc, deps)
+    this.fileRegistry = createFileRegistry(this.fileRegistry, fileRegistry)
+    return this.fileRegistry.getFile(fileDesc.name)
+  }
+
+  toMessageDescFromName(fd: DescFile, msgName: string): DescMessage {
+    for (let i = 0; i < fd.messages.length; i++) {
+      if (fd.messages[i].typeName === msgName) {
+        return fd.messages[i]
+      }
+    }
+    throw new SerializationError('message descriptor not found')
   }
 }
 
@@ -302,7 +333,7 @@ export type ProtobufDeserializerConfig = DeserializerConfig
  * ProtobufDeserializer is a deserializer for Protobuf messages.
  */
 export class ProtobufDeserializer extends Deserializer implements ProtobufSerde {
-  registry: FileRegistry
+  fileRegistry: FileRegistry
   schemaToDescCache: LRUCache<string, DescFile>
 
   /**
@@ -314,7 +345,7 @@ export class ProtobufDeserializer extends Deserializer implements ProtobufSerde
    */
   constructor(client: Client, serdeType: SerdeType, conf: ProtobufDeserializerConfig, ruleRegistry?: RuleRegistry) {
     super(client, serdeType, conf, ruleRegistry)
-    this.registry = createFileRegistry()
+    this.fileRegistry = createFileRegistry()
     this.schemaToDescCache = new LRUCache<string, DescFile>({ max: this.config().cacheCapacity ?? 1000 } )
     this.fieldTransformer = async (ctx: RuleContext, fieldTransform: FieldTransform, msg: any) => {
       return await this.fieldTransform(ctx, fieldTransform, msg)
@@ -340,7 +371,7 @@ export class ProtobufDeserializer extends Deserializer implements ProtobufSerde
     const info = await this.getSchema(topic, payload, 'serialized')
     const fd = await this.toFileDesc(this.client, info)
     const [bytesRead, msgIndexes] = this.readMessageIndexes(payload.subarray(5))
-    const messageDesc = this.toMessageDesc(fd, msgIndexes)
+    const messageDesc = this.toMessageDescFromIndexes(fd, msgIndexes)
 
     const subject = this.subjectName(topic, info)
     const readerMeta = await this.getReaderSchema(subject, 'serialized')
@@ -361,14 +392,9 @@ export class ProtobufDeserializer extends Deserializer implements ProtobufSerde
   }
 
   async fieldTransform(ctx: RuleContext, fieldTransform: FieldTransform, msg: any): Promise<any> {
+    const fileDesc = await this.toFileDesc(this.client, ctx.target)
     const typeName = msg.$typeName
-    if (typeName == null) {
-      throw new SerializationError('message type name is empty')
-    }
-    const messageDesc = this.registry.getMessage(typeName)
-    if (messageDesc == null) {
-      throw new SerializationError('message descriptor not in registry')
-    }
+    const messageDesc = this.toMessageDescFromName(fileDesc, typeName)
     return await transform(ctx, messageDesc, msg, fieldTransform)
   }
 
@@ -389,26 +415,18 @@ export class ProtobufDeserializer extends Deserializer implements ProtobufSerde
     const deps = new Map<string, string>()
     await this.resolveReferences(client, info, deps, 'serialized')
     const fileDesc = fromBinary(FileDescriptorProtoSchema, Buffer.from(info.schema, 'base64'))
-    const resolve = (depName: string) => {
-      if (isBuiltin(depName)) {
-        const dep = builtinDeps.get(depName)
-        if (dep == null) {
-          throw new SerializationError(`dependency ${depName} not found`)
-        }
-        return dep
-      } else {
-        const dep = deps.get(depName)
-        if (dep == null) {
-          throw new SerializationError(`dependency ${depName} not found`)
-        }
-        const fileDesc = fromBinary(FileDescriptorProtoSchema, Buffer.from(dep, 'base64'))
-        fileDesc.name = depName
-        return fileDesc
+    const fileRegistry = newFileRegistry(fileDesc, deps)
+    this.fileRegistry = createFileRegistry(this.fileRegistry, fileRegistry)
+    return this.fileRegistry.getFile(fileDesc.name)
+  }
+
+  toMessageDescFromName(fd: DescFile, msgName: string): DescMessage {
+    for (let i = 0; i < fd.messages.length; i++) {
+      if (fd.messages[i].typeName === msgName) {
+        return fd.messages[i]
       }
     }
-    const fileRegistry = createFileRegistry(fileDesc, resolve)
-    this.registry = createFileRegistry(this.registry, fileRegistry)
-    return this.registry.getFile(fileDesc.name)
+    throw new SerializationError('message descriptor not found')
   }
 
   readMessageIndexes(payload: Buffer): [number, number[]] {
@@ -421,7 +439,7 @@ export class ProtobufDeserializer extends Deserializer implements ProtobufSerde
     return [bw.pos, msgIndexes]
   }
 
-  toMessageDesc(fd: DescFile, msgIndexes: number[]): DescMessage {
+  toMessageDescFromIndexes(fd: DescFile, msgIndexes: number[]): DescMessage {
     let index = msgIndexes[0]
     if (msgIndexes.length === 1) {
       return fd.messages[index]
@@ -438,6 +456,27 @@ export class ProtobufDeserializer extends Deserializer implements ProtobufSerde
   }
 }
 
+function newFileRegistry(fileDesc: FileDescriptorProto, deps: Map<string, string>): FileRegistry {
+  const resolve = (depName: string) => {
+    if (isBuiltin(depName)) {
+      const dep = builtinDeps.get(depName)
+      if (dep == null) {
+        throw new SerializationError(`dependency ${depName} not found`)
+      }
+      return dep
+    } else {
+      const dep = deps.get(depName)
+      if (dep == null) {
+        throw new SerializationError(`dependency ${depName} not found`)
+      }
+      const fileDesc = fromBinary(FileDescriptorProtoSchema, Buffer.from(dep, 'base64'))
+      fileDesc.name = depName
+      return fileDesc
+    }
+  }
+  return createFileRegistry(fileDesc, resolve)
+}
+
 async function transform(ctx: RuleContext, descriptor: DescMessage, msg: any, fieldTransform: FieldTransform): Promise<any> {
   if (msg == null || descriptor == null) {
     return msg
@@ -470,7 +509,7 @@ async function transform(ctx: RuleContext, descriptor: DescMessage, msg: any, fi
 }
 
 async function transformField(ctx: RuleContext, fd: DescField, desc: DescMessage,
-                        msg: any, fieldTransform: FieldTransform) {
+                              msg: any, fieldTransform: FieldTransform) {
   try {
     ctx.enterField(
       msg,