use crate::prelude::*;
use axum::{
extract::{Query, State},
response::sse::{Event, KeepAlive, Sse},
};
use std::convert::Infallible;
use std::sync::atomic::AtomicBool;
use crate::web::search_sources::{CancelOnDrop, FileSource, SearchEvent, SessionSource};
use crate::web::state::WebState;
#[derive(Deserialize)]
pub struct SearchQuery {
pub q: String,
pub sort: Option<String>,
pub project: Option<String>,
}
pub async fn global_search(
State(state): State<WebState>,
Query(params): Query<SearchQuery>,
) -> Sse<impl futures::Stream<Item = Result<Event, Infallible>>> {
let cancelled = Arc::new(AtomicBool::new(false));
let cancel_guard = CancelOnDrop(cancelled.clone());
let (tx, mut rx) = tokio::sync::mpsc::channel::<SearchEvent>(256);
{
let tx = tx.clone();
let q = params.q.clone();
let c = cancelled.clone();
let db = state.db.clone();
let sort = params.sort.clone();
let project = params.project.clone();
tokio::spawn(async move {
let source = SessionSource::new(db, sort, project);
source.run(q, tx, c).await;
});
}
{
let tx = tx.clone();
let q = params.q.clone();
let c = cancelled.clone();
let source = FileSource::new(state.workspace_root.clone());
tokio::task::spawn_blocking(move || source.run(q, tx, c));
}
drop(tx);
let stream = async_stream::stream! {
let _guard = cancel_guard;
while let Some(event) = rx.recv().await {
match serde_json::to_string(&event) {
Ok(json) => yield Ok(Event::default().event("result").data(json)),
Err(e) => yield Ok(Event::default().event("error").data(e.to_string())),
}
}
yield Ok(Event::default().event("end").data("{}"));
};
Sse::new(stream).keep_alive(KeepAlive::default())
}