gluesql-core 0.20.0

use {
    super::ExprNode,
    crate::{
        ast::{Aggregate, CountArgExpr},
        parse_sql::parse_expr,
        plan::{AggregateExprPlan, AggregateFunctionPlan, CountArgExprPlan},
        result::Result,
        translate::{NO_PARAMS, translate_expr},
    },
};

#[derive(Clone, Debug)]
pub enum AggregateNode<'a> {
    Count(CountArgExprNode<'a>, bool), // second field is distinct
    Sum(ExprNode<'a>, bool),
    Min(ExprNode<'a>, bool),
    Max(ExprNode<'a>, bool),
    Avg(ExprNode<'a>, bool),
    Variance(ExprNode<'a>, bool),
    Stdev(ExprNode<'a>, bool),
}

#[derive(Clone, Debug)]
pub enum CountArgExprNode<'a> {
    Text(String),
    Expr(ExprNode<'a>),
}

impl<'a> From<&'a str> for CountArgExprNode<'a> {
    fn from(count_arg_str: &str) -> Self {
        Self::Text(count_arg_str.to_owned())
    }
}

impl<'a> From<ExprNode<'a>> for CountArgExprNode<'a> {
    fn from(expr_node: ExprNode<'a>) -> Self {
        Self::Expr(expr_node)
    }
}

impl CountArgExprNode<'_> {
    pub(super) fn build_count_arg_expr(self) -> Result<CountArgExpr> {
        match self {
            CountArgExprNode::Text(s) if &s == "*" => Ok(CountArgExpr::Wildcard),
            CountArgExprNode::Text(s) => {
                let expr = parse_expr(s).and_then(|expr| translate_expr(&expr, NO_PARAMS))?;

                Ok(CountArgExpr::Expr(expr))
            }
            CountArgExprNode::Expr(expr_node) => expr_node.build_expr().map(CountArgExpr::Expr),
        }
    }

    pub(super) fn build_count_arg_expr_plan(self) -> Result<CountArgExprPlan> {
        match self {
            CountArgExprNode::Text(s) if &s == "*" => Ok(CountArgExprPlan::Wildcard),
            CountArgExprNode::Text(s) => {
                let expr = parse_expr(s).and_then(|expr| translate_expr(&expr, NO_PARAMS))?;

                Ok(CountArgExprPlan::Expr(expr.into()))
            }
            CountArgExprNode::Expr(expr_node) => {
                expr_node.build_expr_plan().map(CountArgExprPlan::Expr)
            }
        }
    }
}

impl AggregateNode<'_> {
    pub(super) fn build_aggregate(self) -> Result<Aggregate> {
        match self {
            AggregateNode::Count(count_arg_expr_node, distinct) => count_arg_expr_node
                .build_count_arg_expr()
                .map(|expr| Aggregate::count(expr, distinct)),
            AggregateNode::Sum(expr_node, distinct) => expr_node
                .build_expr()
                .map(|expr| Aggregate::sum(expr, distinct)),
            AggregateNode::Min(expr_node, distinct) => expr_node
                .build_expr()
                .map(|expr| Aggregate::min(expr, distinct)),
            AggregateNode::Max(expr_node, distinct) => expr_node
                .build_expr()
                .map(|expr| Aggregate::max(expr, distinct)),
            AggregateNode::Avg(expr_node, distinct) => expr_node
                .build_expr()
                .map(|expr| Aggregate::avg(expr, distinct)),
            AggregateNode::Variance(expr_node, distinct) => expr_node
                .build_expr()
                .map(|expr| Aggregate::variance(expr, distinct)),
            AggregateNode::Stdev(expr_node, distinct) => expr_node
                .build_expr()
                .map(|expr| Aggregate::stdev(expr, distinct)),
        }
    }

    pub(super) fn build_aggregate_expr_plan(self) -> Result<AggregateExprPlan> {
        let (func, distinct) = match self {
            AggregateNode::Count(count_arg_expr_node, distinct) => count_arg_expr_node
                .build_count_arg_expr_plan()
                .map(|expr| (AggregateFunctionPlan::Count(expr), distinct)),
            AggregateNode::Sum(expr_node, distinct) => expr_node
                .build_expr_plan()
                .map(|expr| (AggregateFunctionPlan::Sum(expr), distinct)),
            AggregateNode::Min(expr_node, distinct) => expr_node
                .build_expr_plan()
                .map(|expr| (AggregateFunctionPlan::Min(expr), distinct)),
            AggregateNode::Max(expr_node, distinct) => expr_node
                .build_expr_plan()
                .map(|expr| (AggregateFunctionPlan::Max(expr), distinct)),
            AggregateNode::Avg(expr_node, distinct) => expr_node
                .build_expr_plan()
                .map(|expr| (AggregateFunctionPlan::Avg(expr), distinct)),
            AggregateNode::Variance(expr_node, distinct) => expr_node
                .build_expr_plan()
                .map(|expr| (AggregateFunctionPlan::Variance(expr), distinct)),
            AggregateNode::Stdev(expr_node, distinct) => expr_node
                .build_expr_plan()
                .map(|expr| (AggregateFunctionPlan::Stdev(expr), distinct)),
        }?;

        Ok(AggregateExprPlan {
            func,
            distinct,
            slot: None,
        })
    }
}

impl<'a> ExprNode<'a> {
    #[must_use]
    pub fn count(self) -> ExprNode<'a> {
        ExprNode::Aggregate(Box::new(AggregateNode::Count(self.into(), false)))
    }

    #[must_use]
    pub fn count_distinct(self) -> ExprNode<'a> {
        ExprNode::Aggregate(Box::new(AggregateNode::Count(self.into(), true)))
    }

    #[must_use]
    pub fn sum(self) -> ExprNode<'a> {
        ExprNode::Aggregate(Box::new(AggregateNode::Sum(self, false)))
    }

    #[must_use]
    pub fn sum_distinct(self) -> ExprNode<'a> {
        ExprNode::Aggregate(Box::new(AggregateNode::Sum(self, true)))
    }

    #[must_use]
    pub fn min(self) -> ExprNode<'a> {
        ExprNode::Aggregate(Box::new(AggregateNode::Min(self, false)))
    }

    #[must_use]
    pub fn min_distinct(self) -> ExprNode<'a> {
        ExprNode::Aggregate(Box::new(AggregateNode::Min(self, true)))
    }

    #[must_use]
    pub fn max(self) -> ExprNode<'a> {
        ExprNode::Aggregate(Box::new(AggregateNode::Max(self, false)))
    }

    #[must_use]
    pub fn max_distinct(self) -> ExprNode<'a> {
        ExprNode::Aggregate(Box::new(AggregateNode::Max(self, true)))
    }

    #[must_use]
    pub fn avg(self) -> ExprNode<'a> {
        ExprNode::Aggregate(Box::new(AggregateNode::Avg(self, false)))
    }

    #[must_use]
    pub fn avg_distinct(self) -> ExprNode<'a> {
        ExprNode::Aggregate(Box::new(AggregateNode::Avg(self, true)))
    }

    #[must_use]
    pub fn variance(self) -> ExprNode<'a> {
        ExprNode::Aggregate(Box::new(AggregateNode::Variance(self, false)))
    }

    #[must_use]
    pub fn variance_distinct(self) -> ExprNode<'a> {
        ExprNode::Aggregate(Box::new(AggregateNode::Variance(self, true)))
    }

    #[must_use]
    pub fn stdev(self) -> ExprNode<'a> {
        ExprNode::Aggregate(Box::new(AggregateNode::Stdev(self, false)))
    }

    #[must_use]
    pub fn stdev_distinct(self) -> ExprNode<'a> {
        ExprNode::Aggregate(Box::new(AggregateNode::Stdev(self, true)))
    }
}

