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