micromegas/servers/
ingestion.rs

1use axum::Extension;
2use axum::Router;
3use axum::body::Body;
4use axum::http::{Response, StatusCode};
5use axum::response::IntoResponse;
6use axum::routing::post;
7use micromegas_auth::types::AuthProvider;
8use micromegas_ingestion::data_lake_connection::DataLakeConnection;
9use micromegas_ingestion::web_ingestion_service::{IngestionServiceError, WebIngestionService};
10use micromegas_tracing::prelude::*;
11use std::future::Future;
12use std::net::SocketAddr;
13use std::sync::Arc;
14use std::time::Duration;
15use thiserror::Error;
16
17#[derive(Error, Debug)]
18pub enum IngestionError {
19    #[error("Bad request: {0}")]
20    BadRequest(String),
21
22    #[error("Internal server error: {0}")]
23    Internal(String),
24}
25
26impl IntoResponse for IngestionError {
27    fn into_response(self) -> Response<Body> {
28        let (status, category, detail) = match self {
29            IngestionError::BadRequest(msg) => (StatusCode::BAD_REQUEST, "Bad request", msg),
30            IngestionError::Internal(msg) => (
31                StatusCode::INTERNAL_SERVER_ERROR,
32                "Internal server error",
33                msg,
34            ),
35        };
36        error!("{status}: {detail}");
37        (status, category).into_response()
38    }
39}
40
41impl From<IngestionServiceError> for IngestionError {
42    fn from(err: IngestionServiceError) -> Self {
43        match err {
44            IngestionServiceError::ParseError(msg) => IngestionError::BadRequest(msg),
45            IngestionServiceError::DatabaseError(msg) => IngestionError::Internal(msg),
46            IngestionServiceError::StorageError(msg) => IngestionError::Internal(msg),
47        }
48    }
49}
50
51/// Handles requests to insert process information.
52///
53/// Returns 400 for malformed CBOR, 500 for database errors.
54pub async fn insert_process_request(
55    Extension(service): Extension<Arc<WebIngestionService>>,
56    body: bytes::Bytes,
57) -> Result<(), IngestionError> {
58    service.insert_process(body).await.map_err(Into::into)
59}
60
61/// Handles requests to insert stream information.
62///
63/// Returns 400 for malformed CBOR, 500 for database errors.
64pub async fn insert_stream_request(
65    Extension(service): Extension<Arc<WebIngestionService>>,
66    body: bytes::Bytes,
67) -> Result<(), IngestionError> {
68    service.insert_stream(body).await.map_err(Into::into)
69}
70
71/// Handles requests to insert block information.
72///
73/// Returns 400 for empty body or malformed CBOR, 500 for database/storage errors.
74pub async fn insert_block_request(
75    Extension(service): Extension<Arc<WebIngestionService>>,
76    body: bytes::Bytes,
77) -> Result<(), IngestionError> {
78    if body.is_empty() {
79        return Err(IngestionError::BadRequest("empty body".to_string()));
80    }
81    service.insert_block(body).await.map_err(Into::into)
82}
83
84async fn ready_handler(Extension(service): Extension<Arc<WebIngestionService>>) -> StatusCode {
85    if service.check_ready().await {
86        StatusCode::OK
87    } else {
88        StatusCode::SERVICE_UNAVAILABLE
89    }
90}
91
92/// Registers the ingestion routes with the given Axum `Router`.
93///
94/// This function adds routes for `/ingestion/insert_process`,
95/// `/ingestion/insert_stream`, and `/ingestion/insert_block`.
96pub fn register_routes(router: Router) -> Router {
97    router
98        .route("/ingestion/insert_process", post(insert_process_request))
99        .route("/ingestion/insert_stream", post(insert_stream_request))
100        .route("/ingestion/insert_block", post(insert_block_request))
101}
102
103/// Assemble and serve the HTTP ingestion endpoint.
104///
105/// Binds `listen_addr`, wires the ingestion routes + OTLP routes, applies
106/// the supplied `auth_provider` (or runs open when `None`), and shuts down
107/// gracefully when `shutdown` resolves.
108pub async fn serve_ingestion(
109    listen_addr: SocketAddr,
110    lake: DataLakeConnection,
111    auth_provider: Option<Arc<dyn AuthProvider>>,
112    shutdown: impl Future<Output = ()> + Send + 'static,
113    grace: Duration,
114) -> anyhow::Result<()> {
115    use axum::extract::DefaultBodyLimit;
116    use axum::middleware;
117    use axum::routing::get;
118    use micromegas_auth::axum::auth_middleware;
119    use tower_http::limit::RequestBodyLimitLayer;
120
121    use super::axum_utils::observability_middleware;
122    use super::shutdown::serve_axum_with_graceful_shutdown;
123
124    let service = Arc::new(WebIngestionService::new(lake));
125
126    let health_router = Router::new()
127        .route("/health", get(|| async { axum::http::StatusCode::OK }))
128        .route("/ready", get(ready_handler))
129        .layer(Extension(service.clone()));
130
131    let firehose_auth = auth_provider.clone();
132    let cw_logs_firehose_auth = auth_provider.clone();
133
134    let mut protected_app = register_routes(Router::new())
135        .merge(super::otlp::otlp_router())
136        .merge(super::webhook::webhook_router())
137        .layer(DefaultBodyLimit::disable())
138        .layer(RequestBodyLimitLayer::new(100 * 1024 * 1024))
139        .layer(Extension(service.clone()));
140
141    let auth_enabled = auth_provider.is_some();
142    if let Some(provider) = auth_provider {
143        info!("Ingestion: authentication enabled");
144        protected_app = protected_app.layer(middleware::from_fn(move |req, next| {
145            auth_middleware(provider.clone(), req, next)
146        }));
147    } else {
148        warn!("Ingestion: authentication disabled — development mode only");
149    }
150
151    // The Firehose routes carry their own auth (Firehose can only send its credential via
152    // X-Amz-Firehose-Access-Key, not Authorization: Bearer), so they are merged outside
153    // protected_app and never hit the global Bearer auth_middleware.
154    let firehose_app = super::firehose::firehose_router(service.clone(), firehose_auth);
155    let cw_logs_firehose_app =
156        super::firehose_cloudwatch_logs::firehose_router(service.clone(), cw_logs_firehose_auth);
157
158    let app = health_router
159        .merge(protected_app)
160        .merge(firehose_app)
161        .merge(cw_logs_firehose_app)
162        .layer(middleware::from_fn(observability_middleware));
163
164    let listener = tokio::net::TcpListener::bind(listen_addr)
165        .await
166        .map_err(|e| anyhow::anyhow!("ingestion: binding to {listen_addr}: {e}"))?;
167    info!("Ingestion serving on {listen_addr} authentication={auth_enabled}");
168
169    serve_axum_with_graceful_shutdown(
170        listener,
171        app.into_make_service_with_connect_info::<SocketAddr>(),
172        shutdown,
173        grace,
174    )
175    .await?;
176
177    Ok(())
178}
micromegas/servers/ingestion.rs

micromegas/servers/
ingestion.rs
