@databricks/zerobus-ingest-sdk 1.1.0 → 1.3.0

package/src/headers_provider.ts

@@ -1,69 +1,57 @@
 /**
- * Interface for providing custom headers to Zerobus streams.
+ * Custom headers provider accepted by `createStream()`.
  *
- * Implement this interface to use custom authentication beyond OAuth,
- * such as Personal Access Tokens (PAT) or custom auth tokens.
+ * The native adapter invokes `getHeadersCallback` synchronously once during
+ * stream creation and stores the returned tuples. Returning a Promise, or
+ * passing a class with async `getHeaders()`, is not supported and can terminate
+ * the process. Token refresh is not currently wired through this callback.
  */
 export interface HeadersProvider {
     /**
-     * Returns headers as array of [name, value] tuples.
+     * Returns headers as an array of [name, value] tuples.
      *
      * Required headers:
      * - ["authorization", "Bearer <token>"]
      * - ["x-databricks-zerobus-table-name", "<table_name>"]
-     *
-     * @returns Promise resolving to array of header name-value pairs
      */
-    getHeaders(): Promise<Array<[string, string]>>;
+    getHeadersCallback: () => Array<[string, string]>;
 }
 
 /**
  * OAuth 2.0 Client Credentials headers provider.
  *
- * **IMPORTANT: DO NOT instantiate this class directly.**
+ * Do not instantiate this class or pass it to `createStream()`.
  *
  * OAuth authentication is handled automatically by the Rust SDK when you call
- * `createStream()` with clientId and clientSecret parameters (without providing
- * a headers_provider).
- *
- * This class exists for:
- * 1. Documentation purposes - showing the HeadersProvider pattern
- * 2. API consistency with other Zerobus SDKs (Python, Java, Rust)
+ * `createStream()` with clientId and clientSecret and omit the headers provider.
  *
- * **How to use OAuth authentication:**
+ * How to use OAuth authentication:
  * ```typescript
- * // OAuth is the default - just pass clientId and clientSecret
  * const stream = await sdk.createStream(
  *     tableProperties,
- *     clientId,      // OAuth client ID
- *     clientSecret,  // OAuth client secret
+ *     clientId,
+ *     clientSecret,
  *     options
- *     // No headers_provider parameter = OAuth authentication
  * );
  * ```
  *
- * **How to use custom authentication (PAT, etc.):**
+ * How to use custom authentication (PAT or a static token):
  * ```typescript
- * class CustomHeadersProvider implements HeadersProvider {
- *     async getHeaders() {
- *         return [
- *             ["authorization", `Bearer ${myToken}`],
- *             ["x-databricks-zerobus-table-name", tableName]
- *         ];
- *     }
- * }
- *
- * const provider = new CustomHeadersProvider();
  * const stream = await sdk.createStream(
  *     tableProperties,
- *     '', // ignored
- *     '', // ignored
+ *     '',
+ *     '',
  *     options,
- *     { getHeadersCallback: provider.getHeaders.bind(provider) }
+ *     {
+ *         getHeadersCallback: () => [
+ *             ["authorization", `Bearer ${myToken}`],
+ *             ["x-databricks-zerobus-table-name", tableName]
+ *         ]
+ *     }
  * );
  * ```
  */
