ENG-6268: migrate plateau to axum 0.8

Cargo.toml

@@ -23,7 +23,9 @@ license = "MIT OR Apache-2.0"
 
 [workspace.dependencies]
 
+axum = "0.8"
 chrono = "0.4"
+headers = "0.4"
 reqwest = { version = "0.11", default-features = false, features = [
   "json",
   "rustls-tls",
@@ -32,7 +34,10 @@ reqwest = { version = "0.11", default-features = false, features = [
 ] }
 test-log = { version = "0.2", default-features = false, features = ["trace"] }
 thiserror = "2.0"
+tower-http = { version = "0.7", features = ["trace"] }
 tracing = "0.1"
+utoipa = { version = "5", features = ["axum_extras", "chrono"] }
+utoipa-swagger-ui = { version = "9", features = ["axum"] }
 
 plateau-catalog = { path = "./catalog" }
 plateau-cli = { path = "./cli" }

catalog/Cargo.toml

@@ -10,7 +10,6 @@ authors.workspace = true
 
 [dependencies]
 anyhow = "1"
-axum = { version = "0.6", features = ["headers"] }
 bytes = "1.6"
 bytesize = { version = "1.1.0", features = ["serde"] }
 config = "0.14"

client/src/lib.rs

@@ -304,7 +304,7 @@ impl Client {
 
         Ok(match position.into() {
             Some(position) => base_request.json(&position),
-            None => base_request.json(&{}),
+            None => base_request.json(&TopicIterator::default()),
         })
     }
 

server/Cargo.toml

@@ -10,7 +10,6 @@ authors.workspace = true
 
 [dependencies]
 anyhow = "1"
-axum = { version = "0.6", features = ["headers"] }
 bytes = "1.6"
 bytesize = { version = "1.1.0", features = ["serde"] }
 config = "0.14"
@@ -25,13 +24,14 @@ toml = "0.7"
 tracing = "0.1"
 tokio-stream = { version = "0.1", features = ["signal"] }
 tokio = { version = "1", features = ["full"] }
-tower-http = { version = "0.4", features = ["trace"] }
-# TODO: 0.7.4 adds a deprecation warning that will need to be fixed down the road
-utoipa = { version = "4", features = ["axum_extras"] }
-utoipa-swagger-ui = { version = "4", features = ["axum"] }
 
+axum.workspace = true
 chrono.workspace = true
+headers.workspace = true
 thiserror.workspace = true
+tower-http.workspace = true
+utoipa.workspace = true
+utoipa-swagger-ui.workspace = true
 
 plateau-catalog.workspace = true
 plateau-client = { workspace = true, features = ["replicate"] }

server/src/axum_util/query.rs

@@ -12,7 +12,6 @@ pub enum QueryRejection {
     FailedToDeserializeQueryString,
 }
 
-#[axum::async_trait]
 impl<T, S> extract::FromRequestParts<S> for Query<T>
 where
     T: de::DeserializeOwned,
@@ -41,3 +40,22 @@ impl response::IntoResponse for QueryRejection {
         http::StatusCode::NOT_ACCEPTABLE.into_response()
     }
 }
+
+impl<T, S> extract::OptionalFromRequestParts<S> for Query<T>
+where
+    T: de::DeserializeOwned,
+    S: Send + Sync,
+{
+    type Rejection = QueryRejection;
+
+    async fn from_request_parts(
+        parts: &mut http::request::Parts,
+        state: &S,
+    ) -> Result<Option<Self>, Self::Rejection> {
+        Ok(
+            <Self as extract::FromRequestParts<S>>::from_request_parts(parts, state)
+                .await
+                .ok(),
+        )
+    }
+}

server/src/http.rs

@@ -5,13 +5,17 @@ use std::sync::Arc;
 use std::time::{Duration, SystemTime};
 
 use anyhow::Result;
-use axum::{
-    body::Body,
-    extract::{DefaultBodyLimit, FromRef, Path, State},
-    http::{header::ACCEPT, HeaderMap, Request},
-    routing::{get, post},
-    Json, Router, Server,
-};
+use axum::extract::DefaultBodyLimit;
+use axum::extract::FromRef;
+use axum::extract::Path;
+use axum::extract::Request;
+use axum::extract::State;
+use axum::http::header::ACCEPT;
+use axum::http::HeaderMap;
+use axum::routing::get;
+use axum::routing::post;
+use axum::Json;
+use axum::Router;
 
 use chrono::{DateTime, Utc};
 use futures::{Future, FutureExt};
@@ -136,19 +140,19 @@ pub async fn serve(
         .route("/ok", get(healthcheck))
         .route("/topics", get(get_topics))
         .route(
-            "/topic/:topic_name/partition/:partition_name/records",
+            "/topic/{topic_name}/partition/{partition_name}/records",
             get(partition_get_records),
         )
         .route(
-            "/topic/:topic_name/partition/:partition_name",
+            "/topic/{topic_name}/partition/{partition_name}",
             post(topic_append).layer(DefaultBodyLimit::max(config.http.max_append_bytes)),
         )
-        .route("/topic/:topic_name/records", post(topic_iterate_route))
-        .route("/topic/:topic_name", get(topic_get_info))
+        .route("/topic/{topic_name}/records", post(topic_iterate_route))
+        .route("/topic/{topic_name}", get(topic_get_info))
         .route("/info", get(get_info))
         .layer(
             TraceLayer::new(log_codes.into_make_classifier())
-                .make_span_with(|request: &Request<Body>| {
+                .make_span_with(|request: &Request| {
                     tracing::span!(
                         target: "plateau::http",
                         tracing::Level::INFO,
@@ -166,17 +170,20 @@ pub async fn serve(
         )
         .with_state(AppState(catalog, Arc::clone(&config)));
 
-    let server = Server::bind(&config.http.bind).serve(filter.into_make_service());
-    let addr = server.local_addr();
+    let listener = tokio::net::TcpListener::bind(&config.http.bind)
+        .await
+        .unwrap();
+    let addr = listener.local_addr().unwrap();
 
-    let fut = server.with_graceful_shutdown(FutureExt::map(rx_shutdown, |_| ()));
+    let server =
+        axum::serve(listener, filter).with_graceful_shutdown(FutureExt::map(rx_shutdown, |_| ()));
     let span = tracing::info_span!("Server::run", ?addr);
     tracing::info!(parent: &span, %addr, "listening");
 
     (
         addr,
         tx_shutdown,
-        Box::pin(async move { fut.instrument(span).await.unwrap_or(()) }),
+        Box::pin(async move { server.await.unwrap_or(()) }.instrument(span)),
     )
 }
 
@@ -228,7 +235,7 @@ async fn get_topics(
     responses(
         (status = 200, description = "Span of inserted records", body = Inserted),
     ),
-    request_body(content = SchemaChunk<crate::transport::ArrowSchema>, content_type = "application/vnd.apache.arrow.file"),
+    request_body(content = Vec<u8>, content_type = "application/vnd.apache.arrow.file"),
   )]
 async fn topic_append(
     State(AppState(catalog, _config)): State<AppState>,
@@ -316,24 +323,25 @@ async fn topic_iterate_route(
     position: Option<Json<TopicIterator>>,
 ) -> Result<axum::response::Response, ErrorReply> {
     let max_page = config.http.max_page;
-    topic_iterate(topic_name, query, headers, position, catalog, max_page).await
+    let query = query.map(|Query(query)| query).unwrap_or_default();
+    let accept = headers.get(ACCEPT).and_then(|header| header.to_str().ok());
+    let position = position.map(|Json(value)| value).unwrap_or_default();
+    topic_iterate(topic_name, query, accept, position, catalog, max_page).await
 }
 
+// Takes plain domain types rather than axum extractors so this can be called
+// directly (e.g. from fitzroy) without depending on plateau-server's axum
+// version.
 pub async fn topic_iterate(
     topic_name: String,
-    query: Option<Query<TopicIterationQuery>>,
-    headers: HeaderMap,
-    position: Option<Json<TopicIterator>>,
+    query: TopicIterationQuery,
+    accept: Option<&str>,
+    position: TopicIterator,
     catalog: Arc<Catalog>,
     max_page: RowLimit,
 ) -> Result<axum::response::Response, ErrorReply> {
-    let query = query.map(|Query(query)| query).unwrap_or_default();
-    let content = headers.get(ACCEPT).and_then(|header| header.to_str().ok());
-    let position = position.map(|Json(value)| value);
-
     let topic = catalog.get_topic(&topic_name).await;
     let page_size = RowLimit::records(query.page_size.unwrap_or(1000)).min(max_page);
-    let position = position.unwrap_or_default();
     let partition_filter = query.partition_filter;
     let order: Ordering = query.order.unwrap_or(TopicIterationOrder::Asc).into();
 
@@ -364,7 +372,7 @@ pub async fn topic_iterate(
         );
     }
 
-    chunk::to_reply(content, result.batch, query.data_focus)
+    chunk::to_reply(accept, result.batch, query.data_focus)
 }
 
 #[utoipa::path(
@@ -599,7 +607,6 @@ async fn get_info(
             Inserted,
             Partitions,
             // PartitionFilter,
-            crate::transport::ArrowSchemaChunk,
             Span,
             Topic,
             Topics,

server/src/http/chunk.rs

@@ -1,15 +1,14 @@
-use axum::{
-    async_trait,
-    body::{boxed, Full, HttpBody},
-    extract::{
-        rejection::{BytesRejection, FailedToBufferBody},
-        FromRef, FromRequest,
-    },
-    headers::ContentType,
-    http::{header::CONTENT_TYPE, Request, StatusCode},
-    response::Response,
-    BoxError, RequestExt as _,
-};
+use axum::body::Body;
+use axum::extract::rejection::BytesRejection;
+use axum::extract::rejection::FailedToBufferBody;
+use axum::extract::FromRef;
+use axum::extract::FromRequest;
+use axum::extract::Request;
+use axum::http::header::CONTENT_TYPE;
+use axum::http::StatusCode;
+use axum::response::Response;
+use axum::RequestExt as _;
+use headers::ContentType;
 
 use bytes::Bytes;
 use std::io::{Cursor, Write};
@@ -35,18 +34,14 @@ const CONTENT_TYPE_PANDAS_RECORD: &str = "application/json; format=pandas-record
 
 pub(crate) struct SchemaChunkRequest(pub(crate) SchemaChunk<Schema>);
 
-#[async_trait]
-impl<S, B> FromRequest<S, B> for SchemaChunkRequest
+impl<S> FromRequest<S> for SchemaChunkRequest
 where
-    B: HttpBody + Send + 'static,
-    B::Data: Send,
-    B::Error: Into<BoxError>,
     Config: FromRef<S>,
     S: Send + Sync,
 {
     type Rejection = ErrorReply;
 
-    async fn from_request(req: Request<B>, state: &S) -> Result<Self, Self::Rejection> {
+    async fn from_request(req: Request, state: &S) -> Result<Self, Self::Rejection> {
         let config = Config::from_ref(state);
         let max_append_bytes = config.http.max_append_bytes;
 
@@ -59,20 +54,20 @@ where
             ))?;
 
         if content_type == CONTENT_TYPE_ARROW {
-            let bytes = match req.with_limited_body() {
-                Ok(req) => req.extract::<Bytes, _>(),
-                Err(req) => req.extract::<Bytes, _>(),
-            }
-            .await
-            .map_err(|e| {
-                if let BytesRejection::FailedToBufferBody(FailedToBufferBody::LengthLimitError(_)) =
-                    e
-                {
-                    return ErrorReply::PayloadTooLarge(max_append_bytes);
-                }
-
-                ErrorReply::Arrow(ArrowError::from_external_error(Box::new(e)))
-            })?;
+            let bytes = req
+                .with_limited_body()
+                .extract::<Bytes, _>()
+                .await
+                .map_err(|e| {
+                    if let BytesRejection::FailedToBufferBody(
+                        FailedToBufferBody::LengthLimitError(_),
+                    ) = e
+                    {
+                        return ErrorReply::PayloadTooLarge(max_append_bytes);
+                    }
+
+                    ErrorReply::Arrow(ArrowError::from_external_error(Box::new(e)))
+                })?;
 
             deserialize_request(bytes).await
         } else {
@@ -144,14 +139,14 @@ pub(crate) fn to_reply(
                 Response::builder()
                     .header("Content-Type", CONTENT_TYPE_ARROW)
                     .status(StatusCode::OK)
-                    .body(boxed(Full::new(Bytes::from(bytes))))
+                    .body(Body::from(bytes))
                     .map_err(|_| ErrorReply::Unknown)
             }
             None | Some("*/*") | Some(CONTENT_TYPE_JSON) | Some(CONTENT_TYPE_PANDAS_RECORD) => {
                 Response::builder()
                     .header("Content-Type", CONTENT_TYPE_PANDAS_RECORD)
                     .status(StatusCode::OK)
-                    .body(boxed(Full::new(Bytes::from("[]"))))
+                    .body(Body::from("[]"))
                     .map_err(|_| ErrorReply::Unknown)
             }
             Some(other) => Err(ErrorReply::CannotEmit(other.to_string())),
@@ -196,7 +191,7 @@ pub(crate) fn to_reply(
                         .get("status")
                         .unwrap_or(&"{}".to_string()),
                 )
-                .body(boxed(Full::new(Bytes::from(bytes))))
+                .body(Body::from(bytes))
                 .map_err(|_| ErrorReply::Unknown)
         }
         None | Some("*/*") | Some(CONTENT_TYPE_JSON) | Some(CONTENT_TYPE_PANDAS_RECORD) => {
@@ -233,7 +228,7 @@ pub(crate) fn to_reply(
                         .unwrap_or(&"{}".to_string()),
                 )
                 .status(StatusCode::OK)
-                .body(boxed(Full::new(Bytes::from(bytes))))
+                .body(Body::from(bytes))
                 .map_err(|_| ErrorReply::Unknown)
         }
         Some(other) => Err(ErrorReply::CannotEmit(other.to_string())),

transport/Cargo.toml

@@ -29,10 +29,10 @@ arrow-ipc = "55.2.0"
 strum = { version = "0.26", features = ["derive"] }
 thiserror = "1"
 regex = "1.10"
-utoipa = { version = "4", features = ["axum_extras"] }
 chrono = { version = "0.4", features = ["serde"] }
 
 tracing.workspace = true
+utoipa.workspace = true
 
 [features]
 rweb = ["dep:rweb"]

transport/src/lib.rs

@@ -264,6 +264,7 @@ pub struct RecordQuery {
     pub data_focus: DataFocus,
     #[serde(default)]
     #[cfg_attr(feature = "clap", arg(skip))]
+    #[param(value_type = Vec<String>)]
     pub partition_filter: PartitionFilter,
 }
 
@@ -327,6 +328,7 @@ pub struct TopicIterationQuery {
     pub data_focus: DataFocus,
     #[serde(default)]
     #[cfg_attr(feature = "clap", clap(skip))]
+    #[param(value_type = Vec<String>)]
     pub partition_filter: PartitionFilter,
 }
 
@@ -336,7 +338,7 @@ pub struct TopicIterationQuery {
 /// begin with "regex:" to signify that any partition matching the following string can be converted.
 pub type PartitionFilter = Option<Vec<PartitionSelector>>;
 
-#[derive(Clone, Debug, serde::Serialize, serde::Deserialize, utoipa::ToSchema)]
+#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
 #[serde(from = "String", into = "String")]
 pub enum PartitionSelector {
     /// The exact name of a partition.
@@ -533,8 +535,7 @@ pub const CONTENT_TYPE_JSON: &str = "application/json";
 pub type SegmentChunk = RecordBatch;
 
 /// A [SegmentChunk] packaged with its associated [ArrowSchema].
-#[derive(Debug, Clone, PartialEq, ToSchema)]
-#[aliases(ArrowSchemaChunk = SchemaChunk<ArrowSchema>)]
+#[derive(Debug, Clone, PartialEq)]
 pub struct SchemaChunk<S: Borrow<ArrowSchema> + Clone + PartialEq> {
     pub schema: S,
     pub chunk: SegmentChunk,