Skip to main content

mongreldb_server/
lib.rs

1//! mongreldb-server — a long-lived process holding a multi-table `Database`
2//! open, serving SQL + table-qualified native APIs over HTTP.
3//!
4//! Endpoints:
5//!   GET    /health                    → 200 OK
6//!   GET    /tables                    → ["t1", "t2", ...]
7//!   POST   /tables                    → create table
8//!   DELETE /tables/{name}              → drop table
9//!   POST   /tables/{name}/put          → upsert one row
10//!   POST   /tables/{name}/count        → { "count": N }
11//!   POST   /tables/{name}/commit       → { "epoch": N, "epoch_text": "N" }
12//!   POST   /sql                       → Arrow IPC bytes
13//!   POST   /txn                       → exact atomic commit receipt
14//!
15//! Usage: `mongreldb-server <db_dir> [port]`
16
17use std::sync::atomic::{AtomicBool, Ordering};
18use std::sync::{Arc, Mutex};
19
20use axum::extract::{Path, State};
21use axum::http::header;
22use axum::http::StatusCode;
23use axum::response::{IntoResponse, Response};
24use axum::routing::{get, post};
25use axum::Json;
26use mongreldb_core::schema::{Schema, TypeId};
27use mongreldb_core::{CancellationReason, Database, Value};
28use mongreldb_query::{
29    CancelOutcome, ExternalTableModule, ManagedQueryBatches, MongrelSession, QueryId,
30    RegisteredQueryGuard, RegisteredSqlQuery, SqlQueryOptions, SqlQueryPhase, SqlQueryRegistry,
31    SqlStreamCompletion,
32};
33use serde::{Deserialize, Serialize};
34use serde_json::json;
35use sha2::Digest;
36use zeroize::Zeroizing;
37
38mod audit;
39mod kit;
40mod metrics;
41mod pre_cancel;
42mod procedure;
43mod sessions;
44mod sql_idempotency;
45mod sql_pages;
46mod trigger;
47
48pub use sessions::{spawn_session_reaper, SessionStore};
49
50fn client_closed_request_status() -> StatusCode {
51    StatusCode::from_u16(499).unwrap_or(StatusCode::BAD_REQUEST)
52}
53
54fn cancellation_checkpoint_error(query: &RegisteredSqlQuery) -> mongreldb_query::MongrelQueryError {
55    query.checkpoint().err().unwrap_or_else(|| {
56        mongreldb_query::MongrelQueryError::InvalidQueryState(
57            "cancellation notification observed without a terminal checkpoint".into(),
58        )
59    })
60}
61
62/// Map an engine error to the appropriate HTTP status code for defense-in-depth.
63/// Auth errors get 401/403; everything else stays 500. This ensures that even
64/// after the HTTP auth middleware lets a request through, the storage layer's
65/// permission checks surface as the right status (not a generic 500).
66fn status_for_error(e: &mongreldb_core::MongrelError) -> StatusCode {
67    use mongreldb_core::MongrelError;
68    match e {
69        MongrelError::AuthRequired | MongrelError::InvalidCredentials { .. } => {
70            StatusCode::UNAUTHORIZED
71        }
72        MongrelError::AuthNotRequired => StatusCode::BAD_REQUEST,
73        MongrelError::PermissionDenied { .. } => StatusCode::FORBIDDEN,
74        MongrelError::InvalidArgument(_) => StatusCode::CONFLICT,
75        MongrelError::Conflict(_) => StatusCode::CONFLICT,
76        MongrelError::ReadOnlyReplica => StatusCode::CONFLICT,
77        MongrelError::NotFound(_) => StatusCode::NOT_FOUND,
78        MongrelError::DeadlineExceeded => StatusCode::GATEWAY_TIMEOUT,
79        MongrelError::WorkBudgetExceeded => StatusCode::TOO_MANY_REQUESTS,
80        MongrelError::Cancelled => client_closed_request_status(),
81        MongrelError::CursorStale(_) => StatusCode::CONFLICT,
82        MongrelError::CursorExpired => StatusCode::GONE,
83        _ => StatusCode::INTERNAL_SERVER_ERROR,
84    }
85}
86
87/// Map a query-layer error (which wraps engine errors via `Core(...)`) to the
88/// appropriate HTTP status code.
89fn status_for_query_error(e: &mongreldb_query::MongrelQueryError) -> StatusCode {
90    use mongreldb_query::MongrelQueryError;
91    match e {
92        MongrelQueryError::Core(core) => status_for_error(core),
93        MongrelQueryError::DeadlineExceeded { .. } => StatusCode::GATEWAY_TIMEOUT,
94        MongrelQueryError::QueryCancelled { .. } => client_closed_request_status(),
95        MongrelQueryError::QueryIdConflict { .. } => StatusCode::CONFLICT,
96        MongrelQueryError::QueryRegistryFull => StatusCode::SERVICE_UNAVAILABLE,
97        MongrelQueryError::ResultLimitExceeded { .. } => StatusCode::PAYLOAD_TOO_LARGE,
98        MongrelQueryError::TransactionAborted => StatusCode::CONFLICT,
99        MongrelQueryError::NoSqlTransaction | MongrelQueryError::SavepointNotFound { .. } => {
100            StatusCode::CONFLICT
101        }
102        MongrelQueryError::CommitOutcome { .. } => StatusCode::CONFLICT,
103        MongrelQueryError::OutcomeUnknown { .. } => StatusCode::CONFLICT,
104        _ => StatusCode::INTERNAL_SERVER_ERROR,
105    }
106}
107
108/// Extractor that pulls the authenticated [`mongreldb_core::Principal`] (if the
109/// auth middleware injected one) from request extensions without erroring when
110/// absent (e.g. token-authenticated requests carry no `Principal`).
111struct OptionalPrincipal(Option<mongreldb_core::Principal>);
112
113impl<S> axum::extract::FromRequestParts<S> for OptionalPrincipal
114where
115    S: Send + Sync,
116{
117    type Rejection = std::convert::Infallible;
118
119    async fn from_request_parts(
120        parts: &mut axum::http::request::Parts,
121        _state: &S,
122    ) -> Result<Self, Self::Rejection> {
123        Ok(OptionalPrincipal(
124            parts.extensions.get::<mongreldb_core::Principal>().cloned(),
125        ))
126    }
127}
128
129struct AppState {
130    db: Arc<Database>,
131    idem: kit::IdempotencyStore,
132    external_modules: Vec<Arc<dyn ExternalTableModule>>,
133    auth_token: Option<String>,
134    /// When true, authenticate via catalog users (HTTP Basic auth).
135    user_auth: bool,
136    /// Daemon-wide Prometheus-style counters, shared by all handlers.
137    metrics: Arc<metrics::Metrics>,
138    /// `/sql` requests slower than this are logged as slow queries.
139    slow_query_threshold: std::time::Duration,
140    /// Bounded security audit log (auth + DDL/privilege events).
141    audit: Arc<audit::AuditLog>,
142    /// Token-keyed pool of live sessions for cross-request interactive
143    /// transactions (`X-Session-ID` on `/sql`).
144    sessions: Arc<sessions::SessionStore>,
145    /// Bounds CPU-heavy AI work submitted to Tokio's blocking pool.
146    ai_semaphore: Arc<tokio::sync::Semaphore>,
147    /// Process-wide SQL registry. Cancellation never takes a session lock.
148    query_registry: Arc<SqlQueryRegistry>,
149    /// Serializes registration with cancel-before-registration bookkeeping.
150    query_lifecycle: Mutex<()>,
151    /// Short-lived, owner/session-bound cancellations received before SQL.
152    pre_cancellations: pre_cancel::PreCancelStore,
153    /// Owner- and request-bound durable SQL write receipts.
154    sql_idempotency: Arc<sql_idempotency::SqlIdempotencyStore>,
155    /// Bounded projected results retained for stable SQL continuation cursors.
156    sql_pages: sql_pages::SqlPageStore,
157    /// Admission control for ordinary SQL, separate from AI workers.
158    sql_semaphore: Arc<tokio::sync::Semaphore>,
159    sql_default_timeout: std::time::Duration,
160    sql_max_timeout: std::time::Duration,
161    sql_cancel_grace: std::time::Duration,
162    sql_max_output_bytes: usize,
163    sql_max_output_rows: usize,
164    accepting_sql: Arc<AtomicBool>,
165    /// Lazily generated process-local HMAC key. Restart invalidates cursors.
166    cursor_mac_key: CursorMacKey,
167}
168
169#[derive(Default)]
170struct CursorMacKey {
171    key: Mutex<Option<[u8; 32]>>,
172}
173
174impl CursorMacKey {
175    fn get(&self) -> mongreldb_core::Result<[u8; 32]> {
176        let mut key = match self.key.lock() {
177            Ok(key) => key,
178            Err(poisoned) => poisoned.into_inner(),
179        };
180        if let Some(key) = *key {
181            return Ok(key);
182        }
183        let mut generated = [0u8; 32];
184        mongreldb_core::encryption::fill_random(&mut generated)?;
185        *key = Some(generated);
186        Ok(generated)
187    }
188}
189
190/// Handle retained by the daemon to coordinate graceful SQL shutdown without
191/// coupling signal handling to Axum's router internals.
192#[derive(Clone)]
193pub struct ServerControl {
194    query_registry: Arc<SqlQueryRegistry>,
195    sessions: Arc<sessions::SessionStore>,
196    accepting_sql: Arc<AtomicBool>,
197    cancel_grace: std::time::Duration,
198    metrics: Arc<metrics::Metrics>,
199}
200
201impl ServerControl {
202    pub async fn shutdown(&self) -> usize {
203        self.accepting_sql.store(false, Ordering::Release);
204        self.query_registry
205            .cancel_all(CancellationReason::ServerShutdown);
206        let deadline = tokio::time::Instant::now() + self.cancel_grace;
207        while self.query_registry.active_count() > 0 && tokio::time::Instant::now() < deadline {
208            tokio::time::sleep(std::time::Duration::from_millis(5)).await;
209        }
210        self.sessions.close_all();
211        let stuck_queries = self.query_registry.active_statuses();
212        for status in &stuck_queries {
213            eprintln!(
214                "[sql-cancel-stuck] query_id={} phase={}",
215                status.query_id,
216                query_phase_name(status.phase)
217            );
218        }
219        let stuck = stuck_queries.len();
220        self.metrics.add_sql_stuck_after_cancel(stuck);
221        stuck
222    }
223}
224
225pub fn build_app(db: Arc<Database>) -> axum::Router {
226    build_app_with_config(
227        db,
228        std::iter::empty::<Arc<dyn ExternalTableModule>>(),
229        None,
230        None,
231    )
232}
233
234pub fn build_app_with_external_modules(
235    db: Arc<Database>,
236    external_modules: impl IntoIterator<Item = Arc<dyn ExternalTableModule>>,
237) -> axum::Router {
238    build_app_with_config(db, external_modules, None, None)
239}
240
241/// Build the daemon router with optional auth token and max-connections limit.
242pub fn build_app_with_config(
243    db: Arc<Database>,
244    external_modules: impl IntoIterator<Item = Arc<dyn ExternalTableModule>>,
245    auth_token: Option<String>,
246    max_connections: Option<usize>,
247) -> axum::Router {
248    build_app_full(db, external_modules, auth_token, max_connections, false)
249}
250
251/// Build the daemon router with full auth configuration including user-based auth.
252/// Sessions are enabled with a default capacity (256) and idle timeout (300 s).
253pub fn build_app_full(
254    db: Arc<Database>,
255    external_modules: impl IntoIterator<Item = Arc<dyn ExternalTableModule>>,
256    auth_token: Option<String>,
257    max_connections: Option<usize>,
258    user_auth: bool,
259) -> axum::Router {
260    let sessions = Arc::new(sessions::SessionStore::new(
261        default_max_sessions(),
262        default_session_idle_timeout(),
263    ));
264    build_app_with_sessions(
265        db,
266        external_modules,
267        auth_token,
268        max_connections,
269        user_auth,
270        sessions,
271    )
272}
273
274/// Build the daemon router with an explicit, externally-owned session store.
275/// The caller (typically `main`) keeps the `Arc<SessionStore>` so it can spawn
276/// the idle reaper against the same map the handlers use.
277pub fn build_app_with_sessions(
278    db: Arc<Database>,
279    external_modules: impl IntoIterator<Item = Arc<dyn ExternalTableModule>>,
280    auth_token: Option<String>,
281    max_connections: Option<usize>,
282    user_auth: bool,
283    sessions: Arc<sessions::SessionStore>,
284) -> axum::Router {
285    build_app_with_sessions_and_control(
286        db,
287        external_modules,
288        auth_token,
289        max_connections,
290        user_auth,
291        sessions,
292    )
293    .0
294}
295
296/// Build the daemon router and return a handle for graceful query shutdown.
297pub fn build_app_with_sessions_and_control(
298    db: Arc<Database>,
299    external_modules: impl IntoIterator<Item = Arc<dyn ExternalTableModule>>,
300    auth_token: Option<String>,
301    max_connections: Option<usize>,
302    user_auth: bool,
303    sessions: Arc<sessions::SessionStore>,
304) -> (axum::Router, ServerControl) {
305    db.set_replication_wal_retention_segments(default_replication_wal_segments());
306    if let Err(error) = db.set_history_retention_epochs(default_history_retention_epochs()) {
307        eprintln!("[history] failed to configure retention: {error}");
308    }
309    let max_active_queries = default_sql_max_active_queries();
310    let query_registry = Arc::new(SqlQueryRegistry::new(
311        max_active_queries,
312        max_active_queries.saturating_mul(2),
313        2 * 1024 * 1024,
314        default_sql_finished_query_ttl(),
315    ));
316    let accepting_sql = Arc::new(AtomicBool::new(true));
317    let sql_cancel_grace = default_sql_cancel_grace();
318    let sql_max_timeout = default_sql_max_timeout();
319    let sql_default_timeout = default_sql_default_timeout().min(sql_max_timeout);
320    let metrics = Arc::new(metrics::Metrics::default());
321    let (idempotency_root, idempotency_integrity) =
322        match sql_idempotency::IdempotencyIntegrity::for_database(&db) {
323            Ok((root, integrity)) => (root, Some(integrity)),
324            Err(error) => {
325                eprintln!("[idempotency] durable integrity key unavailable: {error}");
326                (db.durable_root(), None)
327            }
328        };
329    let sql_idempotency = Arc::new(sql_idempotency::SqlIdempotencyStore::new_with_integrity(
330        Arc::clone(&idempotency_root),
331        idempotency_integrity.clone(),
332        default_sql_idempotency_ttl(),
333        default_sql_idempotency_max_entries(),
334    ));
335    let server_control = ServerControl {
336        query_registry: Arc::clone(&query_registry),
337        sessions: Arc::clone(&sessions),
338        accepting_sql: Arc::clone(&accepting_sql),
339        cancel_grace: sql_cancel_grace,
340        metrics: Arc::clone(&metrics),
341    };
342    let state = Arc::new(AppState {
343        idem: kit::IdempotencyStore::new_with_integrity(
344            idempotency_root,
345            idempotency_integrity,
346            default_sql_idempotency_ttl(),
347            default_sql_idempotency_max_entries(),
348        ),
349        db,
350        external_modules: external_modules.into_iter().collect(),
351        auth_token,
352        user_auth,
353        metrics,
354        slow_query_threshold: metrics::slow_query_threshold(),
355        audit: Arc::new(audit::AuditLog::new(8192)),
356        sessions,
357        ai_semaphore: Arc::new(tokio::sync::Semaphore::new(default_ai_max_concurrent())),
358        query_registry,
359        query_lifecycle: Mutex::new(()),
360        pre_cancellations: pre_cancel::PreCancelStore::new(
361            default_sql_pre_cancel_ttl(),
362            default_sql_pre_cancel_max_entries(),
363            default_sql_pre_cancel_max_bytes(),
364            default_sql_pre_cancel_max_entries_per_owner(),
365            default_sql_pre_cancel_rate_window(),
366            default_sql_pre_cancel_rate_per_owner(),
367        ),
368        sql_idempotency,
369        sql_pages: sql_pages::SqlPageStore::new(
370            default_sql_page_ttl(),
371            default_sql_page_max_entries(),
372            default_sql_page_max_bytes(),
373            default_sql_page_max_entries_per_owner(),
374        ),
375        sql_semaphore: Arc::new(tokio::sync::Semaphore::new(default_sql_max_concurrent())),
376        sql_default_timeout,
377        sql_max_timeout,
378        sql_cancel_grace,
379        sql_max_output_bytes: default_sql_max_output_bytes(),
380        sql_max_output_rows: default_sql_max_output_rows(),
381        accepting_sql,
382        cursor_mac_key: CursorMacKey::default(),
383    });
384    let router = axum::Router::new()
385        .route("/health", get(health))
386        .route("/capabilities", get(capabilities))
387        .route(
388            "/history/retention",
389            get(history_retention).put(set_history_retention),
390        )
391        .route("/metrics", get(metrics_handler))
392        .route("/audit", get(audit_handler))
393        .route("/tables", get(list_tables).post(create_table))
394        .route("/tables/{name}", axum::routing::delete(drop_table))
395        .route("/tables/{name}/put", post(put_row))
396        .route("/tables/{name}/count", get(count))
397        .route("/tables/{name}/commit", post(commit))
398        .route("/sql", post(sql))
399        .route("/sql/continue", post(continue_sql_page))
400        .route("/queries/{query_id}", get(query_status))
401        .route("/queries/{query_id}/cancel", post(cancel_query))
402        .route("/txn", post(txn))
403        .route("/sessions", post(create_session))
404        .route("/sessions/{id}", axum::routing::delete(close_session))
405        .route("/sessions/{id}/prepare", post(prepare_statement))
406        .route("/sessions/{id}/execute", post(execute_statement))
407        .route(
408            "/sessions/{id}/statements/{name}",
409            axum::routing::delete(deallocate_statement),
410        )
411        .route("/procedures", get(procedure::list).post(procedure::create))
412        .route(
413            "/procedures/{name}",
414            get(procedure::describe)
415                .put(procedure::replace)
416                .delete(procedure::drop_procedure),
417        )
418        .route("/procedures/{name}/call", post(procedure::call))
419        .route("/triggers", get(trigger::list).post(trigger::create))
420        .route(
421            "/triggers/{name}",
422            get(trigger::describe)
423                .put(trigger::replace)
424                .delete(trigger::drop_trigger),
425        )
426        // Typed Kit-aware surface (authoritative validation + constraints).
427        .route("/kit/schema", get(kit::schema_all))
428        .route("/kit/schema/{table}", get(kit::schema_one))
429        .route("/kit/txn", post(kit::kit_txn))
430        .route("/kit/query", post(kit::kit_query))
431        .route("/kit/retrieve", post(kit::kit_retrieve))
432        .route("/kit/ann_rerank", post(kit::kit_ann_rerank))
433        .route("/kit/ai/metrics", get(kit::kit_ai_metrics))
434        .route("/kit/set_similarity", post(kit::kit_set_similarity))
435        .route("/kit/search", post(kit::kit_search))
436        .route("/kit/create_table", post(kit::kit_create_table))
437        .route("/kit/procedures/{name}/call", post(procedure::kit_call))
438        .route("/compact", post(compact_all))
439        .route("/tables/{name}/compact", post(compact_table))
440        .route("/wal/stream", get(wal_stream))
441        .route("/replication/snapshot", get(replication_snapshot))
442        .route("/events", get(events_stream))
443        .with_state(state.clone());
444
445    // A credential-enforced database must never expose the authenticated
446    // daemon handle when the caller forgot to configure an HTTP auth mode.
447    // With neither mode enabled the middleware rejects every route.
448    let router = if state.auth_token.is_some() || state.user_auth || state.db.require_auth_enabled()
449    {
450        router.layer(axum::middleware::from_fn_with_state(
451            state.clone(),
452            auth_middleware,
453        ))
454    } else {
455        router
456    };
457
458    // Apply connection limit if configured.
459    let router = if let Some(max) = max_connections {
460        router.layer(tower::limit::ConcurrencyLimitLayer::new(max))
461    } else {
462        router
463    };
464    (router, server_control)
465}
466
467/// Auth middleware supporting three modes:
468/// 1. **Token** (`--auth-token <token>`): checks `Authorization: Bearer <token>`.
469/// 2. **User auth** (`--auth-users`): checks `Authorization: Basic <base64(user:pass)>`
470///    against catalog users (Argon2id-verified). Injects a `Principal` into
471///    request extensions.
472/// 3. **Both**: token OR valid user credentials accepted.
473async fn auth_middleware(
474    axum::extract::State(state): axum::extract::State<Arc<AppState>>,
475    mut req: axum::extract::Request,
476    next: axum::middleware::Next,
477) -> Result<axum::response::Response, axum::http::StatusCode> {
478    let header = req
479        .headers()
480        .get("authorization")
481        .and_then(|v| v.to_str().ok())
482        .unwrap_or("");
483
484    // Track the attempted identity + failure reason so EVERY 401 path emits
485    // exactly one `login.fail` audit event (missing, malformed, wrong token,
486    // wrong password are all logged — no unauthenticated probe goes unrecorded).
487    let mut attempted = String::new();
488    let mut fail_reason = "no credentials provided".to_string();
489
490    // Mode 1: Token auth (Bearer).
491    if let Some(token) = &state.auth_token {
492        if let Some(provided) = header.strip_prefix("Bearer ") {
493            attempted = "token".to_string();
494            if provided == token {
495                state
496                    .audit
497                    .record("token", "login.ok", "bearer token accepted");
498                return Ok(next.run(req).await);
499            }
500            fail_reason = "invalid bearer token".to_string();
501        }
502    }
503
504    // Mode 2: User auth (Basic).
505    if state.user_auth {
506        if let Some(encoded) = header.strip_prefix("Basic ") {
507            if let Ok(decoded) = base64_decode(encoded) {
508                let decoded = Zeroizing::new(decoded);
509                if let Ok(creds) = std::str::from_utf8(&decoded) {
510                    if let Some((username, password)) = creds.split_once(':') {
511                        attempted = username.to_string();
512                        if let Ok(Some(principal)) =
513                            state.db.authenticate_principal(username, password)
514                        {
515                            let username = principal.username.clone();
516                            drop(decoded);
517                            state
518                                .audit
519                                .record(&username, "login.ok", "basic credentials accepted");
520                            req.extensions_mut().insert(principal);
521                            return Ok(next.run(req).await);
522                        }
523                        fail_reason = "invalid basic credentials".to_string();
524                    } else {
525                        fail_reason = "malformed basic credentials (no ':')".to_string();
526                    }
527                } else {
528                    fail_reason = "malformed basic credentials (non-utf8)".to_string();
529                }
530            } else {
531                fail_reason = "malformed basic credentials (bad base64)".to_string();
532            }
533        }
534    }
535
536    let who = if attempted.is_empty() {
537        "anonymous"
538    } else {
539        attempted.as_str()
540    };
541    state.audit.record(who, "login.fail", fail_reason);
542    Err(axum::http::StatusCode::UNAUTHORIZED)
543}
544
545/// Minimal Base64 decoder (no extra dep).
546fn base64_decode(input: &str) -> Result<Vec<u8>, ()> {
547    const TABLE: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
548    let input: Vec<u8> = input
549        .bytes()
550        .filter(|&b| b != b'\n' && b != b'\r' && b != b' ')
551        .collect();
552    let mut out = Vec::with_capacity(input.len() * 3 / 4);
553    let mut buf = 0u32;
554    let mut bits = 0u32;
555    for &b in &input {
556        if b == b'=' {
557            break;
558        }
559        let val = TABLE.iter().position(|&t| t == b).ok_or(())? as u32;
560        buf = (buf << 6) | val;
561        bits += 6;
562        if bits >= 8 {
563            bits -= 8;
564            out.push((buf >> bits) as u8);
565        }
566    }
567    Ok(out)
568}
569
570/// `GET /wal/stream?since=<epoch>` — return complete committed WAL
571/// transactions after the follower epoch. A 409 response means retained WAL
572/// cannot close the gap and the follower must fetch `/replication/snapshot`.
573/// Records are newline-delimited JSON.
574/// newline-delimited JSON for replication followers. Each line is a JSON
575/// object `{ "seq": N, "txn_id": N, "op": {...} }`.
576async fn wal_stream(
577    axum::extract::State(state): axum::extract::State<Arc<AppState>>,
578    OptionalPrincipal(principal): OptionalPrincipal,
579    axum::extract::Query(params): axum::extract::Query<WalStreamParams>,
580) -> Result<Response, StatusCode> {
581    state
582        .db
583        .require_for(
584            request_principal(&state, &principal).as_ref(),
585            &mongreldb_core::Permission::Admin,
586        )
587        .map_err(|error| status_for_error(&error))?;
588    let since = params.since.unwrap_or(0);
589    let db = Arc::clone(&state.db);
590    let batch = tokio::task::spawn_blocking(move || db.replication_batch_since(since))
591        .await
592        .map_err(|_e| StatusCode::INTERNAL_SERVER_ERROR)?;
593    let batch = batch.map_err(|e| {
594        eprintln!("wal_stream error: {e}");
595        StatusCode::INTERNAL_SERVER_ERROR
596    })?;
597    if batch.requires_snapshot {
598        let mut response = (
599            StatusCode::CONFLICT,
600            "replication snapshot required: WAL retention gap or spilled run",
601        )
602            .into_response();
603        response.headers_mut().insert(
604            "x-mongreldb-replication-status",
605            "snapshot-required"
606                .parse()
607                .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?,
608        );
609        set_replication_headers(
610            &mut response,
611            &batch.source_id,
612            batch.from_epoch,
613            batch.current_epoch,
614            batch.earliest_epoch,
615            batch.commit_count,
616            &batch.records_sha256,
617        )?;
618        return Ok(response);
619    }
620    let mut body = String::new();
621    for record in &batch.records {
622        let json = serde_json::to_string(record).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
623        body.push_str(&json);
624        body.push('\n');
625    }
626    let mut response = (
627        [
628            (header::CONTENT_TYPE, "application/x-ndjson".to_string()),
629            (header::CACHE_CONTROL, "no-cache".to_string()),
630        ],
631        body,
632    )
633        .into_response();
634    set_replication_headers(
635        &mut response,
636        &batch.source_id,
637        batch.from_epoch,
638        batch.current_epoch,
639        batch.earliest_epoch,
640        batch.commit_count,
641        &batch.records_sha256,
642    )?;
643    Ok(response)
644}
645
646async fn replication_snapshot(
647    axum::extract::State(state): axum::extract::State<Arc<AppState>>,
648    OptionalPrincipal(principal): OptionalPrincipal,
649) -> Response {
650    if let Err(error) = state.db.require_for(
651        request_principal(&state, &principal).as_ref(),
652        &mongreldb_core::Permission::Admin,
653    ) {
654        return (status_for_error(&error), error.to_string()).into_response();
655    }
656    let db = Arc::clone(&state.db);
657    let snapshot = match tokio::task::spawn_blocking(move || db.replication_snapshot()).await {
658        Ok(Ok(snapshot)) => snapshot,
659        Ok(Err(error)) => {
660            return (StatusCode::INTERNAL_SERVER_ERROR, error.to_string()).into_response()
661        }
662        Err(error) => {
663            return (StatusCode::INTERNAL_SERVER_ERROR, error.to_string()).into_response()
664        }
665    };
666    let epoch = snapshot.epoch();
667    // ponytail: bootstrap buffers one image; add framed file streaming when
668    // real snapshot sizes make this memory ceiling measurable.
669    match snapshot.encode() {
670        Ok(bytes) => {
671            let mut response =
672                ([(header::CONTENT_TYPE, "application/octet-stream")], bytes).into_response();
673            let Ok(value) = epoch.to_string().parse() else {
674                return (
675                    StatusCode::INTERNAL_SERVER_ERROR,
676                    "invalid replication epoch response header",
677                )
678                    .into_response();
679            };
680            response
681                .headers_mut()
682                .insert("x-mongreldb-current-epoch", value);
683            let source_id = snapshot
684                .source_id()
685                .iter()
686                .map(|byte| format!("{byte:02x}"))
687                .collect::<String>();
688            let Ok(source_id) = source_id.parse() else {
689                return (
690                    StatusCode::INTERNAL_SERVER_ERROR,
691                    "invalid replication source response header",
692                )
693                    .into_response();
694            };
695            response
696                .headers_mut()
697                .insert("x-mongreldb-source-id", source_id);
698            response
699        }
700        Err(error) => (StatusCode::INTERNAL_SERVER_ERROR, error.to_string()).into_response(),
701    }
702}
703
704fn set_replication_headers(
705    response: &mut Response,
706    source_id: &[u8; 32],
707    from: u64,
708    current: u64,
709    earliest: Option<u64>,
710    commit_count: u64,
711    records_sha256: &[u8; 32],
712) -> Result<(), StatusCode> {
713    let source_id = source_id
714        .iter()
715        .map(|byte| format!("{byte:02x}"))
716        .collect::<String>();
717    response.headers_mut().insert(
718        "x-mongreldb-source-id",
719        source_id
720            .parse()
721            .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?,
722    );
723    response.headers_mut().insert(
724        "x-mongreldb-from-epoch",
725        from.to_string()
726            .parse()
727            .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?,
728    );
729    response.headers_mut().insert(
730        "x-mongreldb-current-epoch",
731        current
732            .to_string()
733            .parse()
734            .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?,
735    );
736    response.headers_mut().insert(
737        "x-mongreldb-commit-count",
738        commit_count
739            .to_string()
740            .parse()
741            .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?,
742    );
743    let digest = records_sha256
744        .iter()
745        .map(|byte| format!("{byte:02x}"))
746        .collect::<String>();
747    response.headers_mut().insert(
748        "x-mongreldb-records-sha256",
749        digest
750            .parse()
751            .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?,
752    );
753    if let Some(earliest) = earliest {
754        response.headers_mut().insert(
755            "x-mongreldb-earliest-epoch",
756            earliest
757                .to_string()
758                .parse()
759                .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?,
760        );
761    }
762    Ok(())
763}
764
765#[derive(serde::Deserialize)]
766struct WalStreamParams {
767    since: Option<u64>,
768}
769
770/// `GET /events` — long-lived SSE for durable WAL-backed row changes plus
771/// ephemeral SQL NOTIFY messages. `Last-Event-ID` resumes from a stable
772/// `<commit_epoch>:<operation_index>` id. A retention gap returns 409 before
773/// the stream starts, or a terminal `gap` SSE if the client falls behind.
774async fn events_stream(
775    axum::extract::State(state): axum::extract::State<Arc<AppState>>,
776    OptionalPrincipal(principal): OptionalPrincipal,
777    headers: axum::http::HeaderMap,
778) -> Result<Response, StatusCode> {
779    use axum::response::sse::{Event, KeepAlive, Sse};
780    use futures::Stream;
781    use std::collections::VecDeque;
782    use std::convert::Infallible;
783
784    state
785        .db
786        .require_for(
787            request_principal(&state, &principal).as_ref(),
788            &mongreldb_core::Permission::Admin,
789        )
790        .map_err(|error| status_for_error(&error))?;
791
792    struct State {
793        db: Arc<Database>,
794        receiver: tokio::sync::broadcast::Receiver<mongreldb_core::ChangeEvent>,
795        change_wake: tokio::sync::broadcast::Receiver<()>,
796        interval: tokio::time::Interval,
797        pending: VecDeque<mongreldb_core::ChangeEvent>,
798        last_id: Option<String>,
799        poll_now: bool,
800        done: bool,
801    }
802
803    fn event(change: mongreldb_core::ChangeEvent) -> Event {
804        let id = change.id.clone();
805        let kind = if change.op == "notify" {
806            "notify"
807        } else {
808            "change"
809        };
810        let mut event = Event::default().event(kind).data(
811            serde_json::to_string(&change)
812                .unwrap_or_else(|error| format!(r#"{{"error":"{error}"}}"#)),
813        );
814        if let Some(id) = id {
815            event = event.id(id);
816        }
817        event
818    }
819
820    let last_id = match headers.get("last-event-id") {
821        Some(value) => Some(
822            value
823                .to_str()
824                .map_err(|_| StatusCode::BAD_REQUEST)?
825                .to_owned(),
826        ),
827        None => None,
828    };
829    let receiver = state.db.subscribe_changes();
830    let change_wake = state.db.subscribe_change_commits();
831    let db = Arc::clone(&state.db);
832    let resume = last_id.clone();
833    let initial = tokio::task::spawn_blocking(move || db.change_events_since(resume.as_deref()))
834        .await
835        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
836        .map_err(|error| match error {
837            mongreldb_core::MongrelError::InvalidArgument(_) => StatusCode::BAD_REQUEST,
838            _ => StatusCode::INTERNAL_SERVER_ERROR,
839        })?;
840    if initial.gap {
841        return Ok((
842            StatusCode::CONFLICT,
843            Json(json!({
844                "error": "cdc retention gap",
845                "earliest_epoch": initial.earliest_epoch,
846                "current_epoch": initial.current_epoch,
847            })),
848        )
849            .into_response());
850    }
851
852    let stream: std::pin::Pin<Box<dyn Stream<Item = Result<Event, Infallible>> + Send>> = Box::pin(
853        futures::stream::unfold(
854            State {
855                db: Arc::clone(&state.db),
856                receiver,
857                change_wake,
858                interval: tokio::time::interval(std::time::Duration::from_millis(250)),
859                pending: initial.events.into(),
860                last_id,
861                poll_now: false,
862                done: false,
863            },
864            |mut stream| async move {
865                if stream.done {
866                    return None;
867                }
868                loop {
869                    if stream.poll_now {
870                        stream.poll_now = false;
871                        let db = Arc::clone(&stream.db);
872                        let last_id = stream.last_id.clone();
873                        match tokio::task::spawn_blocking(move || {
874                            db.change_events_since(last_id.as_deref())
875                        })
876                        .await
877                        {
878                            Ok(Ok(batch)) if batch.gap => {
879                                stream.done = true;
880                                let gap = Event::default().event("gap").data(
881                                    json!({
882                                        "error": "cdc retention gap",
883                                        "earliest_epoch": batch.earliest_epoch,
884                                        "current_epoch": batch.current_epoch,
885                                    })
886                                    .to_string(),
887                                );
888                                return Some((Ok(gap), stream));
889                            }
890                            Ok(Ok(batch)) => stream.pending.extend(batch.events),
891                            Ok(Err(error)) => {
892                                stream.done = true;
893                                return Some((
894                                    Ok(Event::default().event("error").data(error.to_string())),
895                                    stream,
896                                ));
897                            }
898                            Err(error) => {
899                                stream.done = true;
900                                return Some((
901                                    Ok(Event::default().event("error").data(error.to_string())),
902                                    stream,
903                                ));
904                            }
905                        }
906                    }
907                    if let Some(change) = stream.pending.pop_front() {
908                        if let Some(id) = &change.id {
909                            stream.last_id = Some(id.clone());
910                        }
911                        return Some((Ok(event(change)), stream));
912                    }
913                    tokio::select! {
914                        received = stream.receiver.recv() => {
915                            match received {
916                                Ok(change) => return Some((Ok(event(change)), stream)),
917                                Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {},
918                                Err(tokio::sync::broadcast::error::RecvError::Closed) => {
919                                    return None;
920                                }
921                            }
922                        }
923                        received = stream.change_wake.recv() => {
924                            match received {
925                                Ok(()) | Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {
926                                    stream.poll_now = true;
927                                }
928                                Err(tokio::sync::broadcast::error::RecvError::Closed) => {}
929                            }
930                        }
931                        _ = stream.interval.tick() => {
932                            stream.poll_now = true;
933                        }
934                    }
935                }
936            },
937        ),
938    );
939
940    Ok(Sse::new(stream)
941        .keep_alive(
942            KeepAlive::new()
943                .interval(std::time::Duration::from_secs(15))
944                .text("keep-alive"),
945        )
946        .into_response())
947}
948
949/// Launch the §5.9 background auto-compaction sweep (run-count cost trigger).
950/// One OS thread, sleeping `interval` between sweeps; each tick locks each
951/// table individually and calls `Table::maybe_compact`. Best-effort: a
952/// compaction error is logged and never aborts the sweep.
953pub fn spawn_auto_compactor(db: Arc<Database>) {
954    if let Err(error) = std::thread::Builder::new()
955        .name("mongreldb-auto-compact".into())
956        .spawn(move || loop {
957            std::thread::sleep(std::time::Duration::from_secs(30));
958            for name in db.table_names() {
959                let Ok(handle) = db.table(&name) else {
960                    continue;
961                };
962                let mut t = handle.lock();
963                let before = t.run_count();
964                match t.maybe_compact() {
965                    Ok(true) => {
966                        eprintln!(
967                            "[auto-compact] {name}: {} runs -> {}",
968                            before,
969                            t.run_count()
970                        );
971                    }
972                    Ok(false) => {}
973                    Err(e) => {
974                        eprintln!("[auto-compact] {name}: compaction failed: {e}");
975                    }
976                }
977            }
978        })
979    {
980        eprintln!("[auto-compact] failed to start background thread: {error}");
981    }
982}
983
984async fn health() -> StatusCode {
985    StatusCode::OK
986}
987
988#[derive(Debug, Serialize)]
989struct SqlCancellationCapabilities {
990    version: u8,
991    client_query_ids: bool,
992    cancel_endpoint: bool,
993    query_status: bool,
994    pre_registration_cancel: bool,
995    stream_disconnect_cancels: bool,
996}
997
998#[derive(Debug, Serialize)]
999struct SqlIdempotencyCapabilities {
1000    version: u8,
1001    durable_pre_execution_intent: bool,
1002    replay_committed_receipt: bool,
1003    indeterminate_never_reexecutes: bool,
1004}
1005
1006#[derive(Debug, Serialize)]
1007struct SqlPaginationCapabilities {
1008    version: u8,
1009    continuation_endpoint: &'static str,
1010    retained_snapshot: bool,
1011    projection_required: bool,
1012    byte_and_token_hints: bool,
1013}
1014
1015#[derive(Debug, Serialize)]
1016struct CapabilitiesResponse {
1017    sql_cancellation: SqlCancellationCapabilities,
1018    sql_idempotency: SqlIdempotencyCapabilities,
1019    sql_pagination: SqlPaginationCapabilities,
1020}
1021
1022async fn capabilities() -> Json<CapabilitiesResponse> {
1023    Json(CapabilitiesResponse {
1024        sql_cancellation: SqlCancellationCapabilities {
1025            version: 2,
1026            client_query_ids: true,
1027            cancel_endpoint: true,
1028            query_status: true,
1029            pre_registration_cancel: true,
1030            stream_disconnect_cancels: true,
1031        },
1032        sql_idempotency: SqlIdempotencyCapabilities {
1033            version: 1,
1034            durable_pre_execution_intent: true,
1035            replay_committed_receipt: true,
1036            indeterminate_never_reexecutes: true,
1037        },
1038        sql_pagination: SqlPaginationCapabilities {
1039            version: 1,
1040            continuation_endpoint: "/sql/continue",
1041            retained_snapshot: true,
1042            projection_required: true,
1043            byte_and_token_hints: true,
1044        },
1045    })
1046}
1047
1048#[derive(Debug, Deserialize)]
1049struct HistoryRetentionRequest {
1050    #[serde(default)]
1051    history_retention_epochs: serde_json::Value,
1052}
1053
1054#[derive(Debug, Serialize)]
1055struct HistoryRetentionResponse {
1056    history_retention_epochs: u64,
1057    earliest_retained_epoch: u64,
1058}
1059
1060fn history_retention_response(db: &Database) -> HistoryRetentionResponse {
1061    HistoryRetentionResponse {
1062        history_retention_epochs: db.history_retention_epochs(),
1063        earliest_retained_epoch: db.earliest_retained_epoch().0,
1064    }
1065}
1066
1067/// `GET /history/retention` — inspect the durable MVCC history window.
1068async fn history_retention(
1069    State(state): State<Arc<AppState>>,
1070    OptionalPrincipal(principal): OptionalPrincipal,
1071) -> Response {
1072    if let Err(error) = state.db.require_for(
1073        request_principal(&state, &principal).as_ref(),
1074        &mongreldb_core::Permission::Admin,
1075    ) {
1076        return (status_for_error(&error), error.to_string()).into_response();
1077    }
1078    Json(history_retention_response(&state.db)).into_response()
1079}
1080
1081/// `PUT /history/retention` — set the durable MVCC history window.
1082async fn set_history_retention(
1083    State(state): State<Arc<AppState>>,
1084    OptionalPrincipal(principal): OptionalPrincipal,
1085    Json(request): Json<HistoryRetentionRequest>,
1086) -> Response {
1087    if let Err(error) = state.db.require_for(
1088        request_principal(&state, &principal).as_ref(),
1089        &mongreldb_core::Permission::Admin,
1090    ) {
1091        return (status_for_error(&error), error.to_string()).into_response();
1092    }
1093    let Some(epochs) = request.history_retention_epochs.as_u64() else {
1094        return (
1095            StatusCode::BAD_REQUEST,
1096            Json(json!({"error": "history_retention_epochs must be a u64"})),
1097        )
1098            .into_response();
1099    };
1100    match state.db.set_history_retention_epochs(epochs) {
1101        Ok(()) => Json(history_retention_response(&state.db)).into_response(),
1102        Err(error) => (status_for_error(&error), error.to_string()).into_response(),
1103    }
1104}
1105
1106/// `GET /audit` — recent security-audit events (auth + DDL/privilege) as a JSON
1107/// array, oldest-first. Subject to the same auth middleware as every other
1108/// route. This is a best-effort in-memory ring buffer, not a tamper-evident
1109/// log (see `audit` module docs).
1110async fn audit_handler(
1111    State(state): State<Arc<AppState>>,
1112    OptionalPrincipal(principal): OptionalPrincipal,
1113) -> Response {
1114    if let Err(error) = state.db.require_for(
1115        request_principal(&state, &principal).as_ref(),
1116        &mongreldb_core::Permission::Admin,
1117    ) {
1118        return (status_for_error(&error), error.to_string()).into_response();
1119    }
1120    let recent = state.audit.recent();
1121    Json(recent).into_response()
1122}
1123
1124/// Default max live sessions when `--max-sessions` is not given.
1125fn default_max_sessions() -> usize {
1126    std::env::var("MONGRELBL_MAX_SESSIONS")
1127        .ok()
1128        .and_then(|v| v.parse().ok())
1129        .unwrap_or(256)
1130}
1131
1132/// Default idle-session timeout when `--session-idle-timeout` is not given.
1133fn default_session_idle_timeout() -> std::time::Duration {
1134    std::time::Duration::from_secs(
1135        std::env::var("MONGRELBL_SESSION_IDLE_TIMEOUT_SECS")
1136            .ok()
1137            .and_then(|v| v.parse().ok())
1138            .unwrap_or(300),
1139    )
1140}
1141
1142fn default_replication_wal_segments() -> usize {
1143    let replication = std::env::var("MONGRELDB_REPLICATION_WAL_SEGMENTS")
1144        .ok()
1145        .and_then(|value| value.parse().ok())
1146        .unwrap_or(16);
1147    let cdc = std::env::var("MONGRELDB_CDC_WAL_SEGMENTS")
1148        .ok()
1149        .and_then(|value| value.parse().ok())
1150        .unwrap_or(16);
1151    replication.max(cdc)
1152}
1153
1154fn default_history_retention_epochs() -> u64 {
1155    std::env::var("MONGRELDB_HISTORY_RETENTION_EPOCHS")
1156        .ok()
1157        .and_then(|value| value.parse().ok())
1158        .unwrap_or(1024)
1159}
1160
1161fn default_ai_max_concurrent() -> usize {
1162    std::env::var("MONGRELDB_AI_MAX_CONCURRENT")
1163        .ok()
1164        .and_then(|value| value.parse().ok())
1165        .filter(|value| *value > 0)
1166        .unwrap_or(4)
1167}
1168
1169fn positive_env_u64(name: &str, default: u64) -> u64 {
1170    std::env::var(name)
1171        .ok()
1172        .and_then(|value| value.parse().ok())
1173        .filter(|value| *value > 0)
1174        .unwrap_or(default)
1175}
1176
1177fn positive_env_usize(name: &str, default: usize) -> usize {
1178    std::env::var(name)
1179        .ok()
1180        .and_then(|value| value.parse().ok())
1181        .filter(|value| *value > 0)
1182        .unwrap_or(default)
1183}
1184
1185fn default_sql_default_timeout() -> std::time::Duration {
1186    std::time::Duration::from_millis(positive_env_u64("MONGRELDB_SQL_DEFAULT_TIMEOUT_MS", 30_000))
1187}
1188
1189fn default_sql_max_timeout() -> std::time::Duration {
1190    std::time::Duration::from_millis(positive_env_u64("MONGRELDB_SQL_MAX_TIMEOUT_MS", 300_000))
1191}
1192
1193fn default_sql_max_concurrent() -> usize {
1194    positive_env_usize(
1195        "MONGRELDB_SQL_MAX_CONCURRENT",
1196        std::thread::available_parallelism()
1197            .map(usize::from)
1198            .unwrap_or(4),
1199    )
1200}
1201
1202fn default_sql_max_active_queries() -> usize {
1203    positive_env_usize("MONGRELDB_SQL_MAX_ACTIVE_QUERIES", 1_024)
1204}
1205
1206fn default_sql_finished_query_ttl() -> std::time::Duration {
1207    std::time::Duration::from_secs(positive_env_u64(
1208        "MONGRELDB_SQL_FINISHED_QUERY_TTL_SECS",
1209        60,
1210    ))
1211}
1212
1213fn default_sql_pre_cancel_ttl() -> std::time::Duration {
1214    std::time::Duration::from_millis(positive_env_u64("MONGRELDB_SQL_PRE_CANCEL_TTL_MS", 15_000))
1215}
1216
1217fn default_sql_pre_cancel_max_entries() -> usize {
1218    positive_env_usize("MONGRELDB_SQL_PRE_CANCEL_MAX_ENTRIES", 2_048)
1219}
1220
1221fn default_sql_pre_cancel_max_bytes() -> usize {
1222    positive_env_usize("MONGRELDB_SQL_PRE_CANCEL_MAX_BYTES", 1024 * 1024)
1223}
1224
1225fn default_sql_pre_cancel_max_entries_per_owner() -> usize {
1226    positive_env_usize("MONGRELDB_SQL_PRE_CANCEL_MAX_PER_OWNER", 256)
1227}
1228
1229fn default_sql_pre_cancel_rate_window() -> std::time::Duration {
1230    std::time::Duration::from_millis(positive_env_u64(
1231        "MONGRELDB_SQL_PRE_CANCEL_RATE_WINDOW_MS",
1232        1_000,
1233    ))
1234}
1235
1236fn default_sql_pre_cancel_rate_per_owner() -> usize {
1237    positive_env_usize("MONGRELDB_SQL_PRE_CANCEL_RATE_PER_OWNER", 256)
1238}
1239
1240pub(crate) fn default_sql_idempotency_ttl() -> std::time::Duration {
1241    std::time::Duration::from_secs(positive_env_u64(
1242        "MONGRELDB_SQL_IDEMPOTENCY_TTL_SECS",
1243        86_400,
1244    ))
1245}
1246
1247pub(crate) fn default_sql_idempotency_max_entries() -> usize {
1248    positive_env_usize("MONGRELDB_SQL_IDEMPOTENCY_MAX_ENTRIES", 4_096)
1249}
1250
1251fn default_sql_page_ttl() -> std::time::Duration {
1252    std::time::Duration::from_secs(positive_env_u64("MONGRELDB_SQL_PAGE_TTL_SECS", 60))
1253}
1254
1255fn default_sql_page_max_entries() -> usize {
1256    positive_env_usize("MONGRELDB_SQL_PAGE_MAX_ENTRIES", 128)
1257}
1258
1259fn default_sql_page_max_bytes() -> usize {
1260    positive_env_usize("MONGRELDB_SQL_PAGE_MAX_RETAINED_BYTES", 128 * 1024 * 1024)
1261}
1262
1263fn default_sql_page_max_entries_per_owner() -> usize {
1264    positive_env_usize("MONGRELDB_SQL_PAGE_MAX_PER_OWNER", 16)
1265}
1266
1267fn default_sql_cancel_grace() -> std::time::Duration {
1268    std::time::Duration::from_millis(positive_env_u64("MONGRELDB_SQL_CANCEL_GRACE_MS", 1_000))
1269}
1270
1271fn default_sql_max_output_bytes() -> usize {
1272    positive_env_usize("MONGRELDB_SQL_MAX_OUTPUT_BYTES", 64 * 1024 * 1024)
1273}
1274
1275fn default_sql_max_output_rows() -> usize {
1276    positive_env_usize("MONGRELDB_SQL_MAX_OUTPUT_ROWS", 1_000_000)
1277}
1278
1279/// Stable request ownership. Usernames and the literal bearer token are never
1280/// ownership keys, so replacement credentials cannot inherit live resources.
1281fn request_owner(state: &AppState, principal: &Option<mongreldb_core::Principal>) -> String {
1282    if let Some(p) = principal {
1283        return format!("user:{}:{}", p.user_id, p.created_epoch);
1284    }
1285    if let Some(token) = state.auth_token.as_deref() {
1286        let mut digest = sha2::Sha256::new();
1287        digest.update(b"mongreldb-server-bearer-owner-v1\0");
1288        digest.update((token.len() as u64).to_le_bytes());
1289        digest.update(token.as_bytes());
1290        return format!("bearer:{}", sql_idempotency::hex(&digest.finalize()));
1291    }
1292    if state.user_auth || state.db.require_auth_enabled() {
1293        return "unauthenticated".into();
1294    }
1295    "anonymous".into()
1296}
1297
1298fn request_principal(
1299    state: &AppState,
1300    principal: &Option<mongreldb_core::Principal>,
1301) -> Option<mongreldb_core::Principal> {
1302    if let Some(principal) = principal {
1303        return state
1304            .db
1305            .resolve_current_principal(principal)
1306            .or_else(|| Some(principal.clone()));
1307    }
1308    state.auth_token.as_ref().and_then(|_| {
1309        if state.db.require_auth_enabled() {
1310            return state
1311                .db
1312                .principal_snapshot()
1313                .and_then(|principal| state.db.resolve_current_principal(&principal))
1314                .filter(|principal| principal.is_admin);
1315        }
1316        Some(mongreldb_core::Principal {
1317            user_id: 0,
1318            created_epoch: 0,
1319            username: "token".into(),
1320            is_admin: true,
1321            roles: Vec::new(),
1322            permissions: Vec::new(),
1323        })
1324    })
1325}
1326
1327fn current_request_principal(
1328    state: &AppState,
1329    principal: &Option<mongreldb_core::Principal>,
1330) -> Option<mongreldb_core::Principal> {
1331    if let Some(principal) = principal {
1332        return state.db.resolve_current_principal(principal);
1333    }
1334    request_principal(state, principal)
1335}
1336
1337fn request_identity_is_current(
1338    state: &AppState,
1339    principal: &Option<mongreldb_core::Principal>,
1340) -> bool {
1341    if principal.is_some() || state.auth_token.is_some() {
1342        return current_request_principal(state, principal).is_some();
1343    }
1344    !state.user_auth && !state.db.require_auth_enabled()
1345}
1346
1347/// `POST /sessions` — open a long-lived session for cross-request interactive
1348/// transactions. Returns `{"session_id": "..."}`; send `X-Session-ID: <token>`
1349/// on subsequent `/sql` requests to route to it. The session is owned by the
1350/// authenticated principal and auto-expires after the idle timeout.
1351async fn create_session(
1352    State(state): State<Arc<AppState>>,
1353    OptionalPrincipal(principal): OptionalPrincipal,
1354) -> Response {
1355    if !state.accepting_sql.load(Ordering::Acquire) {
1356        return (StatusCode::SERVICE_UNAVAILABLE, "server is shutting down").into_response();
1357    }
1358    if !request_identity_is_current(&state, &principal) {
1359        return StatusCode::UNAUTHORIZED.into_response();
1360    }
1361    let owner = request_owner(&state, &principal);
1362    let session = match MongrelSession::open_with_external_modules_as(
1363        Arc::clone(&state.db),
1364        state.external_modules.iter().cloned(),
1365        request_principal(&state, &principal),
1366    ) {
1367        Ok(session) => session.with_query_registry(Arc::clone(&state.query_registry)),
1368        Err(e) => return (status_for_query_error(&e), e.to_string()).into_response(),
1369    };
1370    match state.sessions.create(session, owner.clone()) {
1371        Some(token) => {
1372            state.audit.record(owner, "session.open", "session created");
1373            Json(json!({ "session_id": token })).into_response()
1374        }
1375        None => (
1376            StatusCode::SERVICE_UNAVAILABLE,
1377            "session limit reached; close an idle session or raise --max-sessions",
1378        )
1379            .into_response(),
1380    }
1381}
1382
1383/// `DELETE /sessions/{id}` — close a session, discarding any open (staged but
1384/// uncommitted) transaction. Only the owning principal may close it.
1385async fn close_session(
1386    State(state): State<Arc<AppState>>,
1387    OptionalPrincipal(principal): OptionalPrincipal,
1388    Path(id): Path<String>,
1389) -> Response {
1390    if !request_identity_is_current(&state, &principal) {
1391        return StatusCode::NOT_FOUND.into_response();
1392    }
1393    let owner = request_owner(&state, &principal);
1394    if let Some(entry) = state.sessions.take_for_close(&id, &owner) {
1395        entry
1396            .session
1397            .query_registry()
1398            .cancel_session(&id, CancellationReason::SessionClosed);
1399        let deadline = tokio::time::Instant::now() + state.sql_cancel_grace;
1400        while entry.session.query_registry().active_for_session(&id) > 0
1401            && tokio::time::Instant::now() < deadline
1402        {
1403            tokio::time::sleep(std::time::Duration::from_millis(5)).await;
1404        }
1405        state.audit.record(owner, "session.close", "session closed");
1406        StatusCode::OK.into_response()
1407    } else {
1408        (
1409            StatusCode::NOT_FOUND,
1410            "session not found or not owned by caller",
1411        )
1412            .into_response()
1413    }
1414}
1415
1416/// Choose a buffered response serialization: `"arrow"` (IPC file) or JSON.
1417/// Streaming IPC is dispatched before query collection.
1418async fn dispatch_buffered_sql_format(
1419    state: &AppState,
1420    format: Option<&str>,
1421    output: ManagedQueryBatches,
1422    query_id: QueryId,
1423    test_hook: Option<mongreldb_query::SqlTestHook>,
1424    output_limits: (usize, usize),
1425) -> Response {
1426    let format = format.unwrap_or("json").to_owned();
1427    let (max_rows, max_bytes) = output_limits;
1428    // Keep the lifecycle guard outside the worker. If the worker panics, the
1429    // join-error path must record a serialization failure, not let dropping a
1430    // moved guard misclassify it as a client disconnect.
1431    let serialization_batches = output.batches().to_vec();
1432    let serialization_query = output.query().clone();
1433    let serialized = tokio::task::spawn_blocking(move || {
1434        serialize_buffered_output(
1435            &format,
1436            &serialization_batches,
1437            &serialization_query,
1438            max_rows,
1439            max_bytes,
1440            test_hook.as_ref(),
1441        )
1442    })
1443    .await;
1444    let result = match serialized {
1445        Ok(result) => result,
1446        Err(error) => {
1447            output
1448                .query()
1449                .record_serialization_failure("SERIALIZATION_WORKER_FAILED");
1450            output.fail();
1451            state.metrics.inc_sql_errors();
1452            return terminal_server_error_response(
1453                state,
1454                query_id,
1455                StatusCode::INTERNAL_SERVER_ERROR,
1456                "SERIALIZATION_WORKER_FAILED",
1457                error.to_string(),
1458            );
1459        }
1460    };
1461    match result {
1462        Ok(serialized) => {
1463            if let Err(error) = output.try_complete() {
1464                state.metrics.inc_sql_errors();
1465                return tracked_query_error_response(state, &error, Some(query_id));
1466            }
1467            state.metrics.add_sql_output_bytes(serialized.bytes.len());
1468            let content_type = if serialized.arrow {
1469                "application/vnd.apache.arrow.file"
1470            } else {
1471                "application/json"
1472            };
1473            with_query_id(
1474                ([(header::CONTENT_TYPE, content_type)], serialized.bytes).into_response(),
1475                query_id,
1476            )
1477        }
1478        Err(BufferedSerializationError::Query(error)) => {
1479            output.fail();
1480            state.metrics.inc_sql_errors();
1481            tracked_query_error_response(state, &error, Some(query_id))
1482        }
1483        Err(BufferedSerializationError::Limit(message)) => {
1484            output.fail_result_limit();
1485            state.metrics.inc_sql_errors();
1486            terminal_server_error_response(
1487                state,
1488                query_id,
1489                StatusCode::PAYLOAD_TOO_LARGE,
1490                "RESULT_LIMIT_EXCEEDED",
1491                message,
1492            )
1493        }
1494        Err(BufferedSerializationError::Encoding(message)) => {
1495            output.fail_serialization();
1496            state.metrics.inc_sql_errors();
1497            terminal_server_error_response(
1498                state,
1499                query_id,
1500                StatusCode::INTERNAL_SERVER_ERROR,
1501                "SERIALIZATION_FAILED",
1502                message,
1503            )
1504        }
1505    }
1506}
1507
1508#[derive(Debug)]
1509enum PaginatedSerializationError {
1510    Query(mongreldb_query::MongrelQueryError),
1511    Limit(String),
1512    Projection(String),
1513    Encoding(String),
1514}
1515
1516struct SerializedPageRows {
1517    rows: Vec<serde_json::Value>,
1518    retained_bytes: usize,
1519}
1520
1521struct PaginatedJsonReader<'a> {
1522    cursor: std::io::Cursor<&'a [u8]>,
1523    query: &'a RegisteredSqlQuery,
1524    test_hook: Option<&'a mongreldb_query::SqlTestHook>,
1525}
1526
1527impl std::io::Read for PaginatedJsonReader<'_> {
1528    fn read(&mut self, buffer: &mut [u8]) -> std::io::Result<usize> {
1529        if let Some(hook) = self.test_hook {
1530            hook(mongreldb_query::SqlTestHookPoint::DuringPaginationDeserialization);
1531        }
1532        self.query
1533            .checkpoint()
1534            .map_err(|error| std::io::Error::other(error.to_string()))?;
1535        // Bound work between cancellation checks even for one very large row.
1536        let length = buffer.len().min(64 * 1024);
1537        std::io::Read::read(&mut self.cursor, &mut buffer[..length])
1538    }
1539}
1540
1541const PAGINATED_VALUE_NODE_BYTES: usize = 64;
1542const PAGINATED_OBJECT_ENTRY_BYTES: usize = 128;
1543const PAGINATED_MEMORY_LIMIT_ERROR: &str = "SQL retained output memory limit exceeded";
1544
1545struct PaginatedDecodeBudget<'a> {
1546    used: usize,
1547    limit: usize,
1548    nodes: usize,
1549    exceeded: bool,
1550    query: &'a RegisteredSqlQuery,
1551    test_hook: Option<&'a mongreldb_query::SqlTestHook>,
1552}
1553
1554impl PaginatedDecodeBudget<'_> {
1555    fn begin_value<E: serde::de::Error>(&mut self) -> Result<(), E> {
1556        self.nodes = self.nodes.saturating_add(1);
1557        if self.nodes & 255 == 0 {
1558            if let Some(hook) = self.test_hook {
1559                hook(mongreldb_query::SqlTestHookPoint::DuringPaginationDeserialization);
1560            }
1561            self.query.checkpoint().map_err(E::custom)?;
1562        }
1563        self.charge::<E>(PAGINATED_VALUE_NODE_BYTES)
1564    }
1565
1566    fn charge<E: serde::de::Error>(&mut self, bytes: usize) -> Result<(), E> {
1567        let next = self.used.saturating_add(bytes);
1568        if next > self.limit {
1569            self.exceeded = true;
1570            return Err(E::custom(PAGINATED_MEMORY_LIMIT_ERROR));
1571        }
1572        self.used = next;
1573        Ok(())
1574    }
1575}
1576
1577struct BudgetedJsonValueSeed<'a, 'query> {
1578    budget: &'a mut PaginatedDecodeBudget<'query>,
1579}
1580
1581impl<'de> serde::de::DeserializeSeed<'de> for BudgetedJsonValueSeed<'_, '_> {
1582    type Value = serde_json::Value;
1583
1584    fn deserialize<D>(self, deserializer: D) -> Result<Self::Value, D::Error>
1585    where
1586        D: serde::Deserializer<'de>,
1587    {
1588        self.budget.begin_value::<D::Error>()?;
1589        deserializer.deserialize_any(BudgetedJsonValueVisitor {
1590            budget: self.budget,
1591        })
1592    }
1593}
1594
1595struct BudgetedJsonValueVisitor<'a, 'query> {
1596    budget: &'a mut PaginatedDecodeBudget<'query>,
1597}
1598
1599impl<'de> serde::de::Visitor<'de> for BudgetedJsonValueVisitor<'_, '_> {
1600    type Value = serde_json::Value;
1601
1602    fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1603        formatter.write_str("a JSON value")
1604    }
1605
1606    fn visit_bool<E>(self, value: bool) -> Result<Self::Value, E> {
1607        Ok(serde_json::Value::Bool(value))
1608    }
1609
1610    fn visit_i64<E>(self, value: i64) -> Result<Self::Value, E> {
1611        Ok(serde_json::Value::Number(value.into()))
1612    }
1613
1614    fn visit_u64<E>(self, value: u64) -> Result<Self::Value, E> {
1615        Ok(serde_json::Value::Number(value.into()))
1616    }
1617
1618    fn visit_f64<E>(self, value: f64) -> Result<Self::Value, E>
1619    where
1620        E: serde::de::Error,
1621    {
1622        serde_json::Number::from_f64(value)
1623            .map(serde_json::Value::Number)
1624            .ok_or_else(|| E::custom("JSON number is not finite"))
1625    }
1626
1627    fn visit_str<E>(self, value: &str) -> Result<Self::Value, E>
1628    where
1629        E: serde::de::Error,
1630    {
1631        self.budget.charge::<E>(value.len())?;
1632        Ok(serde_json::Value::String(value.to_owned()))
1633    }
1634
1635    fn visit_borrowed_str<E>(self, value: &'de str) -> Result<Self::Value, E>
1636    where
1637        E: serde::de::Error,
1638    {
1639        self.visit_str(value)
1640    }
1641
1642    fn visit_string<E>(self, value: String) -> Result<Self::Value, E>
1643    where
1644        E: serde::de::Error,
1645    {
1646        self.budget.charge::<E>(value.capacity())?;
1647        Ok(serde_json::Value::String(value))
1648    }
1649
1650    fn visit_none<E>(self) -> Result<Self::Value, E> {
1651        Ok(serde_json::Value::Null)
1652    }
1653
1654    fn visit_unit<E>(self) -> Result<Self::Value, E> {
1655        Ok(serde_json::Value::Null)
1656    }
1657
1658    fn visit_some<D>(self, deserializer: D) -> Result<Self::Value, D::Error>
1659    where
1660        D: serde::Deserializer<'de>,
1661    {
1662        deserializer.deserialize_any(self)
1663    }
1664
1665    fn visit_newtype_struct<D>(self, deserializer: D) -> Result<Self::Value, D::Error>
1666    where
1667        D: serde::Deserializer<'de>,
1668    {
1669        deserializer.deserialize_any(self)
1670    }
1671
1672    fn visit_seq<A>(self, mut sequence: A) -> Result<Self::Value, A::Error>
1673    where
1674        A: serde::de::SeqAccess<'de>,
1675    {
1676        let mut values = Vec::new();
1677        while let Some(value) = sequence.next_element_seed(BudgetedJsonValueSeed {
1678            budget: self.budget,
1679        })? {
1680            values.push(value);
1681        }
1682        Ok(serde_json::Value::Array(values))
1683    }
1684
1685    fn visit_map<A>(self, mut map: A) -> Result<Self::Value, A::Error>
1686    where
1687        A: serde::de::MapAccess<'de>,
1688    {
1689        let mut values = serde_json::Map::new();
1690        while let Some(key) = map.next_key::<String>()? {
1691            self.budget
1692                .charge::<A::Error>(PAGINATED_OBJECT_ENTRY_BYTES.saturating_add(key.capacity()))?;
1693            let value = map.next_value_seed(BudgetedJsonValueSeed {
1694                budget: self.budget,
1695            })?;
1696            values.insert(key, value);
1697        }
1698        Ok(serde_json::Value::Object(values))
1699    }
1700}
1701
1702struct BudgetedJsonRowsSeed<'a, 'query> {
1703    budget: &'a mut PaginatedDecodeBudget<'query>,
1704}
1705
1706impl<'de> serde::de::DeserializeSeed<'de> for BudgetedJsonRowsSeed<'_, '_> {
1707    type Value = Vec<serde_json::Value>;
1708
1709    fn deserialize<D>(self, deserializer: D) -> Result<Self::Value, D::Error>
1710    where
1711        D: serde::Deserializer<'de>,
1712    {
1713        deserializer.deserialize_seq(BudgetedJsonRowsVisitor {
1714            budget: self.budget,
1715        })
1716    }
1717}
1718
1719struct BudgetedJsonRowsVisitor<'a, 'query> {
1720    budget: &'a mut PaginatedDecodeBudget<'query>,
1721}
1722
1723impl<'de> serde::de::Visitor<'de> for BudgetedJsonRowsVisitor<'_, '_> {
1724    type Value = Vec<serde_json::Value>;
1725
1726    fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1727        formatter.write_str("a JSON array of SQL rows")
1728    }
1729
1730    fn visit_seq<A>(self, mut sequence: A) -> Result<Self::Value, A::Error>
1731    where
1732        A: serde::de::SeqAccess<'de>,
1733    {
1734        let mut rows = Vec::new();
1735        while let Some(row) = sequence.next_element_seed(BudgetedJsonValueSeed {
1736            budget: self.budget,
1737        })? {
1738            rows.push(row);
1739        }
1740        Ok(rows)
1741    }
1742}
1743
1744#[allow(clippy::too_many_arguments)]
1745async fn dispatch_paginated_sql(
1746    state: &AppState,
1747    output: ManagedQueryBatches,
1748    query_id: QueryId,
1749    owner: &str,
1750    pagination: ResolvedSqlPagination,
1751    output_limits: (usize, usize),
1752    test_hook: Option<mongreldb_query::SqlTestHook>,
1753    binding: sql_pages::SqlPageBinding,
1754) -> Response {
1755    let projection = pagination.projection;
1756    let serialization_projection = projection.clone();
1757    let serialization_batches = output.batches().to_vec();
1758    let serialization_query = output.query().clone();
1759    let response_test_hook = test_hook.clone();
1760    let retained_memory_limit = output_limits.1.min(state.sql_pages.max_retained_bytes());
1761    let serialized = tokio::task::spawn_blocking(move || {
1762        serialize_paginated_rows(
1763            &serialization_batches,
1764            &serialization_query,
1765            &serialization_projection,
1766            output_limits.0,
1767            output_limits.1,
1768            retained_memory_limit,
1769            test_hook.as_ref(),
1770        )
1771    })
1772    .await;
1773    let serialized = match serialized {
1774        Ok(result) => result,
1775        Err(error) => {
1776            output
1777                .query()
1778                .record_serialization_failure("SERIALIZATION_WORKER_FAILED");
1779            output.fail();
1780            state.metrics.inc_sql_errors();
1781            return terminal_server_error_response(
1782                state,
1783                query_id,
1784                StatusCode::INTERNAL_SERVER_ERROR,
1785                "SERIALIZATION_WORKER_FAILED",
1786                error.to_string(),
1787            );
1788        }
1789    };
1790    let serialized = match serialized {
1791        Ok(serialized) => serialized,
1792        Err(PaginatedSerializationError::Query(error)) => {
1793            output.fail();
1794            state.metrics.inc_sql_errors();
1795            return tracked_query_error_response(state, &error, Some(query_id));
1796        }
1797        Err(PaginatedSerializationError::Limit(message)) => {
1798            output.fail_result_limit();
1799            state.metrics.inc_sql_errors();
1800            return terminal_server_error_response(
1801                state,
1802                query_id,
1803                StatusCode::PAYLOAD_TOO_LARGE,
1804                "RESULT_LIMIT_EXCEEDED",
1805                message,
1806            );
1807        }
1808        Err(PaginatedSerializationError::Projection(message)) => {
1809            output.fail_with_error(
1810                "INVALID_SQL_PROJECTION",
1811                mongreldb_query::QueryTerminalErrorCategory::Execution,
1812            );
1813            state.metrics.inc_sql_errors();
1814            return terminal_server_error_response(
1815                state,
1816                query_id,
1817                StatusCode::BAD_REQUEST,
1818                "INVALID_SQL_PROJECTION",
1819                message,
1820            );
1821        }
1822        Err(PaginatedSerializationError::Encoding(message)) => {
1823            output.fail_serialization();
1824            state.metrics.inc_sql_errors();
1825            return terminal_server_error_response(
1826                state,
1827                query_id,
1828                StatusCode::INTERNAL_SERVER_ERROR,
1829                "SERIALIZATION_FAILED",
1830                message,
1831            );
1832        }
1833    };
1834    let retained_bytes = serialized.retained_bytes;
1835    let current_binding = sql_pages::SqlPageBinding {
1836        security_version: state.db.security_version(),
1837        catalog_epoch: state.db.catalog_snapshot().db_epoch,
1838    };
1839    if current_binding != binding {
1840        output.fail_with_error(
1841            "SQL_CURSOR_EXPIRED",
1842            mongreldb_query::QueryTerminalErrorCategory::Execution,
1843        );
1844        return sql_cursor_error_response(
1845            StatusCode::CONFLICT,
1846            "SQL_CURSOR_EXPIRED",
1847            "authorization or schema changed while retaining the SQL result",
1848        );
1849    }
1850    let retained = match state.sql_pages.insert(
1851        owner,
1852        serialized.rows,
1853        projection,
1854        pagination.limits,
1855        retained_bytes,
1856        binding,
1857    ) {
1858        Ok(retained) => retained,
1859        Err(sql_pages::InsertError::Full | sql_pages::InsertError::OwnerLimit) => {
1860            output.fail_with_error(
1861                "SQL_PAGE_STORE_FULL",
1862                mongreldb_query::QueryTerminalErrorCategory::Execution,
1863            );
1864            state.metrics.inc_sql_errors();
1865            return terminal_server_error_response(
1866                state,
1867                query_id,
1868                StatusCode::SERVICE_UNAVAILABLE,
1869                "SQL_PAGE_STORE_FULL",
1870                "retained SQL result capacity reached",
1871            );
1872        }
1873        Err(sql_pages::InsertError::EntropyUnavailable) => {
1874            output.fail_with_error(
1875                "ENTROPY_UNAVAILABLE",
1876                mongreldb_query::QueryTerminalErrorCategory::Execution,
1877            );
1878            state.metrics.inc_sql_errors();
1879            return terminal_server_error_response(
1880                state,
1881                query_id,
1882                StatusCode::INTERNAL_SERVER_ERROR,
1883                "ENTROPY_UNAVAILABLE",
1884                "OS CSPRNG unavailable",
1885            );
1886        }
1887    };
1888    let cursor_mac_key = match state.cursor_mac_key.get() {
1889        Ok(key) => key,
1890        Err(_) => {
1891            state.sql_pages.discard(&retained);
1892            output.fail_with_error(
1893                "ENTROPY_UNAVAILABLE",
1894                mongreldb_query::QueryTerminalErrorCategory::Execution,
1895            );
1896            state.metrics.inc_sql_errors();
1897            return terminal_server_error_response(
1898                state,
1899                query_id,
1900                StatusCode::INTERNAL_SERVER_ERROR,
1901                "ENTROPY_UNAVAILABLE",
1902                "OS CSPRNG unavailable",
1903            );
1904        }
1905    };
1906    let page = match sql_pages::SqlPageStore::first_page(&retained, &cursor_mac_key) {
1907        Ok(page) => page,
1908        Err(sql_pages::PageError::RowExceedsLimits) => {
1909            state.sql_pages.discard(&retained);
1910            output.fail_result_limit();
1911            state.metrics.inc_sql_errors();
1912            return terminal_server_error_response(
1913                state,
1914                query_id,
1915                StatusCode::PAYLOAD_TOO_LARGE,
1916                "RESULT_LIMIT_EXCEEDED",
1917                "one projected row exceeds the page byte or token limit",
1918            );
1919        }
1920        Err(sql_pages::PageError::OffsetInvalid) => {
1921            state.sql_pages.discard(&retained);
1922            output.fail_with_error(
1923                "INVALID_PAGE_OFFSET",
1924                mongreldb_query::QueryTerminalErrorCategory::Serialization,
1925            );
1926            state.metrics.inc_sql_errors();
1927            return terminal_server_error_response(
1928                state,
1929                query_id,
1930                StatusCode::INTERNAL_SERVER_ERROR,
1931                "INVALID_PAGE_OFFSET",
1932                "failed to create the first SQL result page",
1933            );
1934        }
1935    };
1936    let discard_after_response = page.next_cursor.is_none();
1937    let page_byte_count = page.byte_count;
1938    let encoded_page = tokio::task::spawn_blocking(move || serialize_sql_page(page)).await;
1939    let encoded_page = match encoded_page {
1940        Ok(Ok(encoded_page)) => encoded_page,
1941        Ok(Err(error)) => {
1942            state.sql_pages.discard(&retained);
1943            output.fail_serialization();
1944            state.metrics.inc_sql_errors();
1945            return terminal_server_error_response(
1946                state,
1947                query_id,
1948                StatusCode::INTERNAL_SERVER_ERROR,
1949                "SERIALIZATION_FAILED",
1950                error.to_string(),
1951            );
1952        }
1953        Err(_) => {
1954            state.sql_pages.discard(&retained);
1955            output.fail_serialization();
1956            state.metrics.inc_sql_errors();
1957            return terminal_server_error_response(
1958                state,
1959                query_id,
1960                StatusCode::INTERNAL_SERVER_ERROR,
1961                "SERIALIZATION_WORKER_FAILED",
1962                "SQL page response serialization worker failed",
1963            );
1964        }
1965    };
1966    if let Some(hook) = response_test_hook {
1967        hook(mongreldb_query::SqlTestHookPoint::AfterPageResponseSerialization);
1968    }
1969    if let Err(error) = output.query().checkpoint() {
1970        state.sql_pages.discard(&retained);
1971        state.metrics.inc_sql_errors();
1972        return tracked_query_error_response(state, &error, Some(query_id));
1973    }
1974    if let Err(error) = output.try_complete() {
1975        state.sql_pages.discard(&retained);
1976        state.metrics.inc_sql_errors();
1977        return tracked_query_error_response(state, &error, Some(query_id));
1978    }
1979    if discard_after_response {
1980        state.sql_pages.discard(&retained);
1981    }
1982    state.metrics.add_sql_output_bytes(page_byte_count);
1983    with_query_id(sql_page_response(encoded_page), query_id)
1984}
1985
1986fn serialize_paginated_rows(
1987    batches: &[arrow::record_batch::RecordBatch],
1988    query: &RegisteredSqlQuery,
1989    projection: &[String],
1990    max_rows: usize,
1991    max_bytes: usize,
1992    retained_memory_limit: usize,
1993    test_hook: Option<&mongreldb_query::SqlTestHook>,
1994) -> Result<SerializedPageRows, PaginatedSerializationError> {
1995    const ROW_CHECKPOINT_INTERVAL: usize = 256;
1996    if batches.is_empty() {
1997        let retained_bytes = sql_pages::accounted_bytes(2, &[], projection, || query.checkpoint())
1998            .map_err(PaginatedSerializationError::Query)?;
1999        if retained_bytes > retained_memory_limit {
2000            return Err(PaginatedSerializationError::Limit(
2001                PAGINATED_MEMORY_LIMIT_ERROR.into(),
2002            ));
2003        }
2004        return Ok(SerializedPageRows {
2005            rows: Vec::new(),
2006            retained_bytes,
2007        });
2008    }
2009    let fields = batches[0].schema();
2010    let mut indices = Vec::with_capacity(projection.len());
2011    for name in projection {
2012        let matches: Vec<_> = fields
2013            .fields()
2014            .iter()
2015            .enumerate()
2016            .filter_map(|(index, field)| (field.name() == name).then_some(index))
2017            .collect();
2018        if matches.len() != 1 {
2019            return Err(PaginatedSerializationError::Projection(format!(
2020                "projected output column {name:?} is missing or ambiguous"
2021            )));
2022        }
2023        indices.push(matches[0]);
2024    }
2025
2026    let mut writer_output = LimitedOutput::new(max_bytes);
2027    let mut rows = 0usize;
2028    let encoding = (|| {
2029        let mut writer = arrow::json::writer::ArrayWriter::new(&mut writer_output);
2030        for batch in batches {
2031            let batch = batch.project(&indices).map_err(|error| error.to_string())?;
2032            for offset in (0..batch.num_rows()).step_by(ROW_CHECKPOINT_INTERVAL) {
2033                if let Some(hook) = test_hook {
2034                    hook(mongreldb_query::SqlTestHookPoint::BeforeSerializationBatch);
2035                }
2036                query.checkpoint().map_err(|error| error.to_string())?;
2037                let length = ROW_CHECKPOINT_INTERVAL.min(batch.num_rows() - offset);
2038                rows = rows.saturating_add(length);
2039                if rows > max_rows {
2040                    return Err("SQL retained output row limit exceeded".into());
2041                }
2042                let slice = batch.slice(offset, length);
2043                writer
2044                    .write_batches(&[&slice])
2045                    .map_err(|error| error.to_string())?;
2046            }
2047        }
2048        writer.finish().map_err(|error| error.to_string())
2049    })();
2050    if let Err(error) = encoding {
2051        if let Err(query_error) = query.checkpoint() {
2052            return Err(PaginatedSerializationError::Query(query_error));
2053        }
2054        if writer_output.exceeded || rows > max_rows {
2055            return Err(PaginatedSerializationError::Limit(error));
2056        }
2057        return Err(PaginatedSerializationError::Encoding(error));
2058    }
2059    let bytes = writer_output.bytes.len();
2060    let retained_base = sql_pages::accounted_bytes(bytes, &[], projection, || query.checkpoint())
2061        .map_err(PaginatedSerializationError::Query)?;
2062    if retained_base > retained_memory_limit {
2063        return Err(PaginatedSerializationError::Limit(
2064            PAGINATED_MEMORY_LIMIT_ERROR.into(),
2065        ));
2066    }
2067    let reader = PaginatedJsonReader {
2068        cursor: std::io::Cursor::new(writer_output.bytes.as_slice()),
2069        query,
2070        test_hook,
2071    };
2072    let mut deserializer =
2073        serde_json::Deserializer::from_reader(std::io::BufReader::with_capacity(64 * 1024, reader));
2074    let mut budget = PaginatedDecodeBudget {
2075        used: retained_base,
2076        limit: retained_memory_limit,
2077        nodes: 0,
2078        exceeded: false,
2079        query,
2080        test_hook,
2081    };
2082    let rows = match serde::de::DeserializeSeed::deserialize(
2083        BudgetedJsonRowsSeed {
2084            budget: &mut budget,
2085        },
2086        &mut deserializer,
2087    ) {
2088        Ok(rows) => {
2089            if let Err(error) = deserializer.end() {
2090                return Err(PaginatedSerializationError::Encoding(error.to_string()));
2091            }
2092            rows
2093        }
2094        Err(error) => {
2095            if let Err(query_error) = query.checkpoint() {
2096                return Err(PaginatedSerializationError::Query(query_error));
2097            }
2098            if budget.exceeded {
2099                return Err(PaginatedSerializationError::Limit(
2100                    PAGINATED_MEMORY_LIMIT_ERROR.into(),
2101                ));
2102            }
2103            return Err(PaginatedSerializationError::Encoding(error.to_string()));
2104        }
2105    };
2106    if let Some(hook) = test_hook {
2107        hook(mongreldb_query::SqlTestHookPoint::AfterSerialization);
2108    }
2109    let retained_bytes =
2110        sql_pages::accounted_bytes(bytes, &rows, projection, || query.checkpoint())
2111            .map_err(PaginatedSerializationError::Query)?;
2112    if retained_bytes > retained_memory_limit {
2113        return Err(PaginatedSerializationError::Limit(
2114            PAGINATED_MEMORY_LIMIT_ERROR.into(),
2115        ));
2116    }
2117    Ok(SerializedPageRows {
2118        rows,
2119        retained_bytes,
2120    })
2121}
2122
2123fn serialize_sql_page(page: sql_pages::SqlPage) -> Result<Vec<u8>, serde_json::Error> {
2124    serde_json::to_vec(&json!({
2125        "status": "completed",
2126        "rows": page.rows,
2127        "next_cursor": page.next_cursor,
2128        "page": {
2129            "offset": page.offset,
2130            "row_count": page.row_count,
2131            "total_rows": page.total_rows,
2132            "byte_count": page.byte_count,
2133            "estimated_tokens": page.estimated_tokens,
2134            "limits": page.limits,
2135            "projection": page.projection,
2136            "expires_at_ms": page.expires_at_ms,
2137            "snapshot": "retained_result",
2138            "token_estimate": "ceil(projected_json_bytes/4)",
2139        }
2140    }))
2141}
2142
2143fn sql_page_response(body: Vec<u8>) -> Response {
2144    ([(header::CONTENT_TYPE, "application/json")], body).into_response()
2145}
2146
2147/// Validate a prepared-statement name: bare identifier only
2148/// (`[A-Za-z_][A-Za-z0-9_]*`). Prevents SQL injection via the name, which is
2149/// interpolated into `PREPARE <name> AS ...` / `EXECUTE <name>(...)`.
2150fn validate_stmt_name(name: &str) -> Result<(), String> {
2151    let mut chars = name.chars();
2152    match chars.next() {
2153        Some(c) if c.is_ascii_alphabetic() || c == '_' => {}
2154        _ => return Err("statement name must start with a letter or underscore".into()),
2155    }
2156    if !chars.all(|c| c.is_ascii_alphanumeric() || c == '_') {
2157        return Err("statement name may contain only letters, digits, or underscore".into());
2158    }
2159    Ok(())
2160}
2161
2162/// Render a JSON value as a safe SQL literal for `EXECUTE` parameter binding.
2163/// Values are escaped so a client cannot inject SQL through a parameter.
2164/// Returns `Err` for non-scalar JSON (arrays/objects) so the caller rejects the
2165/// request with 400 rather than silently binding NULL.
2166fn render_sql_literal(v: &serde_json::Value) -> Result<String, String> {
2167    match v {
2168        serde_json::Value::Null => Ok("NULL".into()),
2169        serde_json::Value::Bool(b) => {
2170            if *b {
2171                Ok("TRUE".into())
2172            } else {
2173                Ok("FALSE".into())
2174            }
2175        }
2176        serde_json::Value::Number(n) => Ok(n.to_string()),
2177        // Single-quote, doubling embedded single quotes (SQL-standard escape).
2178        serde_json::Value::String(s) => {
2179            let mut out = String::with_capacity(s.len() + 2);
2180            out.push('\'');
2181            for c in s.chars() {
2182                if c == '\'' {
2183                    out.push_str("''");
2184                } else {
2185                    out.push(c);
2186                }
2187            }
2188            out.push('\'');
2189            Ok(out)
2190        }
2191        // Arrays/objects are not valid scalar params; reject explicitly.
2192        _ => Err("prepared-statement parameters must be scalar (null/bool/number/string)".into()),
2193    }
2194}
2195
2196#[derive(Deserialize)]
2197struct PrepareRequest {
2198    name: String,
2199    sql: String,
2200    #[serde(default)]
2201    query_id: Option<QueryId>,
2202    #[serde(default)]
2203    timeout_ms: Option<u64>,
2204}
2205
2206/// `POST /sessions/{id}/prepare` — parse+plan `sql` once and store it under
2207/// `name` on the session. Subsequent `EXECUTE name(...)` calls (via this
2208/// endpoint or `EXECUTE` SQL) reuse the cached plan, skipping re-planning.
2209async fn prepare_statement(
2210    State(state): State<Arc<AppState>>,
2211    OptionalPrincipal(principal): OptionalPrincipal,
2212    Path(id): Path<String>,
2213    headers: axum::http::HeaderMap,
2214    Json(req): Json<PrepareRequest>,
2215) -> Response {
2216    if !state.accepting_sql.load(Ordering::Acquire) {
2217        return (StatusCode::SERVICE_UNAVAILABLE, "server is shutting down").into_response();
2218    }
2219    if !request_identity_is_current(&state, &principal) {
2220        return StatusCode::NOT_FOUND.into_response();
2221    }
2222    if let Err(msg) = validate_stmt_name(&req.name) {
2223        return (StatusCode::BAD_REQUEST, msg).into_response();
2224    }
2225    let owner = request_owner(&state, &principal);
2226    let Some(entry) = state.sessions.get(&id, &owner) else {
2227        return (
2228            StatusCode::NOT_FOUND,
2229            "session not found or not owned by caller",
2230        )
2231            .into_response();
2232    };
2233    let (options, query_id) = match resolve_query_options(
2234        &state,
2235        &headers,
2236        req.query_id,
2237        req.timeout_ms,
2238        owner,
2239        Some(id),
2240    ) {
2241        Ok(options) => options,
2242        Err(response) => return *response,
2243    };
2244    let query = match register_controlled_query(&state, &entry.session, options) {
2245        Ok(query) => query,
2246        Err(error) => return tracked_query_error_response(&state, &error, Some(query_id)),
2247    };
2248    let registration = RegisteredQueryGuard::new(query);
2249    if mongreldb_query::contains_boolean_ai_predicate(&req.sql) {
2250        registration.fail();
2251        return with_query_id(remote_boolean_ai_error(), query_id);
2252    }
2253    let _sql_permit = match acquire_sql_permit(&state, &entry.session, registration.query()).await {
2254        Ok(permit) => permit,
2255        Err(error) => return tracked_query_error_response(&state, &error, Some(query_id)),
2256    };
2257    let _guard = tokio::select! {
2258        guard = entry.lock.lock() => guard,
2259        _ = registration.query().control().cancelled() => {
2260            return tracked_query_error_response(
2261                &state,
2262                &cancellation_checkpoint_error(registration.query()),
2263                Some(query_id),
2264            );
2265        }
2266    };
2267    if entry.is_closed() {
2268        return with_query_id(
2269            (StatusCode::NOT_FOUND, "session no longer available").into_response(),
2270            query_id,
2271        );
2272    }
2273    entry.touch();
2274    let sql = format!("PREPARE {} AS {}", req.name, req.sql);
2275    let query = registration.into_query();
2276    match entry.session.run_with_query(&sql, query).await {
2277        Ok(_) => with_query_id(
2278            Json(json!({ "prepared": req.name })).into_response(),
2279            query_id,
2280        ),
2281        Err(error) => tracked_query_error_response(&state, &error, Some(query_id)),
2282    }
2283}
2284
2285#[derive(Deserialize)]
2286struct ExecuteRequest {
2287    name: String,
2288    params: Vec<serde_json::Value>,
2289    #[serde(default)]
2290    format: Option<String>,
2291    #[serde(default)]
2292    query_id: Option<QueryId>,
2293    #[serde(default)]
2294    timeout_ms: Option<u64>,
2295}
2296
2297/// `POST /sessions/{id}/execute` — run a previously-prepared statement with
2298/// typed parameters, reusing its cached plan. Returns the same formats as
2299/// `/sql` (`json` default, `arrow`, `arrow-stream`).
2300async fn execute_statement(
2301    State(state): State<Arc<AppState>>,
2302    OptionalPrincipal(principal): OptionalPrincipal,
2303    Path(id): Path<String>,
2304    headers: axum::http::HeaderMap,
2305    Json(req): Json<ExecuteRequest>,
2306) -> Response {
2307    if !state.accepting_sql.load(Ordering::Acquire) {
2308        return (StatusCode::SERVICE_UNAVAILABLE, "server is shutting down").into_response();
2309    }
2310    if !request_identity_is_current(&state, &principal) {
2311        return StatusCode::NOT_FOUND.into_response();
2312    }
2313    if let Err(msg) = validate_stmt_name(&req.name) {
2314        return (StatusCode::BAD_REQUEST, msg).into_response();
2315    }
2316    let owner = request_owner(&state, &principal);
2317    let Some(entry) = state.sessions.get(&id, &owner) else {
2318        return (
2319            StatusCode::NOT_FOUND,
2320            "session not found or not owned by caller",
2321        )
2322            .into_response();
2323    };
2324    let (options, query_id) = match resolve_query_options(
2325        &state,
2326        &headers,
2327        req.query_id,
2328        req.timeout_ms,
2329        owner,
2330        Some(id),
2331    ) {
2332        Ok(options) => options,
2333        Err(response) => return *response,
2334    };
2335    let query = match register_controlled_query(&state, &entry.session, options) {
2336        Ok(query) => query,
2337        Err(error) => return tracked_query_error_response(&state, &error, Some(query_id)),
2338    };
2339    let registration = RegisteredQueryGuard::new(query);
2340    let sql_permit = match acquire_sql_permit(&state, &entry.session, registration.query()).await {
2341        Ok(permit) => permit,
2342        Err(error) => return tracked_query_error_response(&state, &error, Some(query_id)),
2343    };
2344    let _guard = tokio::select! {
2345        guard = entry.lock.lock() => guard,
2346        _ = registration.query().control().cancelled() => {
2347            return tracked_query_error_response(
2348                &state,
2349                &cancellation_checkpoint_error(registration.query()),
2350                Some(query_id),
2351            );
2352        }
2353    };
2354    if entry.is_closed() {
2355        return with_query_id(
2356            (StatusCode::NOT_FOUND, "session no longer available").into_response(),
2357            query_id,
2358        );
2359    }
2360    entry.touch();
2361    state.metrics.inc_sql_queries();
2362    let literals: Vec<String> = match req
2363        .params
2364        .iter()
2365        .map(render_sql_literal)
2366        .collect::<Result<_, _>>()
2367    {
2368        Ok(v) => v,
2369        Err(msg) => {
2370            state.metrics.inc_sql_errors();
2371            return (StatusCode::BAD_REQUEST, msg).into_response();
2372        }
2373    };
2374    let sql = format!("EXECUTE {}({})", req.name, literals.join(", "));
2375    let start = std::time::Instant::now();
2376    let result = if req.format.as_deref() == Some("arrow-stream") {
2377        let query = registration.into_query();
2378        match entry
2379            .session
2380            .run_stream_with_query_for_serialization(&sql, query)
2381            .await
2382        {
2383            Ok((stream, completion)) => Ok(sql_arrow_stream_response_controlled(
2384                stream,
2385                completion,
2386                sql_permit,
2387                (state.sql_max_output_rows, state.sql_max_output_bytes),
2388                &state,
2389                query_id,
2390                entry.session.sql_test_hook(),
2391            )),
2392            Err(error) => Err(error),
2393        }
2394    } else {
2395        let query = registration.into_query();
2396        match entry
2397            .session
2398            .run_with_query_for_serialization_with_limits(
2399                &sql,
2400                query,
2401                mongreldb_query::SqlCollectionLimits::new(
2402                    state.sql_max_output_rows,
2403                    state.sql_max_output_bytes,
2404                ),
2405            )
2406            .await
2407        {
2408            Ok(output) => Ok(dispatch_buffered_sql_format(
2409                &state,
2410                req.format.as_deref(),
2411                output,
2412                query_id,
2413                entry.session.sql_test_hook(),
2414                (state.sql_max_output_rows, state.sql_max_output_bytes),
2415            )
2416            .await),
2417            Err(error) => Err(error),
2418        }
2419    };
2420    let elapsed = start.elapsed();
2421    if elapsed >= state.slow_query_threshold {
2422        state.metrics.inc_slow_queries();
2423        eprintln!(
2424            "[slow-query] {}\u{00b5}s \u{2014} EXECUTE {}",
2425            elapsed.as_micros(),
2426            req.name
2427        );
2428    }
2429    match result {
2430        Ok(response) => with_query_id(response, query_id),
2431        Err(e) => {
2432            state.metrics.inc_sql_errors();
2433            // A reference to an unprepared/unknown statement is a client error.
2434            let msg = format!("{e}");
2435            let status = if msg.contains("does not exist") {
2436                StatusCode::NOT_FOUND
2437            } else {
2438                status_for_query_error(&e)
2439            };
2440            if status == status_for_query_error(&e) {
2441                tracked_query_error_response(&state, &e, Some(query_id))
2442            } else {
2443                with_query_id(
2444                    (status, format!("{msg} ({}µs)", elapsed.as_micros())).into_response(),
2445                    query_id,
2446                )
2447            }
2448        }
2449    }
2450}
2451
2452/// `DELETE /sessions/{id}/statements/{name}` — drop a prepared statement from
2453/// the session (SQL `DEALLOCATE`).
2454async fn deallocate_statement(
2455    State(state): State<Arc<AppState>>,
2456    OptionalPrincipal(principal): OptionalPrincipal,
2457    Path((id, name)): Path<(String, String)>,
2458    headers: axum::http::HeaderMap,
2459) -> Response {
2460    if !state.accepting_sql.load(Ordering::Acquire) {
2461        return (StatusCode::SERVICE_UNAVAILABLE, "server is shutting down").into_response();
2462    }
2463    if !request_identity_is_current(&state, &principal) {
2464        return StatusCode::NOT_FOUND.into_response();
2465    }
2466    if let Err(msg) = validate_stmt_name(&name) {
2467        return (StatusCode::BAD_REQUEST, msg).into_response();
2468    }
2469    let owner = request_owner(&state, &principal);
2470    let Some(entry) = state.sessions.get(&id, &owner) else {
2471        return (
2472            StatusCode::NOT_FOUND,
2473            "session not found or not owned by caller",
2474        )
2475            .into_response();
2476    };
2477    let (options, query_id) =
2478        match resolve_query_options(&state, &headers, None, None, owner, Some(id)) {
2479            Ok(options) => options,
2480            Err(response) => return *response,
2481        };
2482    let query = match register_controlled_query(&state, &entry.session, options) {
2483        Ok(query) => query,
2484        Err(error) => return tracked_query_error_response(&state, &error, Some(query_id)),
2485    };
2486    let registration = RegisteredQueryGuard::new(query);
2487    let _sql_permit = match acquire_sql_permit(&state, &entry.session, registration.query()).await {
2488        Ok(permit) => permit,
2489        Err(error) => return tracked_query_error_response(&state, &error, Some(query_id)),
2490    };
2491    let _guard = tokio::select! {
2492        guard = entry.lock.lock() => guard,
2493        _ = registration.query().control().cancelled() => {
2494            return tracked_query_error_response(
2495                &state,
2496                &cancellation_checkpoint_error(registration.query()),
2497                Some(query_id),
2498            );
2499        }
2500    };
2501    if entry.is_closed() {
2502        return with_query_id(
2503            (StatusCode::NOT_FOUND, "session no longer available").into_response(),
2504            query_id,
2505        );
2506    }
2507    entry.touch();
2508    state.metrics.inc_sql_queries();
2509    let sql = format!("DEALLOCATE {name}");
2510    let start = std::time::Instant::now();
2511    let result = entry
2512        .session
2513        .run_with_query(&sql, registration.into_query())
2514        .await;
2515    let elapsed = start.elapsed();
2516    if elapsed >= state.slow_query_threshold {
2517        state.metrics.inc_slow_queries();
2518        eprintln!(
2519            "[slow-query] {}\u{00b5}s query_id={} operation=DEALLOCATE",
2520            elapsed.as_micros(),
2521            query_id,
2522        );
2523    }
2524    match result {
2525        Ok(_) => with_query_id(
2526            Json(json!({ "deallocated": name })).into_response(),
2527            query_id,
2528        ),
2529        Err(error) => {
2530            state.metrics.inc_sql_errors();
2531            tracked_query_error_response(&state, &error, Some(query_id))
2532        }
2533    }
2534}
2535
2536/// `mongreldb_tables` gauge. Subject to the same auth middleware as every other
2537/// route (scrape with the configured Bearer token / Basic credentials).
2538async fn metrics_handler(
2539    State(state): State<Arc<AppState>>,
2540    OptionalPrincipal(principal): OptionalPrincipal,
2541) -> Response {
2542    if let Err(error) = state.db.require_for(
2543        request_principal(&state, &principal).as_ref(),
2544        &mongreldb_core::Permission::Admin,
2545    ) {
2546        return (status_for_error(&error), error.to_string()).into_response();
2547    }
2548    let body = state.metrics.prometheus_text(
2549        state.db.table_names().len(),
2550        state.query_registry.active_count(),
2551        state.query_registry.queued_count(),
2552        state.query_registry.entry_count(),
2553        state.query_registry.approximate_bytes(),
2554        (
2555            state.pre_cancellations.len(),
2556            state.pre_cancellations.approximate_bytes(),
2557        ),
2558    );
2559    (
2560        [(
2561            header::CONTENT_TYPE,
2562            "text/plain; version=0.0.4; charset=utf-8".to_string(),
2563        )],
2564        body,
2565    )
2566        .into_response()
2567}
2568
2569/// `POST /compact` — compact all mounted tables.
2570async fn compact_all(
2571    State(state): State<Arc<AppState>>,
2572    OptionalPrincipal(principal): OptionalPrincipal,
2573) -> (StatusCode, Json<serde_json::Value>) {
2574    if let Err(error) = state.db.require_for(
2575        request_principal(&state, &principal).as_ref(),
2576        &mongreldb_core::Permission::Ddl,
2577    ) {
2578        return (
2579            status_for_error(&error),
2580            Json(json!({ "status": "error", "message": error.to_string() })),
2581        );
2582    }
2583    match state.db.compact() {
2584        Ok((compacted, skipped)) => (
2585            StatusCode::OK,
2586            Json(json!({
2587                "status": "ok",
2588                "compacted": compacted,
2589                "skipped": skipped,
2590            })),
2591        ),
2592        Err(e) => (
2593            StatusCode::INTERNAL_SERVER_ERROR,
2594            Json(json!({ "status": "error", "message": format!("{e}") })),
2595        ),
2596    }
2597}
2598
2599/// `POST /tables/{name}/compact` — compact a single table.
2600async fn compact_table(
2601    State(state): State<Arc<AppState>>,
2602    OptionalPrincipal(principal): OptionalPrincipal,
2603    Path(name): Path<String>,
2604) -> (StatusCode, Json<serde_json::Value>) {
2605    if let Err(error) = state.db.require_for(
2606        request_principal(&state, &principal).as_ref(),
2607        &mongreldb_core::Permission::Ddl,
2608    ) {
2609        return (
2610            status_for_error(&error),
2611            Json(json!({ "status": "error", "table": name, "message": error.to_string() })),
2612        );
2613    }
2614    match state.db.compact_table(&name) {
2615        Ok(true) => (
2616            StatusCode::OK,
2617            Json(json!({ "status": "compacted", "table": name })),
2618        ),
2619        Ok(false) => (
2620            StatusCode::OK,
2621            Json(json!({ "status": "skipped", "table": name, "reason": "fewer than 2 runs" })),
2622        ),
2623        Err(e) => (
2624            StatusCode::INTERNAL_SERVER_ERROR,
2625            Json(json!({ "status": "error", "table": name, "message": format!("{e}") })),
2626        ),
2627    }
2628}
2629
2630#[derive(Deserialize)]
2631struct CreateTableRequest {
2632    name: String,
2633    columns: Vec<ColumnDefJson>,
2634}
2635
2636#[derive(Deserialize)]
2637struct ColumnDefJson {
2638    id: u16,
2639    name: String,
2640    ty: String,
2641    primary_key: bool,
2642    #[serde(default)]
2643    nullable: bool,
2644}
2645
2646async fn create_table(
2647    State(state): State<Arc<AppState>>,
2648    OptionalPrincipal(principal): OptionalPrincipal,
2649    Json(req): Json<CreateTableRequest>,
2650) -> Response {
2651    if let Err(error) = state.db.require_for(
2652        request_principal(&state, &principal).as_ref(),
2653        &mongreldb_core::Permission::Ddl,
2654    ) {
2655        return (status_for_error(&error), error.to_string()).into_response();
2656    }
2657    let mut columns = Vec::new();
2658    for c in &req.columns {
2659        let ty = match c.ty.as_str() {
2660            "int64" | "bigint" => TypeId::Int64,
2661            "float64" | "double" => TypeId::Float64,
2662            "bytes" | "varchar" | "text" => TypeId::Bytes,
2663            "bool" => TypeId::Bool,
2664            other => {
2665                return (StatusCode::BAD_REQUEST, format!("unknown type: {other}")).into_response()
2666            }
2667        };
2668        let mut flags = mongreldb_core::schema::ColumnFlags::empty();
2669        if c.primary_key {
2670            flags = flags.with(mongreldb_core::schema::ColumnFlags::PRIMARY_KEY);
2671        }
2672        if c.nullable {
2673            flags = flags.with(mongreldb_core::schema::ColumnFlags::NULLABLE);
2674        }
2675        columns.push(mongreldb_core::schema::ColumnDef {
2676            id: c.id,
2677            name: c.name.clone(),
2678            ty,
2679            flags,
2680            default_value: None,
2681        });
2682    }
2683    let schema = Schema {
2684        schema_id: 0,
2685        columns,
2686        indexes: vec![],
2687        colocation: vec![],
2688        constraints: Default::default(),
2689        clustered: false,
2690    };
2691    if let Err(msg) = validate_table_name(&req.name) {
2692        return (StatusCode::BAD_REQUEST, msg).into_response();
2693    }
2694    match state.db.create_table(&req.name, schema) {
2695        Ok(id) => Json(json!({
2696            "table_id": id,
2697            "table_id_text": id.to_string()
2698        }))
2699        .into_response(),
2700        Err(error) => crate::kit::durable_core_error_response(&error)
2701            .unwrap_or_else(|| (status_for_error(&error), error.to_string()).into_response()),
2702    }
2703}
2704
2705async fn list_tables(
2706    State(state): State<Arc<AppState>>,
2707    OptionalPrincipal(principal): OptionalPrincipal,
2708) -> Json<Vec<String>> {
2709    let principal = request_principal(&state, &principal);
2710    Json(
2711        state
2712            .db
2713            .table_names()
2714            .into_iter()
2715            .filter(|table| {
2716                state
2717                    .db
2718                    .select_column_ids_for(table, principal.as_ref())
2719                    .is_ok()
2720            })
2721            .collect(),
2722    )
2723}
2724
2725async fn drop_table(
2726    State(state): State<Arc<AppState>>,
2727    OptionalPrincipal(principal): OptionalPrincipal,
2728    Path(name): Path<String>,
2729) -> Response {
2730    if let Err(error) = state.db.require_for(
2731        request_principal(&state, &principal).as_ref(),
2732        &mongreldb_core::Permission::Ddl,
2733    ) {
2734        return (status_for_error(&error), error.to_string()).into_response();
2735    }
2736    match state.db.drop_table_with_epoch(&name) {
2737        Ok(epoch) => Json(json!({
2738            "status": "committed",
2739            "epoch": epoch.0,
2740            "epoch_text": epoch.0.to_string()
2741        }))
2742        .into_response(),
2743        Err(error) => crate::kit::durable_core_error_response(&error)
2744            .unwrap_or_else(|| (status_for_error(&error), error.to_string()).into_response()),
2745    }
2746}
2747
2748#[derive(Deserialize)]
2749struct PutRequest {
2750    row: Vec<serde_json::Value>,
2751}
2752
2753pub(crate) fn json_to_value(v: &serde_json::Value, expected: &TypeId) -> Value {
2754    match (v, expected) {
2755        (serde_json::Value::Number(n), TypeId::Float64) => {
2756            n.as_f64().map(Value::Float64).unwrap_or(Value::Null)
2757        }
2758        (serde_json::Value::Number(n), TypeId::Int64) => {
2759            n.as_i64().map(Value::Int64).unwrap_or(Value::Null)
2760        }
2761        (serde_json::Value::String(s), TypeId::Bytes) => Value::Bytes(s.as_bytes().to_vec()),
2762        (serde_json::Value::String(s), TypeId::Enum { variants }) => {
2763            if variants.iter().any(|v| v == s) {
2764                Value::Bytes(s.as_bytes().to_vec())
2765            } else {
2766                Value::Null
2767            }
2768        }
2769        (serde_json::Value::Bool(b), TypeId::Bool) => Value::Bool(*b),
2770        // Embedding input: a JSON array of numbers, validated against the
2771        // declared dimension. Mismatched length or non-numeric elements → Null.
2772        (serde_json::Value::Array(arr), TypeId::Embedding { dim }) => {
2773            if arr.len() as u32 != *dim {
2774                return Value::Null;
2775            }
2776            let vec: Option<Vec<f32>> =
2777                arr.iter().map(|el| el.as_f64().map(|f| f as f32)).collect();
2778            vec.map(Value::Embedding).unwrap_or(Value::Null)
2779        }
2780        (serde_json::Value::Null, _) => Value::Null,
2781        // Lenient fallbacks for unknown/loosely-typed JSON.
2782        (serde_json::Value::Number(n), _) => {
2783            if let Some(i) = n.as_i64() {
2784                Value::Int64(i)
2785            } else if let Some(f) = n.as_f64() {
2786                Value::Float64(f)
2787            } else {
2788                Value::Null
2789            }
2790        }
2791        (serde_json::Value::String(s), _) => Value::Bytes(s.as_bytes().to_vec()),
2792        (serde_json::Value::Bool(b), _) => Value::Bool(*b),
2793        _ => Value::Null,
2794    }
2795}
2796
2797fn legacy_json_to_value(value: &serde_json::Value, expected: &TypeId) -> Result<Value, String> {
2798    if value.is_null() {
2799        return Ok(Value::Null);
2800    }
2801    match expected {
2802        TypeId::Bool => value
2803            .as_bool()
2804            .map(Value::Bool)
2805            .ok_or_else(|| "expected a boolean".into()),
2806        TypeId::Int8
2807        | TypeId::Int16
2808        | TypeId::Int32
2809        | TypeId::Int64
2810        | TypeId::UInt8
2811        | TypeId::UInt16
2812        | TypeId::UInt32
2813        | TypeId::UInt64
2814        | TypeId::TimestampNanos
2815        | TypeId::Date32
2816        | TypeId::Date64
2817        | TypeId::Time64 => value
2818            .as_i64()
2819            .map(Value::Int64)
2820            .ok_or_else(|| "expected a signed 64-bit integer".into()),
2821        TypeId::Float32 | TypeId::Float64 => value
2822            .as_f64()
2823            .filter(|value| value.is_finite())
2824            .map(Value::Float64)
2825            .ok_or_else(|| "expected a finite number".into()),
2826        TypeId::Bytes => match value.as_str() {
2827            Some(value) => Ok(Value::Bytes(value.as_bytes().to_vec())),
2828            None => decode_tagged_hex(value, "bytes").map(Value::Bytes),
2829        },
2830        TypeId::Enum { variants } => {
2831            let bytes = match value.as_str() {
2832                Some(value) => value.as_bytes().to_vec(),
2833                None => decode_tagged_hex(value, "bytes")?,
2834            };
2835            let value =
2836                std::str::from_utf8(&bytes).map_err(|_| "enum variant is not UTF-8".to_string())?;
2837            if !variants.iter().any(|variant| variant == value) {
2838                return Err("expected a declared enum variant".into());
2839            }
2840            Ok(Value::Bytes(bytes))
2841        }
2842        TypeId::Embedding { dim } => {
2843            let values = value
2844                .as_array()
2845                .ok_or_else(|| "expected an embedding array".to_string())?;
2846            if values.len() != *dim as usize {
2847                return Err(format!("expected an embedding with {dim} values"));
2848            }
2849            values
2850                .iter()
2851                .map(|value| {
2852                    value
2853                        .as_f64()
2854                        .map(|value| value as f32)
2855                        .filter(|value| value.is_finite())
2856                        .ok_or_else(|| "embedding values must be finite numbers".to_string())
2857                })
2858                .collect::<Result<Vec<_>, _>>()
2859                .map(Value::Embedding)
2860        }
2861        TypeId::Decimal128 { .. } => {
2862            let object = exact_tagged_object(value, "decimal", &["unscaled"])?;
2863            let text = object["unscaled"]
2864                .as_str()
2865                .ok_or_else(|| "decimal unscaled value must be a string".to_string())?;
2866            let value = text
2867                .parse::<i128>()
2868                .map_err(|_| "decimal unscaled value is invalid".to_string())?;
2869            if value.to_string() != text {
2870                return Err("decimal unscaled value is not canonical".into());
2871            }
2872            Ok(Value::Decimal(value))
2873        }
2874        TypeId::Interval => {
2875            let object = exact_tagged_object(value, "interval", &["months", "days", "nanos"])?;
2876            let months = canonical_i64(&object["months"], "interval months")?;
2877            let days = canonical_i64(&object["days"], "interval days")?
2878                .try_into()
2879                .map_err(|_| "interval days is outside i32 range".to_string())?;
2880            let nanos = canonical_i64(&object["nanos"], "interval nanos")?;
2881            Ok(Value::Interval {
2882                months,
2883                days,
2884                nanos,
2885            })
2886        }
2887        TypeId::Uuid => {
2888            let bytes = decode_tagged_hex(value, "uuid")?;
2889            let bytes: [u8; 16] = bytes
2890                .try_into()
2891                .map_err(|_| "UUID must contain exactly 16 bytes".to_string())?;
2892            Ok(Value::Uuid(bytes))
2893        }
2894        TypeId::Json => {
2895            let bytes = decode_tagged_hex(value, "json")?;
2896            std::str::from_utf8(&bytes).map_err(|_| "JSON value is not UTF-8".to_string())?;
2897            serde_json::from_slice::<serde_json::Value>(&bytes)
2898                .map_err(|error| format!("JSON value is invalid: {error}"))?;
2899            Ok(Value::Json(bytes))
2900        }
2901        TypeId::Array { .. } => Err("legacy put does not support array columns".into()),
2902    }
2903}
2904
2905fn exact_tagged_object<'a>(
2906    value: &'a serde_json::Value,
2907    expected_kind: &str,
2908    fields: &[&str],
2909) -> Result<&'a serde_json::Map<String, serde_json::Value>, String> {
2910    let object = value
2911        .as_object()
2912        .ok_or_else(|| format!("expected tagged {expected_kind} value"))?;
2913    if object.len() != fields.len() + 1
2914        || object
2915            .get("$mongreldb_type")
2916            .and_then(|value| value.as_str())
2917            != Some(expected_kind)
2918        || fields.iter().any(|field| !object.contains_key(*field))
2919    {
2920        return Err(format!("invalid tagged {expected_kind} value"));
2921    }
2922    Ok(object)
2923}
2924
2925fn decode_tagged_hex(value: &serde_json::Value, expected_kind: &str) -> Result<Vec<u8>, String> {
2926    let object = exact_tagged_object(value, expected_kind, &["hex"])?;
2927    let encoded = object["hex"]
2928        .as_str()
2929        .ok_or_else(|| format!("tagged {expected_kind} hex must be a string"))?;
2930    if encoded.len() % 2 != 0 {
2931        return Err(format!("tagged {expected_kind} hex has odd length"));
2932    }
2933    encoded
2934        .as_bytes()
2935        .chunks_exact(2)
2936        .map(|pair| {
2937            let high = hex_nibble(pair[0])?;
2938            let low = hex_nibble(pair[1])?;
2939            Ok((high << 4) | low)
2940        })
2941        .collect()
2942}
2943
2944fn hex_nibble(value: u8) -> Result<u8, String> {
2945    match value {
2946        b'0'..=b'9' => Ok(value - b'0'),
2947        b'a'..=b'f' => Ok(value - b'a' + 10),
2948        _ => Err("hex value must use lowercase ASCII digits".into()),
2949    }
2950}
2951
2952fn canonical_i64(value: &serde_json::Value, field: &str) -> Result<i64, String> {
2953    let text = value
2954        .as_str()
2955        .ok_or_else(|| format!("{field} must be a string"))?;
2956    let value = text
2957        .parse::<i64>()
2958        .map_err(|_| format!("{field} is invalid"))?;
2959    if value.to_string() != text {
2960        return Err(format!("{field} is not canonical"));
2961    }
2962    Ok(value)
2963}
2964
2965#[cfg(test)]
2966mod legacy_wire_tests {
2967    use super::*;
2968
2969    #[test]
2970    fn typed_values_are_exact_and_malformed_values_fail_closed() {
2971        assert_eq!(
2972            legacy_json_to_value(
2973                &json!({"$mongreldb_type": "bytes", "hex": "00ff61"}),
2974                &TypeId::Bytes,
2975            )
2976            .unwrap(),
2977            Value::Bytes(vec![0, 0xff, b'a'])
2978        );
2979        for (value, ty) in [
2980            (
2981                json!({"$mongreldb_type": "bytes", "hex": "00FF"}),
2982                TypeId::Bytes,
2983            ),
2984            (
2985                json!({"$mongreldb_type": "decimal", "unscaled": "01"}),
2986                TypeId::Decimal128 {
2987                    precision: 38,
2988                    scale: 0,
2989                },
2990            ),
2991            (
2992                json!({"$mongreldb_type": "uuid", "hex": "00"}),
2993                TypeId::Uuid,
2994            ),
2995            (
2996                json!({"$mongreldb_type": "json", "hex": "7b"}),
2997                TypeId::Json,
2998            ),
2999            (json!([1.0]), TypeId::Embedding { dim: 2 }),
3000            (json!([1e100, 1.0]), TypeId::Embedding { dim: 2 }),
3001        ] {
3002            assert!(legacy_json_to_value(&value, &ty).is_err(), "{value}");
3003        }
3004    }
3005}
3006
3007/// Parse a flat JSON array `[col_id, val, col_id, val, ...]` into typed cell
3008/// pairs, validating the schema. Returns `Err(message)` on any malformed pair.
3009fn parse_cells(
3010    row: &[serde_json::Value],
3011    schema: &mongreldb_core::schema::Schema,
3012) -> Result<Vec<(u16, Value)>, String> {
3013    if row.len() & 1 != 0 {
3014        return Err("row must be an even-length array of [col_id, value] pairs".into());
3015    }
3016    let mut out = Vec::with_capacity(row.len() / 2);
3017    let mut seen = std::collections::HashSet::new();
3018    for chunk in row.chunks(2) {
3019        let col_id = chunk[0]
3020            .as_u64()
3021            .and_then(|value| u16::try_from(value).ok())
3022            .ok_or("column id must be an unsigned 16-bit integer")?;
3023        if !seen.insert(col_id) {
3024            return Err(format!("duplicate column id {col_id}"));
3025        }
3026        let expected = schema
3027            .columns
3028            .iter()
3029            .find(|c| c.id == col_id)
3030            .map(|c| c.ty.clone())
3031            .ok_or_else(|| format!("unknown column id {col_id}"))?;
3032        let val = legacy_json_to_value(&chunk[1], &expected)?;
3033        out.push((col_id, val));
3034    }
3035    Ok(out)
3036}
3037
3038/// Basic validation for a table name: non-empty and no path separators.
3039pub(crate) fn validate_table_name(name: &str) -> Result<(), String> {
3040    if name.is_empty() {
3041        return Err("table name must not be empty".into());
3042    }
3043    if name.contains('/') || name.contains('\\') || name.contains('\0') {
3044        return Err("table name contains invalid characters".into());
3045    }
3046    Ok(())
3047}
3048
3049async fn put_row(
3050    State(state): State<Arc<AppState>>,
3051    OptionalPrincipal(principal): OptionalPrincipal,
3052    Path(name): Path<String>,
3053    Json(req): Json<PutRequest>,
3054) -> Response {
3055    let handle = match state.db.table(&name) {
3056        Ok(h) => h,
3057        Err(e) => return (StatusCode::NOT_FOUND, e.to_string()).into_response(),
3058    };
3059    let schema = handle.lock().schema().clone();
3060    let row = match parse_cells(&req.row, &schema) {
3061        Ok(r) => r,
3062        Err(msg) => return (StatusCode::BAD_REQUEST, msg).into_response(),
3063    };
3064    state.metrics.inc_puts();
3065    let principal = request_principal(&state, &principal);
3066    match state.db.put_for(&name, row, principal.as_ref()) {
3067        Ok(rid) => Json(json!({ "row_id": rid.0.to_string() })).into_response(),
3068        Err(e) => (status_for_error(&e), e.to_string()).into_response(),
3069    }
3070}
3071
3072async fn count(
3073    State(state): State<Arc<AppState>>,
3074    OptionalPrincipal(principal): OptionalPrincipal,
3075    Path(name): Path<String>,
3076) -> Response {
3077    let principal = request_principal(&state, &principal);
3078    match state.db.count_for(&name, principal.as_ref()) {
3079        Ok(count) => Json(json!({ "count": count })).into_response(),
3080        Err(error) => (status_for_error(&error), error.to_string()).into_response(),
3081    }
3082}
3083
3084async fn commit(
3085    State(state): State<Arc<AppState>>,
3086    OptionalPrincipal(principal): OptionalPrincipal,
3087    Path(name): Path<String>,
3088) -> Response {
3089    if let Err(error) = state.db.require_for(
3090        request_principal(&state, &principal).as_ref(),
3091        &mongreldb_core::Permission::Update {
3092            table: name.clone(),
3093        },
3094    ) {
3095        return (status_for_error(&error), error.to_string()).into_response();
3096    }
3097    let handle = match state.db.table(&name) {
3098        Ok(h) => h,
3099        Err(e) => return (StatusCode::NOT_FOUND, e.to_string()).into_response(),
3100    };
3101    let mut g = handle.lock();
3102    state.metrics.inc_commits();
3103    match g.commit() {
3104        Ok(epoch) => Json(json!({
3105            "epoch": epoch.0,
3106            "epoch_text": epoch.0.to_string()
3107        }))
3108        .into_response(),
3109        Err(error) => crate::kit::durable_core_error_response(&error)
3110            .unwrap_or_else(|| (status_for_error(&error), error.to_string()).into_response()),
3111    }
3112}
3113
3114#[derive(Deserialize)]
3115struct SqlRequest {
3116    sql: String,
3117    /// Output format: `"json"` (the default) for a JSON array of row objects,
3118    /// `"arrow"` for Arrow IPC file bytes.
3119    #[serde(default)]
3120    format: Option<String>,
3121    /// Body values take precedence over the equivalent convenience headers.
3122    #[serde(default)]
3123    query_id: Option<QueryId>,
3124    #[serde(default)]
3125    timeout_ms: Option<u64>,
3126    #[serde(default)]
3127    max_output_rows: Option<u64>,
3128    #[serde(default)]
3129    max_output_bytes: Option<u64>,
3130    #[serde(default)]
3131    idempotency_key: Option<String>,
3132    #[serde(default)]
3133    pagination: Option<SqlPaginationRequest>,
3134}
3135
3136#[derive(Clone, Debug, Deserialize, Serialize)]
3137struct SqlPaginationRequest {
3138    page_size_rows: u64,
3139    projection: Vec<String>,
3140    #[serde(default)]
3141    max_page_bytes: Option<u64>,
3142    #[serde(default)]
3143    max_page_tokens: Option<u64>,
3144}
3145
3146#[derive(Clone)]
3147struct ResolvedSqlPagination {
3148    projection: Vec<String>,
3149    limits: sql_pages::SqlPageLimits,
3150}
3151
3152struct ResolvedSqlRequest {
3153    request: SqlRequest,
3154    output_limits: (usize, usize),
3155    idempotency: Option<sql_idempotency::SqlIdempotencyExecution>,
3156    pagination: Option<ResolvedSqlPagination>,
3157}
3158
3159fn query_error_response(
3160    error: &mongreldb_query::MongrelQueryError,
3161    query_id: Option<QueryId>,
3162) -> Response {
3163    query_error_response_with_status(error, query_id, None)
3164}
3165
3166fn query_error_response_with_status(
3167    error: &mongreldb_query::MongrelQueryError,
3168    query_id: Option<QueryId>,
3169    status: Option<&mongreldb_query::QueryStatus>,
3170) -> Response {
3171    use mongreldb_query::MongrelQueryError;
3172    let (base_code, id) = match error {
3173        MongrelQueryError::QueryCancelled { query_id, .. } => ("QUERY_CANCELLED", Some(*query_id)),
3174        MongrelQueryError::DeadlineExceeded { query_id, .. } => {
3175            ("DEADLINE_EXCEEDED", Some(*query_id))
3176        }
3177        MongrelQueryError::QueryIdConflict { query_id } => ("QUERY_ID_CONFLICT", Some(*query_id)),
3178        MongrelQueryError::QueryRegistryFull => ("QUERY_REGISTRY_FULL", query_id),
3179        MongrelQueryError::ResultLimitExceeded { query_id, .. } => {
3180            ("RESULT_LIMIT_EXCEEDED", Some(*query_id))
3181        }
3182        MongrelQueryError::TransactionAborted => ("TRANSACTION_ABORTED", query_id),
3183        MongrelQueryError::NoSqlTransaction => ("NO_SQL_TRANSACTION", query_id),
3184        MongrelQueryError::SavepointNotFound { .. } => ("SAVEPOINT_NOT_FOUND", query_id),
3185        MongrelQueryError::CommitOutcome { query_id, .. } => ("COMMIT_OUTCOME", Some(*query_id)),
3186        MongrelQueryError::OutcomeUnknown { query_id, .. } => {
3187            ("QUERY_OUTCOME_UNKNOWN", Some(*query_id))
3188        }
3189        _ => ("QUERY_FAILED", query_id),
3190    };
3191    let (
3192        error_committed,
3193        error_committed_statements,
3194        error_last_commit_epoch,
3195        error_first_commit_statement_index,
3196        error_last_commit_statement_index,
3197    ) = match error {
3198        MongrelQueryError::QueryCancelled {
3199            committed,
3200            committed_statements,
3201            last_commit_epoch,
3202            first_commit_statement_index,
3203            last_commit_statement_index,
3204            ..
3205        }
3206        | MongrelQueryError::DeadlineExceeded {
3207            committed,
3208            committed_statements,
3209            last_commit_epoch,
3210            first_commit_statement_index,
3211            last_commit_statement_index,
3212            ..
3213        }
3214        | MongrelQueryError::ResultLimitExceeded {
3215            committed,
3216            committed_statements,
3217            last_commit_epoch,
3218            first_commit_statement_index,
3219            last_commit_statement_index,
3220            ..
3221        } => (
3222            *committed,
3223            *committed_statements,
3224            *last_commit_epoch,
3225            *first_commit_statement_index,
3226            *last_commit_statement_index,
3227        ),
3228        MongrelQueryError::CommitOutcome {
3229            committed,
3230            committed_statements,
3231            last_commit_epoch,
3232            first_commit_statement_index,
3233            last_commit_statement_index,
3234            ..
3235        } => (
3236            *committed,
3237            *committed_statements,
3238            *last_commit_epoch,
3239            *first_commit_statement_index,
3240            *last_commit_statement_index,
3241        ),
3242        _ => (false, 0, None, None, None),
3243    };
3244    let committed = status.map_or_else(
3245        || error_committed,
3246        |status| status.durable_outcome.committed,
3247    );
3248    let outcome_unknown = matches!(error, MongrelQueryError::OutcomeUnknown { .. })
3249        || status.is_some_and(|status| status.outcome_unknown);
3250    let code = status
3251        .and_then(|status| {
3252            status
3253                .terminal_error
3254                .as_ref()
3255                .map(|error| error.code.as_str())
3256        })
3257        .unwrap_or(match (base_code, committed) {
3258            ("QUERY_CANCELLED", true) => "QUERY_CANCELLED_AFTER_COMMIT",
3259            ("DEADLINE_EXCEEDED", true) => "DEADLINE_AFTER_COMMIT",
3260            _ => base_code,
3261        });
3262    let response_status = if outcome_unknown {
3263        "outcome_unknown"
3264    } else {
3265        status
3266            .and_then(mongreldb_query::QueryStatus::terminal_state)
3267            .map(terminal_state_name)
3268            .unwrap_or_else(|| match (error, committed) {
3269                (MongrelQueryError::QueryCancelled { .. }, true) => "cancelled_after_commit",
3270                (MongrelQueryError::QueryCancelled { .. }, false) => "cancelled_before_commit",
3271                (MongrelQueryError::DeadlineExceeded { .. }, true) => "deadline_after_commit",
3272                (MongrelQueryError::DeadlineExceeded { .. }, false) => "deadline_before_commit",
3273                (_, true) => "committed_with_error",
3274                _ => "failed_before_commit",
3275            })
3276    };
3277    let (completed_statements, statement_index) = status.map_or_else(
3278        || match error {
3279            MongrelQueryError::QueryCancelled {
3280                completed_statements,
3281                cancelled_statement_index,
3282                ..
3283            }
3284            | MongrelQueryError::DeadlineExceeded {
3285                completed_statements,
3286                cancelled_statement_index,
3287                ..
3288            } => (*completed_statements, *cancelled_statement_index),
3289            MongrelQueryError::ResultLimitExceeded {
3290                completed_statements,
3291                statement_index,
3292                ..
3293            } => (*completed_statements, *statement_index),
3294            MongrelQueryError::CommitOutcome {
3295                completed_statements,
3296                statement_index,
3297                ..
3298            } => (*completed_statements, *statement_index),
3299            _ => (0, 0),
3300        },
3301        |status| (status.completed_statements, status.statement_index),
3302    );
3303    let committed_statements = status.map_or(error_committed_statements, |status| {
3304        status.durable_outcome.committed_statements
3305    });
3306    let last_commit_epoch = status.map_or(error_last_commit_epoch, |status| {
3307        status.durable_outcome.last_commit_epoch
3308    });
3309    let first_commit_statement_index = status
3310        .map_or(error_first_commit_statement_index, |status| {
3311            status.durable_outcome.first_commit_statement_index
3312        });
3313    let last_commit_statement_index = status.map_or(error_last_commit_statement_index, |status| {
3314        status.durable_outcome.last_commit_statement_index
3315    });
3316    let cancellation_reason = status
3317        .map(|status| status.cancellation_reason)
3318        .or(match error {
3319            MongrelQueryError::QueryCancelled { reason, .. } => Some(*reason),
3320            MongrelQueryError::DeadlineExceeded { .. } => Some(CancellationReason::Deadline),
3321            _ => None,
3322        })
3323        .map(cancellation_reason_name);
3324    let cancel_outcome = match error {
3325        MongrelQueryError::QueryCancelled { .. } | MongrelQueryError::DeadlineExceeded { .. } => {
3326            Some("accepted")
3327        }
3328        _ => status.and_then(query_cancel_outcome),
3329    };
3330    let outcome = if outcome_unknown {
3331        json!({
3332            "committed": null,
3333            "committed_statements": null,
3334            "last_commit_epoch": null,
3335            "last_commit_epoch_text": null,
3336            "first_commit_statement_index": null,
3337            "last_commit_statement_index": null,
3338            "completed_statements": null,
3339            "statement_index": null,
3340            "serialization": "unknown",
3341        })
3342    } else {
3343        status.map_or_else(
3344            || {
3345                json!({
3346                    "committed": committed,
3347                    "committed_statements": committed_statements,
3348                    "last_commit_epoch": last_commit_epoch,
3349                    "last_commit_epoch_text": epoch_text(last_commit_epoch),
3350                    "first_commit_statement_index": first_commit_statement_index,
3351                    "last_commit_statement_index": last_commit_statement_index,
3352                    "completed_statements": completed_statements,
3353                    "statement_index": statement_index,
3354                    "serialization": "unknown",
3355                })
3356            },
3357            |status| query_outcome_json(Some(status)),
3358        )
3359    };
3360    let http_status = match code {
3361        "QUERY_CANCELLED_AFTER_COMMIT" | "DEADLINE_AFTER_COMMIT" => StatusCode::CONFLICT,
3362        "QUERY_CANCELLED" => client_closed_request_status(),
3363        "DEADLINE_EXCEEDED" => StatusCode::GATEWAY_TIMEOUT,
3364        _ => status_for_query_error(error),
3365    };
3366    let mut response = (
3367        http_status,
3368        Json(json!({
3369            "query_id": id.map(|value| value.to_string()),
3370            "status": response_status,
3371            "terminal_state": response_status,
3372            "committed": (!outcome_unknown).then_some(committed),
3373            "committed_statements": (!outcome_unknown).then_some(committed_statements),
3374            "last_commit_epoch": (!outcome_unknown).then_some(last_commit_epoch).flatten(),
3375            "last_commit_epoch_text": (!outcome_unknown).then_some(epoch_text(last_commit_epoch)).flatten(),
3376            "first_commit_statement_index": (!outcome_unknown).then_some(first_commit_statement_index).flatten(),
3377            "last_commit_statement_index": (!outcome_unknown).then_some(last_commit_statement_index).flatten(),
3378            "completed_statements": (!outcome_unknown).then_some(completed_statements),
3379            "statement_index": (!outcome_unknown).then_some(statement_index),
3380            "cancel_outcome": cancel_outcome,
3381            "cancellation_reason": cancellation_reason,
3382            "retryable": matches!(error, MongrelQueryError::QueryRegistryFull),
3383            "server_state": status.map(|status| query_phase_name(status.phase)),
3384            "outcome": outcome,
3385            "error": {
3386                "code": code,
3387                "message": error.to_string(),
3388                "query_id": id.map(|value| value.to_string()),
3389                "committed": (!outcome_unknown).then_some(committed),
3390                "retryable": matches!(error, MongrelQueryError::QueryRegistryFull),
3391            }
3392        })),
3393    )
3394        .into_response();
3395    if let Some(id) = id {
3396        add_query_id_header(&mut response, id);
3397    }
3398    response
3399}
3400
3401fn record_query_error(metrics: &metrics::Metrics, error: &mongreldb_query::MongrelQueryError) {
3402    match error {
3403        mongreldb_query::MongrelQueryError::QueryCancelled { reason, .. } => {
3404            metrics.inc_sql_cancelled(*reason)
3405        }
3406        mongreldb_query::MongrelQueryError::DeadlineExceeded { .. } => {
3407            metrics.inc_sql_deadline_exceeded();
3408            metrics.inc_sql_cancelled(CancellationReason::Deadline);
3409        }
3410        _ => {}
3411    }
3412}
3413
3414fn tracked_query_error_response(
3415    state: &AppState,
3416    error: &mongreldb_query::MongrelQueryError,
3417    query_id: Option<QueryId>,
3418) -> Response {
3419    record_query_error(&state.metrics, error);
3420    if let mongreldb_query::MongrelQueryError::QueryCancelled { query_id, .. } = error {
3421        if let Some(requested_at) = state
3422            .query_registry
3423            .status(*query_id)
3424            .and_then(|status| status.cancel_requested_at)
3425        {
3426            state
3427                .metrics
3428                .observe_sql_cancel_latency(requested_at.elapsed());
3429        }
3430    }
3431    let status = if matches!(
3432        error,
3433        mongreldb_query::MongrelQueryError::QueryIdConflict { .. }
3434            | mongreldb_query::MongrelQueryError::QueryRegistryFull
3435    ) {
3436        None
3437    } else {
3438        query_id
3439            .or(match error {
3440                mongreldb_query::MongrelQueryError::QueryCancelled { query_id, .. }
3441                | mongreldb_query::MongrelQueryError::DeadlineExceeded { query_id, .. }
3442                | mongreldb_query::MongrelQueryError::CommitOutcome { query_id, .. }
3443                | mongreldb_query::MongrelQueryError::OutcomeUnknown { query_id, .. } => {
3444                    Some(*query_id)
3445                }
3446                _ => None,
3447            })
3448            .and_then(|query_id| state.query_registry.status(query_id))
3449    };
3450    query_error_response_with_status(error, query_id, status.as_ref())
3451}
3452
3453fn add_query_id_header(response: &mut Response, query_id: QueryId) {
3454    if let Ok(value) = axum::http::HeaderValue::from_str(&query_id.to_string()) {
3455        response.headers_mut().insert("x-mongreldb-query-id", value);
3456    }
3457}
3458
3459fn with_query_id(mut response: Response, query_id: QueryId) -> Response {
3460    add_query_id_header(&mut response, query_id);
3461    response
3462}
3463
3464fn bad_query_control_request(message: impl Into<String>, query_id: Option<QueryId>) -> Response {
3465    let mut response = (
3466        StatusCode::BAD_REQUEST,
3467        Json(json!({
3468            "query_id": query_id.map(|value| value.to_string()),
3469            "status": "failed_before_commit",
3470            "terminal_state": "failed_before_commit",
3471            "committed": false,
3472            "committed_statements": 0,
3473            "last_commit_epoch": null,
3474            "last_commit_epoch_text": null,
3475            "first_commit_statement_index": null,
3476            "last_commit_statement_index": null,
3477            "completed_statements": 0,
3478            "statement_index": 0,
3479            "cancel_outcome": null,
3480            "cancellation_reason": null,
3481            "retryable": false,
3482            "server_state": "failed",
3483            "outcome": {
3484                "committed": false,
3485                "committed_statements": 0,
3486                "last_commit_epoch": null,
3487                "last_commit_epoch_text": null,
3488                "first_commit_statement_index": null,
3489                "last_commit_statement_index": null,
3490                "completed_statements": 0,
3491                "statement_index": 0,
3492                "serialization": "not_started",
3493            },
3494            "error": {
3495                "code": "INVALID_QUERY_OPTIONS",
3496                "message": message.into(),
3497                "query_id": query_id.map(|value| value.to_string()),
3498                "committed": false,
3499                "retryable": false,
3500            }
3501        })),
3502    )
3503        .into_response();
3504    if let Some(query_id) = query_id {
3505        add_query_id_header(&mut response, query_id);
3506    }
3507    response
3508}
3509
3510fn resolve_query_options(
3511    state: &AppState,
3512    headers: &axum::http::HeaderMap,
3513    body_query_id: Option<QueryId>,
3514    body_timeout_ms: Option<u64>,
3515    owner: String,
3516    session_id: Option<String>,
3517) -> std::result::Result<(SqlQueryOptions, QueryId), Box<Response>> {
3518    let query_id = match body_query_id {
3519        Some(query_id) => query_id,
3520        None => match headers.get("x-mongreldb-query-id") {
3521            Some(value) => {
3522                let value = value.to_str().map_err(|_| {
3523                    Box::new(bad_query_control_request(
3524                        "X-MongrelDB-Query-ID is not valid text",
3525                        None,
3526                    ))
3527                })?;
3528                value
3529                    .parse()
3530                    .map_err(|error: mongreldb_query::MongrelQueryError| {
3531                        Box::new(bad_query_control_request(error.to_string(), None))
3532                    })?
3533            }
3534            None => {
3535                QueryId::random().map_err(|error| Box::new(query_error_response(&error, None)))?
3536            }
3537        },
3538    };
3539    let timeout_ms = match body_timeout_ms {
3540        Some(timeout_ms) => timeout_ms,
3541        None => match headers.get("x-mongreldb-timeout-ms") {
3542            Some(value) => value
3543                .to_str()
3544                .ok()
3545                .and_then(|value| value.parse::<u64>().ok())
3546                .ok_or_else(|| {
3547                    Box::new(bad_query_control_request(
3548                        "X-MongrelDB-Timeout-Ms must be a positive integer",
3549                        Some(query_id),
3550                    ))
3551                })?,
3552            None => state
3553                .sql_default_timeout
3554                .as_millis()
3555                .min(u128::from(u64::MAX)) as u64,
3556        },
3557    };
3558    if timeout_ms == 0 {
3559        return Err(Box::new(bad_query_control_request(
3560            "timeout_ms must be positive",
3561            Some(query_id),
3562        )));
3563    }
3564    let timeout = std::time::Duration::from_millis(timeout_ms);
3565    if timeout > state.sql_max_timeout {
3566        return Err(Box::new(bad_query_control_request(
3567            format!(
3568                "timeout_ms exceeds server maximum of {}",
3569                state.sql_max_timeout.as_millis()
3570            ),
3571            Some(query_id),
3572        )));
3573    }
3574    Ok((
3575        SqlQueryOptions {
3576            query_id: Some(query_id),
3577            timeout: Some(timeout),
3578            owner: Some(owner),
3579            session_id,
3580            parent_control: None,
3581        },
3582        query_id,
3583    ))
3584}
3585
3586fn resolve_sql_output_limits(
3587    state: &AppState,
3588    request: &SqlRequest,
3589    query_id: QueryId,
3590) -> std::result::Result<(usize, usize), Box<Response>> {
3591    fn resolve(
3592        requested: Option<u64>,
3593        configured: usize,
3594        name: &str,
3595        query_id: QueryId,
3596    ) -> std::result::Result<usize, Box<Response>> {
3597        if requested == Some(0) {
3598            return Err(Box::new(bad_query_control_request(
3599                format!("{name} must be positive"),
3600                Some(query_id),
3601            )));
3602        }
3603        let requested = requested
3604            .and_then(|value| usize::try_from(value).ok())
3605            .unwrap_or(usize::MAX);
3606        Ok(requested.min(configured))
3607    }
3608
3609    Ok((
3610        resolve(
3611            request.max_output_rows,
3612            state.sql_max_output_rows,
3613            "max_output_rows",
3614            query_id,
3615        )?,
3616        resolve(
3617            request.max_output_bytes,
3618            state.sql_max_output_bytes,
3619            "max_output_bytes",
3620            query_id,
3621        )?,
3622    ))
3623}
3624
3625fn resolve_sql_pagination(
3626    headers: &axum::http::HeaderMap,
3627    request: &SqlRequest,
3628    output_limits: (usize, usize),
3629    registration: RegisteredQueryGuard,
3630    query_id: QueryId,
3631) -> Result<(RegisteredQueryGuard, Option<ResolvedSqlPagination>), Box<Response>> {
3632    let Some(pagination) = request.pagination.as_ref() else {
3633        return Ok((registration, None));
3634    };
3635    if requested_sql_idempotency_key(headers, request)
3636        .ok()
3637        .flatten()
3638        .is_some()
3639    {
3640        return Err(Box::new(registered_sql_error_response(
3641            registration,
3642            query_id,
3643            StatusCode::BAD_REQUEST,
3644            "INCOMPATIBLE_SQL_CONTROLS",
3645            "idempotency_key cannot be combined with SQL pagination",
3646            false,
3647        )));
3648    }
3649    if request
3650        .format
3651        .as_deref()
3652        .is_some_and(|format| format != "json")
3653    {
3654        return Err(Box::new(registered_sql_error_response(
3655            registration,
3656            query_id,
3657            StatusCode::BAD_REQUEST,
3658            "PAGINATION_REQUIRES_JSON",
3659            "SQL pagination supports JSON responses only",
3660            false,
3661        )));
3662    }
3663    registration.query().set_sql_metadata(&request.sql);
3664    if !mongreldb_query::is_single_read_only_query(&request.sql) {
3665        return Err(Box::new(registered_sql_error_response(
3666            registration,
3667            query_id,
3668            StatusCode::BAD_REQUEST,
3669            "PAGINATION_REQUIRES_SINGLE_READ_QUERY",
3670            "SQL pagination accepts exactly one read-only query statement",
3671            false,
3672        )));
3673    }
3674    let page_size = match usize::try_from(pagination.page_size_rows) {
3675        Ok(0) | Err(_) => {
3676            return Err(Box::new(registered_sql_error_response(
3677                registration,
3678                query_id,
3679                StatusCode::BAD_REQUEST,
3680                "INVALID_PAGINATION_OPTIONS",
3681                "pagination.page_size_rows must be positive",
3682                false,
3683            )))
3684        }
3685        Ok(value) => value.min(output_limits.0),
3686    };
3687    if pagination.projection.is_empty() || pagination.projection.len() > 128 {
3688        return Err(Box::new(registered_sql_error_response(
3689            registration,
3690            query_id,
3691            StatusCode::BAD_REQUEST,
3692            "INVALID_SQL_PROJECTION",
3693            "pagination.projection must contain between 1 and 128 output column names",
3694            false,
3695        )));
3696    }
3697    let mut seen = std::collections::HashSet::new();
3698    let metadata_bytes = pagination
3699        .projection
3700        .iter()
3701        .map(String::len)
3702        .fold(0usize, usize::saturating_add);
3703    if metadata_bytes > 16 * 1024
3704        || pagination.projection.iter().any(|column| {
3705            column.is_empty()
3706                || column == "*"
3707                || column.len() > 256
3708                || !seen.insert(column.as_str())
3709        })
3710    {
3711        return Err(Box::new(registered_sql_error_response(
3712            registration,
3713            query_id,
3714            StatusCode::BAD_REQUEST,
3715            "INVALID_SQL_PROJECTION",
3716            "pagination.projection requires unique explicit output names of at most 256 bytes",
3717            false,
3718        )));
3719    }
3720    let max_page_bytes = match pagination.max_page_bytes {
3721        Some(0) => {
3722            return Err(Box::new(registered_sql_error_response(
3723                registration,
3724                query_id,
3725                StatusCode::BAD_REQUEST,
3726                "INVALID_PAGINATION_OPTIONS",
3727                "pagination.max_page_bytes must be positive",
3728                false,
3729            )))
3730        }
3731        Some(value) => usize::try_from(value)
3732            .unwrap_or(usize::MAX)
3733            .min(output_limits.1),
3734        None => output_limits.1.min(1024 * 1024),
3735    };
3736    let token_cap = (output_limits.1.saturating_add(3) / 4).max(1);
3737    let max_page_tokens = match pagination.max_page_tokens {
3738        Some(0) => {
3739            return Err(Box::new(registered_sql_error_response(
3740                registration,
3741                query_id,
3742                StatusCode::BAD_REQUEST,
3743                "INVALID_PAGINATION_OPTIONS",
3744                "pagination.max_page_tokens must be positive",
3745                false,
3746            )))
3747        }
3748        Some(value) => usize::try_from(value).unwrap_or(usize::MAX).min(token_cap),
3749        None => (max_page_bytes.saturating_add(3) / 4).max(1),
3750    };
3751    Ok((
3752        registration,
3753        Some(ResolvedSqlPagination {
3754            projection: pagination.projection.clone(),
3755            limits: sql_pages::SqlPageLimits {
3756                rows: page_size,
3757                bytes: max_page_bytes,
3758                tokens: max_page_tokens,
3759            },
3760        }),
3761    ))
3762}
3763
3764fn requested_sql_idempotency_key(
3765    headers: &axum::http::HeaderMap,
3766    request: &SqlRequest,
3767) -> Result<Option<String>, &'static str> {
3768    let header = match headers.get("idempotency-key") {
3769        Some(value) => Some(
3770            value
3771                .to_str()
3772                .map_err(|_| "Idempotency-Key must be valid UTF-8")?,
3773        ),
3774        None => None,
3775    };
3776    match (request.idempotency_key.as_deref(), header) {
3777        (Some(body), Some(header)) if body != header => {
3778            Err("body idempotency_key and Idempotency-Key header must match")
3779        }
3780        (Some(body), _) => Ok(Some(body.to_owned())),
3781        (None, Some(header)) => Ok(Some(header.to_owned())),
3782        (None, None) => Ok(None),
3783    }
3784}
3785
3786fn sql_idempotency_binding(
3787    request: &SqlRequest,
3788    output_limits: (usize, usize),
3789    session_id: Option<&str>,
3790    expires_after_ms: u64,
3791) -> Result<sql_idempotency::SqlIdempotencyBinding, serde_json::Error> {
3792    let request_semantics = serde_json::to_vec(&json!({
3793        "format": request.format.as_deref().unwrap_or("json"),
3794        "max_output_rows": output_limits.0,
3795        "max_output_bytes": output_limits.1,
3796        "pagination": request.pagination.as_ref(),
3797    }))?;
3798    let session_semantics = session_id.map_or_else(
3799        || b"ephemeral".to_vec(),
3800        |session_id| {
3801            let mut semantics = b"session\0".to_vec();
3802            semantics.extend_from_slice(session_id.as_bytes());
3803            semantics
3804        },
3805    );
3806    Ok(sql_idempotency::SqlIdempotencyBinding {
3807        sql_fingerprint: mongreldb_query::normalized_sql_fingerprint(&request.sql),
3808        // `/sql` currently has no separate bind-parameter array. SQL literals
3809        // are covered by the normalized SQL fingerprint above.
3810        parameter_hash: sql_idempotency::hash(b"[]"),
3811        request_semantics_hash: sql_idempotency::hash(&request_semantics),
3812        session_semantics_hash: sql_idempotency::hash(&session_semantics),
3813        expires_after_ms,
3814    })
3815}
3816
3817struct SqlIdempotencyContext<'a> {
3818    headers: &'a axum::http::HeaderMap,
3819    request: &'a SqlRequest,
3820    output_limits: (usize, usize),
3821    owner: &'a str,
3822    session_id: Option<&'a str>,
3823    session_in_transaction: bool,
3824    query_id: QueryId,
3825}
3826
3827async fn begin_sql_idempotency(
3828    state: &AppState,
3829    context: SqlIdempotencyContext<'_>,
3830    registration: RegisteredQueryGuard,
3831) -> Result<
3832    (
3833        RegisteredQueryGuard,
3834        Option<sql_idempotency::SqlIdempotencyExecution>,
3835    ),
3836    Response,
3837> {
3838    let SqlIdempotencyContext {
3839        headers,
3840        request,
3841        output_limits,
3842        owner,
3843        session_id,
3844        session_in_transaction,
3845        query_id,
3846    } = context;
3847    let key = match requested_sql_idempotency_key(headers, request) {
3848        Ok(key) => key,
3849        Err(message) => {
3850            return Err(registered_sql_error_response(
3851                registration,
3852                query_id,
3853                StatusCode::BAD_REQUEST,
3854                "INVALID_IDEMPOTENCY_KEY",
3855                message,
3856                false,
3857            ))
3858        }
3859    };
3860    let Some(key) = key else {
3861        return Ok((registration, None));
3862    };
3863    match mongreldb_query::classify_sql_idempotency(&request.sql) {
3864        mongreldb_query::SqlIdempotencyClass::ReadOnly
3865        | mongreldb_query::SqlIdempotencyClass::Unsupported => {
3866            return Err(registered_sql_error_response(
3867                registration,
3868                query_id,
3869                StatusCode::BAD_REQUEST,
3870                "IDEMPOTENCY_REQUIRES_SINGLE_WRITE",
3871                "idempotency_key accepts one non-transaction SQL write statement",
3872                false,
3873            ));
3874        }
3875        mongreldb_query::SqlIdempotencyClass::SingleWrite => {}
3876    }
3877    if let Err(message) = sql_idempotency::SqlIdempotencyStore::validate_key(&key) {
3878        return Err(registered_sql_error_response(
3879            registration,
3880            query_id,
3881            StatusCode::BAD_REQUEST,
3882            "INVALID_IDEMPOTENCY_KEY",
3883            message,
3884            false,
3885        ));
3886    }
3887    if request
3888        .format
3889        .as_deref()
3890        .is_some_and(|format| format != "json")
3891    {
3892        return Err(registered_sql_error_response(
3893            registration,
3894            query_id,
3895            StatusCode::BAD_REQUEST,
3896            "IDEMPOTENCY_REQUIRES_JSON",
3897            "SQL idempotency supports buffered JSON responses only",
3898            false,
3899        ));
3900    }
3901    if session_in_transaction {
3902        return Err(registered_sql_error_response(
3903            registration,
3904            query_id,
3905            StatusCode::CONFLICT,
3906            "IDEMPOTENCY_UNSUPPORTED_IN_TRANSACTION",
3907            "SQL idempotency cannot be used inside an open session transaction",
3908            false,
3909        ));
3910    }
3911    registration.query().set_sql_metadata(&request.sql);
3912    let binding = match sql_idempotency_binding(
3913        request,
3914        output_limits,
3915        session_id,
3916        state.sql_idempotency.expires_after_ms(),
3917    ) {
3918        Ok(binding) => binding,
3919        Err(_) => {
3920            return Err(registered_sql_error_response(
3921                registration,
3922                query_id,
3923                StatusCode::INTERNAL_SERVER_ERROR,
3924                "SERIALIZATION_FAILED",
3925                "failed to serialize SQL idempotency request semantics",
3926                false,
3927            ))
3928        }
3929    };
3930    let begin = tokio::select! {
3931        begin = state.sql_idempotency.begin(owner, &key, binding) => begin,
3932        _ = registration.query().control().cancelled() => {
3933            return Err(tracked_query_error_response(
3934                state,
3935                &cancellation_checkpoint_error(registration.query()),
3936                Some(query_id),
3937            ));
3938        }
3939    };
3940    match begin {
3941        sql_idempotency::BeginResult::Execute(execution) => Ok((registration, Some(execution))),
3942        sql_idempotency::BeginResult::Replay {
3943            receipt,
3944            expires_at_ms,
3945        } => match restore_idempotency_replay(registration, &receipt) {
3946            Ok(()) => Err(sql_idempotency_receipt_response(
3947                query_id,
3948                &receipt,
3949                true,
3950                expires_at_ms,
3951                true,
3952            )),
3953            Err(error) => Err(tracked_query_error_response(state, &error, Some(query_id))),
3954        },
3955        sql_idempotency::BeginResult::Mismatch => Err(registered_sql_error_response(
3956            registration,
3957            query_id,
3958            StatusCode::CONFLICT,
3959            "IDEMPOTENCY_KEY_REUSE_MISMATCH",
3960            "idempotency key was already used with different SQL or request semantics",
3961            false,
3962        )),
3963        sql_idempotency::BeginResult::Indeterminate { created_at_ms } => Err(
3964            sql_idempotency_indeterminate_response(registration, query_id, created_at_ms),
3965        ),
3966        sql_idempotency::BeginResult::Full => Err(registered_sql_error_response(
3967            registration,
3968            query_id,
3969            StatusCode::SERVICE_UNAVAILABLE,
3970            "IDEMPOTENCY_STORE_FULL",
3971            "SQL idempotency receipt store is full",
3972            true,
3973        )),
3974        sql_idempotency::BeginResult::Unavailable => Err(registered_sql_error_response(
3975            registration,
3976            query_id,
3977            StatusCode::SERVICE_UNAVAILABLE,
3978            "IDEMPOTENCY_STORE_UNAVAILABLE",
3979            "could not durably reserve the SQL idempotency key",
3980            true,
3981        )),
3982    }
3983}
3984
3985fn restore_idempotency_replay(
3986    registration: RegisteredQueryGuard,
3987    receipt: &sql_idempotency::SqlDurableReceipt,
3988) -> mongreldb_query::Result<()> {
3989    use mongreldb_query::{
3990        DurableOutcome, QueryTerminalError, QueryTerminalErrorCategory, QueryTerminalState,
3991        SerializationOutcome,
3992    };
3993
3994    let invalid_receipt = |field: &str| {
3995        mongreldb_query::MongrelQueryError::InvalidQueryState(format!(
3996            "durable SQL idempotency receipt has invalid {field}"
3997        ))
3998    };
3999    let terminal_state = match receipt.status.as_str() {
4000        "completed" => QueryTerminalState::Completed,
4001        "failed_before_commit" => QueryTerminalState::FailedBeforeCommit,
4002        "cancelled_before_commit" => QueryTerminalState::CancelledBeforeCommit,
4003        "deadline_before_commit" => QueryTerminalState::DeadlineBeforeCommit,
4004        "committed" => QueryTerminalState::Committed,
4005        "committed_with_error" => QueryTerminalState::CommittedWithError,
4006        "partially_committed" => QueryTerminalState::PartiallyCommitted,
4007        "cancelled_after_commit" => QueryTerminalState::CancelledAfterCommit,
4008        "deadline_after_commit" => QueryTerminalState::DeadlineAfterCommit,
4009        _ => return Err(invalid_receipt("terminal state")),
4010    };
4011    let serialization = match receipt.outcome.serialization.as_str() {
4012        "not_started" => SerializationOutcome::NotStarted,
4013        "in_progress" => SerializationOutcome::InProgress,
4014        "succeeded" => SerializationOutcome::Succeeded,
4015        "failed" => SerializationOutcome::Failed,
4016        _ => return Err(invalid_receipt("serialization state")),
4017    };
4018    let terminal_error = match receipt.terminal_error.as_ref() {
4019        Some(error) => Some(QueryTerminalError {
4020            code: error.code.clone(),
4021            category: match error.category.as_str() {
4022                "cancellation" => QueryTerminalErrorCategory::Cancellation,
4023                "deadline" => QueryTerminalErrorCategory::Deadline,
4024                "result_limit" => QueryTerminalErrorCategory::ResultLimit,
4025                "serialization" => QueryTerminalErrorCategory::Serialization,
4026                "execution" => QueryTerminalErrorCategory::Execution,
4027                _ => return Err(invalid_receipt("terminal error category")),
4028            },
4029        }),
4030        None => None,
4031    };
4032    let cancellation_reason = CancellationReason::from_protocol_str(&receipt.cancellation_reason)
4033        .ok_or_else(|| invalid_receipt("cancellation reason"))?;
4034    let phase = match receipt.server_state.as_str() {
4035        "completed" => SqlQueryPhase::Completed,
4036        "cancelled" => SqlQueryPhase::Cancelled,
4037        "failed" => SqlQueryPhase::Failed,
4038        _ => return Err(invalid_receipt("server state")),
4039    };
4040    let query = registration.into_query();
4041    query.restore_replayed_outcome(
4042        DurableOutcome {
4043            committed: receipt.outcome.committed,
4044            committed_statements: receipt.outcome.committed_statements,
4045            last_commit_epoch: receipt.outcome.last_commit_epoch,
4046            first_commit_statement_index: receipt.outcome.first_commit_statement_index,
4047            last_commit_statement_index: receipt.outcome.last_commit_statement_index,
4048        },
4049        receipt.outcome.completed_statements,
4050        receipt.outcome.statement_index,
4051        serialization,
4052        terminal_error,
4053        terminal_state,
4054        cancellation_reason,
4055        phase,
4056    );
4057    query.try_complete()
4058}
4059
4060fn sql_idempotency_indeterminate_response(
4061    registration: RegisteredQueryGuard,
4062    query_id: QueryId,
4063    created_at_ms: Option<u64>,
4064) -> Response {
4065    registration.query().mark_outcome_unknown();
4066    registration.fail();
4067    with_query_id(
4068        (
4069            StatusCode::CONFLICT,
4070            Json(json!({
4071                "query_id": query_id.to_string(),
4072                "status": "outcome_unknown",
4073                "terminal_state": "outcome_unknown",
4074                "committed": null,
4075                "committed_statements": null,
4076                "last_commit_epoch": null,
4077                "last_commit_epoch_text": null,
4078                "first_commit_statement_index": null,
4079                "last_commit_statement_index": null,
4080                "completed_statements": null,
4081                "statement_index": null,
4082                "cancel_outcome": null,
4083                "cancellation_reason": null,
4084                "retryable": false,
4085                "server_state": "failed",
4086                "idempotency_replayed": true,
4087                "idempotency_intent_created_at_ms": created_at_ms,
4088                "outcome": {
4089                    "committed": null,
4090                    "committed_statements": null,
4091                    "last_commit_epoch": null,
4092                    "last_commit_epoch_text": null,
4093                    "first_commit_statement_index": null,
4094                    "last_commit_statement_index": null,
4095                    "completed_statements": null,
4096                    "statement_index": null,
4097                    "serialization": "unknown",
4098                },
4099                "error": {
4100                    "code": "QUERY_OUTCOME_UNKNOWN",
4101                    "message": "a durable write intent exists without a durable receipt; the SQL was not re-executed",
4102                    "query_id": query_id.to_string(),
4103                    "committed": null,
4104                    "retryable": false,
4105                }
4106            })),
4107        )
4108            .into_response(),
4109        query_id,
4110    )
4111}
4112
4113fn registered_sql_error_response(
4114    registration: RegisteredQueryGuard,
4115    query_id: QueryId,
4116    status: StatusCode,
4117    code: &'static str,
4118    message: impl Into<String>,
4119    retryable: bool,
4120) -> Response {
4121    let message = message.into();
4122    registration
4123        .query()
4124        .record_terminal_error(code, mongreldb_query::QueryTerminalErrorCategory::Execution);
4125    registration.fail();
4126    with_query_id(
4127        (
4128            status,
4129            Json(json!({
4130                "query_id": query_id.to_string(),
4131                "status": "failed_before_commit",
4132                "terminal_state": "failed_before_commit",
4133                "committed": false,
4134                "committed_statements": 0,
4135                "last_commit_epoch": null,
4136                "last_commit_epoch_text": null,
4137                "first_commit_statement_index": null,
4138                "last_commit_statement_index": null,
4139                "completed_statements": 0,
4140                "statement_index": 0,
4141                "cancel_outcome": null,
4142                "cancellation_reason": null,
4143                "retryable": retryable,
4144                "server_state": "failed",
4145                "outcome": {
4146                    "committed": false,
4147                    "committed_statements": 0,
4148                    "last_commit_epoch": null,
4149                    "last_commit_epoch_text": null,
4150                    "first_commit_statement_index": null,
4151                    "last_commit_statement_index": null,
4152                    "completed_statements": 0,
4153                    "statement_index": 0,
4154                    "serialization": "not_started",
4155                },
4156                "error": {
4157                    "code": code,
4158                    "message": message,
4159                    "query_id": query_id.to_string(),
4160                    "committed": false,
4161                    "retryable": retryable,
4162                }
4163            })),
4164        )
4165            .into_response(),
4166        query_id,
4167    )
4168}
4169
4170fn register_controlled_query(
4171    state: &AppState,
4172    session: &MongrelSession,
4173    options: SqlQueryOptions,
4174) -> std::result::Result<RegisteredSqlQuery, mongreldb_query::MongrelQueryError> {
4175    let query_id = options.query_id.ok_or_else(|| {
4176        mongreldb_query::MongrelQueryError::InvalidQueryState(
4177            "server query registration requires a query id".into(),
4178        )
4179    })?;
4180    let owner = options.owner.clone().unwrap_or_default();
4181    let session_id = options.session_id.clone();
4182    let _lifecycle = state
4183        .query_lifecycle
4184        .lock()
4185        .unwrap_or_else(|error| error.into_inner());
4186    let pre_cancel_reason = state
4187        .pre_cancellations
4188        .reason(query_id, &owner, session_id.as_deref());
4189    let query = session.register_query(options)?;
4190    let Some(reason) = pre_cancel_reason else {
4191        return Ok(query);
4192    };
4193    state
4194        .pre_cancellations
4195        .take(query_id, &owner, session_id.as_deref());
4196    query.request_cancel(reason);
4197    let error = query.checkpoint().err().unwrap_or_else(|| {
4198        mongreldb_query::MongrelQueryError::InvalidQueryState(format!(
4199            "pre-cancelled query {query_id} remained runnable"
4200        ))
4201    });
4202    query.fail();
4203    Err(error)
4204}
4205
4206async fn acquire_sql_permit(
4207    state: &AppState,
4208    session: &MongrelSession,
4209    query: &RegisteredSqlQuery,
4210) -> std::result::Result<tokio::sync::OwnedSemaphorePermit, mongreldb_query::MongrelQueryError> {
4211    session.fire_test_hook(mongreldb_query::SqlTestHookPoint::WaitingForSqlPermit);
4212    tokio::select! {
4213        permit = Arc::clone(&state.sql_semaphore).acquire_owned() => permit.map_err(|_| {
4214            mongreldb_query::MongrelQueryError::InvalidQueryState(
4215                "SQL admission semaphore closed".into(),
4216            )
4217        }),
4218        _ = query.control().cancelled() => Err(cancellation_checkpoint_error(query)),
4219    }
4220}
4221
4222fn caller_may_manage_query(
4223    state: &AppState,
4224    principal: &Option<mongreldb_core::Principal>,
4225    owner: Option<&str>,
4226) -> bool {
4227    let current = current_request_principal(state, principal);
4228    if (principal.is_some()
4229        || state.auth_token.is_some()
4230        || state.user_auth
4231        || state.db.require_auth_enabled())
4232        && current.is_none()
4233    {
4234        return false;
4235    }
4236    current.is_some_and(|principal| principal.is_admin)
4237        || owner == Some(request_owner(state, principal).as_str())
4238}
4239
4240fn query_phase_name(phase: SqlQueryPhase) -> &'static str {
4241    match phase {
4242        SqlQueryPhase::Queued => "queued",
4243        SqlQueryPhase::Planning => "planning",
4244        SqlQueryPhase::Executing => "executing",
4245        SqlQueryPhase::Streaming => "streaming",
4246        SqlQueryPhase::Serializing => "serializing",
4247        SqlQueryPhase::CommitCritical => "commit_critical",
4248        SqlQueryPhase::Cancelling => "cancelling",
4249        SqlQueryPhase::Completed => "completed",
4250        SqlQueryPhase::Failed => "failed",
4251        SqlQueryPhase::Cancelled => "cancelled",
4252    }
4253}
4254
4255fn commit_fence_outcome_name(outcome: mongreldb_query::CommitFenceOutcome) -> &'static str {
4256    match outcome {
4257        mongreldb_query::CommitFenceOutcome::NotReached => "not_reached",
4258        mongreldb_query::CommitFenceOutcome::CancelWon => "cancel_won",
4259        mongreldb_query::CommitFenceOutcome::CommitWon => "commit_won",
4260    }
4261}
4262
4263fn terminal_state_name(state: mongreldb_query::QueryTerminalState) -> &'static str {
4264    use mongreldb_query::QueryTerminalState;
4265    match state {
4266        QueryTerminalState::OutcomeUnknown => "outcome_unknown",
4267        QueryTerminalState::Completed => "completed",
4268        QueryTerminalState::FailedBeforeCommit => "failed_before_commit",
4269        QueryTerminalState::CancelledBeforeCommit => "cancelled_before_commit",
4270        QueryTerminalState::DeadlineBeforeCommit => "deadline_before_commit",
4271        QueryTerminalState::Committed => "committed",
4272        QueryTerminalState::CommittedWithError => "committed_with_error",
4273        QueryTerminalState::PartiallyCommitted => "partially_committed",
4274        QueryTerminalState::CancelledAfterCommit => "cancelled_after_commit",
4275        QueryTerminalState::DeadlineAfterCommit => "deadline_after_commit",
4276    }
4277}
4278
4279fn serialization_outcome_name(outcome: mongreldb_query::SerializationOutcome) -> &'static str {
4280    use mongreldb_query::SerializationOutcome;
4281    match outcome {
4282        SerializationOutcome::NotStarted => "not_started",
4283        SerializationOutcome::InProgress => "in_progress",
4284        SerializationOutcome::Succeeded => "succeeded",
4285        SerializationOutcome::Failed => "failed",
4286    }
4287}
4288
4289fn terminal_error_category_name(
4290    category: mongreldb_query::QueryTerminalErrorCategory,
4291) -> &'static str {
4292    use mongreldb_query::QueryTerminalErrorCategory;
4293    match category {
4294        QueryTerminalErrorCategory::Cancellation => "cancellation",
4295        QueryTerminalErrorCategory::Deadline => "deadline",
4296        QueryTerminalErrorCategory::ResultLimit => "result_limit",
4297        QueryTerminalErrorCategory::Serialization => "serialization",
4298        QueryTerminalErrorCategory::Execution => "execution",
4299    }
4300}
4301
4302fn terminal_error_retryable(error: Option<&mongreldb_query::QueryTerminalError>) -> bool {
4303    error.is_some_and(|error| {
4304        matches!(
4305            error.code.as_str(),
4306            "IDEMPOTENCY_STORE_FULL" | "IDEMPOTENCY_STORE_UNAVAILABLE"
4307        )
4308    })
4309}
4310
4311fn epoch_text(epoch: Option<u64>) -> Option<String> {
4312    epoch.map(|epoch| epoch.to_string())
4313}
4314
4315fn cancellation_reason_name(reason: CancellationReason) -> &'static str {
4316    reason.as_str()
4317}
4318
4319fn query_cancel_outcome(status: &mongreldb_query::QueryStatus) -> Option<&'static str> {
4320    match status.phase {
4321        SqlQueryPhase::CommitCritical => Some("too_late"),
4322        SqlQueryPhase::Completed | SqlQueryPhase::Failed | SqlQueryPhase::Cancelled => {
4323            Some("already_finished")
4324        }
4325        SqlQueryPhase::Cancelling => Some("accepted"),
4326        _ => None,
4327    }
4328}
4329
4330fn query_outcome_json(status: Option<&mongreldb_query::QueryStatus>) -> serde_json::Value {
4331    let Some(status) = status else {
4332        return json!({
4333            "committed": false,
4334            "committed_statements": 0,
4335            "last_commit_epoch": null,
4336            "last_commit_epoch_text": null,
4337            "first_commit_statement_index": null,
4338            "last_commit_statement_index": null,
4339            "completed_statements": 0,
4340            "statement_index": 0,
4341            "serialization": "not_started",
4342        });
4343    };
4344    if status.outcome_unknown {
4345        return json!({
4346            "committed": null,
4347            "committed_statements": null,
4348            "last_commit_epoch": null,
4349            "last_commit_epoch_text": null,
4350            "first_commit_statement_index": null,
4351            "last_commit_statement_index": null,
4352            "completed_statements": null,
4353            "statement_index": null,
4354            "serialization": "unknown",
4355        });
4356    }
4357    json!({
4358        "committed": status.durable_outcome.committed,
4359        "committed_statements": (!status.outcome_unknown).then_some(status.durable_outcome.committed_statements),
4360        "last_commit_epoch": (!status.outcome_unknown).then_some(status.durable_outcome.last_commit_epoch).flatten(),
4361        "last_commit_epoch_text": (!status.outcome_unknown).then_some(epoch_text(status.durable_outcome.last_commit_epoch)).flatten(),
4362        "first_commit_statement_index": (!status.outcome_unknown).then_some(status.durable_outcome.first_commit_statement_index).flatten(),
4363        "last_commit_statement_index": (!status.outcome_unknown).then_some(status.durable_outcome.last_commit_statement_index).flatten(),
4364        "completed_statements": (!status.outcome_unknown).then_some(status.completed_statements),
4365        "statement_index": (!status.outcome_unknown).then_some(status.statement_index),
4366        "serialization": serialization_outcome_name(status.serialization_outcome),
4367    })
4368}
4369
4370fn sql_terminal_idempotency_receipt(
4371    status: &mongreldb_query::QueryStatus,
4372) -> Option<sql_idempotency::SqlDurableReceipt> {
4373    if status.outcome_unknown {
4374        return None;
4375    }
4376    let terminal_state = status.terminal_state()?;
4377    if !status.durable_outcome.committed
4378        && terminal_state != mongreldb_query::QueryTerminalState::Completed
4379    {
4380        return None;
4381    }
4382    Some(sql_idempotency::SqlDurableReceipt {
4383        original_query_id: status.query_id.to_string(),
4384        status: status
4385            .terminal_state()
4386            .map(terminal_state_name)
4387            .unwrap_or("committed")
4388            .to_owned(),
4389        server_state: query_phase_name(status.phase).to_owned(),
4390        cancellation_reason: cancellation_reason_name(status.cancellation_reason).to_owned(),
4391        outcome: sql_idempotency::SqlReceiptOutcome {
4392            committed: status.durable_outcome.committed,
4393            committed_statements: status.durable_outcome.committed_statements,
4394            last_commit_epoch: status.durable_outcome.last_commit_epoch,
4395            last_commit_epoch_text: epoch_text(status.durable_outcome.last_commit_epoch),
4396            first_commit_statement_index: status.durable_outcome.first_commit_statement_index,
4397            last_commit_statement_index: status.durable_outcome.last_commit_statement_index,
4398            completed_statements: status.completed_statements,
4399            statement_index: status.statement_index,
4400            serialization: serialization_outcome_name(status.serialization_outcome).to_owned(),
4401        },
4402        terminal_error: status.terminal_error.as_ref().map(|error| {
4403            sql_idempotency::SqlReceiptTerminalError {
4404                code: error.code.clone(),
4405                category: terminal_error_category_name(error.category).to_owned(),
4406            }
4407        }),
4408    })
4409}
4410
4411fn sql_idempotency_receipt_response(
4412    query_id: QueryId,
4413    receipt: &sql_idempotency::SqlDurableReceipt,
4414    replayed: bool,
4415    expires_at_ms: u64,
4416    persisted: bool,
4417) -> Response {
4418    let mut response = Json(json!({
4419        "query_id": query_id.to_string(),
4420        "original_query_id": receipt.original_query_id,
4421        "status": receipt.status,
4422        "terminal_state": receipt.status,
4423        "server_state": receipt.server_state,
4424        "cancel_outcome": "already_finished",
4425        "cancellation_reason": receipt.cancellation_reason,
4426        "committed": receipt.outcome.committed,
4427        "committed_statements": receipt.outcome.committed_statements,
4428        "last_commit_epoch": receipt.outcome.last_commit_epoch,
4429        "last_commit_epoch_text": receipt.outcome.last_commit_epoch_text.as_deref(),
4430        "first_commit_statement_index": receipt.outcome.first_commit_statement_index,
4431        "last_commit_statement_index": receipt.outcome.last_commit_statement_index,
4432        "completed_statements": receipt.outcome.completed_statements,
4433        "statement_index": receipt.outcome.statement_index,
4434        "retryable": false,
4435        "idempotency_replayed": replayed,
4436        "idempotency_persisted": persisted,
4437        "idempotency_expires_at_ms": expires_at_ms,
4438        "outcome": receipt.outcome,
4439        "terminal_error": receipt.terminal_error,
4440    }))
4441    .into_response();
4442    response.headers_mut().insert(
4443        "idempotency-replayed",
4444        axum::http::HeaderValue::from_static(if replayed { "true" } else { "false" }),
4445    );
4446    response.headers_mut().insert(
4447        "idempotency-persisted",
4448        axum::http::HeaderValue::from_static(if persisted { "true" } else { "false" }),
4449    );
4450    if let Ok(value) = axum::http::HeaderValue::from_str(&receipt.original_query_id) {
4451        response
4452            .headers_mut()
4453            .insert("x-mongreldb-original-query-id", value);
4454    }
4455    with_query_id(response, query_id)
4456}
4457
4458fn terminal_server_error_response(
4459    state: &AppState,
4460    query_id: QueryId,
4461    http_status: StatusCode,
4462    base_code: &'static str,
4463    message: impl Into<String>,
4464) -> Response {
4465    let status = state.query_registry.status(query_id);
4466    let committed = status
4467        .as_ref()
4468        .is_some_and(|status| status.durable_outcome.committed);
4469    let code = if committed && base_code.starts_with("SERIALIZATION_") {
4470        "SERIALIZATION_FAILED_AFTER_COMMIT"
4471    } else {
4472        base_code
4473    };
4474    let response_status = status
4475        .as_ref()
4476        .and_then(mongreldb_query::QueryStatus::terminal_state)
4477        .map(terminal_state_name)
4478        .unwrap_or(if committed {
4479            "committed_with_error"
4480        } else {
4481            "failed_before_commit"
4482        });
4483    let outcome = query_outcome_json(status.as_ref());
4484    with_query_id(
4485        (
4486            http_status,
4487            Json(json!({
4488                "query_id": query_id.to_string(),
4489                "status": response_status,
4490                "terminal_state": response_status,
4491                "committed": committed,
4492                "committed_statements": status.as_ref().map_or(0, |status| status.durable_outcome.committed_statements),
4493                "last_commit_epoch": status.as_ref().and_then(|status| status.durable_outcome.last_commit_epoch),
4494                "last_commit_epoch_text": epoch_text(status.as_ref().and_then(|status| status.durable_outcome.last_commit_epoch)),
4495                "first_commit_statement_index": status.as_ref().and_then(|status| status.durable_outcome.first_commit_statement_index),
4496                "last_commit_statement_index": status.as_ref().and_then(|status| status.durable_outcome.last_commit_statement_index),
4497                "completed_statements": status.as_ref().map_or(0, |status| status.completed_statements),
4498                "statement_index": status.as_ref().map_or(0, |status| status.statement_index),
4499                "cancel_outcome": null,
4500                "cancellation_reason": status.as_ref().map(|status| cancellation_reason_name(status.cancellation_reason)),
4501                "retryable": false,
4502                "server_state": status.as_ref().map(|status| query_phase_name(status.phase)),
4503                "outcome": outcome,
4504                "error": {
4505                    "code": code,
4506                    "message": message.into(),
4507                    "query_id": query_id.to_string(),
4508                    "committed": committed,
4509                    "retryable": false,
4510                }
4511            })),
4512        )
4513            .into_response(),
4514        query_id,
4515    )
4516}
4517
4518fn query_not_found_response(query_id: Option<QueryId>) -> Response {
4519    let mut response = (
4520        StatusCode::NOT_FOUND,
4521        Json(json!({
4522            "query_id": query_id.map(|value| value.to_string()),
4523            "status": "unknown",
4524            "terminal_state": null,
4525            "committed": null,
4526            "committed_statements": null,
4527            "last_commit_epoch": null,
4528            "last_commit_epoch_text": null,
4529            "first_commit_statement_index": null,
4530            "last_commit_statement_index": null,
4531            "completed_statements": null,
4532            "statement_index": null,
4533            "cancel_outcome": "not_found",
4534            "cancellation_reason": null,
4535            "retryable": false,
4536            "server_state": "not_found",
4537            "outcome": {
4538                "committed": null,
4539                "committed_statements": null,
4540                "last_commit_epoch": null,
4541                "last_commit_epoch_text": null,
4542                "first_commit_statement_index": null,
4543                "last_commit_statement_index": null,
4544                "completed_statements": null,
4545                "statement_index": null,
4546                "serialization": "unknown",
4547            },
4548            "error": {
4549                "code": "QUERY_NOT_FOUND",
4550                "message": "query not found",
4551                "query_id": query_id.map(|value| value.to_string()),
4552                "committed": null,
4553                "retryable": false,
4554            }
4555        })),
4556    )
4557        .into_response();
4558    if let Some(query_id) = query_id {
4559        add_query_id_header(&mut response, query_id);
4560    }
4561    response
4562}
4563
4564fn query_session_header(
4565    headers: &axum::http::HeaderMap,
4566    query_id: Option<QueryId>,
4567) -> std::result::Result<Option<String>, Box<Response>> {
4568    match headers.get("x-session-id") {
4569        Some(value) => match value.to_str() {
4570            Ok(value) if value.len() <= 256 => Ok(Some(value.to_owned())),
4571            _ => Err(Box::new(bad_query_control_request(
4572                "X-Session-ID must be valid text no longer than 256 bytes",
4573                query_id,
4574            ))),
4575        },
4576        None => Ok(None),
4577    }
4578}
4579
4580fn pre_cancelled_query_response(
4581    query_id: QueryId,
4582    reason: CancellationReason,
4583    status: StatusCode,
4584) -> Response {
4585    with_query_id(
4586        (
4587            status,
4588            Json(json!({
4589                "query_id": query_id.to_string(),
4590                "status": "cancelled_before_start",
4591                "terminal_state": "cancelled_before_start",
4592                "state": "pre_cancelled",
4593                "server_state": "pre_cancelled",
4594                "cancel_outcome": "pre_cancelled",
4595                "committed": false,
4596                "committed_statements": 0,
4597                "last_commit_epoch": null,
4598                "last_commit_epoch_text": null,
4599                "first_commit_statement_index": null,
4600                "last_commit_statement_index": null,
4601                "completed_statements": 0,
4602                "statement_index": 0,
4603                "cancellation_reason": cancellation_reason_name(reason),
4604                "outcome": {
4605                    "committed": false,
4606                    "committed_statements": 0,
4607                    "last_commit_epoch": null,
4608                    "last_commit_epoch_text": null,
4609                    "first_commit_statement_index": null,
4610                    "last_commit_statement_index": null,
4611                    "completed_statements": 0,
4612                    "statement_index": 0,
4613                    "serialization": "not_started",
4614                },
4615                "terminal_error": {
4616                    "code": "QUERY_CANCELLED",
4617                    "category": "cancellation",
4618                },
4619                "retryable": false,
4620            })),
4621        )
4622            .into_response(),
4623        query_id,
4624    )
4625}
4626
4627fn compact_finished_query_response(query_id: QueryId) -> Response {
4628    with_query_id(
4629        Json(json!({
4630            "query_id": query_id.to_string(),
4631            "status": "finished",
4632            "terminal_state": null,
4633            "state": "finished",
4634            "server_state": "finished",
4635            "cancel_outcome": "already_finished",
4636            "code": "QUERY_ALREADY_FINISHED",
4637            "committed": null,
4638            "committed_statements": null,
4639            "last_commit_epoch": null,
4640            "last_commit_epoch_text": null,
4641            "first_commit_statement_index": null,
4642            "last_commit_statement_index": null,
4643            "completed_statements": null,
4644            "statement_index": null,
4645            "cancellation_reason": "none",
4646            "outcome": {
4647                "committed": null,
4648                "committed_statements": null,
4649                "last_commit_epoch": null,
4650                "last_commit_epoch_text": null,
4651                "first_commit_statement_index": null,
4652                "last_commit_statement_index": null,
4653                "completed_statements": null,
4654                "statement_index": null,
4655                "serialization": "unknown",
4656            },
4657            "terminal_error": null,
4658            "retryable": false,
4659        }))
4660        .into_response(),
4661        query_id,
4662    )
4663}
4664
4665async fn query_status(
4666    State(state): State<Arc<AppState>>,
4667    OptionalPrincipal(principal): OptionalPrincipal,
4668    Path(query_id): Path<String>,
4669    headers: axum::http::HeaderMap,
4670) -> Response {
4671    let Ok(query_id) = query_id.parse::<QueryId>() else {
4672        return query_not_found_response(None);
4673    };
4674    if !request_identity_is_current(&state, &principal) {
4675        return query_not_found_response(Some(query_id));
4676    }
4677    let requested_session = match query_session_header(&headers, Some(query_id)) {
4678        Ok(session_id) => session_id,
4679        Err(response) => return *response,
4680    };
4681    let owner = request_owner(&state, &principal);
4682    let is_admin =
4683        current_request_principal(&state, &principal).is_some_and(|principal| principal.is_admin);
4684    let _lifecycle = state
4685        .query_lifecycle
4686        .lock()
4687        .unwrap_or_else(|error| error.into_inner());
4688    let Some(status) = state.query_registry.status(query_id) else {
4689        if let Some((finished_owner, finished_session)) =
4690            state.query_registry.compact_finished_identity(query_id)
4691        {
4692            if !caller_may_manage_query(&state, &principal, finished_owner.as_deref())
4693                || requested_session
4694                    .as_deref()
4695                    .is_some_and(|session| finished_session.as_deref() != Some(session))
4696            {
4697                return query_not_found_response(Some(query_id));
4698            }
4699            return compact_finished_query_response(query_id);
4700        }
4701        let reason = state
4702            .pre_cancellations
4703            .reason(query_id, &owner, requested_session.as_deref())
4704            .or_else(|| {
4705                is_admin.then(|| match requested_session.as_deref() {
4706                    Some(session_id) => state
4707                        .pre_cancellations
4708                        .reason_for_query_in_session(query_id, session_id),
4709                    None => state.pre_cancellations.reason_for_query(query_id),
4710                })?
4711            });
4712        if let Some(reason) = reason {
4713            return pre_cancelled_query_response(query_id, reason, StatusCode::OK);
4714        }
4715        return query_not_found_response(Some(query_id));
4716    };
4717    if !caller_may_manage_query(&state, &principal, status.owner.as_deref())
4718        || requested_session
4719            .as_deref()
4720            .is_some_and(|session| status.session_id.as_deref() != Some(session))
4721    {
4722        return query_not_found_response(Some(query_id));
4723    }
4724    let terminal_status = status.terminal_state().map(terminal_state_name);
4725    let terminal_error = status.terminal_error.as_ref().map(|error| {
4726        json!({
4727            "code": error.code,
4728            "category": terminal_error_category_name(error.category),
4729        })
4730    });
4731    let retryable = terminal_error_retryable(status.terminal_error.as_ref());
4732    let cancel_outcome = query_cancel_outcome(&status);
4733    let outcome = query_outcome_json(Some(&status));
4734    let response = Json(json!({
4735        "query_id": query_id.to_string(),
4736        "status": terminal_status.unwrap_or(if status.durable_outcome.committed {
4737            "committed"
4738        } else {
4739            "running"
4740        }),
4741        "terminal_state": terminal_status,
4742        "state": query_phase_name(status.phase),
4743        "server_state": query_phase_name(status.phase),
4744        "started_ms_ago": status.started_at.elapsed().as_millis(),
4745        "deadline_ms_remaining": status.deadline.map(|deadline| {
4746            deadline.saturating_duration_since(std::time::Instant::now()).as_millis()
4747        }),
4748        "session_id": status.session_id,
4749        "operation": status.operation,
4750        "committed": (!status.outcome_unknown).then_some(status.committed),
4751        "committed_statements": (!status.outcome_unknown).then_some(status.durable_outcome.committed_statements),
4752        "last_commit_epoch": (!status.outcome_unknown).then_some(status.durable_outcome.last_commit_epoch).flatten(),
4753        "last_commit_epoch_text": (!status.outcome_unknown).then_some(epoch_text(status.durable_outcome.last_commit_epoch)).flatten(),
4754        "first_commit_statement_index": (!status.outcome_unknown).then_some(status.durable_outcome.first_commit_statement_index).flatten(),
4755        "last_commit_statement_index": (!status.outcome_unknown).then_some(status.durable_outcome.last_commit_statement_index).flatten(),
4756        "cancellation_reason": cancellation_reason_name(status.cancellation_reason),
4757        "completed_statements": (!status.outcome_unknown).then_some(status.completed_statements),
4758        "statement_index": (!status.outcome_unknown).then_some(status.statement_index),
4759        "cancel_outcome": cancel_outcome,
4760        "retryable": retryable,
4761        "outcome": outcome,
4762        "terminal_error": terminal_error,
4763        "trace": {
4764            "queue_duration_us": status.queue_duration.as_micros(),
4765            "planning_duration_us": status.planning_duration.as_micros(),
4766            "execution_duration_us": status.execution_duration.as_micros(),
4767            "serialization_duration_us": status.serialization_duration.as_micros(),
4768            "cancel_requested_phase": status.cancel_requested_phase.map(query_phase_name),
4769            "cancel_observed_phase": status.cancel_observed_phase.map(query_phase_name),
4770            "commit_fence_outcome": commit_fence_outcome_name(status.commit_fence_outcome),
4771        },
4772    }))
4773    .into_response();
4774    with_query_id(response, query_id)
4775}
4776
4777async fn cancel_query(
4778    State(state): State<Arc<AppState>>,
4779    OptionalPrincipal(principal): OptionalPrincipal,
4780    Path(query_id): Path<String>,
4781    headers: axum::http::HeaderMap,
4782) -> Response {
4783    let Ok(query_id) = query_id.parse::<QueryId>() else {
4784        return query_not_found_response(None);
4785    };
4786    if !request_identity_is_current(&state, &principal) {
4787        return query_not_found_response(Some(query_id));
4788    }
4789    let requested_session = match query_session_header(&headers, Some(query_id)) {
4790        Ok(session_id) => session_id,
4791        Err(response) => return *response,
4792    };
4793    let owner = request_owner(&state, &principal);
4794    let _lifecycle = state
4795        .query_lifecycle
4796        .lock()
4797        .unwrap_or_else(|error| error.into_inner());
4798    state.metrics.inc_sql_cancel_requests();
4799    let Some(status) = state.query_registry.status(query_id) else {
4800        if let Some((finished_owner, finished_session)) =
4801            state.query_registry.compact_finished_identity(query_id)
4802        {
4803            if !caller_may_manage_query(&state, &principal, finished_owner.as_deref())
4804                || requested_session
4805                    .as_deref()
4806                    .is_some_and(|session| finished_session.as_deref() != Some(session))
4807            {
4808                return query_not_found_response(Some(query_id));
4809            }
4810            return compact_finished_query_response(query_id);
4811        }
4812        return match state.pre_cancellations.insert(
4813            query_id,
4814            &owner,
4815            requested_session.as_deref(),
4816            CancellationReason::ClientRequest,
4817        ) {
4818            Ok(()) => {
4819                state.metrics.inc_sql_commit_cancel_winner_cancel();
4820                pre_cancelled_query_response(
4821                    query_id,
4822                    CancellationReason::ClientRequest,
4823                    StatusCode::ACCEPTED,
4824                )
4825            }
4826            Err(pre_cancel::InsertError::MetadataTooLarge) => bad_query_control_request(
4827                "query owner or session metadata exceeds 256 bytes",
4828                Some(query_id),
4829            ),
4830            Err(
4831                pre_cancel::InsertError::Full
4832                | pre_cancel::InsertError::OwnerLimit
4833                | pre_cancel::InsertError::RateLimited,
4834            ) => with_query_id(
4835                (
4836                    StatusCode::TOO_MANY_REQUESTS,
4837                    Json(json!({
4838                        "query_id": query_id.to_string(),
4839                        "status": "failed_before_commit",
4840                        "terminal_state": "failed_before_commit",
4841                        "server_state": "failed",
4842                        "cancel_outcome": null,
4843                        "cancellation_reason": null,
4844                        "committed": false,
4845                        "committed_statements": 0,
4846                        "last_commit_epoch": null,
4847                        "last_commit_epoch_text": null,
4848                        "first_commit_statement_index": null,
4849                        "last_commit_statement_index": null,
4850                        "completed_statements": 0,
4851                        "statement_index": 0,
4852                        "retryable": true,
4853                        "outcome": {
4854                            "committed": false,
4855                            "committed_statements": 0,
4856                            "last_commit_epoch": null,
4857                            "last_commit_epoch_text": null,
4858                            "first_commit_statement_index": null,
4859                            "last_commit_statement_index": null,
4860                            "completed_statements": 0,
4861                            "statement_index": 0,
4862                            "serialization": "not_started",
4863                        },
4864                        "error": {
4865                            "code": "QUERY_REGISTRY_FULL",
4866                            "message": "pre-registration cancellation limit reached",
4867                            "query_id": query_id.to_string(),
4868                            "committed": false,
4869                            "retryable": true,
4870                        }
4871                    })),
4872                )
4873                    .into_response(),
4874                query_id,
4875            ),
4876        };
4877    };
4878    if !caller_may_manage_query(&state, &principal, status.owner.as_deref())
4879        || requested_session
4880            .as_deref()
4881            .is_some_and(|session| status.session_id.as_deref() != Some(session))
4882    {
4883        return query_not_found_response(Some(query_id));
4884    }
4885    let (http_status, mut body) = match state.query_registry.cancel(query_id) {
4886        CancelOutcome::Accepted => (
4887            {
4888                state.metrics.inc_sql_commit_cancel_winner_cancel();
4889                StatusCode::ACCEPTED
4890            },
4891            json!({
4892                "query_id": query_id.to_string(),
4893                "state": "cancellation_requested",
4894                "cancel_outcome": "accepted",
4895            }),
4896        ),
4897        CancelOutcome::AlreadyCancelling => (
4898            StatusCode::OK,
4899            json!({
4900                "query_id": query_id.to_string(),
4901                "state": "cancelling",
4902                "cancel_outcome": "already_cancelling",
4903            }),
4904        ),
4905        CancelOutcome::TooLate => (
4906            {
4907                state.metrics.inc_sql_commit_cancel_winner_commit();
4908                StatusCode::CONFLICT
4909            },
4910            json!({
4911                "query_id": query_id.to_string(),
4912                "state": "commit_critical",
4913                "cancel_outcome": "too_late",
4914                "committed": status.durable_outcome.committed,
4915                "outcome": query_outcome_json(Some(&status)),
4916                "retryable": false,
4917                "error": {
4918                    "code": "CANCEL_TOO_LATE",
4919                    "message": "the query has entered its durable commit phase",
4920                    "committed": status.durable_outcome.committed,
4921                    "retryable": false,
4922                }
4923            }),
4924        ),
4925        CancelOutcome::AlreadyFinished => (
4926            StatusCode::OK,
4927            json!({
4928                "query_id": query_id.to_string(),
4929                "state": "finished",
4930                "status": status.terminal_state().map(terminal_state_name),
4931                "cancel_outcome": "already_finished",
4932                "code": "QUERY_ALREADY_FINISHED",
4933                "committed": status.durable_outcome.committed,
4934                "outcome": query_outcome_json(Some(&status)),
4935                "retryable": false,
4936            }),
4937        ),
4938        CancelOutcome::NotFound => return query_not_found_response(Some(query_id)),
4939    };
4940    let status = state.query_registry.status(query_id).unwrap_or(status);
4941    let response_status = status.terminal_state().map(terminal_state_name).unwrap_or(
4942        if status.durable_outcome.committed {
4943            "committed"
4944        } else {
4945            "running"
4946        },
4947    );
4948    if let Some(body) = body.as_object_mut() {
4949        body.insert("status".into(), json!(response_status));
4950        body.insert(
4951            "terminal_state".into(),
4952            json!(status.terminal_state().map(terminal_state_name)),
4953        );
4954        body.insert(
4955            "committed".into(),
4956            json!((!status.outcome_unknown).then_some(status.durable_outcome.committed)),
4957        );
4958        body.insert(
4959            "committed_statements".into(),
4960            json!((!status.outcome_unknown).then_some(status.durable_outcome.committed_statements)),
4961        );
4962        body.insert(
4963            "last_commit_epoch".into(),
4964            json!((!status.outcome_unknown)
4965                .then_some(status.durable_outcome.last_commit_epoch)
4966                .flatten()),
4967        );
4968        body.insert(
4969            "last_commit_epoch_text".into(),
4970            json!((!status.outcome_unknown)
4971                .then_some(epoch_text(status.durable_outcome.last_commit_epoch))
4972                .flatten()),
4973        );
4974        body.insert(
4975            "first_commit_statement_index".into(),
4976            json!((!status.outcome_unknown)
4977                .then_some(status.durable_outcome.first_commit_statement_index)
4978                .flatten()),
4979        );
4980        body.insert(
4981            "last_commit_statement_index".into(),
4982            json!((!status.outcome_unknown)
4983                .then_some(status.durable_outcome.last_commit_statement_index)
4984                .flatten()),
4985        );
4986        body.insert(
4987            "completed_statements".into(),
4988            json!((!status.outcome_unknown).then_some(status.completed_statements)),
4989        );
4990        body.insert(
4991            "statement_index".into(),
4992            json!((!status.outcome_unknown).then_some(status.statement_index)),
4993        );
4994        body.insert(
4995            "cancellation_reason".into(),
4996            json!(cancellation_reason_name(status.cancellation_reason)),
4997        );
4998        body.insert("retryable".into(), json!(false));
4999        body.insert("server_state".into(), json!(query_phase_name(status.phase)));
5000        body.insert("outcome".into(), query_outcome_json(Some(&status)));
5001    }
5002    with_query_id((http_status, Json(body)).into_response(), query_id)
5003}
5004
5005#[derive(Deserialize)]
5006struct SqlContinuationRequest {
5007    cursor: String,
5008}
5009
5010async fn continue_sql_page(
5011    State(state): State<Arc<AppState>>,
5012    OptionalPrincipal(principal): OptionalPrincipal,
5013    Json(request): Json<SqlContinuationRequest>,
5014) -> Response {
5015    if request.cursor.len() > 2_048 {
5016        return sql_cursor_error_response(
5017            StatusCode::BAD_REQUEST,
5018            "INVALID_SQL_CURSOR",
5019            "invalid SQL continuation cursor",
5020        );
5021    }
5022    if !request_identity_is_current(&state, &principal) {
5023        return sql_cursor_error_response(
5024            StatusCode::NOT_FOUND,
5025            "SQL_CURSOR_NOT_FOUND",
5026            "SQL continuation result is unavailable",
5027        );
5028    }
5029    let owner = request_owner(&state, &principal);
5030    let _permit = match state.sql_semaphore.acquire().await {
5031        Ok(permit) => permit,
5032        Err(_) => {
5033            return sql_cursor_error_response(
5034                StatusCode::SERVICE_UNAVAILABLE,
5035                "SQL_ADMISSION_CLOSED",
5036                "SQL continuation admission is closed",
5037            );
5038        }
5039    };
5040    let cursor_mac_key = match state.cursor_mac_key.get() {
5041        Ok(key) => key,
5042        Err(_) => {
5043            return sql_cursor_error_response(
5044                StatusCode::INTERNAL_SERVER_ERROR,
5045                "ENTROPY_UNAVAILABLE",
5046                "OS CSPRNG unavailable",
5047            );
5048        }
5049    };
5050    match state.sql_pages.continue_page(
5051        &request.cursor,
5052        &owner,
5053        &cursor_mac_key,
5054        sql_pages::SqlPageBinding {
5055            security_version: state.db.security_version(),
5056            catalog_epoch: state.db.catalog_snapshot().db_epoch,
5057        },
5058    ) {
5059        Ok(page) => {
5060            let page_byte_count = page.byte_count;
5061            match tokio::task::spawn_blocking(move || serialize_sql_page(page)).await {
5062                Ok(Ok(body)) => {
5063                    state.metrics.add_sql_output_bytes(page_byte_count);
5064                    sql_page_response(body)
5065                }
5066                Ok(Err(_)) => sql_cursor_error_response(
5067                    StatusCode::INTERNAL_SERVER_ERROR,
5068                    "SERIALIZATION_FAILED",
5069                    "failed to serialize SQL continuation page",
5070                ),
5071                Err(_) => sql_cursor_error_response(
5072                    StatusCode::INTERNAL_SERVER_ERROR,
5073                    "SERIALIZATION_WORKER_FAILED",
5074                    "SQL continuation serialization worker failed",
5075                ),
5076            }
5077        }
5078        Err(sql_pages::CursorError::Invalid) => sql_cursor_error_response(
5079            StatusCode::BAD_REQUEST,
5080            "INVALID_SQL_CURSOR",
5081            "invalid SQL continuation cursor",
5082        ),
5083        Err(sql_pages::CursorError::Expired) => sql_cursor_error_response(
5084            StatusCode::GONE,
5085            "SQL_CURSOR_EXPIRED",
5086            "SQL continuation cursor expired",
5087        ),
5088        Err(sql_pages::CursorError::NotFound) => sql_cursor_error_response(
5089            StatusCode::NOT_FOUND,
5090            "SQL_CURSOR_NOT_FOUND",
5091            "SQL continuation result is unavailable",
5092        ),
5093        Err(sql_pages::CursorError::PageLimit) => sql_cursor_error_response(
5094            StatusCode::PAYLOAD_TOO_LARGE,
5095            "RESULT_LIMIT_EXCEEDED",
5096            "one projected row exceeds the page byte or token limit",
5097        ),
5098    }
5099}
5100
5101fn sql_cursor_error_response(
5102    status: StatusCode,
5103    code: &'static str,
5104    message: &'static str,
5105) -> Response {
5106    (
5107        status,
5108        Json(json!({
5109            "status": "failed_before_commit",
5110            "terminal_state": "failed_before_commit",
5111            "server_state": "failed",
5112            "committed": false,
5113            "committed_statements": 0,
5114            "last_commit_epoch": null,
5115            "last_commit_epoch_text": null,
5116            "first_commit_statement_index": null,
5117            "last_commit_statement_index": null,
5118            "completed_statements": 0,
5119            "statement_index": 0,
5120            "cancel_outcome": null,
5121            "cancellation_reason": null,
5122            "retryable": false,
5123            "outcome": {
5124                "committed": false,
5125                "committed_statements": 0,
5126                "last_commit_epoch": null,
5127                "last_commit_epoch_text": null,
5128                "first_commit_statement_index": null,
5129                "last_commit_statement_index": null,
5130                "completed_statements": 0,
5131                "statement_index": 0,
5132                "serialization": "not_started",
5133            },
5134            "error": {
5135                "code": code,
5136                "message": message,
5137                "committed": false,
5138                "retryable": false,
5139            }
5140        })),
5141    )
5142        .into_response()
5143}
5144
5145async fn sql(
5146    State(state): State<Arc<AppState>>,
5147    OptionalPrincipal(principal): OptionalPrincipal,
5148    headers: axum::http::HeaderMap,
5149    Json(req): Json<SqlRequest>,
5150) -> Response {
5151    if !state.accepting_sql.load(Ordering::Acquire) {
5152        return (StatusCode::SERVICE_UNAVAILABLE, "server is shutting down").into_response();
5153    }
5154    if !request_identity_is_current(&state, &principal) {
5155        return StatusCode::UNAUTHORIZED.into_response();
5156    }
5157    // Session routing: an `X-Session-ID` header routes the request to a pooled
5158    // long-lived session, enabling cross-request `BEGIN`/`INSERT`/`COMMIT`
5159    // transactions. Without the header, a fresh ephemeral session is used
5160    // (the historical behavior).
5161    let session_id = match query_session_header(&headers, None) {
5162        Ok(session_id) => session_id,
5163        Err(response) => return *response,
5164    };
5165
5166    let owner = request_owner(&state, &principal);
5167    if let Some(sid) = session_id {
5168        let Some(entry) = state.sessions.get(&sid, &owner) else {
5169            return (
5170                StatusCode::NOT_FOUND,
5171                "session not found or not owned by caller",
5172            )
5173                .into_response();
5174        };
5175        let (options, query_id) = match resolve_query_options(
5176            &state,
5177            &headers,
5178            req.query_id,
5179            req.timeout_ms,
5180            owner.clone(),
5181            Some(sid.clone()),
5182        ) {
5183            Ok(options) => options,
5184            Err(response) => return *response,
5185        };
5186        let query = match register_controlled_query(&state, &entry.session, options) {
5187            Ok(query) => query,
5188            Err(error) => return tracked_query_error_response(&state, &error, Some(query_id)),
5189        };
5190        let registration = RegisteredQueryGuard::new(query);
5191        if mongreldb_query::contains_boolean_ai_predicate(&req.sql) {
5192            registration.fail();
5193            return with_query_id(remote_boolean_ai_error(), query_id);
5194        }
5195        let output_limits = match resolve_sql_output_limits(&state, &req, query_id) {
5196            Ok(limits) => limits,
5197            Err(response) => {
5198                registration.fail();
5199                return *response;
5200            }
5201        };
5202        let (registration, pagination) =
5203            match resolve_sql_pagination(&headers, &req, output_limits, registration, query_id) {
5204                Ok(resolved) => resolved,
5205                Err(response) => return *response,
5206            };
5207        let sql_permit =
5208            match acquire_sql_permit(&state, &entry.session, registration.query()).await {
5209                Ok(permit) => permit,
5210                Err(error) => {
5211                    return tracked_query_error_response(&state, &error, Some(query_id));
5212                }
5213            };
5214        // Registration and global admission happen before this session lock.
5215        let _guard = tokio::select! {
5216            guard = entry.lock.lock() => guard,
5217            _ = registration.query().control().cancelled() => {
5218                return tracked_query_error_response(
5219                    &state,
5220                    &cancellation_checkpoint_error(registration.query()),
5221                    Some(query_id),
5222                );
5223            }
5224        };
5225        // Re-check closed: the session may have been closed/evicted between
5226        // get() and acquiring the lock.
5227        if entry.is_closed() {
5228            return (StatusCode::NOT_FOUND, "session no longer available").into_response();
5229        }
5230        if req.idempotency_key.is_some() || headers.contains_key("idempotency-key") {
5231            entry
5232                .session
5233                .fire_test_hook(mongreldb_query::SqlTestHookPoint::BeforeServerIdempotencyCheck);
5234        }
5235        let (registration, idempotency) = match begin_sql_idempotency(
5236            &state,
5237            SqlIdempotencyContext {
5238                headers: &headers,
5239                request: &req,
5240                output_limits,
5241                owner: &owner,
5242                session_id: Some(&sid),
5243                session_in_transaction: entry.session.staged_sql_operation_count().is_some(),
5244                query_id,
5245            },
5246            registration,
5247        )
5248        .await
5249        {
5250            Ok(resolved) => resolved,
5251            Err(response) => return response,
5252        };
5253        entry.touch();
5254        let query = registration.into_query();
5255        execute_sql(
5256            &state,
5257            &principal,
5258            &entry.session,
5259            ResolvedSqlRequest {
5260                request: req,
5261                output_limits,
5262                idempotency,
5263                pagination,
5264            },
5265            query,
5266            query_id,
5267            sql_permit,
5268        )
5269        .await
5270    } else {
5271        let session = match MongrelSession::open_with_external_modules_as(
5272            Arc::clone(&state.db),
5273            state.external_modules.iter().cloned(),
5274            request_principal(&state, &principal),
5275        ) {
5276            Ok(session) => session.with_query_registry(Arc::clone(&state.query_registry)),
5277            Err(e) => return (status_for_query_error(&e), e.to_string()).into_response(),
5278        };
5279        let (options, query_id) = match resolve_query_options(
5280            &state,
5281            &headers,
5282            req.query_id,
5283            req.timeout_ms,
5284            owner.clone(),
5285            None,
5286        ) {
5287            Ok(options) => options,
5288            Err(response) => return *response,
5289        };
5290        let query = match register_controlled_query(&state, &session, options) {
5291            Ok(query) => query,
5292            Err(error) => return tracked_query_error_response(&state, &error, Some(query_id)),
5293        };
5294        let registration = RegisteredQueryGuard::new(query);
5295        if mongreldb_query::contains_boolean_ai_predicate(&req.sql) {
5296            registration.fail();
5297            return with_query_id(remote_boolean_ai_error(), query_id);
5298        }
5299        let output_limits = match resolve_sql_output_limits(&state, &req, query_id) {
5300            Ok(limits) => limits,
5301            Err(response) => {
5302                registration.fail();
5303                return *response;
5304            }
5305        };
5306        let (registration, pagination) =
5307            match resolve_sql_pagination(&headers, &req, output_limits, registration, query_id) {
5308                Ok(resolved) => resolved,
5309                Err(response) => return *response,
5310            };
5311        let (registration, idempotency) = match begin_sql_idempotency(
5312            &state,
5313            SqlIdempotencyContext {
5314                headers: &headers,
5315                request: &req,
5316                output_limits,
5317                owner: &owner,
5318                session_id: None,
5319                session_in_transaction: false,
5320                query_id,
5321            },
5322            registration,
5323        )
5324        .await
5325        {
5326            Ok(resolved) => resolved,
5327            Err(response) => return response,
5328        };
5329        let sql_permit = match acquire_sql_permit(&state, &session, registration.query()).await {
5330            Ok(permit) => permit,
5331            Err(error) => {
5332                if let Some(idempotency) = idempotency {
5333                    idempotency.abort();
5334                }
5335                return tracked_query_error_response(&state, &error, Some(query_id));
5336            }
5337        };
5338        let query = registration.into_query();
5339        execute_sql(
5340            &state,
5341            &principal,
5342            &session,
5343            ResolvedSqlRequest {
5344                request: req,
5345                output_limits,
5346                idempotency,
5347                pagination,
5348            },
5349            query,
5350            query_id,
5351            sql_permit,
5352        )
5353        .await
5354    }
5355}
5356
5357/// Run one SQL request against a given session: bump counters, audit DDL
5358/// (after execution, with redaction), log slow queries on both success and
5359/// failure, and dispatch the response format. Shared by the pooled-session and
5360/// fresh-session paths.
5361async fn execute_sql(
5362    state: &AppState,
5363    principal: &Option<mongreldb_core::Principal>,
5364    session: &MongrelSession,
5365    request: ResolvedSqlRequest,
5366    query: RegisteredSqlQuery,
5367    query_id: QueryId,
5368    sql_permit: tokio::sync::OwnedSemaphorePermit,
5369) -> Response {
5370    let ResolvedSqlRequest {
5371        request: req,
5372        output_limits,
5373        idempotency,
5374        pagination,
5375    } = request;
5376    // Keep a direct handle until the durable receipt is persisted. Finished
5377    // tombstone eviction must not make a committed idempotent write ambiguous.
5378    let idempotency_query = idempotency.as_ref().map(|_| query.clone());
5379    state.metrics.inc_sql_queries();
5380    let audited = audit::is_audited_sql(&req.sql);
5381    let actor = request_owner(state, principal);
5382    let page_binding = sql_pages::SqlPageBinding {
5383        security_version: state.db.security_version(),
5384        catalog_epoch: state.db.catalog_snapshot().db_epoch,
5385    };
5386    let start = std::time::Instant::now();
5387    // NOTE: deliberately NOT using `run_sql_traced` here. Its thread-local
5388    // push/pop spans an `.await`, and on a multi-threaded tokio runtime the
5389    // task can resume on a different thread, corrupting the trace stack and
5390    // leaking scopes. Wall-clock timing is sufficient for slow-query detection
5391    // and works across awaits.
5392    let result = if let Some(pagination) = pagination {
5393        match session
5394            .run_with_query_for_serialization_with_limits(
5395                &req.sql,
5396                query,
5397                mongreldb_query::SqlCollectionLimits::new(output_limits.0, output_limits.1),
5398            )
5399            .await
5400        {
5401            Ok(output) => Ok(dispatch_paginated_sql(
5402                state,
5403                output,
5404                query_id,
5405                &actor,
5406                pagination,
5407                output_limits,
5408                session.sql_test_hook(),
5409                page_binding,
5410            )
5411            .await),
5412            Err(error) => Err(error),
5413        }
5414    } else if req.format.as_deref() == Some("arrow-stream") {
5415        match session
5416            .run_stream_with_query_for_serialization(&req.sql, query)
5417            .await
5418        {
5419            Ok((stream, completion)) => Ok(sql_arrow_stream_response_controlled(
5420                stream,
5421                completion,
5422                sql_permit,
5423                output_limits,
5424                state,
5425                query_id,
5426                session.sql_test_hook(),
5427            )),
5428            Err(error) => Err(error),
5429        }
5430    } else {
5431        match session
5432            .run_with_query_for_serialization_with_limits(
5433                &req.sql,
5434                query,
5435                mongreldb_query::SqlCollectionLimits::new(output_limits.0, output_limits.1),
5436            )
5437            .await
5438        {
5439            Ok(output) => Ok(dispatch_buffered_sql_format(
5440                state,
5441                req.format.as_deref(),
5442                output,
5443                query_id,
5444                session.sql_test_hook(),
5445                output_limits,
5446            )
5447            .await),
5448            Err(error) => Err(error),
5449        }
5450    };
5451    let elapsed = start.elapsed();
5452    // Slow-query logging covers BOTH success and failure (the slowest errors
5453    // matter most for diagnosis), checked before branching on the outcome.
5454    if elapsed >= state.slow_query_threshold {
5455        state.metrics.inc_slow_queries();
5456        eprintln!(
5457            "[slow-query] {}\u{00b5}s query_id={} operation={}",
5458            elapsed.as_micros(),
5459            query_id,
5460            safe_sql_operation(&req.sql)
5461        );
5462    }
5463    // Audit DDL/privilege AFTER execution so the outcome (ok/fail) is captured.
5464    // `redacted_ddl_detail` never logs credential literals.
5465    if audited {
5466        let (action, detail) = audit::redacted_ddl_detail(&req.sql, result.is_ok());
5467        state.audit.record(actor, action, detail);
5468    }
5469    let response = match result {
5470        Ok(response) => with_query_id(response, query_id),
5471        Err(e) => {
5472            state.metrics.inc_sql_errors();
5473            tracked_query_error_response(state, &e, Some(query_id))
5474        }
5475    };
5476    let Some(idempotency) = idempotency else {
5477        return response;
5478    };
5479    let status = idempotency_query.map(|query| query.status());
5480    if let Some(receipt) = status.as_ref().and_then(sql_terminal_idempotency_receipt) {
5481        let (expires_at_ms, persisted) = idempotency.commit(receipt.clone());
5482        return sql_idempotency_receipt_response(
5483            query_id,
5484            &receipt,
5485            false,
5486            expires_at_ms,
5487            persisted,
5488        );
5489    }
5490    if status.as_ref().is_some_and(can_abort_idempotency_intent) {
5491        idempotency.abort();
5492    }
5493    response
5494}
5495
5496fn can_abort_idempotency_intent(status: &mongreldb_query::QueryStatus) -> bool {
5497    !status.outcome_unknown
5498        && !status.durable_outcome.committed
5499        && matches!(
5500            status.terminal_state(),
5501            Some(
5502                mongreldb_query::QueryTerminalState::FailedBeforeCommit
5503                    | mongreldb_query::QueryTerminalState::CancelledBeforeCommit
5504                    | mongreldb_query::QueryTerminalState::DeadlineBeforeCommit
5505            )
5506        )
5507}
5508
5509fn safe_sql_operation(sql: &str) -> String {
5510    sql.split_whitespace()
5511        .next()
5512        .unwrap_or("UNKNOWN")
5513        .chars()
5514        .filter(|character| character.is_ascii_alphabetic())
5515        .take(16)
5516        .collect::<String>()
5517        .to_ascii_uppercase()
5518}
5519
5520fn remote_boolean_ai_error() -> Response {
5521    (
5522        StatusCode::BAD_REQUEST,
5523        "Boolean ANN/Sparse SQL is disabled remotely; use scored SQL functions",
5524    )
5525        .into_response()
5526}
5527
5528#[derive(Debug)]
5529struct SerializedOutput {
5530    bytes: Vec<u8>,
5531    arrow: bool,
5532}
5533
5534#[derive(Debug)]
5535enum BufferedSerializationError {
5536    Query(mongreldb_query::MongrelQueryError),
5537    Limit(String),
5538    Encoding(String),
5539}
5540
5541struct LimitedOutput {
5542    bytes: Vec<u8>,
5543    max_bytes: usize,
5544    exceeded: bool,
5545}
5546
5547impl LimitedOutput {
5548    fn new(max_bytes: usize) -> Self {
5549        Self {
5550            bytes: Vec::new(),
5551            max_bytes,
5552            exceeded: false,
5553        }
5554    }
5555}
5556
5557impl std::io::Write for LimitedOutput {
5558    fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
5559        if self.bytes.len().saturating_add(bytes.len()) > self.max_bytes {
5560            self.exceeded = true;
5561            return Err(std::io::Error::other("SQL output byte limit exceeded"));
5562        }
5563        self.bytes.extend_from_slice(bytes);
5564        Ok(bytes.len())
5565    }
5566
5567    fn flush(&mut self) -> std::io::Result<()> {
5568        Ok(())
5569    }
5570}
5571
5572fn serialize_buffered_output(
5573    format: &str,
5574    batches: &[arrow::record_batch::RecordBatch],
5575    query: &RegisteredSqlQuery,
5576    max_rows: usize,
5577    max_bytes: usize,
5578    test_hook: Option<&mongreldb_query::SqlTestHook>,
5579) -> std::result::Result<SerializedOutput, BufferedSerializationError> {
5580    const ROW_CHECKPOINT_INTERVAL: usize = 256;
5581    let mut rows = 0usize;
5582    let mut writer_output = LimitedOutput::new(max_bytes);
5583
5584    if format == "arrow" {
5585        if batches.is_empty() {
5586            if let Some(hook) = test_hook {
5587                hook(mongreldb_query::SqlTestHookPoint::AfterSerialization);
5588            }
5589            return Ok(SerializedOutput {
5590                bytes: Vec::new(),
5591                arrow: true,
5592            });
5593        }
5594        let schema = batches[0].schema();
5595        let encoding_result = (|| {
5596            let mut writer =
5597                arrow::ipc::writer::FileWriter::try_new(&mut writer_output, schema.as_ref())
5598                    .map_err(|error| error.to_string())?;
5599            for batch in batches {
5600                for offset in (0..batch.num_rows()).step_by(ROW_CHECKPOINT_INTERVAL) {
5601                    if let Some(hook) = test_hook {
5602                        hook(mongreldb_query::SqlTestHookPoint::BeforeSerializationBatch);
5603                    }
5604                    query.checkpoint().map_err(|error| error.to_string())?;
5605                    let length = ROW_CHECKPOINT_INTERVAL.min(batch.num_rows() - offset);
5606                    rows = rows.saturating_add(length);
5607                    if rows > max_rows {
5608                        return Err("SQL output row limit exceeded".into());
5609                    }
5610                    writer
5611                        .write(&batch.slice(offset, length))
5612                        .map_err(|error| error.to_string())?;
5613                }
5614            }
5615            writer.finish().map_err(|error| error.to_string())
5616        })();
5617        if let Err(error) = encoding_result {
5618            if let Err(query_error) = query.checkpoint() {
5619                return Err(BufferedSerializationError::Query(query_error));
5620            }
5621            if writer_output.exceeded || rows > max_rows {
5622                return Err(BufferedSerializationError::Limit(error));
5623            }
5624            return Err(BufferedSerializationError::Encoding(error));
5625        }
5626        if let Some(hook) = test_hook {
5627            hook(mongreldb_query::SqlTestHookPoint::AfterSerialization);
5628        }
5629        return Ok(SerializedOutput {
5630            bytes: writer_output.bytes,
5631            arrow: true,
5632        });
5633    }
5634
5635    let encoding_result = (|| {
5636        let mut writer = arrow::json::writer::ArrayWriter::new(&mut writer_output);
5637        for batch in batches {
5638            for offset in (0..batch.num_rows()).step_by(ROW_CHECKPOINT_INTERVAL) {
5639                if let Some(hook) = test_hook {
5640                    hook(mongreldb_query::SqlTestHookPoint::BeforeSerializationBatch);
5641                }
5642                query.checkpoint().map_err(|error| error.to_string())?;
5643                let length = ROW_CHECKPOINT_INTERVAL.min(batch.num_rows() - offset);
5644                rows = rows.saturating_add(length);
5645                if rows > max_rows {
5646                    return Err("SQL output row limit exceeded".into());
5647                }
5648                let slice = batch.slice(offset, length);
5649                writer
5650                    .write_batches(&[&slice])
5651                    .map_err(|error| error.to_string())?;
5652            }
5653        }
5654        writer.finish().map_err(|error| error.to_string())
5655    })();
5656    if let Err(error) = encoding_result {
5657        if let Err(query_error) = query.checkpoint() {
5658            return Err(BufferedSerializationError::Query(query_error));
5659        }
5660        if writer_output.exceeded || rows > max_rows {
5661            return Err(BufferedSerializationError::Limit(error));
5662        }
5663        return Err(BufferedSerializationError::Encoding(error));
5664    }
5665    if let Some(hook) = test_hook {
5666        hook(mongreldb_query::SqlTestHookPoint::AfterSerialization);
5667    }
5668    Ok(SerializedOutput {
5669        bytes: writer_output.bytes,
5670        arrow: false,
5671    })
5672}
5673
5674/// Serialize a DataFusion record-batch stream as Arrow streaming IPC. The body
5675/// holds only the active query batch and one serialized IPC message.
5676#[cfg(test)]
5677fn sql_arrow_stream_response(batches: mongreldb_query::MongrelRecordBatchStream) -> Response {
5678    use futures::{stream, StreamExt};
5679
5680    const STREAM_CT: &str = "application/vnd.apache.arrow.stream";
5681
5682    let schema = batches.schema();
5683    let mut writer = match arrow::ipc::writer::StreamWriter::try_new(Vec::new(), schema.as_ref()) {
5684        Ok(w) => w,
5685        Err(e) => {
5686            return (
5687                StatusCode::INTERNAL_SERVER_ERROR,
5688                format!("arrow stream init error: {e}"),
5689            )
5690                .into_response()
5691        }
5692    };
5693    // `try_new` synchronously writes the schema message; drain it so it becomes
5694    // the first chunk yielded to the client.
5695    let schema_chunk: Vec<u8> = std::mem::take(writer.get_mut());
5696    let batch_stream = stream::unfold(
5697        (batches, Some(writer)),
5698        |(mut batches, writer)| async move {
5699            let mut writer = writer?;
5700            match batches.next().await {
5701                Some(Ok(batch)) => match writer.write(&batch) {
5702                    Ok(()) => {
5703                        let chunk = std::mem::take(writer.get_mut());
5704                        Some((Ok(chunk), (batches, Some(writer))))
5705                    }
5706                    Err(error) => Some((Err(std::io::Error::other(error)), (batches, None))),
5707                },
5708                Some(Err(error)) => Some((Err(std::io::Error::other(error)), (batches, None))),
5709                None => match writer.finish() {
5710                    Ok(()) => {
5711                        let chunk = std::mem::take(writer.get_mut());
5712                        Some((Ok(chunk), (batches, None)))
5713                    }
5714                    Err(error) => Some((Err(std::io::Error::other(error)), (batches, None))),
5715                },
5716            }
5717        },
5718    );
5719
5720    // Schema first, then batches + EOS. Each item is already `Result<Vec<u8>,
5721    // io::Error>` so a mid-stream encode failure surfaces as a body error.
5722    let schema_item: Result<Vec<u8>, std::io::Error> = Ok(schema_chunk);
5723    let full = stream::iter([schema_item]).chain(batch_stream);
5724    let body = axum::body::Body::from_stream(full);
5725    ([(header::CONTENT_TYPE, STREAM_CT)], body).into_response()
5726}
5727
5728fn sql_arrow_stream_response_controlled(
5729    batches: mongreldb_query::MongrelRecordBatchStream,
5730    completion: SqlStreamCompletion,
5731    sql_permit: tokio::sync::OwnedSemaphorePermit,
5732    limits: (usize, usize),
5733    state: &AppState,
5734    query_id: QueryId,
5735    test_hook: Option<mongreldb_query::SqlTestHook>,
5736) -> Response {
5737    use futures::{stream, StreamExt};
5738
5739    const STREAM_CT: &str = "application/vnd.apache.arrow.stream";
5740    let (max_rows, max_bytes) = limits;
5741    let schema = batches.schema();
5742    let mut writer = match arrow::ipc::writer::StreamWriter::try_new(Vec::new(), schema.as_ref()) {
5743        Ok(writer) => writer,
5744        Err(error) => {
5745            completion.fail_serialization();
5746            drop(batches);
5747            return terminal_server_error_response(
5748                state,
5749                query_id,
5750                StatusCode::INTERNAL_SERVER_ERROR,
5751                "SERIALIZATION_FAILED",
5752                format!("arrow stream init error: {error}"),
5753            );
5754        }
5755    };
5756    let schema_chunk = std::mem::take(writer.get_mut());
5757    if schema_chunk.len() > max_bytes {
5758        completion.fail_result_limit();
5759        drop(batches);
5760        return terminal_server_error_response(
5761            state,
5762            query_id,
5763            StatusCode::PAYLOAD_TOO_LARGE,
5764            "RESULT_LIMIT_EXCEEDED",
5765            "SQL output byte limit exceeded",
5766        );
5767    }
5768    let metrics = Arc::clone(&state.metrics);
5769    metrics.add_sql_output_bytes(schema_chunk.len());
5770    let batch_stream = stream::unfold(
5771        (
5772            batches,
5773            Some(writer),
5774            completion,
5775            Some(sql_permit),
5776            0usize,
5777            schema_chunk.len(),
5778            metrics,
5779        ),
5780        move |(mut batches, writer, completion, permit, rows, bytes, metrics)| {
5781            let test_hook = test_hook.clone();
5782            async move {
5783                let mut writer = writer?;
5784                match batches.next().await {
5785                    Some(Ok(batch)) => {
5786                        let next_rows = rows.saturating_add(batch.num_rows());
5787                        if next_rows > max_rows {
5788                            completion.fail_result_limit();
5789                            return Some((
5790                                Err(std::io::Error::other("SQL output row limit exceeded")),
5791                                (batches, None, completion, permit, next_rows, bytes, metrics),
5792                            ));
5793                        }
5794                        match writer.write(&batch) {
5795                            Ok(()) => {
5796                                let chunk = std::mem::take(writer.get_mut());
5797                                let next_bytes = bytes.saturating_add(chunk.len());
5798                                if next_bytes > max_bytes {
5799                                    completion.fail_result_limit();
5800                                    return Some((
5801                                        Err(std::io::Error::other(
5802                                            "SQL output byte limit exceeded",
5803                                        )),
5804                                        (
5805                                            batches, None, completion, permit, next_rows,
5806                                            next_bytes, metrics,
5807                                        ),
5808                                    ));
5809                                }
5810                                metrics.add_sql_output_bytes(chunk.len());
5811                                Some((
5812                                    Ok(chunk),
5813                                    (
5814                                        batches,
5815                                        Some(writer),
5816                                        completion,
5817                                        permit,
5818                                        next_rows,
5819                                        next_bytes,
5820                                        metrics,
5821                                    ),
5822                                ))
5823                            }
5824                            Err(error) => {
5825                                completion.fail_serialization();
5826                                Some((
5827                                    Err(std::io::Error::other(error)),
5828                                    (batches, None, completion, permit, rows, bytes, metrics),
5829                                ))
5830                            }
5831                        }
5832                    }
5833                    Some(Err(error)) => Some((
5834                        Err(std::io::Error::other(error)),
5835                        (batches, None, completion, permit, rows, bytes, metrics),
5836                    )),
5837                    None => match writer.finish() {
5838                        Ok(()) => {
5839                            let chunk = std::mem::take(writer.get_mut());
5840                            let next_bytes = bytes.saturating_add(chunk.len());
5841                            if next_bytes > max_bytes {
5842                                completion.fail_result_limit();
5843                                return Some((
5844                                    Err(std::io::Error::other("SQL output byte limit exceeded")),
5845                                    (batches, None, completion, permit, rows, next_bytes, metrics),
5846                                ));
5847                            }
5848                            if let Some(hook) = test_hook {
5849                                hook(mongreldb_query::SqlTestHookPoint::AfterSerialization);
5850                            }
5851                            match completion.try_complete() {
5852                                Ok(()) => {
5853                                    metrics.add_sql_output_bytes(chunk.len());
5854                                    Some((
5855                                        Ok(chunk),
5856                                        (
5857                                            batches, None, completion, permit, rows, next_bytes,
5858                                            metrics,
5859                                        ),
5860                                    ))
5861                                }
5862                                Err(error) => {
5863                                    metrics.inc_sql_errors();
5864                                    Some((
5865                                        Err(std::io::Error::other(error.to_string())),
5866                                        (batches, None, completion, permit, rows, bytes, metrics),
5867                                    ))
5868                                }
5869                            }
5870                        }
5871                        Err(error) => {
5872                            completion.fail_serialization();
5873                            Some((
5874                                Err(std::io::Error::other(error)),
5875                                (batches, None, completion, permit, rows, bytes, metrics),
5876                            ))
5877                        }
5878                    },
5879                }
5880            }
5881        },
5882    );
5883    let schema_item: Result<Vec<u8>, std::io::Error> = Ok(schema_chunk);
5884    let body = axum::body::Body::from_stream(stream::iter([schema_item]).chain(batch_stream));
5885    ([(header::CONTENT_TYPE, STREAM_CT)], body).into_response()
5886}
5887
5888#[derive(Deserialize)]
5889struct TxnOp {
5890    table: String,
5891    op: String,
5892    cells: Option<Vec<serde_json::Value>>,
5893    row_id: Option<u64>,
5894}
5895
5896#[derive(Deserialize)]
5897struct TxnRequest {
5898    ops: Vec<TxnOp>,
5899}
5900
5901async fn txn(
5902    State(state): State<Arc<AppState>>,
5903    OptionalPrincipal(principal): OptionalPrincipal,
5904    Json(req): Json<TxnRequest>,
5905) -> Response {
5906    // Pre-validate every op against the live schemas before entering the
5907    // transaction, so malformed input returns 400 without consuming an epoch
5908    // or poisoning a txn.
5909    let mut parsed: Vec<(String, TxnAction)> = Vec::with_capacity(req.ops.len());
5910    for op in &req.ops {
5911        match op.op.as_str() {
5912            "put" => {
5913                let cells_json = match op.cells.as_ref() {
5914                    Some(c) if !c.is_empty() => c,
5915                    _ => {
5916                        return (StatusCode::BAD_REQUEST, "put op requires non-empty cells")
5917                            .into_response()
5918                    }
5919                };
5920                let handle = match state.db.table(&op.table) {
5921                    Ok(h) => h,
5922                    Err(e) => return (StatusCode::NOT_FOUND, e.to_string()).into_response(),
5923                };
5924                let schema = handle.lock().schema().clone();
5925                let cells = match parse_cells(cells_json, &schema) {
5926                    Ok(c) => c,
5927                    Err(msg) => return (StatusCode::BAD_REQUEST, msg).into_response(),
5928                };
5929                parsed.push((op.table.clone(), TxnAction::Put(cells)));
5930            }
5931            "delete" => {
5932                let rid = match op.row_id {
5933                    Some(r) => r,
5934                    None => {
5935                        return (StatusCode::BAD_REQUEST, "delete op requires row_id")
5936                            .into_response()
5937                    }
5938                };
5939                parsed.push((op.table.clone(), TxnAction::Delete(rid)));
5940            }
5941            other => {
5942                return (StatusCode::BAD_REQUEST, format!("unknown op: {other}")).into_response()
5943            }
5944        }
5945    }
5946
5947    state.metrics.inc_txns();
5948    let mut transaction = state.db.begin_as(request_principal(&state, &principal));
5949    let result = (|| {
5950        for (table, action) in &parsed {
5951            match action {
5952                TxnAction::Put(cells) => {
5953                    transaction.put(table, cells.clone())?;
5954                }
5955                TxnAction::Delete(rid) => {
5956                    transaction.delete(table, mongreldb_core::RowId(*rid))?;
5957                }
5958            }
5959        }
5960        transaction.commit()
5961    })();
5962    match result {
5963        Ok(epoch) => Json(json!({
5964            "status": "committed",
5965            "epoch": epoch.0,
5966            "epoch_text": epoch.0.to_string()
5967        }))
5968        .into_response(),
5969        Err(error) => crate::kit::durable_core_error_response(&error)
5970            .unwrap_or_else(|| (status_for_error(&error), error.to_string()).into_response()),
5971    }
5972}
5973
5974enum TxnAction {
5975    Put(Vec<(u16, Value)>),
5976    Delete(u64),
5977}
5978
5979#[cfg(test)]
5980mod auth_tests {
5981    use super::*;
5982    use mongreldb_core::Database;
5983    use tempfile::tempdir;
5984
5985    #[test]
5986    fn slow_query_operation_does_not_include_literals() {
5987        let sql = "CREATE USER alice PASSWORD 'never-log-this'";
5988        let operation = safe_sql_operation(sql);
5989        assert_eq!(operation, "CREATE");
5990        assert!(!operation.contains("never-log-this"));
5991    }
5992
5993    #[tokio::test]
5994    async fn auth_rejects_missing_token() {
5995        let dir = tempdir().unwrap();
5996        let db = Arc::new(Database::create(dir.path()).unwrap());
5997        let app = build_app_with_config(db, std::iter::empty(), Some("secret".into()), None);
5998        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
5999        let addr = listener.local_addr().unwrap();
6000        tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
6001        // Give the server a moment to start.
6002        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
6003
6004        let client = reqwest::Client::new();
6005        let resp = client
6006            .get(format!("http://{addr}/health"))
6007            .send()
6008            .await
6009            .unwrap();
6010        assert_eq!(resp.status(), 401);
6011    }
6012
6013    #[tokio::test]
6014    async fn auth_accepts_valid_token() {
6015        let dir = tempdir().unwrap();
6016        let db = Arc::new(Database::create(dir.path()).unwrap());
6017        let app = build_app_with_config(db, std::iter::empty(), Some("secret".into()), None);
6018        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
6019        let addr = listener.local_addr().unwrap();
6020        tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
6021        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
6022
6023        let client = reqwest::Client::new();
6024        let resp = client
6025            .get(format!("http://{addr}/health"))
6026            .header("Authorization", "Bearer secret")
6027            .send()
6028            .await
6029            .unwrap();
6030        assert_eq!(resp.status(), 200);
6031    }
6032
6033    #[tokio::test]
6034    async fn no_auth_when_token_unset() {
6035        let dir = tempdir().unwrap();
6036        let db = Arc::new(Database::create(dir.path()).unwrap());
6037        let app = build_app_with_config(db, std::iter::empty(), None, None);
6038        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
6039        let addr = listener.local_addr().unwrap();
6040        tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
6041        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
6042
6043        let client = reqwest::Client::new();
6044        let resp = client
6045            .get(format!("http://{addr}/health"))
6046            .send()
6047            .await
6048            .unwrap();
6049        assert_eq!(resp.status(), 200);
6050    }
6051
6052    #[tokio::test]
6053    async fn capabilities_advertise_sql_cancellation_v2() {
6054        let dir = tempdir().unwrap();
6055        let db = Arc::new(Database::create(dir.path()).unwrap());
6056        let app = build_app_with_config(db, std::iter::empty(), None, None);
6057        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
6058        let addr = listener.local_addr().unwrap();
6059        tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
6060
6061        let body: serde_json::Value = reqwest::Client::new()
6062            .get(format!("http://{addr}/capabilities"))
6063            .send()
6064            .await
6065            .unwrap()
6066            .json()
6067            .await
6068            .unwrap();
6069        assert_eq!(body["sql_cancellation"]["version"], 2);
6070        assert_eq!(body["sql_cancellation"]["client_query_ids"], true);
6071        assert_eq!(body["sql_cancellation"]["cancel_endpoint"], true);
6072        assert_eq!(body["sql_cancellation"]["query_status"], true);
6073        assert_eq!(body["sql_cancellation"]["pre_registration_cancel"], true);
6074        assert_eq!(body["sql_cancellation"]["stream_disconnect_cancels"], true);
6075        assert_eq!(body["sql_idempotency"]["version"], 1);
6076        assert_eq!(
6077            body["sql_idempotency"]["indeterminate_never_reexecutes"],
6078            true
6079        );
6080        assert_eq!(body["sql_pagination"]["version"], 1);
6081        assert_eq!(
6082            body["sql_pagination"]["continuation_endpoint"],
6083            "/sql/continue"
6084        );
6085    }
6086}
6087
6088#[cfg(test)]
6089mod query_response_tests {
6090    use super::*;
6091
6092    #[test]
6093    fn cancellation_reason_names_are_stable_snake_case() {
6094        assert_eq!(cancellation_reason_name(CancellationReason::None), "none");
6095        assert_eq!(
6096            cancellation_reason_name(CancellationReason::ClientRequest),
6097            "client_request"
6098        );
6099        assert_eq!(
6100            cancellation_reason_name(CancellationReason::ClientDisconnected),
6101            "client_disconnected"
6102        );
6103        assert_eq!(
6104            cancellation_reason_name(CancellationReason::SessionClosed),
6105            "session_closed"
6106        );
6107        assert_eq!(
6108            cancellation_reason_name(CancellationReason::ServerShutdown),
6109            "server_shutdown"
6110        );
6111        assert_eq!(
6112            cancellation_reason_name(CancellationReason::Deadline),
6113            "deadline"
6114        );
6115    }
6116
6117    #[test]
6118    fn unknown_outcome_never_proves_idempotency_intent_safe_to_abort() {
6119        let registry = Arc::new(SqlQueryRegistry::default());
6120        let unknown_id: QueryId = "102132435465768798a9bacbdcedfe0f".parse().unwrap();
6121        let unknown = registry
6122            .register(SqlQueryOptions {
6123                query_id: Some(unknown_id),
6124                ..SqlQueryOptions::default()
6125            })
6126            .unwrap();
6127        unknown.mark_outcome_unknown();
6128        unknown.fail();
6129        assert!(!can_abort_idempotency_intent(
6130            &registry.status(unknown_id).unwrap()
6131        ));
6132
6133        let failed_id: QueryId = "2031425364758697a8b9cadbecfd0e1f".parse().unwrap();
6134        let failed = registry
6135            .register(SqlQueryOptions {
6136                query_id: Some(failed_id),
6137                ..SqlQueryOptions::default()
6138            })
6139            .unwrap();
6140        failed.fail();
6141        assert!(can_abort_idempotency_intent(
6142            &registry.status(failed_id).unwrap()
6143        ));
6144    }
6145
6146    #[test]
6147    fn unknown_outcome_never_becomes_durable_receipt() {
6148        let registry = Arc::new(SqlQueryRegistry::default());
6149        let query = registry.register(SqlQueryOptions::default()).unwrap();
6150        query.record_commit(0, 42);
6151        query.mark_outcome_unknown();
6152        query.fail();
6153        let status = registry.status(query.id()).unwrap();
6154        assert!(status.durable_outcome.committed);
6155        assert!(status.outcome_unknown);
6156        assert!(sql_terminal_idempotency_receipt(&status).is_none());
6157    }
6158
6159    #[test]
6160    fn successful_noop_write_becomes_noncommitting_receipt() {
6161        let registry = Arc::new(SqlQueryRegistry::default());
6162        let query = registry.register(SqlQueryOptions::default()).unwrap();
6163        query
6164            .transition(SqlQueryPhase::Queued, SqlQueryPhase::Executing)
6165            .unwrap();
6166        query.try_complete().unwrap();
6167        let status = registry.status(query.id()).unwrap();
6168        assert!(!status.durable_outcome.committed);
6169        assert!(!can_abort_idempotency_intent(&status));
6170        let receipt = sql_terminal_idempotency_receipt(&status).unwrap();
6171        assert_eq!(receipt.status, "completed");
6172        assert!(!receipt.outcome.committed);
6173        assert_eq!(receipt.outcome.committed_statements, 0);
6174        assert_eq!(receipt.outcome.last_commit_epoch, None);
6175    }
6176
6177    #[test]
6178    fn conflicting_idempotency_key_sources_are_rejected() {
6179        let request = SqlRequest {
6180            sql: "INSERT INTO items VALUES (1)".into(),
6181            format: None,
6182            query_id: None,
6183            timeout_ms: None,
6184            max_output_rows: None,
6185            max_output_bytes: None,
6186            idempotency_key: Some("body-key".into()),
6187            pagination: None,
6188        };
6189        let mut headers = axum::http::HeaderMap::new();
6190        headers.insert("idempotency-key", "header-key".parse().unwrap());
6191        assert_eq!(
6192            requested_sql_idempotency_key(&headers, &request),
6193            Err("body idempotency_key and Idempotency-Key header must match")
6194        );
6195
6196        headers.insert("idempotency-key", "body-key".parse().unwrap());
6197        assert_eq!(
6198            requested_sql_idempotency_key(&headers, &request),
6199            Ok(Some("body-key".into()))
6200        );
6201    }
6202
6203    #[test]
6204    fn idempotency_binding_includes_pagination_semantics() {
6205        let mut request = SqlRequest {
6206            sql: "INSERT INTO items VALUES (1)".into(),
6207            format: None,
6208            query_id: None,
6209            timeout_ms: None,
6210            max_output_rows: None,
6211            max_output_bytes: None,
6212            idempotency_key: Some("key".into()),
6213            pagination: None,
6214        };
6215        let unpaged = sql_idempotency_binding(&request, (100, 1_024), None, 60_000).unwrap();
6216        request.pagination = Some(SqlPaginationRequest {
6217            page_size_rows: 10,
6218            projection: vec!["id".into()],
6219            max_page_bytes: Some(512),
6220            max_page_tokens: Some(128),
6221        });
6222        let paged = sql_idempotency_binding(&request, (100, 1_024), None, 60_000).unwrap();
6223        assert_ne!(unpaged.request_semantics_hash, paged.request_semantics_hash);
6224    }
6225
6226    #[test]
6227    fn paginated_decode_stops_nested_heap_amplification_at_budget() {
6228        let registry = Arc::new(SqlQueryRegistry::default());
6229        let query = registry.register(SqlQueryOptions::default()).unwrap();
6230        let json = format!("[[{}]]", vec!["null"; 10_000].join(","));
6231        let mut deserializer = serde_json::Deserializer::from_slice(json.as_bytes());
6232        let mut budget = PaginatedDecodeBudget {
6233            used: 0,
6234            limit: 4 * 1024,
6235            nodes: 0,
6236            exceeded: false,
6237            query: &query,
6238            test_hook: None,
6239        };
6240        let error = serde::de::DeserializeSeed::deserialize(
6241            BudgetedJsonRowsSeed {
6242                budget: &mut budget,
6243            },
6244            &mut deserializer,
6245        )
6246        .unwrap_err();
6247        assert!(error.to_string().contains(PAGINATED_MEMORY_LIMIT_ERROR));
6248        assert!(budget.exceeded);
6249        assert!(budget.nodes < 100, "decoded {} nodes", budget.nodes);
6250        query.fail();
6251    }
6252
6253    #[tokio::test]
6254    async fn unknown_outcome_response_never_claims_no_commit() {
6255        let registry = Arc::new(SqlQueryRegistry::default());
6256        let query = registry.register(SqlQueryOptions::default()).unwrap();
6257        let query_id = query.id();
6258        query
6259            .transition(SqlQueryPhase::Queued, SqlQueryPhase::Executing)
6260            .unwrap();
6261        let error = query.outcome_unknown_error("fenced maintenance failed");
6262        query.fail();
6263        let status = registry.status(query_id).unwrap();
6264        assert_eq!(
6265            status.terminal_state(),
6266            Some(mongreldb_query::QueryTerminalState::OutcomeUnknown)
6267        );
6268
6269        let response = query_error_response_with_status(&error, Some(query_id), Some(&status));
6270        assert_eq!(response.status(), StatusCode::CONFLICT);
6271        let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
6272            .await
6273            .unwrap();
6274        let body: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
6275        assert_eq!(body["status"], "outcome_unknown");
6276        assert_eq!(body["error"]["code"], "QUERY_OUTCOME_UNKNOWN");
6277        assert!(body["committed"].is_null());
6278        assert!(body["committed_statements"].is_null());
6279        assert!(body["last_commit_epoch"].is_null());
6280        assert!(body["completed_statements"].is_null());
6281        assert!(body["statement_index"].is_null());
6282        assert!(body["outcome"]["committed"].is_null());
6283        assert!(body["error"]["committed"].is_null());
6284    }
6285
6286    #[tokio::test]
6287    async fn retryable_idempotency_error_survives_terminal_status() {
6288        let registry = Arc::new(SqlQueryRegistry::default());
6289        let query = registry.register(SqlQueryOptions::default()).unwrap();
6290        let query_id = query.id();
6291        let response = registered_sql_error_response(
6292            RegisteredQueryGuard::new(query),
6293            query_id,
6294            StatusCode::SERVICE_UNAVAILABLE,
6295            "IDEMPOTENCY_STORE_UNAVAILABLE",
6296            "could not durably reserve the SQL idempotency key",
6297            true,
6298        );
6299        let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
6300            .await
6301            .unwrap();
6302        let body: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
6303        assert_eq!(body["retryable"], true);
6304        assert_eq!(body["error"]["retryable"], true);
6305
6306        let status = registry.status(query_id).unwrap();
6307        assert_eq!(
6308            status.terminal_error.as_ref().unwrap().code,
6309            "IDEMPOTENCY_STORE_UNAVAILABLE"
6310        );
6311        assert!(terminal_error_retryable(status.terminal_error.as_ref()));
6312    }
6313
6314    #[tokio::test]
6315    async fn committed_idempotency_terminal_error_is_a_receipt() {
6316        let query_id: QueryId = "00112233445566778899aabbccddeeff".parse().unwrap();
6317        let receipt = sql_idempotency::SqlDurableReceipt {
6318            original_query_id: query_id.to_string(),
6319            status: "committed_with_error".into(),
6320            server_state: "failed".into(),
6321            cancellation_reason: "client_disconnected".into(),
6322            outcome: sql_idempotency::SqlReceiptOutcome {
6323                committed: true,
6324                committed_statements: 1,
6325                last_commit_epoch: Some(42),
6326                last_commit_epoch_text: Some("42".into()),
6327                first_commit_statement_index: Some(0),
6328                last_commit_statement_index: Some(0),
6329                completed_statements: 1,
6330                statement_index: 0,
6331                serialization: "failed".into(),
6332            },
6333            terminal_error: Some(sql_idempotency::SqlReceiptTerminalError {
6334                code: "SERIALIZATION_FAILED_AFTER_COMMIT".into(),
6335                category: "serialization".into(),
6336            }),
6337        };
6338        let response = sql_idempotency_receipt_response(query_id, &receipt, false, 99, true);
6339        assert_eq!(response.status(), StatusCode::OK);
6340        let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
6341            .await
6342            .unwrap();
6343        let body: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
6344        assert_eq!(body["status"], "committed_with_error");
6345        assert_eq!(body["server_state"], "failed");
6346        assert_eq!(body["cancellation_reason"], "client_disconnected");
6347        assert_eq!(body["first_commit_statement_index"], 0);
6348        assert_eq!(body["last_commit_statement_index"], 0);
6349        assert_eq!(
6350            body["terminal_error"]["code"],
6351            "SERIALIZATION_FAILED_AFTER_COMMIT"
6352        );
6353    }
6354
6355    #[test]
6356    fn durable_replay_restores_terminal_status_parity() {
6357        let registry = Arc::new(SqlQueryRegistry::default());
6358        let query_id: QueryId = "11223344556677889900aabbccddeeff".parse().unwrap();
6359        let query = registry
6360            .register(SqlQueryOptions {
6361                query_id: Some(query_id),
6362                ..SqlQueryOptions::default()
6363            })
6364            .unwrap();
6365        let receipt = sql_idempotency::SqlDurableReceipt {
6366            original_query_id: "00112233445566778899aabbccddeeff".into(),
6367            status: "cancelled_after_commit".into(),
6368            server_state: "cancelled".into(),
6369            cancellation_reason: "client_disconnected".into(),
6370            outcome: sql_idempotency::SqlReceiptOutcome {
6371                committed: true,
6372                committed_statements: 1,
6373                last_commit_epoch: Some(42),
6374                last_commit_epoch_text: Some("42".into()),
6375                first_commit_statement_index: Some(0),
6376                last_commit_statement_index: Some(0),
6377                completed_statements: 1,
6378                statement_index: 1,
6379                serialization: "failed".into(),
6380            },
6381            terminal_error: Some(sql_idempotency::SqlReceiptTerminalError {
6382                code: "QUERY_CANCELLED_AFTER_COMMIT".into(),
6383                category: "cancellation".into(),
6384            }),
6385        };
6386        restore_idempotency_replay(RegisteredQueryGuard::new(query), &receipt).unwrap();
6387        let status = registry.status(query_id).unwrap();
6388        assert_eq!(status.phase, SqlQueryPhase::Cancelled);
6389        assert_eq!(
6390            status.terminal_state(),
6391            Some(mongreldb_query::QueryTerminalState::CancelledAfterCommit)
6392        );
6393        assert_eq!(
6394            status.cancellation_reason,
6395            CancellationReason::ClientDisconnected
6396        );
6397        assert_eq!(status.durable_outcome.committed_statements, 1);
6398        assert_eq!(status.durable_outcome.last_commit_epoch, Some(42));
6399        assert_eq!(
6400            status.terminal_error.unwrap().code,
6401            "QUERY_CANCELLED_AFTER_COMMIT"
6402        );
6403    }
6404
6405    #[test]
6406    fn durable_replay_rejects_invalid_authenticated_state() {
6407        let registry = Arc::new(SqlQueryRegistry::default());
6408        let query = registry.register(SqlQueryOptions::default()).unwrap();
6409        let receipt = sql_idempotency::SqlDurableReceipt {
6410            original_query_id: query.id().to_string(),
6411            status: "invented_terminal_state".into(),
6412            server_state: "completed".into(),
6413            cancellation_reason: "none".into(),
6414            outcome: sql_idempotency::SqlReceiptOutcome {
6415                committed: true,
6416                committed_statements: 1,
6417                last_commit_epoch: Some(42),
6418                last_commit_epoch_text: Some("42".into()),
6419                first_commit_statement_index: Some(0),
6420                last_commit_statement_index: Some(0),
6421                completed_statements: 1,
6422                statement_index: 0,
6423                serialization: "succeeded".into(),
6424            },
6425            terminal_error: None,
6426        };
6427        let error =
6428            restore_idempotency_replay(RegisteredQueryGuard::new(query), &receipt).unwrap_err();
6429        assert!(error.to_string().contains("invalid terminal state"));
6430    }
6431
6432    #[test]
6433    fn durable_replay_cancel_wins_before_receipt_response() {
6434        let registry = Arc::new(SqlQueryRegistry::default());
6435        let query_id: QueryId = "22334455667788990011aabbccddeeff".parse().unwrap();
6436        let query = registry
6437            .register(SqlQueryOptions {
6438                query_id: Some(query_id),
6439                ..SqlQueryOptions::default()
6440            })
6441            .unwrap();
6442        assert_eq!(
6443            query.request_cancel(CancellationReason::ClientRequest),
6444            CancelOutcome::Accepted
6445        );
6446        let receipt = sql_idempotency::SqlDurableReceipt {
6447            original_query_id: "00112233445566778899aabbccddeeff".into(),
6448            status: "completed".into(),
6449            server_state: "completed".into(),
6450            cancellation_reason: "none".into(),
6451            outcome: sql_idempotency::SqlReceiptOutcome {
6452                committed: true,
6453                committed_statements: 1,
6454                last_commit_epoch: Some(42),
6455                last_commit_epoch_text: Some("42".into()),
6456                first_commit_statement_index: Some(0),
6457                last_commit_statement_index: Some(0),
6458                completed_statements: 1,
6459                statement_index: 0,
6460                serialization: "succeeded".into(),
6461            },
6462            terminal_error: None,
6463        };
6464        let error = restore_idempotency_replay(RegisteredQueryGuard::new(query), &receipt)
6465            .expect_err("accepted cancellation must suppress replay success");
6466        assert!(matches!(
6467            error,
6468            mongreldb_query::MongrelQueryError::QueryCancelled { .. }
6469        ));
6470        let status = registry.status(query_id).unwrap();
6471        assert_eq!(status.phase, SqlQueryPhase::Cancelled);
6472        assert!(status.durable_outcome.committed);
6473    }
6474
6475    #[test]
6476    fn direct_query_handle_preserves_receipt_after_tombstone_eviction() {
6477        let registry = Arc::new(SqlQueryRegistry::new(
6478            1,
6479            1,
6480            usize::MAX,
6481            std::time::Duration::from_secs(60),
6482        ));
6483        let first = registry.register(SqlQueryOptions::default()).unwrap();
6484        first
6485            .transition(SqlQueryPhase::Queued, SqlQueryPhase::Executing)
6486            .unwrap();
6487        first.record_commit(0, 42);
6488        first.complete_current_statement();
6489        first.try_complete().unwrap();
6490
6491        let second = registry.register(SqlQueryOptions::default()).unwrap();
6492        second
6493            .transition(SqlQueryPhase::Queued, SqlQueryPhase::Executing)
6494            .unwrap();
6495        second.try_complete().unwrap();
6496
6497        assert!(registry.status(first.id()).is_none());
6498        let receipt = sql_terminal_idempotency_receipt(&first.status()).unwrap();
6499        assert_eq!(receipt.outcome.committed_statements, 1);
6500        assert_eq!(receipt.outcome.last_commit_epoch, Some(42));
6501    }
6502
6503    #[test]
6504    fn cancellation_checkpoint_mismatch_returns_typed_error() {
6505        let registry = Arc::new(SqlQueryRegistry::default());
6506        let query = registry.register(SqlQueryOptions::default()).unwrap();
6507        assert!(matches!(
6508            cancellation_checkpoint_error(&query),
6509            mongreldb_query::MongrelQueryError::InvalidQueryState(_)
6510        ));
6511        assert_eq!(
6512            query.request_cancel(CancellationReason::ClientRequest),
6513            CancelOutcome::Accepted
6514        );
6515        assert!(matches!(
6516            cancellation_checkpoint_error(&query),
6517            mongreldb_query::MongrelQueryError::QueryCancelled { .. }
6518        ));
6519    }
6520
6521    #[tokio::test]
6522    async fn cancellation_after_commit_reports_durable_outcome() {
6523        let registry = Arc::new(SqlQueryRegistry::default());
6524        let query = registry.register(SqlQueryOptions::default()).unwrap();
6525        let query_id = query.id();
6526        query
6527            .transition(SqlQueryPhase::Queued, SqlQueryPhase::Executing)
6528            .unwrap();
6529        query.record_commit(0, 42);
6530        assert_eq!(
6531            query.request_cancel(CancellationReason::ClientRequest),
6532            CancelOutcome::Accepted
6533        );
6534        let error = query.checkpoint().unwrap_err();
6535        query.fail();
6536        let status = registry.status(query_id).unwrap();
6537
6538        let response = query_error_response_with_status(&error, Some(query_id), Some(&status));
6539        assert_eq!(response.status(), StatusCode::CONFLICT);
6540        let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
6541            .await
6542            .unwrap();
6543        let body: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
6544        assert_eq!(body["status"], "cancelled_after_commit");
6545        assert_eq!(body["error"]["code"], "QUERY_CANCELLED_AFTER_COMMIT");
6546        assert_eq!(body["committed"], true);
6547        assert_eq!(body["outcome"]["committed_statements"], 1);
6548        assert_eq!(body["outcome"]["last_commit_epoch"], 42);
6549        assert_eq!(body["outcome"]["last_commit_epoch_text"], "42");
6550        assert_eq!(body["first_commit_statement_index"], 0);
6551        assert_eq!(body["last_commit_statement_index"], 0);
6552        assert_eq!(body["outcome"]["first_commit_statement_index"], 0);
6553        assert_eq!(body["outcome"]["last_commit_statement_index"], 0);
6554
6555        let response = query_error_response_with_status(&error, Some(query_id), None);
6556        assert_eq!(response.status(), StatusCode::CONFLICT);
6557        let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
6558            .await
6559            .unwrap();
6560        let body: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
6561        assert_eq!(body["status"], "cancelled_after_commit");
6562        assert_eq!(body["error"]["code"], "QUERY_CANCELLED_AFTER_COMMIT");
6563        assert_eq!(body["committed"], true);
6564        assert_eq!(body["committed_statements"], 1);
6565        assert_eq!(body["last_commit_epoch"], 42);
6566        assert_eq!(body["last_commit_epoch_text"], "42");
6567        assert_eq!(body["first_commit_statement_index"], 0);
6568        assert_eq!(body["last_commit_statement_index"], 0);
6569        assert_eq!(body["outcome"]["committed"], true);
6570        assert_eq!(body["outcome"]["committed_statements"], 1);
6571        assert_eq!(body["outcome"]["last_commit_epoch"], 42);
6572        assert_eq!(body["outcome"]["last_commit_epoch_text"], "42");
6573        assert_eq!(body["outcome"]["first_commit_statement_index"], 0);
6574        assert_eq!(body["outcome"]["last_commit_statement_index"], 0);
6575    }
6576
6577    #[tokio::test]
6578    async fn commit_outcome_fallback_preserves_exact_progress() {
6579        let query_id: QueryId = "33445566778899001122aabbccddeeff".parse().unwrap();
6580        let error = mongreldb_query::MongrelQueryError::CommitOutcome {
6581            query_id,
6582            committed: true,
6583            committed_statements: 3,
6584            last_commit_epoch: Some(77),
6585            first_commit_statement_index: Some(1),
6586            last_commit_statement_index: Some(4),
6587            completed_statements: 4,
6588            statement_index: 5,
6589            message: "durable outcome retained".into(),
6590        };
6591        let response = query_error_response_with_status(&error, Some(query_id), None);
6592        assert_eq!(response.status(), StatusCode::CONFLICT);
6593        let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
6594            .await
6595            .unwrap();
6596        let body: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
6597        assert_eq!(body["status"], "committed_with_error");
6598        assert_eq!(body["committed"], true);
6599        assert_eq!(body["committed_statements"], 3);
6600        assert_eq!(body["last_commit_epoch"], 77);
6601        assert_eq!(body["last_commit_epoch_text"], "77");
6602        assert_eq!(body["first_commit_statement_index"], 1);
6603        assert_eq!(body["last_commit_statement_index"], 4);
6604        assert_eq!(body["completed_statements"], 4);
6605        assert_eq!(body["statement_index"], 5);
6606        assert_eq!(body["outcome"]["committed_statements"], 3);
6607        assert_eq!(body["outcome"]["last_commit_epoch"], 77);
6608        assert_eq!(body["outcome"]["first_commit_statement_index"], 1);
6609        assert_eq!(body["outcome"]["last_commit_statement_index"], 4);
6610        assert_eq!(body["outcome"]["completed_statements"], 4);
6611        assert_eq!(body["outcome"]["statement_index"], 5);
6612    }
6613
6614    #[test]
6615    fn status_cancel_outcome_matches_cancel_endpoint_state() {
6616        let registry = Arc::new(SqlQueryRegistry::default());
6617        let commit = registry.register(SqlQueryOptions::default()).unwrap();
6618        commit
6619            .transition(SqlQueryPhase::Queued, SqlQueryPhase::Executing)
6620            .unwrap();
6621        commit.enter_commit_critical().unwrap();
6622        assert_eq!(
6623            query_cancel_outcome(&registry.status(commit.id()).unwrap()),
6624            Some("too_late")
6625        );
6626
6627        let completed = registry.register(SqlQueryOptions::default()).unwrap();
6628        completed
6629            .transition(SqlQueryPhase::Queued, SqlQueryPhase::Executing)
6630            .unwrap();
6631        completed.try_complete().unwrap();
6632        assert_eq!(
6633            query_cancel_outcome(&registry.status(completed.id()).unwrap()),
6634            Some("already_finished")
6635        );
6636    }
6637}
6638
6639#[cfg(test)]
6640mod wal_stream_tests {
6641    use super::*;
6642    use mongreldb_client::ReplicationFollower;
6643    use mongreldb_core::Database;
6644    use tempfile::tempdir;
6645
6646    #[tokio::test]
6647    async fn wal_stream_returns_records_after_commit() {
6648        let dir = tempdir().unwrap();
6649        let db = Arc::new(Database::create(dir.path()).unwrap());
6650        let table_schema = mongreldb_core::schema::Schema {
6651            schema_id: 1,
6652            columns: vec![mongreldb_core::schema::ColumnDef {
6653                id: 1,
6654                name: "id".into(),
6655                ty: TypeId::Int64,
6656                flags: mongreldb_core::schema::ColumnFlags::empty()
6657                    .with(mongreldb_core::schema::ColumnFlags::PRIMARY_KEY),
6658                default_value: None,
6659            }],
6660            indexes: vec![],
6661            colocation: vec![],
6662            constraints: Default::default(),
6663            clustered: false,
6664        };
6665        db.create_table("items", table_schema).unwrap();
6666        // Write a row to generate WAL records.
6667        let handle = db.table("items").unwrap();
6668        handle.lock().put(vec![(1, Value::Int64(1))]).unwrap();
6669        handle.lock().flush().unwrap();
6670
6671        let app = build_app(db);
6672        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
6673        let addr = listener.local_addr().unwrap();
6674        tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
6675        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
6676
6677        let resp = reqwest::get(format!("http://{addr}/wal/stream"))
6678            .await
6679            .unwrap();
6680        assert_eq!(resp.status(), 200);
6681        let body = resp.text().await.unwrap();
6682        // Should contain at least one record (the flush commit).
6683        assert!(!body.is_empty(), "wal_stream should return records");
6684        assert!(body.contains("seq"), "response should contain seq field");
6685    }
6686
6687    #[tokio::test]
6688    async fn follower_bootstraps_and_applies_incremental_commit() {
6689        let leader_dir = tempdir().unwrap();
6690        let follower_dir = tempdir().unwrap();
6691        let follower_path = follower_dir.path().join("copy");
6692        let db = Arc::new(Database::create(leader_dir.path()).unwrap());
6693        db.create_table(
6694            "items",
6695            mongreldb_core::schema::Schema {
6696                schema_id: 1,
6697                columns: vec![mongreldb_core::schema::ColumnDef {
6698                    id: 1,
6699                    name: "id".into(),
6700                    ty: TypeId::Int64,
6701                    flags: mongreldb_core::schema::ColumnFlags::empty()
6702                        .with(mongreldb_core::schema::ColumnFlags::PRIMARY_KEY),
6703                    default_value: None,
6704                }],
6705                indexes: vec![],
6706                colocation: vec![],
6707                constraints: Default::default(),
6708                clustered: false,
6709            },
6710        )
6711        .unwrap();
6712        let handle = db.table("items").unwrap();
6713        handle.lock().put(vec![(1, Value::Int64(1))]).unwrap();
6714        handle.lock().commit().unwrap();
6715
6716        let app = build_app_with_config(
6717            Arc::clone(&db),
6718            std::iter::empty::<Arc<dyn ExternalTableModule>>(),
6719            Some("replication-secret".into()),
6720            None,
6721        );
6722        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
6723        let addr = listener.local_addr().unwrap();
6724        tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
6725        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
6726
6727        let leader_url = format!("http://{addr}");
6728        let first_path = follower_path.clone();
6729        let (mut follower, initial) = tokio::task::spawn_blocking(move || {
6730            let mut follower = ReplicationFollower::new(&leader_url, first_path)
6731                .unwrap()
6732                .with_bearer_token("replication-secret");
6733            let applied = follower.sync().unwrap();
6734            (follower, applied)
6735        })
6736        .await
6737        .unwrap();
6738        assert_eq!(initial, 0);
6739
6740        handle.lock().put(vec![(1, Value::Int64(2))]).unwrap();
6741        handle.lock().commit().unwrap();
6742        let applied = tokio::task::spawn_blocking(move || {
6743            let count = follower.sync().unwrap();
6744            (follower, count)
6745        })
6746        .await
6747        .unwrap();
6748        follower = applied.0;
6749        assert!(applied.1 > 0);
6750        assert!(follower.last_epoch() > 0);
6751
6752        let replica = Database::open(&follower_path).unwrap();
6753        assert_eq!(replica.table("items").unwrap().lock().count(), 2);
6754        drop(replica);
6755
6756        db.set_spill_threshold(1);
6757        db.transaction(|txn| {
6758            txn.put("items", vec![(1, Value::Int64(3))])?;
6759            Ok(())
6760        })
6761        .unwrap();
6762        let (follower_after_bootstrap, applied) = tokio::task::spawn_blocking(move || {
6763            let count = follower.sync().unwrap();
6764            (follower, count)
6765        })
6766        .await
6767        .unwrap();
6768        assert_eq!(applied, 0, "spilled run should trigger safe rebootstrap");
6769        assert!(follower_after_bootstrap.last_epoch() > 0);
6770        let replica = Database::open(&follower_path).unwrap();
6771        assert_eq!(replica.table("items").unwrap().lock().count(), 3);
6772    }
6773}
6774
6775#[cfg(test)]
6776mod metrics_tests {
6777    use super::*;
6778    use mongreldb_core::Database;
6779    use tempfile::tempdir;
6780
6781    /// Helper: spin up a daemon over a fresh DB with one `items(id int64 pk)`
6782    /// table pre-created, returning the bound address.
6783    async fn setup() -> (tempfile::TempDir, std::net::SocketAddr) {
6784        let dir = tempdir().unwrap();
6785        let db = Arc::new(Database::create(dir.path()).unwrap());
6786        let table_schema = mongreldb_core::schema::Schema {
6787            schema_id: 1,
6788            columns: vec![mongreldb_core::schema::ColumnDef {
6789                id: 1,
6790                name: "id".into(),
6791                ty: TypeId::Int64,
6792                flags: mongreldb_core::schema::ColumnFlags::empty()
6793                    .with(mongreldb_core::schema::ColumnFlags::PRIMARY_KEY),
6794                default_value: None,
6795            }],
6796            indexes: vec![],
6797            colocation: vec![],
6798            constraints: Default::default(),
6799            clustered: false,
6800        };
6801        db.create_table("items", table_schema).unwrap();
6802        let app = build_app(db);
6803        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
6804        let addr = listener.local_addr().unwrap();
6805        tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
6806        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
6807        (dir, addr)
6808    }
6809
6810    #[tokio::test]
6811    async fn metrics_endpoint_returns_prometheus_text() {
6812        let (_dir, addr) = setup().await;
6813        let client = reqwest::Client::new();
6814
6815        // Exercise a few handlers to bump counters.
6816        let _ = client
6817            .post(format!("http://{addr}/tables/items/put"))
6818            .json(&json!({ "row": [1, 1] }))
6819            .send()
6820            .await
6821            .unwrap();
6822        let _ = client
6823            .post(format!("http://{addr}/sql"))
6824            .json(&json!({ "sql": "SELECT count(*) FROM items" }))
6825            .send()
6826            .await
6827            .unwrap();
6828
6829        let resp = client
6830            .get(format!("http://{addr}/metrics"))
6831            .send()
6832            .await
6833            .unwrap();
6834        assert_eq!(resp.status(), 200);
6835        let ct = resp
6836            .headers()
6837            .get("content-type")
6838            .and_then(|v| v.to_str().ok())
6839            .unwrap_or_default()
6840            .to_string();
6841        assert!(
6842            ct.contains("text/plain"),
6843            "content-type is prometheus text: {ct}"
6844        );
6845        let body = resp.text().await.unwrap();
6846        // Prometheus series + type lines are present.
6847        assert!(body.contains("# TYPE mongreldb_sql_queries_total counter"));
6848        assert!(body.contains("# TYPE mongreldb_puts_total counter"));
6849        assert!(body.contains("# TYPE mongreldb_tables gauge"));
6850        // Counters were bumped: at least one query and one put were served.
6851        assert!(
6852            body.contains("mongreldb_sql_queries_total 1"),
6853            "sql_queries counter should reflect the /sql call: {body}"
6854        );
6855        assert!(
6856            body.contains("mongreldb_puts_total 1"),
6857            "puts counter should reflect the put call: {body}"
6858        );
6859        // The tables gauge reflects the single `items` table.
6860        assert!(body.contains("mongreldb_tables 1"));
6861    }
6862
6863    #[tokio::test]
6864    async fn metrics_error_counter_increments_on_bad_sql() {
6865        let (_dir, addr) = setup().await;
6866        let client = reqwest::Client::new();
6867        // A query against a non-existent table errors at the engine layer.
6868        let _ = client
6869            .post(format!("http://{addr}/sql"))
6870            .json(&json!({ "sql": "SELECT * FROM does_not_exist" }))
6871            .send()
6872            .await
6873            .unwrap();
6874        let body = client
6875            .get(format!("http://{addr}/metrics"))
6876            .send()
6877            .await
6878            .unwrap()
6879            .text()
6880            .await
6881            .unwrap();
6882        assert!(
6883            body.contains("mongreldb_sql_errors_total 1"),
6884            "sql_errors should increment on a failed query: {body}"
6885        );
6886    }
6887
6888    #[tokio::test]
6889    async fn arrow_stream_returns_ipc_stream_bytes() {
6890        let (_dir, addr) = setup().await;
6891        let client = reqwest::Client::new();
6892        // Insert a couple of rows so there are real batches to stream, then
6893        // flush so the rows are durable/visible to a fresh SQL session.
6894        for i in 1..=3 {
6895            let resp = client
6896                .post(format!("http://{addr}/tables/items/put"))
6897                .json(&json!({ "row": [1, i] }))
6898                .send()
6899                .await
6900                .unwrap();
6901            assert_eq!(resp.status(), 200, "put should succeed");
6902        }
6903        let _ = client
6904            .post(format!("http://{addr}/tables/items/commit"))
6905            .send()
6906            .await
6907            .unwrap();
6908        // Sanity: the rows are durable and visible to the table handle.
6909        let count_body = client
6910            .get(format!("http://{addr}/tables/items/count"))
6911            .send()
6912            .await
6913            .unwrap()
6914            .text()
6915            .await
6916            .unwrap();
6917        assert!(
6918            count_body.contains("\"count\":3"),
6919            "expected 3 visible rows, got: {count_body}"
6920        );
6921        let resp = client
6922            .post(format!("http://{addr}/sql"))
6923            .json(&json!({ "sql": "SELECT count(*) FROM items", "format": "arrow-stream" }))
6924            .send()
6925            .await
6926            .unwrap();
6927        assert_eq!(resp.status(), 200, "streaming query should succeed");
6928        let ct = resp
6929            .headers()
6930            .get("content-type")
6931            .and_then(|v| v.to_str().ok())
6932            .unwrap_or_default()
6933            .to_string();
6934        assert!(
6935            ct.contains("application/vnd.apache.arrow.stream"),
6936            "content-type should be the arrow stream format: {ct}"
6937        );
6938        let bytes = resp.bytes().await.unwrap();
6939        // Arrow IPC streams begin with the magic continuation marker
6940        // 0xFFFFFFFF followed by the schema message length. The stream must be
6941        // non-empty (3 rows were written) and end with the EOS marker.
6942        assert!(
6943            !bytes.is_empty(),
6944            "arrow stream body should contain schema + batch + EOS"
6945        );
6946        assert!(
6947            bytes.starts_with(&0xFFFFFFFFu32.to_le_bytes()),
6948            "arrow stream must begin with the IPC continuation marker"
6949        );
6950        assert!(
6951            bytes.ends_with(&[0u8, 0, 0, 0]),
6952            "arrow stream should end with the EOS marker (trailing zero length)"
6953        );
6954    }
6955}
6956
6957#[cfg(test)]
6958mod streaming_tests {
6959    use super::*;
6960    use arrow::array::Int64Array;
6961    use arrow::datatypes::{DataType, Field, Schema};
6962    use arrow::record_batch::RecordBatch;
6963    use datafusion::common::DataFusionError;
6964    use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
6965    use futures::StreamExt;
6966    use std::sync::Arc as StdArc;
6967
6968    fn batch_stream(batches: Vec<RecordBatch>) -> mongreldb_query::MongrelRecordBatchStream {
6969        let schema = batches
6970            .first()
6971            .map(RecordBatch::schema)
6972            .unwrap_or_else(|| StdArc::new(Schema::empty()));
6973        let batches =
6974            futures::stream::iter(batches.into_iter().map(Ok::<RecordBatch, DataFusionError>));
6975        Box::pin(RecordBatchStreamAdapter::new(schema, batches))
6976    }
6977
6978    /// Unit-level check: feed two synthetic batches through the streaming
6979    /// serializer and re-parse the resulting IPC stream end-to-end. This
6980    /// validates the per-message chunking (schema + N batches + EOS) without
6981    /// depending on the engine's scan visibility.
6982    #[tokio::test]
6983    async fn arrow_stream_serializes_multiple_batches_roundtrip() {
6984        let schema = StdArc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)]));
6985        let b1 = RecordBatch::try_new(
6986            schema.clone(),
6987            vec![StdArc::new(Int64Array::from(vec![1, 2]))],
6988        )
6989        .unwrap();
6990        let b2 = RecordBatch::try_new(
6991            schema.clone(),
6992            vec![StdArc::new(Int64Array::from(vec![3, 4, 5]))],
6993        )
6994        .unwrap();
6995
6996        let resp = sql_arrow_stream_response(batch_stream(vec![b1, b2]));
6997        let bytes = axum::body::to_bytes(resp.into_body(), 1 << 20)
6998            .await
6999            .unwrap();
7000
7001        // Begins with the IPC continuation marker.
7002        assert!(bytes.starts_with(&0xFFFFFFFFu32.to_le_bytes()));
7003
7004        // Re-parse the full stream and confirm all rows survived the chunked
7005        // serialization.
7006        let slice: &[u8] = bytes.as_ref();
7007        let mut reader = arrow::ipc::reader::StreamReader::try_new(slice, None).unwrap();
7008        let mut total_rows = 0;
7009        for batch in reader.by_ref() {
7010            let batch = batch.expect("each IPC message should decode");
7011            total_rows += batch.num_rows();
7012        }
7013        assert_eq!(
7014            total_rows, 5,
7015            "all rows should round-trip through the stream"
7016        );
7017    }
7018
7019    #[tokio::test]
7020    async fn arrow_stream_emits_schema_before_first_batch() {
7021        let schema = StdArc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)]));
7022        let pending = futures::stream::pending::<Result<RecordBatch, DataFusionError>>();
7023        let batches = Box::pin(RecordBatchStreamAdapter::new(schema, pending));
7024        let mut body = sql_arrow_stream_response(batches)
7025            .into_body()
7026            .into_data_stream();
7027
7028        let chunk = tokio::time::timeout(std::time::Duration::from_millis(100), body.next())
7029            .await
7030            .expect("schema chunk should not wait for a query batch")
7031            .unwrap()
7032            .unwrap();
7033        assert!(chunk.starts_with(&0xFFFFFFFFu32.to_le_bytes()));
7034    }
7035
7036    #[tokio::test]
7037    async fn arrow_stream_empty_query_is_valid_ipc() {
7038        let resp = sql_arrow_stream_response(batch_stream(Vec::new()));
7039        let bytes = axum::body::to_bytes(resp.into_body(), 1 << 20)
7040            .await
7041            .unwrap();
7042        let slice: &[u8] = bytes.as_ref();
7043        let reader = arrow::ipc::reader::StreamReader::try_new(slice, None).unwrap();
7044        assert_eq!(reader.count(), 0);
7045    }
7046
7047    #[tokio::test]
7048    async fn buffered_output_limits_are_typed() {
7049        let dir = tempfile::tempdir().unwrap();
7050        let db = StdArc::new(mongreldb_core::Database::create(dir.path()).unwrap());
7051        let session = MongrelSession::open(db).unwrap();
7052        let query = session.register_query(SqlQueryOptions::default()).unwrap();
7053        let output = session
7054            .run_with_query_for_serialization("SELECT 1", query)
7055            .await
7056            .unwrap();
7057
7058        let row_error =
7059            serialize_buffered_output("json", output.batches(), output.query(), 0, 1024, None)
7060                .unwrap_err();
7061        assert!(matches!(row_error, BufferedSerializationError::Limit(_)));
7062        let byte_error =
7063            serialize_buffered_output("json", output.batches(), output.query(), 10, 1, None)
7064                .unwrap_err();
7065        assert!(matches!(byte_error, BufferedSerializationError::Limit(_)));
7066        output.fail();
7067    }
7068}
7069
7070#[cfg(test)]
7071mod audit_tests {
7072    use super::*;
7073    use mongreldb_core::Database;
7074    use tempfile::tempdir;
7075
7076    /// Build a daemon with user-auth enabled plus one catalog user `alice`.
7077    async fn auth_setup(password: &str) -> std::net::SocketAddr {
7078        let dir = tempdir().unwrap();
7079        let db = Arc::new(Database::create(dir.path()).unwrap());
7080        db.create_user("alice", password).unwrap();
7081        db.set_user_admin("alice", true).unwrap();
7082        let app = build_app_full(
7083            db,
7084            std::iter::empty::<Arc<dyn ExternalTableModule>>(),
7085            None,
7086            None,
7087            true,
7088        );
7089        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
7090        let addr = listener.local_addr().unwrap();
7091        tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
7092        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
7093        addr
7094    }
7095
7096    #[tokio::test]
7097    async fn audit_records_login_success_and_failure() {
7098        let addr = auth_setup("s3cret").await;
7099        let client = reqwest::Client::new();
7100
7101        // Successful basic auth.
7102        let resp = client
7103            .get(format!("http://{addr}/health"))
7104            .header("Authorization", basic("alice", "s3cret"))
7105            .send()
7106            .await
7107            .unwrap();
7108        assert_eq!(resp.status(), 200);
7109
7110        // Failed basic auth (wrong password).
7111        let resp = client
7112            .get(format!("http://{addr}/health"))
7113            .header("Authorization", basic("alice", "wrong"))
7114            .send()
7115            .await
7116            .unwrap();
7117        assert_eq!(resp.status(), 401);
7118
7119        let body = client
7120            .get(format!("http://{addr}/audit"))
7121            .header("Authorization", basic("alice", "s3cret"))
7122            .send()
7123            .await
7124            .unwrap()
7125            .text()
7126            .await
7127            .unwrap();
7128        assert!(
7129            body.contains("\"action\":\"login.ok\""),
7130            "audit should record the successful login: {body}"
7131        );
7132        assert!(
7133            body.contains("\"action\":\"login.fail\""),
7134            "audit should record the failed login: {body}"
7135        );
7136        assert!(
7137            body.contains("\"principal\":\"alice\""),
7138            "audit should attribute events to alice: {body}"
7139        );
7140    }
7141
7142    #[tokio::test]
7143    async fn audit_records_ddl_sql() {
7144        let dir = tempdir().unwrap();
7145        let db = Arc::new(Database::create(dir.path()).unwrap());
7146        let app = build_app(db);
7147        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
7148        let addr = listener.local_addr().unwrap();
7149        tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
7150        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
7151
7152        let client = reqwest::Client::new();
7153        let _ = client
7154            .post(format!("http://{addr}/sql"))
7155            .json(&json!({ "sql": "CREATE TABLE t (id BIGINT PRIMARY KEY)" }))
7156            .send()
7157            .await
7158            .unwrap();
7159        // A plain SELECT is NOT audited.
7160        let _ = client
7161            .post(format!("http://{addr}/sql"))
7162            .json(&json!({ "sql": "SELECT 1" }))
7163            .send()
7164            .await
7165            .unwrap();
7166
7167        let body = client
7168            .get(format!("http://{addr}/audit"))
7169            .send()
7170            .await
7171            .unwrap()
7172            .text()
7173            .await
7174            .unwrap();
7175        assert!(
7176            body.contains("\"action\":\"ddl.ok\""),
7177            "audit should record the successful DDL statement: {body}"
7178        );
7179        assert!(
7180            body.contains("CREATE TABLE"),
7181            "audit detail should carry the DDL snippet: {body}"
7182        );
7183        // SELECT must not appear.
7184        assert!(
7185            !body.contains("SELECT 1"),
7186            "non-DDL reads should not be audited: {body}"
7187        );
7188    }
7189
7190    #[tokio::test]
7191    async fn audit_redacts_credential_passwords() {
7192        let dir = tempdir().unwrap();
7193        let db = Arc::new(Database::create(dir.path()).unwrap());
7194        let app = build_app(db);
7195        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
7196        let addr = listener.local_addr().unwrap();
7197        tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
7198        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
7199
7200        let client = reqwest::Client::new();
7201        // A CREATE USER carrying a plaintext password must be audited but the
7202        // password must NEVER reach /audit or stderr.
7203        let _ = client
7204            .post(format!("http://{addr}/sql"))
7205            .json(&json!({ "sql": "CREATE USER alice WITH PASSWORD 'topsecret'" }))
7206            .send()
7207            .await
7208            .unwrap();
7209
7210        let body = client
7211            .get(format!("http://{addr}/audit"))
7212            .send()
7213            .await
7214            .unwrap()
7215            .text()
7216            .await
7217            .unwrap();
7218        assert!(
7219            !body.contains("topsecret"),
7220            "password must never appear in the audit log: {body}"
7221        );
7222        assert!(
7223            body.contains("redacted credential statement"),
7224            "credential DDL should be recorded as redacted: {body}"
7225        );
7226    }
7227
7228    fn basic(user: &str, pass: &str) -> String {
7229        let raw = format!("{user}:{pass}");
7230        format!("Basic {}", base64_encode(raw.as_bytes()))
7231    }
7232
7233    fn base64_encode(input: &[u8]) -> String {
7234        const TABLE: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
7235        let mut out = String::new();
7236        let mut buf = 0u32;
7237        let mut bits = 0u32;
7238        for &b in input {
7239            buf = (buf << 8) | b as u32;
7240            bits += 8;
7241            while bits >= 6 {
7242                bits -= 6;
7243                out.push(TABLE[((buf >> bits) & 0x3F) as usize] as char);
7244            }
7245        }
7246        if bits > 0 {
7247            out.push(TABLE[((buf << (6 - bits)) & 0x3F) as usize] as char);
7248        }
7249        while !out.len().is_multiple_of(4) {
7250            out.push('=');
7251        }
7252        out
7253    }
7254}
7255
7256#[cfg(test)]
7257mod session_tests {
7258    use super::*;
7259    use mongreldb_core::Database;
7260    use tempfile::tempdir;
7261
7262    /// Spin up a daemon over a fresh DB with one `items(id int64 pk)` table.
7263    /// Returns the TempDir (must be held alive for the test's duration — dropping
7264    /// it deletes the database directory mid-test) and the bound address.
7265    async fn setup() -> (tempfile::TempDir, std::net::SocketAddr) {
7266        let dir = tempdir().unwrap();
7267        let db = Arc::new(Database::create(dir.path()).unwrap());
7268        let table_schema = mongreldb_core::schema::Schema {
7269            schema_id: 1,
7270            columns: vec![mongreldb_core::schema::ColumnDef {
7271                id: 1,
7272                name: "id".into(),
7273                ty: TypeId::Int64,
7274                flags: mongreldb_core::schema::ColumnFlags::empty()
7275                    .with(mongreldb_core::schema::ColumnFlags::PRIMARY_KEY),
7276                default_value: None,
7277            }],
7278            indexes: vec![],
7279            colocation: vec![],
7280            constraints: Default::default(),
7281            clustered: false,
7282        };
7283        db.create_table("items", table_schema).unwrap();
7284        let app = build_app(db);
7285        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
7286        let addr = listener.local_addr().unwrap();
7287        tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
7288        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
7289        (dir, addr)
7290    }
7291
7292    /// Open a session and return its token.
7293    async fn open_session(client: &reqwest::Client, addr: &std::net::SocketAddr) -> String {
7294        let resp = client
7295            .post(format!("http://{addr}/sessions"))
7296            .send()
7297            .await
7298            .unwrap();
7299        assert_eq!(resp.status(), 200);
7300        resp.json::<serde_json::Value>()
7301            .await
7302            .unwrap()
7303            .get("session_id")
7304            .unwrap()
7305            .as_str()
7306            .unwrap()
7307            .to_string()
7308    }
7309
7310    async fn sql_on(
7311        client: &reqwest::Client,
7312        addr: &std::net::SocketAddr,
7313        session: &str,
7314        sql: &str,
7315    ) -> reqwest::Response {
7316        client
7317            .post(format!("http://{addr}/sql"))
7318            .header("X-Session-ID", session)
7319            .json(&json!({ "sql": sql }))
7320            .send()
7321            .await
7322            .unwrap()
7323    }
7324
7325    async fn count_items(client: &reqwest::Client, addr: &std::net::SocketAddr) -> u64 {
7326        client
7327            .get(format!("http://{addr}/tables/items/count"))
7328            .send()
7329            .await
7330            .unwrap()
7331            .json::<serde_json::Value>()
7332            .await
7333            .unwrap()
7334            .get("count")
7335            .unwrap()
7336            .as_u64()
7337            .unwrap()
7338    }
7339
7340    #[tokio::test]
7341    async fn cross_request_transaction_commits() {
7342        let (_dir, addr) = setup().await;
7343        let client = reqwest::Client::new();
7344        let session = open_session(&client, &addr).await;
7345
7346        // BEGIN, INSERT, COMMIT — each its own HTTP request on the same session.
7347        let r = sql_on(&client, &addr, &session, "BEGIN").await;
7348        assert_eq!(r.status(), 200);
7349        let r = sql_on(
7350            &client,
7351            &addr,
7352            &session,
7353            "INSERT INTO items (id) VALUES (1)",
7354        )
7355        .await;
7356        assert_eq!(r.status(), 200, "INSERT should stage successfully");
7357        // Not yet committed → not visible to other connections.
7358        assert_eq!(count_items(&client, &addr).await, 0);
7359        let r = sql_on(&client, &addr, &session, "COMMIT").await;
7360        assert_eq!(r.status(), 200);
7361
7362        // After COMMIT the row is durable and visible.
7363        assert_eq!(count_items(&client, &addr).await, 1);
7364    }
7365
7366    #[tokio::test]
7367    async fn cross_request_transaction_rolls_back() {
7368        let (_dir, addr) = setup().await;
7369        let client = reqwest::Client::new();
7370        let session = open_session(&client, &addr).await;
7371
7372        sql_on(&client, &addr, &session, "BEGIN").await;
7373        sql_on(
7374            &client,
7375            &addr,
7376            &session,
7377            "INSERT INTO items (id) VALUES (5)",
7378        )
7379        .await;
7380        // ROLLBACK discards the staged insert.
7381        let r = sql_on(&client, &addr, &session, "ROLLBACK").await;
7382        assert_eq!(r.status(), 200);
7383        assert_eq!(
7384            count_items(&client, &addr).await,
7385            0,
7386            "rollback discards staged writes"
7387        );
7388    }
7389
7390    #[tokio::test]
7391    async fn unknown_session_id_is_404() {
7392        let (_dir, addr) = setup().await;
7393        let client = reqwest::Client::new();
7394        let resp = client
7395            .post(format!("http://{addr}/sql"))
7396            .header("X-Session-ID", "does-not-exist")
7397            .json(&json!({ "sql": "SELECT 1" }))
7398            .send()
7399            .await
7400            .unwrap();
7401        assert_eq!(resp.status(), 404);
7402    }
7403
7404    #[tokio::test]
7405    async fn invalid_session_headers_do_not_autocommit() {
7406        let (_dir, addr) = setup().await;
7407        let client = reqwest::Client::new();
7408
7409        let resp = client
7410            .post(format!("http://{addr}/sql"))
7411            .header(
7412                "X-Session-ID",
7413                reqwest::header::HeaderValue::from_bytes(&[0xff]).unwrap(),
7414            )
7415            .json(&json!({ "sql": "INSERT INTO items (id) VALUES (1)" }))
7416            .send()
7417            .await
7418            .unwrap();
7419        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
7420
7421        let resp = client
7422            .post(format!("http://{addr}/sql"))
7423            .header("X-Session-ID", "x".repeat(257))
7424            .json(&json!({ "sql": "INSERT INTO items (id) VALUES (2)" }))
7425            .send()
7426            .await
7427            .unwrap();
7428        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
7429        assert_eq!(count_items(&client, &addr).await, 0);
7430    }
7431
7432    #[tokio::test]
7433    async fn close_session_ends_cross_request_state() {
7434        let (_dir, addr) = setup().await;
7435        let client = reqwest::Client::new();
7436        let session = open_session(&client, &addr).await;
7437
7438        // BEGIN on the session, then close it → staged txn is discarded.
7439        sql_on(&client, &addr, &session, "BEGIN").await;
7440        let r = client
7441            .delete(format!("http://{addr}/sessions/{session}"))
7442            .send()
7443            .await
7444            .unwrap();
7445        assert_eq!(r.status(), 200);
7446
7447        // The session token is now invalid.
7448        let resp = sql_on(&client, &addr, &session, "COMMIT").await;
7449        assert_eq!(resp.status(), 404, "closed session is no longer usable");
7450    }
7451
7452    #[tokio::test]
7453    async fn no_session_header_uses_fresh_ephemeral_session() {
7454        // Without X-Session-ID, BEGIN..COMMIT must still work within a single
7455        // multi-statement /sql body (the historical behavior is preserved).
7456        let (_dir, addr) = setup().await;
7457        let client = reqwest::Client::new();
7458        let resp = client
7459            .post(format!("http://{addr}/sql"))
7460            .json(&json!({ "sql": "INSERT INTO items (id) VALUES (42)" }))
7461            .send()
7462            .await
7463            .unwrap();
7464        assert_eq!(resp.status(), 200);
7465        assert_eq!(count_items(&client, &addr).await, 1);
7466    }
7467
7468    #[tokio::test]
7469    async fn prepared_statement_prepare_execute_and_reuse() {
7470        let (_dir, addr) = setup().await;
7471        let client = reqwest::Client::new();
7472        let session = open_session(&client, &addr).await;
7473        sql_on(
7474            &client,
7475            &addr,
7476            &session,
7477            "INSERT INTO items (id) VALUES (1), (2), (3), (4)",
7478        )
7479        .await;
7480
7481        // Prepare a parameterized query once.
7482        let resp = client
7483            .post(format!("http://{addr}/sessions/{session}/prepare"))
7484            .json(&json!({"name":"gt","sql":"SELECT id FROM items WHERE id > $1"}))
7485            .send()
7486            .await
7487            .unwrap();
7488        assert_eq!(resp.status(), 200);
7489
7490        // Execute with param 2 → ids 3,4.
7491        let resp = client
7492            .post(format!("http://{addr}/sessions/{session}/execute"))
7493            .json(&json!({"name":"gt","params":[2]}))
7494            .send()
7495            .await
7496            .unwrap();
7497        assert_eq!(resp.status(), 200);
7498        let body = resp.json::<serde_json::Value>().await.unwrap();
7499        let arr = body
7500            .as_array()
7501            .expect("execute returns a JSON array of rows");
7502        assert_eq!(arr.len(), 2, "ids > 2 are {{3,4}}: {body}");
7503
7504        // Re-execute with a different param → reuses the cached plan, fewer rows.
7505        let resp = client
7506            .post(format!("http://{addr}/sessions/{session}/execute"))
7507            .json(&json!({"name":"gt","params":[3]}))
7508            .send()
7509            .await
7510            .unwrap();
7511        let body = resp.json::<serde_json::Value>().await.unwrap();
7512        assert_eq!(body.as_array().unwrap().len(), 1, "ids > 3 is {{4}}");
7513    }
7514
7515    #[tokio::test]
7516    async fn prepared_statement_deallocate_then_execute_fails() {
7517        let (_dir, addr) = setup().await;
7518        let client = reqwest::Client::new();
7519        let session = open_session(&client, &addr).await;
7520        let _ = client
7521            .post(format!("http://{addr}/sessions/{session}/prepare"))
7522            .json(&json!({"name":"p","sql":"SELECT $1"}))
7523            .send()
7524            .await
7525            .unwrap();
7526        let deallocate_query_id = "dadadadadadadadadadadadadadadada";
7527        let resp = client
7528            .delete(format!("http://{addr}/sessions/{session}/statements/p"))
7529            .header("X-MongrelDB-Query-ID", deallocate_query_id)
7530            .header("X-MongrelDB-Timeout-Ms", "10000")
7531            .send()
7532            .await
7533            .unwrap();
7534        assert_eq!(resp.status(), 200);
7535        assert_eq!(
7536            resp.headers()
7537                .get("X-MongrelDB-Query-ID")
7538                .unwrap()
7539                .to_str()
7540                .unwrap(),
7541            deallocate_query_id
7542        );
7543        let status = client
7544            .get(format!("http://{addr}/queries/{deallocate_query_id}"))
7545            .send()
7546            .await
7547            .unwrap()
7548            .json::<serde_json::Value>()
7549            .await
7550            .unwrap();
7551        assert_eq!(status["state"], "completed");
7552        assert_eq!(status["operation"], "DEALLOCATE");
7553        // Execute after deallocate must error.
7554        let resp = client
7555            .post(format!("http://{addr}/sessions/{session}/execute"))
7556            .json(&json!({"name":"p","params":[1]}))
7557            .send()
7558            .await
7559            .unwrap();
7560        assert_ne!(resp.status(), 200, "execute after DEALLOCATE must fail");
7561    }
7562
7563    #[tokio::test]
7564    async fn prepared_statement_deallocate_honors_pre_registration_cancel() {
7565        let (_dir, addr) = setup().await;
7566        let client = reqwest::Client::new();
7567        let session = open_session(&client, &addr).await;
7568        let prepared = client
7569            .post(format!("http://{addr}/sessions/{session}/prepare"))
7570            .json(&json!({"name":"p","sql":"SELECT $1"}))
7571            .send()
7572            .await
7573            .unwrap();
7574        assert_eq!(prepared.status(), StatusCode::OK);
7575
7576        let query_id = "dbdbdbdbdbdbdbdbdbdbdbdbdbdbdbdb";
7577        let cancel = client
7578            .post(format!("http://{addr}/queries/{query_id}/cancel"))
7579            .header("X-Session-ID", &session)
7580            .send()
7581            .await
7582            .unwrap();
7583        assert_eq!(cancel.status(), StatusCode::ACCEPTED);
7584        let deallocate = client
7585            .delete(format!("http://{addr}/sessions/{session}/statements/p"))
7586            .header("X-MongrelDB-Query-ID", query_id)
7587            .send()
7588            .await
7589            .unwrap();
7590        assert_eq!(deallocate.status().as_u16(), 499);
7591        assert_eq!(
7592            deallocate.json::<serde_json::Value>().await.unwrap()["error"]["code"],
7593            "QUERY_CANCELLED"
7594        );
7595
7596        let execute = client
7597            .post(format!("http://{addr}/sessions/{session}/execute"))
7598            .json(&json!({"name":"p","params":[1]}))
7599            .send()
7600            .await
7601            .unwrap();
7602        assert_eq!(execute.status(), StatusCode::OK);
7603    }
7604
7605    #[tokio::test]
7606    async fn prepared_statement_rejects_bad_name() {
7607        let (_dir, addr) = setup().await;
7608        let client = reqwest::Client::new();
7609        let session = open_session(&client, &addr).await;
7610        let resp = client
7611            .post(format!("http://{addr}/sessions/{session}/prepare"))
7612            .json(&json!({"name":"1bad","sql":"SELECT 1"}))
7613            .send()
7614            .await
7615            .unwrap();
7616        assert_eq!(
7617            resp.status(),
7618            400,
7619            "statement name starting with a digit must be rejected"
7620        );
7621    }
7622}