razeri-plugin-contract 0.11.0

// Copyright (C) 2026 Gabriel Shea
// Licensed under either of the Apache License, Version 2.0 or the MIT
// license, at your option.

//! The scalar-function grammars.

use super::{
    inline_parameters, note_above, or_replace, plain_name, quoted_name, quoted_text,
    refusal_comment,
};
use crate::routine_form::ScalarFunctionForm;
use crate::routine_spec::{RoutineBody, RoutineReturns, RoutineSpec};

/// **Spell this scalar function in a named grammar.**
///
/// Parameters:
/// - `form`: the grammar to spell it in. `Custom` answers with the
///   `PostgreSQL` form under a comment naming the grammar, so a dialect that is
///   nearly one of these has something to start from;
/// - `routine`: the function, in the shape a `Custom` dialect receives it.
#[must_use]
pub fn render_scalar_function_as(form: ScalarFunctionForm, routine: &RoutineSpec) -> String {
    // No wildcard arm, deliberately: a form added to the enum has to be
    // spelled here or the build stops.
    match form {
        ScalarFunctionForm::PostgreSql => postgres(routine),
        ScalarFunctionForm::Redshift => redshift(routine),
        ScalarFunctionForm::MySql => my_sql(routine),
        ScalarFunctionForm::TSql => t_sql(routine),
        ScalarFunctionForm::Snowflake => snowflake(routine),
        ScalarFunctionForm::Databricks => databricks(routine),
        ScalarFunctionForm::GoogleSql => google_sql(routine),
        ScalarFunctionForm::DuckDb => duckdb(routine),
        ScalarFunctionForm::ClickHouse => clickhouse(routine),
        ScalarFunctionForm::Custom(grammar) => note_above(
            &format!(
                "You asked for {grammar}, which is your dialect's own grammar and not one spelled \
                 here. The PostgreSQL form below is the nearest neutral starting point; the named \
                 forms are reachable by naming one."
            ),
            &postgres(routine),
        ),
        ScalarFunctionForm::Unsupported(orientation) => {
            refusal_comment("function", &routine.name, orientation)
        }
    }
}

/// The type the function returns, or an empty string where the spec carries a
/// shape no scalar grammar can declare.
fn returns(routine: &RoutineSpec) -> String {
    match &routine.returns {
        RoutineReturns::Value(data_type) => data_type.clone(),
        _ => String::new(),
    }
}

/// The one expression or query the function's value is.
fn body(routine: &RoutineSpec) -> String {
    match &routine.body {
        RoutineBody::Expression(expression) => expression.clone(),
        RoutineBody::Query(query) | RoutineBody::Block(query) => query.clone(),
        RoutineBody::Statements(statements) => statements.join("; "),
    }
}

/// `PostgreSQL` and Citus: a `RETURN` clause over a SQL body.
fn postgres(routine: &RoutineSpec) -> String {
    format!(
        "CREATE {}FUNCTION {}({})\nRETURNS {}\nLANGUAGE SQL\nRETURN {};",
        or_replace(routine),
        plain_name(routine),
        inline_parameters(routine),
        returns(routine),
        body(routine)
    )
}

/// `Redshift`: the volatility is mandatory rather than optional, there is no
/// `PARALLEL SAFE`, `LANGUAGE` follows the body rather than preceding it, and the
/// body addresses its parameters as `$1` and never by name.
fn redshift(routine: &RoutineSpec) -> String {
    let mut expression = body(routine);
    for (index, parameter) in routine.parameters.iter().enumerate() {
        if !parameter.name.is_empty() {
            expression = expression.replace(&parameter.name, &format!("${}", index + 1));
        }
    }
    format!(
        "CREATE {}FUNCTION {}({})\nRETURNS {}\nIMMUTABLE\nAS $$\n  SELECT {}\n$$ LANGUAGE sql;",
        or_replace(routine),
        plain_name(routine),
        inline_parameters(routine),
        returns(routine),
        expression
    )
}

/// `MySQL` and `MariaDB`: a replace is a `DROP` in front, and the description is
/// a clause rather than a statement.
fn my_sql(routine: &RoutineSpec) -> String {
    let mut lines = Vec::new();
    if routine.or_replace {
        lines.push(format!("DROP FUNCTION IF EXISTS {};", plain_name(routine)));
    }
    lines.push(format!("CREATE FUNCTION {}({})", plain_name(routine), inline_parameters(routine)));
    lines.push(format!("RETURNS {}", returns(routine)));
    lines.push("READS SQL DATA".to_owned());
    if let Some(comment) = &routine.comment {
        lines.push(format!("COMMENT '{}'", quoted_text(comment)));
    }
    lines.push(format!("RETURN {}", body(routine)));
    lines.join("\n")
}

