Skip to content

ide: add columns to document symbols #778

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 29, 2025
Merged

ide: add columns to document symbols #778

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
226 changes: 177 additions & 49 deletions crates/squawk_ide/src/document_symbols.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,17 +4,24 @@ use squawk_syntax::ast::{self, AstNode};
use crate::binder;
use crate::resolve::{resolve_function_info, resolve_table_info};

#[derive(Debug)]
pub enum DocumentSymbolKind {
Table,
Function,
Column,
}

#[derive(Debug)]
pub struct DocumentSymbol {
pub name: String,
pub detail: Option<String>,
pub kind: DocumentSymbolKind,
pub range: TextRange,
pub selection_range: TextRange,
/// Range used for determining when cursor is inside the symbol for showing
/// in the UI
pub full_range: TextRange,
/// Range selected when symbol is selected
pub focus_range: TextRange,
pub children: Vec<DocumentSymbol>,
}

pub fn document_symbols(file: &ast::SourceFile) -> Vec<DocumentSymbol> {
Expand Down Expand Up @@ -51,15 +58,27 @@ fn create_table_symbol(
let (schema, table_name) = resolve_table_info(binder, &path)?;
let name = format!("{}.{}", schema.0, table_name);

let range = create_table.syntax().text_range();
let selection_range = name_node.syntax().text_range();
let full_range = create_table.syntax().text_range();
let focus_range = name_node.syntax().text_range();

let mut children = vec![];
if let Some(table_arg_list) = create_table.table_arg_list() {
for arg in table_arg_list.args() {
if let ast::TableArg::Column(column) = arg
&& let Some(column_symbol) = create_column_symbol(column)
{
children.push(column_symbol);
}
}
}

Some(DocumentSymbol {
name,
detail: None,
kind: DocumentSymbolKind::Table,
range,
selection_range,
full_range,
focus_range,
children,
})
}

Expand All @@ -74,22 +93,44 @@ fn create_function_symbol(
let (schema, function_name) = resolve_function_info(binder, &path)?;
let name = format!("{}.{}", schema.0, function_name);

let range = create_function.syntax().text_range();
let selection_range = name_node.syntax().text_range();
let full_range = create_function.syntax().text_range();
let focus_range = name_node.syntax().text_range();

Some(DocumentSymbol {
name,
detail: None,
kind: DocumentSymbolKind::Function,
range,
selection_range,
full_range,
focus_range,
children: vec![],
})
}

fn create_column_symbol(column: ast::Column) -> Option<DocumentSymbol> {
let name_node = column.name()?;
let name = name_node.syntax().text().to_string();

let detail = column.ty().map(|t| t.syntax().text().to_string());

let full_range = column.syntax().text_range();
let focus_range = name_node.syntax().text_range();

Some(DocumentSymbol {
name,
detail,
kind: DocumentSymbolKind::Column,
full_range,
focus_range,
children: vec![],
})
}

#[cfg(test)]
mod tests {
use super::*;
use annotate_snippets::{AnnotationKind, Level, Renderer, Snippet, renderer::DecorStyle};
use annotate_snippets::{
AnnotationKind, Group, Level, Renderer, Snippet, renderer::DecorStyle,
};
use insta::assert_snapshot;

fn symbols_not_found(sql: &str) {
Expand All @@ -109,44 +150,110 @@ mod tests {
panic!("No symbols found. If this is expected, use `symbols_not_found` instead.")
}

let mut groups = vec![];
let mut output = vec![];
for symbol in symbols {
let kind = match symbol.kind {
DocumentSymbolKind::Table => "table",
DocumentSymbolKind::Function => "function",
};
let title = format!("{}: {}", kind, symbol.name);
let group = Level::INFO.primary_title(title).element(
Snippet::source(sql)
.fold(true)
let group = symbol_to_group(&symbol, sql);
output.push(group);
}
Renderer::plain()
.decor_style(DecorStyle::Unicode)
.render(&output)
.to_string()
}

fn symbol_to_group<'a>(symbol: &DocumentSymbol, sql: &'a str) -> Group<'a> {
let kind = match symbol.kind {
DocumentSymbolKind::Table => "table",
DocumentSymbolKind::Function => "function",
DocumentSymbolKind::Column => "column",
};

let title = if let Some(detail) = &symbol.detail {
format!("{}: {} {}", kind, symbol.name, detail)
} else {
format!("{}: {}", kind, symbol.name)
};

let snippet = Snippet::source(sql)
.fold(true)
.annotation(
AnnotationKind::Primary
.span(symbol.focus_range.into())
.label("focus range"),
)
.annotation(
AnnotationKind::Context
.span(symbol.full_range.into())
.label("full range"),
);

let mut group = Level::INFO.primary_title(title.clone()).element(snippet);

if !symbol.children.is_empty() {
let child_labels: Vec<String> = symbol
.children
.iter()
.map(|child| {
let kind = match child.kind {
DocumentSymbolKind::Column => "column",
_ => unreachable!("only columns can be children"),
};
let detail = &child.detail.as_ref().unwrap();
format!("{}: {} {}", kind, child.name, detail)
})
.collect();

let mut children_snippet = Snippet::source(sql).fold(true);

for (i, child) in symbol.children.iter().enumerate() {
children_snippet = children_snippet
.annotation(
AnnotationKind::Primary
.span(symbol.selection_range.into())
.label("name"),
AnnotationKind::Context
.span(child.full_range.into())
.label(format!("full range for `{}`", child_labels[i].clone())),
)
.annotation(
AnnotationKind::Context
.span(symbol.range.into())
.label("select range"),
),
);
groups.push(group);
AnnotationKind::Primary
.span(child.focus_range.into())
.label("focus range"),
);
}

group = group.element(children_snippet);
}