pub fn count<'a, T: Into<CountArgExprNode<'a>>>(expr: T) -> ExprNode<'a> {
    ExprNode::Aggregate(Box::new(AggregateNode::Count(expr.into(), false)))
}

pub fn count_distinct<'a, T: Into<CountArgExprNode<'a>>>(expr: T) -> ExprNode<'a> {
    ExprNode::Aggregate(Box::new(AggregateNode::Count(expr.into(), true)))
}

pub fn sum<'a, T: Into<ExprNode<'a>>>(expr: T) -> ExprNode<'a> {
    ExprNode::Aggregate(Box::new(AggregateNode::Sum(expr.into(), false)))
}

pub fn sum_distinct<'a, T: Into<ExprNode<'a>>>(expr: T) -> ExprNode<'a> {
    ExprNode::Aggregate(Box::new(AggregateNode::Sum(expr.into(), true)))
}

pub fn min<'a, T: Into<ExprNode<'a>>>(expr: T) -> ExprNode<'a> {
    ExprNode::Aggregate(Box::new(AggregateNode::Min(expr.into(), false)))
}

pub fn min_distinct<'a, T: Into<ExprNode<'a>>>(expr: T) -> ExprNode<'a> {
    ExprNode::Aggregate(Box::new(AggregateNode::Min(expr.into(), true)))
}

pub fn max<'a, T: Into<ExprNode<'a>>>(expr: T) -> ExprNode<'a> {
    ExprNode::Aggregate(Box::new(AggregateNode::Max(expr.into(), false)))
}

pub fn max_distinct<'a, T: Into<ExprNode<'a>>>(expr: T) -> ExprNode<'a> {
    ExprNode::Aggregate(Box::new(AggregateNode::Max(expr.into(), true)))
}

pub fn avg<'a, T: Into<ExprNode<'a>>>(expr: T) -> ExprNode<'a> {
    ExprNode::Aggregate(Box::new(AggregateNode::Avg(expr.into(), false)))
}

pub fn avg_distinct<'a, T: Into<ExprNode<'a>>>(expr: T) -> ExprNode<'a> {
    ExprNode::Aggregate(Box::new(AggregateNode::Avg(expr.into(), true)))
}

