micromegas/servers/
query_audit.rs

1//! Structured per-query audit record for the FlightSQL service.
2//!
3//! `execute_query` (see `flight_sql_service_impl`) emits one JSON-serialized
4//! [`QueryAuditRecord`] per query, at completion, under the dedicated
5//! `flightsql_query_audit` log target. Unlike the untagged `imetric!` cost
6//! metrics (whose `PropertySet` can't carry high-cardinality values such as
7//! SQL text), a free-text log `msg` has no cardinality constraint, so it can
8//! carry both attribution and cost in one self-contained, queryable record.
9
10use datafusion::physical_plan::ExecutionPlan;
11
12/// Aggregated DataFusion plan metrics for one query, read after the stream drains.
13pub struct ScanMetrics {
14    pub output_rows: Option<u64>,
15    pub bytes_scanned: u64,
16}
17
18/// Walk the physical-plan tree: `output_rows` from the root node (final result
19/// grain), `bytes_scanned` summed across every node (leaf `DataSourceExec`
20/// nodes carry it).
21pub fn aggregate_scan_metrics(plan: &dyn ExecutionPlan) -> ScanMetrics {
22    fn sum_bytes(plan: &dyn ExecutionPlan) -> u64 {
23        let mut total = plan
24            .metrics()
25            .and_then(|m| m.sum_by_name("bytes_scanned"))
26            .map(|v| v.as_usize() as u64)
27            .unwrap_or(0);
28        for child in plan.children() {
29            total += sum_bytes(child.as_ref());
30        }
31        total
32    }
33    ScanMetrics {
34        output_rows: plan
35            .metrics()
36            .and_then(|m| m.output_rows())
37            .map(|r| r as u64),
38        bytes_scanned: sum_bytes(plan),
39    }
40}
41
42/// One structured record per FlightSQL query, emitted as a JSON log line under
43/// the `flightsql_query_audit` target when the query completes, fails, or is
44/// abandoned mid-drain (client disconnect/cancel).
45#[derive(serde::Serialize)]
46pub struct QueryAuditRecord {
47    pub client: String,
48    pub user: String,
49    pub email: String,
50    #[serde(skip_serializing_if = "Option::is_none")]
51    pub name: Option<String>,
52    pub service_account: bool,
53    #[serde(skip_serializing_if = "Option::is_none")]
54    pub service_account_name: Option<String>,
55    pub sql: String,
56    #[serde(skip_serializing_if = "Option::is_none")]
57    pub range_begin: Option<String>, // RFC3339
58    #[serde(skip_serializing_if = "Option::is_none")]
59    pub range_end: Option<String>,
60    #[serde(skip_serializing_if = "Option::is_none")]
61    pub limit: Option<u64>,
62    pub context_init_ms: f64,
63    pub planning_ms: f64,
64    pub execution_ms: f64, // stream construction (matches query_execution_duration semantics)
65    pub setup_ms: f64,     // parse+attribution+context+planning+stream-build (query_setup_duration)
66    pub total_ms: f64,     // end-to-end incl. drain
67    pub status: &'static str, // "ok" | "error" | "incomplete"
68    #[serde(skip_serializing_if = "Option::is_none")]
69    pub error: Option<String>,
70    #[serde(skip_serializing_if = "Option::is_none")]
71    pub output_rows: Option<u64>,
72    pub bytes_scanned: u64,
73}
micromegas/servers/query_audit.rs

micromegas/servers/
query_audit.rs
