@bitofsky/databricks-sql 1.0.0 → 1.0.3

package/dist/index.js

@@ -48,6 +48,8 @@ var RateLimitError = class extends HttpError {
 };
 
 // src/util.ts
+import { Readable } from "stream";
+import { pipeline } from "stream/promises";
 function extractWarehouseId(httpPath) {
   const match = httpPath.match(/\/sql\/\d+\.\d+\/warehouses\/([a-zA-Z0-9]+)/);
   if (!match?.[1])
@@ -97,6 +99,24 @@ function validateSucceededResult(statementResult) {
     );
   return statementResult.manifest;
 }
+function isWebReadableStream(body) {
+  return typeof body.getReader === "function";
+}
+async function pipeUrlToOutput(url, output, signal) {
+  if (signal?.aborted)
+    throw new AbortError("Aborted while streaming");
+  const response = await fetch(url, signal ? { signal } : void 0);
+  if (!response.ok) {
+    throw new Error(
+      `Failed to fetch external link: ${response.status} ${response.statusText}`
+    );
+  }
+  if (!response.body)
+    return void output.end();
+  const body = response.body;
+  const input = isWebReadableStream(body) ? Readable.fromWeb(body) : body;
+  await pipeline(input, output);
+}
 
 // src/http.ts
 var MAX_RETRIES = 3;
@@ -169,6 +189,7 @@ async function httpRequest(auth, options) {
 
 // src/databricks-api.ts
 var BASE_PATH = "/api/2.0/sql/statements";
+var HISTORY_BASE_PATH = "/api/2.0/sql/history/queries";
 async function postStatement(auth, request, signal) {
   return httpRequest(auth, {
     method: "POST",
@@ -198,6 +219,13 @@ async function getChunk(auth, statementId, chunkIndex, signal) {
     ...signal ? { signal } : {}
   });
 }
+async function getQueryMetrics(auth, queryId, signal) {
+  return httpRequest(auth, {
+    method: "GET",
+    path: `${HISTORY_BASE_PATH}/${queryId}?include_metrics=true`,
+    ...signal ? { signal } : {}
+  });
+}
 
 // src/api/executeStatement.ts
 var TERMINAL_STATES = /* @__PURE__ */ new Set([
@@ -206,12 +234,22 @@ var TERMINAL_STATES = /* @__PURE__ */ new Set([
   "CANCELED",
   "CLOSED"
 ]);
-var POLL_INTERVAL_MS = 500;
-var MAX_POLL_INTERVAL_MS = 5e3;
+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;
+  }
+}
 async function executeStatement(query, auth, options = {}) {
   const warehouseId = options.warehouse_id ?? extractWarehouseId(auth.httpPath);
-  const { signal, onProgress } = 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, enableMetrics ? await fetchMetrics(auth, statementId, signal) : void 0) : void 0;
   const request = Object.fromEntries(
     Object.entries({
       warehouse_id: warehouseId,
@@ -219,28 +257,51 @@ async function executeStatement(query, auth, options = {}) {
       byte_limit: options.byte_limit,
       disposition: options.disposition,
       format: options.format,
-      on_wait_timeout: options.on_wait_timeout,
-      wait_timeout: options.wait_timeout,
+      on_wait_timeout: options.on_wait_timeout ?? "CONTINUE",
+      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);
-  let pollInterval = POLL_INTERVAL_MS;
-  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 emitProgress?.(result.statement_id);
+      await delay(POLL_INTERVAL_MS, signal);
+      result = await getStatement(auth, result.statement_id, signal);
+    }
+  } 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");
     }
-    onProgress?.(result.status);
-    await delay(pollInterval, signal);
-    pollInterval = Math.min(pollInterval * 1.5, MAX_POLL_INTERVAL_MS);
-    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);
   }
-  onProgress?.(result.status);
+  await emitProgress?.(result.statement_id);
   if (result.status.state === "SUCCEEDED")
     return result;
   if (result.status.state === "CANCELED")
@@ -547,69 +608,157 @@ function convertBoolean(value) {
 import { PassThrough } from "stream";
 import { mergeStreamsFromUrls } from "@bitofsky/merge-streams";
 function fetchStream(statementResult, auth, options = {}) {
-  const { signal } = 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",
+      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 = () => {
+      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);
-    });
+    output.once("close", () => signal.removeEventListener("abort", onAbort));
   }
-  mergeChunksToStream(statementResult, auth, manifest, format, output, signal).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) {
-  const result = statementResult.result;
-  let urls = result?.external_links?.map((link) => link.external_link) ?? [];
-  if (urls.length === 0 && manifest.total_chunk_count > 0) {
-    for (let i = 0; i < manifest.total_chunk_count; i++) {
-      if (signal?.aborted) throw new AbortError("Aborted while collecting URLs");
-      const chunkData = await getChunk(auth, statementResult.statement_id, i, signal);
-      const chunkUrls = chunkData.external_links?.map((link) => link.external_link) ?? [];
-      urls.push(...chunkUrls);
+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) {
+    logger?.info?.(`fetchStream no external links found for statement ${statementId}.`, baseLog);
+    return void output.end();
+  }
+  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 chunkUrls = /* @__PURE__ */ new Map();
+  addChunkLinks(chunkUrls, statementResult.result?.external_links);
+  if (!manifest.total_chunk_count)
+    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);
+    addChunkLinks(chunkUrls, chunkData.external_links);
+  }
+  return flattenChunkUrls(chunkUrls);
+}
+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]);
     }
   }
-  if (urls.length === 0)
-    return void output.end();
-  await mergeStreamsFromUrls(format, signal ? { urls, output, signal } : { urls, output });
+}
+function flattenChunkUrls(chunkUrls) {
+  if (chunkUrls.size === 0)
+    return [];
+  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;
 }
 
 // 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);
@@ -628,10 +777,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");
     }
@@ -649,34 +802,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 } = 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;
-  const stream = fetchStream(statementResult, auth, signal ? { signal } : {});
+  }
+  if (!forceMerge) {
+    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 } : {},
+    ...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: [