feat(kafka): add logic to handle delimited protobufs (#4071)

Author: dreamorosi

package-lock.json

@@ -51,9 +51,9 @@
     "@standard-schema/spec": "^1.0.0"
   },
   "peerDependencies": {
-    "zod": ">=3.24.0",
+    "arktype": ">=2.0.0",
     "valibot": ">=1.0.0",
-    "arktype": ">=2.0.0"
+    "zod": ">=3.24.0"
   },
   "peerDependenciesMeta": {
     "zod": {
@@ -93,12 +93,12 @@
     },
     "./types": {
       "require": {
-        "types": "./lib/cjs/types/types.d.ts",
-        "default": "./lib/cjs/types/types.js"
+        "types": "./lib/cjs/types/index.d.ts",
+        "default": "./lib/cjs/types/index.js"
       },
       "import": {
-        "types": "./lib/esm/types/types.d.ts",
-        "default": "./lib/esm/types/types.js"
+        "types": "./lib/esm/types/index.d.ts",
+        "default": "./lib/esm/types/index.js"
       }
     }
   },
@@ -109,8 +109,8 @@
         "lib/esm/errors.d.ts"
       ],
       "types": [
-        "lib/cjs/types/types.d.ts",
-        "lib/esm/types/types.d.ts"
+        "lib/cjs/types/index.d.ts",
+        "lib/esm/types/index.d.ts"
       ]
     }
   },

packages/kafka/package.json

