ide: create table as (#893)

Author: sbdchd

crates/squawk_ide/src/binder.rs

@@ -249,6 +249,7 @@ fn bind_file(b: &mut Binder, file: &ast::SourceFile) {
 fn bind_stmt(b: &mut Binder, stmt: ast::Stmt) {
     match stmt {
         ast::Stmt::CreateTable(create_table) => bind_create_table(b, create_table),
+        ast::Stmt::CreateTableAs(create_table_as) => bind_create_table_as(b, create_table_as),
         ast::Stmt::CreateForeignTable(create_foreign_table) => {
             bind_create_table(b, create_foreign_table)
         }
@@ -318,6 +319,41 @@ fn bind_create_table(b: &mut Binder, create_table: impl ast::HasCreateTable) {
     b.scopes[root].insert(table_name, type_id);
 }
 
+fn bind_create_table_as(b: &mut Binder, create_table_as: ast::CreateTableAs) {
+    let Some(path) = create_table_as.path() else {
+        return;
+    };
+    let Some(table_name) = item_name(&path) else {
+        return;
+    };
+    let name_ptr = path_to_ptr(&path);
+    let is_temp =
+        create_table_as.temp_token().is_some() || create_table_as.temporary_token().is_some();
+    let Some(schema) = schema_name(b, &path, is_temp) else {
+        return;
+    };
+
+    let table_id = b.symbols.alloc(Symbol {
+        kind: SymbolKind::Table,
+        ptr: name_ptr,
+        schema: Some(schema.clone()),
+        params: None,
+        table: None,
+    });
+
+    let type_id = b.symbols.alloc(Symbol {
+        kind: SymbolKind::Type,
+        ptr: name_ptr,
+        schema: Some(schema),
+        params: None,
+        table: None,
+    });
+
+    let root = b.root_scope();
+    b.scopes[root].insert(table_name.clone(), table_id);
+    b.scopes[root].insert(table_name, type_id);
+}
+
 fn bind_create_index(b: &mut Binder, create_index: ast::CreateIndex) {
     let Some(name) = create_index.name() else {
         return;

crates/squawk_ide/src/document_symbols.rs

@@ -65,6 +65,11 @@ pub fn document_symbols(file: &ast::SourceFile) -> Vec<DocumentSymbol> {
                     symbols.push(symbol);
                 }
             }
+            ast::Stmt::CreateTableAs(create_table_as) => {
+                if let Some(symbol) = create_table_as_symbol(&binder, create_table_as) {
+                    symbols.push(symbol);
+                }
+            }
             ast::Stmt::CreateForeignTable(create_foreign_table) => {
                 if let Some(symbol) = create_table_symbol(&binder, create_foreign_table) {
                     symbols.push(symbol);
@@ -289,6 +294,36 @@ fn create_table_symbol(
     })
 }
 
+fn create_table_as_symbol(
+    binder: &binder::Binder,
+    create_table_as: ast::CreateTableAs,
+) -> Option<DocumentSymbol> {
+    let path = create_table_as.path()?;
+    let segment = path.segment()?;
+    let name_node = if let Some(name) = segment.name() {
+        name.syntax().clone()
+    } else {
+        return None;
+    };
+
+    let (schema, table_name) = resolve_table_info(binder, &path)?;
+    let name = format!("{}.{}", schema.0, table_name);
+
+    let full_range = create_table_as.syntax().text_range();
+    let focus_range = name_node.text_range();
+
+    Some(DocumentSymbol {
+        name,
+        detail: None,
+        kind: DocumentSymbolKind::Table,
+        full_range,
+        focus_range,
+        // TODO: infer the column names, we need the same for views without
+        // explicit column lists
+        children: vec![],
+    })
+}
+
 fn create_view_symbol(
     binder: &binder::Binder,
     create_view: ast::CreateView,
@@ -959,6 +994,21 @@ create table users (
         ");
     }
 
+    #[test]
+    fn create_table_as() {
+        assert_snapshot!(symbols("
+create table t as select 1 a;
+"), @r"
+        info: table: public.t
+          ╭▸ 
+        2 │ create table t as select 1 a;
+          │ ┬────────────┯──────────────
+          │ │            │
+          │ │            focus range
+          ╰╴full range
+        ");
+    }
+
     #[test]
     fn create_schema() {
         assert_snapshot!(symbols("

crates/squawk_ide/src/goto_definition.rs

@@ -2247,6 +2247,34 @@ select v.a$0 from v;
         ");
     }
 
+    #[test]
+    fn goto_create_table_as_column() {
+        assert_snapshot!(goto("
+create table t as select 1 a;
+select a$0 from t;
+"), @r"
+          ╭▸ 
+        2 │ create table t as select 1 a;
+          │                            ─ 2. destination
+        3 │ select a from t;
+          ╰╴       ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_select_from_create_table_as() {
+        assert_snapshot!(goto("
+create table t as select 1 a;
+select a from t$0;
+"), @r"
+          ╭▸ 
+        2 │ create table t as select 1 a;
+          │              ─ 2. destination
+        3 │ select a from t;
+          ╰╴              ─ 1. source
+        ");
+    }
+
     #[test]
     fn goto_view_with_explicit_column_list() {
         assert_snapshot!(goto("

crates/squawk_ide/src/hover.rs

@@ -385,6 +385,21 @@ fn format_hover_for_column_node(
         {
             return format_view_column(&create_view, column_name, binder);
         }
+
+        if let Some(create_table_as) = ast::CreateTableAs::cast(a.clone()) {
+            let column_name = if let Some(name) = ast::Name::cast(column_name_node.clone()) {
+                Name::from_node(&name)
+            } else {
+                continue;
+            };
+            let path = create_table_as.path()?;
+            let (schema, table_name) = resolve::resolve_table_info(binder, &path)?;
+            return Some(ColumnHover::schema_table_column(
+                &schema.to_string(),
+                &table_name,
+                &column_name.to_string(),
+            ));
+        }
     }
 
     let column = column_name_node.ancestors().find_map(ast::Column::cast)?;
@@ -3614,6 +3629,19 @@ select a$0, b from v;
         ");
     }
 
+    #[test]
+    fn hover_on_select_column_from_create_table_as() {
+        assert_snapshot!(check_hover("
+create table t as select 1 a;
+select a$0 from t;
+"), @r"
+        hover: column public.t.a
+          ╭▸ 
+        3 │ select a from t;
+          ╰╴       ─ hover
+        ");
+    }
+
     #[test]
     fn hover_on_select_from_view_with_schema() {
         assert_snapshot!(check_hover("