Preserve comments

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;

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();

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(),