Embed and expose examples in Rust

Cargo.toml

@@ -8,7 +8,15 @@ version = "0.1.39"
 edition = "2021"
 
 build = "rust/build.rs"
-include = ["/rust/*.rs", "/Cargo.toml", "/schemas/**", "/topics/**", "/README.md", "/LICENSE.md"]
+include = [
+    "/examples/**",
+    "/rust/*.rs",
+    "/schemas/**",
+    "/topics/**",
+    "/Cargo.toml",
+    "/LICENSE.md",
+    "/README.md",
+]
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
@@ -27,10 +35,10 @@ default = ["type_generation"]
 type_generation = ["prettyplease", "schemars", "syn", "typify", "regress"]
 
 [dependencies]
-serde = { version = "1.0", features = ["derive"]}
+serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 serde_yaml = "0.9"
-regress = { version = "0.5.0", optional = true }  # required by generated code
+regress = { version = "0.5.0", optional = true }   # required by generated code
 jsonschema = "0.17.1"
 thiserror = "1.0.49"
 

rust/build.rs

@@ -48,34 +48,75 @@ fn generate_embedded_data() -> String {
     let manifest_root = std::env::var("CARGO_MANIFEST_DIR").unwrap_or_default();
     let topics = enumerate_dir("topics");
     let schemas = enumerate_dir("schemas");
+    let examples = enumerate_dir("examples");
 
     let mut code = String::new();
     use std::fmt::Write;
     let w = &mut code;
 
-    writeln!(w, "const TOPICS: &[(&'static str, &'static str)] = &[").unwrap();
+    writeln!(w, "const TOPICS: &[(&str, &str)] = &[").unwrap();
     for topic in topics {
         let name = topic.strip_suffix(".yaml").unwrap();
         let path = format!("{manifest_root}/topics/{topic}");
         writeln!(w, "    ({name:?}, include_str!({path:?})),").unwrap();
     }
     writeln!(w, "];\n").unwrap();
 
-    writeln!(w, "const SCHEMAS: &[(&'static str, &'static str)] = &[").unwrap();
+    writeln!(w, "const SCHEMAS: &[(&str, &str)] = &[").unwrap();
     for schema in schemas {
         let path = format!("{manifest_root}/schemas/{schema}");
         writeln!(w, "    ({schema:?}, include_str!({path:?})),").unwrap();
     }
     writeln!(w, "];").unwrap();
 
+    let mut last_prefix = None;
+    writeln!(w, "const EXAMPLES: &[(&str, &[&[u8]])] = &[").unwrap();
+    for example in &examples {
+        let last_slash = example.rfind('/').unwrap();
+        let (prefix, _name) = example.split_at(last_slash + 1);
+        match last_prefix {
+            None => {
+                writeln!(w, "    ({prefix:?}, &[").unwrap();
+            }
+            Some(last_prefix) if last_prefix != prefix => {
+                writeln!(w, "    ]),").unwrap();
+                writeln!(w, "    ({prefix:?}, &[").unwrap();
+            }
+            _ => {}
+        }
+        last_prefix = Some(prefix);
+        let path = format!("{manifest_root}/examples/{example}");
+        writeln!(w, "        include_bytes!({path:?}),").unwrap();
+    }
+    writeln!(w, "    ]),").unwrap();
+    writeln!(w, "];").unwrap();
+
     code
 }
 
 fn enumerate_dir(dir: &str) -> Vec<String> {
-    let mut files: Vec<_> = std::fs::read_dir(dir)
-        .unwrap()
-        .map(|entry| entry.unwrap().file_name().into_string().unwrap())
-        .collect();
+    let mut files = vec![];
+
+    fn collect_files(files: &mut Vec<String>, dir: &Path, prefix: &str) {
+        for entry in std::fs::read_dir(dir).unwrap() {
+            let entry = entry.unwrap();
+            let name = entry.file_name().into_string().unwrap();
+            let name = if prefix.is_empty() {
+                name
+            } else {
+                format!("{prefix}/{name}")
+            };
+            let ty = entry.file_type().unwrap();
+            if ty.is_dir() {
+                collect_files(files, &entry.path(), &name)
+            } else if ty.is_file() {
+                files.push(name);
+            }
+        }
+    }
+
+    collect_files(&mut files, dir.as_ref(), "");
+
     files.sort();
     files
 }

