crank/crates/mcpaas-proto/src/convert/to_schema.rs

use std::collections::BTreeMap;

use mcpaas_schema::{Schema, SchemaKind};

use crate::{
    errors::ProtoError,
    model::{
        ProtoEnum, ProtoField, ProtoFieldCardinality, ProtoFieldType, ProtoMapField, ProtoMessage,
        ProtoOneof, ProtoScalarKind,
    },
};

pub fn message_to_schema(message: &ProtoMessage) -> Result<Schema, ProtoError> {
    let mut fields = BTreeMap::new();

    for field in &message.fields {
        let (name, schema) = convert_field(field)?;
        fields.insert(name, schema);
    }

    for oneof in &message.oneofs {
        let (name, schema) = convert_oneof(oneof)?;
        fields.insert(name, schema);
    }

    Ok(object_schema(fields))
}

fn convert_field(field: &ProtoField) -> Result<(String, Schema), ProtoError> {
    let mut schema = convert_field_type(&field.field_type)?;

    schema.required = matches!(field.cardinality, ProtoFieldCardinality::Required);
    schema.nullable = false;

    if matches!(field.cardinality, ProtoFieldCardinality::Repeated) {
        schema = Schema {
            kind: SchemaKind::Array,
            description: None,
            required: false,
            nullable: false,
            default_value: None,
            fields: BTreeMap::new(),
            items: Some(Box::new(schema)),
            enum_values: Vec::new(),
            variants: Vec::new(),
        };
    }

    Ok((field.name.clone(), schema))
}

fn convert_field_type(field_type: &ProtoFieldType) -> Result<Schema, ProtoError> {
    match field_type {
        ProtoFieldType::Scalar { scalar } => Ok(scalar_schema(scalar)),
        ProtoFieldType::Message { message } => message_to_schema(message),
        ProtoFieldType::Enum { enumeration } => Ok(enum_schema(enumeration)),
        ProtoFieldType::Map { map } => convert_map(map),
    }
}

fn convert_map(map: &ProtoMapField) -> Result<Schema, ProtoError> {
    let entry = object_schema(BTreeMap::from([
        ("key".to_owned(), scalar_schema(&map.key)),
        ("value".to_owned(), convert_field_type(&map.value)?),
    ]));

    Ok(Schema {
        kind: SchemaKind::Array,
        description: None,
        required: false,
        nullable: false,
        default_value: None,
        fields: BTreeMap::new(),
        items: Some(Box::new(entry)),
        enum_values: Vec::new(),
        variants: Vec::new(),
    })
}

fn convert_oneof(oneof: &ProtoOneof) -> Result<(String, Schema), ProtoError> {
    if oneof.variants.is_empty() {
        return Err(ProtoError::EmptyOneof {
            oneof_name: oneof.name.clone(),
        });
    }

    let mut variants = Vec::with_capacity(oneof.variants.len());

    for field in &oneof.variants {
        let (field_name, field_schema) = convert_field(field)?;
        variants.push(object_schema(BTreeMap::from([(field_name, field_schema)])));
    }

    Ok((
        oneof.name.clone(),
        Schema {
            kind: SchemaKind::Oneof,
            description: None,
            required: false,
            nullable: false,
            default_value: None,
            fields: BTreeMap::new(),
            items: None,
            enum_values: Vec::new(),
            variants,
        },
    ))
}

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

fn enum_schema(enumeration: &ProtoEnum) -> Schema {
    Schema {
        kind: SchemaKind::Enum,
        description: None,
        required: false,
        nullable: false,
        default_value: None,
        fields: BTreeMap::new(),
        items: None,
        enum_values: enumeration.values.clone(),
        variants: Vec::new(),
    }
}

fn scalar_schema(kind: &ProtoScalarKind) -> Schema {
    let schema_kind = match kind {
        ProtoScalarKind::String | ProtoScalarKind::Bytes => SchemaKind::String,
        ProtoScalarKind::Bool => SchemaKind::Boolean,
        ProtoScalarKind::Int32
        | ProtoScalarKind::Int64
        | ProtoScalarKind::Uint32
        | ProtoScalarKind::Uint64 => SchemaKind::Integer,
        ProtoScalarKind::Float | ProtoScalarKind::Double => SchemaKind::Number,
    };

    Schema {
        kind: schema_kind,
        description: None,
        required: false,
        nullable: false,
        default_value: None,
        fields: BTreeMap::new(),
        items: None,
        enum_values: Vec::new(),
        variants: Vec::new(),
    }
}

#[cfg(test)]
mod tests {
    use mcpaas_schema::SchemaKind;

    use crate::{
        ProtoEnum, ProtoField, ProtoFieldCardinality, ProtoFieldType, ProtoMapField, ProtoMessage,
        ProtoMethod, ProtoOneof, ProtoScalarKind, ProtoService, message_to_schema,
    };

