Skip to main content

squawk_ide/
completion.rs

1use rowan::TextSize;
2use salsa::Database as Db;
3use squawk_syntax::ast::{self, AstNode};
4use squawk_syntax::{SyntaxKind, SyntaxNode, SyntaxToken};
5
6use crate::ast_nav;
7use crate::binder;
8use crate::collect;
9use crate::db::{File, bind, parse};
10use crate::file::InFile;
11use crate::name::{self, Name, Schema};
12use crate::resolve;
13use crate::symbols::SymbolKind;
14use crate::tokens::is_string_or_comment;
15
16const COMPLETION_MARKER: &str = "squawkCompletionMarker";
17
18pub fn completion(db: &dyn Db, position: InFile<TextSize>) -> Vec<CompletionItem> {
19    let file = position.file_id;
20    let offset = position.value;
21    let parse = parse(db, file);
22    let source_file = parse.tree();
23
24    let marker_file = file_with_completion_marker(&source_file, offset);
25    let Some(token) = token_at_offset(&marker_file, offset) else {
26        // empty file
27        return default_completions();
28    };
29    // We don't support completions inside comments since we don't have doc
30    // comments a la JSDoc.
31    // And we don't support enums aka string literal types yet so we bail out
32    // early for strings as well
33    if is_string_or_comment(token.kind()) {
34        return vec![];
35    }
36
37    match completion_context(&token) {
38        CompletionContext::TableOnly => table_completions(&marker_file, &token),
39        CompletionContext::Default => default_completions(),
40        CompletionContext::SelectClause(select_clause) => {
41            select_completions(db, file, select_clause, &token)
42        }
43        CompletionContext::SelectClauses(select) => select_clauses_completions(&select),
44        CompletionContext::SelectExpr(select) => select_expr_completions(db, file, &select, &token),
45        CompletionContext::LimitClause => limit_completions(db, file, &token),
46        CompletionContext::OffsetClause => offset_completions(db, file, &token),
47        CompletionContext::DeleteClauses(delete) => {
48            delete_clauses_completions(&marker_file, &delete, &token)
49        }
50        CompletionContext::DeleteExpr(delete) => {
51            delete_expr_completions(db, file, &marker_file, &delete, &token)
52        }
53    }
54}
55
56fn select_completions(
57    db: &dyn Db,
58    file: File,
59    select_clause: ast::SelectClause,
60    token: &SyntaxToken,
61) -> Vec<CompletionItem> {
62    let binder = bind(db, file);
63    let mut completions = vec![];
64    let schema = schema_qualifier_at_token(token);
65    let position = token.text_range().start();
66
67    completions.extend(function_completions(db, file, schema.as_ref(), position));
68
69    let tables = binder.all_symbols_by_kind(SymbolKind::Table, schema.as_ref());
70    completions.extend(tables.into_iter().map(|name| CompletionItem {
71        label: name.to_string(),
72        kind: CompletionItemKind::Table,
73        detail: None,
74        insert_text: None,
75        insert_text_format: None,
76        trigger_completion_after_insert: false,
77        sort_text: None,
78    }));
79
80    if schema.is_none() {
81        completions.extend(schema_completions(&binder));
82    }
83
84    if let Some(parent) = select_clause.syntax().parent()
85        && let Some(select) = ast::Select::cast(parent)
86    {
87        if let Some(from_clause) = select.from_clause() {
88            completions.push(CompletionItem {
89                label: "*".to_string(),
90                kind: CompletionItemKind::Operator,
91                detail: None,
92                insert_text: None,
93                insert_text_format: None,
94                trigger_completion_after_insert: false,
95                sort_text: None,
96            });
97            completions.extend(column_completions_from_clause(db, file, &from_clause));
98        } else if schema.is_none() {
99            completions.extend(select_clauses_completions(&select));
100        }
101    }
102
103    completions
104}
105
106fn select_clauses_completions(select: &ast::Select) -> Vec<CompletionItem> {
107    let mut completions = vec![];
108
109    if select.from_clause().is_none() {
110        completions.push(CompletionItem {
111            label: "from".to_owned(),
112            kind: CompletionItemKind::Snippet,
113            detail: None,
114            insert_text: Some("from $0".to_owned()),
115            insert_text_format: Some(CompletionInsertTextFormat::Snippet),
116            trigger_completion_after_insert: true,
117            sort_text: None,
118        });
119    }
120
121    if select.where_clause().is_none() {
122        completions.push(CompletionItem {
123            label: "where".to_owned(),
124            kind: CompletionItemKind::Snippet,
125            detail: None,
126            insert_text: Some("where $0".to_owned()),
127            insert_text_format: Some(CompletionInsertTextFormat::Snippet),
128            trigger_completion_after_insert: true,
129            sort_text: None,
130        });
131    }
132
133    if select.group_by_clause().is_none() {
134        completions.push(CompletionItem {
135            label: "group by".to_owned(),
136            kind: CompletionItemKind::Snippet,
137            detail: None,
138            insert_text: Some("group by $0".to_owned()),
139            insert_text_format: Some(CompletionInsertTextFormat::Snippet),
140            trigger_completion_after_insert: true,
141            sort_text: None,
142        });
143    }
144
145    if select.having_clause().is_none() {
146        completions.push(CompletionItem {
147            label: "having".to_owned(),
148            kind: CompletionItemKind::Snippet,
149            detail: None,
150            insert_text: Some("having $0".to_owned()),
151            insert_text_format: Some(CompletionInsertTextFormat::Snippet),
152            trigger_completion_after_insert: true,
153            sort_text: None,
154        });
155    }
156
157    if select.order_by_clause().is_none() {
158        completions.push(CompletionItem {
159            label: "order by".to_owned(),
160            kind: CompletionItemKind::Snippet,
161            detail: None,
162            insert_text: Some("order by $0".to_owned()),
163            insert_text_format: Some(CompletionInsertTextFormat::Snippet),
164            trigger_completion_after_insert: true,
165            sort_text: None,
166        });
167    }
168
169    if select.limit_clause().is_none() {
170        completions.push(CompletionItem {
171            label: "limit".to_owned(),
172            kind: CompletionItemKind::Snippet,
173            detail: None,
174            insert_text: Some("limit $0".to_owned()),
175            insert_text_format: Some(CompletionInsertTextFormat::Snippet),
176            trigger_completion_after_insert: true,
177            sort_text: None,
178        });
179    }
180
181    if select.offset_clause().is_none() {
182        completions.push(CompletionItem {
183            label: "offset".to_owned(),
184            kind: CompletionItemKind::Snippet,
185            detail: None,
186            insert_text: Some("offset $0".to_owned()),
187            insert_text_format: Some(CompletionInsertTextFormat::Snippet),
188            trigger_completion_after_insert: true,
189            sort_text: None,
190        });
191    }
192
193    if select.fetch_clause().is_none() {
194        completions.push(CompletionItem {
195            label: "fetch".to_owned(),
196            kind: CompletionItemKind::Snippet,
197            detail: None,
198            insert_text: Some(
199                "fetch ${1|first,next|} $2 ${3|row,rows|} ${4|only,with ties|}".to_owned(),
200            ),
201            insert_text_format: Some(CompletionInsertTextFormat::Snippet),
202            trigger_completion_after_insert: true,
203            sort_text: None,
204        });
205    }
206
207    if select.locking_clauses().next().is_none() {
208        completions.push(CompletionItem {
209            label: "for".to_owned(),
210            kind: CompletionItemKind::Snippet,
211            detail: None,
212            insert_text: Some("for ${1|update,no key update,share,key share|} $2".to_owned()),
213            insert_text_format: Some(CompletionInsertTextFormat::Snippet),
214            trigger_completion_after_insert: true,
215            sort_text: None,
216        });
217    }
218
219    if select.window_clause().is_none() {
220        completions.push(CompletionItem {
221            label: "window".to_owned(),
222            kind: CompletionItemKind::Snippet,
223            detail: None,
224            insert_text: Some("window $1 as ($0)".to_owned()),
225            insert_text_format: Some(CompletionInsertTextFormat::Snippet),
226            trigger_completion_after_insert: true,
227            sort_text: None,
228        });
229    }
230
231    completions.push(CompletionItem {
232        label: "union".to_owned(),
233        kind: CompletionItemKind::Snippet,
234        detail: None,
235        insert_text: Some("union $0".to_owned()),
236        insert_text_format: Some(CompletionInsertTextFormat::Snippet),
237        trigger_completion_after_insert: true,
238        sort_text: None,
239    });
240    completions.push(CompletionItem {
241        label: "intersect".to_owned(),
242        kind: CompletionItemKind::Snippet,
243        detail: None,
244        insert_text: Some("intersect $0".to_owned()),
245        insert_text_format: Some(CompletionInsertTextFormat::Snippet),
246        trigger_completion_after_insert: true,
247        sort_text: None,
248    });
249    completions.push(CompletionItem {
250        label: "except".to_owned(),
251        kind: CompletionItemKind::Snippet,
252        detail: None,
253        insert_text: Some("except $0".to_owned()),
254        insert_text_format: Some(CompletionInsertTextFormat::Snippet),
255        trigger_completion_after_insert: true,
256        sort_text: None,
257    });
258
259    completions
260}
261
262fn limit_completions(db: &dyn Db, file: File, token: &SyntaxToken) -> Vec<CompletionItem> {
263    let schema = schema_qualifier_at_token(token);
264    let position = token.text_range().start();
265
266    let mut completions = vec![CompletionItem {
267        label: "all".to_owned(),
268        kind: CompletionItemKind::Keyword,
269        detail: None,
270        insert_text: None,
271        insert_text_format: None,
272        trigger_completion_after_insert: false,
273        sort_text: None,
274    }];
275
276    completions.extend(function_completions(db, file, schema.as_ref(), position));
277    completions
278}
279
280fn offset_completions(db: &dyn Db, file: File, token: &SyntaxToken) -> Vec<CompletionItem> {
281    let schema = schema_qualifier_at_token(token);
282    let position = token.text_range().start();
283
284    function_completions(db, file, schema.as_ref(), position)
285}
286
287fn select_expr_completions(
288    db: &dyn Db,
289    file: File,
290    select: &ast::Select,
291    token: &SyntaxToken,
292) -> Vec<CompletionItem> {
293    let mut completions = vec![];
294    let schema = schema_qualifier_at_token(token);
295    let position = token.text_range().start();
296
297    completions.extend(function_completions(db, file, schema.as_ref(), position));
298
299    if let Some(from_clause) = select.from_clause() {
300        for from_item in from_clause.from_items() {
301            if let Some(table_name) = table_name_from_from_item(&from_item) {
302                completions.push(CompletionItem {
303                    label: table_name.to_string(),
304                    kind: CompletionItemKind::Table,
305                    detail: None,
306                    insert_text: None,
307                    insert_text_format: None,
308                    trigger_completion_after_insert: false,
309                    sort_text: None,
310                });
311            }
312        }
313
314        completions.extend(column_completions_from_clause(db, file, &from_clause));
315    }
316
317    completions
318}
319
320fn function_completions(
321    db: &dyn Db,
322    file: File,
323    schema: Option<&Schema>,
324    position: TextSize,
325) -> Vec<CompletionItem> {
326    let binder = bind(db, file);
327    binder
328        .all_symbols_by_kind(SymbolKind::Function, schema)
329        .into_iter()
330        .map(|name| CompletionItem {
331            label: format!("{name}()"),
332            kind: CompletionItemKind::Function,
333            detail: function_detail(db, file, name, schema, position),
334            insert_text: None,
335            insert_text_format: None,
336            trigger_completion_after_insert: false,
337            sort_text: None,
338        })
339        .collect()
340}
341
342fn column_completions_from_clause(
343    db: &dyn Db,
344    file: File,
345    from_clause: &ast::FromClause,
346) -> Vec<CompletionItem> {
347    let mut completions = vec![];
348    let syntax_root = from_clause.syntax().ancestors().last().unwrap();
349    for table_ptr in resolve::table_ptrs_from_clause(db, InFile::new(file, from_clause)) {
350        let table_node = table_ptr.to_node(&syntax_root);
351        match ast_nav::parent_source(&table_node) {
352            Some(ast_nav::ParentSouce::CreateTable(create_table)) => {
353                let columns = collect::table_columns(db, file, &create_table);
354                completions.extend(columns.into_iter().map(|(name, ty)| CompletionItem {
355                    label: name.to_string(),
356                    kind: CompletionItemKind::Column,
357                    detail: ty.map(|t| t.to_string()),
358                    insert_text: None,
359                    insert_text_format: None,
360                    trigger_completion_after_insert: false,
361                    sort_text: None,
362                }));
363            }
364            Some(ast_nav::ParentSouce::WithTable(with_table)) => {
365                let columns = collect::with_table_columns_with_types(db, file, with_table);
366                completions.extend(columns.into_iter().map(|(name, ty)| CompletionItem {
367                    label: name.to_string(),
368                    kind: CompletionItemKind::Column,
369                    detail: ty.map(|t| t.to_string()),
370                    insert_text: None,
371                    insert_text_format: None,
372                    trigger_completion_after_insert: false,
373                    sort_text: None,
374                }));
375            }
376            Some(ast_nav::ParentSouce::CreateTableAs(create_table_as)) => {
377                let columns =
378                    collect::create_table_as_columns_with_types(db, file, &create_table_as);
379                completions.extend(columns.into_iter().map(|(name, ty)| CompletionItem {
380                    label: name.to_string(),
381                    kind: CompletionItemKind::Column,
382                    detail: ty.map(|t| t.to_string()),
383                    insert_text: None,
384                    insert_text_format: None,
385                    trigger_completion_after_insert: false,
386                    sort_text: None,
387                }));
388            }
389            Some(ast_nav::ParentSouce::CreateView(create_view)) => {
390                let columns = collect::view_like_columns_with_types(db, file, &create_view);
391                completions.extend(columns.into_iter().map(|(name, ty)| CompletionItem {
392                    label: name.to_string(),
393                    kind: CompletionItemKind::Column,
394                    detail: ty.map(|t| t.to_string()),
395                    insert_text: None,
396                    insert_text_format: None,
397                    trigger_completion_after_insert: false,
398                    sort_text: None,
399                }));
400            }
401            Some(ast_nav::ParentSouce::Alias(alias)) => {
402                let alias_columns: Vec<Name> = alias
403                    .column_list()
404                    .into_iter()
405                    .flat_map(|column_list| column_list.columns())
406                    .filter_map(|column| column.name().map(|name| Name::from_node(&name)))
407                    .collect();
408
409                let base_columns = alias_base_columns_with_types(db, file, &syntax_root, &alias);
410
411                for (idx, alias_column) in alias_columns.iter().enumerate() {
412                    completions.push(CompletionItem {
413                        label: alias_column.to_string(),
414                        kind: CompletionItemKind::Column,
415                        detail: base_columns.get(idx).and_then(|(_, ty)| ty.clone()),
416                        insert_text: None,
417                        insert_text_format: None,
418                        trigger_completion_after_insert: false,
419                        sort_text: Some(format!("{idx:04}")),
420                    });
421                }
422
423                completions.extend(
424                    base_columns
425                        .into_iter()
426                        .skip(alias_columns.len())
427                        .enumerate()
428                        .map(|(idx, (name, ty))| CompletionItem {
429                            label: name.to_string(),
430                            kind: CompletionItemKind::Column,
431                            detail: ty,
432                            insert_text: None,
433                            insert_text_format: None,
434                            trigger_completion_after_insert: false,
435                            sort_text: Some(format!("{:04}", idx + alias_columns.len())),
436                        }),
437                );
438            }
439            Some(ast_nav::ParentSouce::ParenSelect(paren_select)) => {
440                let columns = collect::paren_select_columns_with_types(db, file, &paren_select);
441                completions.extend(columns.into_iter().map(|(name, ty)| CompletionItem {
442                    label: name.to_string(),
443                    kind: CompletionItemKind::Column,
444                    detail: ty.map(|t| t.to_string()),
445                    insert_text: None,
446                    insert_text_format: None,
447                    trigger_completion_after_insert: false,
448                    sort_text: None,
449                }));
450            }
451            None => {}
452        }
453    }
454    completions
455}
456
457fn alias_base_columns_with_types(
458    db: &dyn Db,
459    file: File,
460    syntax_root: &SyntaxNode,
461    alias: &ast::Alias,
462) -> Vec<(Name, Option<String>)> {
463    let Some(from_item) = alias.syntax().ancestors().find_map(ast::FromItem::cast) else {
464        return vec![];
465    };
466    let Some(table_ptr) = resolve::table_ptr_from_from_item(db, InFile::new(file, &from_item))
467    else {
468        return vec![];
469    };
470
471    let table_node = table_ptr.to_node(syntax_root);
472
473    match ast_nav::parent_source(&table_node) {
474        Some(ast_nav::ParentSouce::CreateTable(create_table)) => {
475            collect::table_columns(db, file, &create_table)
476                .into_iter()
477                .map(|(name, ty)| (name, ty.map(|t| t.to_string())))
478                .collect()
479        }
480        Some(ast_nav::ParentSouce::WithTable(with_table)) => {
481            collect::with_table_columns_with_types(db, file, with_table)
482                .into_iter()
483                .map(|(name, ty)| (name, ty.map(|t| t.to_string())))
484                .collect()
485        }
486        Some(ast_nav::ParentSouce::CreateView(create_view)) => {
487            collect::view_like_columns_with_types(db, file, &create_view)
488                .into_iter()
489                .map(|(name, ty)| (name, ty.map(|t| t.to_string())))
490                .collect()
491        }
492        Some(ast_nav::ParentSouce::ParenSelect(paren_select)) => {
493            collect::paren_select_columns_with_types(db, file, &paren_select)
494                .into_iter()
495                .map(|(name, ty)| (name, ty.map(|t| t.to_string())))
496                .collect()
497        }
498        Some(ast_nav::ParentSouce::CreateTableAs(create_table_as)) => {
499            collect::create_table_as_columns_with_types(db, file, &create_table_as)
500                .into_iter()
501                .map(|(name, ty)| (name, ty.map(|t| t.to_string())))
502                .collect()
503        }
504        Some(ast_nav::ParentSouce::Alias(_)) | None => vec![],
505    }
506}
507
508fn schema_completions(binder: &binder::Binder) -> Vec<CompletionItem> {
509    let builtin_schemas = [
510        "public",
511        "pg_catalog",
512        "pg_temp",
513        "pg_toast",
514        "information_schema",
515    ];
516    let mut completions: Vec<CompletionItem> = builtin_schemas
517        .into_iter()
518        .enumerate()
519        .map(|(i, name)| CompletionItem {
520            label: name.to_string(),
521            kind: CompletionItemKind::Schema,
522            detail: None,
523            insert_text: None,
524            insert_text_format: None,
525            trigger_completion_after_insert: false,
526            sort_text: Some(format!("{i}")),
527        })
528        .collect();
529
530    for name in binder.all_symbols_by_kind(SymbolKind::Schema, None) {
531        completions.push(CompletionItem {
532            label: name.to_string(),
533            kind: CompletionItemKind::Schema,
534            detail: None,
535            insert_text: None,
536            insert_text_format: None,
537            trigger_completion_after_insert: false,
538            sort_text: None,
539        });
540    }
541
542    completions
543}
544
545fn table_completions(file: &ast::SourceFile, token: &SyntaxToken) -> Vec<CompletionItem> {
546    let binder = binder::bind(file);
547    let schema = schema_qualifier_at_token(token);
548    let tables = binder.all_symbols_by_kind(SymbolKind::Table, schema.as_ref());
549    let mut completions: Vec<CompletionItem> = tables
550        .into_iter()
551        .map(|name| CompletionItem {
552            label: name.to_string(),
553            kind: CompletionItemKind::Table,
554            detail: None,
555            insert_text: None,
556            insert_text_format: None,
557            trigger_completion_after_insert: false,
558            sort_text: None,
559        })
560        .collect();
561
562    if schema.is_none() {
563        completions.extend(schema_completions(&binder));
564    }
565
566    completions
567}
568
569fn delete_clauses_completions(
570    file: &ast::SourceFile,
571    delete: &ast::Delete,
572    token: &SyntaxToken,
573) -> Vec<CompletionItem> {
574    let mut completions = vec![];
575
576    // `delete from $0`
577    if token.kind() == SyntaxKind::FROM_KW {
578        return table_completions(file, token);
579    }
580
581    if delete.using_clause().is_none() {
582        completions.push(CompletionItem {
583            label: "using".to_owned(),
584            kind: CompletionItemKind::Snippet,
585            detail: None,
586            insert_text: Some("using $0".to_owned()),
587            insert_text_format: Some(CompletionInsertTextFormat::Snippet),
588            trigger_completion_after_insert: true,
589            sort_text: None,
590        });
591    }
592
593    if delete.where_clause().is_none() {
594        completions.push(CompletionItem {
595            label: "where".to_owned(),
596            kind: CompletionItemKind::Snippet,
597            detail: None,
598            insert_text: Some("where $0".to_owned()),
599            insert_text_format: Some(CompletionInsertTextFormat::Snippet),
600            trigger_completion_after_insert: true,
601            sort_text: None,
602        });
603    }
604
605    if delete.returning_clause().is_none() {
606        completions.push(CompletionItem {
607            label: "returning".to_owned(),
608            kind: CompletionItemKind::Snippet,
609            detail: None,
610            insert_text: Some("returning $0".to_owned()),
611            insert_text_format: Some(CompletionInsertTextFormat::Snippet),
612            trigger_completion_after_insert: true,
613            sort_text: None,
614        });
615    }
616
617    completions
618}
619
620fn delete_expr_completions(
621    db: &dyn Db,
622    file: File,
623    source_file: &ast::SourceFile,
624    delete: &ast::Delete,
625    token: &SyntaxToken,
626) -> Vec<CompletionItem> {
627    let binder = binder::bind(source_file);
628    let mut completions = vec![];
629
630    let Some(path) = delete.relation_name().and_then(|r| r.path()) else {
631        return completions;
632    };
633
634    let Some(delete_table_name) = name::table_name(&path) else {
635        return completions;
636    };
637
638    let has_table_qualifier = qualifier_at_token(token).is_some_and(|q| q == delete_table_name);
639    let schema = schema_qualifier_at_token(token);
640    let position = token.text_range().start();
641
642    if has_table_qualifier {
643        let functions = binder.functions_with_single_param(&delete_table_name);
644        completions.extend(functions.into_iter().map(|name| CompletionItem {
645            label: name.to_string(),
646            kind: CompletionItemKind::Function,
647            detail: function_detail(db, file, name, schema.as_ref(), position),
648            insert_text: None,
649            insert_text_format: None,
650            trigger_completion_after_insert: false,
651            sort_text: None,
652        }));
653    } else {
654        let functions = binder.all_symbols_by_kind(SymbolKind::Function, None);
655        completions.extend(functions.into_iter().map(|name| CompletionItem {
656            label: format!("{name}()"),
657            kind: CompletionItemKind::Function,
658            detail: function_detail(db, file, name, schema.as_ref(), position),
659            insert_text: None,
660            insert_text_format: None,
661            trigger_completion_after_insert: false,
662            sort_text: None,
663        }));
664
665        completions.push(CompletionItem {
666            label: delete_table_name.to_string(),
667            kind: CompletionItemKind::Table,
668            detail: None,
669            insert_text: None,
670            insert_text_format: None,
671            trigger_completion_after_insert: false,
672            sort_text: None,
673        });
674    }
675
676    let schema = name::schema_name(&path);
677    let schemas = binder.resolved_schemas(position, schema.as_ref());
678    if let Some(table_ptr) = binder.lookup_with(&delete_table_name, SymbolKind::Table, &schemas)
679        && let Some(create_table) = table_ptr
680            .to_node(source_file.syntax())
681            .ancestors()
682            .find_map(ast::CreateTableLike::cast)
683    {
684        let columns = collect::table_columns(db, file, &create_table);
685        completions.extend(columns.into_iter().map(|(name, ty)| CompletionItem {
686            label: name.to_string(),
687            kind: CompletionItemKind::Column,
688            detail: ty.map(|t| t.to_string()),
689            insert_text: None,
690            insert_text_format: None,
691            trigger_completion_after_insert: false,
692            sort_text: None,
693        }));
694    }
695
696    completions
697}
698
699fn table_name_from_from_item(from_item: &ast::FromItem) -> Option<Name> {
700    if let Some(alias) = from_item.alias()
701        && let Some(alias_name) = alias.name()
702    {
703        return Some(Name::from_node(&alias_name));
704    }
705    if let Some(name_ref) = from_item.name_ref() {
706        return Some(Name::from_node(&name_ref));
707    }
708    None
709}
710
711fn qualifier_at_token(token: &SyntaxToken) -> Option<Name> {
712    let qualifier_token = if token.kind() == SyntaxKind::DOT {
713        token.prev_token()
714    } else if token.kind() == SyntaxKind::IDENT
715        && let Some(prev) = token.prev_token()
716        && prev.kind() == SyntaxKind::DOT
717    {
718        prev.prev_token()
719    } else {
720        None
721    };
722
723    qualifier_token
724        .filter(|tk| tk.kind() == SyntaxKind::IDENT)
725        .map(|tk| Name::from_string(tk.text().to_string()))
726}
727
728#[derive(Debug)]
729enum CompletionContext {
730    TableOnly,
731    Default,
732    SelectClause(ast::SelectClause),
733    SelectClauses(ast::Select),
734    SelectExpr(ast::Select),
735    LimitClause,
736    OffsetClause,
737    DeleteClauses(ast::Delete),
738    DeleteExpr(ast::Delete),
739}
740
741fn completion_context(token: &SyntaxToken) -> CompletionContext {
742    if let Some(node) = token.parent() {
743        let mut inside_delete_clause = false;
744        let mut inside_from_item = false;
745        let mut inside_paren_expr = false;
746        let mut inside_select_expr_clause = false;
747        let mut inside_limit_clause = false;
748        let mut inside_offset_clause = false;
749        for a in node.ancestors() {
750            if ast::Truncate::can_cast(a.kind()) || ast::Table::can_cast(a.kind()) {
751                return CompletionContext::TableOnly;
752            }
753            if ast::WhereClause::can_cast(a.kind())
754                || ast::UsingClause::can_cast(a.kind())
755                || ast::ReturningClause::can_cast(a.kind())
756            {
757                inside_delete_clause = true;
758            }
759            if ast::LimitClause::can_cast(a.kind()) {
760                inside_limit_clause = true;
761            }
762            if ast::OffsetClause::can_cast(a.kind()) {
763                inside_offset_clause = true;
764            }
765            if ast::WhereClause::can_cast(a.kind())
766                || ast::GroupByClause::can_cast(a.kind())
767                || ast::HavingClause::can_cast(a.kind())
768                || ast::OrderByClause::can_cast(a.kind())
769            {
770                inside_select_expr_clause = true;
771            }
772            if ast::FromItem::can_cast(a.kind()) {
773                inside_from_item = true;
774            }
775            if ast::ParenExpr::can_cast(a.kind()) {
776                inside_paren_expr = true;
777            }
778            if let Some(delete) = ast::Delete::cast(a.clone()) {
779                if inside_delete_clause {
780                    return CompletionContext::DeleteExpr(delete);
781                }
782                if delete.relation_name().is_some() {
783                    return CompletionContext::DeleteClauses(delete);
784                }
785                return CompletionContext::TableOnly;
786            }
787            if let Some(select) = ast::Select::cast(a.clone()) {
788                if inside_limit_clause {
789                    return CompletionContext::LimitClause;
790                }
791                if inside_offset_clause {
792                    return CompletionContext::OffsetClause;
793                }
794                if inside_select_expr_clause {
795                    return CompletionContext::SelectExpr(select);
796                }
797                if inside_from_item && !inside_paren_expr && select.from_clause().is_some() {
798                    return CompletionContext::SelectClauses(select);
799                }
800            }
801            if let Some(select_clause) = ast::SelectClause::cast(a.clone()) {
802                return CompletionContext::SelectClause(select_clause);
803            }
804        }
805    }
806    CompletionContext::Default
807}
808
809fn token_at_offset(file: &ast::SourceFile, offset: TextSize) -> Option<SyntaxToken> {
810    let Some(mut token) = file.syntax().token_at_offset(offset).left_biased() else {
811        // empty file - definitely at top level
812        return None;
813    };
814    while token.kind() == SyntaxKind::WHITESPACE {
815        if let Some(tk) = token.prev_token() {
816            token = tk;
817        }
818    }
819    Some(token)
820}
821
822// In order to make completions, we do something similar to rust analyzer by
823// inserting an ident to make the parse tree parse in more cases.
824// Rust analyzer does fancier things for this, which we can investigate later.
825//
826// This helps us support `select t. from t`, which parses as `select t.from t`.
827// If we insert the ident we get, `select t.c from t`.
828fn file_with_completion_marker(file: &ast::SourceFile, offset: TextSize) -> ast::SourceFile {
829    let mut sql = file.syntax().text().to_string();
830    let offset = u32::from(offset) as usize;
831    let offset = offset.min(sql.len());
832    sql.insert_str(offset, COMPLETION_MARKER);
833    ast::SourceFile::parse(&sql).tree()
834}
835
836fn schema_qualifier_at_token(token: &SyntaxToken) -> Option<Schema> {
837    qualifier_at_token(token).map(Schema)
838}
839
840fn function_detail(
841    db: &dyn Db,
842    file: File,
843    function_name: &Name,
844    schema: Option<&Schema>,
845    position: TextSize,
846) -> Option<String> {
847    let binder = bind(db, file);
848    let source_file = parse(db, file).tree();
849    let schemas = binder.resolved_schemas(position, schema);
850    let create_function = binder
851        .lookup_with(function_name, SymbolKind::Function, &schemas)?
852        .to_node(source_file.syntax())
853        .ancestors()
854        .find_map(ast::CreateFunction::cast)?;
855    let path = create_function.path()?;
856    let (schema, function_name) = resolve::resolve_function_info(db, InFile::new(file, &path))?;
857
858    let param_list = create_function.param_list()?;
859    let params = param_list.syntax().text().to_string();
860
861    let ret_type = create_function.ret_type()?;
862    let return_type = ret_type.syntax().text().to_string();
863
864    Some(format!("{schema}.{function_name}{params} {return_type}"))
865}
866
867fn default_completions() -> Vec<CompletionItem> {
868    ["delete from", "select", "table", "truncate"]
869        .map(|stmt| CompletionItem {
870            label: stmt.to_owned(),
871            kind: CompletionItemKind::Snippet,
872            detail: None,
873            insert_text: Some(format!("{stmt} $0;")),
874            insert_text_format: Some(CompletionInsertTextFormat::Snippet),
875            trigger_completion_after_insert: true,
876            sort_text: None,
877        })
878        .into_iter()
879        .collect()
880}
881
882#[derive(Debug, Clone, Copy, PartialEq, Eq)]
883pub enum CompletionItemKind {
884    Keyword,
885    Table,
886    Column,
887    Function,
888    Schema,
889    Type,
890    Snippet,
891    Operator,
892}
893
894impl CompletionItemKind {
895    fn sort_prefix(self) -> &'static str {
896        match self {
897            Self::Column => "0",
898            Self::Keyword => "1",
899            Self::Table => "1",
900            Self::Type => "1",
901            Self::Snippet => "1",
902            Self::Function => "2",
903            Self::Operator => "8",
904            Self::Schema => "9",
905        }
906    }
907}
908
909impl CompletionItem {
910    pub fn sort_text(&self) -> String {
911        let prefix = self.kind.sort_prefix();
912        let suffix = self.sort_text.as_ref().unwrap_or(&self.label);
913        format!("{prefix}_{suffix}")
914    }
915}
916
917#[derive(Debug, Clone, Copy, PartialEq, Eq)]
918pub enum CompletionInsertTextFormat {
919    PlainText,
920    Snippet,
921}
922
923#[derive(Debug, Clone, PartialEq, Eq)]
924pub struct CompletionItem {
925    pub label: String,
926    pub kind: CompletionItemKind,
927    pub detail: Option<String>,
928    pub insert_text: Option<String>,
929    pub insert_text_format: Option<CompletionInsertTextFormat>,
930    pub trigger_completion_after_insert: bool,
931    pub sort_text: Option<String>,
932}
933
934#[cfg(test)]
935mod tests {
936    use super::completion;
937
938    use crate::test_utils::Fixture;
939    use insta::assert_snapshot;
940    use tabled::builder::Builder;
941    use tabled::settings::Style;
942
943    #[must_use]
944    fn completions(sql: &str) -> String {
945        let fixture = Fixture::new_allow_errors(sql);
946        let offset = fixture.marker().offset();
947        let items = completion(fixture.db(), offset);
948        assert!(
949            !items.is_empty(),
950            "No completions found. If this was intended, use `completions_not_found` instead."
951        );
952        format_items(items)
953    }
954
955    fn completions_not_found(sql: &str) {
956        let fixture = Fixture::new_allow_errors(sql);
957        let offset = fixture.marker().offset();
958        let items = completion(fixture.db(), offset);
959        assert_eq!(
960            items,
961            vec![],
962            "Completions found. If this was unintended, use `completions` instead."
963        )
964    }
965
966    fn format_items(mut items: Vec<super::CompletionItem>) -> String {
967        items.sort_by_key(|a| a.sort_text());
968
969        let rows: Vec<Vec<String>> = items
970            .into_iter()
971            .map(|item| {
972                vec![
973                    item.label,
974                    format!("{:?}", item.kind),
975                    item.detail.unwrap_or_default(),
976                ]
977            })
978            .collect();
979
980        let mut builder = Builder::default();
981        builder.push_record(["label", "kind", "detail"]);
982        for row in rows {
983            builder.push_record(row);
984        }
985
986        let mut table = builder.build();
987        table.with(Style::psql());
988        table.to_string()
989    }
990
991    #[test]
992    fn completion_at_start() {
993        assert_snapshot!(completions("$0"), @r"
994         label       | kind    | detail 
995        -------------+---------+--------
996         delete from | Snippet |        
997         select      | Snippet |        
998         table       | Snippet |        
999         truncate    | Snippet |
1000        ");
1001    }
1002
1003    #[test]
1004    fn completion_at_top_level() {
1005        assert_snapshot!(completions("
1006create table t(a int);
1007$0
1008"), @r"
1009         label       | kind    | detail 
1010        -------------+---------+--------
1011         delete from | Snippet |        
1012         select      | Snippet |        
1013         table       | Snippet |        
1014         truncate    | Snippet |
1015        ");
1016    }
1017
1018    #[test]
1019    fn completion_in_string() {
1020        completions_not_found("select '$0';");
1021    }
1022
1023    #[test]
1024    fn completion_in_comment() {
1025        completions_not_found("-- $0 ");
1026    }
1027
1028    #[test]
1029    fn completion_after_truncate() {
1030        assert_snapshot!(completions("
1031create table users (id int);
1032truncate $0;
1033"), @r"
1034         label              | kind   | detail 
1035        --------------------+--------+--------
1036         users              | Table  |        
1037         public             | Schema |        
1038         pg_catalog         | Schema |        
1039         pg_temp            | Schema |        
1040         pg_toast           | Schema |        
1041         information_schema | Schema |
1042        ");
1043    }
1044
1045    #[test]
1046    fn completion_table_at_top_level() {
1047        assert_snapshot!(completions("$0"), @r"
1048         label       | kind    | detail 
1049        -------------+---------+--------
1050         delete from | Snippet |        
1051         select      | Snippet |        
1052         table       | Snippet |        
1053         truncate    | Snippet |
1054        ");
1055    }
1056
1057    #[test]
1058    fn completion_table_nested() {
1059        assert_snapshot!(completions("select * from ($0)"), @r"
1060         label       | kind    | detail 
1061        -------------+---------+--------
1062         delete from | Snippet |        
1063         select      | Snippet |        
1064         table       | Snippet |        
1065         truncate    | Snippet |
1066        ");
1067    }
1068
1069    #[test]
1070    fn completion_after_table() {
1071        assert_snapshot!(completions("
1072create table users (id int);
1073table $0;
1074"), @r"
1075         label              | kind   | detail 
1076        --------------------+--------+--------
1077         users              | Table  |        
1078         public             | Schema |        
1079         pg_catalog         | Schema |        
1080         pg_temp            | Schema |        
1081         pg_toast           | Schema |        
1082         information_schema | Schema |
1083        ");
1084    }
1085
1086    #[test]
1087    fn completion_select_without_from() {
1088        assert_snapshot!(completions("
1089create table t (a int);
1090select $0;
1091"), @r"
1092         label              | kind    | detail 
1093        --------------------+---------+--------
1094         except             | Snippet |        
1095         fetch              | Snippet |        
1096         for                | Snippet |        
1097         from               | Snippet |        
1098         group by           | Snippet |        
1099         having             | Snippet |        
1100         intersect          | Snippet |        
1101         limit              | Snippet |        
1102         offset             | Snippet |        
1103         order by           | Snippet |        
1104         t                  | Table   |        
1105         union              | Snippet |        
1106         where              | Snippet |        
1107         window             | Snippet |        
1108         public             | Schema  |        
1109         pg_catalog         | Schema  |        
1110         pg_temp            | Schema  |        
1111         pg_toast           | Schema  |        
1112         information_schema | Schema  |
1113        ");
1114    }
1115
1116    #[test]
1117    fn completion_after_select() {
1118        assert_snapshot!(completions("
1119create table t(a text, b int);
1120create function f() returns text as 'select 1::text' language sql;
1121select $0 from t;
1122"), @r"
1123         label              | kind     | detail                  
1124        --------------------+----------+-------------------------
1125         a                  | Column   | text                    
1126         b                  | Column   | int                     
1127         t                  | Table    |                         
1128         f()                | Function | public.f() returns text 
1129         *                  | Operator |                         
1130         public             | Schema   |                         
1131         pg_catalog         | Schema   |                         
1132         pg_temp            | Schema   |                         
1133         pg_toast           | Schema   |                         
1134         information_schema | Schema   |
1135        ");
1136    }
1137
1138    #[test]
1139    fn completion_after_select_create_table_inherits_builtin() {
1140        assert_snapshot!(completions("
1141-- include-builtins
1142create table t ()
1143inherits (information_schema.sql_features);
1144select $0 from t;
1145"), @"
1146         label              | kind     | detail         
1147        --------------------+----------+----------------
1148         comments           | Column   | character_data 
1149         feature_id         | Column   | character_data 
1150         feature_name       | Column   | character_data 
1151         is_supported       | Column   | yes_or_no      
1152         is_verified_by     | Column   | character_data 
1153         sub_feature_id     | Column   | character_data 
1154         sub_feature_name   | Column   | character_data 
1155         t                  | Table    |                
1156         *                  | Operator |                
1157         public             | Schema   |                
1158         pg_catalog         | Schema   |                
1159         pg_temp            | Schema   |                
1160         pg_toast           | Schema   |                
1161         information_schema | Schema   |
1162        ");
1163    }
1164
1165    #[test]
1166    fn completion_after_select_create_table_inherits_create_table_as() {
1167        assert_snapshot!(completions("
1168create table parent as select 1 a, 'x'::text b;
1169create table child (c int) inherits (parent);
1170select $0 from child;
1171"), @"
1172         label              | kind     | detail  
1173        --------------------+----------+---------
1174         a                  | Column   | integer 
1175         b                  | Column   | text    
1176         c                  | Column   | int     
1177         child              | Table    |         
1178         parent             | Table    |         
1179         *                  | Operator |         
1180         public             | Schema   |         
1181         pg_catalog         | Schema   |         
1182         pg_temp            | Schema   |         
1183         pg_toast           | Schema   |         
1184         information_schema | Schema   |
1185        ");
1186    }
1187
1188    #[test]
1189    fn completion_after_select_create_table_like_select_into() {
1190        assert_snapshot!(completions("
1191select 1 a, 'x'::text b into parent;
1192create table child (like parent);
1193select $0 from child;
1194"), @"
1195         label              | kind     | detail  
1196        --------------------+----------+---------
1197         a                  | Column   | integer 
1198         b                  | Column   | text    
1199         child              | Table    |         
1200         parent             | Table    |         
1201         *                  | Operator |         
1202         public             | Schema   |         
1203         pg_catalog         | Schema   |         
1204         pg_temp            | Schema   |         
1205         pg_toast           | Schema   |         
1206         information_schema | Schema   |
1207        ");
1208    }
1209
1210    #[test]
1211    fn completion_select_table_qualified() {
1212        assert_snapshot!(completions("
1213create table t (c int);
1214select t.$0 from t;
1215"), @r"
1216         label | kind     | detail 
1217        -------+----------+--------
1218         c     | Column   | int    
1219         *     | Operator |
1220        ");
1221    }
1222
1223    #[test]
1224    fn completion_after_select_with_cte() {
1225        assert_snapshot!(completions("
1226with t as (select 1 a)
1227select $0 from t;
1228"), @r"
1229         label              | kind     | detail  
1230        --------------------+----------+---------
1231         a                  | Column   | integer 
1232         *                  | Operator |         
1233         public             | Schema   |         
1234         pg_catalog         | Schema   |         
1235         pg_temp            | Schema   |         
1236         pg_toast           | Schema   |         
1237         information_schema | Schema   |
1238        ");
1239    }
1240
1241    #[test]
1242    fn completion_after_select_with_cte_alias_column_list() {
1243        assert_snapshot!(completions("
1244with t as (select 1 a, 2 b, 3 c)
1245select $0 from t as u(x, y);
1246"), @r"
1247         label              | kind     | detail  
1248        --------------------+----------+---------
1249         x                  | Column   | integer 
1250         y                  | Column   | integer 
1251         c                  | Column   | integer 
1252         *                  | Operator |         
1253         public             | Schema   |         
1254         pg_catalog         | Schema   |         
1255         pg_temp            | Schema   |         
1256         pg_toast           | Schema   |         
1257         information_schema | Schema   |
1258        ");
1259    }
1260
1261    #[test]
1262    fn completion_values_cte() {
1263        assert_snapshot!(completions("
1264with t as (values (1, 'foo', false))
1265select $0 from t;
1266"), @r"
1267         label              | kind     | detail  
1268        --------------------+----------+---------
1269         column1            | Column   | integer 
1270         column2            | Column   | text    
1271         column3            | Column   | boolean 
1272         *                  | Operator |         
1273         public             | Schema   |         
1274         pg_catalog         | Schema   |         
1275         pg_temp            | Schema   |         
1276         pg_toast           | Schema   |         
1277         information_schema | Schema   |
1278        ");
1279    }
1280
1281    #[test]
1282    fn completion_values_subquery() {
1283        assert_snapshot!(completions("
1284select $0 from (values (1, 'foo', 1.5, false));
1285"), @r"
1286         label              | kind     | detail  
1287        --------------------+----------+---------
1288         column1            | Column   | integer 
1289         column2            | Column   | text    
1290         column3            | Column   | numeric 
1291         column4            | Column   | boolean 
1292         *                  | Operator |         
1293         public             | Schema   |         
1294         pg_catalog         | Schema   |         
1295         pg_temp            | Schema   |         
1296         pg_toast           | Schema   |         
1297         information_schema | Schema   |
1298        ");
1299    }
1300
1301    #[test]
1302    fn completion_with_schema_qualifier() {
1303        assert_snapshot!(completions("
1304create function f() returns int8 as 'select 1' language sql;
1305create function foo.b() returns int8 as 'select 2' language sql;
1306select public.$0;
1307"), @r"
1308         label | kind     | detail                  
1309        -------+----------+-------------------------
1310         f()   | Function | public.f() returns int8
1311        ");
1312    }
1313
1314    #[test]
1315    fn completion_truncate_with_schema_qualifier() {
1316        assert_snapshot!(completions("
1317create table users (id int);
1318truncate public.$0;
1319"), @r"
1320         label | kind  | detail 
1321        -------+-------+--------
1322         users | Table |
1323        ");
1324    }
1325
1326    #[test]
1327    fn completion_after_delete_from() {
1328        assert_snapshot!(completions("
1329create table users (id int);
1330delete from $0;
1331"), @r"
1332         label              | kind   | detail 
1333        --------------------+--------+--------
1334         users              | Table  |        
1335         public             | Schema |        
1336         pg_catalog         | Schema |        
1337         pg_temp            | Schema |        
1338         pg_toast           | Schema |        
1339         information_schema | Schema |
1340        ");
1341    }
1342
1343    #[test]
1344    fn completion_delete_clauses() {
1345        assert_snapshot!(completions("
1346create table t (id int);
1347delete from t $0;
1348"), @r"
1349         label     | kind    | detail 
1350        -----------+---------+--------
1351         returning | Snippet |        
1352         using     | Snippet |        
1353         where     | Snippet |
1354        ");
1355    }
1356
1357    #[test]
1358    fn completion_delete_where_expr() {
1359        assert_snapshot!(completions("
1360create table t (id int, name text);
1361create function is_active() returns bool as 'select true' language sql;
1362delete from t where $0;
1363"), @r"
1364         label       | kind     | detail                          
1365        -------------+----------+---------------------------------
1366         id          | Column   | int                             
1367         name        | Column   | text                            
1368         t           | Table    |                                 
1369         is_active() | Function | public.is_active() returns bool
1370        ")
1371    }
1372
1373    #[test]
1374    fn completion_delete_returning_expr() {
1375        assert_snapshot!(completions("
1376create table t (id int, name text);
1377delete from t returning $0;
1378"), @r"
1379         label | kind   | detail 
1380        -------+--------+--------
1381         id    | Column | int    
1382         name  | Column | text   
1383         t     | Table  |
1384        ");
1385    }
1386
1387    #[test]
1388    fn completion_delete_where_qualified() {
1389        assert_snapshot!(completions("
1390-- different type than the table, so we shouldn't show this
1391create function b(diff_type) returns int8
1392  as 'select 1'
1393  language sql;
1394create function f(t) returns int8
1395  as 'select 1'
1396  language sql;
1397create table t (a int, b text);
1398delete from t where t.$0;
1399"), @r"
1400         label | kind     | detail 
1401        -------+----------+--------
1402         a     | Column   | int    
1403         b     | Column   | text   
1404         f     | Function |
1405        ");
1406    }
1407
1408    #[test]
1409    fn completion_select_clauses() {
1410        assert_snapshot!(completions("
1411with t as (select 1 a)
1412select a from t $0;
1413"), @r"
1414         label     | kind    | detail 
1415        -----------+---------+--------
1416         except    | Snippet |        
1417         fetch     | Snippet |        
1418         for       | Snippet |        
1419         group by  | Snippet |        
1420         having    | Snippet |        
1421         intersect | Snippet |        
1422         limit     | Snippet |        
1423         offset    | Snippet |        
1424         order by  | Snippet |        
1425         union     | Snippet |        
1426         where     | Snippet |        
1427         window    | Snippet |
1428        ");
1429    }
1430
1431    #[test]
1432    fn completion_select_clauses_simple() {
1433        assert_snapshot!(completions("
1434select 1 from t $0;
1435"), @r"
1436         label     | kind    | detail 
1437        -----------+---------+--------
1438         except    | Snippet |        
1439         fetch     | Snippet |        
1440         for       | Snippet |        
1441         group by  | Snippet |        
1442         having    | Snippet |        
1443         intersect | Snippet |        
1444         limit     | Snippet |        
1445         offset    | Snippet |        
1446         order by  | Snippet |        
1447         union     | Snippet |        
1448         where     | Snippet |        
1449         window    | Snippet |
1450        ");
1451    }
1452
1453    #[test]
1454    fn completion_select_group_by_expr() {
1455        assert_snapshot!(completions("
1456with t as (select 1 a)
1457select a from t group by $0;
1458"), @r"
1459         label | kind   | detail  
1460        -------+--------+---------
1461         a     | Column | integer 
1462         t     | Table  |
1463        ");
1464    }
1465
1466    #[test]
1467    fn completion_select_where_expr() {
1468        assert_snapshot!(completions("
1469create table t (id int, name text);
1470select * from t where $0;
1471"), @r"
1472         label | kind   | detail 
1473        -------+--------+--------
1474         id    | Column | int    
1475         name  | Column | text   
1476         t     | Table  |
1477        ");
1478    }
1479
1480    #[test]
1481    fn completion_select_limit() {
1482        assert_snapshot!(completions("
1483create function get_limit() returns int as 'select 10' language sql;
1484select 1 from t limit $0;
1485"), @r"
1486         label       | kind     | detail                         
1487        -------------+----------+--------------------------------
1488         all         | Keyword  |                                
1489         get_limit() | Function | public.get_limit() returns int
1490        ");
1491    }
1492
1493    #[test]
1494    fn completion_select_offset() {
1495        assert_snapshot!(completions("
1496create function get_offset() returns int as 'select 10' language sql;
1497select 1 from t offset $0;
1498"), @r"
1499         label        | kind     | detail                          
1500        --------------+----------+---------------------------------
1501         get_offset() | Function | public.get_offset() returns int
1502        ");
1503    }
1504}