let renderer = Renderer::plain().decor_style(DecorStyle::Unicode);
renderer.render(&groups).to_string()
group
}

#[test]
fn create_table() {
assert_snapshot!(symbols("create table users (id int);"), @r"
assert_snapshot!(symbols("
create table users (
id int,
email citext
);"), @r"
info: table: public.users
╭▸
1 │ create table users (id int);
│ ┬────────────┯━━━━─────────
│ │ │
│ │ name
╰╴select range
2 │ create table users (
│ │ ━━━━━ focus range
│ ┌─┘
│ │
3 │ │ id int,
4 │ │ email citext
5 │ │ );
│ └─┘ full range
│
⸬
3 │ id int,
│ ┯━────
│ │
│ full range for `column: id int`
│ focus range
4 │ email citext
│ ┯━━━━───────
│ │
│ full range for `column: email citext`
╰╴ focus range
");
}

Expand All @@ -160,8 +267,8 @@ mod tests {
1 │ create function hello() returns void as $$ select 1; $$ language sql;
│ ┬───────────────┯━━━━───────────────────────────────────────────────
│ │ │
│ │ name
╰╴select range
│ │ focus range
╰╴full range
"
);
}
Expand All @@ -178,23 +285,37 @@ create function get_user(user_id int) returns void as $$ select 1; $$ language s
2 │ create table users (id int);
│ ┬────────────┯━━━━─────────
│ │ │
│ │ name
│ select range
│ │ focus range
│ full range
│
⸬
2 │ create table users (id int);
│ ┯━────
│ │
│ full range for `column: id int`
│ focus range
╰╴
info: table: public.posts
╭▸
3 │ create table posts (id int);
│ ┬────────────┯━━━━─────────
│ │ │
│ │ name
╰╴select range
│ │ focus range
│ full range
│
⸬
3 │ create table posts (id int);
│ ┯━────
│ │
│ full range for `column: id int`
╰╴ focus range
info: function: public.get_user
╭▸
4 │ create function get_user(user_id int) returns void as $$ select 1; $$ language sql;
│ ┬───────────────┯━━━━━━━──────────────────────────────────────────────────────────
│ │ │
│ │ name
╰╴select range
│ │ focus range
╰╴full range
");
}

Expand All @@ -209,16 +330,23 @@ create function my_schema.hello() returns void as $$ select 1; $$ language sql;
2 │ create table public.users (id int);
│ ┬───────────────────┯━━━━─────────
│ │ │
│ │ name
│ select range
│ │ focus range
│ full range
│
⸬
2 │ create table public.users (id int);
│ ┯━────
│ │
│ full range for `column: id int`
│ focus range
╰╴
info: function: my_schema.hello
╭▸
3 │ create function my_schema.hello() returns void as $$ select 1; $$ language sql;
│ ┬─────────────────────────┯━━━━───────────────────────────────────────────────
│ │ │
│ │ name
╰╴select range
│ │ focus range
╰╴full range
");
}

Expand Down
52 changes: 33 additions & 19 deletions crates/squawk_server/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -316,27 +316,41 @@ fn handle_document_symbol(

let symbols = document_symbols(&file);

fn convert_symbol(
sym: squawk_ide::document_symbols::DocumentSymbol,
line_index: &LineIndex,
) -> DocumentSymbol {
let range = lsp_utils::range(line_index, sym.full_range);
let selection_range = lsp_utils::range(line_index, sym.focus_range);

let children = sym
.children
.into_iter()
.map(|child| convert_symbol(child, line_index))
.collect::<Vec<_>>();

let children = (!children.is_empty()).then_some(children);

DocumentSymbol {
name: sym.name,
detail: sym.detail,
kind: match sym.kind {
DocumentSymbolKind::Table => SymbolKind::STRUCT,
DocumentSymbolKind::Function => SymbolKind::FUNCTION,
DocumentSymbolKind::Column => SymbolKind::FIELD,
},
tags: None,
range,
selection_range,
children,
#[allow(deprecated)]
deprecated: None,
}
}

let lsp_symbols: Vec<DocumentSymbol> = symbols
.into_iter()
.map(|sym| {
let range = lsp_utils::range(&line_index, sym.range);
let selection_range = lsp_utils::range(&line_index, sym.selection_range);

DocumentSymbol {
name: sym.name,
detail: sym.detail,
kind: match sym.kind {
DocumentSymbolKind::Table => SymbolKind::STRUCT,
DocumentSymbolKind::Function => SymbolKind::FUNCTION,
},
tags: None,
range,
selection_range,
children: None,
#[allow(deprecated)]
deprecated: None,
}
})
.map(|sym| convert_symbol(sym, &line_index))
.collect();

let resp = Response {
Expand Down
Loading