diff --git a/crates/pgls_workspace/src/workspace/server.rs b/crates/pgls_workspace/src/workspace/server.rs index 123702450..ec20cf15c 100644 --- a/crates/pgls_workspace/src/workspace/server.rs +++ b/crates/pgls_workspace/src/workspace/server.rs @@ -949,7 +949,7 @@ impl Workspace for WorkspaceServer { // Only format SQL function bodies if skip_fn_bodies is false if !settings.formatter.skip_fn_bodies { - for (id, _stmt_range, _text, ast_result) in &format_candidates { + for (id, _stmt_range, text, ast_result) in &format_candidates { if !id.is_child() { continue; } @@ -962,6 +962,13 @@ impl Workspace for WorkspaceServer { continue; }; + // Comments are not represented in the AST, so reformatting a function + // body that contains comments would silently drop them. Leave the + // original body untouched in that case. + if statement_contains_comment(text) { + continue; + } + let Ok(result) = pgls_pretty_print::format_statement(ast, &config) else { continue; }; @@ -970,17 +977,37 @@ impl Workspace for WorkspaceServer { } } + // Reconstruct the file by replacing each statement's source span with its + // formatted output while copying the gaps between statements (whitespace and + // standalone comments) verbatim. libpg_query drops comments from the AST, so + // anything that lives outside a statement's range would otherwise be lost. + let source = doc.get_document_content(); + let mut last_end = 0usize; + for (id, stmt_range, text, ast_result) in format_candidates { if id.is_child() { continue; } + let stmt_start = usize::from(stmt_range.start()); + let stmt_end = usize::from(stmt_range.end()); + + // Preserve leading text, blank lines, and standalone comments that precede + // this statement. + formatted_output.push_str(&source[last_end..stmt_start]); + last_end = stmt_end; + + // Leave statements outside the requested range untouched. if let Some(filter_range) = params.range && stmt_range.intersect(filter_range).is_none() { - if !formatted_output.is_empty() { - formatted_output.push_str("\n\n"); - } + formatted_output.push_str(&text); + continue; + } + + // A comment inside the statement cannot survive a round-trip through the + // AST, so keep the original text rather than dropping the comment. + if statement_contains_comment(&text) { formatted_output.push_str(&text); continue; } @@ -1001,9 +1028,6 @@ impl Workspace for WorkspaceServer { range: stmt_range, }); } - if !formatted_output.is_empty() { - formatted_output.push_str("\n\n"); - } formatted_output.push_str(&result.formatted); } Err(err) => { @@ -1015,9 +1039,6 @@ impl Workspace for WorkspaceServer { .with_file_span(stmt_range), )); - if !formatted_output.is_empty() { - formatted_output.push_str("\n\n"); - } formatted_output.push_str(&text); } } @@ -1027,14 +1048,14 @@ impl Workspace for WorkspaceServer { pgls_diagnostics::Error::from(syntax_err).with_file_path(&path_str), )); - if !formatted_output.is_empty() { - formatted_output.push_str("\n\n"); - } formatted_output.push_str(&text); } } } + // Preserve any trailing text or comments after the last statement. + formatted_output.push_str(&source[last_end..]); + Ok(PullFormattingResult { original: doc.get_document_content().to_string(), formatted: formatted_output, @@ -1169,6 +1190,30 @@ fn is_dir(path: &Path) -> bool { path.is_dir() || (path.is_symlink() && fs::read_link(path).is_ok_and(|path| path.is_dir())) } +/// Returns `true` if the SQL `statement` contains a line (`--`) or block (`/* */`) +/// comment. +/// +/// libpg_query strips comments while building the AST, so the formatter (which +/// renders from the AST) cannot reproduce them. We use the scanner, which exposes +/// comments as dedicated tokens, to detect them and fall back to the original text. +/// Comments inside string literals (including dollar-quoted bodies) are part of the +/// string token and are correctly not reported here. +fn statement_contains_comment(statement: &str) -> bool { + use pgls_query::protobuf::Token; + + match pgls_query::scan(statement) { + Ok(scan) => scan.tokens.iter().any(|token| { + matches!( + Token::try_from(token.token), + Ok(Token::SqlComment | Token::CComment) + ) + }), + // If scanning fails we cannot reason about the statement; let the regular + // formatting path (which will likely fail to parse too) handle it. + Err(_) => false, + } +} + #[cfg(all(test, feature = "db"))] #[path = "server.tests.rs"] mod tests; diff --git a/crates/pgls_workspace/src/workspace/server.tests.rs b/crates/pgls_workspace/src/workspace/server.tests.rs index 792db080e..e4a05a1a6 100644 --- a/crates/pgls_workspace/src/workspace/server.tests.rs +++ b/crates/pgls_workspace/src/workspace/server.tests.rs @@ -778,6 +778,78 @@ async fn test_format_keeps_sql_function_body_intact() { ); } +fn format_content(content: &str) -> String { + let mut conf = PartialConfiguration::init(); + conf.merge_with(PartialConfiguration { + format: Some(PartialFormatConfiguration { + enabled: Some(true), + ..Default::default() + }), + ..Default::default() + }); + + let workspace = get_test_workspace(Some(conf)).expect("Unable to create test workspace"); + + let path = PgLSPath::new("test.sql"); + workspace + .open_file(OpenFileParams { + path: path.clone(), + content: content.into(), + version: 1, + }) + .expect("Unable to open test file"); + + workspace + .pull_file_formatting(PullFileFormattingParams { path, range: None }) + .expect("Unable to pull formatting") + .formatted +} + +#[tokio::test] +async fn test_format_preserves_leading_comment() { + let formatted = format_content("-- a header comment\nselect 1;"); + assert_eq!(formatted, "-- a header comment\nselect 1;"); +} + +#[tokio::test] +async fn test_format_preserves_trailing_comment() { + let formatted = format_content("select amount from customers; -- selects amount"); + assert_eq!(formatted, "select amount from customers; -- selects amount"); +} + +#[tokio::test] +async fn test_format_preserves_between_statement_comment() { + let formatted = format_content("select 1;\n\n-- in between\n\nselect 2;"); + assert_eq!(formatted, "select 1;\n\n-- in between\n\nselect 2;"); +} + +#[tokio::test] +async fn test_format_preserves_interior_comment() { + // A comment wedged between tokens of a statement cannot survive the AST + // round-trip, so the statement is left untouched instead of dropping it. + let content = "select amount -- the amount\nfrom customers;"; + let formatted = format_content(content); + assert_eq!(formatted, content); +} + +#[tokio::test] +async fn test_format_preserves_block_comment() { + let content = "select 1 /* keep me */;"; + let formatted = format_content(content); + assert_eq!(formatted, content); +} + +#[tokio::test] +async fn test_format_preserves_comment_in_sql_function_body() { + let content = + "create function f() returns int as $$\n select 1; -- inner comment\n$$ language sql;"; + let formatted = format_content(content); + assert!( + formatted.contains("-- inner comment"), + "comment inside the function body should be preserved:\n{formatted}", + ); +} + #[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] async fn test_cstyle_comments(test_db: PgPool) { let mut conf = PartialConfiguration::init(); diff --git a/test-db/seed.sql b/test-db/seed.sql index 97f641b91..5b9eb0f0d 100644 --- a/test-db/seed.sql +++ b/test-db/seed.sql @@ -1,3 +1,5 @@ +-- Help line + create table public.contact ( id serial primary key not null, created_at timestamp with time zone not null default now(),