Merge branch 'main' into sibling-relative-import

Author: sitaowang1998

Cargo.lock

@@ -4,7 +4,7 @@ version = "0.1.0"
 edition = "2024"
 
 [lib]
-name = "clp_client"
+name = "api_server"
 path = "src/lib.rs"
 
 [[bin]]
@@ -29,6 +29,9 @@ sqlx = { version = "0.8.6", features = ["runtime-tokio", "mysql"] }
 thiserror = "2.0.17"
 tokio = { version = "1.48.0", features = ["full"] }
 tokio-stream = { version = "0.1.17", features = ["fs"] }
+tower-http = { version = "0.6.7", features = ["cors"] }
 tracing = "0.1.41"
 tracing-appender = "0.2.3"
 tracing-subscriber = { version = "0.3.20", features = ["json", "env-filter", "fmt", "std"] }
+utoipa = { version = "5.4.0", features = ["axum_extras"] }
+utoipa-axum = "0.2.0"

components/api-server/Cargo.toml

@@ -1,20 +1,6 @@
 use anyhow::Context;
-use axum::{
-    Json,
-    Router,
-    extract::{Path, State},
-    http::StatusCode,
-    response::{
-        IntoResponse,
-        Sse,
-        sse::{Event, KeepAlive},
-    },
-    routing::{get, post},
-};
 use clap::Parser;
 use clp_rust_utils::{clp_config::package, serde::yaml};
-use futures::{Stream, StreamExt};
-use thiserror::Error;
 use tracing_appender::rolling::{RollingFileAppender, Rotation};
 use tracing_subscriber::{self, fmt::writer::MakeWriterExt};
 
@@ -95,75 +81,15 @@ async fn main() -> anyhow::Result<()> {
         .await
         .context(format!("Cannot listen to {addr}"))?;
 
-    let client = clp_client::Client::connect(&config, &credentials)
+    let client = api_server::client::Client::connect(&config, &credentials)
         .await
         .context("Cannot connect to CLP")?;
 
-    let app = Router::new()
-        .route("/", get(health))
-        .route("/health", get(health))
-        .route("/query", post(query))
-        .route("/query_results/{search_job_id}", get(query_results))
-        .with_state(client);
+    let router = api_server::routes::from_client(client)?;
 
     tracing::info!("Server started at {addr}");
-    axum::serve(listener, app)
+    axum::serve(listener, router)
         .with_graceful_shutdown(shutdown_signal())
         .await?;
     Ok(())
 }
-
-async fn health() -> String {
-    "API server is running".to_owned()
-}
-
-async fn query(
-    State(client): State<clp_client::Client>,
-    Json(query_config): Json<clp_client::QueryConfig>,
-) -> Result<Json<serde_json::Value>, HandlerError> {
-    let search_job_id = client.submit_query(query_config).await?;
-    let uri = format!("/query_results/{search_job_id}");
-    Ok(Json(serde_json::json!({"query_results_uri": uri})))
-}
-
-async fn query_results(
-    State(client): State<clp_client::Client>,
-    Path(search_job_id): Path<u64>,
-) -> Result<Sse<impl Stream<Item = Result<Event, HandlerError>>>, HandlerError> {
-    let results_stream = client.fetch_results(search_job_id).await?;
-    let event_stream = results_stream.map(|res| {
-        let message = res?;
-        let trimmed_message = message.trim();
-        if trimmed_message.lines().count() != 1 {
-            return Err(HandlerError::InternalServer);
-        }
-        Ok(Event::default().data(trimmed_message))
-    });
-    Ok(Sse::new(event_stream).keep_alive(KeepAlive::default()))
-}
-
-/// Generic errors for request handlers.
-#[derive(Error, Debug)]
-enum HandlerError {
-    #[error("Internal server error")]
-    InternalServer,
-}
-
-trait IntoHandlerError {}
-
-impl IntoHandlerError for axum::Error {}
-
-impl IntoHandlerError for clp_client::ClientError {}
-
-impl<T: IntoHandlerError> From<T> for HandlerError {
-    fn from(_: T) -> Self {
-        Self::InternalServer
-    }
-}
-
-/// Converts [`HandlerError`] into an HTTP response.
-impl IntoResponse for HandlerError {
-    fn into_response(self) -> axum::response::Response {
-        StatusCode::INTERNAL_SERVER_ERROR.into_response()
-    }
-}

components/api-server/src/bin/api_server.rs

@@ -13,24 +13,39 @@ use futures::{Stream, StreamExt};
 use pin_project_lite::pin_project;
 use serde::{Deserialize, Serialize};
 use sqlx::Row;
+use utoipa::ToSchema;
 
 pub use crate::error::ClientError;
 
 /// Defines the request configuration for submitting a search query.
-#[derive(Clone, Serialize, Deserialize)]
+#[derive(Clone, Serialize, Deserialize, ToSchema)]
 #[serde(deny_unknown_fields)]
 pub struct QueryConfig {
+    /// The search query as a KQL string.
     pub query_string: String,
+
+    /// The dataset to search within. If not provided, only `default` dataset will be searched.
     #[serde(default)]
     pub dataset: Option<String>,
+
+    /// The maximum number of results to return. Set to `0` for no limit.
     #[serde(default)]
     pub max_num_results: u32,
+
+    /// The beginning timestamp (in epoch milliseconds) for the search range (inclusive).
     #[serde(default)]
     pub begin_timestamp: Option<i64>,
+
+    /// The ending timestamp (in epoch milliseconds) for the search range (inclusive).
     #[serde(default)]
     pub end_timestamp: Option<i64>,
+
+    /// Whether the string match should be case-insensitive.
     #[serde(default)]
     pub ignore_case: bool,
+
+    /// Whether to write the search results to files. If `false`, results will be stored in
+    /// `MongoDB`.
     #[serde(default)]
     pub write_to_file: bool,
 }

components/api-server/src/client.rs

@@ -1,3 +1,3 @@
-mod client;
+pub mod client;
 mod error;
-pub use client::*;
+pub mod routes;