apalis-board-api 1.0.0

use std::{
    fmt::Display,
    marker::PhantomData,
    str::FromStr,
    time::{SystemTime, UNIX_EPOCH},
};

use actix_web::{
    Error, FromRequest, HttpRequest, HttpResponse, Responder, Scope,
    dev::Payload,
    error::ErrorBadRequest,
    http::header::HeaderMap,
    web::{self, Data},
};
use apalis_core::{
    backend::{
        Backend, BackendConfig, FetchById, Filter, ListAllTasks, ListQueues, ListTasks,
        ListWorkers, Metrics, TaskSink, WireFormatBackend, codec::Codec,
    },
    task::{Task, builder::TaskBuilder, metadata::MetadataStore},
};
use futures::future::LocalBoxFuture;
use serde::{Serialize, de::DeserializeOwned};
use tokio::sync::RwLock;

use crate::{
    META_PREFIX, RawTask, fetch_queues,
    framework::{ApiBuilder, RegisterRoute, TaskRejection},
    get_all_tasks, get_all_workers, get_task_by_id, get_tasks, get_workers, overview, push_task,
    stats_by_queue,
};

#[cfg(feature = "ui")]
use crate::ui::ServeUI;

/// Handler struct for Actix web routes.
#[derive(Debug, Clone)]
pub struct Handler<S, T, Compact> {
    _phantom: PhantomData<(S, T, Compact)>,
}

impl<S, T, Compact> Handler<S, T, Compact> {
    /// Get tasks for a specific queue.
    pub async fn get_tasks(
        storage: web::Data<RwLock<S>>,
        query: web::Query<Filter>,
    ) -> impl Responder
    where
        T: Serialize + DeserializeOwned + 'static,
        S: ListTasks + Send + 'static,
        S::Codec: Codec<T, Compact = S::Compact>,
        <S::Codec as Codec<T>>::Error: std::error::Error,
    {
        let storage = storage.into_inner();
        let filter = query.into_inner();

        match get_tasks::<S, T>(storage, filter).await {
            Ok(tasks) => HttpResponse::Ok().json(
                tasks
                    .into_iter()
                    .map(|task| RawTask(task))
                    .collect::<Vec<RawTask<T>>>(),
            ),
            Err(e) => HttpResponse::InternalServerError().json(e),
        }
    }

    /// Get statistics for a specific queue.
    pub async fn stats_by_queue(storage: web::Data<RwLock<S>>) -> impl Responder
    where
        S::Error: std::error::Error,
        S: Metrics,
    {
        let storage = storage.into_inner();

        match stats_by_queue::<S>(storage).await {
            Ok(stats) => HttpResponse::Ok().json(stats),
            Err(e) => HttpResponse::InternalServerError().json(e),
        }
    }

    /// Get workers for a specific queue.
    pub async fn get_workers(storage: web::Data<RwLock<S>>) -> impl Responder
    where
        S: ListWorkers,
        S::Error: std::error::Error,
    {
        let storage = storage.into_inner();

        match get_workers::<S>(storage).await {
            Ok(workers) => HttpResponse::Ok().json(workers),
            Err(e) => HttpResponse::InternalServerError().json(e),
        }
    }

    /// Push a new task to the specified queue.
    pub async fn push_task(task: RawTask<T>, storage: Data<RwLock<S>>) -> impl Responder
    where
        S: BackendConfig,
        T: Serialize + DeserializeOwned + 'static,
        S: TaskSink<T, S::Kind> + Send + Sync + 'static,
    {
        match push_task(task, storage.into_inner()).await {
            Ok(_) => HttpResponse::Created().finish(),
            Err(e) => HttpResponse::InternalServerError().json(e),
        }
    }

    /// Get a task by its ID.
    pub async fn get_task_by_id(
        task_id: web::Path<String>,
        storage: web::Data<RwLock<S>>,
    ) -> impl Responder
    where
        S: FetchById + 'static,
        S::Error: std::error::Error,
        S: WireFormatBackend<Compact = Compact> + Backend<Task = Task<Compact>>,
        S::Codec: Codec<T, Compact = Compact>,
        <S::Codec as Codec<T>>::Error: std::error::Error,
        T: Serialize + DeserializeOwned + 'static,
    {
        let task_id = task_id.into_inner();
        let storage = storage.into_inner();

        match get_task_by_id::<S, T>(task_id, storage).await {
            Ok(Some(task)) => HttpResponse::Ok().json(RawTask(task)),
            Ok(None) => HttpResponse::NotFound().finish(),
            Err(e) => HttpResponse::InternalServerError().json(e),
        }
    }

