apalis_board_api/framework/
axum.rs

1use apalis_board_types::ApiError;
2use apalis_core::{
3    backend::{
4        Backend, BackendExt, FetchById, Filter, ListAllTasks, ListQueues, ListTasks, ListWorkers,
5        Metrics, QueueInfo, RunningWorker, Statistic, TaskSink, codec::Codec,
6    },
7    task::Task,
8};
9use axum::{
10    Extension, Json, Router,
11    extract::{Path, Query, rejection::JsonRejection},
12    http::StatusCode,
13    response::{IntoResponse, Response},
14    routing::{get, put},
15};
16
17use serde::{Serialize, de::DeserializeOwned};
18use std::{str::FromStr, sync::Arc};
19use tokio::sync::RwLock;
20
21use crate::framework::{ApiBuilder, RegisterRoute};
22
23/// An enumeration of possible application errors.
24#[derive(Debug, thiserror::Error)]
25pub enum AppError {
26    /// The request body contained invalid JSON
27    #[error("JSON Rejection: {0}")]
28    JsonRejection(JsonRejection),
29
30    /// An error occurred in the API
31    #[error("API Error: {0}")]
32    ApiError(ApiError),
33    /// Resource not found
34    #[error("Resource not found")]
35    NotFound,
36
37    /// Missing application state
38    #[error("Missing application state")]
39    MissingState,
40}
41
42impl IntoResponse for AppError {
43    fn into_response(self) -> Response {
44        match self {
45            Self::JsonRejection(rejection) => {
46                // This error is caused by bad user input so don't log it
47                (rejection.status(), rejection.body_text()).into_response()
48            }
49            Self::ApiError(err) => {
50                // These errors are unexpected and should be logged
51                (StatusCode::INTERNAL_SERVER_ERROR, Json(err).into_response()).into_response()
52            }
53            Self::NotFound => (StatusCode::NOT_FOUND, ()).into_response(),
54            Self::MissingState => (
55                StatusCode::INTERNAL_SERVER_ERROR,
56                "Missing application state",
57            )
58                .into_response(),
59        }
60    }
61}
62
63/// Type alias for application state extension.
64pub type State<B> = Extension<Arc<RwLock<B>>>;
65
66/// Fetch all tasks from the backend storage.
67pub async fn get_tasks<S, T, Compact>(
68    query: Query<Filter>,
69    storage: State<S>,
70) -> Result<Json<Vec<Task<T, S::Context, S::IdType>>>, AppError>
71where
72    T: Serialize + DeserializeOwned + 'static,
73    S: ListTasks<T> + Send + 'static + BackendExt,
74    S::Context: Serialize + 'static,
75    S::IdType: Serialize + 'static,
76    <S as Backend>::Error: std::error::Error + 'static,
77    S::Codec: Codec<T, Compact = Compact> + 'static,
78    Compact: 'static,
79{
80    let storage = storage.0;
81    let filter = query.0;
82
83    crate::get_tasks::<S, T, Compact>(storage, filter)
84        .await
85        .map(Json)
86        .map_err(AppError::ApiError)
87}
88
89/// Fetch statistics for a specific queue from the backend storage.
90pub async fn stats_by_queue<S>(storage: State<S>) -> Result<Json<Vec<Statistic>>, AppError>
91where
92    S::Error: std::error::Error,
93    S: Metrics + BackendExt,
94{
95    let storage = storage.0;
96
97    match crate::stats_by_queue::<S>(storage).await {
98        Ok(stats) => Ok(Json(stats)),
99        Err(e) => Err(AppError::ApiError(e)),
100    }
101}
102
103/// Fetch all workers from the backend storage.
104pub async fn get_workers<S>(storage: State<S>) -> Result<Json<Vec<RunningWorker>>, AppError>
105where
106    S: ListWorkers + BackendExt,
107    S::Error: std::error::Error,
108{
109    let storage = storage.0;
110
111    match crate::get_workers::<S>(storage).await {
112        Ok(workers) => Ok(Json(workers)),
113        Err(e) => Err(AppError::ApiError(e)),
114    }
115}
116
117/// Push a new task to the backend storage.
118pub async fn push_task<S, T, Compact>(
119    storage: State<S>,
120    task: Json<T>,
121) -> Result<Json<()>, AppError>
122where
123    T: Serialize + DeserializeOwned + 'static + Send,
124    S: TaskSink<T> + 'static + Send + BackendExt,
125    S::Error: std::error::Error,
126    S::Codec: Codec<T, Compact = Compact>,
127    <<S as BackendExt>::Codec as Codec<T>>::Error: std::error::Error,
128{
129    match crate::push_task(task.0, storage.0).await {
130        Ok(_) => Ok(Json(())),
131        Err(e) => Err(AppError::ApiError(e)),
132    }
133}
134
135/// Fetch a task by its ID from the backend storage.
136pub async fn get_task_by_id<S, T>(
137    Path(task_id): Path<String>,
138    storage: State<S>,
139) -> Result<Json<Task<T, S::Context, S::IdType>>, AppError>
140where
141    T: Serialize + DeserializeOwned + 'static + Send,
142    S: FetchById<T> + Send + 'static,
143    S::Context: Serialize + 'static + Send,
144    S::IdType: Serialize + 'static + Send,
145    S::Error: std::error::Error,
146    S::IdType: FromStr + 'static + Send,
147    <<S as Backend>::IdType as FromStr>::Err: std::error::Error,
148{
149    let task_id = task_id.clone();
150    let storage = storage.0;
151
152    match crate::get_task_by_id::<S, T>(task_id, storage).await {
153        Ok(Some(task)) => Ok(Json(task)),
154        Ok(None) => Err(AppError::NotFound),
155        Err(e) => Err(AppError::ApiError(e)),
156    }
157}
158
159/// Fetch all tasks from the backend storage.
160pub async fn get_all_tasks<S>(
161    query: Query<Filter>,
162    storage: State<S>,
163) -> Result<Json<Vec<Task<S::Compact, S::Context, S::IdType>>>, AppError>
164where
165    S: ListAllTasks + BackendExt + Send + 'static,
166    S::Context: Serialize,
167    S::IdType: Serialize,
168    S::Compact: Serialize,
169    <S as Backend>::Error: std::error::Error,
170    <<S as BackendExt>::Codec as Codec<<S as Backend>::Args>>::Error: std::error::Error,
171{
172    let storage = storage.0;
173    let filter = query.0;
174
175    match crate::get_all_tasks::<S>(storage, filter).await {
176        Ok(tasks) => Ok(Json(tasks)),
177        Err(e) => Err(AppError::ApiError(e)),
178    }
179}
180
181/// Fetch all workers from the backend storage.
182pub async fn get_all_workers<S>(storage: State<S>) -> Result<Json<Vec<RunningWorker>>, AppError>
183where
184    S: ListWorkers + 'static,
185    S::Error: std::error::Error,
186{
187    let storage = storage.0;
188
189    match crate::get_all_workers::<S>(storage).await {
190        Ok(workers) => Ok(Json(workers)),
191        Err(e) => Err(AppError::ApiError(e)),
192    }
193}
194
195/// Fetch all queues from the backend storage.
196pub async fn fetch_queues<S>(storage: State<S>) -> Result<Json<Vec<QueueInfo>>, AppError>
197where
198    S::Error: std::error::Error,
199    S: ListQueues,
200{
201    let storage = storage.0;
202
203    crate::fetch_queues::<S>(storage)
204        .await
205        .map_err(AppError::ApiError)
206        .map(Json)
207}
208
209/// Get an overview of statistics across all queues.
210pub async fn overview<S>(storage: State<S>) -> Result<Json<Vec<Statistic>>, AppError>
211where
212    S::Error: std::error::Error,
213    S: Metrics + 'static,
214{
215    let storage = storage.0;
216
217    let tasks = crate::overview::<S>(storage)
218        .await
219        .map_err(AppError::ApiError)?;
220
221    Ok(Json(tasks))
222}
223
224impl<B, T, Compact> RegisterRoute<B, T> for ApiBuilder<Router>
225where
226    B: Metrics + ListWorkers + ListAllTasks + ListQueues,
227    B::Context: Serialize,
228    B::IdType: Serialize,
229    <B as Backend>::Error: std::error::Error,
230    B::IdType: FromStr + 'static + Send,
231    <<B as Backend>::IdType as FromStr>::Err: std::error::Error,
232    Compact: Serialize + 'static + Send,
233    B::Compact: Serialize + 'static + Send,
234    B::Context: Serialize + 'static + Send,
235    <B as Backend>::Error: std::error::Error,
236    <<B as BackendExt>::Codec as Codec<<B as Backend>::Args>>::Error: std::error::Error,
237    T: Serialize + DeserializeOwned + 'static + Send,
238    B: ListTasks<T> + FetchById<T>,
239    B::Codec: Codec<T, Compact = Compact>,
240    <<B as BackendExt>::Codec as Codec<T>>::Error: std::error::Error,
241    B: TaskSink<T> + BackendExt + Send + Sync + 'static,
242{
243    fn register(mut self, backend: B) -> Self {
244        let queue = backend.get_queue();
245        let backend = Arc::new(RwLock::new(backend));
246        if self.root {
247            #[allow(unused_mut)]
248            let mut r = self
249                .router
250                .route("/queues", get(fetch_queues::<B>))
251                .route("/tasks", get(get_all_tasks::<B>))
252                .route("/workers", get(get_all_workers::<B>))
253                .route("/overview", get(overview::<B>));
254
255            #[cfg(feature = "sse")]
256            {
257                r = r.route("/events", get(sse::new_client));
258            }
259
260            self.router = r.layer(Extension(backend.clone()));
261        }
262        let scope = self.router.nest(
263            &format!("/queues/{queue}"),
264            Router::new()
265                .route("/tasks", get(get_tasks::<B, T, Compact>))
266                .route("/stats", get(stats_by_queue::<B>))
267                .route("/workers", get(get_workers::<B>))
268                .route("/tasks", put(push_task::<B, T, Compact>))
269                .route("/tasks/{task_id}", get(get_task_by_id::<B, T>))
270                .layer(Extension(queue))
271                .layer(Extension(backend)),
272        );
273
274        Self {
275            router: scope,
276            root: false,
277        }
278    }
279}
280
281#[cfg(feature = "ui")]
282mod ui {
283    use std::{
284        convert::Infallible,
285        task::{Context, Poll},
286    };
287
288    use apalis_core::layers::Service;
289    use axum::{
290        body::Body,
291        http::{Request, Response, StatusCode},
292    };
293
294    use crate::ui::ServeUI;
295
296    impl Service<Request<Body>> for ServeUI {
297        type Response = Response<Body>;
298        type Error = Infallible;
299        type Future = std::future::Ready<Result<Self::Response, Self::Error>>;
300
301        fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
302            Poll::Ready(Ok(()))
303        }
304
305        fn call(&mut self, req: Request<Body>) -> Self::Future {
306            let path = req.uri().path();
307            let mut file = Self::get_file(path);
308
309            // If no matching file, fall back to index.html
310            if file.is_none() {
311                file = Self::get_file("index.html");
312            }
313
314            let response = match file {
315                Some(file) => {
316                    let path_str = file.path().to_str().unwrap_or("");
317                    let content_type = Self::content_type(path_str);
318                    let mut builder = Response::builder()
319                        .status(StatusCode::OK)
320                        .header("Content-Type", content_type);
321
322                    if let Some(cache) = Self::cache_control(path_str) {
323                        builder = builder.header("Cache-Control", cache);
324                    }
325
326                    builder.body(file.contents().to_vec().into()).unwrap()
327                }
328                None => Response::builder()
329                    .status(StatusCode::NOT_FOUND)
330                    .body(Vec::new().into())
331                    .unwrap(),
332            };
333
334            std::future::ready(Ok(response))
335        }
336    }
337}
338
339/// Expose Server-Sent Events (SSE) functionality.
340#[cfg(feature = "sse")]
341pub mod sse {
342
343    use std::{sync::Mutex, time::Duration};
344
345    use axum::response::{Sse, sse::Event};
346    use futures::{Stream, StreamExt, channel::mpsc::TryRecvError};
347
348    use crate::sse::TracingBroadcaster;
349
350    use super::*;
351
352    /// Create a new SSE client and register it with the broadcaster.
353    pub async fn new_client(
354        broadcaster: Extension<Arc<Mutex<TracingBroadcaster>>>,
355    ) -> Sse<impl Stream<Item = Result<Event, TryRecvError>>> {
356        let rx = broadcaster.lock().unwrap().new_client();
357        let stream = rx
358            .filter(|s| futures::future::ready(s.as_ref().is_ok_and(|e| e.span.is_some())))
359            .map(|entry| Ok(Event::default().json_data(entry?).unwrap()));
360
361        Sse::new(stream).keep_alive(
362            axum::response::sse::KeepAlive::new()
363                .interval(Duration::from_secs(1))
364                .text("keep-alive-text"),
365        )
366    }
367}
apalis_board_api/framework/axum.rs

apalis_board_api/framework/
axum.rs