/// T-SQL's scalar UDF: bracketed names, `@`-prefixed parameters, and a `BEGIN
/// … END` block whose `RETURN` is the value.
fn t_sql(routine: &RoutineSpec) -> String {
    let create = if routine.or_replace { "CREATE OR ALTER FUNCTION" } else { "CREATE FUNCTION" };
    let mut lines = vec![format!("{create} {}(", quoted_name(routine, '[', ']'))];
    let last = routine.parameters.len().saturating_sub(1);
    for (index, parameter) in routine.parameters.iter().enumerate() {
        let comma = if index == last { "" } else { "," };
        lines.push(format!("    @{} {}{}", parameter.name, parameter.data_type, comma));
    }
    lines.push(")".to_owned());
    lines.push(format!("RETURNS {}", returns(routine)));
    lines.push("AS".to_owned());
    lines.push("BEGIN".to_owned());
    lines.push(format!("    RETURN {}", body(routine)));
    lines.push("END".to_owned());
    lines.join("\n")
}

/// Snowflake's UDF, with the description as a clause.
fn snowflake(routine: &RoutineSpec) -> String {
    let mut lines = vec![
        format!(
            "CREATE {}FUNCTION {}({})",
            or_replace(routine),
            plain_name(routine),
            inline_parameters(routine)
        ),
        format!("RETURNS {}", returns(routine)),
        "LANGUAGE SQL".to_owned(),
    ];
    if let Some(comment) = &routine.comment {
        lines.push(format!("COMMENT = '{}'", quoted_text(comment)));
    }
    lines.push("AS $$".to_owned());
    lines.push(format!("RETURN {};", body(routine)));
    lines.push("$$;".to_owned());
    lines.join("\n")
}

/// Databricks SQL's UDF: a `RETURN` clause and no `LANGUAGE`.
fn databricks(routine: &RoutineSpec) -> String {
    let mut lines = vec![
        format!(
            "CREATE {}FUNCTION {}({})",
            or_replace(routine),
            plain_name(routine),
            inline_parameters(routine)
        ),
        format!("RETURNS {}", returns(routine)),
    ];
    if let Some(comment) = &routine.comment {
        lines.push(format!("COMMENT '{}'", quoted_text(comment)));
    }
    lines.push(format!("RETURN {}", body(routine)));
    lines.join("\n")
}

/// `GoogleSQL`'s `CREATE FUNCTION`: the value is a parenthesised expression.
fn google_sql(routine: &RoutineSpec) -> String {
    format!(
        "CREATE {}FUNCTION {}({})\nRETURNS {}\nAS (\n{}\n);",
        or_replace(routine),
        quoted_name(routine, '`', '`'),
        inline_parameters(routine),
        returns(routine),
        body(routine)
    )
}

/// `DuckDB`'s `CREATE MACRO`. A macro declares no types at all — not on its
/// parameters and not on its result — so the types the spec carries have nowhere
/// to go here.
fn duckdb(routine: &RoutineSpec) -> String {
    let parameters =
        routine.parameters.iter().map(|p| p.name.clone()).collect::<Vec<_>>().join(", ");
    format!(
        "CREATE {}MACRO {}({}) AS (\n{}\n);",
        or_replace(routine),
        plain_name(routine),
        parameters,
        body(routine)
    )
}

