hamelin_lib 0.15.5

use crate::err::TranslationErrors;
use crate::sql::expression::identifier::{Identifier, SimpleIdentifier};
use crate::sql::expression::literal::ColumnReference;
use crate::sql::expression::SQLExpression;
use crate::sql::query::cte::CTE;
use crate::sql::query::projection::{Binding, ColumnProjection, Projection};
use crate::sql::query::set::SetOperation;
use crate::sql::query::window::{WindowSpecification, Windows};
use crate::write_utils::{maybe_newline, maybe_pad, newline_or_space, Indent};
use derive_more::{From, TryUnwrap};
use std::fmt;
use std::fmt::{Debug, Display, Formatter};
use std::str::FromStr;
use thiserror::Error;

use super::expression::OrderByExpression;

pub mod cte;
pub mod dml;
pub mod projection;
pub mod set;
pub mod window;

#[derive(Debug, Clone, From, TryUnwrap, PartialEq, Eq)]
pub enum SQLQueryExpression {
    SetOperation(SetOperation),
    SQLQuery(SQLQuery),
}

impl SQLQueryExpression {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: Indent) -> fmt::Result {
        match self {
            SQLQueryExpression::SetOperation(set) => set.fmt_indented(f, indentation),
            SQLQueryExpression::SQLQuery(query) => query.fmt_indented(f, indentation),
        }
    }

    pub fn get_table_references(&self) -> Vec<TableReference> {
        match self {
            SQLQueryExpression::SetOperation(set) => {
                let cte_refs = set
                    .cte
                    .get_names()
                    .into_iter()
                    .map(|n| Identifier::from(n.clone()))
                    .collect::<Vec<Identifier>>();
                set.left
                    .get_table_references()
                    .into_iter()
                    .chain(set.right.get_table_references())
                    .chain(set.cte.get_table_references())
                    .filter(|t| !cte_refs.contains(&t.name))
                    .collect()
            }
            SQLQueryExpression::SQLQuery(query) => {
                let cte_refs = query
                    .cte
                    .get_names()
                    .into_iter()
                    .map(|n| Identifier::from(n.clone()))
                    .collect::<Vec<Identifier>>();
                query
                    .from
                    .clone()
                    .into_iter()
                    .flat_map(|t| t.get_table_references())
                    .chain(query.cte.get_table_references())
                    .filter(|t| !cte_refs.contains(&t.name))
                    .collect()
            }
        }
    }
}

impl Display for SQLQueryExpression {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> std::fmt::Result {
        self.fmt_indented(f, Indent::default())
    }
}

#[derive(Debug, Clone, From, TryUnwrap, PartialEq, Eq)]
pub enum TableExpression {
    TableReference(TableReference),
    TableFunctionApply(TableFunctionApply),
    SubQuery(SubQuery),
    Join(Join),
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TableFunctionApply {
    pub function_name: String,
    pub arguments: Vec<SQLExpression>,
    pub alias: Option<TableAlias>,
}

impl TableFunctionApply {
    pub fn new(function_name: String) -> Self {
        Self {
            function_name,
            arguments: Vec::new(),
            alias: None,
        }
    }

    pub fn with_argument(mut self, argument: SQLExpression) -> Self {
        self.arguments.push(argument);
        self
    }

    pub fn with_alias(mut self, alias: TableAlias) -> Self {
        self.alias = Some(alias);
        self
    }

    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: Indent) -> fmt::Result {
        write!(f, "{}(", self.function_name)?;
        for (i, arg) in self.arguments.iter().enumerate() {
            if i > 0 {
                write!(f, ", ")?;
            }
            arg.fmt_indented(f, indentation.nested())?;
        }
        write!(f, ")")?;
        if let Some(alias) = &self.alias {
            write!(f, " AS ")?;
            alias.table_alias.fmt_indented(f, indentation.nested())?;
            if let Some(column_aliases) = &alias.field_aliases {
                write!(f, "(")?;
                for (i, column_alias) in column_aliases.iter().enumerate() {
                    if i > 0 {
                        write!(f, ", ")?;
                    }
                    column_alias.fmt_indented(f, indentation.nested())?;
                }
                write!(f, ")")?;
            }
        }
        Ok(())
    }
}

