ide: add code action for making inferred column alias explicit (#781)

Author: sbdchd

crates/squawk_ide/src/code_actions.rs

@@ -1,11 +1,15 @@
 use rowan::TextSize;
 use squawk_linter::Edit;
 use squawk_syntax::{
-    SyntaxKind, SyntaxNode,
+    SyntaxKind,
     ast::{self, AstNode},
 };
 
-use crate::{generated::keywords::RESERVED_KEYWORDS, offsets::token_from_offset};
+use crate::{
+    column_name::ColumnName,
+    offsets::token_from_offset,
+    quote::{quote_column_alias, unquote_ident},
+};
 
 #[derive(Debug, Clone)]
 pub enum ActionKind {
@@ -29,6 +33,7 @@ pub fn code_actions(file: ast::SourceFile, offset: TextSize) -> Option<Vec<CodeA
     rewrite_select_as_table(&mut actions, &file, offset);
     quote_identifier(&mut actions, &file, offset);
     unquote_identifier(&mut actions, &file, offset);
+    add_explicit_alias(&mut actions, &file, offset);
     Some(actions)
 }
 
@@ -346,7 +351,7 @@ fn unquote_identifier(
         return None;
     };
 
-    let unquoted = unquote(&name_node)?;
+    let unquoted = unquote_ident(&name_node)?;
 
     actions.push(CodeAction {
         title: "Unquote identifier".to_owned(),
@@ -357,45 +362,37 @@ fn unquote_identifier(
     Some(())
 }
 
-fn unquote(node: &SyntaxNode) -> Option<String> {
-    let text = node.text().to_string();
-
-    if !text.starts_with('"') || !text.ends_with('"') {
-        return None;
-    }
-
-    let text = &text[1..text.len() - 1];
-
-    if is_reserved_word(text) {
-        return None;
-    }
+// Postgres docs call these output names.
+// Postgres' parser calls this a column label.
+// Third-party docs call these aliases, so going with that.
+fn add_explicit_alias(
+    actions: &mut Vec<CodeAction>,
+    file: &ast::SourceFile,
+    offset: TextSize,
+) -> Option<()> {
+    let token = token_from_offset(file, offset)?;
+    let target = token.parent_ancestors().find_map(ast::Target::cast)?;
 
-    if text.is_empty() {
+    if target.as_name().is_some() {
         return None;
     }
 
-    let mut chars = text.chars();
+    let alias = ColumnName::from_target(target.clone()).and_then(|c| c.0.to_string())?;
 
-    // see: https://www.postgresql.org/docs/18/sql-syntax-lexical.html#SQL-SYNTAX-IDENTIFIERS
-    match chars.next() {
-        Some(c) if c.is_lowercase() || c == '_' => {}
-        _ => return None,
-    }
+    let expr_end = target.expr().map(|e| e.syntax().text_range().end())?;
 
-    for c in chars {
-        if c.is_lowercase() || c.is_ascii_digit() || c == '_' || c == '$' {
-            continue;
-        }
-        return None;
-    }
+    let quoted_alias = quote_column_alias(&alias);
+    // Postgres docs recommend either using `as` or quoting the name. I think
+    // `as` looks a bit nicer.
+    let replacement = format!(" as {}", quoted_alias);
 
-    Some(text.to_string())
-}
+    actions.push(CodeAction {
+        title: "Add explicit alias".to_owned(),
+        edits: vec![Edit::insert(replacement, expr_end)],
+        kind: ActionKind::RefactorRewrite,
+    });
 
-fn is_reserved_word(text: &str) -> bool {
-    RESERVED_KEYWORDS
-        .binary_search(&text.to_lowercase().as_str())
-        .is_ok()
+    Some(())
 }
 
 #[cfg(test)]
@@ -933,4 +930,119 @@ mod test {
             @"create table T(x int);"
         );
     }
+
+    #[test]
+    fn add_explicit_alias_simple_column() {
+        assert_snapshot!(apply_code_action(
+            add_explicit_alias,
+            "select col_na$0me from t;"),
+            @"select col_name as col_name from t;"
+        );
+    }
+
+    #[test]
+    fn add_explicit_alias_quoted_identifier() {
+        assert_snapshot!(apply_code_action(
+            add_explicit_alias,
+            r#"select "b"$0 from t;"#),
+            @r#"select "b" as b from t;"#
+        );
+    }
+
+    #[test]
+    fn add_explicit_alias_field_expr() {
+        assert_snapshot!(apply_code_action(
+            add_explicit_alias,
+            "select t.col$0umn from t;"),
+            @"select t.column as column from t;"
+        );
+    }
+
+    #[test]
+    fn add_explicit_alias_function_call() {
+        assert_snapshot!(apply_code_action(
+            add_explicit_alias,
+            "select cou$0nt(*) from t;"),
+            @"select count(*) as count from t;"
+        );
+    }
+
+    #[test]
+    fn add_explicit_alias_cast_to_type() {
+        assert_snapshot!(apply_code_action(
+            add_explicit_alias,
+            "select '1'::bigi$0nt from t;"),
+            @"select '1'::bigint as int8 from t;"
+        );
+    }
+
+    #[test]
+    fn add_explicit_alias_cast_column() {
+        assert_snapshot!(apply_code_action(
+            add_explicit_alias,
+            "select col_na$0me::text from t;"),
+            @"select col_name::text as col_name from t;"
+        );
+    }
+
+    #[test]
+    fn add_explicit_alias_case_expr() {
+        assert_snapshot!(apply_code_action(
+            add_explicit_alias,
+            "select ca$0se when true then 'a' end from t;"),
+            @"select case when true then 'a' end as case from t;"
+        );
+    }
+
+    #[test]
+    fn add_explicit_alias_case_with_else() {
+        assert_snapshot!(apply_code_action(
+            add_explicit_alias,
+            "select ca$0se when true then 'a' else now()::text end from t;"),
+            @"select case when true then 'a' else now()::text end as now from t;"
+        );
+    }
+
+    #[test]
+    fn add_explicit_alias_array() {
+        assert_snapshot!(apply_code_action(
+            add_explicit_alias,
+            "select arr$0ay[1, 2, 3] from t;"),
+            @"select array[1, 2, 3] as array from t;"
+        );
+    }
+
+    #[test]
+    fn add_explicit_alias_not_applicable_already_has_alias() {
+        assert!(code_action_not_applicable(
+            add_explicit_alias,
+            "select col_name$0 as foo from t;"
+        ));
+    }
+
+    #[test]
+    fn add_explicit_alias_unknown_column() {
+        assert_snapshot!(apply_code_action(
+            add_explicit_alias,
+            "select 1 $0+ 2 from t;"),
+            @r#"select 1 + 2 as "?column?" from t;"#
+        );
+    }
+
+    #[test]
+    fn add_explicit_alias_not_applicable_star() {
+        assert!(code_action_not_applicable(
+            add_explicit_alias,
+            "select $0* from t;"
+        ));
+    }
+
+    #[test]
+    fn add_explicit_alias_literal() {
+        assert_snapshot!(apply_code_action(
+            add_explicit_alias,
+            "select 'foo$0' from t;"),
+            @r#"select 'foo' as "?column?" from t;"#
+        );
+    }
 }

crates/squawk_ide/src/column_name.rs

@@ -30,7 +30,6 @@ pub(crate) enum ColumnName {
 }
 
 impl ColumnName {
-    #[allow(dead_code)]
     pub(crate) fn from_target(target: ast::Target) -> Option<(ColumnName, SyntaxNode)> {
         if let Some(as_name) = target.as_name()
             && let Some(name_node) = as_name.name()
@@ -55,6 +54,16 @@ impl ColumnName {
             ColumnName::Column(name)
         }
     }
+
+    pub(crate) fn to_string(&self) -> Option<String> {
+        match self {
+            ColumnName::Column(string) => Some(string.to_string()),
+            ColumnName::Star => None,
+            ColumnName::UnknownColumn(c) => {
+                Some(c.clone().unwrap_or_else(|| "?column?".to_string()))
+            }
+        }
+    }
 }
 
 fn name_from_type(ty: ast::Type, unknown_column: bool) -> Option<(ColumnName, SyntaxNode)> {
@@ -174,10 +183,9 @@ fn name_from_name_ref(name_ref: ast::NameRef, in_type: bool) -> Option<(ColumnNa
             }
         }
     }
-    return Some((
-        ColumnName::Column(name_ref.text().to_string()),
-        name_ref.syntax().clone(),
-    ));
+    let text = name_ref.text();
+    let normalized = normalize_identifier(&text);
+    return Some((ColumnName::Column(normalized), name_ref.syntax().clone()));
 }
 
 /*
@@ -418,11 +426,7 @@ fn examples() {
             .unwrap();
 
         ColumnName::from_target(target)
-            .map(|x| match x.0 {
-                ColumnName::Column(string) => string,
-                ColumnName::Star => unreachable!(),
-                ColumnName::UnknownColumn(c) => c.unwrap_or_else(|| "?column?".to_string()),
-            })
+            .and_then(|x| x.0.to_string())
             .unwrap()
     }
 }

crates/squawk_ide/src/lib.rs

@@ -9,6 +9,7 @@ pub mod goto_definition;
 pub mod hover;
 pub mod inlay_hints;
 mod offsets;
+mod quote;
 mod resolve;
 mod scope;
 mod symbols;

crates/squawk_ide/src/quote.rs

@@ -0,0 +1,105 @@
+use squawk_syntax::SyntaxNode;
+
+use crate::generated::keywords::RESERVED_KEYWORDS;
+
+pub(crate) fn quote_column_alias(text: &str) -> String {
+    if needs_quoting(text) {
+        format!(r#""{}""#, text.replace('"', r#""""#))
+    } else {
+        text.to_string()
+    }
+}
+
+pub(crate) fn unquote_ident(node: &SyntaxNode) -> Option<String> {
+    let text = node.text().to_string();
+
+    if !text.starts_with('"') || !text.ends_with('"') {
+        return None;
+    }
+
+    let text = &text[1..text.len() - 1];
+
+    if is_reserved_word(text) {
+        return None;
+    }
+
+    if text.is_empty() {
+        return None;
+    }
+
+    let mut chars = text.chars();
+
+    // see: https://www.postgresql.org/docs/18/sql-syntax-lexical.html#SQL-SYNTAX-IDENTIFIERS
+    match chars.next() {
+        Some(c) if c.is_lowercase() || c == '_' => {}
+        _ => return None,
+    }
+
+    for c in chars {
+        if c.is_lowercase() || c.is_ascii_digit() || c == '_' || c == '$' {
+            continue;
+        }
+        return None;
+    }
+
+    Some(text.to_string())
+}
+
+fn needs_quoting(text: &str) -> bool {
+    if text.is_empty() {
+        return true;
+    }
+
+    // Column labels in AS clauses allow all keywords, so we don't need to check
+    // for reserved words. See PostgreSQL grammar:
+    // ColLabel: IDENT | unreserved_keyword | col_name_keyword | type_func_name_keyword | reserved_keyword
+
+    let mut chars = text.chars();
+
+    match chars.next() {
+        Some(c) if c.is_lowercase() || c == '_' => {}
+        _ => return true,
+    }
+
+    for c in chars {
+        if c.is_lowercase() || c.is_ascii_digit() || c == '_' || c == '$' {
+            continue;
+        }
+        return true;
+    }
+
+    false
+}
+
+pub(crate) fn is_reserved_word(text: &str) -> bool {
+    RESERVED_KEYWORDS
+        .binary_search(&text.to_lowercase().as_str())
+        .is_ok()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn quote_column_alias_handles_embedded_quotes() {
+        assert_eq!(quote_column_alias(r#"foo"bar"#), r#""foo""bar""#);
+    }
+
+    #[test]
+    fn quote_column_alias_doesnt_quote_reserved_words() {
+        // Keywords are allowed as column labels in AS clauses
+        assert_eq!(quote_column_alias("case"), "case");
+        assert_eq!(quote_column_alias("array"), "array");
+    }
+
+    #[test]
+    fn quote_column_alias_doesnt_quote_simple_identifiers() {
+        assert_eq!(quote_column_alias("col_name"), "col_name");
+    }
+
+    #[test]
+    fn quote_column_alias_handles_special_column_name() {
+        assert_eq!(quote_column_alias("?column?"), r#""?column?""#);
+    }
+}