Skip to main content

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 { name, .. } => Some((
39            crate::parse_regobject_name(name)
40                .and_then(|mut names| names.pop())
41                .unwrap_or_else(|| name.clone()),
42            true,
43        )),
44        Expr::Cast { expr, ty } => {
45            let inner = expression_name(expr);
46            if inner.as_ref().is_some_and(|(_, strong)| *strong) {
47                inner
48            } else {
49                Some((
50                    crate::parse_regtype_name(ty)
51                        .ok()
52                        .flatten()
53                        .and_then(|mut name| name.names.pop())
54                        .unwrap_or_else(|| ty.clone()),
55                    false,
56                ))
57            }
58        }
59        Expr::Case { else_branch, .. } => {
60            let inner = else_branch.as_deref().and_then(expression_name);
61            Some(
62                inner
63                    .filter(|(_, strong)| *strong)
64                    .unwrap_or_else(|| ("case".into(), false)),
65            )
66        }
67        Expr::Array(_) => Some(("array".into(), true)),
68        Expr::Row(_) => Some(("row".into(), true)),
69        _ => None,
70    }
71}
72
73pub fn require_column_key<'a>(key: &'a IndexKey, method: &str) -> Result<&'a str, SQLError> {
74    key.column().ok_or_else(|| {
75        SQLError::Unsupported(format!(
76            "expression keys for access method `{method}` are not implemented"
77        ))
78    })
79}
80
81pub fn prepare_index_keys(
82    context: &SchemaBindingContext<'_, '_>,
83    statement: &mut CreateIndex,
84) -> Result<Vec<ColumnType>, SQLError> {
85    let expressions = bind_index_keys(context, statement)?;
86    finish_index_keys(context, statement, &expressions)
87}
88
89fn bind_index_keys(
90    context: &SchemaBindingContext<'_, '_>,
91    statement: &mut CreateIndex,
92) -> Result<Vec<Option<ColumnType>>, SQLError> {
93    statement
94        .columns
95        .iter_mut()
96        .map(|key| {
97            let IndexKey::Expression(expression) = key else {
98                return Ok(None);
99            };
100            let ty = super::bind_index_expression(
101                context.catalog,
102                context.binding,
103                &statement.table,
104                expression,
105                false,
106            )?;
107            if let Some(ty) = ty {
108                return Ok(Some(ty));
109            }
110            **expression = Expr::Cast {
111                expr: expression.clone(),
112                ty: "text".into(),
113            };
114            Ok(Some(ColumnType::Text))
115        })
116        .collect()
117}
118
119fn finish_index_keys(
120    context: &SchemaBindingContext<'_, '_>,
121    statement: &mut CreateIndex,
122    expressions: &[Option<ColumnType>],
123) -> Result<Vec<ColumnType>, SQLError> {
124    let definitions = context
125        .catalog
126        .schema_expression_columns(&statement.table)?
127        .ok_or_else(|| SQLError::UnknownTable(statement.table.clone()))?;
128    let mut types = Vec::with_capacity(statement.columns.len());
129    for (key, analyzed) in statement.columns.iter_mut().zip(expressions) {
130        match key {
131            IndexKey::Column(name) => {
132                let Some(column) = definitions.iter().find(|column| column.name == *name) else {
133                    if definitions.is_empty() {
134                        types.push(ColumnType::Text);
135                        continue;
136                    }
137                    return Err(SQLError::UnknownColumn(name.clone()));
138                };
139                reject_virtual_generated_column(column)?;
140                types.push(column.ty.clone());
141            }
142            IndexKey::Expression(expression) => {
143                for column in &definitions {
144                    if column
145                        .generated
146                        .as_ref()
147                        .is_some_and(|generated| generated.kind == GeneratedColumnKind::Virtual)
148                        && crate::schema::dependencies::schema_expr_references_column(
149                            expression,
150                            &column.name,
151                        )
152                    {
153                        return Err(SQLError::Unsupported(format!(
154                            "index expressions cannot use virtual generated column `{}`",
155                            column.name
156                        )));
157                    }
158                }
159                super::validate_index_expression_immutability(
160                    context.catalog,
161                    &statement.table,
162                    expression,
163                    false,
164                )?;
165                let ty = analyzed
166                    .as_ref()
167                    .expect("analyzed index expression")
168                    .clone();
169                let column = match expression.as_ref() {
170                    crate::ast::Expr::Column(name) => Some(name.clone()),
171                    crate::ast::Expr::Cast { expr, .. } => {
172                        if let crate::ast::Expr::Column(name) = expr.as_ref() {
173                            definitions
174                                .iter()
175                                .any(|column| column.name == *name && column.ty == ty)
176                                .then(|| name.clone())
177                        } else {
178                            None
179                        }
180                    }
181                    _ => None,
182                };
183                if let Some(column) = column {
184                    *key = IndexKey::Column(column);
185                }
186                types.push(ty);
187            }
188        }
189    }
190    // As in `PostgreSQL`, an included column may repeat and may also be a key column.
191    for name in &statement.included_columns {
192        match definitions.iter().find(|column| column.name == *name) {
193            Some(column) => reject_virtual_generated_column(column)?,
194            None if definitions.is_empty() => {}
195            None => return Err(SQLError::UnknownColumn(name.clone())),
196        }
197    }
198    reject_unsupported_included_columns(statement)?;
199    Ok(types)
200}
201
202fn reject_virtual_generated_column(column: &crate::ast::ColumnDef) -> Result<(), SQLError> {
203    if column
204        .generated
205        .as_ref()
206        .is_some_and(|generated| generated.kind == GeneratedColumnKind::Virtual)
207    {
208        return Err(SQLError::Unsupported(
209            "indexes on virtual generated columns are not supported".into(),
210        ));
211    }
212    Ok(())
213}
214
215/// Only a btree index carries columns beside its key.
216fn reject_unsupported_included_columns(statement: &CreateIndex) -> Result<(), SQLError> {
217    let method = statement.access_method.to_ascii_lowercase();
218    if statement.included_columns.is_empty() || matches!(method.as_str(), "" | "btree") {
219        return Ok(());
220    }
221    Err(SQLError::Unsupported(format!(
222        "access method \"{method}\" does not support included columns"
223    )))
224}
225
226/// Bind keys and predicates and retain the public attribute names assigned before expression simplification.
227pub fn prepare_index_definition(
228    catalog: &dyn crate::schema::SchemaExpressionCatalog,
229    bindings: &dyn crate::semantics::conflict::InferenceBindingScope,
230    c: &mut CreateIndex,
231) -> Result<crate::catalog::index::IndexDefinition, SQLError> {
232    let attribute_keys = c
233        .columns
234        .iter()
235        .cloned()
236        .chain(
237            c.included_columns
238                .iter()
239                .cloned()
240                .map(crate::ast::IndexKey::Column),
241        )
242        .collect::<Vec<_>>();
243    let key_names = key_names(&attribute_keys);
244    let binding = bindings.binding_scope()?;
245    let scope = binding.context();
246    let context = SchemaBindingContext {
247        catalog,
248        binding: &scope,
249    };
250    let expressions = bind_index_keys(&context, c)?;
251    if let Some(predicate) = c.predicate.as_deref_mut() {
252        super::bind_index_predicate(catalog, &scope, &c.table, predicate)?;
253    }
254    super::options::index_access_method(c)?;
255    if c.unique {
256        super::unique::validate_unique_index_method(c)?;
257    }
258    reject_unsupported_included_columns(c)?;
259    if let Some(predicate) = c.predicate.as_deref_mut() {
260        super::validate_index_expression_immutability(catalog, &c.table, predicate, true)?;
261    }
262    super::options::validate_index_options(c)?;
263    let key_types = finish_index_keys(&context, c, &expressions)?;
264    Ok(crate::catalog::index::IndexDefinition {
265        catalog: None,
266        relationships: crate::catalog::index::IndexRelationships::default(),
267        key_names,
268        key_types,
269        included_columns: c.included_columns.clone(),
270        column_order: c.column_order.clone(),
271        predicate: c.predicate.clone(),
272        unique: c.unique,
273        nulls_not_distinct: c.nulls_not_distinct,
274    })
275}