uqa-sql 0.5.4

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

//! Catalog and name analysis for execution-free query schema binding.

mod calls;
mod functions;
mod query_sources;
mod references;
mod windows;

pub(super) use calls::validate_call;
pub(super) use functions::validate_window_function;
pub(super) use query_sources::{with_projected_open_columns, with_query_source_columns};
pub(super) use references::require_scalar_boolean_test;
pub(super) use windows::references_local_column;

use super::{QueryBlockPlan, QueryPlan, SQLError, SQLParam, ScalarExpr, SchemaScope};
use crate::ast::{ColumnType, FunctionBinding};
use crate::routines::RoutineResolution;
use crate::semantics::{META_DOC_ID_COLUMN, META_QUALIFIER, META_SCORE_COLUMN};
use crate::{ColumnIdentity, FunctionTypeResolver, RowSchema};

pub(super) struct TableFunctionSourceValidation<'a> {
    pub(super) name: &'a str,
    pub(super) binding: Option<&'a FunctionBinding>,
    pub(super) args: &'a [ScalarExpr],
    pub(super) subqueries: &'a [QueryPlan],
    pub(super) input: &'a RowSchema,
    pub(super) params: &'a [SQLParam],
}

/// A LIMIT or OFFSET argument with the scopes it is validated against.
struct LimitArgument<'a> {
    expression: &'a ScalarExpr,
    construct: &'static str,
    /// The block's own columns over the enclosing scope.
    source: &'a RowSchema,
    enclosing: &'a RowSchema,
    subqueries: &'a [QueryPlan],
}

struct AliasReferenceScope<'a> {
    primary: &'a RowSchema,
    fallback: &'a RowSchema,
    nested: &'a RowSchema,
    subqueries: &'a [QueryPlan],
    params: &'a [SQLParam],
}

impl SchemaScope {
    /// `DISTINCT ON` and `ORDER BY` items name an output column before a column of the block's source.
    fn validate_output_references(
        &mut self,
        engine: &dyn RoutineResolution,
        block: &QueryBlockPlan,
        output: &RowSchema,
        source: &RowSchema,
        params: &[SQLParam],
    ) -> Result<(), SQLError> {
        for expression in block
            .distinct_on
            .iter()
            .chain(block.order_by.iter().map(|order| &order.expr))
        {
            self.validate_alias_reference(
                engine,
                expression,
                AliasReferenceScope {
                    primary: output,
                    fallback: source,
                    nested: source,
                    subqueries: &block.subqueries,
                    params,
                },
            )?;
        }
        Ok(())
    }

