feat: Background task SSE streaming endpoint
- Added /api/update-center/tasks/{id}/stream SSE endpoint
- Polls DB for task status updates every 1s (up to 2 min)
- Streams status + log events, ends on completion/timeout
- Added async-stream dependency
- 18/18 tests pass
This commit is contained in:
@@ -11,6 +11,7 @@ tower = "0.5"
|
||||
tower-http = { version = "0.6", features = ["cors", "trace", "fs", "compression-gzip"] }
|
||||
|
||||
# Async runtime
|
||||
async-stream = "0.3"
|
||||
tokio = { version = "1", features = ["full"] }
|
||||
|
||||
# Database
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
// Background tasks routes — list and inspect async operations
|
||||
|
||||
use std::convert::Infallible;
|
||||
|
||||
use crate::{
|
||||
db::models::background_task::BackgroundTask,
|
||||
error::AppResult,
|
||||
@@ -7,11 +9,17 @@ use crate::{
|
||||
};
|
||||
use axum::{
|
||||
extract::{Path, Query, State},
|
||||
response::{
|
||||
sse::{Event, Sse},
|
||||
IntoResponse,
|
||||
},
|
||||
routing::get,
|
||||
Json, Router,
|
||||
};
|
||||
use futures::stream::{self, Stream};
|
||||
use serde::Deserialize;
|
||||
use serde_json::{json, Value};
|
||||
use tokio_stream::StreamExt;
|
||||
use tracing::info;
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
@@ -23,6 +31,7 @@ pub fn routes() -> Router<AppState> {
|
||||
Router::new()
|
||||
.route("/api/tasks", get(list_tasks))
|
||||
.route("/api/tasks/{id}", get(get_task))
|
||||
.route("/api/update-center/tasks/{id}/stream", get(stream_task_logs))
|
||||
}
|
||||
|
||||
/// GET /api/tasks — list background tasks
|
||||
@@ -62,3 +71,109 @@ async fn get_task(
|
||||
None => Ok(Json(json!({ "success": false, "message": "task not found" }))),
|
||||
}
|
||||
}
|
||||
|
||||
/// GET /api/update-center/tasks/{id}/stream — SSE stream of task progress
|
||||
async fn stream_task_logs(
|
||||
State(state): State<AppState>,
|
||||
Path(id): Path<String>,
|
||||
) -> Result<Sse<impl Stream<Item = Result<Event, Infallible>>>, axum::response::Response> {
|
||||
info!("Streaming task logs for: {}", id);
|
||||
|
||||
// Check task exists
|
||||
let task: Option<BackgroundTask> = sqlx::query_as::<_, BackgroundTask>(
|
||||
"SELECT * FROM background_tasks WHERE id = ?"
|
||||
)
|
||||
.bind(&id)
|
||||
.fetch_optional(&*state.db)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
(axum::http::StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")).into_response()
|
||||
})?;
|
||||
|
||||
let task = match task {
|
||||
Some(t) => t,
|
||||
None => {
|
||||
return Err((
|
||||
axum::http::StatusCode::NOT_FOUND,
|
||||
json!({ "success": false, "message": "task not found" }).to_string(),
|
||||
).into_response());
|
||||
}
|
||||
};
|
||||
|
||||
// Stream task state with polling
|
||||
let initial_status = task.status.clone();
|
||||
let is_done = initial_status != "pending" && initial_status != "running";
|
||||
|
||||
let state_clone = state.clone();
|
||||
let stream = async_stream::stream! {
|
||||
// Yield current task status
|
||||
yield Ok::<_, Infallible>(Event::default()
|
||||
.event("status")
|
||||
.data(json!({ "status": initial_status }).to_string()));
|
||||
|
||||
yield Ok::<_, Infallible>(Event::default()
|
||||
.event("log")
|
||||
.data(json!({
|
||||
"seq": 1,
|
||||
"message": format!("Task {}: {}", initial_status, task.message.as_deref().unwrap_or("")),
|
||||
"createdAt": task.created_at.as_deref().unwrap_or(""),
|
||||
}).to_string()));
|
||||
|
||||
if is_done {
|
||||
yield Ok::<_, Infallible>(Event::default()
|
||||
.event("done")
|
||||
.data(json!({ "status": initial_status }).to_string()));
|
||||
return;
|
||||
}
|
||||
|
||||
// Poll for updates
|
||||
let poll_id = id.clone();
|
||||
for i in 0..120 { // Poll up to 120 times (2 minutes)
|
||||
tokio::time::sleep(tokio::time::Duration::from_millis(1000)).await;
|
||||
|
||||
let current: Option<BackgroundTask> = sqlx::query_as(
|
||||
"SELECT * FROM background_tasks WHERE id = ?"
|
||||
)
|
||||
.bind(&poll_id)
|
||||
.fetch_optional(&*state_clone.db)
|
||||
.await
|
||||
.ok()
|
||||
.flatten();
|
||||
|
||||
match current {
|
||||
Some(t) => {
|
||||
let status = t.status;
|
||||
let updated_at = t.updated_at.unwrap_or_default();
|
||||
|
||||
yield Ok::<_, Infallible>(Event::default()
|
||||
.event("log")
|
||||
.data(json!({
|
||||
"seq": i + 2,
|
||||
"message": format!("Status: {} — {}", status, t.message.as_deref().unwrap_or("")),
|
||||
"createdAt": updated_at,
|
||||
}).to_string()));
|
||||
|
||||
if status != "pending" && status != "running" {
|
||||
yield Ok::<_, Infallible>(Event::default()
|
||||
.event("done")
|
||||
.data(json!({ "status": status }).to_string()));
|
||||
return;
|
||||
}
|
||||
}
|
||||
None => {
|
||||
yield Ok::<_, Infallible>(Event::default()
|
||||
.event("done")
|
||||
.data(json!({ "status": "unknown", "message": "task removed" }).to_string()));
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Timeout
|
||||
yield Ok::<_, Infallible>(Event::default()
|
||||
.event("done")
|
||||
.data(json!({ "status": "timeout" }).to_string()));
|
||||
};
|
||||
|
||||
Ok(Sse::new(stream))
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user