razeri-plugin-contract 0.5.0

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

//! The table-valued-function grammars.

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

/// **Spell this table-valued 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;
/// - `routine`: the function, in the shape a `Custom` dialect receives it.
#[must_use]
pub fn render_table_function_as(form: TableFunctionForm, routine: &RoutineSpec) -> String {
    // No wildcard arm, deliberately: a form added to the enum has to be
    // spelled here or the build stops.
    match form {
        TableFunctionForm::PostgreSql => postgres(routine),
        TableFunctionForm::TSql => t_sql(routine),
        TableFunctionForm::Snowflake => snowflake(routine),
        TableFunctionForm::Databricks => databricks(routine),
        TableFunctionForm::GoogleSql => google_sql(routine),
        TableFunctionForm::DuckDb => duckdb(routine),
        TableFunctionForm::ClickHouse => clickhouse(routine),
        TableFunctionForm::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),
        ),
        TableFunctionForm::Unsupported(orientation) => {
            refusal_comment("table function", &routine.name, orientation)
        }
    }
}

/// The row shape, spelled the way most grammars spell it.
fn row_shape(routine: &RoutineSpec) -> String {
    match &routine.returns {
        RoutineReturns::Columns(columns) => format!(
            "TABLE({})",
            columns
                .iter()
                .map(|c| format!("{} {}", c.name, c.data_type))
                .collect::<Vec<_>>()
                .join(", ")
        ),
        RoutineReturns::RowsOf(relation) => format!("SETOF {relation}"),
        RoutineReturns::Value(data_type) => format!("SETOF {data_type}"),
        RoutineReturns::Inferred | RoutineReturns::Nothing => "TABLE".to_owned(),
    }
}

/// The query whose rows the function returns.
fn body(routine: &RoutineSpec) -> String {
    match &routine.body {
        RoutineBody::Query(query) | RoutineBody::Block(query) => query.clone(),
        RoutineBody::Expression(expression) => expression.clone(),
        RoutineBody::Statements(statements) => statements.join(";\n"),
    }
}

/// `PostgreSQL` and Citus: `RETURNS TABLE(…)` or `RETURNS SETOF …` over a
/// dollar-quoted body.
fn postgres(routine: &RoutineSpec) -> String {
    format!(
        "CREATE {}FUNCTION {}({})\nRETURNS {}\nLANGUAGE SQL\nAS $$\n{}\n$$;",
        or_replace(routine),
        plain_name(routine),
        inline_parameters(routine),
        row_shape(routine),
        body(routine)
    )
}

/// T-SQL's inline table-valued function. The row shape is not declared: the
/// engine reads it off the query, and `RETURNS TABLE` is the whole clause.
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("RETURNS TABLE".to_owned());
    lines.push("AS".to_owned());
    lines.push("RETURN (".to_owned());
    lines.push(format!("    {}", body(routine)));
    lines.push(")".to_owned());
    lines.join("\n")
}

/// Snowflake's UDTF, 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 {}", row_shape(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(body(routine));
    lines.push("$$;".to_owned());
    lines.join("\n")
}

/// Databricks SQL: `RETURNS TABLE … RETURN <query>`, 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 {}", row_shape(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 TABLE FUNCTION`. The row shape is read off the query.
fn google_sql(routine: &RoutineSpec) -> String {
    format!(
        "CREATE {}TABLE FUNCTION {}({}) AS (\n{}\n);",
        or_replace(routine),
        quoted_name(routine, '`', '`'),
        inline_parameters(routine),
        body(routine)
    )
}

/// `DuckDB`'s table `MACRO`. There is no `RETURNS` clause at all, which is why
/// riding the `PostgreSQL` form here is a parse error at `RETURNS`.
fn duckdb(routine: &RoutineSpec) -> String {
    let parameters =
        routine.parameters.iter().map(|p| p.name.clone()).collect::<Vec<_>>().join(", ");
    format!(
        "CREATE {}MACRO {}({}) AS TABLE\n{}",
        or_replace(routine),
        plain_name(routine),
        parameters,
        body(routine)
    )
}

