ide: more create table syntax support for goto def (#813)

Author: sbdchd

crates/squawk_ide/src/binder.rs

@@ -88,6 +88,7 @@ fn bind_stmt(b: &mut Binder, stmt: ast::Stmt) {
         ast::Stmt::CreateMaterializedView(create_view) => {
             bind_create_materialized_view(b, create_view)
         }
+        ast::Stmt::CreateSequence(create_sequence) => bind_create_sequence(b, create_sequence),
         ast::Stmt::Set(set) => bind_set(b, set),
         _ => {}
     }
@@ -331,6 +332,34 @@ fn bind_create_materialized_view(b: &mut Binder, create_view: ast::CreateMateria
     b.scopes[root].insert(view_name, view_id);
 }
 
+fn bind_create_sequence(b: &mut Binder, create_sequence: ast::CreateSequence) {
+    let Some(path) = create_sequence.path() else {
+        return;
+    };
+
+    let Some(sequence_name) = item_name(&path) else {
+        return;
+    };
+
+    let name_ptr = path_to_ptr(&path);
+    let is_temp =
+        create_sequence.temp_token().is_some() || create_sequence.temporary_token().is_some();
+
+    let Some(schema) = schema_name(b, &path, is_temp) else {
+        return;
+    };
+
+    let sequence_id = b.symbols.alloc(Symbol {
+        kind: SymbolKind::Sequence,
+        ptr: name_ptr,
+        schema,
+        params: None,
+    });
+
+    let root = b.root_scope();
+    b.scopes[root].insert(sequence_name, sequence_id);
+}
+
 fn item_name(path: &ast::Path) -> Option<Name> {
     let segment = path.segment()?;
 

crates/squawk_ide/src/classify.rs

@@ -8,6 +8,20 @@ pub(crate) enum NameRefClass {
     DropType,
     DropView,
     DropMaterializedView,
+    DropSequence,
+    ForeignKeyTable,
+    ForeignKeyColumn,
+    ForeignKeyLocalColumn,
+    CheckConstraintColumn,
+    GeneratedColumn,
+    UniqueConstraintColumn,
+    PrimaryKeyConstraintColumn,
+    NotNullConstraintColumn,
+    ExcludeConstraintColumn,
+    PartitionByColumn,
+    PartitionOfTable,
+    LikeTable,
+    InheritsTable,
     DropFunction,
     DropAggregate,
     DropProcedure,
@@ -36,13 +50,16 @@ pub(crate) enum NameRefClass {
 }
 
 pub(crate) fn classify_name_ref(name_ref: &ast::NameRef) -> Option<NameRefClass> {
-    let mut in_partition_item = false;
     let mut in_call_expr = false;
     let mut in_arg_list = false;
     let mut in_column_list = false;
     let mut in_where_clause = false;
     let mut in_from_clause = false;
     let mut in_set_clause = false;
+    let mut in_constraint_exclusion_list = false;
+    let mut in_constraint_include_clause = false;
+    let mut in_constraint_where_clause = false;
+    let mut in_partition_item = false;
 
     // TODO: can we combine this if and the one that follows?
     if let Some(parent) = name_ref.syntax().parent()
@@ -170,6 +187,62 @@ pub(crate) fn classify_name_ref(name_ref: &ast::NameRef) -> Option<NameRefClass>
         if ast::DropMaterializedView::can_cast(ancestor.kind()) {
             return Some(NameRefClass::DropMaterializedView);
         }
+        if ast::DropSequence::can_cast(ancestor.kind()) {
+            return Some(NameRefClass::DropSequence);
+        }
+        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
+                // and from columns to differentiate which would help us avoid
+                // this
+                if let Some(to_columns) = foreign_key.to_columns()
+                    && to_columns
+                        .syntax()
+                        .text_range()
+                        .contains_range(name_ref.syntax().text_range())
+                {
+                    return Some(NameRefClass::ForeignKeyColumn);
+                }
+                if let Some(from_columns) = foreign_key.from_columns()
+                    && from_columns
+                        .syntax()
+                        .text_range()
+                        .contains_range(name_ref.syntax().text_range())
+                {
+                    return Some(NameRefClass::ForeignKeyLocalColumn);
+                }
+            } else {
+                return Some(NameRefClass::ForeignKeyTable);
+            }
+        }
+        if ast::CheckConstraint::can_cast(ancestor.kind()) {
+            return Some(NameRefClass::CheckConstraintColumn);
+        }
+        if ast::GeneratedConstraint::can_cast(ancestor.kind()) {
+            return Some(NameRefClass::GeneratedColumn);
+        }
+        if in_column_list && ast::UniqueConstraint::can_cast(ancestor.kind()) {
+            return Some(NameRefClass::UniqueConstraintColumn);
+        }
+        if in_column_list && ast::PrimaryKeyConstraint::can_cast(ancestor.kind()) {
+            return Some(NameRefClass::PrimaryKeyConstraintColumn);
+        }
+        if ast::NotNullConstraint::can_cast(ancestor.kind()) {
+            return Some(NameRefClass::NotNullConstraintColumn);
+        }
+        if (in_constraint_exclusion_list
+            || in_constraint_include_clause
+            || in_constraint_where_clause)
+            && ast::ExcludeConstraint::can_cast(ancestor.kind())
+        {
+            return Some(NameRefClass::ExcludeConstraintColumn);
+        }
+        if ast::LikeClause::can_cast(ancestor.kind()) {
+            return Some(NameRefClass::LikeTable);
+        }
+        if ast::Inherits::can_cast(ancestor.kind()) {
+            return Some(NameRefClass::InheritsTable);
+        }
         if ast::CastExpr::can_cast(ancestor.kind()) && in_type {
             return Some(NameRefClass::TypeReference);
         }
@@ -197,15 +270,18 @@ pub(crate) fn classify_name_ref(name_ref: &ast::NameRef) -> Option<NameRefClass>
         {
             return Some(NameRefClass::CreateSchema);
         }
-        if ast::PartitionItem::can_cast(ancestor.kind()) {
-            in_partition_item = true;
-        }
         if ast::CreateIndex::can_cast(ancestor.kind()) {
             if in_partition_item {
                 return Some(NameRefClass::CreateIndexColumn);
             }
             return Some(NameRefClass::CreateIndex);
         }
+        if in_partition_item && ast::CreateTable::can_cast(ancestor.kind()) {
+            return Some(NameRefClass::PartitionByColumn);
+        }
+        if ast::PartitionOf::can_cast(ancestor.kind()) {
+            return Some(NameRefClass::PartitionOfTable);
+        }
         if ast::ArgList::can_cast(ancestor.kind()) {
             in_arg_list = true;
         }
@@ -229,6 +305,18 @@ pub(crate) fn classify_name_ref(name_ref: &ast::NameRef) -> Option<NameRefClass>
         if ast::ColumnList::can_cast(ancestor.kind()) {
             in_column_list = true;
         }
+        if ast::ConstraintExclusionList::can_cast(ancestor.kind()) {
+            in_constraint_exclusion_list = true;
+        }
+        if ast::ConstraintIncludeClause::can_cast(ancestor.kind()) {
+            in_constraint_include_clause = true;
+        }
+        if ast::WhereConditionClause::can_cast(ancestor.kind()) {
+            in_constraint_where_clause = true;
+        }
+        if ast::PartitionItem::can_cast(ancestor.kind()) {
+            in_partition_item = true;
+        }
         if ast::Insert::can_cast(ancestor.kind()) {
             if in_column_list {
                 return Some(NameRefClass::InsertColumn);
@@ -273,6 +361,7 @@ pub(crate) enum NameClass {
     CreateTable(ast::CreateTable),
     WithTable(ast::WithTable),
     CreateIndex(ast::CreateIndex),
+    CreateSequence(ast::CreateSequence),
     CreateType(ast::CreateType),
     CreateFunction(ast::CreateFunction),
     CreateAggregate(ast::CreateAggregate),
@@ -307,6 +396,9 @@ pub(crate) fn classify_name(name: &ast::Name) -> Option<NameClass> {
         if let Some(create_index) = ast::CreateIndex::cast(ancestor.clone()) {
             return Some(NameClass::CreateIndex(create_index));
         }
+        if let Some(create_sequence) = ast::CreateSequence::cast(ancestor.clone()) {
+            return Some(NameClass::CreateSequence(create_sequence));
+        }
         if let Some(create_type) = ast::CreateType::cast(ancestor.clone()) {
             return Some(NameClass::CreateType(create_type));
         }

crates/squawk_ide/src/document_symbols.rs

@@ -9,6 +9,7 @@ use crate::resolve::{
 
 #[derive(Debug)]
 pub enum DocumentSymbolKind {
+    Schema,
     Table,
     View,
     MaterializedView,
@@ -40,6 +41,11 @@ pub fn document_symbols(file: &ast::SourceFile) -> Vec<DocumentSymbol> {
 
     for stmt in file.stmts() {
         match stmt {
+            ast::Stmt::CreateSchema(create_schema) => {
+                if let Some(symbol) = create_schema_symbol(create_schema) {
+                    symbols.push(symbol);
+                }
+            }
             ast::Stmt::CreateTable(create_table) => {
                 if let Some(symbol) = create_table_symbol(&binder, create_table) {
                     symbols.push(symbol);
@@ -135,6 +141,37 @@ fn create_cte_table_symbol(with_table: ast::WithTable) -> Option<DocumentSymbol>
     })
 }
 
+fn create_schema_symbol(create_schema: ast::CreateSchema) -> Option<DocumentSymbol> {
+    let (name, focus_range) = if let Some(name_node) = create_schema.name() {
+        (
+            name_node.syntax().text().to_string(),
+            name_node.syntax().text_range(),
+        )
+    } else if let Some(schema_name_ref) = create_schema
+        .schema_authorization()
+        .and_then(|authorization| authorization.role())
+        .and_then(|role| role.name_ref())
+    {
+        (
+            schema_name_ref.syntax().text().to_string(),
+            schema_name_ref.syntax().text_range(),
+        )
+    } else {
+        return None;
+    };
+
+    let full_range = create_schema.syntax().text_range();
+
+    Some(DocumentSymbol {
+        name,
+        detail: None,
+        kind: DocumentSymbolKind::Schema,
+        full_range,
+        focus_range,
+        children: vec![],
+    })
+}
+
 fn create_table_symbol(
     binder: &binder::Binder,
     create_table: ast::CreateTable,
@@ -425,6 +462,7 @@ mod tests {
 
     fn symbol_to_group<'a>(symbol: &DocumentSymbol, sql: &'a str) -> Group<'a> {
         let kind = match symbol.kind {
+            DocumentSymbolKind::Schema => "schema",
             DocumentSymbolKind::Table => "table",
             DocumentSymbolKind::View => "view",
             DocumentSymbolKind::MaterializedView => "materialized view",
@@ -530,6 +568,36 @@ create table users (
         ");
     }
 
+    #[test]
+    fn create_schema() {
+        assert_snapshot!(symbols("
+create schema foo;
+"), @r"
+        info: schema: foo
+          ╭▸ 
+        2 │ create schema foo;
+          │ ┬─────────────┯━━
+          │ │             │
+          │ │             focus range
+          ╰╴full range
+        ");
+    }
+
+    #[test]
+    fn create_schema_authorization() {
+        assert_snapshot!(symbols("
+create schema authorization foo;
+"), @r"
+        info: schema: foo
+          ╭▸ 
+        2 │ create schema authorization foo;
+          │ ┬───────────────────────────┯━━
+          │ │                           │
+          │ │                           focus range
+          ╰╴full range
+        ");
+    }
+
     #[test]
     fn create_function() {
         assert_snapshot!(