fix(pretty_print): close a subquery opened on the WHERE or ON line under the keyword

crates/pgls_pretty_print/src/nodes/join_expr.rs

@@ -123,9 +123,7 @@ pub(super) fn emit_join_expr(e: &mut EventEmitter, n: &JoinExpr) {
         e.line(LineType::SoftOrSpace);
         e.token(TokenKind::ON_KW);
         e.space();
-        e.indent_start();
-        super::emit_node(quals, e);
-        e.indent_end();
+        super::emit_condition_on_keyword_line(e, quals);
     } else if matches!(jointype, JoinType::JoinInner) && !n.is_natural {
         // For INNER JOIN without qualifications (converted from CROSS JOIN), add ON TRUE
         // This is semantically equivalent to CROSS JOIN

crates/pgls_pretty_print/src/nodes/mod.rs

@@ -609,13 +609,69 @@ pub(super) fn emit_clause_condition(e: &mut EventEmitter, clause: &Node) {
     // Compact keeps the condition on the keyword line. The indent stays in both cases so that the
     // condition's own continuation lines, the second AND of a chain for instance, sit under it.
     match e.config().clause_body_style {
-        ClauseBodyStyle::Compact => e.space(),
-        ClauseBodyStyle::Break => e.line(LineType::SoftOrSpace),
+        ClauseBodyStyle::Compact => {
+            e.space();
+            emit_condition_on_keyword_line(e, clause);
+        }
+        ClauseBodyStyle::Break => {
+            e.line(LineType::SoftOrSpace);
+            e.indent_start();
+            emit_node(clause, e);
+            e.indent_end();
+        }
     }
+}
 
-    e.indent_start();
-    emit_node(clause, e);
-    e.indent_end();
+/// Emits a condition that starts on its keyword's line, `where` or `on`, indenting its
+/// continuation lines.
+///
+/// A condition that only breaks inside the parentheses of a subquery, `exists (...)` or
+/// `a = (...)`, has no continuation lines of its own. Indenting it would put the subquery two
+/// levels in and the closing parenthesis one level in, so it keeps the keyword's indent and closes
+/// under the keyword.
+pub(super) fn emit_condition_on_keyword_line(e: &mut EventEmitter, condition: &Node) {
+    if breaks_only_inside_a_subquery(condition) {
+        emit_node(condition, e);
+    } else {
+        e.indent_start();
+        emit_node(condition, e);
+        e.indent_end();
+    }
+}
+
+fn breaks_only_inside_a_subquery(node: &Node) -> bool {
+    use pgls_query::protobuf::{AExprKind, BoolExprType, SubLinkType};
+
+    match node.node.as_ref() {
+        Some(NodeEnum::SubLink(sub_link)) => {
+            sub_link.sub_link_type != SubLinkType::MultiexprSublink as i32
+        }
+        Some(NodeEnum::AExpr(expr)) => {
+            expr.kind() == AExprKind::AexprOp
+                && expr.lexpr.as_deref().is_some_and(is_simple_operand)
+                && expr.rexpr.as_deref().is_some_and(|rexpr| {
+                    matches!(
+                        rexpr.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
+                    )
+                })
+        }
+        Some(NodeEnum::BoolExpr(expr)) => {
+            expr.boolop() == BoolExprType::NotExpr
+                && matches!(expr.args.as_slice(), [arg] if breaks_only_inside_a_subquery(arg))
+        }
+        _ => false,
+    }
+}
+
+/// A column reference, parameter or constant, which never breaks.
+fn is_simple_operand(node: &Node) -> bool {
+    matches!(
+        node.node.as_ref(),
+        Some(NodeEnum::ColumnRef(_) | NodeEnum::ParamRef(_) | NodeEnum::AConst(_))
+    )
 }
 
 /// Emits a condition wrapped in parentheses, such as the `USING (...)` of a policy.

crates/pgls_pretty_print/tests/snapshots/multi/tests__compact_select_list_100.snap

@@ -26,10 +26,10 @@ select id into temporary table active_user_ids from users where active;
 select *
 from users
 where exists (
-    select 1
-    from memberships
-    where memberships.user_id = users.id and memberships.organization_id = 42
-  );
+  select 1
+  from memberships
+  where memberships.user_id = users.id and memberships.organization_id = 42
+);
 
 select case
     when users.status = 'active' then 'currently an active member'

crates/pgls_pretty_print/tests/snapshots/multi/tests__compact_select_list_80.snap

@@ -29,10 +29,10 @@ select id into temporary table active_user_ids from users where active;
 select *
 from users
 where exists (
-    select 1
-    from memberships
-    where memberships.user_id = users.id and memberships.organization_id = 42
-  );
+  select 1
+  from memberships
+  where memberships.user_id = users.id and memberships.organization_id = 42
+);
 
 select case
     when users.status = 'active' then 'currently an active member'

crates/pgls_pretty_print/tests/snapshots/multi/tests__compact_select_list_expanded_100.snap

@@ -39,11 +39,11 @@ where active;
 select *
 from users
 where exists (
-    select 1
-    from memberships
-    where memberships.user_id = users.id and
-      memberships.organization_id = 42
-  );
+  select 1
+  from memberships
+  where memberships.user_id = users.id and
+    memberships.organization_id = 42
+);
 
 select case
          when users.status = 'active'

crates/pgls_pretty_print/tests/snapshots/multi/tests__compact_select_list_expanded_80.snap

@@ -39,11 +39,11 @@ where active;
 select *
 from users
 where exists (
-    select 1
-    from memberships
-    where memberships.user_id = users.id and
-      memberships.organization_id = 42
-  );
+  select 1
+  from memberships
+  where memberships.user_id = users.id and
+    memberships.organization_id = 42
+);
 
 select case
          when users.status = 'active'

