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 mut 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    body.push_str(&metrics::hot_lookup_metrics(
4078        &metrics::aggregate_hot_metrics(state.as_ref()),
4079    ));
4080    (
4081        [(
4082            header::CONTENT_TYPE,
4083            "text/plain; version=0.0.4; charset=utf-8".to_string(),
4084        )],
4085        body,
4086    )
4087        .into_response()
4088}
4089
4090/// `POST /compact` — compact all mounted tables.
4091async fn compact_all(
4092    State(state): State<Arc<AppState>>,
4093    OptionalPrincipal(principal): OptionalPrincipal,
4094) -> (StatusCode, Json<serde_json::Value>) {
4095    if let Err(error) = state.db().require_for(
4096        request_principal(&state, &principal).as_ref(),
4097        &mongreldb_core::Permission::Ddl,
4098    ) {
4099        return (
4100            status_for_error(&error),
4101            Json(json!({ "status": "error", "message": error.to_string() })),
4102        );
4103    }
4104    match state.db().compact() {
4105        Ok((compacted, skipped)) => (
4106            StatusCode::OK,
4107            Json(json!({
4108                "status": "ok",
4109                "compacted": compacted,
4110                "skipped": skipped,
4111            })),
4112        ),
4113        Err(e) => (
4114            StatusCode::INTERNAL_SERVER_ERROR,
4115            Json(json!({ "status": "error", "message": format!("{e}") })),
4116        ),
4117    }
4118}
4119
4120/// `POST /tables/{name}/compact` — compact a single table.
4121async fn compact_table(
4122    State(state): State<Arc<AppState>>,
4123    OptionalPrincipal(principal): OptionalPrincipal,
4124    Path(name): Path<String>,
4125) -> (StatusCode, Json<serde_json::Value>) {
4126    if let Err(error) = state.db().require_for(
4127        request_principal(&state, &principal).as_ref(),
4128        &mongreldb_core::Permission::Ddl,
4129    ) {
4130        return (
4131            status_for_error(&error),
4132            Json(json!({ "status": "error", "table": name, "message": error.to_string() })),
4133        );
4134    }
4135    match state.db().compact_table(&name) {
4136        Ok(true) => (
4137            StatusCode::OK,
4138            Json(json!({ "status": "compacted", "table": name })),
4139        ),
4140        Ok(false) => (
4141            StatusCode::OK,
4142            Json(json!({ "status": "skipped", "table": name, "reason": "fewer than 2 runs" })),
4143        ),
4144        Err(e) => (
4145            StatusCode::INTERNAL_SERVER_ERROR,
4146            Json(json!({ "status": "error", "table": name, "message": format!("{e}") })),
4147        ),
4148    }
4149}
4150
4151#[derive(Deserialize)]
4152struct CreateTableRequest {
4153    name: String,
4154    columns: Vec<ColumnDefJson>,
4155}
4156
4157#[derive(Deserialize)]
4158struct ColumnDefJson {
4159    id: u16,
4160    name: String,
4161    ty: String,
4162    primary_key: bool,
4163    #[serde(default)]
4164    nullable: bool,
4165}
4166
4167async fn create_table(
4168    State(state): State<Arc<AppState>>,
4169    OptionalPrincipal(principal): OptionalPrincipal,
4170    Json(req): Json<CreateTableRequest>,
4171) -> Response {
4172    if let Some(response) = require_writes_open(&state) {
4173        return response;
4174    }
4175    if let Err(error) = state.db().require_for(
4176        request_principal(&state, &principal).as_ref(),
4177        &mongreldb_core::Permission::Ddl,
4178    ) {
4179        return (status_for_error(&error), error.to_string()).into_response();
4180    }
4181    let mut columns = Vec::new();
4182    for c in &req.columns {
4183        let ty = match c.ty.as_str() {
4184            "int64" | "bigint" => TypeId::Int64,
4185            "float64" | "double" => TypeId::Float64,
4186            "bytes" | "varchar" | "text" => TypeId::Bytes,
4187            "bool" => TypeId::Bool,
4188            other => {
4189                return (StatusCode::BAD_REQUEST, format!("unknown type: {other}")).into_response()
4190            }
4191        };
4192        let mut flags = mongreldb_core::schema::ColumnFlags::empty();
4193        if c.primary_key {
4194            flags = flags.with(mongreldb_core::schema::ColumnFlags::PRIMARY_KEY);
4195        }
4196        if c.nullable {
4197            flags = flags.with(mongreldb_core::schema::ColumnFlags::NULLABLE);
4198        }
4199        columns.push(mongreldb_core::schema::ColumnDef {
4200            id: c.id,
4201            name: c.name.clone(),
4202            ty,
4203            flags,
4204            default_value: None,
4205            embedding_source: None,
4206        });
4207    }
4208    let schema = Schema {
4209        schema_id: 0,
4210        columns,
4211        indexes: vec![],
4212        colocation: vec![],
4213        constraints: Default::default(),
4214        clustered: false,
4215    };
4216    if let Err(msg) = validate_table_name(&req.name) {
4217        return (StatusCode::BAD_REQUEST, msg).into_response();
4218    }
4219    match state.db().create_table(&req.name, schema) {
4220        Ok(id) => Json(json!({
4221            "table_id": id,
4222            "table_id_text": id.to_string()
4223        }))
4224        .into_response(),
4225        Err(error) => crate::kit::durable_core_error_response(&error)
4226            .unwrap_or_else(|| (status_for_error(&error), error.to_string()).into_response()),
4227    }
4228}
4229
4230async fn list_tables(
4231    State(state): State<Arc<AppState>>,
4232    OptionalPrincipal(principal): OptionalPrincipal,
4233) -> Response {
4234    if let Some(response) = refuse_cluster_standalone_data_plane(&state) {
4235        return response;
4236    }
4237    let principal = request_principal(&state, &principal);
4238    Json(
4239        state
4240            .db()
4241            .table_names()
4242            .into_iter()
4243            .filter(|table| {
4244                state
4245                    .db()
4246                    .select_column_ids_for(table, principal.as_ref())
4247                    .is_ok()
4248            })
4249            .collect::<Vec<_>>(),
4250    )
4251    .into_response()
4252}
4253
4254async fn drop_table(
4255    State(state): State<Arc<AppState>>,
4256    OptionalPrincipal(principal): OptionalPrincipal,
4257    Path(name): Path<String>,
4258) -> Response {
4259    if let Some(response) = require_writes_open(&state) {
4260        return response;
4261    }
4262    if let Err(error) = state.db().require_for(
4263        request_principal(&state, &principal).as_ref(),
4264        &mongreldb_core::Permission::Ddl,
4265    ) {
4266        return (status_for_error(&error), error.to_string()).into_response();
4267    }
4268    match state.db().drop_table_with_epoch(&name) {
4269        Ok(epoch) => Json(json!({
4270            "status": "committed",
4271            "epoch": epoch.0,
4272            "epoch_text": epoch.0.to_string()
4273        }))
4274        .into_response(),
4275        Err(error) => crate::kit::durable_core_error_response(&error)
4276            .unwrap_or_else(|| (status_for_error(&error), error.to_string()).into_response()),
4277    }
4278}
4279
4280#[derive(Deserialize)]
4281struct PutRequest {
4282    row: Vec<serde_json::Value>,
4283}
4284
4285pub(crate) fn json_to_value(v: &serde_json::Value, expected: &TypeId) -> Value {
4286    match (v, expected) {
4287        (serde_json::Value::Number(n), TypeId::Float64) => {
4288            n.as_f64().map(Value::Float64).unwrap_or(Value::Null)
4289        }
4290        (serde_json::Value::Number(n), TypeId::Int64) => {
4291            n.as_i64().map(Value::Int64).unwrap_or(Value::Null)
4292        }
4293        (serde_json::Value::String(s), TypeId::Bytes) => Value::Bytes(s.as_bytes().to_vec()),
4294        (serde_json::Value::String(s), TypeId::Enum { variants }) => {
4295            if variants.iter().any(|v| v == s) {
4296                Value::Bytes(s.as_bytes().to_vec())
4297            } else {
4298                Value::Null
4299            }
4300        }
4301        (serde_json::Value::Bool(b), TypeId::Bool) => Value::Bool(*b),
4302        // Embedding input: a JSON array of numbers, validated against the
4303        // declared dimension. Mismatched length or non-numeric elements → Null.
4304        (serde_json::Value::Array(arr), TypeId::Embedding { dim }) => {
4305            if arr.len() as u32 != *dim {
4306                return Value::Null;
4307            }
4308            let vec: Option<Vec<f32>> =
4309                arr.iter().map(|el| el.as_f64().map(|f| f as f32)).collect();
4310            vec.map(Value::Embedding).unwrap_or(Value::Null)
4311        }
4312        (serde_json::Value::Null, _) => Value::Null,
4313        // Lenient fallbacks for unknown/loosely-typed JSON.
4314        (serde_json::Value::Number(n), _) => {
4315            if let Some(i) = n.as_i64() {
4316                Value::Int64(i)
4317            } else if let Some(f) = n.as_f64() {
4318                Value::Float64(f)
4319            } else {
4320                Value::Null
4321            }
4322        }
4323        (serde_json::Value::String(s), _) => Value::Bytes(s.as_bytes().to_vec()),
4324        (serde_json::Value::Bool(b), _) => Value::Bool(*b),
4325        _ => Value::Null,
4326    }
4327}
4328
4329fn legacy_json_to_value(value: &serde_json::Value, expected: &TypeId) -> Result<Value, String> {
4330    if value.is_null() {
4331        return Ok(Value::Null);
4332    }
4333    match expected {
4334        TypeId::Bool => value
4335            .as_bool()
4336            .map(Value::Bool)
4337            .ok_or_else(|| "expected a boolean".into()),
4338        TypeId::Int8
4339        | TypeId::Int16
4340        | TypeId::Int32
4341        | TypeId::Int64
4342        | TypeId::UInt8
4343        | TypeId::UInt16
4344        | TypeId::UInt32
4345        | TypeId::UInt64
4346        | TypeId::TimestampNanos
4347        | TypeId::Date32
4348        | TypeId::Date64
4349        | TypeId::Time64 => value
4350            .as_i64()
4351            .map(Value::Int64)
4352            .ok_or_else(|| "expected a signed 64-bit integer".into()),
4353        TypeId::Float32 | TypeId::Float64 => value
4354            .as_f64()
4355            .filter(|value| value.is_finite())
4356            .map(Value::Float64)
4357            .ok_or_else(|| "expected a finite number".into()),
4358        TypeId::Bytes => match value.as_str() {
4359            Some(value) => Ok(Value::Bytes(value.as_bytes().to_vec())),
4360            None => decode_tagged_hex(value, "bytes").map(Value::Bytes),
4361        },
4362        TypeId::Enum { variants } => {
4363            let bytes = match value.as_str() {
4364                Some(value) => value.as_bytes().to_vec(),
4365                None => decode_tagged_hex(value, "bytes")?,
4366            };
4367            let value =
4368                std::str::from_utf8(&bytes).map_err(|_| "enum variant is not UTF-8".to_string())?;
4369            if !variants.iter().any(|variant| variant == value) {
4370                return Err("expected a declared enum variant".into());
4371            }
4372            Ok(Value::Bytes(bytes))
4373        }
4374        TypeId::Embedding { dim } => {
4375            let values = value
4376                .as_array()
4377                .ok_or_else(|| "expected an embedding array".to_string())?;
4378            if values.len() != *dim as usize {
4379                return Err(format!("expected an embedding with {dim} values"));
4380            }
4381            values
4382                .iter()
4383                .map(|value| {
4384                    value
4385                        .as_f64()
4386                        .map(|value| value as f32)
4387                        .filter(|value| value.is_finite())
4388                        .ok_or_else(|| "embedding values must be finite numbers".to_string())
4389                })
4390                .collect::<Result<Vec<_>, _>>()
4391                .map(Value::Embedding)
4392        }
4393        TypeId::Decimal128 { .. } => {
4394            let object = exact_tagged_object(value, "decimal", &["unscaled"])?;
4395            let text = object["unscaled"]
4396                .as_str()
4397                .ok_or_else(|| "decimal unscaled value must be a string".to_string())?;
4398            let value = text
4399                .parse::<i128>()
4400                .map_err(|_| "decimal unscaled value is invalid".to_string())?;
4401            if value.to_string() != text {
4402                return Err("decimal unscaled value is not canonical".into());
4403            }
4404            Ok(Value::Decimal(value))
4405        }
4406        TypeId::Interval => {
4407            let object = exact_tagged_object(value, "interval", &["months", "days", "nanos"])?;
4408            let months = canonical_i64(&object["months"], "interval months")?;
4409            let days = canonical_i64(&object["days"], "interval days")?
4410                .try_into()
4411                .map_err(|_| "interval days is outside i32 range".to_string())?;
4412            let nanos = canonical_i64(&object["nanos"], "interval nanos")?;
4413            Ok(Value::Interval {
4414                months,
4415                days,
4416                nanos,
4417            })
4418        }
4419        TypeId::Uuid => {
4420            let bytes = decode_tagged_hex(value, "uuid")?;
4421            let bytes: [u8; 16] = bytes
4422                .try_into()
4423                .map_err(|_| "UUID must contain exactly 16 bytes".to_string())?;
4424            Ok(Value::Uuid(bytes))
4425        }
4426        TypeId::Json => {
4427            let bytes = decode_tagged_hex(value, "json")?;
4428            std::str::from_utf8(&bytes).map_err(|_| "JSON value is not UTF-8".to_string())?;
4429            serde_json::from_slice::<serde_json::Value>(&bytes)
4430                .map_err(|error| format!("JSON value is invalid: {error}"))?;
4431            Ok(Value::Json(bytes))
4432        }
4433        TypeId::Array { .. } => Err("legacy put does not support array columns".into()),
4434    }
4435}
4436
4437fn exact_tagged_object<'a>(
4438    value: &'a serde_json::Value,
4439    expected_kind: &str,
4440    fields: &[&str],
4441) -> Result<&'a serde_json::Map<String, serde_json::Value>, String> {
4442    let object = value
4443        .as_object()
4444        .ok_or_else(|| format!("expected tagged {expected_kind} value"))?;
4445    if object.len() != fields.len() + 1
4446        || object
4447            .get("$mongreldb_type")
4448            .and_then(|value| value.as_str())
4449            != Some(expected_kind)
4450        || fields.iter().any(|field| !object.contains_key(*field))
4451    {
4452        return Err(format!("invalid tagged {expected_kind} value"));
4453    }
4454    Ok(object)
4455}
4456
4457fn decode_tagged_hex(value: &serde_json::Value, expected_kind: &str) -> Result<Vec<u8>, String> {
4458    let object = exact_tagged_object(value, expected_kind, &["hex"])?;
4459    let encoded = object["hex"]
4460        .as_str()
4461        .ok_or_else(|| format!("tagged {expected_kind} hex must be a string"))?;
4462    if encoded.len() % 2 != 0 {
4463        return Err(format!("tagged {expected_kind} hex has odd length"));
4464    }
4465    encoded
4466        .as_bytes()
4467        .chunks_exact(2)
4468        .map(|pair| {
4469            let high = hex_nibble(pair[0])?;
4470            let low = hex_nibble(pair[1])?;
4471            Ok((high << 4) | low)
4472        })
4473        .collect()
4474}
4475
4476fn hex_nibble(value: u8) -> Result<u8, String> {
4477    match value {
4478        b'0'..=b'9' => Ok(value - b'0'),
4479        b'a'..=b'f' => Ok(value - b'a' + 10),
4480        _ => Err("hex value must use lowercase ASCII digits".into()),
4481    }
4482}
4483
4484fn canonical_i64(value: &serde_json::Value, field: &str) -> Result<i64, String> {
4485    let text = value
4486        .as_str()
4487        .ok_or_else(|| format!("{field} must be a string"))?;
4488    let value = text
4489        .parse::<i64>()
4490        .map_err(|_| format!("{field} is invalid"))?;
4491    if value.to_string() != text {
4492        return Err(format!("{field} is not canonical"));
4493    }
4494    Ok(value)
4495}
4496
4497#[cfg(test)]
4498mod legacy_wire_tests {
4499    use super::*;
4500
4501    #[test]
4502    fn typed_values_are_exact_and_malformed_values_fail_closed() {
4503        assert_eq!(
4504            legacy_json_to_value(
4505                &json!({"$mongreldb_type": "bytes", "hex": "00ff61"}),
4506                &TypeId::Bytes,
4507            )
4508            .unwrap(),
4509            Value::Bytes(vec![0, 0xff, b'a'])
4510        );
4511        for (value, ty) in [
4512            (
4513                json!({"$mongreldb_type": "bytes", "hex": "00FF"}),
4514                TypeId::Bytes,
4515            ),
4516            (
4517                json!({"$mongreldb_type": "decimal", "unscaled": "01"}),
4518                TypeId::Decimal128 {
4519                    precision: 38,
4520                    scale: 0,
4521                },
4522            ),
4523            (
4524                json!({"$mongreldb_type": "uuid", "hex": "00"}),
4525                TypeId::Uuid,
4526            ),
4527            (
4528                json!({"$mongreldb_type": "json", "hex": "7b"}),
4529                TypeId::Json,
4530            ),
4531            (json!([1.0]), TypeId::Embedding { dim: 2 }),
4532            (json!([1e100, 1.0]), TypeId::Embedding { dim: 2 }),
4533        ] {
4534            assert!(legacy_json_to_value(&value, &ty).is_err(), "{value}");
4535        }
4536    }
4537}
4538
4539/// Parse a flat JSON array `[col_id, val, col_id, val, ...]` into typed cell
4540/// pairs, validating the schema. Returns `Err(message)` on any malformed pair.
4541fn parse_cells(
4542    row: &[serde_json::Value],
4543    schema: &mongreldb_core::schema::Schema,
4544) -> Result<Vec<(u16, Value)>, String> {
4545    if row.len() & 1 != 0 {
4546        return Err("row must be an even-length array of [col_id, value] pairs".into());
4547    }
4548    let mut out = Vec::with_capacity(row.len() / 2);
4549    let mut seen = std::collections::HashSet::new();
4550    for chunk in row.chunks(2) {
4551        let col_id = chunk[0]
4552            .as_u64()
4553            .and_then(|value| u16::try_from(value).ok())
4554            .ok_or("column id must be an unsigned 16-bit integer")?;
4555        if !seen.insert(col_id) {
4556            return Err(format!("duplicate column id {col_id}"));
4557        }
4558        let expected = schema
4559            .columns
4560            .iter()
4561            .find(|c| c.id == col_id)
4562            .map(|c| c.ty.clone())
4563            .ok_or_else(|| format!("unknown column id {col_id}"))?;
4564        let val = legacy_json_to_value(&chunk[1], &expected)?;
4565        out.push((col_id, val));
4566    }
4567    Ok(out)
4568}
4569
4570/// Basic validation for a table name: non-empty and no path separators.
4571pub(crate) fn validate_table_name(name: &str) -> Result<(), String> {
4572    if name.is_empty() {
4573        return Err("table name must not be empty".into());
4574    }
4575    if name.contains('/') || name.contains('\\') || name.contains('\0') {
4576        return Err("table name contains invalid characters".into());
4577    }
4578    Ok(())
4579}
4580
4581async fn put_row(
4582    State(state): State<Arc<AppState>>,
4583    OptionalPrincipal(principal): OptionalPrincipal,
4584    Path(name): Path<String>,
4585    Json(req): Json<PutRequest>,
4586) -> Response {
4587    if let Some(response) = require_writes_open(&state) {
4588        return response;
4589    }
4590    let handle = match state.db().table(&name) {
4591        Ok(h) => h,
4592        Err(e) => return (StatusCode::NOT_FOUND, e.to_string()).into_response(),
4593    };
4594    let schema = handle.lock().schema().clone();
4595    let row = match parse_cells(&req.row, &schema) {
4596        Ok(r) => r,
4597        Err(msg) => return (StatusCode::BAD_REQUEST, msg).into_response(),
4598    };
4599    state.metrics.inc_puts();
4600    let principal = request_principal(&state, &principal);
4601    match state.db().put_for(&name, row, principal.as_ref()) {
4602        Ok(rid) => Json(json!({ "row_id": rid.0.to_string() })).into_response(),
4603        Err(e) => (status_for_error(&e), e.to_string()).into_response(),
4604    }
4605}
4606
4607async fn count(
4608    State(state): State<Arc<AppState>>,
4609    OptionalPrincipal(principal): OptionalPrincipal,
4610    Path(name): Path<String>,
4611) -> Response {
4612    if let Some(response) = refuse_cluster_standalone_data_plane(&state) {
4613        return response;
4614    }
4615    let principal = request_principal(&state, &principal);
4616    match state.db().count_for(&name, principal.as_ref()) {
4617        Ok(count) => Json(json!({ "count": count })).into_response(),
4618        Err(error) => (status_for_error(&error), error.to_string()).into_response(),
4619    }
4620}
4621
4622async fn commit(
4623    State(state): State<Arc<AppState>>,
4624    OptionalPrincipal(principal): OptionalPrincipal,
4625    Path(name): Path<String>,
4626) -> Response {
4627    if let Some(response) = require_writes_open(&state) {
4628        return response;
4629    }
4630    if let Err(error) = state.db().require_for(
4631        request_principal(&state, &principal).as_ref(),
4632        &mongreldb_core::Permission::Update {
4633            table: name.clone(),
4634        },
4635    ) {
4636        return (status_for_error(&error), error.to_string()).into_response();
4637    }
4638    let handle = match state.db().table(&name) {
4639        Ok(h) => h,
4640        Err(e) => return (StatusCode::NOT_FOUND, e.to_string()).into_response(),
4641    };
4642    let mut g = handle.lock();
4643    state.metrics.inc_commits();
4644    match g.commit() {
4645        Ok(epoch) => Json(json!({
4646            "epoch": epoch.0,
4647            "epoch_text": epoch.0.to_string()
4648        }))
4649        .into_response(),
4650        Err(error) => crate::kit::durable_core_error_response(&error)
4651            .unwrap_or_else(|| (status_for_error(&error), error.to_string()).into_response()),
4652    }
4653}
4654
4655#[derive(Deserialize)]
4656struct SqlRequest {
4657    sql: String,
4658    /// Output format: `"json"` (the default) for a JSON array of row objects,
4659    /// `"arrow"` for Arrow IPC file bytes.
4660    #[serde(default)]
4661    format: Option<String>,
4662    /// Body values take precedence over the equivalent convenience headers.
4663    #[serde(default)]
4664    query_id: Option<QueryId>,
4665    #[serde(default)]
4666    timeout_ms: Option<u64>,
4667    #[serde(default)]
4668    max_output_rows: Option<u64>,
4669    #[serde(default)]
4670    max_output_bytes: Option<u64>,
4671    #[serde(default)]
4672    idempotency_key: Option<String>,
4673    #[serde(default)]
4674    pagination: Option<SqlPaginationRequest>,
4675}
4676
4677#[derive(Clone, Debug, Deserialize, Serialize)]
4678struct SqlPaginationRequest {
4679    page_size_rows: u64,
4680    projection: Vec<String>,
4681    #[serde(default)]
4682    max_page_bytes: Option<u64>,
4683    #[serde(default)]
4684    max_page_tokens: Option<u64>,
4685}
4686
4687#[derive(Clone)]
4688struct ResolvedSqlPagination {
4689    projection: Vec<String>,
4690    limits: sql_pages::SqlPageLimits,
4691}
4692
4693struct ResolvedSqlRequest {
4694    request: SqlRequest,
4695    output_limits: (usize, usize),
4696    idempotency: Option<sql_idempotency::SqlIdempotencyExecution>,
4697    pagination: Option<ResolvedSqlPagination>,
4698}
4699
4700fn query_error_response(
4701    error: &mongreldb_query::MongrelQueryError,
4702    query_id: Option<QueryId>,
4703) -> Response {
4704    query_error_response_with_status(error, query_id, None)
4705}
4706
4707/// Map a query-layer error onto the stable error taxonomy (spec section 9.7).
4708/// The taxonomy is deliberately coarser than `MongrelQueryError`; non-obvious
4709/// choices follow the `MongrelError::category()` precedents and are
4710/// documented at the arm.
4711fn query_error_category(
4712    error: &mongreldb_query::MongrelQueryError,
4713) -> mongreldb_types::errors::ErrorCategory {
4714    use mongreldb_query::MongrelQueryError;
4715    use mongreldb_types::errors::ErrorCategory;
4716    match error {
4717        MongrelQueryError::Core(error) => error.category(),
4718        MongrelQueryError::QueryCancelled { .. } => ErrorCategory::Cancelled,
4719        MongrelQueryError::DeadlineExceeded { .. } => ErrorCategory::DeadlineExceeded,
4720        // A known-aborted commit may be retried as a fresh transaction; a
4721        // committed-with-error one must never be blindly replayed (the core
4722        // `DurableCommit` precedent).
4723        MongrelQueryError::CommitOutcome { committed, .. } => {
4724            if *committed {
4725                ErrorCategory::CommitOutcomeUnknown
4726            } else {
4727                ErrorCategory::TransactionAborted
4728            }
4729        }
4730        MongrelQueryError::OutcomeUnknown { .. } => ErrorCategory::CommitOutcomeUnknown,
4731        MongrelQueryError::TransactionAborted => ErrorCategory::TransactionAborted,
4732        // The session has no usable transaction; begin a fresh one.
4733        MongrelQueryError::NoSqlTransaction => ErrorCategory::TransactionAborted,
4734        // The caller's view of session state is stale (the core `NotFound`
4735        // precedent).
4736        MongrelQueryError::SavepointNotFound { .. } => ErrorCategory::StaleMetadata,
4737        MongrelQueryError::QueryRegistryFull | MongrelQueryError::ResultLimitExceeded { .. } => {
4738            ErrorCategory::ResourceExhausted
4739        }
4740        // The taxonomy has no invalid-request category: id reuse, invalid
4741        // query state, and SQL parse/plan failures are client/server contract
4742        // disagreements (the core `InvalidArgument` precedent).
4743        MongrelQueryError::QueryIdConflict { .. }
4744        | MongrelQueryError::InvalidQueryState(_)
4745        | MongrelQueryError::Arrow(_)
4746        | MongrelQueryError::DataFusion(_) => ErrorCategory::ClusterVersionMismatch,
4747        MongrelQueryError::Schema(_) => ErrorCategory::SchemaVersionMismatch,
4748        // Uncategorized query-layer failures: the serving replica failed to
4749        // complete the request (the core `Other` precedent).
4750        _ => ErrorCategory::ReplicaUnavailable,
4751    }
4752}
4753
4754fn query_error_response_with_status(
4755    error: &mongreldb_query::MongrelQueryError,
4756    query_id: Option<QueryId>,
4757    status: Option<&mongreldb_query::QueryStatus>,
4758) -> Response {
4759    use mongreldb_query::MongrelQueryError;
4760    let (base_code, id) = match error {
4761        MongrelQueryError::QueryCancelled { query_id, .. } => ("QUERY_CANCELLED", Some(*query_id)),
4762        MongrelQueryError::DeadlineExceeded { query_id, .. } => {
4763            ("DEADLINE_EXCEEDED", Some(*query_id))
4764        }
4765        MongrelQueryError::QueryIdConflict { query_id } => ("QUERY_ID_CONFLICT", Some(*query_id)),
4766        MongrelQueryError::QueryRegistryFull => ("QUERY_REGISTRY_FULL", query_id),
4767        MongrelQueryError::ResultLimitExceeded { query_id, .. } => {
4768            ("RESULT_LIMIT_EXCEEDED", Some(*query_id))
4769        }
4770        MongrelQueryError::TransactionAborted => ("TRANSACTION_ABORTED", query_id),
4771        MongrelQueryError::NoSqlTransaction => ("NO_SQL_TRANSACTION", query_id),
4772        MongrelQueryError::SavepointNotFound { .. } => ("SAVEPOINT_NOT_FOUND", query_id),
4773        MongrelQueryError::CommitOutcome { query_id, .. } => ("COMMIT_OUTCOME", Some(*query_id)),
4774        MongrelQueryError::OutcomeUnknown { query_id, .. } => {
4775            ("QUERY_OUTCOME_UNKNOWN", Some(*query_id))
4776        }
4777        _ => ("QUERY_FAILED", query_id),
4778    };
4779    let (
4780        error_committed,
4781        error_committed_statements,
4782        error_last_commit_epoch,
4783        error_first_commit_statement_index,
4784        error_last_commit_statement_index,
4785    ) = match error {
4786        MongrelQueryError::QueryCancelled {
4787            committed,
4788            committed_statements,
4789            last_commit_epoch,
4790            first_commit_statement_index,
4791            last_commit_statement_index,
4792            ..
4793        }
4794        | MongrelQueryError::DeadlineExceeded {
4795            committed,
4796            committed_statements,
4797            last_commit_epoch,
4798            first_commit_statement_index,
4799            last_commit_statement_index,
4800            ..
4801        }
4802        | MongrelQueryError::ResultLimitExceeded {
4803            committed,
4804            committed_statements,
4805            last_commit_epoch,
4806            first_commit_statement_index,
4807            last_commit_statement_index,
4808            ..
4809        } => (
4810            *committed,
4811            *committed_statements,
4812            *last_commit_epoch,
4813            *first_commit_statement_index,
4814            *last_commit_statement_index,
4815        ),
4816        MongrelQueryError::CommitOutcome {
4817            committed,
4818            committed_statements,
4819            last_commit_epoch,
4820            first_commit_statement_index,
4821            last_commit_statement_index,
4822            ..
4823        } => (
4824            *committed,
4825            *committed_statements,
4826            *last_commit_epoch,
4827            *first_commit_statement_index,
4828            *last_commit_statement_index,
4829        ),
4830        _ => (false, 0, None, None, None),
4831    };
4832    let committed = status.map_or_else(
4833        || error_committed,
4834        |status| status.durable_outcome.committed,
4835    );
4836    let outcome_unknown = matches!(error, MongrelQueryError::OutcomeUnknown { .. })
4837        || status.is_some_and(|status| status.outcome_unknown);
4838    let code = status
4839        .and_then(|status| {
4840            status
4841                .terminal_error
4842                .as_ref()
4843                .map(|error| error.code.as_str())
4844        })
4845        .unwrap_or(match (base_code, committed) {
4846            ("QUERY_CANCELLED", true) => "QUERY_CANCELLED_AFTER_COMMIT",
4847            ("DEADLINE_EXCEEDED", true) => "DEADLINE_AFTER_COMMIT",
4848            _ => base_code,
4849        });
4850    let response_status = if outcome_unknown {
4851        "outcome_unknown"
4852    } else {
4853        status
4854            .and_then(mongreldb_query::QueryStatus::terminal_state)
4855            .map(terminal_state_name)
4856            .unwrap_or_else(|| match (error, committed) {
4857                (MongrelQueryError::QueryCancelled { .. }, true) => "cancelled_after_commit",
4858                (MongrelQueryError::QueryCancelled { .. }, false) => "cancelled_before_commit",
4859                (MongrelQueryError::DeadlineExceeded { .. }, true) => "deadline_after_commit",
4860                (MongrelQueryError::DeadlineExceeded { .. }, false) => "deadline_before_commit",
4861                (_, true) => "committed_with_error",
4862                _ => "failed_before_commit",
4863            })
4864    };
4865    let (completed_statements, statement_index) = status.map_or_else(
4866        || match error {
4867            MongrelQueryError::QueryCancelled {
4868                completed_statements,
4869                cancelled_statement_index,
4870                ..
4871            }
4872            | MongrelQueryError::DeadlineExceeded {
4873                completed_statements,
4874                cancelled_statement_index,
4875                ..
4876            } => (*completed_statements, *cancelled_statement_index),
4877            MongrelQueryError::ResultLimitExceeded {
4878                completed_statements,
4879                statement_index,
4880                ..
4881            } => (*completed_statements, *statement_index),
4882            MongrelQueryError::CommitOutcome {
4883                completed_statements,
4884                statement_index,
4885                ..
4886            } => (*completed_statements, *statement_index),
4887            _ => (0, 0),
4888        },
4889        |status| (status.completed_statements, status.statement_index),
4890    );
4891    let committed_statements = status.map_or(error_committed_statements, |status| {
4892        status.durable_outcome.committed_statements
4893    });
4894    let last_commit_epoch = status.map_or(error_last_commit_epoch, |status| {
4895        status.durable_outcome.last_commit_epoch
4896    });
4897    let first_commit_statement_index = status
4898        .map_or(error_first_commit_statement_index, |status| {
4899            status.durable_outcome.first_commit_statement_index
4900        });
4901    let last_commit_statement_index = status.map_or(error_last_commit_statement_index, |status| {
4902        status.durable_outcome.last_commit_statement_index
4903    });
4904    let cancellation_reason = status
4905        .map(|status| status.cancellation_reason)
4906        .or(match error {
4907            MongrelQueryError::QueryCancelled { reason, .. } => Some(*reason),
4908            MongrelQueryError::DeadlineExceeded { .. } => Some(CancellationReason::Deadline),
4909            _ => None,
4910        })
4911        .map(cancellation_reason_name);
4912    let cancel_outcome = match error {
4913        MongrelQueryError::QueryCancelled { .. } | MongrelQueryError::DeadlineExceeded { .. } => {
4914            Some("accepted")
4915        }
4916        _ => status.and_then(query_cancel_outcome),
4917    };
4918    let outcome = if outcome_unknown {
4919        json!({
4920            "committed": null,
4921            "committed_statements": null,
4922            "last_commit_epoch": null,
4923            "last_commit_epoch_text": null,
4924            "first_commit_statement_index": null,
4925            "last_commit_statement_index": null,
4926            "completed_statements": null,
4927            "statement_index": null,
4928            "serialization": "unknown",
4929        })
4930    } else {
4931        status.map_or_else(
4932            || {
4933                json!({
4934                    "committed": committed,
4935                    "committed_statements": committed_statements,
4936                    "last_commit_epoch": last_commit_epoch,
4937                    "last_commit_epoch_text": epoch_text(last_commit_epoch),
4938                    "first_commit_statement_index": first_commit_statement_index,
4939                    "last_commit_statement_index": last_commit_statement_index,
4940                    "completed_statements": completed_statements,
4941                    "statement_index": statement_index,
4942                    "serialization": "unknown",
4943                })
4944            },
4945            |status| query_outcome_json(Some(status)),
4946        )
4947    };
4948    let http_status = match code {
4949        "QUERY_CANCELLED_AFTER_COMMIT" | "DEADLINE_AFTER_COMMIT" => StatusCode::CONFLICT,
4950        "QUERY_CANCELLED" => client_closed_request_status(),
4951        "DEADLINE_EXCEEDED" => StatusCode::GATEWAY_TIMEOUT,
4952        _ => status_for_query_error(error),
4953    };
4954    // The stable error taxonomy (spec 9.7) rides additively on the existing
4955    // error object: programmatic handling keys off `category` /
4956    // `category_code`, never the message.
4957    let category = query_error_category(error);
4958    let mut response = (
4959        http_status,
4960        Json(json!({
4961            "query_id": id.map(|value| value.to_string()),
4962            "status": response_status,
4963            "terminal_state": response_status,
4964            "committed": (!outcome_unknown).then_some(committed),
4965            "committed_statements": (!outcome_unknown).then_some(committed_statements),
4966            "last_commit_epoch": (!outcome_unknown).then_some(last_commit_epoch).flatten(),
4967            "last_commit_epoch_text": (!outcome_unknown).then_some(epoch_text(last_commit_epoch)).flatten(),
4968            "first_commit_statement_index": (!outcome_unknown).then_some(first_commit_statement_index).flatten(),
4969            "last_commit_statement_index": (!outcome_unknown).then_some(last_commit_statement_index).flatten(),
4970            "completed_statements": (!outcome_unknown).then_some(completed_statements),
4971            "statement_index": (!outcome_unknown).then_some(statement_index),
4972            "cancel_outcome": cancel_outcome,
4973            "cancellation_reason": cancellation_reason,
4974            "retryable": matches!(error, MongrelQueryError::QueryRegistryFull),
4975            "server_state": status.map(|status| query_phase_name(status.phase)),
4976            "outcome": outcome,
4977            "error": {
4978                "code": code,
4979                "message": error.to_string(),
4980                "category": category.to_string(),
4981                "category_code": category.code(),
4982                "query_id": id.map(|value| value.to_string()),
4983                "committed": (!outcome_unknown).then_some(committed),
4984                "retryable": matches!(error, MongrelQueryError::QueryRegistryFull),
4985            }
4986        })),
4987    )
4988        .into_response();
4989    if let Some(id) = id {
4990        add_query_id_header(&mut response, id);
4991    }
4992    response
4993}
4994
4995fn record_query_error(metrics: &metrics::Metrics, error: &mongreldb_query::MongrelQueryError) {
4996    match error {
4997        mongreldb_query::MongrelQueryError::QueryCancelled { reason, .. } => {
4998            metrics.inc_sql_cancelled(*reason)
4999        }
5000        mongreldb_query::MongrelQueryError::DeadlineExceeded { .. } => {
5001            metrics.inc_sql_deadline_exceeded();
5002            metrics.inc_sql_cancelled(CancellationReason::Deadline);
5003        }
5004        _ => {}
5005    }
5006}
5007
5008fn tracked_query_error_response(
5009    state: &AppState,
5010    error: &mongreldb_query::MongrelQueryError,
5011    query_id: Option<QueryId>,
5012) -> Response {
5013    record_query_error(&state.metrics, error);
5014    if let mongreldb_query::MongrelQueryError::QueryCancelled { query_id, .. } = error {
5015        if let Some(requested_at) = state
5016            .query_registry
5017            .status(*query_id)
5018            .and_then(|status| status.cancel_requested_at)
5019        {
5020            state
5021                .metrics
5022                .observe_sql_cancel_latency(requested_at.elapsed());
5023        }
5024    }
5025    let status = if matches!(
5026        error,
5027        mongreldb_query::MongrelQueryError::QueryIdConflict { .. }
5028            | mongreldb_query::MongrelQueryError::QueryRegistryFull
5029    ) {
5030        None
5031    } else {
5032        query_id
5033            .or(match error {
5034                mongreldb_query::MongrelQueryError::QueryCancelled { query_id, .. }
5035                | mongreldb_query::MongrelQueryError::DeadlineExceeded { query_id, .. }
5036                | mongreldb_query::MongrelQueryError::CommitOutcome { query_id, .. }
5037                | mongreldb_query::MongrelQueryError::OutcomeUnknown { query_id, .. } => {
5038                    Some(*query_id)
5039                }
5040                _ => None,
5041            })
5042            .and_then(|query_id| state.query_registry.status(query_id))
5043    };
5044    query_error_response_with_status(error, query_id, status.as_ref())
5045}
5046
5047fn add_query_id_header(response: &mut Response, query_id: QueryId) {
5048    if let Ok(value) = axum::http::HeaderValue::from_str(&query_id.to_string()) {
5049        response.headers_mut().insert("x-mongreldb-query-id", value);
5050    }
5051}
5052
5053fn with_query_id(mut response: Response, query_id: QueryId) -> Response {
5054    add_query_id_header(&mut response, query_id);
5055    response
5056}
5057
5058fn bad_query_control_request(message: impl Into<String>, query_id: Option<QueryId>) -> Response {
5059    let mut response = (
5060        StatusCode::BAD_REQUEST,
5061        Json(json!({
5062            "query_id": query_id.map(|value| value.to_string()),
5063            "status": "failed_before_commit",
5064            "terminal_state": "failed_before_commit",
5065            "committed": false,
5066            "committed_statements": 0,
5067            "last_commit_epoch": null,
5068            "last_commit_epoch_text": null,
5069            "first_commit_statement_index": null,
5070            "last_commit_statement_index": null,
5071            "completed_statements": 0,
5072            "statement_index": 0,
5073            "cancel_outcome": null,
5074            "cancellation_reason": null,
5075            "retryable": false,
5076            "server_state": "failed",
5077            "outcome": {
5078                "committed": false,
5079                "committed_statements": 0,
5080                "last_commit_epoch": null,
5081                "last_commit_epoch_text": null,
5082                "first_commit_statement_index": null,
5083                "last_commit_statement_index": null,
5084                "completed_statements": 0,
5085                "statement_index": 0,
5086                "serialization": "not_started",
5087            },
5088            "error": {
5089                "code": "INVALID_QUERY_OPTIONS",
5090                "message": message.into(),
5091                "query_id": query_id.map(|value| value.to_string()),
5092                "committed": false,
5093                "retryable": false,
5094            }
5095        })),
5096    )
5097        .into_response();
5098    if let Some(query_id) = query_id {
5099        add_query_id_header(&mut response, query_id);
5100    }
5101    response
5102}
5103
5104fn resolve_query_options(
5105    state: &AppState,
5106    headers: &axum::http::HeaderMap,
5107    body_query_id: Option<QueryId>,
5108    body_timeout_ms: Option<u64>,
5109    owner: String,
5110    session_id: Option<String>,
5111) -> std::result::Result<(SqlQueryOptions, QueryId), Box<Response>> {
5112    let query_id = match body_query_id {
5113        Some(query_id) => query_id,
5114        None => match headers.get("x-mongreldb-query-id") {
5115            Some(value) => {
5116                let value = value.to_str().map_err(|_| {
5117                    Box::new(bad_query_control_request(
5118                        "X-MongrelDB-Query-ID is not valid text",
5119                        None,
5120                    ))
5121                })?;
5122                value
5123                    .parse()
5124                    .map_err(|error: mongreldb_query::MongrelQueryError| {
5125                        Box::new(bad_query_control_request(error.to_string(), None))
5126                    })?
5127            }
5128            None => {
5129                QueryId::random().map_err(|error| Box::new(query_error_response(&error, None)))?
5130            }
5131        },
5132    };
5133    let timeout_ms = match body_timeout_ms {
5134        Some(timeout_ms) => timeout_ms,
5135        None => match headers.get("x-mongreldb-timeout-ms") {
5136            Some(value) => value
5137                .to_str()
5138                .ok()
5139                .and_then(|value| value.parse::<u64>().ok())
5140                .ok_or_else(|| {
5141                    Box::new(bad_query_control_request(
5142                        "X-MongrelDB-Timeout-Ms must be a positive integer",
5143                        Some(query_id),
5144                    ))
5145                })?,
5146            None => state.reloadable.sql_default_timeout.as_millis(),
5147        },
5148    };
5149    if timeout_ms == 0 {
5150        return Err(Box::new(bad_query_control_request(
5151            "timeout_ms must be positive",
5152            Some(query_id),
5153        )));
5154    }
5155    let timeout = std::time::Duration::from_millis(timeout_ms);
5156    if timeout > state.reloadable.sql_max_timeout.get() {
5157        return Err(Box::new(bad_query_control_request(
5158            format!(
5159                "timeout_ms exceeds server maximum of {}",
5160                state.reloadable.sql_max_timeout.as_millis()
5161            ),
5162            Some(query_id),
5163        )));
5164    }
5165    Ok((
5166        SqlQueryOptions {
5167            query_id: Some(query_id),
5168            timeout: Some(timeout),
5169            owner: Some(owner),
5170            session_id,
5171            parent_control: None,
5172        },
5173        query_id,
5174    ))
5175}
5176
5177fn resolve_sql_output_limits(
5178    state: &AppState,
5179    request: &SqlRequest,
5180    query_id: QueryId,
5181) -> std::result::Result<(usize, usize), Box<Response>> {
5182    fn resolve(
5183        requested: Option<u64>,
5184        configured: usize,
5185        name: &str,
5186        query_id: QueryId,
5187    ) -> std::result::Result<usize, Box<Response>> {
5188        if requested == Some(0) {
5189            return Err(Box::new(bad_query_control_request(
5190                format!("{name} must be positive"),
5191                Some(query_id),
5192            )));
5193        }
5194        let requested = requested
5195            .and_then(|value| usize::try_from(value).ok())
5196            .unwrap_or(usize::MAX);
5197        Ok(requested.min(configured))
5198    }
5199
5200    Ok((
5201        resolve(
5202            request.max_output_rows,
5203            state.reloadable.sql_max_output_rows.get(),
5204            "max_output_rows",
5205            query_id,
5206        )?,
5207        resolve(
5208            request.max_output_bytes,
5209            state.reloadable.sql_max_output_bytes.get(),
5210            "max_output_bytes",
5211            query_id,
5212        )?,
5213    ))
5214}
5215
5216fn resolve_sql_pagination(
5217    headers: &axum::http::HeaderMap,
5218    request: &SqlRequest,
5219    output_limits: (usize, usize),
5220    registration: RegisteredQueryGuard,
5221    query_id: QueryId,
5222) -> Result<(RegisteredQueryGuard, Option<ResolvedSqlPagination>), Box<Response>> {
5223    let Some(pagination) = request.pagination.as_ref() else {
5224        return Ok((registration, None));
5225    };
5226    if requested_sql_idempotency_key(headers, request)
5227        .ok()
5228        .flatten()
5229        .is_some()
5230    {
5231        return Err(Box::new(registered_sql_error_response(
5232            registration,
5233            query_id,
5234            StatusCode::BAD_REQUEST,
5235            "INCOMPATIBLE_SQL_CONTROLS",
5236            "idempotency_key cannot be combined with SQL pagination",
5237            false,
5238        )));
5239    }
5240    if request
5241        .format
5242        .as_deref()
5243        .is_some_and(|format| format != "json")
5244    {
5245        return Err(Box::new(registered_sql_error_response(
5246            registration,
5247            query_id,
5248            StatusCode::BAD_REQUEST,
5249            "PAGINATION_REQUIRES_JSON",
5250            "SQL pagination supports JSON responses only",
5251            false,
5252        )));
5253    }
5254    registration.query().set_sql_metadata(&request.sql);
5255    if !mongreldb_query::is_single_read_only_query(&request.sql) {
5256        return Err(Box::new(registered_sql_error_response(
5257            registration,
5258            query_id,
5259            StatusCode::BAD_REQUEST,
5260            "PAGINATION_REQUIRES_SINGLE_READ_QUERY",
5261            "SQL pagination accepts exactly one read-only query statement",
5262            false,
5263        )));
5264    }
5265    let page_size = match usize::try_from(pagination.page_size_rows) {
5266        Ok(0) | Err(_) => {
5267            return Err(Box::new(registered_sql_error_response(
5268                registration,
5269                query_id,
5270                StatusCode::BAD_REQUEST,
5271                "INVALID_PAGINATION_OPTIONS",
5272                "pagination.page_size_rows must be positive",
5273                false,
5274            )))
5275        }
5276        Ok(value) => value.min(output_limits.0),
5277    };
5278    if pagination.projection.is_empty() || pagination.projection.len() > 128 {
5279        return Err(Box::new(registered_sql_error_response(
5280            registration,
5281            query_id,
5282            StatusCode::BAD_REQUEST,
5283            "INVALID_SQL_PROJECTION",
5284            "pagination.projection must contain between 1 and 128 output column names",
5285            false,
5286        )));
5287    }
5288    let mut seen = std::collections::HashSet::new();
5289    let metadata_bytes = pagination
5290        .projection
5291        .iter()
5292        .map(String::len)
5293        .fold(0usize, usize::saturating_add);
5294    if metadata_bytes > 16 * 1024
5295        || pagination.projection.iter().any(|column| {
5296            column.is_empty()
5297                || column == "*"
5298                || column.len() > 256
5299                || !seen.insert(column.as_str())
5300        })
5301    {
5302        return Err(Box::new(registered_sql_error_response(
5303            registration,
5304            query_id,
5305            StatusCode::BAD_REQUEST,
5306            "INVALID_SQL_PROJECTION",
5307            "pagination.projection requires unique explicit output names of at most 256 bytes",
5308            false,
5309        )));
5310    }
5311    let max_page_bytes = match pagination.max_page_bytes {
5312        Some(0) => {
5313            return Err(Box::new(registered_sql_error_response(
5314                registration,
5315                query_id,
5316                StatusCode::BAD_REQUEST,
5317                "INVALID_PAGINATION_OPTIONS",
5318                "pagination.max_page_bytes must be positive",
5319                false,
5320            )))
5321        }
5322        Some(value) => usize::try_from(value)
5323            .unwrap_or(usize::MAX)
5324            .min(output_limits.1),
5325        None => output_limits.1.min(1024 * 1024),
5326    };
5327    let token_cap = (output_limits.1.saturating_add(3) / 4).max(1);
5328    let max_page_tokens = match pagination.max_page_tokens {
5329        Some(0) => {
5330            return Err(Box::new(registered_sql_error_response(
5331                registration,
5332                query_id,
5333                StatusCode::BAD_REQUEST,
5334                "INVALID_PAGINATION_OPTIONS",
5335                "pagination.max_page_tokens must be positive",
5336                false,
5337            )))
5338        }
5339        Some(value) => usize::try_from(value).unwrap_or(usize::MAX).min(token_cap),
5340        None => (max_page_bytes.saturating_add(3) / 4).max(1),
5341    };
5342    Ok((
5343        registration,
5344        Some(ResolvedSqlPagination {
5345            projection: pagination.projection.clone(),
5346            limits: sql_pages::SqlPageLimits {
5347                rows: page_size,
5348                bytes: max_page_bytes,
5349                tokens: max_page_tokens,
5350            },
5351        }),
5352    ))
5353}
5354
5355fn requested_sql_idempotency_key(
5356    headers: &axum::http::HeaderMap,
5357    request: &SqlRequest,
5358) -> Result<Option<String>, &'static str> {
5359    let header = match headers.get("idempotency-key") {
5360        Some(value) => Some(
5361            value
5362                .to_str()
5363                .map_err(|_| "Idempotency-Key must be valid UTF-8")?,
5364        ),
5365        None => None,
5366    };
5367    match (request.idempotency_key.as_deref(), header) {
5368        (Some(body), Some(header)) if body != header => {
5369            Err("body idempotency_key and Idempotency-Key header must match")
5370        }
5371        (Some(body), _) => Ok(Some(body.to_owned())),
5372        (None, Some(header)) => Ok(Some(header.to_owned())),
5373        (None, None) => Ok(None),
5374    }
5375}
5376
5377fn sql_idempotency_binding(
5378    request: &SqlRequest,
5379    output_limits: (usize, usize),
5380    session_id: Option<&str>,
5381    expires_after_ms: u64,
5382) -> Result<sql_idempotency::SqlIdempotencyBinding, serde_json::Error> {
5383    let request_semantics = serde_json::to_vec(&json!({
5384        "format": request.format.as_deref().unwrap_or("json"),
5385        "max_output_rows": output_limits.0,
5386        "max_output_bytes": output_limits.1,
5387        "pagination": request.pagination.as_ref(),
5388    }))?;
5389    let session_semantics = session_id.map_or_else(
5390        || b"ephemeral".to_vec(),
5391        |session_id| {
5392            let mut semantics = b"session\0".to_vec();
5393            semantics.extend_from_slice(session_id.as_bytes());
5394            semantics
5395        },
5396    );
5397    Ok(sql_idempotency::SqlIdempotencyBinding {
5398        sql_fingerprint: mongreldb_query::normalized_sql_fingerprint(&request.sql),
5399        // `/sql` currently has no separate bind-parameter array. SQL literals
5400        // are covered by the normalized SQL fingerprint above.
5401        parameter_hash: sql_idempotency::hash(b"[]"),
5402        request_semantics_hash: sql_idempotency::hash(&request_semantics),
5403        session_semantics_hash: sql_idempotency::hash(&session_semantics),
5404        expires_after_ms,
5405    })
5406}
5407
5408struct SqlIdempotencyContext<'a> {
5409    headers: &'a axum::http::HeaderMap,
5410    request: &'a SqlRequest,
5411    output_limits: (usize, usize),
5412    owner: &'a str,
5413    session_id: Option<&'a str>,
5414    session_in_transaction: bool,
5415    query_id: QueryId,
5416}
5417
5418async fn begin_sql_idempotency(
5419    state: &AppState,
5420    context: SqlIdempotencyContext<'_>,
5421    registration: RegisteredQueryGuard,
5422) -> Result<
5423    (
5424        RegisteredQueryGuard,
5425        Option<sql_idempotency::SqlIdempotencyExecution>,
5426    ),
5427    Response,
5428> {
5429    let SqlIdempotencyContext {
5430        headers,
5431        request,
5432        output_limits,
5433        owner,
5434        session_id,
5435        session_in_transaction,
5436        query_id,
5437    } = context;
5438    let key = match requested_sql_idempotency_key(headers, request) {
5439        Ok(key) => key,
5440        Err(message) => {
5441            return Err(registered_sql_error_response(
5442                registration,
5443                query_id,
5444                StatusCode::BAD_REQUEST,
5445                "INVALID_IDEMPOTENCY_KEY",
5446                message,
5447                false,
5448            ))
5449        }
5450    };
5451    let Some(key) = key else {
5452        return Ok((registration, None));
5453    };
5454    match mongreldb_query::classify_sql_idempotency(&request.sql) {
5455        mongreldb_query::SqlIdempotencyClass::ReadOnly
5456        | mongreldb_query::SqlIdempotencyClass::Unsupported => {
5457            return Err(registered_sql_error_response(
5458                registration,
5459                query_id,
5460                StatusCode::BAD_REQUEST,
5461                "IDEMPOTENCY_REQUIRES_SINGLE_WRITE",
5462                "idempotency_key accepts one non-transaction SQL write statement",
5463                false,
5464            ));
5465        }
5466        mongreldb_query::SqlIdempotencyClass::SingleWrite => {}
5467    }
5468    if let Err(message) = sql_idempotency::SqlIdempotencyStore::validate_key(&key) {
5469        return Err(registered_sql_error_response(
5470            registration,
5471            query_id,
5472            StatusCode::BAD_REQUEST,
5473            "INVALID_IDEMPOTENCY_KEY",
5474            message,
5475            false,
5476        ));
5477    }
5478    if request
5479        .format
5480        .as_deref()
5481        .is_some_and(|format| format != "json")
5482    {
5483        return Err(registered_sql_error_response(
5484            registration,
5485            query_id,
5486            StatusCode::BAD_REQUEST,
5487            "IDEMPOTENCY_REQUIRES_JSON",
5488            "SQL idempotency supports buffered JSON responses only",
5489            false,
5490        ));
5491    }
5492    if session_in_transaction {
5493        return Err(registered_sql_error_response(
5494            registration,
5495            query_id,
5496            StatusCode::CONFLICT,
5497            "IDEMPOTENCY_UNSUPPORTED_IN_TRANSACTION",
5498            "SQL idempotency cannot be used inside an open session transaction",
5499            false,
5500        ));
5501    }
5502    registration.query().set_sql_metadata(&request.sql);
5503    let binding = match sql_idempotency_binding(
5504        request,
5505        output_limits,
5506        session_id,
5507        state.sql_idempotency.expires_after_ms(),
5508    ) {
5509        Ok(binding) => binding,
5510        Err(_) => {
5511            return Err(registered_sql_error_response(
5512                registration,
5513                query_id,
5514                StatusCode::INTERNAL_SERVER_ERROR,
5515                "SERIALIZATION_FAILED",
5516                "failed to serialize SQL idempotency request semantics",
5517                false,
5518            ))
5519        }
5520    };
5521    let begin = tokio::select! {
5522        begin = state.sql_idempotency.begin(owner, &key, binding) => begin,
5523        _ = registration.query().control().cancelled() => {
5524            return Err(tracked_query_error_response(
5525                state,
5526                &cancellation_checkpoint_error(registration.query()),
5527                Some(query_id),
5528            ));
5529        }
5530    };
5531    match begin {
5532        sql_idempotency::BeginResult::Execute(execution) => Ok((registration, Some(execution))),
5533        sql_idempotency::BeginResult::Replay {
5534            receipt,
5535            expires_at_ms,
5536        } => match restore_idempotency_replay(registration, &receipt) {
5537            Ok(()) => Err(sql_idempotency_receipt_response(
5538                query_id,
5539                &receipt,
5540                true,
5541                expires_at_ms,
5542                true,
5543            )),
5544            Err(error) => Err(tracked_query_error_response(state, &error, Some(query_id))),
5545        },
5546        sql_idempotency::BeginResult::Mismatch => Err(registered_sql_error_response(
5547            registration,
5548            query_id,
5549            StatusCode::CONFLICT,
5550            "IDEMPOTENCY_KEY_REUSE_MISMATCH",
5551            "idempotency key was already used with different SQL or request semantics",
5552            false,
5553        )),
5554        sql_idempotency::BeginResult::Indeterminate { created_at_ms } => Err(
5555            sql_idempotency_indeterminate_response(registration, query_id, created_at_ms),
5556        ),
5557        sql_idempotency::BeginResult::Full => Err(registered_sql_error_response(
5558            registration,
5559            query_id,
5560            StatusCode::SERVICE_UNAVAILABLE,
5561            "IDEMPOTENCY_STORE_FULL",
5562            "SQL idempotency receipt store is full",
5563            true,
5564        )),
5565        sql_idempotency::BeginResult::Unavailable(_reason) => Err(registered_sql_error_response(
5566            registration,
5567            query_id,
5568            StatusCode::SERVICE_UNAVAILABLE,
5569            "IDEMPOTENCY_STORE_UNAVAILABLE",
5570            "could not durably reserve the SQL idempotency key",
5571            true,
5572        )),
5573    }
5574}
5575
5576fn restore_idempotency_replay(
5577    registration: RegisteredQueryGuard,
5578    receipt: &sql_idempotency::SqlDurableReceipt,
5579) -> mongreldb_query::Result<()> {
5580    use mongreldb_query::{
5581        DurableOutcome, QueryTerminalError, QueryTerminalErrorCategory, QueryTerminalState,
5582        SerializationOutcome,
5583    };
5584
5585    let invalid_receipt = |field: &str| {
5586        mongreldb_query::MongrelQueryError::InvalidQueryState(format!(
5587            "durable SQL idempotency receipt has invalid {field}"
5588        ))
5589    };
5590    let terminal_state = match receipt.status.as_str() {
5591        "completed" => QueryTerminalState::Completed,
5592        "failed_before_commit" => QueryTerminalState::FailedBeforeCommit,
5593        "cancelled_before_commit" => QueryTerminalState::CancelledBeforeCommit,
5594        "deadline_before_commit" => QueryTerminalState::DeadlineBeforeCommit,
5595        "committed" => QueryTerminalState::Committed,
5596        "committed_with_error" => QueryTerminalState::CommittedWithError,
5597        "partially_committed" => QueryTerminalState::PartiallyCommitted,
5598        "cancelled_after_commit" => QueryTerminalState::CancelledAfterCommit,
5599        "deadline_after_commit" => QueryTerminalState::DeadlineAfterCommit,
5600        _ => return Err(invalid_receipt("terminal state")),
5601    };
5602    let serialization = match receipt.outcome.serialization.as_str() {
5603        "not_started" => SerializationOutcome::NotStarted,
5604        "in_progress" => SerializationOutcome::InProgress,
5605        "succeeded" => SerializationOutcome::Succeeded,
5606        "failed" => SerializationOutcome::Failed,
5607        _ => return Err(invalid_receipt("serialization state")),
5608    };
5609    let terminal_error = match receipt.terminal_error.as_ref() {
5610        Some(error) => Some(QueryTerminalError {
5611            code: error.code.clone(),
5612            category: match error.category.as_str() {
5613                "cancellation" => QueryTerminalErrorCategory::Cancellation,
5614                "deadline" => QueryTerminalErrorCategory::Deadline,
5615                "result_limit" => QueryTerminalErrorCategory::ResultLimit,
5616                "serialization" => QueryTerminalErrorCategory::Serialization,
5617                "execution" => QueryTerminalErrorCategory::Execution,
5618                _ => return Err(invalid_receipt("terminal error category")),
5619            },
5620        }),
5621        None => None,
5622    };
5623    let cancellation_reason = CancellationReason::from_protocol_str(&receipt.cancellation_reason)
5624        .ok_or_else(|| invalid_receipt("cancellation reason"))?;
5625    let phase = match receipt.server_state.as_str() {
5626        "completed" => SqlQueryPhase::Completed,
5627        "cancelled" => SqlQueryPhase::Cancelled,
5628        "failed" => SqlQueryPhase::Failed,
5629        _ => return Err(invalid_receipt("server state")),
5630    };
5631    let query = registration.into_query();
5632    query.restore_replayed_outcome(
5633        DurableOutcome {
5634            committed: receipt.outcome.committed,
5635            committed_statements: receipt.outcome.committed_statements,
5636            last_commit_epoch: receipt.outcome.last_commit_epoch,
5637            first_commit_statement_index: receipt.outcome.first_commit_statement_index,
5638            last_commit_statement_index: receipt.outcome.last_commit_statement_index,
5639            // The replayed outcome carries the same literal commit lineage as
5640            // the original write: the core ledger receipt rides the durable
5641            // HTTP receipt (S1B-005).
5642            commit_ts: receipt
5643                .commit_receipt
5644                .as_ref()
5645                .map(sql_idempotency::SqlCommitReceipt::commit_ts),
5646        },
5647        receipt.outcome.completed_statements,
5648        receipt.outcome.statement_index,
5649        serialization,
5650        terminal_error,
5651        terminal_state,
5652        cancellation_reason,
5653        phase,
5654    );
5655    query.try_complete()
5656}
5657
5658fn sql_idempotency_indeterminate_response(
5659    registration: RegisteredQueryGuard,
5660    query_id: QueryId,
5661    created_at_ms: Option<u64>,
5662) -> Response {
5663    registration.query().mark_outcome_unknown();
5664    registration.fail();
5665    with_query_id(
5666        (
5667            StatusCode::CONFLICT,
5668            Json(json!({
5669                "query_id": query_id.to_string(),
5670                "status": "outcome_unknown",
5671                "terminal_state": "outcome_unknown",
5672                "committed": null,
5673                "committed_statements": null,
5674                "last_commit_epoch": null,
5675                "last_commit_epoch_text": null,
5676                "first_commit_statement_index": null,
5677                "last_commit_statement_index": null,
5678                "completed_statements": null,
5679                "statement_index": null,
5680                "cancel_outcome": null,
5681                "cancellation_reason": null,
5682                "retryable": false,
5683                "server_state": "failed",
5684                "idempotency_replayed": true,
5685                "idempotency_intent_created_at_ms": created_at_ms,
5686                "outcome": {
5687                    "committed": null,
5688                    "committed_statements": null,
5689                    "last_commit_epoch": null,
5690                    "last_commit_epoch_text": null,
5691                    "first_commit_statement_index": null,
5692                    "last_commit_statement_index": null,
5693                    "completed_statements": null,
5694                    "statement_index": null,
5695                    "serialization": "unknown",
5696                },
5697                "error": {
5698                    "code": "QUERY_OUTCOME_UNKNOWN",
5699                    "message": "a durable write intent exists without a durable receipt; the SQL was not re-executed",
5700                    "query_id": query_id.to_string(),
5701                    "committed": null,
5702                    "retryable": false,
5703                }
5704            })),
5705        )
5706            .into_response(),
5707        query_id,
5708    )
5709}
5710
5711fn registered_sql_error_response(
5712    registration: RegisteredQueryGuard,
5713    query_id: QueryId,
5714    status: StatusCode,
5715    code: &'static str,
5716    message: impl Into<String>,
5717    retryable: bool,
5718) -> Response {
5719    let message = message.into();
5720    registration
5721        .query()
5722        .record_terminal_error(code, mongreldb_query::QueryTerminalErrorCategory::Execution);
5723    registration.fail();
5724    with_query_id(
5725        (
5726            status,
5727            Json(json!({
5728                "query_id": query_id.to_string(),
5729                "status": "failed_before_commit",
5730                "terminal_state": "failed_before_commit",
5731                "committed": false,
5732                "committed_statements": 0,
5733                "last_commit_epoch": null,
5734                "last_commit_epoch_text": null,
5735                "first_commit_statement_index": null,
5736                "last_commit_statement_index": null,
5737                "completed_statements": 0,
5738                "statement_index": 0,
5739                "cancel_outcome": null,
5740                "cancellation_reason": null,
5741                "retryable": retryable,
5742                "server_state": "failed",
5743                "outcome": {
5744                    "committed": false,
5745                    "committed_statements": 0,
5746                    "last_commit_epoch": null,
5747                    "last_commit_epoch_text": null,
5748                    "first_commit_statement_index": null,
5749                    "last_commit_statement_index": null,
5750                    "completed_statements": 0,
5751                    "statement_index": 0,
5752                    "serialization": "not_started",
5753                },
5754                "error": {
5755                    "code": code,
5756                    "message": message,
5757                    "query_id": query_id.to_string(),
5758                    "committed": false,
5759                    "retryable": retryable,
5760                }
5761            })),
5762        )
5763            .into_response(),
5764        query_id,
5765    )
5766}
5767
5768fn register_controlled_query(
5769    state: &AppState,
5770    session: &MongrelSession,
5771    options: SqlQueryOptions,
5772) -> std::result::Result<RegisteredSqlQuery, mongreldb_query::MongrelQueryError> {
5773    let query_id = options.query_id.ok_or_else(|| {
5774        mongreldb_query::MongrelQueryError::InvalidQueryState(
5775            "server query registration requires a query id".into(),
5776        )
5777    })?;
5778    let owner = options.owner.clone().unwrap_or_default();
5779    let session_id = options.session_id.clone();
5780    let _lifecycle = state
5781        .query_lifecycle
5782        .lock()
5783        .unwrap_or_else(|error| error.into_inner());
5784    let pre_cancel_reason = match state.pre_cancellations.lookup_for_registration(
5785        query_id,
5786        &owner,
5787        session_id.as_deref(),
5788    ) {
5789        pre_cancel::RegistrationLookup::NoReservation => None,
5790        pre_cancel::RegistrationLookup::Matching(reason) => Some(reason),
5791        pre_cancel::RegistrationLookup::ReservedByAnotherIdentity => {
5792            return Err(mongreldb_query::MongrelQueryError::QueryIdConflict { query_id });
5793        }
5794    };
5795    let query = session.register_query(options)?;
5796    let Some(reason) = pre_cancel_reason else {
5797        return Ok(query);
5798    };
5799    state
5800        .pre_cancellations
5801        .take(query_id, &owner, session_id.as_deref());
5802    query.request_cancel(reason);
5803    let error = query.checkpoint().err().unwrap_or_else(|| {
5804        mongreldb_query::MongrelQueryError::InvalidQueryState(format!(
5805            "pre-cancelled query {query_id} remained runnable"
5806        ))
5807    });
5808    query.fail();
5809    Err(error)
5810}
5811
5812/// Acquire outer node semaphore + hierarchical class admission for interactive SQL.
5813///
5814/// Choice (S1E-002): `sql_semaphore` is the absolute node hard cap; the
5815/// hierarchical scheduler enforces class/tenant fairness inside that bound.
5816/// QueueFull / TenantQuota map to ResourceExhausted (HTTP 503). Cancel while
5817/// waiting cancels the scheduler work id. RAII complete frees class slots.
5818/// Scheduler admission runs only with the `cluster` feature; standalone
5819/// builds take the semaphore-only fast path.
5820///
5821/// S4B: best-effort node governor evaluate runs before admission so
5822/// `ReduceAdmission` / cache eviction take effect under pressure.
5823async fn acquire_sql_permit(
5824    state: &AppState,
5825    session: &MongrelSession,
5826    query: &RegisteredSqlQuery,
5827) -> std::result::Result<admission::SqlAdmissionGuard, mongreldb_query::MongrelQueryError> {
5828    refresh_node_pressure(state);
5829
5830    session.fire_test_hook(mongreldb_query::SqlTestHookPoint::WaitingForSqlPermit);
5831
5832    // Outer node hard cap.
5833    let permit = tokio::select! {
5834        permit = Arc::clone(&state.sql_semaphore).acquire_owned() => permit.map_err(|_| {
5835            mongreldb_query::MongrelQueryError::InvalidQueryState(
5836                "SQL admission semaphore closed".into(),
5837            )
5838        })?,
5839        _ = query.control().cancelled() => {
5840            return Err(cancellation_checkpoint_error(query));
5841        }
5842    };
5843
5844    // P1.1: product SQL path admits through NodeAdmissionController::admit_parent
5845    // (hierarchical scheduler + parent memory budget), not a bare scheduler
5846    // clone. Standalone builds (no `cluster` feature) take the fast path: the
5847    // outer semaphore permit only.
5848    #[cfg(feature = "cluster")]
5849    {
5850        let class = mongreldb_core::WorkloadClass::InteractiveSql;
5851        let priority = admission::priority_for_class(&state.resource_groups, class);
5852        let types_query_id = mongreldb_types::ids::QueryId::from_bytes(*query.id().as_bytes());
5853
5854        const SQL_PARENT_BUDGET_BYTES: u64 = 16 * 1024 * 1024;
5855        let parent = match state
5856            .node_admission
5857            .admit_parent(
5858                admission::AdmitRequest {
5859                    tenant: "default",
5860                    class,
5861                    priority,
5862                    deadline: None,
5863                    query_id: Some(types_query_id),
5864                    tag: "sql",
5865                },
5866                mongreldb_core::MemoryClass::QueryExecution,
5867                SQL_PARENT_BUDGET_BYTES,
5868                query.control().cancelled(),
5869            )
5870            .await
5871        {
5872            Ok(parent) => parent,
5873            Err(admission::AdmitError::Cancelled) => {
5874                return Err(cancellation_checkpoint_error(query));
5875            }
5876            Err(error) => {
5877                return Err(admission::admit_error_to_query(error));
5878            }
5879        };
5880
5881        Ok(admission::SqlAdmissionGuard::new(permit, parent))
5882    }
5883    #[cfg(not(feature = "cluster"))]
5884    {
5885        Ok(admission::SqlAdmissionGuard::from_permit_only(permit))
5886    }
5887}
5888
5889/// Best-effort S4B evaluate: live DB reservations, AI semaphore saturation,
5890/// optional process RSS. Poisoned locks / missing metrics never fail the request.
5891fn refresh_node_pressure(state: &AppState) {
5892    let Ok(mut governor) = state.node_governor.lock() else {
5893        return;
5894    };
5895    let ai_capacity = default_ai_max_concurrent();
5896    let (db_reserved, db_max, db_gov) = match state.try_db() {
5897        Ok(db) => {
5898            let gov = db.memory_governor();
5899            (gov.total_used(), gov.max_bytes(), Some(gov))
5900        }
5901        Err(_) => (0, governor.governor.max_bytes(), None),
5902    };
5903    let inputs = admission::build_pressure_inputs(&admission::PressureInputSources {
5904        db_reserved_bytes: db_reserved,
5905        db_max_bytes: db_max,
5906        node_configured_max_bytes: governor.governor.max_bytes(),
5907        tablet_reserved_bytes: governor.tablet_reserved_bytes(),
5908        ai_capacity,
5909        ai_available: state.ai_semaphore.available_permits(),
5910        process_rss_bytes: admission::process_rss_bytes(),
5911    });
5912    admission::refresh_pressure(&mut governor, &inputs, &state.scheduler, db_gov);
5913}
5914
5915fn caller_may_manage_query(
5916    state: &AppState,
5917    principal: &Option<mongreldb_core::Principal>,
5918    owner: Option<&str>,
5919) -> bool {
5920    let current = current_request_principal(state, principal);
5921    let catalog_requires_auth = state
5922        .try_db()
5923        .map(|db| db.require_auth_enabled())
5924        .unwrap_or(false);
5925    if (principal.is_some()
5926        || state.auth_token.is_some()
5927        || state.user_auth
5928        || catalog_requires_auth)
5929        && current.is_none()
5930    {
5931        return false;
5932    }
5933    current.is_some_and(|principal| principal.is_admin)
5934        || owner == Some(request_owner(state, principal).as_str())
5935}
5936
5937fn query_phase_name(phase: SqlQueryPhase) -> &'static str {
5938    match phase {
5939        SqlQueryPhase::Queued => "queued",
5940        SqlQueryPhase::Planning => "planning",
5941        SqlQueryPhase::Executing => "executing",
5942        SqlQueryPhase::Streaming => "streaming",
5943        SqlQueryPhase::Serializing => "serializing",
5944        SqlQueryPhase::CommitCritical => "commit_critical",
5945        SqlQueryPhase::Cancelling => "cancelling",
5946        SqlQueryPhase::Completed => "completed",
5947        SqlQueryPhase::Failed => "failed",
5948        SqlQueryPhase::Cancelled => "cancelled",
5949    }
5950}
5951
5952fn commit_fence_outcome_name(outcome: mongreldb_query::CommitFenceOutcome) -> &'static str {
5953    match outcome {
5954        mongreldb_query::CommitFenceOutcome::NotReached => "not_reached",
5955        mongreldb_query::CommitFenceOutcome::CancelWon => "cancel_won",
5956        mongreldb_query::CommitFenceOutcome::CommitWon => "commit_won",
5957    }
5958}
5959
5960fn terminal_state_name(state: mongreldb_query::QueryTerminalState) -> &'static str {
5961    use mongreldb_query::QueryTerminalState;
5962    match state {
5963        QueryTerminalState::OutcomeUnknown => "outcome_unknown",
5964        QueryTerminalState::Completed => "completed",
5965        QueryTerminalState::FailedBeforeCommit => "failed_before_commit",
5966        QueryTerminalState::CancelledBeforeCommit => "cancelled_before_commit",
5967        QueryTerminalState::DeadlineBeforeCommit => "deadline_before_commit",
5968        QueryTerminalState::Committed => "committed",
5969        QueryTerminalState::CommittedWithError => "committed_with_error",
5970        QueryTerminalState::PartiallyCommitted => "partially_committed",
5971        QueryTerminalState::CancelledAfterCommit => "cancelled_after_commit",
5972        QueryTerminalState::DeadlineAfterCommit => "deadline_after_commit",
5973    }
5974}
5975
5976fn serialization_outcome_name(outcome: mongreldb_query::SerializationOutcome) -> &'static str {
5977    use mongreldb_query::SerializationOutcome;
5978    match outcome {
5979        SerializationOutcome::NotStarted => "not_started",
5980        SerializationOutcome::InProgress => "in_progress",
5981        SerializationOutcome::Succeeded => "succeeded",
5982        SerializationOutcome::Failed => "failed",
5983    }
5984}
5985
5986fn terminal_error_category_name(
5987    category: mongreldb_query::QueryTerminalErrorCategory,
5988) -> &'static str {
5989    use mongreldb_query::QueryTerminalErrorCategory;
5990    match category {
5991        QueryTerminalErrorCategory::Cancellation => "cancellation",
5992        QueryTerminalErrorCategory::Deadline => "deadline",
5993        QueryTerminalErrorCategory::ResultLimit => "result_limit",
5994        QueryTerminalErrorCategory::Serialization => "serialization",
5995        QueryTerminalErrorCategory::Execution => "execution",
5996    }
5997}
5998
5999fn terminal_error_retryable(error: Option<&mongreldb_query::QueryTerminalError>) -> bool {
6000    error.is_some_and(|error| {
6001        matches!(
6002            error.code.as_str(),
6003            "IDEMPOTENCY_STORE_FULL" | "IDEMPOTENCY_STORE_UNAVAILABLE"
6004        )
6005    })
6006}
6007
6008fn epoch_text(epoch: Option<u64>) -> Option<String> {
6009    epoch.map(|epoch| epoch.to_string())
6010}
6011
6012fn cancellation_reason_name(reason: CancellationReason) -> &'static str {
6013    reason.as_str()
6014}
6015
6016fn query_cancel_outcome(status: &mongreldb_query::QueryStatus) -> Option<&'static str> {
6017    match status.phase {
6018        SqlQueryPhase::CommitCritical => Some("too_late"),
6019        SqlQueryPhase::Completed | SqlQueryPhase::Failed | SqlQueryPhase::Cancelled => {
6020            Some("already_finished")
6021        }
6022        SqlQueryPhase::Cancelling => Some("accepted"),
6023        _ => None,
6024    }
6025}
6026
6027/// Structural HLC fields for HTTP durable recovery (P0.6 / 0.64 client surface).
6028fn hlc_json(ts: Option<mongreldb_types::hlc::HlcTimestamp>) -> serde_json::Value {
6029    match ts {
6030        Some(ts) if ts != mongreldb_types::hlc::HlcTimestamp::ZERO => json!({
6031            "physical_micros": ts.physical_micros,
6032            "logical": ts.logical,
6033            "node_tiebreaker": ts.node_tiebreaker,
6034        }),
6035        _ => serde_json::Value::Null,
6036    }
6037}
6038
6039fn query_outcome_json(status: Option<&mongreldb_query::QueryStatus>) -> serde_json::Value {
6040    let Some(status) = status else {
6041        return json!({
6042            "committed": false,
6043            "committed_statements": 0,
6044            "last_commit_epoch": null,
6045            "last_commit_epoch_text": null,
6046            "last_commit_hlc": null,
6047            "first_commit_statement_index": null,
6048            "last_commit_statement_index": null,
6049            "completed_statements": 0,
6050            "statement_index": 0,
6051            "serialization": "not_started",
6052            "serialization_state": "not_started",
6053            "terminal_state": null,
6054        });
6055    };
6056    if status.outcome_unknown {
6057        return json!({
6058            "committed": null,
6059            "committed_statements": null,
6060            "last_commit_epoch": null,
6061            "last_commit_epoch_text": null,
6062            "last_commit_hlc": null,
6063            "first_commit_statement_index": null,
6064            "last_commit_statement_index": null,
6065            "completed_statements": null,
6066            "statement_index": null,
6067            "serialization": "unknown",
6068            "serialization_state": "unknown",
6069            "terminal_state": null,
6070        });
6071    }
6072    let serialization = serialization_outcome_name(status.serialization_outcome);
6073    json!({
6074        "committed": status.durable_outcome.committed,
6075        "committed_statements": status.durable_outcome.committed_statements,
6076        "last_commit_epoch": status.durable_outcome.last_commit_epoch,
6077        "last_commit_epoch_text": epoch_text(status.durable_outcome.last_commit_epoch),
6078        "last_commit_hlc": hlc_json(status.durable_outcome.commit_ts),
6079        "first_commit_statement_index": status.durable_outcome.first_commit_statement_index,
6080        "last_commit_statement_index": status.durable_outcome.last_commit_statement_index,
6081        "completed_statements": status.completed_statements,
6082        "statement_index": status.statement_index,
6083        "serialization": serialization,
6084        "serialization_state": serialization,
6085        "terminal_state": status.terminal_state().map(terminal_state_name),
6086    })
6087}
6088
6089fn sql_terminal_idempotency_receipt(
6090    status: &mongreldb_query::QueryStatus,
6091) -> Option<sql_idempotency::SqlDurableReceipt> {
6092    if status.outcome_unknown {
6093        return None;
6094    }
6095    let terminal_state = status.terminal_state()?;
6096    if !status.durable_outcome.committed
6097        && terminal_state != mongreldb_query::QueryTerminalState::Completed
6098    {
6099        return None;
6100    }
6101    Some(sql_idempotency::SqlDurableReceipt {
6102        original_query_id: status.query_id.to_string(),
6103        status: status
6104            .terminal_state()
6105            .map(terminal_state_name)
6106            .unwrap_or("committed")
6107            .to_owned(),
6108        server_state: query_phase_name(status.phase).to_owned(),
6109        cancellation_reason: cancellation_reason_name(status.cancellation_reason).to_owned(),
6110        outcome: sql_idempotency::SqlReceiptOutcome {
6111            committed: status.durable_outcome.committed,
6112            committed_statements: status.durable_outcome.committed_statements,
6113            last_commit_epoch: status.durable_outcome.last_commit_epoch,
6114            last_commit_epoch_text: epoch_text(status.durable_outcome.last_commit_epoch),
6115            first_commit_statement_index: status.durable_outcome.first_commit_statement_index,
6116            last_commit_statement_index: status.durable_outcome.last_commit_statement_index,
6117            completed_statements: status.completed_statements,
6118            statement_index: status.statement_index,
6119            serialization: serialization_outcome_name(status.serialization_outcome).to_owned(),
6120            // Mirrors of the query-status outcome surface (0.64). The HLC is
6121            // attached with the core ledger receipt
6122            // (`SqlDurableReceipt::attach_commit_receipt`) so the replayed
6123            // response and the restored query status report the same stamp.
6124            last_commit_hlc: None,
6125            serialization_state: Some(
6126                serialization_outcome_name(status.serialization_outcome).to_owned(),
6127            ),
6128            terminal_state: status
6129                .terminal_state()
6130                .map(|state| terminal_state_name(state).to_owned()),
6131        },
6132        terminal_error: status.terminal_error.as_ref().map(|error| {
6133            sql_idempotency::SqlReceiptTerminalError {
6134                code: error.code.clone(),
6135                category: terminal_error_category_name(error.category).to_owned(),
6136            }
6137        }),
6138        commit_receipt: None,
6139    })
6140}
6141
6142fn sql_idempotency_receipt_response(
6143    query_id: QueryId,
6144    receipt: &sql_idempotency::SqlDurableReceipt,
6145    replayed: bool,
6146    expires_at_ms: u64,
6147    persisted: bool,
6148) -> Response {
6149    let mut body = json!({
6150        "query_id": query_id.to_string(),
6151        "original_query_id": receipt.original_query_id,
6152        "status": receipt.status,
6153        "terminal_state": receipt.status,
6154        "server_state": receipt.server_state,
6155        "cancel_outcome": "already_finished",
6156        "cancellation_reason": receipt.cancellation_reason,
6157        "committed": receipt.outcome.committed,
6158        "committed_statements": receipt.outcome.committed_statements,
6159        "last_commit_epoch": receipt.outcome.last_commit_epoch,
6160        "last_commit_epoch_text": receipt.outcome.last_commit_epoch_text.as_deref(),
6161        "first_commit_statement_index": receipt.outcome.first_commit_statement_index,
6162        "last_commit_statement_index": receipt.outcome.last_commit_statement_index,
6163        "completed_statements": receipt.outcome.completed_statements,
6164        "statement_index": receipt.outcome.statement_index,
6165        "retryable": false,
6166        "idempotency_replayed": replayed,
6167        "idempotency_persisted": persisted,
6168        "idempotency_expires_at_ms": expires_at_ms,
6169        "outcome": receipt.outcome,
6170        "terminal_error": receipt.terminal_error,
6171    });
6172    // S1B-005 additive: the core commit log's receipt for the committed
6173    // write, when the durable `TXN_IDEMPOTENCY` ledger recorded one. Absent
6174    // (not null) for receipts predating the unification and for no-op writes
6175    // that never committed.
6176    if let Some(commit_receipt) = &receipt.commit_receipt {
6177        body["commit_receipt"] = json!(commit_receipt);
6178    }
6179    let mut response = Json(body).into_response();
6180    response.headers_mut().insert(
6181        "idempotency-replayed",
6182        axum::http::HeaderValue::from_static(if replayed { "true" } else { "false" }),
6183    );
6184    response.headers_mut().insert(
6185        "idempotency-persisted",
6186        axum::http::HeaderValue::from_static(if persisted { "true" } else { "false" }),
6187    );
6188    if let Ok(value) = axum::http::HeaderValue::from_str(&receipt.original_query_id) {
6189        response
6190            .headers_mut()
6191            .insert("x-mongreldb-original-query-id", value);
6192    }
6193    with_query_id(response, query_id)
6194}
6195
6196fn terminal_server_error_response(
6197    state: &AppState,
6198    query_id: QueryId,
6199    http_status: StatusCode,
6200    base_code: &'static str,
6201    message: impl Into<String>,
6202) -> Response {
6203    let status = state.query_registry.status(query_id);
6204    let committed = status
6205        .as_ref()
6206        .is_some_and(|status| status.durable_outcome.committed);
6207    let code = if committed && base_code.starts_with("SERIALIZATION_") {
6208        "SERIALIZATION_FAILED_AFTER_COMMIT"
6209    } else {
6210        base_code
6211    };
6212    // Stable taxonomy category for this terminal code (spec 9.7); mappings
6213    // follow the `MongrelError::category()` precedents.
6214    let category = {
6215        use mongreldb_types::errors::ErrorCategory;
6216        match code {
6217            "RESULT_LIMIT_EXCEEDED" | "SQL_PAGE_STORE_FULL" | "ENTROPY_UNAVAILABLE" => {
6218                ErrorCategory::ResourceExhausted
6219            }
6220            // Client/supplied-shape disagreements and codec failures are
6221            // contract disagreements (the core `InvalidArgument` /
6222            // `Serialization` precedents).
6223            "INVALID_SQL_PROJECTION" | "INVALID_PAGE_OFFSET" => {
6224                ErrorCategory::ClusterVersionMismatch
6225            }
6226            _ if code.starts_with("SERIALIZATION_") => ErrorCategory::ClusterVersionMismatch,
6227            _ => ErrorCategory::ReplicaUnavailable,
6228        }
6229    };
6230    let response_status = status
6231        .as_ref()
6232        .and_then(mongreldb_query::QueryStatus::terminal_state)
6233        .map(terminal_state_name)
6234        .unwrap_or(if committed {
6235            "committed_with_error"
6236        } else {
6237            "failed_before_commit"
6238        });
6239    let outcome = query_outcome_json(status.as_ref());
6240    with_query_id(
6241        (
6242            http_status,
6243            Json(json!({
6244                "query_id": query_id.to_string(),
6245                "status": response_status,
6246                "terminal_state": response_status,
6247                "committed": committed,
6248                "committed_statements": status.as_ref().map_or(0, |status| status.durable_outcome.committed_statements),
6249                "last_commit_epoch": status.as_ref().and_then(|status| status.durable_outcome.last_commit_epoch),
6250                "last_commit_epoch_text": epoch_text(status.as_ref().and_then(|status| status.durable_outcome.last_commit_epoch)),
6251                "first_commit_statement_index": status.as_ref().and_then(|status| status.durable_outcome.first_commit_statement_index),
6252                "last_commit_statement_index": status.as_ref().and_then(|status| status.durable_outcome.last_commit_statement_index),
6253                "completed_statements": status.as_ref().map_or(0, |status| status.completed_statements),
6254                "statement_index": status.as_ref().map_or(0, |status| status.statement_index),
6255                "cancel_outcome": null,
6256                "cancellation_reason": status.as_ref().map(|status| cancellation_reason_name(status.cancellation_reason)),
6257                "retryable": false,
6258                "server_state": status.as_ref().map(|status| query_phase_name(status.phase)),
6259                "outcome": outcome,
6260                "error": {
6261                    "code": code,
6262                    "message": message.into(),
6263                    "category": category.to_string(),
6264                    "category_code": category.code(),
6265                    "query_id": query_id.to_string(),
6266                    "committed": committed,
6267                    "retryable": false,
6268                }
6269            })),
6270        )
6271            .into_response(),
6272        query_id,
6273    )
6274}
6275
6276fn query_not_found_response(query_id: Option<QueryId>) -> Response {
6277    let mut response = (
6278        StatusCode::NOT_FOUND,
6279        Json(json!({
6280            "query_id": query_id.map(|value| value.to_string()),
6281            "status": "unknown",
6282            "terminal_state": null,
6283            "committed": null,
6284            "committed_statements": null,
6285            "last_commit_epoch": null,
6286            "last_commit_epoch_text": null,
6287            "first_commit_statement_index": null,
6288            "last_commit_statement_index": null,
6289            "completed_statements": null,
6290            "statement_index": null,
6291            "cancel_outcome": "not_found",
6292            "cancellation_reason": null,
6293            "retryable": false,
6294            "server_state": "not_found",
6295            "outcome": {
6296                "committed": null,
6297                "committed_statements": null,
6298                "last_commit_epoch": null,
6299                "last_commit_epoch_text": null,
6300                "first_commit_statement_index": null,
6301                "last_commit_statement_index": null,
6302                "completed_statements": null,
6303                "statement_index": null,
6304                "serialization": "unknown",
6305            },
6306            "error": {
6307                "code": "QUERY_NOT_FOUND",
6308                "message": "query not found",
6309                "query_id": query_id.map(|value| value.to_string()),
6310                "committed": null,
6311                "retryable": false,
6312            }
6313        })),
6314    )
6315        .into_response();
6316    if let Some(query_id) = query_id {
6317        add_query_id_header(&mut response, query_id);
6318    }
6319    response
6320}
6321
6322fn query_session_header(
6323    headers: &axum::http::HeaderMap,
6324    query_id: Option<QueryId>,
6325) -> std::result::Result<Option<String>, Box<Response>> {
6326    match headers.get("x-session-id") {
6327        Some(value) => match value.to_str() {
6328            Ok(value) if value.len() <= 256 => Ok(Some(value.to_owned())),
6329            _ => Err(Box::new(bad_query_control_request(
6330                "X-Session-ID must be valid text no longer than 256 bytes",
6331                query_id,
6332            ))),
6333        },
6334        None => Ok(None),
6335    }
6336}
6337
6338fn pre_cancelled_query_response(
6339    query_id: QueryId,
6340    reason: CancellationReason,
6341    status: StatusCode,
6342) -> Response {
6343    with_query_id(
6344        (
6345            status,
6346            Json(json!({
6347                "query_id": query_id.to_string(),
6348                "status": "cancelled_before_start",
6349                "terminal_state": "cancelled_before_start",
6350                "state": "pre_cancelled",
6351                "server_state": "pre_cancelled",
6352                "cancel_outcome": "pre_cancelled",
6353                "committed": false,
6354                "committed_statements": 0,
6355                "last_commit_epoch": null,
6356                "last_commit_epoch_text": null,
6357                "first_commit_statement_index": null,
6358                "last_commit_statement_index": null,
6359                "completed_statements": 0,
6360                "statement_index": 0,
6361                "cancellation_reason": cancellation_reason_name(reason),
6362                "outcome": {
6363                    "committed": false,
6364                    "committed_statements": 0,
6365                    "last_commit_epoch": null,
6366                    "last_commit_epoch_text": null,
6367                    "first_commit_statement_index": null,
6368                    "last_commit_statement_index": null,
6369                    "completed_statements": 0,
6370                    "statement_index": 0,
6371                    "serialization": "not_started",
6372                },
6373                "terminal_error": {
6374                    "code": "QUERY_CANCELLED",
6375                    "category": "cancellation",
6376                },
6377                "retryable": false,
6378            })),
6379        )
6380            .into_response(),
6381        query_id,
6382    )
6383}
6384
6385fn compact_finished_query_response(status: &CompactFinishedQuery) -> Response {
6386    let query_id = status.query_id;
6387    let durable = &status.durable_outcome;
6388    let outcome_unknown =
6389        status.terminal_state == mongreldb_query::QueryTerminalState::OutcomeUnknown;
6390    let terminal_error = status.terminal_error.as_ref().map(|error| {
6391        json!({
6392            "code": error.code,
6393            "category": terminal_error_category_name(error.category),
6394        })
6395    });
6396    with_query_id(
6397        Json(json!({
6398            "detail": "compact",
6399            "query_id": query_id.to_string(),
6400            "status": terminal_state_name(status.terminal_state),
6401            "terminal_state": terminal_state_name(status.terminal_state),
6402            "state": query_phase_name(status.phase),
6403            "server_state": query_phase_name(status.phase),
6404            "cancel_outcome": "already_finished",
6405            "code": "QUERY_ALREADY_FINISHED",
6406            "committed": (!outcome_unknown).then_some(durable.committed),
6407            "committed_statements": (!outcome_unknown).then_some(durable.committed_statements),
6408            "last_commit_epoch": (!outcome_unknown).then_some(durable.last_commit_epoch).flatten(),
6409            "last_commit_epoch_text": (!outcome_unknown).then_some(epoch_text(durable.last_commit_epoch)).flatten(),
6410            "first_commit_statement_index": (!outcome_unknown).then_some(durable.first_commit_statement_index).flatten(),
6411            "last_commit_statement_index": (!outcome_unknown).then_some(durable.last_commit_statement_index).flatten(),
6412            "completed_statements": (!outcome_unknown).then_some(status.completed_statements),
6413            "statement_index": (!outcome_unknown).then_some(status.statement_index),
6414            "cancellation_reason": cancellation_reason_name(status.cancellation_reason),
6415            "outcome": {
6416                "committed": (!outcome_unknown).then_some(durable.committed),
6417                "committed_statements": (!outcome_unknown).then_some(durable.committed_statements),
6418                "last_commit_epoch": (!outcome_unknown).then_some(durable.last_commit_epoch).flatten(),
6419                "last_commit_epoch_text": (!outcome_unknown).then_some(epoch_text(durable.last_commit_epoch)).flatten(),
6420                "first_commit_statement_index": (!outcome_unknown).then_some(durable.first_commit_statement_index).flatten(),
6421                "last_commit_statement_index": (!outcome_unknown).then_some(durable.last_commit_statement_index).flatten(),
6422                "completed_statements": (!outcome_unknown).then_some(status.completed_statements),
6423                "statement_index": (!outcome_unknown).then_some(status.statement_index),
6424                "serialization": serialization_outcome_name(status.serialization_outcome),
6425            },
6426            "terminal_error": terminal_error,
6427            "retryable": terminal_error_retryable(status.terminal_error.as_ref()),
6428        }))
6429        .into_response(),
6430        query_id,
6431    )
6432}
6433
6434async fn query_status(
6435    State(state): State<Arc<AppState>>,
6436    OptionalPrincipal(principal): OptionalPrincipal,
6437    Path(query_id): Path<String>,
6438    headers: axum::http::HeaderMap,
6439) -> Response {
6440    let Ok(query_id) = query_id.parse::<QueryId>() else {
6441        return query_not_found_response(None);
6442    };
6443    if !request_identity_is_current(&state, &principal) {
6444        return query_not_found_response(Some(query_id));
6445    }
6446    let requested_session = match query_session_header(&headers, Some(query_id)) {
6447        Ok(session_id) => session_id,
6448        Err(response) => return *response,
6449    };
6450    let owner = request_owner(&state, &principal);
6451    let is_admin =
6452        current_request_principal(&state, &principal).is_some_and(|principal| principal.is_admin);
6453    let _lifecycle = state
6454        .query_lifecycle
6455        .lock()
6456        .unwrap_or_else(|error| error.into_inner());
6457    let Some(status) = state.query_registry.status(query_id) else {
6458        if let Some(finished) = state.query_registry.compact_finished_status(query_id) {
6459            if !caller_may_manage_query(&state, &principal, finished.owner.as_deref())
6460                || requested_session
6461                    .as_deref()
6462                    .is_some_and(|session| finished.session_id.as_deref() != Some(session))
6463            {
6464                return query_not_found_response(Some(query_id));
6465            }
6466            return compact_finished_query_response(&finished);
6467        }
6468        let reason = state
6469            .pre_cancellations
6470            .reason(query_id, &owner, requested_session.as_deref())
6471            .or_else(|| {
6472                is_admin.then(|| match requested_session.as_deref() {
6473                    Some(session_id) => state
6474                        .pre_cancellations
6475                        .reason_for_query_in_session(query_id, session_id),
6476                    None => state.pre_cancellations.reason_for_query(query_id),
6477                })?
6478            });
6479        if let Some(reason) = reason {
6480            return pre_cancelled_query_response(query_id, reason, StatusCode::OK);
6481        }
6482        return query_not_found_response(Some(query_id));
6483    };
6484    if !caller_may_manage_query(&state, &principal, status.owner.as_deref())
6485        || requested_session
6486            .as_deref()
6487            .is_some_and(|session| status.session_id.as_deref() != Some(session))
6488    {
6489        return query_not_found_response(Some(query_id));
6490    }
6491    let terminal_status = status.terminal_state().map(terminal_state_name);
6492    let terminal_error = status.terminal_error.as_ref().map(|error| {
6493        json!({
6494            "code": error.code,
6495            "category": terminal_error_category_name(error.category),
6496        })
6497    });
6498    let retryable = terminal_error_retryable(status.terminal_error.as_ref());
6499    let cancel_outcome = query_cancel_outcome(&status);
6500    let outcome = query_outcome_json(Some(&status));
6501    let response = Json(json!({
6502        "query_id": query_id.to_string(),
6503        "status": terminal_status.unwrap_or(if status.durable_outcome.committed {
6504            "committed"
6505        } else {
6506            "running"
6507        }),
6508        "terminal_state": terminal_status,
6509        "state": query_phase_name(status.phase),
6510        "server_state": query_phase_name(status.phase),
6511        "started_ms_ago": status.started_at.elapsed().as_millis(),
6512        "deadline_ms_remaining": status.deadline.map(|deadline| {
6513            deadline.saturating_duration_since(std::time::Instant::now()).as_millis()
6514        }),
6515        "session_id": status.session_id,
6516        "operation": status.operation,
6517        "committed": (!status.outcome_unknown).then_some(status.committed),
6518        "committed_statements": (!status.outcome_unknown).then_some(status.durable_outcome.committed_statements),
6519        "last_commit_epoch": (!status.outcome_unknown).then_some(status.durable_outcome.last_commit_epoch).flatten(),
6520        "last_commit_epoch_text": (!status.outcome_unknown).then_some(epoch_text(status.durable_outcome.last_commit_epoch)).flatten(),
6521        "last_commit_hlc": if !status.outcome_unknown {
6522            hlc_json(status.durable_outcome.commit_ts)
6523        } else {
6524            serde_json::Value::Null
6525        },
6526        "first_commit_statement_index": (!status.outcome_unknown).then_some(status.durable_outcome.first_commit_statement_index).flatten(),
6527        "last_commit_statement_index": (!status.outcome_unknown).then_some(status.durable_outcome.last_commit_statement_index).flatten(),
6528        "cancellation_reason": cancellation_reason_name(status.cancellation_reason),
6529        "completed_statements": (!status.outcome_unknown).then_some(status.completed_statements),
6530        "statement_index": (!status.outcome_unknown).then_some(status.statement_index),
6531        "cancel_outcome": cancel_outcome,
6532        "retryable": retryable,
6533        "outcome": outcome,
6534        // Nested durable object (parity with native DurableOutcome) for clients
6535        // that prefer a single structural recovery payload (0.64).
6536        "durable": outcome,
6537        "terminal_error": terminal_error,
6538        "trace": {
6539            "queue_duration_us": status.queue_duration.as_micros(),
6540            "planning_duration_us": status.planning_duration.as_micros(),
6541            "execution_duration_us": status.execution_duration.as_micros(),
6542            "serialization_duration_us": status.serialization_duration.as_micros(),
6543            "cancel_requested_phase": status.cancel_requested_phase.map(query_phase_name),
6544            "cancel_observed_phase": status.cancel_observed_phase.map(query_phase_name),
6545            "commit_fence_outcome": commit_fence_outcome_name(status.commit_fence_outcome),
6546        },
6547    }))
6548    .into_response();
6549    with_query_id(response, query_id)
6550}
6551
6552async fn cancel_query(
6553    State(state): State<Arc<AppState>>,
6554    OptionalPrincipal(principal): OptionalPrincipal,
6555    Path(query_id): Path<String>,
6556    headers: axum::http::HeaderMap,
6557) -> Response {
6558    let Ok(query_id) = query_id.parse::<QueryId>() else {
6559        return query_not_found_response(None);
6560    };
6561    if !request_identity_is_current(&state, &principal) {
6562        return query_not_found_response(Some(query_id));
6563    }
6564    let requested_session = match query_session_header(&headers, Some(query_id)) {
6565        Ok(session_id) => session_id,
6566        Err(response) => return *response,
6567    };
6568    let owner = request_owner(&state, &principal);
6569    let _lifecycle = state
6570        .query_lifecycle
6571        .lock()
6572        .unwrap_or_else(|error| error.into_inner());
6573    state.metrics.inc_sql_cancel_requests();
6574    let Some(status) = state.query_registry.status(query_id) else {
6575        if let Some(finished) = state.query_registry.compact_finished_status(query_id) {
6576            if !caller_may_manage_query(&state, &principal, finished.owner.as_deref())
6577                || requested_session
6578                    .as_deref()
6579                    .is_some_and(|session| finished.session_id.as_deref() != Some(session))
6580            {
6581                return query_not_found_response(Some(query_id));
6582            }
6583            return compact_finished_query_response(&finished);
6584        }
6585        return match state.pre_cancellations.insert(
6586            query_id,
6587            &owner,
6588            requested_session.as_deref(),
6589            CancellationReason::ClientRequest,
6590        ) {
6591            Ok(()) => {
6592                state.metrics.inc_sql_commit_cancel_winner_cancel();
6593                pre_cancelled_query_response(
6594                    query_id,
6595                    CancellationReason::ClientRequest,
6596                    StatusCode::ACCEPTED,
6597                )
6598            }
6599            Err(pre_cancel::InsertError::MetadataTooLarge) => bad_query_control_request(
6600                "query owner or session metadata exceeds 256 bytes",
6601                Some(query_id),
6602            ),
6603            Err(
6604                pre_cancel::InsertError::Full
6605                | pre_cancel::InsertError::OwnerLimit
6606                | pre_cancel::InsertError::RateLimited,
6607            ) => with_query_id(
6608                (
6609                    StatusCode::TOO_MANY_REQUESTS,
6610                    Json(json!({
6611                        "query_id": query_id.to_string(),
6612                        "status": "failed_before_commit",
6613                        "terminal_state": "failed_before_commit",
6614                        "server_state": "failed",
6615                        "cancel_outcome": null,
6616                        "cancellation_reason": null,
6617                        "committed": false,
6618                        "committed_statements": 0,
6619                        "last_commit_epoch": null,
6620                        "last_commit_epoch_text": null,
6621                        "first_commit_statement_index": null,
6622                        "last_commit_statement_index": null,
6623                        "completed_statements": 0,
6624                        "statement_index": 0,
6625                        "retryable": true,
6626                        "outcome": {
6627                            "committed": false,
6628                            "committed_statements": 0,
6629                            "last_commit_epoch": null,
6630                            "last_commit_epoch_text": null,
6631                            "first_commit_statement_index": null,
6632                            "last_commit_statement_index": null,
6633                            "completed_statements": 0,
6634                            "statement_index": 0,
6635                            "serialization": "not_started",
6636                        },
6637                        "error": {
6638                            "code": "QUERY_REGISTRY_FULL",
6639                            "message": "pre-registration cancellation limit reached",
6640                            "query_id": query_id.to_string(),
6641                            "committed": false,
6642                            "retryable": true,
6643                        }
6644                    })),
6645                )
6646                    .into_response(),
6647                query_id,
6648            ),
6649        };
6650    };
6651    if !caller_may_manage_query(&state, &principal, status.owner.as_deref())
6652        || requested_session
6653            .as_deref()
6654            .is_some_and(|session| status.session_id.as_deref() != Some(session))
6655    {
6656        return query_not_found_response(Some(query_id));
6657    }
6658    let (http_status, mut body) = match state.query_registry.cancel(query_id) {
6659        CancelOutcome::Accepted => (
6660            {
6661                state.metrics.inc_sql_commit_cancel_winner_cancel();
6662                StatusCode::ACCEPTED
6663            },
6664            json!({
6665                "query_id": query_id.to_string(),
6666                "state": "cancellation_requested",
6667                "cancel_outcome": "accepted",
6668            }),
6669        ),
6670        CancelOutcome::AlreadyCancelling => (
6671            StatusCode::OK,
6672            json!({
6673                "query_id": query_id.to_string(),
6674                "state": "cancelling",
6675                "cancel_outcome": "already_cancelling",
6676            }),
6677        ),
6678        CancelOutcome::TooLate => (
6679            {
6680                state.metrics.inc_sql_commit_cancel_winner_commit();
6681                StatusCode::CONFLICT
6682            },
6683            json!({
6684                "query_id": query_id.to_string(),
6685                "state": "commit_critical",
6686                "cancel_outcome": "too_late",
6687                "committed": status.durable_outcome.committed,
6688                "outcome": query_outcome_json(Some(&status)),
6689                "retryable": false,
6690                "error": {
6691                    "code": "CANCEL_TOO_LATE",
6692                    "message": "the query has entered its durable commit phase",
6693                    "committed": status.durable_outcome.committed,
6694                    "retryable": false,
6695                }
6696            }),
6697        ),
6698        CancelOutcome::AlreadyFinished => (
6699            StatusCode::OK,
6700            json!({
6701                "query_id": query_id.to_string(),
6702                "state": "finished",
6703                "status": status.terminal_state().map(terminal_state_name),
6704                "cancel_outcome": "already_finished",
6705                "code": "QUERY_ALREADY_FINISHED",
6706                "committed": status.durable_outcome.committed,
6707                "outcome": query_outcome_json(Some(&status)),
6708                "retryable": false,
6709            }),
6710        ),
6711        CancelOutcome::NotFound => return query_not_found_response(Some(query_id)),
6712    };
6713    let status = state.query_registry.status(query_id).unwrap_or(status);
6714    let response_status = status.terminal_state().map(terminal_state_name).unwrap_or(
6715        if status.durable_outcome.committed {
6716            "committed"
6717        } else {
6718            "running"
6719        },
6720    );
6721    if let Some(body) = body.as_object_mut() {
6722        body.insert("status".into(), json!(response_status));
6723        body.insert(
6724            "terminal_state".into(),
6725            json!(status.terminal_state().map(terminal_state_name)),
6726        );
6727        body.insert(
6728            "committed".into(),
6729            json!((!status.outcome_unknown).then_some(status.durable_outcome.committed)),
6730        );
6731        body.insert(
6732            "committed_statements".into(),
6733            json!((!status.outcome_unknown).then_some(status.durable_outcome.committed_statements)),
6734        );
6735        body.insert(
6736            "last_commit_epoch".into(),
6737            json!((!status.outcome_unknown)
6738                .then_some(status.durable_outcome.last_commit_epoch)
6739                .flatten()),
6740        );
6741        body.insert(
6742            "last_commit_epoch_text".into(),
6743            json!((!status.outcome_unknown)
6744                .then_some(epoch_text(status.durable_outcome.last_commit_epoch))
6745                .flatten()),
6746        );
6747        body.insert(
6748            "first_commit_statement_index".into(),
6749            json!((!status.outcome_unknown)
6750                .then_some(status.durable_outcome.first_commit_statement_index)
6751                .flatten()),
6752        );
6753        body.insert(
6754            "last_commit_statement_index".into(),
6755            json!((!status.outcome_unknown)
6756                .then_some(status.durable_outcome.last_commit_statement_index)
6757                .flatten()),
6758        );
6759        body.insert(
6760            "completed_statements".into(),
6761            json!((!status.outcome_unknown).then_some(status.completed_statements)),
6762        );
6763        body.insert(
6764            "statement_index".into(),
6765            json!((!status.outcome_unknown).then_some(status.statement_index)),
6766        );
6767        body.insert(
6768            "cancellation_reason".into(),
6769            json!(cancellation_reason_name(status.cancellation_reason)),
6770        );
6771        body.insert("retryable".into(), json!(false));
6772        body.insert("server_state".into(), json!(query_phase_name(status.phase)));
6773        body.insert("outcome".into(), query_outcome_json(Some(&status)));
6774    }
6775    with_query_id((http_status, Json(body)).into_response(), query_id)
6776}
6777
6778#[derive(Deserialize)]
6779#[serde(deny_unknown_fields)]
6780struct SqlContinuationRequest {
6781    cursor: String,
6782    #[serde(default)]
6783    operation_id: Option<QueryId>,
6784    #[serde(default)]
6785    timeout_ms: Option<u64>,
6786}
6787
6788fn register_page_operation(
6789    state: &AppState,
6790    options: SqlQueryOptions,
6791) -> mongreldb_query::Result<RegisteredSqlQuery> {
6792    let query_id = options.query_id.ok_or_else(|| {
6793        mongreldb_query::MongrelQueryError::InvalidQueryState(
6794            "page operation registration requires an operation id".into(),
6795        )
6796    })?;
6797    let owner = options.owner.clone().unwrap_or_default();
6798    let session_id = options.session_id.clone();
6799    let _lifecycle = state
6800        .query_lifecycle
6801        .lock()
6802        .unwrap_or_else(|error| error.into_inner());
6803    let reason = match state.pre_cancellations.lookup_for_registration(
6804        query_id,
6805        &owner,
6806        session_id.as_deref(),
6807    ) {
6808        pre_cancel::RegistrationLookup::NoReservation => None,
6809        pre_cancel::RegistrationLookup::Matching(reason) => Some(reason),
6810        pre_cancel::RegistrationLookup::ReservedByAnotherIdentity => {
6811            return Err(mongreldb_query::MongrelQueryError::QueryIdConflict { query_id });
6812        }
6813    };
6814    let query = state.query_registry.register(options)?;
6815    if let Some(reason) = reason {
6816        state
6817            .pre_cancellations
6818            .take(query_id, &owner, session_id.as_deref());
6819        query.request_cancel(reason);
6820        let error = cancellation_checkpoint_error(&query);
6821        query.fail();
6822        return Err(error);
6823    }
6824    Ok(query)
6825}
6826
6827async fn continue_sql_page(
6828    State(state): State<Arc<AppState>>,
6829    OptionalPrincipal(principal): OptionalPrincipal,
6830    headers: axum::http::HeaderMap,
6831    Json(request): Json<SqlContinuationRequest>,
6832) -> Response {
6833    let query_id = match request.operation_id {
6834        Some(query_id) => query_id,
6835        None => match QueryId::random() {
6836            Ok(query_id) => query_id,
6837            Err(error) => return query_error_response(&error, None),
6838        },
6839    };
6840    if !request_identity_is_current(&state, &principal) {
6841        return with_query_id(
6842            sql_cursor_error_response(
6843                StatusCode::NOT_FOUND,
6844                "SQL_CURSOR_NOT_FOUND",
6845                "SQL continuation result is unavailable",
6846                query_id,
6847            ),
6848            query_id,
6849        );
6850    }
6851    let owner = request_owner(&state, &principal);
6852    let session_id = match query_session_header(&headers, Some(query_id)) {
6853        Ok(session_id) => session_id,
6854        Err(response) => return *response,
6855    };
6856    let timeout_ms = request.timeout_ms.unwrap_or_else(|| {
6857        state
6858            .sql_page_default_timeout
6859            .as_millis()
6860            .min(u128::from(u64::MAX)) as u64
6861    });
6862    if timeout_ms == 0 || Duration::from_millis(timeout_ms) > state.sql_page_max_timeout {
6863        return bad_query_control_request(
6864            format!(
6865                "timeout_ms must be positive and no greater than {}",
6866                state.sql_page_max_timeout.as_millis()
6867            ),
6868            Some(query_id),
6869        );
6870    }
6871    let query = match register_page_operation(
6872        &state,
6873        SqlQueryOptions {
6874            query_id: Some(query_id),
6875            timeout: Some(Duration::from_millis(timeout_ms)),
6876            owner: Some(owner.clone()),
6877            session_id,
6878            parent_control: None,
6879        },
6880    ) {
6881        Ok(query) => query,
6882        Err(error) => return tracked_query_error_response(&state, &error, Some(query_id)),
6883    };
6884    query.set_sql_metadata("CONTINUE SQL PAGE");
6885    let _permit = match tokio::select! {
6886        permit = Arc::clone(&state.sql_page_semaphore).acquire_owned() => permit.map_err(|_| {
6887            mongreldb_query::MongrelQueryError::InvalidQueryState(
6888                "SQL page admission semaphore closed".into(),
6889            )
6890        }),
6891        _ = query.control().cancelled() => Err(cancellation_checkpoint_error(&query)),
6892    } {
6893        Ok(permit) => permit,
6894        Err(error) => {
6895            query.fail();
6896            return tracked_query_error_response(&state, &error, Some(query_id));
6897        }
6898    };
6899    if let Err(error) = query.transition(SqlQueryPhase::Queued, SqlQueryPhase::Executing) {
6900        query.fail();
6901        return tracked_query_error_response(&state, &error, Some(query_id));
6902    }
6903    let fail = |status, code, message| {
6904        query.record_terminal_error(code, mongreldb_query::QueryTerminalErrorCategory::Execution);
6905        query.fail();
6906        with_query_id(
6907            sql_cursor_error_response(status, code, message, query_id),
6908            query_id,
6909        )
6910    };
6911    if request.cursor.is_empty() || request.cursor.len() > 2_048 {
6912        return fail(
6913            StatusCode::BAD_REQUEST,
6914            "INVALID_SQL_CURSOR",
6915            "invalid SQL continuation cursor",
6916        );
6917    }
6918    if let Err(error) = query.checkpoint() {
6919        query.fail();
6920        return tracked_query_error_response(&state, &error, Some(query_id));
6921    }
6922    let cursor_mac_key = match state.cursor_mac_key.get() {
6923        Ok(key) => key,
6924        Err(_) => {
6925            return fail(
6926                StatusCode::INTERNAL_SERVER_ERROR,
6927                "ENTROPY_UNAVAILABLE",
6928                "OS CSPRNG unavailable",
6929            );
6930        }
6931    };
6932    match state.sql_pages.continue_page_with_control(
6933        &request.cursor,
6934        &owner,
6935        &cursor_mac_key,
6936        sql_pages::SqlPageBinding {
6937            security_version: state.db().security_version(),
6938            catalog_epoch: state.db().catalog_snapshot().db_epoch,
6939        },
6940        &query,
6941    ) {
6942        Ok(page) => {
6943            let page_byte_count = page.byte_count;
6944            if let Err(error) = query.begin_serialization() {
6945                query.fail();
6946                return tracked_query_error_response(&state, &error, Some(query_id));
6947            }
6948            let serialization_query = query.clone();
6949            match tokio::task::spawn_blocking(move || {
6950                serialize_sql_page_controlled(page, &serialization_query)
6951            })
6952            .await
6953            {
6954                Ok(Ok(body)) => {
6955                    if let Err(error) = query.try_complete() {
6956                        return tracked_query_error_response(&state, &error, Some(query_id));
6957                    }
6958                    state.metrics.add_sql_output_bytes(page_byte_count);
6959                    with_query_id(sql_page_response(body), query_id)
6960                }
6961                Ok(Err(ControlledPageSerializationError::Query(error))) => {
6962                    query.fail();
6963                    tracked_query_error_response(&state, &error, Some(query_id))
6964                }
6965                Ok(Err(ControlledPageSerializationError::Encoding)) => {
6966                    if let Err(error) = query.checkpoint() {
6967                        query.fail();
6968                        return tracked_query_error_response(&state, &error, Some(query_id));
6969                    }
6970                    fail(
6971                        StatusCode::INTERNAL_SERVER_ERROR,
6972                        "SERIALIZATION_FAILED",
6973                        "failed to serialize SQL continuation page",
6974                    )
6975                }
6976                Err(_) => {
6977                    if let Err(error) = query.checkpoint() {
6978                        query.fail();
6979                        return tracked_query_error_response(&state, &error, Some(query_id));
6980                    }
6981                    fail(
6982                        StatusCode::INTERNAL_SERVER_ERROR,
6983                        "SERIALIZATION_WORKER_FAILED",
6984                        "SQL continuation serialization worker failed",
6985                    )
6986                }
6987            }
6988        }
6989        Err(sql_pages::CursorError::Cancelled) => {
6990            let error = cancellation_checkpoint_error(&query);
6991            query.fail();
6992            tracked_query_error_response(&state, &error, Some(query_id))
6993        }
6994        Err(sql_pages::CursorError::Invalid) => fail(
6995            StatusCode::BAD_REQUEST,
6996            "INVALID_SQL_CURSOR",
6997            "invalid SQL continuation cursor",
6998        ),
6999        Err(sql_pages::CursorError::Expired) => fail(
7000            StatusCode::GONE,
7001            "SQL_CURSOR_EXPIRED",
7002            "SQL continuation cursor expired",
7003        ),
7004        Err(sql_pages::CursorError::NotFound) => fail(
7005            StatusCode::NOT_FOUND,
7006            "SQL_CURSOR_NOT_FOUND",
7007            "SQL continuation result is unavailable",
7008        ),
7009        Err(sql_pages::CursorError::PageLimit) => fail(
7010            StatusCode::PAYLOAD_TOO_LARGE,
7011            "RESULT_LIMIT_EXCEEDED",
7012            "one projected row exceeds the page byte or token limit",
7013        ),
7014    }
7015}
7016
7017fn sql_cursor_error_response(
7018    status: StatusCode,
7019    code: &'static str,
7020    message: &'static str,
7021    query_id: QueryId,
7022) -> Response {
7023    // Stable taxonomy category for this cursor code (spec 9.7); mappings
7024    // follow the `MongrelError::category()` precedents.
7025    let category = {
7026        use mongreldb_types::errors::ErrorCategory;
7027        match code {
7028            // Cursor staleness/expiry is stale server-side state (the core
7029            // `CursorStale` / `CursorExpired` precedent).
7030            "SQL_CURSOR_NOT_FOUND" | "SQL_CURSOR_EXPIRED" => ErrorCategory::StaleMetadata,
7031            "RESULT_LIMIT_EXCEEDED" | "ENTROPY_UNAVAILABLE" => ErrorCategory::ResourceExhausted,
7032            _ => ErrorCategory::ClusterVersionMismatch,
7033        }
7034    };
7035    (
7036        status,
7037        Json(json!({
7038            "query_id": query_id.to_string(),
7039            "status": "failed_before_commit",
7040            "terminal_state": "failed_before_commit",
7041            "server_state": "failed",
7042            "committed": false,
7043            "committed_statements": 0,
7044            "last_commit_epoch": null,
7045            "last_commit_epoch_text": null,
7046            "first_commit_statement_index": null,
7047            "last_commit_statement_index": null,
7048            "completed_statements": 0,
7049            "statement_index": 0,
7050            "cancel_outcome": "already_finished",
7051            "cancellation_reason": "none",
7052            "retryable": false,
7053            "outcome": {
7054                "committed": false,
7055                "committed_statements": 0,
7056                "last_commit_epoch": null,
7057                "last_commit_epoch_text": null,
7058                "first_commit_statement_index": null,
7059                "last_commit_statement_index": null,
7060                "completed_statements": 0,
7061                "statement_index": 0,
7062                "serialization": "not_started",
7063            },
7064            "error": {
7065                "code": code,
7066                "message": message,
7067                "category": category.to_string(),
7068                "category_code": category.code(),
7069                "query_id": query_id.to_string(),
7070                "committed": false,
7071                "retryable": false,
7072            }
7073        })),
7074    )
7075        .into_response()
7076}
7077
7078async fn sql(
7079    State(state): State<Arc<AppState>>,
7080    OptionalPrincipal(principal): OptionalPrincipal,
7081    headers: axum::http::HeaderMap,
7082    Json(req): Json<SqlRequest>,
7083) -> Response {
7084    if !state.accepting_sql.load(Ordering::Acquire) {
7085        return (StatusCode::SERVICE_UNAVAILABLE, "server is shutting down").into_response();
7086    }
7087    if !request_identity_is_current(&state, &principal) {
7088        return StatusCode::UNAUTHORIZED.into_response();
7089    }
7090    // P0.2: handle §15 admin SQL before opening any standalone session, so
7091    // cluster mode never touches AppState.db for SHOW CLUSTER / TRANSFER LEADER.
7092    #[cfg(feature = "cluster")]
7093    if let Some(response) = cluster_admin::try_admin_sql(&state, &principal, &req.sql).await {
7094        return response;
7095    }
7096    // P0.2: ordinary public SQL in cluster mode routes through the tablet
7097    // data plane (Raft propose / tablet_rows). Unsupported statements and a
7098    // missing runtime still fail closed — never AppState.db.
7099    #[cfg(feature = "cluster")]
7100    if state.is_cluster_mode() {
7101        if let Some(response) = cluster_data_plane::try_execute_sql(&state, &req.sql).await {
7102            return response;
7103        }
7104        if let Some(response) = refuse_cluster_standalone_data_plane(&state) {
7105            return response;
7106        }
7107    }
7108    // Session routing: an `X-Session-ID` header routes the request to a pooled
7109    // long-lived session, enabling cross-request `BEGIN`/`INSERT`/`COMMIT`
7110    // transactions. Without the header, a fresh ephemeral session is used
7111    // (the historical behavior).
7112    let session_id = match query_session_header(&headers, None) {
7113        Ok(session_id) => session_id,
7114        Err(response) => return *response,
7115    };
7116
7117    let owner = request_owner(&state, &principal);
7118    if let Some(sid) = session_id {
7119        let Some(entry) = state.sessions.get(&sid, &owner) else {
7120            return (
7121                StatusCode::NOT_FOUND,
7122                "session not found or not owned by caller",
7123            )
7124                .into_response();
7125        };
7126        let (options, query_id) = match resolve_query_options(
7127            &state,
7128            &headers,
7129            req.query_id,
7130            req.timeout_ms,
7131            owner.clone(),
7132            Some(sid.clone()),
7133        ) {
7134            Ok(options) => options,
7135            Err(response) => return *response,
7136        };
7137        let query = match register_controlled_query(&state, &entry.session(), options) {
7138            Ok(query) => query,
7139            Err(error) => return tracked_query_error_response(&state, &error, Some(query_id)),
7140        };
7141        let registration = RegisteredQueryGuard::new(query);
7142        if mongreldb_query::contains_boolean_ai_predicate(&req.sql) {
7143            registration.fail();
7144            return with_query_id(remote_boolean_ai_error(), query_id);
7145        }
7146        let output_limits = match resolve_sql_output_limits(&state, &req, query_id) {
7147            Ok(limits) => limits,
7148            Err(response) => {
7149                registration.fail();
7150                return *response;
7151            }
7152        };
7153        let (registration, pagination) =
7154            match resolve_sql_pagination(&headers, &req, output_limits, registration, query_id) {
7155                Ok(resolved) => resolved,
7156                Err(response) => return *response,
7157            };
7158        let sql_permit =
7159            match acquire_sql_permit(&state, &entry.session(), registration.query()).await {
7160                Ok(permit) => permit,
7161                Err(error) => {
7162                    return tracked_query_error_response(&state, &error, Some(query_id));
7163                }
7164            };
7165        // Registration and global admission happen before this session lock.
7166        let _guard = tokio::select! {
7167            guard = entry.lock.lock() => guard,
7168            _ = registration.query().control().cancelled() => {
7169                return tracked_query_error_response(
7170                    &state,
7171                    &cancellation_checkpoint_error(registration.query()),
7172                    Some(query_id),
7173                );
7174            }
7175        };
7176        // Re-check closed: the session may have been closed/evicted between
7177        // get() and acquiring the lock.
7178        if entry.is_closed() {
7179            return (StatusCode::NOT_FOUND, "session no longer available").into_response();
7180        }
7181        if req.idempotency_key.is_some() || headers.contains_key("idempotency-key") {
7182            entry
7183                .session()
7184                .fire_test_hook(mongreldb_query::SqlTestHookPoint::BeforeServerIdempotencyCheck);
7185        }
7186        let (registration, idempotency) = match begin_sql_idempotency(
7187            &state,
7188            SqlIdempotencyContext {
7189                headers: &headers,
7190                request: &req,
7191                output_limits,
7192                owner: &owner,
7193                session_id: Some(&sid),
7194                session_in_transaction: entry.session().staged_sql_operation_count().is_some(),
7195                query_id,
7196            },
7197            registration,
7198        )
7199        .await
7200        {
7201            Ok(resolved) => resolved,
7202            Err(response) => return response,
7203        };
7204        entry.touch();
7205        let query = registration.into_query();
7206        let (response, idempotent_commit_ts) = execute_sql(
7207            &state,
7208            &principal,
7209            &entry.session(),
7210            ResolvedSqlRequest {
7211                request: req,
7212                output_limits,
7213                idempotency,
7214                pagination,
7215            },
7216            query,
7217            query_id,
7218            sql_permit,
7219        )
7220        .await;
7221        // Keep the canonical S1D-004 record (transaction state,
7222        // read-your-writes token, last activity) in step with the session.
7223        // The token carries the write's literal commit-timestamp lineage when
7224        // one is known: the core commit log's receipt recorded for an
7225        // idempotent commit (S1B-005), else the exact commit timestamp the
7226        // query layer sourced from core into the durable outcome, else the
7227        // per-open epoch→commit-ts ledger (`Database::commit_ts_for_epoch`).
7228        // Only when none of those resolve does it fall back to the node's HLC
7229        // timestamp captured after the commit became visible, which the
7230        // single HLC authority orders after the write's commit timestamp
7231        // (spec §8.2).
7232        let durable = state
7233            .query_registry
7234            .status(query_id)
7235            .map(|status| status.durable_outcome);
7236        entry.sync_record_after_request(match durable {
7237            Some(outcome) if outcome.committed => Some(
7238                idempotent_commit_ts
7239                    .or(outcome.commit_ts)
7240                    .or_else(|| {
7241                        outcome.last_commit_epoch.and_then(|epoch| {
7242                            state.db().commit_ts_for_epoch(mongreldb_core::Epoch(epoch))
7243                        })
7244                    })
7245                    .unwrap_or_else(|| ryw_commit_timestamp(state.db())),
7246            ),
7247            _ => None,
7248        });
7249        response
7250    } else {
7251        let session = match MongrelSession::open_with_external_modules_as(
7252            Arc::clone(state.db()),
7253            state.external_modules.iter().cloned(),
7254            request_principal(&state, &principal),
7255        ) {
7256            Ok(session) => session.with_query_registry(Arc::clone(&state.query_registry)),
7257            Err(e) => return (status_for_query_error(&e), e.to_string()).into_response(),
7258        };
7259        let (options, query_id) = match resolve_query_options(
7260            &state,
7261            &headers,
7262            req.query_id,
7263            req.timeout_ms,
7264            owner.clone(),
7265            None,
7266        ) {
7267            Ok(options) => options,
7268            Err(response) => return *response,
7269        };
7270        let query = match register_controlled_query(&state, &session, options) {
7271            Ok(query) => query,
7272            Err(error) => return tracked_query_error_response(&state, &error, Some(query_id)),
7273        };
7274        let registration = RegisteredQueryGuard::new(query);
7275        if mongreldb_query::contains_boolean_ai_predicate(&req.sql) {
7276            registration.fail();
7277            return with_query_id(remote_boolean_ai_error(), query_id);
7278        }
7279        let output_limits = match resolve_sql_output_limits(&state, &req, query_id) {
7280            Ok(limits) => limits,
7281            Err(response) => {
7282                registration.fail();
7283                return *response;
7284            }
7285        };
7286        let (registration, pagination) =
7287            match resolve_sql_pagination(&headers, &req, output_limits, registration, query_id) {
7288                Ok(resolved) => resolved,
7289                Err(response) => return *response,
7290            };
7291        let (registration, idempotency) = match begin_sql_idempotency(
7292            &state,
7293            SqlIdempotencyContext {
7294                headers: &headers,
7295                request: &req,
7296                output_limits,
7297                owner: &owner,
7298                session_id: None,
7299                session_in_transaction: false,
7300                query_id,
7301            },
7302            registration,
7303        )
7304        .await
7305        {
7306            Ok(resolved) => resolved,
7307            Err(response) => return response,
7308        };
7309        let sql_permit = match acquire_sql_permit(&state, &session, registration.query()).await {
7310            Ok(permit) => permit,
7311            Err(error) => {
7312                if let Some(idempotency) = idempotency {
7313                    idempotency.abort();
7314                }
7315                return tracked_query_error_response(&state, &error, Some(query_id));
7316            }
7317        };
7318        let query = registration.into_query();
7319        execute_sql(
7320            &state,
7321            &principal,
7322            &session,
7323            ResolvedSqlRequest {
7324                request: req,
7325                output_limits,
7326                idempotency,
7327                pagination,
7328            },
7329            query,
7330            query_id,
7331            sql_permit,
7332        )
7333        .await
7334        .0
7335    }
7336}
7337
7338/// The read-your-writes fallback token for a just-committed request whose
7339/// literal commit timestamp did not resolve (no idempotency receipt, no
7340/// query-layer `commit_ts`, no per-open ledger entry): the node's HLC
7341/// timestamp captured at a fresh `begin` after the commit became visible. The
7342/// single HLC authority (spec §8.2) orders it after every commit timestamp it
7343/// has already issued — including the commit this request observed — so a
7344/// later read presenting the token sees the session's own write. Falls back
7345/// to wall-clock micros when clock allocation is rejected (skew gate).
7346fn ryw_commit_timestamp(db: &mongreldb_core::Database) -> mongreldb_types::hlc::HlcTimestamp {
7347    if let Some(ts) = db.begin().read_ts() {
7348        return ts;
7349    }
7350    mongreldb_types::hlc::HlcTimestamp {
7351        physical_micros: sessions::now_unix_micros(),
7352        logical: 0,
7353        node_tiebreaker: 0,
7354    }
7355}
7356
7357/// Run one SQL request against a given session: bump counters, audit DDL
7358/// (after execution, with redaction), log slow queries on both success and
7359/// failure, and dispatch the response format. Shared by the pooled-session and
7360/// fresh-session paths.
7361///
7362/// The second tuple element is the HLC commit timestamp of the core receipt
7363/// recorded for an idempotent committed write (S1B-005), when one exists —
7364/// the pooled-session path folds it into the session's read-your-writes token.
7365async fn execute_sql(
7366    state: &AppState,
7367    principal: &Option<mongreldb_core::Principal>,
7368    session: &MongrelSession,
7369    request: ResolvedSqlRequest,
7370    query: RegisteredSqlQuery,
7371    query_id: QueryId,
7372    sql_permit: admission::SqlAdmissionGuard,
7373) -> (Response, Option<mongreldb_types::hlc::HlcTimestamp>) {
7374    let ResolvedSqlRequest {
7375        request: req,
7376        output_limits,
7377        idempotency,
7378        pagination,
7379    } = request;
7380    // §15 admin SQL surface: intercept before the ordinary SQL path so
7381    // SHOW CLUSTER / ALTER NODE DRAIN / … never open tablet files from the
7382    // query gateway (spec §12.10 / §15).
7383    // Note: cluster_admin::try_admin_sql and cluster_data_plane checks are
7384    // already performed by the `sql` handler before calling `execute_sql`.
7385    // They are not repeated here to avoid per-request double-parsing overhead.
7386    // If `execute_sql` is ever called from a path that did NOT pre-filter,
7387    // re-add the guard here.
7388    //
7389    // Keep a direct handle until the durable receipt is persisted. Finished
7390    // tombstone eviction must not make a committed idempotent write ambiguous.
7391    let idempotency_query = idempotency.as_ref().map(|_| query.clone());
7392    state.metrics.inc_sql_queries();
7393    let audited = audit::is_audited_sql(&req.sql);
7394    let actor = request_owner(state, principal);
7395    let page_binding = sql_pages::SqlPageBinding {
7396        security_version: state.db().security_version(),
7397        catalog_epoch: state.db().catalog_snapshot().db_epoch,
7398    };
7399    let start = std::time::Instant::now();
7400    // NOTE: deliberately NOT using `run_sql_traced` here. Its thread-local
7401    // push/pop spans an `.await`, and on a multi-threaded tokio runtime the
7402    // task can resume on a different thread, corrupting the trace stack and
7403    // leaking scopes. Wall-clock timing is sufficient for slow-query detection
7404    // and works across awaits.
7405    let result = if let Some(pagination) = pagination {
7406        match session
7407            .run_with_query_for_serialization_with_limits(
7408                &req.sql,
7409                query,
7410                mongreldb_query::SqlCollectionLimits::new(output_limits.0, output_limits.1),
7411            )
7412            .await
7413        {
7414            Ok(output) => Ok(dispatch_paginated_sql(
7415                state,
7416                output,
7417                query_id,
7418                &actor,
7419                pagination,
7420                output_limits,
7421                session.sql_test_hook(),
7422                page_binding,
7423            )
7424            .await),
7425            Err(error) => Err(error),
7426        }
7427    } else if req.format.as_deref() == Some("arrow-stream") {
7428        match session
7429            .run_stream_with_query_for_serialization(&req.sql, query)
7430            .await
7431        {
7432            Ok((stream, completion)) => Ok(sql_arrow_stream_response_controlled(
7433                stream,
7434                completion,
7435                sql_permit,
7436                output_limits,
7437                state,
7438                query_id,
7439                session.sql_test_hook(),
7440            )),
7441            Err(error) => Err(error),
7442        }
7443    } else {
7444        match session
7445            .run_with_query_for_serialization_with_limits(
7446                &req.sql,
7447                query,
7448                mongreldb_query::SqlCollectionLimits::new(output_limits.0, output_limits.1),
7449            )
7450            .await
7451        {
7452            Ok(output) => Ok(dispatch_buffered_sql_format(
7453                state,
7454                req.format.as_deref(),
7455                output,
7456                query_id,
7457                session.sql_test_hook(),
7458                output_limits,
7459            )
7460            .await),
7461            Err(error) => Err(error),
7462        }
7463    };
7464    let elapsed = start.elapsed();
7465    // Slow-query logging covers BOTH success and failure (the slowest errors
7466    // matter most for diagnosis), checked before branching on the outcome.
7467    if elapsed >= state.reloadable.slow_query_threshold.get() {
7468        state.metrics.inc_slow_queries();
7469        eprintln!(
7470            "[slow-query] {}\u{00b5}s query_id={} operation={}",
7471            elapsed.as_micros(),
7472            query_id,
7473            safe_sql_operation(&req.sql)
7474        );
7475    }
7476    // Audit DDL/privilege AFTER execution so the outcome (ok/fail) is captured.
7477    // `redacted_ddl_detail` never logs credential literals.
7478    if audited {
7479        let (action, detail) = audit::redacted_ddl_detail(&req.sql, result.is_ok());
7480        state.audit.record(actor, action, detail);
7481    }
7482    let response = match result {
7483        Ok(response) => with_query_id(response, query_id),
7484        Err(e) => {
7485            state.metrics.inc_sql_errors();
7486            tracked_query_error_response(state, &e, Some(query_id))
7487        }
7488    };
7489    let Some(idempotency) = idempotency else {
7490        return (response, None);
7491    };
7492    let status = idempotency_query.map(|query| query.status());
7493    if let Some(receipt) = status.as_ref().and_then(sql_terminal_idempotency_receipt) {
7494        let mut receipt = receipt;
7495        // S1B-005: a write that durably committed also records its key in the
7496        // core's durable TXN_IDEMPOTENCY ledger, so the key → original-receipt
7497        // binding survives restart under the commit log's authority (identical
7498        // replay → original receipt; different fingerprint → Conflict). The
7499        // ledger's receipt rides the durable HTTP receipt additively; the HTTP
7500        // store remains the wire authority if the bookkeeping fails.
7501        if receipt.outcome.committed {
7502            let db = Arc::clone(state.db());
7503            let owner = idempotency.owner().to_owned();
7504            let key = idempotency.key().to_owned();
7505            let binding = idempotency.binding().clone();
7506            let ttl = idempotency.ttl();
7507            let commit_receipt = tokio::task::spawn_blocking(move || {
7508                sql_idempotency::record_core_idempotency_commit(&db, &owner, &key, &binding, ttl)
7509            })
7510            .await
7511            .unwrap_or_else(|error| {
7512                eprintln!("[idempotency] core ledger record task failed: {error}");
7513                None
7514            });
7515            receipt.attach_commit_receipt(commit_receipt);
7516        }
7517        let commit_ts = receipt
7518            .commit_receipt
7519            .as_ref()
7520            .map(sql_idempotency::SqlCommitReceipt::commit_ts);
7521        let (expires_at_ms, persisted) = idempotency.commit(receipt.clone());
7522        return (
7523            sql_idempotency_receipt_response(query_id, &receipt, false, expires_at_ms, persisted),
7524            commit_ts,
7525        );
7526    }
7527    if status.as_ref().is_some_and(can_abort_idempotency_intent) {
7528        idempotency.abort();
7529    }
7530    (response, None)
7531}
7532
7533fn can_abort_idempotency_intent(status: &mongreldb_query::QueryStatus) -> bool {
7534    !status.outcome_unknown
7535        && !status.durable_outcome.committed
7536        && matches!(
7537            status.terminal_state(),
7538            Some(
7539                mongreldb_query::QueryTerminalState::FailedBeforeCommit
7540                    | mongreldb_query::QueryTerminalState::CancelledBeforeCommit
7541                    | mongreldb_query::QueryTerminalState::DeadlineBeforeCommit
7542            )
7543        )
7544}
7545
7546fn safe_sql_operation(sql: &str) -> String {
7547    sql.split_whitespace()
7548        .next()
7549        .unwrap_or("UNKNOWN")
7550        .chars()
7551        .filter(|character| character.is_ascii_alphabetic())
7552        .take(16)
7553        .collect::<String>()
7554        .to_ascii_uppercase()
7555}
7556
7557fn remote_boolean_ai_error() -> Response {
7558    (
7559        StatusCode::BAD_REQUEST,
7560        "Boolean ANN/Sparse SQL is disabled remotely; use scored SQL functions",
7561    )
7562        .into_response()
7563}
7564
7565#[derive(Debug)]
7566struct SerializedOutput {
7567    bytes: Vec<u8>,
7568    arrow: bool,
7569}
7570
7571#[derive(Debug)]
7572enum BufferedSerializationError {
7573    Query(mongreldb_query::MongrelQueryError),
7574    Limit(String),
7575    Encoding(String),
7576}
7577
7578struct LimitedOutput {
7579    bytes: Vec<u8>,
7580    max_bytes: usize,
7581    exceeded: bool,
7582}
7583
7584impl LimitedOutput {
7585    fn new(max_bytes: usize) -> Self {
7586        Self {
7587            bytes: Vec::new(),
7588            max_bytes,
7589            exceeded: false,
7590        }
7591    }
7592}
7593
7594impl std::io::Write for LimitedOutput {
7595    fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
7596        if self.bytes.len().saturating_add(bytes.len()) > self.max_bytes {
7597            self.exceeded = true;
7598            return Err(std::io::Error::other("SQL output byte limit exceeded"));
7599        }
7600        self.bytes.extend_from_slice(bytes);
7601        Ok(bytes.len())
7602    }
7603
7604    fn flush(&mut self) -> std::io::Result<()> {
7605        Ok(())
7606    }
7607}
7608
7609fn serialize_buffered_output(
7610    format: &str,
7611    batches: &[arrow::record_batch::RecordBatch],
7612    query: &RegisteredSqlQuery,
7613    max_rows: usize,
7614    max_bytes: usize,
7615    test_hook: Option<&mongreldb_query::SqlTestHook>,
7616) -> std::result::Result<SerializedOutput, BufferedSerializationError> {
7617    const ROW_CHECKPOINT_INTERVAL: usize = 256;
7618    let mut rows = 0usize;
7619    let mut writer_output = LimitedOutput::new(max_bytes);
7620
7621    if format == "arrow" {
7622        if batches.is_empty() {
7623            if let Some(hook) = test_hook {
7624                hook(mongreldb_query::SqlTestHookPoint::AfterSerialization);
7625            }
7626            return Ok(SerializedOutput {
7627                bytes: Vec::new(),
7628                arrow: true,
7629            });
7630        }
7631        let schema = batches[0].schema();
7632        let encoding_result = (|| {
7633            let mut writer =
7634                arrow::ipc::writer::FileWriter::try_new(&mut writer_output, schema.as_ref())
7635                    .map_err(|error| error.to_string())?;
7636            for batch in batches {
7637                for offset in (0..batch.num_rows()).step_by(ROW_CHECKPOINT_INTERVAL) {
7638                    if let Some(hook) = test_hook {
7639                        hook(mongreldb_query::SqlTestHookPoint::BeforeSerializationBatch);
7640                    }
7641                    query.checkpoint().map_err(|error| error.to_string())?;
7642                    let length = ROW_CHECKPOINT_INTERVAL.min(batch.num_rows() - offset);
7643                    rows = rows.saturating_add(length);
7644                    if rows > max_rows {
7645                        return Err("SQL output row limit exceeded".into());
7646                    }
7647                    writer
7648                        .write(&batch.slice(offset, length))
7649                        .map_err(|error| error.to_string())?;
7650                }
7651            }
7652            writer.finish().map_err(|error| error.to_string())
7653        })();
7654        if let Err(error) = encoding_result {
7655            if let Err(query_error) = query.checkpoint() {
7656                return Err(BufferedSerializationError::Query(query_error));
7657            }
7658            if writer_output.exceeded || rows > max_rows {
7659                return Err(BufferedSerializationError::Limit(error));
7660            }
7661            return Err(BufferedSerializationError::Encoding(error));
7662        }
7663        if let Some(hook) = test_hook {
7664            hook(mongreldb_query::SqlTestHookPoint::AfterSerialization);
7665        }
7666        return Ok(SerializedOutput {
7667            bytes: writer_output.bytes,
7668            arrow: true,
7669        });
7670    }
7671
7672    let encoding_result = (|| {
7673        let mut writer = arrow::json::writer::ArrayWriter::new(&mut writer_output);
7674        for batch in batches {
7675            for offset in (0..batch.num_rows()).step_by(ROW_CHECKPOINT_INTERVAL) {
7676                if let Some(hook) = test_hook {
7677                    hook(mongreldb_query::SqlTestHookPoint::BeforeSerializationBatch);
7678                }
7679                query.checkpoint().map_err(|error| error.to_string())?;
7680                let length = ROW_CHECKPOINT_INTERVAL.min(batch.num_rows() - offset);
7681                rows = rows.saturating_add(length);
7682                if rows > max_rows {
7683                    return Err("SQL output row limit exceeded".into());
7684                }
7685                let slice = batch.slice(offset, length);
7686                writer
7687                    .write_batches(&[&slice])
7688                    .map_err(|error| error.to_string())?;
7689            }
7690        }
7691        writer.finish().map_err(|error| error.to_string())
7692    })();
7693    if let Err(error) = encoding_result {
7694        if let Err(query_error) = query.checkpoint() {
7695            return Err(BufferedSerializationError::Query(query_error));
7696        }
7697        if writer_output.exceeded || rows > max_rows {
7698            return Err(BufferedSerializationError::Limit(error));
7699        }
7700        return Err(BufferedSerializationError::Encoding(error));
7701    }
7702    if let Some(hook) = test_hook {
7703        hook(mongreldb_query::SqlTestHookPoint::AfterSerialization);
7704    }
7705    Ok(SerializedOutput {
7706        bytes: writer_output.bytes,
7707        arrow: false,
7708    })
7709}
7710
7711/// Serialize a DataFusion record-batch stream as Arrow streaming IPC. The body
7712/// holds only the active query batch and one serialized IPC message.
7713#[cfg(test)]
7714fn sql_arrow_stream_response(batches: mongreldb_query::MongrelRecordBatchStream) -> Response {
7715    use futures::{stream, StreamExt};
7716
7717    const STREAM_CT: &str = "application/vnd.apache.arrow.stream";
7718
7719    let schema = batches.schema();
7720    let mut writer = match arrow::ipc::writer::StreamWriter::try_new(Vec::new(), schema.as_ref()) {
7721        Ok(w) => w,
7722        Err(e) => {
7723            return (
7724                StatusCode::INTERNAL_SERVER_ERROR,
7725                format!("arrow stream init error: {e}"),
7726            )
7727                .into_response()
7728        }
7729    };
7730    // `try_new` synchronously writes the schema message; drain it so it becomes
7731    // the first chunk yielded to the client.
7732    let schema_chunk: Vec<u8> = std::mem::take(writer.get_mut());
7733    let batch_stream = stream::unfold(
7734        (batches, Some(writer)),
7735        |(mut batches, writer)| async move {
7736            let mut writer = writer?;
7737            match batches.next().await {
7738                Some(Ok(batch)) => match writer.write(&batch) {
7739                    Ok(()) => {
7740                        let chunk = std::mem::take(writer.get_mut());
7741                        Some((Ok(chunk), (batches, Some(writer))))
7742                    }
7743                    Err(error) => Some((Err(std::io::Error::other(error)), (batches, None))),
7744                },
7745                Some(Err(error)) => Some((Err(std::io::Error::other(error)), (batches, None))),
7746                None => match writer.finish() {
7747                    Ok(()) => {
7748                        let chunk = std::mem::take(writer.get_mut());
7749                        Some((Ok(chunk), (batches, None)))
7750                    }
7751                    Err(error) => Some((Err(std::io::Error::other(error)), (batches, None))),
7752                },
7753            }
7754        },
7755    );
7756
7757    // Schema first, then batches + EOS. Each item is already `Result<Vec<u8>,
7758    // io::Error>` so a mid-stream encode failure surfaces as a body error.
7759    let schema_item: Result<Vec<u8>, std::io::Error> = Ok(schema_chunk);
7760    let full = stream::iter([schema_item]).chain(batch_stream);
7761    let body = axum::body::Body::from_stream(full);
7762    ([(header::CONTENT_TYPE, STREAM_CT)], body).into_response()
7763}
7764
7765fn sql_arrow_stream_response_controlled(
7766    batches: mongreldb_query::MongrelRecordBatchStream,
7767    completion: SqlStreamCompletion,
7768    sql_permit: admission::SqlAdmissionGuard,
7769    limits: (usize, usize),
7770    state: &AppState,
7771    query_id: QueryId,
7772    test_hook: Option<mongreldb_query::SqlTestHook>,
7773) -> Response {
7774    use futures::{stream, StreamExt};
7775
7776    const STREAM_CT: &str = "application/vnd.apache.arrow.stream";
7777    let (max_rows, max_bytes) = limits;
7778    let schema = batches.schema();
7779    let mut writer = match arrow::ipc::writer::StreamWriter::try_new(Vec::new(), schema.as_ref()) {
7780        Ok(writer) => writer,
7781        Err(error) => {
7782            completion.fail_serialization();
7783            drop(batches);
7784            return terminal_server_error_response(
7785                state,
7786                query_id,
7787                StatusCode::INTERNAL_SERVER_ERROR,
7788                "SERIALIZATION_FAILED",
7789                format!("arrow stream init error: {error}"),
7790            );
7791        }
7792    };
7793    let schema_chunk = std::mem::take(writer.get_mut());
7794    if schema_chunk.len() > max_bytes {
7795        completion.fail_result_limit();
7796        drop(batches);
7797        return terminal_server_error_response(
7798            state,
7799            query_id,
7800            StatusCode::PAYLOAD_TOO_LARGE,
7801            "RESULT_LIMIT_EXCEEDED",
7802            "SQL output byte limit exceeded",
7803        );
7804    }
7805    let metrics = Arc::clone(&state.metrics);
7806    metrics.add_sql_output_bytes(schema_chunk.len());
7807    let batch_stream = stream::unfold(
7808        (
7809            batches,
7810            Some(writer),
7811            completion,
7812            Some(sql_permit),
7813            0usize,
7814            schema_chunk.len(),
7815            metrics,
7816        ),
7817        move |(mut batches, writer, completion, permit, rows, bytes, metrics)| {
7818            let test_hook = test_hook.clone();
7819            async move {
7820                let mut writer = writer?;
7821                match batches.next().await {
7822                    Some(Ok(batch)) => {
7823                        let next_rows = rows.saturating_add(batch.num_rows());
7824                        if next_rows > max_rows {
7825                            completion.fail_result_limit();
7826                            return Some((
7827                                Err(std::io::Error::other("SQL output row limit exceeded")),
7828                                (batches, None, completion, permit, next_rows, bytes, metrics),
7829                            ));
7830                        }
7831                        match writer.write(&batch) {
7832                            Ok(()) => {
7833                                let chunk = std::mem::take(writer.get_mut());
7834                                let next_bytes = bytes.saturating_add(chunk.len());
7835                                if next_bytes > max_bytes {
7836                                    completion.fail_result_limit();
7837                                    return Some((
7838                                        Err(std::io::Error::other(
7839                                            "SQL output byte limit exceeded",
7840                                        )),
7841                                        (
7842                                            batches, None, completion, permit, next_rows,
7843                                            next_bytes, metrics,
7844                                        ),
7845                                    ));
7846                                }
7847                                metrics.add_sql_output_bytes(chunk.len());
7848                                Some((
7849                                    Ok(chunk),
7850                                    (
7851                                        batches,
7852                                        Some(writer),
7853                                        completion,
7854                                        permit,
7855                                        next_rows,
7856                                        next_bytes,
7857                                        metrics,
7858                                    ),
7859                                ))
7860                            }
7861                            Err(error) => {
7862                                completion.fail_serialization();
7863                                Some((
7864                                    Err(std::io::Error::other(error)),
7865                                    (batches, None, completion, permit, rows, bytes, metrics),
7866                                ))
7867                            }
7868                        }
7869                    }
7870                    Some(Err(error)) => Some((
7871                        Err(std::io::Error::other(error)),
7872                        (batches, None, completion, permit, rows, bytes, metrics),
7873                    )),
7874                    None => match writer.finish() {
7875                        Ok(()) => {
7876                            let chunk = std::mem::take(writer.get_mut());
7877                            let next_bytes = bytes.saturating_add(chunk.len());
7878                            if next_bytes > max_bytes {
7879                                completion.fail_result_limit();
7880                                return Some((
7881                                    Err(std::io::Error::other("SQL output byte limit exceeded")),
7882                                    (batches, None, completion, permit, rows, next_bytes, metrics),
7883                                ));
7884                            }
7885                            if let Some(hook) = test_hook {
7886                                hook(mongreldb_query::SqlTestHookPoint::AfterSerialization);
7887                            }
7888                            match completion.try_complete() {
7889                                Ok(()) => {
7890                                    metrics.add_sql_output_bytes(chunk.len());
7891                                    Some((
7892                                        Ok(chunk),
7893                                        (
7894                                            batches, None, completion, permit, rows, next_bytes,
7895                                            metrics,
7896                                        ),
7897                                    ))
7898                                }
7899                                Err(error) => {
7900                                    metrics.inc_sql_errors();
7901                                    Some((
7902                                        Err(std::io::Error::other(error.to_string())),
7903                                        (batches, None, completion, permit, rows, bytes, metrics),
7904                                    ))
7905                                }
7906                            }
7907                        }
7908                        Err(error) => {
7909                            completion.fail_serialization();
7910                            Some((
7911                                Err(std::io::Error::other(error)),
7912                                (batches, None, completion, permit, rows, bytes, metrics),
7913                            ))
7914                        }
7915                    },
7916                }
7917            }
7918        },
7919    );
7920    let schema_item: Result<Vec<u8>, std::io::Error> = Ok(schema_chunk);
7921    let body = axum::body::Body::from_stream(stream::iter([schema_item]).chain(batch_stream));
7922    ([(header::CONTENT_TYPE, STREAM_CT)], body).into_response()
7923}
7924
7925#[derive(Deserialize)]
7926struct TxnOp {
7927    table: String,
7928    op: String,
7929    cells: Option<Vec<serde_json::Value>>,
7930    row_id: Option<u64>,
7931}
7932
7933#[derive(Deserialize)]
7934struct TxnRequest {
7935    ops: Vec<TxnOp>,
7936}
7937
7938async fn txn(
7939    State(state): State<Arc<AppState>>,
7940    OptionalPrincipal(principal): OptionalPrincipal,
7941    Json(req): Json<TxnRequest>,
7942) -> Response {
7943    if let Some(response) = require_writes_open(&state) {
7944        return response;
7945    }
7946    // Pre-validate every op against the live schemas before entering the
7947    // transaction, so malformed input returns 400 without consuming an epoch
7948    // or poisoning a txn.
7949    let mut parsed: Vec<(String, TxnAction)> = Vec::with_capacity(req.ops.len());
7950    for op in &req.ops {
7951        match op.op.as_str() {
7952            "put" => {
7953                let cells_json = match op.cells.as_ref() {
7954                    Some(c) if !c.is_empty() => c,
7955                    _ => {
7956                        return (StatusCode::BAD_REQUEST, "put op requires non-empty cells")
7957                            .into_response()
7958                    }
7959                };
7960                let handle = match state.db().table(&op.table) {
7961                    Ok(h) => h,
7962                    Err(e) => return (StatusCode::NOT_FOUND, e.to_string()).into_response(),
7963                };
7964                let schema = handle.lock().schema().clone();
7965                let cells = match parse_cells(cells_json, &schema) {
7966                    Ok(c) => c,
7967                    Err(msg) => return (StatusCode::BAD_REQUEST, msg).into_response(),
7968                };
7969                parsed.push((op.table.clone(), TxnAction::Put(cells)));
7970            }
7971            "delete" => {
7972                let rid = match op.row_id {
7973                    Some(r) => r,
7974                    None => {
7975                        return (StatusCode::BAD_REQUEST, "delete op requires row_id")
7976                            .into_response()
7977                    }
7978                };
7979                parsed.push((op.table.clone(), TxnAction::Delete(rid)));
7980            }
7981            other => {
7982                return (StatusCode::BAD_REQUEST, format!("unknown op: {other}")).into_response()
7983            }
7984        }
7985    }
7986
7987    state.metrics.inc_txns();
7988    let mut transaction = state.db().begin_as(request_principal(&state, &principal));
7989    let result = (|| {
7990        for (table, action) in &parsed {
7991            match action {
7992                TxnAction::Put(cells) => {
7993                    transaction.put(table, cells.clone())?;
7994                }
7995                TxnAction::Delete(rid) => {
7996                    transaction.delete(table, mongreldb_core::RowId(*rid))?;
7997                }
7998            }
7999        }
8000        transaction.commit()
8001    })();
8002    match result {
8003        Ok(epoch) => Json(json!({
8004            "status": "committed",
8005            "epoch": epoch.0,
8006            "epoch_text": epoch.0.to_string()
8007        }))
8008        .into_response(),
8009        Err(error) => crate::kit::durable_core_error_response(&error)
8010            .unwrap_or_else(|| (status_for_error(&error), error.to_string()).into_response()),
8011    }
8012}
8013
8014enum TxnAction {
8015    Put(Vec<(u16, Value)>),
8016    Delete(u64),
8017}
8018
8019#[cfg(test)]
8020mod auth_tests {
8021    use super::*;
8022    use mongreldb_core::Database;
8023    use tempfile::tempdir;
8024
8025    #[test]
8026    fn slow_query_operation_does_not_include_literals() {
8027        let sql = "CREATE USER alice PASSWORD 'never-log-this'";
8028        let operation = safe_sql_operation(sql);
8029        assert_eq!(operation, "CREATE");
8030        assert!(!operation.contains("never-log-this"));
8031    }
8032
8033    #[tokio::test]
8034    async fn auth_rejects_missing_token() {
8035        let dir = tempdir().unwrap();
8036        let db = Arc::new(Database::create(dir.path()).unwrap());
8037        let app = build_app_with_config(db, std::iter::empty(), Some("secret".into()), None);
8038        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
8039        let addr = listener.local_addr().unwrap();
8040        tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
8041        // Give the server a moment to start.
8042        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
8043
8044        let client = reqwest::Client::new();
8045        let resp = client
8046            .get(format!("http://{addr}/health"))
8047            .send()
8048            .await
8049            .unwrap();
8050        assert_eq!(resp.status(), 401);
8051    }
8052
8053    #[tokio::test]
8054    async fn auth_accepts_valid_token() {
8055        let dir = tempdir().unwrap();
8056        let db = Arc::new(Database::create(dir.path()).unwrap());
8057        let app = build_app_with_config(db, std::iter::empty(), Some("secret".into()), None);
8058        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
8059        let addr = listener.local_addr().unwrap();
8060        tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
8061        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
8062
8063        let client = reqwest::Client::new();
8064        let resp = client
8065            .get(format!("http://{addr}/health"))
8066            .header("Authorization", "Bearer secret")
8067            .send()
8068            .await
8069            .unwrap();
8070        assert_eq!(resp.status(), 200);
8071    }
8072
8073    #[tokio::test]
8074    async fn no_auth_when_token_unset() {
8075        let dir = tempdir().unwrap();
8076        let db = Arc::new(Database::create(dir.path()).unwrap());
8077        let app = build_app_with_config(db, std::iter::empty(), None, None);
8078        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
8079        let addr = listener.local_addr().unwrap();
8080        tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
8081        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
8082
8083        let client = reqwest::Client::new();
8084        let resp = client
8085            .get(format!("http://{addr}/health"))
8086            .send()
8087            .await
8088            .unwrap();
8089        assert_eq!(resp.status(), 200);
8090    }
8091
8092    #[tokio::test]
8093    async fn capabilities_advertise_sql_cancellation_v2() {
8094        let dir = tempdir().unwrap();
8095        let db = Arc::new(Database::create(dir.path()).unwrap());
8096        let app = build_app_with_config(db, std::iter::empty(), None, None);
8097        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
8098        let addr = listener.local_addr().unwrap();
8099        tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
8100
8101        let body: serde_json::Value = reqwest::Client::new()
8102            .get(format!("http://{addr}/capabilities"))
8103            .send()
8104            .await
8105            .unwrap()
8106            .json()
8107            .await
8108            .unwrap();
8109        assert_eq!(body["sql_cancellation"]["version"], 2);
8110        assert_eq!(body["sql_cancellation"]["client_query_ids"], true);
8111        assert_eq!(body["sql_cancellation"]["cancel_endpoint"], true);
8112        assert_eq!(body["sql_cancellation"]["query_status"], true);
8113        assert_eq!(body["sql_cancellation"]["pre_registration_cancel"], true);
8114        assert_eq!(body["sql_cancellation"]["stream_disconnect_cancels"], true);
8115        assert_eq!(body["sql_idempotency"]["version"], 1);
8116        assert_eq!(
8117            body["sql_idempotency"]["indeterminate_never_reexecutes"],
8118            true
8119        );
8120        assert_eq!(body["sql_pagination"]["version"], 1);
8121        assert_eq!(
8122            body["sql_pagination"]["continuation_endpoint"],
8123            "/sql/continue"
8124        );
8125    }
8126}
8127
8128#[cfg(test)]
8129mod query_response_tests {
8130    use super::*;
8131
8132    #[test]
8133    fn cancellation_reason_names_are_stable_snake_case() {
8134        assert_eq!(cancellation_reason_name(CancellationReason::None), "none");
8135        assert_eq!(
8136            cancellation_reason_name(CancellationReason::ClientRequest),
8137            "client_request"
8138        );
8139        assert_eq!(
8140            cancellation_reason_name(CancellationReason::ClientDisconnected),
8141            "client_disconnected"
8142        );
8143        assert_eq!(
8144            cancellation_reason_name(CancellationReason::SessionClosed),
8145            "session_closed"
8146        );
8147        assert_eq!(
8148            cancellation_reason_name(CancellationReason::ServerShutdown),
8149            "server_shutdown"
8150        );
8151        assert_eq!(
8152            cancellation_reason_name(CancellationReason::Deadline),
8153            "deadline"
8154        );
8155    }
8156
8157    #[test]
8158    fn unknown_outcome_never_proves_idempotency_intent_safe_to_abort() {
8159        let registry = Arc::new(SqlQueryRegistry::default());
8160        let unknown_id: QueryId = "102132435465768798a9bacbdcedfe0f".parse().unwrap();
8161        let unknown = registry
8162            .register(SqlQueryOptions {
8163                query_id: Some(unknown_id),
8164                ..SqlQueryOptions::default()
8165            })
8166            .unwrap();
8167        unknown.mark_outcome_unknown();
8168        unknown.fail();
8169        assert!(!can_abort_idempotency_intent(
8170            &registry.status(unknown_id).unwrap()
8171        ));
8172
8173        let failed_id: QueryId = "2031425364758697a8b9cadbecfd0e1f".parse().unwrap();
8174        let failed = registry
8175            .register(SqlQueryOptions {
8176                query_id: Some(failed_id),
8177                ..SqlQueryOptions::default()
8178            })
8179            .unwrap();
8180        failed.fail();
8181        assert!(can_abort_idempotency_intent(
8182            &registry.status(failed_id).unwrap()
8183        ));
8184    }
8185
8186    #[test]
8187    fn unknown_outcome_never_becomes_durable_receipt() {
8188        let registry = Arc::new(SqlQueryRegistry::default());
8189        let query = registry.register(SqlQueryOptions::default()).unwrap();
8190        query.record_commit(0, 42);
8191        query.mark_outcome_unknown();
8192        query.fail();
8193        let status = registry.status(query.id()).unwrap();
8194        assert!(status.durable_outcome.committed);
8195        assert!(status.outcome_unknown);
8196        assert!(sql_terminal_idempotency_receipt(&status).is_none());
8197    }
8198
8199    #[test]
8200    fn successful_noop_write_becomes_noncommitting_receipt() {
8201        let registry = Arc::new(SqlQueryRegistry::default());
8202        let query = registry.register(SqlQueryOptions::default()).unwrap();
8203        query
8204            .transition(SqlQueryPhase::Queued, SqlQueryPhase::Executing)
8205            .unwrap();
8206        query.try_complete().unwrap();
8207        let status = registry.status(query.id()).unwrap();
8208        assert!(!status.durable_outcome.committed);
8209        assert!(!can_abort_idempotency_intent(&status));
8210        let receipt = sql_terminal_idempotency_receipt(&status).unwrap();
8211        assert_eq!(receipt.status, "completed");
8212        assert!(!receipt.outcome.committed);
8213        assert_eq!(receipt.outcome.committed_statements, 0);
8214        assert_eq!(receipt.outcome.last_commit_epoch, None);
8215    }
8216
8217    #[test]
8218    fn conflicting_idempotency_key_sources_are_rejected() {
8219        let request = SqlRequest {
8220            sql: "INSERT INTO items VALUES (1)".into(),
8221            format: None,
8222            query_id: None,
8223            timeout_ms: None,
8224            max_output_rows: None,
8225            max_output_bytes: None,
8226            idempotency_key: Some("body-key".into()),
8227            pagination: None,
8228        };
8229        let mut headers = axum::http::HeaderMap::new();
8230        headers.insert("idempotency-key", "header-key".parse().unwrap());
8231        assert_eq!(
8232            requested_sql_idempotency_key(&headers, &request),
8233            Err("body idempotency_key and Idempotency-Key header must match")
8234        );
8235
8236        headers.insert("idempotency-key", "body-key".parse().unwrap());
8237        assert_eq!(
8238            requested_sql_idempotency_key(&headers, &request),
8239            Ok(Some("body-key".into()))
8240        );
8241    }
8242
8243    #[test]
8244    fn idempotency_binding_includes_pagination_semantics() {
8245        let mut request = SqlRequest {
8246            sql: "INSERT INTO items VALUES (1)".into(),
8247            format: None,
8248            query_id: None,
8249            timeout_ms: None,
8250            max_output_rows: None,
8251            max_output_bytes: None,
8252            idempotency_key: Some("key".into()),
8253            pagination: None,
8254        };
8255        let unpaged = sql_idempotency_binding(&request, (100, 1_024), None, 60_000).unwrap();
8256        request.pagination = Some(SqlPaginationRequest {
8257            page_size_rows: 10,
8258            projection: vec!["id".into()],
8259            max_page_bytes: Some(512),
8260            max_page_tokens: Some(128),
8261        });
8262        let paged = sql_idempotency_binding(&request, (100, 1_024), None, 60_000).unwrap();
8263        assert_ne!(unpaged.request_semantics_hash, paged.request_semantics_hash);
8264    }
8265
8266    #[test]
8267    fn paginated_decode_stops_nested_heap_amplification_at_budget() {
8268        let registry = Arc::new(SqlQueryRegistry::default());
8269        let query = registry.register(SqlQueryOptions::default()).unwrap();
8270        let json = format!("[[{}]]", vec!["null"; 10_000].join(","));
8271        let mut deserializer = serde_json::Deserializer::from_slice(json.as_bytes());
8272        let mut budget = PaginatedDecodeBudget {
8273            used: 0,
8274            limit: 4 * 1024,
8275            nodes: 0,
8276            exceeded: false,
8277            query: &query,
8278            test_hook: None,
8279        };
8280        let error = serde::de::DeserializeSeed::deserialize(
8281            BudgetedJsonRowsSeed {
8282                budget: &mut budget,
8283            },
8284            &mut deserializer,
8285        )
8286        .unwrap_err();
8287        assert!(error.to_string().contains(PAGINATED_MEMORY_LIMIT_ERROR));
8288        assert!(budget.exceeded);
8289        assert!(budget.nodes < 100, "decoded {} nodes", budget.nodes);
8290        query.fail();
8291    }
8292
8293    #[tokio::test]
8294    async fn unknown_outcome_response_never_claims_no_commit() {
8295        let registry = Arc::new(SqlQueryRegistry::default());
8296        let query = registry.register(SqlQueryOptions::default()).unwrap();
8297        let query_id = query.id();
8298        query
8299            .transition(SqlQueryPhase::Queued, SqlQueryPhase::Executing)
8300            .unwrap();
8301        let error = query.outcome_unknown_error("fenced maintenance failed");
8302        query.fail();
8303        let status = registry.status(query_id).unwrap();
8304        assert_eq!(
8305            status.terminal_state(),
8306            Some(mongreldb_query::QueryTerminalState::OutcomeUnknown)
8307        );
8308
8309        let response = query_error_response_with_status(&error, Some(query_id), Some(&status));
8310        assert_eq!(response.status(), StatusCode::CONFLICT);
8311        let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
8312            .await
8313            .unwrap();
8314        let body: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
8315        assert_eq!(body["status"], "outcome_unknown");
8316        assert_eq!(body["error"]["code"], "QUERY_OUTCOME_UNKNOWN");
8317        assert!(body["committed"].is_null());
8318        assert!(body["committed_statements"].is_null());
8319        assert!(body["last_commit_epoch"].is_null());
8320        assert!(body["completed_statements"].is_null());
8321        assert!(body["statement_index"].is_null());
8322        assert!(body["outcome"]["committed"].is_null());
8323        assert!(body["error"]["committed"].is_null());
8324    }
8325
8326    #[tokio::test]
8327    async fn retryable_idempotency_error_survives_terminal_status() {
8328        let registry = Arc::new(SqlQueryRegistry::default());
8329        let query = registry.register(SqlQueryOptions::default()).unwrap();
8330        let query_id = query.id();
8331        let response = registered_sql_error_response(
8332            RegisteredQueryGuard::new(query),
8333            query_id,
8334            StatusCode::SERVICE_UNAVAILABLE,
8335            "IDEMPOTENCY_STORE_UNAVAILABLE",
8336            "could not durably reserve the SQL idempotency key",
8337            true,
8338        );
8339        let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
8340            .await
8341            .unwrap();
8342        let body: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
8343        assert_eq!(body["retryable"], true);
8344        assert_eq!(body["error"]["retryable"], true);
8345
8346        let status = registry.status(query_id).unwrap();
8347        assert_eq!(
8348            status.terminal_error.as_ref().unwrap().code,
8349            "IDEMPOTENCY_STORE_UNAVAILABLE"
8350        );
8351        assert!(terminal_error_retryable(status.terminal_error.as_ref()));
8352    }
8353
8354    #[tokio::test]
8355    async fn committed_idempotency_terminal_error_is_a_receipt() {
8356        let query_id: QueryId = "00112233445566778899aabbccddeeff".parse().unwrap();
8357        let receipt = sql_idempotency::SqlDurableReceipt {
8358            original_query_id: query_id.to_string(),
8359            status: "committed_with_error".into(),
8360            server_state: "failed".into(),
8361            cancellation_reason: "client_disconnected".into(),
8362            outcome: sql_idempotency::SqlReceiptOutcome {
8363                committed: true,
8364                committed_statements: 1,
8365                last_commit_epoch: Some(42),
8366                last_commit_epoch_text: Some("42".into()),
8367                first_commit_statement_index: Some(0),
8368                last_commit_statement_index: Some(0),
8369                completed_statements: 1,
8370                statement_index: 0,
8371                serialization: "failed".into(),
8372                last_commit_hlc: None,
8373                serialization_state: None,
8374                terminal_state: None,
8375            },
8376            terminal_error: Some(sql_idempotency::SqlReceiptTerminalError {
8377                code: "SERIALIZATION_FAILED_AFTER_COMMIT".into(),
8378                category: "serialization".into(),
8379            }),
8380            commit_receipt: None,
8381        };
8382        let response = sql_idempotency_receipt_response(query_id, &receipt, false, 99, true);
8383        assert_eq!(response.status(), StatusCode::OK);
8384        let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
8385            .await
8386            .unwrap();
8387        let body: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
8388        assert_eq!(body["status"], "committed_with_error");
8389        assert_eq!(body["server_state"], "failed");
8390        assert_eq!(body["cancellation_reason"], "client_disconnected");
8391        assert_eq!(body["first_commit_statement_index"], 0);
8392        assert_eq!(body["last_commit_statement_index"], 0);
8393        assert_eq!(
8394            body["terminal_error"]["code"],
8395            "SERIALIZATION_FAILED_AFTER_COMMIT"
8396        );
8397    }
8398
8399    #[test]
8400    fn durable_replay_restores_terminal_status_parity() {
8401        let registry = Arc::new(SqlQueryRegistry::default());
8402        let query_id: QueryId = "11223344556677889900aabbccddeeff".parse().unwrap();
8403        let query = registry
8404            .register(SqlQueryOptions {
8405                query_id: Some(query_id),
8406                ..SqlQueryOptions::default()
8407            })
8408            .unwrap();
8409        let receipt = sql_idempotency::SqlDurableReceipt {
8410            original_query_id: "00112233445566778899aabbccddeeff".into(),
8411            status: "cancelled_after_commit".into(),
8412            server_state: "cancelled".into(),
8413            cancellation_reason: "client_disconnected".into(),
8414            outcome: sql_idempotency::SqlReceiptOutcome {
8415                committed: true,
8416                committed_statements: 1,
8417                last_commit_epoch: Some(42),
8418                last_commit_epoch_text: Some("42".into()),
8419                first_commit_statement_index: Some(0),
8420                last_commit_statement_index: Some(0),
8421                completed_statements: 1,
8422                statement_index: 1,
8423                serialization: "failed".into(),
8424                last_commit_hlc: None,
8425                serialization_state: None,
8426                terminal_state: None,
8427            },
8428            terminal_error: Some(sql_idempotency::SqlReceiptTerminalError {
8429                code: "QUERY_CANCELLED_AFTER_COMMIT".into(),
8430                category: "cancellation".into(),
8431            }),
8432            commit_receipt: None,
8433        };
8434        restore_idempotency_replay(RegisteredQueryGuard::new(query), &receipt).unwrap();
8435        let status = registry.status(query_id).unwrap();
8436        assert_eq!(status.phase, SqlQueryPhase::Cancelled);
8437        assert_eq!(
8438            status.terminal_state(),
8439            Some(mongreldb_query::QueryTerminalState::CancelledAfterCommit)
8440        );
8441        assert_eq!(
8442            status.cancellation_reason,
8443            CancellationReason::ClientDisconnected
8444        );
8445        assert_eq!(status.durable_outcome.committed_statements, 1);
8446        assert_eq!(status.durable_outcome.last_commit_epoch, Some(42));
8447        assert_eq!(
8448            status.terminal_error.unwrap().code,
8449            "QUERY_CANCELLED_AFTER_COMMIT"
8450        );
8451    }
8452
8453    #[test]
8454    fn durable_replay_rejects_invalid_authenticated_state() {
8455        let registry = Arc::new(SqlQueryRegistry::default());
8456        let query = registry.register(SqlQueryOptions::default()).unwrap();
8457        let receipt = sql_idempotency::SqlDurableReceipt {
8458            original_query_id: query.id().to_string(),
8459            status: "invented_terminal_state".into(),
8460            server_state: "completed".into(),
8461            cancellation_reason: "none".into(),
8462            outcome: sql_idempotency::SqlReceiptOutcome {
8463                committed: true,
8464                committed_statements: 1,
8465                last_commit_epoch: Some(42),
8466                last_commit_epoch_text: Some("42".into()),
8467                first_commit_statement_index: Some(0),
8468                last_commit_statement_index: Some(0),
8469                completed_statements: 1,
8470                statement_index: 0,
8471                serialization: "succeeded".into(),
8472                last_commit_hlc: None,
8473                serialization_state: None,
8474                terminal_state: None,
8475            },
8476            terminal_error: None,
8477            commit_receipt: None,
8478        };
8479        let error =
8480            restore_idempotency_replay(RegisteredQueryGuard::new(query), &receipt).unwrap_err();
8481        assert!(error.to_string().contains("invalid terminal state"));
8482    }
8483
8484    #[test]
8485    fn durable_replay_cancel_wins_before_receipt_response() {
8486        let registry = Arc::new(SqlQueryRegistry::default());
8487        let query_id: QueryId = "22334455667788990011aabbccddeeff".parse().unwrap();
8488        let query = registry
8489            .register(SqlQueryOptions {
8490                query_id: Some(query_id),
8491                ..SqlQueryOptions::default()
8492            })
8493            .unwrap();
8494        assert_eq!(
8495            query.request_cancel(CancellationReason::ClientRequest),
8496            CancelOutcome::Accepted
8497        );
8498        let receipt = sql_idempotency::SqlDurableReceipt {
8499            original_query_id: "00112233445566778899aabbccddeeff".into(),
8500            status: "completed".into(),
8501            server_state: "completed".into(),
8502            cancellation_reason: "none".into(),
8503            outcome: sql_idempotency::SqlReceiptOutcome {
8504                committed: true,
8505                committed_statements: 1,
8506                last_commit_epoch: Some(42),
8507                last_commit_epoch_text: Some("42".into()),
8508                first_commit_statement_index: Some(0),
8509                last_commit_statement_index: Some(0),
8510                completed_statements: 1,
8511                statement_index: 0,
8512                serialization: "succeeded".into(),
8513                last_commit_hlc: None,
8514                serialization_state: None,
8515                terminal_state: None,
8516            },
8517            terminal_error: None,
8518            commit_receipt: None,
8519        };
8520        let error = restore_idempotency_replay(RegisteredQueryGuard::new(query), &receipt)
8521            .expect_err("accepted cancellation must suppress replay success");
8522        assert!(matches!(
8523            error,
8524            mongreldb_query::MongrelQueryError::QueryCancelled { .. }
8525        ));
8526        let status = registry.status(query_id).unwrap();
8527        assert_eq!(status.phase, SqlQueryPhase::Cancelled);
8528        assert!(status.durable_outcome.committed);
8529    }
8530
8531    #[test]
8532    fn direct_query_handle_preserves_receipt_after_tombstone_eviction() {
8533        let registry = Arc::new(SqlQueryRegistry::new(
8534            1,
8535            1,
8536            usize::MAX,
8537            std::time::Duration::from_secs(60),
8538        ));
8539        let first = registry.register(SqlQueryOptions::default()).unwrap();
8540        first
8541            .transition(SqlQueryPhase::Queued, SqlQueryPhase::Executing)
8542            .unwrap();
8543        first.record_commit(0, 42);
8544        first.complete_current_statement();
8545        first.try_complete().unwrap();
8546
8547        let second = registry.register(SqlQueryOptions::default()).unwrap();
8548        second
8549            .transition(SqlQueryPhase::Queued, SqlQueryPhase::Executing)
8550            .unwrap();
8551        second.try_complete().unwrap();
8552
8553        assert!(registry.status(first.id()).is_none());
8554        let receipt = sql_terminal_idempotency_receipt(&first.status()).unwrap();
8555        assert_eq!(receipt.outcome.committed_statements, 1);
8556        assert_eq!(receipt.outcome.last_commit_epoch, Some(42));
8557    }
8558
8559    #[test]
8560    fn cancellation_checkpoint_mismatch_returns_typed_error() {
8561        let registry = Arc::new(SqlQueryRegistry::default());
8562        let query = registry.register(SqlQueryOptions::default()).unwrap();
8563        assert!(matches!(
8564            cancellation_checkpoint_error(&query),
8565            mongreldb_query::MongrelQueryError::InvalidQueryState(_)
8566        ));
8567        assert_eq!(
8568            query.request_cancel(CancellationReason::ClientRequest),
8569            CancelOutcome::Accepted
8570        );
8571        assert!(matches!(
8572            cancellation_checkpoint_error(&query),
8573            mongreldb_query::MongrelQueryError::QueryCancelled { .. }
8574        ));
8575    }
8576
8577    #[tokio::test]
8578    async fn cancellation_after_commit_reports_durable_outcome() {
8579        let registry = Arc::new(SqlQueryRegistry::default());
8580        let query = registry.register(SqlQueryOptions::default()).unwrap();
8581        let query_id = query.id();
8582        query
8583            .transition(SqlQueryPhase::Queued, SqlQueryPhase::Executing)
8584            .unwrap();
8585        query.record_commit(0, 42);
8586        assert_eq!(
8587            query.request_cancel(CancellationReason::ClientRequest),
8588            CancelOutcome::Accepted
8589        );
8590        let error = query.checkpoint().unwrap_err();
8591        query.fail();
8592        let status = registry.status(query_id).unwrap();
8593
8594        let response = query_error_response_with_status(&error, Some(query_id), Some(&status));
8595        assert_eq!(response.status(), StatusCode::CONFLICT);
8596        let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
8597            .await
8598            .unwrap();
8599        let body: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
8600        assert_eq!(body["status"], "cancelled_after_commit");
8601        assert_eq!(body["error"]["code"], "QUERY_CANCELLED_AFTER_COMMIT");
8602        assert_eq!(body["committed"], true);
8603        assert_eq!(body["outcome"]["committed_statements"], 1);
8604        assert_eq!(body["outcome"]["last_commit_epoch"], 42);
8605        assert_eq!(body["outcome"]["last_commit_epoch_text"], "42");
8606        assert_eq!(body["first_commit_statement_index"], 0);
8607        assert_eq!(body["last_commit_statement_index"], 0);
8608        assert_eq!(body["outcome"]["first_commit_statement_index"], 0);
8609        assert_eq!(body["outcome"]["last_commit_statement_index"], 0);
8610
8611        let response = query_error_response_with_status(&error, Some(query_id), None);
8612        assert_eq!(response.status(), StatusCode::CONFLICT);
8613        let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
8614            .await
8615            .unwrap();
8616        let body: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
8617        assert_eq!(body["status"], "cancelled_after_commit");
8618        assert_eq!(body["error"]["code"], "QUERY_CANCELLED_AFTER_COMMIT");
8619        assert_eq!(body["committed"], true);
8620        assert_eq!(body["committed_statements"], 1);
8621        assert_eq!(body["last_commit_epoch"], 42);
8622        assert_eq!(body["last_commit_epoch_text"], "42");
8623        assert_eq!(body["first_commit_statement_index"], 0);
8624        assert_eq!(body["last_commit_statement_index"], 0);
8625        assert_eq!(body["outcome"]["committed"], true);
8626        assert_eq!(body["outcome"]["committed_statements"], 1);
8627        assert_eq!(body["outcome"]["last_commit_epoch"], 42);
8628        assert_eq!(body["outcome"]["last_commit_epoch_text"], "42");
8629        assert_eq!(body["outcome"]["first_commit_statement_index"], 0);
8630        assert_eq!(body["outcome"]["last_commit_statement_index"], 0);
8631    }
8632
8633    #[tokio::test]
8634    async fn commit_outcome_fallback_preserves_exact_progress() {
8635        let query_id: QueryId = "33445566778899001122aabbccddeeff".parse().unwrap();
8636        let error = mongreldb_query::MongrelQueryError::CommitOutcome {
8637            query_id,
8638            committed: true,
8639            committed_statements: 3,
8640            last_commit_epoch: Some(77),
8641            first_commit_statement_index: Some(1),
8642            last_commit_statement_index: Some(4),
8643            completed_statements: 4,
8644            statement_index: 5,
8645            message: "durable outcome retained".into(),
8646        };
8647        let response = query_error_response_with_status(&error, Some(query_id), None);
8648        assert_eq!(response.status(), StatusCode::CONFLICT);
8649        let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
8650            .await
8651            .unwrap();
8652        let body: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
8653        assert_eq!(body["status"], "committed_with_error");
8654        assert_eq!(body["committed"], true);
8655        assert_eq!(body["committed_statements"], 3);
8656        assert_eq!(body["last_commit_epoch"], 77);
8657        assert_eq!(body["last_commit_epoch_text"], "77");
8658        assert_eq!(body["first_commit_statement_index"], 1);
8659        assert_eq!(body["last_commit_statement_index"], 4);
8660        assert_eq!(body["completed_statements"], 4);
8661        assert_eq!(body["statement_index"], 5);
8662        assert_eq!(body["outcome"]["committed_statements"], 3);
8663        assert_eq!(body["outcome"]["last_commit_epoch"], 77);
8664        assert_eq!(body["outcome"]["first_commit_statement_index"], 1);
8665        assert_eq!(body["outcome"]["last_commit_statement_index"], 4);
8666        assert_eq!(body["outcome"]["completed_statements"], 4);
8667        assert_eq!(body["outcome"]["statement_index"], 5);
8668    }
8669
8670    #[test]
8671    fn status_cancel_outcome_matches_cancel_endpoint_state() {
8672        let registry = Arc::new(SqlQueryRegistry::default());
8673        let commit = registry.register(SqlQueryOptions::default()).unwrap();
8674        commit
8675            .transition(SqlQueryPhase::Queued, SqlQueryPhase::Executing)
8676            .unwrap();
8677        commit.enter_commit_critical().unwrap();
8678        assert_eq!(
8679            query_cancel_outcome(&registry.status(commit.id()).unwrap()),
8680            Some("too_late")
8681        );
8682
8683        let completed = registry.register(SqlQueryOptions::default()).unwrap();
8684        completed
8685            .transition(SqlQueryPhase::Queued, SqlQueryPhase::Executing)
8686            .unwrap();
8687        completed.try_complete().unwrap();
8688        assert_eq!(
8689            query_cancel_outcome(&registry.status(completed.id()).unwrap()),
8690            Some("already_finished")
8691        );
8692    }
8693}
8694
8695#[cfg(test)]
8696mod wal_stream_tests {
8697    use super::*;
8698    use mongreldb_client::ReplicationFollower;
8699    use mongreldb_core::Database;
8700    use tempfile::tempdir;
8701
8702    #[tokio::test]
8703    async fn wal_stream_returns_records_after_commit() {
8704        let dir = tempdir().unwrap();
8705        let db = Arc::new(Database::create(dir.path()).unwrap());
8706        let table_schema = mongreldb_core::schema::Schema {
8707            schema_id: 1,
8708            columns: vec![mongreldb_core::schema::ColumnDef {
8709                id: 1,
8710                name: "id".into(),
8711                ty: TypeId::Int64,
8712                flags: mongreldb_core::schema::ColumnFlags::empty()
8713                    .with(mongreldb_core::schema::ColumnFlags::PRIMARY_KEY),
8714                default_value: None,
8715                embedding_source: None,
8716            }],
8717            indexes: vec![],
8718            colocation: vec![],
8719            constraints: Default::default(),
8720            clustered: false,
8721        };
8722        db.create_table("items", table_schema).unwrap();
8723        // Write a row to generate WAL records.
8724        let handle = db.table("items").unwrap();
8725        handle.lock().put(vec![(1, Value::Int64(1))]).unwrap();
8726        handle.lock().flush().unwrap();
8727
8728        let app = build_app(db);
8729        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
8730        let addr = listener.local_addr().unwrap();
8731        tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
8732        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
8733
8734        let resp = reqwest::get(format!("http://{addr}/wal/stream"))
8735            .await
8736            .unwrap();
8737        assert_eq!(resp.status(), 200);
8738        let body = resp.text().await.unwrap();
8739        // Should contain at least one record (the flush commit).
8740        assert!(!body.is_empty(), "wal_stream should return records");
8741        assert!(body.contains("seq"), "response should contain seq field");
8742    }
8743
8744    #[tokio::test]
8745    async fn follower_bootstraps_and_applies_incremental_commit() {
8746        let leader_dir = tempdir().unwrap();
8747        let follower_dir = tempdir().unwrap();
8748        let follower_path = follower_dir.path().join("copy");
8749        let db = Arc::new(Database::create(leader_dir.path()).unwrap());
8750        db.create_table(
8751            "items",
8752            mongreldb_core::schema::Schema {
8753                schema_id: 1,
8754                columns: vec![mongreldb_core::schema::ColumnDef {
8755                    id: 1,
8756                    name: "id".into(),
8757                    ty: TypeId::Int64,
8758                    flags: mongreldb_core::schema::ColumnFlags::empty()
8759                        .with(mongreldb_core::schema::ColumnFlags::PRIMARY_KEY),
8760                    default_value: None,
8761                    embedding_source: None,
8762                }],
8763                indexes: vec![],
8764                colocation: vec![],
8765                constraints: Default::default(),
8766                clustered: false,
8767            },
8768        )
8769        .unwrap();
8770        let handle = db.table("items").unwrap();
8771        handle.lock().put(vec![(1, Value::Int64(1))]).unwrap();
8772        handle.lock().commit().unwrap();
8773
8774        let app = build_app_with_config(
8775            Arc::clone(&db),
8776            std::iter::empty::<Arc<dyn ExternalTableModule>>(),
8777            Some("replication-secret".into()),
8778            None,
8779        );
8780        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
8781        let addr = listener.local_addr().unwrap();
8782        tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
8783        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
8784
8785        let leader_url = format!("http://{addr}");
8786        let first_path = follower_path.clone();
8787        let (mut follower, initial) = tokio::task::spawn_blocking(move || {
8788            let mut follower = ReplicationFollower::new(&leader_url, first_path)
8789                .unwrap()
8790                .with_bearer_token("replication-secret");
8791            let applied = follower.sync().unwrap();
8792            (follower, applied)
8793        })
8794        .await
8795        .unwrap();
8796        assert_eq!(initial, 0);
8797
8798        handle.lock().put(vec![(1, Value::Int64(2))]).unwrap();
8799        handle.lock().commit().unwrap();
8800        let applied = tokio::task::spawn_blocking(move || {
8801            let count = follower.sync().unwrap();
8802            (follower, count)
8803        })
8804        .await
8805        .unwrap();
8806        follower = applied.0;
8807        assert!(applied.1 > 0);
8808        assert!(follower.last_epoch() > 0);
8809
8810        let replica = Database::open(&follower_path).unwrap();
8811        assert_eq!(replica.table("items").unwrap().lock().count(), 2);
8812        drop(replica);
8813
8814        db.set_spill_threshold(1);
8815        db.transaction(|txn| {
8816            txn.put("items", vec![(1, Value::Int64(3))])?;
8817            Ok(())
8818        })
8819        .unwrap();
8820        let (follower_after_bootstrap, applied) = tokio::task::spawn_blocking(move || {
8821            let count = follower.sync().unwrap();
8822            (follower, count)
8823        })
8824        .await
8825        .unwrap();
8826        assert_eq!(applied, 0, "spilled run should trigger safe rebootstrap");
8827        assert!(follower_after_bootstrap.last_epoch() > 0);
8828        let replica = Database::open(&follower_path).unwrap();
8829        assert_eq!(replica.table("items").unwrap().lock().count(), 3);
8830    }
8831}
8832
8833#[cfg(test)]
8834mod metrics_tests {
8835    use super::*;
8836    use mongreldb_core::Database;
8837    use tempfile::tempdir;
8838
8839    /// Helper: spin up a daemon over a fresh DB with one `items(id int64 pk)`
8840    /// table pre-created, returning the bound address.
8841    async fn setup() -> (tempfile::TempDir, std::net::SocketAddr) {
8842        let dir = tempdir().unwrap();
8843        let db = Arc::new(Database::create(dir.path()).unwrap());
8844        let table_schema = mongreldb_core::schema::Schema {
8845            schema_id: 1,
8846            columns: vec![mongreldb_core::schema::ColumnDef {
8847                id: 1,
8848                name: "id".into(),
8849                ty: TypeId::Int64,
8850                flags: mongreldb_core::schema::ColumnFlags::empty()
8851                    .with(mongreldb_core::schema::ColumnFlags::PRIMARY_KEY),
8852                default_value: None,
8853                embedding_source: None,
8854            }],
8855            indexes: vec![],
8856            colocation: vec![],
8857            constraints: Default::default(),
8858            clustered: false,
8859        };
8860        db.create_table("items", table_schema).unwrap();
8861        let app = build_app(db);
8862        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
8863        let addr = listener.local_addr().unwrap();
8864        tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
8865        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
8866        (dir, addr)
8867    }
8868
8869    #[tokio::test]
8870    async fn metrics_endpoint_returns_prometheus_text() {
8871        let (_dir, addr) = setup().await;
8872        let client = reqwest::Client::new();
8873
8874        // Exercise a few handlers to bump counters.
8875        let _ = client
8876            .post(format!("http://{addr}/tables/items/put"))
8877            .json(&json!({ "row": [1, 1] }))
8878            .send()
8879            .await
8880            .unwrap();
8881        let _ = client
8882            .post(format!("http://{addr}/sql"))
8883            .json(&json!({ "sql": "SELECT count(*) FROM items" }))
8884            .send()
8885            .await
8886            .unwrap();
8887
8888        let resp = client
8889            .get(format!("http://{addr}/metrics"))
8890            .send()
8891            .await
8892            .unwrap();
8893        assert_eq!(resp.status(), 200);
8894        let ct = resp
8895            .headers()
8896            .get("content-type")
8897            .and_then(|v| v.to_str().ok())
8898            .unwrap_or_default()
8899            .to_string();
8900        assert!(
8901            ct.contains("text/plain"),
8902            "content-type is prometheus text: {ct}"
8903        );
8904        let body = resp.text().await.unwrap();
8905        // Prometheus series + type lines are present.
8906        assert!(body.contains("# TYPE mongreldb_sql_queries_total counter"));
8907        assert!(body.contains("# TYPE mongreldb_puts_total counter"));
8908        assert!(body.contains("# TYPE mongreldb_tables gauge"));
8909        // Counters were bumped: at least one query and one put were served.
8910        assert!(
8911            body.contains("mongreldb_sql_queries_total 1"),
8912            "sql_queries counter should reflect the /sql call: {body}"
8913        );
8914        assert!(
8915            body.contains("mongreldb_puts_total 1"),
8916            "puts counter should reflect the put call: {body}"
8917        );
8918        // The tables gauge reflects the single `items` table.
8919        assert!(body.contains("mongreldb_tables 1"));
8920    }
8921
8922    #[tokio::test]
8923    async fn metrics_error_counter_increments_on_bad_sql() {
8924        let (_dir, addr) = setup().await;
8925        let client = reqwest::Client::new();
8926        // A query against a non-existent table errors at the engine layer.
8927        let _ = client
8928            .post(format!("http://{addr}/sql"))
8929            .json(&json!({ "sql": "SELECT * FROM does_not_exist" }))
8930            .send()
8931            .await
8932            .unwrap();
8933        let body = client
8934            .get(format!("http://{addr}/metrics"))
8935            .send()
8936            .await
8937            .unwrap()
8938            .text()
8939            .await
8940            .unwrap();
8941        assert!(
8942            body.contains("mongreldb_sql_errors_total 1"),
8943            "sql_errors should increment on a failed query: {body}"
8944        );
8945    }
8946
8947    #[tokio::test]
8948    async fn arrow_stream_returns_ipc_stream_bytes() {
8949        let (_dir, addr) = setup().await;
8950        let client = reqwest::Client::new();
8951        // Insert a couple of rows so there are real batches to stream, then
8952        // flush so the rows are durable/visible to a fresh SQL session.
8953        for i in 1..=3 {
8954            let resp = client
8955                .post(format!("http://{addr}/tables/items/put"))
8956                .json(&json!({ "row": [1, i] }))
8957                .send()
8958                .await
8959                .unwrap();
8960            assert_eq!(resp.status(), 200, "put should succeed");
8961        }
8962        let _ = client
8963            .post(format!("http://{addr}/tables/items/commit"))
8964            .send()
8965            .await
8966            .unwrap();
8967        // Sanity: the rows are durable and visible to the table handle.
8968        let count_body = client
8969            .get(format!("http://{addr}/tables/items/count"))
8970            .send()
8971            .await
8972            .unwrap()
8973            .text()
8974            .await
8975            .unwrap();
8976        assert!(
8977            count_body.contains("\"count\":3"),
8978            "expected 3 visible rows, got: {count_body}"
8979        );
8980        let resp = client
8981            .post(format!("http://{addr}/sql"))
8982            .json(&json!({ "sql": "SELECT count(*) FROM items", "format": "arrow-stream" }))
8983            .send()
8984            .await
8985            .unwrap();
8986        assert_eq!(resp.status(), 200, "streaming query should succeed");
8987        let ct = resp
8988            .headers()
8989            .get("content-type")
8990            .and_then(|v| v.to_str().ok())
8991            .unwrap_or_default()
8992            .to_string();
8993        assert!(
8994            ct.contains("application/vnd.apache.arrow.stream"),
8995            "content-type should be the arrow stream format: {ct}"
8996        );
8997        let bytes = resp.bytes().await.unwrap();
8998        // Arrow IPC streams begin with the magic continuation marker
8999        // 0xFFFFFFFF followed by the schema message length. The stream must be
9000        // non-empty (3 rows were written) and end with the EOS marker.
9001        assert!(
9002            !bytes.is_empty(),
9003            "arrow stream body should contain schema + batch + EOS"
9004        );
9005        assert!(
9006            bytes.starts_with(&0xFFFFFFFFu32.to_le_bytes()),
9007            "arrow stream must begin with the IPC continuation marker"
9008        );
9009        assert!(
9010            bytes.ends_with(&[0u8, 0, 0, 0]),
9011            "arrow stream should end with the EOS marker (trailing zero length)"
9012        );
9013    }
9014}
9015
9016#[cfg(test)]
9017mod streaming_tests {
9018    use super::*;
9019    use arrow::array::Int64Array;
9020    use arrow::datatypes::{DataType, Field, Schema};
9021    use arrow::record_batch::RecordBatch;
9022    use datafusion::common::DataFusionError;
9023    use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
9024    use futures::StreamExt;
9025    use std::sync::Arc as StdArc;
9026
9027    fn batch_stream(batches: Vec<RecordBatch>) -> mongreldb_query::MongrelRecordBatchStream {
9028        let schema = batches
9029            .first()
9030            .map(RecordBatch::schema)
9031            .unwrap_or_else(|| StdArc::new(Schema::empty()));
9032        let batches =
9033            futures::stream::iter(batches.into_iter().map(Ok::<RecordBatch, DataFusionError>));
9034        Box::pin(RecordBatchStreamAdapter::new(schema, batches))
9035    }
9036
9037    /// Unit-level check: feed two synthetic batches through the streaming
9038    /// serializer and re-parse the resulting IPC stream end-to-end. This
9039    /// validates the per-message chunking (schema + N batches + EOS) without
9040    /// depending on the engine's scan visibility.
9041    #[tokio::test]
9042    async fn arrow_stream_serializes_multiple_batches_roundtrip() {
9043        let schema = StdArc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)]));
9044        let b1 = RecordBatch::try_new(
9045            schema.clone(),
9046            vec![StdArc::new(Int64Array::from(vec![1, 2]))],
9047        )
9048        .unwrap();
9049        let b2 = RecordBatch::try_new(
9050            schema.clone(),
9051            vec![StdArc::new(Int64Array::from(vec![3, 4, 5]))],
9052        )
9053        .unwrap();
9054
9055        let resp = sql_arrow_stream_response(batch_stream(vec![b1, b2]));
9056        let bytes = axum::body::to_bytes(resp.into_body(), 1 << 20)
9057            .await
9058            .unwrap();
9059
9060        // Begins with the IPC continuation marker.
9061        assert!(bytes.starts_with(&0xFFFFFFFFu32.to_le_bytes()));
9062
9063        // Re-parse the full stream and confirm all rows survived the chunked
9064        // serialization.
9065        let slice: &[u8] = bytes.as_ref();
9066        let mut reader = arrow::ipc::reader::StreamReader::try_new(slice, None).unwrap();
9067        let mut total_rows = 0;
9068        for batch in reader.by_ref() {
9069            let batch = batch.expect("each IPC message should decode");
9070            total_rows += batch.num_rows();
9071        }
9072        assert_eq!(
9073            total_rows, 5,
9074            "all rows should round-trip through the stream"
9075        );
9076    }
9077
9078    #[tokio::test]
9079    async fn arrow_stream_emits_schema_before_first_batch() {
9080        let schema = StdArc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)]));
9081        let pending = futures::stream::pending::<Result<RecordBatch, DataFusionError>>();
9082        let batches = Box::pin(RecordBatchStreamAdapter::new(schema, pending));
9083        let mut body = sql_arrow_stream_response(batches)
9084            .into_body()
9085            .into_data_stream();
9086
9087        let chunk = tokio::time::timeout(std::time::Duration::from_millis(100), body.next())
9088            .await
9089            .expect("schema chunk should not wait for a query batch")
9090            .unwrap()
9091            .unwrap();
9092        assert!(chunk.starts_with(&0xFFFFFFFFu32.to_le_bytes()));
9093    }
9094
9095    #[tokio::test]
9096    async fn arrow_stream_empty_query_is_valid_ipc() {
9097        let resp = sql_arrow_stream_response(batch_stream(Vec::new()));
9098        let bytes = axum::body::to_bytes(resp.into_body(), 1 << 20)
9099            .await
9100            .unwrap();
9101        let slice: &[u8] = bytes.as_ref();
9102        let reader = arrow::ipc::reader::StreamReader::try_new(slice, None).unwrap();
9103        assert_eq!(reader.count(), 0);
9104    }
9105
9106    #[tokio::test]
9107    async fn buffered_output_limits_are_typed() {
9108        let dir = tempfile::tempdir().unwrap();
9109        let db = StdArc::new(mongreldb_core::Database::create(dir.path()).unwrap());
9110        let session = MongrelSession::open(db).unwrap();
9111        let query = session.register_query(SqlQueryOptions::default()).unwrap();
9112        let output = session
9113            .run_with_query_for_serialization("SELECT 1", query)
9114            .await
9115            .unwrap();
9116
9117        let row_error =
9118            serialize_buffered_output("json", output.batches(), output.query(), 0, 1024, None)
9119                .unwrap_err();
9120        assert!(matches!(row_error, BufferedSerializationError::Limit(_)));
9121        let byte_error =
9122            serialize_buffered_output("json", output.batches(), output.query(), 10, 1, None)
9123                .unwrap_err();
9124        assert!(matches!(byte_error, BufferedSerializationError::Limit(_)));
9125        output.fail();
9126    }
9127}
9128
9129#[cfg(test)]
9130mod audit_tests {
9131    use super::*;
9132    use mongreldb_core::Database;
9133    use tempfile::tempdir;
9134
9135    /// Build a daemon with user-auth enabled plus one catalog user `alice`.
9136    async fn auth_setup(password: &str) -> std::net::SocketAddr {
9137        let dir = tempdir().unwrap();
9138        let db = Arc::new(Database::create(dir.path()).unwrap());
9139        db.create_user("alice", password).unwrap();
9140        db.set_user_admin("alice", true).unwrap();
9141        let app = build_app_full(
9142            db,
9143            std::iter::empty::<Arc<dyn ExternalTableModule>>(),
9144            None,
9145            None,
9146            true,
9147        );
9148        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
9149        let addr = listener.local_addr().unwrap();
9150        tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
9151        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
9152        addr
9153    }
9154
9155    #[tokio::test]
9156    async fn audit_records_login_success_and_failure() {
9157        let addr = auth_setup("s3cret").await;
9158        let client = reqwest::Client::new();
9159
9160        // Successful basic auth.
9161        let resp = client
9162            .get(format!("http://{addr}/health"))
9163            .header("Authorization", basic("alice", "s3cret"))
9164            .send()
9165            .await
9166            .unwrap();
9167        assert_eq!(resp.status(), 200);
9168
9169        // Failed basic auth (wrong password).
9170        let resp = client
9171            .get(format!("http://{addr}/health"))
9172            .header("Authorization", basic("alice", "wrong"))
9173            .send()
9174            .await
9175            .unwrap();
9176        assert_eq!(resp.status(), 401);
9177
9178        let body = client
9179            .get(format!("http://{addr}/audit"))
9180            .header("Authorization", basic("alice", "s3cret"))
9181            .send()
9182            .await
9183            .unwrap()
9184            .text()
9185            .await
9186            .unwrap();
9187        assert!(
9188            body.contains("\"action\":\"login.ok\""),
9189            "audit should record the successful login: {body}"
9190        );
9191        assert!(
9192            body.contains("\"action\":\"login.fail\""),
9193            "audit should record the failed login: {body}"
9194        );
9195        assert!(
9196            body.contains("\"principal\":\"alice\""),
9197            "audit should attribute events to alice: {body}"
9198        );
9199    }
9200
9201    #[tokio::test]
9202    async fn audit_records_ddl_sql() {
9203        let dir = tempdir().unwrap();
9204        let db = Arc::new(Database::create(dir.path()).unwrap());
9205        let app = build_app(db);
9206        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
9207        let addr = listener.local_addr().unwrap();
9208        tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
9209        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
9210
9211        let client = reqwest::Client::new();
9212        let _ = client
9213            .post(format!("http://{addr}/sql"))
9214            .json(&json!({ "sql": "CREATE TABLE t (id BIGINT PRIMARY KEY)" }))
9215            .send()
9216            .await
9217            .unwrap();
9218        // A plain SELECT is NOT audited.
9219        let _ = client
9220            .post(format!("http://{addr}/sql"))
9221            .json(&json!({ "sql": "SELECT 1" }))
9222            .send()
9223            .await
9224            .unwrap();
9225
9226        let body = client
9227            .get(format!("http://{addr}/audit"))
9228            .send()
9229            .await
9230            .unwrap()
9231            .text()
9232            .await
9233            .unwrap();
9234        assert!(
9235            body.contains("\"action\":\"ddl.ok\""),
9236            "audit should record the successful DDL statement: {body}"
9237        );
9238        assert!(
9239            body.contains("CREATE TABLE"),
9240            "audit detail should carry the DDL snippet: {body}"
9241        );
9242        // SELECT must not appear.
9243        assert!(
9244            !body.contains("SELECT 1"),
9245            "non-DDL reads should not be audited: {body}"
9246        );
9247    }
9248
9249    #[tokio::test]
9250    async fn audit_redacts_credential_passwords() {
9251        let dir = tempdir().unwrap();
9252        let db = Arc::new(Database::create(dir.path()).unwrap());
9253        let app = build_app(db);
9254        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
9255        let addr = listener.local_addr().unwrap();
9256        tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
9257        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
9258
9259        let client = reqwest::Client::new();
9260        // A CREATE USER carrying a plaintext password must be audited but the
9261        // password must NEVER reach /audit or stderr.
9262        let _ = client
9263            .post(format!("http://{addr}/sql"))
9264            .json(&json!({ "sql": "CREATE USER alice WITH PASSWORD 'topsecret'" }))
9265            .send()
9266            .await
9267            .unwrap();
9268
9269        let body = client
9270            .get(format!("http://{addr}/audit"))
9271            .send()
9272            .await
9273            .unwrap()
9274            .text()
9275            .await
9276            .unwrap();
9277        assert!(
9278            !body.contains("topsecret"),
9279            "password must never appear in the audit log: {body}"
9280        );
9281        assert!(
9282            body.contains("redacted credential statement"),
9283            "credential DDL should be recorded as redacted: {body}"
9284        );
9285    }
9286
9287    fn basic(user: &str, pass: &str) -> String {
9288        let raw = format!("{user}:{pass}");
9289        format!("Basic {}", base64_encode(raw.as_bytes()))
9290    }
9291
9292    fn base64_encode(input: &[u8]) -> String {
9293        const TABLE: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
9294        let mut out = String::new();
9295        let mut buf = 0u32;
9296        let mut bits = 0u32;
9297        for &b in input {
9298            buf = (buf << 8) | b as u32;
9299            bits += 8;
9300            while bits >= 6 {
9301                bits -= 6;
9302                out.push(TABLE[((buf >> bits) & 0x3F) as usize] as char);
9303            }
9304        }
9305        if bits > 0 {
9306            out.push(TABLE[((buf << (6 - bits)) & 0x3F) as usize] as char);
9307        }
9308        while !out.len().is_multiple_of(4) {
9309            out.push('=');
9310        }
9311        out
9312    }
9313}
9314
9315#[cfg(test)]
9316mod session_tests {
9317    use super::*;
9318    use mongreldb_core::Database;
9319    use tempfile::tempdir;
9320
9321    /// Spin up a daemon over a fresh DB with one `items(id int64 pk)` table.
9322    /// Returns the TempDir (must be held alive for the test's duration — dropping
9323    /// it deletes the database directory mid-test) and the bound address.
9324    async fn setup() -> (tempfile::TempDir, std::net::SocketAddr) {
9325        let dir = tempdir().unwrap();
9326        let db = Arc::new(Database::create(dir.path()).unwrap());
9327        let table_schema = mongreldb_core::schema::Schema {
9328            schema_id: 1,
9329            columns: vec![mongreldb_core::schema::ColumnDef {
9330                id: 1,
9331                name: "id".into(),
9332                ty: TypeId::Int64,
9333                flags: mongreldb_core::schema::ColumnFlags::empty()
9334                    .with(mongreldb_core::schema::ColumnFlags::PRIMARY_KEY),
9335                default_value: None,
9336                embedding_source: None,
9337            }],
9338            indexes: vec![],
9339            colocation: vec![],
9340            constraints: Default::default(),
9341            clustered: false,
9342        };
9343        db.create_table("items", table_schema).unwrap();
9344        let app = build_app(db);
9345        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
9346        let addr = listener.local_addr().unwrap();
9347        tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
9348        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
9349        (dir, addr)
9350    }
9351
9352    /// Open a session and return its token.
9353    async fn open_session(client: &reqwest::Client, addr: &std::net::SocketAddr) -> String {
9354        let resp = client
9355            .post(format!("http://{addr}/sessions"))
9356            .send()
9357            .await
9358            .unwrap();
9359        assert_eq!(resp.status(), 200);
9360        resp.json::<serde_json::Value>()
9361            .await
9362            .unwrap()
9363            .get("session_id")
9364            .unwrap()
9365            .as_str()
9366            .unwrap()
9367            .to_string()
9368    }
9369
9370    async fn sql_on(
9371        client: &reqwest::Client,
9372        addr: &std::net::SocketAddr,
9373        session: &str,
9374        sql: &str,
9375    ) -> reqwest::Response {
9376        client
9377            .post(format!("http://{addr}/sql"))
9378            .header("X-Session-ID", session)
9379            .json(&json!({ "sql": sql }))
9380            .send()
9381            .await
9382            .unwrap()
9383    }
9384
9385    async fn count_items(client: &reqwest::Client, addr: &std::net::SocketAddr) -> u64 {
9386        client
9387            .get(format!("http://{addr}/tables/items/count"))
9388            .send()
9389            .await
9390            .unwrap()
9391            .json::<serde_json::Value>()
9392            .await
9393            .unwrap()
9394            .get("count")
9395            .unwrap()
9396            .as_u64()
9397            .unwrap()
9398    }
9399
9400    #[tokio::test]
9401    async fn cross_request_transaction_commits() {
9402        let (_dir, addr) = setup().await;
9403        let client = reqwest::Client::new();
9404        let session = open_session(&client, &addr).await;
9405
9406        // BEGIN, INSERT, COMMIT — each its own HTTP request on the same session.
9407        let r = sql_on(&client, &addr, &session, "BEGIN").await;
9408        assert_eq!(r.status(), 200);
9409        let r = sql_on(
9410            &client,
9411            &addr,
9412            &session,
9413            "INSERT INTO items (id) VALUES (1)",
9414        )
9415        .await;
9416        assert_eq!(r.status(), 200, "INSERT should stage successfully");
9417        // Not yet committed → not visible to other connections.
9418        assert_eq!(count_items(&client, &addr).await, 0);
9419        let r = sql_on(&client, &addr, &session, "COMMIT").await;
9420        assert_eq!(r.status(), 200);
9421
9422        // After COMMIT the row is durable and visible.
9423        assert_eq!(count_items(&client, &addr).await, 1);
9424    }
9425
9426    #[tokio::test]
9427    async fn cross_request_transaction_rolls_back() {
9428        let (_dir, addr) = setup().await;
9429        let client = reqwest::Client::new();
9430        let session = open_session(&client, &addr).await;
9431
9432        sql_on(&client, &addr, &session, "BEGIN").await;
9433        sql_on(
9434            &client,
9435            &addr,
9436            &session,
9437            "INSERT INTO items (id) VALUES (5)",
9438        )
9439        .await;
9440        // ROLLBACK discards the staged insert.
9441        let r = sql_on(&client, &addr, &session, "ROLLBACK").await;
9442        assert_eq!(r.status(), 200);
9443        assert_eq!(
9444            count_items(&client, &addr).await,
9445            0,
9446            "rollback discards staged writes"
9447        );
9448    }
9449
9450    #[tokio::test]
9451    async fn unknown_session_id_is_404() {
9452        let (_dir, addr) = setup().await;
9453        let client = reqwest::Client::new();
9454        let resp = client
9455            .post(format!("http://{addr}/sql"))
9456            .header("X-Session-ID", "does-not-exist")
9457            .json(&json!({ "sql": "SELECT 1" }))
9458            .send()
9459            .await
9460            .unwrap();
9461        assert_eq!(resp.status(), 404);
9462    }
9463
9464    #[tokio::test]
9465    async fn invalid_session_headers_do_not_autocommit() {
9466        let (_dir, addr) = setup().await;
9467        let client = reqwest::Client::new();
9468
9469        let resp = client
9470            .post(format!("http://{addr}/sql"))
9471            .header(
9472                "X-Session-ID",
9473                reqwest::header::HeaderValue::from_bytes(&[0xff]).unwrap(),
9474            )
9475            .json(&json!({ "sql": "INSERT INTO items (id) VALUES (1)" }))
9476            .send()
9477            .await
9478            .unwrap();
9479        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
9480
9481        let resp = client
9482            .post(format!("http://{addr}/sql"))
9483            .header("X-Session-ID", "x".repeat(257))
9484            .json(&json!({ "sql": "INSERT INTO items (id) VALUES (2)" }))
9485            .send()
9486            .await
9487            .unwrap();
9488        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
9489        assert_eq!(count_items(&client, &addr).await, 0);
9490    }
9491
9492    #[tokio::test]
9493    async fn close_session_ends_cross_request_state() {
9494        let (_dir, addr) = setup().await;
9495        let client = reqwest::Client::new();
9496        let session = open_session(&client, &addr).await;
9497
9498        // BEGIN on the session, then close it → staged txn is discarded.
9499        sql_on(&client, &addr, &session, "BEGIN").await;
9500        let r = client
9501            .delete(format!("http://{addr}/sessions/{session}"))
9502            .send()
9503            .await
9504            .unwrap();
9505        assert_eq!(r.status(), 200);
9506
9507        // The session token is now invalid.
9508        let resp = sql_on(&client, &addr, &session, "COMMIT").await;
9509        assert_eq!(resp.status(), 404, "closed session is no longer usable");
9510    }
9511
9512    #[tokio::test]
9513    async fn no_session_header_uses_fresh_ephemeral_session() {
9514        // Without X-Session-ID, BEGIN..COMMIT must still work within a single
9515        // multi-statement /sql body (the historical behavior is preserved).
9516        let (_dir, addr) = setup().await;
9517        let client = reqwest::Client::new();
9518        let resp = client
9519            .post(format!("http://{addr}/sql"))
9520            .json(&json!({ "sql": "INSERT INTO items (id) VALUES (42)" }))
9521            .send()
9522            .await
9523            .unwrap();
9524        assert_eq!(resp.status(), 200);
9525        assert_eq!(count_items(&client, &addr).await, 1);
9526    }
9527
9528    #[tokio::test]
9529    async fn prepared_statement_prepare_execute_and_reuse() {
9530        let (_dir, addr) = setup().await;
9531        let client = reqwest::Client::new();
9532        let session = open_session(&client, &addr).await;
9533        sql_on(
9534            &client,
9535            &addr,
9536            &session,
9537            "INSERT INTO items (id) VALUES (1), (2), (3), (4)",
9538        )
9539        .await;
9540
9541        // Prepare a parameterized query once.
9542        let resp = client
9543            .post(format!("http://{addr}/sessions/{session}/prepare"))
9544            .json(&json!({"name":"gt","sql":"SELECT id FROM items WHERE id > $1"}))
9545            .send()
9546            .await
9547            .unwrap();
9548        assert_eq!(resp.status(), 200);
9549
9550        // Execute with param 2 → ids 3,4.
9551        let resp = client
9552            .post(format!("http://{addr}/sessions/{session}/execute"))
9553            .json(&json!({"name":"gt","params":[2]}))
9554            .send()
9555            .await
9556            .unwrap();
9557        assert_eq!(resp.status(), 200);
9558        let body = resp.json::<serde_json::Value>().await.unwrap();
9559        let arr = body
9560            .as_array()
9561            .expect("execute returns a JSON array of rows");
9562        assert_eq!(arr.len(), 2, "ids > 2 are {{3,4}}: {body}");
9563
9564        // Re-execute with a different param → reuses the cached plan, fewer rows.
9565        let resp = client
9566            .post(format!("http://{addr}/sessions/{session}/execute"))
9567            .json(&json!({"name":"gt","params":[3]}))
9568            .send()
9569            .await
9570            .unwrap();
9571        let body = resp.json::<serde_json::Value>().await.unwrap();
9572        assert_eq!(body.as_array().unwrap().len(), 1, "ids > 3 is {{4}}");
9573    }
9574
9575    #[tokio::test]
9576    async fn prepared_statement_deallocate_then_execute_fails() {
9577        let (_dir, addr) = setup().await;
9578        let client = reqwest::Client::new();
9579        let session = open_session(&client, &addr).await;
9580        let _ = client
9581            .post(format!("http://{addr}/sessions/{session}/prepare"))
9582            .json(&json!({"name":"p","sql":"SELECT $1"}))
9583            .send()
9584            .await
9585            .unwrap();
9586        let deallocate_query_id = "dadadadadadadadadadadadadadadada";
9587        let resp = client
9588            .delete(format!("http://{addr}/sessions/{session}/statements/p"))
9589            .header("X-MongrelDB-Query-ID", deallocate_query_id)
9590            .header("X-MongrelDB-Timeout-Ms", "10000")
9591            .send()
9592            .await
9593            .unwrap();
9594        assert_eq!(resp.status(), 200);
9595        assert_eq!(
9596            resp.headers()
9597                .get("X-MongrelDB-Query-ID")
9598                .unwrap()
9599                .to_str()
9600                .unwrap(),
9601            deallocate_query_id
9602        );
9603        let status = client
9604            .get(format!("http://{addr}/queries/{deallocate_query_id}"))
9605            .send()
9606            .await
9607            .unwrap()
9608            .json::<serde_json::Value>()
9609            .await
9610            .unwrap();
9611        assert_eq!(status["state"], "completed");
9612        assert_eq!(status["operation"], "DEALLOCATE");
9613        // Execute after deallocate must error.
9614        let resp = client
9615            .post(format!("http://{addr}/sessions/{session}/execute"))
9616            .json(&json!({"name":"p","params":[1]}))
9617            .send()
9618            .await
9619            .unwrap();
9620        assert_ne!(resp.status(), 200, "execute after DEALLOCATE must fail");
9621    }
9622
9623    #[tokio::test]
9624    async fn prepared_statement_deallocate_honors_pre_registration_cancel() {
9625        let (_dir, addr) = setup().await;
9626        let client = reqwest::Client::new();
9627        let session = open_session(&client, &addr).await;
9628        let prepared = client
9629            .post(format!("http://{addr}/sessions/{session}/prepare"))
9630            .json(&json!({"name":"p","sql":"SELECT $1"}))
9631            .send()
9632            .await
9633            .unwrap();
9634        assert_eq!(prepared.status(), StatusCode::OK);
9635
9636        let query_id = "dbdbdbdbdbdbdbdbdbdbdbdbdbdbdbdb";
9637        let cancel = client
9638            .post(format!("http://{addr}/queries/{query_id}/cancel"))
9639            .header("X-Session-ID", &session)
9640            .send()
9641            .await
9642            .unwrap();
9643        assert_eq!(cancel.status(), StatusCode::ACCEPTED);
9644        let deallocate = client
9645            .delete(format!("http://{addr}/sessions/{session}/statements/p"))
9646            .header("X-MongrelDB-Query-ID", query_id)
9647            .send()
9648            .await
9649            .unwrap();
9650        assert_eq!(deallocate.status().as_u16(), 499);
9651        assert_eq!(
9652            deallocate.json::<serde_json::Value>().await.unwrap()["error"]["code"],
9653            "QUERY_CANCELLED"
9654        );
9655
9656        let execute = client
9657            .post(format!("http://{addr}/sessions/{session}/execute"))
9658            .json(&json!({"name":"p","params":[1]}))
9659            .send()
9660            .await
9661            .unwrap();
9662        assert_eq!(execute.status(), StatusCode::OK);
9663    }
9664
9665    #[tokio::test]
9666    async fn prepared_statement_rejects_bad_name() {
9667        let (_dir, addr) = setup().await;
9668        let client = reqwest::Client::new();
9669        let session = open_session(&client, &addr).await;
9670        let resp = client
9671            .post(format!("http://{addr}/sessions/{session}/prepare"))
9672            .json(&json!({"name":"1bad","sql":"SELECT 1"}))
9673            .send()
9674            .await
9675            .unwrap();
9676        assert_eq!(
9677            resp.status(),
9678            400,
9679            "statement name starting with a digit must be rejected"
9680        );
9681    }
9682}