featherbit/plugins/native/
clickhouse_logger.rs1use std::collections::HashMap;
22use std::sync::atomic::{AtomicUsize, Ordering};
23use std::sync::Arc;
24use std::time::Duration;
25
26use async_trait::async_trait;
27use bytes::Bytes;
28use serde_json::Value;
29
30use crate::batch::{BatchConfig, BatchFlusher, BatchSink, FlushError};
31use crate::context::Context;
32use crate::outbound::{OutboundClient, OutboundRequest};
33use crate::plugins::resources::PluginResources;
34use crate::plugins::util::log_entry::{build_entry, parse_log_format};
35use crate::plugins::{Plugin, PluginOutput, PluginResult};
36
37pub struct ClickhouseLoggerPlugin {
39 sink: BatchSink,
40 log_format: Option<HashMap<String, Value>>,
41 include_req_body: bool,
42 include_resp_body: bool,
43}
44
45struct ClickhouseFlusher {
47 client: Arc<OutboundClient>,
48 endpoints: Vec<String>,
50 cursor: AtomicUsize,
51 database: String,
52 logtable: String,
53 user: String,
54 password: String,
55 ssl_verify: bool,
56 timeout: Duration,
57}
58
59impl ClickhouseLoggerPlugin {
60 pub fn from_config(
89 config: &HashMap<String, Value>,
90 resources: &Arc<PluginResources>,
91 ) -> Result<Self, String> {
92 let endpoints = collect_endpoints(config)?;
93
94 let logtable = required_string(config, "logtable")?;
95 let database = required_string(config, "database")?;
96 let user = config
97 .get("user")
98 .and_then(|v| v.as_str())
99 .unwrap_or("")
100 .to_string();
101 let password = config
102 .get("password")
103 .and_then(|v| v.as_str())
104 .unwrap_or("")
105 .to_string();
106
107 let ssl_verify = config
108 .get("ssl_verify")
109 .and_then(|v| v.as_bool())
110 .unwrap_or(true);
111 let timeout =
112 Duration::from_secs(config.get("timeout").and_then(|v| v.as_u64()).unwrap_or(3));
113
114 let include_req_body = config
115 .get("include_req_body")
116 .and_then(|v| v.as_bool())
117 .unwrap_or(false);
118 let include_resp_body = config
119 .get("include_resp_body")
120 .and_then(|v| v.as_bool())
121 .unwrap_or(false);
122 let log_format = parse_log_format(config)?;
123
124 let batch_cfg = BatchConfig::from_config(config)?;
125 let flusher = Arc::new(ClickhouseFlusher {
126 client: resources.outbound.clone(),
127 endpoints,
128 cursor: AtomicUsize::new(0),
129 database,
130 logtable,
131 user,
132 password,
133 ssl_verify,
134 timeout,
135 });
136 let sink = BatchSink::spawn("clickhouse-logger", batch_cfg, flusher);
137
138 Ok(Self {
139 sink,
140 log_format,
141 include_req_body,
142 include_resp_body,
143 })
144 }
145}
146
147fn required_string(config: &HashMap<String, Value>, key: &str) -> Result<String, String> {
148 config
149 .get(key)
150 .and_then(|v| v.as_str())
151 .filter(|s| !s.is_empty())
152 .map(|s| s.to_string())
153 .ok_or_else(|| format!("clickhouse-logger requires '{key}'"))
154}
155
156fn collect_endpoints(config: &HashMap<String, Value>) -> Result<Vec<String>, String> {
157 let mut endpoints = Vec::new();
158 if let Some(s) = config.get("endpoint_addr").and_then(|v| v.as_str()) {
159 endpoints.push(s.to_string());
160 }
161 if let Some(arr) = config.get("endpoint_addrs").and_then(|v| v.as_array()) {
162 for v in arr {
163 if let Some(s) = v.as_str() {
164 endpoints.push(s.to_string());
165 }
166 }
167 }
168 if endpoints.is_empty() {
169 return Err("clickhouse-logger requires 'endpoint_addr' or 'endpoint_addrs'".to_string());
170 }
171 Ok(endpoints)
172}
173
174fn build_insert_body(logtable: &str, entries: &[Value]) -> String {
178 let rows: Vec<String> = entries
179 .iter()
180 .map(|e| serde_json::to_string(e).unwrap_or_else(|_| "{}".to_string()))
181 .collect();
182 format!(
183 "INSERT INTO {logtable} FORMAT JSONEachRow {}",
184 rows.join("\n")
185 )
186}
187
188#[async_trait]
189impl BatchFlusher for ClickhouseFlusher {
190 async fn flush(&self, entries: &[Value]) -> Result<(), FlushError> {
191 let idx = self.cursor.fetch_add(1, Ordering::Relaxed) % self.endpoints.len();
192 let url = self.endpoints[idx].clone();
193 let body = build_insert_body(&self.logtable, entries);
194
195 let headers = vec![
196 ("Content-Type".to_string(), "application/json".to_string()),
197 ("X-ClickHouse-User".to_string(), self.user.clone()),
198 ("X-ClickHouse-Key".to_string(), self.password.clone()),
199 ("X-ClickHouse-Database".to_string(), self.database.clone()),
200 ];
201
202 let req = OutboundRequest {
203 method: http::Method::POST,
204 url,
205 headers,
206 body: Bytes::from(body),
207 timeout: self.timeout,
208 ssl_verify: self.ssl_verify,
209 tls: None,
210 };
211
212 match self.client.request(req).await {
213 Ok(resp) if resp.status < 400 => Ok(()),
215 Ok(resp) => Err(FlushError {
216 message: format!(
217 "clickhouse returned status {}: {}",
218 resp.status,
219 String::from_utf8_lossy(&resp.body)
220 ),
221 first_fail: None,
222 }),
223 Err(e) => Err(FlushError {
224 message: format!("clickhouse callout failed: {e}"),
225 first_fail: None,
226 }),
227 }
228 }
229}
230
231#[async_trait]
232impl Plugin for ClickhouseLoggerPlugin {
233 fn plugin_type(&self) -> &str {
234 "clickhouse-logger"
235 }
236
237 async fn execute(&self, ctx: Context, _named_inputs: &HashMap<String, Value>) -> PluginResult {
238 let entry = build_entry(
239 &ctx,
240 self.log_format.as_ref(),
241 self.include_req_body,
242 self.include_resp_body,
243 );
244 self.sink.push(entry);
245 Ok(PluginOutput {
246 context: ctx,
247 named_outputs: HashMap::new(),
248 })
249 }
250}
251
252#[cfg(test)]
253mod tests {
254 use super::*;
255 use serde_json::json;
256
257 fn cfg(v: Value) -> HashMap<String, Value> {
258 serde_json::from_value(v).unwrap()
259 }
260
261 fn full_cfg() -> Value {
262 json!({
263 "endpoint_addr": "http://clickhouse:8123",
264 "database": "default",
265 "logtable": "logs"
266 })
267 }
268
269 #[tokio::test]
270 async fn requires_endpoint_database_and_table() {
271 assert!(ClickhouseLoggerPlugin::from_config(
272 &cfg(json!({ "database": "d", "logtable": "t" })),
273 &PluginResources::empty()
274 )
275 .is_err());
276 assert!(ClickhouseLoggerPlugin::from_config(
277 &cfg(json!({ "endpoint_addr": "http://c:8123", "logtable": "t" })),
278 &PluginResources::empty()
279 )
280 .is_err());
281 assert!(ClickhouseLoggerPlugin::from_config(
282 &cfg(json!({ "endpoint_addr": "http://c:8123", "database": "d" })),
283 &PluginResources::empty()
284 )
285 .is_err());
286 assert!(
287 ClickhouseLoggerPlugin::from_config(&cfg(full_cfg()), &PluginResources::empty())
288 .is_ok()
289 );
290 }
291
292 #[test]
293 fn insert_body_shape() {
294 let entries = vec![json!({ "a": 1 }), json!({ "b": 2 })];
295 let body = build_insert_body("logs", &entries);
296 assert_eq!(
297 body,
298 "INSERT INTO logs FORMAT JSONEachRow {\"a\":1}\n{\"b\":2}"
299 );
300 }
301
302 #[test]
303 fn insert_body_single_entry() {
304 let body = build_insert_body("t", &[json!({ "x": true })]);
305 assert_eq!(body, "INSERT INTO t FORMAT JSONEachRow {\"x\":true}");
306 }
307}