dbkit/remote/
clickhouse.rs

1//! ClickHouse remote source over the HTTP interface.
2//!
3//! Queries are sent with `FORMAT ArrowStream`, so ClickHouse returns an Arrow
4//! IPC stream that is parsed straight into `RecordBatch`es — no row-by-row
5//! decoding.
6
7use crate::DbkitError;
8use crate::analytical::RecordBatch;
9use crate::analytical::arrow::ipc::reader::StreamReader;
10use crate::analytical::arrow::ipc::writer::StreamWriter;
11use crate::remote::{RemoteSink, RemoteSource};
12use crate::value::DbValue;
13use async_trait::async_trait;
14use std::io::Cursor;
15
16/// A handle to a ClickHouse server queried over HTTP.
17///
18/// ```ignore
19/// let ch = ClickHouseSource::new("http://localhost:8123")
20///     .database("analytics")
21///     .credentials("default", "");
22/// let batches = ch.query_arrow("SELECT count() FROM events", &[]).await?;
23/// ```
24pub struct ClickHouseSource {
25    client: reqwest::Client,
26    url: String,
27    database: Option<String>,
28    user: Option<String>,
29    password: Option<String>,
30}
31
32impl ClickHouseSource {
33    /// Point at a ClickHouse HTTP endpoint, e.g. `http://localhost:8123`.
34    pub fn new(url: impl Into<String>) -> Self {
35        Self {
36            client: reqwest::Client::new(),
37            url: url.into(),
38            database: None,
39            user: None,
40            password: None,
41        }
42    }
43
44    /// Set the default database for queries.
45    pub fn database(mut self, database: impl Into<String>) -> Self {
46        self.database = Some(database.into());
47        self
48    }
49
50    /// Set HTTP basic-auth credentials.
51    pub fn credentials(mut self, user: impl Into<String>, password: impl Into<String>) -> Self {
52        self.user = Some(user.into());
53        self.password = Some(password.into());
54        self
55    }
56}
57
58#[async_trait]
59impl RemoteSource for ClickHouseSource {
60    async fn query_arrow(
61        &self,
62        sql: &str,
63        params: &[DbValue],
64    ) -> Result<Vec<RecordBatch>, DbkitError> {
65        if !params.is_empty() {
66            return Err(DbkitError::Remote(
67                "the ClickHouse source does not support bind parameters yet; \
68                 inline values into the query"
69                    .into(),
70            ));
71        }
72
73        // ArrowStream yields an Arrow IPC stream we can decode directly.
74        let body = format!("{sql}\nFORMAT ArrowStream");
75
76        let mut req = self.client.post(&self.url).body(body);
77        if let Some(db) = &self.database {
78            req = req.query(&[("database", db)]);
79        }
80        if let Some(user) = &self.user {
81            req = req.basic_auth(user, self.password.as_deref());
82        }
83
84        let resp = req
85            .send()
86            .await
87            .map_err(|e| DbkitError::Remote(e.to_string()))?;
88
89        // ClickHouse reports errors in the response body, so surface it.
90        if !resp.status().is_success() {
91            let status = resp.status();
92            let text = resp.text().await.unwrap_or_default();
93            return Err(DbkitError::Remote(format!(
94                "ClickHouse returned {status}: {}",
95                text.trim()
96            )));
97        }
98
99        let bytes = resp
100            .bytes()
101            .await
102            .map_err(|e| DbkitError::Remote(e.to_string()))?;
103
104        let reader = StreamReader::try_new(Cursor::new(bytes), None)
105            .map_err(|e| DbkitError::Remote(format!("arrow decode: {e}")))?;
106        reader
107            .collect::<Result<Vec<_>, _>>()
108            .map_err(|e| DbkitError::Remote(format!("arrow decode: {e}")))
109    }
110}
111
112#[async_trait]
113impl RemoteSink for ClickHouseSource {
114    async fn write_arrow(
115        &self,
116        table: &str,
117        batches: Vec<RecordBatch>,
118    ) -> Result<(), DbkitError> {
119        if batches.is_empty() {
120            return Ok(());
121        }
122
123        // Encode the batches as an Arrow IPC stream for `FORMAT ArrowStream`.
124        let schema = batches[0].schema();
125        let mut buf = Vec::new();
126        {
127            let mut writer = StreamWriter::try_new(&mut buf, &schema)
128                .map_err(|e| DbkitError::Remote(format!("arrow encode: {e}")))?;
129            for batch in &batches {
130                writer
131                    .write(batch)
132                    .map_err(|e| DbkitError::Remote(format!("arrow encode: {e}")))?;
133            }
134            writer
135                .finish()
136                .map_err(|e| DbkitError::Remote(format!("arrow encode: {e}")))?;
137        }
138
139        let mut req = self
140            .client
141            .post(&self.url)
142            .query(&[("query", format!("INSERT INTO {table} FORMAT ArrowStream"))])
143            .body(buf);
144        if let Some(db) = &self.database {
145            req = req.query(&[("database", db)]);
146        }
147        if let Some(user) = &self.user {
148            req = req.basic_auth(user, self.password.as_deref());
149        }
150
151        let resp = req
152            .send()
153            .await
154            .map_err(|e| DbkitError::Remote(e.to_string()))?;
155
156        if !resp.status().is_success() {
157            let status = resp.status();
158            let text = resp.text().await.unwrap_or_default();
159            return Err(DbkitError::Remote(format!(
160                "ClickHouse returned {status}: {}",
161                text.trim()
162            )));
163        }
164        Ok(())
165    }
166}
dbkit/remote/clickhouse.rs

dbkit/remote/
clickhouse.rs
