keelson_psql/

extras.rs

use std::borrow::Cow;

use keelson_core::clause::{GroupingSet, GroupingSetKind};
use keelson_core::expr::{Expr, IntoExprList};
use keelson_core::{Expression, Query, SqlWriter};

/// `SELECT` **`DISTINCT`** or `SELECT` **`DISTINCT ON (a, b)`**.
///
/// From PostgreSQL 17: `SELECT [ ALL | DISTINCT [ ON ( expression [, ...] ) ] ]`.
/// `ALL` is the default and adds nothing, so it is not representable; the absence
/// of a `Distinct` is what `ALL` means.
///
/// A `SelectQuery` stores this as an `Option`, because `DISTINCT` with an empty
/// `ON` list is a different statement from no `DISTINCT` at all — which is exactly
/// the distinction bob loses by keying off `On != nil`.
#[derive(Debug, Clone, Default)]
pub struct Distinct {
    /// The `ON` expressions. Empty is a plain `DISTINCT`.
    pub on: Vec<Expr>,
}

impl Expression for Distinct {
    fn write_sql(&self, w: &mut SqlWriter<'_>) {
        w.push_str("DISTINCT");
        w.write_slice(&self.on, " ON (", ", ", ")");
    }
}

/// `OVERRIDING { SYSTEM | USER } VALUE`, an `INSERT`'s treatment of an identity
/// column.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Overriding {
    /// `OVERRIDING SYSTEM VALUE` — write the supplied value into a
    /// `GENERATED ALWAYS` identity column.
    System,
    /// `OVERRIDING USER VALUE` — ignore the supplied value and use the sequence,
    /// for a `GENERATED BY DEFAULT` column.
    User,
}

impl Overriding {
    /// The keyword, as written between `OVERRIDING` and `VALUE`.
    pub fn as_str(self) -> &'static str {
        match self {
            Overriding::System => "SYSTEM",
            Overriding::User => "USER",
        }
    }
}

/// `TABLESAMPLE method ( args ) [ REPEATABLE ( seed ) ]`.
#[derive(Debug, Clone)]
pub(crate) struct Sample {
    pub(crate) method: Cow<'static, str>,
    pub(crate) args: Vec<Expr>,
    pub(crate) repeatable: Option<Expr>,
}

/// A from-item with a `TABLESAMPLE` clause, alias and all.
///
/// PostgreSQL's `gram.y` puts the sampling clause *after* the alias:
///
/// ```text
/// table_ref: relation_expr opt_alias_clause tablesample_clause
/// ```
///
/// [`TableRef`](keelson_core::clause::TableRef) has no slot in that position — it
/// writes the alias last of the pre-join decorations — so the alias, the column
/// aliases and the sampling clause are folded into one expression here and put in
/// `TableRef::expression`. `ONLY` and `LATERAL` stay on the `TableRef`, because
/// they precede the table name, and so do the joins, because they follow
/// everything.
#[derive(Debug)]
pub(crate) struct SampledTable {
    pub(crate) table: Expr,
    pub(crate) alias: Option<Cow<'static, str>>,
    pub(crate) columns: Vec<Cow<'static, str>>,
    pub(crate) sample: Sample,
}

impl Expression for SampledTable {
    fn write_sql(&self, w: &mut SqlWriter<'_>) {
        w.write_expr(&self.table);

        if let Some(alias) = &self.alias {
            w.push_str(" AS ");
            w.push_quoted(&[alias]);
        }
        if !self.columns.is_empty() {
            w.push_str(" (");
            for (i, column) in self.columns.iter().enumerate() {
                if i > 0 {
                    w.push_str(", ");
                }
                w.push_quoted(&[column]);
            }
            w.push_str(")");
        }

        w.push_str(" TABLESAMPLE ");
        w.push_str(&self.sample.method);
        // Unconditional parentheses: every sampling method takes at least the
        // percentage, so an empty list is a caller error rather than a shape.
        w.push_str(" (");
        w.write_slice(&self.sample.args, "", ", ", "");
        w.push_str(")");

        if let Some(seed) = &self.sample.repeatable {
            w.push_str(" REPEATABLE (");
            w.write_expr(seed);
            w.push_str(")");
        }
    }
}

/// A whole query standing in an expression slot, rendered in **its own** dialect.
///
/// bob's `BaseQuery.WriteSQL` ignores the dialect it is handed and uses the one it
/// was built with; [`SqlWriter::write_with_dialect`] is how that is done here, and
/// it keeps one shared argument list and placeholder counter, so a sub-query
/// re-indexes into its container for free.
#[derive(Debug)]
struct QueryExpr<Q>(Q);

impl<Q: Query> Expression for QueryExpr<Q> {
    fn write_sql(&self, w: &mut SqlWriter<'_>) {
        w.write_with_dialect(self.0.dialect(), &self.0);
    }
}

/// A query as an expression, **not** parenthesised.
///
/// This is the form for slots that supply their own parentheses — a `WITH` body,
/// a set-operation operand, `IN (…)`, `INSERT … SELECT`. Use [`subquery`] where
/// the parentheses are part of the sub-query itself, as in a `FROM` item.
pub fn query(q: impl Query + 'static) -> Expr {
    Expr::custom(QueryExpr(q))
}

