ide: goto def support for tablespace (#814)

Author: sbdchd

crates/squawk_ide/src/binder.rs

@@ -89,6 +89,9 @@ fn bind_stmt(b: &mut Binder, stmt: ast::Stmt) {
             bind_create_materialized_view(b, create_view)
         }
         ast::Stmt::CreateSequence(create_sequence) => bind_create_sequence(b, create_sequence),
+        ast::Stmt::CreateTablespace(create_tablespace) => {
+            bind_create_tablespace(b, create_tablespace)
+        }
         ast::Stmt::Set(set) => bind_set(b, set),
         _ => {}
     }
@@ -360,6 +363,27 @@ fn bind_create_sequence(b: &mut Binder, create_sequence: ast::CreateSequence) {
     b.scopes[root].insert(sequence_name, sequence_id);
 }
 
+fn bind_create_tablespace(b: &mut Binder, create_tablespace: ast::CreateTablespace) {
+    let Some(name) = create_tablespace.name() else {
+        return;
+    };
+
+    let tablespace_name = Name::from_node(&name);
+    let name_ptr = SyntaxNodePtr::new(name.syntax());
+
+    let tablespace_id = b.symbols.alloc(Symbol {
+        kind: SymbolKind::Tablespace,
+        ptr: name_ptr,
+        // TODO: tablespaces don't actually have schemas so we should think
+        // about this more
+        schema: Schema::new("pg_tablespace"),
+        params: None,
+    });
+
+    let root = b.root_scope();
+    b.scopes[root].insert(tablespace_name, tablespace_id);
+}
+
 fn item_name(path: &ast::Path) -> Option<Name> {
     let segment = path.segment()?;
 

crates/squawk_ide/src/classify.rs

@@ -9,6 +9,7 @@ pub(crate) enum NameRefClass {
     DropView,
     DropMaterializedView,
     DropSequence,
+    Tablespace,
     ForeignKeyTable,
     ForeignKeyColumn,
     ForeignKeyLocalColumn,
@@ -190,6 +191,13 @@ pub(crate) fn classify_name_ref(name_ref: &ast::NameRef) -> Option<NameRefClass>
         if ast::DropSequence::can_cast(ancestor.kind()) {
             return Some(NameRefClass::DropSequence);
         }
+        if ast::DropTablespace::can_cast(ancestor.kind())
+            || ast::Tablespace::can_cast(ancestor.kind())
+            || ast::SetTablespace::can_cast(ancestor.kind())
+            || ast::ConstraintIndexTablespace::can_cast(ancestor.kind())
+        {
+            return Some(NameRefClass::Tablespace);
+        }
         if let Some(foreign_key) = ast::ForeignKeyConstraint::cast(ancestor.clone()) {
             if in_column_list {
                 // TODO: ast is too "flat" here, we need a unique node for to
@@ -362,6 +370,7 @@ pub(crate) enum NameClass {
     WithTable(ast::WithTable),
     CreateIndex(ast::CreateIndex),
     CreateSequence(ast::CreateSequence),
+    CreateTablespace(ast::CreateTablespace),
     CreateType(ast::CreateType),
     CreateFunction(ast::CreateFunction),
     CreateAggregate(ast::CreateAggregate),
@@ -399,6 +408,9 @@ pub(crate) fn classify_name(name: &ast::Name) -> Option<NameClass> {
         if let Some(create_sequence) = ast::CreateSequence::cast(ancestor.clone()) {
             return Some(NameClass::CreateSequence(create_sequence));
         }
+        if let Some(create_tablespace) = ast::CreateTablespace::cast(ancestor.clone()) {
+            return Some(NameClass::CreateTablespace(create_tablespace));
+        }
         if let Some(create_type) = ast::CreateType::cast(ancestor.clone()) {
             return Some(NameClass::CreateType(create_type));
         }

crates/squawk_ide/src/goto_definition.rs

@@ -246,6 +246,34 @@ drop sequence s$0;
         ");
     }
 
+    #[test]
+    fn goto_drop_tablespace() {
+        assert_snapshot!(goto("
+create tablespace ts location '/tmp/ts';
+drop tablespace ts$0;
+"), @r"
+          ╭▸ 
+        2 │ create tablespace ts location '/tmp/ts';
+          │                   ── 2. destination
+        3 │ drop tablespace ts;
+          ╰╴                 ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_create_table_tablespace() {
+        assert_snapshot!(goto("
+create tablespace bar location '/tmp/ts';
+create table t (a int) tablespace b$0ar;
+"), @r"
+          ╭▸ 
+        2 │ create tablespace bar location '/tmp/ts';
+          │                   ─── 2. destination
+        3 │ create table t (a int) tablespace bar;
+          ╰╴                                  ─ 1. source
+        ");
+    }
+
     #[test]
     fn goto_drop_sequence_with_schema() {
         assert_snapshot!(goto("

crates/squawk_ide/src/hover.rs

@@ -67,6 +67,7 @@ pub fn hover(file: &ast::SourceFile, offset: TextSize) -> Option<String> {
                 return hover_column(file, &name_ref, &binder);
             }
             NameRefClass::DropSequence => return hover_sequence(file, &name_ref, &binder),
+            NameRefClass::Tablespace => return hover_tablespace(file, &name_ref, &binder),
             NameRefClass::DropIndex => return hover_index(file, &name_ref, &binder),
             NameRefClass::DropFunction => return hover_function(file, &name_ref, &binder),
             NameRefClass::DropAggregate => return hover_aggregate(file, &name_ref, &binder),
@@ -107,6 +108,9 @@ pub fn hover(file: &ast::SourceFile, offset: TextSize) -> Option<String> {
             NameClass::CreateSequence(create_sequence) => {
                 return format_create_sequence(&create_sequence, &binder);
             }
+            NameClass::CreateTablespace(create_tablespace) => {
+                return format_create_tablespace(&create_tablespace);
+            }
             NameClass::CreateType(create_type) => {
                 return format_create_type(&create_type, &binder);
             }
@@ -322,6 +326,17 @@ fn hover_sequence(
     format_create_sequence(&create_sequence, binder)
 }
 
+fn hover_tablespace(
+    file: &ast::SourceFile,
+    name_ref: &ast::NameRef,
+    binder: &binder::Binder,
+) -> Option<String> {
+    let tablespace_ptr = resolve::resolve_name_ref(binder, name_ref)?;
+    let root = file.syntax();
+    let tablespace_name_node = tablespace_ptr.to_node(root);
+    Some(format!("tablespace {}", tablespace_name_node.text()))
+}
+
 fn hover_type(
     file: &ast::SourceFile,
     name_ref: &ast::NameRef,
@@ -474,6 +489,11 @@ fn format_create_sequence(
     Some(format!("sequence {}.{}", schema, sequence_name))
 }
 
+fn format_create_tablespace(create_tablespace: &ast::CreateTablespace) -> Option<String> {
+    let name = create_tablespace.name()?.syntax().text().to_string();
+    Some(format!("tablespace {}", name))
+}
+
 fn index_schema(create_index: &ast::CreateIndex, binder: &binder::Binder) -> Option<String> {
     let position = create_index.syntax().text_range().start();
     let search_path = binder.search_path_at(position);

crates/squawk_ide/src/resolve.rs

@@ -98,6 +98,10 @@ pub(crate) fn resolve_name_ref(binder: &Binder, name_ref: &ast::NameRef) -> Opti
             let position = name_ref.syntax().text_range().start();
             resolve_sequence(binder, &sequence_name, &schema, position)
         }
+        NameRefClass::Tablespace => {
+            let tablespace_name = Name::from_node(name_ref);
+            resolve_tablespace(binder, &tablespace_name)
+        }
         NameRefClass::ForeignKeyTable => {
             let foreign_key = name_ref
                 .syntax()
@@ -420,6 +424,15 @@ fn resolve_sequence(
     )
 }
 
+fn resolve_tablespace(binder: &Binder, tablespace_name: &Name) -> Option<SyntaxNodePtr> {
+    let symbols = binder.scopes[binder.root_scope()].get(tablespace_name)?;
+    let symbol_id = symbols.iter().copied().find(|id| {
+        let symbol = &binder.symbols[*id];
+        symbol.kind == SymbolKind::Tablespace
+    })?;
+    Some(binder.symbols[symbol_id].ptr)
+}
+
 fn resolve_for_kind(
     binder: &Binder,
     name: &Name,

crates/squawk_ide/src/symbols.rs

@@ -53,6 +53,7 @@ pub(crate) enum SymbolKind {
     Type,
     View,
     Sequence,
+    Tablespace,
 }
 
 #[derive(Clone, Debug)]