Skip to content

Commit 1086965

Browse files
sbdchdsbdchd
committed
ide: support set search_path 
1 parent 1850c99 commit 1086965

File tree

10 files changed

+394
-27
lines changed

10 files changed

+394
-27
lines changed

‎crates/squawk_ide/src/binder.rs‎

Lines changed: 103 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,21 @@
11
/// Loosely based on TypeScript's binder
22
/// see: typescript-go/internal/binder/binder.go
33
use la_arena::Arena;
4+
use rowan::TextSize;
45
use squawk_syntax::{SyntaxNodePtr, ast, ast::AstNode};
56

67
use crate::scope::{Scope, ScopeId};
78
use crate::symbols::{Name, Schema, Symbol, SymbolKind};
89

10+
pub(crate) struct SearchPathChange {
11+
position: TextSize,
12+
search_path: Vec<Schema>,
13+
}
14+
915
pub(crate) struct Binder {
1016
pub(crate) scopes: Arena<Scope>,
1117
pub(crate) symbols: Arena<Symbol>,
18+
pub(crate) search_path_changes: Vec<SearchPathChange>,
1219
}
1320

1421
impl Binder {
@@ -18,6 +25,10 @@ impl Binder {
1825
Binder {
1926
scopes,
2027
symbols: Arena::new(),
28+
search_path_changes: vec![SearchPathChange {
29+
position: TextSize::from(0),
30+
search_path: vec![Schema::new("pg_temp"), Schema::new("public")],
31+
}],
2132
}
2233
}
2334

@@ -28,6 +39,18 @@ impl Binder {
2839
.map(|(id, _)| id)
2940
.expect("root scope must exist")
3041
}
42+
43+
pub(crate) fn search_path_at(&self, position: TextSize) -> &[Schema] {
44+
// We're assuming people don't actually use `set search_path` that much,
45+
// so linear search is fine
46+
for change in self.search_path_changes.iter().rev() {
47+
if change.position <= position {
48+
return &change.search_path;
49+
}
50+
}
51+
// default search path
52+
&self.search_path_changes[0].search_path
53+
}
3154
}
3255

