runtime: propagate request context through adapters
This commit is contained in:
@@ -33,7 +33,8 @@ use crank_registry::{
|
||||
UsageSummary, UsageTimelinePoint, WorkspaceMembershipRecord, WorkspaceRecord,
|
||||
};
|
||||
use crank_runtime::{
|
||||
PreparedRequest, ResolvedAuth, RuntimeError, RuntimeExecutor, RuntimeOperation, SecretCrypto,
|
||||
PreparedRequest, ResolvedAuth, RuntimeError, RuntimeExecutor, RuntimeOperation,
|
||||
RuntimeRequestContext, SecretCrypto,
|
||||
};
|
||||
use crank_schema::Schema;
|
||||
use serde::{Deserialize, Serialize};
|
||||
@@ -2096,6 +2097,7 @@ impl AdminService {
|
||||
payload: TestRunPayload,
|
||||
request_id: &str,
|
||||
) -> Result<TestRunResult, ApiError> {
|
||||
let runtime_request_context = RuntimeRequestContext::from_request_id(request_id);
|
||||
let record = self
|
||||
.get_operation_version(workspace_id, operation_id, payload.version)
|
||||
.await?;
|
||||
@@ -2149,12 +2151,22 @@ impl AdminService {
|
||||
Ok(resolved_auth) => match mode {
|
||||
ExecutionMode::Unary => self
|
||||
.runtime
|
||||
.execute_with_auth(&runtime, &payload.input, resolved_auth.as_ref())
|
||||
.execute_with_auth_and_context(
|
||||
&runtime,
|
||||
&payload.input,
|
||||
resolved_auth.as_ref(),
|
||||
Some(&runtime_request_context),
|
||||
)
|
||||
.await
|
||||
.map(|output| TestRunOutcome::Unary { output }),
|
||||
ExecutionMode::Window => self
|
||||
.runtime
|
||||
.execute_window_with_auth(&runtime, &payload.input, resolved_auth.as_ref())
|
||||
.execute_window_with_auth_and_context(
|
||||
&runtime,
|
||||
&payload.input,
|
||||
resolved_auth.as_ref(),
|
||||
Some(&runtime_request_context),
|
||||
)
|
||||
.await
|
||||
.map(|output| TestRunOutcome::Window { output }),
|
||||
ExecutionMode::Session => self
|
||||
@@ -2164,11 +2176,18 @@ impl AdminService {
|
||||
&runtime,
|
||||
&payload.input,
|
||||
resolved_auth.as_ref(),
|
||||
&runtime_request_context,
|
||||
)
|
||||
.await
|
||||
.map(|output| TestRunOutcome::Session { output }),
|
||||
ExecutionMode::AsyncJob => self
|
||||
.start_async_job_test(workspace_id, &record.snapshot, &runtime, &payload.input)
|
||||
.start_async_job_test(
|
||||
workspace_id,
|
||||
&record.snapshot,
|
||||
&runtime,
|
||||
&payload.input,
|
||||
&runtime_request_context,
|
||||
)
|
||||
.await
|
||||
.map(|output| TestRunOutcome::AsyncJob { output }),
|
||||
},
|
||||
@@ -2272,6 +2291,7 @@ impl AdminService {
|
||||
runtime: &RuntimeOperation,
|
||||
input: &Value,
|
||||
resolved_auth: Option<&ResolvedAuth>,
|
||||
request_context: &RuntimeRequestContext,
|
||||
) -> Result<StreamSessionStartView, RuntimeError> {
|
||||
let streaming = runtime.execution_config.streaming.as_ref().ok_or_else(|| {
|
||||
RuntimeError::MissingStreamingConfig {
|
||||
@@ -2280,7 +2300,12 @@ impl AdminService {
|
||||
})?;
|
||||
let seed = self
|
||||
.runtime
|
||||
.execute_session_seed_with_auth(runtime, input, resolved_auth)
|
||||
.execute_session_seed_with_auth_and_context(
|
||||
runtime,
|
||||
input,
|
||||
resolved_auth,
|
||||
Some(request_context),
|
||||
)
|
||||
.await?;
|
||||
let batch_size = streaming.max_items.unwrap_or(10).max(1) as usize;
|
||||
let preview_count = seed.items.len().min(batch_size);
|
||||
@@ -2342,6 +2367,7 @@ impl AdminService {
|
||||
operation: &RegistryOperation,
|
||||
runtime: &RuntimeOperation,
|
||||
input: &Value,
|
||||
request_context: &RuntimeRequestContext,
|
||||
) -> Result<AsyncJobStartView, RuntimeError> {
|
||||
let created_at = OffsetDateTime::now_utc();
|
||||
let streaming = runtime.execution_config.streaming.as_ref().ok_or_else(|| {
|
||||
@@ -2385,6 +2411,7 @@ impl AdminService {
|
||||
let workspace_for_task = workspace_id.clone();
|
||||
let operation_for_task = runtime.clone();
|
||||
let input_for_task = input.clone();
|
||||
let request_context_for_task = request_context.clone();
|
||||
let job_id = job.id.clone();
|
||||
tokio::spawn(async move {
|
||||
let resolved_auth = resolve_runtime_auth_for_task(
|
||||
@@ -2397,10 +2424,11 @@ impl AdminService {
|
||||
let result = match resolved_auth {
|
||||
Ok(resolved_auth) => {
|
||||
runtime_for_task
|
||||
.execute_with_auth(
|
||||
.execute_with_auth_and_context(
|
||||
&operation_for_task,
|
||||
&input_for_task,
|
||||
resolved_auth.as_ref(),
|
||||
Some(&request_context_for_task),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user