solidb 2.0.2

use super::documents::get_transaction_id;
use super::system::AppState;
use crate::{
    error::DbError,
    sdbql::{ast::Expression, BodyClause, Query, QueryExecutor},
    server::response::ApiResponse,
    storage::{query_cache, StorageEngine},
};
use axum::{
    extract::{Path, State},
    http::{HeaderMap, StatusCode},
    response::Json,
};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::sync::Arc;

// ==================== Constants ====================

/// Default query execution timeout (30 seconds)
const QUERY_TIMEOUT_SECS: u64 = 30;

/// Default slow query threshold in milliseconds (100ms)
const SLOW_QUERY_THRESHOLD_MS: f64 = 100.0;

/// Maximum batch size to prevent memory exhaustion
const MAX_BATCH_SIZE: usize = 10_000;

// ==================== Structs ====================

#[derive(Debug, Deserialize)]
pub struct ExecuteQueryRequest {
    pub query: String,
    #[serde(default, alias = "bindVars")]
    pub bind_vars: std::collections::HashMap<String, Value>,
    #[serde(default = "default_batch_size", alias = "batchSize")]
    pub batch_size: usize,
    #[serde(default = "default_cache")]
    pub cache: bool,
}

fn default_cache() -> bool {
    true
}

fn default_batch_size() -> usize {
    1000
}

/// The rows of a query response: owned, or shared with the query cache.
///
/// A cache hit used to deep-clone every cached document into the response,
/// only to serialize and drop the copy — on a repeated 50-row read that clone
/// and its drop were most of what was left after the state clones.
/// Serializes exactly like the `Vec` it wraps.
#[derive(Debug)]
pub enum QueryRows {
    Owned(Vec<Value>),
    Shared(std::sync::Arc<Vec<Value>>),
}

impl QueryRows {
    pub fn as_slice(&self) -> &[Value] {
        match self {
            QueryRows::Owned(rows) => rows,
            QueryRows::Shared(rows) => rows,
        }
    }
}

impl From<Vec<Value>> for QueryRows {
    fn from(rows: Vec<Value>) -> Self {
        QueryRows::Owned(rows)
    }
}

impl Serialize for QueryRows {
    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        self.as_slice().serialize(serializer)
    }
}

#[derive(Debug, Serialize)]
pub struct ExecuteQueryResponse {
    pub result: QueryRows,
    pub count: usize,
    pub has_more: bool,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub id: Option<String>,
    pub cached: bool,
    #[serde(rename = "executionTimeMs")]
    pub execution_time_ms: f64,
    #[serde(rename = "inserted")]
    pub documents_inserted: usize,
    #[serde(rename = "updated")]
    pub documents_updated: usize,
    #[serde(rename = "deleted")]
    pub documents_removed: usize,
}

// ==================== Helper Functions ====================

/// Check if a query is potentially long-running (contains mutations or range iterations)
///
/// Looks into set-operation operands and CTE bodies too: an operand's `FOR`
/// scan is just as blocking as a top-level one.
pub(crate) fn is_long_running_query(query: &Query) -> bool {
    let own = query.body_clauses.iter().any(|clause| match clause {
        BodyClause::Insert(_)
        | BodyClause::Update(_)
        | BodyClause::Remove(_)
        | BodyClause::Upsert(_) => true,
        // All FOR loops should use spawn_blocking because:
        // 1. Range expressions (source_expression.is_some()) can be large
        // 2. Collection scans might trigger scatter-gather with blocking HTTP calls
        BodyClause::For(_) => true,
        _ => false,
    });

    own || query
        .set_operations
        .iter()
        .any(|op| is_long_running_query(&op.query))
        || query.with_clause.as_ref().is_some_and(|with| {
            with.ctes
                .iter()
                .any(|cte| is_long_running_query(&cte.query))
        })
        // Pre-FOR `LET x = (FOR ...)` lives outside `body_clauses`, and a
        // bare `RETURN EMBED(...)` / `RETURN NEIGHBORS(...)` has none at all;
        // both ran inline on a tokio worker for up to the whole timeout.
        || query
            .let_clauses
            .iter()
            .any(|l| expression_is_heavy(&l.expression))
        || query
            .return_clause
            .as_ref()
            .is_some_and(|r| expression_is_heavy(&r.expression))
}

/// Largest collection a lone `FOR` may scan inline, on the async worker.
const INLINE_SCAN_MAX_DOCS: usize = 256;