impl Display for TableFunctionApply {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        self.fmt_indented(f, Indent::default())
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct TableReference {
    pub name: Identifier,
    pub alias: Option<Identifier>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SubQuery {
    pub query: Box<SQLQueryExpression>,
    pub alias: Option<Identifier>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TableAlias {
    pub table_alias: Identifier,
    pub field_aliases: Option<Vec<SimpleIdentifier>>,
}

impl TableAlias {
    pub fn new(table_alias: Identifier) -> Self {
        Self {
            table_alias,
            field_aliases: None,
        }
    }

    pub fn with_field_alias(mut self, field_alias: SimpleIdentifier) -> Self {
        if let Some(ref mut aliases) = self.field_aliases {
            aliases.push(field_alias);
        } else {
            self.field_aliases = Some(vec![field_alias]);
        }

        self
    }
}

impl SubQuery {
    pub fn new(sub: SQLQueryExpression) -> Self {
        Self {
            query: Box::new(sub),
            alias: None,
        }
    }

    pub fn alias(mut self, alias: Identifier) -> Self {
        self.alias = Some(alias);
        self
    }

    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: Indent) -> fmt::Result {
        write!(f, "(")?;
        maybe_newline(f, indentation)?;
        maybe_pad(f, indentation.nested())?;
        self.query.fmt_indented(f, indentation.nested())?;
        maybe_newline(f, indentation)?;
        maybe_pad(f, indentation)?;
        write!(f, ")")?;
        if let Some(alias) = &self.alias {
            write!(f, " AS ")?;
            alias.fmt_indented(f, indentation)?;
        }
        Ok(())
    }
}

impl Display for SubQuery {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        self.fmt_indented(f, Indent::default())
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Join {
    pub left: Box<TableExpression>,
    pub clauses: Vec<JoinClause>,
}

impl Join {
    pub fn new(left: TableExpression) -> Self {
        Self {
            left: Box::new(left),
            clauses: Vec::new(),
        }
    }

    pub fn with_clause(mut self, clause: JoinClause) -> Self {
        self.clauses.push(clause);
        self
    }

    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: Indent) -> fmt::Result {
        self.left.fmt_indented(f, indentation)?;
        for clause in &self.clauses {
            newline_or_space(f, indentation)?;
            maybe_pad(f, indentation.nested())?;
            clause.fmt_indented(f, indentation.nested())?;
        }
        Ok(())
    }
}

impl Display for Join {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        self.fmt_indented(f, Indent::default())
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum JoinType {
    INNER,
    LEFT,
    RIGHT,
    CROSS,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct JoinClause {
    pub table: TableExpression,
    pub join_type: JoinType,
    pub condition: Option<SQLExpression>,
}

/// A pattern variable used in MATCH commands.
/// Pattern variables must have "clean" names (no special characters requiring backticks).
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PatternVariable {
    pub name: String,
}

impl PatternVariable {
    pub fn new(name: String) -> Self {
        Self { name }
    }
}

impl Display for PatternVariable {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.name)
    }
}

#[derive(Debug, Clone, Error, PartialEq, Eq)]
#[error("Pattern variable {name} must be clean. (You shouldn't need a backtick to ref it.)")]
pub struct BadPatternVariable {
    pub name: SimpleIdentifier,
}

impl TryFrom<SimpleIdentifier> for PatternVariable {
    type Error = BadPatternVariable;

    fn try_from(value: SimpleIdentifier) -> Result<Self, Self::Error> {
        if value.is_clean() {
            Ok(Self { name: value.name })
        } else {
            Err(BadPatternVariable { name: value })
        }
    }
}

impl TableReference {
    pub fn new(name: Identifier) -> Self {
        Self { name, alias: None }
    }

    pub fn alias(mut self, alias: Identifier) -> Self {
        self.alias = Some(alias);
        self
    }

    pub fn to_hamelin(&self) -> String {
        self.name.to_hamelin()
    }

    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: Indent) -> fmt::Result {
        self.name.fmt_indented(f, indentation.nested())?;
        if let Some(alias) = &self.alias {
            write!(f, " AS ")?;
            alias.fmt_indented(f, indentation.nested())?;
        }

        Ok(())
    }
}

impl Display for TableReference {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        self.fmt_indented(f, Indent::default())
    }
}

impl FromStr for TableReference {
    type Err = TranslationErrors;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        Ok(TableReference::new(s.parse()?))
    }
}

/// A SQL query. Part of the SQL AST
///
/// It does have a slight bit of magic associated with it. It can "push down" the current set of
/// projections into a subquery when it needs to. It does this when it identifies that a column
/// reference in a newly added expression references a column defined in its currently top-level
/// projection list. The magic methods that have this behavior are `prependProjections(List,
/// Environment)`, `replaceProjections(List)`, and `addFilterCondition(SQLExpression,
/// Environment)`

#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct SQLQuery {
    pub projections: Vec<Projection>,
    pub from: Option<TableExpression>,
    pub where_: Option<SQLExpression>,
    pub group_by: Option<Vec<SQLExpression>>,
    pub order_by: Option<Vec<OrderByExpression>>,
    pub limit: Option<SQLExpression>,
    pub cte: CTE,
    pub windows: Windows,
}

