Skip to content

Commit 113f45b

Browse files
authored
fix: Send column schema and empty data for empty /cubesql result sets instead of empty response (#9798)
* fix: Send column schema and empty data for empty handle_sql_query result sets * Linter * Add E2E test
1 parent 3917c50 commit 113f45b

File tree

3 files changed

+213
-46
lines changed

3 files changed

+213
-46
lines changed

packages/cubejs-backend-native/src/node_export.rs

Lines changed: 73 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
use cubesql::compile::parser::parse_sql_to_statement;
22
use cubesql::compile::{convert_statement_to_cube_query, get_df_batches};
33
use cubesql::config::processing_loop::ShutdownMode;
4+
use cubesql::sql::dataframe::{arrow_to_column_type, Column};
5+
use cubesql::sql::ColumnFlags;
46
use cubesql::transport::{SpanId, TransportService};
57
use futures::StreamExt;
68

@@ -192,6 +194,32 @@ fn shutdown_interface(mut cx: FunctionContext) -> JsResult<JsPromise> {
192194

193195
const CHUNK_DELIM: &str = "\n";
194196

197+
async fn write_jsonl_message(
198+
channel: Arc<Channel>,
199+
write_fn: Arc<Root<JsFunction>>,
200+
stream: Arc<Root<JsObject>>,
201+
value: serde_json::Value,
202+
) -> Result<bool, CubeError> {
203+
let message = format!("{}{}", serde_json::to_string(&value)?, CHUNK_DELIM);
204+
205+
call_js_fn(
206+
channel,
207+
write_fn,
208+
Box::new(move |cx| {
209+
let arg = cx.string(message).upcast::<JsValue>();
210+
Ok(vec![arg.upcast::<JsValue>()])
211+
}),
212+
Box::new(|cx, v| match v.downcast_or_throw::<JsBoolean, _>(cx) {
213+
Ok(v) => Ok(v.value(cx)),
214+
Err(_) => Err(CubeError::internal(
215+
"Failed to downcast write response".to_string(),
216+
)),
217+
}),
218+
stream,
219+
)
220+
.await
221+
}
222+
195223
async fn handle_sql_query(
196224
services: Arc<NodeCubeServices>,
197225
native_auth_ctx: Arc<NativeSQLAuthContext>,
@@ -262,59 +290,44 @@ async fn handle_sql_query(
262290

263291
drain_handler.handle(stream_methods.on.clone()).await?;
264292

265-
let mut is_first_batch = true;
266-
while let Some(batch) = stream.next().await {
267-
let (columns, data) = batch_to_rows(batch?)?;
293+
// Get schema from stream and convert to DataFrame columns format
294+
let stream_schema = stream.schema();
295+
let mut columns = Vec::with_capacity(stream_schema.fields().len());
296+
for field in stream_schema.fields().iter() {
297+
columns.push(Column::new(
298+
field.name().clone(),
299+
arrow_to_column_type(field.data_type().clone())?,
300+
ColumnFlags::empty(),
301+
));
302+
}
268303

269-
if is_first_batch {
270-
let mut schema = Map::new();
271-
schema.insert("schema".into(), columns);
272-
let columns = format!(
273-
"{}{}",
274-
serde_json::to_string(&serde_json::Value::Object(schema))?,
275-
CHUNK_DELIM
276-
);
277-
is_first_batch = false;
304+
// Send schema first
305+
let columns_json = serde_json::to_value(&columns)?;
306+
let mut schema_response = Map::new();
307+
schema_response.insert("schema".into(), columns_json);
278308

279-
call_js_fn(
280-
channel.clone(),
281-
stream_methods.write.clone(),
282-
Box::new(|cx| {
283-
let arg = cx.string(columns).upcast::<JsValue>();
309+
write_jsonl_message(
310+
channel.clone(),
311+
stream_methods.write.clone(),
312+
stream_methods.stream.clone(),
313+
serde_json::Value::Object(schema_response),
314+
)
315+
.await?;
284316

285-
Ok(vec![arg.upcast::<JsValue>()])
286-
}),
287-
Box::new(|cx, v| match v.downcast_or_throw::<JsBoolean, _>(cx) {
288-
Ok(v) => Ok(v.value(cx)),
289-
Err(_) => Err(CubeError::internal(
290-
"Failed to downcast write response".to_string(),
291-
)),
292-
}),
293-
stream_methods.stream.clone(),
294-
)
295-
.await?;
296-
}
317+
// Process all batches
318+
let mut has_data = false;
319+
while let Some(batch) = stream.next().await {
320+
let (_, data) = batch_to_rows(batch?)?;
321+
has_data = true;
297322

298323
let mut rows = Map::new();
299324
rows.insert("data".into(), serde_json::Value::Array(data));
300-
let data = format!("{}{}", serde_json::to_string(&rows)?, CHUNK_DELIM);
301-
let js_stream_write_fn = stream_methods.write.clone();
302325

303-
let should_pause = !call_js_fn(
326+
let should_pause = !write_jsonl_message(
304327
channel.clone(),
305-
js_stream_write_fn,
306-
Box::new(|cx| {
307-
let arg = cx.string(data).upcast::<JsValue>();
308-
309-
Ok(vec![arg.upcast::<JsValue>()])
310-
}),
311-
Box::new(|cx, v| match v.downcast_or_throw::<JsBoolean, _>(cx) {
312-
Ok(v) => Ok(v.value(cx)),
313-
Err(_) => Err(CubeError::internal(
314-
"Failed to downcast write response".to_string(),
315-
)),
316-
}),
328+
stream_methods.write.clone(),
317329
stream_methods.stream.clone(),
330+
serde_json::Value::Object(rows),
318331
)
319332
.await?;
320333

@@ -324,6 +337,20 @@ async fn handle_sql_query(
324337
}
325338
}
326339

340+
// If no data was processed, send empty data
341+
if !has_data {
342+
let mut rows = Map::new();
343+
rows.insert("data".into(), serde_json::Value::Array(vec![]));
344+
345+
write_jsonl_message(
346+
channel.clone(),
347+
stream_methods.write.clone(),
348+
stream_methods.stream.clone(),
349+
serde_json::Value::Object(rows),
350+
)
351+
.await?;
352+
}
353+
327354
Ok::<(), CubeError>(())
328355
};
329356

@@ -465,13 +492,13 @@ fn exec_sql(mut cx: FunctionContext) -> JsResult<JsValue> {
465492
Err(err) => {
466493
let mut error_response = Map::new();
467494
error_response.insert("error".into(), err.to_string().into());
468-
let error_response = format!(
495+
let error_message = format!(
469496
"{}{}",
470497
serde_json::to_string(&serde_json::Value::Object(error_response))
471498
.expect("Failed to serialize error response to JSON"),
472499
CHUNK_DELIM
473500
);
474-
let arg = cx.string(error_response).upcast::<JsValue>();
501+
let arg = cx.string(error_message).upcast::<JsValue>();
475502

476503
vec![arg]
477504
}

packages/cubejs-backend-native/test/sql.test.ts

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -372,4 +372,74 @@ describe('SQLInterface', () => {
372372
expect(process.env.CUBESQL_STREAM_MODE).toBeFalsy();
373373
}
374374
});
375+
376+
test('schema from stream and empty data when no batches', async () => {
377+
const interfaceMethods_ = interfaceMethods();
378+
const instance = await native.registerInterface({
379+
...interfaceMethods_,
380+
canSwitchUserForSession: (_payload) => true,
381+
});
382+
383+
let schemaReceived = false;
384+
let dataReceived = false;
385+
let emptyDataReceived = false;
386+
let buf = '';
387+
388+
const write = jest.fn((chunk, _, callback) => {
389+
const lines = (buf + chunk.toString('utf-8')).split('\n');
390+
buf = lines.pop() || '';
391+
392+
lines
393+
.filter((it) => it.trim().length)
394+
.forEach((line) => {
395+
const json = JSON.parse(line);
396+
397+
if (json.error) {
398+
// Ignore errors for this test
399+
return;
400+
}
401+
402+
if (json.schema) {
403+
schemaReceived = true;
404+
expect(json.schema).toBeDefined();
405+
expect(Array.isArray(json.schema)).toBe(true);
406+
expect(json.data).toBeUndefined();
407+
} else if (json.data) {
408+
dataReceived = true;
409+
// Check if it's empty data
410+
if (Array.isArray(json.data) && json.data.length === 0) {
411+
emptyDataReceived = true;
412+
}
413+
}
414+
});
415+
416+
callback();
417+
});
418+
419+
const cubeSqlStream = new Writable({
420+
write,
421+
});
422+
423+
try {
424+
// Use LIMIT 0 to test the real case where SQL produces no results
425+
await native.execSql(
426+
instance,
427+
'SELECT order_date FROM KibanaSampleDataEcommerce LIMIT 0;',
428+
cubeSqlStream
429+
);
430+
431+
// Verify schema was sent and empty data was sent for LIMIT 0 query
432+
expect(schemaReceived).toBe(true);
433+
expect(dataReceived).toBe(true);
434+
expect(emptyDataReceived).toBe(true);
435+
} catch (error) {
436+
// Even if query fails, we should get schema
437+
console.log('Query error (expected in test):', error);
438+
if (schemaReceived) {
439+
expect(schemaReceived).toBe(true);
440+
}
441+
}
442+
443+
await native.shutdownInterface(instance, 'fast');
444+
});
375445
});

packages/cubejs-testing/test/smoke-cubesql.test.ts

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -148,6 +148,76 @@ describe('SQL API', () => {
148148
expect(rows).toBe(ROWS_LIMIT);
149149
});
150150

151+
it('streams schema and empty data with LIMIT 0', async () => {
152+
const response = await fetch(`${birdbox.configuration.apiUrl}/cubesql`, {
153+
method: 'POST',
154+
headers: {
155+
'Content-Type': 'application/json',
156+
Authorization: token,
157+
},
158+
body: JSON.stringify({
159+
query: `SELECT orderDate FROM ECommerce LIMIT 0;`,
160+
}),
161+
});
162+
163+
const reader = response.body;
164+
let isFirstChunk = true;
165+
let schemaReceived = false;
166+
let emptyDataReceived = false;
167+
168+
let data = '';
169+
const execute = () => new Promise<void>((resolve, reject) => {
170+
const onData = jest.fn((chunk: Buffer) => {
171+
const chunkStr = chunk.toString('utf-8');
172+
173+
if (isFirstChunk) {
174+
isFirstChunk = false;
175+
const json = JSON.parse(chunkStr);
176+
expect(json.schema).toEqual([
177+
{
178+
name: 'orderDate',
179+
column_type: 'Timestamp',
180+
},
181+
]);
182+
schemaReceived = true;
183+
} else {
184+
data += chunkStr;
185+
const json = JSON.parse(chunkStr);
186+
if (json.data && Array.isArray(json.data) && json.data.length === 0) {
187+
emptyDataReceived = true;
188+
}
189+
}
190+
});
191+
reader.on('data', onData);
192+
193+
const onError = jest.fn(() => reject(new Error('Stream error')));
194+
reader.on('error', onError);
195+
196+
const onEnd = jest.fn(() => {
197+
resolve();
198+
});
199+
200+
reader.on('end', onEnd);
201+
});
202+
203+
await execute();
204+
205+
// Verify schema was sent first
206+
expect(schemaReceived).toBe(true);
207+
208+
// Verify empty data was sent
209+
expect(emptyDataReceived).toBe(true);
210+
211+
// Verify no actual rows were returned
212+
const dataLines = data.split('\n').filter((it) => it.trim());
213+
if (dataLines.length > 0) {
214+
const rows = dataLines
215+
.map((it) => JSON.parse(it).data?.length || 0)
216+
.reduce((a, b) => a + b, 0);
217+
expect(rows).toBe(0);
218+
}
219+
});
220+
151221
describe('sql4sql', () => {
152222
async function generateSql(query: string, disablePostPprocessing: boolean = false) {
153223
const response = await fetch(`${birdbox.configuration.apiUrl}/sql`, {

0 commit comments

Comments
 (0)