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(_) | Expr::CompositeRow { .. } => 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                implicit: true,
112                expr: expression.clone(),
113                ty: "text".into(),
114            };
115            Ok(Some(ColumnType::Text))
116        })
117        .collect()
118}
119
120fn finish_index_keys(
121    context: &SchemaBindingContext<'_, '_>,
122    statement: &mut CreateIndex,
123    expressions: &[Option<ColumnType>],
124) -> Result<Vec<ColumnType>, SQLError> {
125    let definitions = context
126        .catalog
127        .schema_expression_columns(&statement.table)?
128        .ok_or_else(|| SQLError::UnknownTable(statement.table.clone()))?;
129    let mut types = Vec::with_capacity(statement.columns.len());
130    for (key, analyzed) in statement.columns.iter_mut().zip(expressions) {
131        match key {
132            IndexKey::Column(name) => {
133                let Some(column) = definitions.iter().find(|column| column.name == *name) else {
134                    if definitions.is_empty() {
135                        types.push(ColumnType::Text);
136                        continue;
137                    }
138                    return Err(SQLError::UnknownColumn(name.clone()));
139                };
140                reject_virtual_generated_column(column)?;
141                types.push(column.ty.clone());
142            }
143            IndexKey::Expression(expression) => {
144                for column in &definitions {
145                    if column
146                        .generated
147                        .as_ref()
148                        .is_some_and(|generated| generated.kind == GeneratedColumnKind::Virtual)
149                        && crate::schema::dependencies::schema_expr_references_column(
150                            expression,
151                            &column.name,
152                        )
153                    {
154                        return Err(SQLError::Unsupported(format!(
155                            "index expressions cannot use virtual generated column `{}`",
156                            column.name
157                        )));
158                    }
159                }
160                super::validate_index_expression_immutability(
161                    context.catalog,
162                    &statement.table,
163                    expression,
164                    false,
165                )?;
166                let ty = analyzed
167                    .as_ref()
168                    .expect("analyzed index expression")
169                    .clone();
170                let column = match expression.as_ref() {
171                    crate::ast::Expr::Column(name) => Some(name.clone()),
172                    crate::ast::Expr::Cast { expr, .. } => {
173                        if let crate::ast::Expr::Column(name) = expr.as_ref() {
174                            definitions
175                                .iter()
176                                .any(|column| column.name == *name && column.ty == ty)
177                                .then(|| name.clone())
178                        } else {
179                            None
180                        }
181                    }
182                    _ => None,
183                };
184                if let Some(column) = column {
185                    *key = IndexKey::Column(column);
186                }
187                types.push(ty);
188            }
189        }
190    }
191    // As in `PostgreSQL`, an included column may repeat and may also be a key column.
192    for name in &statement.included_columns {
193        match definitions.iter().find(|column| column.name == *name) {
194            Some(column) => reject_virtual_generated_column(column)?,
195            None if definitions.is_empty() => {}
196            None => return Err(SQLError::UnknownColumn(name.clone())),
197        }
198    }
199    reject_unsupported_included_columns(statement)?;
200    Ok(types)
201}
202
203fn reject_virtual_generated_column(column: &crate::ast::ColumnDef) -> Result<(), SQLError> {
204    if column
205        .generated
206        .as_ref()
207        .is_some_and(|generated| generated.kind == GeneratedColumnKind::Virtual)
208    {
209        return Err(SQLError::Unsupported(
210            "indexes on virtual generated columns are not supported".into(),
211        ));
212    }
213    Ok(())
214}
215
216/// Only a btree index carries columns beside its key.
217fn reject_unsupported_included_columns(statement: &CreateIndex) -> Result<(), SQLError> {
218    let method = statement.access_method.to_ascii_lowercase();
219    if statement.included_columns.is_empty() || matches!(method.as_str(), "" | "btree") {
220        return Ok(());
221    }
222    Err(SQLError::Unsupported(format!(
223        "access method \"{method}\" does not support included columns"
224    )))
225}
226
227/// Bind keys and predicates and retain the public attribute names assigned before expression simplification.
228pub fn prepare_index_definition(
229    catalog: &dyn crate::schema::SchemaExpressionCatalog,
230    bindings: &dyn crate::semantics::conflict::InferenceBindingScope,
231    c: &mut CreateIndex,
232) -> Result<crate::catalog::index::IndexDefinition, SQLError> {
233    let attribute_keys = c
234        .columns
235        .iter()
236        .cloned()
237        .chain(
238            c.included_columns
239                .iter()
240                .cloned()
241                .map(crate::ast::IndexKey::Column),
242        )
243        .collect::<Vec<_>>();
244    let key_names = key_names(&attribute_keys);
245    let binding = bindings.binding_scope()?;
246    let scope = binding.context();
247    let context = SchemaBindingContext {
248        catalog,
249        binding: &scope,
250    };
251    let expressions = bind_index_keys(&context, c)?;
252    if let Some(predicate) = c.predicate.as_deref_mut() {
253        super::bind_index_predicate(catalog, &scope, &c.table, predicate)?;
254    }
255    super::options::index_access_method(c)?;
256    if c.unique {
257        super::unique::validate_unique_index_method(c)?;
258    }
259    reject_unsupported_included_columns(c)?;
260    if let Some(predicate) = c.predicate.as_deref_mut() {
261        super::validate_index_expression_immutability(catalog, &c.table, predicate, true)?;
262    }
263    super::options::validate_index_options(c)?;
264    let key_types = finish_index_keys(&context, c, &expressions)?;
265    Ok(crate::catalog::index::IndexDefinition {
266        catalog: None,
267        relationships: crate::catalog::index::IndexRelationships::default(),
268        key_names,
269        key_types,
270        included_columns: c.included_columns.clone(),
271        column_order: c.column_order.clone(),
272        predicate: c.predicate.clone(),
273        unique: c.unique,
274        nulls_not_distinct: c.nulls_not_distinct,
275    })
276}