ide: fix goto def for composite type subfield

crates/squawk_ide/src/classify.rs

@@ -476,6 +476,7 @@ pub(crate) fn classify_name_ref(node: &SyntaxNode) -> Option<NameRefClass> {
         let mut in_on_clause = false;
         let mut in_returning_clause = false;
         let mut in_set_clause = false;
+        let mut in_set_expr = false;
         let mut in_where_clause = false;
         let mut in_when_clause = false;
         let mut in_when_condition = false;
@@ -501,13 +502,21 @@ pub(crate) fn classify_name_ref(node: &SyntaxNode) -> Option<NameRefClass> {
             if ast::SetClause::can_cast(ancestor.kind()) {
                 in_set_clause = true;
             }
+            if ast::SetExpr::can_cast(ancestor.kind()) {
+                in_set_expr = true;
+            }
             if ast::WhereClause::can_cast(ancestor.kind()) {
                 in_where_clause = true;
             }
             if ast::MergeWhenClause::can_cast(ancestor.kind()) {
                 in_when_clause = true;
             }
             if ast::Update::can_cast(ancestor.kind()) {
+                // a set target can't be qualified with the relation name, so
+                // `a` in `update t set a.b = 1` is a composite type column
+                if in_set_clause && !in_set_expr {
+                    return Some(NameRefClass::UpdateColumn);
+                }
                 if in_returning_clause || in_set_clause || in_where_clause || in_from_clause {
                     if is_function_call || is_schema_table_col {
                         return Some(NameRefClass::Schema);
@@ -604,6 +613,8 @@ pub(crate) fn classify_name_ref(node: &SyntaxNode) -> Option<NameRefClass> {
         let mut in_when_clause = false;
         let mut in_when_condition = false;
         let mut in_returning_clause = false;
+        let mut in_set_clause = false;
+        let mut in_set_expr = false;
         for ancestor in parent.ancestors() {
             if ast::OnClause::can_cast(ancestor.kind()) {
                 in_on_clause = true;
@@ -626,6 +637,17 @@ pub(crate) fn classify_name_ref(node: &SyntaxNode) -> Option<NameRefClass> {
             if ast::ReturningClause::can_cast(ancestor.kind()) {
                 in_returning_clause = true;
             }
+            if ast::SetClause::can_cast(ancestor.kind()) {
+                in_set_clause = true;
+            }
+            if ast::SetExpr::can_cast(ancestor.kind()) {
+                in_set_expr = true;
+            }
+            // `x` in `update t set a.x = 1` is a field of the composite type of
+            // the column `a`
+            if ast::Update::can_cast(ancestor.kind()) && in_set_clause && !in_set_expr {
+                return Some(NameRefClass::CompositeTypeField);
+            }
             if ast::Merge::can_cast(ancestor.kind())
                 && (in_on_clause || in_when_clause || in_returning_clause)
             {

crates/squawk_ide/src/goto_definition.rs

@@ -12666,15 +12666,70 @@ returning t$0.*;"
 
     #[test]
     fn goto_update_alias_in_set_clause() {
-        assert_snapshot!(goto("
+        // Query 1: ERROR:  column "f" of relation "t" does not exist
+        // LINE 1: update t as f set f.a = 10;
+        //                           ^
+        // HINT:  SET target columns cannot be qualified with the relation name.
+        goto_not_found(
+            "
 create table t(a int, b int);
-update t as f set f$0.a = 10;"
-        ), @r"
-          ╭▸ 
-        3 │ update t as f set f.a = 10;
-          │             ┬     ─ 1. source
-          │             │
-          ╰╴            2. destination
+update t as f set f$0.a = 10;",
+        );
+    }
+
+    #[test]
+    fn goto_update_set_clause_subfield() {
+        // point isn't a composite type so we can't update it like this
+        goto_not_found(
+            "
+create table t(a point, b int);
+update t set a.x$0 = 10;",
+        );
+    }
+
+    #[test]
+    fn goto_update_set_clause_composite_type_subfield() {
+        assert_snapshot!(goto(
+            "
+create type coordinate as (
+  x double precision,
+  y double precision
+); 
+create table t (
+  a coordinate,
+  b integer
+);
+update t set a.x$0 = 10;",
+        ), @"
+           ╭▸ 
+         3 │   x double precision,
+           │   ─ 2. destination
+           ‡
+        10 │ update t set a.x = 10;
+           ╰╴               ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_update_set_clause_composite_type_col_name() {
+        assert_snapshot!(goto(
+            "
+create type coordinate as (
+  x double precision,
+  y double precision
+); 
+create table t (
+  a coordinate,
+  b integer
+);
+update t set a$0.x = 10;",
+        ), @"
+           ╭▸ 
+         7 │   a coordinate,
+           │   ─ 2. destination
+           ‡
+        10 │ update t set a.x = 10;
+           ╰╴             ─ 1. source
         ");
     }