    /// Get all tasks across all queues.
    pub async fn get_all_tasks(
        storage: web::Data<RwLock<S>>,
        query: web::Query<Filter>,
    ) -> impl Responder
    where
        S: ListAllTasks + Send,
        S::Compact: Serialize,
    {
        let storage = storage.into_inner();
        let filter = query.into_inner();

        match get_all_tasks::<S>(storage, filter).await {
            Ok(tasks) => HttpResponse::Ok().json(
                tasks
                    .into_iter()
                    .map(RawTask)
                    .collect::<Vec<RawTask<S::Compact>>>(),
            ),
            Err(e) => HttpResponse::InternalServerError().json(e),
        }
    }

    /// Get all workers across all queues.
    pub async fn get_all_workers(storage: web::Data<RwLock<S>>) -> impl Responder
    where
        S: ListWorkers,
        S::Error: std::error::Error,
    {
        let storage = storage.into_inner();

        match get_all_workers::<S>(storage).await {
            Ok(workers) => HttpResponse::Ok().json(workers),
            Err(e) => HttpResponse::InternalServerError().json(e),
        }
    }

    /// Fetch all queues.
    pub async fn fetch_queues(storage: web::Data<RwLock<S>>) -> impl Responder
    where
        S::Error: std::error::Error,
        S: ListQueues,
    {
        let storage = storage.into_inner();

        match fetch_queues::<S>(storage).await {
            Ok(queues) => HttpResponse::Ok().json(queues),
            Err(e) => HttpResponse::InternalServerError().json(e),
        }
    }

    /// Get an overview of statistics.
    pub async fn overview(storage: web::Data<RwLock<S>>) -> impl Responder
    where
        S::Error: std::error::Error,
        S: Metrics,
    {
        let storage = storage.into_inner();

        match overview::<S>(storage).await {
            Ok(stats) => HttpResponse::Ok().json(stats),
            Err(e) => HttpResponse::InternalServerError().json(e),
        }
    }
}

impl<B, T, Compact> RegisterRoute<B, T> for ApiBuilder<Scope>
where
    B: Metrics + ListWorkers + ListAllTasks + ListQueues + BackendConfig + Send + 'static,
    <B as Backend>::Error: std::error::Error,
    Compact: Serialize + 'static,
    B::Compact: Serialize,
    <B as Backend>::Error: std::error::Error,
    T: Serialize + DeserializeOwned + 'static,
    B: ListTasks + FetchById,
    B::Codec: Codec<T, Compact = Compact>,
    B: TaskSink<T, B::Kind> + Send + Sync + 'static,
    B: WireFormatBackend<Compact = Compact> + Backend<Task = Task<Compact>>,
    <B::Codec as Codec<T>>::Error: std::error::Error,
{
    fn register(mut self, queue: &str, backend: B) -> Self {
        let backend = web::Data::new(RwLock::new(backend));
        if self.root {
            #[allow(unused_mut)]
            let mut router = self
                .router
                .app_data(backend.clone())
                .route(
                    "/queues",
                    web::get().to(Handler::<B, (), Compact>::fetch_queues),
                )
                .route(
                    "/tasks",
                    web::get().to(Handler::<B, (), Compact>::get_all_tasks),
                )
                .route(
                    "/workers",
                    web::get().to(Handler::<B, (), Compact>::get_all_workers),
                )
                .route(
                    "/overview",
                    web::get().to(Handler::<B, (), Compact>::overview),
                );

            #[cfg(feature = "sse")]
            {
                router = router.route("/events", web::get().to(sse::new_client));
            }

            self.router = router;
        }
        let scope = self.router.service(
            Scope::new(&format!("/queues/{queue}"))
                .app_data(backend)
                .route("/tasks", web::get().to(Handler::<B, T, Compact>::get_tasks))
                .route(
                    "/stats",
                    web::get().to(Handler::<B, T, Compact>::stats_by_queue),
                )
                .route(
                    "/workers",
                    web::get().to(Handler::<B, T, Compact>::get_workers),
                )
                .route("/tasks", web::put().to(Handler::<B, T, Compact>::push_task)) // Allow add jobs via api
                .route(
                    "/tasks/{id}",
                    web::get().to(Handler::<B, T, B::Compact>::get_task_by_id),
                ),
        );

        Self {
            router: scope,
            root: false,
        }
    }
}

#[cfg(feature = "ui")]
mod ui {
    use super::ServeUI;
    use actix_web::{
        HttpRequest, HttpResponse, HttpResponseBuilder,
        dev::HttpServiceFactory,
        http::{StatusCode, header},
    };
    impl ServeUI {
        fn serve_file(path: &str) -> HttpResponse {
            let mut file = Self::get_file(path);
            if file.is_none() {
                // Try fallback to index.html for unknown routes
                file = Self::get_file("index.html");
            }

            match file {
                Some(f) => {
                    let path_str = f.path().to_str().unwrap_or("");
                    let mut builder = HttpResponse::Ok();
                    let mut builder =
                        builder.insert_header((header::CONTENT_TYPE, Self::content_type(path_str)));

                    if let Some(cache) = Self::cache_control(path_str) {
                        builder = builder.insert_header((header::CACHE_CONTROL, cache));
                    }

                    builder.body(f.contents().to_vec())
                }
                None => HttpResponseBuilder::new(StatusCode::NOT_FOUND).finish(),
            }
        }
    }
    impl HttpServiceFactory for ServeUI {
        fn register(self, config: &mut actix_web::dev::AppService) {
            let resource = actix_web::Resource::new("/{tail:.*}").route(actix_web::web::get().to(
                move |req: HttpRequest| async move {
                    let path = req.match_info().query("tail");

                    Self::serve_file(path)
                },
            ));
            resource.register(config);
        }
    }
}

