linkrs_server/http/handlers/
query.rs

1use crate::graph_service::GraphServiceError;
2use crate::value::{from_json as json_value_to_core, to_json as value_to_json};
3use axum::{
4    extract::{Json, State},
5    response::Json as JsonResponse,
6};
7use linkrs_metrics::{ErrorInfo, ErrorType, QueryMetrics, QueryPhase, QueryProfile, StatsManager};
8use linkrs_wire::query::{
9    BatchQueryRequest, BatchQueryResponse, ExplainRequest, QueryData, QueryMetadata, QueryRequest,
10    QueryResponse, QueryStageTimings, ValidateRequest, ValidateResponse,
11};
12use std::sync::Arc;
13use std::time::Instant;
14
15use crate::http::{error::HttpError, state::AppState};
16use crate::storage::{
17    StorageClient, StorageOperationContextOps, StorageSchemaContextOps, StorageSyncContextOps,
18};
19
20#[utoipa::path(
21    post,
22    operation_id = "post_v1_query",
23    path = "/v1/query",
24    tag = "Query",
25    request_body = QueryRequest,
26    responses(
27        (status = 200, body = QueryResponse, description = "Query executed"),
28        (status = 500, description = "Internal error")
29    )
30)]
31pub async fn execute<
32    S: StorageClient
33        + StorageSchemaContextOps
34        + StorageSyncContextOps
35        + StorageOperationContextOps
36        + Clone
37        + Send
38        + Sync
39        + 'static,
40>(
41    State(state): State<AppState<S>>,
42    Json(request): Json<QueryRequest>,
43) -> Result<JsonResponse<QueryResponse>, HttpError> {
44    let graph_service = state.server.get_graph_service();
45
46    let parameters = json_params_to_core(&request.parameters);
47    let session_variables = json_params_to_core(&request.session_variables);
48
49    // Map consistency from wire to API consistency
50    let consistency = match request.consistency.as_deref() {
51        Some(s)
52            if s.eq_ignore_ascii_case("read_your_writes")
53                || s.eq_ignore_ascii_case("ryw")
54                || s.eq_ignore_ascii_case("read-your-writes") =>
55        {
56            linkrs_api::api_core::types::ConsistencyLevel::ReadYourWrites(
57                linkrs_core::types::ReadYourWritesConfig {
58                    timeout_ms: request.consistency_timeout_ms.unwrap_or(2000),
59                    minimum_lsn: request.minimum_lsn.map(linkrs_core::types::CommitLsn::new),
60                },
61            )
62        }
63        _ if request.consistency_timeout_ms.is_some() => {
64            linkrs_api::api_core::types::ConsistencyLevel::ReadYourWrites(
65                linkrs_core::types::ReadYourWritesConfig {
66                    timeout_ms: request.consistency_timeout_ms.unwrap(),
67                    minimum_lsn: request.minimum_lsn.map(linkrs_core::types::CommitLsn::new),
68                },
69            )
70        }
71        _ => linkrs_api::api_core::types::ConsistencyLevel::Eventual,
72    };
73
74    // Executing Queries with GraphService
75    let row_limit = result_size_limit(&state);
76    let stats = state.server.get_stats_manager().clone();
77    let trace_id = new_trace_id();
78    let start = Instant::now();
79    let result = match graph_service
80        .execute_with_consistency(
81            request.session_id,
82            &request.query,
83            parameters,
84            session_variables,
85            consistency,
86        )
87        .await
88    {
89        Ok(exec_result) => {
90            // Configuration intents (`UPDATE CONFIGS` / `SHOW CONFIGS`)
91            // resolve against the live store here; anything else passes
92            // through to the standard response mapping.
93            let store = state.server.config_store();
94            let config_path = state.server.get_config_path();
95            match crate::http::handlers::config::resolve_query_config_intent(
96                exec_result,
97                &store,
98                config_path.as_deref(),
99            ) {
100                Ok(resolved) => {
101                    let elapsed_us = start.elapsed().as_micros() as u64;
102                    let response =
103                        query_result_to_response(resolved, row_limit, &trace_id, elapsed_us);
104                    record_success_profile(
105                        &stats,
106                        request.session_id,
107                        &request.query,
108                        trace_id.clone(),
109                        elapsed_us,
110                        response.metadata.rows_returned,
111                    );
112                    Ok::<_, HttpError>(response)
113                }
114                Err(e) => {
115                    let elapsed_us = start.elapsed().as_micros() as u64;
116                    record_failure_profile(
117                        &stats,
118                        request.session_id,
119                        &request.query,
120                        trace_id.clone(),
121                        elapsed_us,
122                        &e,
123                    );
124                    Ok::<_, HttpError>(error_with_trace(
125                        "CONFIG_ERROR".to_string(),
126                        GraphServiceError::new(e),
127                        trace_id,
128                        elapsed_us,
129                    ))
130                }
131            }
132        }
133        Err(e) => {
134            let elapsed_us = start.elapsed().as_micros() as u64;
135            record_failure_profile(
136                &stats,
137                request.session_id,
138                &request.query,
139                trace_id.clone(),
140                elapsed_us,
141                e.message(),
142            );
143            Ok::<_, HttpError>(error_with_trace(
144                "QUERY_ERROR".to_string(),
145                e,
146                trace_id,
147                elapsed_us,
148            ))
149        }
150    };
151
152    Ok(JsonResponse(result?))
153}
154
155#[utoipa::path(
156    post,
157    path = "/v1/query/batch",
158    tag = "Query",
159    request_body = BatchQueryRequest,
160    responses(
161        (status = 200, body = BatchQueryResponse, description = "Batch queries executed"),
162        (status = 500, description = "Internal error")
163    )
164)]
165pub async fn execute_batch<
166    S: StorageClient
167        + StorageSchemaContextOps
168        + StorageSyncContextOps
169        + StorageOperationContextOps
170        + crate::storage::AutoCommitBatchOps
171        + Clone
172        + Send
173        + Sync
174        + 'static,
175>(
176    State(state): State<AppState<S>>,
177    Json(request): Json<BatchQueryRequest>,
178) -> Result<JsonResponse<BatchQueryResponse>, HttpError> {
179    let graph_service = state.server.get_graph_service();
180
181    let parameters = json_params_to_core(&request.parameters);
182    let session_variables = json_params_to_core(&request.session_variables);
183
184    let outcomes = graph_service
185        .execute_batch(
186            request.session_id,
187            &request.statements,
188            parameters,
189            session_variables,
190        )
191        .await;
192
193    let store = state.server.config_store();
194    let config_path = state.server.get_config_path();
195    let row_limit = result_size_limit(&state);
196    let stats = state.server.get_stats_manager().clone();
197    let batch_start = Instant::now();
198    let results = outcomes
199        .into_iter()
200        .enumerate()
201        .map(|(index, outcome)| {
202            let trace_id = new_trace_id();
203            let elapsed_us = batch_start.elapsed().as_micros() as u64;
204            let query_text = request.statements.get(index).cloned().unwrap_or_default();
205            match outcome {
206                Ok(exec_result) => {
207                    match crate::http::handlers::config::resolve_query_config_intent(
208                        exec_result,
209                        &store,
210                        config_path.as_deref(),
211                    ) {
212                        Ok(resolved) => {
213                            let response = query_result_to_response(
214                                resolved, row_limit, &trace_id, elapsed_us,
215                            );
216                            record_success_profile(
217                                &stats,
218                                request.session_id,
219                                &query_text,
220                                trace_id,
221                                elapsed_us,
222                                response.metadata.rows_returned,
223                            );
224                            response
225                        }
226                        Err(e) => {
227                            record_failure_profile(
228                                &stats,
229                                request.session_id,
230                                &query_text,
231                                trace_id.clone(),
232                                elapsed_us,
233                                &e,
234                            );
235                            error_with_trace(
236                                "CONFIG_ERROR".into(),
237                                GraphServiceError::new(e),
238                                trace_id,
239                                elapsed_us,
240                            )
241                        }
242                    }
243                }
244                Err(e) => {
245                    record_failure_profile(
246                        &stats,
247                        request.session_id,
248                        &query_text,
249                        trace_id.clone(),
250                        elapsed_us,
251                        e.message(),
252                    );
253                    error_with_trace("QUERY_ERROR".into(), e, trace_id, elapsed_us)
254                }
255            }
256        })
257        .collect();
258
259    Ok(JsonResponse(BatchQueryResponse { results }))
260}
261
262#[utoipa::path(
263    post,
264    path = "/v1/query/validate",
265    tag = "Query",
266    request_body = ValidateRequest,
267    responses(
268        (status = 200, body = ValidateResponse, description = "Query validation result"),
269        (status = 500, description = "Internal error")
270    )
271)]
272pub async fn validate<
273    S: StorageClient
274        + StorageSchemaContextOps
275        + StorageSyncContextOps
276        + StorageOperationContextOps
277        + Clone
278        + Send
279        + Sync
280        + 'static,
281>(
282    State(state): State<AppState<S>>,
283    Json(request): Json<ValidateRequest>,
284) -> Result<JsonResponse<ValidateResponse>, HttpError> {
285    // Real validation: parse plus binder name resolution, without executing.
286    match validate_gql(&request.query) {
287        Ok(_) => {
288            // Advisory row estimate for automatic routing. Command-like
289            // statements never stream, so they carry no estimate.
290            let estimated_rows = if !request.need_estimate
291                || crate::graph_service::GraphService::<S>::is_command_like(&request.query)
292            {
293                None
294            } else {
295                state
296                    .server
297                    .get_graph_service()
298                    .estimate_rows(request.session_id, &request.query)
299                    .await
300            };
301            Ok(JsonResponse(ValidateResponse {
302                valid: true,
303                message: "Query is valid".to_string(),
304                estimated_rows,
305            }))
306        }
307        Err(e) => Ok(JsonResponse(ValidateResponse {
308            valid: false,
309            message: e,
310            estimated_rows: None,
311        })),
312    }
313}
314
315#[utoipa::path(
316    post,
317    operation_id = "post_v1_query_explain",
318    path = "/v1/query/explain",
319    tag = "Query",
320    request_body = ExplainRequest,
321    responses(
322        (status = 200, body = QueryResponse, description = "Query plan"),
323        (status = 500, description = "Internal error")
324    )
325)]
326pub async fn explain<
327    S: StorageClient
328        + StorageSchemaContextOps
329        + StorageSyncContextOps
330        + StorageOperationContextOps
331        + Clone
332        + Send
333        + Sync
334        + 'static,
335>(
336    State(state): State<AppState<S>>,
337    Json(request): Json<ExplainRequest>,
338) -> Result<JsonResponse<QueryResponse>, HttpError> {
339    let graph_service = state.server.get_graph_service();
340    let parameters = json_params_to_core(&request.parameters);
341    let session_variables = json_params_to_core(&request.session_variables);
342    let row_limit = result_size_limit(&state);
343    let trace_id = new_trace_id();
344    let start = Instant::now();
345    match graph_service
346        .explain(
347            request.session_id,
348            &request.query,
349            parameters,
350            session_variables,
351        )
352        .await
353    {
354        Ok(exec_result) => {
355            let elapsed_us = start.elapsed().as_micros() as u64;
356            Ok(JsonResponse(query_result_to_response(
357                exec_result,
358                row_limit,
359                &trace_id,
360                elapsed_us,
361            )))
362        }
363        Err(e) => {
364            let elapsed_us = start.elapsed().as_micros() as u64;
365            record_failure_profile(
366                &state.server.get_stats_manager().clone(),
367                request.session_id,
368                &request.query,
369                trace_id.clone(),
370                elapsed_us,
371                e.message(),
372            );
373            Ok(JsonResponse(error_with_trace(
374                "QUERY_ERROR".to_string(),
375                e,
376                trace_id,
377                elapsed_us,
378            )))
379        }
380    }
381}
382
383/// Validate query text by parsing and binding it, without executing.
384///
385/// Returns the declared parameter names on success, or a human-readable
386/// reason on failure. Binding runs without a schema manager, so checks
387/// that need live schema (e.g. tag existence) are skipped while syntax,
388/// name resolution, and function resolution still apply.
389pub(crate) fn validate_gql(query: &str) -> Result<Vec<String>, String> {
390    if query.trim().is_empty() {
391        return Err("query must not be empty".to_string());
392    }
393    let mut parser = crate::query::parser::Parser::new(query);
394    let result = parser.parse().map_err(|e| format!("syntax error: {e}"))?;
395    if parser.has_errors() {
396        return Err(format!("syntax error: {}", parser.take_errors()));
397    }
398    crate::query::binder::Binder::new()
399        .bind(result.ast)
400        .map_err(|e| format!("semantic error: {e}"))?;
401    Ok(extract_parameter_names(query))
402}
403
404/// Collect `@name` parameter references from query text.
405pub(crate) fn extract_parameter_names(query: &str) -> Vec<String> {
406    let mut names = Vec::new();
407    let bytes = query.as_bytes();
408    let mut i = 0;
409    while i < bytes.len() {
410        if bytes[i] == b'@' {
411            let mut j = i + 1;
412            while j < bytes.len() && (bytes[j].is_ascii_alphanumeric() || bytes[j] == b'_') {
413                j += 1;
414            }
415            if j > i + 1 {
416                let name = &query[i + 1..j];
417                if !names.iter().any(|n: &String| n == name) {
418                    names.push(name.to_string());
419                }
420            }
421            i = j;
422        } else {
423            i += 1;
424        }
425    }
426    names
427}
428
429fn new_trace_id() -> String {
430    uuid::Uuid::new_v4().to_string()
431}
432
433fn stages_for_elapsed_us(elapsed_us: u64) -> QueryStageTimings {
434    QueryStageTimings {
435        parse_ms: 0.0,
436        validate_ms: 0.0,
437        plan_ms: 0.0,
438        optimize_ms: 0.0,
439        execute_ms: elapsed_us as f64 / 1000.0,
440    }
441}
442
443fn error_with_trace(
444    code: String,
445    error: GraphServiceError,
446    trace_id: String,
447    elapsed_us: u64,
448) -> QueryResponse {
449    let position = error.position().map(|p| linkrs_wire::query::ErrorPosition {
450        line: p.line,
451        column: p.column,
452    });
453    let mut response =
454        QueryResponse::error_with_position(code, error.message().to_string(), None, position);
455    response.metadata.trace_id = Some(trace_id);
456    response.metadata.stages = Some(stages_for_elapsed_us(elapsed_us));
457    response.metadata.execution_time_ms = elapsed_us / 1000;
458    response
459}
460
461fn record_success_profile(
462    stats: &Arc<StatsManager>,
463    session_id: i64,
464    query_text: &str,
465    trace_id: String,
466    elapsed_us: u64,
467    rows_returned: usize,
468) {
469    let mut profile = QueryProfile::new(session_id, query_text.to_string());
470    profile.trace_id = trace_id;
471    profile.total_duration_us = elapsed_us;
472    profile.stages.execute_us = elapsed_us;
473    profile.result_count = rows_returned;
474    let mut metrics = QueryMetrics::new();
475    metrics.execute_time_us = elapsed_us;
476    metrics.total_time_us = elapsed_us;
477    metrics.result_row_count = rows_returned;
478    stats.record_query_metrics(&metrics);
479    stats.record_query_profile(profile);
480}
481
482fn record_failure_profile(
483    stats: &Arc<StatsManager>,
484    session_id: i64,
485    query_text: &str,
486    trace_id: String,
487    elapsed_us: u64,
488    message: &str,
489) {
490    let mut profile = QueryProfile::new(session_id, query_text.to_string());
491    profile.trace_id = trace_id;
492    profile.total_duration_us = elapsed_us;
493    profile.stages.execute_us = elapsed_us;
494    let mut metrics = QueryMetrics::new();
495    metrics.execute_time_us = elapsed_us;
496    metrics.total_time_us = elapsed_us;
497    stats.record_query_metrics(&metrics);
498    let error_info = ErrorInfo::new(
499        ErrorType::ExecutionError,
500        QueryPhase::Execute,
501        message.to_string(),
502    );
503    stats.record_failed_query(profile, error_info);
504}
505
506/// Convert a core-layer [`QueryResult`] into the wire `QueryResponse`.
507/// The core result carries the engine `ExecutionResult` unchanged; each
508/// variant is rendered here into the JSON wire shape (rows stay in column
509/// order, no intermediate map conversion). When `row_limit` is set, rows
510/// past the ceiling are dropped and the response is marked truncated.
511fn query_result_to_response(
512    result: linkrs_api::api_core::QueryResult,
513    row_limit: Option<usize>,
514    trace_id: &str,
515    elapsed_us: u64,
516) -> QueryResponse {
517    // `metadata.space_id` surfaces the switched-to space. The engine executes
518    // USE as a DataSet with a `space_id` column (the `SpaceSwitched` variant
519    // is never produced); `QueryResult::space_summary` recognizes both.
520    let space_id = result.space_summary().map(|s| s.id);
521    let (columns, mut rows): (
522        Vec<String>,
523        Vec<std::collections::HashMap<String, serde_json::Value>>,
524    ) = match result.execution {
525        crate::query::executor::base::ExecutionResult::DataSet { data } => {
526            let rows = data
527                .rows
528                .iter()
529                .map(|row| {
530                    data.col_names
531                        .iter()
532                        .zip(row.iter())
533                        .map(|(col, value)| (col.clone(), value_to_json(value.clone())))
534                        .collect()
535                })
536                .collect();
537            (data.col_names, rows)
538        }
539        crate::query::executor::base::ExecutionResult::SpaceSwitched(summary) => {
540            let row = std::collections::HashMap::from([
541                (
542                    "space_name".to_string(),
543                    serde_json::Value::String(summary.name.clone()),
544                ),
545                (
546                    "space_id".to_string(),
547                    serde_json::Value::Number(summary.id.into()),
548                ),
549                (
550                    "vid_type".to_string(),
551                    serde_json::Value::String(summary.vid_type.to_string()),
552                ),
553            ]);
554            (
555                vec![
556                    "space_name".to_string(),
557                    "space_id".to_string(),
558                    "vid_type".to_string(),
559                ],
560                vec![row],
561            )
562        }
563        _ => (vec![], vec![]),
564    };
565    let truncated = match row_limit {
566        Some(limit) if rows.len() > limit => {
567            rows.truncate(limit);
568            true
569        }
570        _ => false,
571    };
572    let row_count = rows.len();
573    let execution_time_ms = if result.metadata.execution_time_ms == 0 {
574        elapsed_us / 1000
575    } else {
576        result.metadata.execution_time_ms
577    };
578
579    QueryResponse::success(
580        QueryData::new(columns, rows),
581        QueryMetadata {
582            execution_time_ms,
583            rows_scanned: result.metadata.rows_scanned,
584            rows_returned: row_count,
585            space_id,
586            truncated,
587            trace_id: Some(trace_id.to_string()),
588            stages: Some(stages_for_elapsed_us(elapsed_us)),
589            plan_node_count: None,
590            result_row_count: Some(row_count),
591        },
592    )
593}
594
595/// Read the configured result row ceiling (`None` means unlimited).
596/// Shared by the materialized, streaming, and export paths so one
597/// configuration value guards every result shape.
598pub(crate) fn result_size_limit<
599    S: StorageClient
600        + StorageSchemaContextOps
601        + StorageSyncContextOps
602        + StorageOperationContextOps
603        + Clone
604        + Send
605        + Sync
606        + 'static,
607>(
608    state: &AppState<S>,
609) -> Option<usize> {
610    let store = state.server.config_store();
611    let guard = store.read();
612    let resources = guard.query_resource();
613    if resources.has_result_size_limit() {
614        Some(resources.max_result_size)
615    } else {
616        None
617    }
618}
619
620/// Convert an HTTP request's JSON parameter map to core `Value` bindings.
621/// Empty maps are passed through as `None` so the core sees no bindings.
622pub(crate) fn json_params_to_core(
623    params: &std::collections::HashMap<String, serde_json::Value>,
624) -> Option<std::collections::HashMap<String, linkrs_core::Value>> {
625    if params.is_empty() {
626        return None;
627    }
628    Some(
629        params
630            .iter()
631            .map(|(k, v)| (k.clone(), json_value_to_core(v)))
632            .collect(),
633    )
634}
linkrs_server/http/handlers/query.rs

linkrs_server/http/handlers/
query.rs
