Skip to content

ide: support set search_path #747

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
kodiakhq merged 1 commit into from
Dec 18, 2025
Merged

ide: support set search_path #747

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
105 changes: 103 additions & 2 deletions crates/squawk_ide/src/binder.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,14 +1,21 @@
/// Loosely based on TypeScript's binder
/// see: typescript-go/internal/binder/binder.go
use la_arena::Arena;
use rowan::TextSize;
use squawk_syntax::{SyntaxNodePtr, ast, ast::AstNode};

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

pub(crate) struct SearchPathChange {
position: TextSize,
search_path: Vec<Schema>,
}

pub(crate) struct Binder {
pub(crate) scopes: Arena<Scope>,
pub(crate) symbols: Arena<Symbol>,
pub(crate) search_path_changes: Vec<SearchPathChange>,
}

impl Binder {
Expand All @@ -18,6 +25,10 @@ impl Binder {
Binder {
scopes,
symbols: Arena::new(),
search_path_changes: vec![SearchPathChange {
position: TextSize::from(0),
search_path: vec![Schema::new("pg_temp"), Schema::new("public")],
}],
}
}

Expand All @@ -28,6 +39,18 @@ impl Binder {
.map(|(id, _)| id)
.expect("root scope must exist")
}

pub(crate) fn search_path_at(&self, position: TextSize) -> &[Schema] {
// We're assuming people don't actually use `set search_path` that much,
// so linear search is fine
for change in self.search_path_changes.iter().rev() {
if change.position <= position {
return &change.search_path;
}
}
// default search path
&self.search_path_changes[0].search_path
}
}

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

fn bind_stmt(b: &mut Binder, stmt: ast::Stmt) {
if let ast::Stmt::CreateTable(create_table) = stmt {
bind_create_table(b, create_table)
match stmt {
ast::Stmt::CreateTable(create_table) => bind_create_table(b, create_table),
ast::Stmt::Set(set) => bind_set(b, set),
_ => {}
}
}

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

Schema(schema_name)
}

fn bind_set(b: &mut Binder, set: ast::Set) {
let position = set.syntax().text_range().start();

// `set schema` is an alternative to `set search_path`
if set.schema_token().is_some() {
if let Some(literal) = set.literal() {
if let Some(string_value) = extract_string_literal(&literal) {
b.search_path_changes.push(SearchPathChange {
position,
search_path: vec![Schema::new(string_value)],
});
}
}
return;
}

let Some(path) = set.path() else { return };

if path.qualifier().is_some() {
return;
}

let Some(segment) = path.segment() else {
return;
};

let param_name = if let Some(name_ref) = segment.name_ref() {
name_ref.syntax().text().to_string()
} else {
return;
};

if !param_name.eq_ignore_ascii_case("search_path") {
return;
}

// `set search_path`
if set.default_token().is_some() {
b.search_path_changes.push(SearchPathChange {
position,
search_path: vec![Schema::new("pg_temp"), Schema::new("public")],
});
} else {
let mut search_path = vec![];
for config_value in set.config_values() {
match config_value {
ast::ConfigValue::Literal(literal) => {
if let Some(string_value) = extract_string_literal(&literal) {
if !string_value.is_empty() {
search_path.push(Schema::new(string_value));
}
}
}
ast::ConfigValue::NameRef(name_ref) => {
let schema_name = name_ref.syntax().text().to_string();
search_path.push(Schema::new(schema_name));
}
}
}
b.search_path_changes.push(SearchPathChange {
position,
search_path,
});
}
}

fn extract_string_literal(literal: &ast::Literal) -> Option<String> {
let text = literal.syntax().text().to_string();

if text.starts_with('\'') && text.ends_with('\'') && text.len() >= 2 {
Some(text[1..text.len() - 1].to_string())
} else {
None
}
}
163 changes: 163 additions & 0 deletions crates/squawk_ide/src/goto_definition.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -440,4 +440,167 @@ commit;
╰╴────── 2. destination
");
}

