datasynth-server 5.34.0

gRPC and REST server for synthetic data generation
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85

//! Request ID middleware.
//!
//! Generates or preserves a unique request ID for each request.

use axum::{
    body::Body,
    http::{header::HeaderValue, Request},
    middleware::Next,
    response::Response,
};
use uuid::Uuid;

const REQUEST_ID_HEADER: &str = "x-request-id";

/// Request ID middleware.
///
/// If the request already has an `X-Request-Id` header, it is preserved.
/// Otherwise, a new UUID v4 is generated and added to both the request
/// extension and the response headers.
pub async fn request_id_middleware(mut request: Request<Body>, next: Next) -> Response {
    let request_id = request
        .headers()
        .get(REQUEST_ID_HEADER)
        .and_then(|v| v.to_str().ok())
        .map(String::from)
        .unwrap_or_else(|| Uuid::new_v4().to_string());

    // Store in request extensions for logging
    request
        .extensions_mut()
        .insert(RequestId(request_id.clone()));

    let mut response = next.run(request).await;
    if let Ok(header_val) = HeaderValue::try_from(&request_id) {
        response.headers_mut().insert(REQUEST_ID_HEADER, header_val);
    }

    response
}

/// Request ID stored in request extensions.
#[derive(Clone, Debug)]
pub struct RequestId(pub String);

#[cfg(test)]
mod tests {
    use super::*;
    use axum::{routing::get, Router};
    use tower::ServiceExt;

    async fn ok_handler() -> &'static str {
        "ok"
    }

    fn test_router() -> Router {
        Router::new()
            .route("/test", get(ok_handler))
            .layer(axum::middleware::from_fn(request_id_middleware))
    }

    #[tokio::test]
    async fn test_generates_request_id() {
        let router = test_router();
        let request = Request::builder().uri("/test").body(Body::empty()).unwrap();

        let response = router.oneshot(request).await.unwrap();
        assert!(response.headers().get("x-request-id").is_some());
    }

    #[tokio::test]
    async fn test_preserves_client_request_id() {
        let router = test_router();
        let request = Request::builder()
            .uri("/test")
            .header("x-request-id", "client-123")
            .body(Body::empty())
            .unwrap();

        let response = router.oneshot(request).await.unwrap();
        assert_eq!(
            response.headers().get("x-request-id").unwrap(),
            "client-123"
        );
    }
}