1use serde_json::Value as JsonValue;
15
16use crate::executor::ExecutionResult;
17
18const EXPORT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
20
21#[derive(Clone, Copy, Debug, Eq, PartialEq, clap::ValueEnum)]
22pub enum ExportFormat {
23 Influx,
24 Loki,
25 Prometheus,
26}
27
28#[derive(Debug, thiserror::Error)]
29pub enum ExportError {
30 #[error("HTTP request failed: {0}")]
31 Http(#[from] reqwest::Error),
32 #[error("serialization error: {0}")]
33 Serialize(#[from] serde_json::Error),
34 #[error("export target returned {status}: {body}")]
35 BadStatus { status: u16, body: String },
36}
37
38pub async fn push_to_influxdb(url: &str, result: &ExecutionResult) -> Result<(), ExportError> {
44 let body = format_as_influx(result, "containers");
45 let client = reqwest::Client::builder().timeout(EXPORT_TIMEOUT).build()?;
46 let resp = client
47 .post(url)
48 .header("Content-Type", "application/octet-stream")
49 .body(body)
50 .send()
51 .await?;
52 let status = resp.status();
53 if !status.is_success() {
54 let body = resp.text().await.unwrap_or_default();
55 return Err(ExportError::BadStatus {
56 status: status.as_u16(),
57 body,
58 });
59 }
60 Ok(())
61}
62
63pub async fn push_to_loki(url: &str, result: &ExecutionResult) -> Result<(), ExportError> {
68 let body = format_as_loki(result)?;
69 let push_url = format!("{}/loki/api/v1/push", url.trim_end_matches('/'));
70 let client = reqwest::Client::builder().timeout(EXPORT_TIMEOUT).build()?;
71 let resp = client
72 .post(&push_url)
73 .header("Content-Type", "application/json")
74 .body(body)
75 .send()
76 .await?;
77 let status = resp.status();
78 if !status.is_success() {
79 let body = resp.text().await.unwrap_or_default();
80 return Err(ExportError::BadStatus {
81 status: status.as_u16(),
82 body,
83 });
84 }
85 Ok(())
86}
87
88pub async fn push_to_prometheus(url: &str, result: &ExecutionResult) -> Result<(), ExportError> {
93 let body = format_as_prometheus(result);
94 let push_url = format!("{}/metrics/job/dol", url.trim_end_matches('/'));
95 let client = reqwest::Client::builder().timeout(EXPORT_TIMEOUT).build()?;
96 let resp = client
97 .put(&push_url)
98 .header("Content-Type", "text/plain; version=0.0.4")
99 .body(body)
100 .send()
101 .await?;
102 let status = resp.status();
103 if !status.is_success() {
104 let body = resp.text().await.unwrap_or_default();
105 return Err(ExportError::BadStatus {
106 status: status.as_u16(),
107 body,
108 });
109 }
110 Ok(())
111}
112
113#[must_use]
118pub fn format_as_influx(result: &ExecutionResult, measurement: &str) -> String {
119 let mut lines = Vec::new();
120 for row in &result.rows {
121 let name = row
122 .fields
123 .get("name")
124 .and_then(|v| v.as_str())
125 .unwrap_or("unknown");
126 let image = row
127 .fields
128 .get("image")
129 .and_then(|v| v.as_str())
130 .unwrap_or("unknown");
131 let state = row
132 .fields
133 .get("state")
134 .and_then(|v| v.as_str())
135 .unwrap_or("unknown");
136
137 let mut fields = Vec::new();
138 let mut tags = vec![
139 format!("name={}", escape_influx_tag(name)),
140 format!("image={}", escape_influx_tag(image)),
141 format!("state={}", escape_influx_tag(state)),
142 ];
143
144 for (key, val) in &row.fields {
146 match key.as_str() {
147 "name" | "image" | "state" | "id" | "status" | "ports" => {
148 if let Some(s) = val.as_str() {
149 tags.push(format!("{key}={}", escape_influx_tag(s)));
150 }
151 }
152 "cpu" | "memory" | "memory_limit" | "restart_count" | "network_rx"
153 | "network_tx" | "disk_read" | "disk_write" | "size" | "count" => {
154 if let Some(n) = val.as_f64() {
155 fields.push(format!("{key}={n}"));
156 }
157 }
158 _ => {}
159 }
160 }
161
162 if fields.is_empty() {
163 fields.push("_value=1".to_owned());
164 }
165
166 let line = format!(
167 "{measurement},{tags} {fields}",
168 tags = tags.join(","),
169 fields = fields.join(",")
170 );
171 lines.push(line);
172 }
173 lines.join("\n") + "\n"
174}
175
176pub fn format_as_loki(result: &ExecutionResult) -> Result<String, serde_json::Error> {
178 let mut entries: Vec<Vec<JsonValue>> = Vec::new();
179 for row in &result.rows {
180 let ts = std::time::SystemTime::now()
181 .duration_since(std::time::UNIX_EPOCH)
182 .unwrap_or_default()
183 .as_nanos()
184 .to_string();
185 let log_line = serde_json::to_string(&row.fields)?;
186 entries.push(vec![JsonValue::String(ts), JsonValue::String(log_line)]);
187 }
188
189 let payload = serde_json::json!({
190 "streams": [{
191 "stream": {
192 "app": "dol",
193 "source": "docker"
194 },
195 "values": entries
196 }]
197 });
198 serde_json::to_string(&payload)
199}
200
201#[must_use]
206pub fn format_as_prometheus(result: &ExecutionResult) -> String {
207 let mut lines = Vec::new();
208 lines.push("# HELP dol_query Docker query results from DOL".to_owned());
209 lines.push("# TYPE dol_query gauge".to_owned());
210
211 for row in &result.rows {
212 let name = row
213 .fields
214 .get("name")
215 .and_then(|v| v.as_str())
216 .unwrap_or("unknown");
217 let image = row
218 .fields
219 .get("image")
220 .and_then(|v| v.as_str())
221 .unwrap_or("unknown");
222 let state = row
223 .fields
224 .get("state")
225 .and_then(|v| v.as_str())
226 .unwrap_or("unknown");
227
228 let labels = format!(
229 "container=\"{}\",image=\"{}\",state=\"{}\"",
230 escape_prom_label(name),
231 escape_prom_label(image),
232 escape_prom_label(state),
233 );
234
235 for (key, val) in &row.fields {
236 if let Some(n) = val.as_f64() {
237 lines.push(format!("dol_{key}{{{labels}}} {n}"));
238 } else if let Some(i) = val.as_i64() {
239 lines.push(format!("dol_{key}{{{labels}}} {i}"));
240 }
241 }
242 }
243
244 lines.join("\n") + "\n"
245}
246
247fn escape_influx_tag(s: &str) -> String {
248 s.replace('\\', "\\\\")
249 .replace(',', "\\,")
250 .replace('=', "\\=")
251 .replace(' ', "\\ ")
252}
253
254fn escape_prom_label(s: &str) -> String {
255 s.replace('\\', "\\\\")
256 .replace('"', "\\\"")
257 .replace('\n', "\\n")
258}
259
260#[cfg(test)]
261mod tests {
262 use super::*;
263 use crate::executor::Row;
264
265 fn sample_result() -> ExecutionResult {
266 ExecutionResult {
267 rows: vec![
268 Row {
269 fields: std::collections::BTreeMap::from([
270 ("name".to_owned(), JsonValue::String("web".to_owned())),
271 (
272 "image".to_owned(),
273 JsonValue::String("nginx:latest".to_owned()),
274 ),
275 ("state".to_owned(), JsonValue::String("running".to_owned())),
276 (
277 "cpu".to_owned(),
278 JsonValue::Number(serde_json::Number::from_f64(12.5).unwrap()),
279 ),
280 (
281 "memory".to_owned(),
282 JsonValue::Number(serde_json::Number::from_f64(64_000_000.0).unwrap()),
283 ),
284 (
285 "restart_count".to_owned(),
286 JsonValue::Number(serde_json::Number::from(0)),
287 ),
288 ]),
289 },
290 Row {
291 fields: std::collections::BTreeMap::from([
292 ("name".to_owned(), JsonValue::String("db".to_owned())),
293 (
294 "image".to_owned(),
295 JsonValue::String("postgres:16".to_owned()),
296 ),
297 ("state".to_owned(), JsonValue::String("exited".to_owned())),
298 (
299 "cpu".to_owned(),
300 JsonValue::Number(serde_json::Number::from_f64(0.0).unwrap()),
301 ),
302 (
303 "memory".to_owned(),
304 JsonValue::Number(serde_json::Number::from_f64(256_000_000.0).unwrap()),
305 ),
306 (
307 "restart_count".to_owned(),
308 JsonValue::Number(serde_json::Number::from(3)),
309 ),
310 ]),
311 },
312 ],
313 }
314 }
315
316 #[test]
317 fn formats_influx_line_protocol() {
318 let result = sample_result();
319 let output = format_as_influx(&result, "containers");
320 assert!(output.contains("containers,"));
321 assert!(output.contains("name=web"));
322 assert!(output.contains("name=db"));
323 assert!(output.contains("cpu=12.5"));
324 assert!(output.contains("restart_count=3"));
325 assert!(output.ends_with('\n'));
326 }
327
328 #[test]
329 fn formats_loki_json() {
330 let result = sample_result();
331 let output = format_as_loki(&result).unwrap();
332 assert!(output.contains(r#""app":"dol""#));
333 assert!(output.contains(r#""streams""#));
334 assert!(output.contains(r#""values""#));
335 }
336
337 #[test]
338 fn formats_prometheus_exposition() {
339 let result = sample_result();
340 let output = format_as_prometheus(&result);
341 assert!(output.contains("dol_cpu{"));
342 assert!(output.contains("dol_memory{"));
343 assert!(output.contains("container=\"web\""));
344 assert!(output.contains("container=\"db\""));
345 assert!(output.contains("# HELP dol_query"));
346 }
347}