Skip to main content

dol/
export.rs

1//! Export query results to external monitoring systems.
2//!
3//! Supports three formats:
4//! - **InfluxDB** — line protocol via HTTP POST
5//! - **Grafana Loki** — JSON push API
6//! - **Prometheus** — exposition format via Pushgateway
7//!
8//! # Example
9//!
10//! ```ignore
11//! export::push_to_influxdb("http://localhost:8086/write?db=mydb", &result).await?;
12//! ```
13
14use serde_json::Value as JsonValue;
15
16use crate::executor::ExecutionResult;
17
18/// HTTP request timeout for export push operations.
19const 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
38/// Push a query result to `InfluxDB` v1/v2 HTTP write API.
39///
40/// `url` should be the full write endpoint, e.g.:
41///   `http://localhost:8086/write?db=mydb` (`InfluxDB` v1)
42///   `http://localhost:8086/api/v2/write?org=myorg&bucket=mybucket` (v2)
43pub 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
63/// Push a query result to Grafana Loki HTTP push API.
64///
65/// `url` should be the base Loki URL, e.g. `http://localhost:3100`.
66/// The push endpoint `/loki/api/v1/push` is appended automatically.
67pub 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
88/// Push a query result to Prometheus Pushgateway.
89///
90/// `url` should be the Pushgateway URL, e.g. `http://localhost:9091`.
91/// The job name is fixed as `dol`.
92pub 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/// Format execution result as `InfluxDB` line protocol.
114///
115/// Each row becomes a line in the format:
116///   `<measurement>,name=<name>,image=<image>,state=<state> <field_kv>`
117#[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        // Also add other string fields as tags
145        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
176/// Format execution result as Loki JSON push payload.
177pub 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/// Format execution result as Prometheus exposition format.
202///
203/// Numeric fields become gauge metrics:
204///   `dol_<field>{container="<name>",image="<image>",state="<state>"} <value>`
205#[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}