@@ -14,6 +14,7 @@ import {
 import type {
   ConsumerRecord,
   ConsumerRecords,
+  DeserializeOptions,
   Deserializer,
   Record as KafkaRecord,
   MSKEvent,
@@ -74,11 +75,12 @@ const deserializeHeaders = (headers: Record<string, number[]>[] | null) => {
  * @param config - The schema configuration to use for deserialization. See {@link SchemaConfigValue | `SchemaConfigValue`}.
  *   If not provided, the value is decoded as a UTF-8 string.
  */
-const deserialize = (
-  value: string,
-  deserializer: Deserializer,
-  config?: SchemaConfigValue
-) => {
+const deserialize = ({
+  value,
+  deserializer,
+  config,
+  schemaMetadata,
+}: DeserializeOptions) => {
   if (config === undefined) {
     return deserializer(value);
   }
@@ -100,7 +102,7 @@ const deserialize = (
         'Schema string is required for protobuf deserialization'
       );
     }
-    return deserializer(value, config.schema);
+    return deserializer(value, config.schema, schemaMetadata);
   }
 };
 
@@ -162,7 +164,14 @@ const deserializeRecord = async (
   record: KafkaRecord,
   config?: SchemaConfig
 ) => {
-  const { key, value, headers, ...rest } = record;
+  const {
+    key,
+    value,
+    headers,
+    valueSchemaMetadata,
+    keySchemaMetadata,
+    ...rest
+  } = record;
   const { key: keyConfig, value: valueConfig } = config || {};
 
   const deserializerKey = await getDeserializer(keyConfig?.type);
@@ -175,19 +184,25 @@ const deserializeRecord = async (
         return undefined;
       }
       if (isNull(key)) return null;
-      const deserializedKey = deserialize(key, deserializerKey, keyConfig);
+      const deserializedKey = deserialize({
+        value: key,
+        deserializer: deserializerKey,
+        config: keyConfig,
+        schemaMetadata: keySchemaMetadata,
+      });
 
       return keyConfig?.parserSchema
         ? parseSchema(deserializedKey, keyConfig.parserSchema)
         : deserializedKey;
     },
     originalKey: key,
     get value() {
-      const deserializedValue = deserialize(
-        value,
-        deserializerValue,
-        valueConfig
-      );
+      const deserializedValue = deserialize({
+        value: value,
+        deserializer: deserializerValue,
+        config: valueConfig,
+        schemaMetadata: valueSchemaMetadata,
+      });
 
       return valueConfig?.parserSchema
         ? parseSchema(deserializedValue, valueConfig.parserSchema)

packages/kafka/src/consumer.ts

@@ -1,6 +1,11 @@
-import type { Message } from 'protobufjs';
+import { BufferReader, type Message } from 'protobufjs';
 import { KafkaConsumerDeserializationError } from '../errors.js';
-import type { ProtobufMessage } from '../types/types.js';
+import type { ProtobufMessage, SchemaMetadata } from '../types/types.js';
+
+/**
+ * Default order of varint types used in Protobuf to attempt deserializing Confluent Schema Registry messages.
+ */
+const varintOrder: Array<'int32' | 'sint32'> = ['int32', 'sint32'];
 
 /**
  * Deserialize a Protobuf message from a base64-encoded string.
@@ -10,15 +15,80 @@ import type { ProtobufMessage } from '../types/types.js';
  * @param data - The base64-encoded string representing the Protobuf binary data.
  * @param messageType - The Protobuf message type definition - see {@link Message | `Message`} from {@link https://www.npmjs.com/package/protobufjs | `protobufjs`}.
  */
-const deserialize = <T>(data: string, messageType: ProtobufMessage<T>): T => {
+const deserialize = <T>(
+  data: string,
+  messageType: ProtobufMessage<T>,
+  schemaMetadata: SchemaMetadata
+): T => {
+  const buffer = Buffer.from(data, 'base64');
   try {
-    const buffer = Buffer.from(data, 'base64');
-    return messageType.decode(buffer, buffer.length);
+    if (schemaMetadata.schemaId === undefined) {
+      return messageType.decode(buffer, buffer.length);
+    }
+    /**
+     * If `schemaId` is longer than 10 chars, it's an UUID, otherwise it's a numeric ID.
+     *
+     * When this is the case, we know the schema is coming from Glue Schema Registry,
+     * and the first byte of the buffer is a magic byte that we need to remove before
+     * decoding the message.
+     */
+    if (schemaMetadata.schemaId.length > 10) {
+      // remove the first byte from the buffer
+      const reader = new BufferReader(buffer);
+      reader.uint32();
+      return messageType.decode(reader);
+    }
   } catch (error) {
     throw new KafkaConsumerDeserializationError(
       `Failed to deserialize Protobuf message: ${error}, message: ${data}, messageType: ${messageType}`
     );
   }
+
+  /**
+   * If schemaId is numeric, inferred from its length, we know it's coming from Confluent Schema Registry,
+   * so we need to remove the MessageIndex bytes.
+   * We don't know the type of the index, so we try both `int32` and `sint32`. If both fail, we throw an error.
+   */
+  try {
+    const newBuffer = clipConfluentSchemaRegistryBuffer(buffer, varintOrder[0]);
+    return messageType.decode(newBuffer);
+  } catch (error) {
+    try {
+      const newBuffer = clipConfluentSchemaRegistryBuffer(
+        buffer,
+        varintOrder[1]
+      );
+      const decoded = messageType.decode(newBuffer);
+      // swap varint order if the first attempt failed so we can use the correct one for subsequent messages
+      varintOrder.reverse();
+      return decoded;
+    } catch {
+      throw new KafkaConsumerDeserializationError(
+        `Failed to deserialize Protobuf message: ${error}, message: ${data}, messageType: ${messageType}`
+      );
+    }
+  }
+};
+
+/**
+ * Clip the Confluent Schema Registry buffer to remove the index bytes.
+ *
+ * @param buffer - The buffer to clip.
+ * @param intType - The type of the integer to read from the buffer, either 'int32' or 'sint32'.
+ */
+const clipConfluentSchemaRegistryBuffer = (
+  buffer: Buffer,
+  intType: 'int32' | 'sint32'
+) => {
+  const reader = new BufferReader(buffer);
+  /**
+   * Read the first varint byte to get the index count or 0.
+   * Doing so, also advances the reader position to the next byte after the index count.
+   */
+  const indexCount = intType === 'int32' ? reader.int32() : reader.sint32();
+  // Skip the index bytes
+  reader.skip(indexCount);
+  return reader;
 };
 
 export { deserialize };

packages/kafka/src/deserializer/protobuf.ts

@@ -0,0 +1,12 @@
+export type {
+  ConsumerRecord,
+  ConsumerRecords,
+  MSKEvent,
+  ProtobufMessage,
+  Record,
+  RecordHeader,
+  SchemaType,
+  SchemaConfig,
+  SchemaConfigValue,
+  SchemaMetadata,
+} from './types.js';