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