1use super::{constants::*, reference::Reference};
2use serde_json::Value;
3use std::collections::HashMap;
4
5#[derive(thiserror::Error, Debug)]
6pub enum Error {
7 #[error("Missing schema ID")]
8 MissingId(Value),
9 #[error("Invalid sub-schema ID")]
10 InvalidId(String),
11 #[error("Invalid reference")]
12 InvalidRef(#[from] super::reference::Error),
13 #[error("Missing $defs in base schema")]
14 MissingDefs(Value),
15}
16
17pub fn compose(base: &Value, sub_schemas: &[(Option<&str>, Value)]) -> Result<Value, Error> {
18 let mut result = base.clone();
19 let defs = result
20 .get_mut(DEFS_KEY)
21 .and_then(|value| value.as_object_mut())
22 .ok_or_else(|| Error::MissingDefs(base.clone()))?;
23
24 let mut prefixes = HashMap::new();
25
26 for (prefix, sub_schema) in sub_schemas {
27 let id = get_id(sub_schema)?;
28 prefixes.insert(id, prefix.map(str::to_string));
29
30 let (path_prefix, path_name) = match id.parse::<Reference>() {
31 Ok(Reference::PathOnly {
32 path_prefix,
33 path_name,
34 }) => Ok((path_prefix, path_name)),
35 _ => Err(Error::InvalidId(id.to_string())),
36 }?;
37
38 let mut new_sub_schema = sub_schema.clone();
39
40 modify_references(&mut new_sub_schema, &|old_reference| {
41 Ok(match old_reference {
42 Reference::FragmentOnly { fragment_name } => Some(Reference::new(
43 path_prefix.clone(),
44 path_name.clone(),
45 fragment_name.clone(),
46 )),
47 _ => None,
48 })
49 })?;
50
51 if let Some(top_level_def) = get_top_level_def(&new_sub_schema) {
52 defs.insert(
53 format!("{}{}", prefix.unwrap_or_default(), path_name),
54 top_level_def,
55 );
56 }
57
58 if let Some(fields) = new_sub_schema
59 .get(DEFS_KEY)
60 .and_then(|value| value.as_object())
61 {
62 for (key, value) in fields {
63 defs.insert(
64 format!("{}{}", prefix.unwrap_or_default(), key),
65 value.clone(),
66 );
67 }
68 }
69 }
70
71 modify_references(&mut result, &|old_reference| {
72 old_reference
73 .path()
74 .map(|value| {
75 let prefix = prefixes
76 .get(&value.as_ref())
77 .ok_or_else(|| Error::InvalidId(old_reference.to_string()))?;
78
79 Ok(Reference::from_fragment_name(format!(
80 "{}{}",
81 prefix
82 .as_ref()
83 .map(|value| value.as_str())
84 .unwrap_or_default(),
85 old_reference.name()
86 )))
87 })
88 .map_or(Ok(None), |value| value.map(Some))
89 })?;
90
91 Ok(result)
92}
93
94fn get_id(value: &Value) -> Result<&str, Error> {
95 value
96 .get(ID_KEY)
97 .and_then(|value| value.as_str())
98 .ok_or_else(|| Error::MissingId(value.clone()))
99}
100
101fn get_top_level_def(value: &Value) -> Option<Value> {
102 if let Some(fields) = value.as_object() {
103 if fields.keys().any(|key| key != ID_KEY && key != DEFS_KEY) {
104 let mut result = value.clone();
105 let fields = result.as_object_mut().unwrap();
106 fields.remove(DEFS_KEY);
107
108 Some(result)
109 } else {
110 None
111 }
112 } else {
113 None
114 }
115}
116
117fn modify_references<F: Fn(&Reference) -> Result<Option<Reference>, Error>>(
118 value: &mut Value,
119 f: &F,
120) -> Result<(), Error> {
121 if let Some(values) = value.as_array_mut() {
122 for value in values {
123 modify_references(value, f)?;
124 }
125 } else if let Some(fields) = value.as_object_mut() {
126 if let Some(reference) = fields.get_mut(REF_KEY) {
127 if let Some(previous_value) = reference.as_str() {
128 let previous_reference = previous_value.parse::<Reference>()?;
129
130 if let Some(new_reference) = f(&previous_reference)? {
131 *reference = Value::String(new_reference.to_string());
132 }
133 }
134 }
135
136 for value in fields.values_mut() {
137 modify_references(value, f)?;
138 }
139 }
140
141 Ok(())
142}
143
144#[cfg(test)]
145mod tests {
146 use super::*;
147 use serde_json::Value;
148
149 #[test]
150 fn test_compose() {
151 let base_schema = serde_json::from_str::<Value>(
152 r###"
153 {
154 "type": "object",
155 "properties": {
156 "foo": {
157 "$ref": "/schemas/bar"
158 },
159 "baz": {
160 "$ref": "/schemas/qux#/$defs/oof"
161 },
162 "p": {
163 "$ref": "/schemas/prefixed#/$defs/prefixed_thing"
164 }
165 },
166 "$defs": {}
167 }
168 "###,
169 )
170 .unwrap();
171
172 let sub_schema_bar = serde_json::from_str::<Value>(
173 r###"
174 {
175 "$id": "/schemas/bar",
176 "type": "integer"
177 }
178 "###,
179 )
180 .unwrap();
181
182 let sub_schema_qux = serde_json::from_str::<Value>(
183 r###"
184 {
185 "$id": "/schemas/qux",
186 "$defs": {
187 "oof": {
188 "enum": ["ABC", "DEF"]
189 }
190 }
191 }
192 "###,
193 )
194 .unwrap();
195
196 let sub_schema_top_level_and_defs = serde_json::from_str::<Value>(
197 r###"
198 {
199 "$id": "/schemas/top_level",
200 "enum": [1, 2, 3],
201 "$defs": {
202 "top-level_stuff": {
203 "type": "boolean"
204 }
205 }
206 }
207 "###,
208 )
209 .unwrap();
210
211 let sub_schema_prefixed = serde_json::from_str::<Value>(
212 r###"
213 {
214 "$id": "/schemas/prefixed",
215 "$defs": {
216 "prefixed_thing": {
217 "type": "array",
218 "items": "number"
219 }
220 }
221 }
222 "###,
223 )
224 .unwrap();
225
226 let expected = serde_json::from_str::<Value>(
227 r###"
228 {
229 "type": "object",
230 "properties": {
231 "foo": {
232 "$ref": "#/$defs/bar"
233 },
234 "baz": {
235 "$ref": "#/$defs/oof"
236 },
237 "p": {
238 "$ref": "#/$defs/abcd_prefixed_thing"
239 }
240 },
241 "$defs": {
242 "bar": {
243 "$id": "/schemas/bar",
244 "type": "integer"
245 },
246 "oof": {
247 "enum": ["ABC", "DEF"]
248 },
249 "top_level": {
250 "$id": "/schemas/top_level",
251 "enum": [1, 2, 3]
252 },
253 "top-level_stuff": {
254 "type": "boolean"
255 },
256 "abcd_prefixed_thing": {
257 "type": "array",
258 "items": "number"
259 }
260 }
261 }
262 "###,
263 )
264 .unwrap();
265
266 let composed = compose(
267 &base_schema,
268 &vec![
269 (None, sub_schema_bar),
270 (None, sub_schema_qux),
271 (None, sub_schema_top_level_and_defs),
272 (Some("abcd_"), sub_schema_prefixed),
273 ],
274 )
275 .unwrap();
276
277 assert_eq!(composed, expected);
278 }
279}