micromegas/servers/
ingestion.rs

1use super::write_audience::resolve_write_audience;
2use axum::Extension;
3use axum::Router;
4use axum::body::Body;
5use axum::http::{Response, StatusCode};
6use axum::response::IntoResponse;
7use axum::routing::post;
8use micromegas_auth::types::{AuthContext, AuthProvider};
9use micromegas_ingestion::data_lake_connection::DataLakeConnection;
10use micromegas_ingestion::web_ingestion_service::{IngestionServiceError, WebIngestionService};
11use micromegas_ingestion::write_audience::WriteAudience;
12use micromegas_tracing::prelude::*;
13use std::future::Future;
14use std::net::SocketAddr;
15use std::sync::Arc;
16use std::time::Duration;
17use thiserror::Error;
18
19#[derive(Error, Debug)]
20pub enum IngestionError {
21    #[error("Bad request: {0}")]
22    BadRequest(String),
23
24    #[error("Internal server error: {0}")]
25    Internal(String),
26
27    /// A conflicting `insert_process` re-registration under a different audience. Maps to 403.
28    #[error("Forbidden: {0}")]
29    Forbidden(String),
30}
31
32impl IntoResponse for IngestionError {
33    fn into_response(self) -> Response<Body> {
34        let (status, category, detail) = match self {
35            IngestionError::BadRequest(msg) => (StatusCode::BAD_REQUEST, "Bad request", msg),
36            IngestionError::Internal(msg) => (
37                StatusCode::INTERNAL_SERVER_ERROR,
38                "Internal server error",
39                msg,
40            ),
41            IngestionError::Forbidden(msg) => (StatusCode::FORBIDDEN, "Forbidden", msg),
42        };
43        error!("{status}: {detail}");
44        (status, category).into_response()
45    }
46}
47
48impl From<IngestionServiceError> for IngestionError {
49    fn from(err: IngestionServiceError) -> Self {
50        match err {
51            IngestionServiceError::ParseError(msg) => IngestionError::BadRequest(msg),
52            IngestionServiceError::DatabaseError(msg) => IngestionError::Internal(msg),
53            IngestionServiceError::StorageError(msg) => IngestionError::Internal(msg),
54            IngestionServiceError::AudienceConflict { .. }
55            | IngestionServiceError::StreamAudienceConflict { .. } => {
56                IngestionError::Forbidden(err.to_string())
57            }
58        }
59    }
60}
61
62/// Handles requests to insert process information.
63///
64/// Returns 403 for a conflicting re-registration under a different audience, 400 for malformed
65/// CBOR, 500 for database errors.
66pub async fn insert_process_request(
67    Extension(service): Extension<Arc<WebIngestionService>>,
68    ctx: Option<Extension<AuthContext>>,
69    body: bytes::Bytes,
70) -> Result<(), IngestionError> {
71    let audience = resolve_write_audience(ctx.as_ref(), service.default_audience());
72    service
73        .insert_process(body, &audience)
74        .await
75        .map_err(Into::into)
76}
77
78/// Handles requests to insert stream information.
79///
80/// Returns 400 for malformed CBOR, 500 for database errors.
81pub async fn insert_stream_request(
82    Extension(service): Extension<Arc<WebIngestionService>>,
83    ctx: Option<Extension<AuthContext>>,
84    body: bytes::Bytes,
85) -> Result<(), IngestionError> {
86    let audience = resolve_write_audience(ctx.as_ref(), service.default_audience());
87    service
88        .insert_stream(body, &audience)
89        .await
90        .map_err(Into::into)
91}
92
93/// Handles requests to insert block information.
94///
95/// Returns 400 for empty body or malformed CBOR, 500 for database/storage errors.
96pub async fn insert_block_request(
97    Extension(service): Extension<Arc<WebIngestionService>>,
98    ctx: Option<Extension<AuthContext>>,
99    body: bytes::Bytes,
100) -> Result<(), IngestionError> {
101    if body.is_empty() {
102        return Err(IngestionError::BadRequest("empty body".to_string()));
103    }
104    let audience = resolve_write_audience(ctx.as_ref(), service.default_audience());
105    service
106        .insert_block(body, &audience)
107        .await
108        .map_err(Into::into)
109}
110
111async fn ready_handler(Extension(service): Extension<Arc<WebIngestionService>>) -> StatusCode {
112    if service.check_ready().await {
113        StatusCode::OK
114    } else {
115        StatusCode::SERVICE_UNAVAILABLE
116    }
117}
118
119/// Registers the ingestion routes with the given Axum `Router`.
120///
121/// This function adds routes for `/ingestion/insert_process`,
122/// `/ingestion/insert_stream`, and `/ingestion/insert_block`.
123pub fn register_routes(router: Router) -> Router {
124    router
125        .route("/ingestion/insert_process", post(insert_process_request))
126        .route("/ingestion/insert_stream", post(insert_stream_request))
127        .route("/ingestion/insert_block", post(insert_block_request))
128}
129
130/// Assemble and serve the HTTP ingestion endpoint.
131///
132/// Binds `listen_addr`, wires the ingestion routes + OTLP routes, applies
133/// the supplied `auth_provider` (or runs open when `None`), and shuts down
134/// gracefully when `shutdown` resolves.
135///
136/// Ingestion exposes no key-management HTTP surface of its own — keys for both
137/// `ingestion_api_keys` and `analytics_api_keys` are administered exclusively through
138/// `analytics-web-srv`'s own routes. Ingestion still *validates* incoming API keys via
139/// whichever `auth_provider` it was built with (including a `DbApiKeyAuthProvider`), it just
140/// does not expose a way to mint/list/revoke them.
141pub async fn serve_ingestion(
142    listen_addr: SocketAddr,
143    lake: DataLakeConnection,
144    auth_provider: Option<Arc<dyn AuthProvider>>,
145    shutdown: impl Future<Output = ()> + Send + 'static,
146    grace: Duration,
147) -> anyhow::Result<()> {
148    use axum::extract::DefaultBodyLimit;
149    use axum::middleware;
150    use axum::routing::get;
151    use micromegas_auth::axum::auth_middleware;
152    use tower_http::limit::RequestBodyLimitLayer;
153
154    use super::axum_utils::observability_middleware;
155    use super::shutdown::serve_axum_with_graceful_shutdown;
156
157    // The deployment's default audience, resolved once at startup so a typo fails
158    // startup rather than surfacing per-request -- matching `LakehouseContext::new`'s handling
159    // of the same variable. Unprefixed (`""`): the ingestion edge must read exactly the
160    // unprefixed `MICROMEGAS_DEFAULT_AUDIENCE` name the lakehouse roles read, never a
161    // `{prefix}_DEFAULT_AUDIENCE` override, or the write side and the read side could be
162    // configured to disagree -- the exact failure this resolution exists to remove. The `?`
163    // here is redundant-by-design, not a reachable error path: `default_audience_from_env`
164    // already validates against the same charset predicate `WriteAudience::new` applies.
165    let default_audience =
166        WriteAudience::new(&micromegas_auth::policy::default_audience_from_env("")?)?;
167    let service = Arc::new(WebIngestionService::new(lake, default_audience));
168
169    let health_router = Router::new()
170        .route("/health", get(|| async { axum::http::StatusCode::OK }))
171        .route("/ready", get(ready_handler))
172        .layer(Extension(service.clone()));
173
174    let firehose_auth = auth_provider.clone();
175    let cw_logs_firehose_auth = auth_provider.clone();
176
177    let mut protected_app = register_routes(Router::new())
178        .merge(super::otlp::otlp_router())
179        .merge(super::webhook::webhook_router())
180        .layer(DefaultBodyLimit::disable())
181        .layer(RequestBodyLimitLayer::new(100 * 1024 * 1024))
182        .layer(Extension(service.clone()));
183
184    let auth_enabled = auth_provider.is_some();
185    if let Some(provider) = auth_provider {
186        info!("Ingestion: authentication enabled");
187        protected_app = protected_app.layer(middleware::from_fn(move |req, next| {
188            auth_middleware(provider.clone(), req, next)
189        }));
190    } else {
191        warn!("Ingestion: authentication disabled — development mode only");
192    }
193
194    // The Firehose routes carry their own auth (Firehose can only send its credential via
195    // X-Amz-Firehose-Access-Key, not Authorization: Bearer), so they are merged outside
196    // protected_app and never hit the global Bearer auth_middleware.
197    let firehose_app = super::firehose::firehose_router(service.clone(), firehose_auth);
198    let cw_logs_firehose_app =
199        super::firehose_cloudwatch_logs::firehose_router(service.clone(), cw_logs_firehose_auth);
200
201    let app = health_router
202        .merge(protected_app)
203        .merge(firehose_app)
204        .merge(cw_logs_firehose_app)
205        .layer(middleware::from_fn(observability_middleware));
206
207    let listener = tokio::net::TcpListener::bind(listen_addr)
208        .await
209        .map_err(|e| anyhow::anyhow!("ingestion: binding to {listen_addr}: {e}"))?;
210    info!("Ingestion serving on {listen_addr} authentication={auth_enabled}");
211
212    serve_axum_with_graceful_shutdown(
213        listener,
214        app.into_make_service_with_connect_info::<SocketAddr>(),
215        shutdown,
216        grace,
217    )
218    .await?;
219
220    Ok(())
221}
micromegas/servers/ingestion.rs

micromegas/servers/
ingestion.rs
