vector_config/schema/visitors/
scoped_visit.rs

1use std::collections::VecDeque;
2
3use vector_config_common::schema::{visit::Visitor, *};
4
5/// A schema reference which can refer to either a schema definition or the root schema itself.
6#[derive(Clone, Debug, Eq, Hash, PartialEq)]
7pub enum SchemaReference {
8    /// A defined schema.
9    Definition(String),
10
11    /// The root schema itself.
12    Root,
13}
14
15impl std::fmt::Display for SchemaReference {
16    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
17        match self {
18            Self::Definition(name) => write!(f, "{name}"),
19            Self::Root => write!(f, "<root>"),
20        }
21    }
22}
23
24impl<'a, T> From<&'a T> for SchemaReference
25where
26    T: AsRef<str> + ?Sized,
27{
28    fn from(value: &'a T) -> Self {
29        Self::Definition(value.as_ref().to_string())
30    }
31}
32
33impl AsRef<str> for SchemaReference {
34    fn as_ref(&self) -> &str {
35        match self {
36            Self::Definition(name) => name.as_str(),
37            Self::Root => "<root>",
38        }
39    }
40}
41
42/// The schema scope stack is used to understand where in the visiting of a root schema the visitor
43/// currently is. All visiting will inherently start at the root, and continue to exist in the root
44/// scope until a schema reference is resolved, and the visitor visits the resolved schema. Once
45/// that happens, the resolved schema scope is pushed onto the stack, such that the scope is now the
46/// schema reference. This can happen multiple times as subsequent schema references are resolved.
47/// Once the visitor recurses back out of the resolved schema, it is popped from the stack.
48#[derive(Debug, Default)]
49pub struct SchemaScopeStack {
50    stack: VecDeque<SchemaReference>,
51}
52
53impl SchemaScopeStack {
54    pub fn push<S: Into<SchemaReference>>(&mut self, scope: S) {
55        self.stack.push_front(scope.into());
56    }
57
58    pub fn pop(&mut self) -> Option<SchemaReference> {
59        self.stack.pop_front()
60    }
61
62    pub fn current(&self) -> Option<&SchemaReference> {
63        self.stack.front()
64    }
65}
66
67pub trait ScopedVisitor: Visitor {
68    fn push_schema_scope<S: Into<SchemaReference>>(&mut self, scope: S);
69
70    fn pop_schema_scope(&mut self);
71
72    fn get_current_schema_scope(&self) -> &SchemaReference;
73}
74
75pub fn visit_schema_object_scoped<SV: ScopedVisitor + ?Sized>(
76    sv: &mut SV,
77    definitions: &mut Map<String, Schema>,
78    schema: &mut SchemaObject,
79) {
80    if let Some(reference) = schema.reference.as_ref() {
81        let schema_def_key = get_cleaned_schema_reference(reference);
82        let mut referenced_schema = definitions
83            .get(schema_def_key)
84            .cloned()
85            .expect("schema reference should exist");
86
87        if let Schema::Object(referenced_schema) = &mut referenced_schema {
88            sv.push_schema_scope(schema_def_key);
89
90            sv.visit_schema_object(definitions, referenced_schema);
91
92            sv.pop_schema_scope();
93        }
94
95        definitions.insert(schema_def_key.to_string(), referenced_schema);
96    }
97
98    if let Some(sub) = &mut schema.subschemas {
99        visit_vec_scoped(sv, definitions, &mut sub.all_of);
100        visit_vec_scoped(sv, definitions, &mut sub.any_of);
101        visit_vec_scoped(sv, definitions, &mut sub.one_of);
102        visit_box_scoped(sv, definitions, &mut sub.not);
103        visit_box_scoped(sv, definitions, &mut sub.if_schema);
104        visit_box_scoped(sv, definitions, &mut sub.then_schema);
105        visit_box_scoped(sv, definitions, &mut sub.else_schema);
106    }
107
108    if let Some(arr) = &mut schema.array {
109        visit_single_or_vec_scoped(sv, definitions, &mut arr.items);
110        visit_box_scoped(sv, definitions, &mut arr.additional_items);
111        visit_box_scoped(sv, definitions, &mut arr.contains);
112    }
113
114    if let Some(obj) = &mut schema.object {
115        visit_map_values_scoped(sv, definitions, &mut obj.properties);
116        visit_map_values_scoped(sv, definitions, &mut obj.pattern_properties);
117        visit_box_scoped(sv, definitions, &mut obj.additional_properties);
118        visit_box_scoped(sv, definitions, &mut obj.property_names);
119    }
120}
121
122fn visit_box_scoped<SV: ScopedVisitor + ?Sized>(
123    sv: &mut SV,
124    definitions: &mut Map<String, Schema>,
125    target: &mut Option<Box<Schema>>,
126) {
127    if let Some(s) = target {
128        if let Schema::Object(s) = s.as_mut() {
129            sv.visit_schema_object(definitions, s);
130        }
131    }
132}
133
134fn visit_vec_scoped<SV: ScopedVisitor + ?Sized>(
135    sv: &mut SV,
136    definitions: &mut Map<String, Schema>,
137    target: &mut Option<Vec<Schema>>,
138) {
139    if let Some(vec) = target {
140        for s in vec {
141            if let Schema::Object(s) = s {
142                sv.visit_schema_object(definitions, s);
143            }
144        }
145    }
146}
147
148fn visit_map_values_scoped<SV: ScopedVisitor + ?Sized>(
149    sv: &mut SV,
150    definitions: &mut Map<String, Schema>,
151    target: &mut Map<String, Schema>,
152) {
153    for s in target.values_mut() {
154        if let Schema::Object(s) = s {
155            sv.visit_schema_object(definitions, s);
156        }
157    }
158}
159
160fn visit_single_or_vec_scoped<SV: ScopedVisitor + ?Sized>(
161    sv: &mut SV,
162    definitions: &mut Map<String, Schema>,
163    target: &mut Option<SingleOrVec<Schema>>,
164) {
165    match target {
166        None => {}
167        Some(SingleOrVec::Single(s)) => {
168            if let Schema::Object(s) = s.as_mut() {
169                sv.visit_schema_object(definitions, s);
170            }
171        }
172        Some(SingleOrVec::Vec(vec)) => {
173            for s in vec {
174                if let Schema::Object(s) = s {
175                    sv.visit_schema_object(definitions, s);
176                }
177            }
178        }
179    }
180}