uqa-sql 0.5.4

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

//! Projection row-type binding and reference validation.

use crate::ast::ColumnType;
use crate::plan::{ProjectionPlan, QueryBlockPlan, QueryPlan};
use crate::RowSchema;
use crate::{SQLError, SQLParam};

use super::{
    bind_query_plan_schema, projection_columns, BindingContext, QueryFunctionTypeResolver,
    ScalarExpr, SchemaScope,
};
use crate::routines::RoutineResolution;

type ProjectionStarColumn = (String, Option<ColumnType>, Option<u32>);

pub(super) struct ProjectionMetadata {
    pub ty: Option<ColumnType>,
    pub collation: Option<u32>,
    pub record: Option<crate::schema::RecordMetadata>,
}

/// Keep only selections that differ from their output type's default while the output schema is being assembled.
#[derive(Default)]
pub(super) struct OutputCollations(Vec<(usize, u32)>);

impl OutputCollations {
    pub(super) fn push(&mut self, position: usize, ty: Option<&ColumnType>, oid: Option<u32>) {
        if let Some((ty, oid)) = ty.zip(oid) {
            if crate::schema::composites::type_is_collatable(ty)
                && crate::catalog::type_metadata::pg_type_collation_oid(ty) != i64::from(oid)
            {
                self.0.push((position, oid));
            }
        }
    }

    pub(super) fn apply(self, mut schema: RowSchema) -> RowSchema {
        for (position, oid) in self.0 {
            schema = schema
                .with_collation_override(crate::schema::CollationBinding::Column(position), oid)
                .expect("selected collation belongs to its typed output");
        }
        schema
    }
}

/// Bind a projection against an already-declared input schema. `star_schema` identifies the relation expanded by bare `*`; `expression_schema` may also contain joined sources and hidden lookup aliases used by scalar expressions.
pub fn bind_projection_output_schema(
    routines: &dyn RoutineResolution,
    projections: &[ProjectionPlan],
    expression_schema: &RowSchema,
    star_schema: &RowSchema,
    subqueries: &[QueryPlan],
    params: &[SQLParam],
    ctes: &BindingContext,
) -> Result<RowSchema, SQLError> {
    projection_output_schema(
        SchemaScope::from_context(ctes)?,
        routines,
        projections,
        expression_schema,
        star_schema,
        subqueries,
        params,
    )
}

/// Derive and validate a projection's exact output row type without executing it.
pub fn analyze_projection_output_schema(
    routines: &dyn RoutineResolution,
    projections: &[ProjectionPlan],
    expression_schema: &RowSchema,
    star_schema: &RowSchema,
    subqueries: &[QueryPlan],
    params: &[SQLParam],
    ctes: &BindingContext,
) -> Result<RowSchema, SQLError> {
    projection_output_schema(
        SchemaScope::for_analysis(ctes)?,
        routines,
        projections,
        expression_schema,
        star_schema,
        subqueries,
        params,
    )
}

fn projection_output_schema(
    mut scope: SchemaScope,
    routines: &dyn RoutineResolution,
    projections: &[ProjectionPlan],
    expression_schema: &RowSchema,
    star_schema: &RowSchema,
    subqueries: &[QueryPlan],
    params: &[SQLParam],
) -> Result<RowSchema, SQLError> {
    let labels = projection_columns(projections);
    let mut columns = Vec::new();
    let mut types = Vec::new();
    let mut records = Vec::new();
    let mut collations = OutputCollations::default();
    for (position, projection) in projections.iter().enumerate() {
        let expansion_schema = match projection.expr {
            ScalarExpr::QualifiedStar(_) => expression_schema,
            _ => star_schema,
        };
        if let Some(star_columns) = projection_star_columns(&projection.expr, expansion_schema)? {
            records.extend(
                super::record_fields::star_metadata(&projection.expr, expansion_schema)
                    .unwrap_or_default(),
            );
            for (column, ty, collation) in star_columns {
                collations.push(types.len(), ty.as_ref(), collation);
                columns.push(column);
                types.push(ty);
            }
            continue;
        }
        columns.push(labels[position].clone());
        let metadata = scope.bind_projection_metadata(
            routines,
            &projection.expr,
            expression_schema,
            subqueries,
            params,
        )?;
        collations.push(types.len(), metadata.ty.as_ref(), metadata.collation);
        types.push(metadata.ty);
        records.push(metadata.record);
    }
    for projection in projections {
        scope.collations.validate_expression(&projection.expr)?;
    }
    Ok(collations
        .apply(RowSchema::with_types(columns, types))
        .with_record_metadata(
            records
                .into_iter()
                .enumerate()
                .filter_map(|(index, fields)| fields.map(|fields| (index, fields))),
        ))
}