crates/pgls_pretty_print/tests/snapshots/multi/tests__compact_select_list_leading_100.snap

@@ -26,10 +26,10 @@ select id into temporary table active_user_ids from users where active;
 select *
 from users
 where exists (
-    select 1
-    from memberships
-    where memberships.user_id = users.id and memberships.organization_id = 42
-  );
+  select 1
+  from memberships
+  where memberships.user_id = users.id and memberships.organization_id = 42
+);
 
 select case
     when users.status = 'active' then 'currently an active member'

crates/pgls_pretty_print/tests/snapshots/multi/tests__compact_select_list_leading_80.snap

@@ -29,10 +29,10 @@ select id into temporary table active_user_ids from users where active;
 select *
 from users
 where exists (
-    select 1
-    from memberships
-    where memberships.user_id = users.id and memberships.organization_id = 42
-  );
+  select 1
+  from memberships
+  where memberships.user_id = users.id and memberships.organization_id = 42
+);
 
 select case
     when users.status = 'active' then 'currently an active member'

crates/pgls_pretty_print/tests/snapshots/multi/tests__predicate_100.snap

@@ -71,16 +71,16 @@ from
   left outer join
     pred_tab as t2
   on exists (
-      select
-        1
-      from
-        pred_tab as t3,
-        pred_tab as t4,
-        pred_tab as t5,
-        pred_tab as t6
-      where
-        t1.a = t3.a and t6.a is not null
-    );
+    select
+      1
+    from
+      pred_tab as t3,
+      pred_tab as t4,
+      pred_tab as t5,
+      pred_tab as t6
+    where
+      t1.a = t3.a and t6.a is not null
+  );
 
 select
   *
@@ -89,16 +89,16 @@ from
   left outer join
     pred_tab as t2
   on exists (
-      select
-        1
-      from
-        pred_tab as t3,
-        pred_tab as t4,
-        pred_tab as t5,
-        pred_tab as t6
-      where
-        t1.a = t3.a and t6.a is null
-    );
+    select
+      1
+    from
+      pred_tab as t3,
+      pred_tab as t4,
+      pred_tab as t5,
+      pred_tab as t6
+    where
+      t1.a = t3.a and t6.a is null
+  );
 
 drop table pred_tab;
 

crates/pgls_pretty_print/tests/snapshots/multi/tests__predicate_80.snap

@@ -77,16 +77,16 @@ from
   left outer join
     pred_tab as t2
   on exists (
-      select
-        1
-      from
-        pred_tab as t3,
-        pred_tab as t4,
-        pred_tab as t5,
-        pred_tab as t6
-      where
-        t1.a = t3.a and t6.a is not null
-    );
+    select
+      1
+    from
+      pred_tab as t3,
+      pred_tab as t4,
+      pred_tab as t5,
+      pred_tab as t6
+    where
+      t1.a = t3.a and t6.a is not null
+  );
 
 select
   *
@@ -95,16 +95,16 @@ from
   left outer join
     pred_tab as t2
   on exists (
-      select
-        1
-      from
-        pred_tab as t3,
-        pred_tab as t4,
-        pred_tab as t5,
-        pred_tab as t6
-      where
-        t1.a = t3.a and t6.a is null
-    );
+    select
+      1
+    from
+      pred_tab as t3,
+      pred_tab as t4,
+      pred_tab as t5,
+      pred_tab as t6
+    where
+      t1.a = t3.a and t6.a is null
+  );
 
 drop table pred_tab;
 

crates/pgls_pretty_print/tests/snapshots/multi/tests__scalar_subquery_expanded_100.snap

@@ -18,34 +18,34 @@ using (
 select *
 from documents as d
 where d.owner_id = (
-    select m.user_id
-    from memberships as m
-    where m.organization_id = d.organization_id
-    limit 1
-  );
+  select m.user_id
+  from memberships as m
+  where m.organization_id = d.organization_id
+  limit 1
+);
 
 select *
 from documents
 where tags = array(
-    select tag
-    from document_tags
-    where document_tags.document_id = documents.id
-    order by tag
-  );
+  select tag
+  from document_tags
+  where document_tags.document_id = documents.id
+  order by tag
+);
 
 select *
 from t
 where a = (
-    select coalesce(
-        (
-          select max(x)
-          from y
-          where y.id = t.id
-            and y.kind = 'some kind of value'
-        ),
-        0
-      )
-  );
+  select coalesce(
+      (
+        select max(x)
+        from y
+        where y.id = t.id
+          and y.kind = 'some kind of value'
+      ),
+      0
+    )
+);
 
 select *
 from t

crates/pgls_pretty_print/tests/snapshots/multi/tests__scalar_subquery_expanded_80.snap

@@ -18,34 +18,34 @@ using (
 select *
 from documents as d
 where d.owner_id = (
-    select m.user_id
-    from memberships as m
-    where m.organization_id = d.organization_id
-    limit 1
-  );
+  select m.user_id
+  from memberships as m
+  where m.organization_id = d.organization_id
+  limit 1
+);
 
 select *
 from documents
 where tags = array(
-    select tag
-    from document_tags
-    where document_tags.document_id = documents.id
-    order by tag
-  );
+  select tag
+  from document_tags
+  where document_tags.document_id = documents.id
+  order by tag
+);
 
 select *
 from t
 where a = (
-    select coalesce(
-        (
-          select max(x)
-          from y
-          where y.id = t.id
-            and y.kind = 'some kind of value'
-        ),
-        0
-      )
-  );
+  select coalesce(
+      (
+        select max(x)
+        from y
+        where y.id = t.id
+          and y.kind = 'some kind of value'
+      ),
+      0
+    )
+);
 
 select *
 from t