use std::collections::hash_map::Entry;
use std::collections::{BTreeSet, HashMap};
use std::sync::Arc;
use std::vec;
use crate::query_planner::ast::operation::SubgraphFetchOperation;
use crate::query_planner::planner::plan_nodes::{CustomScalarPaths, FetchNode, FlattenNode};
use crate::query_planner::planner::query_plan::QUERY_PLAN_KIND;
use crate::query_planner::{
ast::operation::OperationDefinition,
planner::plan_nodes::{
ConditionNode, EntityBatch, EntityBatchAlias, FetchRewrite, FlattenNodePath, PlanNode,
QueryPlan, SequenceNode,
},
state::supergraph_state::OperationKind,
};
use crate::telemetry::logging::{summary, targets};
use crate::telemetry::metrics::graphql_metrics::GraphQLErrorMetricsRecorder;
use crate::telemetry::traces::spans::graphql::{
GraphQLOperationSpan, GraphQLSpanOperationIdentity, GraphQLSubgraphOperationSpan,
};
use crate::telemetry::traces::spans::observed_error::ObservedError;
use ahash::{HashMap as AHashMap, HashMapExt};
use bytes::BufMut;
use futures::TryFutureExt;
use futures::{
future::BoxFuture,
stream::{BoxStream, FuturesUnordered},
FutureExt, StreamExt,
};
use http::{HeaderMap, StatusCode};
use serde::Serialize;
use sonic_rs::{JsonValueTrait, ValueRef};
use tracing::{debug, error, warn, Instrument};
use crate::executor::execution::client_request_details::OperationDetails;
use crate::executor::execution::demand_control::DemandControlExecutionContext;
use crate::executor::execution::error_masking::ErrorMaskingRuntime;
use crate::executor::execution::operation_name::OperationNameFactory;
use crate::executor::executors::common::ConnectionFingerprint;
use crate::executor::headers::cache_control;
use crate::executor::{
execution::{
client_request_details::ClientRequestDetails,
error::{IntoPlanExecutionError, LazyPlanContext, PlanExecutionError},
jwt_forward::JwtAuthForwardingPlan,
rewrites::FetchRewriteExt,
},
execution_context::ExecutionContext,
executors::{common::SubgraphExecutionRequest, map::SubgraphExecutorMap},
extensions::{aggregator::apply_subgraph_extensions, plan::ExtensionsPlan},
headers::{
plan::HeaderRulesPlan,
request::modify_subgraph_request_headers,
response::{apply_subgraph_response_headers, ResponseHeaderAggregator, ResponseHeaderSink},
},
hooks::{
on_execute::{
DemandControlCost, DemandControlEstimatedCost, OnExecuteEndHookPayload,
OnExecuteResponse, OnExecuteStartHookPayload,
},
on_graphql_error::handle_graphql_errors_with_plugins,
},
introspection::{
resolve::{resolve_introspection, IntrospectionContext},
schema::SchemaMetadata,
},
plugin_context::PluginRequestState,
plugin_trait::{EarlyHTTPResponse, EndControlFlow, StartControlFlow},
plugins::hooks,
projection::{
plan::FieldProjectionPlan, request::project_requires, response::project_by_operation,
},
response::{
graphql_error::{GraphQLError, GraphQLErrorPath, GraphQLErrorPathSegment},
merge::deep_merge,
subgraph_response::SubgraphResponse,
value::Value,
},
utils::{
consts::{CLOSE_BRACKET, OPEN_BRACKET},
traverse::{traverse_and_callback, traverse_and_callback_mut},
},
};
pub type VariablesMap = HashMap<String, sonic_rs::Value>;
#[derive(Debug, Default, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct ExecutionResultExtensions<'exec> {
#[serde(skip_serializing_if = "Option::is_none")]
pub query_plan: Option<&'exec QueryPlan>,
#[serde(flatten)]
pub extensions: HashMap<String, sonic_rs::Value>,
}
impl ExecutionResultExtensions<'_> {
pub fn is_empty(&self) -> bool {
self.query_plan.is_none() && self.extensions.is_empty()
}
}
#[derive(Clone, Debug, Default)]
pub struct CoerceVariablesPayload {
pub variables_map: Option<VariablesMap>,
}
impl CoerceVariablesPayload {
pub fn variable_equals_true(&self, name: &str) -> bool {
self.variables_map
.as_ref()
.and_then(|vars| vars.get(name))
.and_then(|value| value.as_bool())
.unwrap_or(false)
}
}
pub struct QueryPlanExecutionOpts<'exec> {
pub query_plan: &'exec QueryPlan,
pub operation_for_plan: Arc<OperationDefinition>,
pub projection_plan: Arc<Vec<FieldProjectionPlan>>,
pub headers_plan: Arc<HeaderRulesPlan>,
pub extensions_plan: Arc<ExtensionsPlan>,
pub variable_values: Arc<CoerceVariablesPayload>,
pub extensions: ExecutionResultExtensions<'exec>,
pub client_request: Arc<ClientRequestDetails<'exec>>,
pub introspection_context: Arc<IntrospectionContext>,
pub operation_type_name: String,
pub operation_kind: OperationKind,
pub executors: Arc<SubgraphExecutorMap>,
pub jwt_auth_forwarding: Option<Arc<JwtAuthForwardingPlan>>,
pub graphql_error_recorder: Option<GraphQLErrorMetricsRecorder>,
pub demand_control_context: Option<Arc<DemandControlExecutionContext>>,
pub initial_errors: Vec<GraphQLError>,
pub span: GraphQLOperationSpan,
pub plugin_req_state: Option<PluginRequestState<'exec>>,
pub operation_name_factory: OperationNameFactory,
pub response_header_sink: ResponseHeaderSink,
pub error_masking_runtime: Arc<Option<ErrorMaskingRuntime>>,
pub connection_fingerprint: Option<ConnectionFingerprint>,
}
pub struct PlanSubscriptionOutput {
pub body: BoxStream<'static, Vec<u8>>,
pub error_count: usize,
}
pub enum QueryPlanExecutionResult {
Single(PlanExecutionOutput),
Stream(PlanSubscriptionOutput),
}
#[derive(Default)]
pub struct PlanExecutionOutput {
pub body: Vec<u8>,
pub error_count: usize,
pub status_code: StatusCode,
}
#[derive(Serialize)]
pub struct FailedExecutionResult {
pub errors: Vec<GraphQLError>,
}
impl FailedExecutionResult {
pub fn serialize(&self) -> Vec<u8> {
sonic_rs::to_vec(&self).unwrap_or_else(|err| {
error!(target: targets::CORE, error = ?err, "Failed to serialize pipeline error to response");
sonic_rs::to_vec(&FailedExecutionResult {
errors: vec![GraphQLError::from_message_and_code(
"Failed to serialize error response",
"INTERNAL_SERVER_ERROR",
)],
})
.unwrap()
})
}
}
#[inline]
fn serialize_masked_failure(
mut error: GraphQLError,
error_masking_runtime: &Option<ErrorMaskingRuntime>,
) -> Vec<u8> {
if let Some(runtime) = error_masking_runtime {
runtime.apply(&mut error);
}
FailedExecutionResult {
errors: vec![error],
}
.serialize()
}
fn early_http_response_into_execution_output(
response: EarlyHTTPResponse,
response_header_sink: &ResponseHeaderSink,
) -> PlanExecutionOutput {
if !response.headers.is_empty() {
response_header_sink.store(ResponseHeaderAggregator::from_http_headers(
&response.headers,
));
}
PlanExecutionOutput {
body: response.body,
error_count: 0,
status_code: response.status_code,
}
}
#[inline]
fn log_plan_execution_error(error: &PlanExecutionError) {
if let Some(subgraph_name) = error.subgraph_name() {
error!(
target: targets::GRAPHQL_EXECUTION,
subgraph = ?subgraph_name,
error = ?error,
"Error executing step",
);
} else {
error!(
target: targets::GRAPHQL_EXECUTION,
error = ?error,
"Error executing step",
);
}
}
pub async fn execute_query_plan<'exec>(
opts: QueryPlanExecutionOpts<'exec>,
) -> Result<QueryPlanExecutionResult, PlanExecutionError> {
let (subscription_node, remaining_nodes) = match &opts.query_plan.node {
Some(PlanNode::Subscription(sub)) => (Some(sub), None),
Some(PlanNode::Sequence(seq)) => match seq.nodes.first() {
Some(PlanNode::Subscription(sub)) => {
let remaining = if seq.nodes.len() > 1 {
Some(seq.nodes[1..].to_vec())
} else {
None
};
(Some(sub), remaining)
}
_ => (None, None),
},
_ => (None, None),
};
if let Some(sub) = subscription_node {
let fetch_node = &sub.primary;
let query_plan: Arc<QueryPlan> = Arc::new(QueryPlan {
kind: QUERY_PLAN_KIND,
node: remaining_nodes.map(|nodes| {
if nodes.len() == 1 {
nodes.into_iter().next().unwrap()
} else {
PlanNode::Sequence(SequenceNode { nodes })
}
}),
});
let mut headers_map = HeaderMap::new();
modify_subgraph_request_headers(
&opts.headers_plan,
&fetch_node.service_name,
&opts.client_request,
&mut headers_map,
)
.with_plan_context(LazyPlanContext {
subgraph_name: || Some(fetch_node.service_name.to_string()),
affected_path: || None,
})?;
let variable_refs = select_fetch_variables(
&opts.variable_values.variables_map,
fetch_node.variable_usages.as_ref(),
);
let mut subgraph_request = SubgraphExecutionRequest {
query: fetch_node.operation.document_str.as_str(),
document_name_write_pos: fetch_node.operation.name_write_position,
dedupe: false,
operation_name: opts
.operation_name_factory
.generate(&fetch_node.service_name, fetch_node.id),
variables: variable_refs,
headers: headers_map,
raw_variable_values: None,
extensions: None,
custom_scalar_paths: fetch_node.custom_scalar_paths.as_ref(),
connection_fingerprint: opts.connection_fingerprint,
};
if let Some(jwt_forwarding_plan) = &opts.jwt_auth_forwarding {
subgraph_request.add_request_extensions_field(
jwt_forwarding_plan.extension_field_name.clone(),
jwt_forwarding_plan.extension_field_value.clone(),
);
}
let mut response_stream = opts
.executors
.subscribe(
&fetch_node.service_name,
subgraph_request,
&opts.client_request,
)
.await
.with_plan_context(LazyPlanContext {
subgraph_name: || Some(fetch_node.service_name.to_string()),
affected_path: || None,
})?;
let subgraph_name: String = fetch_node.service_name.clone();
let client_method = opts.client_request.method.clone();
let client_url = opts.client_request.url.clone();
let client_headers = opts.client_request.headers.clone();
let client_operation_name = opts.client_request.operation.name.map(|s| s.to_string());
let client_operation_query = opts.client_request.operation.query.to_string();
let client_operation_kind = opts.client_request.operation.kind;
let client_jwt = opts.client_request.jwt.clone();
let client_path_params = opts.client_request.path_params.into_owned();
let response_header_sink = opts.response_header_sink.clone();
let error_masking_runtime = opts.error_masking_runtime.clone();
let operation_name_factory = opts.operation_name_factory.clone();
let body_stream = Box::pin(async_stream::stream! {
while let Some(stream_result) = response_stream.next().await {
let response = match stream_result.with_plan_context(LazyPlanContext {
subgraph_name: || Some(subgraph_name.to_string()),
affected_path: || None,
}) {
Ok(response) => response,
Err(ref err) => {
log_plan_execution_error(err);
yield serialize_masked_failure(err.into(), error_masking_runtime.as_ref());
return;
}
};
let mut initial_errors = opts.initial_errors.clone();
if let Some(new_errors) = response.errors {
initial_errors.extend(new_errors);
}
let opts = QueryPlanExecutionOpts {
query_plan: &query_plan,
operation_for_plan: opts.operation_for_plan.clone(),
projection_plan: opts.projection_plan.clone(),
headers_plan: opts.headers_plan.clone(),
extensions_plan: opts.extensions_plan.clone(),
variable_values: opts.variable_values.clone(),
extensions: ExecutionResultExtensions::default(),
client_request: ClientRequestDetails {
method: &client_method,
url: &client_url,
headers: client_headers.clone(),
operation: OperationDetails {
query: &client_operation_query,
name: client_operation_name.as_deref(),
kind: client_operation_kind,
},
jwt: client_jwt.clone(),
path_params: client_path_params.clone(),
}.into(),
introspection_context: opts.introspection_context.clone(),
operation_type_name: opts.operation_type_name.clone(),
operation_kind: opts.operation_kind.clone(),
executors: opts.executors.clone(),
jwt_auth_forwarding: opts.jwt_auth_forwarding.clone(),
initial_errors,
span: GraphQLOperationSpan { span: opts.span.clone() },
plugin_req_state: None,
graphql_error_recorder: None,
operation_name_factory: operation_name_factory.clone(),
demand_control_context: opts.demand_control_context.clone(),
response_header_sink: response_header_sink.clone(),
error_masking_runtime: opts.error_masking_runtime.clone(),
connection_fingerprint: opts.connection_fingerprint,
};
match execute_query_plan_with_data(response.data, opts).await {
Ok(result) => yield result.body,
Err(ref err) => {
log_plan_execution_error(err);
yield serialize_masked_failure(err.into(), error_masking_runtime.as_ref());
return;
}
}
}
});
return Ok(QueryPlanExecutionResult::Stream(PlanSubscriptionOutput {
body: body_stream,
error_count: 0, }));
}
let introspection_context_clone = Arc::clone(&opts.introspection_context);
let data = if let Some(introspection_query) = &introspection_context_clone.query {
resolve_introspection(introspection_query, &introspection_context_clone)
} else if opts.projection_plan.is_empty() {
Value::Null
} else {
Value::Object(Vec::new())
};
let output = execute_query_plan_with_data(data, opts).await?;
Ok(QueryPlanExecutionResult::Single(output))
}
async fn execute_query_plan_with_data<'exec>(
mut data: Value<'exec>,
mut opts: QueryPlanExecutionOpts<'exec>,
) -> Result<PlanExecutionOutput, PlanExecutionError> {
let mut errors = opts.initial_errors;
let dedupe_subgraph_requests = opts.operation_kind.is_query();
let mut on_end_callbacks = vec![];
if let Some(plugin_req_state) = opts.plugin_req_state.as_ref() {
let mut start_payload = OnExecuteStartHookPayload {
router_http_request: &plugin_req_state.router_http_request,
context: &plugin_req_state.context,
request_context: plugin_req_state
.request_context
.for_plugin::<hooks::OnExecute>(),
query_plan: opts.query_plan,
operation_for_plan: &opts.operation_for_plan,
data,
errors,
extensions: opts.extensions.extensions,
variable_values: &opts.variable_values.variables_map,
dedupe_subgraph_requests,
demand_control_estimate: opts.demand_control_context.as_ref().map(|dc| {
DemandControlEstimatedCost {
estimated: dc.evaluation.estimated_cost,
max: dc.operation.operation_max_cost,
}
}),
};
for plugin in plugin_req_state.plugins.as_ref() {
let result = plugin.on_execute(start_payload).await;
start_payload = result.payload;
match result.control_flow {
StartControlFlow::Proceed => { }
StartControlFlow::EndWithResponse(response) => match response {
OnExecuteResponse::Output(response) => return Ok(response),
OnExecuteResponse::EarlyResponse(response) => {
return Ok(early_http_response_into_execution_output(
response,
&opts.response_header_sink,
));
}
},
StartControlFlow::OnEnd(callback) => {
on_end_callbacks.push(callback);
}
}
}
data = start_payload.data;
errors = start_payload.errors;
opts.extensions.extensions = start_payload.extensions;
}
let mut exec_ctx = ExecutionContext::new(data, errors);
let executor = Executor {
variable_values: &opts.variable_values.variables_map,
schema_metadata: &opts.introspection_context.metadata,
executors: &opts.executors,
client_request: &opts.client_request,
headers_plan: &opts.headers_plan,
extensions_plan: &opts.extensions_plan,
jwt_forwarding_plan: opts.jwt_auth_forwarding,
dedupe_subgraph_requests,
demand_control_context: opts.demand_control_context.clone(),
plugin_req_state: opts.plugin_req_state.as_ref(),
operation_name_factory: &opts.operation_name_factory,
connection_fingerprint: opts.connection_fingerprint,
};
if let Some(node) = &opts.query_plan.node {
executor.execute_plan_node(&mut exec_ctx, node).await;
}
let error_count = exec_ctx.errors.len();
if error_count > 0 {
opts.span.record_error_count(error_count);
opts.span
.record_errors(|| exec_ctx.errors.iter().map(|e| e.into()).collect());
if let Some(error_recorder) = opts.graphql_error_recorder.as_ref() {
error_recorder.record_errors(|| {
exec_ctx
.errors
.iter()
.map(|err| err.extensions.code.as_deref())
});
}
}
let mut data = exec_ctx.data;
let mut errors = exec_ctx.errors;
summary::record(|s| s.set_partial_response(error_count > 0 && !data.is_null()));
let mut response_size_estimate = exec_ctx.response_storage.estimate_final_response_size();
let mut demand_control_cost = None;
if let Some(demand_control) = executor.demand_control_context {
let actual = demand_control.calculate_actual_cost(
&data,
&opts.variable_values.variables_map,
&exec_ctx.subgraph_response_cost_tracker,
);
demand_control_cost = Some(DemandControlCost {
estimated: demand_control.evaluation.estimated_cost,
max: demand_control.operation.operation_max_cost,
actual,
});
if actual > demand_control.operation.operation_max_cost {
warn!(
target: targets::DEMAND_CONTROL,
operation_name = ?opts.operation_for_plan.name.as_deref(),
actual_cost = actual,
estimated_cost = demand_control.evaluation.estimated_cost,
max_cost = demand_control.operation.operation_max_cost,
"actual cost exceeds max cost (not enforced)"
);
}
demand_control.report_telemetry(
actual,
opts.operation_for_plan.name.as_deref(),
&opts.span,
);
demand_control.apply_expose_headers(&mut exec_ctx.response_headers_aggregator, actual);
}
cache_control::finalize(
&mut exec_ctx.response_headers_aggregator,
opts.operation_kind.is_mutation() || !errors.is_empty(),
exec_ctx.response_storage.len(),
);
opts.response_header_sink
.store(std::mem::take(&mut exec_ctx.response_headers_aggregator));
std::mem::take(&mut exec_ctx.extensions_aggregator).merge_into(&mut opts.extensions.extensions);
if !on_end_callbacks.is_empty() {
let mut end_payload = OnExecuteEndHookPayload {
data,
errors,
extensions: opts.extensions.extensions,
response_size_estimate,
request_context: opts
.plugin_req_state
.as_ref()
.map(|state| state.request_context.for_plugin::<hooks::OnExecute>())
.expect("plugin state not available, but on_end_callbacks are present"),
demand_control_cost,
};
for callback in on_end_callbacks {
let result = callback(end_payload);
end_payload = result.payload;
match result.control_flow {
EndControlFlow::Proceed => { }
EndControlFlow::EndWithResponse(response) => match response {
OnExecuteResponse::Output(response) => return Ok(response),
OnExecuteResponse::EarlyResponse(response) => {
return Ok(early_http_response_into_execution_output(
response,
&opts.response_header_sink,
));
}
},
}
}
data = end_payload.data;
errors = end_payload.errors;
opts.extensions.extensions = end_payload.extensions;
response_size_estimate = end_payload.response_size_estimate;
}
let mut status_code = StatusCode::OK;
if !errors.is_empty() {
if let Some(plugin_req_state) = opts.plugin_req_state.as_ref() {
let (new_errors, new_status_code) = handle_graphql_errors_with_plugins(
plugin_req_state.plugins.as_ref(),
plugin_req_state.context.as_ref(),
&plugin_req_state.request_context,
errors,
status_code,
);
errors = new_errors;
status_code = new_status_code;
}
}
if let Some(error_masking_runtime) = opts.error_masking_runtime.as_ref() {
for error in &mut errors {
error_masking_runtime.apply(error);
}
}
let body = project_by_operation(
&data,
errors,
&opts.extensions,
opts.operation_type_name.as_str(),
&opts.projection_plan,
&opts.variable_values.variables_map,
response_size_estimate,
&opts.introspection_context.metadata,
)
.with_plan_context(LazyPlanContext {
subgraph_name: || None,
affected_path: || None,
})?;
Ok(PlanExecutionOutput {
body,
error_count,
status_code,
})
}
pub struct Executor<'exec> {
pub variable_values: &'exec Option<VariablesMap>,
pub schema_metadata: &'exec SchemaMetadata,
pub executors: &'exec SubgraphExecutorMap,
pub client_request: &'exec ClientRequestDetails<'exec>,
pub headers_plan: &'exec HeaderRulesPlan,
pub extensions_plan: &'exec ExtensionsPlan,
pub jwt_forwarding_plan: Option<Arc<JwtAuthForwardingPlan>>,
pub dedupe_subgraph_requests: bool,
pub demand_control_context: Option<Arc<DemandControlExecutionContext>>,
pub plugin_req_state: Option<&'exec PluginRequestState<'exec>>,
pub operation_name_factory: &'exec OperationNameFactory,
pub connection_fingerprint: Option<ConnectionFingerprint>,
}
pub enum ExecutionJob<'exec> {
Fetch {
subgraph_name: &'exec str,
operation: &'exec SubgraphFetchOperation,
response: SubgraphResponse<'exec>,
output_rewrites: Option<&'exec [FetchRewrite]>,
merge_path: Option<&'exec FlattenNodePath>,
},
FlattenFetch {
subgraph_name: &'exec str,
operation: &'exec SubgraphFetchOperation,
response: SubgraphResponse<'exec>,
flatten_node_path: &'exec FlattenNodePath,
representation_hashes: Vec<Option<u64>>,
representation_hash_to_index: AHashMap<u64, usize>,
output_rewrites: Option<&'exec [FetchRewrite]>,
},
BatchFetch {
subgraph_name: &'exec str,
operation: &'exec SubgraphFetchOperation,
response: SubgraphResponse<'exec>,
aliases: Vec<AliasBatchState<'exec>>,
},
}
pub struct AliasBatchState<'exec> {
alias_spec: &'exec EntityBatchAlias,
representation_hash_to_index: AHashMap<u64, usize>,
paths: Vec<AliasPathState<'exec>>,
}
struct AliasPathState<'exec> {
merge_path: &'exec FlattenNodePath,
representation_hashes: Arc<Vec<Option<u64>>>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
struct AliasIndex(usize);
#[derive(Default)]
struct BatchFetchErrors {
by_alias_index: AHashMap<AliasIndex, Vec<GraphQLError>>,
unmatched: Vec<GraphQLError>,
}
impl<'exec> ExecutionJob<'exec> {
fn response(self) -> SubgraphResponse<'exec> {
match self {
ExecutionJob::Fetch { response, .. } => response,
ExecutionJob::FlattenFetch { response, .. } => response,
ExecutionJob::BatchFetch { response, .. } => response,
}
}
pub fn response_ref(&self) -> &SubgraphResponse<'exec> {
match self {
ExecutionJob::Fetch { response, .. } => response,
ExecutionJob::FlattenFetch { response, .. } => response,
ExecutionJob::BatchFetch { response, .. } => response,
}
}
pub fn subgraph_name(&self) -> &'exec str {
match self {
ExecutionJob::Fetch { subgraph_name, .. } => subgraph_name,
ExecutionJob::FlattenFetch { subgraph_name, .. } => subgraph_name,
ExecutionJob::BatchFetch { subgraph_name, .. } => subgraph_name,
}
}
pub fn operation(&self) -> &'exec SubgraphFetchOperation {
match self {
ExecutionJob::Fetch { operation, .. } => operation,
ExecutionJob::FlattenFetch { operation, .. } => operation,
ExecutionJob::BatchFetch { operation, .. } => operation,
}
}
fn affected_path(&self) -> Option<&'exec FlattenNodePath> {
match self {
ExecutionJob::Fetch { merge_path, .. } => *merge_path,
ExecutionJob::FlattenFetch {
flatten_node_path, ..
} => Some(flatten_node_path),
ExecutionJob::BatchFetch { .. } => None,
}
}
}
struct PrepareExecutionJobOpts<'exec> {
subgraph_name: &'exec str,
variable_usages: Option<&'exec BTreeSet<String>>,
operation_name: Option<String>,
operation_kind: Option<&'exec OperationKind>,
operation: &'exec SubgraphFetchOperation,
output_rewrites: Option<&'exec [FetchRewrite]>,
custom_scalar_paths: Option<&'exec CustomScalarPaths>,
raw_variable_values: Option<Vec<(&'exec str, Vec<u8>)>>,
affected_path: Option<&'exec FlattenNodePath>,
}
impl<'exec> Executor<'exec> {
async fn execute_plan_node(&self, ctx: &mut ExecutionContext<'exec>, node: &'exec PlanNode) {
match node {
PlanNode::Parallel(parallel_node) => {
let mut scope = FuturesUnordered::new();
for child in ¶llel_node.nodes {
if let Some(fut) = self.prepare_job_future(child, &ctx.data) {
scope.push(fut);
}
}
while let Some(job) = scope.next().await {
self.process_job_result(ctx, job);
}
}
PlanNode::Sequence(sequence_node) => {
for child in &sequence_node.nodes {
Box::pin(self.execute_plan_node(ctx, child)).await;
}
}
PlanNode::Condition(condition_node) => {
if let Some(next_node) =
condition_node_by_variables(condition_node, self.variable_values)
{
Box::pin(self.execute_plan_node(ctx, next_node)).await;
}
}
node => {
if let Some(fut) = self.prepare_job_future(node, &ctx.data) {
let job = fut.await;
self.process_job_result(ctx, job);
}
}
}
}
fn prepare_root_reentry_fetch_job<'wave>(
&'wave self,
flatten_node: &'exec FlattenNode,
fetch_node: &'exec FetchNode,
) -> BoxFuture<'wave, Result<ExecutionJob<'exec>, PlanExecutionError>> {
self.prepare_execution_job(PrepareExecutionJobOpts {
subgraph_name: &fetch_node.service_name,
variable_usages: fetch_node.variable_usages.as_ref(),
operation_name: self
.operation_name_factory
.generate(&fetch_node.service_name, fetch_node.id),
operation_kind: fetch_node.operation_kind.as_ref(),
operation: &fetch_node.operation,
output_rewrites: fetch_node.output_rewrites.as_deref(),
custom_scalar_paths: fetch_node.custom_scalar_paths.as_ref(),
raw_variable_values: None,
affected_path: Some(&flatten_node.path),
})
.boxed()
}
fn prepare_job_future<'wave>(
&'wave self,
node: &'exec PlanNode,
data: &Value<'exec>,
) -> Option<BoxFuture<'wave, Result<ExecutionJob<'exec>, PlanExecutionError>>> {
match node {
PlanNode::Fetch(fetch_node) => Some(
self.prepare_execution_job(PrepareExecutionJobOpts {
subgraph_name: &fetch_node.service_name,
variable_usages: fetch_node.variable_usages.as_ref(),
operation_name: self
.operation_name_factory
.generate(&fetch_node.service_name, fetch_node.id),
operation_kind: fetch_node.operation_kind.as_ref(),
operation: &fetch_node.operation,
output_rewrites: fetch_node.output_rewrites.as_deref(),
custom_scalar_paths: fetch_node.custom_scalar_paths.as_ref(),
raw_variable_values: None,
affected_path: None,
})
.boxed(),
),
PlanNode::BatchFetch(batch_fetch_node) => {
let (raw_variable_values, aliases) =
self.prepare_batch_fetch_job_state(&batch_fetch_node.entity_batch, data);
if aliases
.iter()
.all(|alias| alias.representation_hash_to_index.is_empty())
{
debug!(
target: targets::GRAPHQL_EXECUTION,
alias_count = aliases.len(),
"Skipping batched entity fetch with no representations"
);
return None;
}
Some(
self.prepare_execution_job(PrepareExecutionJobOpts {
subgraph_name: &batch_fetch_node.service_name,
variable_usages: batch_fetch_node.variable_usages.as_ref(),
operation_name: self
.operation_name_factory
.generate(&batch_fetch_node.service_name, batch_fetch_node.id),
operation_kind: batch_fetch_node.operation_kind.as_ref(),
operation: &batch_fetch_node.operation,
output_rewrites: None,
custom_scalar_paths: batch_fetch_node.custom_scalar_paths.as_ref(),
raw_variable_values: Some(raw_variable_values),
affected_path: None,
})
.map_ok(|fetch_job| ExecutionJob::BatchFetch {
operation: fetch_job.operation(),
subgraph_name: fetch_job.subgraph_name(),
response: fetch_job.response(),
aliases,
})
.boxed(),
)
}
PlanNode::Flatten(flatten_node) => {
let fetch_node = match flatten_node.node.as_ref() {
PlanNode::Fetch(fetch_node) => fetch_node,
_ => return None,
};
let Some(requires_nodes) = fetch_node.requires.as_ref() else {
return Some(self.prepare_root_reentry_fetch_job(flatten_node, fetch_node));
};
let mut index = 0;
let normalized_path = flatten_node.path.as_slice();
let mut filtered_representations = Vec::new();
filtered_representations.put(OPEN_BRACKET);
let possible_types = &self.schema_metadata.possible_types;
let mut representation_hashes: Vec<Option<u64>> = Vec::new();
let mut representation_hash_to_index: AHashMap<u64, usize> = AHashMap::new();
let arena = bumpalo::Bump::new();
traverse_and_callback(
data,
normalized_path,
&self.schema_metadata.possible_types,
&mut |entity| {
if entity.is_null() {
representation_hashes.push(None);
return;
}
let hash = entity.to_hash(&requires_nodes.items, possible_types);
representation_hashes.push(Some(hash));
let is_first_representation = representation_hash_to_index.is_empty();
let vacant_entry = match representation_hash_to_index.entry(hash) {
Entry::Occupied(_) => return,
Entry::Vacant(vacant_entry) => vacant_entry,
};
let entity = if let Some(input_rewrites) = &fetch_node.input_rewrites {
let new_entity = arena.alloc(entity.clone());
for input_rewrite in input_rewrites {
input_rewrite
.rewrite(&self.schema_metadata.possible_types, new_entity);
}
new_entity
} else {
entity
};
let is_projected = project_requires(
possible_types,
&requires_nodes.items,
entity,
&mut filtered_representations,
is_first_representation,
None,
);
if is_projected {
vacant_entry.insert(index);
index += 1;
}
},
);
filtered_representations.put(CLOSE_BRACKET);
if representation_hash_to_index.is_empty() {
return None;
}
Some(
self.prepare_execution_job(PrepareExecutionJobOpts {
subgraph_name: &fetch_node.service_name,
variable_usages: fetch_node.variable_usages.as_ref(),
operation_name: self
.operation_name_factory
.generate(&fetch_node.service_name, fetch_node.id),
operation_kind: fetch_node.operation_kind.as_ref(),
operation: &fetch_node.operation,
output_rewrites: fetch_node.output_rewrites.as_deref(),
custom_scalar_paths: fetch_node.custom_scalar_paths.as_ref(),
raw_variable_values: Some(vec![(
"representations",
filtered_representations,
)]),
affected_path: Some(&flatten_node.path),
})
.map_ok(|fetch_job| ExecutionJob::FlattenFetch {
operation: fetch_job.operation(),
flatten_node_path: &flatten_node.path,
response: fetch_job.response(),
subgraph_name: fetch_node.service_name.as_str(),
representation_hashes,
representation_hash_to_index,
output_rewrites: fetch_node.output_rewrites.as_deref(),
})
.boxed(),
)
}
PlanNode::Condition(node) => condition_node_by_variables(node, self.variable_values)
.and_then(|node| self.prepare_job_future(node, data)),
_ => None,
}
}
fn process_job_result(
&self,
ctx: &mut ExecutionContext<'exec>,
job: Result<ExecutionJob<'exec>, PlanExecutionError>,
) {
match job {
Err(ref err) => {
if let (Some(subgraph_name), Some(subgraph_headers)) = (
err.subgraph_name().as_deref(),
err.subgraph_response_headers(),
) {
if let Err(ref propagation_err) = apply_subgraph_response_headers(
self.headers_plan,
subgraph_name,
subgraph_headers,
self.client_request,
&mut ctx.response_headers_aggregator,
)
.with_plan_context(LazyPlanContext {
subgraph_name: || Some(subgraph_name.to_string()),
affected_path: || err.affected_path().clone(),
}) {
self.log_error(propagation_err);
ctx.errors.push(propagation_err.into());
}
}
self.log_error(err);
ctx.errors.push(err.into());
}
Ok(job) => {
let subgraph_name = job.subgraph_name();
let affected_path = job.affected_path();
if let Some(demand_control) = &self.demand_control_context {
demand_control.record_subgraph_response_cost(
&mut ctx.subgraph_response_cost_tracker,
&job,
self.variable_values,
);
}
if let Some(ref subgraph_headers) = job.response_ref().headers {
if let Err(ref err) = apply_subgraph_response_headers(
self.headers_plan,
job.subgraph_name(),
subgraph_headers,
self.client_request,
&mut ctx.response_headers_aggregator,
)
.with_plan_context(LazyPlanContext {
subgraph_name: || Some(subgraph_name.to_string()),
affected_path: || affected_path.map(|p| p.to_string()),
}) {
self.log_error(err);
ctx.errors.push(err.into());
}
}
if let Some(ref subgraph_extensions) = job.response_ref().extensions {
apply_subgraph_extensions(
self.extensions_plan,
subgraph_extensions,
&mut ctx.extensions_aggregator,
);
}
match job {
ExecutionJob::Fetch {
mut response,
output_rewrites,
merge_path,
..
} => {
if let Some(response_bytes) = response.bytes {
ctx.response_storage.add_response(response_bytes);
}
if let Some(output_rewrites) = output_rewrites {
for output_rewrite in output_rewrites {
output_rewrite.rewrite(
&self.schema_metadata.possible_types,
&mut response.data,
);
}
}
match merge_path {
None => deep_merge(&mut ctx.data, response.data),
Some(merge_path) => {
let source = response.data;
traverse_and_callback_mut(
&mut ctx.data,
merge_path.as_slice(),
self.schema_metadata,
None,
&mut |target, _error_path| {
deep_merge(target, source.clone());
},
);
}
}
ctx.handle_errors(subgraph_name, affected_path, response.errors, None);
}
ExecutionJob::FlattenFetch {
mut response,
flatten_node_path,
representation_hashes,
ref representation_hash_to_index,
output_rewrites,
..
} => {
if let Some(response_bytes) = response.bytes {
ctx.response_storage.add_response(response_bytes);
}
if let Some(mut entities) = response.data.take_entities() {
if let Some(output_rewrites) = output_rewrites {
for output_rewrite in output_rewrites {
for entity in &mut entities {
output_rewrite
.rewrite(&self.schema_metadata.possible_types, entity);
}
}
}
let mut index = 0;
let normalized_path = flatten_node_path.as_slice();
let initial_error_path = response.errors.as_ref().map(|_| {
GraphQLErrorPath::with_capacity(normalized_path.len() + 2)
});
let mut entity_index_error_map = response
.errors
.as_ref()
.map(|_| HashMap::with_capacity(entities.len()));
traverse_and_callback_mut(
&mut ctx.data,
normalized_path,
self.schema_metadata,
initial_error_path,
&mut |target, error_path| {
let hash = representation_hashes.get(index).copied();
index += 1;
let Some(Some(hash)) = hash else {
return;
};
if let Some(entity_index) =
representation_hash_to_index.get(&hash)
{
if let (Some(error_path), Some(entity_index_error_map)) =
(error_path, entity_index_error_map.as_mut())
{
let error_paths = entity_index_error_map
.entry(entity_index)
.or_insert_with(Vec::new);
error_paths.push(error_path);
}
if let Some(entity) = entities.get(*entity_index) {
let new_val: Value<'_> =
unsafe { std::mem::transmute(entity.clone()) };
deep_merge(target, new_val);
}
}
},
);
ctx.handle_errors(
subgraph_name,
affected_path,
response.errors,
entity_index_error_map,
);
} else {
ctx.handle_errors(subgraph_name, affected_path, response.errors, None);
}
}
ExecutionJob::BatchFetch {
mut response,
aliases,
..
} => {
if let Some(response_bytes) = response.bytes {
ctx.response_storage.add_response(response_bytes);
}
let mut errors =
self.partition_batch_errors_by_alias(&aliases, response.errors.take());
let mut entities_by_alias =
Self::collect_batched_entities_by_alias(&mut response.data, &aliases);
for (alias_index, alias_state) in aliases.iter().enumerate() {
let alias_index = AliasIndex(alias_index);
let mut alias_errors = errors.by_alias_index.remove(&alias_index);
let entity_index_error_map = self.merge_batch_alias_entities(
ctx,
alias_state,
entities_by_alias.get_mut(&alias_index),
alias_errors.as_deref(),
);
let affected_path = if alias_state.paths.len() == 1 {
Some(alias_state.paths[0].merge_path)
} else {
None
};
ctx.handle_errors(
subgraph_name,
affected_path,
alias_errors.take(),
entity_index_error_map,
);
}
if !errors.unmatched.is_empty() {
ctx.handle_errors(subgraph_name, None, Some(errors.unmatched), None);
}
tracing::trace!(
target: targets::GRAPHQL_EXECUTION,
alias_count = aliases.len(),
"Patched entity batch alias results"
);
}
}
}
}
}
fn log_error(&self, error: &PlanExecutionError) {
log_plan_execution_error(error);
}
fn partition_batch_errors_by_alias(
&self,
aliases: &[AliasBatchState<'exec>],
response_errors: Option<Vec<GraphQLError>>,
) -> BatchFetchErrors {
let mut alias_index_by_name: AHashMap<&str, AliasIndex> =
AHashMap::with_capacity(aliases.len());
for (alias_index, alias_state) in aliases.iter().enumerate() {
alias_index_by_name.insert(
alias_state.alias_spec.alias.as_str(),
AliasIndex(alias_index),
);
}
let mut errors_by_alias_index: AHashMap<AliasIndex, Vec<GraphQLError>> = AHashMap::new();
let mut unmatched_errors: Vec<GraphQLError> = Vec::new();
let Some(response_errors) = response_errors else {
return BatchFetchErrors::default();
};
for mut error in response_errors {
let maybe_alias = error.path.as_ref().and_then(|path| {
path.segments.first().and_then(|segment| match segment {
GraphQLErrorPathSegment::String(alias) => Some(alias.as_str()),
_ => None,
})
});
let Some(alias) = maybe_alias else {
unmatched_errors.push(error);
continue;
};
let Some(alias_index) = alias_index_by_name.get(alias) else {
unmatched_errors.push(error);
continue;
};
if let Some(path) = error.path.as_mut() {
if let Some(GraphQLErrorPathSegment::String(first)) = path.segments.first_mut() {
*first = "_entities".to_string();
}
}
errors_by_alias_index
.entry(*alias_index)
.or_default()
.push(error);
}
BatchFetchErrors {
by_alias_index: errors_by_alias_index,
unmatched: unmatched_errors,
}
}
fn collect_batched_entities_by_alias(
response_data: &mut Value<'exec>,
aliases: &[AliasBatchState<'exec>],
) -> AHashMap<AliasIndex, Vec<Value<'exec>>> {
let mut entities_by_alias: AHashMap<AliasIndex, Vec<Value<'exec>>> =
AHashMap::with_capacity(aliases.len());
for (alias_index, alias_state) in aliases.iter().enumerate() {
let Some(entities) =
response_data.take_entities_by_key(alias_state.alias_spec.alias.as_str())
else {
continue;
};
entities_by_alias.insert(AliasIndex(alias_index), entities);
}
entities_by_alias
}
fn merge_batch_alias_entities<'alias>(
&self,
ctx: &mut ExecutionContext<'exec>,
alias_state: &'alias AliasBatchState<'exec>,
entities: Option<&mut Vec<Value<'exec>>>,
alias_errors: Option<&[GraphQLError]>,
) -> Option<HashMap<&'alias usize, Vec<GraphQLErrorPath>>> {
let has_alias_errors = alias_errors.is_some();
let mut entity_index_error_map = has_alias_errors.then(HashMap::new);
let Some(entities) = entities else {
return entity_index_error_map;
};
if let Some(output_rewrites) = alias_state.alias_spec.output_rewrites.as_ref() {
for output_rewrite in output_rewrites {
for entity in entities.iter_mut() {
output_rewrite.rewrite(&self.schema_metadata.possible_types, entity);
}
}
}
if alias_state.representation_hash_to_index.is_empty() {
return entity_index_error_map;
}
for path_state in &alias_state.paths {
let mut index = 0;
let normalized_path = path_state.merge_path.as_slice();
let initial_error_path = has_alias_errors
.then(|| GraphQLErrorPath::with_capacity(normalized_path.len() + 2));
traverse_and_callback_mut(
&mut ctx.data,
normalized_path,
self.schema_metadata,
initial_error_path,
&mut |target_data, error_path| {
let hash = path_state.representation_hashes.get(index).copied();
index += 1;
let Some(Some(hash)) = hash else {
return;
};
if let Some(entity_index) = alias_state.representation_hash_to_index.get(&hash)
{
if let (Some(error_path), Some(entity_index_error_map)) =
(error_path, entity_index_error_map.as_mut())
{
let error_paths = entity_index_error_map
.entry(entity_index)
.or_insert_with(Vec::new);
error_paths.push(error_path);
}
if let Some(entity) = entities.get(*entity_index) {
let new_val: Value<'_> = unsafe { std::mem::transmute(entity.clone()) };
deep_merge(target_data, new_val);
}
}
},
);
}
entity_index_error_map
}
fn prepare_batch_fetch_job_state(
&self,
entity_batch: &'exec EntityBatch,
data: &Value<'exec>,
) -> (Vec<(&'exec str, Vec<u8>)>, Vec<AliasBatchState<'exec>>) {
let mut raw_variable_values: Vec<(&'exec str, Vec<u8>)> =
Vec::with_capacity(entity_batch.aliases.len());
let mut raw_variable_indices_by_name: AHashMap<&'exec str, usize> =
AHashMap::with_capacity(entity_batch.aliases.len());
let mut aliases = Vec::with_capacity(entity_batch.aliases.len());
let possible_types = &self.schema_metadata.possible_types;
for alias_spec in &entity_batch.aliases {
let mut index = 0;
let mut filtered_representations = Vec::new();
filtered_representations.put(OPEN_BRACKET);
let mut representation_hash_to_index: AHashMap<u64, usize> = AHashMap::new();
let arena = bumpalo::Bump::new();
let mut path_hashes_by_index: Vec<Option<Arc<Vec<Option<u64>>>>> =
vec![None; alias_spec.merge_paths.len()];
let mut path_groups: Vec<(&FlattenNodePath, Vec<usize>)> =
Vec::with_capacity(alias_spec.merge_paths.len());
for (path_index, merge_path) in alias_spec.merge_paths.iter().enumerate() {
if let Some((_, target_indices)) =
path_groups.iter_mut().find(|(path, _)| *path == merge_path)
{
target_indices.push(path_index);
} else {
path_groups.push((merge_path, vec![path_index]));
}
}
for (merge_path, grouped_target_indices) in path_groups {
let mut representation_hashes: Vec<Option<u64>> = Vec::new();
traverse_and_callback(data, merge_path.as_slice(), possible_types, &mut |entity| {
if entity.is_null() {
representation_hashes.push(None);
return;
}
let hash = entity.to_hash(&alias_spec.requires.items, possible_types);
representation_hashes.push(Some(hash));
let is_first_representation = representation_hash_to_index.is_empty();
let vacant_entry = match representation_hash_to_index.entry(hash) {
Entry::Occupied(_) => return,
Entry::Vacant(vacant_entry) => vacant_entry,
};
let entity = if let Some(input_rewrites) = &alias_spec.input_rewrites {
let new_entity = arena.alloc(entity.clone());
for input_rewrite in input_rewrites {
input_rewrite.rewrite(&self.schema_metadata.possible_types, new_entity);
}
new_entity
} else {
entity
};
let is_projected = project_requires(
possible_types,
&alias_spec.requires.items,
entity,
&mut filtered_representations,
is_first_representation,
None,
);
if is_projected {
vacant_entry.insert(index);
index += 1;
}
});
let representation_hashes = Arc::new(representation_hashes);
for path_index in grouped_target_indices {
path_hashes_by_index[path_index] = Some(Arc::clone(&representation_hashes));
}
}
filtered_representations.put(CLOSE_BRACKET);
let mut paths = Vec::with_capacity(alias_spec.merge_paths.len());
for (path_index, merge_path) in alias_spec.merge_paths.iter().enumerate() {
paths.push(AliasPathState {
merge_path,
representation_hashes: path_hashes_by_index[path_index]
.take()
.unwrap_or_else(|| Arc::new(Vec::new())),
});
}
let variable_name = alias_spec.representations_variable_name.as_str();
if !raw_variable_indices_by_name.contains_key(variable_name) {
raw_variable_indices_by_name.insert(variable_name, raw_variable_values.len());
raw_variable_values.push((variable_name, filtered_representations));
}
aliases.push(AliasBatchState {
alias_spec,
representation_hash_to_index,
paths,
});
}
(raw_variable_values, aliases)
}
async fn prepare_execution_job(
&self,
opts: PrepareExecutionJobOpts<'exec>,
) -> Result<ExecutionJob<'exec>, PlanExecutionError> {
let subgraph_operation_span =
GraphQLSubgraphOperationSpan::new(opts.subgraph_name, &opts.operation.document_str);
async {
let mut headers_map = HeaderMap::new();
let subgraph_name_factory = || Some(opts.subgraph_name.to_string());
let affected_path_factory = || opts.affected_path.map(|p| p.to_string());
modify_subgraph_request_headers(
self.headers_plan,
opts.subgraph_name,
self.client_request,
&mut headers_map,
)
.with_plan_context(LazyPlanContext {
subgraph_name: subgraph_name_factory,
affected_path: affected_path_factory,
})?;
let variable_refs = select_fetch_variables(self.variable_values, opts.variable_usages);
let mut subgraph_request = SubgraphExecutionRequest {
query: &opts.operation.document_str,
document_name_write_pos: opts.operation.name_write_position,
dedupe: self.dedupe_subgraph_requests,
operation_name: opts.operation_name,
variables: variable_refs,
raw_variable_values: opts.raw_variable_values,
headers: headers_map,
extensions: None,
custom_scalar_paths: opts.custom_scalar_paths,
connection_fingerprint: self.connection_fingerprint,
};
let client_document_hash_str = opts.operation.hash.to_string();
subgraph_operation_span.record_operation_identity(GraphQLSpanOperationIdentity {
name: subgraph_request.operation_name.as_deref(),
operation_type: match opts.operation_kind {
Some(OperationKind::Query) | None => "query",
Some(OperationKind::Mutation) => "mutation",
Some(OperationKind::Subscription) => "subscription",
},
client_document_hash: &client_document_hash_str,
});
if let Some(jwt_forwarding_plan) = &self.jwt_forwarding_plan {
subgraph_request.add_request_extensions_field(
jwt_forwarding_plan.extension_field_name.clone(),
jwt_forwarding_plan.extension_field_value.clone(),
);
}
let mut response = self
.executors
.execute(
opts.subgraph_name,
subgraph_request,
self.client_request,
self.plugin_req_state,
self.demand_control_context.as_deref(),
)
.await
.with_plan_context(LazyPlanContext {
subgraph_name: subgraph_name_factory,
affected_path: affected_path_factory,
})?;
if let Some(errors) = &response.errors {
if !errors.is_empty() {
subgraph_operation_span.record_error_count(errors.len());
subgraph_operation_span
.record_errors(|| errors.iter().map(|e| e.into()).collect());
}
}
if let Some(status) = &response.status {
if !status.is_success() {
response.append_error(GraphQLError::create_from_status_code(
status,
opts.subgraph_name,
));
}
}
Ok(ExecutionJob::Fetch {
subgraph_name: opts.subgraph_name,
operation: opts.operation,
response,
output_rewrites: opts.output_rewrites,
merge_path: opts.affected_path,
})
}
.inspect_err(|err: &PlanExecutionError| {
subgraph_operation_span.record_error_count(1);
subgraph_operation_span
.record_errors(|| vec![ObservedError::from(&GraphQLError::from(err))]);
})
.instrument(subgraph_operation_span.clone())
.await
}
}
fn condition_node_by_variables<'a>(
condition_node: &'a ConditionNode,
variable_values: &'a Option<VariablesMap>,
) -> Option<&'a PlanNode> {
let vars = variable_values.as_ref()?;
let value = vars.get(&condition_node.condition)?;
let condition_met = matches!(value.as_ref(), ValueRef::Bool(true));
if condition_met {
condition_node.if_clause.as_deref()
} else {
condition_node.else_clause.as_deref()
}
}
fn select_fetch_variables<'a>(
variable_values: &'a Option<VariablesMap>,
variable_usages: Option<&BTreeSet<String>>,
) -> Option<HashMap<&'a str, &'a sonic_rs::Value>> {
let values = variable_values.as_ref()?;
variable_usages.map(|variable_usages| {
variable_usages
.iter()
.filter_map(|var_name| {
values
.get_key_value(var_name.as_str())
.map(|(key, value)| (key.as_str(), value))
})
.collect()
})
}
#[cfg(test)]
mod tests {
use crate::executor::{
execution::{
client_request_details::{ClientRequestDetails, JwtRequestDetails, OperationDetails},
operation_name::OperationNameFactory,
plan::Executor,
},
execution_context::ExecutionContext,
extensions::plan::ExtensionsPlan,
headers::plan::HeaderRulesPlan,
introspection::schema::SchemaMetadata,
response::{
graphql_error::{GraphQLErrorExtensions, GraphQLErrorPath},
value::Value as ResponseValue,
},
SubgraphExecutorMap,
};
use super::select_fetch_variables;
use crate::query_planner::{
ast::{document::Document, operation::SubgraphFetchOperation},
planner::plan_nodes::{EntityBatch, EntityBatchAlias, FetchNode, ParallelNode, PlanNode},
utils::parsing::parse_operation,
};
use crate::telemetry::TelemetryContext;
use dashmap::DashMap;
use graphql_tools::parser::query::{self, Definition};
use ntex::http::HeaderMap;
use sonic_rs::Value;
use std::{
collections::{BTreeSet, HashMap},
sync::{mpsc::channel, Arc},
time::Duration,
vec,
};
fn value_from_number(n: i32) -> Value {
sonic_rs::from_str(&n.to_string()).unwrap()
}
fn parse_document(query: &str) -> Document {
let document = parse_operation(query);
let mut operation = None;
let mut fragments = Vec::new();
for definition in document.definitions {
match definition {
Definition::Operation(current_operation) => {
if operation.is_none() {
operation = Some(current_operation.into());
}
}
Definition::Fragment(fragment) => fragments.push(fragment.into()),
}
}
Document {
operation: operation.expect("operation definition should exist"),
fragments,
}
}
#[test]
fn select_fetch_variables_only_used_variables() {
let mut variable_values_map = HashMap::new();
variable_values_map.insert("used".to_string(), value_from_number(1));
variable_values_map.insert("unused".to_string(), value_from_number(2));
let variable_values = Some(variable_values_map);
let mut usages = BTreeSet::new();
usages.insert("used".to_string());
let selected = select_fetch_variables(&variable_values, Some(&usages)).unwrap();
assert_eq!(selected.len(), 1);
assert!(selected.contains_key("used"));
assert!(!selected.contains_key("unused"));
}
#[test]
fn select_fetch_variables_ignores_missing_usage_entries() {
let mut variable_values_map = HashMap::new();
variable_values_map.insert("present".to_string(), value_from_number(3));
let variable_values = Some(variable_values_map);
let mut usages = BTreeSet::new();
usages.insert("present".to_string());
usages.insert("missing".to_string());
let selected = select_fetch_variables(&variable_values, Some(&usages)).unwrap();
assert_eq!(selected.len(), 1);
assert!(selected.contains_key("present"));
assert!(!selected.contains_key("missing"));
}
#[test]
fn select_fetch_variables_for_no_usage_entries() {
let mut variable_values_map = HashMap::new();
variable_values_map.insert("unused_1".to_string(), value_from_number(1));
variable_values_map.insert("unused_2".to_string(), value_from_number(2));
let variable_values = Some(variable_values_map);
let selected = select_fetch_variables(&variable_values, None);
assert!(selected.is_none());
}
#[test]
fn normalize_entity_errors_correctly() {
use crate::executor::response::graphql_error::{GraphQLError, GraphQLErrorPathSegment};
use std::collections::HashMap;
let mut ctx = ExecutionContext::default();
let mut entity_index_error_map: HashMap<&usize, Vec<GraphQLErrorPath>> = HashMap::new();
entity_index_error_map.insert(
&0,
vec![
GraphQLErrorPath {
segments: vec![
GraphQLErrorPathSegment::String("a".to_string()),
GraphQLErrorPathSegment::Index(0),
],
},
GraphQLErrorPath {
segments: vec![
GraphQLErrorPathSegment::String("b".to_string()),
GraphQLErrorPathSegment::Index(1),
],
},
],
);
let response_errors = vec![GraphQLError {
message: "Error 1".to_string(),
locations: None,
path: Some(GraphQLErrorPath {
segments: vec![
GraphQLErrorPathSegment::String("_entities".to_string()),
GraphQLErrorPathSegment::Index(0),
GraphQLErrorPathSegment::String("field1".to_string()),
],
}),
extensions: GraphQLErrorExtensions::default(),
}];
ctx.handle_errors(
"subgraph_a",
None,
Some(response_errors),
Some(entity_index_error_map),
);
assert_eq!(ctx.errors.len(), 2);
assert_eq!(ctx.errors[0].message, "Error 1");
assert_eq!(
ctx.errors[0].path.as_ref().unwrap().segments,
vec![
GraphQLErrorPathSegment::String("a".to_string()),
GraphQLErrorPathSegment::Index(0),
GraphQLErrorPathSegment::String("field1".to_string())
]
);
assert_eq!(ctx.errors[1].message, "Error 1");
assert_eq!(
ctx.errors[1].path.as_ref().unwrap().segments,
vec![
GraphQLErrorPathSegment::String("b".to_string()),
GraphQLErrorPathSegment::Index(1),
GraphQLErrorPathSegment::String("field1".to_string())
]
);
}
#[test]
fn prepare_batch_fetch_job_state_deduplicates_shared_variable_payloads() {
let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
let subgraph_endpoint_map = HashMap::from([(
"inventory".to_string(),
"http://example.com/graphql".parse().unwrap(),
)]);
let executors = SubgraphExecutorMap::from_http_endpoint_map(
&subgraph_endpoint_map,
Default::default(),
Default::default(),
Default::default(),
None,
Arc::new(TelemetryContext::from_propagation_config(
&Default::default(),
&Default::default(),
)),
Arc::new(DashMap::new()),
)
.unwrap();
let executor = Executor {
variable_values: &None,
schema_metadata: &SchemaMetadata::default(),
executors: &executors,
client_request: &ClientRequestDetails {
method: &http::Method::POST,
url: &"http://example.com".parse().unwrap(),
headers: HeaderMap::new().into(),
operation: OperationDetails {
name: None,
query: "{ products { upc } }",
kind: "query",
},
jwt: JwtRequestDetails::Unauthenticated.into(),
path_params: Default::default(),
},
headers_plan: &HeaderRulesPlan::default(),
extensions_plan: &ExtensionsPlan::default(),
jwt_forwarding_plan: None,
dedupe_subgraph_requests: false,
demand_control_context: None,
plugin_req_state: None,
operation_name_factory: &OperationNameFactory::default(),
connection_fingerprint: None,
};
let data: ResponseValue = sonic_rs::from_str(
r#"{
"products": [
{"__typename": "Product", "upc": "1"},
{"__typename": "Product", "upc": "2"}
]
}"#,
)
.unwrap();
fn document_into_selection<'a>(
doc: query::Document<'a, String>,
) -> query::SelectionSet<'a, String> {
doc.definitions
.iter()
.find_map(|def| {
let query::Definition::Operation(op) = def else {
return None;
};
match op {
query::OperationDefinition::SelectionSet(sel) => Some(sel),
query::OperationDefinition::Query(q) => Some(&q.selection_set),
query::OperationDefinition::Mutation(m) => Some(&m.selection_set),
query::OperationDefinition::Subscription(s) => Some(&s.selection_set),
}
})
.unwrap()
.clone()
}
let requires_query = parse_operation("{ ... on Product { upc } }");
let requires_selection = document_into_selection(requires_query);
let shared_var = "__batch_reps_0".to_string();
let entity_batch = EntityBatch {
aliases: vec![
EntityBatchAlias {
alias: "_e0".to_string(),
representations_variable_name: shared_var.clone(),
merge_paths: vec![],
requires: requires_selection.clone().into(),
input_rewrites: None,
output_rewrites: None,
},
EntityBatchAlias {
alias: "_e1".to_string(),
representations_variable_name: shared_var,
merge_paths: vec![],
requires: requires_selection.into(),
input_rewrites: None,
output_rewrites: None,
},
],
};
let (raw_variable_values, aliases) =
executor.prepare_batch_fetch_job_state(&entity_batch, &data);
assert_eq!(aliases.len(), 2);
assert_eq!(raw_variable_values.len(), 1);
assert_eq!(raw_variable_values[0].0, "__batch_reps_0");
}
#[tokio::test]
async fn runs_parallel_jobs_in_parallel() {
let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
let mut subgraph_a = mockito::Server::new_async().await;
let mut subgraph_b = mockito::Server::new_async().await;
let data = crate::executor::response::value::Value::Null;
let subgraph_endpoint_map = HashMap::from([
(
"subgraph_a".to_string(),
format!("http://{}/graphql", subgraph_a.host_with_port())
.parse()
.unwrap(),
),
(
"subgraph_b".to_string(),
format!("http://{}/graphql", subgraph_b.host_with_port())
.parse()
.unwrap(),
),
]);
let executor = Executor {
variable_values: &None,
schema_metadata: &SchemaMetadata::default(),
executors: &SubgraphExecutorMap::from_http_endpoint_map(
&subgraph_endpoint_map,
Default::default(),
Default::default(),
Default::default(),
None,
Arc::new(TelemetryContext::from_propagation_config(
&Default::default(),
&Default::default(),
)),
Arc::new(DashMap::new()),
)
.unwrap(),
client_request: &ClientRequestDetails {
method: &http::Method::POST,
url: &"http://example.com".parse().unwrap(),
headers: HeaderMap::new().into(),
operation: OperationDetails {
name: None,
query: "{ from_a from_b }",
kind: "query",
},
jwt: JwtRequestDetails::Unauthenticated.into(),
path_params: Default::default(),
},
headers_plan: &HeaderRulesPlan::default(),
extensions_plan: &ExtensionsPlan::default(),
jwt_forwarding_plan: None,
dedupe_subgraph_requests: false,
demand_control_context: None,
plugin_req_state: None,
operation_name_factory: &OperationNameFactory::default(),
connection_fingerprint: None,
};
let mock_a = subgraph_a
.mock("POST", "/graphql")
.with_status(200)
.with_header("content-type", "application/json")
.with_body(r#"{"data":{"from_a":"value_a"}}"#)
.create();
let mut exec_ctx = ExecutionContext {
data,
..Default::default()
};
let data_ref: &'static crate::executor::response::value::Value<'static> =
unsafe { std::mem::transmute(&exec_ctx.data) };
let (sender, receiver) = channel();
let mock_b = subgraph_b
.mock("POST", "/graphql")
.with_status(200)
.with_header("content-type", "application/json")
.with_chunked_body(move |writer| {
std::thread::sleep(Duration::from_millis(1000));
if let Some(data) = data_ref.as_object() {
let from_a_index = data.iter().position(|(k, _)| k == &"from_a");
let from_a_value = from_a_index
.and_then(|index| data.get(index))
.and_then(|(_, v)| v.as_str());
if let Some(from_a_value) = from_a_value {
sender
.send(from_a_value.to_string())
.expect("Failed to send from_a value through channel");
}
}
writer.write_fmt(format_args!(r#"{{"data":{{"from_b":"value_b"}}}}"#))
})
.create();
executor
.execute_plan_node(
&mut exec_ctx,
&PlanNode::Parallel(ParallelNode {
nodes: vec![
PlanNode::Fetch(FetchNode {
id: 1,
service_name: "subgraph_a".to_string(),
operation: SubgraphFetchOperation::from_anonymous_operation(
parse_document("{ from_a }"),
),
custom_scalar_paths: None,
requires: None,
input_rewrites: None,
output_rewrites: None,
variable_usages: None,
operation_kind: None,
}),
PlanNode::Fetch(FetchNode {
id: 2,
service_name: "subgraph_b".to_string(),
operation: SubgraphFetchOperation::from_anonymous_operation(
parse_document("{ from_b }"),
),
custom_scalar_paths: None,
requires: None,
input_rewrites: None,
output_rewrites: None,
variable_usages: None,
operation_kind: None,
}),
],
}),
)
.await;
mock_a.assert();
mock_b.assert();
let from_a_value = receiver
.recv()
.expect("Failed to receive from_a value through channel");
assert_eq!(from_a_value, "value_a");
}
}