@databricks/zerobus-ingest-sdk 1.1.0 → 1.2.0

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>,
 }
 
@@ -220,8 +221,8 @@ fn convert_js_to_record_payload(env: &Env, payload: Unknown) -> Result<RustRecor
 ///
 /// ```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 +401,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 +420,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 +465,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 +483,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 +531,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 +550,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 +583,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 +674,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 +712,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 +746,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 +830,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 +850,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 +879,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 +903,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)))?;
 
@@ -927,7 +968,7 @@ impl ZerobusSdk {
     ///   "", // ignored
     ///   undefined,
     ///   {
-    ///     getHeadersCallback: async () => [
+    ///     getHeadersCallback: () => [
     ///       ["authorization", `Bearer ${myToken}`],
     ///       ["x-databricks-zerobus-table-name", tableName]
     ///     ]
@@ -947,19 +988,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 +1117,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 +1134,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]
@@ -1438,9 +1491,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()?;
@@ -1686,7 +1737,8 @@ impl ZerobusSdk {
     ///
     /// # 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) {