/// A read whose only loop is one `FOR` over a small local collection.
///
/// [`is_long_running_query`] sends every `FOR` to the blocking pool, because
/// a scan can be arbitrarily large. When the collection holds at most
/// [`INLINE_SCAN_MAX_DOCS`] documents the scan is tens of microseconds, less
/// than the handoff itself (a futex wake, two context switches, and the
/// timeout machinery): measured on a 50-document collection it is several
/// percent of the request. Anything that could grow the work — a second
/// loop, a join, a mutation, a subquery or heavy function, a sharded or
/// unknown collection — keeps the blocking pool.
pub(crate) fn is_small_local_scan(query: &Query, storage: &StorageEngine, db_name: &str) -> bool {
    if !query.set_operations.is_empty()
        || !query.join_clauses.is_empty()
        || query.with_clause.is_some()
        || query.window_clause.is_some()
        || query.has_mutations()
    {
        return false;
    }
    let Some((BodyClause::For(for_clause), rest)) = query.body_clauses.split_first() else {
        return false;
    };
    if for_clause.source_expression.is_some()
        || for_clause.system_time.is_some()
        || for_clause.valid_time.is_some()
        || for_clause
            .source_variable
            .as_ref()
            .is_some_and(|s| s != &for_clause.collection)
    {
        return false;
    }
    let light = |e: &Expression| !expression_is_heavy(e);
    let clauses_light = rest.iter().all(|c| match c {
        BodyClause::Filter(f) => light(&f.expression),
        BodyClause::Let(l) => light(&l.expression),
        _ => false,
    });
    if !clauses_light
        || !query.let_clauses.iter().all(|l| light(&l.expression))
        || !query.post_limit_lets.iter().all(|l| light(&l.expression))
        || !query
            .sort_clause
            .as_ref()
            .is_none_or(|s| s.fields.iter().all(|(e, _)| light(e)))
        || !query
            .return_clause
            .as_ref()
            .is_none_or(|r| light(&r.expression))
    {
        return false;
    }
    let Ok(collection) = storage
        .get_database(db_name)
        .and_then(|db| db.get_collection(&for_clause.collection))
    else {
        return false;
    };
    collection
        .get_shard_config()
        .is_none_or(|c| c.num_shards == 0)
        && collection.count() <= INLINE_SCAN_MAX_DOCS
}

/// Functions that scan a collection, call out over the network, or run a
/// model — none of which belong on an async worker thread.
const HEAVY_FUNCTIONS: &[&str] = &[
    "SAMPLE",
    "LEVENSHTEIN",
    "FUZZY_MATCH",
    "HYBRID_SEARCH",
    "SEARCH_INDEX",
    "SNAPSHOT_DIFF",
    "EMBED",
    "EMBED_BATCH",
    "RERANK",
    "LLM",
    "CHAT",
    "GRAPH_RAG",
    "GRAPH_RAG_SEARCH",
    "NEIGHBORS",
    "PAGERANK",
    "DEGREE_CENTRALITY",
    "COMMUNITY_SEARCH",
    "SHORTEST_PATH",
    "K_PATHS",
    "FULLTEXT",
    "VECTOR_SEARCH",
];

/// Whether evaluating this expression can do work proportional to a
/// collection or wait on I/O: a subquery anywhere inside it, or a call to a
/// function in `HEAVY_FUNCTIONS`.
fn expression_is_heavy(expr: &Expression) -> bool {
    use crate::sdbql::ast::Expression as E;
    match expr {
        E::Subquery(q) => {
            !q.body_clauses.is_empty()
                || is_long_running_query(q)
                || q.return_clause
                    .as_ref()
                    .is_some_and(|r| expression_is_heavy(&r.expression))
        }
        E::FunctionCall { name, args } => {
            HEAVY_FUNCTIONS.contains(&name.to_ascii_uppercase().as_str())
                || args.iter().any(expression_is_heavy)
        }
        other => {
            let mut heavy = false;
            other.for_each_child(&mut |child| heavy |= expression_is_heavy(child));
            heavy
        }
    }
}

/// Get collection names affected by mutation clauses for targeted cache invalidation.
/// Collections a query mutates, for cache invalidation.
///
/// `pub(crate)` so the native-driver query handler can invalidate on the same
/// terms as this one — it used to skip invalidation entirely because it had no
/// access to this.
/// Drop cached results for the collections a mutation touched, or the whole
/// cache when the set could not be determined.
pub(crate) fn invalidate_collections(collections: &[String]) {
    if collections.is_empty() {
        query_cache::get_query_cache().invalidate_all();
    } else {
        for collection in collections {
            query_cache::get_query_cache().invalidate_collection(collection);
        }
    }
}

pub(crate) fn mutated_collections(query: &Query) -> std::collections::HashSet<String> {
    // Every collection the query names, from the same AST walk the cache is
    // keyed on (audit P2). That is a superset of the written ones, and it sees
    // a mutation inside a subquery or CTE that a clause-only scan missed.
    // Empty — which makes `invalidate_collections` drop everything — when the
    // set cannot be determined.
    match query_cache::cacheable_collections(query) {
        Some(refs) => refs.required.into_iter().chain(refs.maybe).collect(),
        None => std::collections::HashSet::new(),
    }
}

