Expose summary in request (#397)

Author: alrevuelta

Cargo.toml

@@ -126,6 +126,7 @@ clickhouse-macros = { version = "0.3.0", path = "macros" }
 clickhouse-types = { version = "0.1.2", path = "types" }
 
 serde = { version = "1.0.106", features = ["derive"] }
+serde_json = "1"
 
 thiserror = "2.0"
 bytes = { version = "1.5.0", features = ["serde"] }

src/cursors/bytes.rs

@@ -1,4 +1,4 @@
-use crate::{cursors::RawCursor, error::Result, response::Response};
+use crate::{cursors::RawCursor, error::Result, query_summary::QuerySummary, response::Response};
 use bytes::{Buf, Bytes, BytesMut};
 use std::{
     io::Result as IoResult,
@@ -123,6 +123,16 @@ impl BytesCursor {
     pub fn decoded_bytes(&self) -> u64 {
         self.raw.decoded_bytes()
     }
+
+    /// Returns the parsed `X-ClickHouse-Summary` response header, if
+    /// present. Available once the response headers have been received.
+    ///
+    /// Note: the summary values may be incomplete unless the query was
+    /// executed with `wait_end_of_query=1`.
+    #[inline]
+    pub fn summary(&self) -> Option<&QuerySummary> {
+        self.raw.summary()
+    }
 }
 
 impl AsyncRead for BytesCursor {

src/cursors/raw.rs

@@ -1,5 +1,6 @@
 use crate::{
     error::Result,
+    query_summary::QuerySummary,
     response::{Chunks, Response, ResponseFuture},
 };
 use bytes::Bytes;
@@ -20,6 +21,7 @@ enum RawCursorState {
 
 struct RawCursorLoading {
     chunks: Chunks,
+    summary: Option<Box<QuerySummary>>,
     net_size: u64,
     data_size: u64,
 }
@@ -63,10 +65,15 @@ impl RawCursor {
         // in order to provide proper fused behavior of the cursor.
         let res = ready!(future.as_mut().poll(cx));
         let mut chunks = Chunks::empty();
-        let res = res.map(|c| chunks = c);
+        let mut summary = None;
+        let res = res.map(|(c, s)| {
+            chunks = c;
+            summary = s;
+        });
 
         self.0 = RawCursorState::Loading(RawCursorLoading {
             chunks,
+            summary,
             net_size: 0,
             data_size: 0,
         });
@@ -88,6 +95,13 @@ impl RawCursor {
         }
     }
 
+    pub(crate) fn summary(&self) -> Option<&QuerySummary> {
+        match &self.0 {
+            RawCursorState::Loading(state) => state.summary.as_deref(),
+            RawCursorState::Waiting(_) => None,
+        }
+    }
+
     #[cfg(feature = "futures03")]
     pub(crate) fn is_terminated(&self) -> bool {
         match &self.0 {

src/cursors/row.rs

@@ -6,6 +6,7 @@ use crate::{
     bytes_ext::BytesExt,
     cursors::RawCursor,
     error::{Error, Result},
+    query_summary::QuerySummary,
     response::Response,
     rowbinary,
 };
@@ -159,6 +160,16 @@ impl<T> RowCursor<T> {
     pub fn decoded_bytes(&self) -> u64 {
         self.raw.decoded_bytes()
     }
+
+    /// Returns the parsed `X-ClickHouse-Summary` response header, if
+    /// present. Available once the response headers have been received.
+    ///
+    /// Note: the summary values may be incomplete unless the query was
+    /// executed with `wait_end_of_query=1`.
+    #[inline]
+    pub fn summary(&self) -> Option<&QuerySummary> {
+        self.raw.summary()
+    }
 }
 
 #[cfg(feature = "futures03")]

src/lib.rs

@@ -3,6 +3,7 @@
 
 pub use self::{
     compression::Compression,
+    query_summary::QuerySummary,
     row::{Row, RowOwned, RowRead, RowWrite},
 };
 use self::{error::Result, http_client::HttpClient};
@@ -35,6 +36,7 @@ mod compression;
 mod cursors;
 mod headers;
 mod http_client;
+mod query_summary;
 mod request_body;
 mod response;
 mod row;

src/query_summary.rs

@@ -0,0 +1,78 @@
+use std::collections::HashMap;
+
+/// Parsed representation of the `X-ClickHouse-Summary` HTTP response header.
+///
+/// Provides typed getters for known fields and a generic [`get`](Self::get)
+/// fallback for forward-compatibility with future ClickHouse versions.
+///
+/// All getters return `Option<u64>`: `None` if the field is absent or cannot
+/// be parsed. This ensures deserialization never fails, even if ClickHouse
+/// renames, removes, or adds fields.
+///
+/// Note: the summary values may be incomplete unless the query was executed
+/// with `wait_end_of_query=1`, because ClickHouse sends this header before
+/// the response body and the values reflect progress at that point.
+#[derive(Debug, Clone)]
+pub struct QuerySummary {
+    fields: HashMap<String, String>,
+}
+
+impl QuerySummary {
+    /// Returns the raw string value for the given key, if present.
+    ///
+    /// Use this to access fields that are not yet covered by typed getters,
+    /// e.g. fields added in newer ClickHouse versions.
+    pub fn get(&self, key: &str) -> Option<&str> {
+        self.fields.get(key).map(String::as_str)
+    }
+
+    pub fn read_rows(&self) -> Option<u64> {
+        self.get_u64("read_rows")
+    }
+
+    pub fn read_bytes(&self) -> Option<u64> {
+        self.get_u64("read_bytes")
+    }
+
+    pub fn written_rows(&self) -> Option<u64> {
+        self.get_u64("written_rows")
+    }
+
+    pub fn written_bytes(&self) -> Option<u64> {
+        self.get_u64("written_bytes")
+    }
+
+    pub fn total_rows_to_read(&self) -> Option<u64> {
+        self.get_u64("total_rows_to_read")
+    }
+
+    pub fn result_rows(&self) -> Option<u64> {
+        self.get_u64("result_rows")
+    }
+
+    pub fn result_bytes(&self) -> Option<u64> {
+        self.get_u64("result_bytes")
+    }
+
+    pub fn elapsed_ns(&self) -> Option<u64> {
+        self.get_u64("elapsed_ns")
+    }
+
+    pub fn memory_usage(&self) -> Option<u64> {
+        self.get_u64("memory_usage")
+    }
+
+    fn get_u64(&self, key: &str) -> Option<u64> {
+        self.fields.get(key)?.parse().ok()
+    }
+
+    /// Parses the raw header value into a `QuerySummary`.
+    ///
+    /// Returns `None` if the value is not valid JSON or not an object with
+    /// string values. This matches ClickHouse's encoding, where all values
+    /// are JSON strings (e.g. `"1000"` instead of `1000`).
+    pub(crate) fn from_header(raw: &str) -> Option<Self> {
+        let fields: HashMap<String, String> = serde_json::from_str(raw).ok()?;
+        Some(Self { fields })
+    }
+}

src/response.rs

@@ -18,6 +18,7 @@ use crate::compression::lz4::Lz4Decoder;
 use crate::{
     compression::Compression,
     error::{Error, Result},
+    query_summary::QuerySummary,
 };
 
 // === Response ===
@@ -30,7 +31,8 @@ pub(crate) enum Response {
     Loading(Chunks),
 }
 
-pub(crate) type ResponseFuture = Pin<Box<dyn Future<Output = Result<Chunks>> + Send>>;
+pub(crate) type ResponseFuture =
+    Pin<Box<dyn Future<Output = Result<(Chunks, Option<Box<QuerySummary>>)>> + Send>>;
 
 impl Response {
     pub(crate) fn new(response: HyperResponseFuture, compression: Compression) -> Self {
@@ -46,9 +48,16 @@ impl Response {
                     .get("X-ClickHouse-Exception-Tag")
                     .map(|value| value.as_bytes().into());
 
+                let summary = response
+                    .headers()
+                    .get("X-ClickHouse-Summary")
+                    .and_then(|v| v.to_str().ok())
+                    .and_then(QuerySummary::from_header)
+                    .map(Box::new);
+
                 // More likely to be successful, start streaming.
                 // It still can fail, but we'll handle it in `DetectDbException`.
-                Ok(Chunks::new(response.into_body(), compression, tag))
+                Ok((Chunks::new(response.into_body(), compression, tag), summary))
             } else {
                 // An instantly failed request.
                 Err(collect_bad_response(
@@ -74,7 +83,10 @@ impl Response {
     pub(crate) async fn finish(&mut self) -> Result<()> {
         let chunks = loop {
             match self {
-                Self::Waiting(future) => *self = Self::Loading(future.await?),
+                Self::Waiting(future) => {
+                    let (chunks, _summary) = future.await?;
+                    *self = Self::Loading(chunks);
+                }
                 Self::Loading(chunks) => break chunks,
             }
         };
@@ -157,7 +169,9 @@ pub(crate) struct Chunk {
 
 // * Uses `Option<_>` to make this stream fused.
 // * Uses `Box<_>` in order to reduce the size of cursors.
-pub(crate) struct Chunks(Option<Box<DetectDbException<Decompress<IncomingStream>>>>);
+pub(crate) struct Chunks {
+    inner: Option<Box<DetectDbException<Decompress<IncomingStream>>>>,
+}
 
 impl Chunks {
     fn new(stream: Incoming, compression: Compression, exception_tag: Option<Box<[u8]>>) -> Self {
@@ -167,16 +181,18 @@ impl Chunks {
             stream,
             exception_tag,
         };
-        Self(Some(Box::new(stream)))
+        Self {
+            inner: Some(Box::new(stream)),
+        }
     }
 
     pub(crate) fn empty() -> Self {
-        Self(None)
+        Self { inner: None }
     }
 
     #[cfg(feature = "futures03")]
     pub(crate) fn is_terminated(&self) -> bool {
-        self.0.is_none()
+        self.inner.is_none()
     }
 }
 
@@ -185,11 +201,11 @@ impl Stream for Chunks {
 
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         // We use `take()` to make the stream fused, including the case of panics.
-        if let Some(mut stream) = self.0.take() {
+        if let Some(mut stream) = self.inner.take() {
             let res = Pin::new(&mut stream).poll_next(cx);
 
             if matches!(res, Poll::Pending | Poll::Ready(Some(Ok(_)))) {
-                self.0 = Some(stream);
+                self.inner = Some(stream);
             }
 
             res

src/test/handlers.rs

@@ -60,6 +60,26 @@ where
     Thunk(Response::new(buffer.into()))
 }
 
+// === provide_with_summary ===
+
+/// Like [`provide`], but includes an `X-ClickHouse-Summary` response header.
+#[track_caller]
+pub fn provide_with_summary<T>(rows: impl IntoIterator<Item = T>, summary: &str) -> impl Handler
+where
+    T: Serialize + Row,
+{
+    let mut buffer = Vec::with_capacity(BUFFER_INITIAL_CAPACITY);
+    for row in rows {
+        rowbinary::serialize_row_binary(&mut buffer, &row).expect("failed to serialize");
+    }
+    Thunk(
+        Response::builder()
+            .header("X-ClickHouse-Summary", summary)
+            .body(Bytes::from(buffer))
+            .expect("invalid builder"),
+    )
+}
+
 // === record ===
 
 struct RecordHandler<T>(PhantomData<T>);

tests/it/main.rs

@@ -264,6 +264,7 @@ mod mock;
 mod nested;
 mod query;
 mod query_readonly;
+mod query_summary;
 mod query_syntax;
 mod rbwnat_header;
 mod rbwnat_smoke;