impl TableExpression {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: Indent) -> fmt::Result {
        match self {
            TableExpression::TableReference(table_reference) => {
                table_reference.fmt_indented(f, indentation)
            }
            TableExpression::TableFunctionApply(table_function_apply) => {
                table_function_apply.fmt_indented(f, indentation)
            }
            TableExpression::SubQuery(sub_query) => sub_query.fmt_indented(f, indentation),
            TableExpression::Join(join) => join.fmt_indented(f, indentation),
        }
    }

    pub fn get_table_references(&self) -> Vec<TableReference> {
        match self {
            TableExpression::SubQuery(q) => q.query.get_table_references(),
            TableExpression::TableReference(t) => vec![t.clone()],
            TableExpression::Join(j) => j
                .left
                .get_table_references()
                .into_iter()
                .chain(
                    j.clauses
                        .iter()
                        .flat_map(|clause| clause.table.get_table_references())
                        .collect::<Vec<_>>(),
                )
                .collect(),
            TableExpression::TableFunctionApply(_) => {
                // A table function apply generates a table without referring to one.
                vec![]
            }
        }
    }
}

impl Display for TableExpression {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        self.fmt_indented(f, Indent::default())
    }
}

impl JoinClause {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: Indent) -> fmt::Result {
        match self.join_type {
            JoinType::INNER => write!(f, "INNER JOIN ")?,
            JoinType::LEFT => write!(f, "LEFT JOIN ")?,
            JoinType::RIGHT => write!(f, "RIGHT JOIN ")?,
            JoinType::CROSS => write!(f, "CROSS JOIN ")?,
        }
        self.table.fmt_indented(f, indentation.nested())?;
        if let Some(condition) = &self.condition {
            write!(f, " ON")?;
            newline_or_space(f, indentation)?;
            maybe_pad(f, indentation.nested())?;
            condition.fmt_indented(f, indentation.nested())?;
        }
        Ok(())
    }
}

impl SQLQuery {
    pub fn from(mut self, from: TableExpression) -> Self {
        self.from = Some(from);
        self
    }

    pub fn select(mut self, projections: Vec<Projection>) -> Self {
        self.projections = projections;
        self
    }

    pub fn limit(mut self, limit: SQLExpression) -> Self {
        self.limit = Some(limit);
        self
    }

    pub fn where_(mut self, where_: SQLExpression) -> Self {
        self.where_ = Some(where_);
        self
    }

    pub fn group_by(mut self, group_by: Vec<SQLExpression>) -> Self {
        self.group_by = Some(group_by);
        self
    }

    pub fn order_by(mut self, order_by: Vec<OrderByExpression>) -> Self {
        self.order_by = Some(order_by);
        self
    }

    pub fn cte(mut self, cte: CTE) -> Self {
        self.cte = cte;
        self
    }

    pub fn with_cte(mut self, name: SimpleIdentifier, query: SQLQuery) -> Self {
        self.cte = self.cte.with(name, query);
        self
    }

    pub fn windows(mut self, windows: Windows) -> Self {
        self.windows = windows;
        self
    }

    pub fn with_window(
        mut self,
        identifier: SimpleIdentifier,
        window: WindowSpecification,
    ) -> Self {
        self.windows = self.windows.with(identifier, window);
        self
    }

    pub fn replace_projections(&self, projections: Vec<Projection>) -> Self {
        let mut ret = self.clone();
        if ret.references_column_in_projections(&projections) {
            ret = ret.push_down();
        }
        ret.select(projections)
    }

    pub fn remove_projections(&self, columns_to_drop: &[Identifier]) -> Self {
        let mut ret = self.clone();
        for column in columns_to_drop {
            ret.projections.retain(|p| match p {
                Projection::Binding(Binding { name, .. }) => {
                    Identifier::from(name.clone()).last() != column.last()
                }
                Projection::ColumnProjection(ColumnProjection {
                    identifier: column_reference,
                }) => column_reference.last() != column.last(),
            });
        }
        ret
    }

    pub fn remove_simple_projections(&self, columns_to_drop: &[SimpleIdentifier]) -> Self {
        let mut ret = self.clone();
        for column in columns_to_drop {
            ret.projections.retain(|p| match p {
                Projection::Binding(Binding { name, .. }) => *name != *column,
                Projection::ColumnProjection(ColumnProjection {
                    identifier: column_reference,
                }) => *column_reference.last() != *column,
            });
        }
        ret
    }

    pub fn add_group_expression(self, group: Vec<SQLExpression>) -> Self {
        if group.is_empty() {
            return self;
        }
        let push_down_needed = group
            .iter()
            .map(|e| e.get_column_references())
            .any(|c| self.references_columns_in_column_refs(&c[..]));
        let ret_query = if push_down_needed {
            self.push_down()
        } else {
            self
        };
        ret_query.group_by(group)
    }