/// Expose Server-Sent Events (SSE) functionality.
#[cfg(feature = "sse")]
pub mod sse {
    use std::{sync::Arc, time::Duration};

    use crate::sse::TracingBroadcaster;
    use actix_web::web::*;
    use actix_web_lab::sse::Event;
    use futures::StreamExt;
    use std::sync::Mutex;

    /// Create a new SSE client connection.
    pub async fn new_client(
        broadcaster: Data<Arc<Mutex<TracingBroadcaster>>>,
    ) -> impl actix_web::Responder {
        let rx = broadcaster.lock().unwrap().new_client();

        actix_web_lab::sse::Sse::from_stream(
            rx.filter(|s| futures::future::ready(s.as_ref().is_ok_and(|e| e.span.is_some())))
                .map(|entry| {
                    match actix_web_lab::sse::Data::new_json(
                        entry.map_err(actix_web::error::ErrorInternalServerError)?,
                    ) {
                        Ok(data) => Ok(Event::Data(data)),
                        Err(e) => Err(actix_web::error::ErrorInternalServerError(e)),
                    }
                }),
        )
        .with_keep_alive(Duration::from_secs(60 * 5))
    }
}

impl<T> FromRequest for RawTask<T>
where
    T: DeserializeOwned + 'static,
{
    type Error = Error;
    type Future = LocalBoxFuture<'static, Result<Self, Self::Error>>;

    fn from_request(req: &HttpRequest, payload: &mut Payload) -> Self::Future {
        let task = match extract_task(req) {
            Ok(ctx) => ctx,
            Err(err) => return Box::pin(async move { Err(ErrorBadRequest(err)) }),
        };

        let body = web::Json::<T>::from_request(req, payload);

        Box::pin(async move {
            let args = body.await?.into_inner();
            let task = task.map_args(|_| args);
            Ok(Self(task))
        })
    }
}

fn extract_task(req: &HttpRequest) -> Result<Task<()>, TaskRejection> {
    let headers = req.headers();
    let mut task = TaskBuilder::new(());

    if let Some(id) = parse_header(headers, "x-task-id")? {
        task = task.task_id(id);
    }

    if let Some(attempt) = parse_header::<usize>(headers, "x-task-attempt")? {
        task = task.attempt(attempt);
    }

    if let Some(status) = parse_header(headers, "x-task-status")? {
        task = task.status(status);
    }

    if let Some(run_at) = parse_header(headers, "x-task-run-at")? {
        task = task.run_at_timestamp(run_at);
    } else {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .expect("Time went backwards");
        task = task.run_at_timestamp(now.as_secs())
    }

    task = task.done_at(parse_header(headers, "x-task-done-at")?);

    task = task.lock_at(parse_header(headers, "x-task-lock-at")?);

    task = task.lock_by(parse_header(headers, "x-task-lock-by")?);

    if let Some(idempotency_key) = parse_header::<String>(headers, "x-task-idempotency-key")? {
        task = task.idempotency_key(idempotency_key);
    }

    if let Some(max_attempts) = parse_header(headers, "x-task-max-attempts")? {
        task = task.max_attempts(max_attempts);
    }

    if let Some(priority) = parse_header(headers, "x-task-priority")? {
        task = task.priority(priority);
    }

    let mut metadata = MetadataStore::default();
    for (name, value) in headers.iter() {
        if let Some(key) = name.as_str().strip_prefix(META_PREFIX) {
            let value = value.to_str().map_err(|e| TaskRejection::BadHeader {
                name: name.to_string(),
                reason: e.to_string(),
            })?;

            metadata.insert(key.to_owned(), value.to_owned()).unwrap();
        }
    }

    task = task.with_metadata(metadata);

    Ok(task.build())
}

fn parse_header<T>(headers: &HeaderMap, name: &str) -> Result<Option<T>, TaskRejection>
where
    T: FromStr,
    T::Err: Display,
{
    let Some(raw) = headers.get(name) else {
        return Ok(None);
    };
    let bad = |reason: String| TaskRejection::BadHeader {
        name: name.into(),
        reason,
    };
    let s = raw.to_str().map_err(|e| bad(e.to_string()))?;
    s.parse::<T>().map(Some).map_err(|e| bad(e.to_string()))
}