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