From 07955ea20f0b5f3c2ddc3e9506085495e0181fc9 Mon Sep 17 00:00:00 2001 From: psteinroe Date: Thu, 22 Oct 2026 16:48:52 +0000 Subject: [PATCH 1/2] feat: code action to expand * into the column list Co-Authored-By: Claude Opus 5.5 --- crates/pgls_catalog/src/resolve/mod.rs | 25 +++- .../src/resolve/nodes/res_target.rs | 111 ++++++++++++++- .../src/resolve/nodes/select_stmt.rs | 8 +- .../src/resolve/reference_tests.rs | 78 +++++++++++ crates/pgls_catalog/src/resolve/resolver.rs | 5 + crates/pgls_catalog/src/resolve/span.rs | 48 ++++++- crates/pgls_lsp/src/handlers/code_actions.rs | 9 +- crates/pgls_lsp/tests/server.rs | 108 +++++++++++++++ .../src/features/code_actions.rs | 3 +- crates/pgls_workspace/src/workspace/server.rs | 63 +++++++-- .../src/workspace/server.tests.rs | 110 ++++++++++++++- .../src/workspace/server/expand_star.rs | 127 ++++++++++++++++++ docs/features/editor_features.md | 14 ++ 13 files changed, 689 insertions(+), 20 deletions(-) create mode 100644 crates/pgls_workspace/src/workspace/server/expand_star.rs diff --git a/crates/pgls_catalog/src/resolve/mod.rs b/crates/pgls_catalog/src/resolve/mod.rs index 44abe6d6e..b8653bff1 100644 --- a/crates/pgls_catalog/src/resolve/mod.rs +++ b/crates/pgls_catalog/src/resolve/mod.rs @@ -204,6 +204,25 @@ pub enum Reference { span: Option, relation: crate::RelationInfo, }, + /// A `*` or `t.*` in a select list or RETURNING list whose columns are all known, and each + /// can be referenced by a name that resolves to it. + Star { + /// The location of the `*`, or of the qualifier of `t.*`. + location: i32, + /// The span of the `*`, including its qualifier. + span: Option, + /// The columns that the `*` expands to, in order. + columns: Vec, + }, +} + +/// A column that a `*` expands to, and the name that references it from where the `*` is. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct StarColumn { + /// The qualifier of the column, e.g. `["u"]` for `u.name`. Empty if the column name alone + /// references it. + pub qualifier: Vec, + pub name: String, } impl Reference { @@ -211,7 +230,8 @@ impl Reference { match self { Self::Function { location, .. } | Self::Column { location, .. } - | Self::Relation { location, .. } => *location, + | Self::Relation { location, .. } + | Self::Star { location, .. } => *location, } } @@ -219,7 +239,8 @@ impl Reference { match self { Self::Function { span, .. } | Self::Column { span, .. } - | Self::Relation { span, .. } => *span, + | Self::Relation { span, .. } + | Self::Star { span, .. } => *span, } } } diff --git a/crates/pgls_catalog/src/resolve/nodes/res_target.rs b/crates/pgls_catalog/src/resolve/nodes/res_target.rs index 8a41a6604..24c9d1344 100644 --- a/crates/pgls_catalog/src/resolve/nodes/res_target.rs +++ b/crates/pgls_catalog/src/resolve/nodes/res_target.rs @@ -2,10 +2,13 @@ use pgls_query::{Node, NodeEnum}; -use super::{Resolver, resolve_list, string::string_value}; +use super::{ + Resolver, resolve_list, + string::{string_value, string_values}, +}; use crate::resolve::column_name::figure_column_name; use crate::resolve::{ - FindingKind, + FindingKind, Reference, StarColumn, scope::{Columns, Item, Level, find_item}, }; @@ -51,11 +54,115 @@ pub(super) fn target_list_columns(r: &Resolver, targets: &[Node]) -> Columns { Some(columns) } +/// Records the columns that each `*` and `t.*` of a select list or RETURNING list expands to, +/// when references are collected. The FROM items of the list must be the innermost level. +pub(super) fn record_star_references(r: &mut Resolver, targets: &[Node]) { + if !r.collects_references() { + return; + } + for target in targets { + let Some(NodeEnum::ResTarget(target)) = target.node.as_ref() else { + continue; + }; + let Some(NodeEnum::ColumnRef(column)) = target.val.as_deref().and_then(|v| v.node.as_ref()) + else { + continue; + }; + let Some((last, qualifier)) = column.fields.split_last() else { + continue; + }; + if !matches!(last.node, Some(NodeEnum::AStar(_))) { + continue; + } + let columns = string_values(qualifier).and_then(|qualifier| match qualifier.as_slice() { + [] => r.levels.last().and_then(star_columns), + [name] => find_item(&r.levels, name).and_then(|item| item_columns(item, &qualifier)), + [schema, name] => r + .levels + .iter() + .rev() + .find_map(|level| { + level.items.iter().find(|item| { + item.schema.as_deref() == Some(schema.as_str()) + && item.name.as_deref() == Some(name.as_str()) + }) + }) + .and_then(|item| item_columns(item, &qualifier)), + _ => None, + }); + if let Some(columns) = columns { + r.record(Reference::Star { + location: column.location, + span: r.star_span(column.location), + columns, + }); + } + } +} + +/// The columns of `*`, like [`ExpandAllTables`]. A column is qualified with the name of its FROM +/// item when there are several. `None` if a column is not known, or if a column name doesn't +/// reference the column with certainty. +/// +/// Columns merged by `JOIN ... USING` and `NATURAL` joins are not supported: they come first in +/// the output of their join, which the level doesn't record, and the bare name of a merged +/// column can be ambiguous. +/// +/// [`ExpandAllTables`]: https://github.com/postgres/postgres/blob/REL_18_6/src/backend/parser/parse_target.c#L1296 +fn star_columns(level: &Level) -> Option> { + if level.items.is_empty() + || level.has_natural_join + || level.has_opaque_items + || !level.merged_columns.is_empty() + { + return None; + } + if let [item] = level.items.as_slice() { + return item_columns(item, &[]); + } + let mut columns = Vec::new(); + for (i, item) in level.items.iter().enumerate() { + let name = item.name.as_ref()?; + // The qualifier must reference only this item. + if level.items[..i] + .iter() + .any(|other| other.name.as_ref() == Some(name)) + { + return None; + } + columns.extend(item_columns(item, std::slice::from_ref(name))?); + } + Some(columns) +} + +/// The columns of `item`, qualified with `qualifier`. `None` if they are not known, or if two +/// have the same name, which then doesn't reference either. +fn item_columns(item: &Item, qualifier: &[String]) -> Option> { + let names = item.columns.as_ref()?; + if names + .iter() + .enumerate() + .any(|(i, name)| names[..i].contains(name)) + { + return None; + } + Some( + names + .iter() + .map(|name| StarColumn { + qualifier: qualifier.to_vec(), + name: name.clone(), + }) + .collect(), + ) +} + /// Resolves a RETURNING list against `level` and returns its columns. pub(super) fn resolve_returning_list(r: &mut Resolver, list: &[Node], mut level: Level) -> Columns { add_old_and_new(r, &mut level); r.enter_level(level); resolve_list(r, list); + record_star_references(r, list); let columns = target_list_columns(r, list); r.output = columns.as_ref().map(|names| { let mut typed = Vec::new(); diff --git a/crates/pgls_catalog/src/resolve/nodes/select_stmt.rs b/crates/pgls_catalog/src/resolve/nodes/select_stmt.rs index 02ea16060..977276c78 100644 --- a/crates/pgls_catalog/src/resolve/nodes/select_stmt.rs +++ b/crates/pgls_catalog/src/resolve/nodes/select_stmt.rs @@ -6,8 +6,11 @@ use pgls_query::{ }; use super::{ - Resolver, from_clause::resolve_from_clause, res_target::target_list_columns, resolve_list, - resolve_node, with_clause::resolve_with_clause, + Resolver, + from_clause::resolve_from_clause, + res_target::{record_star_references, target_list_columns}, + resolve_list, resolve_node, + with_clause::resolve_with_clause, }; use crate::resolve::column_name::figure_column_name; use crate::resolve::scope::Columns; @@ -136,6 +139,7 @@ fn resolve_select(r: &mut Resolver, n: &SelectStmt) -> Columns { r.enter_level(level); resolve_list(r, &n.target_list); + record_star_references(r, &n.target_list); for node in n.where_clause.iter().chain(&n.having_clause) { resolve_node(r, node); if let Some(expr) = node.node.as_ref() { diff --git a/crates/pgls_catalog/src/resolve/reference_tests.rs b/crates/pgls_catalog/src/resolve/reference_tests.rs index 4434b2b61..5fb0289a2 100644 --- a/crates/pgls_catalog/src/resolve/reference_tests.rs +++ b/crates/pgls_catalog/src/resolve/reference_tests.rs @@ -111,6 +111,22 @@ fn references(catalog: &Catalog, search_path: &[String], sql: &str) -> Vec { format!("relation {}.{}", relation.schema, relation.name) } + Reference::Star { columns, .. } => { + let columns = columns + .iter() + .map(|column| { + column + .qualifier + .iter() + .chain(std::iter::once(&column.name)) + .cloned() + .collect::>() + .join(".") + }) + .collect::>() + .join(", "); + format!("star {columns}") + } }; format!("{text} => {target}") }) @@ -289,6 +305,68 @@ async fn relation_references(test_db: PgPool) { ); } +/// The `*` references of `sql`. +fn stars_of(catalog: &Catalog, search_path: &[String], sql: &str) -> Vec { + references(catalog, search_path, sql) + .into_iter() + .filter(|reference| reference.contains("=> star ")) + .collect() +} + +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn star_references(test_db: PgPool) { + let (catalog, search_path) = catalog(&test_db).await; + let stars = |sql| stars_of(&catalog, &search_path, sql); + + assert_eq!(stars("select * from users"), ["* => star id, name"]); + // Columns are qualified when there are several FROM items. + assert_eq!( + stars("select *, p.* from public.users join posts p on p.user_id = users.id"), + [ + "* => star users.id, users.name, p.id, p.user_id, p.title", + "p.* => star p.id, p.user_id, p.title", + ] + ); + assert_eq!( + stars("select public.users.* from public.users"), + ["public.users.* => star public.users.id, public.users.name"] + ); + // Aliases rename the columns. + assert_eq!(stars("select * from users u(uid)"), ["* => star uid, name"]); + // CTEs, subqueries, and RETURNING lists. + assert_eq!( + stars( + "with c as (select id, name as n from users) \ + select c.*, s.* from c, (select count(*) as total from posts) s" + ), + ["c.* => star c.id, c.n", "s.* => star s.total",] + ); + assert_eq!( + stars("delete from users returning *"), + ["* => star id, name"] + ); + // The `*` of a subquery expands to its own FROM items. + assert_eq!( + stars("select exists (select * from posts where user_id = users.id) from users"), + ["* => star id, user_id, title"] + ); + // Unknown columns, duplicate column names, merged columns, and items without a name have + // no reference. + assert_eq!(stars("select * from nope"), [""; 0]); + assert_eq!(stars("select * from users, nope"), [""; 0]); + assert_eq!( + stars("select j.* from (users join posts on true) j"), + [""; 0] + ); + assert_eq!(stars("select * from users join posts using (id)"), [""; 0]); + assert_eq!(stars("select * from users natural join posts"), [""; 0]); + assert_eq!(stars("select * from users, (select 1 as one)"), [""; 0]); + assert_eq!( + stars("select 1 from users where (users.*) is not null"), + [""; 0] + ); +} + #[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] async fn references_are_only_collected_on_request(test_db: PgPool) { let (catalog, search_path) = catalog(&test_db).await; diff --git a/crates/pgls_catalog/src/resolve/resolver.rs b/crates/pgls_catalog/src/resolve/resolver.rs index eb4fd48d0..82f67c27b 100644 --- a/crates/pgls_catalog/src/resolve/resolver.rs +++ b/crates/pgls_catalog/src/resolve/resolver.rs @@ -196,6 +196,11 @@ impl<'a> Resolver<'a> { self.span_to(location, span::reference_end) } + /// The span of the `*` or `t.*` starting at `location`. + pub fn star_span(&self, location: i32) -> Option { + self.span_to(location, span::star_end) + } + /// The span of the operator token starting at `location`. pub fn operator_span(&self, location: i32) -> Option { self.span_to(location, span::operator_end) diff --git a/crates/pgls_catalog/src/resolve/span.rs b/crates/pgls_catalog/src/resolve/span.rs index 93f93d8e2..aaa1ed2bc 100644 --- a/crates/pgls_catalog/src/resolve/span.rs +++ b/crates/pgls_catalog/src/resolve/span.rs @@ -50,6 +50,39 @@ pub(super) fn reference_end(sql: &str, start: usize) -> Option { Some(cursor) } +/// The end of `*`, or of a qualified `t.*` or `"s"."t".*`, starting at `start`. +pub(super) fn star_end(sql: &str, start: usize) -> Option { + let bytes = sql.as_bytes(); + let mut cursor = start; + loop { + match *bytes.get(cursor)? { + b'*' => return Some(cursor + 1), + b'.' => cursor += 1, + byte if byte.is_ascii_whitespace() => cursor += 1, + b'"' => { + cursor += 1; + loop { + if *bytes.get(cursor)? == b'"' { + cursor += 1; + if bytes.get(cursor) != Some(&b'"') { + break; + } + } + cursor += 1; + } + } + byte if byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'$') || byte >= 0x80 => { + while bytes.get(cursor).is_some_and(|byte| { + byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'$') || *byte >= 0x80 + }) { + cursor += 1; + } + } + _ => return None, + } + } +} + pub(super) fn token_end(sql: &str, start: usize) -> Option { let bytes = sql.as_bytes(); let mut cursor = start; @@ -156,7 +189,7 @@ pub(super) fn operator_end(sql: &str, start: usize) -> Option { #[cfg(test)] mod tests { - use super::{cast_end, operator_end, reference_end}; + use super::{cast_end, operator_end, reference_end, star_end}; fn text(sql: &str, start: usize, end: Option) -> &str { &sql[start..end.expect("span")] @@ -187,4 +220,17 @@ mod tests { let start = sql.find("date '").unwrap(); assert_eq!(text(sql, start, cast_end(sql, start)), "date 'today'"); } + + #[test] + fn star_spans_cover_the_qualifier() { + let sql = r#"select *, t.*, "My ""T""" . *, s.t.* from t"#; + let start = sql.find('*').unwrap(); + assert_eq!(text(sql, start, star_end(sql, start)), "*"); + let start = sql.find("t.*").unwrap(); + assert_eq!(text(sql, start, star_end(sql, start)), "t.*"); + let start = sql.find("\"My").unwrap(); + assert_eq!(text(sql, start, star_end(sql, start)), r#""My ""T""" . *"#); + let start = sql.find("s.t").unwrap(); + assert_eq!(text(sql, start, star_end(sql, start)), "s.t.*"); + } } diff --git a/crates/pgls_lsp/src/handlers/code_actions.rs b/crates/pgls_lsp/src/handlers/code_actions.rs index b3c2f6750..1219ec538 100644 --- a/crates/pgls_lsp/src/handlers/code_actions.rs +++ b/crates/pgls_lsp/src/handlers/code_actions.rs @@ -99,11 +99,14 @@ pub fn get_actions( .diagnostics .iter() .filter(|diagnostic| { - diagnostic.code == Some(ls_types::NumberOrString::String(edit.rule.clone())) + edit.rule.as_ref().is_some_and(|rule| { + diagnostic.code == Some(ls_types::NumberOrString::String(rule.clone())) + }) }) .cloned() .collect(); - let is_suppression = edit.category.matches(SUPPRESSION_ACTION_CATEGORY); + let is_fix = + edit.rule.is_some() && !edit.category.matches(SUPPRESSION_ACTION_CATEGORY); Some(CodeAction { title: action.title, @@ -115,7 +118,7 @@ pub fn get_actions( changes: Some(HashMap::from([(url.clone(), edits)])), ..Default::default() }), - is_preferred: (!is_suppression && edit.applicability == Applicability::Always) + is_preferred: (is_fix && edit.applicability == Applicability::Always) .then_some(true), disabled: action .disabled_reason diff --git a/crates/pgls_lsp/tests/server.rs b/crates/pgls_lsp/tests/server.rs index c58197149..d21f29e68 100644 --- a/crates/pgls_lsp/tests/server.rs +++ b/crates/pgls_lsp/tests/server.rs @@ -1571,6 +1571,114 @@ async fn test_lint_fix_code_actions(test_db: PgPool) -> Result<()> { Ok(()) } +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn test_expand_star_code_action(test_db: PgPool) -> Result<()> { + let factory = ServerFactory::default(); + let mut fs = MemoryFileSystem::default(); + + test_db + .execute("create table star_users (id int, name text);") + .await?; + + let database = test_db + .connect_options() + .get_database() + .unwrap() + .to_string(); + let host = test_db.connect_options().get_host().to_string(); + + let mut conf = PartialConfiguration::init(); + conf.merge_with(PartialConfiguration { + db: Some(PartialDatabaseConfiguration { + database: Some(database), + host: Some(host), + ..Default::default() + }), + ..Default::default() + }); + + fs.insert( + url!("postgres-language-server.jsonc") + .to_file_path() + .unwrap() + .into_owned(), + serde_json::to_string_pretty(&conf).unwrap(), + ); + + let (service, client) = factory + .create_with_fs(None, DynRef::Owned(Box::new(fs))) + .into_inner(); + + let (stream, sink) = client.split(); + let mut server = Server::new(service); + + let (sender, _) = channel(CHANNEL_BUFFER_SIZE); + let reader = tokio::spawn(client_handler(stream, sink, sender)); + + server.initialize().await?; + server.initialized().await?; + + server.load_configuration().await?; + + let doc_content = "select * from star_users;"; + let doc_url = url!("test.sql"); + server + .open_named_document(doc_content.to_string(), doc_url.clone(), "sql") + .await?; + + let code_actions_response = server + .request::( + "textDocument/codeAction", + "_code_action", + CodeActionParams { + text_document: TextDocumentIdentifier { + uri: doc_url.clone(), + }, + range: Range { + start: Position::new(0, 7), + end: Position::new(0, 7), + }, + context: CodeActionContext::default(), + partial_result_params: PartialResultParams::default(), + work_done_progress_params: WorkDoneProgressParams::default(), + }, + ) + .await? + .unwrap(); + + let action = code_actions_response + .into_iter() + .find_map(|action| match action { + lsp::CodeActionOrCommand::CodeAction(action) + if action.title == "Expand `*` into its columns" => + { + Some(action) + } + _ => None, + }) + .expect("expected the expand star action"); + assert_eq!( + action.kind, + Some(lsp::CodeActionKind::from("refactor.rewrite.pgls")) + ); + assert_eq!(action.is_preferred, None); + let edits = &action.edit.as_ref().unwrap().changes.as_ref().unwrap()[&doc_url]; + assert_eq!(edits.len(), 1); + assert_eq!(edits[0].new_text, "id, name"); + assert_eq!( + edits[0].range, + Range { + start: Position::new(0, 7), + end: Position::new(0, 8), + } + ); + + server.shutdown().await?; + reader.abort(); + + Ok(()) +} + #[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] async fn test_invalidate_schema_cache(test_db: PgPool) -> Result<()> { let factory = ServerFactory::default(); diff --git a/crates/pgls_workspace/src/features/code_actions.rs b/crates/pgls_workspace/src/features/code_actions.rs index d47ecbaa3..ddeb8a62d 100644 --- a/crates/pgls_workspace/src/features/code_actions.rs +++ b/crates/pgls_workspace/src/features/code_actions.rs @@ -39,7 +39,8 @@ pub enum CodeActionKind { #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] pub struct EditAction { /// The category of the diagnostic the action fixes or suppresses, e.g. `lint/banDropTable`. - pub rule: String, + /// `None` for refactors. + pub rule: Option, pub category: ActionCategory, pub applicability: Applicability, /// Non-overlapping edits of the file, ordered by their position. diff --git a/crates/pgls_workspace/src/workspace/server.rs b/crates/pgls_workspace/src/workspace/server.rs index b186db9f2..895258914 100644 --- a/crates/pgls_workspace/src/workspace/server.rs +++ b/crates/pgls_workspace/src/workspace/server.rs @@ -90,6 +90,7 @@ mod connection_key; #[cfg(feature = "db")] mod connection_manager; pub(crate) mod document; +mod expand_star; mod migration; mod pg_query; #[cfg(feature = "db")] @@ -394,7 +395,7 @@ impl WorkspaceServer { actions.push(CodeAction { title: fix.message.clone(), kind: CodeActionKind::Edit(EditAction { - rule: category.to_string(), + rule: Some(category.to_string()), category: ActionCategory::QuickFix(Cow::Owned(rule.to_string())), applicability: fix.applicability, edits: fix @@ -422,7 +423,7 @@ impl WorkspaceServer { actions.push(CodeAction { title: format!("Suppress rule {rule} for this statement"), kind: CodeActionKind::Edit(EditAction { - rule: category.to_string(), + rule: Some(category.to_string()), category: ActionCategory::Other(Cow::Borrowed(SUPPRESSION_ACTION_CATEGORY)), applicability: Applicability::MaybeIncorrect, edits: vec![TextEdit { @@ -436,6 +437,45 @@ impl WorkspaceServer { actions } + /// The refactors offered at `position`: expanding a `*` into its columns. + fn refactor_code_actions(&self, doc: &Document, position: TextSize) -> Vec { + // Code actions are requested whenever the cursor moves, so don't resolve the file unless + // there is a `*` on the line. + let content = doc.get_document_content(); + let bytes = content.as_bytes(); + let offset = usize::from(position).min(bytes.len()); + let line_start = bytes[..offset] + .iter() + .rposition(|byte| *byte == b'\n') + .map_or(0, |i| i + 1); + let line_end = bytes[offset..] + .iter() + .position(|byte| *byte == b'\n') + .map_or(bytes.len(), |i| offset + i); + if !bytes[line_start..line_end].contains(&b'*') { + return Vec::new(); + } + + #[cfg(feature = "db")] + let snapshot = self.workspaces().settings().and_then(|settings| { + self.connection + .with_pool(&settings.db, |pool| self.snapshot.load(pool)) + .and_then(Result::ok) + }); + + #[cfg(not(feature = "db"))] + let snapshot = self.snapshot.get(); + + let Some(snapshot) = snapshot else { + return Vec::new(); + }; + let (stmt_start, references) = self.references_at(doc, &snapshot, position); + references + .iter() + .filter_map(|reference| expand_star::expand_star_action(reference, content, stmt_start)) + .collect() + } + /// The lint diagnostics of a document that are not suppressed. fn lint_diagnostics( &self, @@ -546,13 +586,14 @@ impl WorkspaceServer { } /// The references whose span contains `position`, in the innermost statement at `position` - /// (the body of a SQL function rather than its `CREATE FUNCTION`). + /// (the body of a SQL function rather than its `CREATE FUNCTION`). Their spans are relative + /// to the start of that statement, which is returned with them. fn references_at( &self, doc: &Document, snapshot: &Arc, position: TextSize, - ) -> Vec { + ) -> (TextSize, Vec) { let mut innermost: Option<(TextRange, Vec)> = None; self.walk_statements( doc, @@ -582,17 +623,18 @@ impl WorkspaceServer { }, ); let Some((range, references)) = innermost else { - return Vec::new(); + return (TextSize::default(), Vec::new()); }; let position = position - range.start(); - references + let references = references .into_iter() .filter(|reference| { reference .span() .is_some_and(|span| span.contains_inclusive(position)) }) - .collect() + .collect(); + (range.start(), references) } /// The indexed database snapshot, built once per snapshot. @@ -1026,6 +1068,11 @@ impl Workspace for WorkspaceServer { }) .collect(); + actions.splice( + 0..0, + self.refactor_code_actions(parser, params.cursor_position), + ); + if let Some(settings) = settings { let mut lint_actions = self.lint_code_actions(settings, ¶ms, parser); lint_actions.append(&mut actions); @@ -1767,7 +1814,7 @@ impl Workspace for WorkspaceServer { { Some((stmt_id, range, ts_tree, maybe_ast)) => { let position_in_stmt = params.position - range.start(); - let references = self.references_at(doc, &snapshot, params.position); + let (_, references) = self.references_at(doc, &snapshot, params.position); let snapshot = self.snapshot_before(doc, range.start(), snapshot); let markdown_blocks = pgls_hover::on_hover(pgls_hover::OnHoverParams { diff --git a/crates/pgls_workspace/src/workspace/server.tests.rs b/crates/pgls_workspace/src/workspace/server.tests.rs index 8f0e857a0..6c60374f5 100644 --- a/crates/pgls_workspace/src/workspace/server.tests.rs +++ b/crates/pgls_workspace/src/workspace/server.tests.rs @@ -3,7 +3,7 @@ use std::{ time::{Duration, Instant}, }; -use pgls_analyse::{Applicability, RuleCategories}; +use pgls_analyse::{ActionCategory, Applicability, RefactorKind, RuleCategories}; use pgls_configuration::{Merge, StringSet}; use pgls_configuration::{ PartialConfiguration, PartialFormatConfiguration, PartialTypecheckConfiguration, @@ -2122,6 +2122,114 @@ async fn test_hover_shows_resolved_references(test_db: PgPool) { assert_eq!(blocks.len(), 2, "{blocks:?}"); } +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn test_code_actions_expand_star(test_db: PgPool) { + test_db + .execute(r#"create table star_items (id int, "Label" text);"#) + .await + .expect("Failed to setup test database"); + + let mut conf = PartialConfiguration::init(); + conf.merge_with(PartialConfiguration { + db: Some(PartialDatabaseConfiguration { + database: Some( + test_db + .connect_options() + .get_database() + .unwrap() + .to_string(), + ), + ..Default::default() + }), + ..Default::default() + }); + let workspace = get_test_workspace(Some(conf)).expect("Unable to create test workspace"); + + let path = PgLSPath::new("test.sql"); + let content = "create table star_tags (id int, item_id int, \"select\" text);\n\ + select * from star_items i join star_tags t on t.item_id = i.id;\n\ + select t.* from star_tags t where id > 0;\n\ + create function all_items() returns setof star_items language sql\n\ + as 'select * from star_items';\n"; + workspace + .open_file(OpenFileParams { + path: path.clone(), + content: content.into(), + version: 1, + }) + .expect("Unable to open test file"); + + let expand_actions = |needle: &str| { + let offset = content.find(needle).expect("needle in content"); + workspace + .pull_code_actions(crate::features::code_actions::CodeActionsParams { + path: path.clone(), + cursor_position: u32::try_from(offset).unwrap().into(), + only: vec![], + skip: vec![], + }) + .expect("Unable to pull code actions") + .actions + .into_iter() + .filter_map(|action| match action.kind { + CodeActionKind::Edit(edit) if edit.rule.is_none() => Some((action.title, edit)), + _ => None, + }) + .collect::>() + }; + let range_of = |needle: &str| { + let start = content.find(needle).expect("needle in content"); + TextRange::at( + u32::try_from(start).unwrap().into(), + u32::try_from(needle.len()).unwrap().into(), + ) + }; + + // The columns of a table of the database and of one created in the file, qualified since + // there are two FROM items. + let actions = expand_actions("* from star_items"); + assert_eq!(actions.len(), 1, "{actions:?}"); + let (title, edit) = &actions[0]; + assert_eq!(title, "Expand `*` into its columns"); + assert_eq!( + edit.category, + ActionCategory::Refactor(RefactorKind::Rewrite) + ); + assert_eq!( + edit.edits, + vec![TextEdit { + range: range_of("*"), + new_text: r#"i.id, i."Label", t.id, t.item_id, t."select""#.into(), + }] + ); + + // `t.*`, with the cursor on its qualifier. + let actions = expand_actions("t.* from"); + assert_eq!(actions.len(), 1, "{actions:?}"); + assert_eq!(actions[0].0, "Expand `t.*` into its columns"); + assert_eq!( + actions[0].1.edits, + vec![TextEdit { + range: range_of("t.*"), + new_text: r#"t.id, t.item_id, t."select""#.into(), + }] + ); + + // In the body of a SQL function. + let actions = expand_actions("* from star_items'"); + assert_eq!(actions.len(), 1, "{actions:?}"); + assert_eq!( + actions[0].1.edits, + vec![TextEdit { + range: TextRange::at(range_of("* from star_items'").start(), 1.into()), + new_text: r#"id, "Label""#.into(), + }] + ); + + // Nothing to expand elsewhere. + assert!(expand_actions("where id").is_empty()); +} + #[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] async fn test_inlay_hints_follow_migration_sections(test_db: PgPool) { let mut conf = PartialConfiguration::init(); diff --git a/crates/pgls_workspace/src/workspace/server/expand_star.rs b/crates/pgls_workspace/src/workspace/server/expand_star.rs new file mode 100644 index 000000000..9a29e01da --- /dev/null +++ b/crates/pgls_workspace/src/workspace/server/expand_star.rs @@ -0,0 +1,127 @@ +//! The refactor that replaces a `*` in a select list with the columns it expands to. + +use pgls_analyse::{ActionCategory, Applicability, RefactorKind}; +use pgls_catalog::resolve::{Reference, StarColumn}; +use pgls_query::protobuf::KeywordKind; +use pgls_text_size::TextSize; + +use crate::features::code_actions::{CodeAction, CodeActionKind, EditAction, TextEdit}; + +/// The action that replaces the `*` or `t.*` of `reference` with its columns. `content` is the +/// text of the file, and the statement of the reference starts at `stmt_start`. +pub(super) fn expand_star_action( + reference: &Reference, + content: &str, + stmt_start: TextSize, +) -> Option { + let Reference::Star { span, columns, .. } = reference else { + return None; + }; + let range = (*span)? + stmt_start; + let star = content.get(usize::from(range.start())..usize::from(range.end()))?; + let new_text = column_list(columns); + // In the body of a function quoted with `'`, a quote in a column name would end the body. + if new_text.contains('\'') + && content + .get(..usize::from(stmt_start)) + .is_some_and(|before| before.ends_with('\'')) + { + return None; + } + Some(CodeAction { + title: format!("Expand `{star}` into its columns"), + kind: CodeActionKind::Edit(EditAction { + rule: None, + category: ActionCategory::Refactor(RefactorKind::Rewrite), + applicability: Applicability::Always, + edits: vec![TextEdit { range, new_text }], + }), + disabled_reason: None, + }) +} + +/// The columns as a select list, e.g. `u.id, u."Name"`. +fn column_list(columns: &[StarColumn]) -> String { + columns + .iter() + .map(|column| { + column + .qualifier + .iter() + .chain(std::iter::once(&column.name)) + .map(|name| quote_identifier(name)) + .collect::>() + .join(".") + }) + .collect::>() + .join(", ") +} + +/// `name`, quoted unless it reads as the same name without quotes: lower case, and not a keyword +/// other than an unreserved one. +fn quote_identifier(name: &str) -> String { + let plain = name + .chars() + .next() + .is_some_and(|first| first.is_ascii_lowercase() || first == '_') + && name + .chars() + .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || matches!(c, '_' | '$')); + let keyword = || { + pgls_query::scan(name).map_or(true, |scan| { + scan.tokens.first().is_none_or(|token| { + !matches!( + token.keyword_kind(), + KeywordKind::NoKeyword | KeywordKind::UnreservedKeyword + ) + }) + }) + }; + if plain && !keyword() { + name.to_owned() + } else { + format!("\"{}\"", name.replace('"', "\"\"")) + } +} + +#[cfg(test)] +mod tests { + use super::{StarColumn, column_list, quote_identifier}; + + #[test] + fn quotes_names_that_need_it() { + assert_eq!(quote_identifier("user_id"), "user_id"); + assert_eq!(quote_identifier("_x$1"), "_x$1"); + // Unreserved keywords can be column names. + assert_eq!(quote_identifier("name"), "name"); + assert_eq!(quote_identifier("Name"), "\"Name\""); + assert_eq!(quote_identifier("1st"), "\"1st\""); + assert_eq!(quote_identifier("my col"), "\"my col\""); + assert_eq!(quote_identifier("a\"b"), "\"a\"\"b\""); + assert_eq!(quote_identifier("select"), "\"select\""); + assert_eq!(quote_identifier("between"), "\"between\""); + assert_eq!(quote_identifier("left"), "\"left\""); + } + + #[test] + fn qualifies_columns() { + let columns = [ + StarColumn { + qualifier: vec!["u".into()], + name: "id".into(), + }, + StarColumn { + qualifier: vec!["public".into(), "Posts".into()], + name: "user".into(), + }, + StarColumn { + qualifier: vec![], + name: "title".into(), + }, + ]; + assert_eq!( + column_list(&columns), + r#"u.id, public."Posts"."user", title"# + ); + } +} diff --git a/docs/features/editor_features.md b/docs/features/editor_features.md index 2d8c598c0..e8184a4a9 100644 --- a/docs/features/editor_features.md +++ b/docs/features/editor_features.md @@ -43,3 +43,17 @@ The language server caches schema information on startup. Function calls show the names of the parameters their positional arguments are passed to, as in `make_point(x: 1, y: 2)`. Hints are only shown when the call resolves to exactly one function, including functions created earlier in the same file. Named arguments (`f(a => 1)`), unnamed parameters, and arguments that already carry the parameter's name (`f(t.x)` for the parameter `x`) get no hint. Of the elements passed to a `VARIADIC` parameter, only the first is labelled. Inlay hints require a database connection. Enable them in your editor, e.g. with `editor.inlayHints.enabled` in VS Code. + +## Expand `*` + +With the cursor on a `*` or `t.*` in a select list or `RETURNING` list, the refactor code action "Expand `*` into its columns" replaces it with the columns it stands for: + +```sql +select * from users u join posts p on p.user_id = u.id; +-- becomes +select u.id, u.name, p.id, p.user_id, p.title from users u join posts p on p.user_id = u.id; +``` + +Columns are qualified with the name or alias of their table when there are several in `FROM`, and quoted where needed. The columns of CTEs, subqueries, and tables created earlier in the file are expanded too. The action is only offered when every column is known and can be referenced unambiguously, so not for joins with `USING` or `NATURAL`, tables whose columns are unknown, or sources with two columns of the same name. + +Like hover, it requires a database connection. From 6d177b373806f756fcffbb6eafd28f023b1e7129 Mon Sep 17 00:00:00 2001 From: psteinroe Date: Sat, 24 Oct 2026 07:38:22 +0000 Subject: [PATCH 2/2] fix: expand * only on real star tokens and quote non-ASCII names minimally The code action resolved the file whenever the cursor's line contained a *, including /* comments, count(*) and multiplications. It now scans the line and only resolves when the cursor is on a * of a select list or on its qualifier, also inside a one-line function body. Names with characters that aren't ASCII are no longer quoted: Postgres accepts them unquoted and doesn't fold them. --- crates/pgls_workspace/src/workspace/server.rs | 14 +- .../src/workspace/server.tests.rs | 12 +- .../src/workspace/server/expand_star.rs | 140 +++++++++++++++++- 3 files changed, 146 insertions(+), 20 deletions(-) diff --git a/crates/pgls_workspace/src/workspace/server.rs b/crates/pgls_workspace/src/workspace/server.rs index 895258914..d4590d679 100644 --- a/crates/pgls_workspace/src/workspace/server.rs +++ b/crates/pgls_workspace/src/workspace/server.rs @@ -440,19 +440,9 @@ impl WorkspaceServer { /// The refactors offered at `position`: expanding a `*` into its columns. fn refactor_code_actions(&self, doc: &Document, position: TextSize) -> Vec { // Code actions are requested whenever the cursor moves, so don't resolve the file unless - // there is a `*` on the line. + // the cursor is on a `*` of a select list. let content = doc.get_document_content(); - let bytes = content.as_bytes(); - let offset = usize::from(position).min(bytes.len()); - let line_start = bytes[..offset] - .iter() - .rposition(|byte| *byte == b'\n') - .map_or(0, |i| i + 1); - let line_end = bytes[offset..] - .iter() - .position(|byte| *byte == b'\n') - .map_or(bytes.len(), |i| offset + i); - if !bytes[line_start..line_end].contains(&b'*') { + if !expand_star::is_at_star(content, position) { return Vec::new(); } diff --git a/crates/pgls_workspace/src/workspace/server.tests.rs b/crates/pgls_workspace/src/workspace/server.tests.rs index 6c60374f5..90375eefc 100644 --- a/crates/pgls_workspace/src/workspace/server.tests.rs +++ b/crates/pgls_workspace/src/workspace/server.tests.rs @@ -2150,7 +2150,10 @@ async fn test_code_actions_expand_star(test_db: PgPool) { select * from star_items i join star_tags t on t.item_id = i.id;\n\ select t.* from star_tags t where id > 0;\n\ create function all_items() returns setof star_items language sql\n\ - as 'select * from star_items';\n"; + as 'select * from star_items';\n\ + create table größen (größe int, \"Ölstand\" int, \"Café\" int);\n\ + select * from größen;\n\ + select count(*) /* * */ from größen;\n"; workspace .open_file(OpenFileParams { path: path.clone(), @@ -2226,8 +2229,15 @@ async fn test_code_actions_expand_star(test_db: PgPool) { }] ); + // Names that aren't ASCII aren't quoted, unless they have upper case ASCII letters. + let actions = expand_actions("* from größen;"); + assert_eq!(actions.len(), 1, "{actions:?}"); + assert_eq!(actions[0].1.edits[0].new_text, r#"größe, Ölstand, "Café""#); + // Nothing to expand elsewhere. assert!(expand_actions("where id").is_empty()); + assert!(expand_actions("*) /*").is_empty()); + assert!(expand_actions("* */").is_empty()); } #[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] diff --git a/crates/pgls_workspace/src/workspace/server/expand_star.rs b/crates/pgls_workspace/src/workspace/server/expand_star.rs index 9a29e01da..2b7807087 100644 --- a/crates/pgls_workspace/src/workspace/server/expand_star.rs +++ b/crates/pgls_workspace/src/workspace/server/expand_star.rs @@ -2,7 +2,7 @@ use pgls_analyse::{ActionCategory, Applicability, RefactorKind}; use pgls_catalog::resolve::{Reference, StarColumn}; -use pgls_query::protobuf::KeywordKind; +use pgls_query::protobuf::{KeywordKind, ScanToken, Token}; use pgls_text_size::TextSize; use crate::features::code_actions::{CodeAction, CodeActionKind, EditAction, TextEdit}; @@ -40,6 +40,91 @@ pub(super) fn expand_star_action( }) } +/// Whether `position` is on a `*` that may stand for columns, or on its qualifier, as in +/// `select *` or `select u.*`. The line of the cursor is scanned on its own: `*` in comments +/// and strings, `count(*)`, and multiplications don't count. If the line doesn't scan, like +/// a line inside a multi-line string, any `*` on it counts. +pub(super) fn is_at_star(content: &str, position: TextSize) -> bool { + let offset = usize::from(position).min(content.len()); + let line_start = content[..offset].rfind('\n').map_or(0, |i| i + 1); + let line_end = content[offset..] + .find('\n') + .map_or(content.len(), |i| offset + i); + is_at_star_in(&content[line_start..line_end], offset - line_start) +} + +/// Whether `cursor` is on a `*` of a select list in `text`, a line or a string constant on it. +fn is_at_star_in(text: &str, cursor: usize) -> bool { + if !text.contains('*') { + return false; + } + let Ok(scan) = pgls_query::scan(text) else { + return true; + }; + // In the body of a function on one line, like `as 'select * from t'`. + if let Some([_, string]) = scan.tokens.windows(2).find(|pair| { + pair[0].token() == Token::As + && pair[1].token() == Token::Sconst + && usize::try_from(pair[1].start).is_ok_and(|start| start < cursor) + && usize::try_from(pair[1].end).is_ok_and(|end| cursor < end) + }) { + let (Ok(start), Ok(end)) = (usize::try_from(string.start), usize::try_from(string.end)) + else { + return true; + }; + let constant = &text[start..end]; + let quote = if constant.starts_with('\'') { + 1 + } else if let Some(after_dollar) = constant.strip_prefix('$') { + after_dollar.find('$').map_or(1, |tag| tag + 2) + } else { + return false; + }; + let inner = constant + .get(quote..constant.len().saturating_sub(quote)) + .unwrap_or_default(); + return cursor >= start + quote && is_at_star_in(inner, cursor - start - quote); + } + let tokens: Vec<&ScanToken> = scan + .tokens + .iter() + .filter(|token| !matches!(token.token(), Token::SqlComment | Token::CComment)) + .collect(); + tokens.iter().enumerate().any(|(i, star)| { + if star.token() != Token::Ascii42 { + return false; + } + // What comes before a `*` of a select list: nothing on this line, `SELECT`, `,`, `.`, + // `RETURNING`, `DISTINCT`, `ALL`, or the `)` of `DISTINCT ON (...)`. + let mut first = i; + while first >= 2 + && tokens[first - 1].token() == Token::Ascii46 + && is_name(tokens[first - 2]) + { + first -= 2; + } + let star_like = first < i + || first == 0 + || matches!( + tokens[first - 1].token(), + Token::Select + | Token::Ascii44 + | Token::Returning + | Token::Distinct + | Token::All + | Token::Ascii41 + ); + let start = usize::try_from(tokens[first].start).unwrap_or(usize::MAX); + let end = usize::try_from(star.end).unwrap_or(0); + star_like && (start..=end).contains(&cursor) + }) +} + +fn is_name(token: &ScanToken) -> bool { + matches!(token.token(), Token::Ident | Token::Uident) + || token.keyword_kind() != KeywordKind::NoKeyword +} + /// The columns as a select list, e.g. `u.id, u."Name"`. fn column_list(columns: &[StarColumn]) -> String { columns @@ -57,16 +142,19 @@ fn column_list(columns: &[StarColumn]) -> String { .join(", ") } -/// `name`, quoted unless it reads as the same name without quotes: lower case, and not a keyword -/// other than an unreserved one. +/// `name`, quoted unless it reads as the same name without quotes: no upper case ASCII letters, +/// and not a keyword other than an unreserved one. Like Postgres' lexer, any character that isn't +/// ASCII can be part of an unquoted name, and isn't folded to lower case ([`scan.l`]). +/// +/// [`scan.l`]: https://github.com/postgres/postgres/blob/REL_18_6/src/backend/parser/scan.l#L332 fn quote_identifier(name: &str) -> String { let plain = name .chars() .next() - .is_some_and(|first| first.is_ascii_lowercase() || first == '_') - && name - .chars() - .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || matches!(c, '_' | '$')); + .is_some_and(|first| first.is_ascii_lowercase() || first == '_' || !first.is_ascii()) + && name.chars().all(|c| { + c.is_ascii_lowercase() || c.is_ascii_digit() || matches!(c, '_' | '$') || !c.is_ascii() + }); let keyword = || { pgls_query::scan(name).map_or(true, |scan| { scan.tokens.first().is_none_or(|token| { @@ -86,6 +174,8 @@ fn quote_identifier(name: &str) -> String { #[cfg(test)] mod tests { + use pgls_text_size::TextSize; + use super::{StarColumn, column_list, quote_identifier}; #[test] @@ -101,6 +191,42 @@ mod tests { assert_eq!(quote_identifier("select"), "\"select\""); assert_eq!(quote_identifier("between"), "\"between\""); assert_eq!(quote_identifier("left"), "\"left\""); + // Characters that aren't ASCII need no quotes, and aren't folded. + assert_eq!(quote_identifier("größe"), "größe"); + assert_eq!(quote_identifier("Ölstand"), "Ölstand"); + assert_eq!(quote_identifier("Café"), "\"Café\""); + } + + #[test] + fn only_stars_of_select_lists_count() { + let at = |line: &str| { + let position = line.find('|').unwrap(); + let content = line.replace('|', ""); + super::is_at_star(&content, TextSize::try_from(position).unwrap()) + }; + assert!(at("select |* from users")); + assert!(at("select *| from users")); + assert!(at("select id, |* from users")); + assert!(at("select |u.* from users u")); + assert!(at("select public.us|ers.* from users")); + assert!(at("insert into t values (1) returning |*")); + assert!(at(" |*\n")); + assert!(!at("select * |from users")); + assert!(!at("select count(|*) from users")); + assert!(!at("select a |* b from t")); + assert!(!at("select 1 /* |comment */")); + assert!(!at("select '|*' from t")); + assert!(!at("select 1| -- *")); + // In the body of a function on the line. + assert!(at( + "create function f() returns setof t as 'select |* from t' language sql" + )); + assert!(at( + "create function f() returns setof t as $f$ select |* from t $f$" + )); + assert!(!at( + "create function f() returns setof t as 'select |count(*) from t'" + )); } #[test]