/// `ClickHouse`'s lambda. It has no schema for a function and no declared types,
/// so the name is the bare one.
fn clickhouse(routine: &RoutineSpec) -> String {
    let parameters =
        routine.parameters.iter().map(|p| p.name.clone()).collect::<Vec<_>>().join(", ");
    format!(
        "CREATE {}FUNCTION {} AS ({}) -> {}",
        or_replace(routine),
        routine.name,
        parameters,
        body(routine)
    )
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::routine_spec::{RoutineKind, RoutineParameter};

    fn a_function() -> RoutineSpec {
        let mut spec = RoutineSpec::new(RoutineKind::ScalarFunction, "busy_day");
        spec.schema = Some("retail_core".to_owned());
        spec.or_replace = true;
        spec.comment = Some("Whether the day was busy.".to_owned());
        spec.parameters = vec![
            RoutineParameter {
                name: "p_rows".to_owned(),
                data_type: "NUMERIC(19, 4)".to_owned(),
                declared_type: None,
                direction: crate::routine_spec::ParameterDirection::In,
                default: None,
            },
            RoutineParameter {
                name: "p_threshold".to_owned(),
                data_type: "BIGINT".to_owned(),
                declared_type: None,
                direction: crate::routine_spec::ParameterDirection::In,
                default: None,
            },
        ];
        spec.returns = RoutineReturns::Value("BOOLEAN".to_owned());
        spec.body = RoutineBody::Expression("p_rows > p_threshold".to_owned());
        spec
    }

    /// `PostgreSQL`, whole.
    #[test]
    fn the_postgres_form_returns_the_expression() {
        assert_eq!(
            render_scalar_function_as(ScalarFunctionForm::PostgreSql, &a_function()),
            "CREATE OR REPLACE FUNCTION retail_core.busy_day(p_rows NUMERIC(19, 4), p_threshold \
             BIGINT)\nRETURNS BOOLEAN\nLANGUAGE SQL\nRETURN p_rows > p_threshold;"
        );
    }

    /// **`Redshift` is not `PostgreSQL` here, and this is the difference that
    /// costs a day.** The volatility is mandatory, `LANGUAGE` follows the body,
    /// and the body addresses its parameters by position.
    #[test]
    fn the_redshift_form_is_positional_and_puts_language_last() {
        assert_eq!(
            render_scalar_function_as(ScalarFunctionForm::Redshift, &a_function()),
            "CREATE OR REPLACE FUNCTION retail_core.busy_day(p_rows NUMERIC(19, 4), p_threshold \
             BIGINT)\nRETURNS BOOLEAN\nIMMUTABLE\nAS $$\n  SELECT $1 > $2\n$$ LANGUAGE sql;"
        );
    }

    /// T-SQL, whole.
    #[test]
    fn the_t_sql_form_brackets_the_name_and_returns_from_a_block() {
        assert_eq!(
            render_scalar_function_as(ScalarFunctionForm::TSql, &a_function()),
            "CREATE OR ALTER FUNCTION [retail_core].[busy_day](\n    @p_rows NUMERIC(19, 4),\n    \
             @p_threshold BIGINT\n)\nRETURNS BOOLEAN\nAS\nBEGIN\n    RETURN p_rows > \
             p_threshold\nEND"
        );
    }

    /// `MySQL` drops first and carries the description as a clause.
    #[test]
    fn the_mysql_form_drops_first() {
        let rendered = render_scalar_function_as(ScalarFunctionForm::MySql, &a_function());
        assert!(rendered.starts_with("DROP FUNCTION IF EXISTS retail_core.busy_day;\n"));
        assert!(rendered.contains("\nREADS SQL DATA\nCOMMENT 'Whether the day was busy.'\n"));
    }

    /// **A macro declares no types.** `DuckDB` and `ClickHouse` both drop them,
    /// and a renderer that emitted them writes a parse error.
    #[test]
    fn the_macro_forms_declare_no_types() {
        let duck = render_scalar_function_as(ScalarFunctionForm::DuckDb, &a_function());
        assert_eq!(
            duck,
            "CREATE OR REPLACE MACRO retail_core.busy_day(p_rows, p_threshold) AS (\np_rows > \
             p_threshold\n);"
        );
        let click = render_scalar_function_as(ScalarFunctionForm::ClickHouse, &a_function());
        assert_eq!(
            click,
            "CREATE OR REPLACE FUNCTION busy_day AS (p_rows, p_threshold) -> p_rows > p_threshold"
        );
        for rendered in [duck, click] {
            assert!(!rendered.contains("NUMERIC"), "{rendered}");
        }
    }

    /// Snowflake and `GoogleSQL`, whole.
    #[test]
    fn the_snowflake_and_google_forms_carry_their_own_clauses() {
        assert_eq!(
            render_scalar_function_as(ScalarFunctionForm::Snowflake, &a_function()),
            "CREATE OR REPLACE FUNCTION retail_core.busy_day(p_rows NUMERIC(19, 4), p_threshold \
             BIGINT)\nRETURNS BOOLEAN\nLANGUAGE SQL\nCOMMENT = 'Whether the day was \
             busy.'\nAS $$\nRETURN p_rows > p_threshold;\n$$;"
        );
        assert_eq!(
            render_scalar_function_as(ScalarFunctionForm::GoogleSql, &a_function()),
            "CREATE OR REPLACE FUNCTION `retail_core`.`busy_day`(p_rows NUMERIC(19, 4), \
             p_threshold BIGINT)\nRETURNS BOOLEAN\nAS (\np_rows > p_threshold\n);"
        );
    }

    /// A refusal is the dialect's own sentence in a comment, and never a
    /// statement.
    #[test]
    fn an_unsupported_form_spells_no_statement() {
        let rendered = render_scalar_function_as(
            ScalarFunctionForm::Unsupported("compute the value in the SELECT that needs it"),
            &a_function(),
        );
        assert!(rendered.contains("compute the value in the SELECT that needs it"));
        assert!(!rendered.contains("CREATE"), "{rendered}");
    }
}