From 0b0e8ade1771d6081f99adc5cb1b2273857ae379 Mon Sep 17 00:00:00 2001 From: ZhuchkaTriplesix Date: Tue, 29 Sep 2026 10:04:55 +0300 Subject: [PATCH] perf(driver): stream ClickHouse HTTP responses instead of full-text buffering. Closes #49 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit db.query's tabular path called reqwest::Response::text(), forcing the entire HTTP response body into one contiguous String before any parsing started. For a large analytical SELECT (tens of MB of JSON lines), the raw text and the parsed QueryResult rows coexisted in memory at once, risking OOM under the 256 MB ClickHouse Sandbox ceiling. - New ClickHouseClient::post_sql_response returns the raw streaming reqwest::Response on success (reusing the existing readonly-retry logic) instead of buffering it into a String. - New src/driver/streaming::stream_compact_output reads that response's bytes_stream() through a LinesCodec-based FramedRead and parses rows incrementally as they arrive, stopping — without reading the rest of the network response — once `limit` rows are parsed or a 200 MB safety byte cap is hit even when no limit was given. - Extracted parse_columns/parse_and_normalize_row out of parse_compact_output so both the buffered (schema introspection, mock mode, tree building) and streaming (large query results) parsers share identical row normalization. - QueryResult gained isTruncated, set whenever limit or the safety cap cut a result short, so callers can tell rows isn't the complete set. - handle_query's tabular branch now goes through post_sql_response + stream_compact_output; the non-tabular (mutation/DDL) branch is unchanged, since those responses are always small. --- Cargo.lock | 1 + Cargo.toml | 3 +- src/driver/client.rs | 48 ++++++++-- src/driver/mod.rs | 1 + src/driver/streaming.rs | 179 ++++++++++++++++++++++++++++++++++++++ src/mapper/row_compact.rs | 152 +++++++++++++++++++------------- src/rpc/handlers/query.rs | 33 +++++-- 7 files changed, 342 insertions(+), 75 deletions(-) create mode 100644 src/driver/streaming.rs diff --git a/Cargo.lock b/Cargo.lock index f5053a7..9f3fd05 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -85,6 +85,7 @@ name = "clickhouse-query-ext" version = "1.0.1" dependencies = [ "anyhow", + "bytes", "futures", "libc", "regex", diff --git a/Cargo.toml b/Cargo.toml index 0753610..eedfb22 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,8 +8,9 @@ description = "High-performance ClickHouse driver for Querya Desktop" [dependencies] # Async runtime & I/O tokio = { version = "1", features = ["rt-multi-thread", "io-std", "sync", "macros", "time", "net"] } -tokio-util = { version = "0.7", features = ["codec"] } +tokio-util = { version = "0.7", features = ["codec", "io"] } futures = "0.3" +bytes = "1" # HTTP & TLS for ClickHouse API reqwest = { version = "0.12", default-features = false, features = ["json", "stream", "rustls-tls"] } diff --git a/src/driver/client.rs b/src/driver/client.rs index f429e5e..eeaa1cc 100644 --- a/src/driver/client.rs +++ b/src/driver/client.rs @@ -165,14 +165,15 @@ impl ClickHouseClient { req } + async fn error_from_response(resp: reqwest::Response) -> DriverError { + let status = resp.status(); + let text = resp.text().await.unwrap_or_default(); + DriverError::Client(format!("ClickHouse HTTP error {}: {}", status, text)) + } + async fn read_response(resp: reqwest::Response) -> Result { if !resp.status().is_success() { - let status = resp.status(); - let text = resp.text().await.unwrap_or_default(); - return Err(DriverError::Client(format!( - "ClickHouse HTTP error {}: {}", - status, text - ))); + return Err(Self::error_from_response(resp).await); } Ok(resp.text().await?) } @@ -248,6 +249,41 @@ impl ClickHouseClient { } } + /// Like `post_sql`, but on success returns the raw streaming `reqwest::Response` + /// instead of buffering the whole body into a `String`, so a large analytical + /// result set can be parsed incrementally as it arrives over the network via + /// `crate::driver::streaming::stream_compact_output` (issue #49). Callers must + /// gate on mock/test connections themselves, same as before calling `post_sql`. + pub async fn post_sql_response( + &self, + sql: &str, + mut extra_params: impl FnMut(&mut Url), + ) -> Result { + let sql = sql.to_string(); + let mut omit = self.omit_readonly_setting(); + loop { + let mut url = Url::parse(&self.base_url)?; + url.query_pairs_mut() + .append_pair("database", &self.database); + extra_params(&mut url); + self.append_safe_mode_settings_with(&mut url, omit); + let req = self + .apply_auth(self.http_client.post(url)) + .body(sql.clone()); + let resp = req.send().await?; + if resp.status().is_success() { + return Ok(resp); + } + let err = Self::error_from_response(resp).await; + if self.readonly && !omit && Self::is_readonly_setting_conflict(&err) { + self.mark_server_readonly_enforced(); + omit = true; + continue; + } + return Err(err); + } + } + /// Check connection health by executing `SELECT version()` against ClickHouse. pub async fn ping_connection(&self) -> Result { let version = self.get_with_query("SELECT version()").await?; diff --git a/src/driver/mod.rs b/src/driver/mod.rs index f691764..a5b9f42 100644 --- a/src/driver/mod.rs +++ b/src/driver/mod.rs @@ -1,3 +1,4 @@ //! ClickHouse HTTP client and session management. pub mod client; pub mod pool; +pub mod streaming; diff --git a/src/driver/streaming.rs b/src/driver/streaming.rs new file mode 100644 index 0000000..58fe3c1 --- /dev/null +++ b/src/driver/streaming.rs @@ -0,0 +1,179 @@ +//! Incremental parsing of ClickHouse's `FORMAT JSONCompactEachRowWithNamesAndTypes` +//! HTTP response, row by row, as bytes arrive over the network. +//! +//! `db.query`'s previous implementation buffered the entire HTTP response body +//! into a single `String` (`reqwest::Response::text()`) before parsing began, +//! so a large analytical result (tens of megabytes of JSON lines) held both +//! the raw text and the parsed rows in memory simultaneously — a real OOM risk +//! under the 256 MB ClickHouse Sandbox ceiling. This module instead reads the +//! response as a line stream and stops pulling further bytes off the network +//! connection as soon as either the caller's `limit` or a safety byte cap is +//! reached (see issue #49). + +use crate::error::DriverError; +use crate::mapper::row_compact::{ + QueryResult, QueryStatistics, parse_and_normalize_row, parse_columns, +}; +use bytes::Bytes; +use futures::Stream; +use futures::StreamExt; +use tokio_util::codec::{FramedRead, LinesCodec}; +use tokio_util::io::StreamReader; + +/// Independent of any client-supplied `limit`, stop consuming further row +/// bytes once this many have been read, so an unbounded query (no `limit` +/// sent) still can't grow the result past a safe ceiling. +const MAX_ROW_BYTES: usize = 200 * 1024 * 1024; + +/// ClickHouse lines (JSON arrays of column values) are expected to be well +/// under this; it exists only to bound a single corrupt/adversarial line's +/// buffered length instead of growing unbounded. +const MAX_LINE_BYTES: usize = 64 * 1024 * 1024; + +/// Reads a byte stream (a ClickHouse HTTP response body) as a stream of +/// lines and incrementally parses it as `FORMAT JSONCompactEachRowWithNamesAndTypes` +/// output, stopping (without reading the rest of the stream) once `limit` +/// rows have been parsed or the `MAX_ROW_BYTES` safety cap is reached. +pub async fn stream_compact_output( + byte_stream: S, + limit: Option, +) -> Result +where + S: Stream> + Unpin, +{ + let stream_reader = StreamReader::new(byte_stream); + let mut lines = FramedRead::new( + stream_reader, + LinesCodec::new_with_max_length(MAX_LINE_BYTES), + ); + + let mut bytes_read: usize = 0; + + let Some(names_line) = lines.next().await.transpose().map_err(line_err)? else { + return Ok(QueryResult { + columns: vec![], + rows: vec![], + statistics: QueryStatistics { + rows_read: 0, + bytes_read: 0, + elapsed_ms: 0, + }, + query_id: None, + is_truncated: false, + }); + }; + bytes_read += names_line.len() + 1; + + let Some(types_line) = lines.next().await.transpose().map_err(line_err)? else { + return Err(DriverError::Client( + "Malformed JSONCompactEachRowWithNamesAndTypes output: missing names or types row" + .to_string(), + )); + }; + bytes_read += types_line.len() + 1; + + let columns = parse_columns(&names_line, &types_line)?; + + let mut rows = Vec::with_capacity(limit.unwrap_or(16).min(1024)); + let mut is_truncated = false; + while let Some(line) = lines.next().await.transpose().map_err(line_err)? { + bytes_read += line.len() + 1; + + if limit.is_some_and(|limit| rows.len() >= limit) || bytes_read > MAX_ROW_BYTES { + is_truncated = true; + break; + } + + rows.push(parse_and_normalize_row(&line, &columns)?); + } + // Drop the frame reader (and the underlying HTTP connection) now instead of + // reading any remaining body bytes off the network when we stopped early. + drop(lines); + + let rows_read = rows.len(); + Ok(QueryResult { + columns, + rows, + statistics: QueryStatistics { + rows_read, + bytes_read, + elapsed_ms: 0, + }, + query_id: None, + is_truncated, + }) +} + +fn line_err(e: tokio_util::codec::LinesCodecError) -> DriverError { + DriverError::Client(format!("Failed to read ClickHouse response stream: {}", e)) +} + +#[cfg(test)] +mod tests { + use super::*; + use futures::stream; + + /// Builds an in-memory byte stream out of a static string, split into + /// arbitrary chunks, to exercise `stream_compact_output` without a live + /// HTTP connection. Splitting mid-line (not just at newlines) verifies + /// the line codec correctly reassembles frames split across chunks. + fn chunked_stream( + body: &'static str, + chunk_size: usize, + ) -> impl Stream> { + let chunks: Vec> = body + .as_bytes() + .chunks(chunk_size.max(1)) + .map(|c| Ok(Bytes::copy_from_slice(c))) + .collect(); + stream::iter(chunks) + } + + #[tokio::test] + async fn test_stream_compact_output_parses_rows() { + let body = "[\"id\"]\n[\"UInt64\"]\n[1]\n[2]\n[3]\n"; + let result = stream_compact_output(chunked_stream(body, 1024), None) + .await + .unwrap(); + assert_eq!(result.rows.len(), 3); + assert!(!result.is_truncated); + } + + #[tokio::test] + async fn test_stream_compact_output_reassembles_lines_split_across_chunks() { + let body = "[\"id\"]\n[\"UInt64\"]\n[1]\n[2]\n[3]\n"; + // Force byte-at-a-time delivery so every line boundary falls mid-chunk. + let result = stream_compact_output(chunked_stream(body, 1), None) + .await + .unwrap(); + assert_eq!(result.rows.len(), 3); + } + + #[tokio::test] + async fn test_stream_compact_output_enforces_limit() { + let body = "[\"id\"]\n[\"UInt64\"]\n[1]\n[2]\n[3]\n"; + let result = stream_compact_output(chunked_stream(body, 1024), Some(2)) + .await + .unwrap(); + assert_eq!(result.rows.len(), 2); + assert!(result.is_truncated); + } + + #[tokio::test] + async fn test_stream_compact_output_empty_body() { + let result = stream_compact_output(chunked_stream("", 1024), None) + .await + .unwrap(); + assert!(result.rows.is_empty()); + assert!(result.columns.is_empty()); + } + + #[tokio::test] + async fn test_stream_compact_output_missing_types_row() { + let body = "[\"id\"]\n"; + let err = stream_compact_output(chunked_stream(body, 1024), None) + .await + .unwrap_err(); + assert!(err.to_string().contains("missing names or types row")); + } +} diff --git a/src/mapper/row_compact.rs b/src/mapper/row_compact.rs index df68901..c134da0 100644 --- a/src/mapper/row_compact.rs +++ b/src/mapper/row_compact.rs @@ -19,6 +19,88 @@ pub struct QueryResult { pub statistics: QueryStatistics, #[serde(skip_serializing_if = "Option::is_none")] pub query_id: Option, + /// `true` when `limit` (or the streaming safety byte cap) cut the result + /// short of what ClickHouse actually had to offer, so the caller knows + /// `rows` isn't the complete result set. + #[serde(skip_serializing_if = "std::ops::Not::not")] + pub is_truncated: bool, +} + +/// Parses the `["name", ...]` / `["Type", ...]` header lines of +/// `FORMAT JSONCompactEachRowWithNamesAndTypes` output into `ColumnSchema`s. +/// Shared by both the buffered (`parse_compact_output`) and streaming +/// (`crate::driver::streaming::stream_compact_output`) parsers. +pub(crate) fn parse_columns( + names_line: &str, + types_line: &str, +) -> Result, DriverError> { + let names: Vec = serde_json::from_str(names_line).map_err(|e| { + DriverError::Client(format!( + "Failed to parse column names from ClickHouse output: {}", + e + )) + })?; + let types: Vec = serde_json::from_str(types_line).map_err(|e| { + DriverError::Client(format!( + "Failed to parse column types from ClickHouse output: {}", + e + )) + })?; + + if names.len() != types.len() { + return Err(DriverError::Client(format!( + "Column names count ({}) does not match types count ({})", + names.len(), + types.len() + ))); + } + + Ok(names + .into_iter() + .zip(types) + .map(|(name, ch_type)| ColumnSchema::new(name, ch_type)) + .collect()) +} + +/// Parses a single `[value, value, ...]` data row line and normalizes it +/// according to each column's `mapped_type` (e.g. converting 64-bit numbers +/// and Decimals into JSON strings to prevent 53-bit float overflow in +/// JS/Flutter). Shared by both the buffered and streaming parsers. +pub(crate) fn parse_and_normalize_row( + line: &str, + columns: &[ColumnSchema], +) -> Result, DriverError> { + let mut raw_row: Vec = serde_json::from_str(line).map_err(|e| { + DriverError::Client(format!( + "Failed to parse data row JSON array from ClickHouse: {}", + e + )) + })?; + + for (i, col) in columns.iter().enumerate() { + if let Some(val) = raw_row.get_mut(i) { + if val.is_null() { + continue; + } + match col.mapped_type { + "string" => { + // 64-bit/large integers and Decimals may arrive as JSON numbers from ClickHouse + if val.is_number() { + *val = Value::String(val.to_string()); + } + } + "integer" => { + if let Some(s) = val.as_str() + && let Ok(n) = s.parse::() + { + *val = Value::Number(serde_json::Number::from(n)); + } + } + _ => {} + } + } + } + Ok(raw_row) } /// Parses the output of ClickHouse `FORMAT JSONCompactEachRowWithNamesAndTypes`. @@ -51,6 +133,7 @@ pub fn parse_compact_output( elapsed_ms, }, query_id: None, + is_truncated: false, }); }; @@ -61,70 +144,16 @@ pub fn parse_compact_output( )); }; - let names: Vec = serde_json::from_str(names_line).map_err(|e| { - DriverError::Client(format!( - "Failed to parse column names from ClickHouse output: {}", - e - )) - })?; - let types: Vec = serde_json::from_str(types_line).map_err(|e| { - DriverError::Client(format!( - "Failed to parse column types from ClickHouse output: {}", - e - )) - })?; - - if names.len() != types.len() { - return Err(DriverError::Client(format!( - "Column names count ({}) does not match types count ({})", - names.len(), - types.len() - ))); - } - - let mut columns = Vec::with_capacity(names.len()); - for (name, ch_type) in names.into_iter().zip(types) { - columns.push(ColumnSchema::new(name, ch_type)); - } + let columns = parse_columns(names_line, types_line)?; let mut rows = Vec::with_capacity(limit.unwrap_or(16).min(1024)); + let mut is_truncated = false; for line in lines { if limit.is_some_and(|limit| rows.len() >= limit) { + is_truncated = true; break; } - - let mut raw_row: Vec = serde_json::from_str(line).map_err(|e| { - DriverError::Client(format!( - "Failed to parse data row JSON array from ClickHouse: {}", - e - )) - })?; - - // Normalize values according to Querya schema mapped_type - for (i, col) in columns.iter().enumerate() { - if let Some(val) = raw_row.get_mut(i) { - if val.is_null() { - continue; - } - match col.mapped_type { - "string" => { - // 64-bit/large integers and Decimals may arrive as JSON numbers from ClickHouse - if val.is_number() { - *val = Value::String(val.to_string()); - } - } - "integer" => { - if let Some(s) = val.as_str() - && let Ok(n) = s.parse::() - { - *val = Value::Number(serde_json::Number::from(n)); - } - } - _ => {} - } - } - } - rows.push(raw_row); + rows.push(parse_and_normalize_row(line, &columns)?); } let rows_read = rows.len(); @@ -137,6 +166,7 @@ pub fn parse_compact_output( elapsed_ms, }, query_id: None, + is_truncated, }) } @@ -218,6 +248,7 @@ mod tests { let unlimited = parse_compact_output(raw_output, 0, None).unwrap(); assert_eq!(unlimited.rows.len(), 5); assert_eq!(unlimited.statistics.rows_read, 5); + assert!(!unlimited.is_truncated); let limited = parse_compact_output(raw_output, 0, Some(2)).unwrap(); assert_eq!(limited.rows.len(), 2); @@ -225,10 +256,13 @@ mod tests { assert_eq!(limited.rows[0][0], json!("1")); assert_eq!(limited.rows[1][0], json!("2")); assert_eq!(limited.statistics.rows_read, 2); + assert!(limited.is_truncated); - // A limit larger than the actual row count is a no-op. + // A limit larger than the actual row count is a no-op, and isn't + // reported as truncated since nothing was actually cut off. let generous_limit = parse_compact_output(raw_output, 0, Some(100)).unwrap(); assert_eq!(generous_limit.rows.len(), 5); + assert!(!generous_limit.is_truncated); // A zero limit returns no rows at all, without erroring. let zero_limit = parse_compact_output(raw_output, 0, Some(0)).unwrap(); diff --git a/src/rpc/handlers/query.rs b/src/rpc/handlers/query.rs index 73199da..9d11c0f 100644 --- a/src/rpc/handlers/query.rs +++ b/src/rpc/handlers/query.rs @@ -2,6 +2,7 @@ use crate::driver::pool::ConnectionPool; use crate::error::DriverError; use crate::mapper::row_compact::parse_compact_output; use crate::utils::secret_guard::ConnectionSecretsPool; +use futures::StreamExt; use serde::Deserialize; use serde_json::{Value, json}; use std::sync::atomic::{AtomicU64, Ordering}; @@ -360,19 +361,33 @@ pub async fn handle_query(params: Option) -> Result { // 3. Real ClickHouse HTTP request let actual_query_id_for_url = actual_query_id.clone(); - let text = client - .post_sql(&sql_to_run, |url| { - url.query_pairs_mut() - .append_pair("query_id", &actual_query_id_for_url); - }) - .await?; - let elapsed = start_time.elapsed().as_millis() as u64; - if is_tabular_query { - let mut result = parse_compact_output(&text, elapsed, query_params.limit)?; + // Stream and parse the response row-by-row instead of buffering the + // whole body into a String first, bounding peak memory for large + // analytical result sets (issue #49). + let response = client + .post_sql_response(&sql_to_run, |url| { + url.query_pairs_mut() + .append_pair("query_id", &actual_query_id_for_url); + }) + .await?; + let byte_stream = response + .bytes_stream() + .map(|chunk| chunk.map_err(std::io::Error::other)); + let mut result = + crate::driver::streaming::stream_compact_output(byte_stream, query_params.limit) + .await?; + result.statistics.elapsed_ms = start_time.elapsed().as_millis() as u64; result.query_id = Some(actual_query_id); Ok(serde_json::to_value(result)?) } else { + let text = client + .post_sql(&sql_to_run, |url| { + url.query_pairs_mut() + .append_pair("query_id", &actual_query_id_for_url); + }) + .await?; + let elapsed = start_time.elapsed().as_millis() as u64; Ok(build_non_tabular_result( &upper_sql, elapsed,