uqa-sql 0.5.2

//
// Unified Query Algebra
//
// Copyright (c) 2023-2026 Cognica, Inc.
//

use super::{
    Deserialize, Expr, FunctionParallel, RoutineAclEntry, RoutineAttributeClauses,
    RoutineConfigAction, RoutineSecurityAttributes, Serialize, Statement,
};

/// Parameter mode of a `CREATE FUNCTION` / `CREATE PROCEDURE`
/// argument. Mirrors `PostgreSQL`'s `FunctionParameterMode`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum FunctionParamMode {
    /// `IN` (also the default when no mode is written).
    In,
    /// `OUT` - shapes the result row, not part of a function's call
    /// signature (but part of a procedure's).
    Out,
    /// `INOUT` - accepted as input and returned in the result row.
    InOut,
    /// `VARIADIC` - a trailing array parameter that accepts either expanded element arguments or one explicit `VARIADIC` array argument.
    Variadic,
    /// `RETURNS TABLE (col type, ...)` column. Behaves like an `OUT`
    /// parameter of a set-returning function.
    Table,
}

/// An input parameter as a SQL-standard body refers to it: the name the body uses, which follows `PostgreSQL`'s positional naming, and where the parameter's value arrives in a call.
#[derive(Debug, Clone, Copy)]
pub struct SQLBodyParameter<'a> {
    /// The name the body uses, empty when the declared parameter at this position is unnamed.
    pub name: &'a str,
    pub parameter: &'a FunctionParam,
    /// The parameter's 1-based position among the call's arguments, which include a procedure's output parameters.
    pub call_position: usize,
}

/// One declared parameter of a user-defined function or procedure.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FunctionParam {
    /// Parameter name. Empty for unnamed parameters (`f(integer)`),
    /// which are only addressable as `$n`.
    pub name: String,
    /// Raw type name as written (last segment, lower-cased by the
    /// compiler; e.g. `int4`, `text`, `numeric`).
    pub type_name: String,
    /// Parsed relation and column identity for `%TYPE`; ordinary types have no reference.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub type_reference: Option<RoutineColumnTypeReference>,
    /// The type as the statement wrote it, as `TypeNameToString` spells it, which names a missing type until registration resolves the type.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub written_type: Option<String>,
    pub mode: FunctionParamMode,
    /// `DEFAULT <expr>` for trailing input parameters.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub default: Option<Expr>,
    /// The analyzed default expression type, used when an omitted polymorphic input supplies call typing.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub default_type: Option<RoutineDefaultType>,
}

/// The type of a stored routine default. Some NULL inputs retain a polymorphic
/// pseudo-type without claiming that it is a concrete SQL value carrier.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum RoutineDefaultType {
    Concrete(super::ColumnType),
    Polymorphic(String),
}

/// Structured relation-column identity carried by a routine `%TYPE` declaration until catalog binding resolves it to a concrete SQL type.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RoutineColumnTypeReference {
    pub schema: Option<String>,
    pub relation: String,
    pub column: String,
}

impl RoutineColumnTypeReference {
    pub fn new(schema: Option<String>, relation: String, column: String) -> Self {
        Self {
            schema,
            relation,
            column,
        }
    }

    pub fn relation_reference(&self) -> String {
        match self.schema.as_deref() {
            Some(schema) => format!(
                "{}.{}",
                render_identifier_component(schema),
                render_identifier_component(&self.relation)
            ),
            None => render_identifier_component(&self.relation),
        }
    }

    pub fn type_reference(&self) -> String {
        format!(
            "{}.{}%type",
            self.relation_reference(),
            render_identifier_component(&self.column)
        )
    }
}

fn render_identifier_component(component: &str) -> String {
    let can_render_bare = component
        .bytes()
        .enumerate()
        .all(|(index, byte)| match byte {
            b'a'..=b'z' | b'_' => true,
            b'0'..=b'9' | b'$' => index != 0,
            _ => false,
        });
    if can_render_bare && !component.is_empty() {
        component.to_string()
    } else {
        format!("\"{}\"", component.replace('"', "\"\""))
    }
}

/// Declared result shape of a user-defined function.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum FunctionReturns {
    /// Procedures and functions whose result is shaped purely by
    /// `OUT` parameters carry no explicit `RETURNS` clause.
    None,
    /// `RETURNS <type>` - includes `RETURNS void` and `RETURNS record`.
    Scalar { type_name: String },
    /// `RETURNS SETOF <type>`.
    SetOf { type_name: String },
    /// `RETURNS TABLE (...)`. The column list lives in
    /// [`CreateFunction::params`] as [`FunctionParamMode::Table`]
    /// entries; this variant just records the set-returning shape.
    Table,
}

