fix(pretty_print): keep with admin option portable

crates/pgls_pretty_print/src/comments.rs

@@ -243,6 +243,15 @@ pub(crate) fn first_node_location(ast: &NodeEnum) -> Option<i32> {
         .min()
 }
 
+/// The location of the node of `ast` that starts last in the source. A trailing comment written
+/// after `ast` is attached there.
+pub(crate) fn last_node_location(ast: &NodeEnum) -> Option<i32> {
+    ast.iter()
+        .filter_map(|node| node_location(&node))
+        .filter(|location| *location >= 0)
+        .max()
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

crates/pgls_pretty_print/src/emitter.rs

@@ -59,12 +59,15 @@ pub struct EventEmitter {
     leading_line_comments_require_break: bool,
     /// Line comments written after a list comma, waiting for the emitter to reach that comma.
     comments_after_separator: Vec<Comment>,
+    /// `config.layout` as configured, while `with_layout` overrides it for part of a statement.
+    configured_layout: crate::Layout,
 }
 
 impl EventEmitter {
     pub fn new(config: FormatConfig) -> Self {
         Self {
             events: Vec::new(),
+            configured_layout: config.layout.clone(),
             config,
             leading_comments: HashMap::new(),
             trailing_comments: HashMap::new(),
@@ -79,6 +82,11 @@ impl EventEmitter {
         &self.config
     }
 
+    /// The layout from the configuration, which `with_layout` does not change.
+    pub fn configured_layout(&self) -> crate::Layout {
+        self.configured_layout.clone()
+    }
+
     /// Emits `body` with `layout` in place of the configured one.
     pub fn with_layout(&mut self, layout: crate::Layout, body: impl FnOnce(&mut Self)) {
         let previous = std::mem::replace(&mut self.config.layout, layout);
@@ -93,6 +101,7 @@ impl EventEmitter {
     ) -> Self {
         Self {
             events: Vec::new(),
+            configured_layout: config.layout.clone(),
             config,
             leading_comments,
             trailing_comments,
@@ -203,10 +212,25 @@ impl EventEmitter {
 
     /// Emits and consumes comments that follow `location`, if any.
     pub fn take_trailing_comments_at(&mut self, location: i32) {
-        let Some(comments) = self.trailing_comments.remove(&location) else {
-            return;
-        };
+        let comments = self.remove_trailing_comments_at(location);
+        self.emit_trailing_comments(comments);
+    }
+
+    /// Consumes the comments that follow `location` without emitting them, so that the caller
+    /// can emit them somewhere else with `emit_trailing_comments`.
+    pub fn remove_trailing_comments_at(&mut self, location: i32) -> Vec<Comment> {
+        self.trailing_comments.remove(&location).unwrap_or_default()
+    }
+
+    /// Puts back comments taken with `remove_trailing_comments_at`.
+    pub fn restore_trailing_comments_at(&mut self, location: i32, comments: Vec<Comment>) {
+        if !comments.is_empty() {
+            self.trailing_comments.insert(location, comments);
+        }
+    }
 
+    /// Emits comments that follow the node just emitted.
+    pub fn emit_trailing_comments(&mut self, comments: Vec<Comment>) {
         for comment in comments {
             if comment.after_separator {
                 self.comments_after_separator.push(comment);

crates/pgls_pretty_print/src/nodes/a_array_expr.rs

@@ -2,19 +2,27 @@ use pgls_query::protobuf::AArrayExpr;
 
 use crate::{
     TokenKind,
-    emitter::{EventEmitter, GroupKind},
+    emitter::{EventEmitter, GroupKind, LineType},
 };
 
-use super::node_list::emit_comma_separated_list;
+use super::node_list::emit_fill_comma_separated_list;
 
 pub(super) fn emit_a_array_expr(e: &mut EventEmitter, n: &AArrayExpr) {
     e.group_start(GroupKind::AArrayExpr);
 
     e.token(TokenKind::ARRAY_KW);
     e.token(TokenKind::L_BRACK);
 
+    // Laid out like an `IN (...)` list: on one line when it fits, otherwise the elements fill
+    // indented lines and `]` goes back to the line the array started on.
     if !n.elements.is_empty() {
-        emit_comma_separated_list(e, &n.elements, super::emit_node);
+        e.indent_start();
+        e.line(LineType::Soft);
+        e.group_start(GroupKind::List);
+        emit_fill_comma_separated_list(e, &n.elements, super::emit_node);
+        e.group_end();
+        e.indent_end();
+        e.line(LineType::Soft);
     }
 
     e.token(TokenKind::R_BRACK);

crates/pgls_pretty_print/src/nodes/a_expr.rs

@@ -72,7 +72,13 @@ fn emit_aexpr_op(e: &mut EventEmitter, n: &AExpr) {
                 }
                 e.space();
                 emit_operator(e, &n.name);
-                e.line(LineType::SoftOrSpace);
+                // A subquery breaks inside its parentheses, so `= (` stays on the operator line,
+                // like `in (` and `exists (`.
+                if starts_with_subquery_parenthesis(rexpr) {
+                    e.space();
+                } else {
+                    e.line(LineType::SoftOrSpace);
+                }
                 emit_operand_with_parens(e, rexpr, parent_info, OperandSide::Right);
                 if grouped_suffix {
                     e.group_end();
@@ -598,6 +604,18 @@ fn operator_needs_space(op: &str) -> bool {
     op.chars().any(|c| c.is_alphabetic())
 }
 
+/// Whether `node` is printed as `(subquery)` or `array(subquery)`.
+fn starts_with_subquery_parenthesis(node: &Node) -> bool {
+    use pgls_query::protobuf::SubLinkType;
+
+    matches!(
+        node.node.as_ref(),
+        Some(NodeEnum::SubLink(sub_link))
+            if sub_link.sub_link_type == SubLinkType::ExprSublink as i32
+                || sub_link.sub_link_type == SubLinkType::ArraySublink as i32
+    )
+}
+
 fn operator_prefers_line_break(name: &[Node]) -> bool {
     match extract_simple_operator(name) {
         // Comparison operators

crates/pgls_pretty_print/src/nodes/case_when.rs

@@ -10,9 +10,13 @@ pub(super) fn emit_case_when(e: &mut EventEmitter, n: &CaseWhen) {
 
     e.token(TokenKind::WHEN_KW);
 
+    // A condition that breaks, an AND chain for instance, continues one level deeper than WHEN, so
+    // that its operators don't read as siblings of WHEN.
     if let Some(ref expr) = n.expr {
         e.space();
+        e.indent_start();
         super::emit_node(expr, e);
+        e.indent_end();
     }
 
     let expanded = matches!(e.config().layout, crate::Layout::Expanded);

crates/pgls_pretty_print/src/nodes/common_table_expr.rs

@@ -22,11 +22,8 @@ pub(super) fn emit_common_table_expr(e: &mut EventEmitter, n: &CommonTableExpr)
         emit_parenthesized_list(e, &n.aliascolnames, super::emit_node);
     }
 
-    if matches!(e.config().layout, crate::Layout::Expanded) {
-        e.space();
-    } else {
-        e.line(LineType::SoftOrSpace);
-    }
+    // `name as (` stays on one line; a body that doesn't fit breaks after the parenthesis.
+    e.space();
     e.token(TokenKind::AS_KW);
     e.space();
 

crates/pgls_pretty_print/src/nodes/constraint.rs

@@ -312,9 +312,7 @@ pub(super) fn emit_constraint(e: &mut EventEmitter, n: &Constraint) {
                 e.space();
                 e.token(TokenKind::WHERE_KW);
                 e.space();
-                e.token(TokenKind::L_PAREN);
-                super::emit_clause_condition(e, where_clause);
-                e.token(TokenKind::R_PAREN);
+                super::emit_parenthesized_condition(e, GroupKind::Constraint, where_clause);
             }
 
             emit_deferrable(e, n, &mut continued);

crates/pgls_pretty_print/src/nodes/create_foreign_table_stmt.rs

@@ -48,10 +48,10 @@ pub(super) fn emit_create_foreign_table_stmt(e: &mut EventEmitter, n: &CreateFor
             e.token(TokenKind::L_PAREN);
             if !base.table_elts.is_empty() {
                 e.indent_start();
-                e.line(crate::emitter::LineType::SoftOrSpace);
+                e.line(crate::emitter::LineType::Soft);
                 emit_comma_separated_list(e, &base.table_elts, super::emit_node);
                 e.indent_end();
-                e.line(crate::emitter::LineType::SoftOrSpace);
+                e.line(crate::emitter::LineType::Soft);
             }
             e.token(TokenKind::R_PAREN);
         }

crates/pgls_pretty_print/src/nodes/create_stmt.rs

@@ -171,7 +171,7 @@ fn emit_create_stmt_impl(e: &mut EventEmitter, n: &CreateStmt, with_semicolon: b
             }
 
             e.indent_end();
-            e.line(LineType::SoftOrSpace);
+            e.line(LineType::Soft);
             e.token(TokenKind::R_PAREN);
         }
     } else {
@@ -182,10 +182,11 @@ fn emit_create_stmt_impl(e: &mut EventEmitter, n: &CreateStmt, with_semicolon: b
         e.space();
         e.group_start(GroupKind::List);
         e.token(TokenKind::L_PAREN);
-        // Force break for tables with multiple columns for readability
+        // Force break for tables with multiple columns for readability. A table without columns
+        // keeps `()` together.
         if multiple_items {
             e.line(LineType::Hard);
-        } else {
+        } else if has_content {
             e.line(LineType::Soft);
         }
 
@@ -223,7 +224,7 @@ fn emit_create_stmt_impl(e: &mut EventEmitter, n: &CreateStmt, with_semicolon: b
 
         if multiple_items {
             e.line(LineType::Hard);
-        } else {
+        } else if has_content {
             e.line(LineType::Soft);
         }
         e.token(TokenKind::R_PAREN);

crates/pgls_pretty_print/src/nodes/def_elem.rs

@@ -105,7 +105,7 @@ pub(super) fn emit_options_def_elem(e: &mut EventEmitter, n: &DefElem) {
 }
 
 /// Emit role options for GRANT...WITH syntax
-/// Uses `NAME TRUE` / `NAME FALSE` syntax for INHERIT, SET, ADMIN
+/// Uses `NAME TRUE` / `NAME FALSE` syntax for INHERIT, SET, ADMIN, except `ADMIN OPTION`.
 pub(super) fn emit_grant_role_option(e: &mut EventEmitter, n: &DefElem) {
     e.group_start(GroupKind::DefElem);
 
@@ -128,7 +128,11 @@ pub(super) fn emit_grant_role_option(e: &mut EventEmitter, n: &DefElem) {
     // Emit TRUE or FALSE
     if let Some(val) = is_true {
         e.space();
-        if val {
+        // `admin option` and `admin true` parse the same, but Postgres 15 and older only accept
+        // `with admin option`.
+        if val && defname.eq_ignore_ascii_case("admin") {
+            e.token(TokenKind::OPTION_KW);
+        } else if val {
             e.token(TokenKind::TRUE_KW);
         } else {
             e.token(TokenKind::FALSE_KW);

crates/pgls_pretty_print/src/nodes/func_call.rs

@@ -102,7 +102,10 @@ fn emit_grouped_comma(
 
 pub(super) fn emit_func_call(e: &mut EventEmitter, n: &FuncCall) {
     e.group_start(GroupKind::FuncCall);
-    if n.over.is_some() {
+    // The name and arguments have a group of their own when a clause follows them, so that
+    // `count(*)` stays together and the line breaks before FILTER, WITHIN GROUP or OVER first.
+    let has_suffix = n.over.is_some() || n.agg_filter.is_some() || n.agg_within_group;
+    if has_suffix {
         e.group_start(GroupKind::List);
     }
 
@@ -115,7 +118,7 @@ pub(super) fn emit_func_call(e: &mut EventEmitter, n: &FuncCall) {
         emit_standard_function(e, n);
     }
 
-    if n.over.is_some() {
+    if has_suffix {
         e.group_end();
     }
 
@@ -125,6 +128,8 @@ pub(super) fn emit_func_call(e: &mut EventEmitter, n: &FuncCall) {
             "ordered-set aggregate is missing ORDER BY list"
         );
 
+        // A clause that breaks onto its own line continues the call, one level in.
+        e.indent_start();
         e.line(LineType::SoftOrSpace);
         e.token(TokenKind::WITHIN_KW);
         e.space();
@@ -144,16 +149,30 @@ pub(super) fn emit_func_call(e: &mut EventEmitter, n: &FuncCall) {
         e.line(LineType::Soft);
         e.token(TokenKind::R_PAREN);
         e.group_end();
+        e.indent_end();
     }
 
     if let Some(ref filter) = n.agg_filter {
+        e.indent_start();
         e.line(LineType::SoftOrSpace);
         e.token(TokenKind::FILTER_KW);
         e.space();
+        // `filter (where ...)` stays on one line when it fits. Otherwise the condition goes on
+        // its own lines inside the parentheses, like a policy's `using (...)`.
+        e.group_start(GroupKind::List);
         e.token(TokenKind::L_PAREN);
+        e.indent_start();
+        e.line(LineType::Soft);
         e.token(TokenKind::WHERE_KW);
-        super::emit_clause_condition(e, filter);
+        e.space();
+        e.indent_start();
+        super::emit_node(filter, e);
+        e.indent_end();
+        e.indent_end();
+        e.line(LineType::Soft);
         e.token(TokenKind::R_PAREN);
+        e.group_end();
+        e.indent_end();
     }
 
     // Handle OVER clause (window functions)

crates/pgls_pretty_print/src/nodes/join_expr.rs

@@ -3,7 +3,7 @@ use pgls_query::protobuf::{JoinExpr, JoinType};
 use crate::TokenKind;
 use crate::emitter::{EventEmitter, GroupKind, LineType};
 
-use super::node_list::{emit_comma_separated_list, emit_parenthesized_list};
+use super::node_list::emit_parenthesized_list;
 use super::string::emit_identifier_maybe_quoted;
 
 pub(super) fn emit_join_expr(e: &mut EventEmitter, n: &JoinExpr) {
@@ -102,29 +102,14 @@ pub(super) fn emit_join_expr(e: &mut EventEmitter, n: &JoinExpr) {
         e.line(LineType::SoftOrSpace);
         e.token(TokenKind::USING_KW);
         e.space();
-        e.token(TokenKind::L_PAREN);
-        if n.using_clause.len() > 1 {
-            e.indent_start();
-            e.line(LineType::SoftOrSpace);
-            emit_comma_separated_list(e, &n.using_clause, |node, e| {
-                // For USING clause, String nodes should be identifiers
-                if let Some(pgls_query::NodeEnum::String(s)) = node.node.as_ref() {
-                    emit_identifier_maybe_quoted(e, &s.sval);
-                } else {
-                    super::emit_node(node, e);
-                }
-            });
-            e.indent_end();
-        } else {
-            emit_comma_separated_list(e, &n.using_clause, |node, e| {
-                if let Some(pgls_query::NodeEnum::String(s)) = node.node.as_ref() {
-                    emit_identifier_maybe_quoted(e, &s.sval);
-                } else {
-                    super::emit_node(node, e);
-                }
-            });
-        }
-        e.token(TokenKind::R_PAREN);
+        emit_parenthesized_list(e, &n.using_clause, |node, e| {
+            // For USING clause, String nodes should be identifiers
+            if let Some(pgls_query::NodeEnum::String(s)) = node.node.as_ref() {
+                emit_identifier_maybe_quoted(e, &s.sval);
+            } else {
+                super::emit_node(node, e);
+            }
+        });
 
         // Emit USING alias if present (PostgreSQL 14+)
         // Format: USING (col1, col2) AS alias