runtime: propagate request context through adapters

This commit is contained in:
a.tolmachev
2026-05-01 11:44:43 +00:00
parent 3b1a7c6993
commit 7be6bed347
6 changed files with 530 additions and 24 deletions
+369 -13
View File
@@ -10,7 +10,7 @@ use serde_json::{Map, Value, json};
use crate::{
AdapterResponse, PreparedRequest, ResolvedAuth, RuntimeError, RuntimeOperation,
WindowExecutionResult,
RuntimeRequestContext, WindowExecutionResult,
};
#[derive(Clone, Debug)]
@@ -44,7 +44,8 @@ impl RuntimeExecutor {
operation: &RuntimeOperation,
input: &Value,
) -> Result<Value, RuntimeError> {
self.execute_with_auth(operation, input, None).await
self.execute_with_auth_and_context(operation, input, None, None)
.await
}
pub async fn execute_with_auth(
@@ -52,10 +53,32 @@ impl RuntimeExecutor {
operation: &RuntimeOperation,
input: &Value,
resolved_auth: Option<&ResolvedAuth>,
) -> Result<Value, RuntimeError> {
self.execute_with_auth_and_context(operation, input, resolved_auth, None)
.await
}
pub async fn execute_with_context(
&self,
operation: &RuntimeOperation,
input: &Value,
request_context: Option<&RuntimeRequestContext>,
) -> Result<Value, RuntimeError> {
self.execute_with_auth_and_context(operation, input, None, request_context)
.await
}
pub async fn execute_with_auth_and_context(
&self,
operation: &RuntimeOperation,
input: &Value,
resolved_auth: Option<&ResolvedAuth>,
request_context: Option<&RuntimeRequestContext>,
) -> Result<Value, RuntimeError> {
let prepared_request = self.prepare_request(operation, input)?;
let prepared_request = apply_resolved_auth(prepared_request, resolved_auth);
self.execute_prepared(operation, prepared_request).await
self.execute_prepared(operation, prepared_request, request_context)
.await
}
pub async fn execute_window(
@@ -63,7 +86,8 @@ impl RuntimeExecutor {
operation: &RuntimeOperation,
input: &Value,
) -> Result<WindowExecutionResult, RuntimeError> {
self.execute_window_with_auth(operation, input, None).await
self.execute_window_with_auth_and_context(operation, input, None, None)
.await
}
pub async fn execute_window_with_auth(
@@ -71,6 +95,27 @@ impl RuntimeExecutor {
operation: &RuntimeOperation,
input: &Value,
resolved_auth: Option<&ResolvedAuth>,
) -> Result<WindowExecutionResult, RuntimeError> {
self.execute_window_with_auth_and_context(operation, input, resolved_auth, None)
.await
}
pub async fn execute_window_with_context(
&self,
operation: &RuntimeOperation,
input: &Value,
request_context: Option<&RuntimeRequestContext>,
) -> Result<WindowExecutionResult, RuntimeError> {
self.execute_window_with_auth_and_context(operation, input, None, request_context)
.await
}
pub async fn execute_window_with_auth_and_context(
&self,
operation: &RuntimeOperation,
input: &Value,
resolved_auth: Option<&ResolvedAuth>,
request_context: Option<&RuntimeRequestContext>,
) -> Result<WindowExecutionResult, RuntimeError> {
let Some(streaming) = operation.execution_config.streaming.as_ref() else {
return Err(RuntimeError::MissingStreamingConfig {
@@ -94,10 +139,11 @@ impl RuntimeExecutor {
operation.target,
Target::Rest(_) | Target::Grpc(_) | Target::Websocket(_)
) {
self.execute_window_adapter(operation, prepared_request)
self.execute_window_adapter(operation, prepared_request, request_context)
.await?
} else {
self.execute_adapter(operation, prepared_request).await?
self.execute_adapter(operation, prepared_request, request_context)
.await?
};
crate::aggregation::collect_window_result(&adapter_response.body, streaming)
@@ -108,7 +154,7 @@ impl RuntimeExecutor {
operation: &RuntimeOperation,
input: &Value,
) -> Result<WindowExecutionResult, RuntimeError> {
self.execute_session_seed_with_auth(operation, input, None)
self.execute_session_seed_with_auth_and_context(operation, input, None, None)
.await
}
@@ -117,6 +163,27 @@ impl RuntimeExecutor {
operation: &RuntimeOperation,
input: &Value,
resolved_auth: Option<&ResolvedAuth>,
) -> Result<WindowExecutionResult, RuntimeError> {
self.execute_session_seed_with_auth_and_context(operation, input, resolved_auth, None)
.await
}
pub async fn execute_session_seed_with_context(
&self,
operation: &RuntimeOperation,
input: &Value,
request_context: Option<&RuntimeRequestContext>,
) -> Result<WindowExecutionResult, RuntimeError> {
self.execute_session_seed_with_auth_and_context(operation, input, None, request_context)
.await
}
pub async fn execute_session_seed_with_auth_and_context(
&self,
operation: &RuntimeOperation,
input: &Value,
resolved_auth: Option<&ResolvedAuth>,
request_context: Option<&RuntimeRequestContext>,
) -> Result<WindowExecutionResult, RuntimeError> {
let Some(streaming) = operation.execution_config.streaming.as_ref() else {
return Err(RuntimeError::MissingStreamingConfig {
@@ -144,8 +211,13 @@ impl RuntimeExecutor {
config.max_items = Some(seed_limit);
}
self.execute_window_with_auth(&seeded_operation, input, resolved_auth)
.await
self.execute_window_with_auth_and_context(
&seeded_operation,
input,
resolved_auth,
request_context,
)
.await
}
pub fn prepare_request(
@@ -163,8 +235,11 @@ impl RuntimeExecutor {
&self,
operation: &RuntimeOperation,
prepared_request: PreparedRequest,
request_context: Option<&RuntimeRequestContext>,
) -> Result<Value, RuntimeError> {
let adapter_response = self.execute_adapter(operation, prepared_request).await?;
let adapter_response = self
.execute_adapter(operation, prepared_request, request_context)
.await?;
let finalized_output = finalize_output(operation, &adapter_response)?;
operation.output_schema.validate_shape(&finalized_output)?;
@@ -176,7 +251,9 @@ impl RuntimeExecutor {
&self,
operation: &RuntimeOperation,
prepared_request: PreparedRequest,
request_context: Option<&RuntimeRequestContext>,
) -> Result<AdapterResponse, RuntimeError> {
let context_headers = runtime_context_headers(request_context);
match &operation.target {
Target::Grpc(target) => {
let request = GrpcRequest {
@@ -184,6 +261,7 @@ impl RuntimeExecutor {
&BTreeMap::new(),
&operation.execution_config.headers,
&prepared_request.headers,
&context_headers,
),
body: prepared_request
.grpc
@@ -206,6 +284,7 @@ impl RuntimeExecutor {
&BTreeMap::new(),
&operation.execution_config.headers,
&prepared_request.headers,
&context_headers,
),
variables: prepared_request.variables.clone(),
timeout_ms: operation.execution_config.timeout_ms,
@@ -227,6 +306,7 @@ impl RuntimeExecutor {
&target.static_headers,
&operation.execution_config.headers,
&prepared_request.headers,
&context_headers,
),
body: prepared_request.body.clone(),
timeout_ms: operation.execution_config.timeout_ms,
@@ -246,6 +326,7 @@ impl RuntimeExecutor {
&BTreeMap::new(),
&operation.execution_config.headers,
&prepared_request.headers,
&context_headers,
),
body: prepared_request
.body
@@ -273,7 +354,9 @@ impl RuntimeExecutor {
&self,
operation: &RuntimeOperation,
prepared_request: PreparedRequest,
request_context: Option<&RuntimeRequestContext>,
) -> Result<AdapterResponse, RuntimeError> {
let context_headers = runtime_context_headers(request_context);
match &operation.target {
Target::Rest(target) => {
let Some(streaming) = operation.execution_config.streaming.as_ref() else {
@@ -289,6 +372,7 @@ impl RuntimeExecutor {
&target.static_headers,
&operation.execution_config.headers,
&prepared_request.headers,
&context_headers,
),
body: prepared_request.body.clone(),
timeout_ms: streaming
@@ -319,6 +403,7 @@ impl RuntimeExecutor {
&BTreeMap::new(),
&operation.execution_config.headers,
&prepared_request.headers,
&context_headers,
),
body: prepared_request
.grpc
@@ -356,6 +441,7 @@ impl RuntimeExecutor {
&BTreeMap::new(),
&operation.execution_config.headers,
&prepared_request.headers,
&context_headers,
),
window_duration_ms: streaming.window_duration_ms.unwrap_or_default(),
max_items: streaming.max_items.map(|value| value.saturating_add(1)),
@@ -384,7 +470,10 @@ impl RuntimeExecutor {
operation_id: operation.operation_id.as_str().to_owned(),
mode: ExecutionMode::Window,
}),
_ => self.execute_adapter(operation, prepared_request).await,
_ => {
self.execute_adapter(operation, prepared_request, request_context)
.await
}
}
}
}
@@ -444,13 +533,23 @@ fn merge_headers(
static_headers: &BTreeMap<String, String>,
execution_headers: &BTreeMap<String, String>,
request_headers: &BTreeMap<String, String>,
context_headers: &BTreeMap<String, String>,
) -> BTreeMap<String, String> {
let mut headers = static_headers.clone();
headers.extend(execution_headers.clone());
headers.extend(request_headers.clone());
headers.extend(context_headers.clone());
headers
}
fn runtime_context_headers(
request_context: Option<&RuntimeRequestContext>,
) -> BTreeMap<String, String> {
request_context
.map(RuntimeRequestContext::outbound_headers)
.unwrap_or_default()
}
fn read_string_map(
value: Option<&Value>,
field_name: &str,
@@ -510,6 +609,7 @@ mod tests {
use axum::{
Json, Router,
http::HeaderMap,
response::sse::{Event, KeepAlive, Sse},
routing::post,
};
@@ -528,9 +628,14 @@ mod tests {
use serde_json::{Value, json};
use time::{OffsetDateTime, format_description::well_known::Rfc3339};
use tokio::net::TcpListener;
use tokio_tungstenite::accept_async;
use tokio_tungstenite::{
accept_async, accept_hdr_async,
tungstenite::handshake::server::{Request, Response},
};
use crate::{PreparedRequest, RuntimeError, RuntimeExecutor, RuntimeOperation};
use crate::{
PreparedRequest, RuntimeError, RuntimeExecutor, RuntimeOperation, RuntimeRequestContext,
};
fn timestamp(value: &str) -> OffsetDateTime {
OffsetDateTime::parse(value, &Rfc3339).unwrap()
@@ -550,6 +655,29 @@ mod tests {
assert_eq!(output, json!({ "id": "lead_123" }));
}
#[tokio::test]
async fn propagates_request_context_to_rest_headers() {
let base_url = spawn_runtime_server().await;
let executor = RuntimeExecutor::new();
let mut operation = test_rest_operation(&base_url, false, false);
let context = RuntimeRequestContext::new("req_rest_123", "corr_rest_123");
if let Target::Rest(target) = &mut operation.target {
target.path_template = "/leads-context".to_owned();
}
let output = executor
.execute_with_context(
&operation,
&json!({ "email": "user@example.com" }),
Some(&context),
)
.await
.unwrap();
assert_eq!(output, json!({ "id": "req_rest_123|corr_rest_123" }));
}
#[tokio::test]
async fn executes_graphql_operation_end_to_end() {
let endpoint = spawn_graphql_server().await;
@@ -564,6 +692,25 @@ mod tests {
assert_eq!(output, json!({ "id": "lead_123" }));
}
#[tokio::test]
async fn propagates_request_context_to_graphql_headers() {
let endpoint = spawn_graphql_context_server().await;
let executor = RuntimeExecutor::new();
let operation = test_graphql_operation(&endpoint, false);
let context = RuntimeRequestContext::new("req_graphql_123", "corr_graphql_123");
let output = executor
.execute_with_context(
&operation,
&json!({ "email": "user@example.com" }),
Some(&context),
)
.await
.unwrap();
assert_eq!(output, json!({ "id": "req_graphql_123|corr_graphql_123" }));
}
#[tokio::test]
async fn executes_grpc_operation_end_to_end() {
let server_addr = grpc_test_support::spawn_unary_echo_server().await;
@@ -578,6 +725,24 @@ mod tests {
assert_eq!(output, json!({ "message": "hello" }));
}
#[tokio::test]
async fn propagates_request_context_to_grpc_metadata() {
let server_addr = grpc_test_support::spawn_metadata_echo_server().await;
let executor = RuntimeExecutor::new();
let operation = test_grpc_operation(&server_addr);
let context = RuntimeRequestContext::new("req_grpc_123", "corr_grpc_123");
let output = executor
.execute_with_context(&operation, &json!({ "message": "hello" }), Some(&context))
.await
.unwrap();
assert_eq!(
output,
json!({ "message": "hello|req_grpc_123|corr_grpc_123" })
);
}
#[tokio::test]
async fn executes_soap_operation_end_to_end() {
let endpoint = spawn_soap_server().await;
@@ -592,6 +757,25 @@ mod tests {
assert_eq!(output, json!({ "id": "lead_123" }));
}
#[tokio::test]
async fn propagates_request_context_to_soap_headers() {
let endpoint = spawn_soap_context_server().await;
let executor = RuntimeExecutor::new();
let operation = test_soap_operation(&endpoint);
let context = RuntimeRequestContext::new("req_soap_123", "corr_soap_123");
let output = executor
.execute_with_context(
&operation,
&json!({ "email": "user@example.com" }),
Some(&context),
)
.await
.unwrap();
assert_eq!(output, json!({ "id": "req_soap_123|corr_soap_123" }));
}
#[tokio::test]
async fn executes_grpc_window_mode_with_server_stream() {
let server_addr = grpc_test_support::spawn_unary_echo_server().await;
@@ -632,6 +816,23 @@ mod tests {
assert!(result.has_more);
}
#[tokio::test]
async fn propagates_request_context_to_websocket_headers() {
let target_url = spawn_request_context_websocket_server().await;
let executor = RuntimeExecutor::new();
let operation =
test_websocket_window_operation(&target_url, AggregationMode::RawItems, Some(2));
let context = RuntimeRequestContext::new("req_ws_123", "corr_ws_123");
let result = executor
.execute_window_with_context(&operation, &json!({}), Some(&context))
.await
.unwrap();
assert_eq!(result.items.len(), 2);
assert_eq!(result.items[0], json!({ "seq": 1, "value": 101 }));
}
#[tokio::test]
async fn rejects_invalid_input_shape() {
let base_url = spawn_runtime_server().await;
@@ -807,6 +1008,7 @@ mod tests {
async fn spawn_runtime_server() -> String {
let app = Router::new()
.route("/leads", post(create_lead))
.route("/leads-context", post(create_lead_with_context))
.route("/events", post(events_window))
.route("/slow-events", post(slow_events_window))
.route("/events-sse", post(events_stream))
@@ -833,6 +1035,18 @@ mod tests {
format!("http://{}", address)
}
async fn spawn_graphql_context_server() -> String {
let app = Router::new().route("/", post(graphql_context_handler));
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
tokio::spawn(async move {
axum::serve(listener, app).await.unwrap();
});
format!("http://{}", address)
}
async fn spawn_soap_server() -> String {
let app = Router::new().route("/", post(soap_handler));
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
@@ -845,6 +1059,18 @@ mod tests {
format!("http://{}", address)
}
async fn spawn_soap_context_server() -> String {
let app = Router::new().route("/", post(soap_context_handler));
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
tokio::spawn(async move {
axum::serve(listener, app).await.unwrap();
});
format!("http://{}", address)
}
async fn spawn_websocket_server() -> String {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
@@ -883,6 +1109,61 @@ mod tests {
format!("ws://{}", address)
}
async fn spawn_request_context_websocket_server() -> String {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
tokio::spawn(async move {
loop {
let (stream, _) = listener.accept().await.unwrap();
tokio::spawn(async move {
let websocket =
accept_hdr_async(stream, |request: &Request, response: Response| {
assert_eq!(
request
.headers()
.get("x-request-id")
.and_then(|value| value.to_str().ok()),
Some("req_ws_123")
);
assert_eq!(
request
.headers()
.get("x-correlation-id")
.and_then(|value| value.to_str().ok()),
Some("corr_ws_123")
);
Ok(response)
})
.await
.unwrap();
let (mut sink, mut source) = websocket.split();
if let Some(message) = source.next().await {
let message = message.unwrap();
if !matches!(message, tokio_tungstenite::tungstenite::Message::Text(_)) {
return;
}
}
for payload in [
json!({ "seq": 1, "value": 101 }),
json!({ "seq": 2, "value": 102 }),
] {
sink.send(tokio_tungstenite::tungstenite::Message::Text(
payload.to_string().into(),
))
.await
.unwrap();
}
let _ = sink.close().await;
});
}
});
format!("ws://{}", address)
}
async fn create_lead(Json(payload): Json<Value>) -> (axum::http::StatusCode, Json<Value>) {
let should_fail = payload
.get("fail")
@@ -902,6 +1183,28 @@ mod tests {
)
}
async fn create_lead_with_context(
headers: HeaderMap,
Json(_payload): Json<Value>,
) -> (axum::http::StatusCode, Json<Value>) {
(
axum::http::StatusCode::OK,
Json(json!({
"id": format!(
"{}|{}",
headers
.get("x-request-id")
.and_then(|value| value.to_str().ok())
.unwrap_or_default(),
headers
.get("x-correlation-id")
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
),
})),
)
}
async fn events_window() -> (axum::http::StatusCode, Json<Value>) {
(
axum::http::StatusCode::OK,
@@ -970,6 +1273,30 @@ mod tests {
}))
}
async fn graphql_context_handler(
headers: HeaderMap,
Json(_payload): Json<Value>,
) -> Json<Value> {
Json(json!({
"data": {
"createLead": {
"id": format!(
"{}|{}",
headers
.get("x-request-id")
.and_then(|value| value.to_str().ok())
.unwrap_or_default(),
headers
.get("x-correlation-id")
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
),
"status": "created"
}
}
}))
}
async fn soap_handler(body: String) -> (axum::http::StatusCode, String) {
assert!(body.contains("<email>user@example.com</email>"));
@@ -987,6 +1314,35 @@ mod tests {
)
}
async fn soap_context_handler(
headers: HeaderMap,
body: String,
) -> (axum::http::StatusCode, String) {
assert!(body.contains("<email>user@example.com</email>"));
let request_id = headers
.get("x-request-id")
.and_then(|value| value.to_str().ok())
.unwrap_or_default();
let correlation_id = headers
.get("x-correlation-id")
.and_then(|value| value.to_str().ok())
.unwrap_or_default();
(
axum::http::StatusCode::OK,
format!(
r#"<soap:Envelope xmlns:soap="http://schemas.xmlsoap.org/soap/envelope/">
<soap:Body>
<CreateLeadResponse>
<id>{request_id}|{correlation_id}</id>
<status>created</status>
</CreateLeadResponse>
</soap:Body>
</soap:Envelope>"#
),
)
}
fn test_rest_operation(
base_url: &str,
should_fail: bool,