/// `IMMUTABLE` / `STABLE` / `VOLATILE` marker.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
pub enum FunctionVolatility {
    Immutable,
    Stable,
    #[default]
    Volatile,
}

/// Body of a user-defined routine.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum FunctionBody {
    /// `AS $$ ... $$` - raw source text, parsed per language at
    /// registration time.
    Source(String),
    /// SQL-standard body (`BEGIN ATOMIC ... END` / `RETURN expr`)
    /// compiled straight to statements.
    Statements(Vec<Statement>),
}

/// How a SQL-standard body was written; `pg_get_function_sqlbody` reproduces it.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum SQLBodyForm {
    /// `RETURN expr`, stored as one `SELECT expr`.
    Return,
    /// `BEGIN ATOMIC stmt; ... END`.
    Atomic,
}

/// `CREATE [OR REPLACE] FUNCTION | PROCEDURE`.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CreateFunction {
    /// Stable catalog identity. The engine assigns this once when the routine is created and preserves it across replacement and rename.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub object_id: Option<[u8; 16]>,
    /// The catalog tuple's replacement identity. A replacement changes this token while keeping `object_id`; older definitions use their object identity until first replacement.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub catalog_revision: Option<[u8; 16]>,
    /// The public OID allocated when the routine was created; replacement and rename keep it. Routines created before OIDs were recorded derive it from their identity.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub catalog_oid: Option<u32>,
    pub name: String,
    pub or_replace: bool,
    pub is_procedure: bool,
    pub params: Vec<FunctionParam>,
    pub returns: FunctionReturns,
    /// Parsed `%TYPE` identity for a scalar or set return declaration until registration resolves it.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub return_type_reference: Option<RoutineColumnTypeReference>,
    /// The declared result type as the statement wrote it, without SETOF, until registration resolves it.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub return_written_type: Option<String>,
    /// Lower-cased language name (`plpgsql`, `sql`).
    pub language: String,
    pub body: FunctionBody,
    /// The written form of a SQL-standard body. Definitions stored before the form was recorded omit it.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub sql_body_form: Option<SQLBodyForm>,
    /// Effective schema search path captured when a SQL-standard body or parameter default is catalog-bound. String and PL/pgSQL bodies keep dynamic lookup, but their parameter defaults still use this captured path.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub creation_search_path: Vec<String>,
    pub volatility: FunctionVolatility,
    /// `STRICT` / `RETURNS NULL ON NULL INPUT` - the function is not
    /// invoked when any input argument is NULL; the result is NULL.
    pub strict: bool,
    /// Catalog owner incarnation. Parsed declarations and anonymous blocks are unbound; registration binds the effective current role before publication.
    #[serde(
        default,
        deserialize_with = "super::routine_security::deserialize_routine_owner"
    )]
    pub owner: Option<uqa_core::catalog_role::RoleIdentity>,
    /// Execution identity and leakproofness, flattened to retain the catalog-definition wire shape.
    #[serde(default, flatten)]
    pub security: RoutineSecurityAttributes,
    /// Parallel-safety classification.
    #[serde(default)]
    pub parallel: FunctionParallel,
    /// Optional planner support routine identity.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub support: Option<String>,
    /// `COST`, the estimated execution cost in units of `cpu_operator_cost` that `pg_proc.procost` reports, a `float4` as `PostgreSQL` stores it; without one the language's default applies.
    #[serde(
        default,
        skip_serializing_if = "Option::is_none",
        with = "super::routine_estimate"
    )]
    pub cost: Option<f32>,
    /// `ROWS`, the estimated number of rows a set-returning routine returns, which `pg_proc.prorows` reports; without one the default applies.
    #[serde(
        default,
        skip_serializing_if = "Option::is_none",
        with = "super::routine_estimate"
    )]
    pub rows: Option<f32>,
    /// Effective per-routine configuration as `name=value` pairs in declaration order.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub config: Vec<(String, String)>,
    /// Creation-time configuration actions awaiting engine/session resolution. Registration consumes this list before persistence.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub config_actions: Vec<RoutineConfigAction>,
    /// The statement's attribute clauses in written order, which registration checks and consumes before persistence.
    #[serde(default, skip_serializing_if = "RoutineAttributeClauses::is_empty")]
    pub attribute_clauses: RoutineAttributeClauses,
    /// Explicit execution privileges, including the owner's revocable EXECUTE. `None` means the `PostgreSQL` default (PUBLIC and owner EXECUTE); ownership always retains implicit grant options.
    #[serde(default)]
    pub execute_acl: Option<Vec<RoutineAclEntry>>,
}