    pub fn references_column_in_projections(&self, projections: &[Projection]) -> bool {
        let column_references = projections
            .iter()
            .flat_map(|p| match p {
                Projection::Binding(Binding { expression, .. }) => {
                    expression.get_column_references()
                }
                Projection::ColumnProjection(ColumnProjection {
                    identifier: column_reference,
                }) => {
                    vec![ColumnReference::new(column_reference.clone())]
                }
            })
            .collect::<Vec<_>>();
        self.references_columns_in_column_refs(&column_references[..])
    }

    pub fn references_columns_in_column_refs(&self, column_references: &[ColumnReference]) -> bool {
        column_references
            .iter()
            .map(|column_reference| match column_reference {
                ColumnReference {
                    identifier: Identifier::Simple(s),
                } => s,
                ColumnReference {
                    identifier: Identifier::Compound(c),
                } => c.first(),
            })
            .any(|expression_identifier| {
                self.projections.iter().any(|projection| match projection {
                    Projection::Binding(Binding { name, .. }) => name == expression_identifier,
                    Projection::ColumnProjection(_) => {
                        // A column projection is a "lift" from a lower query.
                        // That means it doesn't count as being defined in this query.
                        false
                    }
                })
            })
    }

    pub fn push_down(self) -> Self {
        let sqe: SQLQueryExpression = self.into();
        let sq = SubQuery::new(sqe);
        SQLQuery::default().from(sq.into())
    }

    pub fn get_table_references(&self) -> Vec<TableReference> {
        self.from
            .as_ref()
            .map(|f| f.get_table_references())
            .unwrap_or_default()
            .into_iter()
            .chain(self.cte.get_table_references().into_iter())
            .filter(|tr| match &tr.name {
                Identifier::Simple(simple_identifier) => {
                    !self.cte.expressions.contains_key(simple_identifier)
                }
                Identifier::Compound(_) => true,
            })
            .collect()
    }

    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: Indent) -> fmt::Result {
        self.cte.clone().fmt_indented(f, indentation)?;
        write!(f, "SELECT")?;
        newline_or_space(f, indentation)?;
        for (i, projection) in self.projections.iter().enumerate() {
            maybe_pad(f, indentation.nested())?;
            projection.fmt_indented(f, indentation.nested())?;

            if i != self.projections.len() - 1 {
                write!(f, ",")?;
                newline_or_space(f, indentation)?;
            }
        }

        if let Some(from) = &self.from {
            newline_or_space(f, indentation)?;
            maybe_pad(f, indentation)?;

            // This is a hack and I hate it but oh well.
            // Newlines are different for tr's because I have peculiar aesthetic prefs
            if let TableExpression::TableReference(_) = from {
                write!(f, "FROM")?;
                newline_or_space(f, indentation)?;
                maybe_pad(f, indentation.nested())?;
            } else {
                write!(f, "FROM ")?;
            }
            from.fmt_indented(f, indentation)?;
        }

        if let Some(where_) = &self.where_ {
            newline_or_space(f, indentation)?;
            maybe_pad(f, indentation)?;
            write!(f, "WHERE")?;
            newline_or_space(f, indentation)?;
            maybe_pad(f, indentation.nested())?;
            where_.fmt_indented(f, indentation.nested())?;
        }

        if let Some(group_bys) = &self.group_by {
            newline_or_space(f, indentation)?;
            maybe_pad(f, indentation)?;
            write!(f, "GROUP BY")?;
            newline_or_space(f, indentation)?;
            for (i, group_by) in group_bys.iter().enumerate() {
                maybe_pad(f, indentation.nested())?;
                if i != 0 {
                    write!(f, ", ")?;
                }
                group_by.fmt_indented(f, indentation.nested())?;
            }
        }

        if !self.windows.windows.is_empty() {
            self.windows.fmt_indented(f, indentation)?;
        }

        if let Some(order_bys) = &self.order_by {
            newline_or_space(f, indentation)?;
            maybe_pad(f, indentation)?;
            write!(f, "ORDER BY")?;
            newline_or_space(f, indentation)?;
            for (i, order_by) in order_bys.iter().enumerate() {
                maybe_pad(f, indentation.nested())?;
                if i != 0 {
                    write!(f, ", ")?;
                }
                order_by.fmt_indented(f, indentation.nested())?;
            }
        }

        if let Some(limit) = &self.limit {
            newline_or_space(f, indentation)?;
            maybe_pad(f, indentation)?;
            write!(f, "LIMIT")?;
            newline_or_space(f, indentation)?;
            maybe_pad(f, indentation.nested())?;
            limit.fmt_indented(f, indentation)?;
        };
        Ok(())
    }
}

impl Display for SQLQuery {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        self.fmt_indented(f, Indent::default())
    }
}