Skip to main content

squawk_ide/
code_actions.rs

1use itertools::Itertools;
2use rowan::{TextRange, TextSize};
3use squawk_linter::Edit;
4use squawk_syntax::{
5    SyntaxKind, SyntaxToken,
6    ast::{self, AstNode},
7};
8use std::iter;
9
10use crate::{
11    binder,
12    column_name::ColumnName,
13    offsets::token_from_offset,
14    quote::{quote_column_alias, unquote_ident},
15    symbols::Name,
16};
17
18#[derive(Debug, Clone)]
19pub enum ActionKind {
20    QuickFix,
21    RefactorRewrite,
22}
23
24#[derive(Debug, Clone)]
25pub struct CodeAction {
26    pub title: String,
27    pub edits: Vec<Edit>,
28    pub kind: ActionKind,
29}
30
31pub fn code_actions(file: ast::SourceFile, offset: TextSize) -> Option<Vec<CodeAction>> {
32    let mut actions = vec![];
33    rewrite_as_regular_string(&mut actions, &file, offset);
34    rewrite_as_dollar_quoted_string(&mut actions, &file, offset);
35    remove_else_clause(&mut actions, &file, offset);
36    rewrite_table_as_select(&mut actions, &file, offset);
37    rewrite_select_as_table(&mut actions, &file, offset);
38    rewrite_values_as_select(&mut actions, &file, offset);
39    rewrite_select_as_values(&mut actions, &file, offset);
40    add_schema(&mut actions, &file, offset);
41    quote_identifier(&mut actions, &file, offset);
42    unquote_identifier(&mut actions, &file, offset);
43    add_explicit_alias(&mut actions, &file, offset);
44    remove_redundant_alias(&mut actions, &file, offset);
45    rewrite_cast_to_double_colon(&mut actions, &file, offset);
46    rewrite_double_colon_to_cast(&mut actions, &file, offset);
47    Some(actions)
48}
49
50fn rewrite_as_regular_string(
51    actions: &mut Vec<CodeAction>,
52    file: &ast::SourceFile,
53    offset: TextSize,
54) -> Option<()> {
55    let dollar_string = file
56        .syntax()
57        .token_at_offset(offset)
58        .find(|token| token.kind() == SyntaxKind::DOLLAR_QUOTED_STRING)?;
59
60    let replacement = dollar_quoted_to_string(dollar_string.text())?;
61    actions.push(CodeAction {
62        title: "Rewrite as regular string".to_owned(),
63        edits: vec![Edit::replace(dollar_string.text_range(), replacement)],
64        kind: ActionKind::RefactorRewrite,
65    });
66
67    Some(())
68}
69
70fn rewrite_as_dollar_quoted_string(
71    actions: &mut Vec<CodeAction>,
72    file: &ast::SourceFile,
73    offset: TextSize,
74) -> Option<()> {
75    let string = file
76        .syntax()
77        .token_at_offset(offset)
78        .find(|token| token.kind() == SyntaxKind::STRING)?;
79
80    let replacement = string_to_dollar_quoted(string.text())?;
81    actions.push(CodeAction {
82        title: "Rewrite as dollar-quoted string".to_owned(),
83        edits: vec![Edit::replace(string.text_range(), replacement)],
84        kind: ActionKind::RefactorRewrite,
85    });
86
87    Some(())
88}
89
90fn string_to_dollar_quoted(text: &str) -> Option<String> {
91    let normalized = normalize_single_quoted_string(text)?;
92    let delimiter = dollar_delimiter(&normalized)?;
93    let boundary = format!("${}$", delimiter);
94    Some(format!("{boundary}{normalized}{boundary}"))
95}
96
97fn dollar_quoted_to_string(text: &str) -> Option<String> {
98    debug_assert!(text.starts_with('$'));
99    let (delimiter, content) = split_dollar_quoted(text)?;
100    let boundary = format!("${}$", delimiter);
101
102    if !text.starts_with(&boundary) || !text.ends_with(&boundary) {
103        return None;
104    }
105
106    // quotes are escaped by using two of them in Postgres
107    let escaped = content.replace('\'', "''");
108    Some(format!("'{}'", escaped))
109}
110
111fn split_dollar_quoted(text: &str) -> Option<(String, &str)> {
112    debug_assert!(text.starts_with('$'));
113    let second_dollar = text[1..].find('$')?;
114    // the `foo` in `select $foo$bar$foo$`
115    let delimiter = &text[1..=second_dollar];
116    let boundary = format!("${}$", delimiter);
117
118    if !text.ends_with(&boundary) {
119        return None;
120    }
121
122    let start = boundary.len();
123    let end = text.len().checked_sub(boundary.len())?;
124    let content = text.get(start..end)?;
125    Some((delimiter.to_owned(), content))
126}
127
128fn normalize_single_quoted_string(text: &str) -> Option<String> {
129    let body = text.strip_prefix('\'')?.strip_suffix('\'')?;
130    return Some(body.replace("''", "'"));
131}
132
133fn dollar_delimiter(content: &str) -> Option<String> {
134    // We can't safely transform a trailing `$` i.e., `select 'foo $'` with an
135    // empty delim, because we'll  `select $$foo $$$` which isn't valid.
136    if !content.contains("$$") && !content.ends_with('$') {
137        return Some("".to_owned());
138    }
139
140    let mut delim = "q".to_owned();
141    // don't want to just loop forever
142    for idx in 0..10 {
143        if !content.contains(&format!("${}$", delim)) {
144            return Some(delim);
145        }
146        delim.push_str(&idx.to_string());
147    }
148    None
149}
150
151fn remove_else_clause(
152    actions: &mut Vec<CodeAction>,
153    file: &ast::SourceFile,
154    offset: TextSize,
155) -> Option<()> {
156    let else_token = file
157        .syntax()
158        .token_at_offset(offset)
159        .find(|x| x.kind() == SyntaxKind::ELSE_KW)?;
160    let parent = else_token.parent()?;
161    let else_clause = ast::ElseClause::cast(parent)?;
162
163    let mut edits = vec![];
164    edits.push(Edit::delete(else_clause.syntax().text_range()));
165    if let Some(token) = else_token.prev_token() {
166        if token.kind() == SyntaxKind::WHITESPACE {
167            edits.push(Edit::delete(token.text_range()));
168        }
169    }
170
171    actions.push(CodeAction {
172        title: "Remove `else` clause".to_owned(),
173        edits,
174        kind: ActionKind::RefactorRewrite,
175    });
176    Some(())
177}
178
179fn rewrite_table_as_select(
180    actions: &mut Vec<CodeAction>,
181    file: &ast::SourceFile,
182    offset: TextSize,
183) -> Option<()> {
184    let token = token_from_offset(file, offset)?;
185    let table = token.parent_ancestors().find_map(ast::Table::cast)?;
186
187    let relation_name = table.relation_name()?;
188    let table_name = relation_name.syntax().text();
189
190    let replacement = format!("select * from {}", table_name);
191
192    actions.push(CodeAction {
193        title: "Rewrite as `select`".to_owned(),
194        edits: vec![Edit::replace(table.syntax().text_range(), replacement)],
195        kind: ActionKind::RefactorRewrite,
196    });
197
198    Some(())
199}
200
201fn rewrite_select_as_table(
202    actions: &mut Vec<CodeAction>,
203    file: &ast::SourceFile,
204    offset: TextSize,
205) -> Option<()> {
206    let token = token_from_offset(file, offset)?;
207    let select = token.parent_ancestors().find_map(ast::Select::cast)?;
208
209    if !can_transform_select_to_table(&select) {
210        return None;
211    }
212
213    let from_clause = select.from_clause()?;
214    let from_item = from_clause.from_items().next()?;
215
216    let table_name = if let Some(name_ref) = from_item.name_ref() {
217        name_ref.syntax().text().to_string()
218    } else if let Some(field_expr) = from_item.field_expr() {
219        field_expr.syntax().text().to_string()
220    } else {
221        return None;
222    };
223
224    let replacement = format!("table {}", table_name);
225
226    actions.push(CodeAction {
227        title: "Rewrite as `table`".to_owned(),
228        edits: vec![Edit::replace(select.syntax().text_range(), replacement)],
229        kind: ActionKind::RefactorRewrite,
230    });
231
232    Some(())
233}
234
235/// Returns true if a `select` statement can be safely rewritten as a `table` statement.
236///
237/// We can only do this when there are no clauses besides the `select` and
238/// `from` clause. Additionally, we can only have a table reference in the
239/// `from` clause.
240/// The `select`'s target list must only be a `*`.
241fn can_transform_select_to_table(select: &ast::Select) -> bool {
242    if select.with_clause().is_some()
243        || select.where_clause().is_some()
244        || select.group_by_clause().is_some()
245        || select.having_clause().is_some()
246        || select.window_clause().is_some()
247        || select.order_by_clause().is_some()
248        || select.limit_clause().is_some()
249        || select.fetch_clause().is_some()
250        || select.offset_clause().is_some()
251        || select.filter_clause().is_some()
252        || select.locking_clauses().next().is_some()
253    {
254        return false;
255    }
256
257    let Some(select_clause) = select.select_clause() else {
258        return false;
259    };
260
261    if select_clause.distinct_clause().is_some() {
262        return false;
263    }
264
265    let Some(target_list) = select_clause.target_list() else {
266        return false;
267    };
268
269    let mut targets = target_list.targets();
270    let Some(target) = targets.next() else {
271        return false;
272    };
273
274    if targets.next().is_some() {
275        return false;
276    }
277
278    // only want to support: `select *`
279    if target.expr().is_some() || target.star_token().is_none() {
280        return false;
281    }
282
283    let Some(from_clause) = select.from_clause() else {
284        return false;
285    };
286
287    let mut from_items = from_clause.from_items();
288    let Some(from_item) = from_items.next() else {
289        return false;
290    };
291
292    // only can have one from item & no join exprs
293    if from_items.next().is_some() || from_clause.join_exprs().next().is_some() {
294        return false;
295    }
296
297    if from_item.alias().is_some()
298        || from_item.tablesample_clause().is_some()
299        || from_item.only_token().is_some()
300        || from_item.lateral_token().is_some()
301        || from_item.star_token().is_some()
302        || from_item.call_expr().is_some()
303        || from_item.paren_select().is_some()
304        || from_item.json_table().is_some()
305        || from_item.xml_table().is_some()
306        || from_item.cast_expr().is_some()
307    {
308        return false;
309    }
310
311    // only want table refs
312    from_item.name_ref().is_some() || from_item.field_expr().is_some()
313}
314
315fn quote_identifier(
316    actions: &mut Vec<CodeAction>,
317    file: &ast::SourceFile,
318    offset: TextSize,
319) -> Option<()> {
320    let token = token_from_offset(file, offset)?;
321    let parent = token.parent()?;
322
323    let name_node = if let Some(name) = ast::Name::cast(parent.clone()) {
324        name.syntax().clone()
325    } else if let Some(name_ref) = ast::NameRef::cast(parent) {
326        name_ref.syntax().clone()
327    } else {
328        return None;
329    };
330
331    let text = name_node.text().to_string();
332
333    if text.starts_with('"') {
334        return None;
335    }
336
337    let quoted = format!(r#""{}""#, text.to_lowercase());
338
339    actions.push(CodeAction {
340        title: "Quote identifier".to_owned(),
341        edits: vec![Edit::replace(name_node.text_range(), quoted)],
342        kind: ActionKind::RefactorRewrite,
343    });
344
345    Some(())
346}
347
348fn unquote_identifier(
349    actions: &mut Vec<CodeAction>,
350    file: &ast::SourceFile,
351    offset: TextSize,
352) -> Option<()> {
353    let token = token_from_offset(file, offset)?;
354    let parent = token.parent()?;
355
356    let name_node = if let Some(name) = ast::Name::cast(parent.clone()) {
357        name.syntax().clone()
358    } else if let Some(name_ref) = ast::NameRef::cast(parent) {
359        name_ref.syntax().clone()
360    } else {
361        return None;
362    };
363
364    let unquoted = unquote_ident(&name_node)?;
365
366    actions.push(CodeAction {
367        title: "Unquote identifier".to_owned(),
368        edits: vec![Edit::replace(name_node.text_range(), unquoted)],
369        kind: ActionKind::RefactorRewrite,
370    });
371
372    Some(())
373}
374
375// Postgres docs call these output names.
376// Postgres' parser calls this a column label.
377// Third-party docs call these aliases, so going with that.
378fn add_explicit_alias(
379    actions: &mut Vec<CodeAction>,
380    file: &ast::SourceFile,
381    offset: TextSize,
382) -> Option<()> {
383    let token = token_from_offset(file, offset)?;
384    let target = token.parent_ancestors().find_map(ast::Target::cast)?;
385
386    if target.as_name().is_some() {
387        return None;
388    }
389
390    if let Some(ast::Expr::FieldExpr(field_expr)) = target.expr()
391        && field_expr.star_token().is_some()
392    {
393        return None;
394    }
395
396    let alias = ColumnName::from_target(target.clone()).and_then(|c| c.0.to_string())?;
397
398    let expr_end = target.expr().map(|e| e.syntax().text_range().end())?;
399
400    let quoted_alias = quote_column_alias(&alias);
401    // Postgres docs recommend either using `as` or quoting the name. I think
402    // `as` looks a bit nicer.
403    let replacement = format!(" as {}", quoted_alias);
404
405    actions.push(CodeAction {
406        title: "Add explicit alias".to_owned(),
407        edits: vec![Edit::insert(replacement, expr_end)],
408        kind: ActionKind::RefactorRewrite,
409    });
410
411    Some(())
412}
413
414fn remove_redundant_alias(
415    actions: &mut Vec<CodeAction>,
416    file: &ast::SourceFile,
417    offset: TextSize,
418) -> Option<()> {
419    let token = token_from_offset(file, offset)?;
420    let target = token.parent_ancestors().find_map(ast::Target::cast)?;
421
422    let as_name = target.as_name()?;
423    let (inferred_column, _) = ColumnName::inferred_from_target(target.clone())?;
424    let inferred_column_alias = inferred_column.to_string()?;
425
426    let alias = as_name.name()?;
427
428    if Name::from_node(&alias) != Name::from_string(inferred_column_alias) {
429        return None;
430    }
431
432    // TODO:
433    // This lets use remove any whitespace so we don't end up with:
434    //   select x as x, b from t;
435    // becoming
436    //   select x , b from t;
437    // but we probably want a better way to express this.
438    // Maybe a "Remove preceding whitespace" style option for edits.
439    let expr_end = target.expr()?.syntax().text_range().end();
440    let alias_end = as_name.syntax().text_range().end();
441
442    actions.push(CodeAction {
443        title: "Remove redundant alias".to_owned(),
444        edits: vec![Edit::delete(TextRange::new(expr_end, alias_end))],
445        kind: ActionKind::QuickFix,
446    });
447
448    Some(())
449}
450
451fn add_schema(
452    actions: &mut Vec<CodeAction>,
453    file: &ast::SourceFile,
454    offset: TextSize,
455) -> Option<()> {
456    let token = token_from_offset(file, offset)?;
457    let range = token.parent_ancestors().find_map(|node| {
458        if let Some(path) = ast::Path::cast(node.clone()) {
459            if path.qualifier().is_some() {
460                return None;
461            }
462            return Some(path.syntax().text_range());
463        }
464        if let Some(from_item) = ast::FromItem::cast(node.clone()) {
465            let name_ref = from_item.name_ref()?;
466            return Some(name_ref.syntax().text_range());
467        }
468        if let Some(call_expr) = ast::CallExpr::cast(node) {
469            let ast::Expr::NameRef(name_ref) = call_expr.expr()? else {
470                return None;
471            };
472            return Some(name_ref.syntax().text_range());
473        }
474        None
475    })?;
476
477    if !range.contains(offset) {
478        return None;
479    }
480
481    let position = token.text_range().start();
482    let binder = binder::bind(file);
483    let schema = binder.search_path_at(position).first()?.to_string();
484    let replacement = format!("{}.", schema);
485
486    actions.push(CodeAction {
487        title: "Add schema".to_owned(),
488        edits: vec![Edit::insert(replacement, position)],
489        kind: ActionKind::RefactorRewrite,
490    });
491
492    Some(())
493}
494
495fn rewrite_cast_to_double_colon(
496    actions: &mut Vec<CodeAction>,
497    file: &ast::SourceFile,
498    offset: TextSize,
499) -> Option<()> {
500    let token = token_from_offset(file, offset)?;
501    let cast_expr = token.parent_ancestors().find_map(ast::CastExpr::cast)?;
502
503    if cast_expr.colon_colon().is_some() {
504        return None;
505    }
506
507    let expr = cast_expr.expr()?;
508    let ty = cast_expr.ty()?;
509
510    let expr_text = expr.syntax().text();
511    let type_text = ty.syntax().text();
512
513    let replacement = format!("{}::{}", expr_text, type_text);
514
515    actions.push(CodeAction {
516        title: "Rewrite as cast operator `::`".to_owned(),
517        edits: vec![Edit::replace(cast_expr.syntax().text_range(), replacement)],
518        kind: ActionKind::RefactorRewrite,
519    });
520
521    Some(())
522}
523
524fn rewrite_double_colon_to_cast(
525    actions: &mut Vec<CodeAction>,
526    file: &ast::SourceFile,
527    offset: TextSize,
528) -> Option<()> {
529    let token = token_from_offset(file, offset)?;
530    let cast_expr = token.parent_ancestors().find_map(ast::CastExpr::cast)?;
531
532    if cast_expr.cast_token().is_some() {
533        return None;
534    }
535
536    let expr = cast_expr.expr()?;
537    let ty = cast_expr.ty()?;
538
539    let expr_text = expr.syntax().text();
540    let type_text = ty.syntax().text();
541
542    let replacement = format!("cast({} as {})", expr_text, type_text);
543
544    actions.push(CodeAction {
545        title: "Rewrite as cast function `cast()`".to_owned(),
546        edits: vec![Edit::replace(cast_expr.syntax().text_range(), replacement)],
547        kind: ActionKind::RefactorRewrite,
548    });
549
550    Some(())
551}
552
553fn rewrite_values_as_select(
554    actions: &mut Vec<CodeAction>,
555    file: &ast::SourceFile,
556    offset: TextSize,
557) -> Option<()> {
558    let token = token_from_offset(file, offset)?;
559    let values = token.parent_ancestors().find_map(ast::Values::cast)?;
560
561    let value_token_start = values.values_token().map(|x| x.text_range().start())?;
562    let values_end = values.syntax().text_range().end();
563    // `values` but we skip over the possibly preceeding CTE
564    let values_range = TextRange::new(value_token_start, values_end);
565
566    let mut rows = values.row_list()?.rows();
567
568    let first_targets: Vec<_> = rows
569        .next()?
570        .exprs()
571        .enumerate()
572        .map(|(idx, expr)| format!("{} as column{}", expr.syntax().text(), idx + 1))
573        .collect();
574
575    if first_targets.is_empty() {
576        return None;
577    }
578
579    let mut select_parts = vec![format!("select {}", first_targets.join(", "))];
580
581    for row in rows {
582        let row_targets = row
583            .exprs()
584            .map(|e| e.syntax().text().to_string())
585            .join(", ");
586        if row_targets.is_empty() {
587            return None;
588        }
589        select_parts.push(format!("union all\nselect {}", row_targets));
590    }
591
592    let select_stmt = select_parts.join("\n");
593
594    actions.push(CodeAction {
595        title: "Rewrite as `select`".to_owned(),
596        edits: vec![Edit::replace(values_range, select_stmt)],
597        kind: ActionKind::RefactorRewrite,
598    });
599
600    Some(())
601}
602
603fn is_values_row_column_name(target: &ast::Target, idx: usize) -> bool {
604    let Some(as_name) = target.as_name() else {
605        return false;
606    };
607    let Some(name) = as_name.name() else {
608        return false;
609    };
610    let expected = format!("column{}", idx + 1);
611    if Name::from_node(&name) != Name::from_string(expected) {
612        return false;
613    }
614    true
615}
616
617enum SelectContext {
618    Compound(ast::CompoundSelect),
619    Single(ast::Select),
620}
621
622impl SelectContext {
623    fn iter(&self) -> Option<Box<dyn Iterator<Item = ast::Select>>> {
624        // Ideally we'd have something like Python's `yield` and `yield from`
625        // but instead we have to do all of this to avoid creating some temp
626        // vecs
627        fn variant_iter(
628            variant: ast::SelectVariant,
629        ) -> Option<Box<dyn Iterator<Item = ast::Select>>> {
630            match variant {
631                ast::SelectVariant::Select(select) => Some(Box::new(iter::once(select))),
632                ast::SelectVariant::CompoundSelect(compound) => compound_iter(&compound),
633                ast::SelectVariant::ParenSelect(_)
634                | ast::SelectVariant::SelectInto(_)
635                | ast::SelectVariant::Table(_)
636                | ast::SelectVariant::Values(_) => None,
637            }
638        }
639
640        fn compound_iter(
641            node: &ast::CompoundSelect,
642        ) -> Option<Box<dyn Iterator<Item = ast::Select>>> {
643            let lhs_iter = node
644                .lhs()
645                .map(variant_iter)
646                .unwrap_or_else(|| Some(Box::new(iter::empty())))?;
647            let rhs_iter = node
648                .rhs()
649                .map(variant_iter)
650                .unwrap_or_else(|| Some(Box::new(iter::empty())))?;
651            Some(Box::new(lhs_iter.chain(rhs_iter)))
652        }
653
654        match self {
655            SelectContext::Compound(compound) => compound_iter(compound),
656            SelectContext::Single(select) => Some(Box::new(iter::once(select.clone()))),
657        }
658    }
659}
660
661fn rewrite_select_as_values(
662    actions: &mut Vec<CodeAction>,
663    file: &ast::SourceFile,
664    offset: TextSize,
665) -> Option<()> {
666    let token = token_from_offset(file, offset)?;
667
668    let parent = find_select_parent(token)?;
669
670    let mut selects = parent.iter()?.peekable();
671    let select_token_start = selects
672        .peek()?
673        .select_clause()
674        .and_then(|x| x.select_token())
675        .map(|x| x.text_range().start())?;
676
677    let mut rows = vec![];
678    for (idx, select) in selects.enumerate() {
679        let exprs: Vec<String> = select
680            .select_clause()?
681            .target_list()?
682            .targets()
683            .enumerate()
684            .map(|(i, t)| {
685                if idx != 0 || is_values_row_column_name(&t, i) {
686                    t.expr().map(|expr| expr.syntax().text().to_string())
687                } else {
688                    None
689                }
690            })
691            .collect::<Option<_>>()?;
692
693        if exprs.is_empty() {
694            return None;
695        }
696
697        rows.push(format!("({})", exprs.join(", ")));
698    }
699
700    let values_stmt = format!("values {}", rows.join(", "));
701
702    let select_end = match &parent {
703        SelectContext::Compound(compound) => compound.syntax().text_range().end(),
704        SelectContext::Single(select) => select.syntax().text_range().end(),
705    };
706    let select_range = TextRange::new(select_token_start, select_end);
707
708    actions.push(CodeAction {
709        title: "Rewrite as `values`".to_owned(),
710        edits: vec![Edit::replace(select_range, values_stmt)],
711        kind: ActionKind::RefactorRewrite,
712    });
713
714    Some(())
715}
716
717fn find_select_parent(token: SyntaxToken) -> Option<SelectContext> {
718    let mut found_select = None;
719    let mut found_compound = None;
720    for node in token.parent_ancestors() {
721        if let Some(compound_select) = ast::CompoundSelect::cast(node.clone()) {
722            if compound_select.union_token().is_some() && compound_select.all_token().is_some() {
723                found_compound = Some(SelectContext::Compound(compound_select));
724            } else {
725                break;
726            }
727        }
728        if found_select.is_none()
729            && let Some(select) = ast::Select::cast(node)
730        {
731            found_select = Some(SelectContext::Single(select));
732        }
733    }
734    found_compound.or(found_select)
735}
736
737#[cfg(test)]
738mod test {
739    use super::*;
740    use crate::test_utils::fixture;
741    use insta::assert_snapshot;
742    use rowan::TextSize;
743    use squawk_syntax::ast;
744
745    fn apply_code_action(
746        f: impl Fn(&mut Vec<CodeAction>, &ast::SourceFile, TextSize) -> Option<()>,
747        sql: &str,
748    ) -> String {
749        let (mut offset, sql) = fixture(sql);
750        let parse = ast::SourceFile::parse(&sql);
751        assert_eq!(parse.errors(), vec![]);
752        let file: ast::SourceFile = parse.tree();
753
754        offset = offset.checked_sub(1.into()).unwrap_or_default();
755
756        let mut actions = vec![];
757        f(&mut actions, &file, offset);
758
759        assert!(
760            !actions.is_empty(),
761            "We should always have actions for `apply_code_action`. If you want to ensure there are no actions, use `code_action_not_applicable` instead."
762        );
763
764        let action = &actions[0];
765        let mut result = sql.clone();
766
767        let mut edits = action.edits.clone();
768        edits.sort_by_key(|e| e.text_range.start());
769        check_overlap(&edits);
770        edits.reverse();
771
772        for edit in edits {
773            let start: usize = edit.text_range.start().into();
774            let end: usize = edit.text_range.end().into();
775            let replacement = edit.text.as_deref().unwrap_or("");
776            result.replace_range(start..end, replacement);
777        }
778
779        let reparse = ast::SourceFile::parse(&result);
780        assert_eq!(
781            reparse.errors(),
782            vec![],
783            "Code actions shouldn't cause syntax errors"
784        );
785
786        result
787    }
788
789    // There's an invariant where the edits can't overlap.
790    // For example, if we have an edit that deletes the full `else clause` and
791    // another edit that deletes the `else` keyword and they overlap, then
792    // vscode doesn't surface the code action.
793    fn check_overlap(edits: &[Edit]) {
794        for (edit_i, edit_j) in edits.iter().zip(edits.iter().skip(1)) {
795            if let Some(intersection) = edit_i.text_range.intersect(edit_j.text_range) {
796                assert!(
797                    intersection.is_empty(),
798                    "Edit ranges must not overlap: {:?} and {:?} intersect at {:?}",
799                    edit_i.text_range,
800                    edit_j.text_range,
801                    intersection
802                );
803            }
804        }
805    }
806
807    fn code_action_not_applicable(
808        f: impl Fn(&mut Vec<CodeAction>, &ast::SourceFile, TextSize) -> Option<()>,
809        sql: &str,
810    ) -> bool {
811        let (offset, sql) = fixture(sql);
812        let parse = ast::SourceFile::parse(&sql);
813        assert_eq!(parse.errors(), vec![]);
814        let file: ast::SourceFile = parse.tree();
815
816        let mut actions = vec![];
817        f(&mut actions, &file, offset);
818        actions.is_empty()
819    }
820
821    #[test]
822    fn remove_else_clause_() {
823        assert_snapshot!(apply_code_action(
824            remove_else_clause,
825            "select case x when true then 1 else$0 2 end;"),
826            @"select case x when true then 1 end;"
827        );
828    }
829
830    #[test]
831    fn remove_else_clause_before_token() {
832        assert_snapshot!(apply_code_action(
833            remove_else_clause,
834            "select case x when true then 1 e$0lse 2 end;"),
835            @"select case x when true then 1 end;"
836        );
837    }
838
839    #[test]
840    fn remove_else_clause_not_applicable() {
841        assert!(code_action_not_applicable(
842            remove_else_clause,
843            "select case x when true then 1 else 2 end$0;"
844        ));
845    }
846
847    #[test]
848    fn rewrite_string() {
849        assert_snapshot!(apply_code_action(
850            rewrite_as_dollar_quoted_string,
851            "select 'fo$0o';"),
852            @"select $$foo$$;"
853        );
854    }
855
856    #[test]
857    fn rewrite_string_with_single_quote() {
858        assert_snapshot!(apply_code_action(
859            rewrite_as_dollar_quoted_string,
860            "select 'it''s$0 nice';"),
861            @"select $$it's nice$$;"
862        );
863    }
864
865    #[test]
866    fn rewrite_string_with_dollar_signs() {
867        assert_snapshot!(apply_code_action(
868            rewrite_as_dollar_quoted_string,
869            "select 'foo $$ ba$0r';"),
870            @"select $q$foo $$ bar$q$;"
871        );
872    }
873
874    #[test]
875    fn rewrite_string_when_trailing_dollar() {
876        assert_snapshot!(apply_code_action(
877            rewrite_as_dollar_quoted_string,
878            "select 'foo $'$0;"),
879            @"select $q$foo $$q$;"
880        );
881    }
882
883    #[test]
884    fn rewrite_string_not_applicable() {
885        assert!(code_action_not_applicable(
886            rewrite_as_dollar_quoted_string,
887            "select 1 + $0 2;"
888        ));
889    }
890
891    #[test]
892    fn rewrite_prefix_string_not_applicable() {
893        assert!(code_action_not_applicable(
894            rewrite_as_dollar_quoted_string,
895            "select b'foo$0';"
896        ));
897    }
898
899    #[test]
900    fn rewrite_dollar_string() {
901        assert_snapshot!(apply_code_action(
902            rewrite_as_regular_string,
903            "select $$fo$0o$$;"),
904            @"select 'foo';"
905        );
906    }
907
908    #[test]
909    fn rewrite_dollar_string_with_tag() {
910        assert_snapshot!(apply_code_action(
911            rewrite_as_regular_string,
912            "select $tag$fo$0o$tag$;"),
913            @"select 'foo';"
914        );
915    }
916
917    #[test]
918    fn rewrite_dollar_string_with_quote() {
919        assert_snapshot!(apply_code_action(
920            rewrite_as_regular_string,
921            "select $$it'$0s fine$$;"),
922            @"select 'it''s fine';"
923        );
924    }
925
926    #[test]
927    fn rewrite_dollar_string_not_applicable() {
928        assert!(code_action_not_applicable(
929            rewrite_as_regular_string,
930            "select 'foo$0';"
931        ));
932    }
933
934    #[test]
935    fn rewrite_table_as_select_simple() {
936        assert_snapshot!(apply_code_action(
937            rewrite_table_as_select,
938            "tab$0le foo;"),
939            @"select * from foo;"
940        );
941    }
942
943    #[test]
944    fn rewrite_table_as_select_qualified() {
945        assert_snapshot!(apply_code_action(
946            rewrite_table_as_select,
947            "ta$0ble schema.foo;"),
948            @"select * from schema.foo;"
949        );
950    }
951
952    #[test]
953    fn rewrite_table_as_select_after_keyword() {
954        assert_snapshot!(apply_code_action(
955            rewrite_table_as_select,
956            "table$0 bar;"),
957            @"select * from bar;"
958        );
959    }
960
961    #[test]
962    fn rewrite_table_as_select_on_table_name() {
963        assert_snapshot!(apply_code_action(
964            rewrite_table_as_select,
965            "table fo$0o;"),
966            @"select * from foo;"
967        );
968    }
969
970    #[test]
971    fn rewrite_table_as_select_not_applicable() {
972        assert!(code_action_not_applicable(
973            rewrite_table_as_select,
974            "select * from foo$0;"
975        ));
976    }
977
978    #[test]
979    fn rewrite_select_as_table_simple() {
980        assert_snapshot!(apply_code_action(
981            rewrite_select_as_table,
982            "sel$0ect * from foo;"),
983            @"table foo;"
984        );
985    }
986
987    #[test]
988    fn rewrite_select_as_table_qualified() {
989        assert_snapshot!(apply_code_action(
990            rewrite_select_as_table,
991            "select * from sch$0ema.foo;"),
992            @"table schema.foo;"
993        );
994    }
995
996    #[test]
997    fn rewrite_select_as_table_on_star() {
998        assert_snapshot!(apply_code_action(
999            rewrite_select_as_table,
1000            "select $0* from bar;"),
1001            @"table bar;"
1002        );
1003    }
1004
1005    #[test]
1006    fn rewrite_select_as_table_on_from() {
1007        assert_snapshot!(apply_code_action(
1008            rewrite_select_as_table,
1009            "select * fr$0om baz;"),
1010            @"table baz;"
1011        );
1012    }
1013
1014    #[test]
1015    fn rewrite_select_as_table_not_applicable_with_where() {
1016        assert!(code_action_not_applicable(
1017            rewrite_select_as_table,
1018            "select * from foo$0 where x = 1;"
1019        ));
1020    }
1021
1022    #[test]
1023    fn rewrite_select_as_table_not_applicable_with_order_by() {
1024        assert!(code_action_not_applicable(
1025            rewrite_select_as_table,
1026            "select * from foo$0 order by x;"
1027        ));
1028    }
1029
1030    #[test]
1031    fn rewrite_select_as_table_not_applicable_with_limit() {
1032        assert!(code_action_not_applicable(
1033            rewrite_select_as_table,
1034            "select * from foo$0 limit 10;"
1035        ));
1036    }
1037
1038    #[test]
1039    fn add_schema_simple() {
1040        assert_snapshot!(apply_code_action(
1041            add_schema,
1042            "create table t$0(a text, b int);"),
1043            @"create table public.t(a text, b int);"
1044        );
1045    }
1046
1047    #[test]
1048    fn add_schema_create_foreign_table() {
1049        assert_snapshot!(apply_code_action(
1050            add_schema,
1051            "create foreign table t$0(a text, b int) server foo;"),
1052            @"create foreign table public.t(a text, b int) server foo;"
1053        );
1054    }
1055
1056    #[test]
1057    fn add_schema_create_function() {
1058        assert_snapshot!(apply_code_action(
1059            add_schema,
1060            "create function f$0() returns int8\n  as 'select 1'\n  language sql;"),
1061            @"create function public.f() returns int8
1062  as 'select 1'
1063  language sql;"
1064        );
1065    }
1066
1067    #[test]
1068    fn add_schema_create_type() {
1069        assert_snapshot!(apply_code_action(
1070            add_schema,
1071            "create type t$0 as enum ();"),
1072            @"create type public.t as enum ();"
1073        );
1074    }
1075
1076    #[test]
1077    fn add_schema_table_stmt() {
1078        assert_snapshot!(apply_code_action(
1079            add_schema,
1080            "table t$0;"),
1081            @"table public.t;"
1082        );
1083    }
1084
1085    #[test]
1086    fn add_schema_select_from() {
1087        assert_snapshot!(apply_code_action(
1088            add_schema,
1089            "create table t(a text, b int);
1090        select t from t$0;"),
1091            @"create table t(a text, b int);
1092        select t from public.t;"
1093        );
1094    }
1095
1096    #[test]
1097    fn add_schema_select_table_value() {
1098        // we can't insert the schema here because:
1099        // `select public.t from t` isn't valid
1100        assert!(code_action_not_applicable(
1101            add_schema,
1102            "create table t(a text, b int);
1103        select t$0 from t;"
1104        ));
1105    }
1106
1107    #[test]
1108    fn add_schema_select_unqualified_column() {
1109        // not applicable since we don't have the table name set
1110        // we'll have another quick action to insert table names
1111        assert!(code_action_not_applicable(
1112            add_schema,
1113            "create table t(a text, b int);
1114        select a$0 from t;"
1115        ));
1116    }
1117
1118    #[test]
1119    fn add_schema_select_qualified_column() {
1120        // not valid because we haven't specified the schema on the table name
1121        // `select public.t.c from t` isn't valid sql
1122        assert!(code_action_not_applicable(
1123            add_schema,
1124            "create table t(c text);
1125        select t$0.c from t;"
1126        ));
1127    }
1128
1129    #[test]
1130    fn add_schema_with_search_path() {
1131        assert_snapshot!(
1132            apply_code_action(
1133                add_schema,
1134                "
1135set search_path to myschema;
1136create table t$0(a text, b int);"
1137            ),
1138            @"
1139set search_path to myschema;
1140create table myschema.t(a text, b int);"
1141        );
1142    }
1143
1144    #[test]
1145    fn add_schema_not_applicable_with_schema() {
1146        assert!(code_action_not_applicable(
1147            add_schema,
1148            "create table myschema.t$0(a text, b int);"
1149        ));
1150    }
1151
1152    #[test]
1153    fn add_schema_function_call() {
1154        assert_snapshot!(apply_code_action(
1155            add_schema,
1156            "
1157create function f() returns int8
1158  as 'select 1'
1159  language sql;
1160
1161select f$0();"),
1162            @"
1163create function f() returns int8
1164  as 'select 1'
1165  language sql;
1166
1167select public.f();"
1168        );
1169    }
1170
1171    #[test]
1172    fn add_schema_function_call_not_applicable_with_schema() {
1173        assert!(code_action_not_applicable(
1174            add_schema,
1175            "
1176create function f() returns int8 as 'select 1' language sql;
1177select myschema.f$0();"
1178        ));
1179    }
1180
1181    #[test]
1182    fn rewrite_select_as_table_not_applicable_with_distinct() {
1183        assert!(code_action_not_applicable(
1184            rewrite_select_as_table,
1185            "select distinct * from foo$0;"
1186        ));
1187    }
1188
1189    #[test]
1190    fn rewrite_select_as_table_not_applicable_with_columns() {
1191        assert!(code_action_not_applicable(
1192            rewrite_select_as_table,
1193            "select id, name from foo$0;"
1194        ));
1195    }
1196
1197    #[test]
1198    fn rewrite_select_as_table_not_applicable_with_join() {
1199        assert!(code_action_not_applicable(
1200            rewrite_select_as_table,
1201            "select * from foo$0 join bar on foo.id = bar.id;"
1202        ));
1203    }
1204
1205    #[test]
1206    fn rewrite_select_as_table_not_applicable_with_alias() {
1207        assert!(code_action_not_applicable(
1208            rewrite_select_as_table,
1209            "select * from foo$0 f;"
1210        ));
1211    }
1212
1213    #[test]
1214    fn rewrite_select_as_table_not_applicable_with_multiple_tables() {
1215        assert!(code_action_not_applicable(
1216            rewrite_select_as_table,
1217            "select * from foo$0, bar;"
1218        ));
1219    }
1220
1221    #[test]
1222    fn rewrite_select_as_table_not_applicable_on_table() {
1223        assert!(code_action_not_applicable(
1224            rewrite_select_as_table,
1225            "table foo$0;"
1226        ));
1227    }
1228
1229    #[test]
1230    fn quote_identifier_on_name_ref() {
1231        assert_snapshot!(apply_code_action(
1232            quote_identifier,
1233            "select x$0 from t;"),
1234            @r#"select "x" from t;"#
1235        );
1236    }
1237
1238    #[test]
1239    fn quote_identifier_on_name() {
1240        assert_snapshot!(apply_code_action(
1241            quote_identifier,
1242            "create table T(X$0 int);"),
1243            @r#"create table T("x" int);"#
1244        );
1245    }
1246
1247    #[test]
1248    fn quote_identifier_lowercases() {
1249        assert_snapshot!(apply_code_action(
1250            quote_identifier,
1251            "create table T(COL$0 int);"),
1252            @r#"create table T("col" int);"#
1253        );
1254    }
1255
1256    #[test]
1257    fn quote_identifier_not_applicable_when_already_quoted() {
1258        assert!(code_action_not_applicable(
1259            quote_identifier,
1260            r#"select "x"$0 from t;"#
1261        ));
1262    }
1263
1264    #[test]
1265    fn quote_identifier_not_applicable_on_select_keyword() {
1266        assert!(code_action_not_applicable(
1267            quote_identifier,
1268            "sel$0ect x from t;"
1269        ));
1270    }
1271
1272    #[test]
1273    fn quote_identifier_on_keyword_column_name() {
1274        assert_snapshot!(apply_code_action(
1275            quote_identifier,
1276            "select te$0xt from t;"),
1277            @r#"select "text" from t;"#
1278        );
1279    }
1280
1281    #[test]
1282    fn quote_identifier_example_select() {
1283        assert_snapshot!(apply_code_action(
1284            quote_identifier,
1285            "select x$0 from t;"),
1286            @r#"select "x" from t;"#
1287        );
1288    }
1289
1290    #[test]
1291    fn quote_identifier_example_create_table() {
1292        assert_snapshot!(apply_code_action(
1293            quote_identifier,
1294            "create table T(X$0 int);"),
1295            @r#"create table T("x" int);"#
1296        );
1297    }
1298
1299    #[test]
1300    fn unquote_identifier_simple() {
1301        assert_snapshot!(apply_code_action(
1302            unquote_identifier,
1303            r#"select "x"$0 from t;"#),
1304            @"select x from t;"
1305        );
1306    }
1307
1308    #[test]
1309    fn unquote_identifier_with_underscore() {
1310        assert_snapshot!(apply_code_action(
1311            unquote_identifier,
1312            r#"select "user_id"$0 from t;"#),
1313            @"select user_id from t;"
1314        );
1315    }
1316
1317    #[test]
1318    fn unquote_identifier_with_digits() {
1319        assert_snapshot!(apply_code_action(
1320            unquote_identifier,
1321            r#"select "x123"$0 from t;"#),
1322            @"select x123 from t;"
1323        );
1324    }
1325
1326    #[test]
1327    fn unquote_identifier_with_dollar() {
1328        assert_snapshot!(apply_code_action(
1329            unquote_identifier,
1330            r#"select "my_table$1"$0 from t;"#),
1331            @"select my_table$1 from t;"
1332        );
1333    }
1334
1335    #[test]
1336    fn unquote_identifier_starts_with_underscore() {
1337        assert_snapshot!(apply_code_action(
1338            unquote_identifier,
1339            r#"select "_col"$0 from t;"#),
1340            @"select _col from t;"
1341        );
1342    }
1343
1344    #[test]
1345    fn unquote_identifier_starts_with_unicode() {
1346        assert_snapshot!(apply_code_action(
1347            unquote_identifier,
1348            r#"select "é"$0 from t;"#),
1349            @"select é from t;"
1350        );
1351    }
1352
1353    #[test]
1354    fn unquote_identifier_not_applicable() {
1355        // upper case
1356        assert!(code_action_not_applicable(
1357            unquote_identifier,
1358            r#"select "X"$0 from t;"#
1359        ));
1360        // upper case
1361        assert!(code_action_not_applicable(
1362            unquote_identifier,
1363            r#"select "Foo"$0 from t;"#
1364        ));
1365        // dash
1366        assert!(code_action_not_applicable(
1367            unquote_identifier,
1368            r#"select "my-col"$0 from t;"#
1369        ));
1370        // leading digits
1371        assert!(code_action_not_applicable(
1372            unquote_identifier,
1373            r#"select "123"$0 from t;"#
1374        ));
1375        // space
1376        assert!(code_action_not_applicable(
1377            unquote_identifier,
1378            r#"select "foo bar"$0 from t;"#
1379        ));
1380        // quotes
1381        assert!(code_action_not_applicable(
1382            unquote_identifier,
1383            r#"select "foo""bar"$0 from t;"#
1384        ));
1385        // already unquoted
1386        assert!(code_action_not_applicable(
1387            unquote_identifier,
1388            "select x$0 from t;"
1389        ));
1390        // brackets
1391        assert!(code_action_not_applicable(
1392            unquote_identifier,
1393            r#"select "my[col]"$0 from t;"#
1394        ));
1395        // curly brackets
1396        assert!(code_action_not_applicable(
1397            unquote_identifier,
1398            r#"select "my{}"$0 from t;"#
1399        ));
1400        // reserved word
1401        assert!(code_action_not_applicable(
1402            unquote_identifier,
1403            r#"select "select"$0 from t;"#
1404        ));
1405    }
1406
1407    #[test]
1408    fn unquote_identifier_on_name() {
1409        assert_snapshot!(apply_code_action(
1410            unquote_identifier,
1411            r#"create table T("x"$0 int);"#),
1412            @"create table T(x int);"
1413        );
1414    }
1415
1416    #[test]
1417    fn add_explicit_alias_simple_column() {
1418        assert_snapshot!(apply_code_action(
1419            add_explicit_alias,
1420            "select col_na$0me from t;"),
1421            @"select col_name as col_name from t;"
1422        );
1423    }
1424
1425    #[test]
1426    fn add_explicit_alias_quoted_identifier() {
1427        assert_snapshot!(apply_code_action(
1428            add_explicit_alias,
1429            r#"select "b"$0 from t;"#),
1430            @r#"select "b" as b from t;"#
1431        );
1432    }
1433
1434    #[test]
1435    fn add_explicit_alias_field_expr() {
1436        assert_snapshot!(apply_code_action(
1437            add_explicit_alias,
1438            "select t.col$0umn from t;"),
1439            @"select t.column as column from t;"
1440        );
1441    }
1442
1443    #[test]
1444    fn add_explicit_alias_function_call() {
1445        assert_snapshot!(apply_code_action(
1446            add_explicit_alias,
1447            "select cou$0nt(*) from t;"),
1448            @"select count(*) as count from t;"
1449        );
1450    }
1451
1452    #[test]
1453    fn add_explicit_alias_cast_to_type() {
1454        assert_snapshot!(apply_code_action(
1455            add_explicit_alias,
1456            "select '1'::bigi$0nt from t;"),
1457            @"select '1'::bigint as int8 from t;"
1458        );
1459    }
1460
1461    #[test]
1462    fn add_explicit_alias_cast_column() {
1463        assert_snapshot!(apply_code_action(
1464            add_explicit_alias,
1465            "select col_na$0me::text from t;"),
1466            @"select col_name::text as col_name from t;"
1467        );
1468    }
1469
1470    #[test]
1471    fn add_explicit_alias_case_expr() {
1472        assert_snapshot!(apply_code_action(
1473            add_explicit_alias,
1474            "select ca$0se when true then 'a' end from t;"),
1475            @"select case when true then 'a' end as case from t;"
1476        );
1477    }
1478
1479    #[test]
1480    fn add_explicit_alias_case_with_else() {
1481        assert_snapshot!(apply_code_action(
1482            add_explicit_alias,
1483            "select ca$0se when true then 'a' else now()::text end from t;"),
1484            @"select case when true then 'a' else now()::text end as now from t;"
1485        );
1486    }
1487
1488    #[test]
1489    fn add_explicit_alias_array() {
1490        assert_snapshot!(apply_code_action(
1491            add_explicit_alias,
1492            "select arr$0ay[1, 2, 3] from t;"),
1493            @"select array[1, 2, 3] as array from t;"
1494        );
1495    }
1496
1497    #[test]
1498    fn add_explicit_alias_not_applicable_already_has_alias() {
1499        assert!(code_action_not_applicable(
1500            add_explicit_alias,
1501            "select col_name$0 as foo from t;"
1502        ));
1503    }
1504
1505    #[test]
1506    fn add_explicit_alias_unknown_column() {
1507        assert_snapshot!(apply_code_action(
1508            add_explicit_alias,
1509            "select 1 $0+ 2 from t;"),
1510            @r#"select 1 + 2 as "?column?" from t;"#
1511        );
1512    }
1513
1514    #[test]
1515    fn add_explicit_alias_not_applicable_star() {
1516        assert!(code_action_not_applicable(
1517            add_explicit_alias,
1518            "select $0* from t;"
1519        ));
1520    }
1521
1522    #[test]
1523    fn add_explicit_alias_not_applicable_qualified_star() {
1524        assert!(code_action_not_applicable(
1525            add_explicit_alias,
1526            "with t as (select 1 a) select t.*$0 from t;"
1527        ));
1528    }
1529
1530    #[test]
1531    fn add_explicit_alias_literal() {
1532        assert_snapshot!(apply_code_action(
1533            add_explicit_alias,
1534            "select 'foo$0' from t;"),
1535            @r#"select 'foo' as "?column?" from t;"#
1536        );
1537    }
1538
1539    #[test]
1540    fn remove_redundant_alias_simple() {
1541        assert_snapshot!(apply_code_action(
1542            remove_redundant_alias,
1543            "select col_name as col_na$0me from t;"),
1544            @"select col_name from t;"
1545        );
1546    }
1547
1548    #[test]
1549    fn remove_redundant_alias_quoted() {
1550        assert_snapshot!(apply_code_action(
1551            remove_redundant_alias,
1552            r#"select "x"$0 as x from t;"#),
1553            @r#"select "x" from t;"#
1554        );
1555    }
1556
1557    #[test]
1558    fn remove_redundant_alias_case_insensitive() {
1559        assert_snapshot!(apply_code_action(
1560            remove_redundant_alias,
1561            "select col_name$0 as COL_NAME from t;"),
1562            @"select col_name from t;"
1563        );
1564    }
1565
1566    #[test]
1567    fn remove_redundant_alias_function() {
1568        assert_snapshot!(apply_code_action(
1569            remove_redundant_alias,
1570            "select count(*)$0 as count from t;"),
1571            @"select count(*) from t;"
1572        );
1573    }
1574
1575    #[test]
1576    fn remove_redundant_alias_field_expr() {
1577        assert_snapshot!(apply_code_action(
1578            remove_redundant_alias,
1579            "select t.col$0umn as column from t;"),
1580            @"select t.column from t;"
1581        );
1582    }
1583
1584    #[test]
1585    fn remove_redundant_alias_not_applicable_different_name() {
1586        assert!(code_action_not_applicable(
1587            remove_redundant_alias,
1588            "select col_name$0 as foo from t;"
1589        ));
1590    }
1591
1592    #[test]
1593    fn remove_redundant_alias_not_applicable_no_alias() {
1594        assert!(code_action_not_applicable(
1595            remove_redundant_alias,
1596            "select col_name$0 from t;"
1597        ));
1598    }
1599
1600    #[test]
1601    fn rewrite_cast_to_double_colon_simple() {
1602        assert_snapshot!(apply_code_action(
1603            rewrite_cast_to_double_colon,
1604            "select ca$0st(foo as text) from t;"),
1605            @"select foo::text from t;"
1606        );
1607    }
1608
1609    #[test]
1610    fn rewrite_cast_to_double_colon_on_column() {
1611        assert_snapshot!(apply_code_action(
1612            rewrite_cast_to_double_colon,
1613            "select cast(col_na$0me as int) from t;"),
1614            @"select col_name::int from t;"
1615        );
1616    }
1617
1618    #[test]
1619    fn rewrite_cast_to_double_colon_on_type() {
1620        assert_snapshot!(apply_code_action(
1621            rewrite_cast_to_double_colon,
1622            "select cast(x as bigi$0nt) from t;"),
1623            @"select x::bigint from t;"
1624        );
1625    }
1626
1627    #[test]
1628    fn rewrite_cast_to_double_colon_qualified_type() {
1629        assert_snapshot!(apply_code_action(
1630            rewrite_cast_to_double_colon,
1631            "select cast(x as pg_cata$0log.text) from t;"),
1632            @"select x::pg_catalog.text from t;"
1633        );
1634    }
1635
1636    #[test]
1637    fn rewrite_cast_to_double_colon_expression() {
1638        assert_snapshot!(apply_code_action(
1639            rewrite_cast_to_double_colon,
1640            "select ca$0st(1 + 2 as bigint) from t;"),
1641            @"select 1 + 2::bigint from t;"
1642        );
1643    }
1644
1645    #[test]
1646    fn rewrite_cast_to_double_colon_type_first_syntax() {
1647        assert_snapshot!(apply_code_action(
1648            rewrite_cast_to_double_colon,
1649            "select in$0t '1';"),
1650            @"select '1'::int;"
1651        );
1652    }
1653
1654    #[test]
1655    fn rewrite_cast_to_double_colon_type_first_qualified() {
1656        assert_snapshot!(apply_code_action(
1657            rewrite_cast_to_double_colon,
1658            "select pg_catalog.int$04 '1';"),
1659            @"select '1'::pg_catalog.int4;"
1660        );
1661    }
1662
1663    #[test]
1664    fn rewrite_cast_to_double_colon_not_applicable_already_double_colon() {
1665        assert!(code_action_not_applicable(
1666            rewrite_cast_to_double_colon,
1667            "select foo::te$0xt from t;"
1668        ));
1669    }
1670
1671    #[test]
1672    fn rewrite_cast_to_double_colon_not_applicable_outside_cast() {
1673        assert!(code_action_not_applicable(
1674            rewrite_cast_to_double_colon,
1675            "select fo$0o from t;"
1676        ));
1677    }
1678
1679    #[test]
1680    fn rewrite_double_colon_to_cast_simple() {
1681        assert_snapshot!(apply_code_action(
1682            rewrite_double_colon_to_cast,
1683            "select foo::te$0xt from t;"),
1684            @"select cast(foo as text) from t;"
1685        );
1686    }
1687
1688    #[test]
1689    fn rewrite_double_colon_to_cast_on_column() {
1690        assert_snapshot!(apply_code_action(
1691            rewrite_double_colon_to_cast,
1692            "select col_na$0me::int from t;"),
1693            @"select cast(col_name as int) from t;"
1694        );
1695    }
1696
1697    #[test]
1698    fn rewrite_double_colon_to_cast_on_type() {
1699        assert_snapshot!(apply_code_action(
1700            rewrite_double_colon_to_cast,
1701            "select x::bigi$0nt from t;"),
1702            @"select cast(x as bigint) from t;"
1703        );
1704    }
1705
1706    #[test]
1707    fn rewrite_double_colon_to_cast_qualified_type() {
1708        assert_snapshot!(apply_code_action(
1709            rewrite_double_colon_to_cast,
1710            "select x::pg_cata$0log.text from t;"),
1711            @"select cast(x as pg_catalog.text) from t;"
1712        );
1713    }
1714
1715    #[test]
1716    fn rewrite_double_colon_to_cast_expression() {
1717        assert_snapshot!(apply_code_action(
1718            rewrite_double_colon_to_cast,
1719            "select 1 + 2::bigi$0nt from t;"),
1720            @"select 1 + cast(2 as bigint) from t;"
1721        );
1722    }
1723
1724    #[test]
1725    fn rewrite_type_literal_syntax_to_cast() {
1726        assert_snapshot!(apply_code_action(
1727            rewrite_double_colon_to_cast,
1728            "select in$0t '1';"),
1729            @"select cast('1' as int);"
1730        );
1731    }
1732
1733    #[test]
1734    fn rewrite_qualified_type_literal_syntax_to_cast() {
1735        assert_snapshot!(apply_code_action(
1736            rewrite_double_colon_to_cast,
1737            "select pg_catalog.int$04 '1';"),
1738            @"select cast('1' as pg_catalog.int4);"
1739        );
1740    }
1741
1742    #[test]
1743    fn rewrite_double_colon_to_cast_not_applicable_already_cast() {
1744        assert!(code_action_not_applicable(
1745            rewrite_double_colon_to_cast,
1746            "select ca$0st(foo as text) from t;"
1747        ));
1748    }
1749
1750    #[test]
1751    fn rewrite_double_colon_to_cast_not_applicable_outside_cast() {
1752        assert!(code_action_not_applicable(
1753            rewrite_double_colon_to_cast,
1754            "select fo$0o from t;"
1755        ));
1756    }
1757
1758    #[test]
1759    fn rewrite_values_as_select_simple() {
1760        assert_snapshot!(
1761            apply_code_action(rewrite_values_as_select, "valu$0es (1, 'one'), (2, 'two');"),
1762            @r"
1763        select 1 as column1, 'one' as column2
1764        union all
1765        select 2, 'two';
1766        "
1767        );
1768    }
1769
1770    #[test]
1771    fn rewrite_values_as_select_single_row() {
1772        assert_snapshot!(
1773            apply_code_action(rewrite_values_as_select, "val$0ues (1, 2, 3);"),
1774            @"select 1 as column1, 2 as column2, 3 as column3;"
1775        );
1776    }
1777
1778    #[test]
1779    fn rewrite_values_as_select_single_column() {
1780        assert_snapshot!(
1781            apply_code_action(rewrite_values_as_select, "values$0 (1);"),
1782            @"select 1 as column1;"
1783        );
1784    }
1785
1786    #[test]
1787    fn rewrite_values_as_select_multiple_rows() {
1788        assert_snapshot!(
1789            apply_code_action(rewrite_values_as_select, "values (1, 2), (3, 4), (5, 6$0);"),
1790            @r"
1791        select 1 as column1, 2 as column2
1792        union all
1793        select 3, 4
1794        union all
1795        select 5, 6;
1796        "
1797        );
1798    }
1799
1800    #[test]
1801    fn rewrite_values_as_select_with_clause() {
1802        assert_snapshot!(
1803            apply_code_action(
1804                rewrite_values_as_select,
1805                "with cte as (select 1) val$0ues (1, 'one'), (2, 'two');"
1806            ),
1807            @r"
1808        with cte as (select 1) select 1 as column1, 'one' as column2
1809        union all
1810        select 2, 'two';
1811        "
1812        );
1813    }
1814
1815    #[test]
1816    fn rewrite_values_as_select_complex_expressions() {
1817        assert_snapshot!(
1818            apply_code_action(
1819                rewrite_values_as_select,
1820                "values (1 + 2, 'test'::text$0, array[1,2]);"
1821            ),
1822            @"select 1 + 2 as column1, 'test'::text as column2, array[1,2] as column3;"
1823        );
1824    }
1825
1826    #[test]
1827    fn rewrite_values_as_select_on_values_keyword() {
1828        assert_snapshot!(
1829            apply_code_action(rewrite_values_as_select, "val$0ues (1, 2);"),
1830            @"select 1 as column1, 2 as column2;"
1831        );
1832    }
1833
1834    #[test]
1835    fn rewrite_values_as_select_on_row_content() {
1836        assert_snapshot!(
1837            apply_code_action(rewrite_values_as_select, "values (1$0, 2), (3, 4);"),
1838            @r"
1839        select 1 as column1, 2 as column2
1840        union all
1841        select 3, 4;
1842        "
1843        );
1844    }
1845
1846    #[test]
1847    fn rewrite_values_as_select_not_applicable_on_select() {
1848        assert!(code_action_not_applicable(
1849            rewrite_values_as_select,
1850            "sel$0ect 1;"
1851        ));
1852    }
1853
1854    #[test]
1855    fn rewrite_select_as_values_simple() {
1856        assert_snapshot!(
1857            apply_code_action(
1858                rewrite_select_as_values,
1859                "select 1 as column1, 'one' as column2 union all$0 select 2, 'two';"
1860            ),
1861            @"values (1, 'one'), (2, 'two');"
1862        );
1863    }
1864
1865    #[test]
1866    fn rewrite_select_as_values_multiple_rows() {
1867        assert_snapshot!(
1868            apply_code_action(
1869                rewrite_select_as_values,
1870                "select 1 as column1, 2 as column2 union$0 all select 3, 4 union all select 5, 6;"
1871            ),
1872            @"values (1, 2), (3, 4), (5, 6);"
1873        );
1874    }
1875
1876    #[test]
1877    fn rewrite_select_as_values_multiple_rows_cursor_on_second_union() {
1878        assert_snapshot!(
1879            apply_code_action(
1880                rewrite_select_as_values,
1881                "select 1 as column1, 2 as column2 union all select 3, 4 union$0 all select 5, 6;"
1882            ),
1883            @"values (1, 2), (3, 4), (5, 6);"
1884        );
1885    }
1886
1887    #[test]
1888    fn rewrite_select_as_values_single_column() {
1889        assert_snapshot!(
1890            apply_code_action(
1891                rewrite_select_as_values,
1892                "select 1 as column1$0 union all select 2;"
1893            ),
1894            @"values (1), (2);"
1895        );
1896    }
1897
1898    #[test]
1899    fn rewrite_select_as_values_with_clause() {
1900        assert_snapshot!(
1901            apply_code_action(
1902                rewrite_select_as_values,
1903                "with cte as (select 1) select 1 as column1, 'one' as column2 uni$0on all select 2, 'two';"
1904            ),
1905            @"with cte as (select 1) values (1, 'one'), (2, 'two');"
1906        );
1907    }
1908
1909    #[test]
1910    fn rewrite_select_as_values_complex_expressions() {
1911        assert_snapshot!(
1912            apply_code_action(
1913                rewrite_select_as_values,
1914                "select 1 + 2 as column1, 'test'::text as column2$0 union all select 3 * 4, array[1,2]::text;"
1915            ),
1916            @"values (1 + 2, 'test'::text), (3 * 4, array[1,2]::text);"
1917        );
1918    }
1919
1920    #[test]
1921    fn rewrite_select_as_values_single_select() {
1922        assert_snapshot!(
1923            apply_code_action(
1924                rewrite_select_as_values,
1925                "select 1 as column1, 2 as column2$0;"
1926            ),
1927            @"values (1, 2);"
1928        );
1929    }
1930
1931    #[test]
1932    fn rewrite_select_as_values_single_select_with_clause() {
1933        assert_snapshot!(
1934            apply_code_action(
1935                rewrite_select_as_values,
1936                "with cte as (select 1) select 1 as column1$0, 'test' as column2;"
1937            ),
1938            @"with cte as (select 1) values (1, 'test');"
1939        );
1940    }
1941
1942    #[test]
1943    fn rewrite_select_as_values_not_applicable_union_without_all() {
1944        assert!(code_action_not_applicable(
1945            rewrite_select_as_values,
1946            "select 1 as column1 union$0 select 2;"
1947        ));
1948    }
1949
1950    #[test]
1951    fn rewrite_select_as_values_not_applicable_wrong_column_names() {
1952        assert!(code_action_not_applicable(
1953            rewrite_select_as_values,
1954            "select 1 as foo, 2 as bar union all$0 select 3, 4;"
1955        ));
1956    }
1957
1958    #[test]
1959    fn rewrite_select_as_values_not_applicable_missing_aliases() {
1960        assert!(code_action_not_applicable(
1961            rewrite_select_as_values,
1962            "select 1, 2 union all$0 select 3, 4;"
1963        ));
1964    }
1965
1966    #[test]
1967    fn rewrite_select_as_values_case_insensitive_column_names() {
1968        assert_snapshot!(
1969            apply_code_action(
1970                rewrite_select_as_values,
1971                "select 1 as COLUMN1, 2 as CoLuMn2 union all$0 select 3, 4;"
1972            ),
1973            @"values (1, 2), (3, 4);"
1974        );
1975    }
1976
1977    #[test]
1978    fn rewrite_select_as_values_not_applicable_with_values() {
1979        assert!(code_action_not_applicable(
1980            rewrite_select_as_values,
1981            "select 1 as column1, 2 as column2 union all$0 values (3, 4);"
1982        ));
1983    }
1984
1985    #[test]
1986    fn rewrite_select_as_values_not_applicable_with_table() {
1987        assert!(code_action_not_applicable(
1988            rewrite_select_as_values,
1989            "select 1 as column1, 2 as column2 union all$0 table foo;"
1990        ));
1991    }
1992
1993    #[test]
1994    fn rewrite_select_as_values_not_applicable_intersect() {
1995        assert!(code_action_not_applicable(
1996            rewrite_select_as_values,
1997            "select 1 as column1, 2 as column2 inter$0sect select 3, 4;"
1998        ));
1999    }
2000
2001    #[test]
2002    fn rewrite_select_as_values_not_applicable_except() {
2003        assert!(code_action_not_applicable(
2004            rewrite_select_as_values,
2005            "select 1 as column1, 2 as column2 exc$0ept select 3, 4;"
2006        ));
2007    }
2008}