vector_config_common/schema/
visit.rs

1use tracing::debug;
2
3use super::{get_cleaned_schema_reference, Map, RootSchema, Schema, SchemaObject, SingleOrVec};
4
5/// Trait used to recursively modify a constructed schema and its subschemas.
6pub trait Visitor: std::fmt::Debug {
7    /// Override this method to modify a [`RootSchema`] and (optionally) its subschemas.
8    ///
9    /// When overriding this method, you will usually want to call the [`visit_root_schema`] function to visit subschemas.
10    fn visit_root_schema(&mut self, root: &mut RootSchema) {
11        visit_root_schema(self, root);
12    }
13
14    /// Override this method to modify a [`Schema`] and (optionally) its subschemas.
15    ///
16    /// When overriding this method, you will usually want to call the [`visit_schema`] function to visit subschemas.
17    fn visit_schema(&mut self, definitions: &mut Map<String, Schema>, schema: &mut Schema) {
18        visit_schema(self, definitions, schema);
19    }
20
21    /// Override this method to modify a [`SchemaObject`] and (optionally) its subschemas.
22    ///
23    /// When overriding this method, you will usually want to call the [`visit_schema_object`] function to visit subschemas.
24    fn visit_schema_object(
25        &mut self,
26        definitions: &mut Map<String, Schema>,
27        schema: &mut SchemaObject,
28    ) {
29        visit_schema_object(self, definitions, schema);
30    }
31}
32
33/// Visits all subschemas of the [`RootSchema`].
34pub fn visit_root_schema<V: Visitor + ?Sized>(v: &mut V, root: &mut RootSchema) {
35    v.visit_schema_object(&mut root.definitions, &mut root.schema);
36}
37
38/// Visits all subschemas of the [`Schema`].
39pub fn visit_schema<V: Visitor + ?Sized>(
40    v: &mut V,
41    definitions: &mut Map<String, Schema>,
42    schema: &mut Schema,
43) {
44    if let Schema::Object(schema) = schema {
45        v.visit_schema_object(definitions, schema);
46    }
47}
48
49/// Visits all subschemas of the [`SchemaObject`].
50pub fn visit_schema_object<V: Visitor + ?Sized>(
51    v: &mut V,
52    definitions: &mut Map<String, Schema>,
53    schema: &mut SchemaObject,
54) {
55    if schema.reference.is_some() {
56        with_resolved_schema_reference(
57            definitions,
58            schema,
59            |defs, schema_ref, referenced_schema| {
60                debug!(referent = schema_ref, "Visiting schema reference.");
61
62                v.visit_schema(defs, referenced_schema);
63            },
64        )
65    }
66
67    if let Some(sub) = &mut schema.subschemas {
68        visit_vec(v, definitions, &mut sub.all_of);
69        visit_vec(v, definitions, &mut sub.any_of);
70        visit_vec(v, definitions, &mut sub.one_of);
71        visit_box(v, definitions, &mut sub.not);
72        visit_box(v, definitions, &mut sub.if_schema);
73        visit_box(v, definitions, &mut sub.then_schema);
74        visit_box(v, definitions, &mut sub.else_schema);
75    }
76
77    if let Some(arr) = &mut schema.array {
78        visit_single_or_vec(v, definitions, &mut arr.items);
79        visit_box(v, definitions, &mut arr.additional_items);
80        visit_box(v, definitions, &mut arr.contains);
81    }
82
83    if let Some(obj) = &mut schema.object {
84        visit_map_values(v, definitions, &mut obj.properties);
85        visit_map_values(v, definitions, &mut obj.pattern_properties);
86        visit_box(v, definitions, &mut obj.additional_properties);
87        visit_box(v, definitions, &mut obj.property_names);
88    }
89}
90
91fn visit_box<V: Visitor + ?Sized>(
92    v: &mut V,
93    definitions: &mut Map<String, Schema>,
94    target: &mut Option<Box<Schema>>,
95) {
96    if let Some(s) = target {
97        v.visit_schema(definitions, s);
98    }
99}
100
101fn visit_vec<V: Visitor + ?Sized>(
102    v: &mut V,
103    definitions: &mut Map<String, Schema>,
104    target: &mut Option<Vec<Schema>>,
105) {
106    if let Some(vec) = target {
107        for s in vec {
108            v.visit_schema(definitions, s);
109        }
110    }
111}
112
113fn visit_map_values<V: Visitor + ?Sized>(
114    v: &mut V,
115    definitions: &mut Map<String, Schema>,
116    target: &mut Map<String, Schema>,
117) {
118    for s in target.values_mut() {
119        v.visit_schema(definitions, s);
120    }
121}
122
123fn visit_single_or_vec<V: Visitor + ?Sized>(
124    v: &mut V,
125    definitions: &mut Map<String, Schema>,
126    target: &mut Option<SingleOrVec<Schema>>,
127) {
128    match target {
129        None => {}
130        Some(SingleOrVec::Single(s)) => v.visit_schema(definitions, s),
131        Some(SingleOrVec::Vec(vec)) => {
132            for s in vec {
133                v.visit_schema(definitions, s);
134            }
135        }
136    }
137}
138
139pub fn with_resolved_schema_reference<F>(
140    definitions: &mut Map<String, Schema>,
141    schema: &mut SchemaObject,
142    f: F,
143) where
144    F: FnOnce(&mut Map<String, Schema>, &str, &mut Schema),
145{
146    if let Some(reference) = schema.reference.as_ref() {
147        let schema_def_key = get_cleaned_schema_reference(reference);
148        let mut referenced_schema = definitions
149            .get(schema_def_key)
150            .cloned()
151            .expect("schema reference should exist");
152
153        f(definitions, schema_def_key, &mut referenced_schema);
154
155        definitions.insert(schema_def_key.to_string(), referenced_schema);
156    }
157}