faucet_common_clickhouse/
lib.rs1#![cfg_attr(docsrs, feature(doc_cfg))]
2
3use faucet_core::FaucetError;
33use schemars::JsonSchema;
34use serde::{Deserialize, Serialize};
35use serde_json::Value;
36
37pub const DEFAULT_HTTP_PORT: u16 = 8123;
39pub const DEFAULT_DATABASE: &str = "default";
41
42fn default_database() -> String {
43 DEFAULT_DATABASE.to_string()
44}
45
46#[derive(Clone, Serialize, Deserialize, JsonSchema)]
52pub struct ClickHouseConnection {
53 #[serde(default, skip_serializing_if = "Option::is_none")]
57 pub url: Option<String>,
58 #[serde(default, skip_serializing_if = "Option::is_none")]
62 pub host: Option<String>,
63 #[serde(default, skip_serializing_if = "Option::is_none")]
66 pub http_port: Option<u16>,
67 #[serde(default)]
71 pub tls: bool,
72 #[serde(default = "default_database")]
74 pub database: String,
75 #[serde(default, skip_serializing_if = "Option::is_none")]
77 pub user: Option<String>,
78 #[serde(default, skip_serializing_if = "Option::is_none")]
80 pub password: Option<String>,
81}
82
83impl Default for ClickHouseConnection {
84 fn default() -> Self {
85 Self {
86 url: None,
87 host: None,
88 http_port: None,
89 tls: false,
90 database: default_database(),
91 user: None,
92 password: None,
93 }
94 }
95}
96
97impl std::fmt::Debug for ClickHouseConnection {
98 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
99 f.debug_struct("ClickHouseConnection")
100 .field("url", &self.url)
101 .field("host", &self.host)
102 .field("http_port", &self.http_port)
103 .field("tls", &self.tls)
104 .field("database", &self.database)
105 .field("user", &self.user)
106 .field("password", &self.password.as_ref().map(|_| "***"))
107 .finish()
108 }
109}
110
111impl ClickHouseConnection {
112 pub fn from_url(url: impl Into<String>) -> Self {
114 Self {
115 url: Some(url.into()),
116 ..Default::default()
117 }
118 }
119
120 pub fn validate(&self) -> Result<(), FaucetError> {
122 match (&self.url, &self.host) {
123 (Some(_), Some(_)) => Err(FaucetError::Config(
124 "ClickHouse config sets both `url` and `host`; set exactly one".into(),
125 )),
126 (None, None) => Err(FaucetError::Config(
127 "ClickHouse config requires either `url` or `host`".into(),
128 )),
129 _ => Ok(()),
130 }
131 }
132
133 pub fn base_url(&self) -> Result<String, FaucetError> {
138 if let Some(url) = &self.url {
139 return Ok(url.trim_end_matches('/').to_string());
140 }
141 if let Some(host) = &self.host {
142 let scheme = if self.tls { "https" } else { "http" };
143 let port = self.http_port.unwrap_or(DEFAULT_HTTP_PORT);
144 return Ok(format!("{scheme}://{host}:{port}"));
145 }
146 Err(FaucetError::Config(
147 "ClickHouse config requires either `url` or `host`".into(),
148 ))
149 }
150}
151
152pub fn build_client(_conn: &ClickHouseConnection) -> Result<reqwest::Client, FaucetError> {
156 reqwest::Client::builder()
157 .build()
158 .map_err(FaucetError::Http)
159}
160
161pub fn query_params(database: &str, settings: &[(&str, &str)]) -> Vec<(String, String)> {
167 let mut params = Vec::with_capacity(1 + settings.len());
168 params.push(("database".to_string(), database.to_string()));
169 for (k, v) in settings {
170 params.push((k.to_string(), v.to_string()));
171 }
172 params
173}
174
175pub fn apply_auth(
180 mut req: reqwest::RequestBuilder,
181 conn: &ClickHouseConnection,
182) -> reqwest::RequestBuilder {
183 if let Some(user) = &conn.user {
184 req = req.header("X-ClickHouse-User", user);
185 }
186 if let Some(password) = &conn.password {
187 req = req.header("X-ClickHouse-Key", password);
188 }
189 req
190}
191
192pub fn parse_json_each_row(body: &str) -> Result<Vec<Value>, FaucetError> {
197 let mut out = Vec::new();
198 for (idx, line) in body.lines().enumerate() {
199 let trimmed = line.trim();
200 if trimmed.is_empty() {
201 continue;
202 }
203 let value: Value = serde_json::from_str(trimmed).map_err(|e| {
204 FaucetError::Source(format!(
205 "ClickHouse: failed to parse JSONEachRow line {}: {e}",
206 idx + 1
207 ))
208 })?;
209 out.push(value);
210 }
211 Ok(out)
212}
213
214pub fn build_json_each_row(records: &[Value]) -> Result<String, FaucetError> {
220 let mut body = String::new();
221 for record in records {
222 let line = serde_json::to_string(record).map_err(|e| {
223 FaucetError::Sink(format!("ClickHouse: failed to serialize record: {e}"))
224 })?;
225 body.push_str(&line);
226 body.push('\n');
227 }
228 Ok(body)
229}
230
231pub fn sql_literal(value: &Value) -> String {
240 match value {
241 Value::Null => "NULL".to_string(),
242 Value::Bool(b) => {
243 if *b {
244 "1".to_string()
245 } else {
246 "0".to_string()
247 }
248 }
249 Value::Number(n) => n.to_string(),
250 Value::String(s) => quote_string(s),
251 other => quote_string(&other.to_string()),
252 }
253}
254
255fn quote_string(s: &str) -> String {
256 let escaped = s.replace('\\', "\\\\").replace('\'', "\\'");
257 format!("'{escaped}'")
258}
259
260#[cfg(test)]
261mod tests {
262 use super::*;
263 use serde_json::json;
264
265 #[test]
266 fn base_url_from_url_trims_trailing_slash() {
267 let conn = ClickHouseConnection::from_url("http://localhost:8123/");
268 assert_eq!(conn.base_url().unwrap(), "http://localhost:8123");
269 }
270
271 #[test]
272 fn base_url_from_host_defaults_port_and_scheme() {
273 let conn = ClickHouseConnection {
274 host: Some("db.example.com".into()),
275 ..Default::default()
276 };
277 assert_eq!(conn.base_url().unwrap(), "http://db.example.com:8123");
278 }
279
280 #[test]
281 fn base_url_from_host_honors_tls_and_port() {
282 let conn = ClickHouseConnection {
283 host: Some("db.example.com".into()),
284 http_port: Some(8443),
285 tls: true,
286 ..Default::default()
287 };
288 assert_eq!(conn.base_url().unwrap(), "https://db.example.com:8443");
289 }
290
291 #[test]
292 fn base_url_requires_url_or_host() {
293 let conn = ClickHouseConnection::default();
294 assert!(conn.base_url().is_err());
295 }
296
297 #[test]
298 fn validate_rejects_both_and_neither() {
299 let both = ClickHouseConnection {
300 url: Some("http://h:8123".into()),
301 host: Some("h".into()),
302 ..Default::default()
303 };
304 assert!(both.validate().is_err());
305 assert!(ClickHouseConnection::default().validate().is_err());
306 }
307
308 #[test]
309 fn validate_accepts_exactly_one() {
310 assert!(
311 ClickHouseConnection::from_url("http://h:8123")
312 .validate()
313 .is_ok()
314 );
315 let host_only = ClickHouseConnection {
316 host: Some("h".into()),
317 ..Default::default()
318 };
319 assert!(host_only.validate().is_ok());
320 }
321
322 #[test]
323 fn debug_masks_password() {
324 let conn = ClickHouseConnection {
325 url: Some("http://h:8123".into()),
326 user: Some("alice".into()),
327 password: Some("s3cret".into()),
328 ..Default::default()
329 };
330 let dbg = format!("{conn:?}");
331 assert!(dbg.contains("alice"));
332 assert!(dbg.contains("***"));
333 assert!(!dbg.contains("s3cret"));
334 }
335
336 #[test]
337 fn database_defaults_when_missing() {
338 let conn: ClickHouseConnection =
339 serde_json::from_value(json!({ "url": "http://h:8123" })).unwrap();
340 assert_eq!(conn.database, "default");
341 }
342
343 #[test]
344 fn query_params_puts_database_first_then_settings() {
345 let params = query_params("analytics", &[("default_format", "JSONEachRow")]);
346 assert_eq!(
347 params,
348 vec![
349 ("database".to_string(), "analytics".to_string()),
350 ("default_format".to_string(), "JSONEachRow".to_string()),
351 ]
352 );
353 }
354
355 #[test]
356 fn query_params_async_insert_on_and_off() {
357 let on = query_params(
358 "db",
359 &[("async_insert", "1"), ("wait_for_async_insert", "1")],
360 );
361 assert!(on.contains(&("async_insert".to_string(), "1".to_string())));
362 let off = query_params("db", &[]);
363 assert_eq!(off.len(), 1, "only the database param when no settings");
364 }
365
366 #[test]
367 fn parse_json_each_row_multiple_rows() {
368 let body = "{\"a\":1}\n{\"a\":2}\n{\"a\":3}\n";
369 let rows = parse_json_each_row(body).unwrap();
370 assert_eq!(rows.len(), 3);
371 assert_eq!(rows[2]["a"], 3);
372 }
373
374 #[test]
375 fn parse_json_each_row_empty_result_is_empty() {
376 assert!(parse_json_each_row("").unwrap().is_empty());
377 assert!(parse_json_each_row("\n\n").unwrap().is_empty());
378 }
379
380 #[test]
381 fn parse_json_each_row_skips_blank_lines_between_rows() {
382 let rows = parse_json_each_row("{\"a\":1}\n\n{\"a\":2}\n").unwrap();
383 assert_eq!(rows.len(), 2);
384 }
385
386 #[test]
387 fn parse_json_each_row_malformed_line_is_typed_error() {
388 let err = parse_json_each_row("{\"a\":1}\nnot-json\n").unwrap_err();
389 match err {
390 FaucetError::Source(m) => assert!(m.contains("line 2"), "got: {m}"),
391 other => panic!("expected Source error, got {other:?}"),
392 }
393 }
394
395 #[test]
396 fn build_json_each_row_exact_ndjson() {
397 let page = vec![json!({"id": 1, "v": "a"}), json!({"id": 2, "v": "b"})];
398 let body = build_json_each_row(&page).unwrap();
399 assert_eq!(body, "{\"id\":1,\"v\":\"a\"}\n{\"id\":2,\"v\":\"b\"}\n");
400 }
401
402 #[test]
403 fn build_json_each_row_empty_is_empty_string() {
404 assert_eq!(build_json_each_row(&[]).unwrap(), "");
405 }
406
407 #[test]
408 fn build_and_parse_round_trip() {
409 let page = vec![json!({"id": 1, "s": "héllo"}), json!({"id": 2, "s": "x"})];
410 let body = build_json_each_row(&page).unwrap();
411 let back = parse_json_each_row(&body).unwrap();
412 assert_eq!(back, page);
413 }
414
415 #[test]
416 fn sql_literal_scalars() {
417 assert_eq!(sql_literal(&Value::Null), "NULL");
418 assert_eq!(sql_literal(&json!(true)), "1");
419 assert_eq!(sql_literal(&json!(false)), "0");
420 assert_eq!(sql_literal(&json!(42)), "42");
421 assert_eq!(sql_literal(&json!(-1.5)), "-1.5");
422 assert_eq!(sql_literal(&json!("2024-01-01")), "'2024-01-01'");
423 }
424
425 #[test]
426 fn sql_literal_escapes_quote_and_backslash() {
427 assert_eq!(sql_literal(&json!("O'Brien")), "'O\\'Brien'");
428 assert_eq!(sql_literal(&json!("a\\b")), "'a\\\\b'");
429 assert_eq!(
431 sql_literal(&json!("x' OR '1'='1")),
432 "'x\\' OR \\'1\\'=\\'1'"
433 );
434 }
435
436 #[test]
437 fn build_client_succeeds() {
438 assert!(build_client(&ClickHouseConnection::from_url("http://h:8123")).is_ok());
439 }
440}