@databricks/zerobus-ingest-sdk 0.0.1 → 0.1.1

package/utils/descriptor.ts

@@ -0,0 +1,103 @@
+/**
+ * Utility functions for working with Protocol Buffer descriptors.
+ *
+ * These utilities help extract message descriptors from FileDescriptorSets
+ * generated by protoc, providing flexibility in choosing which message to use.
+ */
+
+import * as fs from 'fs';
+import * as protobuf from 'protobufjs';
+import descriptor from 'protobufjs/ext/descriptor';
+
+/**
+ * Options for loading a descriptor from a FileDescriptorSet.
+ */
+export interface LoadDescriptorOptions {
+    /**
+     * Path to the descriptor file generated by protoc.
+     * Example: 'schemas/my_schema_descriptor.pb'
+     */
+    descriptorPath: string;
+
+    /**
+     * Name of the proto file within the FileDescriptorSet.
+     * This should match the filename in your .proto file.
+     * Example: 'schemas/air_quality.proto' or 'air_quality.proto'
+     */
+    protoFileName: string;
+
+    /**
+     * Name of the message type to extract.
+     * Example: 'AirQuality'
+     */
+    messageName: string;
+}
+
+/**
+ * Loads a specific message's DescriptorProto from a FileDescriptorSet.
+ *
+ * This function:
+ * 1. Reads the binary FileDescriptorSet file generated by protoc
+ * 2. Parses it using protobufjs
+ * 3. Finds the specified proto file and message within the FileDescriptorSet
+ * 4. Extracts just that message's DescriptorProto
+ * 5. Encodes it to base64 for use with the Zerobus SDK
+ *
+ * @param options - Configuration specifying which message to extract
+ * @returns Base64-encoded DescriptorProto string for the specified message
+ * @throws Error if the descriptor file, proto file, or message cannot be found
+ *
+ * @example
+ * ```typescript
+ * import { loadDescriptorProto } from '@databricks/zerobus-sdk/utils/descriptor';
+ *
+ * const descriptorBase64 = loadDescriptorProto({
+ *     descriptorPath: 'schemas/air_quality_descriptor.pb',
+ *     protoFileName: 'air_quality.proto',
+ *     messageName: 'AirQuality'
+ * });
+ *
+ * const tableProperties = {
+ *     tableName: 'catalog.schema.table',
+ *     descriptorProto: descriptorBase64
+ * };
+ * ```
+ */
+export function loadDescriptorProto(options: LoadDescriptorOptions): string {
+    const { descriptorPath, protoFileName, messageName } = options;
+
+    // Read the binary descriptor file
+    const descriptorBytes = fs.readFileSync(descriptorPath);
+
+    // Parse as FileDescriptorSet using the descriptor extension
+    const fileDescriptorSet = descriptor.FileDescriptorSet.decode(descriptorBytes) as any;
+
+    // Find the file descriptor matching the proto file name
+    const fileDescriptor = fileDescriptorSet.file.find((f: any) =>
+        f.name === protoFileName || f.name.endsWith('/' + protoFileName) || f.name.endsWith(protoFileName)
+    );
+
+    if (!fileDescriptor) {
+        throw new Error(
+            `Proto file '${protoFileName}' not found in descriptor. Available files: ${
+                fileDescriptorSet.file.map((f: any) => f.name).join(', ')
+            }`
+        );
+    }
+
+    // Find the message type
+    const messageType = fileDescriptor.messageType?.find((m: any) => m.name === messageName);
+
+    if (!messageType) {
+        const availableMessages = fileDescriptor.messageType?.map((m: any) => m.name).join(', ') || 'none';
+        throw new Error(
+            `Message '${messageName}' not found in ${protoFileName}. Available messages: ${availableMessages}`
+        );
+    }
+
+    // Encode the DescriptorProto to binary
+    const descriptorProtoBytes = descriptor.DescriptorProto.encode(messageType).finish();
+
+    // Convert to base64
+    return Buffer.from(descriptorProtoBytes).toString('base64');
+}