uqa_sql/schema/indexes/
keys.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Typed index-key preparation and column names assigned at index creation.
8
9use crate::ast::{Expr, GeneratedColumnKind, IndexKey};
10use crate::schema::SchemaBindingContext;
11use crate::{ast::CreateIndex, ColumnType, SQLError};
12
13pub fn key_names(keys: &[IndexKey]) -> Vec<String> {
14    let mut names = Vec::with_capacity(keys.len());
15    for key in keys {
16        let label = match key {
17            IndexKey::Column(column) => column.clone(),
18            IndexKey::Expression(expression) => {
19                expression_name(expression).map_or_else(|| "expr".into(), |(name, _)| name)
20            }
21        };
22        let mut name = label.clone();
23        let mut suffix = 1_u64;
24        while names.contains(&name) {
25            name = format!("{label}{suffix}");
26            suffix += 1;
27        }
28        names.push(name);
29    }
30    names
31}
32
33fn expression_name(expression: &Expr) -> Option<(String, bool)> {
34    match expression {
35        Expr::Column(name) | Expr::QualifiedColumn { column: name, .. } => {
36            Some((name.clone(), true))
37        }
38        Expr::Func {
39            binding: Some(binding),
40            args,
41            ..
42        } if binding.dispatch == Some(crate::ast::FunctionDispatch::FieldSelect) => {
43            match args.as_slice() {
44                [_, Expr::Literal(uqa_core::Value::Str(name))] => Some((name.clone(), true)),
45                _ => None,
46            }
47        }
48        Expr::Func { name, .. } => Some((
49            crate::parse_regobject_name(name)
50                .and_then(|mut names| names.pop())
51                .unwrap_or_else(|| name.clone()),
52            true,
53        )),
54        Expr::Cast { expr, ty, .. } => {
55            let inner = expression_name(expr);
56            if inner.as_ref().is_some_and(|(_, strong)| *strong) {
57                inner
58            } else {
59                Some((
60                    crate::parse_regtype_name(ty)
61                        .ok()
62                        .flatten()
63                        .and_then(|mut name| name.names.pop())
64                        .unwrap_or_else(|| ty.clone()),
65                    false,
66                ))
67            }
68        }
69        Expr::Case { else_branch, .. } => {
70            let inner = else_branch.as_deref().and_then(expression_name);
71            Some(
72                inner
73                    .filter(|(_, strong)| *strong)
74                    .unwrap_or_else(|| ("case".into(), false)),
75            )
76        }
77        Expr::Array(_) => Some(("array".into(), true)),
78        Expr::Row(_) | Expr::CompositeRow { .. } => Some(("row".into(), true)),
79        _ => None,
80    }
81}
82
83pub fn require_column_key<'a>(key: &'a IndexKey, method: &str) -> Result<&'a str, SQLError> {
84    key.column().ok_or_else(|| {
85        SQLError::Unsupported(format!(
86            "expression keys for access method `{method}` are not implemented"
87        ))
88    })
89}
90
91pub fn prepare_index_keys(
92    context: &SchemaBindingContext<'_, '_>,
93    statement: &mut CreateIndex,
94) -> Result<Vec<ColumnType>, SQLError> {
95    let expressions = bind_index_keys(context, statement)?;
96    finish_index_keys(context, statement, &expressions)
97}
98
99fn bind_index_keys(
100    context: &SchemaBindingContext<'_, '_>,
101    statement: &mut CreateIndex,
102) -> Result<Vec<Option<ColumnType>>, SQLError> {
103    statement
104        .columns
105        .iter_mut()
106        .map(|key| {
107            let IndexKey::Expression(expression) = key else {
108                return Ok(None);
109            };
110            let ty = super::bind_index_expression(
111                context.catalog,
112                context.binding,
113                &statement.table,
114                expression,
115                false,
116            )?;
117            if let Some(ty) = ty {
118                return Ok(Some(ty));
119            }
120            **expression = Expr::Cast {
121                implicit: true,
122                expr: expression.clone(),
123                ty: "text".into(),
124            };
125            Ok(Some(ColumnType::Text))
126        })
127        .collect()
128}
129
130fn finish_index_keys(
131    context: &SchemaBindingContext<'_, '_>,
132    statement: &mut CreateIndex,
133    expressions: &[Option<ColumnType>],
134) -> Result<Vec<ColumnType>, SQLError> {
135    let definitions = context
136        .catalog
137        .schema_expression_columns(&statement.table)?
138        .ok_or_else(|| SQLError::UnknownTable(statement.table.clone()))?;
139    let mut types = Vec::with_capacity(statement.columns.len());
140    for (key, analyzed) in statement.columns.iter_mut().zip(expressions) {
141        match key {
142            IndexKey::Column(name) => {
143                let Some(column) = definitions.iter().find(|column| column.name == *name) else {
144                    if definitions.is_empty() {
145                        types.push(ColumnType::Text);
146                        continue;
147                    }
148                    return Err(SQLError::UnknownColumn(name.clone()));
149                };
150                reject_virtual_generated_column(column)?;
151                types.push(column.ty.clone());
152            }
153            IndexKey::Expression(expression) => {
154                for column in &definitions {
155                    if column
156                        .generated
157                        .as_ref()
158                        .is_some_and(|generated| generated.kind == GeneratedColumnKind::Virtual)
159                        && crate::schema::dependencies::schema_expr_references_column(
160                            expression,
161                            &column.name,
162                        )
163                    {
164                        return Err(SQLError::Unsupported(format!(
165                            "index expressions cannot use virtual generated column `{}`",
166                            column.name
167                        )));
168                    }
169                }
170                super::validate_index_expression_immutability(
171                    context.catalog,
172                    &statement.table,
173                    expression,
174                    false,
175                )?;
176                let ty = analyzed
177                    .as_ref()
178                    .expect("analyzed index expression")
179                    .clone();
180                let column = match expression.as_ref() {
181                    crate::ast::Expr::Column(name) => Some(name.clone()),
182                    crate::ast::Expr::Cast { expr, .. } => {
183                        if let crate::ast::Expr::Column(name) = expr.as_ref() {
184                            definitions
185                                .iter()
186                                .any(|column| column.name == *name && column.ty == ty)
187                                .then(|| name.clone())
188                        } else {
189                            None
190                        }
191                    }
192                    _ => None,
193                };
194                if let Some(column) = column {
195                    *key = IndexKey::Column(column);
196                }
197                types.push(ty);
198            }
199        }
200    }
201    // As in `PostgreSQL`, an included column may repeat and may also be a key column.
202    for name in &statement.included_columns {
203        match definitions.iter().find(|column| column.name == *name) {
204            Some(column) => reject_virtual_generated_column(column)?,
205            None if definitions.is_empty() => {}
206            None => return Err(SQLError::UnknownColumn(name.clone())),
207        }
208    }
209    reject_unsupported_included_columns(statement)?;
210    Ok(types)
211}
212
213fn reject_virtual_generated_column(column: &crate::ast::ColumnDef) -> Result<(), SQLError> {
214    if column
215        .generated
216        .as_ref()
217        .is_some_and(|generated| generated.kind == GeneratedColumnKind::Virtual)
218    {
219        return Err(SQLError::Unsupported(
220            "indexes on virtual generated columns are not supported".into(),
221        ));
222    }
223    Ok(())
224}
225
226/// Only a btree index carries columns beside its key.
227fn reject_unsupported_included_columns(statement: &CreateIndex) -> Result<(), SQLError> {
228    let method = statement.access_method.to_ascii_lowercase();
229    if statement.included_columns.is_empty() || matches!(method.as_str(), "" | "btree") {
230        return Ok(());
231    }
232    Err(SQLError::Unsupported(format!(
233        "access method \"{method}\" does not support included columns"
234    )))
235}
236
237/// Bind keys and predicates and retain the public attribute names assigned before expression simplification.
238pub fn prepare_index_definition(
239    catalog: &dyn crate::schema::SchemaExpressionCatalog,
240    bindings: &dyn crate::semantics::conflict::InferenceBindingScope,
241    c: &mut CreateIndex,
242) -> Result<crate::catalog::index::IndexDefinition, SQLError> {
243    let attribute_keys = c
244        .columns
245        .iter()
246        .cloned()
247        .chain(
248            c.included_columns
249                .iter()
250                .cloned()
251                .map(crate::ast::IndexKey::Column),
252        )
253        .collect::<Vec<_>>();
254    let key_names = key_names(&attribute_keys);
255    let binding = bindings.binding_scope()?;
256    let scope = binding.context();
257    let context = SchemaBindingContext {
258        catalog,
259        binding: &scope,
260    };
261    let expressions = bind_index_keys(&context, c)?;
262    if let Some(predicate) = c.predicate.as_deref_mut() {
263        super::bind_index_predicate(catalog, &scope, &c.table, predicate)?;
264    }
265    super::options::index_access_method(c)?;
266    if c.unique {
267        super::unique::validate_unique_index_method(c)?;
268    }
269    reject_unsupported_included_columns(c)?;
270    if let Some(predicate) = c.predicate.as_deref_mut() {
271        super::validate_index_expression_immutability(catalog, &c.table, predicate, true)?;
272    }
273    super::options::validate_index_options(c)?;
274    let key_types = finish_index_keys(&context, c, &expressions)?;
275    Ok(crate::catalog::index::IndexDefinition {
276        catalog: None,
277        relationships: crate::catalog::index::IndexRelationships::default(),
278        key_names,
279        key_types,
280        included_columns: c.included_columns.clone(),
281        column_order: c.column_order.clone(),
282        predicate: c.predicate.clone(),
283        unique: c.unique,
284        nulls_not_distinct: c.nulls_not_distinct,
285    })
286}
287
288#[cfg(test)]
289mod tests {
290    use super::*;
291
292    #[test]
293    fn field_selection_names_use_the_selected_attribute_through_casts() {
294        let crate::ast::Statement::CreateIndex(statement) =
295            crate::compile("CREATE INDEX ix ON t (((ROW(id, 2)).f1), (((ROW(id, 2)).f1)::bigint))")
296                .unwrap()
297                .remove(0)
298        else {
299            panic!("index")
300        };
301        assert_eq!(key_names(&statement.columns), ["f1", "f11"]);
302    }
303}
uqa_sql/schema/indexes/keys.rs

uqa_sql/schema/indexes/
keys.rs
