@bitofsky/databricks-sql 1.0.0

package/dist/index.js

@@ -0,0 +1,718 @@
+// src/errors.ts
+var DatabricksSqlError = class _DatabricksSqlError extends Error {
+  code;
+  statementId;
+  constructor(message, code, statementId) {
+    super(message);
+    this.name = "DatabricksSqlError";
+    this.code = code ?? "UNKNOWN_ERROR";
+    this.statementId = statementId;
+    Error.captureStackTrace?.(this, _DatabricksSqlError);
+  }
+};
+var StatementCancelledError = class extends DatabricksSqlError {
+  constructor(statementId) {
+    super(`Statement ${statementId} was cancelled`, "CANCELLED", statementId);
+    this.name = "StatementCancelledError";
+  }
+};
+var AbortError = class extends DatabricksSqlError {
+  constructor(message = "Operation was aborted") {
+    super(message, "ABORTED");
+    this.name = "AbortError";
+  }
+};
+var HttpError = class extends DatabricksSqlError {
+  status;
+  statusText;
+  constructor(status, statusText, message) {
+    super(message ?? `HTTP ${status}: ${statusText}`, `HTTP_${status}`);
+    this.name = "HttpError";
+    this.status = status;
+    this.statusText = statusText;
+  }
+};
+var AuthenticationError = class extends HttpError {
+  constructor() {
+    super(401, "Unauthorized", "Authentication failed. Check your token.");
+    this.name = "AuthenticationError";
+  }
+};
+var RateLimitError = class extends HttpError {
+  retryAfter;
+  constructor(retryAfter) {
+    super(429, "Too Many Requests", "Rate limit exceeded");
+    this.name = "RateLimitError";
+    this.retryAfter = retryAfter;
+  }
+};
+
+// src/util.ts
+function extractWarehouseId(httpPath) {
+  const match = httpPath.match(/\/sql\/\d+\.\d+\/warehouses\/([a-zA-Z0-9]+)/);
+  if (!match?.[1])
+    throw new Error(`Cannot extract warehouse_id from httpPath: ${httpPath}`);
+  return match[1];
+}
+function throwIfAborted(signal, context) {
+  if (signal?.aborted)
+    throw new AbortError(`[${context}] Aborted`);
+}
+async function delay(ms, signal) {
+  return new Promise((resolve, reject) => {
+    if (signal?.aborted)
+      return reject(new AbortError("Aborted before delay"));
+    let settled = false;
+    const onAbort = () => {
+      if (settled) return;
+      settled = true;
+      clearTimeout(timer);
+      reject(new AbortError("Aborted during delay"));
+    };
+    const timer = setTimeout(() => {
+      if (settled) return;
+      settled = true;
+      signal?.removeEventListener("abort", onAbort);
+      resolve();
+    }, ms);
+    signal?.addEventListener("abort", onAbort, { once: true });
+  });
+}
+function buildUrl(host, path) {
+  const base = host.startsWith("https://") ? host : `https://${host}`;
+  return new URL(path, base).href;
+}
+function validateSucceededResult(statementResult) {
+  if (statementResult.status.state !== "SUCCEEDED")
+    throw new DatabricksSqlError(
+      `Cannot fetch from non-succeeded statement: ${statementResult.status.state}`,
+      "INVALID_STATE",
+      statementResult.statement_id
+    );
+  if (!statementResult.manifest)
+    throw new DatabricksSqlError(
+      "Statement result has no manifest",
+      "MISSING_MANIFEST",
+      statementResult.statement_id
+    );
+  return statementResult.manifest;
+}
+
+// src/http.ts
+var MAX_RETRIES = 3;
+var INITIAL_RETRY_DELAY_MS = 1e3;
+async function httpRequest(auth, options) {
+  const { method, path, body, signal } = options;
+  const url = buildUrl(auth.host, path);
+  let lastError;
+  let retryDelay = INITIAL_RETRY_DELAY_MS;
+  for (let attempt = 0; attempt <= MAX_RETRIES; attempt++) {
+    if (signal?.aborted)
+      throw new AbortError();
+    try {
+      const fetchInit = Object.fromEntries(
+        Object.entries({
+          method,
+          headers: {
+            Authorization: `Bearer ${auth.token}`,
+            "Content-Type": "application/json",
+            Accept: "application/json"
+          },
+          body: body ? JSON.stringify(body) : void 0,
+          signal
+        }).filter(([, v]) => v !== void 0)
+      );
+      const response = await fetch(url, fetchInit);
+      if (response.ok)
+        return await response.json();
+      if (response.status === 401)
+        throw new AuthenticationError();
+      if (response.status === 429) {
+        const retryAfterHeader = response.headers.get("Retry-After");
+        const retryAfter = retryAfterHeader ? parseInt(retryAfterHeader, 10) : void 0;
+        const error = new RateLimitError(
+          isNaN(retryAfter) ? void 0 : retryAfter
+        );
+        if (error.retryAfter && attempt < MAX_RETRIES) {
+          await delay(error.retryAfter * 1e3, signal);
+          continue;
+        }
+        throw error;
+      }
+      if (response.status >= 500) {
+        const errorBody2 = await response.text().catch(() => "");
+        lastError = new HttpError(response.status, response.statusText, errorBody2);
+        if (attempt < MAX_RETRIES) {
+          await delay(retryDelay, signal);
+          retryDelay *= 2;
+          continue;
+        }
+      }
+      const errorBody = await response.text().catch(() => "");
+      throw new HttpError(response.status, response.statusText, errorBody);
+    } catch (err) {
+      if (err instanceof AbortError || err instanceof AuthenticationError || err instanceof HttpError)
+        throw err;
+      if (err instanceof TypeError && err.message.includes("fetch")) {
+        lastError = err;
+        if (attempt < MAX_RETRIES) {
+          await delay(retryDelay, signal);
+          retryDelay *= 2;
+          continue;
+        }
+      }
+      throw err;
+    }
+  }
+  throw lastError ?? new Error("Request failed after retries");
+}
+
+// src/databricks-api.ts
+var BASE_PATH = "/api/2.0/sql/statements";
+async function postStatement(auth, request, signal) {
+  return httpRequest(auth, {
+    method: "POST",
+    path: BASE_PATH,
+    body: request,
+    ...signal ? { signal } : {}
+  });
+}
+async function getStatement(auth, statementId, signal) {
+  return httpRequest(auth, {
+    method: "GET",
+    path: `${BASE_PATH}/${statementId}`,
+    ...signal ? { signal } : {}
+  });
+}
+async function cancelStatement(auth, statementId, signal) {
+  await httpRequest(auth, {
+    method: "POST",
+    path: `${BASE_PATH}/${statementId}/cancel`,
+    ...signal ? { signal } : {}
+  });
+}
+async function getChunk(auth, statementId, chunkIndex, signal) {
+  return httpRequest(auth, {
+    method: "GET",
+    path: `${BASE_PATH}/${statementId}/result/chunks/${chunkIndex}`,
+    ...signal ? { signal } : {}
+  });
+}
+
+// src/api/executeStatement.ts
+var TERMINAL_STATES = /* @__PURE__ */ new Set([
+  "SUCCEEDED",
+  "FAILED",
+  "CANCELED",
+  "CLOSED"
+]);
+var POLL_INTERVAL_MS = 500;
+var MAX_POLL_INTERVAL_MS = 5e3;
+async function executeStatement(query, auth, options = {}) {
+  const warehouseId = options.warehouse_id ?? extractWarehouseId(auth.httpPath);
+  const { signal, onProgress } = options;
+  throwIfAborted(signal, "executeStatement");
+  const request = Object.fromEntries(
+    Object.entries({
+      warehouse_id: warehouseId,
+      statement: query,
+      byte_limit: options.byte_limit,
+      disposition: options.disposition,
+      format: options.format,
+      on_wait_timeout: options.on_wait_timeout,
+      wait_timeout: options.wait_timeout,
+      row_limit: options.row_limit,
+      catalog: options.catalog,
+      schema: options.schema,
+      parameters: options.parameters
+    }).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);
+    result = await getStatement(auth, result.statement_id, signal);
+  }
+  onProgress?.(result.status);
+  if (result.status.state === "SUCCEEDED")
+    return result;
+  if (result.status.state === "CANCELED")
+    throw new StatementCancelledError(result.statement_id);
+  throw new DatabricksSqlError(
+    result.status.error?.message ?? "Statement execution failed",
+    result.status.error?.error_code,
+    result.statement_id
+  );
+}
+
+// src/api/fetchRow.ts
+import { parser } from "stream-json";
+import { streamArray } from "stream-json/streamers/StreamArray";
+
+// src/createRowMapper.ts
+var INTEGER_TYPES = /* @__PURE__ */ new Set(["TINYINT", "SMALLINT", "INT"]);
+var BIGINT_TYPES = /* @__PURE__ */ new Set(["BIGINT", "LONG"]);
+var FLOAT_TYPES = /* @__PURE__ */ new Set(["FLOAT", "DOUBLE"]);
+var BOOLEAN_TYPES = /* @__PURE__ */ new Set(["BOOLEAN"]);
+var STRING_TYPES = /* @__PURE__ */ new Set([
+  "STRING",
+  "DATE",
+  "TIMESTAMP",
+  "TIMESTAMP_NTZ",
+  "TIMESTAMP_LTZ",
+  "TIME"
+]);
+function createRowMapper(manifest, format) {
+  if (format !== "JSON_OBJECT")
+    return (row) => row;
+  const columnConverters = manifest.schema.columns.map((column) => ({
+    name: column.name,
+    convert: createColumnConverter(column)
+  }));
+  return (row) => {
+    const mapped = {};
+    for (let index = 0; index < columnConverters.length; index++) {
+      const converter = columnConverters[index];
+      if (!converter)
+        continue;
+      const { name, convert } = converter;
+      if (name)
+        mapped[name] = convert(row[index]);
+    }
+    return mapped;
+  };
+}
+function createColumnConverter(column) {
+  const descriptor = parseColumnType(column);
+  return (value) => convertValue(descriptor, value);
+}
+function parseColumnType(column) {
+  if (column.type_name === "STRUCT" || column.type_name === "ARRAY" || column.type_name === "MAP")
+    return parseTypeDescriptor(column.type_text);
+  if (column.type_name === "DECIMAL")
+    return createDecimalDescriptor({
+      typeName: column.type_name,
+      typeText: column.type_text
+    }, column.type_precision, column.type_scale);
+  return {
+    typeName: column.type_name,
+    typeText: column.type_text
+  };
+}
+function parseTypeDescriptor(typeText) {
+  const trimmed = typeText.trim();
+  const typeName = getTypeName(trimmed);
+  if (typeName === "STRUCT")
+    return {
+      typeName,
+      typeText: trimmed,
+      fields: parseStructFields(trimmed)
+    };
+  if (typeName === "ARRAY") {
+    const elementTypeText = parseSingleTypeArgument(trimmed);
+    const descriptor = {
+      typeName,
+      typeText: trimmed
+    };
+    if (elementTypeText)
+      descriptor.elementType = parseTypeDescriptor(elementTypeText);
+    return descriptor;
+  }
+  if (typeName === "MAP") {
+    const [keyTypeText, valueTypeText] = parseTypeArguments(trimmed, 2);
+    const descriptor = {
+      typeName,
+      typeText: trimmed
+    };
+    if (keyTypeText)
+      descriptor.keyType = parseTypeDescriptor(keyTypeText);
+    if (valueTypeText)
+      descriptor.valueType = parseTypeDescriptor(valueTypeText);
+    return descriptor;
+  }
+  if (typeName === "DECIMAL") {
+    const { precision, scale } = parseDecimalInfo(trimmed);
+    return createDecimalDescriptor({ typeName, typeText: trimmed }, precision, scale);
+  }
+  return {
+    typeName,
+    typeText: trimmed
+  };
+}
+function getTypeName(typeText) {
+  return typeText.match(/^[A-Z_]+/)?.[0] ?? typeText;
+}
+function parseDecimalInfo(typeText) {
+  const match = typeText.match(/DECIMAL\((\d+),\s*(\d+)\)/);
+  if (!match)
+    return {};
+  return {
+    precision: Number(match[1]),
+    scale: Number(match[2])
+  };
+}
+function createDecimalDescriptor(base, precision, scale) {
+  const descriptor = { ...base };
+  if (precision !== void 0)
+    descriptor.precision = precision;
+  if (scale !== void 0)
+    descriptor.scale = scale;
+  return descriptor;
+}
+function parseStructFields(typeText) {
+  const start = typeText.indexOf("<");
+  const end = typeText.lastIndexOf(">");
+  if (start === -1 || end === -1 || end <= start)
+    return [];
+  const inner = typeText.slice(start + 1, end);
+  const parts = splitTopLevel(inner);
+  const fields = [];
+  for (const part of parts) {
+    const separatorIndex = part.indexOf(":");
+    if (separatorIndex === -1)
+      continue;
+    const name = part.slice(0, separatorIndex).trim();
+    let fieldTypeText = part.slice(separatorIndex + 1).trim();
+    fieldTypeText = stripNotNull(fieldTypeText);
+    if (!name)
+      continue;
+    fields.push({
+      name,
+      type: parseTypeDescriptor(fieldTypeText)
+    });
+  }
+  return fields;
+}
+function parseSingleTypeArgument(typeText) {
+  const [arg] = parseTypeArguments(typeText, 1);
+  return arg ?? null;
+}
+function parseTypeArguments(typeText, expectedCount) {
+  const start = typeText.indexOf("<");
+  const end = typeText.lastIndexOf(">");
+  if (start === -1 || end === -1 || end <= start)
+    return [];
+  const inner = typeText.slice(start + 1, end);
+  const parts = splitTopLevel(inner);
+  if (parts.length < expectedCount)
+    return parts;
+  return parts.slice(0, expectedCount).map((part) => stripNotNull(part.trim()));
+}
+function splitTopLevel(value) {
+  const result = [];
+  let current = "";
+  let angleDepth = 0;
+  let parenDepth = 0;
+  for (const char of value) {
+    if (char === "<") angleDepth++;
+    if (char === ">") angleDepth--;
+    if (char === "(") parenDepth++;
+    if (char === ")") parenDepth--;
+    if (char === "," && angleDepth === 0 && parenDepth === 0) {
+      result.push(current.trim());
+      current = "";
+      continue;
+    }
+    current += char;
+  }
+  if (current.trim().length > 0)
+    result.push(current.trim());
+  return result;
+}
+function stripNotNull(typeText) {
+  let trimmed = typeText.trim();
+  while (trimmed.endsWith("NOT NULL"))
+    trimmed = trimmed.slice(0, -"NOT NULL".length).trim();
+  return trimmed;
+}
+function convertValue(descriptor, value) {
+  if (value === null || value === void 0)
+    return value;
+  if (descriptor.typeName === "STRUCT" && descriptor.fields)
+    return convertStructValue(descriptor.fields, value);
+  if (descriptor.typeName === "ARRAY" && descriptor.elementType)
+    return convertArrayValue(descriptor.elementType, value);
+  if (descriptor.typeName === "MAP" && descriptor.keyType && descriptor.valueType)
+    return convertMapValue(descriptor.keyType, descriptor.valueType, value);
+  if (descriptor.typeName === "DECIMAL")
+    return convertNumber(value);
+  if (INTEGER_TYPES.has(descriptor.typeName))
+    return convertNumber(value);
+  if (BIGINT_TYPES.has(descriptor.typeName))
+    return convertInteger(value);
+  if (FLOAT_TYPES.has(descriptor.typeName))
+    return convertNumber(value);
+  if (BOOLEAN_TYPES.has(descriptor.typeName))
+    return convertBoolean(value);
+  if (STRING_TYPES.has(descriptor.typeName))
+    return value;
+  return value;
+}
+function convertStructValue(fields, value) {
+  const raw = parseStructValue(value);
+  if (!raw || typeof raw !== "object" || Array.isArray(raw))
+    return value;
+  const mapped = {};
+  for (const field of fields)
+    mapped[field.name] = convertValue(field.type, raw[field.name]);
+  return mapped;
+}
+function convertArrayValue(elementType, value) {
+  const raw = parseJsonValue(value);
+  if (!Array.isArray(raw))
+    return value;
+  return raw.map((entry) => convertValue(elementType, entry));
+}
+function convertMapValue(keyType, valueType, value) {
+  const raw = parseJsonValue(value);
+  if (!raw || typeof raw !== "object")
+    return value;
+  if (Array.isArray(raw)) {
+    const mapped2 = {};
+    for (const entry of raw) {
+      if (!Array.isArray(entry) || entry.length < 2)
+        continue;
+      const convertedKey = convertValue(keyType, entry[0]);
+      mapped2[String(convertedKey)] = convertValue(valueType, entry[1]);
+    }
+    return mapped2;
+  }
+  const mapped = {};
+  for (const [key, entryValue] of Object.entries(raw)) {
+    const convertedKey = convertValue(keyType, key);
+    mapped[String(convertedKey)] = convertValue(valueType, entryValue);
+  }
+  return mapped;
+}
+function parseStructValue(value) {
+  const parsed = parseJsonValue(value);
+  if (parsed && typeof parsed === "object" && !Array.isArray(parsed))
+    return parsed;
+  return parsed;
+}
+function parseJsonValue(value) {
+  if (typeof value === "string") {
+    try {
+      return JSON.parse(value);
+    } catch {
+      throw new DatabricksSqlError("Failed to parse JSON value", "INVALID_JSON");
+    }
+  }
+  return value;
+}
+function convertNumber(value) {
+  if (typeof value === "number")
+    return value;
+  if (typeof value === "string") {
+    const parsed = Number(value);
+    return Number.isNaN(parsed) ? value : parsed;
+  }
+  return value;
+}
+function convertInteger(value) {
+  if (typeof value === "bigint")
+    return value;
+  if (typeof value === "number") {
+    if (Number.isInteger(value))
+      return BigInt(value);
+    return value;
+  }
+  if (typeof value === "string") {
+    try {
+      return BigInt(value);
+    } catch {
+      return value;
+    }
+  }
+  return value;
+}
+function convertBoolean(value) {
+  if (typeof value === "boolean")
+    return value;
+  if (typeof value === "string") {
+    if (value === "true") return true;
+    if (value === "false") return false;
+  }
+  return value;
+}
+
+// src/api/fetchStream.ts
+import { PassThrough } from "stream";
+import { mergeStreamsFromUrls } from "@bitofsky/merge-streams";
+function fetchStream(statementResult, auth, options = {}) {
+  const { signal } = options;
+  const manifest = validateSucceededResult(statementResult);
+  const format = manifest.format;
+  const output = new PassThrough();
+  if (signal) {
+    const onAbort = () => {
+      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).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);
+    }
+  }
+  if (urls.length === 0)
+    return void output.end();
+  await mergeStreamsFromUrls(format, signal ? { urls, output, signal } : { urls, output });
+}
+
+// src/api/fetchRow.ts
+async function fetchRow(statementResult, auth, options = {}) {
+  const { signal, onEachRow, format } = options;
+  const manifest = validateSucceededResult(statementResult);
+  const mapRow = createRowMapper(manifest, format);
+  if (statementResult.result?.external_links) {
+    if (manifest.format !== "JSON_ARRAY") {
+      throw new DatabricksSqlError(
+        `fetchRow only supports JSON_ARRAY for external_links. Received: ${manifest.format}`,
+        "UNSUPPORTED_FORMAT",
+        statementResult.statement_id
+      );
+    }
+    const stream = fetchStream(statementResult, auth, signal ? { signal } : {});
+    await consumeJsonArrayStream(stream, mapRow, onEachRow, signal);
+    return;
+  }
+  const totalChunks = manifest.total_chunk_count;
+  const dataArray = statementResult.result?.data_array;
+  if (dataArray) {
+    for (const row of dataArray) {
+      if (signal?.aborted) throw new AbortError("Aborted");
+      onEachRow?.(mapRow(row));
+    }
+  }
+  if (totalChunks > 1) {
+    const statementId = statementResult.statement_id;
+    for (let chunkIndex = 1; chunkIndex < totalChunks; chunkIndex++) {
+      if (signal?.aborted) throw new AbortError("Aborted");
+      const chunk = await getChunk(auth, statementId, chunkIndex, signal);
+      if (chunk.external_links)
+        throw new DatabricksSqlError(
+          "fetchRow only supports INLINE results. Chunk contains external_links.",
+          "UNSUPPORTED_FORMAT",
+          statementId
+        );
+      if (chunk.data_array) {
+        for (const row of chunk.data_array) {
+          if (signal?.aborted) throw new AbortError("Aborted");
+          onEachRow?.(mapRow(row));
+        }
+      }
+    }
+  }
+}
+async function consumeJsonArrayStream(stream, mapRow, onEachRow, signal) {
+  const jsonStream = stream.pipe(parser()).pipe(streamArray());
+  for await (const item of jsonStream) {
+    if (signal?.aborted) {
+      stream.destroy(new AbortError("Aborted"));
+      throw new AbortError("Aborted");
+    }
+    const row = item.value;
+    if (!Array.isArray(row)) {
+      throw new DatabricksSqlError(
+        "Expected JSON_ARRAY rows to be arrays",
+        "INVALID_FORMAT"
+      );
+    }
+    onEachRow?.(mapRow(row));
+  }
+}
+
+// src/api/fetchAll.ts
+async function fetchAll(statementResult, auth, options = {}) {
+  const rows = [];
+  const fetchOptions = {
+    // Collect rows as they are streamed in.
+    onEachRow: (row) => {
+      rows.push(row);
+    }
+  };
+  if (options.signal)
+    fetchOptions.signal = options.signal;
+  if (options.format)
+    fetchOptions.format = options.format;
+  await fetchRow(statementResult, auth, fetchOptions);
+  return rows;
+}
+
+// src/api/mergeExternalLinks.ts
+async function mergeExternalLinks(statementResult, auth, options) {
+  const { signal, mergeStreamToExternalLink } = options;
+  if (!statementResult.result?.external_links)
+    return statementResult;
+  const stream = fetchStream(statementResult, auth, signal ? { signal } : {});
+  const uploadResult = await mergeStreamToExternalLink(stream);
+  const manifest = validateSucceededResult(statementResult);
+  const totalRowCount = manifest.total_row_count ?? 0;
+  return {
+    statement_id: statementResult.statement_id,
+    status: statementResult.status,
+    manifest: {
+      ...manifest,
+      total_chunk_count: 1,
+      total_byte_count: uploadResult.byte_count,
+      chunks: [
+        {
+          chunk_index: 0,
+          row_offset: 0,
+          row_count: totalRowCount,
+          byte_count: uploadResult.byte_count
+        }
+      ]
+    },
+    result: {
+      external_links: [
+        {
+          chunk_index: 0,
+          row_offset: 0,
+          row_count: totalRowCount,
+          byte_count: uploadResult.byte_count,
+          external_link: uploadResult.externalLink,
+          expiration: uploadResult.expiration
+        }
+      ]
+    }
+  };
+}
+export {
+  AbortError,
+  AuthenticationError,
+  DatabricksSqlError,
+  HttpError,
+  RateLimitError,
+  StatementCancelledError,
+  executeStatement,
+  fetchAll,
+  fetchRow,
+  fetchStream,
+  mergeExternalLinks
+};
+//# sourceMappingURL=index.js.map