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//! [`crate::servers::query_audit::QueryAuditRecord`] per query, at completion, under the dedicated
5//! `flightsql_query_audit` log target. Unlike the untagged `imetric!` cost
6//! metrics (whose properties are meant for bounded-cardinality dimensions, not
7//! values such as 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//!
10//! `peak_memory_bytes`/`spilled_bytes`/`spill_count` complement the cost
11//! fields above: the peak comes from the query's `ScopedMemoryPool` (see
12//! `micromegas_analytics::lakehouse::scoped_memory_pool`), while the spill
13//! counters are summed from the physical plan tree, same as `bytes_scanned`.
14
15use datafusion::physical_plan::ExecutionPlan;
16
17/// Aggregated DataFusion plan metrics for one query, read after the stream drains.
18pub struct ScanMetrics {
19    pub output_rows: Option<u64>,
20    pub bytes_scanned: u64,
21    pub spilled_bytes: u64,
22    pub spill_count: u64,
23}
24
25/// Walk the physical-plan tree: `output_rows` from the root node (final result
26/// grain), `bytes_scanned`/`spilled_bytes`/`spill_count` summed across every
27/// node (leaf `DataSourceExec` nodes carry `bytes_scanned`; spilling operators
28/// such as `ExternalSorter` carry the other two). `sum_by_name` cannot be used
29/// for the spill counters: `MetricsSet::sum_by_name` explicitly returns
30/// `false` for `MetricValue::SpillCount`/`SpilledBytes`, so it would silently
31/// report zero; the dedicated `MetricsSet::spill_count()`/`spilled_bytes()`
32/// accessors are used instead.
33pub fn aggregate_scan_metrics(plan: &dyn ExecutionPlan) -> ScanMetrics {
34    fn sum_bytes(plan: &dyn ExecutionPlan) -> u64 {
35        let mut total = plan
36            .metrics()
37            .and_then(|m| m.sum_by_name("bytes_scanned"))
38            .map(|v| v.as_usize() as u64)
39            .unwrap_or(0);
40        for child in plan.children() {
41            total += sum_bytes(child.as_ref());
42        }
43        total
44    }
45    fn sum_spills(plan: &dyn ExecutionPlan) -> (u64, u64) {
46        let metrics = plan.metrics();
47        let mut spilled_bytes = metrics
48            .as_ref()
49            .and_then(|m| m.spilled_bytes())
50            .map(|v| v as u64)
51            .unwrap_or(0);
52        let mut spill_count = metrics
53            .as_ref()
54            .and_then(|m| m.spill_count())
55            .map(|v| v as u64)
56            .unwrap_or(0);
57        for child in plan.children() {
58            let (child_bytes, child_count) = sum_spills(child.as_ref());
59            spilled_bytes += child_bytes;
60            spill_count += child_count;
61        }
62        (spilled_bytes, spill_count)
63    }
64    let (spilled_bytes, spill_count) = sum_spills(plan);
65    ScanMetrics {
66        output_rows: plan
67            .metrics()
68            .and_then(|m| m.output_rows())
69            .map(|r| r as u64),
70        bytes_scanned: sum_bytes(plan),
71        spilled_bytes,
72        spill_count,
73    }
74}
75
76/// One structured record per FlightSQL query, emitted as a JSON log line under
77/// the `flightsql_query_audit` target when the query completes, fails, or is
78/// abandoned mid-drain (client disconnect/cancel).
79#[derive(serde::Serialize)]
80pub struct QueryAuditRecord {
81    /// Minted as the first statement of `execute_query`, before any fallible
82    /// step -- always present, so a failure's server-log line (which also
83    /// carries it) and this record can be correlated by grepping the id.
84    pub query_id: String,
85    /// The rightmost `X-Forwarded-For` entry (the address the trusted proxy in front of this
86    /// service observed), falling back to `X-Real-IP` and then the gRPC peer address for a
87    /// direct connection -- see `http_utils::get_client_ip`. Unlike `client`/`agent`/`entrypoint`
88    /// below, this comes from the network / trusted-proxy layer rather than from
89    /// client-controlled attribution headers, so it's grouped with `query_id` rather than with
90    /// the self-reported fields -- but that guarantee holds only for requests that actually
91    /// traversed the trusted proxy; for a direct connection with no proxy in front, both
92    /// `X-Forwarded-For` and `X-Real-IP` are just as caller-chosen as the self-reported fields.
93    pub client_ip: String,
94    pub client: String,
95    pub agent: String,
96    pub entrypoint: String,
97    #[serde(skip_serializing_if = "Option::is_none")]
98    pub session: Option<String>,
99    /// Originating notebook name, if the query was issued from a notebook cell.
100    #[serde(skip_serializing_if = "Option::is_none")]
101    pub notebook: Option<String>,
102    /// Originating cell name within the notebook, if the query was issued from a notebook cell.
103    #[serde(skip_serializing_if = "Option::is_none")]
104    pub cell: Option<String>,
105    pub user: String,
106    pub email: String,
107    #[serde(skip_serializing_if = "Option::is_none")]
108    pub name: Option<String>,
109    pub service_account: bool,
110    #[serde(skip_serializing_if = "Option::is_none")]
111    pub service_account_name: Option<String>,
112    pub sql: String,
113    #[serde(skip_serializing_if = "Option::is_none")]
114    pub range_begin: Option<String>, // RFC3339
115    #[serde(skip_serializing_if = "Option::is_none")]
116    pub range_end: Option<String>,
117    #[serde(skip_serializing_if = "Option::is_none")]
118    pub limit: Option<u64>,
119    pub context_init_ms: f64,
120    pub planning_ms: f64,
121    pub execution_ms: f64, // stream construction (matches query_execution_duration semantics)
122    pub setup_ms: f64,     // parse+attribution+context+planning+stream-build (query_setup_duration)
123    pub total_ms: f64,     // end-to-end incl. drain
124    pub status: &'static str, // "ok" | "error" | "incomplete"
125    #[serde(skip_serializing_if = "Option::is_none")]
126    pub error: Option<String>,
127    /// "user" | "resource" | "internal" | "denied", derived from the gRPC status code
128    /// (see `flight_sql_service_impl::error_class`) -- "denied" is stamped directly by the
129    /// query-deny-list check rather than derived from a `DataFusionError`. Omitted when
130    /// `status == "ok"`, matching the `error` field's own "on error" convention.
131    #[serde(skip_serializing_if = "Option::is_none")]
132    pub error_class: Option<&'static str>,
133    #[serde(skip_serializing_if = "Option::is_none")]
134    pub output_rows: Option<u64>,
135    pub bytes_scanned: u64,
136    pub peak_memory_bytes: u64,
137    pub spilled_bytes: u64,
138    pub spill_count: u64,
139    /// Normalized SQL fingerprint -- the first 16 hex chars of the SHA-256 of the
140    /// literal-stripped, whitespace-collapsed token stream. Computed once
141    /// per query and emitted on every terminal path regardless of whether the deny list is in
142    /// active use, since it's what an operator pastes into `deny_queries` after finding an
143    /// offender in the audit log. Appended last so existing JSON consumers are unaffected.
144    pub sql_hash: String,
145}
micromegas/servers/query_audit.rs

micromegas/servers/
query_audit.rs
