Skip to main content

uqa_sql/schema/
generated.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! `PostgreSQL` 18 generated-column validation and row computation.
8
9use super::SchemaExpressionCatalog;
10use crate::ast::ForeignKey;
11use crate::ast::{ColumnDef, Expr, GeneratedColumnKind, TableKeyConstraint};
12use crate::{
13    assignment::conversion::convert_value_to_column_type, semantics::aggregates, ColumnType,
14    SQLError,
15};
16
17pub(super) mod typing;
18
19pub fn prepare_generated_columns(
20    engine: &dyn SchemaExpressionCatalog,
21    qualifier: &str,
22    columns: &mut [ColumnDef],
23    key_constraints: &[TableKeyConstraint],
24    foreign_keys: &[ForeignKey],
25) -> Result<(), SQLError> {
26    let snapshot = columns.to_vec();
27    for (index, column) in snapshot.iter().enumerate() {
28        let Some(generated) = column.generated.as_ref() else {
29            continue;
30        };
31        if column.default.is_some() {
32            return Err(SQLError::TypeMismatch(format!(
33                "both default and generation expression specified for column `{}`",
34                column.name
35            )));
36        }
37        if column.auto_increment.is_some() {
38            return Err(SQLError::TypeMismatch(format!(
39                "both identity and generation expression specified for column `{}`",
40                column.name
41            )));
42        }
43        if generated.kind == GeneratedColumnKind::Virtual {
44            validate_virtual_column_envelope(column, key_constraints, foreign_keys)?;
45        }
46        let plan = crate::plan::ExpressionPlan::lower((*generated.expression).clone());
47        if !plan.subqueries.is_empty() {
48            return Err(SQLError::TypeMismatch(
49                "cannot use subquery in column generation expression".into(),
50            ));
51        }
52        if aggregates::contains_aggregate(engine, &plan.scalar) {
53            return Err(SQLError::TypeMismatch(
54                "aggregate functions are not allowed in column generation expressions".into(),
55            ));
56        }
57        validate_generation_expression(
58            engine,
59            qualifier,
60            &snapshot,
61            &generated.expression,
62            generated.kind,
63        )?;
64        let prepared = columns[index]
65            .generated
66            .as_mut()
67            .ok_or_else(|| SQLError::Internal("generated column disappeared".into()))?;
68        bind_schema_column_references(&mut prepared.expression, qualifier);
69        let (expression_type, function_dependencies) =
70            typing::infer_generation_expression(engine, &snapshot, &mut prepared.expression)?;
71        crate::catalog::regrole_dependencies::reject_stored_regrole_constants(
72            engine,
73            &prepared.expression,
74            Some(&column.ty),
75        )?;
76        if let typing::GenerationType::UnknownLiteral(value) = &expression_type {
77            convert_value_to_column_type(uqa_core::Value::Str(value.clone()), &column.ty)?;
78        } else if !typing::generation_type_assignable_to(&expression_type, &column.ty) {
79            return Err(SQLError::TypeMismatch(format!(
80                "column `{}` has type {} but generation expression has type {}",
81                column.name,
82                crate::catalog::type_metadata::column_type_name(&column.ty),
83                typing::generation_type_name(&expression_type)
84            )));
85        }
86        prepared.function_dependencies = function_dependencies;
87    }
88    Ok(())
89}
90
91fn validate_virtual_column_envelope(
92    column: &ColumnDef,
93    key_constraints: &[TableKeyConstraint],
94    foreign_keys: &[ForeignKey],
95) -> Result<(), SQLError> {
96    if contains_engine_defined_type(&column.ty) {
97        return Err(SQLError::TypeMismatch(format!(
98            "virtual generated column `{}` cannot use a user-defined type",
99            column.name
100        )));
101    }
102    if column.primary_key
103        || key_constraints.iter().any(|constraint| {
104            constraint.kind == crate::ast::TableKeyConstraintKind::PrimaryKey
105                && constraint.columns.iter().any(|name| name == &column.name)
106        })
107    {
108        return Err(SQLError::TypeMismatch(
109            "primary keys on virtual generated columns are not supported".into(),
110        ));
111    }
112    if column.unique
113        || key_constraints.iter().any(|constraint| {
114            constraint.kind == crate::ast::TableKeyConstraintKind::Unique
115                && constraint.columns.iter().any(|name| name == &column.name)
116        })
117    {
118        return Err(SQLError::TypeMismatch(
119            "unique constraints on virtual generated columns are not supported".into(),
120        ));
121    }
122    if column.references.is_some()
123        || foreign_keys.iter().any(|foreign_key| {
124            foreign_key
125                .local_columns
126                .iter()
127                .any(|name| name == &column.name)
128        })
129    {
130        return Err(SQLError::TypeMismatch(
131            "foreign key constraints on virtual generated columns are not supported".into(),
132        ));
133    }
134    Ok(())
135}
136
137fn contains_engine_defined_type(ty: &ColumnType) -> bool {
138    match ty {
139        ColumnType::Vector(_) | ColumnType::Tensor(_) => true,
140        ColumnType::Array(element) => contains_engine_defined_type(element),
141        _ => false,
142    }
143}
144
145#[expect(
146    clippy::too_many_lines,
147    reason = "preserves generated coercion diagnostics"
148)]
149fn validate_generation_expression(
150    engine: &dyn SchemaExpressionCatalog,
151    qualifier: &str,
152    columns: &[ColumnDef],
153    expression: &Expr,
154    kind: GeneratedColumnKind,
155) -> Result<(), SQLError> {
156    match expression {
157        Expr::Column(name) => validate_generation_column_reference(columns, name),
158        Expr::QualifiedColumn {
159            qualifier: expression_qualifier,
160            column,
161            ..
162        } => {
163            if expression_qualifier != qualifier {
164                return Err(SQLError::UnknownTable(expression_qualifier.clone()));
165            }
166            validate_generation_column_reference(columns, column)
167        }
168        Expr::Func {
169            name,
170            args,
171            distinct,
172            order_by,
173            filter,
174            ..
175        } => {
176            if *distinct || !order_by.is_empty() || filter.is_some() {
177                return Err(SQLError::TypeMismatch(
178                    "aggregate syntax is not allowed in column generation expressions".into(),
179                ));
180            }
181            if kind == GeneratedColumnKind::Virtual
182                && (engine
183                    .registered_runtime_function_volatility(name)
184                    .is_some()
185                    || engine.lookup_visible_sql_functions(name)?.is_some())
186            {
187                return Err(SQLError::TypeMismatch(
188                    "generation expression uses user-defined function; virtual generated columns cannot use user-defined functions"
189                        .into(),
190                ));
191            }
192            for argument in args {
193                validate_generation_expression(engine, qualifier, columns, argument, kind)?;
194            }
195            Ok(())
196        }
197        Expr::Array(items) | Expr::Row(items) | Expr::And(items) | Expr::Or(items) => {
198            for item in items {
199                validate_generation_expression(engine, qualifier, columns, item, kind)?;
200            }
201            Ok(())
202        }
203        Expr::Binary { lhs, rhs, .. } => {
204            validate_generation_expression(engine, qualifier, columns, lhs, kind)?;
205            validate_generation_expression(engine, qualifier, columns, rhs, kind)
206        }
207        Expr::Not(inner)
208        | Expr::UnaryMinus(inner)
209        | Expr::IsNull { expr: inner, .. }
210        | Expr::Cast { expr: inner, .. } => {
211            validate_generation_expression(engine, qualifier, columns, inner, kind)
212        }
213        Expr::Between { expr, low, high } => {
214            validate_generation_expression(engine, qualifier, columns, expr, kind)?;
215            validate_generation_expression(engine, qualifier, columns, low, kind)?;
216            validate_generation_expression(engine, qualifier, columns, high, kind)
217        }
218        Expr::InList { expr, list, .. } => {
219            validate_generation_expression(engine, qualifier, columns, expr, kind)?;
220            for item in list {
221                validate_generation_expression(engine, qualifier, columns, item, kind)?;
222            }
223            Ok(())
224        }
225        Expr::Case {
226            base,
227            when,
228            else_branch,
229        } => {
230            if let Some(base) = base {
231                validate_generation_expression(engine, qualifier, columns, base, kind)?;
232            }
233            for (condition, result) in when {
234                validate_generation_expression(engine, qualifier, columns, condition, kind)?;
235                validate_generation_expression(engine, qualifier, columns, result, kind)?;
236            }
237            if let Some(else_branch) = else_branch {
238                validate_generation_expression(engine, qualifier, columns, else_branch, kind)?;
239            }
240            Ok(())
241        }
242        Expr::Default | Expr::Param(_) => Err(SQLError::TypeMismatch(
243            "parameters and DEFAULT are not allowed in column generation expressions".into(),
244        )),
245        Expr::Star | Expr::QualifiedStar(_) => Err(SQLError::TypeMismatch(
246            "whole-row references are not allowed in column generation expressions".into(),
247        )),
248        Expr::InternalColumn(_) => Err(SQLError::Internal(
249            "executor-only column reached generation expression validation".into(),
250        )),
251        Expr::WindowCall { .. } => Err(SQLError::TypeMismatch(
252            "window functions are not allowed in column generation expressions".into(),
253        )),
254        Expr::ScalarSubquery(_) | Expr::Exists { .. } | Expr::InSubquery { .. } => Err(
255            SQLError::TypeMismatch("cannot use subquery in column generation expression".into()),
256        ),
257        Expr::Literal(_) | Expr::TypedLiteral { .. } => Ok(()),
258    }
259}
260
261pub fn bind_schema_column_references(expression: &mut Expr, qualifier: &str) {
262    if let Expr::QualifiedColumn {
263        qualifier: expression_qualifier,
264        column,
265    } = expression
266    {
267        if expression_qualifier == qualifier {
268            *expression = Expr::Column(column.clone());
269        }
270        return;
271    }
272    match expression {
273        Expr::Func {
274            args,
275            order_by,
276            filter,
277            ..
278        } => {
279            for argument in args {
280                bind_schema_column_references(argument, qualifier);
281            }
282            for order in order_by {
283                bind_schema_column_references(&mut order.expr, qualifier);
284            }
285            if let Some(filter) = filter {
286                bind_schema_column_references(filter, qualifier);
287            }
288        }
289        Expr::Array(items) | Expr::Row(items) | Expr::And(items) | Expr::Or(items) => {
290            for item in items {
291                bind_schema_column_references(item, qualifier);
292            }
293        }
294        Expr::Binary { lhs, rhs, .. } => {
295            bind_schema_column_references(lhs, qualifier);
296            bind_schema_column_references(rhs, qualifier);
297        }
298        Expr::Not(inner)
299        | Expr::UnaryMinus(inner)
300        | Expr::IsNull { expr: inner, .. }
301        | Expr::Cast { expr: inner, .. } => {
302            bind_schema_column_references(inner, qualifier);
303        }
304        Expr::Between { expr, low, high } => {
305            bind_schema_column_references(expr, qualifier);
306            bind_schema_column_references(low, qualifier);
307            bind_schema_column_references(high, qualifier);
308        }
309        Expr::InList { expr, list, .. } => {
310            bind_schema_column_references(expr, qualifier);
311            for item in list {
312                bind_schema_column_references(item, qualifier);
313            }
314        }
315        Expr::Case {
316            base,
317            when,
318            else_branch,
319        } => {
320            if let Some(base) = base {
321                bind_schema_column_references(base, qualifier);
322            }
323            for (condition, result) in when {
324                bind_schema_column_references(condition, qualifier);
325                bind_schema_column_references(result, qualifier);
326            }
327            if let Some(else_branch) = else_branch {
328                bind_schema_column_references(else_branch, qualifier);
329            }
330        }
331        Expr::Star
332        | Expr::QualifiedStar(_)
333        | Expr::Default
334        | Expr::Column(_)
335        | Expr::QualifiedColumn { .. }
336        | Expr::InternalColumn(_)
337        | Expr::Literal(_)
338        | Expr::TypedLiteral { .. }
339        | Expr::Param(_)
340        | Expr::WindowCall { .. }
341        | Expr::ScalarSubquery(_)
342        | Expr::Exists { .. }
343        | Expr::InSubquery { .. } => {}
344    }
345}
346
347fn validate_generation_column_reference(columns: &[ColumnDef], name: &str) -> Result<(), SQLError> {
348    let Some(column) = columns.iter().find(|column| column.name == name) else {
349        return Err(SQLError::UnknownColumn(name.to_string()));
350    };
351    if column.generated.is_some() {
352        return Err(SQLError::TypeMismatch(format!(
353            "cannot use generated column `{name}` in column generation expression"
354        )));
355    }
356    Ok(())
357}