/// Log slow query to _slow_queries collection (async, non-blocking)
#[allow(clippy::too_many_arguments)]
fn log_slow_query(
    storage: Arc<StorageEngine>,
    db_name: String,
    query_text: String,
    execution_time_ms: f64,
    results_count: usize,
    documents_inserted: usize,
    documents_updated: usize,
    documents_removed: usize,
    origin: Option<String>,
    cf_ops_during: u64,
    cf_ops_ms_during: f64,
) {
    // Only log if query exceeds threshold
    if execution_time_ms < SLOW_QUERY_THRESHOLD_MS {
        return;
    }

    // Never log queries that touch the slow-query log itself: when CF churn
    // (or load) pushes them past the threshold they re-enter the log right
    // after a clear — a feedback loop that makes clearing unreliable.
    if query_text.contains("_slow_queries") {
        return;
    }

    // Spawn background task to avoid blocking the response
    tokio::spawn(async move {
        let slow_query_coll = format!("{}:_slow_queries", db_name);

        // Get or create the _slow_queries collection.
        // Concurrent slow queries may race to create it, so retry lookup briefly.
        let collection = match storage.get_collection(&slow_query_coll) {
            Ok(coll) => coll,
            Err(_) => {
                if let Ok(db) = storage.get_database(&db_name) {
                    let _ = db.create_collection("_slow_queries".to_string(), None);
                } else {
                    return;
                }

                let mut last_err = None;
                let mut resolved = None;
                for _ in 0..10 {
                    match storage.get_collection(&slow_query_coll) {
                        Ok(coll) => {
                            resolved = Some(coll);
                            break;
                        }
                        Err(e) => {
                            last_err = Some(e);
                            tokio::time::sleep(std::time::Duration::from_millis(10)).await;
                        }
                    }
                }

                if let Some(coll) = resolved {
                    coll
                } else {
                    tracing::warn!(
                        "Failed to get _slow_queries collection after retries: {}",
                        last_err
                            .map(|e| e.to_string())
                            .unwrap_or_else(|| "unknown error".to_string())
                    );
                    return;
                }
            }
        };

        // Create log entry. `cf_ops_during` / `cf_ops_ms_during` record the
        // column-family create/drop activity that overlapped this query —
        // CF ops hold the RocksDB DB mutex for their full OPTIONS rewrite,
        // so a high value means the query was queueing behind CF churn
        // (e.g. a test suite), not doing slow work itself.
        let log_entry = serde_json::json!({
            "query": query_text,
            "execution_time_ms": execution_time_ms,
            "timestamp": chrono::Utc::now().to_rfc3339(),
            "results_count": results_count,
            "documents_inserted": documents_inserted,
            "documents_updated": documents_updated,
            "documents_removed": documents_removed,
            "origin": origin,
            "cf_ops_during": cf_ops_during,
            "cf_ops_ms_during": cf_ops_ms_during
        });

        if let Err(e) = collection.insert(log_entry) {
            tracing::warn!("Failed to log slow query: {}", e);
        }
    });
}

/// The write identity of an HTTP request, for the storage write getters.
///
/// A handler whose route is not behind the auth layer has no claims; that is
/// a non-admin client, never the server.
pub fn write_actor_from_claims(
    claims: Option<&crate::server::auth::Claims>,
) -> crate::storage::WriteActor {
    match claims {
        Some(c) => crate::storage::WriteActor::client(principal_from_claims(c).can_admin),
        None => crate::storage::WriteActor::client(false),
    }
}

/// Whether this request is a coordinator's scatter-gather sub-query, which
/// must run locally instead of fanning out again.
///
/// Audit C1: the header alone used to be honoured, so any client could send
/// `X-Scatter-Gather` and read only the local shard of a sharded collection
/// (skipping the coordinator's routing). Only the cluster-internal principal,
/// which `auth_middleware` grants after verifying the cluster secret, may.
fn is_scatter_gather_subquery(headers: &HeaderMap, claims: &crate::server::auth::Claims) -> bool {
    headers.contains_key("X-Scatter-Gather") && crate::server::auth::is_cluster_internal(claims)
}

pub fn principal_from_claims(claims: &crate::server::auth::Claims) -> crate::sdbql::QueryPrincipal {
    crate::sdbql::QueryPrincipal::from_roles(
        claims.sub.clone(),
        claims.roles.clone().unwrap_or_default(),
    )
}

// ==================== Handlers ====================

