use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use axum::extract::{Path, State};
use axum::http::header;
use axum::http::StatusCode;
use axum::response::{IntoResponse, Response};
use axum::routing::{get, post};
use axum::Json;
use mongreldb_core::schema::{Schema, TypeId};
use mongreldb_core::{CancellationReason, Database, Value};
use mongreldb_query::{
CancelOutcome, CompactFinishedQuery, ExternalTableModule, ManagedQueryBatches, MongrelSession,
QueryId, RegisteredQueryGuard, RegisteredSqlQuery, SqlQueryOptions, SqlQueryPhase,
SqlQueryRegistry, SqlStreamCompletion,
};
use serde::{Deserialize, Serialize};
use serde_json::json;
use sha2::Digest;
use zeroize::Zeroizing;
mod audit;
mod kit;
mod metrics;
mod pre_cancel;
mod procedure;
mod sessions;
mod sql_idempotency;
mod sql_pages;
mod trigger;
pub use sessions::{spawn_session_reaper, SessionStore};
fn client_closed_request_status() -> StatusCode {
StatusCode::from_u16(499).unwrap_or(StatusCode::BAD_REQUEST)
}
fn cancellation_checkpoint_error(query: &RegisteredSqlQuery) -> mongreldb_query::MongrelQueryError {
query.checkpoint().err().unwrap_or_else(|| {
mongreldb_query::MongrelQueryError::InvalidQueryState(
"cancellation notification observed without a terminal checkpoint".into(),
)
})
}
fn status_for_error(e: &mongreldb_core::MongrelError) -> StatusCode {
use mongreldb_core::MongrelError;
match e {
MongrelError::AuthRequired | MongrelError::InvalidCredentials { .. } => {
StatusCode::UNAUTHORIZED
}
MongrelError::AuthNotRequired => StatusCode::BAD_REQUEST,
MongrelError::PermissionDenied { .. } => StatusCode::FORBIDDEN,
MongrelError::InvalidArgument(_) => StatusCode::CONFLICT,
MongrelError::Conflict(_) => StatusCode::CONFLICT,
MongrelError::ReadOnlyReplica => StatusCode::CONFLICT,
MongrelError::NotFound(_) => StatusCode::NOT_FOUND,
MongrelError::DeadlineExceeded => StatusCode::GATEWAY_TIMEOUT,
MongrelError::WorkBudgetExceeded => StatusCode::TOO_MANY_REQUESTS,
MongrelError::Cancelled => client_closed_request_status(),
MongrelError::CursorStale(_) => StatusCode::CONFLICT,
MongrelError::CursorExpired => StatusCode::GONE,
_ => StatusCode::INTERNAL_SERVER_ERROR,
}
}
fn status_for_query_error(e: &mongreldb_query::MongrelQueryError) -> StatusCode {
use mongreldb_query::MongrelQueryError;
match e {
MongrelQueryError::Core(core) => status_for_error(core),
MongrelQueryError::DeadlineExceeded { .. } => StatusCode::GATEWAY_TIMEOUT,
MongrelQueryError::QueryCancelled { .. } => client_closed_request_status(),
MongrelQueryError::QueryIdConflict { .. } => StatusCode::CONFLICT,
MongrelQueryError::QueryRegistryFull => StatusCode::SERVICE_UNAVAILABLE,
MongrelQueryError::ResultLimitExceeded { .. } => StatusCode::PAYLOAD_TOO_LARGE,
MongrelQueryError::TransactionAborted => StatusCode::CONFLICT,
MongrelQueryError::NoSqlTransaction | MongrelQueryError::SavepointNotFound { .. } => {
StatusCode::CONFLICT
}
MongrelQueryError::CommitOutcome { .. } => StatusCode::CONFLICT,
MongrelQueryError::OutcomeUnknown { .. } => StatusCode::CONFLICT,
_ => StatusCode::INTERNAL_SERVER_ERROR,
}
}
struct OptionalPrincipal(Option<mongreldb_core::Principal>);
impl<S> axum::extract::FromRequestParts<S> for OptionalPrincipal
where
S: Send + Sync,
{
type Rejection = std::convert::Infallible;
async fn from_request_parts(
parts: &mut axum::http::request::Parts,
_state: &S,
) -> Result<Self, Self::Rejection> {
Ok(OptionalPrincipal(
parts.extensions.get::<mongreldb_core::Principal>().cloned(),
))
}
}
struct AppState {
db: Arc<Database>,
idem: kit::IdempotencyStore,
external_modules: Vec<Arc<dyn ExternalTableModule>>,
auth_token: Option<String>,
user_auth: bool,
metrics: Arc<metrics::Metrics>,
slow_query_threshold: std::time::Duration,
audit: Arc<audit::AuditLog>,
sessions: Arc<sessions::SessionStore>,
ai_semaphore: Arc<tokio::sync::Semaphore>,
query_registry: Arc<SqlQueryRegistry>,
query_lifecycle: Mutex<()>,
pre_cancellations: pre_cancel::PreCancelStore,
sql_idempotency: Arc<sql_idempotency::SqlIdempotencyStore>,
sql_pages: sql_pages::SqlPageStore,
sql_semaphore: Arc<tokio::sync::Semaphore>,
sql_page_semaphore: Arc<tokio::sync::Semaphore>,
sql_page_default_timeout: std::time::Duration,
sql_page_max_timeout: std::time::Duration,
sql_default_timeout: std::time::Duration,
sql_max_timeout: std::time::Duration,
sql_cancel_grace: std::time::Duration,
sql_max_output_bytes: usize,
sql_max_output_rows: usize,
accepting_sql: Arc<AtomicBool>,
cursor_mac_key: CursorMacKey,
}
#[derive(Default)]
struct CursorMacKey {
key: Mutex<Option<[u8; 32]>>,
}
impl CursorMacKey {
fn get(&self) -> mongreldb_core::Result<[u8; 32]> {
let mut key = match self.key.lock() {
Ok(key) => key,
Err(poisoned) => poisoned.into_inner(),
};
if let Some(key) = *key {
return Ok(key);
}
let mut generated = [0u8; 32];
mongreldb_core::encryption::fill_random(&mut generated)?;
*key = Some(generated);
Ok(generated)
}
}
#[derive(Clone)]
pub struct ServerControl {
query_registry: Arc<SqlQueryRegistry>,
sessions: Arc<sessions::SessionStore>,
accepting_sql: Arc<AtomicBool>,
cancel_grace: std::time::Duration,
metrics: Arc<metrics::Metrics>,
}
impl ServerControl {
pub async fn shutdown(&self) -> usize {
self.accepting_sql.store(false, Ordering::Release);
self.query_registry
.cancel_all(CancellationReason::ServerShutdown);
let deadline = tokio::time::Instant::now() + self.cancel_grace;
while self.query_registry.active_count() > 0 && tokio::time::Instant::now() < deadline {
tokio::time::sleep(std::time::Duration::from_millis(5)).await;
}
self.sessions.close_all();
let stuck_queries = self.query_registry.active_statuses();
for status in &stuck_queries {
eprintln!(
"[sql-cancel-stuck] query_id={} phase={}",
status.query_id,
query_phase_name(status.phase)
);
}
let stuck = stuck_queries.len();
self.metrics.add_sql_stuck_after_cancel(stuck);
stuck
}
}
pub fn build_app(db: Arc<Database>) -> axum::Router {
build_app_with_config(
db,
std::iter::empty::<Arc<dyn ExternalTableModule>>(),
None,
None,
)
}
pub fn build_app_with_external_modules(
db: Arc<Database>,
external_modules: impl IntoIterator<Item = Arc<dyn ExternalTableModule>>,
) -> axum::Router {
build_app_with_config(db, external_modules, None, None)
}
pub fn build_app_with_config(
db: Arc<Database>,
external_modules: impl IntoIterator<Item = Arc<dyn ExternalTableModule>>,
auth_token: Option<String>,
max_connections: Option<usize>,
) -> axum::Router {
build_app_full(db, external_modules, auth_token, max_connections, false)
}
pub fn build_app_full(
db: Arc<Database>,
external_modules: impl IntoIterator<Item = Arc<dyn ExternalTableModule>>,
auth_token: Option<String>,
max_connections: Option<usize>,
user_auth: bool,
) -> axum::Router {
let sessions = Arc::new(sessions::SessionStore::new(
default_max_sessions(),
default_session_idle_timeout(),
));
build_app_with_sessions(
db,
external_modules,
auth_token,
max_connections,
user_auth,
sessions,
)
}
pub fn build_app_with_sessions(
db: Arc<Database>,
external_modules: impl IntoIterator<Item = Arc<dyn ExternalTableModule>>,
auth_token: Option<String>,
max_connections: Option<usize>,
user_auth: bool,
sessions: Arc<sessions::SessionStore>,
) -> axum::Router {
build_app_with_sessions_and_control(
db,
external_modules,
auth_token,
max_connections,
user_auth,
sessions,
)
.0
}
pub fn build_app_with_sessions_and_control(
db: Arc<Database>,
external_modules: impl IntoIterator<Item = Arc<dyn ExternalTableModule>>,
auth_token: Option<String>,
max_connections: Option<usize>,
user_auth: bool,
sessions: Arc<sessions::SessionStore>,
) -> (axum::Router, ServerControl) {
db.set_replication_wal_retention_segments(default_replication_wal_segments());
if let Err(error) = db.set_history_retention_epochs(default_history_retention_epochs()) {
eprintln!("[history] failed to configure retention: {error}");
}
let max_active_queries = default_sql_max_active_queries();
let query_registry = Arc::new(SqlQueryRegistry::new_with_limits(
max_active_queries,
default_sql_finished_detail_max_entries(),
default_sql_finished_detail_max_bytes(),
default_sql_finished_compact_max_entries(),
default_sql_finished_compact_max_bytes(),
default_sql_finished_query_ttl(),
));
let accepting_sql = Arc::new(AtomicBool::new(true));
let sql_cancel_grace = default_sql_cancel_grace();
let sql_max_timeout = default_sql_max_timeout();
let sql_default_timeout = default_sql_default_timeout().min(sql_max_timeout);
let metrics = Arc::new(metrics::Metrics::default());
let (idempotency_root, idempotency_integrity) =
match sql_idempotency::IdempotencyIntegrity::for_database(&db) {
Ok((root, integrity)) => (root, Some(integrity)),
Err(error) => {
eprintln!("[idempotency] durable integrity key unavailable: {error}");
(db.durable_root(), None)
}
};
let sql_idempotency = Arc::new(sql_idempotency::SqlIdempotencyStore::new_with_integrity(
Arc::clone(&idempotency_root),
idempotency_integrity.clone(),
default_sql_idempotency_ttl(),
default_sql_idempotency_max_entries(),
));
let server_control = ServerControl {
query_registry: Arc::clone(&query_registry),
sessions: Arc::clone(&sessions),
accepting_sql: Arc::clone(&accepting_sql),
cancel_grace: sql_cancel_grace,
metrics: Arc::clone(&metrics),
};
let state = Arc::new(AppState {
idem: kit::IdempotencyStore::new_with_integrity(
idempotency_root,
idempotency_integrity,
default_sql_idempotency_ttl(),
default_sql_idempotency_max_entries(),
),
db,
external_modules: external_modules.into_iter().collect(),
auth_token,
user_auth,
metrics,
slow_query_threshold: metrics::slow_query_threshold(),
audit: Arc::new(audit::AuditLog::new(8192)),
sessions,
ai_semaphore: Arc::new(tokio::sync::Semaphore::new(default_ai_max_concurrent())),
query_registry,
query_lifecycle: Mutex::new(()),
pre_cancellations: pre_cancel::PreCancelStore::new(
default_sql_pre_cancel_ttl(),
default_sql_pre_cancel_max_entries(),
default_sql_pre_cancel_max_bytes(),
default_sql_pre_cancel_max_entries_per_owner(),
default_sql_pre_cancel_rate_window(),
default_sql_pre_cancel_rate_per_owner(),
),
sql_idempotency,
sql_pages: sql_pages::SqlPageStore::new(
default_sql_page_ttl(),
default_sql_page_max_entries(),
default_sql_page_max_bytes(),
default_sql_page_max_entries_per_owner(),
),
sql_semaphore: Arc::new(tokio::sync::Semaphore::new(default_sql_max_concurrent())),
sql_page_semaphore: Arc::new(tokio::sync::Semaphore::new(
default_sql_page_max_concurrent(),
)),
sql_page_default_timeout: default_sql_page_default_timeout()
.min(default_sql_page_max_timeout()),
sql_page_max_timeout: default_sql_page_max_timeout(),
sql_default_timeout,
sql_max_timeout,
sql_cancel_grace,
sql_max_output_bytes: default_sql_max_output_bytes(),
sql_max_output_rows: default_sql_max_output_rows(),
accepting_sql,
cursor_mac_key: CursorMacKey::default(),
});
let router = axum::Router::new()
.route("/health", get(health))
.route("/build-info", get(build_info))
.route("/capabilities", get(capabilities))
.route(
"/history/retention",
get(history_retention).put(set_history_retention),
)
.route("/metrics", get(metrics_handler))
.route("/audit", get(audit_handler))
.route("/tables", get(list_tables).post(create_table))
.route("/tables/{name}", axum::routing::delete(drop_table))
.route("/tables/{name}/put", post(put_row))
.route("/tables/{name}/count", get(count))
.route("/tables/{name}/commit", post(commit))
.route("/sql", post(sql))
.route("/sql/continue", post(continue_sql_page))
.route("/queries/{query_id}", get(query_status))
.route("/queries/{query_id}/cancel", post(cancel_query))
.route("/txn", post(txn))
.route("/sessions", post(create_session))
.route("/sessions/{id}", axum::routing::delete(close_session))
.route("/sessions/{id}/prepare", post(prepare_statement))
.route("/sessions/{id}/execute", post(execute_statement))
.route(
"/sessions/{id}/statements/{name}",
axum::routing::delete(deallocate_statement),
)
.route("/procedures", get(procedure::list).post(procedure::create))
.route(
"/procedures/{name}",
get(procedure::describe)
.put(procedure::replace)
.delete(procedure::drop_procedure),
)
.route("/procedures/{name}/call", post(procedure::call))
.route("/triggers", get(trigger::list).post(trigger::create))
.route(
"/triggers/{name}",
get(trigger::describe)
.put(trigger::replace)
.delete(trigger::drop_trigger),
)
.route("/kit/schema", get(kit::schema_all))
.route("/kit/schema/{table}", get(kit::schema_one))
.route("/kit/txn", post(kit::kit_txn))
.route("/kit/query", post(kit::kit_query))
.route("/kit/retrieve", post(kit::kit_retrieve))
.route("/kit/ann_rerank", post(kit::kit_ann_rerank))
.route("/kit/ai/metrics", get(kit::kit_ai_metrics))
.route("/kit/set_similarity", post(kit::kit_set_similarity))
.route("/kit/search", post(kit::kit_search))
.route("/kit/create_table", post(kit::kit_create_table))
.route("/kit/procedures/{name}/call", post(procedure::kit_call))
.route("/compact", post(compact_all))
.route("/tables/{name}/compact", post(compact_table))
.route("/wal/stream", get(wal_stream))
.route("/replication/snapshot", get(replication_snapshot))
.route("/events", get(events_stream))
.with_state(state.clone());
let router = if state.auth_token.is_some() || state.user_auth || state.db.require_auth_enabled()
{
router.layer(axum::middleware::from_fn_with_state(
state.clone(),
auth_middleware,
))
} else {
router
};
let router = if let Some(max) = max_connections {
router.layer(tower::limit::ConcurrencyLimitLayer::new(max))
} else {
router
};
(router, server_control)
}
async fn auth_middleware(
axum::extract::State(state): axum::extract::State<Arc<AppState>>,
mut req: axum::extract::Request,
next: axum::middleware::Next,
) -> Result<axum::response::Response, axum::http::StatusCode> {
let header = req
.headers()
.get("authorization")
.and_then(|v| v.to_str().ok())
.unwrap_or("");
let mut attempted = String::new();
let mut fail_reason = "no credentials provided".to_string();
if let Some(token) = &state.auth_token {
if let Some(provided) = header.strip_prefix("Bearer ") {
attempted = "token".to_string();
if provided == token {
state
.audit
.record("token", "login.ok", "bearer token accepted");
return Ok(next.run(req).await);
}
fail_reason = "invalid bearer token".to_string();
}
}
if state.user_auth {
if let Some(encoded) = header.strip_prefix("Basic ") {
if let Ok(decoded) = base64_decode(encoded) {
let decoded = Zeroizing::new(decoded);
if let Ok(creds) = std::str::from_utf8(&decoded) {
if let Some((username, password)) = creds.split_once(':') {
attempted = username.to_string();
if let Ok(Some(principal)) =
state.db.authenticate_principal(username, password)
{
let username = principal.username.clone();
drop(decoded);
state
.audit
.record(&username, "login.ok", "basic credentials accepted");
req.extensions_mut().insert(principal);
return Ok(next.run(req).await);
}
fail_reason = "invalid basic credentials".to_string();
} else {
fail_reason = "malformed basic credentials (no ':')".to_string();
}
} else {
fail_reason = "malformed basic credentials (non-utf8)".to_string();
}
} else {
fail_reason = "malformed basic credentials (bad base64)".to_string();
}
}
}
let who = if attempted.is_empty() {
"anonymous"
} else {
attempted.as_str()
};
state.audit.record(who, "login.fail", fail_reason);
Err(axum::http::StatusCode::UNAUTHORIZED)
}
fn base64_decode(input: &str) -> Result<Vec<u8>, ()> {
const TABLE: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
let input: Vec<u8> = input
.bytes()
.filter(|&b| b != b'\n' && b != b'\r' && b != b' ')
.collect();
let mut out = Vec::with_capacity(input.len() * 3 / 4);
let mut buf = 0u32;
let mut bits = 0u32;
for &b in &input {
if b == b'=' {
break;
}
let val = TABLE.iter().position(|&t| t == b).ok_or(())? as u32;
buf = (buf << 6) | val;
bits += 6;
if bits >= 8 {
bits -= 8;
out.push((buf >> bits) as u8);
}
}
Ok(out)
}
async fn wal_stream(
axum::extract::State(state): axum::extract::State<Arc<AppState>>,
OptionalPrincipal(principal): OptionalPrincipal,
axum::extract::Query(params): axum::extract::Query<WalStreamParams>,
) -> Result<Response, StatusCode> {
state
.db
.require_for(
request_principal(&state, &principal).as_ref(),
&mongreldb_core::Permission::Admin,
)
.map_err(|error| status_for_error(&error))?;
let since = params.since.unwrap_or(0);
let db = Arc::clone(&state.db);
let batch = tokio::task::spawn_blocking(move || db.replication_batch_since(since))
.await
.map_err(|_e| StatusCode::INTERNAL_SERVER_ERROR)?;
let batch = batch.map_err(|e| {
eprintln!("wal_stream error: {e}");
StatusCode::INTERNAL_SERVER_ERROR
})?;
if batch.requires_snapshot {
let mut response = (
StatusCode::CONFLICT,
"replication snapshot required: WAL retention gap or spilled run",
)
.into_response();
response.headers_mut().insert(
"x-mongreldb-replication-status",
"snapshot-required"
.parse()
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?,
);
set_replication_headers(
&mut response,
&batch.source_id,
batch.from_epoch,
batch.current_epoch,
batch.earliest_epoch,
batch.commit_count,
&batch.records_sha256,
)?;
return Ok(response);
}
let mut body = String::new();
for record in &batch.records {
let json = serde_json::to_string(record).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
body.push_str(&json);
body.push('\n');
}
let mut response = (
[
(header::CONTENT_TYPE, "application/x-ndjson".to_string()),
(header::CACHE_CONTROL, "no-cache".to_string()),
],
body,
)
.into_response();
set_replication_headers(
&mut response,
&batch.source_id,
batch.from_epoch,
batch.current_epoch,
batch.earliest_epoch,
batch.commit_count,
&batch.records_sha256,
)?;
Ok(response)
}
async fn replication_snapshot(
axum::extract::State(state): axum::extract::State<Arc<AppState>>,
OptionalPrincipal(principal): OptionalPrincipal,
) -> Response {
if let Err(error) = state.db.require_for(
request_principal(&state, &principal).as_ref(),
&mongreldb_core::Permission::Admin,
) {
return (status_for_error(&error), error.to_string()).into_response();
}
let db = Arc::clone(&state.db);
let snapshot = match tokio::task::spawn_blocking(move || db.replication_snapshot()).await {
Ok(Ok(snapshot)) => snapshot,
Ok(Err(error)) => {
return (StatusCode::INTERNAL_SERVER_ERROR, error.to_string()).into_response()
}
Err(error) => {
return (StatusCode::INTERNAL_SERVER_ERROR, error.to_string()).into_response()
}
};
let epoch = snapshot.epoch();
match snapshot.encode() {
Ok(bytes) => {
let mut response =
([(header::CONTENT_TYPE, "application/octet-stream")], bytes).into_response();
let Ok(value) = epoch.to_string().parse() else {
return (
StatusCode::INTERNAL_SERVER_ERROR,
"invalid replication epoch response header",
)
.into_response();
};
response
.headers_mut()
.insert("x-mongreldb-current-epoch", value);
let source_id = snapshot
.source_id()
.iter()
.map(|byte| format!("{byte:02x}"))
.collect::<String>();
let Ok(source_id) = source_id.parse() else {
return (
StatusCode::INTERNAL_SERVER_ERROR,
"invalid replication source response header",
)
.into_response();
};
response
.headers_mut()
.insert("x-mongreldb-source-id", source_id);
response
}
Err(error) => (StatusCode::INTERNAL_SERVER_ERROR, error.to_string()).into_response(),
}
}
fn set_replication_headers(
response: &mut Response,
source_id: &[u8; 32],
from: u64,
current: u64,
earliest: Option<u64>,
commit_count: u64,
records_sha256: &[u8; 32],
) -> Result<(), StatusCode> {
let source_id = source_id
.iter()
.map(|byte| format!("{byte:02x}"))
.collect::<String>();
response.headers_mut().insert(
"x-mongreldb-source-id",
source_id
.parse()
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?,
);
response.headers_mut().insert(
"x-mongreldb-from-epoch",
from.to_string()
.parse()
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?,
);
response.headers_mut().insert(
"x-mongreldb-current-epoch",
current
.to_string()
.parse()
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?,
);
response.headers_mut().insert(
"x-mongreldb-commit-count",
commit_count
.to_string()
.parse()
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?,
);
let digest = records_sha256
.iter()
.map(|byte| format!("{byte:02x}"))
.collect::<String>();
response.headers_mut().insert(
"x-mongreldb-records-sha256",
digest
.parse()
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?,
);
if let Some(earliest) = earliest {
response.headers_mut().insert(
"x-mongreldb-earliest-epoch",
earliest
.to_string()
.parse()
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?,
);
}
Ok(())
}
#[derive(serde::Deserialize)]
struct WalStreamParams {
since: Option<u64>,
}
async fn events_stream(
axum::extract::State(state): axum::extract::State<Arc<AppState>>,
OptionalPrincipal(principal): OptionalPrincipal,
headers: axum::http::HeaderMap,
) -> Result<Response, StatusCode> {
use axum::response::sse::{Event, KeepAlive, Sse};
use futures::Stream;
use std::collections::VecDeque;
use std::convert::Infallible;
state
.db
.require_for(
request_principal(&state, &principal).as_ref(),
&mongreldb_core::Permission::Admin,
)
.map_err(|error| status_for_error(&error))?;
struct State {
db: Arc<Database>,
receiver: tokio::sync::broadcast::Receiver<mongreldb_core::ChangeEvent>,
change_wake: tokio::sync::broadcast::Receiver<()>,
interval: tokio::time::Interval,
pending: VecDeque<mongreldb_core::ChangeEvent>,
last_id: Option<String>,
poll_now: bool,
done: bool,
}
fn event(change: mongreldb_core::ChangeEvent) -> Event {
let id = change.id.clone();
let kind = if change.op == "notify" {
"notify"
} else {
"change"
};
let mut event = Event::default().event(kind).data(
serde_json::to_string(&change)
.unwrap_or_else(|error| format!(r#"{{"error":"{error}"}}"#)),
);
if let Some(id) = id {
event = event.id(id);
}
event
}
let last_id = match headers.get("last-event-id") {
Some(value) => Some(
value
.to_str()
.map_err(|_| StatusCode::BAD_REQUEST)?
.to_owned(),
),
None => None,
};
let receiver = state.db.subscribe_changes();
let change_wake = state.db.subscribe_change_commits();
let db = Arc::clone(&state.db);
let resume = last_id.clone();
let initial = tokio::task::spawn_blocking(move || db.change_events_since(resume.as_deref()))
.await
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
.map_err(|error| match error {
mongreldb_core::MongrelError::InvalidArgument(_) => StatusCode::BAD_REQUEST,
_ => StatusCode::INTERNAL_SERVER_ERROR,
})?;
if initial.gap {
return Ok((
StatusCode::CONFLICT,
Json(json!({
"error": "cdc retention gap",
"earliest_epoch": initial.earliest_epoch,
"current_epoch": initial.current_epoch,
})),
)
.into_response());
}
let stream: std::pin::Pin<Box<dyn Stream<Item = Result<Event, Infallible>> + Send>> = Box::pin(
futures::stream::unfold(
State {
db: Arc::clone(&state.db),
receiver,
change_wake,
interval: tokio::time::interval(std::time::Duration::from_millis(250)),
pending: initial.events.into(),
last_id,
poll_now: false,
done: false,
},
|mut stream| async move {
if stream.done {
return None;
}
loop {
if stream.poll_now {
stream.poll_now = false;
let db = Arc::clone(&stream.db);
let last_id = stream.last_id.clone();
match tokio::task::spawn_blocking(move || {
db.change_events_since(last_id.as_deref())
})
.await
{
Ok(Ok(batch)) if batch.gap => {
stream.done = true;
let gap = Event::default().event("gap").data(
json!({
"error": "cdc retention gap",
"earliest_epoch": batch.earliest_epoch,
"current_epoch": batch.current_epoch,
})
.to_string(),
);
return Some((Ok(gap), stream));
}
Ok(Ok(batch)) => stream.pending.extend(batch.events),
Ok(Err(error)) => {
stream.done = true;
return Some((
Ok(Event::default().event("error").data(error.to_string())),
stream,
));
}
Err(error) => {
stream.done = true;
return Some((
Ok(Event::default().event("error").data(error.to_string())),
stream,
));
}
}
}
if let Some(change) = stream.pending.pop_front() {
if let Some(id) = &change.id {
stream.last_id = Some(id.clone());
}
return Some((Ok(event(change)), stream));
}
tokio::select! {
received = stream.receiver.recv() => {
match received {
Ok(change) => return Some((Ok(event(change)), stream)),
Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {},
Err(tokio::sync::broadcast::error::RecvError::Closed) => {
return None;
}
}
}
received = stream.change_wake.recv() => {
match received {
Ok(()) | Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {
stream.poll_now = true;
}
Err(tokio::sync::broadcast::error::RecvError::Closed) => {}
}
}
_ = stream.interval.tick() => {
stream.poll_now = true;
}
}
}
},
),
);
Ok(Sse::new(stream)
.keep_alive(
KeepAlive::new()
.interval(std::time::Duration::from_secs(15))
.text("keep-alive"),
)
.into_response())
}
pub fn spawn_auto_compactor(db: Arc<Database>) {
if let Err(error) = std::thread::Builder::new()
.name("mongreldb-auto-compact".into())
.spawn(move || loop {
std::thread::sleep(std::time::Duration::from_secs(30));
for name in db.table_names() {
let Ok(handle) = db.table(&name) else {
continue;
};
let mut t = handle.lock();
let before = t.run_count();
match t.maybe_compact() {
Ok(true) => {
eprintln!(
"[auto-compact] {name}: {} runs -> {}",
before,
t.run_count()
);
}
Ok(false) => {}
Err(e) => {
eprintln!("[auto-compact] {name}: compaction failed: {e}");
}
}
}
})
{
eprintln!("[auto-compact] failed to start background thread: {error}");
}
}
async fn health() -> StatusCode {
StatusCode::OK
}
#[derive(Debug, Serialize)]
struct SqlCancellationCapabilities {
version: u8,
client_query_ids: bool,
cancel_endpoint: bool,
query_status: bool,
pre_registration_cancel: bool,
stream_disconnect_cancels: bool,
}
#[derive(Debug, Serialize)]
struct SqlIdempotencyCapabilities {
version: u8,
durable_pre_execution_intent: bool,
replay_committed_receipt: bool,
indeterminate_never_reexecutes: bool,
}
#[derive(Debug, Serialize)]
struct SqlPaginationCapabilities {
version: u8,
continuation_endpoint: &'static str,
retained_snapshot: bool,
projection_required: bool,
byte_and_token_hints: bool,
}
#[derive(Debug, Serialize)]
struct CapabilitiesResponse {
sql_cancellation: SqlCancellationCapabilities,
sql_idempotency: SqlIdempotencyCapabilities,
sql_pagination: SqlPaginationCapabilities,
}
async fn capabilities() -> Json<CapabilitiesResponse> {
Json(CapabilitiesResponse {
sql_cancellation: SqlCancellationCapabilities {
version: 2,
client_query_ids: true,
cancel_endpoint: true,
query_status: true,
pre_registration_cancel: true,
stream_disconnect_cancels: true,
},
sql_idempotency: SqlIdempotencyCapabilities {
version: 1,
durable_pre_execution_intent: true,
replay_committed_receipt: true,
indeterminate_never_reexecutes: true,
},
sql_pagination: SqlPaginationCapabilities {
version: 1,
continuation_endpoint: "/sql/continue",
retained_snapshot: true,
projection_required: true,
byte_and_token_hints: true,
},
})
}
async fn build_info() -> Json<mongreldb_query::BuildInfo> {
Json(mongreldb_query::build_info())
}
#[derive(Debug, Deserialize)]
struct HistoryRetentionRequest {
#[serde(default)]
history_retention_epochs: serde_json::Value,
}
#[derive(Debug, Serialize)]
struct HistoryRetentionResponse {
history_retention_epochs: u64,
earliest_retained_epoch: u64,
}
fn history_retention_response(db: &Database) -> HistoryRetentionResponse {
HistoryRetentionResponse {
history_retention_epochs: db.history_retention_epochs(),
earliest_retained_epoch: db.earliest_retained_epoch().0,
}
}
async fn history_retention(
State(state): State<Arc<AppState>>,
OptionalPrincipal(principal): OptionalPrincipal,
) -> Response {
if let Err(error) = state.db.require_for(
request_principal(&state, &principal).as_ref(),
&mongreldb_core::Permission::Admin,
) {
return (status_for_error(&error), error.to_string()).into_response();
}
Json(history_retention_response(&state.db)).into_response()
}
async fn set_history_retention(
State(state): State<Arc<AppState>>,
OptionalPrincipal(principal): OptionalPrincipal,
Json(request): Json<HistoryRetentionRequest>,
) -> Response {
if let Err(error) = state.db.require_for(
request_principal(&state, &principal).as_ref(),
&mongreldb_core::Permission::Admin,
) {
return (status_for_error(&error), error.to_string()).into_response();
}
let Some(epochs) = request.history_retention_epochs.as_u64() else {
return (
StatusCode::BAD_REQUEST,
Json(json!({"error": "history_retention_epochs must be a u64"})),
)
.into_response();
};
match state.db.set_history_retention_epochs(epochs) {
Ok(()) => Json(history_retention_response(&state.db)).into_response(),
Err(error) => (status_for_error(&error), error.to_string()).into_response(),
}
}
async fn audit_handler(
State(state): State<Arc<AppState>>,
OptionalPrincipal(principal): OptionalPrincipal,
) -> Response {
if let Err(error) = state.db.require_for(
request_principal(&state, &principal).as_ref(),
&mongreldb_core::Permission::Admin,
) {
return (status_for_error(&error), error.to_string()).into_response();
}
let recent = state.audit.recent();
Json(recent).into_response()
}
fn default_max_sessions() -> usize {
std::env::var("MONGRELBL_MAX_SESSIONS")
.ok()
.and_then(|v| v.parse().ok())
.unwrap_or(256)
}
fn default_session_idle_timeout() -> std::time::Duration {
std::time::Duration::from_secs(
std::env::var("MONGRELBL_SESSION_IDLE_TIMEOUT_SECS")
.ok()
.and_then(|v| v.parse().ok())
.unwrap_or(300),
)
}
fn default_replication_wal_segments() -> usize {
let replication = std::env::var("MONGRELDB_REPLICATION_WAL_SEGMENTS")
.ok()
.and_then(|value| value.parse().ok())
.unwrap_or(16);
let cdc = std::env::var("MONGRELDB_CDC_WAL_SEGMENTS")
.ok()
.and_then(|value| value.parse().ok())
.unwrap_or(16);
replication.max(cdc)
}
fn default_history_retention_epochs() -> u64 {
std::env::var("MONGRELDB_HISTORY_RETENTION_EPOCHS")
.ok()
.and_then(|value| value.parse().ok())
.unwrap_or(1024)
}
fn default_ai_max_concurrent() -> usize {
std::env::var("MONGRELDB_AI_MAX_CONCURRENT")
.ok()
.and_then(|value| value.parse().ok())
.filter(|value| *value > 0)
.unwrap_or(4)
}
fn positive_env_u64(name: &str, default: u64) -> u64 {
std::env::var(name)
.ok()
.and_then(|value| value.parse().ok())
.filter(|value| *value > 0)
.unwrap_or(default)
}
fn positive_env_usize(name: &str, default: usize) -> usize {
std::env::var(name)
.ok()
.and_then(|value| value.parse().ok())
.filter(|value| *value > 0)
.unwrap_or(default)
}
fn default_sql_default_timeout() -> std::time::Duration {
std::time::Duration::from_millis(positive_env_u64("MONGRELDB_SQL_DEFAULT_TIMEOUT_MS", 30_000))
}
fn default_sql_max_timeout() -> std::time::Duration {
std::time::Duration::from_millis(positive_env_u64("MONGRELDB_SQL_MAX_TIMEOUT_MS", 300_000))
}
fn default_sql_max_concurrent() -> usize {
positive_env_usize(
"MONGRELDB_SQL_MAX_CONCURRENT",
std::thread::available_parallelism()
.map(usize::from)
.unwrap_or(4),
)
}
fn default_sql_max_active_queries() -> usize {
positive_env_usize("MONGRELDB_SQL_MAX_ACTIVE_QUERIES", 1_024)
}
fn default_sql_page_max_concurrent() -> usize {
positive_env_usize("MONGRELDB_SQL_PAGE_MAX_CONCURRENT", 16)
}
fn default_sql_page_default_timeout() -> std::time::Duration {
std::time::Duration::from_millis(positive_env_u64(
"MONGRELDB_SQL_PAGE_DEFAULT_TIMEOUT_MS",
5_000,
))
}
fn default_sql_page_max_timeout() -> std::time::Duration {
std::time::Duration::from_millis(positive_env_u64(
"MONGRELDB_SQL_PAGE_MAX_TIMEOUT_MS",
30_000,
))
}
fn default_sql_finished_query_ttl() -> std::time::Duration {
std::time::Duration::from_secs(positive_env_u64(
"MONGRELDB_SQL_FINISHED_QUERY_TTL_SECS",
60,
))
}
fn default_sql_finished_detail_max_entries() -> usize {
positive_env_usize("MONGRELDB_SQL_FINISHED_DETAIL_MAX_ENTRIES", 2_048)
}
fn default_sql_finished_detail_max_bytes() -> usize {
positive_env_usize("MONGRELDB_SQL_FINISHED_DETAIL_MAX_BYTES", 8 * 1024 * 1024)
}
fn default_sql_finished_compact_max_entries() -> usize {
positive_env_usize("MONGRELDB_SQL_FINISHED_COMPACT_MAX_ENTRIES", 100_000)
}
fn default_sql_finished_compact_max_bytes() -> usize {
positive_env_usize("MONGRELDB_SQL_FINISHED_COMPACT_MAX_BYTES", 32 * 1024 * 1024)
}
fn default_sql_pre_cancel_ttl() -> std::time::Duration {
std::time::Duration::from_millis(positive_env_u64("MONGRELDB_SQL_PRE_CANCEL_TTL_MS", 15_000))
}
fn default_sql_pre_cancel_max_entries() -> usize {
positive_env_usize("MONGRELDB_SQL_PRE_CANCEL_MAX_ENTRIES", 2_048)
}
fn default_sql_pre_cancel_max_bytes() -> usize {
positive_env_usize("MONGRELDB_SQL_PRE_CANCEL_MAX_BYTES", 1024 * 1024)
}
fn default_sql_pre_cancel_max_entries_per_owner() -> usize {
positive_env_usize("MONGRELDB_SQL_PRE_CANCEL_MAX_PER_OWNER", 256)
}
fn default_sql_pre_cancel_rate_window() -> std::time::Duration {
std::time::Duration::from_millis(positive_env_u64(
"MONGRELDB_SQL_PRE_CANCEL_RATE_WINDOW_MS",
1_000,
))
}
fn default_sql_pre_cancel_rate_per_owner() -> usize {
positive_env_usize("MONGRELDB_SQL_PRE_CANCEL_RATE_PER_OWNER", 256)
}
pub(crate) fn default_sql_idempotency_ttl() -> std::time::Duration {
std::time::Duration::from_secs(positive_env_u64(
"MONGRELDB_SQL_IDEMPOTENCY_TTL_SECS",
86_400,
))
}
pub(crate) fn default_sql_idempotency_max_entries() -> usize {
positive_env_usize("MONGRELDB_SQL_IDEMPOTENCY_MAX_ENTRIES", 4_096)
}
fn default_sql_page_ttl() -> std::time::Duration {
std::time::Duration::from_secs(positive_env_u64("MONGRELDB_SQL_PAGE_TTL_SECS", 60))
}
fn default_sql_page_max_entries() -> usize {
positive_env_usize("MONGRELDB_SQL_PAGE_MAX_ENTRIES", 128)
}
fn default_sql_page_max_bytes() -> usize {
positive_env_usize("MONGRELDB_SQL_PAGE_MAX_RETAINED_BYTES", 128 * 1024 * 1024)
}
fn default_sql_page_max_entries_per_owner() -> usize {
positive_env_usize("MONGRELDB_SQL_PAGE_MAX_PER_OWNER", 16)
}
fn default_sql_cancel_grace() -> std::time::Duration {
std::time::Duration::from_millis(positive_env_u64("MONGRELDB_SQL_CANCEL_GRACE_MS", 1_000))
}
fn default_sql_max_output_bytes() -> usize {
positive_env_usize("MONGRELDB_SQL_MAX_OUTPUT_BYTES", 64 * 1024 * 1024)
}
fn default_sql_max_output_rows() -> usize {
positive_env_usize("MONGRELDB_SQL_MAX_OUTPUT_ROWS", 1_000_000)
}
fn request_owner(state: &AppState, principal: &Option<mongreldb_core::Principal>) -> String {
if let Some(p) = principal {
return format!("user:{}:{}", p.user_id, p.created_epoch);
}
if let Some(token) = state.auth_token.as_deref() {
let mut digest = sha2::Sha256::new();
digest.update(b"mongreldb-server-bearer-owner-v1\0");
digest.update((token.len() as u64).to_le_bytes());
digest.update(token.as_bytes());
return format!("bearer:{}", sql_idempotency::hex(&digest.finalize()));
}
if state.user_auth || state.db.require_auth_enabled() {
return "unauthenticated".into();
}
"anonymous".into()
}
fn request_principal(
state: &AppState,
principal: &Option<mongreldb_core::Principal>,
) -> Option<mongreldb_core::Principal> {
if let Some(principal) = principal {
return state
.db
.resolve_current_principal(principal)
.or_else(|| Some(principal.clone()));
}
state.auth_token.as_ref().and_then(|_| {
if state.db.require_auth_enabled() {
return state
.db
.principal_snapshot()
.and_then(|principal| state.db.resolve_current_principal(&principal))
.filter(|principal| principal.is_admin);
}
Some(mongreldb_core::Principal {
user_id: 0,
created_epoch: 0,
username: "token".into(),
is_admin: true,
roles: Vec::new(),
permissions: Vec::new(),
})
})
}
fn current_request_principal(
state: &AppState,
principal: &Option<mongreldb_core::Principal>,
) -> Option<mongreldb_core::Principal> {
if let Some(principal) = principal {
return state.db.resolve_current_principal(principal);
}
request_principal(state, principal)
}
fn request_identity_is_current(
state: &AppState,
principal: &Option<mongreldb_core::Principal>,
) -> bool {
if principal.is_some() || state.auth_token.is_some() {
return current_request_principal(state, principal).is_some();
}
!state.user_auth && !state.db.require_auth_enabled()
}
async fn create_session(
State(state): State<Arc<AppState>>,
OptionalPrincipal(principal): OptionalPrincipal,
) -> Response {
if !state.accepting_sql.load(Ordering::Acquire) {
return (StatusCode::SERVICE_UNAVAILABLE, "server is shutting down").into_response();
}
if !request_identity_is_current(&state, &principal) {
return StatusCode::UNAUTHORIZED.into_response();
}
let owner = request_owner(&state, &principal);
let session = match MongrelSession::open_with_external_modules_as(
Arc::clone(&state.db),
state.external_modules.iter().cloned(),
request_principal(&state, &principal),
) {
Ok(session) => session.with_query_registry(Arc::clone(&state.query_registry)),
Err(e) => return (status_for_query_error(&e), e.to_string()).into_response(),
};
match state.sessions.create(session, owner.clone()) {
Some(token) => {
state.audit.record(owner, "session.open", "session created");
Json(json!({ "session_id": token })).into_response()
}
None => (
StatusCode::SERVICE_UNAVAILABLE,
"session limit reached; close an idle session or raise --max-sessions",
)
.into_response(),
}
}
async fn close_session(
State(state): State<Arc<AppState>>,
OptionalPrincipal(principal): OptionalPrincipal,
Path(id): Path<String>,
) -> Response {
if !request_identity_is_current(&state, &principal) {
return StatusCode::NOT_FOUND.into_response();
}
let owner = request_owner(&state, &principal);
if let Some(entry) = state.sessions.take_for_close(&id, &owner) {
entry
.session
.query_registry()
.cancel_session(&id, CancellationReason::SessionClosed);
let deadline = tokio::time::Instant::now() + state.sql_cancel_grace;
while entry.session.query_registry().active_for_session(&id) > 0
&& tokio::time::Instant::now() < deadline
{
tokio::time::sleep(std::time::Duration::from_millis(5)).await;
}
state.audit.record(owner, "session.close", "session closed");
StatusCode::OK.into_response()
} else {
(
StatusCode::NOT_FOUND,
"session not found or not owned by caller",
)
.into_response()
}
}
async fn dispatch_buffered_sql_format(
state: &AppState,
format: Option<&str>,
output: ManagedQueryBatches,
query_id: QueryId,
test_hook: Option<mongreldb_query::SqlTestHook>,
output_limits: (usize, usize),
) -> Response {
let format = format.unwrap_or("json").to_owned();
let (max_rows, max_bytes) = output_limits;
let serialization_batches = output.batches().to_vec();
let serialization_query = output.query().clone();
let serialized = tokio::task::spawn_blocking(move || {
serialize_buffered_output(
&format,
&serialization_batches,
&serialization_query,
max_rows,
max_bytes,
test_hook.as_ref(),
)
})
.await;
let result = match serialized {
Ok(result) => result,
Err(error) => {
output
.query()
.record_serialization_failure("SERIALIZATION_WORKER_FAILED");
output.fail();
state.metrics.inc_sql_errors();
return terminal_server_error_response(
state,
query_id,
StatusCode::INTERNAL_SERVER_ERROR,
"SERIALIZATION_WORKER_FAILED",
error.to_string(),
);
}
};
match result {
Ok(serialized) => {
if let Err(error) = output.try_complete() {
state.metrics.inc_sql_errors();
return tracked_query_error_response(state, &error, Some(query_id));
}
state.metrics.add_sql_output_bytes(serialized.bytes.len());
let content_type = if serialized.arrow {
"application/vnd.apache.arrow.file"
} else {
"application/json"
};
with_query_id(
([(header::CONTENT_TYPE, content_type)], serialized.bytes).into_response(),
query_id,
)
}
Err(BufferedSerializationError::Query(error)) => {
output.fail();
state.metrics.inc_sql_errors();
tracked_query_error_response(state, &error, Some(query_id))
}
Err(BufferedSerializationError::Limit(message)) => {
output.fail_result_limit();
state.metrics.inc_sql_errors();
terminal_server_error_response(
state,
query_id,
StatusCode::PAYLOAD_TOO_LARGE,
"RESULT_LIMIT_EXCEEDED",
message,
)
}
Err(BufferedSerializationError::Encoding(message)) => {
output.fail_serialization();
state.metrics.inc_sql_errors();
terminal_server_error_response(
state,
query_id,
StatusCode::INTERNAL_SERVER_ERROR,
"SERIALIZATION_FAILED",
message,
)
}
}
}
#[derive(Debug)]
enum PaginatedSerializationError {
Query(mongreldb_query::MongrelQueryError),
Limit(String),
Projection(String),
Encoding(String),
}
struct SerializedPageRows {
rows: Vec<serde_json::Value>,
retained_bytes: usize,
}
struct PaginatedJsonReader<'a> {
cursor: std::io::Cursor<&'a [u8]>,
query: &'a RegisteredSqlQuery,
test_hook: Option<&'a mongreldb_query::SqlTestHook>,
}
impl std::io::Read for PaginatedJsonReader<'_> {
fn read(&mut self, buffer: &mut [u8]) -> std::io::Result<usize> {
if let Some(hook) = self.test_hook {
hook(mongreldb_query::SqlTestHookPoint::DuringPaginationDeserialization);
}
self.query
.checkpoint()
.map_err(|error| std::io::Error::other(error.to_string()))?;
let length = buffer.len().min(64 * 1024);
std::io::Read::read(&mut self.cursor, &mut buffer[..length])
}
}
const PAGINATED_VALUE_NODE_BYTES: usize = 64;
const PAGINATED_OBJECT_ENTRY_BYTES: usize = 128;
const PAGINATED_MEMORY_LIMIT_ERROR: &str = "SQL retained output memory limit exceeded";
struct PaginatedDecodeBudget<'a> {
used: usize,
limit: usize,
nodes: usize,
exceeded: bool,
query: &'a RegisteredSqlQuery,
test_hook: Option<&'a mongreldb_query::SqlTestHook>,
}
impl PaginatedDecodeBudget<'_> {
fn begin_value<E: serde::de::Error>(&mut self) -> Result<(), E> {
self.nodes = self.nodes.saturating_add(1);
if self.nodes & 255 == 0 {
if let Some(hook) = self.test_hook {
hook(mongreldb_query::SqlTestHookPoint::DuringPaginationDeserialization);
}
self.query.checkpoint().map_err(E::custom)?;
}
self.charge::<E>(PAGINATED_VALUE_NODE_BYTES)
}
fn charge<E: serde::de::Error>(&mut self, bytes: usize) -> Result<(), E> {
let next = self.used.saturating_add(bytes);
if next > self.limit {
self.exceeded = true;
return Err(E::custom(PAGINATED_MEMORY_LIMIT_ERROR));
}
self.used = next;
Ok(())
}
}
struct BudgetedJsonValueSeed<'a, 'query> {
budget: &'a mut PaginatedDecodeBudget<'query>,
}
impl<'de> serde::de::DeserializeSeed<'de> for BudgetedJsonValueSeed<'_, '_> {
type Value = serde_json::Value;
fn deserialize<D>(self, deserializer: D) -> Result<Self::Value, D::Error>
where
D: serde::Deserializer<'de>,
{
self.budget.begin_value::<D::Error>()?;
deserializer.deserialize_any(BudgetedJsonValueVisitor {
budget: self.budget,
})
}
}
struct BudgetedJsonValueVisitor<'a, 'query> {
budget: &'a mut PaginatedDecodeBudget<'query>,
}
impl<'de> serde::de::Visitor<'de> for BudgetedJsonValueVisitor<'_, '_> {
type Value = serde_json::Value;
fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.write_str("a JSON value")
}
fn visit_bool<E>(self, value: bool) -> Result<Self::Value, E> {
Ok(serde_json::Value::Bool(value))
}
fn visit_i64<E>(self, value: i64) -> Result<Self::Value, E> {
Ok(serde_json::Value::Number(value.into()))
}
fn visit_u64<E>(self, value: u64) -> Result<Self::Value, E> {
Ok(serde_json::Value::Number(value.into()))
}
fn visit_f64<E>(self, value: f64) -> Result<Self::Value, E>
where
E: serde::de::Error,
{
serde_json::Number::from_f64(value)
.map(serde_json::Value::Number)
.ok_or_else(|| E::custom("JSON number is not finite"))
}
fn visit_str<E>(self, value: &str) -> Result<Self::Value, E>
where
E: serde::de::Error,
{
self.budget.charge::<E>(value.len())?;
Ok(serde_json::Value::String(value.to_owned()))
}
fn visit_borrowed_str<E>(self, value: &'de str) -> Result<Self::Value, E>
where
E: serde::de::Error,
{
self.visit_str(value)
}
fn visit_string<E>(self, value: String) -> Result<Self::Value, E>
where
E: serde::de::Error,
{
self.budget.charge::<E>(value.capacity())?;
Ok(serde_json::Value::String(value))
}
fn visit_none<E>(self) -> Result<Self::Value, E> {
Ok(serde_json::Value::Null)
}
fn visit_unit<E>(self) -> Result<Self::Value, E> {
Ok(serde_json::Value::Null)
}
fn visit_some<D>(self, deserializer: D) -> Result<Self::Value, D::Error>
where
D: serde::Deserializer<'de>,
{
deserializer.deserialize_any(self)
}
fn visit_newtype_struct<D>(self, deserializer: D) -> Result<Self::Value, D::Error>
where
D: serde::Deserializer<'de>,
{
deserializer.deserialize_any(self)
}
fn visit_seq<A>(self, mut sequence: A) -> Result<Self::Value, A::Error>
where
A: serde::de::SeqAccess<'de>,
{
let mut values = Vec::new();
while let Some(value) = sequence.next_element_seed(BudgetedJsonValueSeed {
budget: self.budget,
})? {
values.push(value);
}
Ok(serde_json::Value::Array(values))
}
fn visit_map<A>(self, mut map: A) -> Result<Self::Value, A::Error>
where
A: serde::de::MapAccess<'de>,
{
let mut values = serde_json::Map::new();
while let Some(key) = map.next_key::<String>()? {
self.budget
.charge::<A::Error>(PAGINATED_OBJECT_ENTRY_BYTES.saturating_add(key.capacity()))?;
let value = map.next_value_seed(BudgetedJsonValueSeed {
budget: self.budget,
})?;
values.insert(key, value);
}
Ok(serde_json::Value::Object(values))
}
}
struct BudgetedJsonRowsSeed<'a, 'query> {
budget: &'a mut PaginatedDecodeBudget<'query>,
}
impl<'de> serde::de::DeserializeSeed<'de> for BudgetedJsonRowsSeed<'_, '_> {
type Value = Vec<serde_json::Value>;
fn deserialize<D>(self, deserializer: D) -> Result<Self::Value, D::Error>
where
D: serde::Deserializer<'de>,
{
deserializer.deserialize_seq(BudgetedJsonRowsVisitor {
budget: self.budget,
})
}
}
struct BudgetedJsonRowsVisitor<'a, 'query> {
budget: &'a mut PaginatedDecodeBudget<'query>,
}
impl<'de> serde::de::Visitor<'de> for BudgetedJsonRowsVisitor<'_, '_> {
type Value = Vec<serde_json::Value>;
fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.write_str("a JSON array of SQL rows")
}
fn visit_seq<A>(self, mut sequence: A) -> Result<Self::Value, A::Error>
where
A: serde::de::SeqAccess<'de>,
{
let mut rows = Vec::new();
while let Some(row) = sequence.next_element_seed(BudgetedJsonValueSeed {
budget: self.budget,
})? {
rows.push(row);
}
Ok(rows)
}
}
#[allow(clippy::too_many_arguments)]
async fn dispatch_paginated_sql(
state: &AppState,
output: ManagedQueryBatches,
query_id: QueryId,
owner: &str,
pagination: ResolvedSqlPagination,
output_limits: (usize, usize),
test_hook: Option<mongreldb_query::SqlTestHook>,
binding: sql_pages::SqlPageBinding,
) -> Response {
let projection = pagination.projection;
let serialization_projection = projection.clone();
let serialization_batches = output.batches().to_vec();
let serialization_query = output.query().clone();
let response_test_hook = test_hook.clone();
let retained_memory_limit = output_limits.1.min(state.sql_pages.max_retained_bytes());
let serialized = tokio::task::spawn_blocking(move || {
serialize_paginated_rows(
&serialization_batches,
&serialization_query,
&serialization_projection,
output_limits.0,
output_limits.1,
retained_memory_limit,
test_hook.as_ref(),
)
})
.await;
let serialized = match serialized {
Ok(result) => result,
Err(error) => {
output
.query()
.record_serialization_failure("SERIALIZATION_WORKER_FAILED");
output.fail();
state.metrics.inc_sql_errors();
return terminal_server_error_response(
state,
query_id,
StatusCode::INTERNAL_SERVER_ERROR,
"SERIALIZATION_WORKER_FAILED",
error.to_string(),
);
}
};
let serialized = match serialized {
Ok(serialized) => serialized,
Err(PaginatedSerializationError::Query(error)) => {
output.fail();
state.metrics.inc_sql_errors();
return tracked_query_error_response(state, &error, Some(query_id));
}
Err(PaginatedSerializationError::Limit(message)) => {
output.fail_result_limit();
state.metrics.inc_sql_errors();
return terminal_server_error_response(
state,
query_id,
StatusCode::PAYLOAD_TOO_LARGE,
"RESULT_LIMIT_EXCEEDED",
message,
);
}
Err(PaginatedSerializationError::Projection(message)) => {
output.fail_with_error(
"INVALID_SQL_PROJECTION",
mongreldb_query::QueryTerminalErrorCategory::Execution,
);
state.metrics.inc_sql_errors();
return terminal_server_error_response(
state,
query_id,
StatusCode::BAD_REQUEST,
"INVALID_SQL_PROJECTION",
message,
);
}
Err(PaginatedSerializationError::Encoding(message)) => {
output.fail_serialization();
state.metrics.inc_sql_errors();
return terminal_server_error_response(
state,
query_id,
StatusCode::INTERNAL_SERVER_ERROR,
"SERIALIZATION_FAILED",
message,
);
}
};
let retained_bytes = serialized.retained_bytes;
let current_binding = sql_pages::SqlPageBinding {
security_version: state.db.security_version(),
catalog_epoch: state.db.catalog_snapshot().db_epoch,
};
if current_binding != binding {
output.fail_with_error(
"SQL_CURSOR_EXPIRED",
mongreldb_query::QueryTerminalErrorCategory::Execution,
);
return with_query_id(
sql_cursor_error_response(
StatusCode::CONFLICT,
"SQL_CURSOR_EXPIRED",
"authorization or schema changed while retaining the SQL result",
query_id,
),
query_id,
);
}
let retained = match state.sql_pages.insert(
owner,
serialized.rows,
projection,
pagination.limits,
retained_bytes,
binding,
) {
Ok(retained) => retained,
Err(sql_pages::InsertError::Full | sql_pages::InsertError::OwnerLimit) => {
output.fail_with_error(
"SQL_PAGE_STORE_FULL",
mongreldb_query::QueryTerminalErrorCategory::Execution,
);
state.metrics.inc_sql_errors();
return terminal_server_error_response(
state,
query_id,
StatusCode::SERVICE_UNAVAILABLE,
"SQL_PAGE_STORE_FULL",
"retained SQL result capacity reached",
);
}
Err(sql_pages::InsertError::EntropyUnavailable) => {
output.fail_with_error(
"ENTROPY_UNAVAILABLE",
mongreldb_query::QueryTerminalErrorCategory::Execution,
);
state.metrics.inc_sql_errors();
return terminal_server_error_response(
state,
query_id,
StatusCode::INTERNAL_SERVER_ERROR,
"ENTROPY_UNAVAILABLE",
"OS CSPRNG unavailable",
);
}
};
let cursor_mac_key = match state.cursor_mac_key.get() {
Ok(key) => key,
Err(_) => {
state.sql_pages.discard(&retained);
output.fail_with_error(
"ENTROPY_UNAVAILABLE",
mongreldb_query::QueryTerminalErrorCategory::Execution,
);
state.metrics.inc_sql_errors();
return terminal_server_error_response(
state,
query_id,
StatusCode::INTERNAL_SERVER_ERROR,
"ENTROPY_UNAVAILABLE",
"OS CSPRNG unavailable",
);
}
};
let page = match sql_pages::SqlPageStore::first_page(&retained, &cursor_mac_key) {
Ok(page) => page,
Err(sql_pages::PageError::RowExceedsLimits) => {
state.sql_pages.discard(&retained);
output.fail_result_limit();
state.metrics.inc_sql_errors();
return terminal_server_error_response(
state,
query_id,
StatusCode::PAYLOAD_TOO_LARGE,
"RESULT_LIMIT_EXCEEDED",
"one projected row exceeds the page byte or token limit",
);
}
Err(sql_pages::PageError::OffsetInvalid) => {
state.sql_pages.discard(&retained);
output.fail_with_error(
"INVALID_PAGE_OFFSET",
mongreldb_query::QueryTerminalErrorCategory::Serialization,
);
state.metrics.inc_sql_errors();
return terminal_server_error_response(
state,
query_id,
StatusCode::INTERNAL_SERVER_ERROR,
"INVALID_PAGE_OFFSET",
"failed to create the first SQL result page",
);
}
Err(sql_pages::PageError::Cancelled) => unreachable!("first-page rendering has no control"),
};
let discard_after_response = page.next_cursor.is_none();
let page_byte_count = page.byte_count;
let encoded_page = tokio::task::spawn_blocking(move || serialize_sql_page(page)).await;
let encoded_page = match encoded_page {
Ok(Ok(encoded_page)) => encoded_page,
Ok(Err(error)) => {
state.sql_pages.discard(&retained);
output.fail_serialization();
state.metrics.inc_sql_errors();
return terminal_server_error_response(
state,
query_id,
StatusCode::INTERNAL_SERVER_ERROR,
"SERIALIZATION_FAILED",
error.to_string(),
);
}
Err(_) => {
state.sql_pages.discard(&retained);
output.fail_serialization();
state.metrics.inc_sql_errors();
return terminal_server_error_response(
state,
query_id,
StatusCode::INTERNAL_SERVER_ERROR,
"SERIALIZATION_WORKER_FAILED",
"SQL page response serialization worker failed",
);
}
};
if let Some(hook) = response_test_hook {
hook(mongreldb_query::SqlTestHookPoint::AfterPageResponseSerialization);
}
if let Err(error) = output.query().checkpoint() {
state.sql_pages.discard(&retained);
state.metrics.inc_sql_errors();
return tracked_query_error_response(state, &error, Some(query_id));
}
if let Err(error) = output.try_complete() {
state.sql_pages.discard(&retained);
state.metrics.inc_sql_errors();
return tracked_query_error_response(state, &error, Some(query_id));
}
if discard_after_response {
state.sql_pages.discard(&retained);
}
state.metrics.add_sql_output_bytes(page_byte_count);
with_query_id(sql_page_response(encoded_page), query_id)
}
fn serialize_paginated_rows(
batches: &[arrow::record_batch::RecordBatch],
query: &RegisteredSqlQuery,
projection: &[String],
max_rows: usize,
max_bytes: usize,
retained_memory_limit: usize,
test_hook: Option<&mongreldb_query::SqlTestHook>,
) -> Result<SerializedPageRows, PaginatedSerializationError> {
const ROW_CHECKPOINT_INTERVAL: usize = 256;
if batches.is_empty() {
let retained_bytes = sql_pages::accounted_bytes(2, &[], projection, || query.checkpoint())
.map_err(PaginatedSerializationError::Query)?;
if retained_bytes > retained_memory_limit {
return Err(PaginatedSerializationError::Limit(
PAGINATED_MEMORY_LIMIT_ERROR.into(),
));
}
return Ok(SerializedPageRows {
rows: Vec::new(),
retained_bytes,
});
}
let fields = batches[0].schema();
let mut indices = Vec::with_capacity(projection.len());
for name in projection {
let matches: Vec<_> = fields
.fields()
.iter()
.enumerate()
.filter_map(|(index, field)| (field.name() == name).then_some(index))
.collect();
if matches.len() != 1 {
return Err(PaginatedSerializationError::Projection(format!(
"projected output column {name:?} is missing or ambiguous"
)));
}
indices.push(matches[0]);
}
let mut writer_output = LimitedOutput::new(max_bytes);
let mut rows = 0usize;
let encoding = (|| {
let mut writer = arrow::json::writer::ArrayWriter::new(&mut writer_output);
for batch in batches {
let batch = batch.project(&indices).map_err(|error| error.to_string())?;
for offset in (0..batch.num_rows()).step_by(ROW_CHECKPOINT_INTERVAL) {
if let Some(hook) = test_hook {
hook(mongreldb_query::SqlTestHookPoint::BeforeSerializationBatch);
}
query.checkpoint().map_err(|error| error.to_string())?;
let length = ROW_CHECKPOINT_INTERVAL.min(batch.num_rows() - offset);
rows = rows.saturating_add(length);
if rows > max_rows {
return Err("SQL retained output row limit exceeded".into());
}
let slice = batch.slice(offset, length);
writer
.write_batches(&[&slice])
.map_err(|error| error.to_string())?;
}
}
writer.finish().map_err(|error| error.to_string())
})();
if let Err(error) = encoding {
if let Err(query_error) = query.checkpoint() {
return Err(PaginatedSerializationError::Query(query_error));
}
if writer_output.exceeded || rows > max_rows {
return Err(PaginatedSerializationError::Limit(error));
}
return Err(PaginatedSerializationError::Encoding(error));
}
let bytes = writer_output.bytes.len();
let retained_base = sql_pages::accounted_bytes(bytes, &[], projection, || query.checkpoint())
.map_err(PaginatedSerializationError::Query)?;
if retained_base > retained_memory_limit {
return Err(PaginatedSerializationError::Limit(
PAGINATED_MEMORY_LIMIT_ERROR.into(),
));
}
let reader = PaginatedJsonReader {
cursor: std::io::Cursor::new(writer_output.bytes.as_slice()),
query,
test_hook,
};
let mut deserializer =
serde_json::Deserializer::from_reader(std::io::BufReader::with_capacity(64 * 1024, reader));
let mut budget = PaginatedDecodeBudget {
used: retained_base,
limit: retained_memory_limit,
nodes: 0,
exceeded: false,
query,
test_hook,
};
let rows = match serde::de::DeserializeSeed::deserialize(
BudgetedJsonRowsSeed {
budget: &mut budget,
},
&mut deserializer,
) {
Ok(rows) => {
if let Err(error) = deserializer.end() {
return Err(PaginatedSerializationError::Encoding(error.to_string()));
}
rows
}
Err(error) => {
if let Err(query_error) = query.checkpoint() {
return Err(PaginatedSerializationError::Query(query_error));
}
if budget.exceeded {
return Err(PaginatedSerializationError::Limit(
PAGINATED_MEMORY_LIMIT_ERROR.into(),
));
}
return Err(PaginatedSerializationError::Encoding(error.to_string()));
}
};
if let Some(hook) = test_hook {
hook(mongreldb_query::SqlTestHookPoint::AfterSerialization);
}
let retained_bytes =
sql_pages::accounted_bytes(bytes, &rows, projection, || query.checkpoint())
.map_err(PaginatedSerializationError::Query)?;
if retained_bytes > retained_memory_limit {
return Err(PaginatedSerializationError::Limit(
PAGINATED_MEMORY_LIMIT_ERROR.into(),
));
}
Ok(SerializedPageRows {
rows,
retained_bytes,
})
}
fn serialize_sql_page(page: sql_pages::SqlPage) -> Result<Vec<u8>, serde_json::Error> {
serde_json::to_vec(&json!({
"status": "completed",
"rows": page.rows,
"next_cursor": page.next_cursor,
"page": {
"offset": page.offset,
"row_count": page.row_count,
"total_rows": page.total_rows,
"byte_count": page.byte_count,
"estimated_tokens": page.estimated_tokens,
"limits": page.limits,
"projection": page.projection,
"expires_at_ms": page.expires_at_ms,
"snapshot": "retained_result",
"token_estimate": "ceil(projected_json_bytes/4)",
}
}))
}
enum ControlledPageSerializationError {
Query(mongreldb_query::MongrelQueryError),
Encoding,
}
struct ControlledPageWriter<'a> {
bytes: Vec<u8>,
query: &'a RegisteredSqlQuery,
}
impl std::io::Write for ControlledPageWriter<'_> {
fn write(&mut self, buffer: &[u8]) -> std::io::Result<usize> {
self.query
.checkpoint()
.map_err(|error| std::io::Error::other(error.to_string()))?;
self.bytes.extend_from_slice(buffer);
Ok(buffer.len())
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
fn serialize_sql_page_controlled(
page: sql_pages::SqlPage,
query: &RegisteredSqlQuery,
) -> Result<Vec<u8>, ControlledPageSerializationError> {
query
.checkpoint()
.map_err(ControlledPageSerializationError::Query)?;
let value = json!({
"status": "completed",
"rows": page.rows,
"next_cursor": page.next_cursor,
"page": {
"offset": page.offset,
"row_count": page.row_count,
"total_rows": page.total_rows,
"byte_count": page.byte_count,
"estimated_tokens": page.estimated_tokens,
"limits": page.limits,
"projection": page.projection,
"expires_at_ms": page.expires_at_ms,
"snapshot": "retained_result",
"token_estimate": "ceil(projected_json_bytes/4)",
}
});
let mut writer = ControlledPageWriter {
bytes: Vec::new(),
query,
};
if let Err(error) = serde_json::to_writer(&mut writer, &value) {
return match query.checkpoint() {
Err(error) => Err(ControlledPageSerializationError::Query(error)),
Ok(()) => {
let _ = error;
Err(ControlledPageSerializationError::Encoding)
}
};
}
query
.checkpoint()
.map_err(ControlledPageSerializationError::Query)?;
Ok(writer.bytes)
}
fn sql_page_response(body: Vec<u8>) -> Response {
([(header::CONTENT_TYPE, "application/json")], body).into_response()
}
fn validate_stmt_name(name: &str) -> Result<(), String> {
let mut chars = name.chars();
match chars.next() {
Some(c) if c.is_ascii_alphabetic() || c == '_' => {}
_ => return Err("statement name must start with a letter or underscore".into()),
}
if !chars.all(|c| c.is_ascii_alphanumeric() || c == '_') {
return Err("statement name may contain only letters, digits, or underscore".into());
}
Ok(())
}
fn render_sql_literal(v: &serde_json::Value) -> Result<String, String> {
match v {
serde_json::Value::Null => Ok("NULL".into()),
serde_json::Value::Bool(b) => {
if *b {
Ok("TRUE".into())
} else {
Ok("FALSE".into())
}
}
serde_json::Value::Number(n) => Ok(n.to_string()),
serde_json::Value::String(s) => {
let mut out = String::with_capacity(s.len() + 2);
out.push('\'');
for c in s.chars() {
if c == '\'' {
out.push_str("''");
} else {
out.push(c);
}
}
out.push('\'');
Ok(out)
}
_ => Err("prepared-statement parameters must be scalar (null/bool/number/string)".into()),
}
}
#[derive(Deserialize)]
struct PrepareRequest {
name: String,
sql: String,
#[serde(default)]
query_id: Option<QueryId>,
#[serde(default)]
timeout_ms: Option<u64>,
}
async fn prepare_statement(
State(state): State<Arc<AppState>>,
OptionalPrincipal(principal): OptionalPrincipal,
Path(id): Path<String>,
headers: axum::http::HeaderMap,
Json(req): Json<PrepareRequest>,
) -> Response {
if !state.accepting_sql.load(Ordering::Acquire) {
return (StatusCode::SERVICE_UNAVAILABLE, "server is shutting down").into_response();
}
if !request_identity_is_current(&state, &principal) {
return StatusCode::NOT_FOUND.into_response();
}
if let Err(msg) = validate_stmt_name(&req.name) {
return (StatusCode::BAD_REQUEST, msg).into_response();
}
let owner = request_owner(&state, &principal);
let Some(entry) = state.sessions.get(&id, &owner) else {
return (
StatusCode::NOT_FOUND,
"session not found or not owned by caller",
)
.into_response();
};
let (options, query_id) = match resolve_query_options(
&state,
&headers,
req.query_id,
req.timeout_ms,
owner,
Some(id),
) {
Ok(options) => options,
Err(response) => return *response,
};
let query = match register_controlled_query(&state, &entry.session, options) {
Ok(query) => query,
Err(error) => return tracked_query_error_response(&state, &error, Some(query_id)),
};
let registration = RegisteredQueryGuard::new(query);
if mongreldb_query::contains_boolean_ai_predicate(&req.sql) {
registration.fail();
return with_query_id(remote_boolean_ai_error(), query_id);
}
let _sql_permit = match acquire_sql_permit(&state, &entry.session, registration.query()).await {
Ok(permit) => permit,
Err(error) => return tracked_query_error_response(&state, &error, Some(query_id)),
};
let _guard = tokio::select! {
guard = entry.lock.lock() => guard,
_ = registration.query().control().cancelled() => {
return tracked_query_error_response(
&state,
&cancellation_checkpoint_error(registration.query()),
Some(query_id),
);
}
};
if entry.is_closed() {
return with_query_id(
(StatusCode::NOT_FOUND, "session no longer available").into_response(),
query_id,
);
}
entry.touch();
let sql = format!("PREPARE {} AS {}", req.name, req.sql);
let query = registration.into_query();
match entry.session.run_with_query(&sql, query).await {
Ok(_) => with_query_id(
Json(json!({ "prepared": req.name })).into_response(),
query_id,
),
Err(error) => tracked_query_error_response(&state, &error, Some(query_id)),
}
}
#[derive(Deserialize)]
struct ExecuteRequest {
name: String,
params: Vec<serde_json::Value>,
#[serde(default)]
format: Option<String>,
#[serde(default)]
query_id: Option<QueryId>,
#[serde(default)]
timeout_ms: Option<u64>,
}
async fn execute_statement(
State(state): State<Arc<AppState>>,
OptionalPrincipal(principal): OptionalPrincipal,
Path(id): Path<String>,
headers: axum::http::HeaderMap,
Json(req): Json<ExecuteRequest>,
) -> Response {
if !state.accepting_sql.load(Ordering::Acquire) {
return (StatusCode::SERVICE_UNAVAILABLE, "server is shutting down").into_response();
}
if !request_identity_is_current(&state, &principal) {
return StatusCode::NOT_FOUND.into_response();
}
if let Err(msg) = validate_stmt_name(&req.name) {
return (StatusCode::BAD_REQUEST, msg).into_response();
}
let owner = request_owner(&state, &principal);
let Some(entry) = state.sessions.get(&id, &owner) else {
return (
StatusCode::NOT_FOUND,
"session not found or not owned by caller",
)
.into_response();
};
let (options, query_id) = match resolve_query_options(
&state,
&headers,
req.query_id,
req.timeout_ms,
owner,
Some(id),
) {
Ok(options) => options,
Err(response) => return *response,
};
let query = match register_controlled_query(&state, &entry.session, options) {
Ok(query) => query,
Err(error) => return tracked_query_error_response(&state, &error, Some(query_id)),
};
let registration = RegisteredQueryGuard::new(query);
let sql_permit = match acquire_sql_permit(&state, &entry.session, registration.query()).await {
Ok(permit) => permit,
Err(error) => return tracked_query_error_response(&state, &error, Some(query_id)),
};
let _guard = tokio::select! {
guard = entry.lock.lock() => guard,
_ = registration.query().control().cancelled() => {
return tracked_query_error_response(
&state,
&cancellation_checkpoint_error(registration.query()),
Some(query_id),
);
}
};
if entry.is_closed() {
return with_query_id(
(StatusCode::NOT_FOUND, "session no longer available").into_response(),
query_id,
);
}
entry.touch();
state.metrics.inc_sql_queries();
let literals: Vec<String> = match req
.params
.iter()
.map(render_sql_literal)
.collect::<Result<_, _>>()
{
Ok(v) => v,
Err(msg) => {
state.metrics.inc_sql_errors();
return (StatusCode::BAD_REQUEST, msg).into_response();
}
};
let sql = format!("EXECUTE {}({})", req.name, literals.join(", "));
let start = std::time::Instant::now();
let result = if req.format.as_deref() == Some("arrow-stream") {
let query = registration.into_query();
match entry
.session
.run_stream_with_query_for_serialization(&sql, query)
.await
{
Ok((stream, completion)) => Ok(sql_arrow_stream_response_controlled(
stream,
completion,
sql_permit,
(state.sql_max_output_rows, state.sql_max_output_bytes),
&state,
query_id,
entry.session.sql_test_hook(),
)),
Err(error) => Err(error),
}
} else {
let query = registration.into_query();
match entry
.session
.run_with_query_for_serialization_with_limits(
&sql,
query,
mongreldb_query::SqlCollectionLimits::new(
state.sql_max_output_rows,
state.sql_max_output_bytes,
),
)
.await
{
Ok(output) => Ok(dispatch_buffered_sql_format(
&state,
req.format.as_deref(),
output,
query_id,
entry.session.sql_test_hook(),
(state.sql_max_output_rows, state.sql_max_output_bytes),
)
.await),
Err(error) => Err(error),
}
};
let elapsed = start.elapsed();
if elapsed >= state.slow_query_threshold {
state.metrics.inc_slow_queries();
eprintln!(
"[slow-query] {}\u{00b5}s \u{2014} EXECUTE {}",
elapsed.as_micros(),
req.name
);
}
match result {
Ok(response) => with_query_id(response, query_id),
Err(e) => {
state.metrics.inc_sql_errors();
let msg = format!("{e}");
let status = if msg.contains("does not exist") {
StatusCode::NOT_FOUND
} else {
status_for_query_error(&e)
};
if status == status_for_query_error(&e) {
tracked_query_error_response(&state, &e, Some(query_id))
} else {
with_query_id(
(status, format!("{msg} ({}µs)", elapsed.as_micros())).into_response(),
query_id,
)
}
}
}
}
async fn deallocate_statement(
State(state): State<Arc<AppState>>,
OptionalPrincipal(principal): OptionalPrincipal,
Path((id, name)): Path<(String, String)>,
headers: axum::http::HeaderMap,
) -> Response {
if !state.accepting_sql.load(Ordering::Acquire) {
return (StatusCode::SERVICE_UNAVAILABLE, "server is shutting down").into_response();
}
if !request_identity_is_current(&state, &principal) {
return StatusCode::NOT_FOUND.into_response();
}
if let Err(msg) = validate_stmt_name(&name) {
return (StatusCode::BAD_REQUEST, msg).into_response();
}
let owner = request_owner(&state, &principal);
let Some(entry) = state.sessions.get(&id, &owner) else {
return (
StatusCode::NOT_FOUND,
"session not found or not owned by caller",
)
.into_response();
};
let (options, query_id) =
match resolve_query_options(&state, &headers, None, None, owner, Some(id)) {
Ok(options) => options,
Err(response) => return *response,
};
let query = match register_controlled_query(&state, &entry.session, options) {
Ok(query) => query,
Err(error) => return tracked_query_error_response(&state, &error, Some(query_id)),
};
let registration = RegisteredQueryGuard::new(query);
let _sql_permit = match acquire_sql_permit(&state, &entry.session, registration.query()).await {
Ok(permit) => permit,
Err(error) => return tracked_query_error_response(&state, &error, Some(query_id)),
};
let _guard = tokio::select! {
guard = entry.lock.lock() => guard,
_ = registration.query().control().cancelled() => {
return tracked_query_error_response(
&state,
&cancellation_checkpoint_error(registration.query()),
Some(query_id),
);
}
};
if entry.is_closed() {
return with_query_id(
(StatusCode::NOT_FOUND, "session no longer available").into_response(),
query_id,
);
}
entry.touch();
state.metrics.inc_sql_queries();
let sql = format!("DEALLOCATE {name}");
let start = std::time::Instant::now();
let result = entry
.session
.run_with_query(&sql, registration.into_query())
.await;
let elapsed = start.elapsed();
if elapsed >= state.slow_query_threshold {
state.metrics.inc_slow_queries();
eprintln!(
"[slow-query] {}\u{00b5}s query_id={} operation=DEALLOCATE",
elapsed.as_micros(),
query_id,
);
}
match result {
Ok(_) => with_query_id(
Json(json!({ "deallocated": name })).into_response(),
query_id,
),
Err(error) => {
state.metrics.inc_sql_errors();
tracked_query_error_response(&state, &error, Some(query_id))
}
}
}
async fn metrics_handler(
State(state): State<Arc<AppState>>,
OptionalPrincipal(principal): OptionalPrincipal,
) -> Response {
if let Err(error) = state.db.require_for(
request_principal(&state, &principal).as_ref(),
&mongreldb_core::Permission::Admin,
) {
return (status_for_error(&error), error.to_string()).into_response();
}
let body = state.metrics.prometheus_text(
state.db.table_names().len(),
state.query_registry.stats(),
(
state.pre_cancellations.len(),
state.pre_cancellations.approximate_bytes(),
),
);
(
[(
header::CONTENT_TYPE,
"text/plain; version=0.0.4; charset=utf-8".to_string(),
)],
body,
)
.into_response()
}
async fn compact_all(
State(state): State<Arc<AppState>>,
OptionalPrincipal(principal): OptionalPrincipal,
) -> (StatusCode, Json<serde_json::Value>) {
if let Err(error) = state.db.require_for(
request_principal(&state, &principal).as_ref(),
&mongreldb_core::Permission::Ddl,
) {
return (
status_for_error(&error),
Json(json!({ "status": "error", "message": error.to_string() })),
);
}
match state.db.compact() {
Ok((compacted, skipped)) => (
StatusCode::OK,
Json(json!({
"status": "ok",
"compacted": compacted,
"skipped": skipped,
})),
),
Err(e) => (
StatusCode::INTERNAL_SERVER_ERROR,
Json(json!({ "status": "error", "message": format!("{e}") })),
),
}
}
async fn compact_table(
State(state): State<Arc<AppState>>,
OptionalPrincipal(principal): OptionalPrincipal,
Path(name): Path<String>,
) -> (StatusCode, Json<serde_json::Value>) {
if let Err(error) = state.db.require_for(
request_principal(&state, &principal).as_ref(),
&mongreldb_core::Permission::Ddl,
) {
return (
status_for_error(&error),
Json(json!({ "status": "error", "table": name, "message": error.to_string() })),
);
}
match state.db.compact_table(&name) {
Ok(true) => (
StatusCode::OK,
Json(json!({ "status": "compacted", "table": name })),
),
Ok(false) => (
StatusCode::OK,
Json(json!({ "status": "skipped", "table": name, "reason": "fewer than 2 runs" })),
),
Err(e) => (
StatusCode::INTERNAL_SERVER_ERROR,
Json(json!({ "status": "error", "table": name, "message": format!("{e}") })),
),
}
}
#[derive(Deserialize)]
struct CreateTableRequest {
name: String,
columns: Vec<ColumnDefJson>,
}
#[derive(Deserialize)]
struct ColumnDefJson {
id: u16,
name: String,
ty: String,
primary_key: bool,
#[serde(default)]
nullable: bool,
}
async fn create_table(
State(state): State<Arc<AppState>>,
OptionalPrincipal(principal): OptionalPrincipal,
Json(req): Json<CreateTableRequest>,
) -> Response {
if let Err(error) = state.db.require_for(
request_principal(&state, &principal).as_ref(),
&mongreldb_core::Permission::Ddl,
) {
return (status_for_error(&error), error.to_string()).into_response();
}
let mut columns = Vec::new();
for c in &req.columns {
let ty = match c.ty.as_str() {
"int64" | "bigint" => TypeId::Int64,
"float64" | "double" => TypeId::Float64,
"bytes" | "varchar" | "text" => TypeId::Bytes,
"bool" => TypeId::Bool,
other => {
return (StatusCode::BAD_REQUEST, format!("unknown type: {other}")).into_response()
}
};
let mut flags = mongreldb_core::schema::ColumnFlags::empty();
if c.primary_key {
flags = flags.with(mongreldb_core::schema::ColumnFlags::PRIMARY_KEY);
}
if c.nullable {
flags = flags.with(mongreldb_core::schema::ColumnFlags::NULLABLE);
}
columns.push(mongreldb_core::schema::ColumnDef {
id: c.id,
name: c.name.clone(),
ty,
flags,
default_value: None,
});
}
let schema = Schema {
schema_id: 0,
columns,
indexes: vec![],
colocation: vec![],
constraints: Default::default(),
clustered: false,
};
if let Err(msg) = validate_table_name(&req.name) {
return (StatusCode::BAD_REQUEST, msg).into_response();
}
match state.db.create_table(&req.name, schema) {
Ok(id) => Json(json!({
"table_id": id,
"table_id_text": id.to_string()
}))
.into_response(),
Err(error) => crate::kit::durable_core_error_response(&error)
.unwrap_or_else(|| (status_for_error(&error), error.to_string()).into_response()),
}
}
async fn list_tables(
State(state): State<Arc<AppState>>,
OptionalPrincipal(principal): OptionalPrincipal,
) -> Json<Vec<String>> {
let principal = request_principal(&state, &principal);
Json(
state
.db
.table_names()
.into_iter()
.filter(|table| {
state
.db
.select_column_ids_for(table, principal.as_ref())
.is_ok()
})
.collect(),
)
}
async fn drop_table(
State(state): State<Arc<AppState>>,
OptionalPrincipal(principal): OptionalPrincipal,
Path(name): Path<String>,
) -> Response {
if let Err(error) = state.db.require_for(
request_principal(&state, &principal).as_ref(),
&mongreldb_core::Permission::Ddl,
) {
return (status_for_error(&error), error.to_string()).into_response();
}
match state.db.drop_table_with_epoch(&name) {
Ok(epoch) => Json(json!({
"status": "committed",
"epoch": epoch.0,
"epoch_text": epoch.0.to_string()
}))
.into_response(),
Err(error) => crate::kit::durable_core_error_response(&error)
.unwrap_or_else(|| (status_for_error(&error), error.to_string()).into_response()),
}
}
#[derive(Deserialize)]
struct PutRequest {
row: Vec<serde_json::Value>,
}
pub(crate) fn json_to_value(v: &serde_json::Value, expected: &TypeId) -> Value {
match (v, expected) {
(serde_json::Value::Number(n), TypeId::Float64) => {
n.as_f64().map(Value::Float64).unwrap_or(Value::Null)
}
(serde_json::Value::Number(n), TypeId::Int64) => {
n.as_i64().map(Value::Int64).unwrap_or(Value::Null)
}
(serde_json::Value::String(s), TypeId::Bytes) => Value::Bytes(s.as_bytes().to_vec()),
(serde_json::Value::String(s), TypeId::Enum { variants }) => {
if variants.iter().any(|v| v == s) {
Value::Bytes(s.as_bytes().to_vec())
} else {
Value::Null
}
}
(serde_json::Value::Bool(b), TypeId::Bool) => Value::Bool(*b),
(serde_json::Value::Array(arr), TypeId::Embedding { dim }) => {
if arr.len() as u32 != *dim {
return Value::Null;
}
let vec: Option<Vec<f32>> =
arr.iter().map(|el| el.as_f64().map(|f| f as f32)).collect();
vec.map(Value::Embedding).unwrap_or(Value::Null)
}
(serde_json::Value::Null, _) => Value::Null,
(serde_json::Value::Number(n), _) => {
if let Some(i) = n.as_i64() {
Value::Int64(i)
} else if let Some(f) = n.as_f64() {
Value::Float64(f)
} else {
Value::Null
}
}
(serde_json::Value::String(s), _) => Value::Bytes(s.as_bytes().to_vec()),
(serde_json::Value::Bool(b), _) => Value::Bool(*b),
_ => Value::Null,
}
}
fn legacy_json_to_value(value: &serde_json::Value, expected: &TypeId) -> Result<Value, String> {
if value.is_null() {
return Ok(Value::Null);
}
match expected {
TypeId::Bool => value
.as_bool()
.map(Value::Bool)
.ok_or_else(|| "expected a boolean".into()),
TypeId::Int8
| TypeId::Int16
| TypeId::Int32
| TypeId::Int64
| TypeId::UInt8
| TypeId::UInt16
| TypeId::UInt32
| TypeId::UInt64
| TypeId::TimestampNanos
| TypeId::Date32
| TypeId::Date64
| TypeId::Time64 => value
.as_i64()
.map(Value::Int64)
.ok_or_else(|| "expected a signed 64-bit integer".into()),
TypeId::Float32 | TypeId::Float64 => value
.as_f64()
.filter(|value| value.is_finite())
.map(Value::Float64)
.ok_or_else(|| "expected a finite number".into()),
TypeId::Bytes => match value.as_str() {
Some(value) => Ok(Value::Bytes(value.as_bytes().to_vec())),
None => decode_tagged_hex(value, "bytes").map(Value::Bytes),
},
TypeId::Enum { variants } => {
let bytes = match value.as_str() {
Some(value) => value.as_bytes().to_vec(),
None => decode_tagged_hex(value, "bytes")?,
};
let value =
std::str::from_utf8(&bytes).map_err(|_| "enum variant is not UTF-8".to_string())?;
if !variants.iter().any(|variant| variant == value) {
return Err("expected a declared enum variant".into());
}
Ok(Value::Bytes(bytes))
}
TypeId::Embedding { dim } => {
let values = value
.as_array()
.ok_or_else(|| "expected an embedding array".to_string())?;
if values.len() != *dim as usize {
return Err(format!("expected an embedding with {dim} values"));
}
values
.iter()
.map(|value| {
value
.as_f64()
.map(|value| value as f32)
.filter(|value| value.is_finite())
.ok_or_else(|| "embedding values must be finite numbers".to_string())
})
.collect::<Result<Vec<_>, _>>()
.map(Value::Embedding)
}
TypeId::Decimal128 { .. } => {
let object = exact_tagged_object(value, "decimal", &["unscaled"])?;
let text = object["unscaled"]
.as_str()
.ok_or_else(|| "decimal unscaled value must be a string".to_string())?;
let value = text
.parse::<i128>()
.map_err(|_| "decimal unscaled value is invalid".to_string())?;
if value.to_string() != text {
return Err("decimal unscaled value is not canonical".into());
}
Ok(Value::Decimal(value))
}
TypeId::Interval => {
let object = exact_tagged_object(value, "interval", &["months", "days", "nanos"])?;
let months = canonical_i64(&object["months"], "interval months")?;
let days = canonical_i64(&object["days"], "interval days")?
.try_into()
.map_err(|_| "interval days is outside i32 range".to_string())?;
let nanos = canonical_i64(&object["nanos"], "interval nanos")?;
Ok(Value::Interval {
months,
days,
nanos,
})
}
TypeId::Uuid => {
let bytes = decode_tagged_hex(value, "uuid")?;
let bytes: [u8; 16] = bytes
.try_into()
.map_err(|_| "UUID must contain exactly 16 bytes".to_string())?;
Ok(Value::Uuid(bytes))
}
TypeId::Json => {
let bytes = decode_tagged_hex(value, "json")?;
std::str::from_utf8(&bytes).map_err(|_| "JSON value is not UTF-8".to_string())?;
serde_json::from_slice::<serde_json::Value>(&bytes)
.map_err(|error| format!("JSON value is invalid: {error}"))?;
Ok(Value::Json(bytes))
}
TypeId::Array { .. } => Err("legacy put does not support array columns".into()),
}
}
fn exact_tagged_object<'a>(
value: &'a serde_json::Value,
expected_kind: &str,
fields: &[&str],
) -> Result<&'a serde_json::Map<String, serde_json::Value>, String> {
let object = value
.as_object()
.ok_or_else(|| format!("expected tagged {expected_kind} value"))?;
if object.len() != fields.len() + 1
|| object
.get("$mongreldb_type")
.and_then(|value| value.as_str())
!= Some(expected_kind)
|| fields.iter().any(|field| !object.contains_key(*field))
{
return Err(format!("invalid tagged {expected_kind} value"));
}
Ok(object)
}
fn decode_tagged_hex(value: &serde_json::Value, expected_kind: &str) -> Result<Vec<u8>, String> {
let object = exact_tagged_object(value, expected_kind, &["hex"])?;
let encoded = object["hex"]
.as_str()
.ok_or_else(|| format!("tagged {expected_kind} hex must be a string"))?;
if encoded.len() % 2 != 0 {
return Err(format!("tagged {expected_kind} hex has odd length"));
}
encoded
.as_bytes()
.chunks_exact(2)
.map(|pair| {
let high = hex_nibble(pair[0])?;
let low = hex_nibble(pair[1])?;
Ok((high << 4) | low)
})
.collect()
}
fn hex_nibble(value: u8) -> Result<u8, String> {
match value {
b'0'..=b'9' => Ok(value - b'0'),
b'a'..=b'f' => Ok(value - b'a' + 10),
_ => Err("hex value must use lowercase ASCII digits".into()),
}
}
fn canonical_i64(value: &serde_json::Value, field: &str) -> Result<i64, String> {
let text = value
.as_str()
.ok_or_else(|| format!("{field} must be a string"))?;
let value = text
.parse::<i64>()
.map_err(|_| format!("{field} is invalid"))?;
if value.to_string() != text {
return Err(format!("{field} is not canonical"));
}
Ok(value)
}
#[cfg(test)]
mod legacy_wire_tests {
use super::*;
#[test]
fn typed_values_are_exact_and_malformed_values_fail_closed() {
assert_eq!(
legacy_json_to_value(
&json!({"$mongreldb_type": "bytes", "hex": "00ff61"}),
&TypeId::Bytes,
)
.unwrap(),
Value::Bytes(vec![0, 0xff, b'a'])
);
for (value, ty) in [
(
json!({"$mongreldb_type": "bytes", "hex": "00FF"}),
TypeId::Bytes,
),
(
json!({"$mongreldb_type": "decimal", "unscaled": "01"}),
TypeId::Decimal128 {
precision: 38,
scale: 0,
},
),
(
json!({"$mongreldb_type": "uuid", "hex": "00"}),
TypeId::Uuid,
),
(
json!({"$mongreldb_type": "json", "hex": "7b"}),
TypeId::Json,
),
(json!([1.0]), TypeId::Embedding { dim: 2 }),
(json!([1e100, 1.0]), TypeId::Embedding { dim: 2 }),
] {
assert!(legacy_json_to_value(&value, &ty).is_err(), "{value}");
}
}
}
fn parse_cells(
row: &[serde_json::Value],
schema: &mongreldb_core::schema::Schema,
) -> Result<Vec<(u16, Value)>, String> {
if row.len() & 1 != 0 {
return Err("row must be an even-length array of [col_id, value] pairs".into());
}
let mut out = Vec::with_capacity(row.len() / 2);
let mut seen = std::collections::HashSet::new();
for chunk in row.chunks(2) {
let col_id = chunk[0]
.as_u64()
.and_then(|value| u16::try_from(value).ok())
.ok_or("column id must be an unsigned 16-bit integer")?;
if !seen.insert(col_id) {
return Err(format!("duplicate column id {col_id}"));
}
let expected = schema
.columns
.iter()
.find(|c| c.id == col_id)
.map(|c| c.ty.clone())
.ok_or_else(|| format!("unknown column id {col_id}"))?;
let val = legacy_json_to_value(&chunk[1], &expected)?;
out.push((col_id, val));
}
Ok(out)
}
pub(crate) fn validate_table_name(name: &str) -> Result<(), String> {
if name.is_empty() {
return Err("table name must not be empty".into());
}
if name.contains('/') || name.contains('\\') || name.contains('\0') {
return Err("table name contains invalid characters".into());
}
Ok(())
}
async fn put_row(
State(state): State<Arc<AppState>>,
OptionalPrincipal(principal): OptionalPrincipal,
Path(name): Path<String>,
Json(req): Json<PutRequest>,
) -> Response {
let handle = match state.db.table(&name) {
Ok(h) => h,
Err(e) => return (StatusCode::NOT_FOUND, e.to_string()).into_response(),
};
let schema = handle.lock().schema().clone();
let row = match parse_cells(&req.row, &schema) {
Ok(r) => r,
Err(msg) => return (StatusCode::BAD_REQUEST, msg).into_response(),
};
state.metrics.inc_puts();
let principal = request_principal(&state, &principal);
match state.db.put_for(&name, row, principal.as_ref()) {
Ok(rid) => Json(json!({ "row_id": rid.0.to_string() })).into_response(),
Err(e) => (status_for_error(&e), e.to_string()).into_response(),
}
}
async fn count(
State(state): State<Arc<AppState>>,
OptionalPrincipal(principal): OptionalPrincipal,
Path(name): Path<String>,
) -> Response {
let principal = request_principal(&state, &principal);
match state.db.count_for(&name, principal.as_ref()) {
Ok(count) => Json(json!({ "count": count })).into_response(),
Err(error) => (status_for_error(&error), error.to_string()).into_response(),
}
}
async fn commit(
State(state): State<Arc<AppState>>,
OptionalPrincipal(principal): OptionalPrincipal,
Path(name): Path<String>,
) -> Response {
if let Err(error) = state.db.require_for(
request_principal(&state, &principal).as_ref(),
&mongreldb_core::Permission::Update {
table: name.clone(),
},
) {
return (status_for_error(&error), error.to_string()).into_response();
}
let handle = match state.db.table(&name) {
Ok(h) => h,
Err(e) => return (StatusCode::NOT_FOUND, e.to_string()).into_response(),
};
let mut g = handle.lock();
state.metrics.inc_commits();
match g.commit() {
Ok(epoch) => Json(json!({
"epoch": epoch.0,
"epoch_text": epoch.0.to_string()
}))
.into_response(),
Err(error) => crate::kit::durable_core_error_response(&error)
.unwrap_or_else(|| (status_for_error(&error), error.to_string()).into_response()),
}
}
#[derive(Deserialize)]
struct SqlRequest {
sql: String,
#[serde(default)]
format: Option<String>,
#[serde(default)]
query_id: Option<QueryId>,
#[serde(default)]
timeout_ms: Option<u64>,
#[serde(default)]
max_output_rows: Option<u64>,
#[serde(default)]
max_output_bytes: Option<u64>,
#[serde(default)]
idempotency_key: Option<String>,
#[serde(default)]
pagination: Option<SqlPaginationRequest>,
}
#[derive(Clone, Debug, Deserialize, Serialize)]
struct SqlPaginationRequest {
page_size_rows: u64,
projection: Vec<String>,
#[serde(default)]
max_page_bytes: Option<u64>,
#[serde(default)]
max_page_tokens: Option<u64>,
}
#[derive(Clone)]
struct ResolvedSqlPagination {
projection: Vec<String>,
limits: sql_pages::SqlPageLimits,
}
struct ResolvedSqlRequest {
request: SqlRequest,
output_limits: (usize, usize),
idempotency: Option<sql_idempotency::SqlIdempotencyExecution>,
pagination: Option<ResolvedSqlPagination>,
}
fn query_error_response(
error: &mongreldb_query::MongrelQueryError,
query_id: Option<QueryId>,
) -> Response {
query_error_response_with_status(error, query_id, None)
}
fn query_error_response_with_status(
error: &mongreldb_query::MongrelQueryError,
query_id: Option<QueryId>,
status: Option<&mongreldb_query::QueryStatus>,
) -> Response {
use mongreldb_query::MongrelQueryError;
let (base_code, id) = match error {
MongrelQueryError::QueryCancelled { query_id, .. } => ("QUERY_CANCELLED", Some(*query_id)),
MongrelQueryError::DeadlineExceeded { query_id, .. } => {
("DEADLINE_EXCEEDED", Some(*query_id))
}
MongrelQueryError::QueryIdConflict { query_id } => ("QUERY_ID_CONFLICT", Some(*query_id)),
MongrelQueryError::QueryRegistryFull => ("QUERY_REGISTRY_FULL", query_id),
MongrelQueryError::ResultLimitExceeded { query_id, .. } => {
("RESULT_LIMIT_EXCEEDED", Some(*query_id))
}
MongrelQueryError::TransactionAborted => ("TRANSACTION_ABORTED", query_id),
MongrelQueryError::NoSqlTransaction => ("NO_SQL_TRANSACTION", query_id),
MongrelQueryError::SavepointNotFound { .. } => ("SAVEPOINT_NOT_FOUND", query_id),
MongrelQueryError::CommitOutcome { query_id, .. } => ("COMMIT_OUTCOME", Some(*query_id)),
MongrelQueryError::OutcomeUnknown { query_id, .. } => {
("QUERY_OUTCOME_UNKNOWN", Some(*query_id))
}
_ => ("QUERY_FAILED", query_id),
};
let (
error_committed,
error_committed_statements,
error_last_commit_epoch,
error_first_commit_statement_index,
error_last_commit_statement_index,
) = match error {
MongrelQueryError::QueryCancelled {
committed,
committed_statements,
last_commit_epoch,
first_commit_statement_index,
last_commit_statement_index,
..
}
| MongrelQueryError::DeadlineExceeded {
committed,
committed_statements,
last_commit_epoch,
first_commit_statement_index,
last_commit_statement_index,
..
}
| MongrelQueryError::ResultLimitExceeded {
committed,
committed_statements,
last_commit_epoch,
first_commit_statement_index,
last_commit_statement_index,
..
} => (
*committed,
*committed_statements,
*last_commit_epoch,
*first_commit_statement_index,
*last_commit_statement_index,
),
MongrelQueryError::CommitOutcome {
committed,
committed_statements,
last_commit_epoch,
first_commit_statement_index,
last_commit_statement_index,
..
} => (
*committed,
*committed_statements,
*last_commit_epoch,
*first_commit_statement_index,
*last_commit_statement_index,
),
_ => (false, 0, None, None, None),
};
let committed = status.map_or_else(
|| error_committed,
|status| status.durable_outcome.committed,
);
let outcome_unknown = matches!(error, MongrelQueryError::OutcomeUnknown { .. })
|| status.is_some_and(|status| status.outcome_unknown);
let code = status
.and_then(|status| {
status
.terminal_error
.as_ref()
.map(|error| error.code.as_str())
})
.unwrap_or(match (base_code, committed) {
("QUERY_CANCELLED", true) => "QUERY_CANCELLED_AFTER_COMMIT",
("DEADLINE_EXCEEDED", true) => "DEADLINE_AFTER_COMMIT",
_ => base_code,
});
let response_status = if outcome_unknown {
"outcome_unknown"
} else {
status
.and_then(mongreldb_query::QueryStatus::terminal_state)
.map(terminal_state_name)
.unwrap_or_else(|| match (error, committed) {
(MongrelQueryError::QueryCancelled { .. }, true) => "cancelled_after_commit",
(MongrelQueryError::QueryCancelled { .. }, false) => "cancelled_before_commit",
(MongrelQueryError::DeadlineExceeded { .. }, true) => "deadline_after_commit",
(MongrelQueryError::DeadlineExceeded { .. }, false) => "deadline_before_commit",
(_, true) => "committed_with_error",
_ => "failed_before_commit",
})
};
let (completed_statements, statement_index) = status.map_or_else(
|| match error {
MongrelQueryError::QueryCancelled {
completed_statements,
cancelled_statement_index,
..
}
| MongrelQueryError::DeadlineExceeded {
completed_statements,
cancelled_statement_index,
..
} => (*completed_statements, *cancelled_statement_index),
MongrelQueryError::ResultLimitExceeded {
completed_statements,
statement_index,
..
} => (*completed_statements, *statement_index),
MongrelQueryError::CommitOutcome {
completed_statements,
statement_index,
..
} => (*completed_statements, *statement_index),
_ => (0, 0),
},
|status| (status.completed_statements, status.statement_index),
);
let committed_statements = status.map_or(error_committed_statements, |status| {
status.durable_outcome.committed_statements
});
let last_commit_epoch = status.map_or(error_last_commit_epoch, |status| {
status.durable_outcome.last_commit_epoch
});
let first_commit_statement_index = status
.map_or(error_first_commit_statement_index, |status| {
status.durable_outcome.first_commit_statement_index
});
let last_commit_statement_index = status.map_or(error_last_commit_statement_index, |status| {
status.durable_outcome.last_commit_statement_index
});
let cancellation_reason = status
.map(|status| status.cancellation_reason)
.or(match error {
MongrelQueryError::QueryCancelled { reason, .. } => Some(*reason),
MongrelQueryError::DeadlineExceeded { .. } => Some(CancellationReason::Deadline),
_ => None,
})
.map(cancellation_reason_name);
let cancel_outcome = match error {
MongrelQueryError::QueryCancelled { .. } | MongrelQueryError::DeadlineExceeded { .. } => {
Some("accepted")
}
_ => status.and_then(query_cancel_outcome),
};
let outcome = if outcome_unknown {
json!({
"committed": null,
"committed_statements": null,
"last_commit_epoch": null,
"last_commit_epoch_text": null,
"first_commit_statement_index": null,
"last_commit_statement_index": null,
"completed_statements": null,
"statement_index": null,
"serialization": "unknown",
})
} else {
status.map_or_else(
|| {
json!({
"committed": committed,
"committed_statements": committed_statements,
"last_commit_epoch": last_commit_epoch,
"last_commit_epoch_text": epoch_text(last_commit_epoch),
"first_commit_statement_index": first_commit_statement_index,
"last_commit_statement_index": last_commit_statement_index,
"completed_statements": completed_statements,
"statement_index": statement_index,
"serialization": "unknown",
})
},
|status| query_outcome_json(Some(status)),
)
};
let http_status = match code {
"QUERY_CANCELLED_AFTER_COMMIT" | "DEADLINE_AFTER_COMMIT" => StatusCode::CONFLICT,
"QUERY_CANCELLED" => client_closed_request_status(),
"DEADLINE_EXCEEDED" => StatusCode::GATEWAY_TIMEOUT,
_ => status_for_query_error(error),
};
let mut response = (
http_status,
Json(json!({
"query_id": id.map(|value| value.to_string()),
"status": response_status,
"terminal_state": response_status,
"committed": (!outcome_unknown).then_some(committed),
"committed_statements": (!outcome_unknown).then_some(committed_statements),
"last_commit_epoch": (!outcome_unknown).then_some(last_commit_epoch).flatten(),
"last_commit_epoch_text": (!outcome_unknown).then_some(epoch_text(last_commit_epoch)).flatten(),
"first_commit_statement_index": (!outcome_unknown).then_some(first_commit_statement_index).flatten(),
"last_commit_statement_index": (!outcome_unknown).then_some(last_commit_statement_index).flatten(),
"completed_statements": (!outcome_unknown).then_some(completed_statements),
"statement_index": (!outcome_unknown).then_some(statement_index),
"cancel_outcome": cancel_outcome,
"cancellation_reason": cancellation_reason,
"retryable": matches!(error, MongrelQueryError::QueryRegistryFull),
"server_state": status.map(|status| query_phase_name(status.phase)),
"outcome": outcome,
"error": {
"code": code,
"message": error.to_string(),
"query_id": id.map(|value| value.to_string()),
"committed": (!outcome_unknown).then_some(committed),
"retryable": matches!(error, MongrelQueryError::QueryRegistryFull),
}
})),
)
.into_response();
if let Some(id) = id {
add_query_id_header(&mut response, id);
}
response
}
fn record_query_error(metrics: &metrics::Metrics, error: &mongreldb_query::MongrelQueryError) {
match error {
mongreldb_query::MongrelQueryError::QueryCancelled { reason, .. } => {
metrics.inc_sql_cancelled(*reason)
}
mongreldb_query::MongrelQueryError::DeadlineExceeded { .. } => {
metrics.inc_sql_deadline_exceeded();
metrics.inc_sql_cancelled(CancellationReason::Deadline);
}
_ => {}
}
}
fn tracked_query_error_response(
state: &AppState,
error: &mongreldb_query::MongrelQueryError,
query_id: Option<QueryId>,
) -> Response {
record_query_error(&state.metrics, error);
if let mongreldb_query::MongrelQueryError::QueryCancelled { query_id, .. } = error {
if let Some(requested_at) = state
.query_registry
.status(*query_id)
.and_then(|status| status.cancel_requested_at)
{
state
.metrics
.observe_sql_cancel_latency(requested_at.elapsed());
}
}
let status = if matches!(
error,
mongreldb_query::MongrelQueryError::QueryIdConflict { .. }
| mongreldb_query::MongrelQueryError::QueryRegistryFull
) {
None
} else {
query_id
.or(match error {
mongreldb_query::MongrelQueryError::QueryCancelled { query_id, .. }
| mongreldb_query::MongrelQueryError::DeadlineExceeded { query_id, .. }
| mongreldb_query::MongrelQueryError::CommitOutcome { query_id, .. }
| mongreldb_query::MongrelQueryError::OutcomeUnknown { query_id, .. } => {
Some(*query_id)
}
_ => None,
})
.and_then(|query_id| state.query_registry.status(query_id))
};
query_error_response_with_status(error, query_id, status.as_ref())
}
fn add_query_id_header(response: &mut Response, query_id: QueryId) {
if let Ok(value) = axum::http::HeaderValue::from_str(&query_id.to_string()) {
response.headers_mut().insert("x-mongreldb-query-id", value);
}
}
fn with_query_id(mut response: Response, query_id: QueryId) -> Response {
add_query_id_header(&mut response, query_id);
response
}
fn bad_query_control_request(message: impl Into<String>, query_id: Option<QueryId>) -> Response {
let mut response = (
StatusCode::BAD_REQUEST,
Json(json!({
"query_id": query_id.map(|value| value.to_string()),
"status": "failed_before_commit",
"terminal_state": "failed_before_commit",
"committed": false,
"committed_statements": 0,
"last_commit_epoch": null,
"last_commit_epoch_text": null,
"first_commit_statement_index": null,
"last_commit_statement_index": null,
"completed_statements": 0,
"statement_index": 0,
"cancel_outcome": null,
"cancellation_reason": null,
"retryable": false,
"server_state": "failed",
"outcome": {
"committed": false,
"committed_statements": 0,
"last_commit_epoch": null,
"last_commit_epoch_text": null,
"first_commit_statement_index": null,
"last_commit_statement_index": null,
"completed_statements": 0,
"statement_index": 0,
"serialization": "not_started",
},
"error": {
"code": "INVALID_QUERY_OPTIONS",
"message": message.into(),
"query_id": query_id.map(|value| value.to_string()),
"committed": false,
"retryable": false,
}
})),
)
.into_response();
if let Some(query_id) = query_id {
add_query_id_header(&mut response, query_id);
}
response
}
fn resolve_query_options(
state: &AppState,
headers: &axum::http::HeaderMap,
body_query_id: Option<QueryId>,
body_timeout_ms: Option<u64>,
owner: String,
session_id: Option<String>,
) -> std::result::Result<(SqlQueryOptions, QueryId), Box<Response>> {
let query_id = match body_query_id {
Some(query_id) => query_id,
None => match headers.get("x-mongreldb-query-id") {
Some(value) => {
let value = value.to_str().map_err(|_| {
Box::new(bad_query_control_request(
"X-MongrelDB-Query-ID is not valid text",
None,
))
})?;
value
.parse()
.map_err(|error: mongreldb_query::MongrelQueryError| {
Box::new(bad_query_control_request(error.to_string(), None))
})?
}
None => {
QueryId::random().map_err(|error| Box::new(query_error_response(&error, None)))?
}
},
};
let timeout_ms = match body_timeout_ms {
Some(timeout_ms) => timeout_ms,
None => match headers.get("x-mongreldb-timeout-ms") {
Some(value) => value
.to_str()
.ok()
.and_then(|value| value.parse::<u64>().ok())
.ok_or_else(|| {
Box::new(bad_query_control_request(
"X-MongrelDB-Timeout-Ms must be a positive integer",
Some(query_id),
))
})?,
None => state
.sql_default_timeout
.as_millis()
.min(u128::from(u64::MAX)) as u64,
},
};
if timeout_ms == 0 {
return Err(Box::new(bad_query_control_request(
"timeout_ms must be positive",
Some(query_id),
)));
}
let timeout = std::time::Duration::from_millis(timeout_ms);
if timeout > state.sql_max_timeout {
return Err(Box::new(bad_query_control_request(
format!(
"timeout_ms exceeds server maximum of {}",
state.sql_max_timeout.as_millis()
),
Some(query_id),
)));
}
Ok((
SqlQueryOptions {
query_id: Some(query_id),
timeout: Some(timeout),
owner: Some(owner),
session_id,
parent_control: None,
},
query_id,
))
}
fn resolve_sql_output_limits(
state: &AppState,
request: &SqlRequest,
query_id: QueryId,
) -> std::result::Result<(usize, usize), Box<Response>> {
fn resolve(
requested: Option<u64>,
configured: usize,
name: &str,
query_id: QueryId,
) -> std::result::Result<usize, Box<Response>> {
if requested == Some(0) {
return Err(Box::new(bad_query_control_request(
format!("{name} must be positive"),
Some(query_id),
)));
}
let requested = requested
.and_then(|value| usize::try_from(value).ok())
.unwrap_or(usize::MAX);
Ok(requested.min(configured))
}
Ok((
resolve(
request.max_output_rows,
state.sql_max_output_rows,
"max_output_rows",
query_id,
)?,
resolve(
request.max_output_bytes,
state.sql_max_output_bytes,
"max_output_bytes",
query_id,
)?,
))
}
fn resolve_sql_pagination(
headers: &axum::http::HeaderMap,
request: &SqlRequest,
output_limits: (usize, usize),
registration: RegisteredQueryGuard,
query_id: QueryId,
) -> Result<(RegisteredQueryGuard, Option<ResolvedSqlPagination>), Box<Response>> {
let Some(pagination) = request.pagination.as_ref() else {
return Ok((registration, None));
};
if requested_sql_idempotency_key(headers, request)
.ok()
.flatten()
.is_some()
{
return Err(Box::new(registered_sql_error_response(
registration,
query_id,
StatusCode::BAD_REQUEST,
"INCOMPATIBLE_SQL_CONTROLS",
"idempotency_key cannot be combined with SQL pagination",
false,
)));
}
if request
.format
.as_deref()
.is_some_and(|format| format != "json")
{
return Err(Box::new(registered_sql_error_response(
registration,
query_id,
StatusCode::BAD_REQUEST,
"PAGINATION_REQUIRES_JSON",
"SQL pagination supports JSON responses only",
false,
)));
}
registration.query().set_sql_metadata(&request.sql);
if !mongreldb_query::is_single_read_only_query(&request.sql) {
return Err(Box::new(registered_sql_error_response(
registration,
query_id,
StatusCode::BAD_REQUEST,
"PAGINATION_REQUIRES_SINGLE_READ_QUERY",
"SQL pagination accepts exactly one read-only query statement",
false,
)));
}
let page_size = match usize::try_from(pagination.page_size_rows) {
Ok(0) | Err(_) => {
return Err(Box::new(registered_sql_error_response(
registration,
query_id,
StatusCode::BAD_REQUEST,
"INVALID_PAGINATION_OPTIONS",
"pagination.page_size_rows must be positive",
false,
)))
}
Ok(value) => value.min(output_limits.0),
};
if pagination.projection.is_empty() || pagination.projection.len() > 128 {
return Err(Box::new(registered_sql_error_response(
registration,
query_id,
StatusCode::BAD_REQUEST,
"INVALID_SQL_PROJECTION",
"pagination.projection must contain between 1 and 128 output column names",
false,
)));
}
let mut seen = std::collections::HashSet::new();
let metadata_bytes = pagination
.projection
.iter()
.map(String::len)
.fold(0usize, usize::saturating_add);
if metadata_bytes > 16 * 1024
|| pagination.projection.iter().any(|column| {
column.is_empty()
|| column == "*"
|| column.len() > 256
|| !seen.insert(column.as_str())
})
{
return Err(Box::new(registered_sql_error_response(
registration,
query_id,
StatusCode::BAD_REQUEST,
"INVALID_SQL_PROJECTION",
"pagination.projection requires unique explicit output names of at most 256 bytes",
false,
)));
}
let max_page_bytes = match pagination.max_page_bytes {
Some(0) => {
return Err(Box::new(registered_sql_error_response(
registration,
query_id,
StatusCode::BAD_REQUEST,
"INVALID_PAGINATION_OPTIONS",
"pagination.max_page_bytes must be positive",
false,
)))
}
Some(value) => usize::try_from(value)
.unwrap_or(usize::MAX)
.min(output_limits.1),
None => output_limits.1.min(1024 * 1024),
};
let token_cap = (output_limits.1.saturating_add(3) / 4).max(1);
let max_page_tokens = match pagination.max_page_tokens {
Some(0) => {
return Err(Box::new(registered_sql_error_response(
registration,
query_id,
StatusCode::BAD_REQUEST,
"INVALID_PAGINATION_OPTIONS",
"pagination.max_page_tokens must be positive",
false,
)))
}
Some(value) => usize::try_from(value).unwrap_or(usize::MAX).min(token_cap),
None => (max_page_bytes.saturating_add(3) / 4).max(1),
};
Ok((
registration,
Some(ResolvedSqlPagination {
projection: pagination.projection.clone(),
limits: sql_pages::SqlPageLimits {
rows: page_size,
bytes: max_page_bytes,
tokens: max_page_tokens,
},
}),
))
}
fn requested_sql_idempotency_key(
headers: &axum::http::HeaderMap,
request: &SqlRequest,
) -> Result<Option<String>, &'static str> {
let header = match headers.get("idempotency-key") {
Some(value) => Some(
value
.to_str()
.map_err(|_| "Idempotency-Key must be valid UTF-8")?,
),
None => None,
};
match (request.idempotency_key.as_deref(), header) {
(Some(body), Some(header)) if body != header => {
Err("body idempotency_key and Idempotency-Key header must match")
}
(Some(body), _) => Ok(Some(body.to_owned())),
(None, Some(header)) => Ok(Some(header.to_owned())),
(None, None) => Ok(None),
}
}
fn sql_idempotency_binding(
request: &SqlRequest,
output_limits: (usize, usize),
session_id: Option<&str>,
expires_after_ms: u64,
) -> Result<sql_idempotency::SqlIdempotencyBinding, serde_json::Error> {
let request_semantics = serde_json::to_vec(&json!({
"format": request.format.as_deref().unwrap_or("json"),
"max_output_rows": output_limits.0,
"max_output_bytes": output_limits.1,
"pagination": request.pagination.as_ref(),
}))?;
let session_semantics = session_id.map_or_else(
|| b"ephemeral".to_vec(),
|session_id| {
let mut semantics = b"session\0".to_vec();
semantics.extend_from_slice(session_id.as_bytes());
semantics
},
);
Ok(sql_idempotency::SqlIdempotencyBinding {
sql_fingerprint: mongreldb_query::normalized_sql_fingerprint(&request.sql),
parameter_hash: sql_idempotency::hash(b"[]"),
request_semantics_hash: sql_idempotency::hash(&request_semantics),
session_semantics_hash: sql_idempotency::hash(&session_semantics),
expires_after_ms,
})
}
struct SqlIdempotencyContext<'a> {
headers: &'a axum::http::HeaderMap,
request: &'a SqlRequest,
output_limits: (usize, usize),
owner: &'a str,
session_id: Option<&'a str>,
session_in_transaction: bool,
query_id: QueryId,
}
async fn begin_sql_idempotency(
state: &AppState,
context: SqlIdempotencyContext<'_>,
registration: RegisteredQueryGuard,
) -> Result<
(
RegisteredQueryGuard,
Option<sql_idempotency::SqlIdempotencyExecution>,
),
Response,
> {
let SqlIdempotencyContext {
headers,
request,
output_limits,
owner,
session_id,
session_in_transaction,
query_id,
} = context;
let key = match requested_sql_idempotency_key(headers, request) {
Ok(key) => key,
Err(message) => {
return Err(registered_sql_error_response(
registration,
query_id,
StatusCode::BAD_REQUEST,
"INVALID_IDEMPOTENCY_KEY",
message,
false,
))
}
};
let Some(key) = key else {
return Ok((registration, None));
};
match mongreldb_query::classify_sql_idempotency(&request.sql) {
mongreldb_query::SqlIdempotencyClass::ReadOnly
| mongreldb_query::SqlIdempotencyClass::Unsupported => {
return Err(registered_sql_error_response(
registration,
query_id,
StatusCode::BAD_REQUEST,
"IDEMPOTENCY_REQUIRES_SINGLE_WRITE",
"idempotency_key accepts one non-transaction SQL write statement",
false,
));
}
mongreldb_query::SqlIdempotencyClass::SingleWrite => {}
}
if let Err(message) = sql_idempotency::SqlIdempotencyStore::validate_key(&key) {
return Err(registered_sql_error_response(
registration,
query_id,
StatusCode::BAD_REQUEST,
"INVALID_IDEMPOTENCY_KEY",
message,
false,
));
}
if request
.format
.as_deref()
.is_some_and(|format| format != "json")
{
return Err(registered_sql_error_response(
registration,
query_id,
StatusCode::BAD_REQUEST,
"IDEMPOTENCY_REQUIRES_JSON",
"SQL idempotency supports buffered JSON responses only",
false,
));
}
if session_in_transaction {
return Err(registered_sql_error_response(
registration,
query_id,
StatusCode::CONFLICT,
"IDEMPOTENCY_UNSUPPORTED_IN_TRANSACTION",
"SQL idempotency cannot be used inside an open session transaction",
false,
));
}
registration.query().set_sql_metadata(&request.sql);
let binding = match sql_idempotency_binding(
request,
output_limits,
session_id,
state.sql_idempotency.expires_after_ms(),
) {
Ok(binding) => binding,
Err(_) => {
return Err(registered_sql_error_response(
registration,
query_id,
StatusCode::INTERNAL_SERVER_ERROR,
"SERIALIZATION_FAILED",
"failed to serialize SQL idempotency request semantics",
false,
))
}
};
let begin = tokio::select! {
begin = state.sql_idempotency.begin(owner, &key, binding) => begin,
_ = registration.query().control().cancelled() => {
return Err(tracked_query_error_response(
state,
&cancellation_checkpoint_error(registration.query()),
Some(query_id),
));
}
};
match begin {
sql_idempotency::BeginResult::Execute(execution) => Ok((registration, Some(execution))),
sql_idempotency::BeginResult::Replay {
receipt,
expires_at_ms,
} => match restore_idempotency_replay(registration, &receipt) {
Ok(()) => Err(sql_idempotency_receipt_response(
query_id,
&receipt,
true,
expires_at_ms,
true,
)),
Err(error) => Err(tracked_query_error_response(state, &error, Some(query_id))),
},
sql_idempotency::BeginResult::Mismatch => Err(registered_sql_error_response(
registration,
query_id,
StatusCode::CONFLICT,
"IDEMPOTENCY_KEY_REUSE_MISMATCH",
"idempotency key was already used with different SQL or request semantics",
false,
)),
sql_idempotency::BeginResult::Indeterminate { created_at_ms } => Err(
sql_idempotency_indeterminate_response(registration, query_id, created_at_ms),
),
sql_idempotency::BeginResult::Full => Err(registered_sql_error_response(
registration,
query_id,
StatusCode::SERVICE_UNAVAILABLE,
"IDEMPOTENCY_STORE_FULL",
"SQL idempotency receipt store is full",
true,
)),
sql_idempotency::BeginResult::Unavailable => Err(registered_sql_error_response(
registration,
query_id,
StatusCode::SERVICE_UNAVAILABLE,
"IDEMPOTENCY_STORE_UNAVAILABLE",
"could not durably reserve the SQL idempotency key",
true,
)),
}
}
fn restore_idempotency_replay(
registration: RegisteredQueryGuard,
receipt: &sql_idempotency::SqlDurableReceipt,
) -> mongreldb_query::Result<()> {
use mongreldb_query::{
DurableOutcome, QueryTerminalError, QueryTerminalErrorCategory, QueryTerminalState,
SerializationOutcome,
};
let invalid_receipt = |field: &str| {
mongreldb_query::MongrelQueryError::InvalidQueryState(format!(
"durable SQL idempotency receipt has invalid {field}"
))
};
let terminal_state = match receipt.status.as_str() {
"completed" => QueryTerminalState::Completed,
"failed_before_commit" => QueryTerminalState::FailedBeforeCommit,
"cancelled_before_commit" => QueryTerminalState::CancelledBeforeCommit,
"deadline_before_commit" => QueryTerminalState::DeadlineBeforeCommit,
"committed" => QueryTerminalState::Committed,
"committed_with_error" => QueryTerminalState::CommittedWithError,
"partially_committed" => QueryTerminalState::PartiallyCommitted,
"cancelled_after_commit" => QueryTerminalState::CancelledAfterCommit,
"deadline_after_commit" => QueryTerminalState::DeadlineAfterCommit,
_ => return Err(invalid_receipt("terminal state")),
};
let serialization = match receipt.outcome.serialization.as_str() {
"not_started" => SerializationOutcome::NotStarted,
"in_progress" => SerializationOutcome::InProgress,
"succeeded" => SerializationOutcome::Succeeded,
"failed" => SerializationOutcome::Failed,
_ => return Err(invalid_receipt("serialization state")),
};
let terminal_error = match receipt.terminal_error.as_ref() {
Some(error) => Some(QueryTerminalError {
code: error.code.clone(),
category: match error.category.as_str() {
"cancellation" => QueryTerminalErrorCategory::Cancellation,
"deadline" => QueryTerminalErrorCategory::Deadline,
"result_limit" => QueryTerminalErrorCategory::ResultLimit,
"serialization" => QueryTerminalErrorCategory::Serialization,
"execution" => QueryTerminalErrorCategory::Execution,
_ => return Err(invalid_receipt("terminal error category")),
},
}),
None => None,
};
let cancellation_reason = CancellationReason::from_protocol_str(&receipt.cancellation_reason)
.ok_or_else(|| invalid_receipt("cancellation reason"))?;
let phase = match receipt.server_state.as_str() {
"completed" => SqlQueryPhase::Completed,
"cancelled" => SqlQueryPhase::Cancelled,
"failed" => SqlQueryPhase::Failed,
_ => return Err(invalid_receipt("server state")),
};
let query = registration.into_query();
query.restore_replayed_outcome(
DurableOutcome {
committed: receipt.outcome.committed,
committed_statements: receipt.outcome.committed_statements,
last_commit_epoch: receipt.outcome.last_commit_epoch,
first_commit_statement_index: receipt.outcome.first_commit_statement_index,
last_commit_statement_index: receipt.outcome.last_commit_statement_index,
},
receipt.outcome.completed_statements,
receipt.outcome.statement_index,
serialization,
terminal_error,
terminal_state,
cancellation_reason,
phase,
);
query.try_complete()
}
fn sql_idempotency_indeterminate_response(
registration: RegisteredQueryGuard,
query_id: QueryId,
created_at_ms: Option<u64>,
) -> Response {
registration.query().mark_outcome_unknown();
registration.fail();
with_query_id(
(
StatusCode::CONFLICT,
Json(json!({
"query_id": query_id.to_string(),
"status": "outcome_unknown",
"terminal_state": "outcome_unknown",
"committed": null,
"committed_statements": null,
"last_commit_epoch": null,
"last_commit_epoch_text": null,
"first_commit_statement_index": null,
"last_commit_statement_index": null,
"completed_statements": null,
"statement_index": null,
"cancel_outcome": null,
"cancellation_reason": null,
"retryable": false,
"server_state": "failed",
"idempotency_replayed": true,
"idempotency_intent_created_at_ms": created_at_ms,
"outcome": {
"committed": null,
"committed_statements": null,
"last_commit_epoch": null,
"last_commit_epoch_text": null,
"first_commit_statement_index": null,
"last_commit_statement_index": null,
"completed_statements": null,
"statement_index": null,
"serialization": "unknown",
},
"error": {
"code": "QUERY_OUTCOME_UNKNOWN",
"message": "a durable write intent exists without a durable receipt; the SQL was not re-executed",
"query_id": query_id.to_string(),
"committed": null,
"retryable": false,
}
})),
)
.into_response(),
query_id,
)
}
fn registered_sql_error_response(
registration: RegisteredQueryGuard,
query_id: QueryId,
status: StatusCode,
code: &'static str,
message: impl Into<String>,
retryable: bool,
) -> Response {
let message = message.into();
registration
.query()
.record_terminal_error(code, mongreldb_query::QueryTerminalErrorCategory::Execution);
registration.fail();
with_query_id(
(
status,
Json(json!({
"query_id": query_id.to_string(),
"status": "failed_before_commit",
"terminal_state": "failed_before_commit",
"committed": false,
"committed_statements": 0,
"last_commit_epoch": null,
"last_commit_epoch_text": null,
"first_commit_statement_index": null,
"last_commit_statement_index": null,
"completed_statements": 0,
"statement_index": 0,
"cancel_outcome": null,
"cancellation_reason": null,
"retryable": retryable,
"server_state": "failed",
"outcome": {
"committed": false,
"committed_statements": 0,
"last_commit_epoch": null,
"last_commit_epoch_text": null,
"first_commit_statement_index": null,
"last_commit_statement_index": null,
"completed_statements": 0,
"statement_index": 0,
"serialization": "not_started",
},
"error": {
"code": code,
"message": message,
"query_id": query_id.to_string(),
"committed": false,
"retryable": retryable,
}
})),
)
.into_response(),
query_id,
)
}
fn register_controlled_query(
state: &AppState,
session: &MongrelSession,
options: SqlQueryOptions,
) -> std::result::Result<RegisteredSqlQuery, mongreldb_query::MongrelQueryError> {
let query_id = options.query_id.ok_or_else(|| {
mongreldb_query::MongrelQueryError::InvalidQueryState(
"server query registration requires a query id".into(),
)
})?;
let owner = options.owner.clone().unwrap_or_default();
let session_id = options.session_id.clone();
let _lifecycle = state
.query_lifecycle
.lock()
.unwrap_or_else(|error| error.into_inner());
let pre_cancel_reason = match state.pre_cancellations.lookup_for_registration(
query_id,
&owner,
session_id.as_deref(),
) {
pre_cancel::RegistrationLookup::NoReservation => None,
pre_cancel::RegistrationLookup::Matching(reason) => Some(reason),
pre_cancel::RegistrationLookup::ReservedByAnotherIdentity => {
return Err(mongreldb_query::MongrelQueryError::QueryIdConflict { query_id });
}
};
let query = session.register_query(options)?;
let Some(reason) = pre_cancel_reason else {
return Ok(query);
};
state
.pre_cancellations
.take(query_id, &owner, session_id.as_deref());
query.request_cancel(reason);
let error = query.checkpoint().err().unwrap_or_else(|| {
mongreldb_query::MongrelQueryError::InvalidQueryState(format!(
"pre-cancelled query {query_id} remained runnable"
))
});
query.fail();
Err(error)
}
async fn acquire_sql_permit(
state: &AppState,
session: &MongrelSession,
query: &RegisteredSqlQuery,
) -> std::result::Result<tokio::sync::OwnedSemaphorePermit, mongreldb_query::MongrelQueryError> {
session.fire_test_hook(mongreldb_query::SqlTestHookPoint::WaitingForSqlPermit);
tokio::select! {
permit = Arc::clone(&state.sql_semaphore).acquire_owned() => permit.map_err(|_| {
mongreldb_query::MongrelQueryError::InvalidQueryState(
"SQL admission semaphore closed".into(),
)
}),
_ = query.control().cancelled() => Err(cancellation_checkpoint_error(query)),
}
}
fn caller_may_manage_query(
state: &AppState,
principal: &Option<mongreldb_core::Principal>,
owner: Option<&str>,
) -> bool {
let current = current_request_principal(state, principal);
if (principal.is_some()
|| state.auth_token.is_some()
|| state.user_auth
|| state.db.require_auth_enabled())
&& current.is_none()
{
return false;
}
current.is_some_and(|principal| principal.is_admin)
|| owner == Some(request_owner(state, principal).as_str())
}
fn query_phase_name(phase: SqlQueryPhase) -> &'static str {
match phase {
SqlQueryPhase::Queued => "queued",
SqlQueryPhase::Planning => "planning",
SqlQueryPhase::Executing => "executing",
SqlQueryPhase::Streaming => "streaming",
SqlQueryPhase::Serializing => "serializing",
SqlQueryPhase::CommitCritical => "commit_critical",
SqlQueryPhase::Cancelling => "cancelling",
SqlQueryPhase::Completed => "completed",
SqlQueryPhase::Failed => "failed",
SqlQueryPhase::Cancelled => "cancelled",
}
}
fn commit_fence_outcome_name(outcome: mongreldb_query::CommitFenceOutcome) -> &'static str {
match outcome {
mongreldb_query::CommitFenceOutcome::NotReached => "not_reached",
mongreldb_query::CommitFenceOutcome::CancelWon => "cancel_won",
mongreldb_query::CommitFenceOutcome::CommitWon => "commit_won",
}
}
fn terminal_state_name(state: mongreldb_query::QueryTerminalState) -> &'static str {
use mongreldb_query::QueryTerminalState;
match state {
QueryTerminalState::OutcomeUnknown => "outcome_unknown",
QueryTerminalState::Completed => "completed",
QueryTerminalState::FailedBeforeCommit => "failed_before_commit",
QueryTerminalState::CancelledBeforeCommit => "cancelled_before_commit",
QueryTerminalState::DeadlineBeforeCommit => "deadline_before_commit",
QueryTerminalState::Committed => "committed",
QueryTerminalState::CommittedWithError => "committed_with_error",
QueryTerminalState::PartiallyCommitted => "partially_committed",
QueryTerminalState::CancelledAfterCommit => "cancelled_after_commit",
QueryTerminalState::DeadlineAfterCommit => "deadline_after_commit",
}
}
fn serialization_outcome_name(outcome: mongreldb_query::SerializationOutcome) -> &'static str {
use mongreldb_query::SerializationOutcome;
match outcome {
SerializationOutcome::NotStarted => "not_started",
SerializationOutcome::InProgress => "in_progress",
SerializationOutcome::Succeeded => "succeeded",
SerializationOutcome::Failed => "failed",
}
}
fn terminal_error_category_name(
category: mongreldb_query::QueryTerminalErrorCategory,
) -> &'static str {
use mongreldb_query::QueryTerminalErrorCategory;
match category {
QueryTerminalErrorCategory::Cancellation => "cancellation",
QueryTerminalErrorCategory::Deadline => "deadline",
QueryTerminalErrorCategory::ResultLimit => "result_limit",
QueryTerminalErrorCategory::Serialization => "serialization",
QueryTerminalErrorCategory::Execution => "execution",
}
}
fn terminal_error_retryable(error: Option<&mongreldb_query::QueryTerminalError>) -> bool {
error.is_some_and(|error| {
matches!(
error.code.as_str(),
"IDEMPOTENCY_STORE_FULL" | "IDEMPOTENCY_STORE_UNAVAILABLE"
)
})
}
fn epoch_text(epoch: Option<u64>) -> Option<String> {
epoch.map(|epoch| epoch.to_string())
}
fn cancellation_reason_name(reason: CancellationReason) -> &'static str {
reason.as_str()
}
fn query_cancel_outcome(status: &mongreldb_query::QueryStatus) -> Option<&'static str> {
match status.phase {
SqlQueryPhase::CommitCritical => Some("too_late"),
SqlQueryPhase::Completed | SqlQueryPhase::Failed | SqlQueryPhase::Cancelled => {
Some("already_finished")
}
SqlQueryPhase::Cancelling => Some("accepted"),
_ => None,
}
}
fn query_outcome_json(status: Option<&mongreldb_query::QueryStatus>) -> serde_json::Value {
let Some(status) = status else {
return json!({
"committed": false,
"committed_statements": 0,
"last_commit_epoch": null,
"last_commit_epoch_text": null,
"first_commit_statement_index": null,
"last_commit_statement_index": null,
"completed_statements": 0,
"statement_index": 0,
"serialization": "not_started",
});
};
if status.outcome_unknown {
return json!({
"committed": null,
"committed_statements": null,
"last_commit_epoch": null,
"last_commit_epoch_text": null,
"first_commit_statement_index": null,
"last_commit_statement_index": null,
"completed_statements": null,
"statement_index": null,
"serialization": "unknown",
});
}
json!({
"committed": status.durable_outcome.committed,
"committed_statements": (!status.outcome_unknown).then_some(status.durable_outcome.committed_statements),
"last_commit_epoch": (!status.outcome_unknown).then_some(status.durable_outcome.last_commit_epoch).flatten(),
"last_commit_epoch_text": (!status.outcome_unknown).then_some(epoch_text(status.durable_outcome.last_commit_epoch)).flatten(),
"first_commit_statement_index": (!status.outcome_unknown).then_some(status.durable_outcome.first_commit_statement_index).flatten(),
"last_commit_statement_index": (!status.outcome_unknown).then_some(status.durable_outcome.last_commit_statement_index).flatten(),
"completed_statements": (!status.outcome_unknown).then_some(status.completed_statements),
"statement_index": (!status.outcome_unknown).then_some(status.statement_index),
"serialization": serialization_outcome_name(status.serialization_outcome),
})
}
fn sql_terminal_idempotency_receipt(
status: &mongreldb_query::QueryStatus,
) -> Option<sql_idempotency::SqlDurableReceipt> {
if status.outcome_unknown {
return None;
}
let terminal_state = status.terminal_state()?;
if !status.durable_outcome.committed
&& terminal_state != mongreldb_query::QueryTerminalState::Completed
{
return None;
}
Some(sql_idempotency::SqlDurableReceipt {
original_query_id: status.query_id.to_string(),
status: status
.terminal_state()
.map(terminal_state_name)
.unwrap_or("committed")
.to_owned(),
server_state: query_phase_name(status.phase).to_owned(),
cancellation_reason: cancellation_reason_name(status.cancellation_reason).to_owned(),
outcome: sql_idempotency::SqlReceiptOutcome {
committed: status.durable_outcome.committed,
committed_statements: status.durable_outcome.committed_statements,
last_commit_epoch: status.durable_outcome.last_commit_epoch,
last_commit_epoch_text: epoch_text(status.durable_outcome.last_commit_epoch),
first_commit_statement_index: status.durable_outcome.first_commit_statement_index,
last_commit_statement_index: status.durable_outcome.last_commit_statement_index,
completed_statements: status.completed_statements,
statement_index: status.statement_index,
serialization: serialization_outcome_name(status.serialization_outcome).to_owned(),
},
terminal_error: status.terminal_error.as_ref().map(|error| {
sql_idempotency::SqlReceiptTerminalError {
code: error.code.clone(),
category: terminal_error_category_name(error.category).to_owned(),
}
}),
})
}
fn sql_idempotency_receipt_response(
query_id: QueryId,
receipt: &sql_idempotency::SqlDurableReceipt,
replayed: bool,
expires_at_ms: u64,
persisted: bool,
) -> Response {
let mut response = Json(json!({
"query_id": query_id.to_string(),
"original_query_id": receipt.original_query_id,
"status": receipt.status,
"terminal_state": receipt.status,
"server_state": receipt.server_state,
"cancel_outcome": "already_finished",
"cancellation_reason": receipt.cancellation_reason,
"committed": receipt.outcome.committed,
"committed_statements": receipt.outcome.committed_statements,
"last_commit_epoch": receipt.outcome.last_commit_epoch,
"last_commit_epoch_text": receipt.outcome.last_commit_epoch_text.as_deref(),
"first_commit_statement_index": receipt.outcome.first_commit_statement_index,
"last_commit_statement_index": receipt.outcome.last_commit_statement_index,
"completed_statements": receipt.outcome.completed_statements,
"statement_index": receipt.outcome.statement_index,
"retryable": false,
"idempotency_replayed": replayed,
"idempotency_persisted": persisted,
"idempotency_expires_at_ms": expires_at_ms,
"outcome": receipt.outcome,
"terminal_error": receipt.terminal_error,
}))
.into_response();
response.headers_mut().insert(
"idempotency-replayed",
axum::http::HeaderValue::from_static(if replayed { "true" } else { "false" }),
);
response.headers_mut().insert(
"idempotency-persisted",
axum::http::HeaderValue::from_static(if persisted { "true" } else { "false" }),
);
if let Ok(value) = axum::http::HeaderValue::from_str(&receipt.original_query_id) {
response
.headers_mut()
.insert("x-mongreldb-original-query-id", value);
}
with_query_id(response, query_id)
}
fn terminal_server_error_response(
state: &AppState,
query_id: QueryId,
http_status: StatusCode,
base_code: &'static str,
message: impl Into<String>,
) -> Response {
let status = state.query_registry.status(query_id);
let committed = status
.as_ref()
.is_some_and(|status| status.durable_outcome.committed);
let code = if committed && base_code.starts_with("SERIALIZATION_") {
"SERIALIZATION_FAILED_AFTER_COMMIT"
} else {
base_code
};
let response_status = status
.as_ref()
.and_then(mongreldb_query::QueryStatus::terminal_state)
.map(terminal_state_name)
.unwrap_or(if committed {
"committed_with_error"
} else {
"failed_before_commit"
});
let outcome = query_outcome_json(status.as_ref());
with_query_id(
(
http_status,
Json(json!({
"query_id": query_id.to_string(),
"status": response_status,
"terminal_state": response_status,
"committed": committed,
"committed_statements": status.as_ref().map_or(0, |status| status.durable_outcome.committed_statements),
"last_commit_epoch": status.as_ref().and_then(|status| status.durable_outcome.last_commit_epoch),
"last_commit_epoch_text": epoch_text(status.as_ref().and_then(|status| status.durable_outcome.last_commit_epoch)),
"first_commit_statement_index": status.as_ref().and_then(|status| status.durable_outcome.first_commit_statement_index),
"last_commit_statement_index": status.as_ref().and_then(|status| status.durable_outcome.last_commit_statement_index),
"completed_statements": status.as_ref().map_or(0, |status| status.completed_statements),
"statement_index": status.as_ref().map_or(0, |status| status.statement_index),
"cancel_outcome": null,
"cancellation_reason": status.as_ref().map(|status| cancellation_reason_name(status.cancellation_reason)),
"retryable": false,
"server_state": status.as_ref().map(|status| query_phase_name(status.phase)),
"outcome": outcome,
"error": {
"code": code,
"message": message.into(),
"query_id": query_id.to_string(),
"committed": committed,
"retryable": false,
}
})),
)
.into_response(),
query_id,
)
}
fn query_not_found_response(query_id: Option<QueryId>) -> Response {
let mut response = (
StatusCode::NOT_FOUND,
Json(json!({
"query_id": query_id.map(|value| value.to_string()),
"status": "unknown",
"terminal_state": null,
"committed": null,
"committed_statements": null,
"last_commit_epoch": null,
"last_commit_epoch_text": null,
"first_commit_statement_index": null,
"last_commit_statement_index": null,
"completed_statements": null,
"statement_index": null,
"cancel_outcome": "not_found",
"cancellation_reason": null,
"retryable": false,
"server_state": "not_found",
"outcome": {
"committed": null,
"committed_statements": null,
"last_commit_epoch": null,
"last_commit_epoch_text": null,
"first_commit_statement_index": null,
"last_commit_statement_index": null,
"completed_statements": null,
"statement_index": null,
"serialization": "unknown",
},
"error": {
"code": "QUERY_NOT_FOUND",
"message": "query not found",
"query_id": query_id.map(|value| value.to_string()),
"committed": null,
"retryable": false,
}
})),
)
.into_response();
if let Some(query_id) = query_id {
add_query_id_header(&mut response, query_id);
}
response
}
fn query_session_header(
headers: &axum::http::HeaderMap,
query_id: Option<QueryId>,
) -> std::result::Result<Option<String>, Box<Response>> {
match headers.get("x-session-id") {
Some(value) => match value.to_str() {
Ok(value) if value.len() <= 256 => Ok(Some(value.to_owned())),
_ => Err(Box::new(bad_query_control_request(
"X-Session-ID must be valid text no longer than 256 bytes",
query_id,
))),
},
None => Ok(None),
}
}
fn pre_cancelled_query_response(
query_id: QueryId,
reason: CancellationReason,
status: StatusCode,
) -> Response {
with_query_id(
(
status,
Json(json!({
"query_id": query_id.to_string(),
"status": "cancelled_before_start",
"terminal_state": "cancelled_before_start",
"state": "pre_cancelled",
"server_state": "pre_cancelled",
"cancel_outcome": "pre_cancelled",
"committed": false,
"committed_statements": 0,
"last_commit_epoch": null,
"last_commit_epoch_text": null,
"first_commit_statement_index": null,
"last_commit_statement_index": null,
"completed_statements": 0,
"statement_index": 0,
"cancellation_reason": cancellation_reason_name(reason),
"outcome": {
"committed": false,
"committed_statements": 0,
"last_commit_epoch": null,
"last_commit_epoch_text": null,
"first_commit_statement_index": null,
"last_commit_statement_index": null,
"completed_statements": 0,
"statement_index": 0,
"serialization": "not_started",
},
"terminal_error": {
"code": "QUERY_CANCELLED",
"category": "cancellation",
},
"retryable": false,
})),
)
.into_response(),
query_id,
)
}
fn compact_finished_query_response(status: &CompactFinishedQuery) -> Response {
let query_id = status.query_id;
let durable = &status.durable_outcome;
let outcome_unknown =
status.terminal_state == mongreldb_query::QueryTerminalState::OutcomeUnknown;
let terminal_error = status.terminal_error.as_ref().map(|error| {
json!({
"code": error.code,
"category": terminal_error_category_name(error.category),
})
});
with_query_id(
Json(json!({
"detail": "compact",
"query_id": query_id.to_string(),
"status": terminal_state_name(status.terminal_state),
"terminal_state": terminal_state_name(status.terminal_state),
"state": query_phase_name(status.phase),
"server_state": query_phase_name(status.phase),
"cancel_outcome": "already_finished",
"code": "QUERY_ALREADY_FINISHED",
"committed": (!outcome_unknown).then_some(durable.committed),
"committed_statements": (!outcome_unknown).then_some(durable.committed_statements),
"last_commit_epoch": (!outcome_unknown).then_some(durable.last_commit_epoch).flatten(),
"last_commit_epoch_text": (!outcome_unknown).then_some(epoch_text(durable.last_commit_epoch)).flatten(),
"first_commit_statement_index": (!outcome_unknown).then_some(durable.first_commit_statement_index).flatten(),
"last_commit_statement_index": (!outcome_unknown).then_some(durable.last_commit_statement_index).flatten(),
"completed_statements": (!outcome_unknown).then_some(status.completed_statements),
"statement_index": (!outcome_unknown).then_some(status.statement_index),
"cancellation_reason": cancellation_reason_name(status.cancellation_reason),
"outcome": {
"committed": (!outcome_unknown).then_some(durable.committed),
"committed_statements": (!outcome_unknown).then_some(durable.committed_statements),
"last_commit_epoch": (!outcome_unknown).then_some(durable.last_commit_epoch).flatten(),
"last_commit_epoch_text": (!outcome_unknown).then_some(epoch_text(durable.last_commit_epoch)).flatten(),
"first_commit_statement_index": (!outcome_unknown).then_some(durable.first_commit_statement_index).flatten(),
"last_commit_statement_index": (!outcome_unknown).then_some(durable.last_commit_statement_index).flatten(),
"completed_statements": (!outcome_unknown).then_some(status.completed_statements),
"statement_index": (!outcome_unknown).then_some(status.statement_index),
"serialization": serialization_outcome_name(status.serialization_outcome),
},
"terminal_error": terminal_error,
"retryable": terminal_error_retryable(status.terminal_error.as_ref()),
}))
.into_response(),
query_id,
)
}
async fn query_status(
State(state): State<Arc<AppState>>,
OptionalPrincipal(principal): OptionalPrincipal,
Path(query_id): Path<String>,
headers: axum::http::HeaderMap,
) -> Response {
let Ok(query_id) = query_id.parse::<QueryId>() else {
return query_not_found_response(None);
};
if !request_identity_is_current(&state, &principal) {
return query_not_found_response(Some(query_id));
}
let requested_session = match query_session_header(&headers, Some(query_id)) {
Ok(session_id) => session_id,
Err(response) => return *response,
};
let owner = request_owner(&state, &principal);
let is_admin =
current_request_principal(&state, &principal).is_some_and(|principal| principal.is_admin);
let _lifecycle = state
.query_lifecycle
.lock()
.unwrap_or_else(|error| error.into_inner());
let Some(status) = state.query_registry.status(query_id) else {
if let Some(finished) = state.query_registry.compact_finished_status(query_id) {
if !caller_may_manage_query(&state, &principal, finished.owner.as_deref())
|| requested_session
.as_deref()
.is_some_and(|session| finished.session_id.as_deref() != Some(session))
{
return query_not_found_response(Some(query_id));
}
return compact_finished_query_response(&finished);
}
let reason = state
.pre_cancellations
.reason(query_id, &owner, requested_session.as_deref())
.or_else(|| {
is_admin.then(|| match requested_session.as_deref() {
Some(session_id) => state
.pre_cancellations
.reason_for_query_in_session(query_id, session_id),
None => state.pre_cancellations.reason_for_query(query_id),
})?
});
if let Some(reason) = reason {
return pre_cancelled_query_response(query_id, reason, StatusCode::OK);
}
return query_not_found_response(Some(query_id));
};
if !caller_may_manage_query(&state, &principal, status.owner.as_deref())
|| requested_session
.as_deref()
.is_some_and(|session| status.session_id.as_deref() != Some(session))
{
return query_not_found_response(Some(query_id));
}
let terminal_status = status.terminal_state().map(terminal_state_name);
let terminal_error = status.terminal_error.as_ref().map(|error| {
json!({
"code": error.code,
"category": terminal_error_category_name(error.category),
})
});
let retryable = terminal_error_retryable(status.terminal_error.as_ref());
let cancel_outcome = query_cancel_outcome(&status);
let outcome = query_outcome_json(Some(&status));
let response = Json(json!({
"query_id": query_id.to_string(),
"status": terminal_status.unwrap_or(if status.durable_outcome.committed {
"committed"
} else {
"running"
}),
"terminal_state": terminal_status,
"state": query_phase_name(status.phase),
"server_state": query_phase_name(status.phase),
"started_ms_ago": status.started_at.elapsed().as_millis(),
"deadline_ms_remaining": status.deadline.map(|deadline| {
deadline.saturating_duration_since(std::time::Instant::now()).as_millis()
}),
"session_id": status.session_id,
"operation": status.operation,
"committed": (!status.outcome_unknown).then_some(status.committed),
"committed_statements": (!status.outcome_unknown).then_some(status.durable_outcome.committed_statements),
"last_commit_epoch": (!status.outcome_unknown).then_some(status.durable_outcome.last_commit_epoch).flatten(),
"last_commit_epoch_text": (!status.outcome_unknown).then_some(epoch_text(status.durable_outcome.last_commit_epoch)).flatten(),
"first_commit_statement_index": (!status.outcome_unknown).then_some(status.durable_outcome.first_commit_statement_index).flatten(),
"last_commit_statement_index": (!status.outcome_unknown).then_some(status.durable_outcome.last_commit_statement_index).flatten(),
"cancellation_reason": cancellation_reason_name(status.cancellation_reason),
"completed_statements": (!status.outcome_unknown).then_some(status.completed_statements),
"statement_index": (!status.outcome_unknown).then_some(status.statement_index),
"cancel_outcome": cancel_outcome,
"retryable": retryable,
"outcome": outcome,
"terminal_error": terminal_error,
"trace": {
"queue_duration_us": status.queue_duration.as_micros(),
"planning_duration_us": status.planning_duration.as_micros(),
"execution_duration_us": status.execution_duration.as_micros(),
"serialization_duration_us": status.serialization_duration.as_micros(),
"cancel_requested_phase": status.cancel_requested_phase.map(query_phase_name),
"cancel_observed_phase": status.cancel_observed_phase.map(query_phase_name),
"commit_fence_outcome": commit_fence_outcome_name(status.commit_fence_outcome),
},
}))
.into_response();
with_query_id(response, query_id)
}
async fn cancel_query(
State(state): State<Arc<AppState>>,
OptionalPrincipal(principal): OptionalPrincipal,
Path(query_id): Path<String>,
headers: axum::http::HeaderMap,
) -> Response {
let Ok(query_id) = query_id.parse::<QueryId>() else {
return query_not_found_response(None);
};
if !request_identity_is_current(&state, &principal) {
return query_not_found_response(Some(query_id));
}
let requested_session = match query_session_header(&headers, Some(query_id)) {
Ok(session_id) => session_id,
Err(response) => return *response,
};
let owner = request_owner(&state, &principal);
let _lifecycle = state
.query_lifecycle
.lock()
.unwrap_or_else(|error| error.into_inner());
state.metrics.inc_sql_cancel_requests();
let Some(status) = state.query_registry.status(query_id) else {
if let Some(finished) = state.query_registry.compact_finished_status(query_id) {
if !caller_may_manage_query(&state, &principal, finished.owner.as_deref())
|| requested_session
.as_deref()
.is_some_and(|session| finished.session_id.as_deref() != Some(session))
{
return query_not_found_response(Some(query_id));
}
return compact_finished_query_response(&finished);
}
return match state.pre_cancellations.insert(
query_id,
&owner,
requested_session.as_deref(),
CancellationReason::ClientRequest,
) {
Ok(()) => {
state.metrics.inc_sql_commit_cancel_winner_cancel();
pre_cancelled_query_response(
query_id,
CancellationReason::ClientRequest,
StatusCode::ACCEPTED,
)
}
Err(pre_cancel::InsertError::MetadataTooLarge) => bad_query_control_request(
"query owner or session metadata exceeds 256 bytes",
Some(query_id),
),
Err(
pre_cancel::InsertError::Full
| pre_cancel::InsertError::OwnerLimit
| pre_cancel::InsertError::RateLimited,
) => with_query_id(
(
StatusCode::TOO_MANY_REQUESTS,
Json(json!({
"query_id": query_id.to_string(),
"status": "failed_before_commit",
"terminal_state": "failed_before_commit",
"server_state": "failed",
"cancel_outcome": null,
"cancellation_reason": null,
"committed": false,
"committed_statements": 0,
"last_commit_epoch": null,
"last_commit_epoch_text": null,
"first_commit_statement_index": null,
"last_commit_statement_index": null,
"completed_statements": 0,
"statement_index": 0,
"retryable": true,
"outcome": {
"committed": false,
"committed_statements": 0,
"last_commit_epoch": null,
"last_commit_epoch_text": null,
"first_commit_statement_index": null,
"last_commit_statement_index": null,
"completed_statements": 0,
"statement_index": 0,
"serialization": "not_started",
},
"error": {
"code": "QUERY_REGISTRY_FULL",
"message": "pre-registration cancellation limit reached",
"query_id": query_id.to_string(),
"committed": false,
"retryable": true,
}
})),
)
.into_response(),
query_id,
),
};
};
if !caller_may_manage_query(&state, &principal, status.owner.as_deref())
|| requested_session
.as_deref()
.is_some_and(|session| status.session_id.as_deref() != Some(session))
{
return query_not_found_response(Some(query_id));
}
let (http_status, mut body) = match state.query_registry.cancel(query_id) {
CancelOutcome::Accepted => (
{
state.metrics.inc_sql_commit_cancel_winner_cancel();
StatusCode::ACCEPTED
},
json!({
"query_id": query_id.to_string(),
"state": "cancellation_requested",
"cancel_outcome": "accepted",
}),
),
CancelOutcome::AlreadyCancelling => (
StatusCode::OK,
json!({
"query_id": query_id.to_string(),
"state": "cancelling",
"cancel_outcome": "already_cancelling",
}),
),
CancelOutcome::TooLate => (
{
state.metrics.inc_sql_commit_cancel_winner_commit();
StatusCode::CONFLICT
},
json!({
"query_id": query_id.to_string(),
"state": "commit_critical",
"cancel_outcome": "too_late",
"committed": status.durable_outcome.committed,
"outcome": query_outcome_json(Some(&status)),
"retryable": false,
"error": {
"code": "CANCEL_TOO_LATE",
"message": "the query has entered its durable commit phase",
"committed": status.durable_outcome.committed,
"retryable": false,
}
}),
),
CancelOutcome::AlreadyFinished => (
StatusCode::OK,
json!({
"query_id": query_id.to_string(),
"state": "finished",
"status": status.terminal_state().map(terminal_state_name),
"cancel_outcome": "already_finished",
"code": "QUERY_ALREADY_FINISHED",
"committed": status.durable_outcome.committed,
"outcome": query_outcome_json(Some(&status)),
"retryable": false,
}),
),
CancelOutcome::NotFound => return query_not_found_response(Some(query_id)),
};
let status = state.query_registry.status(query_id).unwrap_or(status);
let response_status = status.terminal_state().map(terminal_state_name).unwrap_or(
if status.durable_outcome.committed {
"committed"
} else {
"running"
},
);
if let Some(body) = body.as_object_mut() {
body.insert("status".into(), json!(response_status));
body.insert(
"terminal_state".into(),
json!(status.terminal_state().map(terminal_state_name)),
);
body.insert(
"committed".into(),
json!((!status.outcome_unknown).then_some(status.durable_outcome.committed)),
);
body.insert(
"committed_statements".into(),
json!((!status.outcome_unknown).then_some(status.durable_outcome.committed_statements)),
);
body.insert(
"last_commit_epoch".into(),
json!((!status.outcome_unknown)
.then_some(status.durable_outcome.last_commit_epoch)
.flatten()),
);
body.insert(
"last_commit_epoch_text".into(),
json!((!status.outcome_unknown)
.then_some(epoch_text(status.durable_outcome.last_commit_epoch))
.flatten()),
);
body.insert(
"first_commit_statement_index".into(),
json!((!status.outcome_unknown)
.then_some(status.durable_outcome.first_commit_statement_index)
.flatten()),
);
body.insert(
"last_commit_statement_index".into(),
json!((!status.outcome_unknown)
.then_some(status.durable_outcome.last_commit_statement_index)
.flatten()),
);
body.insert(
"completed_statements".into(),
json!((!status.outcome_unknown).then_some(status.completed_statements)),
);
body.insert(
"statement_index".into(),
json!((!status.outcome_unknown).then_some(status.statement_index)),
);
body.insert(
"cancellation_reason".into(),
json!(cancellation_reason_name(status.cancellation_reason)),
);
body.insert("retryable".into(), json!(false));
body.insert("server_state".into(), json!(query_phase_name(status.phase)));
body.insert("outcome".into(), query_outcome_json(Some(&status)));
}
with_query_id((http_status, Json(body)).into_response(), query_id)
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct SqlContinuationRequest {
cursor: String,
#[serde(default)]
operation_id: Option<QueryId>,
#[serde(default)]
timeout_ms: Option<u64>,
}
fn register_page_operation(
state: &AppState,
options: SqlQueryOptions,
) -> mongreldb_query::Result<RegisteredSqlQuery> {
let query_id = options.query_id.ok_or_else(|| {
mongreldb_query::MongrelQueryError::InvalidQueryState(
"page operation registration requires an operation id".into(),
)
})?;
let owner = options.owner.clone().unwrap_or_default();
let session_id = options.session_id.clone();
let _lifecycle = state
.query_lifecycle
.lock()
.unwrap_or_else(|error| error.into_inner());
let reason = match state.pre_cancellations.lookup_for_registration(
query_id,
&owner,
session_id.as_deref(),
) {
pre_cancel::RegistrationLookup::NoReservation => None,
pre_cancel::RegistrationLookup::Matching(reason) => Some(reason),
pre_cancel::RegistrationLookup::ReservedByAnotherIdentity => {
return Err(mongreldb_query::MongrelQueryError::QueryIdConflict { query_id });
}
};
let query = state.query_registry.register(options)?;
if let Some(reason) = reason {
state
.pre_cancellations
.take(query_id, &owner, session_id.as_deref());
query.request_cancel(reason);
let error = cancellation_checkpoint_error(&query);
query.fail();
return Err(error);
}
Ok(query)
}
async fn continue_sql_page(
State(state): State<Arc<AppState>>,
OptionalPrincipal(principal): OptionalPrincipal,
headers: axum::http::HeaderMap,
Json(request): Json<SqlContinuationRequest>,
) -> Response {
let query_id = match request.operation_id {
Some(query_id) => query_id,
None => match QueryId::random() {
Ok(query_id) => query_id,
Err(error) => return query_error_response(&error, None),
},
};
if !request_identity_is_current(&state, &principal) {
return with_query_id(
sql_cursor_error_response(
StatusCode::NOT_FOUND,
"SQL_CURSOR_NOT_FOUND",
"SQL continuation result is unavailable",
query_id,
),
query_id,
);
}
let owner = request_owner(&state, &principal);
let session_id = match query_session_header(&headers, Some(query_id)) {
Ok(session_id) => session_id,
Err(response) => return *response,
};
let timeout_ms = request.timeout_ms.unwrap_or_else(|| {
state
.sql_page_default_timeout
.as_millis()
.min(u128::from(u64::MAX)) as u64
});
if timeout_ms == 0 || Duration::from_millis(timeout_ms) > state.sql_page_max_timeout {
return bad_query_control_request(
format!(
"timeout_ms must be positive and no greater than {}",
state.sql_page_max_timeout.as_millis()
),
Some(query_id),
);
}
let query = match register_page_operation(
&state,
SqlQueryOptions {
query_id: Some(query_id),
timeout: Some(Duration::from_millis(timeout_ms)),
owner: Some(owner.clone()),
session_id,
parent_control: None,
},
) {
Ok(query) => query,
Err(error) => return tracked_query_error_response(&state, &error, Some(query_id)),
};
query.set_sql_metadata("CONTINUE SQL PAGE");
let _permit = match tokio::select! {
permit = Arc::clone(&state.sql_page_semaphore).acquire_owned() => permit.map_err(|_| {
mongreldb_query::MongrelQueryError::InvalidQueryState(
"SQL page admission semaphore closed".into(),
)
}),
_ = query.control().cancelled() => Err(cancellation_checkpoint_error(&query)),
} {
Ok(permit) => permit,
Err(error) => {
query.fail();
return tracked_query_error_response(&state, &error, Some(query_id));
}
};
if let Err(error) = query.transition(SqlQueryPhase::Queued, SqlQueryPhase::Executing) {
query.fail();
return tracked_query_error_response(&state, &error, Some(query_id));
}
let fail = |status, code, message| {
query.record_terminal_error(code, mongreldb_query::QueryTerminalErrorCategory::Execution);
query.fail();
with_query_id(
sql_cursor_error_response(status, code, message, query_id),
query_id,
)
};
if request.cursor.is_empty() || request.cursor.len() > 2_048 {
return fail(
StatusCode::BAD_REQUEST,
"INVALID_SQL_CURSOR",
"invalid SQL continuation cursor",
);
}
if let Err(error) = query.checkpoint() {
query.fail();
return tracked_query_error_response(&state, &error, Some(query_id));
}
let cursor_mac_key = match state.cursor_mac_key.get() {
Ok(key) => key,
Err(_) => {
return fail(
StatusCode::INTERNAL_SERVER_ERROR,
"ENTROPY_UNAVAILABLE",
"OS CSPRNG unavailable",
);
}
};
match state.sql_pages.continue_page_with_control(
&request.cursor,
&owner,
&cursor_mac_key,
sql_pages::SqlPageBinding {
security_version: state.db.security_version(),
catalog_epoch: state.db.catalog_snapshot().db_epoch,
},
&query,
) {
Ok(page) => {
let page_byte_count = page.byte_count;
if let Err(error) = query.begin_serialization() {
query.fail();
return tracked_query_error_response(&state, &error, Some(query_id));
}
let serialization_query = query.clone();
match tokio::task::spawn_blocking(move || {
serialize_sql_page_controlled(page, &serialization_query)
})
.await
{
Ok(Ok(body)) => {
if let Err(error) = query.try_complete() {
return tracked_query_error_response(&state, &error, Some(query_id));
}
state.metrics.add_sql_output_bytes(page_byte_count);
with_query_id(sql_page_response(body), query_id)
}
Ok(Err(ControlledPageSerializationError::Query(error))) => {
query.fail();
tracked_query_error_response(&state, &error, Some(query_id))
}
Ok(Err(ControlledPageSerializationError::Encoding)) => fail(
StatusCode::INTERNAL_SERVER_ERROR,
"SERIALIZATION_FAILED",
"failed to serialize SQL continuation page",
),
Err(_) => fail(
StatusCode::INTERNAL_SERVER_ERROR,
"SERIALIZATION_WORKER_FAILED",
"SQL continuation serialization worker failed",
),
}
}
Err(sql_pages::CursorError::Cancelled) => {
let error = cancellation_checkpoint_error(&query);
query.fail();
tracked_query_error_response(&state, &error, Some(query_id))
}
Err(sql_pages::CursorError::Invalid) => fail(
StatusCode::BAD_REQUEST,
"INVALID_SQL_CURSOR",
"invalid SQL continuation cursor",
),
Err(sql_pages::CursorError::Expired) => fail(
StatusCode::GONE,
"SQL_CURSOR_EXPIRED",
"SQL continuation cursor expired",
),
Err(sql_pages::CursorError::NotFound) => fail(
StatusCode::NOT_FOUND,
"SQL_CURSOR_NOT_FOUND",
"SQL continuation result is unavailable",
),
Err(sql_pages::CursorError::PageLimit) => fail(
StatusCode::PAYLOAD_TOO_LARGE,
"RESULT_LIMIT_EXCEEDED",
"one projected row exceeds the page byte or token limit",
),
}
}
fn sql_cursor_error_response(
status: StatusCode,
code: &'static str,
message: &'static str,
query_id: QueryId,
) -> Response {
(
status,
Json(json!({
"query_id": query_id.to_string(),
"status": "failed_before_commit",
"terminal_state": "failed_before_commit",
"server_state": "failed",
"committed": false,
"committed_statements": 0,
"last_commit_epoch": null,
"last_commit_epoch_text": null,
"first_commit_statement_index": null,
"last_commit_statement_index": null,
"completed_statements": 0,
"statement_index": 0,
"cancel_outcome": "already_finished",
"cancellation_reason": "none",
"retryable": false,
"outcome": {
"committed": false,
"committed_statements": 0,
"last_commit_epoch": null,
"last_commit_epoch_text": null,
"first_commit_statement_index": null,
"last_commit_statement_index": null,
"completed_statements": 0,
"statement_index": 0,
"serialization": "not_started",
},
"error": {
"code": code,
"message": message,
"query_id": query_id.to_string(),
"committed": false,
"retryable": false,
}
})),
)
.into_response()
}
async fn sql(
State(state): State<Arc<AppState>>,
OptionalPrincipal(principal): OptionalPrincipal,
headers: axum::http::HeaderMap,
Json(req): Json<SqlRequest>,
) -> Response {
if !state.accepting_sql.load(Ordering::Acquire) {
return (StatusCode::SERVICE_UNAVAILABLE, "server is shutting down").into_response();
}
if !request_identity_is_current(&state, &principal) {
return StatusCode::UNAUTHORIZED.into_response();
}
let session_id = match query_session_header(&headers, None) {
Ok(session_id) => session_id,
Err(response) => return *response,
};
let owner = request_owner(&state, &principal);
if let Some(sid) = session_id {
let Some(entry) = state.sessions.get(&sid, &owner) else {
return (
StatusCode::NOT_FOUND,
"session not found or not owned by caller",
)
.into_response();
};
let (options, query_id) = match resolve_query_options(
&state,
&headers,
req.query_id,
req.timeout_ms,
owner.clone(),
Some(sid.clone()),
) {
Ok(options) => options,
Err(response) => return *response,
};
let query = match register_controlled_query(&state, &entry.session, options) {
Ok(query) => query,
Err(error) => return tracked_query_error_response(&state, &error, Some(query_id)),
};
let registration = RegisteredQueryGuard::new(query);
if mongreldb_query::contains_boolean_ai_predicate(&req.sql) {
registration.fail();
return with_query_id(remote_boolean_ai_error(), query_id);
}
let output_limits = match resolve_sql_output_limits(&state, &req, query_id) {
Ok(limits) => limits,
Err(response) => {
registration.fail();
return *response;
}
};
let (registration, pagination) =
match resolve_sql_pagination(&headers, &req, output_limits, registration, query_id) {
Ok(resolved) => resolved,
Err(response) => return *response,
};
let sql_permit =
match acquire_sql_permit(&state, &entry.session, registration.query()).await {
Ok(permit) => permit,
Err(error) => {
return tracked_query_error_response(&state, &error, Some(query_id));
}
};
let _guard = tokio::select! {
guard = entry.lock.lock() => guard,
_ = registration.query().control().cancelled() => {
return tracked_query_error_response(
&state,
&cancellation_checkpoint_error(registration.query()),
Some(query_id),
);
}
};
if entry.is_closed() {
return (StatusCode::NOT_FOUND, "session no longer available").into_response();
}
if req.idempotency_key.is_some() || headers.contains_key("idempotency-key") {
entry
.session
.fire_test_hook(mongreldb_query::SqlTestHookPoint::BeforeServerIdempotencyCheck);
}
let (registration, idempotency) = match begin_sql_idempotency(
&state,
SqlIdempotencyContext {
headers: &headers,
request: &req,
output_limits,
owner: &owner,
session_id: Some(&sid),
session_in_transaction: entry.session.staged_sql_operation_count().is_some(),
query_id,
},
registration,
)
.await
{
Ok(resolved) => resolved,
Err(response) => return response,
};
entry.touch();
let query = registration.into_query();
execute_sql(
&state,
&principal,
&entry.session,
ResolvedSqlRequest {
request: req,
output_limits,
idempotency,
pagination,
},
query,
query_id,
sql_permit,
)
.await
} else {
let session = match MongrelSession::open_with_external_modules_as(
Arc::clone(&state.db),
state.external_modules.iter().cloned(),
request_principal(&state, &principal),
) {
Ok(session) => session.with_query_registry(Arc::clone(&state.query_registry)),
Err(e) => return (status_for_query_error(&e), e.to_string()).into_response(),
};
let (options, query_id) = match resolve_query_options(
&state,
&headers,
req.query_id,
req.timeout_ms,
owner.clone(),
None,
) {
Ok(options) => options,
Err(response) => return *response,
};
let query = match register_controlled_query(&state, &session, options) {
Ok(query) => query,
Err(error) => return tracked_query_error_response(&state, &error, Some(query_id)),
};
let registration = RegisteredQueryGuard::new(query);
if mongreldb_query::contains_boolean_ai_predicate(&req.sql) {
registration.fail();
return with_query_id(remote_boolean_ai_error(), query_id);
}
let output_limits = match resolve_sql_output_limits(&state, &req, query_id) {
Ok(limits) => limits,
Err(response) => {
registration.fail();
return *response;
}
};
let (registration, pagination) =
match resolve_sql_pagination(&headers, &req, output_limits, registration, query_id) {
Ok(resolved) => resolved,
Err(response) => return *response,
};
let (registration, idempotency) = match begin_sql_idempotency(
&state,
SqlIdempotencyContext {
headers: &headers,
request: &req,
output_limits,
owner: &owner,
session_id: None,
session_in_transaction: false,
query_id,
},
registration,
)
.await
{
Ok(resolved) => resolved,
Err(response) => return response,
};
let sql_permit = match acquire_sql_permit(&state, &session, registration.query()).await {
Ok(permit) => permit,
Err(error) => {
if let Some(idempotency) = idempotency {
idempotency.abort();
}
return tracked_query_error_response(&state, &error, Some(query_id));
}
};
let query = registration.into_query();
execute_sql(
&state,
&principal,
&session,
ResolvedSqlRequest {
request: req,
output_limits,
idempotency,
pagination,
},
query,
query_id,
sql_permit,
)
.await
}
}
async fn execute_sql(
state: &AppState,
principal: &Option<mongreldb_core::Principal>,
session: &MongrelSession,
request: ResolvedSqlRequest,
query: RegisteredSqlQuery,
query_id: QueryId,
sql_permit: tokio::sync::OwnedSemaphorePermit,
) -> Response {
let ResolvedSqlRequest {
request: req,
output_limits,
idempotency,
pagination,
} = request;
let idempotency_query = idempotency.as_ref().map(|_| query.clone());
state.metrics.inc_sql_queries();
let audited = audit::is_audited_sql(&req.sql);
let actor = request_owner(state, principal);
let page_binding = sql_pages::SqlPageBinding {
security_version: state.db.security_version(),
catalog_epoch: state.db.catalog_snapshot().db_epoch,
};
let start = std::time::Instant::now();
let result = if let Some(pagination) = pagination {
match session
.run_with_query_for_serialization_with_limits(
&req.sql,
query,
mongreldb_query::SqlCollectionLimits::new(output_limits.0, output_limits.1),
)
.await
{
Ok(output) => Ok(dispatch_paginated_sql(
state,
output,
query_id,
&actor,
pagination,
output_limits,
session.sql_test_hook(),
page_binding,
)
.await),
Err(error) => Err(error),
}
} else if req.format.as_deref() == Some("arrow-stream") {
match session
.run_stream_with_query_for_serialization(&req.sql, query)
.await
{
Ok((stream, completion)) => Ok(sql_arrow_stream_response_controlled(
stream,
completion,
sql_permit,
output_limits,
state,
query_id,
session.sql_test_hook(),
)),
Err(error) => Err(error),
}
} else {
match session
.run_with_query_for_serialization_with_limits(
&req.sql,
query,
mongreldb_query::SqlCollectionLimits::new(output_limits.0, output_limits.1),
)
.await
{
Ok(output) => Ok(dispatch_buffered_sql_format(
state,
req.format.as_deref(),
output,
query_id,
session.sql_test_hook(),
output_limits,
)
.await),
Err(error) => Err(error),
}
};
let elapsed = start.elapsed();
if elapsed >= state.slow_query_threshold {
state.metrics.inc_slow_queries();
eprintln!(
"[slow-query] {}\u{00b5}s query_id={} operation={}",
elapsed.as_micros(),
query_id,
safe_sql_operation(&req.sql)
);
}
if audited {
let (action, detail) = audit::redacted_ddl_detail(&req.sql, result.is_ok());
state.audit.record(actor, action, detail);
}
let response = match result {
Ok(response) => with_query_id(response, query_id),
Err(e) => {
state.metrics.inc_sql_errors();
tracked_query_error_response(state, &e, Some(query_id))
}
};
let Some(idempotency) = idempotency else {
return response;
};
let status = idempotency_query.map(|query| query.status());
if let Some(receipt) = status.as_ref().and_then(sql_terminal_idempotency_receipt) {
let (expires_at_ms, persisted) = idempotency.commit(receipt.clone());
return sql_idempotency_receipt_response(
query_id,
&receipt,
false,
expires_at_ms,
persisted,
);
}
if status.as_ref().is_some_and(can_abort_idempotency_intent) {
idempotency.abort();
}
response
}
fn can_abort_idempotency_intent(status: &mongreldb_query::QueryStatus) -> bool {
!status.outcome_unknown
&& !status.durable_outcome.committed
&& matches!(
status.terminal_state(),
Some(
mongreldb_query::QueryTerminalState::FailedBeforeCommit
| mongreldb_query::QueryTerminalState::CancelledBeforeCommit
| mongreldb_query::QueryTerminalState::DeadlineBeforeCommit
)
)
}
fn safe_sql_operation(sql: &str) -> String {
sql.split_whitespace()
.next()
.unwrap_or("UNKNOWN")
.chars()
.filter(|character| character.is_ascii_alphabetic())
.take(16)
.collect::<String>()
.to_ascii_uppercase()
}
fn remote_boolean_ai_error() -> Response {
(
StatusCode::BAD_REQUEST,
"Boolean ANN/Sparse SQL is disabled remotely; use scored SQL functions",
)
.into_response()
}
#[derive(Debug)]
struct SerializedOutput {
bytes: Vec<u8>,
arrow: bool,
}
#[derive(Debug)]
enum BufferedSerializationError {
Query(mongreldb_query::MongrelQueryError),
Limit(String),
Encoding(String),
}
struct LimitedOutput {
bytes: Vec<u8>,
max_bytes: usize,
exceeded: bool,
}
impl LimitedOutput {
fn new(max_bytes: usize) -> Self {
Self {
bytes: Vec::new(),
max_bytes,
exceeded: false,
}
}
}
impl std::io::Write for LimitedOutput {
fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
if self.bytes.len().saturating_add(bytes.len()) > self.max_bytes {
self.exceeded = true;
return Err(std::io::Error::other("SQL output byte limit exceeded"));
}
self.bytes.extend_from_slice(bytes);
Ok(bytes.len())
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
fn serialize_buffered_output(
format: &str,
batches: &[arrow::record_batch::RecordBatch],
query: &RegisteredSqlQuery,
max_rows: usize,
max_bytes: usize,
test_hook: Option<&mongreldb_query::SqlTestHook>,
) -> std::result::Result<SerializedOutput, BufferedSerializationError> {
const ROW_CHECKPOINT_INTERVAL: usize = 256;
let mut rows = 0usize;
let mut writer_output = LimitedOutput::new(max_bytes);
if format == "arrow" {
if batches.is_empty() {
if let Some(hook) = test_hook {
hook(mongreldb_query::SqlTestHookPoint::AfterSerialization);
}
return Ok(SerializedOutput {
bytes: Vec::new(),
arrow: true,
});
}
let schema = batches[0].schema();
let encoding_result = (|| {
let mut writer =
arrow::ipc::writer::FileWriter::try_new(&mut writer_output, schema.as_ref())
.map_err(|error| error.to_string())?;
for batch in batches {
for offset in (0..batch.num_rows()).step_by(ROW_CHECKPOINT_INTERVAL) {
if let Some(hook) = test_hook {
hook(mongreldb_query::SqlTestHookPoint::BeforeSerializationBatch);
}
query.checkpoint().map_err(|error| error.to_string())?;
let length = ROW_CHECKPOINT_INTERVAL.min(batch.num_rows() - offset);
rows = rows.saturating_add(length);
if rows > max_rows {
return Err("SQL output row limit exceeded".into());
}
writer
.write(&batch.slice(offset, length))
.map_err(|error| error.to_string())?;
}
}
writer.finish().map_err(|error| error.to_string())
})();
if let Err(error) = encoding_result {
if let Err(query_error) = query.checkpoint() {
return Err(BufferedSerializationError::Query(query_error));
}
if writer_output.exceeded || rows > max_rows {
return Err(BufferedSerializationError::Limit(error));
}
return Err(BufferedSerializationError::Encoding(error));
}
if let Some(hook) = test_hook {
hook(mongreldb_query::SqlTestHookPoint::AfterSerialization);
}
return Ok(SerializedOutput {
bytes: writer_output.bytes,
arrow: true,
});
}
let encoding_result = (|| {
let mut writer = arrow::json::writer::ArrayWriter::new(&mut writer_output);
for batch in batches {
for offset in (0..batch.num_rows()).step_by(ROW_CHECKPOINT_INTERVAL) {
if let Some(hook) = test_hook {
hook(mongreldb_query::SqlTestHookPoint::BeforeSerializationBatch);
}
query.checkpoint().map_err(|error| error.to_string())?;
let length = ROW_CHECKPOINT_INTERVAL.min(batch.num_rows() - offset);
rows = rows.saturating_add(length);
if rows > max_rows {
return Err("SQL output row limit exceeded".into());
}
let slice = batch.slice(offset, length);
writer
.write_batches(&[&slice])
.map_err(|error| error.to_string())?;
}
}
writer.finish().map_err(|error| error.to_string())
})();
if let Err(error) = encoding_result {
if let Err(query_error) = query.checkpoint() {
return Err(BufferedSerializationError::Query(query_error));
}
if writer_output.exceeded || rows > max_rows {
return Err(BufferedSerializationError::Limit(error));
}
return Err(BufferedSerializationError::Encoding(error));
}
if let Some(hook) = test_hook {
hook(mongreldb_query::SqlTestHookPoint::AfterSerialization);
}
Ok(SerializedOutput {
bytes: writer_output.bytes,
arrow: false,
})
}
#[cfg(test)]
fn sql_arrow_stream_response(batches: mongreldb_query::MongrelRecordBatchStream) -> Response {
use futures::{stream, StreamExt};
const STREAM_CT: &str = "application/vnd.apache.arrow.stream";
let schema = batches.schema();
let mut writer = match arrow::ipc::writer::StreamWriter::try_new(Vec::new(), schema.as_ref()) {
Ok(w) => w,
Err(e) => {
return (
StatusCode::INTERNAL_SERVER_ERROR,
format!("arrow stream init error: {e}"),
)
.into_response()
}
};
let schema_chunk: Vec<u8> = std::mem::take(writer.get_mut());
let batch_stream = stream::unfold(
(batches, Some(writer)),
|(mut batches, writer)| async move {
let mut writer = writer?;
match batches.next().await {
Some(Ok(batch)) => match writer.write(&batch) {
Ok(()) => {
let chunk = std::mem::take(writer.get_mut());
Some((Ok(chunk), (batches, Some(writer))))
}
Err(error) => Some((Err(std::io::Error::other(error)), (batches, None))),
},
Some(Err(error)) => Some((Err(std::io::Error::other(error)), (batches, None))),
None => match writer.finish() {
Ok(()) => {
let chunk = std::mem::take(writer.get_mut());
Some((Ok(chunk), (batches, None)))
}
Err(error) => Some((Err(std::io::Error::other(error)), (batches, None))),
},
}
},
);
let schema_item: Result<Vec<u8>, std::io::Error> = Ok(schema_chunk);
let full = stream::iter([schema_item]).chain(batch_stream);
let body = axum::body::Body::from_stream(full);
([(header::CONTENT_TYPE, STREAM_CT)], body).into_response()
}
fn sql_arrow_stream_response_controlled(
batches: mongreldb_query::MongrelRecordBatchStream,
completion: SqlStreamCompletion,
sql_permit: tokio::sync::OwnedSemaphorePermit,
limits: (usize, usize),
state: &AppState,
query_id: QueryId,
test_hook: Option<mongreldb_query::SqlTestHook>,
) -> Response {
use futures::{stream, StreamExt};
const STREAM_CT: &str = "application/vnd.apache.arrow.stream";
let (max_rows, max_bytes) = limits;
let schema = batches.schema();
let mut writer = match arrow::ipc::writer::StreamWriter::try_new(Vec::new(), schema.as_ref()) {
Ok(writer) => writer,
Err(error) => {
completion.fail_serialization();
drop(batches);
return terminal_server_error_response(
state,
query_id,
StatusCode::INTERNAL_SERVER_ERROR,
"SERIALIZATION_FAILED",
format!("arrow stream init error: {error}"),
);
}
};
let schema_chunk = std::mem::take(writer.get_mut());
if schema_chunk.len() > max_bytes {
completion.fail_result_limit();
drop(batches);
return terminal_server_error_response(
state,
query_id,
StatusCode::PAYLOAD_TOO_LARGE,
"RESULT_LIMIT_EXCEEDED",
"SQL output byte limit exceeded",
);
}
let metrics = Arc::clone(&state.metrics);
metrics.add_sql_output_bytes(schema_chunk.len());
let batch_stream = stream::unfold(
(
batches,
Some(writer),
completion,
Some(sql_permit),
0usize,
schema_chunk.len(),
metrics,
),
move |(mut batches, writer, completion, permit, rows, bytes, metrics)| {
let test_hook = test_hook.clone();
async move {
let mut writer = writer?;
match batches.next().await {
Some(Ok(batch)) => {
let next_rows = rows.saturating_add(batch.num_rows());
if next_rows > max_rows {
completion.fail_result_limit();
return Some((
Err(std::io::Error::other("SQL output row limit exceeded")),
(batches, None, completion, permit, next_rows, bytes, metrics),
));
}
match writer.write(&batch) {
Ok(()) => {
let chunk = std::mem::take(writer.get_mut());
let next_bytes = bytes.saturating_add(chunk.len());
if next_bytes > max_bytes {
completion.fail_result_limit();
return Some((
Err(std::io::Error::other(
"SQL output byte limit exceeded",
)),
(
batches, None, completion, permit, next_rows,
next_bytes, metrics,
),
));
}
metrics.add_sql_output_bytes(chunk.len());
Some((
Ok(chunk),
(
batches,
Some(writer),
completion,
permit,
next_rows,
next_bytes,
metrics,
),
))
}
Err(error) => {
completion.fail_serialization();
Some((
Err(std::io::Error::other(error)),
(batches, None, completion, permit, rows, bytes, metrics),
))
}
}
}
Some(Err(error)) => Some((
Err(std::io::Error::other(error)),
(batches, None, completion, permit, rows, bytes, metrics),
)),
None => match writer.finish() {
Ok(()) => {
let chunk = std::mem::take(writer.get_mut());
let next_bytes = bytes.saturating_add(chunk.len());
if next_bytes > max_bytes {
completion.fail_result_limit();
return Some((
Err(std::io::Error::other("SQL output byte limit exceeded")),
(batches, None, completion, permit, rows, next_bytes, metrics),
));
}
if let Some(hook) = test_hook {
hook(mongreldb_query::SqlTestHookPoint::AfterSerialization);
}
match completion.try_complete() {
Ok(()) => {
metrics.add_sql_output_bytes(chunk.len());
Some((
Ok(chunk),
(
batches, None, completion, permit, rows, next_bytes,
metrics,
),
))
}
Err(error) => {
metrics.inc_sql_errors();
Some((
Err(std::io::Error::other(error.to_string())),
(batches, None, completion, permit, rows, bytes, metrics),
))
}
}
}
Err(error) => {
completion.fail_serialization();
Some((
Err(std::io::Error::other(error)),
(batches, None, completion, permit, rows, bytes, metrics),
))
}
},
}
}
},
);
let schema_item: Result<Vec<u8>, std::io::Error> = Ok(schema_chunk);
let body = axum::body::Body::from_stream(stream::iter([schema_item]).chain(batch_stream));
([(header::CONTENT_TYPE, STREAM_CT)], body).into_response()
}
#[derive(Deserialize)]
struct TxnOp {
table: String,
op: String,
cells: Option<Vec<serde_json::Value>>,
row_id: Option<u64>,
}
#[derive(Deserialize)]
struct TxnRequest {
ops: Vec<TxnOp>,
}
async fn txn(
State(state): State<Arc<AppState>>,
OptionalPrincipal(principal): OptionalPrincipal,
Json(req): Json<TxnRequest>,
) -> Response {
let mut parsed: Vec<(String, TxnAction)> = Vec::with_capacity(req.ops.len());
for op in &req.ops {
match op.op.as_str() {
"put" => {
let cells_json = match op.cells.as_ref() {
Some(c) if !c.is_empty() => c,
_ => {
return (StatusCode::BAD_REQUEST, "put op requires non-empty cells")
.into_response()
}
};
let handle = match state.db.table(&op.table) {
Ok(h) => h,
Err(e) => return (StatusCode::NOT_FOUND, e.to_string()).into_response(),
};
let schema = handle.lock().schema().clone();
let cells = match parse_cells(cells_json, &schema) {
Ok(c) => c,
Err(msg) => return (StatusCode::BAD_REQUEST, msg).into_response(),
};
parsed.push((op.table.clone(), TxnAction::Put(cells)));
}
"delete" => {
let rid = match op.row_id {
Some(r) => r,
None => {
return (StatusCode::BAD_REQUEST, "delete op requires row_id")
.into_response()
}
};
parsed.push((op.table.clone(), TxnAction::Delete(rid)));
}
other => {
return (StatusCode::BAD_REQUEST, format!("unknown op: {other}")).into_response()
}
}
}
state.metrics.inc_txns();
let mut transaction = state.db.begin_as(request_principal(&state, &principal));
let result = (|| {
for (table, action) in &parsed {
match action {
TxnAction::Put(cells) => {
transaction.put(table, cells.clone())?;
}
TxnAction::Delete(rid) => {
transaction.delete(table, mongreldb_core::RowId(*rid))?;
}
}
}
transaction.commit()
})();
match result {
Ok(epoch) => Json(json!({
"status": "committed",
"epoch": epoch.0,
"epoch_text": epoch.0.to_string()
}))
.into_response(),
Err(error) => crate::kit::durable_core_error_response(&error)
.unwrap_or_else(|| (status_for_error(&error), error.to_string()).into_response()),
}
}
enum TxnAction {
Put(Vec<(u16, Value)>),
Delete(u64),
}
#[cfg(test)]
mod auth_tests {
use super::*;
use mongreldb_core::Database;
use tempfile::tempdir;
#[test]
fn slow_query_operation_does_not_include_literals() {
let sql = "CREATE USER alice PASSWORD 'never-log-this'";
let operation = safe_sql_operation(sql);
assert_eq!(operation, "CREATE");
assert!(!operation.contains("never-log-this"));
}
#[tokio::test]
async fn auth_rejects_missing_token() {
let dir = tempdir().unwrap();
let db = Arc::new(Database::create(dir.path()).unwrap());
let app = build_app_with_config(db, std::iter::empty(), Some("secret".into()), None);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
let client = reqwest::Client::new();
let resp = client
.get(format!("http://{addr}/health"))
.send()
.await
.unwrap();
assert_eq!(resp.status(), 401);
}
#[tokio::test]
async fn auth_accepts_valid_token() {
let dir = tempdir().unwrap();
let db = Arc::new(Database::create(dir.path()).unwrap());
let app = build_app_with_config(db, std::iter::empty(), Some("secret".into()), None);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
let client = reqwest::Client::new();
let resp = client
.get(format!("http://{addr}/health"))
.header("Authorization", "Bearer secret")
.send()
.await
.unwrap();
assert_eq!(resp.status(), 200);
}
#[tokio::test]
async fn no_auth_when_token_unset() {
let dir = tempdir().unwrap();
let db = Arc::new(Database::create(dir.path()).unwrap());
let app = build_app_with_config(db, std::iter::empty(), None, None);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
let client = reqwest::Client::new();
let resp = client
.get(format!("http://{addr}/health"))
.send()
.await
.unwrap();
assert_eq!(resp.status(), 200);
}
#[tokio::test]
async fn capabilities_advertise_sql_cancellation_v2() {
let dir = tempdir().unwrap();
let db = Arc::new(Database::create(dir.path()).unwrap());
let app = build_app_with_config(db, std::iter::empty(), None, None);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
let body: serde_json::Value = reqwest::Client::new()
.get(format!("http://{addr}/capabilities"))
.send()
.await
.unwrap()
.json()
.await
.unwrap();
assert_eq!(body["sql_cancellation"]["version"], 2);
assert_eq!(body["sql_cancellation"]["client_query_ids"], true);
assert_eq!(body["sql_cancellation"]["cancel_endpoint"], true);
assert_eq!(body["sql_cancellation"]["query_status"], true);
assert_eq!(body["sql_cancellation"]["pre_registration_cancel"], true);
assert_eq!(body["sql_cancellation"]["stream_disconnect_cancels"], true);
assert_eq!(body["sql_idempotency"]["version"], 1);
assert_eq!(
body["sql_idempotency"]["indeterminate_never_reexecutes"],
true
);
assert_eq!(body["sql_pagination"]["version"], 1);
assert_eq!(
body["sql_pagination"]["continuation_endpoint"],
"/sql/continue"
);
}
}
#[cfg(test)]
mod query_response_tests {
use super::*;
#[test]
fn cancellation_reason_names_are_stable_snake_case() {
assert_eq!(cancellation_reason_name(CancellationReason::None), "none");
assert_eq!(
cancellation_reason_name(CancellationReason::ClientRequest),
"client_request"
);
assert_eq!(
cancellation_reason_name(CancellationReason::ClientDisconnected),
"client_disconnected"
);
assert_eq!(
cancellation_reason_name(CancellationReason::SessionClosed),
"session_closed"
);
assert_eq!(
cancellation_reason_name(CancellationReason::ServerShutdown),
"server_shutdown"
);
assert_eq!(
cancellation_reason_name(CancellationReason::Deadline),
"deadline"
);
}
#[test]
fn unknown_outcome_never_proves_idempotency_intent_safe_to_abort() {
let registry = Arc::new(SqlQueryRegistry::default());
let unknown_id: QueryId = "102132435465768798a9bacbdcedfe0f".parse().unwrap();
let unknown = registry
.register(SqlQueryOptions {
query_id: Some(unknown_id),
..SqlQueryOptions::default()
})
.unwrap();
unknown.mark_outcome_unknown();
unknown.fail();
assert!(!can_abort_idempotency_intent(
®istry.status(unknown_id).unwrap()
));
let failed_id: QueryId = "2031425364758697a8b9cadbecfd0e1f".parse().unwrap();
let failed = registry
.register(SqlQueryOptions {
query_id: Some(failed_id),
..SqlQueryOptions::default()
})
.unwrap();
failed.fail();
assert!(can_abort_idempotency_intent(
®istry.status(failed_id).unwrap()
));
}
#[test]
fn unknown_outcome_never_becomes_durable_receipt() {
let registry = Arc::new(SqlQueryRegistry::default());
let query = registry.register(SqlQueryOptions::default()).unwrap();
query.record_commit(0, 42);
query.mark_outcome_unknown();
query.fail();
let status = registry.status(query.id()).unwrap();
assert!(status.durable_outcome.committed);
assert!(status.outcome_unknown);
assert!(sql_terminal_idempotency_receipt(&status).is_none());
}
#[test]
fn successful_noop_write_becomes_noncommitting_receipt() {
let registry = Arc::new(SqlQueryRegistry::default());
let query = registry.register(SqlQueryOptions::default()).unwrap();
query
.transition(SqlQueryPhase::Queued, SqlQueryPhase::Executing)
.unwrap();
query.try_complete().unwrap();
let status = registry.status(query.id()).unwrap();
assert!(!status.durable_outcome.committed);
assert!(!can_abort_idempotency_intent(&status));
let receipt = sql_terminal_idempotency_receipt(&status).unwrap();
assert_eq!(receipt.status, "completed");
assert!(!receipt.outcome.committed);
assert_eq!(receipt.outcome.committed_statements, 0);
assert_eq!(receipt.outcome.last_commit_epoch, None);
}
#[test]
fn conflicting_idempotency_key_sources_are_rejected() {
let request = SqlRequest {
sql: "INSERT INTO items VALUES (1)".into(),
format: None,
query_id: None,
timeout_ms: None,
max_output_rows: None,
max_output_bytes: None,
idempotency_key: Some("body-key".into()),
pagination: None,
};
let mut headers = axum::http::HeaderMap::new();
headers.insert("idempotency-key", "header-key".parse().unwrap());
assert_eq!(
requested_sql_idempotency_key(&headers, &request),
Err("body idempotency_key and Idempotency-Key header must match")
);
headers.insert("idempotency-key", "body-key".parse().unwrap());
assert_eq!(
requested_sql_idempotency_key(&headers, &request),
Ok(Some("body-key".into()))
);
}
#[test]
fn idempotency_binding_includes_pagination_semantics() {
let mut request = SqlRequest {
sql: "INSERT INTO items VALUES (1)".into(),
format: None,
query_id: None,
timeout_ms: None,
max_output_rows: None,
max_output_bytes: None,
idempotency_key: Some("key".into()),
pagination: None,
};
let unpaged = sql_idempotency_binding(&request, (100, 1_024), None, 60_000).unwrap();
request.pagination = Some(SqlPaginationRequest {
page_size_rows: 10,
projection: vec!["id".into()],
max_page_bytes: Some(512),
max_page_tokens: Some(128),
});
let paged = sql_idempotency_binding(&request, (100, 1_024), None, 60_000).unwrap();
assert_ne!(unpaged.request_semantics_hash, paged.request_semantics_hash);
}
#[test]
fn paginated_decode_stops_nested_heap_amplification_at_budget() {
let registry = Arc::new(SqlQueryRegistry::default());
let query = registry.register(SqlQueryOptions::default()).unwrap();
let json = format!("[[{}]]", vec!["null"; 10_000].join(","));
let mut deserializer = serde_json::Deserializer::from_slice(json.as_bytes());
let mut budget = PaginatedDecodeBudget {
used: 0,
limit: 4 * 1024,
nodes: 0,
exceeded: false,
query: &query,
test_hook: None,
};
let error = serde::de::DeserializeSeed::deserialize(
BudgetedJsonRowsSeed {
budget: &mut budget,
},
&mut deserializer,
)
.unwrap_err();
assert!(error.to_string().contains(PAGINATED_MEMORY_LIMIT_ERROR));
assert!(budget.exceeded);
assert!(budget.nodes < 100, "decoded {} nodes", budget.nodes);
query.fail();
}
#[tokio::test]
async fn unknown_outcome_response_never_claims_no_commit() {
let registry = Arc::new(SqlQueryRegistry::default());
let query = registry.register(SqlQueryOptions::default()).unwrap();
let query_id = query.id();
query
.transition(SqlQueryPhase::Queued, SqlQueryPhase::Executing)
.unwrap();
let error = query.outcome_unknown_error("fenced maintenance failed");
query.fail();
let status = registry.status(query_id).unwrap();
assert_eq!(
status.terminal_state(),
Some(mongreldb_query::QueryTerminalState::OutcomeUnknown)
);
let response = query_error_response_with_status(&error, Some(query_id), Some(&status));
assert_eq!(response.status(), StatusCode::CONFLICT);
let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
.await
.unwrap();
let body: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
assert_eq!(body["status"], "outcome_unknown");
assert_eq!(body["error"]["code"], "QUERY_OUTCOME_UNKNOWN");
assert!(body["committed"].is_null());
assert!(body["committed_statements"].is_null());
assert!(body["last_commit_epoch"].is_null());
assert!(body["completed_statements"].is_null());
assert!(body["statement_index"].is_null());
assert!(body["outcome"]["committed"].is_null());
assert!(body["error"]["committed"].is_null());
}
#[tokio::test]
async fn retryable_idempotency_error_survives_terminal_status() {
let registry = Arc::new(SqlQueryRegistry::default());
let query = registry.register(SqlQueryOptions::default()).unwrap();
let query_id = query.id();
let response = registered_sql_error_response(
RegisteredQueryGuard::new(query),
query_id,
StatusCode::SERVICE_UNAVAILABLE,
"IDEMPOTENCY_STORE_UNAVAILABLE",
"could not durably reserve the SQL idempotency key",
true,
);
let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
.await
.unwrap();
let body: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
assert_eq!(body["retryable"], true);
assert_eq!(body["error"]["retryable"], true);
let status = registry.status(query_id).unwrap();
assert_eq!(
status.terminal_error.as_ref().unwrap().code,
"IDEMPOTENCY_STORE_UNAVAILABLE"
);
assert!(terminal_error_retryable(status.terminal_error.as_ref()));
}
#[tokio::test]
async fn committed_idempotency_terminal_error_is_a_receipt() {
let query_id: QueryId = "00112233445566778899aabbccddeeff".parse().unwrap();
let receipt = sql_idempotency::SqlDurableReceipt {
original_query_id: query_id.to_string(),
status: "committed_with_error".into(),
server_state: "failed".into(),
cancellation_reason: "client_disconnected".into(),
outcome: sql_idempotency::SqlReceiptOutcome {
committed: true,
committed_statements: 1,
last_commit_epoch: Some(42),
last_commit_epoch_text: Some("42".into()),
first_commit_statement_index: Some(0),
last_commit_statement_index: Some(0),
completed_statements: 1,
statement_index: 0,
serialization: "failed".into(),
},
terminal_error: Some(sql_idempotency::SqlReceiptTerminalError {
code: "SERIALIZATION_FAILED_AFTER_COMMIT".into(),
category: "serialization".into(),
}),
};
let response = sql_idempotency_receipt_response(query_id, &receipt, false, 99, true);
assert_eq!(response.status(), StatusCode::OK);
let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
.await
.unwrap();
let body: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
assert_eq!(body["status"], "committed_with_error");
assert_eq!(body["server_state"], "failed");
assert_eq!(body["cancellation_reason"], "client_disconnected");
assert_eq!(body["first_commit_statement_index"], 0);
assert_eq!(body["last_commit_statement_index"], 0);
assert_eq!(
body["terminal_error"]["code"],
"SERIALIZATION_FAILED_AFTER_COMMIT"
);
}
#[test]
fn durable_replay_restores_terminal_status_parity() {
let registry = Arc::new(SqlQueryRegistry::default());
let query_id: QueryId = "11223344556677889900aabbccddeeff".parse().unwrap();
let query = registry
.register(SqlQueryOptions {
query_id: Some(query_id),
..SqlQueryOptions::default()
})
.unwrap();
let receipt = sql_idempotency::SqlDurableReceipt {
original_query_id: "00112233445566778899aabbccddeeff".into(),
status: "cancelled_after_commit".into(),
server_state: "cancelled".into(),
cancellation_reason: "client_disconnected".into(),
outcome: sql_idempotency::SqlReceiptOutcome {
committed: true,
committed_statements: 1,
last_commit_epoch: Some(42),
last_commit_epoch_text: Some("42".into()),
first_commit_statement_index: Some(0),
last_commit_statement_index: Some(0),
completed_statements: 1,
statement_index: 1,
serialization: "failed".into(),
},
terminal_error: Some(sql_idempotency::SqlReceiptTerminalError {
code: "QUERY_CANCELLED_AFTER_COMMIT".into(),
category: "cancellation".into(),
}),
};
restore_idempotency_replay(RegisteredQueryGuard::new(query), &receipt).unwrap();
let status = registry.status(query_id).unwrap();
assert_eq!(status.phase, SqlQueryPhase::Cancelled);
assert_eq!(
status.terminal_state(),
Some(mongreldb_query::QueryTerminalState::CancelledAfterCommit)
);
assert_eq!(
status.cancellation_reason,
CancellationReason::ClientDisconnected
);
assert_eq!(status.durable_outcome.committed_statements, 1);
assert_eq!(status.durable_outcome.last_commit_epoch, Some(42));
assert_eq!(
status.terminal_error.unwrap().code,
"QUERY_CANCELLED_AFTER_COMMIT"
);
}
#[test]
fn durable_replay_rejects_invalid_authenticated_state() {
let registry = Arc::new(SqlQueryRegistry::default());
let query = registry.register(SqlQueryOptions::default()).unwrap();
let receipt = sql_idempotency::SqlDurableReceipt {
original_query_id: query.id().to_string(),
status: "invented_terminal_state".into(),
server_state: "completed".into(),
cancellation_reason: "none".into(),
outcome: sql_idempotency::SqlReceiptOutcome {
committed: true,
committed_statements: 1,
last_commit_epoch: Some(42),
last_commit_epoch_text: Some("42".into()),
first_commit_statement_index: Some(0),
last_commit_statement_index: Some(0),
completed_statements: 1,
statement_index: 0,
serialization: "succeeded".into(),
},
terminal_error: None,
};
let error =
restore_idempotency_replay(RegisteredQueryGuard::new(query), &receipt).unwrap_err();
assert!(error.to_string().contains("invalid terminal state"));
}
#[test]
fn durable_replay_cancel_wins_before_receipt_response() {
let registry = Arc::new(SqlQueryRegistry::default());
let query_id: QueryId = "22334455667788990011aabbccddeeff".parse().unwrap();
let query = registry
.register(SqlQueryOptions {
query_id: Some(query_id),
..SqlQueryOptions::default()
})
.unwrap();
assert_eq!(
query.request_cancel(CancellationReason::ClientRequest),
CancelOutcome::Accepted
);
let receipt = sql_idempotency::SqlDurableReceipt {
original_query_id: "00112233445566778899aabbccddeeff".into(),
status: "completed".into(),
server_state: "completed".into(),
cancellation_reason: "none".into(),
outcome: sql_idempotency::SqlReceiptOutcome {
committed: true,
committed_statements: 1,
last_commit_epoch: Some(42),
last_commit_epoch_text: Some("42".into()),
first_commit_statement_index: Some(0),
last_commit_statement_index: Some(0),
completed_statements: 1,
statement_index: 0,
serialization: "succeeded".into(),
},
terminal_error: None,
};
let error = restore_idempotency_replay(RegisteredQueryGuard::new(query), &receipt)
.expect_err("accepted cancellation must suppress replay success");
assert!(matches!(
error,
mongreldb_query::MongrelQueryError::QueryCancelled { .. }
));
let status = registry.status(query_id).unwrap();
assert_eq!(status.phase, SqlQueryPhase::Cancelled);
assert!(status.durable_outcome.committed);
}
#[test]
fn direct_query_handle_preserves_receipt_after_tombstone_eviction() {
let registry = Arc::new(SqlQueryRegistry::new(
1,
1,
usize::MAX,
std::time::Duration::from_secs(60),
));
let first = registry.register(SqlQueryOptions::default()).unwrap();
first
.transition(SqlQueryPhase::Queued, SqlQueryPhase::Executing)
.unwrap();
first.record_commit(0, 42);
first.complete_current_statement();
first.try_complete().unwrap();
let second = registry.register(SqlQueryOptions::default()).unwrap();
second
.transition(SqlQueryPhase::Queued, SqlQueryPhase::Executing)
.unwrap();
second.try_complete().unwrap();
assert!(registry.status(first.id()).is_none());
let receipt = sql_terminal_idempotency_receipt(&first.status()).unwrap();
assert_eq!(receipt.outcome.committed_statements, 1);
assert_eq!(receipt.outcome.last_commit_epoch, Some(42));
}
#[test]
fn cancellation_checkpoint_mismatch_returns_typed_error() {
let registry = Arc::new(SqlQueryRegistry::default());
let query = registry.register(SqlQueryOptions::default()).unwrap();
assert!(matches!(
cancellation_checkpoint_error(&query),
mongreldb_query::MongrelQueryError::InvalidQueryState(_)
));
assert_eq!(
query.request_cancel(CancellationReason::ClientRequest),
CancelOutcome::Accepted
);
assert!(matches!(
cancellation_checkpoint_error(&query),
mongreldb_query::MongrelQueryError::QueryCancelled { .. }
));
}
#[tokio::test]
async fn cancellation_after_commit_reports_durable_outcome() {
let registry = Arc::new(SqlQueryRegistry::default());
let query = registry.register(SqlQueryOptions::default()).unwrap();
let query_id = query.id();
query
.transition(SqlQueryPhase::Queued, SqlQueryPhase::Executing)
.unwrap();
query.record_commit(0, 42);
assert_eq!(
query.request_cancel(CancellationReason::ClientRequest),
CancelOutcome::Accepted
);
let error = query.checkpoint().unwrap_err();
query.fail();
let status = registry.status(query_id).unwrap();
let response = query_error_response_with_status(&error, Some(query_id), Some(&status));
assert_eq!(response.status(), StatusCode::CONFLICT);
let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
.await
.unwrap();
let body: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
assert_eq!(body["status"], "cancelled_after_commit");
assert_eq!(body["error"]["code"], "QUERY_CANCELLED_AFTER_COMMIT");
assert_eq!(body["committed"], true);
assert_eq!(body["outcome"]["committed_statements"], 1);
assert_eq!(body["outcome"]["last_commit_epoch"], 42);
assert_eq!(body["outcome"]["last_commit_epoch_text"], "42");
assert_eq!(body["first_commit_statement_index"], 0);
assert_eq!(body["last_commit_statement_index"], 0);
assert_eq!(body["outcome"]["first_commit_statement_index"], 0);
assert_eq!(body["outcome"]["last_commit_statement_index"], 0);
let response = query_error_response_with_status(&error, Some(query_id), None);
assert_eq!(response.status(), StatusCode::CONFLICT);
let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
.await
.unwrap();
let body: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
assert_eq!(body["status"], "cancelled_after_commit");
assert_eq!(body["error"]["code"], "QUERY_CANCELLED_AFTER_COMMIT");
assert_eq!(body["committed"], true);
assert_eq!(body["committed_statements"], 1);
assert_eq!(body["last_commit_epoch"], 42);
assert_eq!(body["last_commit_epoch_text"], "42");
assert_eq!(body["first_commit_statement_index"], 0);
assert_eq!(body["last_commit_statement_index"], 0);
assert_eq!(body["outcome"]["committed"], true);
assert_eq!(body["outcome"]["committed_statements"], 1);
assert_eq!(body["outcome"]["last_commit_epoch"], 42);
assert_eq!(body["outcome"]["last_commit_epoch_text"], "42");
assert_eq!(body["outcome"]["first_commit_statement_index"], 0);
assert_eq!(body["outcome"]["last_commit_statement_index"], 0);
}
#[tokio::test]
async fn commit_outcome_fallback_preserves_exact_progress() {
let query_id: QueryId = "33445566778899001122aabbccddeeff".parse().unwrap();
let error = mongreldb_query::MongrelQueryError::CommitOutcome {
query_id,
committed: true,
committed_statements: 3,
last_commit_epoch: Some(77),
first_commit_statement_index: Some(1),
last_commit_statement_index: Some(4),
completed_statements: 4,
statement_index: 5,
message: "durable outcome retained".into(),
};
let response = query_error_response_with_status(&error, Some(query_id), None);
assert_eq!(response.status(), StatusCode::CONFLICT);
let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
.await
.unwrap();
let body: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
assert_eq!(body["status"], "committed_with_error");
assert_eq!(body["committed"], true);
assert_eq!(body["committed_statements"], 3);
assert_eq!(body["last_commit_epoch"], 77);
assert_eq!(body["last_commit_epoch_text"], "77");
assert_eq!(body["first_commit_statement_index"], 1);
assert_eq!(body["last_commit_statement_index"], 4);
assert_eq!(body["completed_statements"], 4);
assert_eq!(body["statement_index"], 5);
assert_eq!(body["outcome"]["committed_statements"], 3);
assert_eq!(body["outcome"]["last_commit_epoch"], 77);
assert_eq!(body["outcome"]["first_commit_statement_index"], 1);
assert_eq!(body["outcome"]["last_commit_statement_index"], 4);
assert_eq!(body["outcome"]["completed_statements"], 4);
assert_eq!(body["outcome"]["statement_index"], 5);
}
#[test]
fn status_cancel_outcome_matches_cancel_endpoint_state() {
let registry = Arc::new(SqlQueryRegistry::default());
let commit = registry.register(SqlQueryOptions::default()).unwrap();
commit
.transition(SqlQueryPhase::Queued, SqlQueryPhase::Executing)
.unwrap();
commit.enter_commit_critical().unwrap();
assert_eq!(
query_cancel_outcome(®istry.status(commit.id()).unwrap()),
Some("too_late")
);
let completed = registry.register(SqlQueryOptions::default()).unwrap();
completed
.transition(SqlQueryPhase::Queued, SqlQueryPhase::Executing)
.unwrap();
completed.try_complete().unwrap();
assert_eq!(
query_cancel_outcome(®istry.status(completed.id()).unwrap()),
Some("already_finished")
);
}
}
#[cfg(test)]
mod wal_stream_tests {
use super::*;
use mongreldb_client::ReplicationFollower;
use mongreldb_core::Database;
use tempfile::tempdir;
#[tokio::test]
async fn wal_stream_returns_records_after_commit() {
let dir = tempdir().unwrap();
let db = Arc::new(Database::create(dir.path()).unwrap());
let table_schema = mongreldb_core::schema::Schema {
schema_id: 1,
columns: vec![mongreldb_core::schema::ColumnDef {
id: 1,
name: "id".into(),
ty: TypeId::Int64,
flags: mongreldb_core::schema::ColumnFlags::empty()
.with(mongreldb_core::schema::ColumnFlags::PRIMARY_KEY),
default_value: None,
}],
indexes: vec![],
colocation: vec![],
constraints: Default::default(),
clustered: false,
};
db.create_table("items", table_schema).unwrap();
let handle = db.table("items").unwrap();
handle.lock().put(vec![(1, Value::Int64(1))]).unwrap();
handle.lock().flush().unwrap();
let app = build_app(db);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
let resp = reqwest::get(format!("http://{addr}/wal/stream"))
.await
.unwrap();
assert_eq!(resp.status(), 200);
let body = resp.text().await.unwrap();
assert!(!body.is_empty(), "wal_stream should return records");
assert!(body.contains("seq"), "response should contain seq field");
}
#[tokio::test]
async fn follower_bootstraps_and_applies_incremental_commit() {
let leader_dir = tempdir().unwrap();
let follower_dir = tempdir().unwrap();
let follower_path = follower_dir.path().join("copy");
let db = Arc::new(Database::create(leader_dir.path()).unwrap());
db.create_table(
"items",
mongreldb_core::schema::Schema {
schema_id: 1,
columns: vec![mongreldb_core::schema::ColumnDef {
id: 1,
name: "id".into(),
ty: TypeId::Int64,
flags: mongreldb_core::schema::ColumnFlags::empty()
.with(mongreldb_core::schema::ColumnFlags::PRIMARY_KEY),
default_value: None,
}],
indexes: vec![],
colocation: vec![],
constraints: Default::default(),
clustered: false,
},
)
.unwrap();
let handle = db.table("items").unwrap();
handle.lock().put(vec![(1, Value::Int64(1))]).unwrap();
handle.lock().commit().unwrap();
let app = build_app_with_config(
Arc::clone(&db),
std::iter::empty::<Arc<dyn ExternalTableModule>>(),
Some("replication-secret".into()),
None,
);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
let leader_url = format!("http://{addr}");
let first_path = follower_path.clone();
let (mut follower, initial) = tokio::task::spawn_blocking(move || {
let mut follower = ReplicationFollower::new(&leader_url, first_path)
.unwrap()
.with_bearer_token("replication-secret");
let applied = follower.sync().unwrap();
(follower, applied)
})
.await
.unwrap();
assert_eq!(initial, 0);
handle.lock().put(vec![(1, Value::Int64(2))]).unwrap();
handle.lock().commit().unwrap();
let applied = tokio::task::spawn_blocking(move || {
let count = follower.sync().unwrap();
(follower, count)
})
.await
.unwrap();
follower = applied.0;
assert!(applied.1 > 0);
assert!(follower.last_epoch() > 0);
let replica = Database::open(&follower_path).unwrap();
assert_eq!(replica.table("items").unwrap().lock().count(), 2);
drop(replica);
db.set_spill_threshold(1);
db.transaction(|txn| {
txn.put("items", vec![(1, Value::Int64(3))])?;
Ok(())
})
.unwrap();
let (follower_after_bootstrap, applied) = tokio::task::spawn_blocking(move || {
let count = follower.sync().unwrap();
(follower, count)
})
.await
.unwrap();
assert_eq!(applied, 0, "spilled run should trigger safe rebootstrap");
assert!(follower_after_bootstrap.last_epoch() > 0);
let replica = Database::open(&follower_path).unwrap();
assert_eq!(replica.table("items").unwrap().lock().count(), 3);
}
}
#[cfg(test)]
mod metrics_tests {
use super::*;
use mongreldb_core::Database;
use tempfile::tempdir;
async fn setup() -> (tempfile::TempDir, std::net::SocketAddr) {
let dir = tempdir().unwrap();
let db = Arc::new(Database::create(dir.path()).unwrap());
let table_schema = mongreldb_core::schema::Schema {
schema_id: 1,
columns: vec![mongreldb_core::schema::ColumnDef {
id: 1,
name: "id".into(),
ty: TypeId::Int64,
flags: mongreldb_core::schema::ColumnFlags::empty()
.with(mongreldb_core::schema::ColumnFlags::PRIMARY_KEY),
default_value: None,
}],
indexes: vec![],
colocation: vec![],
constraints: Default::default(),
clustered: false,
};
db.create_table("items", table_schema).unwrap();
let app = build_app(db);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
(dir, addr)
}
#[tokio::test]
async fn metrics_endpoint_returns_prometheus_text() {
let (_dir, addr) = setup().await;
let client = reqwest::Client::new();
let _ = client
.post(format!("http://{addr}/tables/items/put"))
.json(&json!({ "row": [1, 1] }))
.send()
.await
.unwrap();
let _ = client
.post(format!("http://{addr}/sql"))
.json(&json!({ "sql": "SELECT count(*) FROM items" }))
.send()
.await
.unwrap();
let resp = client
.get(format!("http://{addr}/metrics"))
.send()
.await
.unwrap();
assert_eq!(resp.status(), 200);
let ct = resp
.headers()
.get("content-type")
.and_then(|v| v.to_str().ok())
.unwrap_or_default()
.to_string();
assert!(
ct.contains("text/plain"),
"content-type is prometheus text: {ct}"
);
let body = resp.text().await.unwrap();
assert!(body.contains("# TYPE mongreldb_sql_queries_total counter"));
assert!(body.contains("# TYPE mongreldb_puts_total counter"));
assert!(body.contains("# TYPE mongreldb_tables gauge"));
assert!(
body.contains("mongreldb_sql_queries_total 1"),
"sql_queries counter should reflect the /sql call: {body}"
);
assert!(
body.contains("mongreldb_puts_total 1"),
"puts counter should reflect the put call: {body}"
);
assert!(body.contains("mongreldb_tables 1"));
}
#[tokio::test]
async fn metrics_error_counter_increments_on_bad_sql() {
let (_dir, addr) = setup().await;
let client = reqwest::Client::new();
let _ = client
.post(format!("http://{addr}/sql"))
.json(&json!({ "sql": "SELECT * FROM does_not_exist" }))
.send()
.await
.unwrap();
let body = client
.get(format!("http://{addr}/metrics"))
.send()
.await
.unwrap()
.text()
.await
.unwrap();
assert!(
body.contains("mongreldb_sql_errors_total 1"),
"sql_errors should increment on a failed query: {body}"
);
}
#[tokio::test]
async fn arrow_stream_returns_ipc_stream_bytes() {
let (_dir, addr) = setup().await;
let client = reqwest::Client::new();
for i in 1..=3 {
let resp = client
.post(format!("http://{addr}/tables/items/put"))
.json(&json!({ "row": [1, i] }))
.send()
.await
.unwrap();
assert_eq!(resp.status(), 200, "put should succeed");
}
let _ = client
.post(format!("http://{addr}/tables/items/commit"))
.send()
.await
.unwrap();
let count_body = client
.get(format!("http://{addr}/tables/items/count"))
.send()
.await
.unwrap()
.text()
.await
.unwrap();
assert!(
count_body.contains("\"count\":3"),
"expected 3 visible rows, got: {count_body}"
);
let resp = client
.post(format!("http://{addr}/sql"))
.json(&json!({ "sql": "SELECT count(*) FROM items", "format": "arrow-stream" }))
.send()
.await
.unwrap();
assert_eq!(resp.status(), 200, "streaming query should succeed");
let ct = resp
.headers()
.get("content-type")
.and_then(|v| v.to_str().ok())
.unwrap_or_default()
.to_string();
assert!(
ct.contains("application/vnd.apache.arrow.stream"),
"content-type should be the arrow stream format: {ct}"
);
let bytes = resp.bytes().await.unwrap();
assert!(
!bytes.is_empty(),
"arrow stream body should contain schema + batch + EOS"
);
assert!(
bytes.starts_with(&0xFFFFFFFFu32.to_le_bytes()),
"arrow stream must begin with the IPC continuation marker"
);
assert!(
bytes.ends_with(&[0u8, 0, 0, 0]),
"arrow stream should end with the EOS marker (trailing zero length)"
);
}
}
#[cfg(test)]
mod streaming_tests {
use super::*;
use arrow::array::Int64Array;
use arrow::datatypes::{DataType, Field, Schema};
use arrow::record_batch::RecordBatch;
use datafusion::common::DataFusionError;
use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
use futures::StreamExt;
use std::sync::Arc as StdArc;
fn batch_stream(batches: Vec<RecordBatch>) -> mongreldb_query::MongrelRecordBatchStream {
let schema = batches
.first()
.map(RecordBatch::schema)
.unwrap_or_else(|| StdArc::new(Schema::empty()));
let batches =
futures::stream::iter(batches.into_iter().map(Ok::<RecordBatch, DataFusionError>));
Box::pin(RecordBatchStreamAdapter::new(schema, batches))
}
#[tokio::test]
async fn arrow_stream_serializes_multiple_batches_roundtrip() {
let schema = StdArc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)]));
let b1 = RecordBatch::try_new(
schema.clone(),
vec![StdArc::new(Int64Array::from(vec![1, 2]))],
)
.unwrap();
let b2 = RecordBatch::try_new(
schema.clone(),
vec![StdArc::new(Int64Array::from(vec![3, 4, 5]))],
)
.unwrap();
let resp = sql_arrow_stream_response(batch_stream(vec![b1, b2]));
let bytes = axum::body::to_bytes(resp.into_body(), 1 << 20)
.await
.unwrap();
assert!(bytes.starts_with(&0xFFFFFFFFu32.to_le_bytes()));
let slice: &[u8] = bytes.as_ref();
let mut reader = arrow::ipc::reader::StreamReader::try_new(slice, None).unwrap();
let mut total_rows = 0;
for batch in reader.by_ref() {
let batch = batch.expect("each IPC message should decode");
total_rows += batch.num_rows();
}
assert_eq!(
total_rows, 5,
"all rows should round-trip through the stream"
);
}
#[tokio::test]
async fn arrow_stream_emits_schema_before_first_batch() {
let schema = StdArc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)]));
let pending = futures::stream::pending::<Result<RecordBatch, DataFusionError>>();
let batches = Box::pin(RecordBatchStreamAdapter::new(schema, pending));
let mut body = sql_arrow_stream_response(batches)
.into_body()
.into_data_stream();
let chunk = tokio::time::timeout(std::time::Duration::from_millis(100), body.next())
.await
.expect("schema chunk should not wait for a query batch")
.unwrap()
.unwrap();
assert!(chunk.starts_with(&0xFFFFFFFFu32.to_le_bytes()));
}
#[tokio::test]
async fn arrow_stream_empty_query_is_valid_ipc() {
let resp = sql_arrow_stream_response(batch_stream(Vec::new()));
let bytes = axum::body::to_bytes(resp.into_body(), 1 << 20)
.await
.unwrap();
let slice: &[u8] = bytes.as_ref();
let reader = arrow::ipc::reader::StreamReader::try_new(slice, None).unwrap();
assert_eq!(reader.count(), 0);
}
#[tokio::test]
async fn buffered_output_limits_are_typed() {
let dir = tempfile::tempdir().unwrap();
let db = StdArc::new(mongreldb_core::Database::create(dir.path()).unwrap());
let session = MongrelSession::open(db).unwrap();
let query = session.register_query(SqlQueryOptions::default()).unwrap();
let output = session
.run_with_query_for_serialization("SELECT 1", query)
.await
.unwrap();
let row_error =
serialize_buffered_output("json", output.batches(), output.query(), 0, 1024, None)
.unwrap_err();
assert!(matches!(row_error, BufferedSerializationError::Limit(_)));
let byte_error =
serialize_buffered_output("json", output.batches(), output.query(), 10, 1, None)
.unwrap_err();
assert!(matches!(byte_error, BufferedSerializationError::Limit(_)));
output.fail();
}
}
#[cfg(test)]
mod audit_tests {
use super::*;
use mongreldb_core::Database;
use tempfile::tempdir;
async fn auth_setup(password: &str) -> std::net::SocketAddr {
let dir = tempdir().unwrap();
let db = Arc::new(Database::create(dir.path()).unwrap());
db.create_user("alice", password).unwrap();
db.set_user_admin("alice", true).unwrap();
let app = build_app_full(
db,
std::iter::empty::<Arc<dyn ExternalTableModule>>(),
None,
None,
true,
);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
addr
}
#[tokio::test]
async fn audit_records_login_success_and_failure() {
let addr = auth_setup("s3cret").await;
let client = reqwest::Client::new();
let resp = client
.get(format!("http://{addr}/health"))
.header("Authorization", basic("alice", "s3cret"))
.send()
.await
.unwrap();
assert_eq!(resp.status(), 200);
let resp = client
.get(format!("http://{addr}/health"))
.header("Authorization", basic("alice", "wrong"))
.send()
.await
.unwrap();
assert_eq!(resp.status(), 401);
let body = client
.get(format!("http://{addr}/audit"))
.header("Authorization", basic("alice", "s3cret"))
.send()
.await
.unwrap()
.text()
.await
.unwrap();
assert!(
body.contains("\"action\":\"login.ok\""),
"audit should record the successful login: {body}"
);
assert!(
body.contains("\"action\":\"login.fail\""),
"audit should record the failed login: {body}"
);
assert!(
body.contains("\"principal\":\"alice\""),
"audit should attribute events to alice: {body}"
);
}
#[tokio::test]
async fn audit_records_ddl_sql() {
let dir = tempdir().unwrap();
let db = Arc::new(Database::create(dir.path()).unwrap());
let app = build_app(db);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
let client = reqwest::Client::new();
let _ = client
.post(format!("http://{addr}/sql"))
.json(&json!({ "sql": "CREATE TABLE t (id BIGINT PRIMARY KEY)" }))
.send()
.await
.unwrap();
let _ = client
.post(format!("http://{addr}/sql"))
.json(&json!({ "sql": "SELECT 1" }))
.send()
.await
.unwrap();
let body = client
.get(format!("http://{addr}/audit"))
.send()
.await
.unwrap()
.text()
.await
.unwrap();
assert!(
body.contains("\"action\":\"ddl.ok\""),
"audit should record the successful DDL statement: {body}"
);
assert!(
body.contains("CREATE TABLE"),
"audit detail should carry the DDL snippet: {body}"
);
assert!(
!body.contains("SELECT 1"),
"non-DDL reads should not be audited: {body}"
);
}
#[tokio::test]
async fn audit_redacts_credential_passwords() {
let dir = tempdir().unwrap();
let db = Arc::new(Database::create(dir.path()).unwrap());
let app = build_app(db);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
let client = reqwest::Client::new();
let _ = client
.post(format!("http://{addr}/sql"))
.json(&json!({ "sql": "CREATE USER alice WITH PASSWORD 'topsecret'" }))
.send()
.await
.unwrap();
let body = client
.get(format!("http://{addr}/audit"))
.send()
.await
.unwrap()
.text()
.await
.unwrap();
assert!(
!body.contains("topsecret"),
"password must never appear in the audit log: {body}"
);
assert!(
body.contains("redacted credential statement"),
"credential DDL should be recorded as redacted: {body}"
);
}
fn basic(user: &str, pass: &str) -> String {
let raw = format!("{user}:{pass}");
format!("Basic {}", base64_encode(raw.as_bytes()))
}
fn base64_encode(input: &[u8]) -> String {
const TABLE: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
let mut out = String::new();
let mut buf = 0u32;
let mut bits = 0u32;
for &b in input {
buf = (buf << 8) | b as u32;
bits += 8;
while bits >= 6 {
bits -= 6;
out.push(TABLE[((buf >> bits) & 0x3F) as usize] as char);
}
}
if bits > 0 {
out.push(TABLE[((buf << (6 - bits)) & 0x3F) as usize] as char);
}
while !out.len().is_multiple_of(4) {
out.push('=');
}
out
}
}
#[cfg(test)]
mod session_tests {
use super::*;
use mongreldb_core::Database;
use tempfile::tempdir;
async fn setup() -> (tempfile::TempDir, std::net::SocketAddr) {
let dir = tempdir().unwrap();
let db = Arc::new(Database::create(dir.path()).unwrap());
let table_schema = mongreldb_core::schema::Schema {
schema_id: 1,
columns: vec![mongreldb_core::schema::ColumnDef {
id: 1,
name: "id".into(),
ty: TypeId::Int64,
flags: mongreldb_core::schema::ColumnFlags::empty()
.with(mongreldb_core::schema::ColumnFlags::PRIMARY_KEY),
default_value: None,
}],
indexes: vec![],
colocation: vec![],
constraints: Default::default(),
clustered: false,
};
db.create_table("items", table_schema).unwrap();
let app = build_app(db);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
(dir, addr)
}
async fn open_session(client: &reqwest::Client, addr: &std::net::SocketAddr) -> String {
let resp = client
.post(format!("http://{addr}/sessions"))
.send()
.await
.unwrap();
assert_eq!(resp.status(), 200);
resp.json::<serde_json::Value>()
.await
.unwrap()
.get("session_id")
.unwrap()
.as_str()
.unwrap()
.to_string()
}
async fn sql_on(
client: &reqwest::Client,
addr: &std::net::SocketAddr,
session: &str,
sql: &str,
) -> reqwest::Response {
client
.post(format!("http://{addr}/sql"))
.header("X-Session-ID", session)
.json(&json!({ "sql": sql }))
.send()
.await
.unwrap()
}
async fn count_items(client: &reqwest::Client, addr: &std::net::SocketAddr) -> u64 {
client
.get(format!("http://{addr}/tables/items/count"))
.send()
.await
.unwrap()
.json::<serde_json::Value>()
.await
.unwrap()
.get("count")
.unwrap()
.as_u64()
.unwrap()
}
#[tokio::test]
async fn cross_request_transaction_commits() {
let (_dir, addr) = setup().await;
let client = reqwest::Client::new();
let session = open_session(&client, &addr).await;
let r = sql_on(&client, &addr, &session, "BEGIN").await;
assert_eq!(r.status(), 200);
let r = sql_on(
&client,
&addr,
&session,
"INSERT INTO items (id) VALUES (1)",
)
.await;
assert_eq!(r.status(), 200, "INSERT should stage successfully");
assert_eq!(count_items(&client, &addr).await, 0);
let r = sql_on(&client, &addr, &session, "COMMIT").await;
assert_eq!(r.status(), 200);
assert_eq!(count_items(&client, &addr).await, 1);
}
#[tokio::test]
async fn cross_request_transaction_rolls_back() {
let (_dir, addr) = setup().await;
let client = reqwest::Client::new();
let session = open_session(&client, &addr).await;
sql_on(&client, &addr, &session, "BEGIN").await;
sql_on(
&client,
&addr,
&session,
"INSERT INTO items (id) VALUES (5)",
)
.await;
let r = sql_on(&client, &addr, &session, "ROLLBACK").await;
assert_eq!(r.status(), 200);
assert_eq!(
count_items(&client, &addr).await,
0,
"rollback discards staged writes"
);
}
#[tokio::test]
async fn unknown_session_id_is_404() {
let (_dir, addr) = setup().await;
let client = reqwest::Client::new();
let resp = client
.post(format!("http://{addr}/sql"))
.header("X-Session-ID", "does-not-exist")
.json(&json!({ "sql": "SELECT 1" }))
.send()
.await
.unwrap();
assert_eq!(resp.status(), 404);
}
#[tokio::test]
async fn invalid_session_headers_do_not_autocommit() {
let (_dir, addr) = setup().await;
let client = reqwest::Client::new();
let resp = client
.post(format!("http://{addr}/sql"))
.header(
"X-Session-ID",
reqwest::header::HeaderValue::from_bytes(&[0xff]).unwrap(),
)
.json(&json!({ "sql": "INSERT INTO items (id) VALUES (1)" }))
.send()
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
let resp = client
.post(format!("http://{addr}/sql"))
.header("X-Session-ID", "x".repeat(257))
.json(&json!({ "sql": "INSERT INTO items (id) VALUES (2)" }))
.send()
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
assert_eq!(count_items(&client, &addr).await, 0);
}
#[tokio::test]
async fn close_session_ends_cross_request_state() {
let (_dir, addr) = setup().await;
let client = reqwest::Client::new();
let session = open_session(&client, &addr).await;
sql_on(&client, &addr, &session, "BEGIN").await;
let r = client
.delete(format!("http://{addr}/sessions/{session}"))
.send()
.await
.unwrap();
assert_eq!(r.status(), 200);
let resp = sql_on(&client, &addr, &session, "COMMIT").await;
assert_eq!(resp.status(), 404, "closed session is no longer usable");
}
#[tokio::test]
async fn no_session_header_uses_fresh_ephemeral_session() {
let (_dir, addr) = setup().await;
let client = reqwest::Client::new();
let resp = client
.post(format!("http://{addr}/sql"))
.json(&json!({ "sql": "INSERT INTO items (id) VALUES (42)" }))
.send()
.await
.unwrap();
assert_eq!(resp.status(), 200);
assert_eq!(count_items(&client, &addr).await, 1);
}
#[tokio::test]
async fn prepared_statement_prepare_execute_and_reuse() {
let (_dir, addr) = setup().await;
let client = reqwest::Client::new();
let session = open_session(&client, &addr).await;
sql_on(
&client,
&addr,
&session,
"INSERT INTO items (id) VALUES (1), (2), (3), (4)",
)
.await;
let resp = client
.post(format!("http://{addr}/sessions/{session}/prepare"))
.json(&json!({"name":"gt","sql":"SELECT id FROM items WHERE id > $1"}))
.send()
.await
.unwrap();
assert_eq!(resp.status(), 200);
let resp = client
.post(format!("http://{addr}/sessions/{session}/execute"))
.json(&json!({"name":"gt","params":[2]}))
.send()
.await
.unwrap();
assert_eq!(resp.status(), 200);
let body = resp.json::<serde_json::Value>().await.unwrap();
let arr = body
.as_array()
.expect("execute returns a JSON array of rows");
assert_eq!(arr.len(), 2, "ids > 2 are {{3,4}}: {body}");
let resp = client
.post(format!("http://{addr}/sessions/{session}/execute"))
.json(&json!({"name":"gt","params":[3]}))
.send()
.await
.unwrap();
let body = resp.json::<serde_json::Value>().await.unwrap();
assert_eq!(body.as_array().unwrap().len(), 1, "ids > 3 is {{4}}");
}
#[tokio::test]
async fn prepared_statement_deallocate_then_execute_fails() {
let (_dir, addr) = setup().await;
let client = reqwest::Client::new();
let session = open_session(&client, &addr).await;
let _ = client
.post(format!("http://{addr}/sessions/{session}/prepare"))
.json(&json!({"name":"p","sql":"SELECT $1"}))
.send()
.await
.unwrap();
let deallocate_query_id = "dadadadadadadadadadadadadadadada";
let resp = client
.delete(format!("http://{addr}/sessions/{session}/statements/p"))
.header("X-MongrelDB-Query-ID", deallocate_query_id)
.header("X-MongrelDB-Timeout-Ms", "10000")
.send()
.await
.unwrap();
assert_eq!(resp.status(), 200);
assert_eq!(
resp.headers()
.get("X-MongrelDB-Query-ID")
.unwrap()
.to_str()
.unwrap(),
deallocate_query_id
);
let status = client
.get(format!("http://{addr}/queries/{deallocate_query_id}"))
.send()
.await
.unwrap()
.json::<serde_json::Value>()
.await
.unwrap();
assert_eq!(status["state"], "completed");
assert_eq!(status["operation"], "DEALLOCATE");
let resp = client
.post(format!("http://{addr}/sessions/{session}/execute"))
.json(&json!({"name":"p","params":[1]}))
.send()
.await
.unwrap();
assert_ne!(resp.status(), 200, "execute after DEALLOCATE must fail");
}
#[tokio::test]
async fn prepared_statement_deallocate_honors_pre_registration_cancel() {
let (_dir, addr) = setup().await;
let client = reqwest::Client::new();
let session = open_session(&client, &addr).await;
let prepared = client
.post(format!("http://{addr}/sessions/{session}/prepare"))
.json(&json!({"name":"p","sql":"SELECT $1"}))
.send()
.await
.unwrap();
assert_eq!(prepared.status(), StatusCode::OK);
let query_id = "dbdbdbdbdbdbdbdbdbdbdbdbdbdbdbdb";
let cancel = client
.post(format!("http://{addr}/queries/{query_id}/cancel"))
.header("X-Session-ID", &session)
.send()
.await
.unwrap();
assert_eq!(cancel.status(), StatusCode::ACCEPTED);
let deallocate = client
.delete(format!("http://{addr}/sessions/{session}/statements/p"))
.header("X-MongrelDB-Query-ID", query_id)
.send()
.await
.unwrap();
assert_eq!(deallocate.status().as_u16(), 499);
assert_eq!(
deallocate.json::<serde_json::Value>().await.unwrap()["error"]["code"],
"QUERY_CANCELLED"
);
let execute = client
.post(format!("http://{addr}/sessions/{session}/execute"))
.json(&json!({"name":"p","params":[1]}))
.send()
.await
.unwrap();
assert_eq!(execute.status(), StatusCode::OK);
}
#[tokio::test]
async fn prepared_statement_rejects_bad_name() {
let (_dir, addr) = setup().await;
let client = reqwest::Client::new();
let session = open_session(&client, &addr).await;
let resp = client
.post(format!("http://{addr}/sessions/{session}/prepare"))
.json(&json!({"name":"1bad","sql":"SELECT 1"}))
.send()
.await
.unwrap();
assert_eq!(
resp.status(),
400,
"statement name starting with a digit must be rejected"
);
}
}