diff --git a/crates/client-api/src/routes/database.rs b/crates/client-api/src/routes/database.rs index 0a52272153e..a690458d0b0 100644 --- a/crates/client-api/src/routes/database.rs +++ b/crates/client-api/src/routes/database.rs @@ -90,7 +90,6 @@ fn allow_creation(auth: &SpacetimeAuth) -> Result<(), ErrorResponse> { } #[derive(Deserialize)] pub struct CallParams { - name_or_identity: NameOrIdentity, reducer: String, } @@ -141,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 { @@ -156,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 @@ -215,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. @@ -258,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, @@ -266,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 { @@ -461,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)] @@ -502,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, @@ -521,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.). @@ -561,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, @@ -581,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, @@ -610,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 @@ -620,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) @@ -687,11 +654,6 @@ pub(crate) async fn worker_ctx_find_database( .map_err(log_and_500) } -#[derive(Deserialize)] -pub struct SqlParams { - pub name_or_identity: NameOrIdentity, -} - #[derive(Deserialize)] pub struct SqlQueryParams { /// If `true`, return the query result only after its transaction offset @@ -742,7 +704,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, @@ -751,7 +713,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| { @@ -773,7 +735,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, @@ -783,7 +745,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); @@ -798,11 +760,6 @@ pub struct DNSParams { name_or_identity: NameOrIdentity, } -#[derive(Deserialize)] -pub struct ReverseDNSParams { - name_or_identity: NameOrIdentity, -} - #[derive(Deserialize)] pub struct DNSQueryParams {} @@ -817,9 +774,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) @@ -847,6 +804,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, @@ -855,10 +813,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?; @@ -965,12 +920,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(), @@ -979,6 +930,7 @@ pub async fn publish( return self::reset( State(ctx), + Extension(ResolvedDatabase(database)), Path(ResetDatabaseParams { name_or_identity: name_or_identity.clone(), }), @@ -1228,11 +1180,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)] @@ -1243,13 +1190,13 @@ 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> { // 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?; + let database_identity = resolve_and_authenticate(&ctx, database, &auth).await?; let style = match style { PrettyPrintStyle::NoColor => AutoMigratePrettyPrintStyle::NoColor, PrettyPrintStyle::AnsiColor => AutoMigratePrettyPrintStyle::AnsiColor, @@ -1311,17 +1258,13 @@ pub async fn pre_publish .map(axum::Json) } -/// Resolves the [`NameOrIdentity`] to a database identity and checks if the -/// `auth` identity owns the database. +/// Checks if the `auth` identity owns the middleware-resolved database. async fn resolve_and_authenticate( ctx: &S, - name_or_identity: &NameOrIdentity, + database: Database, 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)?; + let database_identity = database.database_identity; ctx.authorize_action(auth.claims.identity, database.database_identity, Action::UpdateDatabase) .await?; @@ -1364,13 +1307,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?; @@ -1384,13 +1324,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?; @@ -1402,19 +1339,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) @@ -1432,14 +1364,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 { @@ -1450,18 +1377,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 @@ -1509,30 +1425,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()) } @@ -1607,7 +1508,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::), @@ -1622,17 +1523,16 @@ where S: NodeDelegate + ControlStateDelegate + Authorization + Clone + 'static, { pub fn into_router(self, ctx: S) -> axum::Router { - let egress_metrics_middleware = - axum::middleware::from_fn_with_state(ctx.clone(), http_response_egress_metrics_middleware::); + 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_put) .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("/call/:reducer", self.call_reducer_procedure_post) .route("/schema", self.schema_get) .route("/logs", self.logs_get) @@ -1643,12 +1543,24 @@ where .route("/reset", self.db_reset) .route("/lock", self.lock_post) .route("/unlock", self.unlock_post) - .route_layer(egress_metrics_middleware.clone()); + .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); - // Add the subscribe route after `egress_metrics_middleware` + // 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 = counted_db_router.route("/subscribe", self.subscribe_get); + let db_router = db_router.route("/subscribe", self.subscribe_get); let authed_root_router = axum::Router::new().route( "/", @@ -1671,7 +1583,7 @@ where .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_layer(egress_metrics_middleware); + .route_layer(resolving_egress_metrics_middleware); axum::Router::new() .merge(authed_root_router) @@ -1680,36 +1592,31 @@ where } } -/// Middleware which counts response bytes in the metric `spacetime_http_response_size_bytes_total`. +/// Resolves an existing database, attaches it as [`ResolvedDatabase`], +/// and counts response bytes in the metric `spacetime_http_response_size_bytes_total`. /// -/// This middleware is intended to be supplied to all HTTP routes which apply to a particular database, -/// *except* the WebSocket connection route (confusingly named `subscribe`), whose egress is measured separately. -async fn http_response_egress_metrics_middleware( +/// 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, - request: Request, + mut request: Request, next: axum::middleware::Next, ) -> axum::response::Response where S: ControlStateDelegate + Clone + Send + Sync + 'static, { - let Ok(Ok(database_identity)) = name_or_identity.try_resolve(&worker_ctx).await else { - // The provided name doesn't map to an `Identity`. - // Run the route unchanged (as opposed to returning an error from the middleware) - // to preserve the error-handling behavior of the route. - return next.run(request).await; + let database_identity = match name_or_identity.resolve(&worker_ctx).await { + Ok(database_identity) => database_identity, + Err(response) => return Err::<(), ErrorResponse>(response).into_response(), }; - - if !matches!( - worker_ctx_find_database(&worker_ctx, &database_identity).await, - Ok(Some(_)) - ) { - // We have what appears to be an `Identity`, but it doesn't name any database. - // Don't create a metrics label for it, - // and run the route unchanged (as opposed to returning an error from the middleware) - // to preserve the error-handling behavior of the route. - return next.run(request).await; - } + let database = match worker_ctx_find_database(&worker_ctx, &database_identity).await { + Ok(Some(database)) => database, + Ok(None) => return NO_SUCH_DATABASE.into_response(), + Err(response) => return Err::<(), ErrorResponse>(response).into_response(), + }; + request.extensions_mut().insert(ResolvedDatabase(database.clone())); let response = next.run(request).await; let (parts, body) = response.into_parts(); @@ -1727,7 +1634,7 @@ where let counter = DB_METRICS .http_response_size_bytes - .with_label_values(&database_identity); + .with_label_values(&database.database_identity); counter.inc_by(header_bytes as u64); // `/logs?follow=true` can stream indefinitely. @@ -1767,6 +1674,7 @@ mod tests { use spacetimedb_paths::FromPathUnchecked; use spacetimedb_schema::auto_migrate::{MigrationPolicy, PrettyPrintStyle}; use std::collections::HashMap; + use std::sync::atomic::{AtomicUsize, Ordering}; use tower::util::ServiceExt; #[derive(Clone, Default)] @@ -1819,6 +1727,7 @@ mod tests { module_logs_dir: ModuleLogsDir, databases: std::sync::Arc>, dns: std::sync::Arc>, + dns_lookups: std::sync::Arc, } impl DummyState { @@ -1831,6 +1740,7 @@ mod tests { 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)), } } @@ -1847,6 +1757,10 @@ mod tests { 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 { @@ -1988,6 +1902,7 @@ mod tests { 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> { @@ -2113,9 +2028,8 @@ 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"); } @@ -2191,7 +2105,7 @@ mod tests { } #[tokio::test] - async fn http_response_egress_metric_skips_unresolved_and_nonexistent_databases() { + 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); @@ -2205,25 +2119,29 @@ mod tests { .into_router(state.clone()) .with_state(state); - let _ = app + let response = app .clone() .oneshot(Request::builder().uri("/unresolved-name").body(Body::empty()).unwrap()) .await - .unwrap() - .into_body() - .collect() - .await .unwrap(); - let _ = app + 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() - .into_body() - .collect() - .await .unwrap(); - let _ = app + 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}")) @@ -2231,11 +2149,12 @@ mod tests { .unwrap(), ) .await - .unwrap() - .into_body() - .collect() - .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); @@ -2244,6 +2163,95 @@ mod tests { remove_http_response_size_metric(arbitrary_identity); } + #[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); diff --git a/crates/client-api/src/routes/mcp.rs b/crates/client-api/src/routes/mcp.rs index 9f60644735b..fab0aa610af 100644 --- a/crates/client-api/src/routes/mcp.rs +++ b/crates/client-api/src/routes/mcp.rs @@ -1,23 +1,22 @@ use std::time::Duration; -use axum::extract::{Path, State}; +use axum::extract::State; use axum::response::{ErrorResponse, IntoResponse, Response}; use axum::{Extension, Json}; use http::StatusCode; -use serde::Deserialize; use serde_json::{json, Value}; use spacetimedb::auth::identity::ConnectionAuthCtx; use spacetimedb::host::{FunctionArgs, ReducerOutcome}; +use spacetimedb::messages::control_db::Database; 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, map_reducer_error, sql_direct, ResolvedDatabase, SqlQueryParams, }; use crate::auth::SpacetimeAuth; use crate::routes::subscribe::generate_random_connection_id; -use crate::util::NameOrIdentity; use crate::{log_and_500, Authorization, ControlStateDelegate, NodeDelegate}; const PROTOCOL_VERSION: &str = "2025-06-18"; @@ -36,15 +35,10 @@ const MAX_ERROR_BODY_BYTES: usize = 64 * 1024; type RpcError = (i64, String); -#[derive(Deserialize)] -pub struct McpParams { - name_or_identity: NameOrIdentity, -} - /// handle MCP JSON-RPC request pub async fn mcp( State(ctx): State, - Path(McpParams { name_or_identity }): Path, + Extension(ResolvedDatabase(database)): Extension, Extension(auth): Extension, Json(request): Json, ) -> axum::response::Result @@ -65,7 +59,7 @@ where // protocol ping, distinct from the ping tool "ping" => jsonrpc_result(&id, json!({})), "tools/list" => jsonrpc_result(&id, tools_list()), - "tools/call" => match tools_call(&ctx, name_or_identity, auth, request.get("params")).await { + "tools/call" => match tools_call(&ctx, database, auth, request.get("params")).await { Ok(result) => jsonrpc_result(&id, result), Err((code, message)) => jsonrpc_error(&id, code, message), }, @@ -143,7 +137,7 @@ fn tools_list() -> Value { async fn tools_call( ctx: &S, - name_or_identity: NameOrIdentity, + database: Database, auth: SpacetimeAuth, params: Option<&Value>, ) -> Result @@ -163,20 +157,20 @@ where Some(message) => format!("pong: {message}"), None => "pong".to_owned(), }), - "get_schema" => tool_get_schema(ctx, name_or_identity).await, + "get_schema" => tool_get_schema(ctx, &database).await, "sql" => { let Some(sql) = arguments.and_then(|a| a.get("sql")).and_then(Value::as_str) else { return Err((INVALID_PARAMS, "sql argument must be a string".to_owned())); }; let confirmed = arguments.and_then(|a| a.get("confirmed")).and_then(Value::as_bool); - tool_sql(ctx, name_or_identity, auth, sql.to_owned(), confirmed).await + tool_sql(ctx, database, auth, sql.to_owned(), confirmed).await } "call" => { let Some(reducer) = arguments.and_then(|a| a.get("reducer")).and_then(Value::as_str) else { return Err((INVALID_PARAMS, "reducer argument must be a string".to_owned())); }; let args_json = reducer_args_json(arguments)?; - tool_call_reducer(ctx, name_or_identity, auth, reducer.to_owned(), args_json).await + tool_call_reducer(ctx, &database, auth, reducer.to_owned(), args_json).await } other => return Err((INVALID_PARAMS, format!("unknown tool: {other}"))), }; @@ -207,11 +201,11 @@ async fn execution_error_to_tool_result(err: ErrorResponse) -> Value { json!({ "content": [ { "type": "text", "text": text } ], "isError": true }) } -async fn tool_get_schema(ctx: &S, name_or_identity: NameOrIdentity) -> axum::response::Result +async fn tool_get_schema(ctx: &S, database: &Database) -> axum::response::Result where S: ControlStateDelegate + NodeDelegate, { - let (leader, _) = find_leader_and_database(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)?; @@ -220,7 +214,7 @@ where async fn tool_sql( ctx: &S, - name_or_identity: NameOrIdentity, + database: Database, auth: SpacetimeAuth, sql: String, confirmed: Option, @@ -232,7 +226,7 @@ where let caller_auth: ConnectionAuthCtx = auth.into(); let rows = sql_direct( ctx.clone(), - SqlParams { name_or_identity }, + database, SqlQueryParams { confirmed }, caller_identity, caller_auth, @@ -245,7 +239,7 @@ where async fn tool_call_reducer( ctx: &S, - name_or_identity: NameOrIdentity, + database: &Database, auth: SpacetimeAuth, reducer: String, args_json: String, @@ -255,7 +249,7 @@ 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 module = find_database_module(ctx, database).await?; let connection_id = generate_random_connection_id(); module diff --git a/crates/pg/src/pg_server.rs b/crates/pg/src/pg_server.rs index 15170a35866..1b3109d506b 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(),