ide: support goto def on using clause in join (#815)

Author: sbdchd

Cargo.lock

@@ -57,6 +57,7 @@ quote = "1.0.40"
 xshell = "0.2.7"
 proc-macro2 = "1.0.95"
 snapbox = { version = "0.6.0", features = ["diff", "term-svg", "cmd"] }
+smallvec = "1.13.2"
 
 # local
 # we have to make the versions explicit otherwise `cargo publish` won't work

Cargo.toml

@@ -20,6 +20,7 @@ annotate-snippets.workspace = true
 log.workspace = true
 smol_str.workspace = true
 la-arena.workspace = true
+smallvec.workspace = true
 
 [dev-dependencies]
 insta.workspace = true

crates/squawk_ide/Cargo.toml

@@ -92,6 +92,7 @@ fn bind_stmt(b: &mut Binder, stmt: ast::Stmt) {
         ast::Stmt::CreateTablespace(create_tablespace) => {
             bind_create_tablespace(b, create_tablespace)
         }
+        ast::Stmt::CreateDatabase(create_database) => bind_create_database(b, create_database),
         ast::Stmt::Set(set) => bind_set(b, set),
         _ => {}
     }
@@ -384,6 +385,25 @@ fn bind_create_tablespace(b: &mut Binder, create_tablespace: ast::CreateTablespa
     b.scopes[root].insert(tablespace_name, tablespace_id);
 }
 
+fn bind_create_database(b: &mut Binder, create_database: ast::CreateDatabase) {
+    let Some(name) = create_database.name() else {
+        return;
+    };
+
+    let database_name = Name::from_node(&name);
+    let name_ptr = SyntaxNodePtr::new(name.syntax());
+
+    let database_id = b.symbols.alloc(Symbol {
+        kind: SymbolKind::Database,
+        ptr: name_ptr,
+        schema: Schema::new("pg_database"),
+        params: None,
+    });
+
+    let root = b.root_scope();
+    b.scopes[root].insert(database_name, database_id);
+}
+
 fn item_name(path: &ast::Path) -> Option<Name> {
     let segment = path.segment()?;
 

crates/squawk_ide/src/binder.rs

@@ -9,7 +9,9 @@ pub(crate) enum NameRefClass {
     DropView,
     DropMaterializedView,
     DropSequence,
+    SequenceOwnedByColumn,
     Tablespace,
+    DropDatabase,
     ForeignKeyTable,
     ForeignKeyColumn,
     ForeignKeyLocalColumn,
@@ -32,6 +34,7 @@ pub(crate) enum NameRefClass {
     CreateSchema,
     CreateIndex,
     CreateIndexColumn,
+    DefaultConstraintFunctionCall,
     SelectFunctionCall,
     SelectFromTable,
     SelectColumn,
@@ -46,6 +49,7 @@ pub(crate) enum NameRefClass {
     UpdateWhereColumn,
     UpdateSetColumn,
     UpdateFromTable,
+    JoinUsingColumn,
     SchemaQualifier,
     TypeReference,
 }
@@ -56,6 +60,7 @@ pub(crate) fn classify_name_ref(name_ref: &ast::NameRef) -> Option<NameRefClass>
     let mut in_column_list = false;
     let mut in_where_clause = false;
     let mut in_from_clause = false;
+    let mut in_on_clause = false;
     let mut in_set_clause = false;
     let mut in_constraint_exclusion_list = false;
     let mut in_constraint_include_clause = false;
@@ -83,11 +88,15 @@ pub(crate) fn classify_name_ref(name_ref: &ast::NameRef) -> Option<NameRefClass>
             .is_some();
 
         let mut in_from_clause = false;
+        let mut in_on_clause = false;
         for ancestor in parent.ancestors() {
+            if ast::OnClause::can_cast(ancestor.kind()) {
+                in_on_clause = true;
+            }
             if ast::FromClause::can_cast(ancestor.kind()) {
                 in_from_clause = true;
             }
-            if ast::Select::can_cast(ancestor.kind()) && !in_from_clause {
+            if ast::Select::can_cast(ancestor.kind()) && (!in_from_clause || in_on_clause) {
                 if is_function_call || is_schema_table_col {
                     return Some(NameRefClass::SchemaQualifier);
                 } else {
@@ -119,15 +128,19 @@ pub(crate) fn classify_name_ref(name_ref: &ast::NameRef) -> Option<NameRefClass>
             .is_some();
 
         let mut in_from_clause = false;
+        let mut in_on_clause = false;
         let mut in_cast_expr = false;
         for ancestor in parent.ancestors() {
+            if ast::OnClause::can_cast(ancestor.kind()) {
+                in_on_clause = true;
+            }
             if ast::CastExpr::can_cast(ancestor.kind()) {
                 in_cast_expr = true;
             }
             if ast::FromClause::can_cast(ancestor.kind()) {
                 in_from_clause = true;
             }
-            if ast::Select::can_cast(ancestor.kind()) && !in_from_clause {
+            if ast::Select::can_cast(ancestor.kind()) && (!in_from_clause || in_on_clause) {
                 if in_cast_expr {
                     return Some(NameRefClass::TypeReference);
                 }
@@ -191,6 +204,15 @@ 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::DropDatabase::can_cast(ancestor.kind()) {
+            return Some(NameRefClass::DropDatabase);
+        }
+        if let Some(sequence_option) = ast::SequenceOption::cast(ancestor.clone())
+            && sequence_option.owned_token().is_some()
+            && sequence_option.by_token().is_some()
+        {
+            return Some(NameRefClass::SequenceOwnedByColumn);
+        }
         if ast::DropTablespace::can_cast(ancestor.kind())
             || ast::Tablespace::can_cast(ancestor.kind())
             || ast::SetTablespace::can_cast(ancestor.kind())
@@ -296,14 +318,20 @@ pub(crate) fn classify_name_ref(name_ref: &ast::NameRef) -> Option<NameRefClass>
         if ast::CallExpr::can_cast(ancestor.kind()) {
             in_call_expr = true;
         }
+        if ast::DefaultConstraint::can_cast(ancestor.kind()) && in_call_expr && !in_arg_list {
+            return Some(NameRefClass::DefaultConstraintFunctionCall);
+        }
+        if ast::OnClause::can_cast(ancestor.kind()) {
+            in_on_clause = true;
+        }
         if ast::FromClause::can_cast(ancestor.kind()) {
             in_from_clause = true;
         }
         if ast::Select::can_cast(ancestor.kind()) {
             if in_call_expr && !in_arg_list {
                 return Some(NameRefClass::SelectFunctionCall);
             }
-            if in_from_clause {
+            if in_from_clause && !in_on_clause {
                 return Some(NameRefClass::SelectFromTable);
             }
             // Classify as SelectColumn for target list, WHERE, ORDER BY, GROUP BY, etc.
@@ -331,6 +359,9 @@ pub(crate) fn classify_name_ref(name_ref: &ast::NameRef) -> Option<NameRefClass>
             }
             return Some(NameRefClass::InsertTable);
         }
+        if ast::JoinUsingClause::can_cast(ancestor.kind()) && in_column_list {
+            return Some(NameRefClass::JoinUsingColumn);
+        }
         if ast::WhereClause::can_cast(ancestor.kind()) {
             in_where_clause = true;
         }
@@ -371,6 +402,7 @@ pub(crate) enum NameClass {
     CreateIndex(ast::CreateIndex),
     CreateSequence(ast::CreateSequence),
     CreateTablespace(ast::CreateTablespace),
+    CreateDatabase(ast::CreateDatabase),
     CreateType(ast::CreateType),
     CreateFunction(ast::CreateFunction),
     CreateAggregate(ast::CreateAggregate),
@@ -411,6 +443,9 @@ pub(crate) fn classify_name(name: &ast::Name) -> Option<NameClass> {
         if let Some(create_tablespace) = ast::CreateTablespace::cast(ancestor.clone()) {
             return Some(NameClass::CreateTablespace(create_tablespace));
         }
+        if let Some(create_database) = ast::CreateDatabase::cast(ancestor.clone()) {
+            return Some(NameClass::CreateDatabase(create_database));
+        }
         if let Some(create_type) = ast::CreateType::cast(ancestor.clone()) {
             return Some(NameClass::CreateType(create_type));
         }

crates/squawk_ide/src/classify.rs

@@ -20,8 +20,8 @@ pub fn find_references(file: &ast::SourceFile, offset: TextSize) -> Vec<TextRang
         match_ast! {
             match node {
                 ast::NameRef(name_ref) => {
-                    if let Some(found) = resolve::resolve_name_ref(&binder, &name_ref)
-                      && found == target
+                    if let Some(found_refs) = resolve::resolve_name_ref(&binder, &name_ref)
+                        && found_refs.contains(&target)
                     {
                         refs.push(name_ref.syntax().text_range());
                     }
@@ -49,10 +49,12 @@ fn find_target(file: &ast::SourceFile, offset: TextSize, binder: &Binder) -> Opt
         return Some(SyntaxNodePtr::new(name.syntax()));
     }
 
-    if let Some(name_ref) = ast::NameRef::cast(parent.clone())
-        && let Some(ptr) = resolve::resolve_name_ref(binder, &name_ref)
-    {
-        return Some(ptr);
+    if let Some(name_ref) = ast::NameRef::cast(parent.clone()) {
+        // TODO: I think we want to return a list of targets so we can support cases like:
+        // select * from t join u using (id);
+        //                               ^ find refs
+        return resolve::resolve_name_ref(binder, &name_ref)
+            .and_then(|ptrs| ptrs.into_iter().next());
     }
 
     None