Skip to main content

velesdb_server/handlers/query/
aggregation.rs

1//! Aggregation query dispatch and execution.
2//!
3//! Handles detection and execution of GROUP BY / aggregate function queries,
4//! routing them to `execute_aggregate` on the appropriate collection.
5
6use axum::{http::StatusCode, response::IntoResponse, Json};
7use std::sync::Arc;
8use velesdb_core::velesql::Query;
9
10use crate::handlers::helpers::notify_query_timing;
11use crate::types::{
12    AggregationResponse, QueryRequest, QueryResponseMeta, VELESQL_CONTRACT_VERSION,
13};
14use crate::AppState;
15
16use super::velesql_helpers::{parse_and_validate, velesql_collection_not_found, velesql_error};
17
18fn aggregation_result_count(result: &serde_json::Value) -> usize {
19    match result {
20        serde_json::Value::Array(rows) => rows.len(),
21        serde_json::Value::Object(_) => 1,
22        _ => 0,
23    }
24}
25
26pub(crate) fn execute_aggregation_query(
27    state: &Arc<AppState>,
28    collection_name: &str,
29    parsed: &Query,
30    params: &std::collections::HashMap<String, serde_json::Value>,
31    start: std::time::Instant,
32) -> axum::response::Response {
33    // Prefer typed vector collection for aggregation.
34    let result = if let Some(vc) = state.db.get_vector_collection(collection_name) {
35        vc.execute_aggregate(parsed, params)
36    } else if let Some(any) = state.db.get_any_collection(collection_name) {
37        any.execute_aggregate(parsed, params)
38    } else {
39        state.operational_metrics.inc_errors();
40        return velesql_collection_not_found(collection_name);
41    };
42
43    let result = match result {
44        Ok(r) => r,
45        Err(e) => {
46            state.operational_metrics.inc_errors();
47            return velesql_error(
48                StatusCode::UNPROCESSABLE_ENTITY,
49                "VELESQL_AGGREGATION_ERROR",
50                &e.to_string(),
51                "Verify GROUP BY/HAVING clauses and aggregate function arguments",
52                None,
53            );
54        }
55    };
56
57    let elapsed = start.elapsed();
58    let timing_ms = elapsed.as_secs_f64() * 1000.0;
59    notify_query_timing(state, collection_name, start);
60    state
61        .query_duration_histogram
62        .observe(elapsed.as_secs_f64());
63    let count = aggregation_result_count(&result);
64
65    Json(AggregationResponse {
66        result,
67        timing_ms,
68        meta: QueryResponseMeta {
69            velesql_contract_version: VELESQL_CONTRACT_VERSION.to_string(),
70            count,
71        },
72    })
73    .into_response()
74}
75
76/// Resolve the collection name for an aggregation query.
77#[allow(clippy::result_large_err)]
78pub(crate) fn resolve_aggregate_collection(
79    parsed: &Query,
80    req: &QueryRequest,
81) -> Result<String, axum::response::Response> {
82    if !parsed.select.from.is_empty() {
83        return Ok(parsed.select.from.clone());
84    }
85    req.collection
86        .as_ref()
87        .filter(|name| !name.is_empty())
88        .cloned()
89        .ok_or_else(|| {
90            velesql_error(
91                StatusCode::UNPROCESSABLE_ENTITY,
92                "VELESQL_MISSING_COLLECTION",
93                "Aggregation query requires a FROM collection or request-body `collection`",
94                "Add FROM <collection> to query or set `collection` in request JSON",
95                Some(serde_json::json!({
96                    "field": "collection",
97                    "endpoint": "/aggregate"
98                })),
99            )
100        })
101}
102
103/// The refuse-or-proceed phase of [`aggregate`], one rule per step: the
104/// query parses, it IS an aggregation (row/search/graph belong on
105/// `/query`), and it names exactly one resolvable collection. Split out so
106/// the handler charges the error metric in ONE place instead of once per
107/// refusal arm.
108#[allow(clippy::result_large_err)]
109fn prepare_aggregation(
110    req: &QueryRequest,
111) -> Result<(velesdb_core::velesql::Query, String), axum::response::Response> {
112    let parsed = parse_and_validate(&req.query)?;
113    if parsed.is_match_query() || !parsed.select.is_aggregation_query() {
114        return Err(velesql_error(
115            StatusCode::UNPROCESSABLE_ENTITY,
116            "VELESQL_AGGREGATION_ERROR",
117            "Only aggregation queries are accepted on /aggregate",
118            "Use /query for row/search/graph queries; use /aggregate for GROUP BY/aggregate workloads.",
119            Some(serde_json::json!({ "endpoint": "/aggregate" })),
120        ));
121    }
122    let collection_name = resolve_aggregate_collection(&parsed, req)?;
123    Ok((parsed, collection_name))
124}
125
126/// Execute an aggregation-only VelesQL query.
127///
128/// This endpoint is explicit and stable for GROUP BY / HAVING / aggregate workloads.
129#[utoipa::path(
130    post,
131    path = "/aggregate",
132    tag = "query",
133    request_body = QueryRequest,
134    responses(
135        (status = 200, description = "Aggregation results", body = AggregationResponse),
136        (status = 400, description = "Query syntax error", body = crate::types::QueryErrorResponse),
137        (status = 422, description = "Aggregation validation/execution error", body = crate::types::VelesqlErrorResponse),
138        (status = 404, description = "Collection not found", body = crate::types::VelesqlErrorResponse)
139    )
140)]
141pub async fn aggregate(
142    axum::extract::State(state): axum::extract::State<Arc<AppState>>,
143    Json(req): Json<QueryRequest>,
144) -> impl IntoResponse {
145    let start = std::time::Instant::now();
146    state.operational_metrics.inc_queries();
147
148    let (parsed, collection_name) = match prepare_aggregation(&req) {
149        Ok(prepared) => prepared,
150        Err(resp) => {
151            state.operational_metrics.inc_errors();
152            return resp;
153        }
154    };
155
156    // The aggregation scan is synchronous core code — run it on the
157    // blocking pool so the async workers stay responsive.
158    let state_clone = Arc::clone(&state);
159    crate::handlers::helpers::run_blocking(move || {
160        execute_aggregation_query(&state_clone, &collection_name, &parsed, &req.params, start)
161    })
162    .await
163    .unwrap_or_else(|resp| resp)
164}