pub async fn execute_query(
    State(state): State<AppState>,
    Path(db_name): Path<String>,
    headers: HeaderMap,
    axum::Extension(claims): axum::Extension<crate::server::auth::Claims>,
    Json(req): Json<ExecuteQueryRequest>,
) -> Result<ApiResponse<ExecuteQueryResponse>, DbError> {
    // Count every query
    state
        .query_counter
        .fetch_add(1, std::sync::atomic::Ordering::Relaxed);

    // Auth is enforced by `auth_middleware` on the route layer (routes.rs).
    // The authz middleware only required Read for this endpoint (queries are
    // reads by default); if the parsed query mutates, upgrade to Write here.
    {
        let prepared = crate::sdbql::get_prepared_statement_cache().parse_if_needed(&req.query)?;
        if prepared.query.has_mutations() {
            crate::server::authz_middleware::enforce(
                &claims,
                &state,
                crate::server::authorization::PermissionAction::Write,
                Some(&db_name),
            )
            .await?;
        }
    }

    // Check for transaction context
    if let Some(tx_id) = get_transaction_id(&headers) {
        // Execute transactional SDBQL query
        use crate::sdbql::ast::BodyClause;

        let prepared = crate::sdbql::get_prepared_statement_cache().parse_if_needed(&req.query)?;
        let query = prepared.query.as_ref();

        // Get transaction manager
        let tx_manager = state.storage.transaction_manager()?;
        let tx_arc = tx_manager.get(tx_id)?;
        let mut tx = tx_arc
            .write()
            .map_err(|_| DbError::InternalError("Transaction lock poisoned".into()))?;
        let wal = tx_manager.wal().clone();
        let lock_manager = tx_manager.lock_manager().clone();

        // Check if query contains mutation operations. `has_mutations()` also
        // sees UPSERT and mutations nested in set-operation operands or CTE
        // bodies — those must not take the read-only path, which bypasses the
        // transaction's WAL and locks.
        let has_mutations = query.has_mutations();

        if !has_mutations {
            // No mutations - just execute normally (read operations)
            // No mutations - just execute normally (read operations)
            let executor = if req.bind_vars.is_empty() {
                QueryExecutor::with_database(&state.storage, db_name)
            } else {
                QueryExecutor::with_database_and_bind_vars(&state.storage, db_name, req.bind_vars)
            }
            .with_principal(principal_from_claims(&claims))
            .with_timeout(std::time::Duration::from_secs(QUERY_TIMEOUT_SECS));

            let results = executor.execute(query)?;
            return Ok(ApiResponse::new(
                ExecuteQueryResponse {
                    count: results.len(),
                    result: results.into(),
                    has_more: false,
                    id: None,
                    cached: false,
                    execution_time_ms: 0.0,
                    documents_inserted: 0,
                    documents_updated: 0,
                    documents_removed: 0,
                },
                &headers,
            ));
        }

        // For mutation operations, execute transactionally
        let executor = if req.bind_vars.is_empty() {
            QueryExecutor::with_database(&state.storage, db_name.clone())
        } else {
            QueryExecutor::with_database_and_bind_vars(
                &state.storage,
                db_name.clone(),
                req.bind_vars.clone(),
            )
        }
        .with_timeout(std::time::Duration::from_secs(QUERY_TIMEOUT_SECS));

        // Execute body clauses manually to intercept mutations
        // Bind variables are read from the executor, never copied into this
        // context: it is cloned once per row below, and a 3 MB `@rows` over
        // 20k rows was 60 GB — the server was OOM-killed.
        let mut initial_bindings = crate::sdbql::executor::types::Context::default();

        // Process LET clauses
        for let_clause in &query.let_clauses {
            let value =
                executor.evaluate_expr_with_context(&let_clause.expression, &initial_bindings)?;
            initial_bindings.insert(let_clause.variable.clone(), value);
        }

        let mut rows: Vec<crate::sdbql::executor::types::Context> = vec![initial_bindings.clone()];
        let mut mutation_count = 0;

        // Process body clauses in order
        for clause in &query.body_clauses {
            match clause {
                BodyClause::For(for_clause) => {
                    let mut new_rows = Vec::new();
                    for ctx in &rows {
                        let docs = if let Some(ref expr) = for_clause.source_expression {
                            let value = executor.evaluate_expr_with_context(expr, ctx)?;
                            match value {
                                Value::Array(arr) => arr,
                                other => vec![other],
                            }
                        } else {
                            let source_name = for_clause
                                .source_variable
                                .as_ref()
                                .unwrap_or(&for_clause.collection);
                            if let Some(value) = ctx.get(source_name) {
                                match value {
                                    Value::Array(arr) => arr.clone(),
                                    other => vec![other.clone()],
                                }
                            } else {
                                // Scan collection - check if sharded
                                let full_coll_name =
                                    format!("{}:{}", db_name, for_clause.collection);
                                let collection = state.storage.get_collection(&full_coll_name)?;
                                let shard_config = collection.get_shard_config();

                                if let (Some(config), Some(coordinator)) =
                                    (shard_config, &state.shard_coordinator)
                                {
                                    // Sharded collection - use scatter-gather
                                    // Execute async operation in blocking context
                                    let coordinator_clone = coordinator.clone();
                                    let db_name_owned = db_name.to_string();
                                    let coll_name_owned = for_clause.collection.clone();
                                    let config_clone = config.clone();

                                    match tokio::task::block_in_place(|| {
                                        tokio::runtime::Handle::current().block_on(async {
                                            coordinator_clone
                                                .scan_all_shards(
                                                    &db_name_owned,
                                                    &coll_name_owned,
                                                    &config_clone,
                                                )
                                                .await
                                        })
                                    }) {
                                        Ok(docs) => {
                                            docs.into_iter().map(|d| d.to_value()).collect()
                                        }
                                        Err(e) => {
                                            eprintln!("Scatter-gather failed: {:?}, using local shards only", e);
                                            collection.scan_values(executor.scan_cap())
                                        }
                                    }
                                } else {
                                    // Non-sharded or no coordinator - local scan
                                    collection.scan_values(executor.scan_cap())
                                }
                            }
                        };

                        for doc in docs {
                            let mut new_ctx = ctx.clone();
                            new_ctx.insert(for_clause.variable.clone(), doc);
                            new_rows.push(new_ctx);
                        }
                        // This hand-rolled pipeline had no ceiling at all.
                        executor.check_budget(new_rows.len())?;
                    }
                    rows = new_rows;
                }
                BodyClause::Let(let_clause) => {
                    for ctx in &mut rows {
                        let value =
                            executor.evaluate_expr_with_context(&let_clause.expression, ctx)?;
                        ctx.insert(let_clause.variable.clone(), value);
                    }
                }
                BodyClause::Filter(filter_clause) | BodyClause::Search(filter_clause) => {
                    rows.retain(|ctx| {
                        executor
                            .evaluate_filter_with_context(&filter_clause.expression, ctx)
                            .unwrap_or(false)
                    });
                }
                BodyClause::Insert(insert_clause) => {
                    let full_coll_name = format!("{}:{}", db_name, insert_clause.collection);
                    let collection = state.storage.get_collection(&full_coll_name)?;

                    for ctx in &rows {
                        let doc_value =
                            executor.evaluate_expr_with_context(&insert_clause.document, ctx)?;
                        collection.insert_tx(&mut tx, &wal, &lock_manager, doc_value)?;
                        mutation_count += 1;
                    }
                }
                BodyClause::Update(update_clause) => {
                    let full_coll_name = format!("{}:{}", db_name, update_clause.collection);
                    let collection = state.storage.get_collection(&full_coll_name)?;

                    for ctx in &rows {
                        let selector_value =
                            executor.evaluate_expr_with_context(&update_clause.selector, ctx)?;
                        let key = match &selector_value {
                            Value::String(s) => s.clone(),
                            Value::Object(obj) => obj.get("_key")
                                .and_then(|v| v.as_str())
                                .map(|s| s.to_string())
                                .ok_or_else(|| DbError::ExecutionError(
                                    "UPDATE: selector object must have a _key field".to_string()
                                ))?,
                            _ => return Err(DbError::ExecutionError(
                                "UPDATE: selector must be a string key or an object with _key field".to_string()
                            )),
                        };

                        let changes_value =
                            executor.evaluate_expr_with_context(&update_clause.changes, ctx)?;
                        collection.update_tx(&mut tx, &wal, &lock_manager, &key, changes_value)?;
                        mutation_count += 1;
                    }
                }
                BodyClause::Remove(remove_clause) => {
                    let full_coll_name = format!("{}:{}", db_name, remove_clause.collection);
                    let collection = state.storage.get_collection(&full_coll_name)?;

                    for ctx in &rows {
                        let selector_value =
                            executor.evaluate_expr_with_context(&remove_clause.selector, ctx)?;
                        let key = match &selector_value {
                            Value::String(s) => s.clone(),
                            Value::Object(obj) => obj.get("_key")
                                .and_then(|v| v.as_str())
                                .map(|s| s.to_string())
                                .ok_or_else(|| DbError::ExecutionError(
                                    "REMOVE: selector object must have a _key field".to_string()
                                ))?,
                            _ => return Err(DbError::ExecutionError(
                                "REMOVE: selector must be a string key or an object with _key field".to_string()
                            )),
                        };

                        collection.delete_tx(&mut tx, &wal, &lock_manager, &key)?;
                        mutation_count += 1;
                    }
                }
                _ => {}
            }
        }

        // Return mutation result
        return Ok(ApiResponse::new(
            ExecuteQueryResponse {
                result: vec![serde_json::json!({
                    "mutationCount": mutation_count,
                    "message": format!("{} operation(s) staged in transaction. Commit to apply changes.", mutation_count)
                })]
                .into(),
                count: 1,
                has_more: false,
                id: None,
                cached: false,
                execution_time_ms: 0.0,
                documents_inserted: 0, // Transactional mutations are not counted until commit
                documents_updated: 0,
                documents_removed: 0,
            },
            &headers,
        ));
    }

    // Non-transactional execution (existing logic)
    // Use prepared statement cache to avoid re-parsing frequently executed queries
    let prepared = crate::sdbql::get_prepared_statement_cache().parse_if_needed(&req.query)?;
    let query = prepared.query.as_ref();

    // Check if query is cacheable (read-only with no mutations). This has to be
    // the deep check: caching a query whose mutation hides in a set-operation
    // operand or a CTE body would serve the cached rows on the next identical
    // request and never run the mutation at all.
    let is_read_only = !query.has_mutations();

    // Try to get cached result for read-only queries (unless cache is disabled).
    // The key covers the database and the caller (audit H1: row policies and
    // CURRENT_USER make rows principal-dependent), and is `None` when the
    // query's collection set cannot be read off the AST (audit P2).
    // The generation is snapshotted first so a write landing mid-query
    // prevents the stale rows from being stored.
    let cache_generation = query_cache::current_generation();
    let cache_key = if is_read_only && req.cache {
        query_cache::cache_key_for(
            &state.storage,
            &db_name,
            &req.query,
            query,
            &req.bind_vars,
            &principal_from_claims(&claims),
        )
    } else {
        None
    };

    // Check cache if query is read-only and caching is enabled
    if let Some(ref key) = cache_key {
        let cached_result = query_cache::get_query_cache().get(key);
        if let Some(result) = cached_result {
            tracing::debug!(
                "Query cache hit for: {}",
                &req.query[..req.query.len().min(50)]
            );
            return Ok(ApiResponse::new(
                ExecuteQueryResponse {
                    count: result.len(),
                    result: QueryRows::Shared(result),
                    has_more: false,
                    id: None,
                    cached: true,
                    execution_time_ms: 0.0,
                    documents_inserted: 0,
                    documents_updated: 0,
                    documents_removed: 0,
                },
                &headers,
            ));
        }
    }

    // Handle CREATE STREAM clause
    if let Some(ref _create_stream) = query.create_stream_clause {
        if let Some(manager) = &state.stream_manager {
            match manager.create_stream_as(
                &db_name,
                (*query).clone(),
                principal_from_claims(&claims),
            ) {
                Ok(_name) => {
                    return Ok(ApiResponse::new(
                        ExecuteQueryResponse {
                            result: Vec::new().into(),
                            count: 0,
                            has_more: false,
                            id: None,
                            cached: false,
                            execution_time_ms: 0.0,
                            documents_inserted: 0,
                            documents_updated: 0,
                            documents_removed: 0,
                        },
                        &headers,
                    ));
                }
                Err(e) => return Err(e),
            }
        } else {
            return Err(DbError::OperationNotSupported(
                "Stream processing not enabled".to_string(),
            ));
        }
    }

    let batch_size = req.batch_size.min(MAX_BATCH_SIZE);

    // Clone db_name and query text for slow query logging (before they're moved)
    let db_name_for_logging = db_name.clone();
    let db_name_for_cursor = db_name.clone();
    let query_text_for_logging = req.query.clone();

    // Snapshot CF-op activity so the slow-query log can tell a genuinely
    // slow query from one that queued behind create_cf/drop_cf churn.
    let cf_ops_before = crate::storage::cf_ops::snapshot();

    // Only use spawn_blocking for potentially long-running queries
    // (mutations or range iterations). Simple reads run directly.
    let mutates = query.has_mutations();
    let (query_result, execution_time_ms) =
        if is_long_running_query(query) && !is_small_local_scan(query, &state.storage, &db_name) {
            let storage = state.storage.clone();
            let bind_vars = req.bind_vars.clone();
            let replication_log = state.replication_log.clone();
            let shard_coordinator = state.shard_coordinator.clone();
            let is_scatter_gather = is_scatter_gather_subquery(&headers, &claims);
            let query = (*query).clone();
            let principal = principal_from_claims(&claims);

            // Collections this query invalidates, resolved before the executor
            // moves out of reach, so the timeout arm below can still drop them.
            let invalidated: Vec<String> = if mutates {
                mutated_collections(&query).into_iter().collect()
            } else {
                Vec::new()
            };

            // Apply timeout to prevent DoS from long-running queries
            let mut task = tokio::task::spawn_blocking(move || {
                let mut executor = if bind_vars.is_empty() {
                    QueryExecutor::with_database(&storage, db_name)
                } else {
                    QueryExecutor::with_database_and_bind_vars(&storage, db_name, bind_vars)
                }
                .with_principal(principal)
                .with_timeout(std::time::Duration::from_secs(QUERY_TIMEOUT_SECS));

                // Add replication service for mutation logging
                if let Some(ref log) = replication_log {
                    executor = executor.with_replication(log);
                }

                // Inject shard coordinator for scatter-gather (if not already a sub-query)
                if !is_scatter_gather {
                    if let Some(coord) = shard_coordinator {
                        executor = executor.with_shard_coordinator(coord);
                    }
                }

                let start = std::time::Instant::now();
                let result = executor.execute_with_stats(&query)?;
                let execution_time_ms = start.elapsed().as_secs_f64() * 1000.0;
                Ok::<_, DbError>((result, execution_time_ms))
            });

            // `&mut task` so the handle survives a timeout and can still be awaited.
            match tokio::time::timeout(
                std::time::Duration::from_secs(QUERY_TIMEOUT_SECS),
                &mut task,
            )
            .await
            {
                Ok(join_result) => join_result
                    .map_err(|e| DbError::InternalError(format!("Task join error: {}", e)))??,
                Err(_) => {
                    // A blocking task is not cancellable: dropping the handle does
                    // not stop the executor, so a mutation that overruns still
                    // commits (and still reaches the replication log). Drop the
                    // cached rows now, and again once the write really lands, or
                    // readers keep being served the pre-mutation result.
                    if mutates {
                        invalidate_collections(&invalidated);
                        tokio::spawn(async move {
                            let _ = task.await;
                            invalidate_collections(&invalidated);
                        });
                    }
                    return Err(DbError::BadRequest(format!(
                        "Query execution timeout: exceeded {} seconds",
                        QUERY_TIMEOUT_SECS
                    )));
                }
            }
        } else {
            let mut executor = if req.bind_vars.is_empty() {
                QueryExecutor::with_database(&state.storage, db_name)
            } else {
                QueryExecutor::with_database_and_bind_vars(&state.storage, db_name, req.bind_vars)
            }
            .with_principal(principal_from_claims(&claims))
            .with_timeout(std::time::Duration::from_secs(QUERY_TIMEOUT_SECS));

            // Add replication service for mutation logging
            if let Some(ref log) = state.replication_log {
                executor = executor.with_replication(log);
            }

            // Inject shard coordinator for scatter-gather (if not already a sub-query)
            if !is_scatter_gather_subquery(&headers, &claims) {
                if let Some(coordinator) = state.shard_coordinator.clone() {
                    executor = executor.with_shard_coordinator(coordinator);
                }
            }

            let start = std::time::Instant::now();
            let result = executor.execute_with_stats(query)?;
            let execution_time_ms = start.elapsed().as_secs_f64() * 1000.0;
            (result, execution_time_ms)
        };

    let total_count = query_result.results.len();
    let mutations = &query_result.mutations;

    // Increment write counter for mutation queries
    if mutations.has_mutations() {
        state
            .write_counter
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    }

    // Cache read-only query results. The cache evicts by entry count, so
    // one large result set per entry is the whole heap; only small ones are
    // worth keeping, and those are the ones a cache hit helps.
    if let Some(key) = cache_key {
        if query_result.results.len() <= query_cache::MAX_CACHED_ROWS {
            let result_clone: Vec<serde_json::Value> = query_result.results.clone();
            query_cache::get_query_cache().put_if_current(key, result_clone, cache_generation);
        }
    }

    // Invalidate query cache when mutations occurred
    if mutations.has_mutations() {
        let collections: Vec<String> = mutated_collections(query).into_iter().collect();
        invalidate_collections(&collections);
    }

    // Log slow query if it exceeds threshold (async, non-blocking)
    let cf_ops_after = crate::storage::cf_ops::snapshot();
    log_slow_query(
        state.storage.clone(),
        db_name_for_logging,
        query_text_for_logging,
        execution_time_ms,
        total_count,
        mutations.documents_inserted,
        mutations.documents_updated,
        mutations.documents_removed,
        Some(claims.sub.clone()),
        cf_ops_before.ops_since(&cf_ops_after),
        cf_ops_before.ms_since(&cf_ops_after),
    );

    let (cursor_id, result_batch, has_more) = state.cursor_store.store_and_get_first_batch(
        db_name_for_cursor,
        query_result.results,
        batch_size,
    );

    Ok(ApiResponse::new(
        ExecuteQueryResponse {
            result: result_batch.into(),
            count: total_count,
            has_more,
            id: cursor_id,
            cached: false,
            execution_time_ms,
            documents_inserted: mutations.documents_inserted,
            documents_updated: mutations.documents_updated,
            documents_removed: mutations.documents_removed,
        },
        &headers,
    ))
}

