hive_router_plan_executor/
context.rs

1use std::collections::HashMap;
2
3use hive_router_query_planner::planner::plan_nodes::{FetchNode, FetchRewrite, QueryPlan};
4
5use crate::{
6    headers::plan::ResponseHeaderAggregator,
7    response::{
8        graphql_error::{GraphQLError, GraphQLErrorPath},
9        storage::ResponsesStorage,
10        value::Value,
11    },
12};
13
14pub struct ExecutionContext<'a> {
15    pub response_storage: ResponsesStorage,
16    pub final_response: Value<'a>,
17    pub errors: Vec<GraphQLError>,
18    pub output_rewrites: OutputRewritesStorage,
19    pub response_headers_aggregator: ResponseHeaderAggregator,
20}
21
22impl<'a> Default for ExecutionContext<'a> {
23    fn default() -> Self {
24        ExecutionContext {
25            response_storage: Default::default(),
26            output_rewrites: Default::default(),
27            errors: Vec::new(),
28            final_response: Value::Null,
29            response_headers_aggregator: Default::default(),
30        }
31    }
32}
33
34impl<'a> ExecutionContext<'a> {
35    pub fn new(
36        query_plan: &QueryPlan,
37        init_final_response: Value<'a>,
38        init_errors: Vec<GraphQLError>,
39    ) -> Self {
40        ExecutionContext {
41            response_storage: ResponsesStorage::new(),
42            output_rewrites: OutputRewritesStorage::from_query_plan(query_plan),
43            errors: init_errors,
44            final_response: init_final_response,
45            response_headers_aggregator: Default::default(),
46        }
47    }
48
49    pub fn handle_errors(
50        &mut self,
51        subgraph_name: &str,
52        affected_path: Option<String>,
53        errors: Option<Vec<GraphQLError>>,
54        entity_index_error_map: Option<HashMap<&usize, Vec<GraphQLErrorPath>>>,
55    ) {
56        if let Some(response_errors) = errors {
57            for response_error in response_errors {
58                let mut processed_error = response_error.add_subgraph_name(subgraph_name);
59
60                if let Some(path) = &affected_path {
61                    processed_error = processed_error.add_affected_path(path.clone());
62                }
63
64                if let Some(entity_index_error_map) = &entity_index_error_map {
65                    let normalized_errors =
66                        processed_error.normalize_entity_error(entity_index_error_map);
67                    self.errors.extend(normalized_errors);
68                } else {
69                    self.errors.push(processed_error);
70                }
71            }
72        }
73    }
74}
75
76#[derive(Default)]
77pub struct OutputRewritesStorage {
78    output_rewrites: HashMap<i64, Vec<FetchRewrite>>,
79}
80
81impl OutputRewritesStorage {
82    pub fn from_query_plan(query_plan: &QueryPlan) -> OutputRewritesStorage {
83        let mut output_rewrites = OutputRewritesStorage {
84            output_rewrites: HashMap::new(),
85        };
86
87        for fetch_node in query_plan.fetch_nodes() {
88            output_rewrites.add_maybe(fetch_node);
89        }
90
91        output_rewrites
92    }
93
94    fn add_maybe(&mut self, fetch_node: &FetchNode) {
95        if let Some(rewrites) = &fetch_node.output_rewrites {
96            self.output_rewrites.insert(fetch_node.id, rewrites.clone());
97        }
98    }
99
100    pub fn get(&self, id: i64) -> Option<&Vec<FetchRewrite>> {
101        self.output_rewrites.get(&id)
102    }
103}