    #[test]
    fn service_filters_unary_methods() {
        let unary = ProtoMethod {
            name: "GetLead".to_owned(),
            input: empty_message("GetLeadRequest"),
            output: empty_message("GetLeadResponse"),
            client_streaming: false,
            server_streaming: false,
        };

        let streaming = ProtoMethod {
            name: "StreamLeads".to_owned(),
            input: empty_message("StreamLeadsRequest"),
            output: empty_message("StreamLeadsResponse"),
            client_streaming: false,
            server_streaming: true,
        };

        let service = ProtoService {
            package: "crm.v1".to_owned(),
            name: "LeadService".to_owned(),
            methods: vec![unary.clone(), streaming],
        };

        let unary_methods = service.unary_methods().collect::<Vec<_>>();

        assert_eq!(unary_methods.len(), 1);
        assert_eq!(unary_methods[0].name, unary.name);
    }

    #[test]
    fn converts_nested_message_to_schema() {
        let request = ProtoMessage {
            name: "CreateLeadRequest".to_owned(),
            fields: vec![ProtoField {
                name: "lead".to_owned(),
                json_name: "lead".to_owned(),
                cardinality: ProtoFieldCardinality::Required,
                field_type: ProtoFieldType::Message {
                    message: Box::new(ProtoMessage {
                        name: "Lead".to_owned(),
                        fields: vec![ProtoField {
                            name: "email".to_owned(),
                            json_name: "email".to_owned(),
                            cardinality: ProtoFieldCardinality::Required,
                            field_type: ProtoFieldType::Scalar {
                                scalar: ProtoScalarKind::String,
                            },
                        }],
                        oneofs: Vec::new(),
                    }),
                },
            }],
            oneofs: Vec::new(),
        };

        let schema = message_to_schema(&request).unwrap();

        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("lead")
                .and_then(|field| field.field("email"))
                .map(|field| field.required),
            Some(true)
        );
    }

    #[test]
    fn converts_enum_map_and_oneof_contracts() {
        let response = ProtoMessage {
            name: "GetLeadResponse".to_owned(),
            fields: vec![
                ProtoField {
                    name: "status".to_owned(),
                    json_name: "status".to_owned(),
                    cardinality: ProtoFieldCardinality::Optional,
                    field_type: ProtoFieldType::Enum {
                        enumeration: ProtoEnum {
                            name: "LeadStatus".to_owned(),
                            values: vec!["OPEN".to_owned(), "CLOSED".to_owned()],
                        },
                    },
                },
                ProtoField {
                    name: "attributes".to_owned(),
                    json_name: "attributes".to_owned(),
                    cardinality: ProtoFieldCardinality::Optional,
                    field_type: ProtoFieldType::Map {
                        map: ProtoMapField {
                            key: ProtoScalarKind::String,
                            value: Box::new(ProtoFieldType::Scalar {
                                scalar: ProtoScalarKind::String,
                            }),
                        },
                    },
                },
            ],
            oneofs: vec![ProtoOneof {
                name: "contact".to_owned(),
                variants: vec![
                    ProtoField {
                        name: "email".to_owned(),
                        json_name: "email".to_owned(),
                        cardinality: ProtoFieldCardinality::Optional,
                        field_type: ProtoFieldType::Scalar {
                            scalar: ProtoScalarKind::String,
                        },
                    },
                    ProtoField {
                        name: "phone".to_owned(),
                        json_name: "phone".to_owned(),
                        cardinality: ProtoFieldCardinality::Optional,
                        field_type: ProtoFieldType::Scalar {
                            scalar: ProtoScalarKind::String,
                        },
                    },
                ],
            }],
        };

        let schema = message_to_schema(&response).unwrap();

        assert_eq!(
            schema.field("status").map(|field| field.kind.clone()),
            Some(SchemaKind::Enum)
        );
        assert_eq!(
            schema
                .field("status")
                .map(|field| field.enum_values.clone())
                .unwrap(),
            vec!["OPEN".to_owned(), "CLOSED".to_owned()]
        );

        let attributes = schema.field("attributes").unwrap();
        assert_eq!(attributes.kind, SchemaKind::Array);
        assert_eq!(
            attributes
                .items
                .as_deref()
                .and_then(|entry| entry.field("key"))
                .map(|field| field.kind.clone()),
            Some(SchemaKind::String)
        );
        assert_eq!(
            attributes
                .items
                .as_deref()
                .and_then(|entry| entry.field("value"))
                .map(|field| field.kind.clone()),
            Some(SchemaKind::String)
        );

        let contact = schema.field("contact").unwrap();
        assert_eq!(contact.kind, SchemaKind::Oneof);
        assert_eq!(contact.variants.len(), 2);
        assert_eq!(
            contact.variants[0]
                .field("email")
                .map(|field| field.kind.clone()),
            Some(SchemaKind::String)
        );
        assert_eq!(
            contact.variants[1]
                .field("phone")
                .map(|field| field.kind.clone()),
            Some(SchemaKind::String)
        );
    }

    #[test]
    fn rejects_empty_oneof() {
        let message = ProtoMessage {
            name: "BrokenMessage".to_owned(),
            fields: Vec::new(),
            oneofs: vec![ProtoOneof {
                name: "selection".to_owned(),
                variants: Vec::new(),
            }],
        };

        let error = message_to_schema(&message).unwrap_err();

        assert_eq!(
            error,
            crate::ProtoError::EmptyOneof {
                oneof_name: "selection".to_owned(),
            }
        );
    }

    fn empty_message(name: &str) -> ProtoMessage {
        ProtoMessage {
            name: name.to_owned(),
            fields: Vec::new(),
            oneofs: Vec::new(),
        }
    }
}