/// A parenthesised sub-query: `(SELECT …)`.
///
/// What a `FROM` item or a scalar sub-expression needs. PostgreSQL additionally
/// requires an alias on a `FROM` sub-query, which is
/// [`select::from(..).as_(..)`](mod@crate::select).
pub fn subquery(q: impl Query + 'static) -> Expr {
    Expr::group(query(q))
}

/// `EXCLUDED."col"` — the proposed row inside `ON CONFLICT DO UPDATE`.
pub fn excluded(column: impl Into<Cow<'static, str>>) -> Expr {
    Expr::join_with("", (Expr::raw("EXCLUDED."), Expr::ident(column.into())))
}

/// A fragment that cannot be rendered, and says so instead of writing nothing.
///
/// The pattern this exists for: a helper is handed an empty list, and the clause
/// that will hold it has *already* committed its keyword by the time the fragment
/// renders. `GroupBy` writes `GROUP BY ` as soon as it has one grouping element, so
/// an element that writes nothing leaves the keyword dangling and
/// [`build`](keelson_core::Query::build) hands back SQL that cannot parse — with no
/// error at all, which is the worst of the available outcomes. Every other
/// unfillable construct in the clause layer records
/// [`Error::Incomplete`](keelson_core::Error::Incomplete); so does this.
#[derive(Debug)]
pub(crate) struct Incomplete(pub(crate) &'static str);

impl Expression for Incomplete {
    fn write_sql(&self, w: &mut SqlWriter<'_>) {
        w.record_error(keelson_core::Error::Incomplete(self.0));
    }
}

/// A from-item that was marked `LATERAL` but is a bare table or CTE name.
///
/// PostgreSQL's grammar puts `LATERAL` only in front of a sub-query or a
/// function item — `JOIN LATERAL "posts"` is a syntax error, and there is
/// nothing for the keyword to mean on a name anyway (a table cannot reference
/// the items before it). The chain methods swap this in when `.lateral()` is
/// called on such an item, so the mistake is caught where it is made; the item
/// still renders, keeping the debug print honest, while `build()` refuses.
///
/// Only [`Expr::Ident`] items are judged. A raw fragment could be anything —
/// progressive enhancement means hand-written SQL is trusted — and sub-queries
/// and function calls arrive as other variants.
#[derive(Debug)]
pub(crate) struct LateralBareName(pub(crate) Expr);

impl Expression for LateralBareName {
    fn write_sql(&self, w: &mut SqlWriter<'_>) {
        w.record_error(keelson_core::Error::other(
            "LATERAL is set on a bare table or CTE name, but LATERAL can precede only a sub-query or a function item",
        ));
        w.write_expr(&self.0);
    }
}

/// Wrap a grouping element, refusing an empty one.
fn grouping_element(kind: GroupingSetKind, groups: impl IntoExprList) -> Expr {
    let set = GroupingSet::new(kind, groups);
    if set.is_empty() {
        // `ROLLUP` with no list is a syntax error, and so is the `GROUP BY ` that
        // would be left in front of it. See `Incomplete`.
        return Expr::custom(Incomplete("the columns of a grouping element"));
    }
    Expr::custom(set)
}

/// `ROLLUP (a, b)` — a grouping element covering every prefix of the list.
pub fn rollup(groups: impl IntoExprList) -> Expr {
    grouping_element(GroupingSetKind::Rollup, groups)
}

/// `CUBE (a, b)` — a grouping element covering every subset of the list.
pub fn cube(groups: impl IntoExprList) -> Expr {
    grouping_element(GroupingSetKind::Cube, groups)
}

/// `GROUPING SETS ((a), (b), ())` — the sets listed explicitly.
///
/// Each element is normally a [`group`](crate::group); the empty set is written
/// [`raw("()")`](crate::raw), because an empty
/// [`Expr::Group`](keelson_core::expr::Expr::Group) renders `(NULL)` — a row of one
/// null, which is a different thing.
pub fn grouping_sets(sets: impl IntoExprList) -> Expr {
    grouping_element(GroupingSetKind::GroupingSets, sets)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{Psql, group, quote, raw};
    use keelson_core::build;

    fn sql(e: impl Expression) -> String {
        build(&Psql, &e).expect("render").0
    }

    #[test]
    fn distinct_renders_with_and_without_an_on_list() {
        assert_eq!(sql(Distinct::default()), "DISTINCT");
        assert_eq!(
            sql(Distinct {
                on: vec![quote("a"), quote("b")]
            }),
            r#"DISTINCT ON ("a", "b")"#
        );
    }

    #[test]
    fn excluded_qualifies_the_column_with_the_pseudo_table() {
        assert_eq!(sql(excluded("email")), r#"EXCLUDED."email""#);
    }

    /// PostgreSQL 17 `sql-select.html`, `grouping_element`.
    #[test]
    fn the_grouping_elements_use_their_own_keywords() {
        assert_eq!(
            sql(rollup((quote("a"), quote("b")))),
            r#"ROLLUP ("a", "b")"#
        );
        assert_eq!(sql(cube((quote("a"), quote("b")))), r#"CUBE ("a", "b")"#);
        assert_eq!(
            sql(grouping_sets((group(quote("a")), raw("()")))),
            r#"GROUPING SETS (("a"), ())"#
        );
    }
}