@bitofsky/databricks-sql 1.0.1 → 1.0.4

package/README.md

@@ -23,6 +23,7 @@ The goal is simple: stream big results with stable memory usage and without forc
 - Optimized polling with server-side wait (up to 50s) before falling back to client polling.
 - Query metrics support via Query History API (`enableMetrics` option).
 - Efficient external link handling: merge chunks into a single stream.
+- Handles partial external link responses by fetching missing chunk metadata.
 - `mergeExternalLinks` supports streaming uploads and returns a new StatementResult with a presigned URL.
 - `fetchRow`/`fetchAll` support `JSON_OBJECT` (schema-based row mapping).
 - External links + JSON_ARRAY are supported for row iteration (streaming JSON parsing).
@@ -48,12 +49,16 @@ console.log(rows) // [{ value: 1 }]
 ```
 
 ## Sample (Streaming + Presigned URL)
-Stream external links into S3, then return a single presigned URL:
+Stream external links into S3 with gzip compression, then return a single presigned URL.
 
 ```ts
 import { executeStatement, mergeExternalLinks } from '@bitofsky/databricks-sql'
-import { GetObjectCommand, PutObjectCommand, S3Client } from '@aws-sdk/client-s3'
+import { GetObjectCommand, HeadObjectCommand, S3Client } from '@aws-sdk/client-s3'
+import { Upload } from '@aws-sdk/lib-storage'
 import { getSignedUrl } from '@aws-sdk/s3-request-presigner'
