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