Skip to main content

squawk_ide/
code_actions.rs

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