kappe/src/compiler.rs

use serde_json::{json, Value};

use std::fmt::{Display, Formatter};

use crate::ast::{Field, KIdentifier, KValue, Select, Where};
use crate::error::KappeError;

#[derive(Debug, PartialEq)]
pub struct Query {
    pub source: KIdentifier,
    pub body: Value,
}

impl Query {
    pub fn new(source: KIdentifier, body: Value) -> Self {
        Self { source, body }
    }
}

impl Display for Query {
    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
        let header = format!("GET /{}/_search", self.source.name);
        let body = match serde_json::to_string_pretty(&self.body) {
            Ok(body) => body,
            Err(_) => return Err(std::fmt::Error {}),
        };
        write!(f, "{}\n{}", header, body)
    }
}

pub fn compile(expr: Select) -> Result<Query, KappeError> {
    let field_names: Vec<Value> = expr
        .fields
        .iter()
        .map(|f| match f {
            Field::Star => Value::String("*".to_string()),
            Field::Named(name) => Value::String(name.clone()),
        })
        .collect();

    let query = if let Some(where_clause) = expr.where_clause {
        compile_where_clause(where_clause)
    } else {
        json!({ "match_all":{} })
    };

    let body = json!({
        "query": query,
        "_source": Value::Array(field_names)
    });

    Ok(Query::new(expr.source, body))
}

pub fn compile_where_clause(where_clause: Where) -> Value {
    let terms: Vec<Value> = where_clause
        .predicates
        .iter()
        .map(|predicate| {
            let mut map = serde_json::Map::new();
            let key = predicate.left.name.clone();
            let value = compile_kvalue(&predicate.right);
            map.insert(key, value);
            json!({ "term": Value::Object(map) })
        })
        .collect();

    json!({
        "bool": {
            "filter": Value::Array(terms)
        }
    })
}

pub fn compile_kvalue(value: &KValue) -> Value {
    match value {
        KValue::Identifier(identifier) => Value::String(identifier.name.to_string()),
        KValue::Number(number) => {
            Value::Number(serde_json::Number::from_f64(number.value).unwrap())
        }
        KValue::String(string) => Value::String(string.value.to_string()),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::lexer::scan;
    use crate::parser::parse;

    fn _compile(input: &str) -> Result<Query, KappeError> {
        let tokens = scan(input).unwrap();
        let select = parse(tokens).unwrap();
        compile(select)
    }

    #[test]
    fn it_compiles_a_simple_select() {
        assert_eq!(
            _compile("SELECT * FROM bar").unwrap(),
            Query::new(
                KIdentifier::new("bar"),
                json!({
                    "query": {
                        "match_all": {}
                    },
                    "_source": ["*"]
                })
            )
        )
    }

    #[test]
    fn it_compiles_a_select_with_specific_fields() {
        assert_eq!(
            _compile("SELECT foo FROM bar").unwrap(),
            Query::new(
                KIdentifier::new("bar"),
                json!({
                    "query": {
                        "match_all": {}
                    },
                    "_source": ["foo"]
                })
            )
        )
    }

    #[test]
    fn it_compiles_a_select_with_a_single_predicate_where_clause() {
        let mut map: serde_json::Map<String, Value> = serde_json::Map::new();
        map.insert("baz".to_string(), Value::String("quux".to_string()));
        assert_eq!(
            _compile("SELECT * FROM bar WHERE baz = \"quux\"").unwrap(),
            Query::new(
                KIdentifier::new("bar"),
                json!({
                    "query": {
                        "bool": {
                            "filter": [
                                json!({ "term": Value::Object(map) })
                            ]
                        }
                    },
                    "_source": ["*"]
                })
            )
        )
    }

    #[test]
    fn it_compiles_a_select_with_a_multiple_predicate_where_claus() {
        let mut map_one: serde_json::Map<String, Value> = serde_json::Map::new();
        map_one.insert("baz".to_string(), Value::String("quux".to_string()));
        let mut map_two: serde_json::Map<String, Value> = serde_json::Map::new();
        map_two.insert(
            "bop".to_string(),
            Value::Number(serde_json::Number::from_f64(5.0).unwrap()),
        );
        assert_eq!(
            _compile("SELECT * FROM bar WHERE baz = \"quux\", bop = 5").unwrap(),
            Query::new(
                KIdentifier::new("bar"),
                json!({
                    "query": {
                        "bool": {
                            "filter": [
                                json!({ "term": Value::Object(map_one) }),
                                json!({ "term": Value::Object(map_two) }),
                            ]
                        }
                    },
                    "_source": ["*"]
                })
            )
        )
    }
}