pub async fn explain_query(
    State(state): State<AppState>,
    Path(db_name): Path<String>,
    headers: HeaderMap,
    axum::Extension(claims): axum::Extension<crate::server::auth::Claims>,
    Json(req): Json<ExecuteQueryRequest>,
) -> Result<Json<crate::sdbql::QueryExplain>, DbError> {
    let prepared = crate::sdbql::get_prepared_statement_cache().parse_if_needed(&req.query)?;
    let query = (*prepared.query).clone();
    let bind_vars = req.bind_vars.clone();
    let storage = state.storage.clone();
    let shard_coordinator = state.shard_coordinator.clone();
    let is_scatter_gather = is_scatter_gather_subquery(&headers, &claims);
    // EXPLAIN reports whether *this* caller's run would auto-index, so it needs
    // the same principal the run would get.
    let principal = principal_from_claims(&claims);

    let explain = {
        let storage = storage.clone();
        match tokio::time::timeout(
            std::time::Duration::from_secs(QUERY_TIMEOUT_SECS),
            tokio::task::spawn_blocking(move || {
                let mut executor = if bind_vars.is_empty() {
                    QueryExecutor::with_database(&storage, db_name)
                } else {
                    QueryExecutor::with_database_and_bind_vars(&storage, db_name, bind_vars)
                }
                .with_principal(principal)
                .with_timeout(std::time::Duration::from_secs(QUERY_TIMEOUT_SECS));

                if !is_scatter_gather {
                    if let Some(coordinator) = shard_coordinator {
                        executor = executor.with_shard_coordinator(coordinator);
                    }
                }

                executor.explain(&query)
            }),
        )
        .await
        {
            Ok(join_result) => join_result
                .map_err(|e| DbError::InternalError(format!("Task join error: {}", e)))??,
            Err(_) => {
                return Err(DbError::BadRequest(format!(
                    "Explain timeout: exceeded {} seconds",
                    QUERY_TIMEOUT_SECS
                )))
            }
        }
    };

    Ok(Json(explain))
}

