ide: fix output column alias resolution in group by / order by

PLAN.md

@@ -540,6 +540,16 @@ select * from x;
 -- LINE 6:     select * from x;
 ```
 
+Note: this should only error if the with query name is used. So we need tracking
+of used / unused.
+
+The following is okay:
+
+```sql
+with t as (delete from k)
+select 1;
+```
+
 ### Rule: dialect: now() to dest
 
 should support various fixes so people can write in one dialect of SQL and have it easily convert to the other one

crates/squawk_ide/src/classify.rs

@@ -48,6 +48,8 @@ pub(crate) enum NameRefClass {
     Schema,
     SelectColumn,
     SelectFunctionCall,
+    SelectGroupByAliasOrColumn,
+    SelectOrderByAliasOrColumn,
     SelectQualifiedColumn,
     SelectQualifiedColumnTable,
     Sequence,
@@ -114,6 +116,8 @@ pub(crate) fn classify_name_ref(node: &SyntaxNode) -> Option<NameRefClass> {
     let mut in_when_clause = false;
     let mut in_special_sql_fn = false;
     let mut in_conflict_target = false;
+    let mut in_group_by_clause = false;
+    let mut in_order_by_clause = false;
 
     // TODO: can we combine this if and the one that follows?
     if let Some(parent) = node.parent()
@@ -663,6 +667,12 @@ pub(crate) fn classify_name_ref(node: &SyntaxNode) -> Option<NameRefClass> {
         if ast::FromClause::can_cast(ancestor.kind()) {
             in_from_clause = true;
         }
+        if ast::GroupByClause::can_cast(ancestor.kind()) {
+            in_group_by_clause = true;
+        }
+        if ast::OrderByClause::can_cast(ancestor.kind()) {
+            in_order_by_clause = true;
+        }
         if ast::Select::can_cast(ancestor.kind()) {
             if in_function_name && !in_special_sql_fn {
                 return Some(NameRefClass::SelectFunctionCall);
@@ -673,7 +683,19 @@ pub(crate) fn classify_name_ref(node: &SyntaxNode) -> Option<NameRefClass> {
                 }
                 return Some(NameRefClass::FromTable);
             }
-            // Classify as SelectColumn for target list, WHERE, ORDER BY, GROUP BY, etc.
+            if in_group_by_clause
+                && let Some(parent) = node.parent()
+                && ast::GroupingExpr::can_cast(parent.kind())
+            {
+                return Some(NameRefClass::SelectGroupByAliasOrColumn);
+            }
+            if in_order_by_clause
+                && let Some(parent) = node.parent()
+                && ast::SortBy::can_cast(parent.kind())
+            {
+                return Some(NameRefClass::SelectOrderByAliasOrColumn);
+            }
+            // Classify as SelectColumn for target list, WHERE, etc.
             // (anything in SELECT except FROM clause)
             return Some(NameRefClass::SelectColumn);
         }

crates/squawk_ide/src/goto_definition.rs

@@ -142,6 +142,7 @@ mod test {
     use rowan::TextRange;
     use rustc_hash::FxHashMap;
 
+    #[must_use]
     #[track_caller]
     fn goto(sql: &str) -> String {
         goto_(sql).expect("should always find a definition")
@@ -4752,6 +4753,95 @@ select * from t group by t.b$0;
         ");
     }
 
+    #[test]
+    fn goto_select_alias_order_by_column_name_conflict() {
+        // If an ORDER BY expression is a simple name that matches both an
+        // output column name and an input column name, ORDER BY will interpret
+        // it as the output column name.
+        assert_snapshot!(goto("
+with t as (select 2 a)
+select 1 a from t
+order by a$0;
+"), @"
+          ╭▸ 
+        3 │ select 1 a from t
+          │          ─ 2. destination
+        4 │ order by a;
+          ╰╴         ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_select_alias_group_by_column_name_conflict() {
+        // If a GROUP BY expression is a simple name that matches both output
+        // column name and an input column name, GROUP BY will interpret it as
+        // the input column name.
+        assert_snapshot!(goto("
+with t as (select 2 a)
+select 1 a from t
+group by a$0;
+"), @"
+          ╭▸ 
+        2 │ with t as (select 2 a)
+          │                     ─ 2. destination
+        3 │ select 1 a from t
+        4 │ group by a;
+          ╰╴         ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_select_alias_in_group_by_with_cte() {
+        assert_snapshot!(goto("
+with t as (select 2 b)
+select 1 a from t
+group by a$0;
+"), @"
+          ╭▸ 
+        3 │ select 1 a from t
+          │          ─ 2. destination
+        4 │ group by a;
+          ╰╴         ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_select_alias_expr_in_order_by_not_found() {
+        goto_not_found(
+            "
+with t as (select 2 b)
+select 1 a from t
+order by a$0 + 1
+",
+        );
+    }
+
+    #[test]
+    fn goto_select_alias_expr_in_group_by_expr_not_found() {
+        goto_not_found(
+            "
+with t as (select 2 b)
+select 1 a from t
+group by a$0 + 1
+",
+        );
+    }
+
+    #[test]
+    fn goto_select_alias_in_order_by_with_cte() {
+        assert_snapshot!(goto("
+with t as (select 2 b)
+select 1 a from t
+order by a$0;
+"), @"
+          ╭▸ 
+        3 │ select 1 a from t
+          │          ─ 2. destination
+        4 │ order by a;
+          ╰╴         ─ 1. source
+        ");
+    }
+
     #[test]
     fn goto_cte_table() {
         assert_snapshot!(goto("

crates/squawk_ide/src/resolve.rs

@@ -638,6 +638,12 @@ pub(crate) fn resolve_name_ref(
         }
         NameRefClass::CreateIndexColumn => resolve_create_index_column_ptr(db, file, name_ref),
         NameRefClass::SelectColumn => resolve_select_column_ptr(db, file, name_ref),
+        NameRefClass::SelectGroupByAliasOrColumn => {
+            resolve_select_group_by_alias_or_column_ptr(db, file, name_ref)
+        }
+        NameRefClass::SelectOrderByAliasOrColumn => {
+            resolve_select_order_by_alias_or_column_ptr(db, file, name_ref)
+        }
         NameRefClass::SelectQualifiedColumnTable => {
             resolve_select_qualified_column_table_name_ptr(db, file, name_ref)
         }
@@ -1728,6 +1734,52 @@ fn resolve_select_column_ptr(
     None
 }
 
+fn resolve_select_group_by_alias_or_column_ptr(
+    db: &dyn Db,
+    file: File,
+    column_name_ref: &ast::NameRef,
+) -> Option<SmallVec<[Location; 1]>> {
+    if let Some(ptr) = resolve_select_column_ptr(db, file, column_name_ref) {
+        return Some(ptr);
+    }
+    resolve_select_target_alias_ptr(file, column_name_ref)
+}
+
+fn resolve_select_order_by_alias_or_column_ptr(
+    db: &dyn Db,
+    file: File,
+    column_name_ref: &ast::NameRef,
+) -> Option<SmallVec<[Location; 1]>> {
+    if let Some(ptr) = resolve_select_target_alias_ptr(file, column_name_ref) {
+        return Some(ptr);
+    }
+    resolve_select_column_ptr(db, file, column_name_ref)
+}
+
+fn resolve_select_target_alias_ptr(
+    file: File,
+    column_name_ref: &ast::NameRef,
+) -> Option<SmallVec<[Location; 1]>> {
+    let select = column_name_ref
+        .syntax()
+        .ancestors()
+        .find_map(ast::Select::cast)?;
+    let column_name = Name::from_node(column_name_ref);
+    let target_list = select.select_clause()?.target_list()?;
+    for target in target_list.targets() {
+        if let Some(alias_name) = target.as_name().and_then(|a| a.name())
+            && Name::from_node(&alias_name) == column_name
+        {
+            return Some(smallvec![Location::new(
+                file,
+                alias_name.syntax().text_range(),
+                LocationKind::Column,
+            )]);
+        }
+    }
+    None
+}
+
 fn resolve_fn_call_column(
     db: &dyn Db,
     file: File,