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