crank/crates/crank-schema/src/model.rs

use std::collections::BTreeMap;

use serde::{Deserialize, Serialize};
use serde_json::Value;

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum SchemaKind {
    Object,
    Array,
    String,
    Integer,
    Number,
    Boolean,
    Enum,
    #[serde(rename = "null")]
    Null,
    Oneof,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct Schema {
    #[serde(rename = "type")]
    pub kind: SchemaKind,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,
    #[serde(default)]
    pub required: bool,
    #[serde(default)]
    pub nullable: bool,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub default_value: Option<Value>,
    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
    pub fields: BTreeMap<String, Schema>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub items: Option<Box<Schema>>,
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub enum_values: Vec<String>,
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub variants: Vec<Schema>,
}

impl Schema {
    pub fn is_object(&self) -> bool {
        self.kind == SchemaKind::Object
    }

    pub fn field(&self, name: &str) -> Option<&Schema> {
        self.fields.get(name)
    }

    pub fn has_required_fields(&self) -> bool {
        self.fields.values().any(|field| field.required)
    }

    pub fn expected_type_name(&self) -> &'static str {
        match self.kind {
            SchemaKind::Object => "object",
            SchemaKind::Array => "array",
            SchemaKind::String => "string",
            SchemaKind::Integer => "integer",
            SchemaKind::Number => "number",
            SchemaKind::Boolean => "boolean",
            SchemaKind::Enum => "enum",
            SchemaKind::Null => "null",
            SchemaKind::Oneof => "oneof",
        }
    }
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeMap;

    use super::{Schema, SchemaKind};

    #[test]
    fn object_schema_exposes_fields() {
        let mut fields = BTreeMap::new();
        fields.insert(
            "email".to_owned(),
            Schema {
                kind: SchemaKind::String,
                description: None,
                required: true,
                nullable: false,
                default_value: None,
                fields: BTreeMap::new(),
                items: None,
                enum_values: Vec::new(),
                variants: Vec::new(),
            },
        );

        let schema = Schema {
            kind: SchemaKind::Object,
            description: Some("User input".to_owned()),
            required: true,
            nullable: false,
            default_value: None,
            fields,
            items: None,
            enum_values: Vec::new(),
            variants: Vec::new(),
        };

        assert!(schema.is_object());
        assert!(schema.has_required_fields());
        assert_eq!(
            schema.field("email").map(|field| field.required),
            Some(true)
        );
    }

    #[test]
    fn schema_serializes_type_field() {
        let schema = Schema {
            kind: SchemaKind::Array,
            description: None,
            required: false,
            nullable: false,
            default_value: None,
            fields: BTreeMap::new(),
            items: Some(Box::new(Schema {
                kind: SchemaKind::String,
                description: None,
                required: false,
                nullable: false,
                default_value: None,
                fields: BTreeMap::new(),
                items: None,
                enum_values: Vec::new(),
                variants: Vec::new(),
            })),
            enum_values: Vec::new(),
            variants: Vec::new(),
        };

        let value = serde_json::to_value(schema).unwrap();

        assert_eq!(value["type"], "array");
        assert_eq!(value["items"]["type"], "string");
    }

    #[test]
    fn schema_roundtrips_through_yaml() {
        let schema = Schema {
            kind: SchemaKind::Object,
            description: Some("Lead output".to_owned()),
            required: true,
            nullable: false,
            default_value: None,
            fields: BTreeMap::from([(
                "id".to_owned(),
                Schema {
                    kind: SchemaKind::String,
                    description: None,
                    required: true,
                    nullable: false,
                    default_value: None,
                    fields: BTreeMap::new(),
                    items: None,
                    enum_values: Vec::new(),
                    variants: Vec::new(),
                },
            )]),
            items: None,
            enum_values: Vec::new(),
            variants: Vec::new(),
        };

        let yaml = serde_yaml::to_string(&schema).unwrap();
        let restored: Schema = serde_yaml::from_str(&yaml).unwrap();

        assert!(yaml.contains("type: object"));
        assert_eq!(restored, schema);
    }
}