Skip to main content

sql_fun_sqlast/sem/
create_index.rs

1use sql_fun_core::IVec;
2
3use crate::{
4    sem::{
5        AnalysisError, AstAndContextPair, ColumnName, FromClause, ParseContext, SemAst,
6        SemScalarExpr, TableName, WithClause, analyze_scaler_expr,
7    },
8    syn::{ListOpt, Opt, ScanToken, SortByDir},
9};
10
11/// Index item sort order
12#[derive(Debug, Clone, Copy)]
13pub enum SortOrder {
14    /// assending
15    Asc,
16    /// dessending
17    Desc,
18}
19
20/// `IndexElement` has column or expression
21#[derive(Debug, Clone)]
22pub enum IndexElement {
23    /// column valiant
24    ColumnName(ColumnName),
25    /// expression variant
26    Expr(Box<SemScalarExpr>),
27}
28
29impl From<&ColumnName> for IndexElement {
30    fn from(value: &ColumnName) -> Self {
31        Self::ColumnName(value.clone())
32    }
33}
34
35impl From<&SemScalarExpr> for IndexElement {
36    fn from(value: &SemScalarExpr) -> Self {
37        Self::Expr(Box::new(value.clone()))
38    }
39}
40
41/// combine column and sort order , `col_name`
42#[derive(Debug, Clone)]
43pub struct IndexElementWithOrder {
44    #[cfg_attr(not(test), expect(dead_code))]
45    element: IndexElement,
46    #[cfg_attr(not(test), expect(dead_code))]
47    col_name: Option<String>,
48    #[cfg_attr(not(test), expect(dead_code))]
49    order: SortOrder,
50}
51
52impl IndexElementWithOrder {
53    /// create instance
54    #[must_use]
55    pub fn new(element: &IndexElement, order: SortOrder, col_name: &Option<String>) -> Self {
56        Self {
57            element: element.clone(),
58            col_name: col_name.clone(),
59            order,
60        }
61    }
62}
63
64#[derive(Debug, Default, Clone)]
65pub struct IndexElementCollection(Vec<IndexElementWithOrder>);
66
67impl IndexElementCollection {
68    pub fn from_column_names(column_names: &IVec<ColumnName>) -> Self {
69        Self(
70            column_names
71                .iter()
72                .map(|v| {
73                    let col = IndexElement::from(v);
74                    IndexElementWithOrder::new(&col, SortOrder::Asc, &None)
75                })
76                .collect(),
77        )
78    }
79
80    fn push(&mut self, item: IndexElementWithOrder) {
81        self.0.push(item);
82    }
83}
84
85#[cfg(test)]
86mod index_element_collection_tests {
87    use super::*;
88
89    #[test]
90    fn from_column_names_uses_asc_order_and_no_alias() {
91        let columns: IVec<ColumnName> =
92            vec![ColumnName::from("id"), ColumnName::from("name")].into();
93
94        let collection = IndexElementCollection::from_column_names(&columns);
95
96        assert_eq!(collection.0.len(), 2);
97        match &collection.0[0] {
98            IndexElementWithOrder {
99                element: IndexElement::ColumnName(name),
100                col_name,
101                order,
102            } => {
103                assert_eq!(name, &ColumnName::from("id"));
104                assert!(col_name.is_none());
105                assert!(matches!(order, SortOrder::Asc));
106            }
107            _ => panic!("expected column name index element"),
108        }
109        match &collection.0[1] {
110            IndexElementWithOrder {
111                element: IndexElement::ColumnName(name),
112                col_name,
113                order,
114            } => {
115                assert_eq!(name, &ColumnName::from("name"));
116                assert!(col_name.is_none());
117                assert!(matches!(order, SortOrder::Asc));
118            }
119            _ => panic!("expected column name index element"),
120        }
121    }
122}
123
124/// analyzed [`crate::syn::IndexStmt`]
125#[derive(Debug, Clone)]
126pub struct CreateIndex {
127    name: String,
128    #[cfg_attr(not(test), expect(dead_code))]
129    table: TableName,
130    #[cfg_attr(not(test), expect(dead_code))]
131    key_columns: IndexElementCollection,
132    #[cfg_attr(not(test), expect(dead_code))]
133    include_columns: IndexElementCollection,
134    #[cfg_attr(not(test), expect(dead_code))]
135    unique: bool,
136}
137
138impl CreateIndex {
139    /// create a instance
140    #[must_use]
141    pub fn new(name: &str, table: &TableName, key_columns: &IndexElementCollection) -> Self {
142        Self {
143            name: name.to_string(),
144            table: table.clone(),
145            key_columns: key_columns.clone(),
146            include_columns: IndexElementCollection::default(),
147            unique: true,
148        }
149    }
150
151    /// create with `include_columns`
152    #[must_use]
153    pub fn new_with_include_column(
154        name: &str,
155        table: &TableName,
156        key_columns: &IndexElementCollection,
157        include_columns: &IndexElementCollection,
158    ) -> Self {
159        Self {
160            name: name.to_string(),
161            table: table.clone(),
162            key_columns: key_columns.clone(),
163            include_columns: include_columns.clone(),
164            unique: true,
165        }
166    }
167
168    /// get index name
169    #[must_use]
170    pub fn name(&self) -> &str {
171        &self.name
172    }
173}
174
175#[cfg(test)]
176mod create_index_tests {
177    use super::*;
178    use crate::sem::FullName;
179
180    #[test]
181    fn new_sets_defaults_and_clones_fields() {
182        let table = TableName::from(FullName::with_schema("public", "items"));
183        let key_columns =
184            IndexElementCollection::from_column_names(&vec![ColumnName::from("id")].into());
185
186        let index = CreateIndex::new("idx_items_id", &table, &key_columns);
187
188        assert_eq!(index.name, "idx_items_id");
189        assert_eq!(index.table, table);
190        assert_eq!(index.key_columns.0.len(), 1);
191        assert_eq!(index.include_columns.0.len(), 0);
192        assert!(index.unique);
193    }
194
195    #[test]
196    fn new_with_include_column_preserves_include_columns() {
197        let table = TableName::from(FullName::with_schema("public", "items"));
198        let key_columns =
199            IndexElementCollection::from_column_names(&vec![ColumnName::from("id")].into());
200        let include_columns =
201            IndexElementCollection::from_column_names(&vec![ColumnName::from("metadata")].into());
202
203        let index = CreateIndex::new_with_include_column(
204            "idx_items_id",
205            &table,
206            &key_columns,
207            &include_columns,
208        );
209
210        assert_eq!(index.key_columns.0.len(), 1);
211        assert_eq!(index.include_columns.0.len(), 1);
212        assert!(index.unique);
213    }
214}
215
216impl CreateIndex {
217    fn sort_order(key: &crate::syn::IndexElem) -> Result<SortOrder, AnalysisError> {
218        let order = if let Some(order) = key.get_ordering().as_inner() {
219            match order {
220                SortByDir::SortbyUsing => {
221                    AnalysisError::raise_unexpected_input("SortbyUsing in IndexElem")?
222                }
223                SortByDir::SortbyDesc => SortOrder::Desc,
224                _ => SortOrder::Asc,
225            }
226        } else {
227            SortOrder::Asc
228        };
229        Ok(order)
230    }
231
232    fn analyze_index_elem<TParseContext>(
233        mut context: TParseContext,
234        keys: Vec<crate::syn::IndexElem>,
235        tokens: &IVec<ScanToken>,
236    ) -> Result<(IndexElementCollection, TParseContext), AnalysisError>
237    where
238        TParseContext: ParseContext,
239    {
240        let mut key_columns = IndexElementCollection::default();
241        for key in keys {
242            let element = if let Some(expr) = key.get_expr().as_inner() {
243                let (scaler_expr, new_context) = analyze_scaler_expr(
244                    context,
245                    &WithClause::default(),
246                    &FromClause::default(),
247                    expr,
248                    tokens,
249                )?;
250                context = new_context;
251                IndexElement::from(&scaler_expr)
252            } else {
253                let name = key.get_name();
254                IndexElement::from(&ColumnName::from(name))
255            };
256
257            let order = CreateIndex::sort_order(&key)?;
258            let col_name = key.get_indexcolname();
259            let col_name = if col_name.is_empty() {
260                None
261            } else {
262                Some(col_name)
263            };
264
265            key_columns.push(IndexElementWithOrder::new(&element, order, &col_name));
266        }
267        Ok((key_columns, context))
268    }
269}
270
271#[cfg(test)]
272mod analyze_index_elem_tests {
273    use super::*;
274    use crate::{
275        sem::ScalarConstExpr,
276        syn::{Node, NodeInner, SortByDir},
277        test_helpers::{SynBuilder, TestParseContext},
278    };
279    use sql_fun_core::IVec;
280
281    #[test]
282    fn analyze_index_elem_collects_column_order_and_alias() {
283        let builder = SynBuilder::new();
284        let keys = vec![
285            builder.index_elem_column("id", SortByDir::SortbyDesc.into(), Some("id_alias")),
286            builder.index_elem_column("name", SortByDir::SortbyDefault.into(), None),
287        ];
288        let tokens: IVec<ScanToken> = IVec::default();
289        let context = TestParseContext::default();
290
291        let (collection, _context) =
292            CreateIndex::analyze_index_elem(context, keys, &tokens).expect("analyze index elem");
293
294        assert_eq!(collection.0.len(), 2);
295
296        let first = &collection.0[0];
297        assert!(matches!(first.order, SortOrder::Desc));
298        assert_eq!(first.col_name.as_deref(), Some("id_alias"));
299        match &first.element {
300            IndexElement::ColumnName(name) => assert_eq!(name, &ColumnName::from("id")),
301            _ => panic!("expected column name"),
302        }
303
304        let second = &collection.0[1];
305        assert!(matches!(second.order, SortOrder::Asc));
306        assert!(second.col_name.is_none());
307        match &second.element {
308            IndexElement::ColumnName(name) => assert_eq!(name, &ColumnName::from("name")),
309            _ => panic!("expected column name"),
310        }
311    }
312
313    #[test]
314    fn analyze_index_elem_collects_expression() {
315        let builder = SynBuilder::new();
316        let expr_node = Node::from(NodeInner::AConst(builder.const_int4(42)));
317        let keys = vec![builder.index_elem_expr(expr_node, SortByDir::SortbyDefault.into(), None)];
318        let tokens: IVec<ScanToken> = IVec::default();
319        let context = TestParseContext::default();
320
321        let (collection, _context) =
322            CreateIndex::analyze_index_elem(context, keys, &tokens).expect("analyze index elem");
323
324        assert_eq!(collection.0.len(), 1);
325        let item = &collection.0[0];
326        assert!(matches!(item.order, SortOrder::Asc));
327        match &item.element {
328            IndexElement::Expr(expr) => match expr.as_ref() {
329                SemScalarExpr::Const(ScalarConstExpr::Integer(value)) => {
330                    assert_eq!(*value, 42);
331                }
332                _ => panic!("expected integer const expr"),
333            },
334            _ => panic!("expected expression element"),
335        }
336    }
337}
338
339/// analyze [`crate::syn::IndexStmt`]
340pub fn analyze_index_stmt<TParseContext>(
341    mut context: TParseContext,
342    _parent_schema: &Option<String>,
343    syn: crate::syn::IndexStmt,
344    tokens: &IVec<ScanToken>,
345) -> Result<AstAndContextPair<TParseContext>, AnalysisError>
346where
347    TParseContext: ParseContext,
348{
349    let Some(table) = syn.get_relation().as_inner() else {
350        AnalysisError::raise_unexpected_none("indexstmt.relation")?
351    };
352    let Some(keys) = syn.get_index_params().map(|v| v.as_index_elem()) else {
353        AnalysisError::raise_unexpected_none("indexstmt.index_params")?
354    };
355
356    let name = syn.get_idxname();
357    let table = TableName::try_from(table)?;
358    let unique = syn.get_unique();
359
360    let (key_columns, new_context) = CreateIndex::analyze_index_elem(context, keys, tokens)?;
361    context = new_context;
362
363    let include_columns =
364        if let Some(includes) = syn.get_index_including_params().map(|v| v.as_index_elem()) {
365            let (cols, new_context) = CreateIndex::analyze_index_elem(context, includes, tokens)?;
366            context = new_context;
367            cols
368        } else {
369            IndexElementCollection::default()
370        };
371
372    let create_index = CreateIndex {
373        name,
374        table,
375        key_columns,
376        include_columns,
377        unique,
378    };
379
380    context = context.apply_create_index(&create_index)?;
381
382    Ok(AstAndContextPair::new(
383        SemAst::CreateIndex(create_index),
384        context,
385    ))
386}
387
388#[cfg(test)]
389mod analyze_index_stmt_tests {
390    use super::*;
391    use crate::{
392        sem::FullName,
393        syn::{AliasOpt, RangeVarOpt, SortByDir},
394        test_helpers::{SynBuilder, TestParseContext, test_context},
395    };
396    use rstest::rstest;
397    use sql_fun_core::IVec;
398
399    #[rstest]
400    fn analyze_index_stmt_collects_keys_including_order_and_includes(
401        test_context: TestParseContext,
402    ) {
403        let builder = SynBuilder::new();
404        let relation = RangeVarOpt::from(
405            builder.range_var(&FullName::with_schema("public", "items"), &AliasOpt::none()),
406        );
407        let index_params = builder.index_elem_list(vec![
408            builder.index_elem_column("id", SortByDir::SortbyDesc.into(), None),
409            builder.index_elem_column("name", SortByDir::SortbyDefault.into(), None),
410        ]);
411        let include_params = builder.index_elem_list(vec![builder.index_elem_column(
412            "metadata",
413            SortByDir::SortbyDefault.into(),
414            None,
415        )]);
416        let index_stmt =
417            builder.index_stmt("idx_items", relation, index_params, include_params, false);
418        let tokens: IVec<ScanToken> = IVec::default();
419
420        let AstAndContextPair(ast, context) =
421            analyze_index_stmt(test_context, &None, index_stmt, &tokens).expect("analyze index");
422
423        let SemAst::CreateIndex(create_index) = ast else {
424            panic!("expected create index ast");
425        };
426
427        assert_eq!(create_index.name, "idx_items");
428        assert!(matches!(
429            create_index.table.full_name().schema(),
430            Some(schema) if schema.as_str() == "public"
431        ));
432        assert!(!create_index.unique);
433        assert_eq!(create_index.key_columns.0.len(), 2);
434        assert_eq!(create_index.include_columns.0.len(), 1);
435
436        let first_key = &create_index.key_columns.0[0];
437        assert!(matches!(first_key.order, SortOrder::Desc));
438        match &first_key.element {
439            IndexElement::ColumnName(name) => assert_eq!(name, &ColumnName::from("id")),
440            _ => panic!("expected column name"),
441        }
442
443        let second_key = &create_index.key_columns.0[1];
444        assert!(matches!(second_key.order, SortOrder::Asc));
445        match &second_key.element {
446            IndexElement::ColumnName(name) => assert_eq!(name, &ColumnName::from("name")),
447            _ => panic!("expected column name"),
448        }
449
450        let include_key = &create_index.include_columns.0[0];
451        match &include_key.element {
452            IndexElement::ColumnName(name) => assert_eq!(name, &ColumnName::from("metadata")),
453            _ => panic!("expected column name"),
454        }
455
456        assert_eq!(context.create_indexes_len(), 1);
457        let applied = context.last_create_index().expect("applied create index");
458        assert_eq!(applied.name, create_index.name);
459    }
460}