diff --git a/services/cubejs/src/routes/loadExport.js b/services/cubejs/src/routes/loadExport.js index 4ea7548d..b3d486d2 100644 --- a/services/cubejs/src/routes/loadExport.js +++ b/services/cubejs/src/routes/loadExport.js @@ -23,6 +23,7 @@ import { writeBinaryChunk, writeRowStreamAsArrow, } from "../utils/arrowSerializer.js"; +import { execClickHouseArrowStream } from "../utils/clickhouseArrow.js"; const prepareAnnotation = typeof prepareAnnotationModule.prepareAnnotation === "function" @@ -54,20 +55,6 @@ function isClickHouseContext(securityContext) { return typeof dbType === "string" && dbType.toLowerCase() === "clickhouse"; } -function removeTrailingSemicolon(query) { - const trimmed = String(query ?? "").trimEnd(); - let lastNonSemiIdx = trimmed.length; - for (let i = lastNonSemiIdx; i > 0; i--) { - if (trimmed[i - 1] !== ";") { - lastNonSemiIdx = i; - break; - } - } - return lastNonSemiIdx !== trimmed.length - ? trimmed.slice(0, lastNonSemiIdx) - : trimmed; -} - function normalizeClickHouseCSVLine(line) { const normalized = line.replace(CLICKHOUSE_NULL_TOKEN_RE, ""); if (normalized.endsWith("\r\n")) return normalized; @@ -426,13 +413,10 @@ async function executeNativeClickHouseCsv( } async function executeNativeClickHouseArrow(res, query, values, driver, signal) { - const result = await driver.client.exec({ - query: `${removeTrailingSemicolon(sqlstring.format(query, values || []))}\nFORMAT ArrowStream`, - clickhouse_settings: { - ...driver.config?.clickhouseSettings, - output_format_arrow_compression_method: "none", - }, - abort_signal: signal, + const result = await execClickHouseArrowStream({ + driver, + sql: sqlstring.format(query, values || []), + signal }); const stream = typeof result.stream === "function" @@ -496,17 +480,26 @@ async function tryHandleLoadExport(req, res, cubejs, query, format) { const driver = await cubejs.options.driverFactory({ securityContext: plan.context.securityContext, }); - res.set(ARROW_HEADERS); - setNativeArrowFieldMappingHeaders(res, plan); - await executeNativeClickHouseArrow( - res, - nativeQuery.query, - nativeQuery.values, - driver, - abortController.signal - ); - res.end(); - return true; + try { + res.set(ARROW_HEADERS); + setNativeArrowFieldMappingHeaders(res, plan); + await executeNativeClickHouseArrow( + res, + nativeQuery.query, + nativeQuery.values, + driver, + abortController.signal + ); + res.end(); + return true; + } catch (err) { + if (abortController.signal.aborted) return true; + if (res.writableEnded) throw err; + console.warn( + "Native ClickHouse Arrow failed; falling back to semantic stream:", + err?.message || err + ); + } } if (!canSemanticStreamLoadExport(format, plan.capabilities)) { diff --git a/services/cubejs/src/routes/runSql.js b/services/cubejs/src/routes/runSql.js index d7ee2dbc..6a406ad2 100644 --- a/services/cubejs/src/routes/runSql.js +++ b/services/cubejs/src/routes/runSql.js @@ -2,6 +2,7 @@ import crypto from "crypto"; import { loadRules } from "../utils/queryRewrite.js"; import { serializeRowsToArrow } from "../utils/arrowSerializer.js"; +import { execClickHouseArrowStream } from "../utils/clickhouseArrow.js"; import { validateFormat } from "../utils/formatValidator.js"; import { writeRowsAsCSV, writeTextChunk } from "../utils/csvSerializer.js"; import { buildJSONStat } from "../utils/jsonstatBuilder.js"; @@ -62,20 +63,6 @@ function addUniqueColumns(target, names) { } } -function removeTrailingSemicolon(query) { - const trimmed = String(query ?? "").trimEnd(); - let lastNonSemiIdx = trimmed.length; - for (let i = lastNonSemiIdx; i > 0; i--) { - if (trimmed[i - 1] !== ";") { - lastNonSemiIdx = i; - break; - } - } - return lastNonSemiIdx !== trimmed.length - ? trimmed.slice(0, lastNonSemiIdx) - : trimmed; -} - function deriveExportColumnsFromRunSql(body, rows) { if (rows.length > 0) { return Object.keys(rows[0]); @@ -213,20 +200,19 @@ export default async (req, res, cubejs) => { } if (format === "arrow" && isClickHouse(securityContext)) { - const clickhouseQuery = `${removeTrailingSemicolon(sql)}\nFORMAT ArrowStream`; - const result = await driver.client.exec({ - query: clickhouseQuery, - clickhouse_settings: { - ...driver.config?.clickhouseSettings, - output_format_arrow_compression_method: "none", - }, - abort_signal: abortController.signal, + const result = await execClickHouseArrowStream({ + driver, + sql, + signal: abortController.signal }); res.set(ARROW_HEADERS); try { - for await (const chunk of result.stream) { + const stream = typeof result.stream === "function" + ? result.stream() + : result.stream; + for await (const chunk of stream) { await writeBinaryChunk(res, chunk, abortController.signal); } } catch (streamErr) { diff --git a/services/cubejs/src/utils/__tests__/clickhouseArrow.test.js b/services/cubejs/src/utils/__tests__/clickhouseArrow.test.js new file mode 100644 index 00000000..6cf5b117 --- /dev/null +++ b/services/cubejs/src/utils/__tests__/clickhouseArrow.test.js @@ -0,0 +1,85 @@ +import { describe, it } from "node:test"; +import assert from "node:assert/strict"; + +import { + execClickHouseArrowStream, + isForbiddenClickHouseArrowCompressionError +} from "../clickhouseArrow.js"; + +describe("clickhouseArrow", () => { + it("detects the readonly Arrow compression SET error", () => { + assert.equal( + isForbiddenClickHouseArrowCompressionError( + new Error( + "Cannot modify 'output_format_arrow_compression_method' setting in readonly mode." + ) + ), + true + ); + assert.equal( + isForbiddenClickHouseArrowCompressionError(new Error("timeout")), + false + ); + }); + + it("retries without the compression override after a readonly SET error", async () => { + const settingsByCall = []; + const driver = { + config: { clickhouseSettings: { max_execution_time: 30 } }, + client: { + exec: async (opts) => { + settingsByCall.push(opts.clickhouse_settings); + if (settingsByCall.length === 1) { + throw new Error( + "Cannot modify 'output_format_arrow_compression_method' setting in readonly mode." + ); + } + return { ok: true, query: opts.query }; + } + } + }; + + const result = await execClickHouseArrowStream({ + driver, + sql: "SELECT 1;", + signal: undefined + }); + + assert.equal(result.ok, true); + assert.match(result.query, /FORMAT ArrowStream/); + assert.equal(settingsByCall.length, 2); + assert.equal(settingsByCall[0].max_execution_time, 30); + assert.equal(settingsByCall[0].output_format_arrow_compression_method, "none"); + assert.equal(settingsByCall[1].max_execution_time, 30); + assert.equal( + settingsByCall[1].output_format_arrow_compression_method, + undefined + ); + }); + + it("skips the compression SET when the driver is readonly", async () => { + const settingsByCall = []; + const driver = { + config: { clickhouseSettings: {} }, + readOnly: () => true, + client: { + exec: async (opts) => { + settingsByCall.push(opts.clickhouse_settings); + return { ok: true, query: opts.query }; + } + } + }; + + await execClickHouseArrowStream({ + driver, + sql: "SELECT 1", + signal: undefined + }); + + assert.equal(settingsByCall.length, 1); + assert.equal( + settingsByCall[0].output_format_arrow_compression_method, + undefined + ); + }); +}); diff --git a/services/cubejs/src/utils/clickhouseArrow.js b/services/cubejs/src/utils/clickhouseArrow.js new file mode 100644 index 00000000..9cec8567 --- /dev/null +++ b/services/cubejs/src/utils/clickhouseArrow.js @@ -0,0 +1,66 @@ +function removeTrailingSemicolon(query) { + const trimmed = String(query ?? "").trimEnd(); + let lastNonSemiIdx = trimmed.length; + for (let i = lastNonSemiIdx; i > 0; i--) { + if (trimmed[i - 1] !== ";") { + lastNonSemiIdx = i; + break; + } + } + return lastNonSemiIdx !== trimmed.length + ? trimmed.slice(0, lastNonSemiIdx) + : trimmed; +} + +function getClickHouseClient(driver) { + if (driver?.client && typeof driver.client.exec === "function") { + return driver.client; + } + if (driver && typeof driver.exec === "function") { + return driver; + } + return null; +} + +export function isForbiddenClickHouseArrowCompressionError(err) { + const msg = String(err?.message || err); + return ( + msg.includes("output_format_arrow_compression_method") + && (msg.includes("readonly") || msg.includes("Cannot modify")) + ); +} + +/** + * Native ClickHouse ArrowStream. Prefer uncompressed IPC so the browser + * apache-arrow decoder can read it. Readonly ClickHouse users cannot SET + * that codec, even to the current value, so skip or retry without it. + */ +export async function execClickHouseArrowStream({ driver, sql, signal }) { + const client = getClickHouseClient(driver); + if (!client) { + throw new Error("ClickHouse driver has no exec client for ArrowStream"); + } + + const query = `${removeTrailingSemicolon(sql)}\nFORMAT ArrowStream`; + const baseSettings = { ...(driver.config?.clickhouseSettings || {}) }; + const readonly = typeof driver.readOnly === "function" && driver.readOnly(); + + const run = (clickhouse_settings) => client.exec({ + query, + clickhouse_settings, + abort_signal: signal + }); + + if (!readonly) { + try { + return await run({ + ...baseSettings, + output_format_arrow_compression_method: "none" + }); + } catch (err) { + if (!isForbiddenClickHouseArrowCompressionError(err)) throw err; + } + } + + return await run(baseSettings); +} diff --git a/services/cubejs/test/loadExport.test.js b/services/cubejs/test/loadExport.test.js index 6450814b..d7e133dc 100644 --- a/services/cubejs/test/loadExport.test.js +++ b/services/cubejs/test/loadExport.test.js @@ -226,6 +226,112 @@ describe("maybeHandleLoadExport", () => { assert.deepEqual(res.binaryOutput, nativeArrowBuffer); }); + it("retries native Arrow without compression override when ClickHouse is readonly", async () => { + const req = createRequest({ + format: "arrow", + query: { + dimensions: ["Orders.city"] + } + }); + const res = new MockResponse(); + const settingsByCall = []; + const nativeArrowBuffer = Buffer.from("arrow-readonly-ok"); + + const cubejs = createMockCube({ + dbType: "clickhouse", + normalizedQuery: req.body.query, + sqlQuery: { + sql: ["SELECT city FROM orders", []], + aliasNameToMember: { + "Orders.city": "Orders.city" + } + }, + metaConfig: createMetaConfig({ + dimensions: [{ name: "Orders.city", type: "string" }] + }), + nativePreAggs: { + preAggregationsTablesToTempTables: [], + values: [] + }, + driver: createClickHouseDriver({ + execImpl: async (opts) => { + settingsByCall.push(opts.clickhouse_settings); + if (settingsByCall.length === 1) { + throw new Error( + "Cannot modify 'output_format_arrow_compression_method' setting in readonly mode." + ); + } + return { + async *stream() { + yield nativeArrowBuffer; + } + }; + } + }) + }); + + await maybeHandleLoadExport(req, res, () => { + throw new Error("next should not be called"); + }, cubejs); + + assert.equal(settingsByCall.length, 2); + assert.equal( + settingsByCall[0].output_format_arrow_compression_method, + "none" + ); + assert.equal( + settingsByCall[1].output_format_arrow_compression_method, + undefined + ); + assert.equal( + res.headers["Content-Type"], + "application/vnd.apache.arrow.stream" + ); + assert.deepEqual(res.binaryOutput, nativeArrowBuffer); + }); + + it("falls back to semantic Arrow when the native ClickHouse client has no exec", async () => { + const req = createRequest({ + format: "arrow", + query: { + dimensions: ["Orders.city"] + } + }); + const res = new MockResponse(); + + const cubejs = createMockCube({ + dbType: "clickhouse", + normalizedQuery: req.body.query, + sqlQuery: { + sql: ["SELECT city FROM orders", []], + aliasNameToMember: { + "Orders.city": "Orders.city" + } + }, + metaConfig: createMetaConfig({ + dimensions: [{ name: "Orders.city", type: "string" }] + }), + nativePreAggs: { + preAggregationsTablesToTempTables: [], + values: [] + }, + streamRows: [{ "Orders.city": "Reykjavik" }], + driver: { + config: { clickhouseSettings: {} } + } + }); + + await maybeHandleLoadExport(req, res, () => { + throw new Error("next should not be called"); + }, cubejs); + + assert.equal( + res.headers["Content-Type"], + "application/vnd.apache.arrow.stream" + ); + assert.ok(res.binaryOutput.length > 0); + }); + it("streams Arrow through the semantic export path when native ClickHouse export is unavailable", async () => { const req = createRequest({ format: "arrow", @@ -523,10 +629,10 @@ function createMockCube(options) { }; } -function createClickHouseDriver({ csvLines = [], arrowChunks = [] }) { +function createClickHouseDriver({ csvLines = [], arrowChunks = [], execImpl } = {}) { return { config: { - clickhouseSettings: {}, + clickhouseSettings: {} }, client: { query: async () => ({ @@ -534,16 +640,19 @@ function createClickHouseDriver({ csvLines = [], arrowChunks = [] }) { for (const line of csvLines) { yield [{ text: line }]; } - }, + } }), - exec: async () => ({ - async *stream() { - for (const chunk of arrowChunks) { - yield chunk; + exec: async (opts) => { + if (execImpl) return execImpl(opts); + return { + async *stream() { + for (const chunk of arrowChunks) { + yield chunk; + } } - }, - }), - }, + }; + } + } }; }