pub fn variance<'a, T: Into<ExprNode<'a>>>(expr: T) -> ExprNode<'a> {
    ExprNode::Aggregate(Box::new(AggregateNode::Variance(expr.into(), false)))
}

pub fn variance_distinct<'a, T: Into<ExprNode<'a>>>(expr: T) -> ExprNode<'a> {
    ExprNode::Aggregate(Box::new(AggregateNode::Variance(expr.into(), true)))
}

pub fn stdev<'a, T: Into<ExprNode<'a>>>(expr: T) -> ExprNode<'a> {
    ExprNode::Aggregate(Box::new(AggregateNode::Stdev(expr.into(), false)))
}

pub fn stdev_distinct<'a, T: Into<ExprNode<'a>>>(expr: T) -> ExprNode<'a> {
    ExprNode::Aggregate(Box::new(AggregateNode::Stdev(expr.into(), true)))
}

#[cfg(test)]
mod tests {
    use crate::{
        query_builder::{
            avg, avg_distinct, col, count, count_distinct, expr, max, max_distinct, min,
            min_distinct, stdev, stdev_distinct, sum, sum_distinct, test_expr, variance,
            variance_distinct,
        },
        result::Error,
    };

    #[test]
    fn aggregate() {
        let actual = col("id").count();
        let expected = "COUNT(id)";
        test_expr(actual, expected);

        let actual = count("id");
        let expected = "COUNT(id)";
        test_expr(actual, expected);

        let actual = count("*");
        let expected = "COUNT(*)";
        test_expr(actual, expected);

        let actual = count_distinct("*");
        let expected = "COUNT(DISTINCT *)";
        test_expr(actual, expected);

        let actual = col("id").count_distinct();
        let expected = "COUNT(DISTINCT id)";
        test_expr(actual, expected);

        let actual = count_distinct("id");
        let expected = "COUNT(DISTINCT id)";
        test_expr(actual, expected);

        let actual = col("amount").sum();
        let expected = "SUM(amount)";
        test_expr(actual, expected);

        let actual = sum("amount");
        let expected = "SUM(amount)";
        test_expr(actual, expected);

        let actual = col("amount").sum_distinct();
        let expected = "SUM(DISTINCT amount)";
        test_expr(actual, expected);

        let actual = sum_distinct("amount");
        let expected = "SUM(DISTINCT amount)";
        test_expr(actual, expected);

        let actual = col("budget").min();
        let expected = "MIN(budget)";
        test_expr(actual, expected);

        let actual = min("budget");
        let expected = "MIN(budget)";
        test_expr(actual, expected);

        let actual = col("budget").min_distinct();
        let expected = "MIN(DISTINCT budget)";
        test_expr(actual, expected);

        let actual = min_distinct("budget");
        let expected = "MIN(DISTINCT budget)";
        test_expr(actual, expected);

        let actual = col("score").max();
        let expected = "MAX(score)";
        test_expr(actual, expected);

        let actual = max("score");
        let expected = "MAX(score)";
        test_expr(actual, expected);

        let actual = col("grade").max_distinct();
        let expected = "MAX(DISTINCT grade)";
        test_expr(actual, expected);

        let actual = max_distinct("grade");
        let expected = "MAX(DISTINCT grade)";
        test_expr(actual, expected);

        let actual = col("grade").avg();
        let expected = "AVG(grade)";
        test_expr(actual, expected);

        let actual = avg("grade");
        let expected = "AVG(grade)";
        test_expr(actual, expected);

        let actual = col("grade").avg_distinct();
        let expected = "AVG(DISTINCT grade)";
        test_expr(actual, expected);

        let actual = avg_distinct("grade");
        let expected = "AVG(DISTINCT grade)";
        test_expr(actual, expected);

        let actual = col("statistic").variance();
        let expected = "VARIANCE(statistic)";
        test_expr(actual, expected);

        let actual = variance("statistic");
        let expected = "VARIANCE(statistic)";
        test_expr(actual, expected);

        let actual = col("statistic").variance_distinct();
        let expected = "VARIANCE(DISTINCT statistic)";
        test_expr(actual, expected);

        let actual = variance_distinct("statistic");
        let expected = "VARIANCE(DISTINCT statistic)";
        test_expr(actual, expected);

        let actual = col("scatterplot").stdev();
        let expected = "STDEV(scatterplot)";
        test_expr(actual, expected);

        let actual = stdev("scatterplot");
        let expected = "STDEV(scatterplot)";
        test_expr(actual, expected);

        let actual = col("scatterplot").stdev_distinct();
        let expected = "STDEV(DISTINCT scatterplot)";
        test_expr(actual, expected);

        let actual = stdev_distinct("scatterplot");
        let expected = "STDEV(DISTINCT scatterplot)";
        test_expr(actual, expected);
    }

    #[test]
    fn aggregate_expr_plan_propagates_expr_error() {
        let actual = sum(expr(")")).build_expr_plan();

        assert!(matches!(actual, Err(Error::Parser(_))));
    }
}