3356
pub(crate) fn bind(file: &ast::SourceFile) -> Binder {
@@ -45,8 +68,10 @@ fn bind_file(b: &mut Binder, file: &ast::SourceFile) {
4568
}
4669

4770
fn bind_stmt(b: &mut Binder, stmt: ast::Stmt) {
48-
if let ast::Stmt::CreateTable(create_table) = stmt {
49-
bind_create_table(b, create_table)
71+
match stmt {
72+
ast::Stmt::CreateTable(create_table) => bind_create_table(b, create_table),
73+
ast::Stmt::Set(set) => bind_set(b, set),
74+
_ => {}
5075
}
5176
}
5277

@@ -113,3 +138,79 @@ fn schema_name(path: &ast::Path, is_temp: bool) -> Schema {
113138

114139
Schema(schema_name)
115140
}
141+
142+
fn bind_set(b: &mut Binder, set: ast::Set) {
143+
let position = set.syntax().text_range().start();
144+
145+
// `set schema` is an alternative to `set search_path`
146+
if set.schema_token().is_some() {
147+
if let Some(literal) = set.literal() {
148+
if let Some(string_value) = extract_string_literal(&literal) {
149+
b.search_path_changes.push(SearchPathChange {
150+
position,
151+
search_path: vec![Schema::new(string_value)],
152+
});
153+
}
154+
}
155+
return;
156+
}
157+
158+
let Some(path) = set.path() else { return };
159+
160+
if path.qualifier().is_some() {
161+
return;
162+
}
163+
164+
let Some(segment) = path.segment() else {
165+
return;
166+
};
167+
168+
let param_name = if let Some(name_ref) = segment.name_ref() {
169+
name_ref.syntax().text().to_string()
170+
} else {
171+
return;
172+
};
173+
174+
if !param_name.eq_ignore_ascii_case("search_path") {
175+
return;
176+
}
177+
178+
// `set search_path`
179+
if set.default_token().is_some() {
180+
b.search_path_changes.push(SearchPathChange {
181+
position,
182+
search_path: vec![Schema::new("pg_temp"), Schema::new("public")],
183+
});
184+
} else {
185+
let mut search_path = vec![];
186+
for config_value in set.config_values() {
187+
match config_value {
188+
ast::ConfigValue::Literal(literal) => {
189+
if let Some(string_value) = extract_string_literal(&literal) {
190+
if !string_value.is_empty() {
191+
search_path.push(Schema::new(string_value));
192+
}
193+
}
194+
}
195+
ast::ConfigValue::NameRef(name_ref) => {
196+
let schema_name = name_ref.syntax().text().to_string();
197+
search_path.push(Schema::new(schema_name));
198+
}
199+
}
200+
}
201+
b.search_path_changes.push(SearchPathChange {
202+
position,
203+
search_path,
204+
});
205+
}
206+
}
207+
208+
fn extract_string_literal(literal: &ast::Literal) -> Option<String> {
209+
let text = literal.syntax().text().to_string();
210+
211+
if text.starts_with('\'') && text.ends_with('\'') && text.len() >= 2 {
212+
Some(text[1..text.len() - 1].to_string())
213+
} else {
214+
None
215+
}
216+
}

‎crates/squawk_ide/src/goto_definition.rs‎

Lines changed: 163 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -440,4 +440,167 @@ commit;
440440
╰╴────── 2. destination
441441
");
442442
}
443+
444+
#[test]
445+
fn goto_with_search_path() {
446+
assert_snapshot!(goto(r#"
447+
set search_path to "foo", public;
448+
create table foo.t();
449+
drop table t$0;
450+
"#), @r"
451+
╭▸
452+
3 │ create table foo.t();
453+
│ ─ 2. destination
454+
4 │ drop table t;
455+
╰╴ ─ 1. source
456+
");
457+
}
458+
459+
#[test]
460+
fn goto_with_search_path_like_variable() {
461+
// not actually search path
462+
goto_not_found(
463+
"
464+
set bar.search_path to foo, public;
465+
create table foo.t();
466+
drop table t$0;
467+
",
468+
)
469+
}
470+
471+
#[test]
472+
fn goto_with_search_path_second_schema() {
473+
assert_snapshot!(goto("
474+
set search_path to foo, bar, public;
475+
create table bar.t();
476+
drop table t$0;
477+
"), @r"
478+
╭▸
479+
3 │ create table bar.t();
480+
│ ─ 2. destination
481+
4 │ drop table t;
482+
╰╴ ─ 1. source
483+
");
484+
}
485+
486+
#[test]
487+
fn goto_with_search_path_skips_first() {
488+
assert_snapshot!(goto("
489+
set search_path to foo, bar, public;
490+
create table foo.t();
491+
create table bar.t();
492+
drop table t$0;
493+
"), @r"
494+
╭▸
495+
3 │ create table foo.t();
496+
│ ─ 2. destination
497+
4 │ create table bar.t();
498+
5 │ drop table t;
499+
╰╴ ─ 1. source
500+
");
501+
}
502+
503+
#[test]
504+
fn goto_without_search_path_uses_default() {
505+
assert_snapshot!(goto("
506+
create table foo.t();
507+
create table public.t();
508+
drop table t$0;
509+
"), @r"
510+
╭▸
511+
3 │ create table public.t();
512+
│ ─ 2. destination
513+
4 │ drop table t;
514+
╰╴ ─ 1. source
515+
");
516+
}
517+
518+
#[test]
519+
fn goto_with_set_schema() {
520+
assert_snapshot!(goto("
521+
set schema 'myschema';
522+
create table myschema.t();
523+
drop table t$0;
524+
"), @r"
525+
╭▸
526+
3 │ create table myschema.t();
527+
│ ─ 2. destination
528+
4 │ drop table t;
529+
╰╴ ─ 1. source
530+
");
531+
}
532+
533+
#[test]
534+
fn goto_with_set_schema_ignores_other_schemas() {
535+
assert_snapshot!(goto("
536+
set schema 'myschema';
537+
create table public.t();
538+
create table myschema.t();
539+
drop table t$0;
540+
"), @r"
541+
╭▸
542+
4 │ create table myschema.t();
543+
│ ─ 2. destination
544+
5 │ drop table t;
545+
╰╴ ─ 1. source
546+
");
547+
}
548+
549+
#[test]
550+
fn goto_with_search_path_changed_twice() {
551+
assert_snapshot!(goto("
552+
set search_path to foo;
553+
create table foo.t();
554+
set search_path to bar;
555+
create table bar.t();
556+
drop table t$0;
557+
"), @r"
558+
╭▸
559+
5 │ create table bar.t();
560+
│ ─ 2. destination
561+
6 │ drop table t;
562+
╰╴ ─ 1. source
563+
");
564+
565+
assert_snapshot!(goto("
566+
set search_path to foo;
567+
create table foo.t();
568+
drop table t$0;
569+
set search_path to bar;
570+
create table bar.t();
571+
drop table t;
572+
"), @r"
573+
╭▸
574+
3 │ create table foo.t();
575+
│ ─ 2. destination
576+
4 │ drop table t;
577+
╰╴ ─ 1. source
578+
");
579+
}
580+
581+
#[test]
582+
fn goto_with_empty_search_path() {
583+
goto_not_found(
584+
"
585+
set search_path to '';
586+
create table public.t();
587+
drop table t$0;
588+
",
589+
)
590+
}
591+
592+
#[test]
593+
fn goto_with_search_path_uppercase() {
594+
assert_snapshot!(goto("
595+
SET SEARCH_PATH TO foo;
596+
create table foo.t();
597+
drop table t$0;
598+
"), @r"
599+
╭▸
600+
3 │ create table foo.t();
601+
│ ─ 2. destination
602+
4 │ drop table t;
603+
╰╴ ─ 1. source
604+
");
605+
}
443606
}

‎crates/squawk_ide/src/resolve.rs‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
use rowan::TextSize;
12
use squawk_syntax::{
23
SyntaxNodePtr,
34
ast::{self, AstNode},
@@ -19,7 +20,8 @@ pub(crate) fn resolve_name_ref(binder: &Binder, name_ref: &ast::NameRef) -> Opti
1920
let path = find_containing_path(name_ref)?;
2021
let table_name = extract_table_name(&path)?;
2122
let schema = extract_schema_name(&path);
22-
resolve_table(binder, &table_name, &schema)
23+
let position = name_ref.syntax().text_range().start();
24+
resolve_table(binder, &table_name, &schema, position)
2325
}
2426
}
2527
}
@@ -38,6 +40,7 @@ fn resolve_table(
3840
binder: &Binder,
3941
table_name: &Name,
4042
schema: &Option<Schema>,
43+
position: TextSize,
4144
) -> Option<SyntaxNodePtr> {
4245
let symbols = binder.scopes[binder.root_scope()].get(table_name)?;
4346

@@ -48,10 +51,11 @@ fn resolve_table(
4851
})?;
4952
return Some(binder.symbols[symbol_id].ptr);
5053
} else {
51-
for search_schema in [Schema::new("pg_temp"), Schema::new("public")] {
54+
let search_path = binder.search_path_at(position);
55+
for search_schema in search_path {
5256
if let Some(symbol_id) = symbols.iter().copied().find(|id| {
5357
let symbol = &binder.symbols[*id];
54-
symbol.kind == SymbolKind::Table && symbol.schema == search_schema
58+
symbol.kind == SymbolKind::Table && &symbol.schema == search_schema
5559
}) {
5660
return Some(binder.symbols[symbol_id].ptr);
5761
}

‎crates/squawk_parser/src/grammar.rs‎

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -13535,14 +13535,10 @@ fn config_value(p: &mut Parser<'_>) -> bool {
1353513535
while !p.at(EOF) {
1353613536
if opt_string_literal(p).is_none()
1353713537
&& opt_numeric_literal(p).is_none()
13538-
&& !opt_ident(p)
13538+
&& opt_name_ref(p).is_none()
1353913539
&& !opt_bool_literal(p)
1354013540
{
13541-
if p.at_ts(BARE_LABEL_KEYWORDS) {
13542-
p.bump_any();
13543-
} else {
13544-
break;
13545-
}
13541+
break;
1354613542
}
1354713543
found_value = true;
1354813544
if !p.eat(COMMA) {

‎crates/squawk_parser/tests/snapshots/tests__alter_database_ok.snap‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -116,7 +116,8 @@ SOURCE_FILE
116116
WHITESPACE " "
117117
TO_KW "to"
118118
WHITESPACE " "
119-
IDENT "v"
119+
NAME_REF
120+
IDENT "v"
120121
SEMICOLON ";"
121122
WHITESPACE "\n"
122123
ALTER_DATABASE
@@ -137,7 +138,8 @@ SOURCE_FILE
137138
WHITESPACE " "
138139
EQ "="
139140
WHITESPACE " "
140-
IDENT "v"
141+
NAME_REF
142+
IDENT "v"
141143
SEMICOLON ";"
142144
WHITESPACE "\n"
143145
ALTER_DATABASE

0 commit comments

Comments
 (0)