wip

lib/vector-config-common/src/schema/json_schema.rs

@@ -1,8 +1,8 @@
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
-use std::ops::Deref;
+use std::{iter, ops::Deref};
 
-use super::{Map, Set};
+use super::{Map, Set, DEFINITIONS_PREFIX};
 
 /// A JSON Schema.
 #[allow(clippy::large_enum_variant)]
@@ -40,6 +40,26 @@ impl Schema {
         }
     }
 
+    /// Gets a reference to the inner schema object if this schema is a JSON Schema object.
+    ///
+    /// Otherwise, `None` is returned.
+    pub fn as_object(&self) -> Option<&SchemaObject> {
+        match self {
+            Schema::Object(schema) => Some(schema),
+            _ => None,
+        }
+    }
+
+    /// Gets a mutable reference to the inner schema object if this schema is a JSON Schema object.
+    ///
+    /// Otherwise, `None` is returned.
+    pub fn as_object_mut(&mut self) -> Option<&mut SchemaObject> {
+        match self {
+            Schema::Object(schema) => Some(schema),
+            _ => None,
+        }
+    }
+
     /// Converts the given schema (if it is a boolean schema) into an equivalent schema object.
     ///
     /// If the given schema is already a schema object, this has no effect.
@@ -539,6 +559,30 @@ pub enum SingleOrVec<T> {
     Vec(Vec<T>),
 }
 
+impl<T: Clone> Extend<T> for SingleOrVec<T> {
+    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
+        match self {
+            Self::Single(item) => {
+                *self = Self::Vec(iter::once(*item.clone()).chain(iter.into_iter()).collect());
+            }
+            Self::Vec(items) => items.extend(iter),
+        }
+    }
+}
+
+impl<'a, T> IntoIterator for &'a SingleOrVec<T> {
+    type Item = &'a T;
+
+    type IntoIter = std::slice::Iter<'a, T>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        match self {
+            SingleOrVec::Single(item) => std::slice::from_ref(item.as_ref()).into_iter(),
+            SingleOrVec::Vec(items) => items.as_slice().into_iter(),
+        }
+    }
+}
+
 impl<T> From<T> for SingleOrVec<T> {
     fn from(single: T) -> Self {
         SingleOrVec::Single(Box::new(single))
@@ -567,3 +611,14 @@ fn is_none_or_default_true(field: &Option<Box<Schema>>) -> bool {
         Some(value) => matches!(value.as_ref(), Schema::Bool(true)),
     }
 }
+
+pub fn get_cleaned_schema_reference(schema_ref: &str) -> &str {
+    if schema_ref.starts_with(DEFINITIONS_PREFIX) {
+        &schema_ref[DEFINITIONS_PREFIX.len()..]
+    } else {
+        panic!(
+            "Tried to clean schema reference that does not start with the definition prefix: {}",
+            schema_ref
+        );
+    }
+}

lib/vector-config-common/src/schema/visit.rs

@@ -1,4 +1,4 @@
-use super::{Map, RootSchema, Schema, SchemaObject, SingleOrVec, DEFINITIONS_PREFIX};
+use super::{get_cleaned_schema_reference, Map, RootSchema, Schema, SchemaObject, SingleOrVec};
 
 /// Trait used to recursively modify a constructed schema and its subschemas.
 pub trait Visitor: std::fmt::Debug {
@@ -136,14 +136,14 @@ pub fn with_resolved_schema_reference<F>(
     F: FnOnce(&mut Map<String, Schema>, &mut Schema),
 {
     if let Some(reference) = schema.reference.as_ref() {
-        let schema_def_key = reference.replace(DEFINITIONS_PREFIX, "");
+        let schema_def_key = get_cleaned_schema_reference(reference);
         let mut referenced_schema = definitions
-            .get(&schema_def_key)
+            .get(schema_def_key)
             .cloned()
             .expect("schema reference should exist");
 
         f(definitions, &mut referenced_schema);
 
-        definitions.insert(schema_def_key, referenced_schema);
+        definitions.insert(schema_def_key.to_string(), referenced_schema);
     }
 }