#[test]
fn goto_with_search_path() {
assert_snapshot!(goto(r#"
set search_path to "foo", public;
create table foo.t();
drop table t$0;
"#), @r"
╭▸
3 │ create table foo.t();
│ ─ 2. destination
4 │ drop table t;
╰╴ ─ 1. source
");
}

#[test]
fn goto_with_search_path_like_variable() {
// not actually search path
goto_not_found(
"
set bar.search_path to foo, public;
create table foo.t();
drop table t$0;
",
)
}

#[test]
fn goto_with_search_path_second_schema() {
assert_snapshot!(goto("
set search_path to foo, bar, public;
create table bar.t();
drop table t$0;
"), @r"
╭▸
3 │ create table bar.t();
│ ─ 2. destination
4 │ drop table t;
╰╴ ─ 1. source
");
}

#[test]
fn goto_with_search_path_skips_first() {
assert_snapshot!(goto("
set search_path to foo, bar, public;
create table foo.t();
create table bar.t();
drop table t$0;
"), @r"
╭▸
3 │ create table foo.t();
│ ─ 2. destination
4 │ create table bar.t();
5 │ drop table t;
╰╴ ─ 1. source
");
}

#[test]
fn goto_without_search_path_uses_default() {
assert_snapshot!(goto("
create table foo.t();
create table public.t();
drop table t$0;
"), @r"
╭▸
3 │ create table public.t();
│ ─ 2. destination
4 │ drop table t;
╰╴ ─ 1. source
");
}

#[test]
fn goto_with_set_schema() {
assert_snapshot!(goto("
set schema 'myschema';
create table myschema.t();
drop table t$0;
"), @r"
╭▸
3 │ create table myschema.t();
│ ─ 2. destination
4 │ drop table t;
╰╴ ─ 1. source
");
}

#[test]
fn goto_with_set_schema_ignores_other_schemas() {
assert_snapshot!(goto("
set schema 'myschema';
create table public.t();
create table myschema.t();
drop table t$0;
"), @r"
╭▸
4 │ create table myschema.t();
│ ─ 2. destination
5 │ drop table t;
╰╴ ─ 1. source
");
}

#[test]
fn goto_with_search_path_changed_twice() {
assert_snapshot!(goto("
set search_path to foo;
create table foo.t();
set search_path to bar;
create table bar.t();
drop table t$0;
"), @r"
╭▸
5 │ create table bar.t();
│ ─ 2. destination
6 │ drop table t;
╰╴ ─ 1. source
");

assert_snapshot!(goto("
set search_path to foo;
create table foo.t();
drop table t$0;
set search_path to bar;
create table bar.t();
drop table t;
"), @r"
╭▸
3 │ create table foo.t();
│ ─ 2. destination
4 │ drop table t;
╰╴ ─ 1. source
");
}

#[test]
fn goto_with_empty_search_path() {
goto_not_found(
"
set search_path to '';
create table public.t();
drop table t$0;
",
)
}

#[test]
fn goto_with_search_path_uppercase() {
assert_snapshot!(goto("
SET SEARCH_PATH TO foo;
create table foo.t();
drop table t$0;
"), @r"
╭▸
3 │ create table foo.t();
│ ─ 2. destination
4 │ drop table t;
╰╴ ─ 1. source
");
}
}
10 changes: 7 additions & 3 deletions crates/squawk_ide/src/resolve.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
use rowan::TextSize;
use squawk_syntax::{
SyntaxNodePtr,
ast::{self, AstNode},
Expand All @@ -19,7 +20,8 @@ pub(crate) fn resolve_name_ref(binder: &Binder, name_ref: &ast::NameRef) -> Opti
let path = find_containing_path(name_ref)?;
let table_name = extract_table_name(&path)?;
let schema = extract_schema_name(&path);
resolve_table(binder, &table_name, &schema)
let position = name_ref.syntax().text_range().start();
resolve_table(binder, &table_name, &schema, position)
}
}
}
Expand All @@ -38,6 +40,7 @@ fn resolve_table(
binder: &Binder,
table_name: &Name,
schema: &Option<Schema>,
position: TextSize,
) -> Option<SyntaxNodePtr> {
let symbols = binder.scopes[binder.root_scope()].get(table_name)?;

Expand All @@ -48,10 +51,11 @@ fn resolve_table(
})?;
return Some(binder.symbols[symbol_id].ptr);
} else {
for search_schema in [Schema::new("pg_temp"), Schema::new("public")] {
let search_path = binder.search_path_at(position);
for search_schema in search_path {
if let Some(symbol_id) = symbols.iter().copied().find(|id| {
let symbol = &binder.symbols[*id];
symbol.kind == SymbolKind::Table && symbol.schema == search_schema
symbol.kind == SymbolKind::Table && &symbol.schema == search_schema
}) {
return Some(binder.symbols[symbol_id].ptr);
}
Expand Down
8 changes: 2 additions & 6 deletions crates/squawk_parser/src/grammar.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13535,14 +13535,10 @@ fn config_value(p: &mut Parser<'_>) -> bool {
while !p.at(EOF) {
if opt_string_literal(p).is_none()
&& opt_numeric_literal(p).is_none()
&& !opt_ident(p)
&& opt_name_ref(p).is_none()
&& !opt_bool_literal(p)
{
if p.at_ts(BARE_LABEL_KEYWORDS) {
p.bump_any();
} else {
break;
}
break;
}
found_value = true;
if !p.eat(COMMA) {
Expand Down
6 changes: 4 additions & 2 deletions crates/squawk_parser/tests/snapshots/tests__alter_database_ok.snap
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -116,7 +116,8 @@ SOURCE_FILE
WHITESPACE " "
TO_KW "to"
WHITESPACE " "
IDENT "v"
NAME_REF
IDENT "v"
SEMICOLON ";"
WHITESPACE "\n"
ALTER_DATABASE
Expand All @@ -137,7 +138,8 @@ SOURCE_FILE
WHITESPACE " "
EQ "="
WHITESPACE " "
IDENT "v"
NAME_REF
IDENT "v"
SEMICOLON ";"
WHITESPACE "\n"
ALTER_DATABASE
Expand Down
Loading
Loading