1use crate::collect;
2use crate::db::{File, parse};
3use crate::file::InFile;
4use crate::goto_definition;
5use crate::resolve;
6use crate::symbols::Name;
7use rowan::{TextRange, TextSize};
8use salsa::Database as Db;
9use squawk_syntax::ast::{self, AstNode};
10
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub enum InlayHintKind {
14 Type,
15 Parameter,
16}
17
18#[derive(Clone, PartialEq, Eq)]
19pub struct InlayHint {
20 pub position: TextSize,
21 pub label: String,
22 pub kind: InlayHintKind,
23 pub target: Option<InFile<TextRange>>,
26}
27
28#[salsa::tracked]
29pub fn inlay_hints(db: &dyn Db, file: File) -> Vec<InlayHint> {
30 let mut hints = vec![];
31 for node in parse(db, file).tree().syntax().descendants() {
32 if let Some(call_expr) = ast::CallExpr::cast(node.clone()) {
33 inlay_hint_call_expr(db, &mut hints, file, call_expr);
34 } else if let Some(insert) = ast::Insert::cast(node) {
35 inlay_hint_insert(db, &mut hints, file, insert);
36 }
37 }
38 hints
39}
40
41fn inlay_hint_call_expr(
42 db: &dyn Db,
43 hints: &mut Vec<InlayHint>,
44 file_id: File,
45 call_expr: ast::CallExpr,
46) -> Option<()> {
47 let arg_list = call_expr.arg_list()?;
48 let expr = call_expr.expr()?;
49
50 let name_ref = if let Some(name_ref) = ast::NameRef::cast(expr.syntax().clone()) {
51 name_ref
52 } else {
53 ast::FieldExpr::cast(expr.syntax().clone())?.field()?
54 };
55
56 let location = goto_definition::goto_definition(
57 db,
58 InFile::new(file_id, name_ref.syntax().text_range().start()),
59 )
60 .into_iter()
61 .next()?;
62
63 let def_file = parse(db, location.file).tree();
64
65 let function_name_node = def_file.syntax().covering_element(location.range);
66
67 if let Some(create_function) = function_name_node
68 .ancestors()
69 .find_map(ast::CreateFunction::cast)
70 && let Some(param_list) = create_function.param_list()
71 {
72 for (param, arg) in param_list.params().zip(arg_list.args()) {
73 if let Some(param_name) = param.name() {
74 let arg_start = arg.syntax().text_range().start();
75 let target = Some(InFile::new(location.file, param_name.syntax().text_range()));
76 hints.push(InlayHint {
77 position: arg_start,
78 label: format!("{}: ", param_name.syntax().text()),
79 kind: InlayHintKind::Parameter,
80 target,
81 });
82 }
83 }
84 };
85
86 Some(())
87}
88
89fn inlay_hint_insert(
90 db: &dyn Db,
91 hints: &mut Vec<InlayHint>,
92 file_id: File,
93 insert: ast::Insert,
94) -> Option<()> {
95 let name_start = insert
96 .relation_name_ref()?
97 .path_ref()?
98 .segment()?
99 .name_ref()?
100 .syntax()
101 .text_range()
102 .start();
103 let location = goto_definition::goto_definition(db, InFile::new(file_id, name_start))
106 .into_iter()
107 .next();
108
109 let def_file = location.as_ref().map(|loc| loc.file).unwrap_or(file_id);
110 let def_tree = parse(db, def_file).tree();
111
112 let create_table = location.as_ref().and_then(|loc| {
113 def_tree
114 .syntax()
115 .covering_element(loc.range)
116 .ancestors()
117 .find_map(ast::CreateTableLike::cast)
118 });
119
120 let columns: Vec<(Name, Option<InFile<TextRange>>)> =
121 if let Some(column_list) = insert.column_ref_list() {
122 column_list
124 .column_refs()
125 .filter_map(|col| {
126 let col_name = col.name_ref().map(|x| Name::from_node(&x))?;
127 let target = create_table
128 .as_ref()
129 .and_then(|x| {
130 resolve::find_column_in_create_table(
131 db,
132 InFile::new(def_file, x),
133 &col_name,
134 )
135 })
136 .and_then(|x| x.into_iter().next())
137 .map(|x| InFile::new(x.file, x.range));
138 Some((col_name, target))
139 })
140 .collect()
141 } else {
142 collect::columns_from_create_table(db, def_file, &create_table?)
144 .into_iter()
145 .map(|(col_name, ptr)| {
146 let target = ptr.map(|ptr| InFile::new(ptr.file_id, ptr.value.text_range()));
147 (col_name, target)
148 })
149 .collect()
150 };
151
152 inlay_hint_insert_select(hints, columns, insert.select_variant()?)
153}
154
155fn inlay_hint_insert_select(
156 hints: &mut Vec<InlayHint>,
157 columns: Vec<(Name, Option<InFile<TextRange>>)>,
158 select_variant: ast::SelectVariant,
159) -> Option<()> {
160 if let ast::SelectVariant::Values(values) = &select_variant {
161 for row in values.row_list()?.rows() {
163 for ((column_name, target), expr) in columns.iter().zip(row.exprs()) {
164 let expr_start = expr.syntax().text_range().start();
165 hints.push(InlayHint {
166 position: expr_start,
167 label: format!("{column_name}: "),
168 kind: InlayHintKind::Parameter,
169 target: *target,
170 });
171 }
172 }
173 return Some(());
174 }
175
176 let target_list = select_variant.target_list()?;
178 for ((column_name, target), target_expr) in columns.iter().zip(target_list.targets()) {
179 let expr = target_expr.expr()?;
180 let expr_start = expr.syntax().text_range().start();
181 hints.push(InlayHint {
182 position: expr_start,
183 label: format!("{column_name}: "),
184 kind: InlayHintKind::Parameter,
185 target: *target,
186 });
187 }
188
189 Some(())
190}
191
192#[cfg(test)]
193mod test {
194 use crate::builtins::builtins_file;
195 use crate::db::{Database, File};
196 use crate::inlay_hints::{InlayHint, inlay_hints};
197 use annotate_snippets::{AnnotationKind, Level, Renderer, Snippet, renderer::DecorStyle};
198 use insta::assert_snapshot;
199 use rustc_hash::FxHashMap;
200 use std::ops::Range;
201
202 #[must_use]
203 #[track_caller]
204 fn check_inlay_hints(sql: &str) -> String {
205 let db = Database::default();
206 let file = File::new(&db, sql.to_string().into());
207
208 assert_eq!(crate::db::parse(&db, file).errors(), vec![]);
209
210 let hints = inlay_hints(&db, file);
211
212 if hints.is_empty() {
213 return String::new();
214 }
215
216 let mut modified_sql = sql.to_string();
217 let mut indexed: Vec<(usize, &InlayHint)> = hints.iter().enumerate().collect();
218 indexed.sort_by_key(|(_, h)| h.position);
219
220 let mut label_annotations: Vec<Range<usize>> = vec![0..0; hints.len()];
221 let mut cumulative = 0;
222 for (i, hint) in &indexed {
223 let pos: usize = hint.position.into();
224 let new_pos = pos + cumulative;
225 modified_sql.insert_str(new_pos, &hint.label);
226 label_annotations[*i] = new_pos..new_pos + hint.label.len();
227 cumulative += hint.label.len();
228 }
229
230 let mut targets_by_file: FxHashMap<File, Vec<(usize, Range<usize>)>> = FxHashMap::default();
231 for (i, hint) in hints.iter().enumerate() {
232 if let Some(target) = &hint.target {
233 let start: usize = target.value.start().into();
234 let end: usize = target.value.end().into();
235 targets_by_file
236 .entry(target.file_id)
237 .or_default()
238 .push((i + 1, start..end));
239 }
240 }
241
242 let mut file_paths: FxHashMap<File, &'static str> = FxHashMap::default();
243 file_paths.insert(file, "current.sql");
244 file_paths.insert(builtins_file(&db), "builtins.sql");
245
246 let mut labels_snippet = Snippet::source(&modified_sql).fold(true);
247 for (i, range) in label_annotations.into_iter().enumerate() {
248 labels_snippet = labels_snippet.annotation(
249 AnnotationKind::Context
250 .span(range)
251 .label(format!("{}. label", i + 1)),
252 );
253 }
254
255 let mut groups = vec![Level::INFO.primary_title("labels").element(labels_snippet)];
256
257 let mut target_entries = targets_by_file.into_iter().collect::<Vec<_>>();
258 target_entries.sort_by_key(|(_, targets)| {
259 targets.iter().map(|(i, _)| *i).min().unwrap_or(usize::MAX)
260 });
261
262 let target_contents = target_entries
263 .into_iter()
264 .map(|(f, targets)| {
265 let path = *file_paths.get(&f).unwrap();
266 (f.content(&db).clone(), path, targets)
267 })
268 .collect::<Vec<_>>();
269
270 for (content, path, targets) in &target_contents {
271 let mut snippet = Snippet::source(content.as_ref()).fold(true).path(*path);
272 for (i, range) in targets {
273 snippet = snippet.annotation(
274 AnnotationKind::Context
275 .span(range.clone())
276 .label(format!("{i}. target")),
277 );
278 }
279 groups.push(Level::INFO.primary_title("targets").element(snippet));
280 }
281
282 let renderer = Renderer::plain().decor_style(DecorStyle::Unicode);
283 renderer
284 .render(&groups)
285 .to_string()
286 .replace("info: labels", "labels:")
287 .replace("info: targets", "targets:")
288 }
289
290 #[test]
291 fn single_param() {
292 assert_snapshot!(check_inlay_hints("
293create function foo(a int) returns int as 'select $$1' language sql;
294select foo(1);
295"), @"
296 labels:
297 ╭▸
298 3 │ select foo(a: 1);
299 │ ─── 1. label
300 ╰╴
301 targets:
302 ╭▸ current.sql:2:21
303 │
304 2 │ create function foo(a int) returns int as 'select $$1' language sql;
305 ╰╴ ─ 1. target
306 ");
307 }
308
309 #[test]
310 fn multiple_params() {
311 assert_snapshot!(check_inlay_hints("
312create function add(a int, b int) returns int as 'select $$1 + $$2' language sql;
313select add(1, 2);
314"), @"
315 labels:
316 ╭▸
317 3 │ select add(a: 1, b: 2);
318 │ ┬── ─── 2. label
319 │ │
320 │ 1. label
321 ╰╴
322 targets:
323 ╭▸ current.sql:2:21
324 │
325 2 │ create function add(a int, b int) returns int as 'select $$1 + $$2' language sql;
326 │ ┬ ─ 2. target
327 │ │
328 ╰╴ 1. target
329 ");
330 }
331
332 #[test]
333 fn no_params() {
334 assert_snapshot!(check_inlay_hints("
335create function foo() returns int as 'select 1' language sql;
336select foo();
337"), @"");
338 }
339
340 #[test]
341 fn with_schema() {
342 assert_snapshot!(check_inlay_hints("
343create function public.foo(x int) returns int as 'select $$1' language sql;
344select public.foo(42);
345"), @"
346 labels:
347 ╭▸
348 3 │ select public.foo(x: 42);
349 │ ─── 1. label
350 ╰╴
351 targets:
352 ╭▸ current.sql:2:28
353 │
354 2 │ create function public.foo(x int) returns int as 'select $$1' language sql;
355 ╰╴ ─ 1. target
356 ");
357 }
358
359 #[test]
360 fn with_search_path() {
361 assert_snapshot!(check_inlay_hints(r#"
362set search_path to myschema;
363create function foo(val int) returns int as 'select $$1' language sql;
364select foo(100);
365"#), @"
366 labels:
367 ╭▸
368 4 │ select foo(val: 100);
369 │ ───── 1. label
370 ╰╴
371 targets:
372 ╭▸ current.sql:3:21
373 │
374 3 │ create function foo(val int) returns int as 'select $$1' language sql;
375 ╰╴ ─── 1. target
376 ");
377 }
378
379 #[test]
380 fn multiple_calls() {
381 assert_snapshot!(check_inlay_hints("
382create function inc(n int) returns int as 'select $$1 + 1' language sql;
383select inc(1), inc(2);
384"), @"
385 labels:
386 ╭▸
387 3 │ select inc(n: 1), inc(n: 2);
388 │ ┬── ─── 2. label
389 │ │
390 │ 1. label
391 ╰╴
392 targets:
393 ╭▸ current.sql:2:21
394 │
395 2 │ create function inc(n int) returns int as 'select $$1 + 1' language sql;
396 │ ┬
397 │ │
398 │ 1. target
399 ╰╴ 2. target
400 ");
401 }
402
403 #[test]
404 fn more_args_than_params() {
405 assert_snapshot!(check_inlay_hints("
406create function foo(a int) returns int as 'select $$1' language sql;
407select foo(1, 2);
408"), @"
409 labels:
410 ╭▸
411 3 │ select foo(a: 1, 2);
412 │ ─── 1. label
413 ╰╴
414 targets:
415 ╭▸ current.sql:2:21
416 │
417 2 │ create function foo(a int) returns int as 'select $$1' language sql;
418 ╰╴ ─ 1. target
419 ");
420 }
421
422 #[test]
423 fn builtin_function() {
424 assert_snapshot!(check_inlay_hints("
425select json_strip_nulls('[1, null]', true);
426"), @"
427 labels:
428 ╭▸
429 2 │ select json_strip_nulls(target: '[1, null]', strip_in_arrays: true);
430 │ ──────── 1. label ───────────────── 2. label
431 ╰╴
432 targets:
433 ╭▸ builtins.sql:9239:45
434 │
435 9239 │ create function pg_catalog.json_strip_nulls(target json, strip_in_arrays boolean DEFAULT false) returns json
436 │ ┬───── ─────────────── 2. target
437 │ │
438 ╰╴ 1. target
439 ");
440 }
441
442 #[test]
443 fn insert_with_column_list() {
444 assert_snapshot!(check_inlay_hints("
445create table t (column_a int, column_b int, column_c text);
446insert into t (column_a, column_c) values (1, 'foo');
447"), @"
448 labels:
449 ╭▸
450 3 │ insert into t (column_a, column_c) values (column_a: 1, column_c: 'foo');
451 │ ┬───────── ────────── 2. label
452 │ │
453 │ 1. label
454 ╰╴
455 targets:
456 ╭▸ current.sql:2:17
457 │
458 2 │ create table t (column_a int, column_b int, column_c text);
459 ╰╴ ──────── 1. target ──────── 2. target
460 ");
461 }
462
463 #[test]
464 fn insert_without_column_list() {
465 assert_snapshot!(check_inlay_hints("
466create table t (column_a int, column_b int, column_c text);
467insert into t values (1, 2, 'foo');
468"), @"
469 labels:
470 ╭▸
471 3 │ insert into t values (column_a: 1, column_b: 2, column_c: 'foo');
472 │ ┬───────── ┬───────── ────────── 3. label
473 │ │ │
474 │ │ 2. label
475 │ 1. label
476 ╰╴
477 targets:
478 ╭▸ current.sql:2:17
479 │
480 2 │ create table t (column_a int, column_b int, column_c text);
481 │ ┬─────── ┬─────── ──────── 3. target
482 │ │ │
483 │ │ 2. target
484 ╰╴ 1. target
485 ");
486 }
487
488 #[test]
489 fn insert_multiple_rows() {
490 assert_snapshot!(check_inlay_hints("
491create table t (x int, y int);
492insert into t values (1, 2), (3, 4);
493"), @"
494 labels:
495 ╭▸
496 3 │ insert into t values (x: 1, y: 2), (x: 3, y: 4);
497 │ ┬── ┬── ┬── ─── 4. label
498 │ │ │ │
499 │ │ │ 3. label
500 │ │ 2. label
501 │ 1. label
502 ╰╴
503 targets:
504 ╭▸ current.sql:2:17
505 │
506 2 │ create table t (x int, y int);
507 │ ┬ ┬
508 │ │ │
509 │ │ 2. target
510 │ │ 4. target
511 │ 1. target
512 ╰╴ 3. target
513 ");
514 }
515
516 #[test]
517 fn insert_no_create_table() {
518 assert_snapshot!(check_inlay_hints("
519insert into t (a, b) values (1, 2);
520"), @"
521 labels:
522 ╭▸
523 2 │ insert into t (a, b) values (a: 1, b: 2);
524 │ ┬── ─── 2. label
525 │ │
526 ╰╴ 1. label
527 ");
528 }
529
530 #[test]
531 fn insert_more_values_than_columns() {
532 assert_snapshot!(check_inlay_hints("
533create table t (a int, b int);
534insert into t values (1, 2, 3);
535"), @"
536 labels:
537 ╭▸
538 3 │ insert into t values (a: 1, b: 2, 3);
539 │ ┬── ─── 2. label
540 │ │
541 │ 1. label
542 ╰╴
543 targets:
544 ╭▸ current.sql:2:17
545 │
546 2 │ create table t (a int, b int);
547 │ ┬ ─ 2. target
548 │ │
549 ╰╴ 1. target
550 ");
551 }
552
553 #[test]
554 fn insert_table_inherits_select() {
555 assert_snapshot!(check_inlay_hints("
556create table t (a int, b int);
557create table u (c int) inherits (t);
558insert into u select 1, 2, 3;
559"), @"
560 labels:
561 ╭▸
562 4 │ insert into u select a: 1, b: 2, c: 3;
563 │ ┬── ┬── ─── 3. label
564 │ │ │
565 │ │ 2. label
566 │ 1. label
567 ╰╴
568 targets:
569 ╭▸ current.sql:2:17
570 │
571 2 │ create table t (a int, b int);
572 │ ┬ ─ 2. target
573 │ │
574 │ 1. target
575 3 │ create table u (c int) inherits (t);
576 ╰╴ ─ 3. target
577 ");
578 }
579
580 #[test]
581 fn insert_table_inherits_builtin_values() {
582 assert_snapshot!(check_inlay_hints("
583create table t ()
584inherits (information_schema.sql_features);
585insert into t values (1, 2, 3, 4, 5, 6, 7);
586"), @"
587 labels:
588 ╭▸
589 4 │ …ues (feature_id: 1, feature_name: 2, sub_feature_id: 3, sub_feature_name: 4, is_supported: 5, is_verified_by: 6, comments: 7);
590 │ ┬─────────── ┬───────────── ┬─────────────── ┬───────────────── ┬───────────── ┬─────────────── ────────── 7. label
591 │ │ │ │ │ │ │
592 │ │ │ │ │ │ 6. label
593 │ │ │ │ │ 5. label
594 │ │ │ │ 4. label
595 │ │ │ 3. label
596 │ │ 2. label
597 │ 1. label
598 ╰╴
599 targets:
600 ╭▸ builtins.sql:436:3
601 │
602 436 │ feature_id information_schema.character_data,
603 │ ────────── 1. target
604 437 │ feature_name information_schema.character_data,
605 │ ──────────── 2. target
606 438 │ sub_feature_id information_schema.character_data,
607 │ ────────────── 3. target
608 439 │ sub_feature_name information_schema.character_data,
609 │ ──────────────── 4. target
610 440 │ is_supported information_schema.yes_or_no,
611 │ ──────────── 5. target
612 441 │ is_verified_by information_schema.character_data,
613 │ ────────────── 6. target
614 442 │ comments information_schema.character_data
615 ╰╴ ──────── 7. target
616 ");
617 }
618
619 #[test]
620 fn insert_table_inherits_create_table_as_values() {
621 assert_snapshot!(check_inlay_hints("
622create table parent as select 1 a, 'x'::text b;
623create table child (c int) inherits (parent);
624insert into child values (1, 2, 3);
625"), @"
626 labels:
627 ╭▸
628 4 │ insert into child values (a: 1, b: 2, c: 3);
629 │ ┬── ┬── ─── 3. label
630 │ │ │
631 │ │ 2. label
632 │ 1. label
633 ╰╴
634 targets:
635 ╭▸ current.sql:3:21
636 │
637 3 │ create table child (c int) inherits (parent);
638 ╰╴ ─ 3. target
639 ");
640 }
641
642 #[test]
643 fn insert_table_inherits_create_table_as_select_star() {
644 assert_snapshot!(check_inlay_hints("
645create table base (a int, b text);
646create table parent as select * from base;
647create table child (c int) inherits (parent);
648insert into child values (1, 2, 3);
649"), @"
650 labels:
651 ╭▸
652 5 │ insert into child values (a: 1, b: 2, c: 3);
653 │ ┬── ┬── ─── 3. label
654 │ │ │
655 │ │ 2. label
656 │ 1. label
657 ╰╴
658 targets:
659 ╭▸ current.sql:4:21
660 │
661 4 │ create table child (c int) inherits (parent);
662 ╰╴ ─ 3. target
663 ");
664 }
665
666 #[test]
667 fn insert_table_like_select() {
668 assert_snapshot!(check_inlay_hints("
669create table x (a int, b int);
670create table y (c int, like x);
671insert into y select 1, 2, 3;
672"), @"
673 labels:
674 ╭▸
675 4 │ insert into y select c: 1, a: 2, b: 3;
676 │ ┬── ┬── ─── 3. label
677 │ │ │
678 │ │ 2. label
679 │ 1. label
680 ╰╴
681 targets:
682 ╭▸ current.sql:2:17
683 │
684 2 │ create table x (a int, b int);
685 │ ┬ ─ 3. target
686 │ │
687 │ 2. target
688 3 │ create table y (c int, like x);
689 ╰╴ ─ 1. target
690 ");
691 }
692
693 #[test]
694 fn insert_select() {
695 assert_snapshot!(check_inlay_hints("
696create table t (a int, b int);
697insert into t select 1, 2;
698"), @"
699 labels:
700 ╭▸
701 3 │ insert into t select a: 1, b: 2;
702 │ ┬── ─── 2. label
703 │ │
704 │ 1. label
705 ╰╴
706 targets:
707 ╭▸ current.sql:2:17
708 │
709 2 │ create table t (a int, b int);
710 │ ┬ ─ 2. target
711 │ │
712 ╰╴ 1. target
713 ");
714 }
715
716 #[test]
717 fn insert_table_like_builtin_values() {
718 assert_snapshot!(check_inlay_hints("
719create table t (like information_schema.sql_features);
720insert into t values (1, 2, 3, 4, 5, 6, 7);
721"), @"
722 labels:
723 ╭▸
724 3 │ …ues (feature_id: 1, feature_name: 2, sub_feature_id: 3, sub_feature_name: 4, is_supported: 5, is_verified_by: 6, comments: 7);
725 │ ┬─────────── ┬───────────── ┬─────────────── ┬───────────────── ┬───────────── ┬─────────────── ────────── 7. label
726 │ │ │ │ │ │ │
727 │ │ │ │ │ │ 6. label
728 │ │ │ │ │ 5. label
729 │ │ │ │ 4. label
730 │ │ │ 3. label
731 │ │ 2. label
732 │ 1. label
733 ╰╴
734 targets:
735 ╭▸ builtins.sql:436:3
736 │
737 436 │ feature_id information_schema.character_data,
738 │ ────────── 1. target
739 437 │ feature_name information_schema.character_data,
740 │ ──────────── 2. target
741 438 │ sub_feature_id information_schema.character_data,
742 │ ────────────── 3. target
743 439 │ sub_feature_name information_schema.character_data,
744 │ ──────────────── 4. target
745 440 │ is_supported information_schema.yes_or_no,
746 │ ──────────── 5. target
747 441 │ is_verified_by information_schema.character_data,
748 │ ────────────── 6. target
749 442 │ comments information_schema.character_data
750 ╰╴ ──────── 7. target
751 ");
752 }
753
754 #[test]
755 fn insert_table_like_select_into_values() {
756 assert_snapshot!(check_inlay_hints("
757select 1 a, 'x'::text b into parent;
758create table child (like parent);
759insert into child values (1, 2);
760"), @"
761 labels:
762 ╭▸
763 4 │ insert into child values (a: 1, b: 2);
764 │ ┬── ─── 2. label
765 │ │
766 ╰╴ 1. label
767 ");
768 }
769}