Skip to main content

hive_router_plan_executor/execution/
rewrites.rs

1use hive_router_query_planner::planner::plan_nodes::{
2    FetchNodePathSegment, FetchRewrite, KeyRenamer, ValueSetter,
3};
4
5use crate::{
6    introspection::schema::PossibleTypes, response::value::Value,
7    utils::consts::TYPENAME_FIELD_NAME,
8};
9
10pub trait FetchRewriteExt {
11    fn rewrite<'a>(&'a self, possible_types: &PossibleTypes, value: &mut Value<'a>);
12}
13
14impl FetchRewriteExt for FetchRewrite {
15    fn rewrite<'a>(&'a self, possible_types: &PossibleTypes, value: &mut Value<'a>) {
16        match self {
17            FetchRewrite::KeyRenamer(key_renamer) => key_renamer.apply(possible_types, value),
18            FetchRewrite::ValueSetter(value_setter) => value_setter.apply(possible_types, value),
19        }
20    }
21}
22
23trait RewriteApplier {
24    fn apply<'a>(&'a self, possible_types: &PossibleTypes, value: &mut Value<'a>);
25    fn apply_path<'a>(
26        &'a self,
27        possible_types: &PossibleTypes,
28        value: &mut Value<'a>,
29        path: &'a [FetchNodePathSegment],
30    );
31}
32
33impl RewriteApplier for KeyRenamer {
34    fn apply<'a>(&'a self, possible_types: &PossibleTypes, value: &mut Value<'a>) {
35        self.apply_path(possible_types, value, &self.path)
36    }
37    fn apply_path<'a>(
38        &'a self,
39        possible_types: &PossibleTypes,
40        value: &mut Value<'a>,
41        path: &'a [FetchNodePathSegment],
42    ) {
43        let current_segment = &path[0];
44        let remaining_path = &path[1..];
45        match value {
46            Value::Array(arr) => {
47                for item in arr {
48                    self.apply_path(possible_types, item, path);
49                }
50            }
51            Value::Object(obj) => match current_segment {
52                FetchNodePathSegment::TypenameEquals(type_condition) => {
53                    let type_name = obj
54                        .iter()
55                        .find(|(key, _)| key == &TYPENAME_FIELD_NAME)
56                        .and_then(|(_, val)| val.as_str())
57                        .unwrap_or(type_condition);
58                    if possible_types.entity_satisfies_type_condition(type_name, type_condition) {
59                        self.apply_path(possible_types, value, remaining_path)
60                    }
61                }
62                FetchNodePathSegment::Key(field_name) => {
63                    if remaining_path.is_empty() {
64                        if field_name != &self.rename_key_to {
65                            if let Some((key, _)) =
66                                obj.iter_mut().find(|(key, _)| key == field_name)
67                            {
68                                *key = self.rename_key_to.as_str()
69                            }
70                        }
71                    } else if let Some(data) = obj.iter_mut().find(|r| r.0 == field_name) {
72                        self.apply_path(possible_types, &mut data.1, remaining_path)
73                    }
74                }
75            },
76            // If the value is not an object or an array, we can't apply the rewrite.
77            _ => (),
78        }
79    }
80}
81
82impl RewriteApplier for ValueSetter {
83    fn apply<'a>(&'a self, possible_types: &PossibleTypes, data: &mut Value<'a>) {
84        self.apply_path(possible_types, data, &self.path)
85    }
86
87    fn apply_path<'a>(
88        &'a self,
89        possible_types: &PossibleTypes,
90        data: &mut Value<'a>,
91        path: &'a [FetchNodePathSegment],
92    ) {
93        if path.is_empty() {
94            *data = Value::String(self.set_value_to.as_str().into());
95            return;
96        }
97
98        match data {
99            Value::Array(arr) => {
100                for data in arr {
101                    self.apply_path(possible_types, data, path);
102                }
103            }
104            Value::Object(map) => {
105                let current_segment = &path[0];
106                let remaining_path = &path[1..];
107
108                match current_segment {
109                    FetchNodePathSegment::TypenameEquals(type_condition) => {
110                        let type_name = map
111                            .iter()
112                            .find(|(key, _)| key == &TYPENAME_FIELD_NAME)
113                            .and_then(|(_, val)| val.as_str())
114                            .unwrap_or(type_condition);
115                        if possible_types.entity_satisfies_type_condition(type_name, type_condition)
116                        {
117                            self.apply_path(possible_types, data, remaining_path)
118                        }
119                    }
120                    FetchNodePathSegment::Key(field_name) => {
121                        if let Some(data) = map.iter_mut().find(|r| r.0 == field_name) {
122                            self.apply_path(possible_types, &mut data.1, remaining_path)
123                        }
124                    }
125                }
126            }
127            // If the value is not an object or an array, we can't apply the rewrite.
128            _ => (),
129        }
130    }
131}