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