diff --git a/crates/squawk_ide/src/document_symbols.rs b/crates/squawk_ide/src/document_symbols.rs index 800c87a4..033747a7 100644 --- a/crates/squawk_ide/src/document_symbols.rs +++ b/crates/squawk_ide/src/document_symbols.rs @@ -47,6 +47,22 @@ pub fn document_symbols(file: &ast::SourceFile) -> Vec { symbols.push(symbol); } } + ast::Stmt::Select(select) => { + symbols.extend(cte_table_symbols(select)); + } + ast::Stmt::SelectInto(select_into) => { + symbols.extend(cte_table_symbols(select_into)); + } + ast::Stmt::Insert(insert) => { + symbols.extend(cte_table_symbols(insert)); + } + ast::Stmt::Update(update) => { + symbols.extend(cte_table_symbols(update)); + } + ast::Stmt::Delete(delete) => { + symbols.extend(cte_table_symbols(delete)); + } + _ => {} } } @@ -54,6 +70,43 @@ pub fn document_symbols(file: &ast::SourceFile) -> Vec { symbols } +fn cte_table_symbols(stmt: impl ast::HasWithClause) -> Vec { + let Some(with_clause) = stmt.with_clause() else { + return vec![]; + }; + + with_clause + .with_tables() + .filter_map(create_cte_table_symbol) + .collect() +} + +fn create_cte_table_symbol(with_table: ast::WithTable) -> Option { + let name_node = with_table.name()?; + let name = name_node.syntax().text().to_string(); + + let full_range = with_table.syntax().text_range(); + let focus_range = name_node.syntax().text_range(); + + let mut children = vec![]; + if let Some(column_list) = with_table.column_list() { + for column in column_list.columns() { + if let Some(column_symbol) = create_column_symbol(column) { + children.push(column_symbol); + } + } + } + + Some(DocumentSymbol { + name, + detail: None, + kind: DocumentSymbolKind::Table, + full_range, + focus_range, + children, + }) +} + fn create_table_symbol( binder: &binder::Binder, create_table: ast::CreateTable, @@ -559,4 +612,64 @@ create function my_schema.hello() returns void as $$ select 1; $$ language sql; fn non_create_statements() { symbols_not_found("select * from users;") } + + #[test] + fn cte_table() { + assert_snapshot!( + symbols(" +with recent_users as ( + select id, email as user_email + from users +) +select * from recent_users; +"), + @r" + info: table: recent_users + ╭▸ + 2 │ with recent_users as ( + │ │━━━━━━━━━━━ + │ │ + │ ┌──────focus range + │ │ + 3 │ │ select id, email as user_email + 4 │ │ from users + 5 │ │ ) + ╰╴└─┘ full range + " + ); + } + + #[test] + fn cte_table_with_column_list() { + assert_snapshot!( + symbols(" +with t(a, b, c) as ( + select 1, 2, 3 +) +select * from t; +"), + @r" + info: table: t + ╭▸ + 2 │ with t(a, b, c) as ( + │ ━ focus range + │ ┌──────┘ + │ │ + 3 │ │ select 1, 2, 3 + 4 │ │ ) + │ └─┘ full range + │ + ⸬ + 2 │ with t(a, b, c) as ( + │ ┯ ┯ ┯ + │ │ │ │ + │ │ │ full range for `column: c` + │ │ │ focus range + │ │ full range for `column: b` + │ │ focus range + │ full range for `column: a` + ╰╴ focus range + " + ); + } } diff --git a/crates/squawk_syntax/src/ast.rs b/crates/squawk_syntax/src/ast.rs index 993196da..aa377541 100644 --- a/crates/squawk_syntax/src/ast.rs +++ b/crates/squawk_syntax/src/ast.rs @@ -55,6 +55,7 @@ pub use self::{ // HasGenericParams, HasLoopBody, HasName, HasParamList, + HasWithClause, NameLike, }, }; diff --git a/crates/squawk_syntax/src/ast/node_ext.rs b/crates/squawk_syntax/src/ast/node_ext.rs index 62fc90c1..29196172 100644 --- a/crates/squawk_syntax/src/ast/node_ext.rs +++ b/crates/squawk_syntax/src/ast/node_ext.rs @@ -231,9 +231,16 @@ pub(crate) fn text_of_first_token(node: &SyntaxNode) -> TokenText<'_> { impl ast::HasParamList for ast::FunctionSig {} impl ast::HasParamList for ast::Aggregate {} + impl ast::NameLike for ast::Name {} impl ast::NameLike for ast::NameRef {} +impl ast::HasWithClause for ast::Select {} +impl ast::HasWithClause for ast::SelectInto {} +impl ast::HasWithClause for ast::Insert {} +impl ast::HasWithClause for ast::Update {} +impl ast::HasWithClause for ast::Delete {} + #[test] fn index_expr() { let source_code = " diff --git a/crates/squawk_syntax/src/ast/traits.rs b/crates/squawk_syntax/src/ast/traits.rs index 087310f9..ed0ba2fb 100644 --- a/crates/squawk_syntax/src/ast/traits.rs +++ b/crates/squawk_syntax/src/ast/traits.rs @@ -11,6 +11,13 @@ pub trait HasName: AstNode { pub trait NameLike: AstNode {} +pub trait HasWithClause: AstNode { + #[inline] + fn with_clause(&self) -> Option { + support::child(self.syntax()) + } +} + pub trait HasArgList: AstNode { fn arg_list(&self) -> Option { support::child(self.syntax())