rust/lib.rs

@@ -58,9 +58,9 @@ struct TopicData {
     schemas: Vec<TopicSchema>,
 }
 
-fn find_entry<'s>(slice: &[(&str, &'s str)], key: &str) -> Option<&'s str> {
+fn find_entry<'s, T>(slice: &'s [(&str, T)], key: &str) -> Option<&'s T> {
     let idx = slice.binary_search_by_key(&key, |&(name, _)| name).ok()?;
-    Some(slice.get(idx)?.1)
+    Some(&slice.get(idx)?.1)
 }
 
 impl TopicData {
@@ -77,6 +77,7 @@ pub struct Schema {
     pub compatibility_mode: CompatibilityMode,
     schema: &'static str,
     compiled_json_schema: JSONSchema,
+    examples: &'static [&'static [u8]],
 }
 
 impl PartialEq for Schema {
@@ -103,6 +104,11 @@ impl Schema {
     pub fn raw_schema(&self) -> &str {
         self.schema
     }
+
+    /// Returns a list of examples for this schema.
+    pub fn examples(&self) -> &[&[u8]] {
+        self.examples
+    }
 }
 
 fn get_topic_schema(topic: &str, version: Option<u16>) -> Result<TopicSchema, SchemaError> {
@@ -146,22 +152,30 @@ pub fn get_schema(topic: &str, version: Option<u16>) -> Result<Schema, SchemaErr
     let s = serde_json::from_str(schema).map_err(|_| SchemaError::InvalidSchema)?;
     let compiled_json_schema = JSONSchema::compile(&s).map_err(|_| SchemaError::InvalidSchema)?;
 
+    // FIXME(swatinem): This assumes that there is only a single `examples` entry in the definition.
+    // If we would want to support multiple, we would have to either merge those in code generation,
+    // or rather use a `fn examples() -> impl Iterator`.
+    let examples = schema_metadata
+        .examples
+        .first()
+        .and_then(|example| find_entry(EXAMPLES, example))
+        .copied()
+        .unwrap_or_default();
+
     Ok(Schema {
         version: schema_metadata.version,
         schema_type: schema_metadata.schema_type,
         compatibility_mode: schema_metadata.compatibility_mode,
         schema,
         compiled_json_schema,
+        examples,
     })
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
 
-    use std::fs::read_to_string;
-    use std::path::Path;
-
     #[test]
     fn test_get_schema() {
         assert!(matches!(
@@ -178,18 +192,16 @@ mod tests {
 
     #[test]
     fn test_validate() {
-        let topic = "snuba-queries";
-        let topic_schema = get_topic_schema(topic, None).unwrap();
-        let schema = get_schema(topic, None).unwrap();
-        let example_dir = topic_schema.examples.first().unwrap();
-        let example_path = Path::new(env!("CARGO_MANIFEST_DIR"))
-            .join(format!("examples/{}{}", example_dir, "snuba-queries1.json"));
-        let example = read_to_string(example_path).unwrap();
-        schema.validate_json(example.as_bytes()).unwrap();
-
-        let invalid = "{}";
+        let schema = get_schema("snuba-queries", None).unwrap();
+
+        let examples = schema.examples();
+        assert!(!examples.is_empty());
+        for example in examples {
+            schema.validate_json(example).unwrap();
+        }
+
         assert!(matches!(
-            schema.validate_json(invalid.as_bytes()),
+            schema.validate_json(b"{}"),
             Err(SchemaError::InvalidMessage)
         ));
     }