ide: add document symbol support for CTEs (#806)

Author: sbdchd

crates/squawk_ide/src/document_symbols.rs

@@ -47,13 +47,66 @@ pub fn document_symbols(file: &ast::SourceFile) -> Vec<DocumentSymbol> {
                     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));
+            }
+
             _ => {}
         }
     }
 
     symbols
 }
 
+fn cte_table_symbols(stmt: impl ast::HasWithClause) -> Vec<DocumentSymbol> {
+    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<DocumentSymbol> {
+    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
+        "
+        );
+    }
 }

crates/squawk_syntax/src/ast.rs

@@ -55,6 +55,7 @@ pub use self::{
         // HasGenericParams, HasLoopBody,
         HasName,
         HasParamList,
+        HasWithClause,
         NameLike,
     },
 };

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 = "

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<ast::WithClause> {
+        support::child(self.syntax())
+    }
+}
+
 pub trait HasArgList: AstNode {
     fn arg_list(&self) -> Option<ast::ArgList> {
         support::child(self.syntax())