impl CreateFunction {
    /// Parameters that define routine identity: `IN` + `INOUT` + `VARIADIC`, in declaration order.
    pub fn identity_params(&self) -> Vec<&FunctionParam> {
        self.params
            .iter()
            .filter(|param| Self::is_identity_param(param))
            .collect()
    }

    /// Number of parameters that define routine identity.
    pub fn identity_arity(&self) -> usize {
        self.params
            .iter()
            .filter(|param| Self::is_identity_param(param))
            .count()
    }

    fn is_identity_param(param: &FunctionParam) -> bool {
        matches!(
            param.mode,
            FunctionParamMode::In | FunctionParamMode::InOut | FunctionParamMode::Variadic
        )
    }

    /// Parameters supplied by a call: identity parameters for functions and every non-`TABLE` parameter for procedures.
    pub fn call_params(&self) -> Vec<&FunctionParam> {
        self.params
            .iter()
            .filter(|param| self.is_call_param(param))
            .collect()
    }

    /// Number of declared call parameters; a variadic parameter can consume multiple actual arguments.
    pub fn call_arity(&self) -> usize {
        self.params
            .iter()
            .filter(|param| self.is_call_param(param))
            .count()
    }

    /// Minimum number of actual arguments for ordinary expanded notation; a variadic parameter accepts zero elements.
    pub fn required_call_arity(&self) -> usize {
        self.params
            .iter()
            .filter(|param| {
                self.is_call_param(param)
                    && param.default.is_none()
                    && param.mode != FunctionParamMode::Variadic
            })
            .count()
    }

    fn is_call_param(&self, param: &FunctionParam) -> bool {
        match param.mode {
            FunctionParamMode::In | FunctionParamMode::InOut | FunctionParamMode::Variadic => true,
            FunctionParamMode::Out => self.is_procedure,
            FunctionParamMode::Table => false,
        }
    }

    /// The parameters a SQL-standard body refers to, as `interpret_AS_clause` gives them to its parser: the input parameters in order, the one at position `n` going by the name of the parameter declared at position `n` among all parameters. When an output parameter precedes an input, the body therefore names that input by the output parameter's name, as `PostgreSQL` does.
    pub fn sql_body_parameters(&self) -> Vec<SQLBodyParameter<'_>> {
        let call_params = self.call_params();
        self.identity_params()
            .into_iter()
            .enumerate()
            .map(|(index, parameter)| SQLBodyParameter {
                name: self
                    .params
                    .get(index)
                    .map_or("", |declared| declared.name.as_str()),
                parameter,
                call_position: call_params
                    .iter()
                    .position(|call| std::ptr::eq(*call, parameter))
                    .map_or(index + 1, |position| position + 1),
            })
            .collect()
    }

    /// Backward-compatible alias for [`Self::call_arity`].
    pub fn signature_arity(&self) -> usize {
        self.call_arity()
    }

    /// Backward-compatible alias for [`Self::required_call_arity`].
    pub fn required_arity(&self) -> usize {
        self.required_call_arity()
    }

    /// Backward-compatible alias for [`Self::call_params`].
    pub fn signature_params(&self) -> Vec<&FunctionParam> {
        self.call_params()
    }

    /// Parameters that shape the result row: `OUT` + `INOUT` +
    /// `RETURNS TABLE` columns, in declaration order.
    pub fn output_params(&self) -> Vec<&FunctionParam> {
        self.params
            .iter()
            .filter(|p| {
                matches!(
                    p.mode,
                    FunctionParamMode::Out | FunctionParamMode::InOut | FunctionParamMode::Table
                )
            })
            .collect()
    }

    /// True when the routine produces a row set (`RETURNS SETOF` /
    /// `RETURNS TABLE`).
    pub fn returns_set(&self) -> bool {
        matches!(
            self.returns,
            FunctionReturns::SetOf { .. } | FunctionReturns::Table
        )
    }
}

/// One `DROP FUNCTION` / `DROP PROCEDURE` target.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DropFunctionItem {
    pub name: String,
    /// `Some(types)` when the statement spelled an argument list
    /// (`DROP FUNCTION f(int, int)` - matched by canonical argument
    /// types); `None` for the bare-name form
    /// (`DROP FUNCTION f`).
    pub arg_types: Option<Vec<String>>,
}

/// `DROP FUNCTION [IF EXISTS] name[(argtypes)] [, ...]` and the
/// `DROP PROCEDURE` equivalent.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DropFunctionStmt {
    pub is_procedure: bool,
    pub if_exists: bool,
    #[serde(default)]
    pub cascade: bool,
    pub items: Vec<DropFunctionItem>,
}