@bitofsky/databricks-sql 1.0.0 → 1.0.1

package/dist/index.d.cts

@@ -19,6 +19,83 @@ type StatementStatus = {
         message: string;
     };
 };
+/**
+ * Query execution metrics from Query History API
+ * @see https://docs.databricks.com/api/workspace/queryhistory/list
+ */
+type QueryMetrics = {
+    /** Total time in milliseconds */
+    total_time_ms?: number;
+    /** Compilation time in milliseconds */
+    compilation_time_ms?: number;
+    /** Execution time in milliseconds */
+    execution_time_ms?: number;
+    /** Result fetch time in milliseconds */
+    result_fetch_time_ms?: number;
+    /** Query execution time in milliseconds */
+    query_execution_time_ms?: number;
+    /** Metadata time in milliseconds */
+    metadata_time_ms?: number;
+    /** Task total time in milliseconds */
+    task_total_time_ms?: number;
+    /** Photon total time in milliseconds */
+    photon_total_time_ms?: number;
+    /** Query compilation start timestamp */
+    query_compilation_start_timestamp?: number;
+    /** Bytes read */
+    read_bytes?: number;
+    /** Remote bytes read */
+    read_remote_bytes?: number;
+    /** Remote bytes written */
+    write_remote_bytes?: number;
+    /** Cache bytes read */
+    read_cache_bytes?: number;
+    /** Bytes spilled to disk */
+    spill_to_disk_bytes?: number;
+    /** Network bytes sent */
+    network_sent_bytes?: number;
+    /** Pruned bytes */
+    pruned_bytes?: number;
+    /** Rows produced count */
+    rows_produced_count?: number;
+    /** Rows read count */
+    rows_read_count?: number;
+    /** Files read count */
+    read_files_count?: number;
+    /** Partitions read count */
+    read_partitions_count?: number;
+    /** Pruned files count */
+    pruned_files_count?: number;
+    /** Whether result is from cache */
+    result_from_cache?: boolean;
+    /** Percentage of bytes read from cache */
+    bytes_read_from_cache_percentage?: number;
+    /** Remote rows written */
+    write_remote_rows?: number;
+    /** Remote files written */
+    write_remote_files?: number;
+};
+/**
+ * Query info from Query History API
+ * @see https://docs.databricks.com/api/workspace/queryhistory/list
+ */
+type QueryInfo = {
+    query_id: string;
+    status: string;
+    query_text: string;
+    query_start_time_ms: number;
+    execution_end_time_ms?: number;
+    query_end_time_ms?: number;
+    user_id: number;
+    user_name: string;
+    endpoint_id: string;
+    warehouse_id: string;
+    rows_produced?: number;
+    metrics?: QueryMetrics;
+    is_final: boolean;
+    duration?: number;
+    statement_type?: string;
+};
 /** Column schema information */
 type ColumnInfo = {
     name: string;
@@ -91,7 +168,9 @@ type StatementParameter = {
  */
 type ExecuteStatementOptions = {
     /** Progress callback (called on each poll) */
-    onProgress?: (status: StatementStatus) => void;
+    onProgress?: (status: StatementStatus, metrics?: QueryMetrics) => void;
+    /** Enable query metrics fetching during polling (default: false) */
+    enableMetrics?: boolean;
     /** Abort signal for cancellation */
     signal?: AbortSignal;
     /** Result byte limit */
@@ -127,7 +206,10 @@ type RowObject = Record<string, unknown>;
 /** Format for fetchRow/fetchAll */
 type FetchRowFormat = 'JSON_ARRAY' | 'JSON_OBJECT';
 /** Options for fetchStream */
-type FetchStreamOptions = SignalOptions;
+type FetchStreamOptions = SignalOptions & {
+    /** Force merge even when there is only a single external link */
+    forceMerge?: boolean;
+};
 /** Options for fetchRow */
 type FetchRowsOptions = SignalOptions & {
     /** Callback for each row */
@@ -153,6 +235,8 @@ type MergeExternalLinksResult = {
 type MergeExternalLinksOptions = SignalOptions & {
     /** Callback to upload merged stream to external link */
     mergeStreamToExternalLink: (stream: Readable) => Promise<MergeExternalLinksResult>;
+    /** Force merge even when there is only a single external link chunk */
+    forceMerge?: boolean;
 };
 /**
  * API request for executeStatement
@@ -247,4 +331,4 @@ declare function fetchStream(statementResult: StatementResult, auth: AuthInfo, o
  */
 declare function mergeExternalLinks(statementResult: StatementResult, auth: AuthInfo, options: MergeExternalLinksOptions): Promise<StatementResult>;
 
-export { AbortError, type AuthInfo, AuthenticationError, type ChunkInfo, type ColumnInfo, DatabricksSqlError, type ExecuteStatementOptions, type ExecuteStatementRequest, type ExternalLinkInfo, type ExternalLinksResultData, type FetchAllOptions, type FetchRowFormat, type FetchRowsOptions, type FetchStreamOptions, type GetChunkResponse, HttpError, type InlineResultData, type MergeExternalLinksOptions, type MergeExternalLinksResult, RateLimitError, type ResultData, type RowArray, type RowObject, type SignalOptions, StatementCancelledError, type StatementManifest, type StatementParameter, type StatementResult, type StatementState, type StatementStatus, executeStatement, fetchAll, fetchRow, fetchStream, mergeExternalLinks };
+export { AbortError, type AuthInfo, AuthenticationError, type ChunkInfo, type ColumnInfo, DatabricksSqlError, type ExecuteStatementOptions, type ExecuteStatementRequest, type ExternalLinkInfo, type ExternalLinksResultData, type FetchAllOptions, type FetchRowFormat, type FetchRowsOptions, type FetchStreamOptions, type GetChunkResponse, HttpError, type InlineResultData, type MergeExternalLinksOptions, type MergeExternalLinksResult, type QueryInfo, type QueryMetrics, RateLimitError, type ResultData, type RowArray, type RowObject, type SignalOptions, StatementCancelledError, type StatementManifest, type StatementParameter, type StatementResult, type StatementState, type StatementStatus, executeStatement, fetchAll, fetchRow, fetchStream, mergeExternalLinks };

package/dist/index.d.ts

@@ -19,6 +19,83 @@ type StatementStatus = {
         message: string;
     };
 };
+/**
+ * Query execution metrics from Query History API
+ * @see https://docs.databricks.com/api/workspace/queryhistory/list
+ */
+type QueryMetrics = {
+    /** Total time in milliseconds */
+    total_time_ms?: number;
+    /** Compilation time in milliseconds */
+    compilation_time_ms?: number;
+    /** Execution time in milliseconds */
+    execution_time_ms?: number;
+    /** Result fetch time in milliseconds */
+    result_fetch_time_ms?: number;
+    /** Query execution time in milliseconds */
+    query_execution_time_ms?: number;
+    /** Metadata time in milliseconds */
+    metadata_time_ms?: number;
+    /** Task total time in milliseconds */
+    task_total_time_ms?: number;
+    /** Photon total time in milliseconds */
+    photon_total_time_ms?: number;
+    /** Query compilation start timestamp */
+    query_compilation_start_timestamp?: number;
+    /** Bytes read */
+    read_bytes?: number;
+    /** Remote bytes read */
+    read_remote_bytes?: number;
+    /** Remote bytes written */
+    write_remote_bytes?: number;
+    /** Cache bytes read */
+    read_cache_bytes?: number;
+    /** Bytes spilled to disk */
+    spill_to_disk_bytes?: number;
+    /** Network bytes sent */
+    network_sent_bytes?: number;
+    /** Pruned bytes */
+    pruned_bytes?: number;
+    /** Rows produced count */
+    rows_produced_count?: number;
+    /** Rows read count */
+    rows_read_count?: number;
+    /** Files read count */
+    read_files_count?: number;
+    /** Partitions read count */
+    read_partitions_count?: number;
+    /** Pruned files count */
+    pruned_files_count?: number;
+    /** Whether result is from cache */
+    result_from_cache?: boolean;
+    /** Percentage of bytes read from cache */
+    bytes_read_from_cache_percentage?: number;
+    /** Remote rows written */
+    write_remote_rows?: number;
+    /** Remote files written */
+    write_remote_files?: number;
+};
+/**
+ * Query info from Query History API
+ * @see https://docs.databricks.com/api/workspace/queryhistory/list
+ */
+type QueryInfo = {
+    query_id: string;
+    status: string;
+    query_text: string;
+    query_start_time_ms: number;
+    execution_end_time_ms?: number;
+    query_end_time_ms?: number;
+    user_id: number;
+    user_name: string;
+    endpoint_id: string;
+    warehouse_id: string;
+    rows_produced?: number;
+    metrics?: QueryMetrics;
+    is_final: boolean;
+    duration?: number;
+    statement_type?: string;
+};
 /** Column schema information */
 type ColumnInfo = {
     name: string;
@@ -91,7 +168,9 @@ type StatementParameter = {
  */
 type ExecuteStatementOptions = {
     /** Progress callback (called on each poll) */
-    onProgress?: (status: StatementStatus) => void;
+    onProgress?: (status: StatementStatus, metrics?: QueryMetrics) => void;
+    /** Enable query metrics fetching during polling (default: false) */
+    enableMetrics?: boolean;
     /** Abort signal for cancellation */
     signal?: AbortSignal;
     /** Result byte limit */
@@ -127,7 +206,10 @@ type RowObject = Record<string, unknown>;
 /** Format for fetchRow/fetchAll */
 type FetchRowFormat = 'JSON_ARRAY' | 'JSON_OBJECT';
 /** Options for fetchStream */
-type FetchStreamOptions = SignalOptions;
+type FetchStreamOptions = SignalOptions & {
+    /** Force merge even when there is only a single external link */
+    forceMerge?: boolean;
+};
 /** Options for fetchRow */
 type FetchRowsOptions = SignalOptions & {
     /** Callback for each row */
@@ -153,6 +235,8 @@ type MergeExternalLinksResult = {
 type MergeExternalLinksOptions = SignalOptions & {
     /** Callback to upload merged stream to external link */
     mergeStreamToExternalLink: (stream: Readable) => Promise<MergeExternalLinksResult>;
+    /** Force merge even when there is only a single external link chunk */
+    forceMerge?: boolean;
 };
 /**
  * API request for executeStatement
@@ -247,4 +331,4 @@ declare function fetchStream(statementResult: StatementResult, auth: AuthInfo, o
  */
 declare function mergeExternalLinks(statementResult: StatementResult, auth: AuthInfo, options: MergeExternalLinksOptions): Promise<StatementResult>;
 
-export { AbortError, type AuthInfo, AuthenticationError, type ChunkInfo, type ColumnInfo, DatabricksSqlError, type ExecuteStatementOptions, type ExecuteStatementRequest, type ExternalLinkInfo, type ExternalLinksResultData, type FetchAllOptions, type FetchRowFormat, type FetchRowsOptions, type FetchStreamOptions, type GetChunkResponse, HttpError, type InlineResultData, type MergeExternalLinksOptions, type MergeExternalLinksResult, RateLimitError, type ResultData, type RowArray, type RowObject, type SignalOptions, StatementCancelledError, type StatementManifest, type StatementParameter, type StatementResult, type StatementState, type StatementStatus, executeStatement, fetchAll, fetchRow, fetchStream, mergeExternalLinks };
+export { AbortError, type AuthInfo, AuthenticationError, type ChunkInfo, type ColumnInfo, DatabricksSqlError, type ExecuteStatementOptions, type ExecuteStatementRequest, type ExternalLinkInfo, type ExternalLinksResultData, type FetchAllOptions, type FetchRowFormat, type FetchRowsOptions, type FetchStreamOptions, type GetChunkResponse, HttpError, type InlineResultData, type MergeExternalLinksOptions, type MergeExternalLinksResult, type QueryInfo, type QueryMetrics, RateLimitError, type ResultData, type RowArray, type RowObject, type SignalOptions, StatementCancelledError, type StatementManifest, type StatementParameter, type StatementResult, type StatementState, type StatementStatus, executeStatement, fetchAll, fetchRow, fetchStream, mergeExternalLinks };

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,20 @@ 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 } = options;
   throwIfAborted(signal, "executeStatement");
+  const emitProgress = onProgress ? async (statementId) => onProgress(result.status, enableMetrics ? await fetchMetrics(auth, statementId, signal) : void 0) : void 0;
   const request = Object.fromEntries(
     Object.entries({
       warehouse_id: warehouseId,
@@ -219,8 +255,8 @@ 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: options.wait_timeout ?? "50s",
       row_limit: options.row_limit,
       catalog: options.catalog,
       schema: options.schema,
@@ -228,19 +264,17 @@ async function executeStatement(query, auth, options = {}) {
     }).filter(([, v]) => v !== void 0)
   );
   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(() => {
       });
       throw new AbortError("Aborted during polling");
     }
-    onProgress?.(result.status);
-    await delay(pollInterval, signal);
-    pollInterval = Math.min(pollInterval * 1.5, MAX_POLL_INTERVAL_MS);
+    await emitProgress?.(result.statement_id);
+    await delay(POLL_INTERVAL_MS, signal);
     result = await getStatement(auth, result.statement_id, signal);
   }
-  onProgress?.(result.status);
+  await emitProgress?.(result.statement_id);
   if (result.status.state === "SUCCEEDED")
     return result;
   if (result.status.state === "CANCELED")
@@ -547,40 +581,55 @@ function convertBoolean(value) {
 import { PassThrough } from "stream";
 import { mergeStreamsFromUrls } from "@bitofsky/merge-streams";
 function fetchStream(statementResult, auth, options = {}) {
-  const { signal } = options;
+  const { signal, forceMerge } = options;
   const manifest = validateSucceededResult(statementResult);
   const format = manifest.format;
+  if (statementResult.result?.data_array) {
+    throw new DatabricksSqlError(
+      "fetchStream only supports EXTERNAL_LINKS results",
+      "UNSUPPORTED_FORMAT",
+      statementResult.statement_id
+    );
+  }
   const output = new PassThrough();
   if (signal) {
-    const onAbort = () => {
-      output.destroy(new AbortError("Stream aborted"));
-    };
+    const onAbort = () => 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);
-    }
-  );
+  mergeChunksToStream(statementResult, auth, manifest, format, output, signal, forceMerge).catch((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) {
+  const urls = await collectExternalUrls(statementResult, auth, manifest, signal);
   if (urls.length === 0)
     return void output.end();
-  await mergeStreamsFromUrls(format, signal ? { urls, output, signal } : { urls, output });
+  if (urls.length === 1 && !forceMerge)
+    return pipeUrlToOutput(urls[0], output, signal);
+  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;
+  if (!manifest.total_chunk_count)
+    return [];
+  const chunkUrls = [];
+  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);
+    chunkUrls.push(...extractExternalLinks(chunkData.external_links));
+  }
+  return chunkUrls;
+}
+function extractExternalLinks(externalLinks) {
+  if (!externalLinks)
+    return [];
+  return externalLinks.map((link) => link.external_link).filter(isNonEmptyString);
+}
+function isNonEmptyString(value) {
+  return typeof value === "string" && value.length > 0;
 }
 
 // src/api/fetchRow.ts
@@ -665,10 +714,20 @@ async function fetchAll(statementResult, auth, options = {}) {
 
 // src/api/mergeExternalLinks.ts
 async function mergeExternalLinks(statementResult, auth, options) {
-  const { signal, mergeStreamToExternalLink } = options;
+  const { signal, mergeStreamToExternalLink, forceMerge } = options;
   if (!statementResult.result?.external_links)
     return statementResult;
-  const stream = fetchStream(statementResult, auth, signal ? { signal } : {});
+  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)
+      return statementResult;
+  }
+  const stream = fetchStream(statementResult, auth, {
+    ...signal ? { signal } : {},
+    ...forceMerge !== void 0 ? { forceMerge } : {}
+  });
   const uploadResult = await mergeStreamToExternalLink(stream);
   const manifest = validateSucceededResult(statementResult);
   const totalRowCount = manifest.total_row_count ?? 0;