-export class OAuthHeadersProvider implements HeadersProvider {
+export class OAuthHeadersProvider {
     constructor(
         private clientId: string,
         private clientSecret: string,
@@ -75,8 +63,8 @@ export class OAuthHeadersProvider implements HeadersProvider {
         throw new Error(
             'OAuthHeadersProvider should not be instantiated directly. ' +
             'OAuth authentication is handled internally by the Rust SDK. ' +
-            'To use OAuth: call createStream(tableProperties, clientId, clientSecret, options) without the headers_provider parameter. ' +
-            'To use custom authentication: implement the HeadersProvider interface.'
+            'To use OAuth: call createStream(tableProperties, clientId, clientSecret, options) without a headers provider. ' +
+            'To use custom authentication: pass { getHeadersCallback: () => [...] } as the headers provider.'
         );
     }
 }

package/src/lib.rs

@@ -50,7 +50,7 @@ pub enum RecordType {
 #[napi(object)]
 pub struct StreamConfigurationOptions {
     /// Maximum number of unacknowledged requests that can be in flight.
-    /// Default: 10,000
+    /// Default: 1,000,000
     pub max_inflight_requests: Option<u32>,
 
     /// Enable automatic stream recovery on transient failures.
@@ -103,7 +103,8 @@ pub struct TableProperties {
     pub table_name: String,
 
     /// Optional Protocol Buffer descriptor as a base64-encoded string.
-    /// If not provided, JSON encoding will be used.
+    /// Omitting this does not select JSON. The stream defaults to Protocol Buffers
+    /// unless `record_type` is set to JSON.
     pub descriptor_proto: Option<String>,
 }
 
@@ -213,15 +214,16 @@ fn convert_js_to_record_payload(env: &Env, payload: Unknown) -> Result<RustRecor
 
 /// A stream for ingesting data into a Databricks Delta table.
 ///
-/// The stream manages a bidirectional gRPC connection, handles acknowledgments,
-/// and provides automatic recovery on transient failures.
+/// The stream manages JSON or Protocol Buffer ingestion over a bidirectional
+/// gRPC connection, handles acknowledgments, and provides automatic recovery
+/// on transient failures.
 ///
 /// # Example
 ///
 /// ```typescript
 /// const stream = await sdk.createStream(tableProps, clientId, clientSecret, options);
-/// const ackPromise = await stream.ingestRecord(Buffer.from([1, 2, 3]));
-/// const offset = await ackPromise;
+/// const offset = await stream.ingestRecordOffset(Buffer.from([1, 2, 3]));
+/// await stream.flush();
 /// await stream.close();
 /// ```
 #[napi]
@@ -400,6 +402,13 @@ impl ZerobusStream {
     /// This is the recommended API for high-throughput scenarios where you want to
     /// decouple record ingestion from acknowledgment tracking.
     ///
+    /// **Acknowledgments:** the idiomatic flow is to ingest in a loop and then `flush()`
+    /// once to confirm everything queued so far. The returned offset, together with
+    /// `waitForOffset()`, lets you confirm a specific record when you need it (acks are
+    /// ordered, so the last offset confirms the whole run) — prefer `flush()` for bulk.
+    /// Avoid calling `waitForOffset()` after every record in a tight loop, since that
+    /// limits throughput to one record per round-trip.
+    ///
     /// # Arguments
     ///
     /// * `payload` - The record data (Buffer, string, protobuf message, or plain object)
@@ -412,11 +421,14 @@ impl ZerobusStream {
     /// # Example
     ///
     /// ```typescript
-    /// // Promise resolves immediately with offset (before server ack)
-    /// const offset1 = await stream.ingestRecordOffset(record1);
-    /// const offset2 = await stream.ingestRecordOffset(record2);
-    /// // Wait for both to be acknowledged
-    /// await stream.waitForOffset(offset2);
+    /// // High-throughput pattern: ingest in a loop, wait once at the end.
+    /// let lastOffset: bigint | null = null;
+    /// for (const record of records) {
+    ///   lastOffset = await stream.ingestRecordOffset(record); // resolves on queue, no round-trip
+    /// }
+    /// // The ack watermark is monotonic: waiting on the last offset confirms all prior records.
+    /// if (lastOffset !== null) await stream.waitForOffset(lastOffset);
+    /// // Or simply: await stream.flush();
     /// ```
     #[napi(ts_return_type = "Promise<bigint>")]
     pub fn ingest_record_offset(&self, env: Env, payload: Unknown) -> Result<JsObject> {
@@ -454,6 +466,12 @@ impl ZerobusStream {
     /// the batch is queued, without waiting for server acknowledgment. Use
     /// `waitForOffset()` to wait for acknowledgment when needed.
     ///
+    /// **Acknowledgments:** the idiomatic flow is to ingest your batches in a loop and
+    /// then `flush()` once to confirm. The returned offset, together with `waitForOffset()`,
+    /// confirms a specific batch when you need it (acks are ordered, so the last offset
+    /// confirms the whole run) — prefer `flush()` for bulk. Avoid calling `waitForOffset()`
+    /// after every batch in a tight loop, since that limits throughput to one round-trip per batch.
+    ///
     /// # Arguments
     ///
     /// * `records` - Array of record data
@@ -466,11 +484,14 @@ impl ZerobusStream {
     /// # Example
     ///
     /// ```typescript
-    /// // Promise resolves immediately with offset (before server ack)
-    /// const offset = await stream.ingestRecordsOffset(batch);
-    /// if (offset !== null) {
-    ///   await stream.waitForOffset(offset);
+    /// // Ingest many batches without waiting, then flush once.
+    /// let lastOffset = null;
+    /// for (const batch of batches) {
+    ///   const offset = await stream.ingestRecordsOffset(batch); // resolves on queue
+    ///   if (offset !== null) lastOffset = offset;
     /// }
+    /// if (lastOffset !== null) await stream.waitForOffset(lastOffset);
+    /// // Or simply: await stream.flush();
     /// ```
     #[napi(ts_return_type = "Promise<bigint | null>")]
     pub fn ingest_records_offset(&self, env: Env, records: Vec<Unknown>) -> Result<JsObject> {
@@ -511,9 +532,12 @@ impl ZerobusStream {
 
     /// Waits for a specific offset to be acknowledged by the server.
     ///
-    /// Use this method with `ingestRecordOffset()` and `ingestRecordsOffset()` to
-    /// selectively wait for acknowledgments. This allows you to ingest many records
-    /// quickly and then wait only for specific offsets when needed.
+    /// Use this method with `ingestRecordOffset()` and `ingestRecordsOffset()` to confirm
+    /// a specific record before continuing. Acks are ordered, so waiting on the LAST offset
+    /// confirms every prior record too — you never need to wait on intermediate offsets.
+    /// For confirming a bulk run, `flush()` is usually simpler; reach for `waitForOffset()`
+    /// when one particular record must be confirmed. Avoid calling it after every record in
+    /// a tight loop, since that limits throughput to one record per round-trip.
     ///
     /// # Arguments
     ///
@@ -527,12 +551,12 @@ impl ZerobusStream {
     /// # Example
     ///
     /// ```typescript
-    /// const offsets = [];
+    /// let lastOffset: bigint | null = null;
     /// for (const record of records) {
-    ///   offsets.push(await stream.ingestRecordOffset(record));
+    ///   lastOffset = await stream.ingestRecordOffset(record); // no per-record wait
     /// }
-    /// // Wait for the last offset (implies all previous are also acknowledged)
-    /// await stream.waitForOffset(offsets[offsets.length - 1]);
+    /// // Wait for the last offset only (implies all previous are also acknowledged).
+    /// if (lastOffset !== null) await stream.waitForOffset(lastOffset);
     /// ```
     #[napi(ts_return_type = "Promise<void>")]
     pub fn wait_for_offset(&self, env: Env, offset_id: BigInt) -> Result<JsObject> {
@@ -560,6 +584,11 @@ impl ZerobusStream {
     /// This method ensures all previously ingested records have been sent to the server
     /// and acknowledged. It's useful for checkpointing or ensuring data durability.
     ///
+    /// This is the idiomatic way to confirm records ingested via `ingestRecordOffset()` /
+    /// `ingestRecordsOffset()`: ingest in a loop, then `flush()` once (for a bounded batch,
+    /// or periodically for a long-running stream). It resolves once everything queued so
+    /// far is acknowledged.
+    ///
     /// # Errors
     ///
     /// - Timeout errors if flush takes longer than configured timeout
@@ -646,15 +675,12 @@ impl ZerobusStream {
     ///
     /// ```typescript
     /// try {
-    ///   await stream.ingestRecords(batch1);
-    ///   await stream.ingestRecords(batch2);
+    ///   await stream.ingestRecordsOffset(batch1);
+    ///   await stream.ingestRecordsOffset(batch2);
+    ///   await stream.flush();
     /// } catch (error) {
     ///   const unackedBatches = await stream.getUnackedBatches();
-    ///
-    ///   // Re-ingest with new stream
-    ///   for (const batch of unackedBatches) {
-    ///     await newStream.ingestRecords(batch);
-    ///   }
+    ///   console.log(`Batches available for recovery: ${unackedBatches.length}`);
     /// }
     /// ```
     #[napi]
@@ -687,10 +713,10 @@ impl ZerobusStream {
 /// JavaScript headers provider callback wrapper.
 ///
 /// Allows TypeScript code to provide custom authentication headers
-/// by implementing a getHeaders() function.
+/// by implementing a getHeadersCallback() function.
 #[napi(object)]
 pub struct JsHeadersProvider {
-    /// JavaScript function: () => Promise<Array<[string, string]>>
+    /// JavaScript function: () => Array<[string, string]>
     pub get_headers_callback: JsFunction,
 }
 
@@ -721,11 +747,6 @@ impl StaticHeadersProvider {
             ));
         }
 
-        // Add TS user agent if not provided
-        if !map.contains_key("user-agent") {
-            map.insert("user-agent", TS_SDK_USER_AGENT.to_string());
-        }
-
         Ok(Self { headers: map })
     }
 }
@@ -810,11 +831,17 @@ impl RustHeadersProvider for TsOAuthHeadersProvider {
         let mut headers = HashMap::new();
         headers.insert("authorization", format!("Bearer {}", token));
         headers.insert("x-databricks-zerobus-table-name", self.table_name.clone());
-        headers.insert("user-agent", TS_SDK_USER_AGENT.to_string());
         Ok(headers)
     }
 }
 
+#[napi(object)]
+#[derive(Default)]
+pub struct ZerobusSdkOptions {
+    /// Identifier appended to the `user-agent` header
+    pub application_name: Option<String>,
+}
+
 /// The main SDK for interacting with the Databricks Zerobus service.
 ///
 /// This is the entry point for creating ingestion streams to Delta tables.
@@ -824,7 +851,8 @@ impl RustHeadersProvider for TsOAuthHeadersProvider {
 /// ```typescript
 /// const sdk = new ZerobusSdk(
 ///   "https://workspace-id.zerobus.region.cloud.databricks.com",
-///   "https://workspace.cloud.databricks.com"
+///   "https://workspace.cloud.databricks.com",
+///   { applicationName: "my-app/1.0" }
 /// );
 ///
 /// const stream = await sdk.createStream(
@@ -852,13 +880,19 @@ impl ZerobusSdk {
     ///   (e.g., "https://workspace-id.zerobus.region.cloud.databricks.com")
     /// * `unity_catalog_url` - The Unity Catalog endpoint URL
     ///   (e.g., "https://workspace.cloud.databricks.com")
+    /// * `options` - Optional SDK configuration (see `ZerobusSdkOptions`),
+    ///   including `applicationName` for server-side attribution.
     ///
     /// # Errors
     ///
     /// - Invalid endpoint URLs
     /// - Failed to extract workspace ID from the endpoint
     #[napi(constructor)]
-    pub fn new(zerobus_endpoint: String, unity_catalog_url: String) -> Result<Self> {
+    pub fn new(
+        zerobus_endpoint: String,
+        unity_catalog_url: String,
+        options: Option<ZerobusSdkOptions>,
+    ) -> Result<Self> {
         let workspace_id = zerobus_endpoint
             .strip_prefix("https://")
             .or_else(|| zerobus_endpoint.strip_prefix("http://"))
@@ -870,9 +904,17 @@ impl ZerobusSdk {
                 )
             })?;
 
-        let inner = RustZerobusSdk::builder()
+        let options = options.unwrap_or_default();
+
+        let builder = RustZerobusSdk::builder()
             .endpoint(&zerobus_endpoint)
             .unity_catalog_url(&unity_catalog_url)
+            .sdk_identifier(TS_SDK_USER_AGENT);
+        let builder = match options.application_name {
+            Some(name) => builder.application_name(name),
+            None => builder,
+        };
+        let inner = builder
             .build()
             .map_err(|e| Error::from_reason(format!("Failed to create SDK: {}", e)))?;
 
@@ -885,7 +927,7 @@ impl ZerobusSdk {
 
     /// Creates a new ingestion stream to a Delta table.
     ///
-    /// This method establishes a bidirectional gRPC connection to the Zerobus service
+    /// This method opens a JSON or Protocol Buffer stream to the Zerobus service
     /// and prepares it for data ingestion. By default, it uses OAuth 2.0 Client Credentials
     /// authentication. For custom authentication (e.g., Personal Access Tokens), provide
     /// a custom headers_provider.
@@ -927,7 +969,7 @@ impl ZerobusSdk {
     ///   "", // ignored
     ///   undefined,
     ///   {
-    ///     getHeadersCallback: async () => [
+    ///     getHeadersCallback: () => [
     ///       ["authorization", `Bearer ${myToken}`],
     ///       ["x-databricks-zerobus-table-name", tableName]
     ///     ]
@@ -947,19 +989,19 @@ impl ZerobusSdk {
         // Decode the optional protobuf descriptor up-front so we can hand it
         // to the builder's `.compiled_proto()` setter; the builder constructs
         // the (now-private) `TableProperties` itself.
-        let descriptor_proto: Option<prost_types::DescriptorProto> =
-            if let Some(ref desc_str) = table_properties.descriptor_proto {
-                let bytes = base64_decode(desc_str).map_err(|e| {
-                    Error::from_reason(format!("Failed to decode descriptor: {}", e))
+        let descriptor_proto: Option<prost_types::DescriptorProto> = if let Some(ref desc_str) =
+            table_properties.descriptor_proto
+        {
+            let bytes = base64_decode(desc_str)
+                .map_err(|e| Error::from_reason(format!("Failed to decode descriptor: {}", e)))?;
+            let dp: prost_types::DescriptorProto =
+                prost::Message::decode(&bytes[..]).map_err(|e| {
+                    Error::from_reason(format!("Failed to parse descriptor proto: {}", e))
                 })?;
-                let dp: prost_types::DescriptorProto = prost::Message::decode(&bytes[..])
-                    .map_err(|e| {
-                        Error::from_reason(format!("Failed to parse descriptor proto: {}", e))
-                    })?;
-                Some(dp)
-            } else {
-                None
-            };
+            Some(dp)
+        } else {
+            None
+        };
 
         let opts = options.unwrap_or(StreamConfigurationOptions {
             max_inflight_requests: None,
@@ -1076,7 +1118,8 @@ impl ZerobusSdk {
     ///
     /// # Arguments
     ///
-    /// * `stream` - The failed or closed stream to recreate
+    /// * `stream` - The terminally failed stream to recreate. The TypeScript wrapper
+    ///   must not have been closed because `close()` releases its native handle.
     ///
     /// # Returns
     ///
@@ -1092,12 +1135,23 @@ impl ZerobusSdk {
     ///
     /// ```typescript
     /// try {
-    ///   await stream.ingestRecords(batch);
+    ///   await stream.ingestRecordsOffset(batch);
+    ///   await stream.flush();
     /// } catch (error) {
-    ///   await stream.close();
-    ///   // Recreate stream with all unacked batches re-ingested
-    ///   const newStream = await sdk.recreateStream(stream);
-    ///   // Continue ingesting with newStream
+    ///   try {
+    ///     const newStream = await sdk.recreateStream(stream);
+    ///     try {
+    ///       await newStream.flush();
+    ///     } finally {
+    ///       await newStream.close();
+    ///     }
+    ///   } finally {
+    ///     try {
+    ///       await stream.close();
+    ///     } catch (closeError) {
+    ///       console.error("Failed stream released:", closeError);
+    ///     }
+    ///   }
     /// }
     /// ```
     #[napi]
@@ -1128,8 +1182,8 @@ fn base64_decode(input: &str) -> std::result::Result<Vec<u8>, String> {
 }
 
 // =============================================================================
-// Arrow Flight Support (Beta)
-// Enabled with feature flag: cargo build --features arrow-flight
+// Arrow Flight support
+// Behind the arrow-flight feature. Enable with: npm run build:arrow
 // =============================================================================
 
 #[cfg(feature = "arrow-flight")]
@@ -1144,8 +1198,6 @@ use databricks_zerobus_ingest_sdk::{
 
 /// IPC compression type for Arrow Flight streams.
 ///
-/// **Beta**: Arrow Flight support is in Beta. The API is stabilising but
-/// may still change before reaching GA.
 #[cfg(feature = "arrow-flight")]
 #[napi]
 pub enum IpcCompressionType {
@@ -1157,8 +1209,6 @@ pub enum IpcCompressionType {
 
 /// Configuration options for Arrow Flight streams.
 ///
-/// **Beta**: Arrow Flight support is in Beta. The API is stabilising but
-/// may still change before reaching GA.
 #[cfg(feature = "arrow-flight")]
 #[napi(object)]
 #[derive(Debug, Clone)]
@@ -1215,7 +1265,6 @@ fn map_ipc_compression(value: Option<i32>) -> Option<arrow_ipc::CompressionType>
 
 /// Arrow data type enum for schema definition.
 ///
-/// **Beta**: Arrow Flight support is in Beta.
 #[cfg(feature = "arrow-flight")]
 #[napi]
 pub enum ArrowDataType {
@@ -1287,7 +1336,6 @@ fn convert_arrow_data_type(dt: i32) -> RustDataType {
 
 /// Arrow field definition for schema.
 ///
-/// **Beta**: Arrow Flight support is in Beta.
 #[cfg(feature = "arrow-flight")]
 #[napi(object)]
 #[derive(Debug, Clone)]
@@ -1305,8 +1353,6 @@ pub struct ArrowField {
 /// Unlike `TableProperties` which uses Protocol Buffers, Arrow Flight streams
 /// require an Arrow schema definition.
 ///
-/// **Beta**: Arrow Flight support is in Beta. The API is stabilising but
-/// may still change before reaching GA.
 #[cfg(feature = "arrow-flight")]
 #[napi(object)]
 #[derive(Debug, Clone)]
@@ -1340,9 +1386,6 @@ fn build_arrow_schema(fields: &[ArrowField]) -> Arc<RustArrowSchema> {
 /// This stream provides a high-performance interface for streaming Arrow data
 /// to Databricks Delta tables using the Arrow Flight protocol.
 ///
-/// **Beta**: Arrow Flight support is in Beta. The API is stabilising but
-/// may still change before reaching GA.
-///
 /// # Lifecycle
 ///
 /// 1. Create a stream via `sdk.createArrowStream()`
@@ -1438,9 +1481,7 @@ impl ZerobusArrowStream {
                 stream_ref
                     .ingest_ipc_batch(Bytes::from(buffer_vec))
                     .await
-                    .map_err(|e| {
-                        napi::Error::from_reason(format!("Failed to ingest batch: {}", e))
-                    })
+                    .map_err(|e| napi::Error::from_reason(format!("Failed to ingest batch: {}", e)))
             },
             |env, offset_id| {
                 let global: JsGlobal = env.get_global()?;
@@ -1566,9 +1607,6 @@ impl ZerobusArrowStream {
 impl ZerobusSdk {
     /// Creates a new Arrow Flight stream to a Delta table.
     ///
-    /// **Beta**: Arrow Flight support is in Beta. The API is stabilising
-    /// but may still change before reaching GA.
-    ///
     /// This method establishes an Arrow Flight connection to the Zerobus service
     /// for high-performance columnar data ingestion.
     ///
@@ -1682,11 +1720,10 @@ impl ZerobusSdk {
 
     /// Recreates an Arrow stream with the same configuration and re-ingests unacknowledged batches.
     ///
-    /// **Beta**: Arrow Flight support is in Beta.
-    ///
     /// # Arguments
     ///
-    /// * `stream` - The failed or closed Arrow stream to recreate
+    /// * `stream` - The terminally failed Arrow stream to recreate. The TypeScript wrapper
+    ///   must not have been closed because `close()` releases its native handle.
     ///
     /// # Returns
     ///

package/utils/descriptor.d.ts

@@ -0,0 +1,32 @@
+/**
+ * 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.
+ *
+ * @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
+ */
+export function loadDescriptorProto(options: LoadDescriptorOptions): string;

package/utils/descriptor.js

@@ -0,0 +1,55 @@
+"use strict";
+
+/**
+ * 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.
+ */
+
+const fs = require("fs");
+const descriptor = require("protobufjs/ext/descriptor");
+
+/**
+ * Loads a specific message's DescriptorProto from a FileDescriptorSet.
+ *
+ * @param {object} options
+ * @param {string} options.descriptorPath Path to the descriptor file generated by protoc.
+ * @param {string} options.protoFileName Name of the proto file within the FileDescriptorSet.
+ * @param {string} options.messageName Name of the message type to extract.
+ * @returns {string} Base64-encoded DescriptorProto string for the specified message.
+ */
+function loadDescriptorProto(options) {
+    const { descriptorPath, protoFileName, messageName } = options;
+
+    const descriptorBytes = fs.readFileSync(descriptorPath);
+    const fileDescriptorSet = descriptor.FileDescriptorSet.decode(descriptorBytes);
+    const fileDescriptor = fileDescriptorSet.file.find((file) =>
+        file.name === protoFileName ||
+        file.name.endsWith("/" + protoFileName)
+    );
+
+    if (!fileDescriptor) {
+        throw new Error(
+            `Proto file '${protoFileName}' not found in descriptor. Available files: ${
+                fileDescriptorSet.file.map((file) => file.name).join(", ")
+            }`
+        );
+    }
+
+    const messageType = fileDescriptor.messageType?.find((message) => message.name === messageName);
+
+    if (!messageType) {
+        const availableMessages = fileDescriptor.messageType?.map((message) => message.name).join(", ") || "none";
+        throw new Error(
+            `Message '${messageName}' not found in ${protoFileName}. Available messages: ${availableMessages}`
+        );
+    }
+
+    const descriptorProtoBytes = descriptor.DescriptorProto.encode(messageType).finish();
+    return Buffer.from(descriptorProtoBytes).toString("base64");
+}
+
+module.exports = {
+    loadDescriptorProto,
+};

package/utils/descriptor.ts

@@ -49,7 +49,7 @@ export interface LoadDescriptorOptions {
  *
  * @example
  * ```typescript
- * import { loadDescriptorProto } from '@databricks/zerobus-sdk/utils/descriptor';
+ * import { loadDescriptorProto } from '@databricks/zerobus-ingest-sdk/utils/descriptor.js';
  *
  * const descriptorBase64 = loadDescriptorProto({
  *     descriptorPath: 'schemas/air_quality_descriptor.pb',
@@ -74,7 +74,7 @@ export function loadDescriptorProto(options: LoadDescriptorOptions): string {
 
     // 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)
+        f.name === protoFileName || f.name.endsWith('/' + protoFileName)
     );
 
     if (!fileDescriptor) {