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,