diff --git a/crates/client-api/src/routes/database.rs b/crates/client-api/src/routes/database.rs index b15f72a1e7c..f45eee60cd1 100644 --- a/crates/client-api/src/routes/database.rs +++ b/crates/client-api/src/routes/database.rs @@ -42,6 +42,7 @@ use spacetimedb_client_api_messages::name::{ self, DatabaseName, DomainName, MigrationPolicy, PrePublishAutoMigrateResult, PrePublishManualMigrateResult, PrePublishResult, PrettyPrintStyle, PublishOp, PublishResult, }; +use spacetimedb_datastore::db_metrics::DB_METRICS; use spacetimedb_lib::db::raw_def::v10::RawModuleDefV10; use spacetimedb_lib::db::raw_def::v9::RawModuleDefV9; use spacetimedb_lib::{http as st_http, ConnectionId}; @@ -89,7 +90,6 @@ fn allow_creation(auth: &SpacetimeAuth) -> Result<(), ErrorResponse> { } #[derive(Deserialize)] pub struct CallParams { - name_or_identity: NameOrIdentity, reducer: String, } @@ -140,10 +140,8 @@ fn map_procedure_error(e: ProcedureCallError, procedure: &str) -> (StatusCode, S pub async fn call( State(worker_ctx): State, Extension(auth): Extension, - Path(CallParams { - name_or_identity, - reducer, - }): Path, + Extension(ResolvedDatabase(database)): Extension, + Path(CallParams { reducer }): Path, TypedHeader(content_type): TypedHeader, ByteStringBody(body): ByteStringBody, ) -> axum::response::Result { @@ -155,7 +153,8 @@ pub async fn call( let caller_auth: ConnectionAuthCtx = auth.into(); - let (module, Database { owner_identity, .. }) = find_module_and_database(&worker_ctx, name_or_identity).await?; + let owner_identity = database.owner_identity; + let module = find_database_module(&worker_ctx, &database).await?; let fut = async move |module: ModuleHost, caller_identity: Identity, connection_id: ConnectionId| { let result = match module @@ -214,42 +213,37 @@ pub async fn call( with_connection(module, caller_auth, caller_identity, fut).await } -#[derive(Deserialize)] -pub struct HttpRouteRootParams { - name_or_identity: NameOrIdentity, -} - #[derive(Deserialize)] pub struct HttpRouteParams { - name_or_identity: NameOrIdentity, path: String, } pub async fn handle_http_route_root( State(worker_ctx): State, - Path(HttpRouteRootParams { name_or_identity }): Path, + Extension(ResolvedDatabase(database)): Extension, OriginalUri(original_uri): OriginalUri, request: Request, ) -> axum::response::Result { - handle_http_route_impl(worker_ctx, name_or_identity, "".to_string(), original_uri, request).await + handle_http_route_impl(worker_ctx, database, "".to_string(), original_uri, request).await } pub async fn handle_http_route_root_slash( State(worker_ctx): State, - Path(HttpRouteRootParams { name_or_identity }): Path, + Extension(ResolvedDatabase(database)): Extension, OriginalUri(original_uri): OriginalUri, request: Request, ) -> axum::response::Result { - handle_http_route_impl(worker_ctx, name_or_identity, "/".to_string(), original_uri, request).await + handle_http_route_impl(worker_ctx, database, "/".to_string(), original_uri, request).await } pub async fn handle_http_route( State(worker_ctx): State, - Path(HttpRouteParams { name_or_identity, path }): Path, + Extension(ResolvedDatabase(database)): Extension, + Path(HttpRouteParams { path }): Path, OriginalUri(original_uri): OriginalUri, request: Request, ) -> axum::response::Result { - handle_http_route_impl(worker_ctx, name_or_identity, format!("/{path}"), original_uri, request).await + handle_http_route_impl(worker_ctx, database, format!("/{path}"), original_uri, request).await } /// Error response body for unknown user-defined HTTP route. @@ -257,7 +251,7 @@ const NO_SUCH_ROUTE: &str = "Database has not registered a handler for this rout async fn handle_http_route_impl( worker_ctx: S, - name_or_identity: NameOrIdentity, + database: Database, handler_path: String, original_uri: http::Uri, request: Request, @@ -265,7 +259,7 @@ async fn handle_http_route_impl( let (parts, body) = request.into_parts(); let st_method = http_method_to_st(&parts.method); - let (module, _database) = find_module_and_database(&worker_ctx, name_or_identity).await?; + let module = find_database_module(&worker_ctx, &database).await?; let module_def = &module.info().module_def; let Some((handler_id, _handler_def, _route_def)) = module_def.match_http_route(&st_method, &handler_path) else { @@ -460,31 +454,23 @@ pub(crate) fn client_disconnected_error_to_response(err: ReducerCallError) -> Er (StatusCode::INTERNAL_SERVER_ERROR, format!("{:#}", anyhow::anyhow!(err))).into() } -pub(crate) async fn find_leader_and_database( +pub(crate) async fn find_database_leader( worker_ctx: &S, - name_or_identity: NameOrIdentity, -) -> axum::response::Result<(Host, Database)> { - let db_identity = name_or_identity.resolve(worker_ctx).await?; - let database = worker_ctx_find_database(worker_ctx, &db_identity) - .await? - .ok_or_else(|| { - log::error!("Could not find database: {}", db_identity.to_hex()); - NO_SUCH_DATABASE - })?; - + database: &Database, +) -> axum::response::Result { let leader = worker_ctx.leader(database.id).await.map_err(Into::into)?; - Ok((leader, database)) + Ok(leader) } -pub(crate) async fn find_module_and_database( +pub(crate) async fn find_database_module( worker_ctx: &S, - name_or_identity: NameOrIdentity, -) -> axum::response::Result<(ModuleHost, Database)> { - let (leader, database) = find_leader_and_database(worker_ctx, name_or_identity).await?; + database: &Database, +) -> axum::response::Result { + let leader = find_database_leader(worker_ctx, database).await?; let module = leader.module().await.map_err(log_and_500)?; - Ok((module, database)) + Ok(module) } #[derive(Debug, derive_more::From)] @@ -501,10 +487,6 @@ fn procedure_outcome_response(return_val: AlgebraicValue) -> (StatusCode, axum:: ) } -#[derive(Deserialize)] -pub struct SchemaParams { - name_or_identity: NameOrIdentity, -} #[derive(Deserialize)] pub struct SchemaQueryParams { version: SchemaVersion, @@ -520,14 +502,14 @@ enum SchemaVersion { pub async fn schema( State(worker_ctx): State, - Path(SchemaParams { name_or_identity }): Path, + Extension(ResolvedDatabase(database)): Extension, Query(SchemaQueryParams { version }): Query, Extension(auth): Extension, ) -> axum::response::Result where S: ControlStateDelegate + NodeDelegate, { - let (leader, _) = find_leader_and_database(&worker_ctx, name_or_identity).await?; + let leader = find_database_leader(&worker_ctx, &database).await?; // Wait for the module to finish loading rather than returning an immediate // 500 error. The database may still be initializing (replaying the log, // running init reducers, etc.). @@ -560,6 +542,9 @@ pub struct DatabaseParam { name_or_identity: NameOrIdentity, } +#[derive(Clone)] +pub struct ResolvedDatabase(pub Database); + #[derive(sats::Serialize)] struct DatabaseResponse { database_identity: Identity, @@ -580,26 +565,12 @@ impl From for DatabaseResponse { } pub async fn db_info( - State(worker_ctx): State, - Path(DatabaseParam { name_or_identity }): Path, + Extension(ResolvedDatabase(database)): Extension, ) -> axum::response::Result { - log::trace!("Trying to resolve database identity: {name_or_identity:?}"); - let database_identity = name_or_identity.resolve(&worker_ctx).await?; - log::trace!("Resolved identity to: {database_identity:?}"); - let database = worker_ctx_find_database(&worker_ctx, &database_identity) - .await? - .ok_or(NO_SUCH_DATABASE)?; - log::trace!("Fetched database from the worker db for database identity: {database_identity:?}"); - let response = DatabaseResponse::from(database); Ok(axum::Json(sats::serde::SerdeWrapper(response))) } -#[derive(Deserialize)] -pub struct LogsParams { - name_or_identity: NameOrIdentity, -} - #[derive(Deserialize)] pub struct LogsQuery { num_lines: Option, @@ -609,7 +580,7 @@ pub struct LogsQuery { pub async fn logs( State(worker_ctx): State, - Path(LogsParams { name_or_identity }): Path, + Extension(ResolvedDatabase(database)): Extension, Query(LogsQuery { num_lines, follow }): Query, Extension(auth): Extension, ) -> axum::response::Result @@ -619,10 +590,7 @@ where // You should not be able to read the logs from a database that you do not own // so, unless you are the owner, this will fail. - let database_identity: Identity = name_or_identity.resolve(&worker_ctx).await?; - let database = worker_ctx_find_database(&worker_ctx, &database_identity) - .await? - .ok_or(NO_SUCH_DATABASE)?; + let database_identity = database.database_identity; worker_ctx .authorize_action(auth.claims.identity, database.database_identity, Action::ViewModuleLogs) @@ -686,9 +654,16 @@ pub(crate) async fn worker_ctx_find_database( .map_err(log_and_500) } -#[derive(Deserialize)] -pub struct SqlParams { - pub name_or_identity: NameOrIdentity, +pub(crate) async fn find_database_or_404( + worker_ctx: &(impl ControlStateDelegate + ?Sized), + name_or_identity: NameOrIdentity, +) -> axum::response::Result { + let identity = name_or_identity.resolve(worker_ctx).await?; + let database = worker_ctx_find_database(worker_ctx, &identity).await?.ok_or_else(|| { + log::debug!("Identity {identity} in HTTP request does not refer to a database"); + NO_SUCH_DATABASE + })?; + Ok(database) } #[derive(Deserialize)] @@ -741,7 +716,7 @@ where pub async fn sql_direct( worker_ctx: S, - SqlParams { name_or_identity }: SqlParams, + database: Database, SqlQueryParams { confirmed }: SqlQueryParams, caller_identity: Identity, caller_auth: ConnectionAuthCtx, @@ -750,7 +725,7 @@ pub async fn sql_direct( where S: NodeDelegate + ControlStateDelegate + Authorization + 'static, { - let (host, database) = find_leader_and_database(&worker_ctx, name_or_identity).await?; + let host = find_database_leader(&worker_ctx, &database).await?; let module = host.module().await.map_err(log_and_500)?; let fut = async move |_module: ModuleHost, caller_identity: Identity, _connection_id: ConnectionId| { @@ -772,7 +747,7 @@ where pub async fn sql( State(worker_ctx): State, - Path(name_or_identity): Path, + Extension(ResolvedDatabase(database)): Extension, Query(params): Query, Extension(auth): Extension, body: String, @@ -782,7 +757,7 @@ where { let caller_identity = auth.claims.identity; let caller_auth: ConnectionAuthCtx = auth.into(); - let json = sql_direct(worker_ctx, name_or_identity, params, caller_identity, caller_auth, body).await?; + let json = sql_direct(worker_ctx, database, params, caller_identity, caller_auth, body).await?; let total_duration = json.iter().fold(0, |acc, x| acc + x.total_duration_micros); @@ -797,11 +772,6 @@ pub struct DNSParams { name_or_identity: NameOrIdentity, } -#[derive(Deserialize)] -pub struct ReverseDNSParams { - name_or_identity: NameOrIdentity, -} - #[derive(Deserialize)] pub struct DNSQueryParams {} @@ -816,9 +786,9 @@ pub async fn get_identity( pub async fn get_names( State(ctx): State, - Path(ReverseDNSParams { name_or_identity }): Path, + Extension(ResolvedDatabase(database)): Extension, ) -> axum::response::Result { - let database_identity = name_or_identity.resolve(&ctx).await?; + let database_identity = database.database_identity; let names = ctx .reverse_lookup(&database_identity) @@ -846,6 +816,7 @@ pub struct ResetDatabaseQueryParams { pub async fn reset( State(ctx): State, + Extension(ResolvedDatabase(database)): Extension, Path(ResetDatabaseParams { name_or_identity }): Path, Query(ResetDatabaseQueryParams { num_replicas, @@ -854,10 +825,7 @@ pub async fn reset( Extension(auth): Extension, program_bytes: Option, ) -> axum::response::Result> { - let database_identity = name_or_identity.resolve(&ctx).await?; - let database = worker_ctx_find_database(&ctx, &database_identity) - .await? - .ok_or(NO_SUCH_DATABASE)?; + let database_identity = database.database_identity; ctx.authorize_action(auth.claims.identity, database.database_identity, Action::ResetDatabase) .await?; @@ -964,12 +932,8 @@ pub async fn publish( .ok_or_else(|| bad_request("Clear database requires database name or identity".into()))?; let database_identity = name_or_identity.try_resolve(&ctx).await.map_err(log_and_500)?; if let Ok(identity) = database_identity { - let exists = ctx - .get_database_by_identity(&identity) - .await - .map_err(log_and_500)? - .is_some(); - if exists { + let database = ctx.get_database_by_identity(&identity).await.map_err(log_and_500)?; + if let Some(database) = database { if parent.is_some() { return Err(bad_request( "Setting the parent of an existing database is not supported".into(), @@ -978,6 +942,7 @@ pub async fn publish( return self::reset( State(ctx), + Extension(ResolvedDatabase(database)), Path(ResetDatabaseParams { name_or_identity: name_or_identity.clone(), }), @@ -1227,11 +1192,6 @@ fn bad_request(message: Cow<'static, str>) -> ErrorResponse { (StatusCode::BAD_REQUEST, message).into() } -#[derive(serde::Deserialize)] -pub struct PrePublishParams { - name_or_identity: NameOrIdentity, -} - #[derive(serde::Deserialize)] pub struct PrePublishQueryParams { #[serde(default)] @@ -1242,13 +1202,17 @@ pub struct PrePublishQueryParams { pub async fn pre_publish( State(ctx): State, - Path(PrePublishParams { name_or_identity }): Path, + Extension(ResolvedDatabase(database)): Extension, Query(PrePublishQueryParams { style, host_type }): Query, Extension(auth): Extension, program_bytes: Bytes, ) -> axum::response::Result> { + let database_identity = database.database_identity; + // User should not be able to print migration plans for a database that they do not own - let database_identity = resolve_and_authenticate(&ctx, &name_or_identity, &auth).await?; + ctx.authorize_action(auth.claims.identity, database_identity, Action::UpdateDatabase) + .await?; + let style = match style { PrettyPrintStyle::NoColor => AutoMigratePrettyPrintStyle::NoColor, PrettyPrintStyle::AnsiColor => AutoMigratePrettyPrintStyle::AnsiColor, @@ -1310,24 +1274,6 @@ pub async fn pre_publish .map(axum::Json) } -/// Resolves the [`NameOrIdentity`] to a database identity and checks if the -/// `auth` identity owns the database. -async fn resolve_and_authenticate( - ctx: &S, - name_or_identity: &NameOrIdentity, - auth: &SpacetimeAuth, -) -> axum::response::Result { - let database_identity = name_or_identity.resolve(ctx).await?; - let database = worker_ctx_find_database(ctx, &database_identity) - .await? - .ok_or(NO_SUCH_DATABASE)?; - - ctx.authorize_action(auth.claims.identity, database.database_identity, Action::UpdateDatabase) - .await?; - - Ok(database_identity) -} - #[derive(Deserialize)] pub struct DeleteDatabaseParams { pub name_or_identity: NameOrIdentity, @@ -1363,13 +1309,10 @@ pub async fn delete_database( pub async fn lock_database( State(ctx): State, - Path(DeleteDatabaseParams { name_or_identity }): Path, + Extension(ResolvedDatabase(database)): Extension, Extension(auth): Extension, ) -> axum::response::Result { - let database_identity = name_or_identity.resolve(&ctx).await?; - let Some(_database) = worker_ctx_find_database(&ctx, &database_identity).await? else { - return Err(StatusCode::NOT_FOUND.into()); - }; + let database_identity = database.database_identity; ctx.authorize_action(auth.claims.identity, database_identity, Action::DeleteDatabase) .await?; @@ -1383,13 +1326,10 @@ pub async fn lock_database( pub async fn unlock_database( State(ctx): State, - Path(DeleteDatabaseParams { name_or_identity }): Path, + Extension(ResolvedDatabase(database)): Extension, Extension(auth): Extension, ) -> axum::response::Result { - let database_identity = name_or_identity.resolve(&ctx).await?; - let Some(_database) = worker_ctx_find_database(&ctx, &database_identity).await? else { - return Err(StatusCode::NOT_FOUND.into()); - }; + let database_identity = database.database_identity; ctx.authorize_action(auth.claims.identity, database_identity, Action::DeleteDatabase) .await?; @@ -1401,19 +1341,14 @@ pub async fn unlock_database( Ok(()) } -#[derive(Deserialize)] -pub struct AddNameParams { - name_or_identity: NameOrIdentity, -} - pub async fn add_name( State(ctx): State, - Path(AddNameParams { name_or_identity }): Path, + Extension(ResolvedDatabase(database)): Extension, Extension(auth): Extension, name: String, ) -> axum::response::Result { let name = DatabaseName::try_from(name).map_err(|err| (StatusCode::BAD_REQUEST, err.to_string()))?; - let database_identity = name_or_identity.resolve(&ctx).await?; + let database_identity = database.database_identity; let response = ctx .create_dns_record(&auth.claims.identity, &name.into(), &database_identity) @@ -1431,14 +1366,9 @@ pub async fn add_name( Ok((code, axum::Json(response))) } -#[derive(Deserialize)] -pub struct SetNamesParams { - name_or_identity: NameOrIdentity, -} - pub async fn set_names( State(ctx): State, - Path(SetNamesParams { name_or_identity }): Path, + Extension(ResolvedDatabase(database)): Extension, Extension(auth): Extension, names: axum::Json>, ) -> axum::response::Result { @@ -1449,18 +1379,7 @@ pub async fn set_names( .collect::, _>>() .map_err(|(input, e)| (StatusCode::BAD_REQUEST, format!("Error parsing `{input}`: {e}")))?; - let database_identity = name_or_identity.resolve(&ctx).await?; - - let database = ctx - .get_database_by_identity(&database_identity) - .await - .map_err(log_and_500)?; - let Some(database) = database else { - return Ok(( - StatusCode::NOT_FOUND, - axum::Json(name::SetDomainsResult::DatabaseNotFound), - )); - }; + let database_identity = database.database_identity; ctx.authorize_action(auth.claims.identity, database.database_identity, Action::RenameDatabase) .await @@ -1508,30 +1427,15 @@ pub async fn set_names( Ok((status, axum::Json(response))) } -#[derive(serde::Deserialize)] -pub struct TimestampParams { - name_or_identity: NameOrIdentity, -} - /// Returns the database's view of the current time, /// as a SATS-JSON encoded [`Timestamp`]. /// -/// Takes a particular database's [`NameOrIdentity`] as an argument +/// Takes a particular database as an argument /// because in a clusterized SpacetimeDB-cloud deployment, /// this request will be routed to the node running the requested database. -async fn get_timestamp( - State(worker_ctx): State, - Path(TimestampParams { name_or_identity }): Path, +async fn get_timestamp( + Extension(ResolvedDatabase(_database)): Extension, ) -> axum::response::Result { - let db_identity = name_or_identity.resolve(&worker_ctx).await?; - - let _database = worker_ctx_find_database(&worker_ctx, &db_identity) - .await? - .ok_or_else(|| { - log::error!("Could not find database: {}", db_identity.to_hex()); - NO_SUCH_DATABASE - })?; - Ok(axum::Json(sats::serde::SerdeWrapper(Timestamp::now())).into_response()) } @@ -1606,7 +1510,7 @@ where mcp_post: post(crate::routes::mcp::mcp::), pre_publish: post(pre_publish::), db_reset: put(reset::), - timestamp_get: get(get_timestamp::), + timestamp_get: get(get_timestamp), lock_post: post(lock_database::), unlock_post: post(unlock_database::), http_route_root: any(handle_http_route_root::), @@ -1621,25 +1525,46 @@ where S: NodeDelegate + ControlStateDelegate + Authorization + Clone + 'static, { pub fn into_router(self, ctx: S) -> axum::Router { - let db_router = axum::Router::::new() - .route("/", self.db_put) + let resolving_egress_metrics_middleware = axum::middleware::from_fn_with_state( + ctx.clone(), + resolve_database_name_and_count_response_egress_middleware::, + ); + + let counted_db_router = axum::Router::::new() .route("/", self.db_get) - .route("/", self.db_delete) .route("/names", self.names_get) .route("/names", self.names_post) .route("/names", self.names_put) - .route("/identity", self.identity_get) - .route("/subscribe", self.subscribe_get) .route("/call/:reducer", self.call_reducer_procedure_post) .route("/schema", self.schema_get) .route("/logs", self.logs_get) .route("/sql", self.sql_post) - .route("/mcp", self.mcp_post) .route("/unstable/timestamp", self.timestamp_get) .route("/pre_publish", self.pre_publish) .route("/reset", self.db_reset) .route("/lock", self.lock_post) - .route("/unlock", self.unlock_post); + .route("/unlock", self.unlock_post) + .route_layer(resolving_egress_metrics_middleware.clone()); + + // Publishing can create a database for a new name, so it bypasses existing-database resolution. + // Publish operations are infrequent and the responses are small, so we don't mind that we don't measure them. + let db_router = counted_db_router.route("/", self.db_put); + + // These routes have different behavior on a non-existent database than a 404, + // and so resolve the database themselves rather than having the middleware do so. + // Like publish, these operations are infrequent and the responses are small, + // so we don't mind that we don't measure them. + let db_router = db_router + .route("/", self.db_delete) + .route("/identity", self.identity_get) + // I (pgoldman 2026-08-07) am actually somewhat concerned that we do care about measuring egress for MCP requests, + // but the MCP handler's name resolution and error handling are significantly incompatible with the middleware. + .route("/mcp", self.mcp_post); + + // Add the subscribe route after `resolving_egress_metrics_middleware` + // so that its egress bytes don't get counted into `http_response_size_bytes`; + // we have different metrics tracking WebSocket message size. + let db_router = db_router.route("/subscribe", self.subscribe_get); let authed_root_router = axum::Router::new().route( "/", @@ -1661,7 +1586,8 @@ where let http_route_router = axum::Router::::new() .route("/:name_or_identity/route", self.http_route_root) .route("/:name_or_identity/route/", self.http_route_root_slash) - .route("/:name_or_identity/route/*path", self.http_route); + .route("/:name_or_identity/route/*path", self.http_route) + .route_layer(resolving_egress_metrics_middleware); axum::Router::new() .merge(authed_root_router) @@ -1670,6 +1596,64 @@ where } } +/// Resolves an existing database, attaches it as [`ResolvedDatabase`], +/// and counts response bytes in the metric `spacetime_http_response_size_bytes_total`. +/// +/// This middleware returns established name and database `404`s before the handler runs. +/// It is intended for HTTP routes that require an existing database, +/// except WebSocket `subscribe`, whose egress is measured separately. +async fn resolve_database_name_and_count_response_egress_middleware( + State(worker_ctx): State, + Path(DatabaseParam { name_or_identity }): Path, + mut request: Request, + next: axum::middleware::Next, +) -> axum::response::Result +where + S: ControlStateDelegate + Clone + Send + Sync + 'static, +{ + let database = match find_database_or_404(&worker_ctx, name_or_identity).await { + Ok(database) => database, + Err(err) => { + // Fully drain the request before responding + // so clients uploading a body receive the HTTP error + // instead of a broken pipe from a connection closed mid-upload. + while let Some(Ok(_)) = request.body_mut().frame().await {} + return Err(err); + } + }; + request.extensions_mut().insert(ResolvedDatabase(database.clone())); + + let response = next.run(request).await; + let (parts, body) = response.into_parts(); + + // Count the number of bytes used by the headers. + // For guest-defined routes bound to HTTP handlers, these may be arbitrarily large and are worth billing for; + // for built-in routes they will be small and it doesn't really matter one way or another whether we do or don't bill. + // N.b. headers installed by other middleware may or may not be counted here, + // depending on the order in which the middleware applies. + let header_bytes: usize = parts + .headers + .iter() + .map(|(name, value)| name.as_str().len() + value.as_bytes().len()) + .sum(); + + let counter = DB_METRICS + .http_response_size_bytes + .with_label_values(&database.database_identity); + counter.inc_by(header_bytes as u64); + + // `/logs?follow=true` can stream indefinitely. + // Counting frames as they are emitted preserves streaming behavior and avoids buffering the response. + let body = body.map_frame(move |frame| { + if let Some(data) = frame.data_ref() { + counter.inc_by(data.len() as u64); + } + frame + }); + + Ok(axum::response::Response::from_parts(parts, Body::new(body))) +} + #[cfg(test)] mod tests { use super::*; @@ -1679,7 +1663,7 @@ mod tests { Action, Authorization, ControlStateReadAccess, ControlStateWriteAccess, MaybeMisdirected, Unauthorized, }; use async_trait::async_trait; - use axum::body::Body; + use axum::body::{Body, Bytes}; use http::Request; use spacetimedb::auth::identity::{JwtError, JwtErrorKind, SpacetimeIdentityClaims}; use spacetimedb::auth::token_validation::{TokenSigner, TokenValidationError, TokenValidator}; @@ -1690,10 +1674,15 @@ mod tests { use spacetimedb_client_api_messages::name::{ DomainName, InsertDomainResult, RegisterTldResult, SetDomainsResult, Tld, }; + use spacetimedb_lib::Hash; use spacetimedb_paths::server::ModuleLogsDir; use spacetimedb_paths::FromPathUnchecked; use spacetimedb_schema::auto_migrate::{MigrationPolicy, PrettyPrintStyle}; + use std::collections::HashMap; + use std::convert::Infallible; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use tower::util::ServiceExt; + #[derive(Clone, Default)] struct DummyValidator; @@ -1710,8 +1699,14 @@ mod tests { } impl TokenSigner for DummyJwtProvider { - fn sign(&self, _claims: &T) -> Result { - Err(JwtError::from(JwtErrorKind::InvalidSignature)) + fn sign(&self, claims: &T) -> Result { + use base64::{engine::general_purpose, Engine}; + + let payload = serde_json::to_vec(claims).map_err(|_| JwtError::from(JwtErrorKind::InvalidSignature))?; + Ok(format!( + "test.{}.signature", + general_purpose::URL_SAFE_NO_PAD.encode(payload) + )) } } @@ -1736,6 +1731,9 @@ mod tests { jwt: DummyJwtProvider, client_actor_index: std::sync::Arc, module_logs_dir: ModuleLogsDir, + databases: std::sync::Arc>, + dns: std::sync::Arc>, + dns_lookups: std::sync::Arc, } impl DummyState { @@ -1746,8 +1744,82 @@ mod tests { }, client_actor_index: std::sync::Arc::new(ClientActorIndex::new()), module_logs_dir: ModuleLogsDir::from_path_unchecked(std::env::temp_dir()), + databases: std::sync::Arc::new(HashMap::new()), + dns: std::sync::Arc::new(HashMap::new()), + dns_lookups: std::sync::Arc::new(AtomicUsize::new(0)), + } + } + + fn with_database(mut self, database_identity: Identity) -> Self { + let mut databases = HashMap::new(); + databases.insert(database_identity, test_database(database_identity)); + self.databases = std::sync::Arc::new(databases); + self + } + + fn with_dns(mut self, name: &str, database_identity: Identity) -> Self { + let mut dns = HashMap::new(); + dns.insert(name.to_owned(), database_identity); + self.dns = std::sync::Arc::new(dns); + self + } + + fn dns_lookups(&self) -> usize { + self.dns_lookups.load(Ordering::Relaxed) + } + } + + fn test_identity(byte: u8) -> Identity { + Identity::from_byte_array([byte; 32]) + } + + fn test_database(database_identity: Identity) -> Database { + Database { + id: u64::from(database_identity.to_byte_array()[0]), + database_identity, + owner_identity: test_identity(254), + host_type: HostType::Wasm, + initial_program: Hash::from_byte_array([0; 32]), + bootstrap_generation: 0, + } + } + + fn http_response_size_metric(database_identity: Identity) -> u64 { + DB_METRICS + .http_response_size_bytes + .with_label_values(&database_identity) + .get() + } + + fn collected_http_response_size_metric(database_identity: Identity) -> Option { + let db_label = database_identity.to_hex(); + for metric_family in prometheus::core::Collector::collect(&DB_METRICS.http_response_size_bytes) { + if metric_family.name() != "spacetime_http_response_size_bytes_total" { + continue; + } + + for metric in metric_family.get_metric() { + let has_db_label = metric + .get_label() + .iter() + .any(|label| label.name() == "db" && label.value() == db_label.as_str()); + if has_db_label { + return Some(metric.get_counter().value() as u64); + } } } + + None + } + + fn remove_http_response_size_metric(database_identity: Identity) { + let _ = DB_METRICS + .http_response_size_bytes + .remove_label_values(&database_identity); + } + + fn text_plain_header_bytes() -> u64 { + "content-type".len() as u64 + "text/plain; charset=utf-8".len() as u64 } impl HasWebSocketOptions for DummyState { @@ -1817,8 +1889,8 @@ mod tests { async fn get_database_by_id(&self, _id: u64) -> anyhow::Result> { Ok(None) } - async fn get_database_by_identity(&self, _database_identity: &Identity) -> anyhow::Result> { - Ok(None) + async fn get_database_by_identity(&self, database_identity: &Identity) -> anyhow::Result> { + Ok(self.databases.get(database_identity).cloned()) } async fn get_databases(&self) -> anyhow::Result> { Ok(Vec::new()) @@ -1835,8 +1907,9 @@ mod tests { async fn get_energy_balance(&self, _identity: &Identity) -> anyhow::Result> { Ok(None) } - async fn lookup_database_identity(&self, _domain: &str) -> anyhow::Result> { - Ok(None) + async fn lookup_database_identity(&self, domain: &str) -> anyhow::Result> { + self.dns_lookups.fetch_add(1, Ordering::Relaxed); + Ok(self.dns.get(domain).copied()) } async fn reverse_lookup(&self, _database_identity: &Identity) -> anyhow::Result> { Ok(Vec::new()) @@ -1961,9 +2034,315 @@ mod tests { assert_eq!(response.status(), StatusCode::NOT_FOUND); let body = response.into_body().collect().await.unwrap().to_bytes(); // We'll get this error message out of the stack: - // - `find_module_and_database` - // - `find_leader_and_database` - // - `name_or_identity.resolve(worker_ctx)` -> `NameOrIdentity::resolve` + // The database-resolution middleware returns this established error + // before the unauthenticated HTTP route handler is invoked. assert_eq!(body, "`not-a-database` not found"); } + + #[tokio::test] + async fn http_response_egress_metric_counts_database_routes() { + let database_identity = test_identity(11); + remove_http_response_size_metric(database_identity); + + let state = DummyState::new().with_database(database_identity); + let app = DatabaseRoutes:: { + db_get: axum::routing::get(|| async { ([(http::header::CONTENT_TYPE, "text/plain")], "hello") }), + ..Default::default() + } + .into_router(state.clone()) + .with_state(state); + + let response = app + .oneshot( + Request::builder() + .uri(format!("/{database_identity}")) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let body = response.into_body().collect().await.unwrap().to_bytes(); + + assert_eq!(body, "hello"); + assert_eq!( + http_response_size_metric(database_identity), + "content-type".len() as u64 + "text/plain".len() as u64 + "hello".len() as u64 + ); + + remove_http_response_size_metric(database_identity); + } + + #[tokio::test] + async fn http_response_egress_metric_counts_user_http_route_headers_and_body() { + let database_identity = test_identity(12); + remove_http_response_size_metric(database_identity); + + let state = DummyState::new() + .with_database(database_identity) + .with_dns("metric-test", database_identity); + let app = DatabaseRoutes:: { + http_route: axum::routing::any(|| async { ([("x-large-test-header", "abcdef")], "route-body") }), + ..Default::default() + } + .into_router(state.clone()) + .with_state(state); + + let response = app + .oneshot( + Request::builder() + .uri("/metric-test/route/health") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let body = response.into_body().collect().await.unwrap().to_bytes(); + + assert_eq!(body, "route-body"); + assert_eq!( + http_response_size_metric(database_identity), + "x-large-test-header".len() as u64 + + "abcdef".len() as u64 + + text_plain_header_bytes() + + "route-body".len() as u64 + ); + + remove_http_response_size_metric(database_identity); + } + + #[tokio::test] + async fn resolving_middleware_returns_not_found_without_creating_metrics_labels() { + let resolved_but_missing_identity = test_identity(13); + let arbitrary_identity = test_identity(14); + remove_http_response_size_metric(resolved_but_missing_identity); + remove_http_response_size_metric(arbitrary_identity); + + let state = DummyState::new().with_dns("missing-db", resolved_but_missing_identity); + let app = DatabaseRoutes:: { + db_get: axum::routing::get(|| async { "not counted" }), + ..Default::default() + } + .into_router(state.clone()) + .with_state(state); + + let response = app + .clone() + .oneshot(Request::builder().uri("/unresolved-name").body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::NOT_FOUND); + assert_eq!( + response.into_body().collect().await.unwrap().to_bytes(), + "`unresolved-name` not found" + ); + + let response = app + .clone() + .oneshot(Request::builder().uri("/missing-db").body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::NOT_FOUND); + assert_eq!( + response.into_body().collect().await.unwrap().to_bytes(), + "No such database." + ); + + let response = app + .oneshot( + Request::builder() + .uri(format!("/{arbitrary_identity}")) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::NOT_FOUND); + assert_eq!( + response.into_body().collect().await.unwrap().to_bytes(), + "No such database." + ); + + assert_eq!(collected_http_response_size_metric(resolved_but_missing_identity), None); + assert_eq!(collected_http_response_size_metric(arbitrary_identity), None); + + remove_http_response_size_metric(resolved_but_missing_identity); + remove_http_response_size_metric(arbitrary_identity); + } + + #[tokio::test] + async fn resolving_middleware_drains_request_body_before_not_found() { + let body_was_polled = std::sync::Arc::new(AtomicBool::new(false)); + let body_was_polled_by_stream = body_was_polled.clone(); + let body = Body::from_stream(futures::stream::once(async move { + body_was_polled_by_stream.store(true, Ordering::Relaxed); + Ok::<_, Infallible>(Bytes::from_static(b"module")) + })); + let state = DummyState::new(); + let app = DatabaseRoutes:: { + db_get: axum::routing::get(|| async { "not reached" }), + ..Default::default() + } + .into_router(state.clone()) + .with_state(state); + + let response = app + .oneshot(Request::builder().uri("/unresolved-name").body(body).unwrap()) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::NOT_FOUND); + assert!(body_was_polled.load(Ordering::Relaxed)); + } + + #[tokio::test] + async fn resolving_middleware_attaches_database_and_resolves_a_name_once() { + let database_identity = test_identity(17); + remove_http_response_size_metric(database_identity); + + let state = DummyState::new() + .with_database(database_identity) + .with_dns("named-database", database_identity); + let app = DatabaseRoutes:: { + db_get: axum::routing::get( + |Extension(ResolvedDatabase(database)): Extension| async move { + database.database_identity.to_string() + }, + ), + ..Default::default() + } + .into_router(state.clone()) + .with_state(state.clone()); + + let response = app + .oneshot(Request::builder().uri("/named-database").body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!( + response.into_body().collect().await.unwrap().to_bytes(), + database_identity.to_string() + ); + assert_eq!(state.dns_lookups(), 1); + + remove_http_response_size_metric(database_identity); + } + + #[tokio::test] + async fn db_info_resolves_a_database_name_once() { + let database_identity = test_identity(19); + remove_http_response_size_metric(database_identity); + + let state = DummyState::new() + .with_database(database_identity) + .with_dns("db-info", database_identity); + let app = DatabaseRoutes::::default() + .into_router(state.clone()) + .with_state(state.clone()); + + let response = app + .oneshot(Request::builder().uri("/db-info").body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(state.dns_lookups(), 1); + + remove_http_response_size_metric(database_identity); + } + + #[tokio::test] + async fn publish_delete_and_identity_bypass_resolving_middleware() { + let missing_identity = test_identity(18); + remove_http_response_size_metric(missing_identity); + + let state = DummyState::new(); + let app = DatabaseRoutes:: { + db_put: axum::routing::put(|| async { "publish" }), + db_delete: axum::routing::delete(|| async { "delete" }), + identity_get: axum::routing::get(|| async { "identity" }), + ..Default::default() + } + .into_router(state.clone()) + .with_state(state); + + let delete_uri = format!("/{missing_identity}"); + let identity_uri = format!("/{missing_identity}/identity"); + for (method, uri, expected) in [ + (http::Method::PUT, "/unregistered-name", "publish"), + (http::Method::DELETE, delete_uri.as_str(), "delete"), + (http::Method::GET, identity_uri.as_str(), "identity"), + ] { + let response = app + .clone() + .oneshot(Request::builder().method(method).uri(uri).body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(response.into_body().collect().await.unwrap().to_bytes(), expected); + } + + assert_eq!(collected_http_response_size_metric(missing_identity), None); + } + + #[tokio::test] + async fn http_response_egress_metric_does_not_count_subscribe() { + let database_identity = test_identity(15); + remove_http_response_size_metric(database_identity); + + let state = DummyState::new().with_database(database_identity); + let app = DatabaseRoutes:: { + subscribe_get: axum::routing::get(|| async { "subscribe response" }), + ..Default::default() + } + .into_router(state.clone()) + .with_state(state); + + let response = app + .oneshot( + Request::builder() + .uri(format!("/{database_identity}/subscribe")) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let body = response.into_body().collect().await.unwrap().to_bytes(); + + assert_eq!(body, "subscribe response"); + assert_eq!(http_response_size_metric(database_identity), 0); + + remove_http_response_size_metric(database_identity); + } + + #[tokio::test] + async fn http_response_egress_metric_counts_error_responses_for_existing_database() { + let database_identity = test_identity(16); + remove_http_response_size_metric(database_identity); + + let state = DummyState::new().with_database(database_identity); + let app = DatabaseRoutes:: { + db_get: axum::routing::get(|| async { (StatusCode::BAD_REQUEST, "bad request body") }), + ..Default::default() + } + .into_router(state.clone()) + .with_state(state); + + let response = app + .oneshot( + Request::builder() + .uri(format!("/{database_identity}")) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let body = response.into_body().collect().await.unwrap().to_bytes(); + + assert_eq!(body, "bad request body"); + assert_eq!( + http_response_size_metric(database_identity), + text_plain_header_bytes() + "bad request body".len() as u64 + ); + + remove_http_response_size_metric(database_identity); + } } diff --git a/crates/client-api/src/routes/mcp.rs b/crates/client-api/src/routes/mcp.rs index 03b2f8d9abf..c8c237fa536 100644 --- a/crates/client-api/src/routes/mcp.rs +++ b/crates/client-api/src/routes/mcp.rs @@ -14,8 +14,8 @@ use spacetimedb_lib::db::raw_def::v9::RawModuleDefV9; use spacetimedb_lib::sats; use super::database::{ - client_connected_error_to_response, client_disconnected_error_to_response, find_leader_and_database, - find_module_and_database, map_reducer_error, sql_direct, SqlParams, SqlQueryParams, + client_connected_error_to_response, client_disconnected_error_to_response, find_database_leader, + find_database_module, find_database_or_404, map_reducer_error, sql_direct, SqlQueryParams, }; use crate::auth::SpacetimeAuth; use crate::routes::subscribe::generate_random_connection_id; @@ -44,6 +44,13 @@ pub struct McpParams { } /// handle MCP JSON-RPC request for the database named in the URL +// +// Due to different name resolution and error handling behavior in different branches, +// this route handler does not use [`super::database::resolve_database_name_and_count_response_egress_middleware`]. +// This is unfortunate, as we probably would like to count egress bytes from MCP calls, +// but I (pgoldman 2026-08-07) do not have the wherewithal +// to significantly rewrite this file in order to make it compatible with the middleware, +// and do not know which of its error-handling behaviors are safe to change. pub async fn mcp( State(ctx): State, Path(McpParams { name_or_identity }): Path, @@ -316,7 +323,8 @@ async fn tool_get_schema(ctx: &S, name_or_identity: NameOrIdentity) -> axum:: where S: ControlStateDelegate + NodeDelegate, { - let (leader, _) = find_leader_and_database(ctx, name_or_identity).await?; + let database = find_database_or_404(ctx, name_or_identity).await?; + let leader = find_database_leader(ctx, &database).await?; let module = leader.wait_for_module(MODULE_WAIT_TIMEOUT).await.map_err(log_and_500)?; let raw = RawModuleDefV9::from(module.info.module_def.as_ref().clone()); let json = serde_json::to_string(&sats::serde::SerdeWrapper(raw)).map_err(log_and_500)?; @@ -335,9 +343,10 @@ where { let caller_identity = auth.claims.identity; let caller_auth: ConnectionAuthCtx = auth.into(); + let database = find_database_or_404(ctx, name_or_identity).await?; let rows = sql_direct( ctx.clone(), - SqlParams { name_or_identity }, + database, SqlQueryParams { confirmed }, caller_identity, caller_auth, @@ -360,7 +369,8 @@ where { let caller_identity = auth.claims.identity; let caller_auth: ConnectionAuthCtx = auth.into(); - let (module, _) = find_module_and_database(ctx, name_or_identity).await?; + let database = find_database_or_404(ctx, name_or_identity).await?; + let module = find_database_module(ctx, &database).await?; let connection_id = generate_random_connection_id(); module diff --git a/crates/core/src/host/host_controller.rs b/crates/core/src/host/host_controller.rs index 60ffb1ba042..9fc3e012fa0 100644 --- a/crates/core/src/host/host_controller.rs +++ b/crates/core/src/host/host_controller.rs @@ -1578,8 +1578,13 @@ pub async fn extract_schema(program_bytes: Box<[u8]>, host_type: HostType) -> an .await } -// Remove all gauges associated with a database. -// This is useful if a database is being deleted. +/// Removes metrics associated with a database. +/// +/// This is called when a database's [`ModuleHost`] exits, +/// including (but not limited to) when a database is deleted. +/// +/// Despite the historical function name, this cleans up per-database metric +/// series even when they are not literally `Gauge`s or `IntGauge`s. pub fn remove_database_gauges<'a, I>(db: &Identity, table_names: I) where I: IntoIterator, @@ -1609,4 +1614,5 @@ where V8HeapMetrics::remove_all_metric_label_values_for_database(db); let _ = WORKER_METRICS.v8_request_queue_length.remove_label_values(db); + let _ = DB_METRICS.http_response_size_bytes.remove_label_values(db); } diff --git a/crates/datastore/src/db_metrics/mod.rs b/crates/datastore/src/db_metrics/mod.rs index b8366034bac..2c7375362b0 100644 --- a/crates/datastore/src/db_metrics/mod.rs +++ b/crates/datastore/src/db_metrics/mod.rs @@ -225,6 +225,13 @@ An individual HTTP response's size in bytes is the sum of the sizes of the heade #[labels(db: Identity)] pub procedure_http_response_size_bytes: IntCounterVec, + #[name = spacetime_http_response_size_bytes_total] + #[help = "Total logical bytes sent in HTTP responses for routes scoped to a specific database. + +An individual HTTP response's size in bytes is the sum of the sizes of the header names, header values and body."] + #[labels(db: Identity)] + pub http_response_size_bytes: IntCounterVec, + #[name = spacetime_procedure_num_http_requests] #[help = "Number of HTTP requests performed by procedures running in databases. diff --git a/crates/pg/src/pg_server.rs b/crates/pg/src/pg_server.rs index b6f848e0e71..e16f13fdf66 100644 --- a/crates/pg/src/pg_server.rs +++ b/crates/pg/src/pg_server.rs @@ -25,7 +25,7 @@ use pgwire::tokio::process_socket; use spacetimedb_auth::identity::ConnectionAuthCtx; use spacetimedb_client_api::auth::validate_token; use spacetimedb_client_api::routes::database; -use spacetimedb_client_api::routes::database::{SqlParams, SqlQueryParams}; +use spacetimedb_client_api::routes::database::SqlQueryParams; use spacetimedb_client_api::{Authorization, ControlStateReadAccess, ControlStateWriteAccess, NodeDelegate}; use spacetimedb_client_api_messages::http::SqlStmtResult; use spacetimedb_client_api_messages::name::DatabaseName; @@ -154,14 +154,25 @@ where { async fn exe_sql(&self, query: String) -> PgWireResult> { let params = self.cached.lock().await.clone().unwrap(); - let db = SqlParams { - name_or_identity: database::NameOrIdentity::Name(DatabaseName(params.database.clone())), - }; + let name_or_identity = database::NameOrIdentity::Name(DatabaseName(params.database.clone())); + let database_identity = response(name_or_identity.resolve(&self.ctx).await, ¶ms.database).await?; + let database = response( + self.ctx + .get_database_by_identity(&database_identity) + .await + .map_err(|err| { + log::warn!("PG: unable to load database {database_identity}: {err:#}"); + (StatusCode::INTERNAL_SERVER_ERROR, "Internal server error").into() + }) + .and_then(|database| database.ok_or_else(|| database::NO_SUCH_DATABASE.into())), + ¶ms.database, + ) + .await?; let sql = match response( database::sql_direct( self.ctx.clone(), - db, + database, SqlQueryParams { confirmed: Some(true) }, params.caller_identity, params.caller_auth.clone(),