|
1 | 1 | import fs from "node:fs/promises"; |
2 | 2 | import path from "node:path"; |
3 | 3 | import { WorkspaceClient } from "@databricks/sdk-experimental"; |
| 4 | +import { tableFromIPC } from "apache-arrow"; |
4 | 5 | import pc from "picocolors"; |
5 | 6 | import { createLogger } from "../logging/logger"; |
6 | 7 | import { CACHE_VERSION, hashSQL, loadCache, saveCache } from "./cache"; |
@@ -78,18 +79,85 @@ function formatParametersType(sql: string): string { |
78 | 79 | : "Record<string, never>"; |
79 | 80 | } |
80 | 81 |
|
| 82 | +/** |
| 83 | + * Map Arrow DataType IDs to Databricks SQL type names. |
| 84 | + * Arrow type IDs come from the Arrow spec (apache-arrow TypeId enum). |
| 85 | + * We only need to cover the types that DESCRIBE QUERY can return. |
| 86 | + */ |
| 87 | +function arrowTypeToSqlName(arrowType: { typeId: number }): string { |
| 88 | + switch (arrowType.typeId) { |
| 89 | + case 1: // Bool |
| 90 | + return "BOOLEAN"; |
| 91 | + case 2: // Int (covers TINYINT, SMALLINT, INT, BIGINT depending on bitWidth) |
| 92 | + return "INT"; |
| 93 | + case 3: // Float (covers FLOAT, DOUBLE) |
| 94 | + return "DOUBLE"; |
| 95 | + case 4: // Decimal |
| 96 | + return "DECIMAL"; |
| 97 | + case 5: // Utf8 |
| 98 | + return "STRING"; |
| 99 | + case 6: // Binary |
| 100 | + return "BINARY"; |
| 101 | + case 7: // FixedSizeBinary |
| 102 | + return "BINARY"; |
| 103 | + case 8: // Date |
| 104 | + return "DATE"; |
| 105 | + case 10: // Timestamp |
| 106 | + return "TIMESTAMP"; |
| 107 | + case 12: // List |
| 108 | + return "ARRAY"; |
| 109 | + case 14: // Struct |
| 110 | + return "STRUCT"; |
| 111 | + case 15: // Map |
| 112 | + return "MAP"; |
| 113 | + default: |
| 114 | + return "STRING"; |
| 115 | + } |
| 116 | +} |
| 117 | + |
| 118 | +/** |
| 119 | + * Decode a base64 Arrow IPC attachment and extract column metadata. |
| 120 | + * Returns the same shape as rows parsed from DESCRIBE QUERY data_array. |
| 121 | + */ |
| 122 | +function columnsFromArrowAttachment( |
| 123 | + attachment: string, |
| 124 | +): Array<{ name: string; type_name: string; comment: string | undefined }> { |
| 125 | + const buf = Buffer.from(attachment, "base64"); |
| 126 | + const table = tableFromIPC(buf); |
| 127 | + return table.schema.fields.map((field) => ({ |
| 128 | + name: field.name, |
| 129 | + type_name: arrowTypeToSqlName(field.type), |
| 130 | + comment: undefined, |
| 131 | + })); |
| 132 | +} |
| 133 | + |
81 | 134 | export function convertToQueryType( |
82 | 135 | result: DatabricksStatementExecutionResponse, |
83 | 136 | sql: string, |
84 | 137 | queryName: string, |
85 | 138 | ): { type: string; hasResults: boolean } { |
86 | 139 | const dataRows = result.result?.data_array || []; |
87 | | - const columns = dataRows.map((row) => ({ |
| 140 | + let columns = dataRows.map((row) => ({ |
88 | 141 | name: row[0] || "", |
89 | 142 | type_name: row[1]?.toUpperCase() || "STRING", |
90 | 143 | comment: row[2] || undefined, |
91 | 144 | })); |
92 | 145 |
|
| 146 | + // Fallback: serverless warehouses may return ARROW_STREAM format with an |
| 147 | + // inline base64 attachment instead of data_array. Decode the Arrow IPC |
| 148 | + // schema to extract column names and types. |
| 149 | + if (columns.length === 0 && result.result?.attachment) { |
| 150 | + logger.debug("data_array empty, decoding Arrow IPC attachment for schema"); |
| 151 | + try { |
| 152 | + columns = columnsFromArrowAttachment(result.result.attachment); |
| 153 | + } catch (err) { |
| 154 | + logger.warn( |
| 155 | + "Failed to decode Arrow IPC attachment: %s", |
| 156 | + err instanceof Error ? err.message : String(err), |
| 157 | + ); |
| 158 | + } |
| 159 | + } |
| 160 | + |
93 | 161 | const paramsType = formatParametersType(sql); |
94 | 162 |
|
95 | 163 | // generate result fields with JSDoc |
@@ -277,10 +345,11 @@ export async function generateQueriesFromDescribe( |
277 | 345 | ); |
278 | 346 |
|
279 | 347 | logger.debug( |
280 | | - "DESCRIBE result for %s: state=%s, rows=%d", |
| 348 | + "DESCRIBE result for %s: state=%s, rows=%d, hasAttachment=%s", |
281 | 349 | queryName, |
282 | 350 | result.status.state, |
283 | 351 | result.result?.data_array?.length ?? 0, |
| 352 | + !!result.result?.attachment, |
284 | 353 | ); |
285 | 354 |
|
286 | 355 | if (result.status.state === "FAILED") { |
|
0 commit comments