/// Validate every scalar expression in a query block while the physical input still carries declared SQL types. This must precede polymorphic rewrites such as `pg_typeof`, because an invalid common type is an error, not an `unknown` result.
pub fn validate_query_block_expression_types(
    routines: &dyn RoutineResolution,
    statement: &QueryBlockPlan,
    schema: &RowSchema,
    params: &[SQLParam],
    ctes: &BindingContext,
) -> Result<(), SQLError> {
    let scalar_subquery_types = statement
        .subqueries
        .iter()
        .map(|plan| {
            bind_query_plan_schema(routines, plan, params, ctes, Some(schema))
                .map(|output| output.column_type(0).cloned())
        })
        .collect::<Result<Vec<_>, _>>()?;
    let resolver = QueryFunctionTypeResolver {
        routines,
        scalar_subquery_types: Some(scalar_subquery_types),
        defer_routine_namespace_errors: true,
    };
    for expression in statement
        .projections
        .iter()
        .map(|projection| &projection.expr)
        .chain(statement.group_by.iter())
        .chain(statement.grouping_sets.iter().flatten())
        .chain(statement.order_by.iter().map(|order| &order.expr))
        .chain(statement.distinct_on.iter())
        .chain(statement.r#where.iter())
        .chain(statement.having.iter())
        .chain(statement.limit.iter())
        .chain(statement.offset.iter())
    {
        crate::scalar_type_with_resolver(expression, schema, params, &resolver)?;
    }
    for expression in statement
        .group_by
        .iter()
        .chain(statement.grouping_sets.iter().flatten())
        .chain(statement.distinct_on.iter())
    {
        if let Some(ty) = crate::scalar_type_with_resolver(expression, schema, params, &resolver)? {
            crate::type_resolution::ComparisonSupport::Equality
                .require_in_catalog(&ty, routines)?;
        }
    }
    for order in &statement.order_by {
        if let Some(ty) = crate::scalar_type_with_resolver(&order.expr, schema, params, &resolver)?
        {
            crate::type_resolution::ComparisonSupport::Ordering
                .require_in_catalog(&ty, routines)?;
        }
    }
    Ok(())
}

/// Validate every query-block reference only after the caller has the authoritative source schema. This preserves registered table-function row shapes and checks recursive argument references before definitive routine namespace lookup. `schema` holds the block's columns over `outer`, the scope of the queries that enclose the block.
pub fn validate_query_block_references(
    routines: &dyn RoutineResolution,
    statement: &QueryBlockPlan,
    schema: &RowSchema,
    params: &[SQLParam],
    ctes: &BindingContext,
    outer: Option<&RowSchema>,
) -> Result<(), SQLError> {
    let output = analyze_projection_output_schema(
        routines,
        &statement.projections,
        schema,
        schema,
        &statement.subqueries,
        params,
        ctes,
    )?;
    if statement.distinct && statement.distinct_on.is_empty() {
        for ty in output.column_types().iter().flatten() {
            crate::type_resolution::ComparisonSupport::Equality.require_in_catalog(ty, routines)?;
        }
    }
    SchemaScope::for_analysis(ctes)?
        .validate_query_block_clauses(routines, statement, schema, &output, params, outer)
}

pub(super) fn projection_star_columns(
    expression: &ScalarExpr,
    schema: &RowSchema,
) -> Result<Option<Vec<ProjectionStarColumn>>, SQLError> {
    match expression {
        ScalarExpr::Star => Ok(Some(
            schema
                .columns()
                .iter()
                .enumerate()
                .map(|(position, column)| {
                    (
                        schema.public_name(position).unwrap_or(column).to_string(),
                        schema.column_type(position).cloned(),
                        schema.column_collation(position),
                    )
                })
                .collect(),
        )),
        ScalarExpr::QualifiedStar(qualifier) => {
            let columns = schema
                .qualified_star_position_layout(qualifier)
                .into_iter()
                .map(|(column, logical, _, ty)| {
                    let collation = logical.map_or_else(
                        || schema.qualified_collation(qualifier, &column),
                        |position| schema.column_collation(position),
                    );
                    (column, ty, collation)
                })
                .collect::<Vec<_>>();
            if columns.is_empty() {
                return Err(SQLError::UnknownTable(qualifier.clone()));
            }
            Ok(Some(columns))
        }
        _ => Ok(None),
    }
}

pub(super) fn rename_schema(
    schema: &RowSchema,
    aliases: &[String],
    qualifier: Option<&str>,
) -> RowSchema {
    let columns = schema
        .columns()
        .iter()
        .enumerate()
        .map(|(position, column)| {
            aliases
                .get(position)
                .cloned()
                .unwrap_or_else(|| schema.public_name(position).unwrap_or(column).to_string())
        })
        .collect();
    let renamed = match qualifier {
        Some(qualifier) => {
            RowSchema::with_qualified_types(qualifier, columns, schema.column_types().to_vec())
        }
        None => RowSchema::with_types(columns, schema.column_types().to_vec()),
    };
    let renamed = renamed
        .with_column_collations((0..schema.len()).map(|position| schema.column_collation(position)))
        .with_record_fields_from(schema);
    let renamed = if schema.columns_are_open(None) {
        RowSchema::with_open_columns(&renamed, qualifier)
    } else {
        renamed
    };
    let mut hidden = Vec::new();
    let mut conflicting = Vec::new();
    for (identity, ty) in schema.typed_virtual_identities() {
        let conflicts = match identity.qualifier() {
            Some(source) => schema.qualified_column_is_ambiguous(source, identity.column()),
            None => schema.column_is_ambiguous(identity.column()),
        };
        let mapped = qualifier.map_or_else(
            || vec![identity.clone()],
            |qualifier| {
                vec![
                    crate::ColumnIdentity::unqualified(identity.column()),
                    crate::ColumnIdentity::qualified(qualifier, identity.column()),
                ]
            },
        );
        for identity in mapped {
            if conflicts {
                conflicting.push((identity, ty.cloned()));
            } else {
                hidden.push((identity, ty.cloned()));
            }
        }
    }
    let renamed = RowSchema::with_typed_virtual_identities(&renamed, &hidden);
    RowSchema::with_typed_conflicting_virtual_identities(&renamed, &conflicting)
}