Skip to main content

hive_router_plan_executor/
context.rs

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