fix(formatter): properly handle child statements (#678)

Author: psteinroe

crates/pgls_workspace/src/workspace/server.rs

@@ -889,7 +889,34 @@ impl Workspace for WorkspaceServer {
         let mut formatted_output = String::new();
         let path_str = params.path.as_path().display().to_string();
 
-        for (_id, stmt_range, text, ast_result) in doc.iter(FormatStatementMapper) {
+        let format_candidates: Vec<_> = doc.iter(FormatStatementMapper).collect();
+        let mut formatted_sql_fn_bodies = HashMap::new();
+
+        for (id, _stmt_range, _text, ast_result) in &format_candidates {
+            if !id.is_child() {
+                continue;
+            }
+
+            let Some(parent_id) = id.parent() else {
+                continue;
+            };
+
+            let Ok(ast) = ast_result else {
+                continue;
+            };
+
+            let Ok(result) = pgls_pretty_print::format_statement(ast, &config) else {
+                continue;
+            };
+
+            formatted_sql_fn_bodies.insert(parent_id, result.formatted);
+        }
+
+        for (id, stmt_range, text, ast_result) in format_candidates {
+            if id.is_child() {
+                continue;
+            }
+
             if let Some(filter_range) = params.range
                 && stmt_range.intersect(filter_range).is_none()
             {
@@ -901,35 +928,42 @@ impl Workspace for WorkspaceServer {
             }
 
             match ast_result {
-                Ok(ast) => match pgls_pretty_print::format_statement(&ast, &config) {
-                    Ok(result) => {
-                        if text != result.formatted {
-                            statements.push(StatementFormatResult {
-                                original: text.clone(),
-                                formatted: result.formatted.clone(),
-                                range: stmt_range,
-                            });
-                        }
-                        if !formatted_output.is_empty() {
-                            formatted_output.push_str("\n\n");
-                        }
-                        formatted_output.push_str(&result.formatted);
+                Ok(ast) => {
+                    let mut ast = ast;
+                    if let Some(formatted_sql_fn_body) = formatted_sql_fn_bodies.get(&id) {
+                        sql_function::set_sql_fn_body(&mut ast, formatted_sql_fn_body);
                     }
-                    Err(err) => {
-                        diagnostics.push(SDiagnostic::new(
-                            pgls_diagnostics::Error::from(WorkspaceError::format_error(
-                                err.to_string(),
-                            ))
-                            .with_file_path(&path_str)
-                            .with_file_span(stmt_range),
-                        ));
-
-                        if !formatted_output.is_empty() {
-                            formatted_output.push_str("\n\n");
+
+                    match pgls_pretty_print::format_statement(&ast, &config) {
+                        Ok(result) => {
+                            if text != result.formatted {
+                                statements.push(StatementFormatResult {
+                                    original: text.clone(),
+                                    formatted: result.formatted.clone(),
+                                    range: stmt_range,
+                                });
+                            }
+                            if !formatted_output.is_empty() {
+                                formatted_output.push_str("\n\n");
+                            }
+                            formatted_output.push_str(&result.formatted);
+                        }
+                        Err(err) => {
+                            diagnostics.push(SDiagnostic::new(
+                                pgls_diagnostics::Error::from(WorkspaceError::format_error(
+                                    err.to_string(),
+                                ))
+                                .with_file_path(&path_str)
+                                .with_file_span(stmt_range),
+                            ));
+
+                            if !formatted_output.is_empty() {
+                                formatted_output.push_str("\n\n");
+                            }
+                            formatted_output.push_str(&text);
                         }
-                        formatted_output.push_str(&text);
                     }
-                },
+                }
                 Err(syntax_err) => {
                     diagnostics.push(SDiagnostic::new(
                         pgls_diagnostics::Error::from(syntax_err).with_file_path(&path_str),

crates/pgls_workspace/src/workspace/server.tests.rs

@@ -3,8 +3,8 @@ use std::sync::Arc;
 use biome_deserialize::{Merge, StringSet};
 use pgls_analyse::RuleCategories;
 use pgls_configuration::{
-    PartialConfiguration, PartialTypecheckConfiguration, database::PartialDatabaseConfiguration,
-    files::PartialFilesConfiguration,
+    PartialConfiguration, PartialFormatConfiguration, PartialTypecheckConfiguration,
+    database::PartialDatabaseConfiguration, files::PartialFilesConfiguration,
 };
 
 #[cfg(not(target_os = "windows"))]
@@ -17,6 +17,7 @@ use sqlx::{Executor, PgPool};
 use crate::{
     Workspace, WorkspaceError,
     features::code_actions::ExecuteStatementResult,
+    features::format::PullFileFormattingParams,
     workspace::{
         OpenFileParams, RegisterProjectFolderParams, StatementId, UpdateSettingsParams,
         server::WorkspaceServer,
@@ -629,6 +630,52 @@ FOR NO KEY UPDATE;
     );
 }
 
+#[tokio::test]
+async fn test_format_keeps_sql_function_body_intact() {
+    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");
+    let content =
+        "create function add(a int, b int) returns int as 'SELECT 424242+$1+$2;' language sql;";
+
+    workspace
+        .open_file(OpenFileParams {
+            path: path.clone(),
+            content: content.into(),
+            version: 1,
+        })
+        .expect("Unable to open test file");
+
+    let result = workspace
+        .pull_file_formatting(PullFileFormattingParams {
+            path: path.clone(),
+            range: None,
+        })
+        .expect("Unable to pull formatting");
+
+    assert_eq!(
+        result.formatted.matches("424242").count(),
+        1,
+        "SQL function body should not be emitted as a second standalone statement:\n{}",
+        result.formatted,
+    );
+
+    assert!(
+        result.formatted.contains("select 424242 + $1 + $2;"),
+        "SQL function body should be formatted inline:\n{}",
+        result.formatted,
+    );
+}
+
 #[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")]
 async fn test_cstyle_comments(test_db: PgPool) {
     let mut conf = PartialConfiguration::init();

crates/pgls_workspace/src/workspace/server/sql_function.rs

@@ -106,6 +106,52 @@ pub fn get_sql_fn_body(ast: &pgls_query::NodeEnum, content: &str) -> Option<SQLF
     })
 }
 
+/// Replaces the SQL function body in CREATE FUNCTION ... AS '<body>' for LANGUAGE SQL functions.
+/// Returns true when a body was replaced.
+pub fn set_sql_fn_body(ast: &mut pgls_query::NodeEnum, new_body: &str) -> bool {
+    let create_fn = match ast {
+        pgls_query::NodeEnum::CreateFunctionStmt(cf) => cf,
+        _ => return false,
+    };
+
+    let language = pgls_query_ext::utils::find_option_value(create_fn, "language");
+    if language.as_deref() != Some("sql") {
+        return false;
+    }
+
+    for option in &mut create_fn.options {
+        let Some(pgls_query::NodeEnum::DefElem(def_elem)) = option.node.as_mut() else {
+            continue;
+        };
+
+        if def_elem.defname != "as" {
+            continue;
+        }
+
+        let Some(arg) = def_elem.arg.as_mut() else {
+            continue;
+        };
+
+        match arg.node.as_mut() {
+            Some(pgls_query::NodeEnum::String(s)) => {
+                s.sval = new_body.to_string();
+                return true;
+            }
+            Some(pgls_query::NodeEnum::List(list)) => {
+                if list.items.len() == 1
+                    && let Some(pgls_query::NodeEnum::String(s)) = list.items[0].node.as_mut()
+                {
+                    s.sval = new_body.to_string();
+                    return true;
+                }
+            }
+            _ => {}
+        }
+    }
+
+    false
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -163,4 +209,25 @@ mod tests {
                 .is_some()
         );
     }
+
+    #[test]
+    fn set_sql_function_body() {
+        let input = "CREATE FUNCTION add(a integer, b integer) RETURNS integer
+    AS 'SELECT $1+$2;'
+    LANGUAGE SQL
+    IMMUTABLE;";
+
+        let mut ast = pgls_query::parse(input).unwrap().into_root().unwrap();
+
+        let changed = set_sql_fn_body(&mut ast, "select $1 + $2;");
+        assert!(changed);
+
+        let create_fn = match &ast {
+            pgls_query::NodeEnum::CreateFunctionStmt(stmt) => stmt,
+            _ => panic!("Expected CreateFunctionStmt"),
+        };
+
+        let body = pgls_query_ext::utils::find_option_value(create_fn, "as");
+        assert_eq!(body, Some("select $1 + $2;".to_string()));
+    }
 }