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#[derive(Debug, Clone, Copy)]
13pub enum SortOrder {
14 Asc,
16 Desc,
18}
19
20#[derive(Debug, Clone)]
22pub enum IndexElement {
23 ColumnName(ColumnName),
25 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#[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 #[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#[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 #[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 #[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 #[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
339pub 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}