pub async fn get_next_batch(
    State(state): State<AppState>,
    Path(cursor_id): Path<String>,
    headers: HeaderMap,
    axum::Extension(claims): axum::Extension<crate::server::auth::Claims>,
) -> Result<ApiResponse<ExecuteQueryResponse>, DbError> {
    // The cursor route has no {db} path param; check read permission against
    // the database the original query ran on.
    if let Some(db) = state.cursor_store.db_name(&cursor_id) {
        crate::server::authz_middleware::enforce(
            &claims,
            &state,
            crate::server::authorization::PermissionAction::Read,
            Some(&db),
        )
        .await?;
    }
    if let Some((batch, has_more)) = state.cursor_store.get_next_batch(&cursor_id) {
        let count = batch.len();
        Ok(ApiResponse::new(
            ExecuteQueryResponse {
                result: batch.into(),
                count,
                has_more,
                id: if has_more { Some(cursor_id) } else { None },
                cached: true,
                execution_time_ms: 0.0, // Cached results, no execution time
                documents_inserted: 0,  // Mutations already counted in first response
                documents_updated: 0,
                documents_removed: 0,
            },
            &headers,
        ))
    } else {
        Err(DbError::DocumentNotFound(format!(
            "Cursor not found or expired: {}",
            cursor_id
        )))
    }
}

