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