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                if column
140                    .generated
141                    .as_ref()
142                    .is_some_and(|generated| generated.kind == GeneratedColumnKind::Virtual)
143                {
144                    return Err(SQLError::Unsupported(format!(
145                        "indexes on virtual generated column `{name}` are not supported"
146                    )));
147                }
148                types.push(column.ty.clone());
149            }
150            IndexKey::Expression(expression) => {
151                for column in &definitions {
152                    if column
153                        .generated
154                        .as_ref()
155                        .is_some_and(|generated| generated.kind == GeneratedColumnKind::Virtual)
156                        && crate::schema::dependencies::schema_expr_references_column(
157                            expression,
158                            &column.name,
159                        )
160                    {
161                        return Err(SQLError::Unsupported(format!(
162                            "index expressions cannot use virtual generated column `{}`",
163                            column.name
164                        )));
165                    }
166                }
167                super::validate_index_expression_immutability(
168                    context.catalog,
169                    &statement.table,
170                    expression,
171                    false,
172                )?;
173                let ty = analyzed
174                    .as_ref()
175                    .expect("analyzed index expression")
176                    .clone();
177                let column = match expression.as_ref() {
178                    crate::ast::Expr::Column(name) => Some(name.clone()),
179                    crate::ast::Expr::Cast { expr, .. } => {
180                        if let crate::ast::Expr::Column(name) = expr.as_ref() {
181                            definitions
182                                .iter()
183                                .any(|column| column.name == *name && column.ty == ty)
184                                .then(|| name.clone())
185                        } else {
186                            None
187                        }
188                    }
189                    _ => None,
190                };
191                if let Some(column) = column {
192                    *key = IndexKey::Column(column);
193                }
194                types.push(ty);
195            }
196        }
197    }
198    let mut included = std::collections::BTreeSet::new();
199    for name in &statement.included_columns {
200        if !definitions.is_empty() && !definitions.iter().any(|column| column.name == *name) {
201            return Err(SQLError::UnknownColumn(name.clone()));
202        }
203        if !included.insert(name)
204            || statement
205                .columns
206                .iter()
207                .any(|key| key.column() == Some(name.as_str()))
208        {
209            return Err(SQLError::Routine {
210                sqlstate: "42701".into(),
211                message: format!("column \"{name}\" included more than once"),
212            });
213        }
214    }
215    if !included.is_empty() && statement.access_method == "gin" {
216        return Err(SQLError::Unsupported(
217            "access method \"gin\" does not support included columns".into(),
218        ));
219    }
220    Ok(types)
221}
222
223/// Bind keys and predicates and retain the public attribute names assigned before expression simplification.
224pub fn prepare_index_definition(
225    catalog: &dyn crate::schema::SchemaExpressionCatalog,
226    bindings: &dyn crate::semantics::conflict::InferenceBindingScope,
227    c: &mut CreateIndex,
228) -> Result<crate::catalog::index::IndexDefinition, SQLError> {
229    let attribute_keys = c
230        .columns
231        .iter()
232        .cloned()
233        .chain(
234            c.included_columns
235                .iter()
236                .cloned()
237                .map(crate::ast::IndexKey::Column),
238        )
239        .collect::<Vec<_>>();
240    let key_names = key_names(&attribute_keys);
241    let binding = bindings.binding_scope()?;
242    let scope = binding.context();
243    let context = SchemaBindingContext {
244        catalog,
245        binding: &scope,
246    };
247    let expressions = bind_index_keys(&context, c)?;
248    if let Some(predicate) = c.predicate.as_deref_mut() {
249        super::bind_index_predicate(catalog, &scope, &c.table, predicate)?;
250    }
251    super::options::index_access_method(c)?;
252    if c.unique {
253        super::unique::validate_unique_index_method(c)?;
254    }
255    if !c.included_columns.is_empty() && c.access_method.eq_ignore_ascii_case("gin") {
256        return Err(SQLError::Unsupported(
257            "access method \"gin\" does not support included columns".into(),
258        ));
259    }
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}