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

use std::collections::BTreeMap;

use serde_json::{Map, Value};

use crate::{Schema, SchemaKind};

impl Schema {
    pub fn from_json_sample(value: &Value) -> Self {
        match value {
            Value::Null => Self {
                kind: SchemaKind::Null,
                description: None,
                required: true,
                nullable: true,
                default_value: None,
                fields: BTreeMap::new(),
                items: None,
                enum_values: Vec::new(),
                variants: Vec::new(),
            },
            Value::Bool(_) => scalar_schema(SchemaKind::Boolean),
            Value::String(_) => scalar_schema(SchemaKind::String),
            Value::Number(number) if number.is_i64() || number.is_u64() => {
                scalar_schema(SchemaKind::Integer)
            }
            Value::Number(_) => scalar_schema(SchemaKind::Number),
            Value::Array(items) => array_schema(items),
            Value::Object(fields) => object_schema(fields),
        }
    }
}

fn scalar_schema(kind: SchemaKind) -> Schema {
    Schema {
        kind,
        description: None,
        required: true,
        nullable: false,
        default_value: None,
        fields: BTreeMap::new(),
        items: None,
        enum_values: Vec::new(),
        variants: Vec::new(),
    }
}

fn array_schema(items: &[Value]) -> Schema {
    Schema {
        kind: SchemaKind::Array,
        description: None,
        required: true,
        nullable: false,
        default_value: None,
        fields: BTreeMap::new(),
        items: items
            .first()
            .map(|item| Box::new(Schema::from_json_sample(item))),
        enum_values: Vec::new(),
        variants: Vec::new(),
    }
}

fn object_schema(fields: &Map<String, Value>) -> Schema {
    Schema {
        kind: SchemaKind::Object,
        description: None,
        required: true,
        nullable: false,
        default_value: None,
        fields: fields
            .iter()
            .map(|(name, value)| (name.clone(), Schema::from_json_sample(value)))
            .collect(),
        items: None,
        enum_values: Vec::new(),
        variants: Vec::new(),
    }
}

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

    use serde_json::json;

    use crate::{Schema, SchemaKind};

    #[test]
    fn infers_object_schema_from_json_sample() {
        let schema = Schema::from_json_sample(&json!({
            "lead": {
                "email": "user@example.com",
                "active": true
            },
            "score": 42
        }));

        assert_eq!(schema.kind, SchemaKind::Object);
        assert_eq!(
            schema.field("lead").map(|field| field.kind.clone()),
            Some(SchemaKind::Object)
        );
        assert_eq!(
            schema
                .field("lead")
                .and_then(|field| field.field("email"))
                .map(|field| field.kind.clone()),
            Some(SchemaKind::String)
        );
        assert_eq!(
            schema.field("score").map(|field| field.kind.clone()),
            Some(SchemaKind::Integer)
        );
    }

    #[test]
    fn infers_array_item_schema_from_first_sample_item() {
        let schema = Schema::from_json_sample(&json!([
            {
                "id": "lead_01",
                "tags": ["warm", "priority"]
            }
        ]));

        assert_eq!(schema.kind, SchemaKind::Array);
        assert_eq!(
            schema
                .items
                .as_deref()
                .and_then(|item| item.field("id"))
                .map(|field| field.kind.clone()),
            Some(SchemaKind::String)
        );
        assert_eq!(
            schema
                .items
                .as_deref()
                .and_then(|item| item.field("tags"))
                .map(|field| field.kind.clone()),
            Some(SchemaKind::Array)
        );
    }

    #[test]
    fn infers_null_schema_for_null_sample() {
        let schema = Schema::from_json_sample(&serde_json::Value::Null);

        assert_eq!(
            schema,
            Schema {
                kind: SchemaKind::Null,
                description: None,
                required: true,
                nullable: true,
                default_value: None,
                fields: BTreeMap::new(),
                items: None,
                enum_values: Vec::new(),
                variants: Vec::new(),
            }
        );
    }
}