pub async fn delete_cursor(
    State(state): State<AppState>,
    Path(cursor_id): Path<String>,
    axum::Extension(claims): axum::Extension<crate::server::auth::Claims>,
) -> Result<StatusCode, DbError> {
    if let Some(db) = state.cursor_store.db_name(&cursor_id) {
        crate::server::authz_middleware::enforce(
            &claims,
            &state,
            crate::server::authorization::PermissionAction::Read,
            Some(&db),
        )
        .await?;
    }
    if state.cursor_store.delete(&cursor_id) {
        Ok(StatusCode::NO_CONTENT)
    } else {
        Err(DbError::DocumentNotFound(format!(
            "Cursor not found: {}",
            cursor_id
        )))
    }
}

#[cfg(test)]
mod long_running_tests {
    use super::is_long_running_query;
    use crate::sdbql::parser::parse;

    #[test]
    fn point_reads_stay_inline() {
        assert!(!is_long_running_query(&parse("RETURN 1 + 1").unwrap()));
        assert!(!is_long_running_query(
            &parse("LET x = 2 RETURN x * 3").unwrap()
        ));
    }

    /// Pre-FOR `LET` lives outside `body_clauses`; a scan hidden in one ran
    /// inline on a tokio worker.
    #[test]
    fn let_subquery_takes_the_blocking_pool() {
        let q = parse("LET xs = (FOR d IN docs RETURN d) RETURN LENGTH(xs)").unwrap();
        assert!(is_long_running_query(&q));
        let q = parse("RETURN (FOR d IN docs RETURN d)").unwrap();
        assert!(is_long_running_query(&q));
    }

