1use 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;
67pub 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#[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
100fn 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 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 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
138fn 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 #[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
213struct 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 storage: ServerStorageRuntime,
238 idem: kit::IdempotencyStore,
239 external_modules: Vec<Arc<dyn ExternalTableModule>>,
240 auth_token: Option<String>,
241 user_auth: bool,
243 metrics: Arc<metrics::Metrics>,
245 audit: Arc<audit::AuditLog>,
247 sessions: Arc<sessions::SessionStore>,
250 ai_semaphore: Arc<tokio::sync::Semaphore>,
252 query_registry: Arc<SqlQueryRegistry>,
254 query_lifecycle: Mutex<()>,
256 pre_cancellations: pre_cancel::PreCancelStore,
258 sql_idempotency: Arc<sql_idempotency::SqlIdempotencyStore>,
260 sql_pages: sql_pages::SqlPageStore,
262 sql_semaphore: Arc<tokio::sync::Semaphore>,
265 sql_page_semaphore: Arc<tokio::sync::Semaphore>,
267 sql_page_default_timeout: std::time::Duration,
268 sql_page_max_timeout: std::time::Duration,
269 max_request_bytes: usize,
273 accepting_sql: Arc<AtomicBool>,
274 cursor_mac_key: CursorMacKey,
276 reloadable: Arc<ReloadableConfig>,
278 drain: Arc<DrainControl>,
280 scheduler: admission::SchedulerAdmission,
284 #[cfg(feature = "cluster")]
288 node_admission: admission::NodeAdmissionController,
289 node_governor: std::sync::Mutex<mongreldb_core::NodeMemoryGovernor>,
291 #[cfg(feature = "cluster")]
294 ai_generations: std::sync::Mutex<mongreldb_core::AiIndexGenerationRegistry>,
295 #[cfg(feature = "cluster")]
297 multi_region: std::sync::Mutex<mongreldb_cluster::multi_region::MultiRegionPolicy>,
298 #[cfg(feature = "cluster")]
301 ops_jobs: std::sync::Mutex<mongreldb_core::OpsJobStore>,
302 resource_groups: mongreldb_core::ResourceGroupRegistry,
305 #[cfg(feature = "cluster")]
309 embedding_providers: mongreldb_core::EmbeddingProviderRegistry,
310}
311
312impl AppState {
313 pub(crate) fn storage(&self) -> &ServerStorageRuntime {
315 &self.storage
316 }
317
318 #[cfg(feature = "cluster")]
320 pub(crate) fn node_admission(&self) -> &admission::NodeAdmissionController {
321 &self.node_admission
322 }
323
324 #[cfg(feature = "cluster")]
326 pub(crate) fn resource_groups(&self) -> &mongreldb_core::ResourceGroupRegistry {
327 &self.resource_groups
328 }
329
330 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 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 #[allow(dead_code)] fn db_arc(&self) -> Arc<Database> {
355 Arc::clone(self.db())
356 }
357
358 #[cfg(feature = "cluster")]
360 pub(crate) fn cluster_runtime(&self) -> Option<&cluster_runtime::ClusterRuntimeHandle> {
361 self.storage.cluster_handle()
362 }
363
364 #[cfg(feature = "cluster")]
367 pub(crate) fn is_cluster_mode(&self) -> bool {
368 self.storage.is_cluster()
369 }
370
371 #[cfg(not(feature = "cluster"))]
374 pub(crate) fn is_cluster_mode(&self) -> bool {
375 false
376 }
377}
378
379fn 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
394pub(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
405struct 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
433struct 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
453struct ReloadableConfig {
464 slow_query_threshold: ReloadableDuration,
466 sql_default_timeout: ReloadableDuration,
469 sql_max_timeout: ReloadableDuration,
471 sql_cancel_grace: ReloadableDuration,
474 sql_max_output_rows: ReloadableUsize,
476 sql_max_output_bytes: ReloadableUsize,
478}
479
480#[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#[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#[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
528fn 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 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
579fn 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#[derive(Default)]
616struct DrainControl {
617 lock: tokio::sync::Mutex<()>,
619 record: std::sync::Mutex<DrainRecord>,
620}
621
622#[derive(Clone, Default, Serialize)]
623struct DrainRecord {
624 initiated: bool,
626 completed: bool,
628 deadline_ms: u64,
630 lifecycle: String,
632 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#[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 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 #[cfg(feature = "native-rpc")]
678 node_admission: admission::NodeAdmissionController,
679 #[cfg(feature = "native-rpc")]
680 sql_priority: u8,
681}
682
683impl ServerControl {
684 pub fn query_registry(&self) -> Arc<SqlQueryRegistry> {
686 Arc::clone(&self.query_registry)
687 }
688
689 pub fn storage(&self) -> &ServerStorageRuntime {
691 &self.storage
692 }
693
694 #[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 #[cfg(feature = "cluster")]
746 pub fn cluster_runtime(&self) -> Option<&cluster_runtime::ClusterRuntimeHandle> {
747 self.storage.cluster_handle()
748 }
749
750 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 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
796pub 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
806pub 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
829pub 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#[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#[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#[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 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
933pub 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 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 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 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 #[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 .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 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 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 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
1227async 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 let mut attempted = String::new();
1248 let mut fail_reason = "no credentials provided".to_string();
1249
1250 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 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
1309fn 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
1334fn 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
1360async 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
1390async 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 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
1590async 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
1769pub 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
1891async 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
1905async 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
1930async 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
1948fn 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 if state.is_cluster_mode() {
1963 if state.user_auth && principal.as_ref().is_none_or(|p| !p.is_admin) {
1964 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
2001fn 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
2028fn 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
2056async 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
2069async 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 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 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 state.accepting_sql.store(false, Ordering::Release);
2155 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 state.sessions.close_all();
2166 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 (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
2211fn default_drain_deadline_ms() -> u64 {
2214 positive_env_u64("MONGRELDB_DRAIN_DEADLINE_MS", 30_000)
2215}
2216
2217async 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 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
2301fn 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
2309fn 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
2490fn default_max_request_bytes() -> usize {
2494 positive_env_usize("MONGRELDB_MAX_REQUEST_BYTES", 2 * 1024 * 1024)
2495}
2496
2497fn 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
2579fn 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
2598async 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
2640async 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
2673async 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 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 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 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 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 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 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
3486fn 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
3501fn 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 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 _ => Err("prepared-statement parameters must be scalar (null/bool/number/string)".into()),
3532 }
3533}
3534
3535fn 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 #[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
3582async 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 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
3699async 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 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 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 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 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 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 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
3969async 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
4057async 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
4090async 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
4120async 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 (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 (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
4539fn 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
4570pub(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 #[serde(default)]
4661 format: Option<String>,
4662 #[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
4707fn 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 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 MongrelQueryError::NoSqlTransaction => ErrorCategory::TransactionAborted,
4734 MongrelQueryError::SavepointNotFound { .. } => ErrorCategory::StaleMetadata,
4737 MongrelQueryError::QueryRegistryFull | MongrelQueryError::ResultLimitExceeded { .. } => {
4738 ErrorCategory::ResourceExhausted
4739 }
4740 MongrelQueryError::QueryIdConflict { .. }
4744 | MongrelQueryError::InvalidQueryState(_)
4745 | MongrelQueryError::Arrow(_)
4746 | MongrelQueryError::DataFusion(_) => ErrorCategory::ClusterVersionMismatch,
4747 MongrelQueryError::Schema(_) => ErrorCategory::SchemaVersionMismatch,
4748 _ => 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 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 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 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
5812async 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 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 #[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
5889fn 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
6027fn 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 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 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 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 "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 "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 let category = {
7026 use mongreldb_types::errors::ErrorCategory;
7027 match code {
7028 "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 #[cfg(feature = "cluster")]
7093 if let Some(response) = cluster_admin::try_admin_sql(&state, &principal, &req.sql).await {
7094 return response;
7095 }
7096 #[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 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 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 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 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
7338fn 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
7357async 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 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 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 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 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 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#[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 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 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 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 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 ®istry.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 ®istry.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(®istry.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(®istry.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 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 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 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 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 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 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 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 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 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 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 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 #[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 assert!(bytes.starts_with(&0xFFFFFFFFu32.to_le_bytes()));
9062
9063 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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}