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 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 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 if !content.contains("$$") && !content.ends_with('$') {
147 return Some("".to_owned());
148 }
149
150 let mut delim = "q".to_owned();
151 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
311fn 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 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 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 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
451fn 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 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 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 let binder = binder::bind(file);
560 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 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 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 }
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 }
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 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 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 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 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 assert!(code_action_not_applicable(
1557 unquote_identifier,
1558 r#"select "X"$0 from t;"#
1559 ));
1560 assert!(code_action_not_applicable(
1562 unquote_identifier,
1563 r#"select "Foo"$0 from t;"#
1564 ));
1565 assert!(code_action_not_applicable(
1567 unquote_identifier,
1568 r#"select "my-col"$0 from t;"#
1569 ));
1570 assert!(code_action_not_applicable(
1572 unquote_identifier,
1573 r#"select "123"$0 from t;"#
1574 ));
1575 assert!(code_action_not_applicable(
1577 unquote_identifier,
1578 r#"select "foo bar"$0 from t;"#
1579 ));
1580 assert!(code_action_not_applicable(
1582 unquote_identifier,
1583 r#"select "foo""bar"$0 from t;"#
1584 ));
1585 assert!(code_action_not_applicable(
1587 unquote_identifier,
1588 "select x$0 from t;"
1589 ));
1590 assert!(code_action_not_applicable(
1592 unquote_identifier,
1593 r#"select "my[col]"$0 from t;"#
1594 ));
1595 assert!(code_action_not_applicable(
1597 unquote_identifier,
1598 r#"select "my{}"$0 from t;"#
1599 ));
1600 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}