phonograph/phono-models/src/expression.rs

use std::fmt::Display;

use phono_backends::escape_identifier;
use serde::{Deserialize, Serialize};

use crate::{datum::Datum, query_builders::QueryFragment};

/// Building block of a syntax tree for a constrained subset of SQL that can be
/// statically analyzed, to validate that user-provided expressions perform only
/// operations that are read-only and otherwise safe to execute.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(tag = "t", content = "c")]
pub enum PgExpressionAny {
    Comparison(PgComparisonExpression),
    Identifier(PgIdentifierExpression),
    Literal(Datum),
    ToJson(PgToJsonExpression),
}

impl PgExpressionAny {
    pub fn into_query_fragment(self) -> QueryFragment {
        match self {
            Self::Comparison(expr) => expr.into_query_fragment(),
            Self::Identifier(expr) => expr.into_query_fragment(),
            Self::Literal(expr) => {
                if expr.is_none() {
                    QueryFragment::from_sql("null")
                } else {
                    QueryFragment::from_param(expr)
                }
            }
            Self::ToJson(expr) => expr.into_query_fragment(),
        }
    }
}

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(tag = "t", content = "c")]
pub enum PgComparisonExpression {
    Infix(PgInfixExpression<PgComparisonOperator>),
    IsNull(PgIsNullExpression),
    IsNotNull(PgIsNotNullExpression),
}

impl PgComparisonExpression {
    fn into_query_fragment(self) -> QueryFragment {
        match self {
            Self::Infix(expr) => expr.into_query_fragment(),
            Self::IsNull(expr) => expr.into_query_fragment(),
            Self::IsNotNull(expr) => expr.into_query_fragment(),
        }
    }
}

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub struct PgInfixExpression<T: Display> {
    pub operator: T,
    pub lhs: Box<PgExpressionAny>,
    pub rhs: Box<PgExpressionAny>,
}

impl<T: Display> PgInfixExpression<T> {
    fn into_query_fragment(self) -> QueryFragment {
        QueryFragment::concat([
            QueryFragment::from_sql("(("),
            self.lhs.into_query_fragment(),
            QueryFragment::from_sql(&format!(") {} (", self.operator)),
            self.rhs.into_query_fragment(),
            QueryFragment::from_sql("))"),
        ])
    }
}

#[derive(Clone, Debug, strum::Display, Deserialize, PartialEq, Serialize)]
pub enum PgComparisonOperator {
    #[strum(to_string = "and")]
    And,
    #[strum(to_string = "=")]
    Eq,
    #[strum(to_string = ">")]
    Gt,
    #[strum(to_string = "<")]
    Lt,
    #[strum(to_string = "<>")]
    Neq,
    #[strum(to_string = "or")]
    Or,
}

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub struct PgIsNullExpression {
    lhs: Box<PgExpressionAny>,
}

impl PgIsNullExpression {
    fn into_query_fragment(self) -> QueryFragment {
        QueryFragment::concat([
            QueryFragment::from_sql("(("),
            self.lhs.into_query_fragment(),
            QueryFragment::from_sql(") is null)"),
        ])
    }
}

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub struct PgIsNotNullExpression {
    lhs: Box<PgExpressionAny>,
}

impl PgIsNotNullExpression {
    fn into_query_fragment(self) -> QueryFragment {
        QueryFragment::concat([
            QueryFragment::from_sql("(("),
            self.lhs.into_query_fragment(),
            QueryFragment::from_sql(") is not null)"),
        ])
    }
}

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub struct PgIdentifierExpression {
    pub parts_raw: Vec<String>,
}

impl PgIdentifierExpression {
    fn into_query_fragment(self) -> QueryFragment {
        QueryFragment::join(
            self.parts_raw
                .iter()
                .map(|part| QueryFragment::from_sql(&escape_identifier(part))),
            QueryFragment::from_sql("."),
        )
    }
}

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub struct PgToJsonExpression {
    entries: Vec<(String, PgExpressionAny)>,
}

impl PgToJsonExpression {
    /// Generates a query fragment to the effect of:
    /// `to_json((select ($expr) as "ident", ($expr2) as "ident2"))`
    fn into_query_fragment(self) -> QueryFragment {
        if self.entries.is_empty() {
            QueryFragment::from_sql("'{}'")
        } else {
            QueryFragment::concat([
                QueryFragment::from_sql("to_json((select "),
                QueryFragment::join(
                    self.entries.into_iter().map(|(key, value)| {
                        QueryFragment::concat([
                            QueryFragment::from_sql("("),
                            value.into_query_fragment(),
                            QueryFragment::from_sql(&format!(") as {}", escape_identifier(&key))),
                        ])
                    }),
                    QueryFragment::from_sql(", "),
                ),
                QueryFragment::from_sql("))"),
            ])
        }
    }
}