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