@qualithm/arrow-flight-sql-js 0.0.1 → 0.1.0

package/dist/client.js

@@ -10,6 +10,7 @@ import { collectToTable, parseFlightData, tryParseSchema } from "./arrow";
 import { AuthenticationError, ConnectionError, FlightSqlError } from "./errors";
 import { getFlightServiceDefinition } from "./generated";
 import { encodeActionClosePreparedStatementRequest, encodeActionCreatePreparedStatementRequest, encodeCommandGetCatalogs, encodeCommandGetDbSchemas, encodeCommandGetExportedKeys, encodeCommandGetImportedKeys, encodeCommandGetPrimaryKeys, encodeCommandGetTables, encodeCommandGetTableTypes, encodeCommandPreparedStatementQuery, encodeCommandStatementQuery, encodeCommandStatementUpdate, getBytesField, parseProtoFields } from "./proto";
+import { SubscriptionMessageType, SubscriptionMode } from "./types";
 // Default configuration values
 const defaultConnectTimeoutMs = 30_000;
 const defaultRequestTimeoutMs = 60_000;
@@ -72,7 +73,7 @@ export class FlightSqlClient {
             const protocolPackage = flightPackage.protocol;
             this.flightService = protocolPackage.FlightService;
             // Create channel credentials
-            const credentials = this.options.credentials !== undefined ? this.options.credentials : this.createCredentials();
+            const credentials = this.options.credentials ?? this.createCredentials();
             // Create the gRPC client
             const address = `${this.options.host}:${String(this.options.port)}`;
             this.grpcClient = new this.flightService(address, credentials, {
@@ -261,7 +262,7 @@ export class FlightSqlClient {
         };
         const flightInfo = await this.getFlightInfo(descriptor);
         return this.fetchCatalogResults(flightInfo, (row) => ({
-            catalogName: row["catalog_name"]
+            catalogName: row.catalog_name
         }));
     }
     /**
@@ -292,8 +293,8 @@ export class FlightSqlClient {
         };
         const flightInfo = await this.getFlightInfo(descriptor);
         return this.fetchCatalogResults(flightInfo, (row) => ({
-            catalogName: row["catalog_name"],
-            schemaName: row["db_schema_name"]
+            catalogName: row.catalog_name,
+            schemaName: row.db_schema_name
         }));
     }
     /**
@@ -333,13 +334,13 @@ export class FlightSqlClient {
         const flightInfo = await this.getFlightInfo(descriptor);
         return this.fetchCatalogResults(flightInfo, (row) => {
             const info = {
-                catalogName: row["catalog_name"],
-                schemaName: row["db_schema_name"],
-                tableName: row["table_name"],
-                tableType: row["table_type"]
+                catalogName: row.catalog_name,
+                schemaName: row.db_schema_name,
+                tableName: row.table_name,
+                tableType: row.table_type
             };
             // If includeSchema was requested, parse the table schema
-            const schemaBytes = row["table_schema"];
+            const schemaBytes = row.table_schema;
             if (schemaBytes && schemaBytes.length > 0) {
                 info.schema = tryParseSchema(schemaBytes) ?? undefined;
             }
@@ -366,7 +367,7 @@ export class FlightSqlClient {
         };
         const flightInfo = await this.getFlightInfo(descriptor);
         return this.fetchCatalogResults(flightInfo, (row) => ({
-            tableType: row["table_type"]
+            tableType: row.table_type
         }));
     }
     /**
@@ -394,12 +395,12 @@ export class FlightSqlClient {
         };
         const flightInfo = await this.getFlightInfo(descriptor);
         return this.fetchCatalogResults(flightInfo, (row) => ({
-            catalogName: row["catalog_name"],
-            schemaName: row["db_schema_name"],
-            tableName: row["table_name"],
-            columnName: row["column_name"],
-            keySequence: row["key_sequence"],
-            keyName: row["key_name"]
+            catalogName: row.catalog_name,
+            schemaName: row.db_schema_name,
+            tableName: row.table_name,
+            columnName: row.column_name,
+            keySequence: row.key_sequence,
+            keyName: row.key_name
         }));
     }
     /**
@@ -462,42 +463,50 @@ export class FlightSqlClient {
         }
         for (const endpoint of flightInfo.endpoints) {
             for await (const flightData of this.doGet(endpoint.ticket)) {
-                if (flightData.dataHeader && flightData.dataBody) {
-                    const batch = parseFlightData(flightData.dataHeader, flightData.dataBody, schema);
-                    if (batch) {
-                        // Convert batch rows to objects
-                        for (let i = 0; i < batch.numRows; i++) {
-                            const row = {};
-                            for (const field of schema.fields) {
-                                const column = batch.getChild(field.name);
-                                row[field.name] = column?.get(i);
-                            }
-                            results.push(mapper(row));
-                        }
-                    }
+                if (!flightData.dataHeader || !flightData.dataBody) {
+                    continue;
+                }
+                const batch = parseFlightData(flightData.dataHeader, flightData.dataBody, schema);
+                if (!batch) {
+                    continue;
                 }
+                // Convert batch rows to objects
+                this.extractRowsFromBatch(batch, schema, mapper, results);
             }
         }
         return results;
     }
+    /**
+     * Helper to extract rows from a batch and map them
+     */
+    extractRowsFromBatch(batch, schema, mapper, results) {
+        for (let i = 0; i < batch.numRows; i++) {
+            const row = {};
+            for (const field of schema.fields) {
+                const column = batch.getChild(field.name);
+                row[field.name] = column?.get(i);
+            }
+            results.push(mapper(row));
+        }
+    }
     /**
      * Helper to fetch foreign key results with the complex schema
      */
     async fetchForeignKeyResults(flightInfo) {
         return this.fetchCatalogResults(flightInfo, (row) => ({
-            pkCatalogName: row["pk_catalog_name"],
-            pkSchemaName: row["pk_db_schema_name"],
-            pkTableName: row["pk_table_name"],
-            pkColumnName: row["pk_column_name"],
-            fkCatalogName: row["fk_catalog_name"],
-            fkSchemaName: row["fk_db_schema_name"],
-            fkTableName: row["fk_table_name"],
-            fkColumnName: row["fk_column_name"],
-            keySequence: row["key_sequence"],
-            fkKeyName: row["fk_key_name"],
-            pkKeyName: row["pk_key_name"],
-            updateRule: row["update_rule"],
-            deleteRule: row["delete_rule"]
+            pkCatalogName: row.pk_catalog_name,
+            pkSchemaName: row.pk_db_schema_name,
+            pkTableName: row.pk_table_name,
+            pkColumnName: row.pk_column_name,
+            fkCatalogName: row.fk_catalog_name,
+            fkSchemaName: row.fk_db_schema_name,
+            fkTableName: row.fk_table_name,
+            fkColumnName: row.fk_column_name,
+            keySequence: row.key_sequence,
+            fkKeyName: row.fk_key_name,
+            pkKeyName: row.pk_key_name,
+            updateRule: row.update_rule,
+            deleteRule: row.delete_rule
         }));
     }
     // ===========================================================================
@@ -622,13 +631,146 @@ export class FlightSqlClient {
             throw error;
         }
     }
+    /**
+     * Open a bidirectional data exchange with the server.
+     *
+     * This is the low-level API for DoExchange. For real-time subscriptions,
+     * use the higher-level `subscribe()` method instead.
+     *
+     * @param descriptor - Flight descriptor for the exchange
+     * @returns An exchange handle for sending and receiving FlightData
+     *
+     * @example
+     * ```typescript
+     * const exchange = client.doExchange({
+     *   type: DescriptorType.CMD,
+     *   cmd: new TextEncoder().encode('SUBSCRIBE:my_topic')
+     * })
+     *
+     * // Send initial request
+     * await exchange.send({ appMetadata: subscribeCommand })
+     *
+     * // Receive responses
+     * for await (const data of exchange) {
+     *   if (data.dataHeader) {
+     *     // Process record batch
+     *   }
+     * }
+     *
+     * // Clean up
+     * await exchange.end()
+     * ```
+     */
+    doExchange(descriptor) {
+        this.ensureConnected();
+        const client = this.grpcClient;
+        const metadata = this.createRequestMetadata();
+        const stream = client.doExchange(metadata);
+        // Use object wrapper to track errors
+        const errorState = { error: null };
+        stream.on("error", (err) => {
+            errorState.error = err;
+        });
+        // Capture references for closure
+        const wrapGrpcError = this.wrapGrpcError.bind(this);
+        const serializeFlightDescriptor = this.serializeFlightDescriptor.bind(this);
+        const wrapStream = this.wrapStream.bind(this);
+        // Create the exchange handle
+        const exchange = {
+            async send(data) {
+                if (errorState.error) {
+                    throw wrapGrpcError(errorState.error);
+                }
+                stream.write({
+                    flightDescriptor: data.flightDescriptor
+                        ? serializeFlightDescriptor(data.flightDescriptor)
+                        : undefined,
+                    dataHeader: data.dataHeader,
+                    dataBody: data.dataBody,
+                    appMetadata: data.appMetadata
+                });
+                return Promise.resolve();
+            },
+            async end() {
+                stream.end();
+                return Promise.resolve();
+            },
+            cancel() {
+                stream.cancel();
+            },
+            async *[Symbol.asyncIterator]() {
+                try {
+                    for await (const data of wrapStream(stream)) {
+                        const flightData = data;
+                        yield {
+                            flightDescriptor: flightData.flightDescriptor
+                                ? {
+                                    type: flightData.flightDescriptor.type,
+                                    cmd: flightData.flightDescriptor.cmd,
+                                    path: flightData.flightDescriptor.path
+                                }
+                                : undefined,
+                            dataHeader: flightData.dataHeader ?? new Uint8Array(),
+                            dataBody: flightData.dataBody ?? new Uint8Array(),
+                            appMetadata: flightData.appMetadata
+                        };
+                    }
+                }
+                catch (error) {
+                    if (errorState.error) {
+                        throw wrapGrpcError(errorState.error);
+                    }
+                    throw error;
+                }
+            }
+        };
+        // Send descriptor in first message (empty data)
+        stream.write({
+            flightDescriptor: this.serializeFlightDescriptor(descriptor)
+        });
+        return exchange;
+    }
+    /**
+     * Subscribe to real-time data updates from a query.
+     *
+     * Returns a Subscription that yields RecordBatches as they arrive from the server.
+     * Automatically handles heartbeats and can reconnect on connection loss.
+     *
+     * @param query - SQL query to subscribe to
+     * @param options - Subscription options
+     * @returns Subscription handle for receiving batches and control
+     *
+     * @example
+     * ```typescript
+     * const subscription = await client.subscribe(
+     *   "SELECT * FROM events WHERE status = 'pending'",
+     *   { mode: 'CHANGES_ONLY', heartbeatMs: 30000 }
+     * )
+     *
+     * for await (const batch of subscription) {
+     *   console.log(`Received ${batch.numRows} rows`)
+     * }
+     *
+     * // Or with AbortController
+     * const controller = new AbortController()
+     * const subscription = await client.subscribe(query, {
+     *   signal: controller.signal
+     * })
+     *
+     * // Later: cancel the subscription
+     * controller.abort()
+     * ```
+     */
+    subscribe(query, options = {}) {
+        return new Subscription(this, query, options);
+    }
     /**
      * Helper to get the first item from an async iterable without consuming the rest
      */
     async getFirstFromIterable(iterable) {
         const iterator = iterable[Symbol.asyncIterator]();
         const result = await iterator.next();
-        return result.done ? undefined : result.value;
+        return result.done === true ? undefined : result.value;
     }
     /**
      * Execute an action on the server.
@@ -752,7 +894,7 @@ export class FlightSqlClient {
     createRequestMetadata() {
         const metadata = new grpc.Metadata();
         // Add auth token if present
-        if (this.authToken) {
+        if (this.authToken !== null && this.authToken !== "") {
             metadata.set("authorization", `Bearer ${this.authToken}`);
         }
         // Add custom metadata from options
@@ -785,7 +927,20 @@ export class FlightSqlClient {
             case grpc.status.UNAVAILABLE:
             case grpc.status.DEADLINE_EXCEEDED:
                 return new ConnectionError(message, { cause: error });
-            default:
+            case grpc.status.OK:
+            case grpc.status.CANCELLED:
+            case grpc.status.UNKNOWN:
+            case grpc.status.INVALID_ARGUMENT:
+            case grpc.status.NOT_FOUND:
+            case grpc.status.ALREADY_EXISTS:
+            case grpc.status.PERMISSION_DENIED:
+            case grpc.status.RESOURCE_EXHAUSTED:
+            case grpc.status.FAILED_PRECONDITION:
+            case grpc.status.ABORTED:
+            case grpc.status.OUT_OF_RANGE:
+            case grpc.status.UNIMPLEMENTED:
+            case grpc.status.INTERNAL:
+            case grpc.status.DATA_LOSS:
                 return new FlightSqlError(message, { cause: error });
         }
     }
@@ -916,7 +1071,7 @@ export class QueryResult {
             const framedParts = [];
             for await (const flightData of this.client.doGet(endpoint.ticket)) {
                 if (flightData.dataHeader && flightData.dataHeader.length > 0) {
-                    const framed = this.frameAsIPC(flightData.dataHeader, flightData.dataBody || new Uint8Array(0));
+                    const framed = this.frameAsIPC(flightData.dataHeader, flightData.dataBody ?? new Uint8Array(0));
                     framedParts.push(framed);
                 }
             }
@@ -1084,4 +1239,289 @@ export class PreparedStatement {
         this.closed = true;
     }
 }
+// ============================================================================
+// Subscription
+// ============================================================================
+/**
+ * Real-time subscription to query results via DoExchange.
+ *
+ * Yields RecordBatches as they arrive from the server. Handles heartbeats
+ * automatically and can reconnect on transient connection failures.
+ *
+ * @example
+ * ```typescript
+ * const subscription = client.subscribe("SELECT * FROM events")
+ *
+ * for await (const batch of subscription) {
+ *   console.log(`Received ${batch.numRows} rows`)
+ * }
+ *
+ * await subscription.unsubscribe()
+ * ```
+ */
+export class Subscription {
+    client;
+    query;
+    options;
+    exchange = null;
+    subscriptionId = null;
+    connectedState = false;
+    batchesReceivedCount = 0;
+    reconnectAttempts = 0;
+    abortedFlag = false;
+    lastHeartbeat = Date.now();
+    iterating = false;
+    /** Check aborted state - method prevents ESLint from static flow analysis across async boundaries */
+    isAborted() {
+        return this.abortedFlag;
+    }
+    constructor(client, query, options = {}) {
+        this.client = client;
+        this.query = query;
+        this.options = {
+            ...options,
+            mode: options.mode ?? SubscriptionMode.ChangesOnly,
+            heartbeatMs: options.heartbeatMs ?? 30_000,
+            autoReconnect: options.autoReconnect ?? true,
+            maxReconnectAttempts: options.maxReconnectAttempts ?? 10,
+            reconnectDelayMs: options.reconnectDelayMs ?? 1_000,
+            maxReconnectDelayMs: options.maxReconnectDelayMs ?? 30_000
+        };
+        // Handle abort signal
+        if (this.options.signal) {
+            this.options.signal.addEventListener("abort", () => {
+                this.abortedFlag = true;
+                this.exchange?.cancel();
+            });
+        }
+    }
+    /**
+     * Unique ID for this subscription (assigned by server after connect)
+     */
+    get id() {
+        return this.subscriptionId ?? "";
+    }
+    /**
+     * Whether the subscription is currently connected
+     */
+    get connected() {
+        return this.connectedState;
+    }
+    /**
+     * Number of batches received
+     */
+    get batchesReceived() {
+        return this.batchesReceivedCount;
+    }
+    /**
+     * Number of reconnection attempts
+     */
+    get reconnectCount() {
+        return this.reconnectAttempts;
+    }
+    /**
+     * Timestamp of the last heartbeat received from the server
+     */
+    get lastHeartbeatTime() {
+        return this.lastHeartbeat;
+    }
+    /**
+     * Time in milliseconds since the last heartbeat
+     */
+    get timeSinceLastHeartbeat() {
+        return Date.now() - this.lastHeartbeat;
+    }
+    /**
+     * Start the subscription and iterate over incoming batches
+     */
+    async *[Symbol.asyncIterator]() {
+        if (this.iterating) {
+            throw new FlightSqlError("Subscription is already being iterated");
+        }
+        this.iterating = true;
+        try {
+            while (!this.isAborted()) {
+                try {
+                    yield* this.streamBatches();
+                    // Stream completed normally
+                    break;
+                }
+                catch (error) {
+                    if (this.isAborted()) {
+                        break;
+                    }
+                    if (!this.options.autoReconnect) {
+                        throw error;
+                    }
+                    if (this.reconnectAttempts >= this.options.maxReconnectAttempts) {
+                        throw new FlightSqlError(`Max reconnection attempts (${String(this.options.maxReconnectAttempts)}) exceeded`, { cause: error instanceof Error ? error : new Error(String(error)) });
+                    }
+                    await this.reconnect();
+                }
+            }
+        }
+        finally {
+            this.iterating = false;
+            await this.cleanup();
+        }
+    }
+    /**
+     * Unsubscribe and close the connection
+     */
+    async unsubscribe() {
+        this.abortedFlag = true;
+        await this.cleanup();
+    }
+    async *streamBatches() {
+        this.initConnection();
+        if (!this.exchange) {
+            throw new FlightSqlError("Exchange not initialized");
+        }
+        for await (const flightData of this.exchange) {
+            if (this.isAborted()) {
+                break;
+            }
+            // Check for metadata messages (heartbeat, control messages)
+            if (flightData.appMetadata && flightData.appMetadata.length > 0) {
+                const metadata = this.parseMetadata(flightData.appMetadata);
+                switch (metadata.type) {
+                    case SubscriptionMessageType.HEARTBEAT:
+                        this.lastHeartbeat = Date.now();
+                        continue;
+                    case SubscriptionMessageType.COMPLETE:
+                        return;
+                    case SubscriptionMessageType.ERROR:
+                        throw new FlightSqlError(metadata.error ?? "Subscription error from server");
+                    case SubscriptionMessageType.DATA:
+                        // Continue to parse data
+                        if (metadata.subscriptionId !== undefined && this.subscriptionId === null) {
+                            this.subscriptionId = metadata.subscriptionId;
+                        }
+                        break;
+                    case SubscriptionMessageType.SUBSCRIBE:
+                    case SubscriptionMessageType.UNSUBSCRIBE:
+                        // These are client-to-server messages, ignore if received
+                        break;
+                }
+            }
+            // Parse data if present
+            if (flightData.dataHeader && flightData.dataHeader.length > 0) {
+                const batches = this.parseFlightDataToBatches(flightData);
+                for (const batch of batches) {
+                    this.batchesReceivedCount++;
+                    yield batch;
+                }
+            }
+        }
+    }
+    initConnection() {
+        const descriptor = {
+            type: 2, // CMD
+            cmd: this.encodeSubscribeCommand()
+        };
+        this.exchange = this.client.doExchange(descriptor);
+        this.connectedState = true;
+        this.lastHeartbeat = Date.now();
+    }
+    async reconnect() {
+        this.reconnectAttempts++;
+        this.connectedState = false;
+        // Calculate backoff delay with jitter
+        const baseDelay = Math.min(this.options.reconnectDelayMs * Math.pow(2, this.reconnectAttempts - 1), this.options.maxReconnectDelayMs);
+        const jitter = Math.random() * baseDelay * 0.1;
+        const delay = baseDelay + jitter;
+        await this.sleep(delay);
+        this.initConnection();
+    }
+    async cleanup() {
+        if (this.exchange && this.connectedState) {
+            try {
+                // Send unsubscribe message
+                await this.exchange.send({
+                    dataHeader: new Uint8Array(),
+                    dataBody: new Uint8Array(),
+                    appMetadata: this.encodeMetadata({
+                        type: SubscriptionMessageType.UNSUBSCRIBE,
+                        subscriptionId: this.subscriptionId ?? undefined,
+                        timestamp: Date.now()
+                    })
+                });
+                await this.exchange.end();
+            }
+            catch {
+                // Ignore cleanup errors
+            }
+        }
+        this.connectedState = false;
+        this.exchange = null;
+    }
+    encodeSubscribeCommand() {
+        // Encode subscription request as JSON in the command
+        const request = {
+            type: SubscriptionMessageType.SUBSCRIBE,
+            query: this.query,
+            mode: this.options.mode,
+            heartbeatMs: this.options.heartbeatMs,
+            metadata: this.options.metadata
+        };
+        return new TextEncoder().encode(JSON.stringify(request));
+    }
+    encodeMetadata(metadata) {
+        return new TextEncoder().encode(JSON.stringify(metadata));
+    }
+    parseMetadata(data) {
+        try {
+            const text = new TextDecoder().decode(data);
+            return JSON.parse(text);
+        }
+        catch {
+            return { type: SubscriptionMessageType.DATA };
+        }
+    }
+    parseFlightDataToBatches(flightData) {
+        const batches = [];
+        if (!flightData.dataHeader || flightData.dataHeader.length === 0) {
+            return batches;
+        }
+        try {
+            // Frame the raw flatbuffer with IPC continuation marker
+            const framed = this.frameAsIPC(flightData.dataHeader, flightData.dataBody ?? new Uint8Array(0));
+            const reader = RecordBatchReader.from(framed);
+            for (const batch of reader) {
+                batches.push(batch);
+            }
+        }
+        catch {
+            // If parsing fails, return empty
+        }
+        return batches;
+    }
+    frameAsIPC(header, body) {
+        const continuationToken = new Uint8Array([0xff, 0xff, 0xff, 0xff]);
+        const metadataLength = new Uint8Array(4);
+        new DataView(metadataLength.buffer).setInt32(0, header.length, true);
+        const headerPadding = (8 - ((header.length + 4 + 4) % 8)) % 8;
+        const padding = new Uint8Array(headerPadding);
+        const totalLength = continuationToken.length +
+            metadataLength.length +
+            header.length +
+            padding.length +
+            body.length;
+        const result = new Uint8Array(totalLength);
+        let offset = 0;
+        result.set(continuationToken, offset);
+        offset += continuationToken.length;
+        result.set(metadataLength, offset);
+        offset += metadataLength.length;
+        result.set(header, offset);
+        offset += header.length;
+        result.set(padding, offset);
+        offset += padding.length;
+        result.set(body, offset);
+        return result;
+    }
+    async sleep(ms) {
+        return new Promise((resolve) => setTimeout(resolve, ms));
+    }
+}
 //# sourceMappingURL=client.js.map