    /// Check a query block's clauses in `transformSelectStmt`'s order after its target list, against `source`, the block's own columns over `outer`, the scope of the queries that enclose it: `WHERE`, `GROUP BY`, `HAVING`, `DISTINCT ON`, `ORDER BY`, `LIMIT` and `OFFSET`, then its window frames.
    pub(super) fn validate_query_block_clauses(
        &mut self,
        engine: &dyn RoutineResolution,
        block: &QueryBlockPlan,
        source: &RowSchema,
        output: &RowSchema,
        params: &[SQLParam],
        outer: Option<&RowSchema>,
    ) -> Result<(), SQLError> {
        if let Some(predicate) = block.r#where.as_ref() {
            self.validate_expression_references(
                engine,
                predicate,
                source,
                None,
                &block.subqueries,
                params,
            )?;
            self.validate_condition(
                engine,
                predicate,
                "WHERE",
                source,
                &block.subqueries,
                params,
            )?;
        }
        for expression in block
            .group_by
            .iter()
            .chain(block.grouping_sets.iter().flatten())
        {
            let expression = crate::semantics::grouping_sets::resolve_grouping_expression(
                engine,
                expression,
                &block.projections,
                source,
                outer,
                params,
            )?;
            self.validate_expression_references(
                engine,
                &expression,
                source,
                None,
                &block.subqueries,
                params,
            )?;
        }
        if let Some(having) = block.having.as_ref() {
            self.validate_expression_references(
                engine,
                having,
                source,
                None,
                &block.subqueries,
                params,
            )?;
            self.validate_condition(engine, having, "HAVING", source, &block.subqueries, params)?;
        }
        self.validate_output_references(engine, block, output, source, params)?;
        let enclosing = outer.cloned().unwrap_or_default();
        for (expression, construct) in block
            .limit
            .iter()
            .map(|expression| (expression, "LIMIT"))
            .chain(block.offset.iter().map(|expression| (expression, "OFFSET")))
        {
            self.validate_limit_argument(
                engine,
                LimitArgument {
                    expression,
                    construct,
                    source,
                    enclosing: &enclosing,
                    subqueries: &block.subqueries,
                },
                params,
            )?;
        }
        self.validate_window_frames(
            engine,
            block,
            &windows::WindowFrameScope {
                source,
                subqueries: &block.subqueries,
                params,
            },
        )?;
        crate::semantics::grouping_sets::validate_grouped_expressions(
            engine, block, source, outer, params,
        )
    }

    /// Validate a LIMIT or OFFSET argument, which can name the columns of the enclosing queries and the parameters but no column of its own query level, as `checkExprIsVarFree` requires.
    fn validate_limit_argument(
        &mut self,
        engine: &dyn RoutineResolution,
        argument: LimitArgument<'_>,
        params: &[SQLParam],
    ) -> Result<(), SQLError> {
        let Err(error) = self.validate_expression_references(
            engine,
            argument.expression,
            argument.enclosing,
            None,
            argument.subqueries,
            params,
        ) else {
            return Ok(());
        };
        if self
            .validate_expression_references(
                engine,
                argument.expression,
                argument.source,
                None,
                argument.subqueries,
                params,
            )
            .is_ok()
        {
            return Err(SQLError::Routine {
                sqlstate: "42P10".into(),
                message: format!(
                    "argument of {} must not contain variables",
                    argument.construct
                ),
            });
        }
        Err(error)
    }

    fn validate_alias_reference(
        &mut self,
        engine: &dyn RoutineResolution,
        expression: &ScalarExpr,
        scope: AliasReferenceScope<'_>,
    ) -> Result<(), SQLError> {
        if matches!(expression, ScalarExpr::Column(_) | ScalarExpr::Position(_)) {
            self.validate_expression_references(
                engine,
                expression,
                scope.primary,
                Some(scope.fallback),
                scope.subqueries,
                scope.params,
            )
        } else {
            self.validate_expression_references(
                engine,
                expression,
                scope.nested,
                None,
                scope.subqueries,
                scope.params,
            )
        }
    }

    pub(super) fn validate_expression_references(
        &mut self,
        engine: &dyn RoutineResolution,
        expression: &ScalarExpr,
        schema: &RowSchema,
        fallback: Option<&RowSchema>,
        subqueries: &[QueryPlan],
        params: &[SQLParam],
    ) -> Result<(), SQLError> {
        let schema = self.with_stored_outer_internal_aliases(schema);
        let schema = &schema;
        let resolver =
            self.query_function_type_resolver(engine, expression, schema, subqueries, params)?;
        Self::validate_expression_references_with_resolver(
            self.rule_condition_event,
            engine,
            expression,
            schema,
            fallback,
            params,
            &resolver,
        )
    }

    /// A clause condition must be boolean, as `transformWhereClause` requires.
    pub(super) fn validate_condition(
        &mut self,
        engine: &dyn RoutineResolution,
        condition: &ScalarExpr,
        construct: &str,
        schema: &RowSchema,
        subqueries: &[QueryPlan],
        params: &[SQLParam],
    ) -> Result<(), SQLError> {
        let schema = self.with_stored_outer_internal_aliases(schema);
        super::preparation::condition_metadata(
            self, engine, condition, construct, &schema, subqueries, params,
        )
    }

    pub(super) fn validate_expression_references_with_resolver(
        rule_event: Option<crate::ast::RuleEvent>,
        engine: &dyn RoutineResolution,
        expression: &ScalarExpr,
        schema: &RowSchema,
        fallback: Option<&RowSchema>,
        params: &[SQLParam],
        resolver: &dyn FunctionTypeResolver,
    ) -> Result<(), SQLError> {
        references::validate_expression(
            rule_event, engine, expression, schema, fallback, params, resolver,
        )
    }

    pub(super) fn validate_table_function_source(
        &mut self,
        engine: &dyn RoutineResolution,
        request: TableFunctionSourceValidation<'_>,
    ) -> Result<Option<crate::semantics::ResolvedUserTableFunction>, SQLError> {
        let TableFunctionSourceValidation {
            name,
            binding,
            args,
            subqueries,
            input,
            params,
        } = request;
        let resolver = self
            .query_function_type_resolver_for_subqueries(engine, args, input, subqueries, params)?;
        for argument in args {
            Self::validate_expression_references_with_resolver(
                self.rule_condition_event,
                engine,
                argument,
                input,
                None,
                params,
                &resolver,
            )?;
        }
        functions::validate_table_function(engine, name, binding, args, input, params, &resolver)
    }
}

pub(super) fn with_table_pseudo_columns(schema: &RowSchema, qualifier: &str) -> RowSchema {
    let columns = [
        (
            ColumnIdentity::qualified(qualifier, "_doc_id"),
            Some(ColumnType::BigInteger),
        ),
        (
            ColumnIdentity::qualified(qualifier, "_score"),
            Some(ColumnType::DoublePrecision),
        ),
        (
            ColumnIdentity::qualified(qualifier, "tableoid"),
            Some(ColumnType::Oid),
        ),
        (
            ColumnIdentity::qualified(qualifier, "xmin"),
            Some(ColumnType::Xid),
        ),
    ];
    RowSchema::with_typed_virtual_identities(schema, &columns)
}

pub(super) fn with_unqualified_table_pseudo_columns(schema: &RowSchema) -> RowSchema {
    let Some(qualifier) = functions::single_pseudo_column_qualifier(schema) else {
        return schema.clone();
    };
    let columns = [
        (
            ColumnIdentity::unqualified("_doc_id"),
            schema.qualified_type(&qualifier, "_doc_id").cloned(),
        ),
        (
            ColumnIdentity::unqualified("_score"),
            schema.qualified_type(&qualifier, "_score").cloned(),
        ),
        (
            ColumnIdentity::unqualified("tableoid"),
            schema.qualified_type(&qualifier, "tableoid").cloned(),
        ),
        (
            ColumnIdentity::unqualified("xmin"),
            schema.qualified_type(&qualifier, "xmin").cloned(),
        ),
    ];
    let schema = RowSchema::with_typed_virtual_identities(schema, &columns);
    if schema.has_qualifier(META_QUALIFIER) {
        return schema;
    }
    RowSchema::with_typed_virtual_identities(
        &schema,
        &[
            (
                ColumnIdentity::qualified(META_QUALIFIER, META_DOC_ID_COLUMN),
                Some(ColumnType::BigInteger),
            ),
            (
                ColumnIdentity::qualified(META_QUALIFIER, META_SCORE_COLUMN),
                Some(ColumnType::DoublePrecision),
            ),
        ],
    )
}