From d841d311718ff56f92bbc125b52315c98eb77c16 Mon Sep 17 00:00:00 2001 From: Steve Dignam Date: Tue, 30 Dec 2025 13:13:17 -0500 Subject: [PATCH] ide: add code action for making inferred column alias explicit --- crates/squawk_ide/src/code_actions.rs | 180 +++++++++++++++++++++----- crates/squawk_ide/src/column_name.rs | 24 ++-- crates/squawk_ide/src/lib.rs | 1 + crates/squawk_ide/src/quote.rs | 105 +++++++++++++++ 4 files changed, 266 insertions(+), 44 deletions(-) create mode 100644 crates/squawk_ide/src/quote.rs diff --git a/crates/squawk_ide/src/code_actions.rs b/crates/squawk_ide/src/code_actions.rs index 0751b153..c3f16990 100644 --- a/crates/squawk_ide/src/code_actions.rs +++ b/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 Option { - 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, + 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;"# + ); + } } diff --git a/crates/squawk_ide/src/column_name.rs b/crates/squawk_ide/src/column_name.rs index 8742cf69..3c1d584d 100644 --- a/crates/squawk_ide/src/column_name.rs +++ b/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 { + 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() } } diff --git a/crates/squawk_ide/src/lib.rs b/crates/squawk_ide/src/lib.rs index cd59c0d1..0399f921 100644 --- a/crates/squawk_ide/src/lib.rs +++ b/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; diff --git a/crates/squawk_ide/src/quote.rs b/crates/squawk_ide/src/quote.rs new file mode 100644 index 00000000..482e0797 --- /dev/null +++ b/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 { + 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?""#); + } +}