faucet_source_clickhouse/
stream.rs

1//! The ClickHouse [`Source`] implementation — HTTP client, query execution,
2//! streaming JSONEachRow decode, and incremental-replication bookkeeping.
3
4use std::collections::HashMap;
5use std::hash::{Hash, Hasher};
6use std::pin::Pin;
7use std::sync::Mutex;
8
9use async_trait::async_trait;
10use faucet_common_clickhouse::{
11    apply_auth, build_client, parse_json_each_row, query_params, sql_literal,
12};
13use faucet_core::check::{CheckContext, CheckReport, Probe};
14use faucet_core::replication::{filter_incremental, max_replication_value, max_value};
15use faucet_core::util::{DEFAULT_ERROR_BODY_MAX_LEN, check_http_response};
16use faucet_core::{FaucetError, Source, Stream, StreamPage};
17use futures::StreamExt;
18use serde_json::Value;
19
20use crate::config::{ClickHouseReplication, ClickHouseSourceConfig};
21
22/// ClickHouse query source (HTTP interface, `JSONEachRow`).
23pub struct ClickHouseSource {
24    config: ClickHouseSourceConfig,
25    client: reqwest::Client,
26    /// Resolved once in [`ClickHouseSource::new`] so the hot path never re-parses.
27    base_url: String,
28    /// Bookmark loaded via [`Source::apply_start_bookmark`]; overrides the
29    /// configured `initial_value` for incremental runs.
30    start_bookmark: Mutex<Option<Value>>,
31}
32
33/// Incremental-replication context resolved for one run.
34#[derive(Debug, Clone, PartialEq)]
35struct IncrementalCtx {
36    column: String,
37    start: Value,
38}
39
40impl ClickHouseSource {
41    /// Validate the config and build the reusable HTTP client.
42    pub fn new(config: ClickHouseSourceConfig) -> Result<Self, FaucetError> {
43        config.validate()?;
44        let base_url = config.connection.base_url()?;
45        let client = build_client(&config.connection)?;
46        Ok(Self {
47            config,
48            client,
49            base_url,
50            start_bookmark: Mutex::new(None),
51        })
52    }
53
54    fn current_start(&self) -> Option<Value> {
55        self.start_bookmark
56            .lock()
57            .expect("start_bookmark mutex poisoned")
58            .clone()
59    }
60
61    /// Build the POST request that runs `query` and returns `JSONEachRow`.
62    fn request(&self, query: String) -> reqwest::RequestBuilder {
63        let params = query_params(
64            &self.config.connection.database,
65            &[("default_format", "JSONEachRow")],
66        );
67        let req = self.client.post(&self.base_url).query(&params).body(query);
68        apply_auth(req, &self.config.connection)
69    }
70
71    /// Run the query and collect all decoded rows plus (for incremental) the new
72    /// bookmark. Used by the non-streaming convenience methods.
73    async fn collect_all(
74        &self,
75        context: &HashMap<String, Value>,
76    ) -> Result<(Vec<Value>, Option<Value>), FaucetError> {
77        let start = self.current_start();
78        let (query, incr) = build_effective_query(&self.config, context, start.as_ref());
79
80        let resp = self.request(query).send().await?;
81        let resp = check_http_response(resp, DEFAULT_ERROR_BODY_MAX_LEN).await?;
82        let body = resp.text().await.map_err(|e| {
83            FaucetError::Source(format!("ClickHouse: reading response failed: {e}"))
84        })?;
85        let records = parse_json_each_row(&body)?;
86
87        let mut running_max: Option<Value> = None;
88        let records = apply_incremental(records, incr.as_ref(), &mut running_max);
89        let bookmark = if incr.is_some() { running_max } else { None };
90        Ok((records, bookmark))
91    }
92}
93
94/// Build the final query string and (for incremental runs) the client-side
95/// filter context. Pure (no client) so it is unit-testable.
96///
97/// Substitution order: parent-context `{key}` tokens (as injection-safe SQL
98/// literals) → the incremental bookmark bound where the user wrote `@bookmark`.
99fn build_effective_query(
100    config: &ClickHouseSourceConfig,
101    context: &HashMap<String, Value>,
102    start_bookmark: Option<&Value>,
103) -> (String, Option<IncrementalCtx>) {
104    let mut query = if context.is_empty() {
105        config.query.clone()
106    } else {
107        substitute_context_sql(&config.query, context)
108    };
109
110    let incremental = match &config.replication {
111        ClickHouseReplication::Full => None,
112        ClickHouseReplication::Incremental {
113            column,
114            initial_value,
115        } => {
116            let start = start_bookmark
117                .cloned()
118                .unwrap_or_else(|| initial_value.clone());
119            // Server-side pushdown: substitute the cursor as an injection-safe
120            // SQL literal where the user wrote `@bookmark`. If absent, only the
121            // client-side filter applies.
122            if query.contains("@bookmark") {
123                query = query.replace("@bookmark", &sql_literal(&start));
124            }
125            Some(IncrementalCtx {
126                column: column.clone(),
127                start,
128            })
129        }
130    };
131
132    (query, incremental)
133}
134
135/// Replace each `{key}` token with the injection-safe SQL literal of the
136/// corresponding context value. Tokens with no matching context entry are left
137/// verbatim.
138fn substitute_context_sql(query: &str, context: &HashMap<String, Value>) -> String {
139    let mut out = query.to_string();
140    for (key, value) in context {
141        out = out.replace(&format!("{{{key}}}"), &sql_literal(value));
142    }
143    out
144}
145
146/// Filter a page for incremental replication and advance `running_max`.
147/// For full replication the page passes through unchanged.
148fn apply_incremental(
149    page: Vec<Value>,
150    incr: Option<&IncrementalCtx>,
151    running_max: &mut Option<Value>,
152) -> Vec<Value> {
153    match incr {
154        None => page,
155        Some(ctx) => {
156            let kept = filter_incremental(page, &ctx.column, &ctx.start);
157            if let Some(m) = max_replication_value(&kept, &ctx.column) {
158                let m = m.clone();
159                *running_max = Some(match running_max.take() {
160                    Some(prev) => max_value(prev, m),
161                    None => m,
162                });
163            }
164            kept
165        }
166    }
167}
168
169/// Drain every complete `\n`-terminated line from `buf`, returning each line's
170/// bytes (without the trailing newline). Splitting on the `0x0A` byte is safe
171/// for UTF-8 because a newline never appears inside a multi-byte sequence, so a
172/// line split across two network chunks is reassembled correctly. Pure and
173/// unit-testable.
174fn split_complete_lines(buf: &mut Vec<u8>) -> Vec<Vec<u8>> {
175    let mut lines = Vec::new();
176    while let Some(pos) = buf.iter().position(|&b| b == b'\n') {
177        let mut line: Vec<u8> = buf.drain(..=pos).collect();
178        line.pop(); // strip the trailing '\n'
179        lines.push(line);
180    }
181    lines
182}
183
184/// Parse one raw JSONEachRow line into a JSON value. Blank lines yield `None`.
185/// Invalid UTF-8 or JSON surfaces as a typed [`FaucetError::Source`].
186fn parse_line(line: &[u8]) -> Result<Option<Value>, FaucetError> {
187    let text = std::str::from_utf8(line)
188        .map_err(|e| FaucetError::Source(format!("ClickHouse: non-UTF-8 response line: {e}")))?
189        .trim();
190    if text.is_empty() {
191        return Ok(None);
192    }
193    let value: Value = serde_json::from_str(text).map_err(|e| {
194        FaucetError::Source(format!("ClickHouse: failed to parse JSONEachRow line: {e}"))
195    })?;
196    Ok(Some(value))
197}
198
199/// Derive a default state-store key from the connection host + a query
200/// fingerprint, stable across runs.
201fn default_state_key(config: &ClickHouseSourceConfig) -> String {
202    let host = config
203        .connection
204        .base_url()
205        .ok()
206        .and_then(|u| url::Url::parse(&u).ok())
207        .and_then(|u| u.host_str().map(str::to_string))
208        .unwrap_or_else(|| "clickhouse".to_string());
209
210    let mut hasher = std::collections::hash_map::DefaultHasher::new();
211    config.query.hash(&mut hasher);
212    let fingerprint = hasher.finish();
213    let host: String = host
214        .chars()
215        .map(|c| {
216            if c.is_ascii_alphanumeric() || matches!(c, '_' | '-' | '.') {
217                c
218            } else {
219                '_'
220            }
221        })
222        .collect();
223    format!("clickhouse:{host}:{fingerprint:016x}")
224}
225
226#[async_trait]
227impl Source for ClickHouseSource {
228    async fn fetch_with_context(
229        &self,
230        context: &HashMap<String, Value>,
231    ) -> Result<Vec<Value>, FaucetError> {
232        Ok(self.collect_all(context).await?.0)
233    }
234
235    async fn fetch_with_context_incremental(
236        &self,
237        context: &HashMap<String, Value>,
238    ) -> Result<(Vec<Value>, Option<Value>), FaucetError> {
239        self.collect_all(context).await
240    }
241
242    /// Stream rows straight off the HTTP response body without buffering the
243    /// whole result set: bytes are accumulated, split into complete
244    /// `JSONEachRow` lines, and yielded in [`ClickHouseSourceConfig::batch_size`]
245    /// pages. The final page carries the incremental bookmark (when replicating
246    /// incrementally) so the pipeline persists only after everything before it
247    /// is written.
248    fn stream_pages<'a>(
249        &'a self,
250        context: &'a HashMap<String, Value>,
251        _batch_size: usize,
252    ) -> Pin<Box<dyn Stream<Item = Result<StreamPage, FaucetError>> + Send + 'a>> {
253        let batch_size = self.config.batch_size;
254        let chunk = if batch_size == 0 {
255            usize::MAX
256        } else {
257            batch_size
258        };
259        let cap = if batch_size == 0 { 1024 } else { batch_size };
260        let start = self.current_start();
261        let (query, incr) = build_effective_query(&self.config, context, start.as_ref());
262
263        Box::pin(async_stream::try_stream! {
264            let resp = self.request(query).send().await?;
265            let resp = check_http_response(resp, DEFAULT_ERROR_BODY_MAX_LEN).await?;
266            let mut body = resp.bytes_stream();
267
268            let mut buf: Vec<u8> = Vec::new();
269            let mut page: Vec<Value> = Vec::with_capacity(cap);
270            let mut running_max: Option<Value> = None;
271            let mut total = 0usize;
272
273            while let Some(chunk_result) = body.next().await {
274                let bytes = chunk_result.map_err(FaucetError::Http)?;
275                buf.extend_from_slice(&bytes);
276                for line in split_complete_lines(&mut buf) {
277                    if let Some(value) = parse_line(&line)? {
278                        page.push(value);
279                        if page.len() >= chunk {
280                            let ready = std::mem::replace(&mut page, Vec::with_capacity(cap));
281                            let kept = apply_incremental(ready, incr.as_ref(), &mut running_max);
282                            total += kept.len();
283                            if !kept.is_empty() {
284                                yield StreamPage { records: kept, bookmark: None };
285                            }
286                        }
287                    }
288                }
289            }
290            // Trailing line without a terminating newline (ClickHouse always
291            // newline-terminates, but be robust).
292            if let Some(value) = parse_line(&buf)? {
293                page.push(value);
294            }
295
296            // Final page carries the bookmark.
297            let kept = apply_incremental(page, incr.as_ref(), &mut running_max);
298            total += kept.len();
299            let bookmark = if incr.is_some() { running_max.clone() } else { None };
300            if !kept.is_empty() || bookmark.is_some() {
301                yield StreamPage { records: kept, bookmark };
302            }
303
304            tracing::info!(rows = total, batch_size, query = %self.config.query, "ClickHouse source stream complete");
305        })
306    }
307
308    fn config_schema(&self) -> Value {
309        serde_json::to_value(faucet_core::schema_for!(ClickHouseSourceConfig))
310            .expect("schema serialization")
311    }
312
313    fn connector_name(&self) -> &'static str {
314        "clickhouse"
315    }
316
317    fn dataset_uri(&self) -> String {
318        format!(
319            "{}?query={}",
320            faucet_core::redact_uri_credentials(&self.base_url),
321            self.config.query
322        )
323    }
324
325    fn state_key(&self) -> Option<String> {
326        match &self.config.replication {
327            ClickHouseReplication::Full => None,
328            ClickHouseReplication::Incremental { .. } => Some(
329                self.config
330                    .state_key
331                    .clone()
332                    .unwrap_or_else(|| default_state_key(&self.config)),
333            ),
334        }
335    }
336
337    async fn apply_start_bookmark(&self, bookmark: Value) -> Result<(), FaucetError> {
338        *self
339            .start_bookmark
340            .lock()
341            .expect("start_bookmark mutex poisoned") = Some(bookmark);
342        Ok(())
343    }
344
345    /// Non-mutating preflight probe (`connect`): runs `SELECT 1` over the HTTP
346    /// interface.
347    async fn check(&self, ctx: &CheckContext) -> Result<CheckReport, FaucetError> {
348        let started = std::time::Instant::now();
349        let hint = "check url / host / database / credentials / that the server is reachable";
350        let req = self.request("SELECT 1".to_string());
351        let probe = match tokio::time::timeout(ctx.timeout, req.send()).await {
352            Ok(Ok(resp)) => match check_http_response(resp, DEFAULT_ERROR_BODY_MAX_LEN).await {
353                Ok(_) => Probe::pass("connect", started.elapsed()),
354                Err(e) => Probe::fail_hint("connect", started.elapsed(), e.to_string(), hint),
355            },
356            Ok(Err(e)) => Probe::fail_hint("connect", started.elapsed(), e.to_string(), hint),
357            Err(_) => Probe::fail_hint("connect", started.elapsed(), "timed out", hint),
358        };
359        Ok(CheckReport::single(probe))
360    }
361}
362
363#[cfg(test)]
364mod tests {
365    use super::*;
366    use serde_json::json;
367
368    fn full_cfg() -> ClickHouseSourceConfig {
369        ClickHouseSourceConfig::new("http://db.example.com:8123", "SELECT * FROM t")
370    }
371
372    #[test]
373    fn build_full_returns_query_unchanged() {
374        let (q, incr) = build_effective_query(&full_cfg(), &HashMap::new(), None);
375        assert_eq!(q, "SELECT * FROM t");
376        assert!(incr.is_none());
377    }
378
379    #[test]
380    fn build_incremental_substitutes_bookmark_literal() {
381        let cfg = ClickHouseSourceConfig::new(
382            "http://h:8123",
383            "SELECT * FROM t WHERE updated_at > @bookmark",
384        )
385        .incremental("updated_at", json!("1970-01-01"));
386        let (q, incr) = build_effective_query(&cfg, &HashMap::new(), None);
387        assert_eq!(q, "SELECT * FROM t WHERE updated_at > '1970-01-01'");
388        assert_eq!(
389            incr,
390            Some(IncrementalCtx {
391                column: "updated_at".into(),
392                start: json!("1970-01-01"),
393            })
394        );
395    }
396
397    #[test]
398    fn build_incremental_prefers_stored_bookmark_over_initial() {
399        let cfg =
400            ClickHouseSourceConfig::new("http://h:8123", "SELECT * FROM t WHERE c > @bookmark")
401                .incremental("c", json!(0));
402        let stored = json!(500);
403        let (q, incr) = build_effective_query(&cfg, &HashMap::new(), Some(&stored));
404        assert_eq!(q, "SELECT * FROM t WHERE c > 500");
405        assert_eq!(incr.unwrap().start, json!(500));
406    }
407
408    #[test]
409    fn build_incremental_without_token_still_returns_filter_ctx() {
410        let cfg = ClickHouseSourceConfig::new("http://h:8123", "SELECT * FROM t")
411            .incremental("c", json!(0));
412        let (q, incr) = build_effective_query(&cfg, &HashMap::new(), None);
413        assert_eq!(q, "SELECT * FROM t");
414        assert!(incr.is_some(), "client-side filter must still run");
415    }
416
417    #[test]
418    fn context_substitution_uses_injection_safe_literals() {
419        let mut ctx = HashMap::new();
420        ctx.insert("tenant".to_string(), json!("ac'me"));
421        ctx.insert("id".to_string(), json!(7));
422        let cfg = ClickHouseSourceConfig::new(
423            "http://h:8123",
424            "SELECT * FROM t WHERE tenant = {tenant} AND id = {id}",
425        );
426        let (q, _incr) = build_effective_query(&cfg, &ctx, None);
427        assert!(q.contains("tenant = 'ac\\'me'"), "got: {q}");
428        assert!(q.contains("id = 7"), "got: {q}");
429    }
430
431    #[test]
432    fn apply_incremental_filters_and_tracks_max() {
433        let ctx = IncrementalCtx {
434            column: "c".into(),
435            start: json!(10),
436        };
437        let mut running = None;
438        let page = vec![json!({"c": 5}), json!({"c": 15}), json!({"c": 20})];
439        let kept = apply_incremental(page, Some(&ctx), &mut running);
440        assert_eq!(kept.len(), 2);
441        assert_eq!(running, Some(json!(20)));
442    }
443
444    #[test]
445    fn apply_incremental_full_passes_through() {
446        let mut running = None;
447        let page = vec![json!({"c": 1}), json!({"c": 2})];
448        let kept = apply_incremental(page, None, &mut running);
449        assert_eq!(kept.len(), 2);
450        assert_eq!(running, None);
451    }
452
453    #[test]
454    fn split_complete_lines_drains_full_lines_only() {
455        let mut buf = b"{\"a\":1}\n{\"a\":2}\n{\"a\":3".to_vec();
456        let lines = split_complete_lines(&mut buf);
457        assert_eq!(lines.len(), 2, "the unterminated tail stays buffered");
458        assert_eq!(lines[0], b"{\"a\":1}");
459        assert_eq!(buf, b"{\"a\":3", "partial line remains in the buffer");
460    }
461
462    #[test]
463    fn split_complete_lines_reassembles_across_chunks() {
464        // A line split across two network chunks is reassembled once the
465        // second chunk (carrying the newline) arrives.
466        let mut buf = b"{\"a\":".to_vec();
467        assert!(split_complete_lines(&mut buf).is_empty());
468        buf.extend_from_slice(b"1}\n");
469        let lines = split_complete_lines(&mut buf);
470        assert_eq!(lines, vec![b"{\"a\":1}".to_vec()]);
471    }
472
473    #[test]
474    fn parse_line_handles_blank_and_valid_and_invalid() {
475        assert_eq!(parse_line(b"").unwrap(), None);
476        assert_eq!(parse_line(b"   ").unwrap(), None);
477        assert_eq!(parse_line(b"{\"a\":1}").unwrap(), Some(json!({"a": 1})));
478        assert!(parse_line(b"not-json").is_err());
479    }
480
481    #[test]
482    fn parse_line_rejects_invalid_utf8() {
483        assert!(parse_line(&[0xff, 0xfe]).is_err());
484    }
485
486    #[test]
487    fn state_key_only_for_incremental_and_is_stable() {
488        let source = ClickHouseSource::new(full_cfg()).unwrap();
489        assert_eq!(source.state_key(), None, "full replication has no bookmark");
490
491        let cfg = ClickHouseSourceConfig::new("http://db.example.com:8123", "SELECT * FROM t")
492            .incremental("c", json!(0));
493        let source = ClickHouseSource::new(cfg).unwrap();
494        let k1 = source.state_key().unwrap();
495        let k2 = source.state_key().unwrap();
496        assert_eq!(k1, k2);
497        assert!(k1.starts_with("clickhouse:db.example.com:"), "got: {k1}");
498        faucet_core::state::validate_state_key(&k1).expect("derived key must be valid");
499    }
500
501    #[test]
502    fn explicit_state_key_overrides_default() {
503        let mut cfg = ClickHouseSourceConfig::new("http://h:8123", "SELECT * FROM t")
504            .incremental("c", json!(0));
505        cfg.state_key = Some("custom-key".into());
506        let source = ClickHouseSource::new(cfg).unwrap();
507        assert_eq!(source.state_key().as_deref(), Some("custom-key"));
508    }
509
510    #[tokio::test]
511    async fn apply_start_bookmark_overrides_initial() {
512        let cfg =
513            ClickHouseSourceConfig::new("http://h:8123", "SELECT * FROM t WHERE c > @bookmark")
514                .incremental("c", json!(0));
515        let source = ClickHouseSource::new(cfg).unwrap();
516        source.apply_start_bookmark(json!(999)).await.unwrap();
517        let (q, incr) = build_effective_query(
518            &source.config,
519            &HashMap::new(),
520            source.current_start().as_ref(),
521        );
522        assert_eq!(q, "SELECT * FROM t WHERE c > 999");
523        assert_eq!(incr.unwrap().start, json!(999));
524    }
525
526    #[test]
527    fn dataset_uri_has_no_credentials() {
528        let source = ClickHouseSource::new(full_cfg()).unwrap();
529        assert_eq!(
530            source.dataset_uri(),
531            "http://db.example.com:8123?query=SELECT * FROM t"
532        );
533    }
534
535    #[test]
536    fn connector_name_is_clickhouse() {
537        let source = ClickHouseSource::new(full_cfg()).unwrap();
538        assert_eq!(source.connector_name(), "clickhouse");
539    }
540
541    #[test]
542    fn config_schema_is_object() {
543        let source = ClickHouseSource::new(full_cfg()).unwrap();
544        assert_eq!(source.config_schema()["type"], "object");
545    }
546
547    #[test]
548    fn new_rejects_invalid_config() {
549        let cfg = ClickHouseSourceConfig::new("http://h:8123", "SELECT 1")
550            .with_batch_size(faucet_core::MAX_BATCH_SIZE + 1);
551        assert!(ClickHouseSource::new(cfg).is_err());
552    }
553
554    #[tokio::test]
555    async fn check_fails_against_unreachable_server() {
556        // Connection-refused on a closed local port exercises the check() I/O
557        // path offline (no external dependency).
558        let source = ClickHouseSource::new(ClickHouseSourceConfig::new(
559            "http://127.0.0.1:1",
560            "SELECT 1",
561        ))
562        .unwrap();
563        let ctx = CheckContext {
564            timeout: std::time::Duration::from_secs(2),
565        };
566        let report = source.check(&ctx).await.unwrap();
567        assert!(
568            matches!(
569                report.probes[0].status,
570                faucet_core::check::ProbeStatus::Fail { .. }
571            ),
572            "unreachable server must fail the connect probe"
573        );
574    }
575}
faucet_source_clickhouse/stream.rs

faucet_source_clickhouse/
stream.rs
