Move ConductorState->AppState conversion to an extractor (#42840)

Introduces a new extractor `MtState`, which acts like the `State` extractor but is overridable. Also, use it in the actions middleware using `middleware::from_fn_with_state` rather than attaching the state using an extension. This is more typesafe than using extensions.

GitOrigin-RevId: bee2c9bee725774ca9c8b3765474c13cdf0fe423

Author: goffrie

crates/common/src/http/extract.rs

@@ -1,7 +1,11 @@
-use std::time::Instant;
+use std::{
+    future,
+    time::Instant,
+};
 
 use axum::{
     extract::{
+        FromRef,
         FromRequest,
         FromRequestParts,
         Request,
@@ -15,9 +19,10 @@ use axum::{
 use bytes::Bytes;
 use errors::ErrorMetadata;
 use fastrace::{
-    future::FutureExt,
+    future::FutureExt as _,
     Span,
 };
+use futures::TryFutureExt as _;
 use http::HeaderMap;
 use serde::{
     de::DeserializeOwned,
@@ -145,3 +150,40 @@ where
         axum::Json(self.0).into_response()
     }
 }
+
+/// Like `axum::extract::State`, but customizable
+pub struct MtState<T>(pub T);
+
+impl<S, T> FromRequestParts<S> for MtState<T>
+where
+    T: FromMtState<S>,
+    S: Send + Sync,
+{
+    type Rejection = HttpResponseError;
+
+    fn from_request_parts(
+        parts: &mut Parts,
+        state: &S,
+    ) -> impl Future<Output = Result<Self, Self::Rejection>> + Send {
+        T::from_request_parts(parts, state).map_ok(MtState)
+    }
+}
+
+pub trait FromMtState<Outer>: Sized {
+    fn from_request_parts(
+        parts: &mut Parts,
+        state: &Outer,
+    ) -> impl Future<Output = Result<Self, HttpResponseError>> + Send;
+}
+
+impl<Outer, T: FromRef<Outer>> FromMtState<Outer> for T
+where
+    T: Send + Sync,
+{
+    fn from_request_parts(
+        _parts: &mut Parts,
+        state: &Outer,
+    ) -> impl Future<Output = Result<Self, HttpResponseError>> + Send {
+        future::ready(Ok(T::from_ref(state)))
+    }
+}

crates/local_backend/src/app_metrics.rs

@@ -1,7 +1,4 @@
-use axum::{
-    extract::State,
-    response::IntoResponse,
-};
+use axum::response::IntoResponse;
 use common::{
     components::{
         ComponentFunctionPath,
@@ -10,6 +7,7 @@ use common::{
     http::{
         extract::{
             Json,
+            MtState,
             Query,
         },
         HttpResponseError,
@@ -39,7 +37,7 @@ pub(crate) struct UdfRateQueryArgs {
 }
 
 pub(crate) async fn udf_rate(
-    State(st): State<LocalAppState>,
+    MtState(st): MtState<LocalAppState>,
     ExtractIdentity(identity): ExtractIdentity,
     Query(UdfRateQueryArgs {
         component_path,
@@ -69,7 +67,7 @@ pub(crate) struct TopKQueryArgs {
 }
 
 pub(crate) async fn failure_percentage_top_k(
-    State(st): State<LocalAppState>,
+    MtState(st): MtState<LocalAppState>,
     ExtractIdentity(identity): ExtractIdentity,
     Query(TopKQueryArgs { window, k }): Query<TopKQueryArgs>,
 ) -> Result<impl IntoResponse, HttpResponseError> {
@@ -85,7 +83,7 @@ pub(crate) async fn failure_percentage_top_k(
 }
 
 pub(crate) async fn cache_hit_percentage_top_k(
-    State(st): State<LocalAppState>,
+    MtState(st): MtState<LocalAppState>,
     ExtractIdentity(identity): ExtractIdentity,
     Query(TopKQueryArgs { window, k }): Query<TopKQueryArgs>,
 ) -> Result<impl IntoResponse, HttpResponseError> {
@@ -110,7 +108,7 @@ pub(crate) struct CacheHitPercentageQueryArgs {
     udf_type: Option<String>,
 }
 pub(crate) async fn cache_hit_percentage(
-    State(st): State<LocalAppState>,
+    MtState(st): MtState<LocalAppState>,
     ExtractIdentity(identity): ExtractIdentity,
     Query(query_args): Query<CacheHitPercentageQueryArgs>,
 ) -> Result<impl IntoResponse, HttpResponseError> {
@@ -140,7 +138,7 @@ pub(crate) struct LatencyPercentilesQueryArgs {
     udf_type: Option<String>,
 }
 pub(crate) async fn latency_percentiles(
-    State(st): State<LocalAppState>,
+    MtState(st): MtState<LocalAppState>,
     ExtractIdentity(identity): ExtractIdentity,
     Query(query_args): Query<LatencyPercentilesQueryArgs>,
 ) -> Result<impl IntoResponse, HttpResponseError> {
@@ -170,7 +168,7 @@ pub(crate) struct TableRateQueryArgs {
     window: String,
 }
 pub(crate) async fn table_rate(
-    State(st): State<LocalAppState>,
+    MtState(st): MtState<LocalAppState>,
     ExtractIdentity(identity): ExtractIdentity,
     Query(query_args): Query<TableRateQueryArgs>,
 ) -> Result<impl IntoResponse, HttpResponseError> {
@@ -224,7 +222,7 @@ pub(crate) struct ScheduledJobLagArgs {
     window: String,
 }
 pub(crate) async fn scheduled_job_lag(
-    State(st): State<LocalAppState>,
+    MtState(st): MtState<LocalAppState>,
     ExtractIdentity(identity): ExtractIdentity,
     Query(query_args): Query<ScheduledJobLagArgs>,
 ) -> Result<impl IntoResponse, HttpResponseError> {

crates/local_backend/src/authentication.rs

@@ -7,15 +7,15 @@ use anyhow::{
 };
 use authentication::extract_bearer_token;
 use axum::{
-    extract::{
-        FromRef,
-        FromRequestParts,
-    },
+    extract::FromRequestParts,
     RequestPartsExt,
 };
 use common::{
     http::{
-        extract::Query,
+        extract::{
+            FromMtState,
+            Query,
+        },
         ExtractRequestId,
         ExtractResolvedHostname,
         HttpResponseError,
@@ -105,7 +105,7 @@ pub struct ExtractIdentity(pub Identity);
 
 impl<S> FromRequestParts<S> for ExtractIdentity
 where
-    LocalAppState: FromRef<S>,
+    LocalAppState: FromMtState<S>,
     S: Send + Sync + Clone + 'static,
 {
     type Rejection = HttpResponseError;
@@ -116,7 +116,7 @@ where
     ) -> Result<Self, Self::Rejection> {
         let token: AuthenticationToken =
             parts.extract::<ExtractAuthenticationToken>().await?.into();
-        let st = LocalAppState::from_ref(st);
+        let st = LocalAppState::from_request_parts(parts, st).await?;
 
         Ok(Self(
             st.application

crates/local_backend/src/canonical_urls.rs

@@ -6,7 +6,10 @@ use axum::{
     response::IntoResponse,
 };
 use common::http::{
-    extract::Json,
+    extract::{
+        Json,
+        MtState,
+    },
     HttpResponseError,
     RequestDestination,
 };
@@ -111,7 +114,7 @@ pub struct GetCanonicalUrlsResponse {
     ),
 )]
 pub async fn get_canonical_urls(
-    State(st): State<LocalAppState>,
+    MtState(st): MtState<LocalAppState>,
     ExtractIdentity(identity): ExtractIdentity,
 ) -> Result<impl IntoResponse, HttpResponseError> {
     must_be_admin(&identity)?;

crates/local_backend/src/dashboard.rs

@@ -5,17 +5,16 @@ use application::{
 };
 use axum::{
     debug_handler,
-    extract::{
-        FromRef,
-        State,
-    },
+    extract::State,
     response::IntoResponse,
 };
 use common::{
     components::ComponentId,
     http::{
         extract::{
+            FromMtState,
             Json,
+            MtState,
             Query,
         },
         ExtractClientVersion,
@@ -94,9 +93,8 @@ pub struct ShapesArgs {
     ),
     responses((status = 200, body = serde_json::Value)),
 )]
-#[debug_handler]
 pub async fn shapes2(
-    State(st): State<LocalAppState>,
+    MtState(st): MtState<LocalAppState>,
     ExtractIdentity(identity): ExtractIdentity,
     Query(ShapesArgs { component }): Query<ShapesArgs>,
 ) -> Result<impl IntoResponse, HttpResponseError> {
@@ -136,9 +134,8 @@ pub async fn shapes2(
     request_body = DeleteTableArgs,
     responses((status = 200)),
 )]
-#[debug_handler]
 pub async fn delete_tables(
-    State(st): State<LocalAppState>,
+    MtState(st): MtState<LocalAppState>,
     ExtractIdentity(identity): ExtractIdentity,
     Json(DeleteTableArgs {
         table_names,
@@ -168,9 +165,8 @@ pub async fn delete_tables(
     request_body = DeleteComponentArgs,
     responses((status = 200)),
 )]
-#[debug_handler]
 pub async fn delete_component(
-    State(st): State<LocalAppState>,
+    MtState(st): MtState<LocalAppState>,
     ExtractIdentity(identity): ExtractIdentity,
     Json(DeleteComponentArgs { component_id }): Json<DeleteComponentArgs>,
 ) -> Result<impl IntoResponse, HttpResponseError> {
@@ -206,9 +202,8 @@ struct GetIndexesResponse {
     ),
     responses((status = 200, body = GetIndexesResponse)),
 )]
-#[debug_handler]
 pub async fn get_indexes(
-    State(st): State<LocalAppState>,
+    MtState(st): MtState<LocalAppState>,
     ExtractIdentity(identity): ExtractIdentity,
     Query(GetIndexesArgs { component_id }): Query<GetIndexesArgs>,
 ) -> Result<impl IntoResponse, HttpResponseError> {
@@ -245,9 +240,8 @@ pub struct GetSourceCodeArgs {
     ),
     responses((status = 200, body = String)),
 )]
-#[debug_handler]
 pub async fn get_source_code(
-    State(st): State<LocalAppState>,
+    MtState(st): MtState<LocalAppState>,
     ExtractIdentity(identity): ExtractIdentity,
     Query(GetSourceCodeArgs { path, component }): Query<GetSourceCodeArgs>,
 ) -> Result<impl IntoResponse, HttpResponseError> {
@@ -351,7 +345,7 @@ pub fn local_only_dashboard_router() -> OpenApiRouter<crate::LocalAppState> {
 // Routes with the same handlers for the local backend + closed source backend
 pub fn common_dashboard_api_router<S>() -> OpenApiRouter<S>
 where
-    LocalAppState: FromRef<S>,
+    LocalAppState: FromMtState<S>,
     S: Clone + Send + Sync + 'static,
 {
     OpenApiRouter::new()

crates/local_backend/src/deploy_config.rs

@@ -16,7 +16,10 @@ use axum::{
 use common::{
     components::ComponentId,
     http::{
-        extract::Json,
+        extract::{
+            Json,
+            MtState,
+        },
         HttpResponseError,
     },
     version::Version,
@@ -182,9 +185,8 @@ pub struct ModuleHashJson {
     environment: Option<String>,
 }
 
-#[debug_handler]
 pub async fn get_config(
-    State(st): State<LocalAppState>,
+    MtState(st): MtState<LocalAppState>,
     Json(req): Json<GetConfigRequest>,
 ) -> Result<impl IntoResponse, HttpResponseError> {
     let identity = must_be_admin_from_key(
@@ -212,9 +214,8 @@ pub async fn get_config(
     }))
 }
 
-#[debug_handler]
 pub async fn get_config_hashes(
-    State(st): State<LocalAppState>,
+    MtState(st): MtState<LocalAppState>,
     Json(req): Json<GetConfigRequest>,
 ) -> Result<impl IntoResponse, HttpResponseError> {
     let identity = must_be_admin_from_key(

crates/local_backend/src/deploy_config2.rs

@@ -22,7 +22,10 @@ use common::{
     },
     bootstrap_model::components::definition::SerializedComponentDefinitionMetadata,
     http::{
-        extract::Json,
+        extract::{
+            Json,
+            MtState,
+        },
         HttpResponseError,
     },
 };
@@ -226,9 +229,8 @@ pub async fn start_push(
 // it won’t start schema validation/index backfill). It can be used to determine
 // what will be the effects of a large push without starting work that can take
 // a long time on large instances.
-#[debug_handler]
 pub async fn evaluate_push(
-    State(st): State<LocalAppState>,
+    MtState(st): MtState<LocalAppState>,
     Json(req): Json<StartPushRequest>,
 ) -> Result<impl IntoResponse, HttpResponseError> {
     let _identity = must_be_admin_from_key_with_write_access(
@@ -258,9 +260,8 @@ pub struct WaitForSchemaRequest {
     timeout_ms: Option<u32>,
 }
 
-#[debug_handler]
 pub async fn wait_for_schema(
-    State(st): State<LocalAppState>,
+    MtState(st): MtState<LocalAppState>,
     Json(req): Json<WaitForSchemaRequest>,
 ) -> Result<impl IntoResponse, HttpResponseError> {
     let identity = must_be_admin_from_key(
@@ -289,9 +290,8 @@ pub struct FinishPushRequest {
     dry_run: bool,
 }
 
-#[debug_handler]
 pub async fn finish_push(
-    State(st): State<LocalAppState>,
+    MtState(st): MtState<LocalAppState>,
     Json(req): Json<FinishPushRequest>,
 ) -> Result<impl IntoResponse, HttpResponseError> {
     let identity = must_be_admin_from_key_with_write_access(

crates/local_backend/src/environment_variables.rs

@@ -7,7 +7,10 @@ use axum::{
     response::IntoResponse,
 };
 use common::http::{
-    extract::Json,
+    extract::{
+        Json,
+        MtState,
+    },
     HttpResponseError,
 };
 use http::StatusCode;
@@ -134,7 +137,7 @@ pub struct ListEnvVarsResponse {
     ),
 )]
 pub async fn list_environment_variables(
-    State(st): State<LocalAppState>,
+    MtState(st): MtState<LocalAppState>,
     ExtractIdentity(identity): ExtractIdentity,
 ) -> Result<impl IntoResponse, HttpResponseError> {
     must_be_admin(&identity)?;