+import { createGzip } from 'zlib'
+import { pipeline } from 'stream/promises'
+import { PassThrough } from 'stream'
 
 const auth = {
   token: process.env.DATABRICKS_TOKEN!,
@@ -72,31 +77,42 @@ const result = await executeStatement(
 
 const merged = await mergeExternalLinks(result, auth, {
   mergeStreamToExternalLink: async (stream) => {
-    const key = `merged-${Date.now()}.csv`
-    await s3.send(
-      new PutObjectCommand({
+    const key = `merged-${Date.now()}.csv.gz`
+    const gzip = createGzip() // Compress with gzip and upload to S3
+    const passThrough = new PassThrough()
+
+    const upload = new Upload({
+      client: s3,
+      params: {
         Bucket: bucket,
         Key: key,
-        Body: stream,
-        ContentType: 'text/csv',
-      })
-    )
-
-    const externalLink = await getSignedUrl(
-      s3,
-      new GetObjectCommand({ Bucket: bucket, Key: key }),
-      { expiresIn: 3600 }
-    )
+        Body: passThrough,
+        ContentType: 'text/csv; charset=utf-8',
+        ContentEncoding: 'gzip',
+      },
+    })
+    const uploadPromise = upload.done()
+
+    await Promise.all([
+      pipeline(stream, gzip, passThrough),
+      uploadPromise,
+    ])
+
+    // Get actual uploaded size via HeadObject
+    const head = await s3.send(new HeadObjectCommand({ Bucket: bucket, Key: key }))
+    // Generate presigned URL valid for 1 hour
+    const externalLink = await getSignedUrl(s3, new GetObjectCommand({ Bucket: bucket, Key: key }),{ expiresIn: 3600 })
 
     return {
-      externalLink,
-      byte_count: 0,
-      expiration: new Date(Date.now() + 3600 * 1000).toISOString(),
+      externalLink, // Presigned URL to merged gzip CSV
+      byte_count: head.ContentLength ?? 0, // Actual compressed size
+      expiration: new Date(Date.now() + 3600 * 1000).toISOString(), // 1 hour from now
     }
   },
 })
 
-console.log(merged.result?.external_links?.[0].external_link) // Presigned URL to merged CSV
+console.log(merged.result?.external_links?.[0].external_link) // Presigned URL to merged gzip CSV
+console.log(merged.result?.external_links?.[0].byte_count)    // Actual compressed size
 ```
 
 ## Sample (Progress with Metrics)
@@ -116,8 +132,8 @@ const result = await executeStatement(
   auth,
   {
     enableMetrics: true,
-    onProgress: (status, metrics) => {
-      console.log(`State: ${status.state}`)
+    onProgress: (result, metrics) => {
+      console.log(`State: ${result.status.state}`)
       if (metrics) {  // metrics is optional, only present when enableMetrics: true
         console.log(`  Execution time: ${metrics.execution_time_ms}ms`)
         console.log(`  Rows produced: ${metrics.rows_produced_count}`)
@@ -178,6 +194,7 @@ function executeStatement(
 ```
 - Calls the Databricks Statement Execution API and polls until completion.
 - Server waits up to 50s (`wait_timeout`) before client-side polling begins.
+- Default `wait_timeout` is `50s`, or `0s` when `onProgress` is provided.
 - Use `options.onProgress` to receive status updates with optional metrics.
 - Set `enableMetrics: true` to fetch query metrics from Query History API on each poll.
 - Throws `DatabricksSqlError` on failure, `StatementCancelledError` on cancel, and `AbortError` on abort.
@@ -193,6 +210,7 @@ function fetchRow(
 - Streams each row to `options.onEachRow`.
 - Use `format: 'JSON_OBJECT'` to map rows into schema-based objects.
 - Supports `INLINE` results or `JSON_ARRAY` formatted `EXTERNAL_LINKS` only.
+- If only a subset of external links is returned, missing chunk metadata is fetched by index.
 
 ### fetchAll(statementResult, auth, options?)
 ```ts
@@ -204,6 +222,7 @@ function fetchAll(
 ```
 - Collects all rows into an array. For large results, prefer `fetchRow`/`fetchStream`.
 - Supports `INLINE` results or `JSON_ARRAY` formatted `EXTERNAL_LINKS` only.
+- If only a subset of external links is returned, missing chunk metadata is fetched by index.
 
 ### fetchStream(statementResult, auth, options?)
 ```ts
@@ -218,6 +237,7 @@ function fetchStream(
 - Throws if the result is `INLINE`.
 - Ends as an empty stream when no external links exist.
 - `forceMerge: true` forces merge even when there is only a single external link.
+- If only a subset of external links is returned, missing chunk metadata is fetched by index.
 
 ### mergeExternalLinks(statementResult, auth, options)
 ```ts
@@ -236,8 +256,9 @@ function mergeExternalLinks(
 ### Options (Summary)
 ```ts
 type ExecuteStatementOptions = {
-  onProgress?: (status: StatementStatus, metrics?: QueryMetrics) => void
+  onProgress?: (result: StatementResult, metrics?: QueryMetrics) => void
   enableMetrics?: boolean      // Fetch metrics from Query History API (default: false)
+  logger?: Logger
   signal?: AbortSignal
   disposition?: 'INLINE' | 'EXTERNAL_LINKS'
   format?: 'JSON_ARRAY' | 'ARROW_STREAM' | 'CSV'
@@ -255,21 +276,25 @@ type FetchRowsOptions = {
   signal?: AbortSignal
   onEachRow?: (row: RowArray | RowObject) => void
   format?: 'JSON_ARRAY' | 'JSON_OBJECT'
+  logger?: Logger
 }
 
 type FetchAllOptions = {
   signal?: AbortSignal
   format?: 'JSON_ARRAY' | 'JSON_OBJECT'
+  logger?: Logger
 }
 
 type FetchStreamOptions = {
   signal?: AbortSignal
   forceMerge?: boolean
+  logger?: Logger
 }
 
 type MergeExternalLinksOptions = {
   signal?: AbortSignal
   forceMerge?: boolean
+  logger?: Logger
   mergeStreamToExternalLink: (stream: Readable) => Promise<{
     externalLink: string
     byte_count: number

package/dist/index.cjs

@@ -272,18 +272,22 @@ var TERMINAL_STATES = /* @__PURE__ */ new Set([
 ]);
 var POLL_INTERVAL_MS = 5e3;
 async function fetchMetrics(auth, statementId, signal) {
-  try {
-    const queryInfo = await getQueryMetrics(auth, statementId, signal);
-    return queryInfo.metrics;
-  } catch {
-    return void 0;
-  }
+  const queryInfo = await getQueryMetrics(auth, statementId, signal);
+  return queryInfo.metrics;
 }
 async function executeStatement(query, auth, options = {}) {
   const warehouseId = options.warehouse_id ?? extractWarehouseId(auth.httpPath);
-  const { signal, onProgress, enableMetrics } = options;
+  const { signal, onProgress, enableMetrics, logger } = options;
+  const waitTimeout = options.wait_timeout ?? (onProgress ? "0s" : "50s");
+  let cancelIssued = false;
   throwIfAborted(signal, "executeStatement");
-  const emitProgress = onProgress ? async (statementId) => onProgress(result.status, enableMetrics ? await fetchMetrics(auth, statementId, signal) : void 0) : void 0;
+  const emitProgress = onProgress ? async () => result ? onProgress(
+    result,
+    enableMetrics ? await fetchMetrics(auth, result.statement_id, signal).catch((e) => {
+      logger?.error?.(`executeStatement Failed to fetch query metrics for statement ${result?.statement_id}: ${String(e)}`, { statementId: result?.statement_id });
+      return void 0;
+    }) : void 0
+  ) : void 0 : void 0;
   const request = Object.fromEntries(
     Object.entries({
       warehouse_id: warehouseId,
@@ -292,25 +296,50 @@ async function executeStatement(query, auth, options = {}) {
       disposition: options.disposition,
       format: options.format,
       on_wait_timeout: options.on_wait_timeout ?? "CONTINUE",
-      wait_timeout: options.wait_timeout ?? "50s",
+      wait_timeout: waitTimeout,
       row_limit: options.row_limit,
       catalog: options.catalog,
       schema: options.schema,
       parameters: options.parameters
     }).filter(([, v]) => v !== void 0)
   );
+  logger?.info?.(`executeStatement Executing statement on warehouse ${warehouseId}...`);
   let result = await postStatement(auth, request, signal);
-  while (!TERMINAL_STATES.has(result.status.state)) {
-    if (signal?.aborted) {
-      await cancelStatement(auth, result.statement_id).catch(() => {
-      });
+  const cancelStatementSafely = async () => {
+    if (cancelIssued) return;
+    logger?.info?.("executeStatement Abort signal received during executeStatement.");
+    cancelIssued = true;
+    await cancelStatement(auth, result.statement_id).catch((err) => {
+      logger?.error?.("executeStatement Failed to cancel statement after abort.", err);
+    });
+  };
+  if (signal?.aborted) {
+    await cancelStatementSafely();
+    throw new AbortError("Aborted during polling");
+  }
+  const onAbort = () => cancelStatementSafely().catch(() => {
+  });
+  try {
+    signal?.addEventListener("abort", onAbort, { once: true });
+    while (!TERMINAL_STATES.has(result.status.state)) {
+      logger?.info?.(`executeStatement Statement ${result.statement_id} in state ${result.status.state}; polling for status...`);
+      await delay(POLL_INTERVAL_MS, signal);
+      result = await getStatement(auth, result.statement_id, signal);
+      await emitProgress?.();
+    }
+  } catch (err) {
+    if (err instanceof AbortError || signal?.aborted) {
+      logger?.info?.("executeStatement Abort detected in executeStatement polling loop.");
+      await cancelStatementSafely();
       throw new AbortError("Aborted during polling");
     }
-    await emitProgress?.(result.statement_id);
-    await delay(POLL_INTERVAL_MS, signal);
-    result = await getStatement(auth, result.statement_id, signal);
+    logger?.error?.(`executeStatement Error during executeStatement polling: ${String(err)}`);
+    throw err;
+  } finally {
+    logger?.info?.(`executeStatement Statement ${result.statement_id} reached final state: ${result.status.state}`);
+    signal?.removeEventListener("abort", onAbort);
   }
-  await emitProgress?.(result.statement_id);
+  await emitProgress?.();
   if (result.status.state === "SUCCEEDED")
     return result;
   if (result.status.state === "CANCELED")
@@ -617,52 +646,110 @@ function convertBoolean(value) {
 var import_node_stream2 = require("stream");
 var import_merge_streams = require("@bitofsky/merge-streams");
 function fetchStream(statementResult, auth, options = {}) {
-  const { signal, forceMerge } = options;
+  const { signal, forceMerge, logger } = options;
   const manifest = validateSucceededResult(statementResult);
   const format = manifest.format;
+  const statementId = statementResult.statement_id;
+  const baseLog = { statementId, manifest, format, forceMerge };
   if (statementResult.result?.data_array) {
+    logger?.error?.(
+      `fetchStream only supports EXTERNAL_LINKS results for statement ${statementId}.`,
+      { ...baseLog, hasDataArray: true }
+    );
     throw new DatabricksSqlError(
       "fetchStream only supports EXTERNAL_LINKS results",
       "UNSUPPORTED_FORMAT",
-      statementResult.statement_id
+      statementId
     );
   }
+  logger?.info?.(`fetchStream creating stream for statement ${statementId}.`, {
+    ...baseLog,
+    hasExternalLinks: Boolean(statementResult.result?.external_links?.length)
+  });
   const output = new import_node_stream2.PassThrough();
   if (signal) {
-    const onAbort = () => output.destroy(new AbortError("Stream aborted"));
+    const onAbort = () => {
+      logger?.info?.(`fetchStream abort signal received while streaming statement ${statementId}.`, baseLog);
+      output.destroy(new AbortError("Stream aborted"));
+    };
     signal.addEventListener("abort", onAbort, { once: true });
     output.once("close", () => signal.removeEventListener("abort", onAbort));
   }
-  mergeChunksToStream(statementResult, auth, manifest, format, output, signal, forceMerge).catch((err) => output.destroy(err));
+  output.on("error", (err) => {
+    if (err instanceof AbortError)
+      return;
+    if (output.listenerCount("error") === 1)
+      throw err;
+  });
+  mergeChunksToStream(statementResult, auth, manifest, format, output, signal, forceMerge, logger).catch((err) => {
+    logger?.error?.(`fetchStream error while streaming statement ${statementId}.`, {
+      ...baseLog,
+      error: err
+    });
+    output.destroy(err);
+  });
   return output;
 }
-async function mergeChunksToStream(statementResult, auth, manifest, format, output, signal, forceMerge) {
+async function mergeChunksToStream(statementResult, auth, manifest, format, output, signal, forceMerge, logger) {
+  const statementId = statementResult.statement_id;
+  const baseLog = { statementId, manifest, format, forceMerge };
+  logger?.info?.(`fetchStream collecting external links for statement ${statementId}.`, baseLog);
   const urls = await collectExternalUrls(statementResult, auth, manifest, signal);
-  if (urls.length === 0)
+  if (urls.length === 0) {
+    logger?.info?.(`fetchStream no external links found for statement ${statementId}.`, baseLog);
     return void output.end();
-  if (urls.length === 1 && !forceMerge)
+  }
+  if (urls.length === 1 && !forceMerge) {
+    logger?.info?.(`fetchStream piping single external link for statement ${statementId}.`, {
+      ...baseLog,
+      urlCount: urls.length
+    });
     return pipeUrlToOutput(urls[0], output, signal);
+  }
+  logger?.info?.(`fetchStream merging ${urls.length} external links for statement ${statementId}.`, {
+    ...baseLog,
+    urlCount: urls.length
+  });
   return (0, import_merge_streams.mergeStreamsFromUrls)(format, signal ? { urls, output, signal } : { urls, output });
 }
 async function collectExternalUrls(statementResult, auth, manifest, signal) {
-  const urls = extractExternalLinks(statementResult.result?.external_links);
-  if (urls.length > 0)
-    return urls;
+  const chunkUrls = /* @__PURE__ */ new Map();
+  addChunkLinks(chunkUrls, statementResult.result?.external_links);
   if (!manifest.total_chunk_count)
-    return [];
-  const chunkUrls = [];
+    return flattenChunkUrls(chunkUrls);
   for (let i = 0; i < manifest.total_chunk_count; i++) {
+    if (chunkUrls.has(i))
+      continue;
     if (signal?.aborted)
       throw new AbortError("Aborted while collecting URLs");
     const chunkData = await getChunk(auth, statementResult.statement_id, i, signal);
-    chunkUrls.push(...extractExternalLinks(chunkData.external_links));
+    addChunkLinks(chunkUrls, chunkData.external_links);
   }
-  return chunkUrls;
+  return flattenChunkUrls(chunkUrls);
 }
-function extractExternalLinks(externalLinks) {
+function addChunkLinks(chunkUrls, externalLinks) {
   if (!externalLinks)
+    return;
+  for (const link of externalLinks) {
+    if (!isNonEmptyString(link.external_link))
+      continue;
+    const existing = chunkUrls.get(link.chunk_index);
+    if (existing) {
+      existing.push(link.external_link);
+    } else {
+      chunkUrls.set(link.chunk_index, [link.external_link]);
+    }
+  }
+}
+function flattenChunkUrls(chunkUrls) {
+  if (chunkUrls.size === 0)
     return [];
-  return externalLinks.map((link) => link.external_link).filter(isNonEmptyString);
+  const sorted = [...chunkUrls.entries()].sort(([a], [b]) => a - b);
+  const urls = [];
+  for (const [, links] of sorted) {
+    urls.push(...links);
+  }
+  return urls;
 }
 function isNonEmptyString(value) {
   return typeof value === "string" && value.length > 0;
@@ -670,31 +757,46 @@ function isNonEmptyString(value) {
 
 // src/api/fetchRow.ts
 async function fetchRow(statementResult, auth, options = {}) {
-  const { signal, onEachRow, format } = options;
+  const { signal, onEachRow, format, logger } = options;
   const manifest = validateSucceededResult(statementResult);
+  const statementId = statementResult.statement_id;
+  const logContext = { statementId, manifest, requestedFormat: format };
   const mapRow = createRowMapper(manifest, format);
+  logger?.info?.(`fetchRow fetching rows for statement ${statementId}.`, {
+    ...logContext,
+    resultType: statementResult.result?.external_links ? "EXTERNAL_LINKS" : "INLINE"
+  });
   if (statementResult.result?.external_links) {
     if (manifest.format !== "JSON_ARRAY") {
+      logger?.error?.(`fetchRow only supports JSON_ARRAY for external_links; got ${manifest.format}.`, logContext);
       throw new DatabricksSqlError(
         `fetchRow only supports JSON_ARRAY for external_links. Received: ${manifest.format}`,
         "UNSUPPORTED_FORMAT",
-        statementResult.statement_id
+        statementId
       );
     }
-    const stream = fetchStream(statementResult, auth, signal ? { signal } : {});
-    await consumeJsonArrayStream(stream, mapRow, onEachRow, signal);
+    logger?.info?.(`fetchRow streaming external links for statement ${statementId}.`, logContext);
+    const stream = fetchStream(statementResult, auth, {
+      ...signal ? { signal } : {},
+      ...logger ? { logger } : {}
+    });
+    await consumeJsonArrayStream(stream, mapRow, onEachRow, signal, logger, logContext);
     return;
   }
   const totalChunks = manifest.total_chunk_count;
   const dataArray = statementResult.result?.data_array;
   if (dataArray) {
+    logger?.info?.(`fetchRow processing inline rows for statement ${statementId}.`, {
+      ...logContext,
+      inlineRows: dataArray.length
+    });
     for (const row of dataArray) {
       if (signal?.aborted) throw new AbortError("Aborted");
       onEachRow?.(mapRow(row));
     }
   }
   if (totalChunks > 1) {
-    const statementId = statementResult.statement_id;
+    logger?.info?.(`fetchRow processing ${totalChunks} chunks for statement ${statementId}.`, logContext);
     for (let chunkIndex = 1; chunkIndex < totalChunks; chunkIndex++) {
       if (signal?.aborted) throw new AbortError("Aborted");
       const chunk = await getChunk(auth, statementId, chunkIndex, signal);
@@ -713,10 +815,14 @@ async function fetchRow(statementResult, auth, options = {}) {
     }
   }
 }
-async function consumeJsonArrayStream(stream, mapRow, onEachRow, signal) {
+async function consumeJsonArrayStream(stream, mapRow, onEachRow, signal, logger, logContext) {
   const jsonStream = stream.pipe((0, import_stream_json.parser)()).pipe((0, import_StreamArray.streamArray)());
   for await (const item of jsonStream) {
     if (signal?.aborted) {
+      logger?.info?.("fetchRow abort detected while streaming JSON_ARRAY rows.", {
+        ...logContext,
+        aborted: signal.aborted
+      });
       stream.destroy(new AbortError("Aborted"));
       throw new AbortError("Aborted");
     }
@@ -734,44 +840,74 @@ async function consumeJsonArrayStream(stream, mapRow, onEachRow, signal) {
 // src/api/fetchAll.ts
 async function fetchAll(statementResult, auth, options = {}) {
   const rows = [];
+  const statementId = statementResult.statement_id;
+  const manifest = statementResult.manifest;
+  const logContext = { statementId, manifest, requestedFormat: options.format };
   const fetchOptions = {
     // Collect rows as they are streamed in.
     onEachRow: (row) => {
       rows.push(row);
     }
   };
+  const { logger } = options;
+  logger?.info?.(`fetchAll fetching all rows for statement ${statementId}.`, logContext);
   if (options.signal)
     fetchOptions.signal = options.signal;
   if (options.format)
     fetchOptions.format = options.format;
+  if (options.logger)
+    fetchOptions.logger = options.logger;
   await fetchRow(statementResult, auth, fetchOptions);
+  logger?.info?.(`fetchAll fetched ${rows.length} rows for statement ${statementId}.`, {
+    ...logContext,
+    rowCount: rows.length,
+    resolvedFormat: options.format ?? manifest?.format
+  });
   return rows;
 }
 
 // src/api/mergeExternalLinks.ts
 async function mergeExternalLinks(statementResult, auth, options) {
-  const { signal, mergeStreamToExternalLink, forceMerge } = options;
-  if (!statementResult.result?.external_links)
+  const { signal, mergeStreamToExternalLink, forceMerge, logger } = options;
+  const statementId = statementResult.statement_id;
+  const manifest = statementResult.manifest;
+  const externalLinks = statementResult.result?.external_links;
+  const totalChunks = manifest?.total_chunk_count ?? 0;
+  const logContext = { statementId, manifest, totalChunks, forceMerge };
+  if (!externalLinks) {
+    logger?.info?.(`mergeExternalLinks no external links to merge for statement ${statementId}.`, logContext);
     return statementResult;
+  }
   if (!forceMerge) {
-    const totalChunks = statementResult.manifest?.total_chunk_count;
-    const externalLinks = statementResult.result.external_links;
-    const isSingleChunk = totalChunks === void 0 ? externalLinks.length <= 1 : totalChunks <= 1;
-    if (isSingleChunk && externalLinks.length <= 1)
+    const isSingleChunk = totalChunks <= 1;
+    if (isSingleChunk) {
+      logger?.info?.(`mergeExternalLinks skipping merge for single external link in statement ${statementId}.`, {
+        ...logContext,
+        totalChunks
+      });
       return statementResult;
+    }
   }
+  logger?.info?.(`mergeExternalLinks merging external links for statement ${statementId}.`, logContext);
   const stream = fetchStream(statementResult, auth, {
     ...signal ? { signal } : {},
-    ...forceMerge !== void 0 ? { forceMerge } : {}
+    ...forceMerge !== void 0 ? { forceMerge } : {},
+    ...logger ? { logger } : {}
   });
+  logger?.info?.(`mergeExternalLinks uploading merged external link for statement ${statementId}.`, logContext);
   const uploadResult = await mergeStreamToExternalLink(stream);
-  const manifest = validateSucceededResult(statementResult);
-  const totalRowCount = manifest.total_row_count ?? 0;
+  logger?.info?.(`mergeExternalLinks uploaded merged external link for statement ${statementId}.`, {
+    ...logContext,
+    byteCount: uploadResult.byte_count,
+    expiration: uploadResult.expiration
+  });
+  const validatedManifest = validateSucceededResult(statementResult);
+  const totalRowCount = validatedManifest.total_row_count ?? 0;
   return {
     statement_id: statementResult.statement_id,
     status: statementResult.status,
     manifest: {
-      ...manifest,
+      ...validatedManifest,
       total_chunk_count: 1,
       total_byte_count: uploadResult.byte_count,
       chunks: [