    #[test]
    fn heavy_functions_take_the_blocking_pool() {
        for q in [
            r#"RETURN EMBED("hello")"#,
            r#"RETURN NEIGHBORS("edges", ["a/1"])"#,
            r#"RETURN {p: PAGERANK("edges")}"#,
        ] {
            assert!(is_long_running_query(&parse(q).unwrap()), "{q}");
        }
    }

    /// A lone FOR over a small local collection skips the blocking pool;
    /// anything that could grow the work does not.
    #[test]
    fn small_local_scans_run_inline() {
        use super::{is_small_local_scan, INLINE_SCAN_MAX_DOCS};
        use crate::storage::StorageEngine;
        let dir = tempfile::tempdir().unwrap();
        let storage = StorageEngine::new(dir.path().to_str().unwrap()).unwrap();
        storage.initialize().unwrap();
        storage.create_database("d".into()).unwrap();
        let db = storage.get_database("d").unwrap();
        db.create_collection("small".into(), None).unwrap();
        db.create_collection("big".into(), None).unwrap();
        let small = db.get_collection("small").unwrap();
        for i in 0..3 {
            small.insert(serde_json::json!({"n": i})).unwrap();
        }
        let big = db.get_collection("big").unwrap();
        for i in 0..=INLINE_SCAN_MAX_DOCS {
            big.insert(serde_json::json!({"n": i})).unwrap();
        }
        let inline = |q: &str| is_small_local_scan(&parse(q).unwrap(), &storage, "d");

        assert!(inline("FOR x IN small RETURN x"));
        assert!(inline(
            "FOR x IN small FILTER x.n > 1 SORT x.n LIMIT 2 RETURN {n: x.n}"
        ));
        assert!(!inline("FOR x IN big RETURN x"));
        assert!(!inline("FOR x IN missing RETURN x"));
        assert!(!inline("FOR x IN small FOR y IN small RETURN [x, y]"));
        assert!(!inline("FOR x IN small RETURN (FOR y IN big RETURN y)"));
        assert!(!inline("FOR x IN small UPDATE x WITH {m: 1} IN small"));
        assert!(!inline("FOR x IN 1..3 RETURN x"));
    }
}