@bitofsky/databricks-sql 1.0.1 → 1.0.4

package/dist/index.js

@@ -236,18 +236,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,
@@ -256,25 +260,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")
@@ -581,52 +610,110 @@ function convertBoolean(value) {
 import { PassThrough } from "stream";
 import { mergeStreamsFromUrls } from "@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 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 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;
@@ -634,31 +721,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);
@@ -677,10 +779,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(parser()).pipe(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");
     }
@@ -698,44 +804,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: [