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