/// `ClickHouse`'s parameterized view. It is a view and not a function, which is
/// why the parameters are not in the statement at all: the query addresses them
/// as `{name:Type}` and the caller supplies them at the call site.
fn clickhouse(routine: &RoutineSpec) -> String {
    format!("CREATE {}VIEW {} AS\n{}", or_replace(routine), plain_name(routine), body(routine))
}

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

    fn a_function() -> RoutineSpec {
        let mut spec = RoutineSpec::new(RoutineKind::TableFunction, "customer_at");
        spec.schema = Some("retail_core".to_owned());
        spec.or_replace = true;
        spec.comment = Some("The customers as they stood.".to_owned());
        spec.parameters = vec![RoutineParameter {
            name: "p_as_of".to_owned(),
            data_type: "DATE".to_owned(),
            declared_type: None,
            direction: ParameterDirection::In,
            default: None,
        }];
        spec.returns = RoutineReturns::Columns(vec![
            RoutineColumn {
                name: "customer_sid".to_owned(),
                data_type: "BIGINT".to_owned(),
                declared_type: None,
            },
            RoutineColumn {
                name: "customer_name".to_owned(),
                data_type: "VARCHAR(200)".to_owned(),
                declared_type: None,
            },
        ]);
        spec.body =
            RoutineBody::Query("SELECT customer_sid, customer_name FROM customer".to_owned());
        spec
    }

    /// `PostgreSQL`, whole.
    #[test]
    fn the_postgres_form_declares_the_row_shape() {
        assert_eq!(
            render_table_function_as(TableFunctionForm::PostgreSql, &a_function()),
            "CREATE OR REPLACE FUNCTION retail_core.customer_at(p_as_of DATE)\nRETURNS \
             TABLE(customer_sid BIGINT, customer_name VARCHAR(200))\nLANGUAGE SQL\nAS \
             $$\nSELECT customer_sid, customer_name FROM customer\n$$;"
        );
    }

    /// **`RETURNS SETOF` is one column, and a routine declaring several has
    /// nowhere to put them.** A spec asking for rows of an existing relation gets
    /// the `SETOF` spelling and nothing else.
    #[test]
    fn rows_of_a_relation_are_setof_and_not_a_column_list() {
        let mut spec = a_function();
        spec.returns = RoutineReturns::RowsOf("customer".to_owned());
        assert!(
            render_table_function_as(TableFunctionForm::PostgreSql, &spec)
                .contains("RETURNS SETOF customer\n")
        );
    }

    /// **T-SQL reads the row shape off the query.** `RETURNS TABLE` is the whole
    /// clause for an inline TVF, and a column list there is a syntax error.
    #[test]
    fn the_t_sql_form_declares_no_columns() {
        let rendered = render_table_function_as(TableFunctionForm::TSql, &a_function());
        assert_eq!(
            rendered,
            "CREATE OR ALTER FUNCTION [retail_core].[customer_at](\n    @p_as_of \
             DATE\n)\nRETURNS TABLE\nAS\nRETURN (\n    SELECT customer_sid, customer_name FROM \
             customer\n)"
        );
        assert!(!rendered.contains("customer_sid BIGINT"), "{rendered}");
    }

    /// **`DuckDB` has no `RETURNS` clause at all**, which is why riding the
    /// `PostgreSQL` form here is a parse error at `RETURNS`.
    #[test]
    fn the_duckdb_macro_has_no_returns_clause() {
        let rendered = render_table_function_as(TableFunctionForm::DuckDb, &a_function());
        assert_eq!(
            rendered,
            "CREATE OR REPLACE MACRO retail_core.customer_at(p_as_of) AS TABLE\nSELECT \
             customer_sid, customer_name FROM customer"
        );
        assert!(!rendered.contains("RETURNS"), "{rendered}");
    }

    /// Snowflake, Databricks and `GoogleSQL`, whole.
    #[test]
    fn the_remaining_forms_carry_their_own_clauses() {
        assert!(
            render_table_function_as(TableFunctionForm::Snowflake, &a_function())
                .contains("\nLANGUAGE SQL\nCOMMENT = 'The customers as they stood.'\nAS $$\n")
        );
        assert!(
            render_table_function_as(TableFunctionForm::Databricks, &a_function())
                .ends_with("RETURN SELECT customer_sid, customer_name FROM customer")
        );
        assert!(
            render_table_function_as(TableFunctionForm::GoogleSql, &a_function())
                .starts_with("CREATE OR REPLACE TABLE FUNCTION `retail_core`.`customer_at`(")
        );
    }

    /// **`ClickHouse` has a parameterized view and not a function**, so the
    /// parameters are not in the statement: the query addresses them and the
    /// caller supplies them at the call site.
    #[test]
    fn the_clickhouse_form_is_a_view_and_carries_no_parameter_list() {
        let rendered = render_table_function_as(TableFunctionForm::ClickHouse, &a_function());
        assert_eq!(
            rendered,
            "CREATE OR REPLACE VIEW retail_core.customer_at AS\nSELECT customer_sid, \
             customer_name FROM customer"
        );
        assert!(!rendered.contains("p_as_of"), "{rendered}");
    }

    /// 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_table_function_as(
            TableFunctionForm::Unsupported(
                "a VIEW returns rows here, with the parameter in the \
                                           caller's own WHERE clause",
            ),
            &a_function(),
        );
        assert!(rendered.contains("a VIEW returns rows here"));
        assert!(!rendered.contains("CREATE"), "{rendered}");
    }
}