ide: support hover & goto def for function call style columns in where (#776)

Author: sbdchd

crates/squawk_ide/src/goto_definition.rs

@@ -1218,6 +1218,36 @@ select * from t where t.b$0 > 0;
         ");
     }
 
+    #[test]
+    fn goto_function_call_style_table_arg_in_where() {
+        assert_snapshot!(goto("
+create table t(a int);
+select * from t where a(t$0) > 2;
+"), @r"
+          ╭▸ 
+        2 │ create table t(a int);
+          │              ─ 2. destination
+        3 │ select * from t where a(t) > 2;
+          ╰╴                        ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_qualified_table_ref_in_where() {
+        assert_snapshot!(goto("
+create table t(a int);
+create function b(t) returns int as 'select 1' language sql;
+select * from t where t$0.b > 2;
+"), @r"
+          ╭▸ 
+        2 │ create table t(a int);
+          │              ─ 2. destination
+        3 │ create function b(t) returns int as 'select 1' language sql;
+        4 │ select * from t where t.b > 2;
+          ╰╴                      ─ 1. source
+        ");
+    }
+
     #[test]
     fn goto_function_call_style_in_order_by() {
         assert_snapshot!(goto("

crates/squawk_ide/src/hover.rs

@@ -359,9 +359,9 @@ fn is_select_from_table(name_ref: &ast::NameRef) -> bool {
 }
 
 fn is_select_column(name_ref: &ast::NameRef) -> bool {
-    let mut in_target_list = false;
     let mut in_call_expr = false;
     let mut in_arg_list = false;
+    let mut in_from_clause = false;
 
     for ancestor in name_ref.syntax().ancestors() {
         if ast::ArgList::can_cast(ancestor.kind()) {
@@ -370,14 +370,19 @@ fn is_select_column(name_ref: &ast::NameRef) -> bool {
         if ast::CallExpr::can_cast(ancestor.kind()) {
             in_call_expr = true;
         }
-        if ast::TargetList::can_cast(ancestor.kind()) {
-            in_target_list = true;
+        if ast::FromClause::can_cast(ancestor.kind()) {
+            in_from_clause = true;
         }
-        if ast::Select::can_cast(ancestor.kind()) && in_target_list {
+        if ast::Select::can_cast(ancestor.kind()) {
             // if we're inside a call expr but not inside an arg list, this is a function call
             if in_call_expr && !in_arg_list {
                 return false;
             }
+            // if we're in FROM clause, this is a table reference, not a column
+            if in_from_clause {
+                return false;
+            }
+            // anything else in SELECT (target list, WHERE, ORDER BY, etc.) is a column
             return true;
         }
     }
@@ -1078,6 +1083,33 @@ select b(t$0) from t;
         ");
     }
 
+    #[test]
+    fn hover_on_function_call_style_table_arg_in_where() {
+        assert_snapshot!(check_hover("
+create table t(a int);
+select * from t where a(t$0) > 2;
+"), @r"
+        hover: table public.t(a int)
+          ╭▸ 
+        3 │ select * from t where a(t) > 2;
+          ╰╴                        ─ hover
+        ");
+    }
+
+    #[test]
+    fn hover_on_qualified_table_ref_in_where() {
+        assert_snapshot!(check_hover("
+create table t(a int);
+create function b(t) returns int as 'select 1' language sql;
+select * from t where t$0.b > 2;
+"), @r"
+        hover: table public.t(a int)
+          ╭▸ 
+        4 │ select * from t where t.b > 2;
+          ╰╴                      ─ hover
+        ");
+    }
+
     #[test]
     fn hover_on_field_style_function_call() {
         assert_snapshot!(check_hover("

crates/squawk_ide/src/resolve.rs

@@ -127,7 +127,6 @@ fn classify_name_ref_context(name_ref: &ast::NameRef) -> Option<NameRefContext>
     let mut in_column_list = false;
     let mut in_where_clause = false;
     let mut in_from_clause = false;
-    let mut in_target_list = false;
 
     // TODO: can we combine this if and the one that follows?
     if let Some(parent) = name_ref.syntax().parent()
@@ -149,8 +148,12 @@ fn classify_name_ref_context(name_ref: &ast::NameRef) -> Option<NameRefContext>
             .and_then(ast::FieldExpr::cast)
             .is_some();
 
+        let mut in_from_clause = false;
         for ancestor in parent.ancestors() {
-            if ast::TargetList::can_cast(ancestor.kind()) {
+            if ast::FromClause::can_cast(ancestor.kind()) {
+                in_from_clause = true;
+            }
+            if ast::Select::can_cast(ancestor.kind()) && !in_from_clause {
                 if is_function_call || is_schema_table_col {
                     return Some(NameRefContext::SchemaQualifier);
                 } else {
@@ -246,19 +249,16 @@ fn classify_name_ref_context(name_ref: &ast::NameRef) -> Option<NameRefContext>
         if ast::FromClause::can_cast(ancestor.kind()) {
             in_from_clause = true;
         }
-        if ast::TargetList::can_cast(ancestor.kind()) {
-            in_target_list = true;
-        }
         if ast::Select::can_cast(ancestor.kind()) {
             if in_call_expr && !in_arg_list {
                 return Some(NameRefContext::SelectFunctionCall);
             }
             if in_from_clause {
                 return Some(NameRefContext::SelectFromTable);
             }
-            if in_target_list {
-                return Some(NameRefContext::SelectColumn);
-            }
+            // Classify as SelectColumn for target list, WHERE, ORDER BY, GROUP BY, etc.
+            // (anything in SELECT except FROM clause)
+            return Some(NameRefContext::SelectColumn);
         }
         if ast::ColumnList::can_cast(ancestor.kind()) {
             in_column_list = true;