fix(editor): avoid quadratic SQL folding scans

apps/desktop/src/lib/editor/__tests__/codemirrorSqlBlockFolding.spec.ts

@@ -3,9 +3,9 @@ import { ensureSyntaxTree, foldable } from "@codemirror/language";
 import { Compartment, EditorState } from "@codemirror/state";
 import { describe, expect, it } from "vitest";
 import { createDbxCodeMirrorSqlDialect } from "@/lib/editor/codemirrorSqlDialect";
-import { sqlBlockFoldService } from "@/lib/editor/codemirrorSqlBlockFolding";
+import { collectUnionBranchFoldRanges, sqlBlockFoldService } from "@/lib/editor/codemirrorSqlBlockFolding";
 
-function stateFor(doc: string, dialectName: "mysql" | "sqlserver" = "mysql"): EditorState {
+function stateFor(doc: string, dialectName: "mysql" | "postgres" | "sqlserver" = "mysql"): EditorState {
   const state = EditorState.create({
     doc,
     extensions: [langSql.sql({ dialect: createDbxCodeMirrorSqlDialect(langSql, dialectName) }), sqlBlockFoldService],
@@ -121,6 +121,99 @@ END`;
     expect(foldedTextAtLine(state, 2)).toBeNull();
   });
 
+  it("folds each CTE and nested FROM subquery in the reported PostgreSQL structure", () => {
+    const sql = `WITH t1 AS (
+    SELECT fee_code
+    FROM fact_activity
+)
+, t2 AS (
+    SELECT fee_code
+    FROM fact_verification
+)
+, t3 AS (
+    SELECT fee_code
+    FROM (
+        SELECT fee_code
+        FROM fact_extract
+    ) extracted
+)
+SELECT *
+FROM t1
+LEFT JOIN t2 ON t1.fee_code = t2.fee_code
+LEFT JOIN t3 ON t1.fee_code = t3.fee_code;`;
+    const state = stateFor(sql, "postgres");
+
+    expect(foldedTextAtLine(state, 1)).toBe("\n    SELECT fee_code\n    FROM fact_activity\n");
+    expect(foldedTextAtLine(state, 5)).toBe("\n    SELECT fee_code\n    FROM fact_verification\n");
+    expect(foldedTextAtLine(state, 9)).toContain("FROM (\n");
+    expect(foldedTextAtLine(state, 11)).toBe("\n        SELECT fee_code\n        FROM fact_extract\n    ");
+  });
+
+  it("folds a LEFT JOIN subquery from its opening line", () => {
+    const sql = `SELECT account.id
+FROM account
+LEFT JOIN (
+  SELECT account_id, SUM(amount) AS total
+  FROM fee
+  GROUP BY account_id
+) summary ON summary.account_id = account.id;`;
+    const state = stateFor(sql, "postgres");
+
+    expect(foldedTextAtLine(state, 3)).toBe("\n  SELECT account_id, SUM(amount) AS total\n  FROM fee\n  GROUP BY account_id\n");
+  });
+
+  it("folds SELECT branches around UNION ALL at the same query level", () => {
+    const sql = `WITH combined AS (
+  SELECT id, amount
+  FROM current_fees
+  UNION ALL
+  SELECT id, amount
+  FROM archived_fees
+)
+SELECT *
+FROM combined
+UNION ALL
+SELECT *
+FROM manual_fees;`;
+    const state = stateFor(sql, "postgres");
+
+    expect(foldedTextAtLine(state, 2)).toBe("\n  FROM current_fees\n  ");
+    expect(foldedTextAtLine(state, 5)).toBe("\n  FROM archived_fees\n");
+    expect(foldedTextAtLine(state, 8)).toBe("\nFROM combined\n");
+    expect(foldedTextAtLine(state, 11)).toBe("\nFROM manual_fees");
+  });
+
+  it("precomputes UNION branch boundaries with one token-position read per token", () => {
+    const branchCount = 1_000;
+    let positionReads = 0;
+    const tokens: Array<{ readonly from: number; readonly keyword: "SELECT" | "UNION" }> = Array.from({ length: branchCount * 2 - 1 }, (_, index) => ({
+      keyword: index % 2 === 0 ? "SELECT" : "UNION",
+      get from() {
+        positionReads++;
+        return index * 10;
+      },
+    }));
+    const scopeEnd = tokens.length * 10;
+
+    const ranges = collectUnionBranchFoldRanges(tokens, scopeEnd);
+
+    expect(ranges).toHaveLength(branchCount);
+    expect(ranges[0]).toEqual({ from: 0, to: 10 });
+    expect(ranges.at(-1)).toEqual({ from: (tokens.length - 1) * 10, to: scopeEnd });
+    expect(positionReads).toBe(tokens.length);
+  });
+
+  it("does not create query folds for keywords in strings, comments, or same-line subqueries", () => {
+    const sql = `SELECT '(SELECT fake)' AS text_value;
+-- LEFT JOIN (SELECT fake)
+SELECT * FROM (SELECT 1 AS value) one_line;`;
+    const state = stateFor(sql, "postgres");
+
+    expect(foldedTextAtLine(state, 1)).toBeNull();
+    expect(foldedTextAtLine(state, 2)).toBeNull();
+    expect(foldedTextAtLine(state, 3)).toBeNull();
+  });
+
   it("regression: a BEGIN...END block entirely on one line must not produce an inverted fold range", () => {
     // `from`/`to` were both derived from the shared line, and `from` (end of line) fell after
     // `to` (start of END) -- CodeMirror's foldable() must not surface that as a range at all.

apps/desktop/src/lib/editor/codemirrorSqlBlockFolding.ts

@@ -12,9 +12,8 @@ interface FoldRange {
 }
 
 // `@codemirror/lang-sql`'s grammar is deliberately shallow (see its `foldNodeProp`, which only
-// covers the flat `Statement` and `BlockComment` nodes): `BEGIN`/`CASE`/`END` are just `Keyword`
-// leaf tokens with no enclosing block node, so there is nothing for tree-based folding to attach
-// to. This foldService reconstructs block structure from the token stream instead.
+// covers the flat `Statement` and `BlockComment` nodes). This service adds structure-aware folds
+// for procedural blocks and query expressions without replacing CodeMirror's native folds.
 //
 // Scoped to `BEGIN...END` and `CASE...END` (issue #6574) -- `BEGIN TRY`/`END TRY` and
 // `BEGIN CATCH`/`END CATCH` fall out of this for free, since `BEGIN` is matched regardless of a
@@ -65,6 +64,96 @@ function isSqlServerTransactionBegin(state: EditorState, tokens: SyntaxNode[], i
 // keep serving ranges computed from a stale or partial parse.
 const rangeCache = new WeakMap<Tree, Map<number, FoldRange>>();
 
+function addMultilineFoldRange(state: EditorState, ranges: Map<number, FoldRange>, openerPosition: number, endPosition: number, overwrite = false) {
+  const openerLine = state.doc.lineAt(openerPosition);
+  const endLine = state.doc.lineAt(endPosition);
+  if (openerLine.number === endLine.number || endPosition <= openerLine.to || (!overwrite && ranges.has(openerLine.number))) return;
+  ranges.set(openerLine.number, { from: openerLine.to, to: endPosition });
+}
+
+function queryParent(node: SyntaxNode): SyntaxNode | null {
+  for (let parent = node.parent; parent; parent = parent.parent) {
+    if (parent.name === "Parens" || parent.name === "Statement") return parent;
+  }
+  return null;
+}
+
+function queryScopeKey(node: SyntaxNode): string {
+  return `${node.name}:${node.from}:${node.to}`;
+}
+
+function queryScopeEnd(state: EditorState, scope: SyntaxNode): number {
+  const lastChild = scope.lastChild;
+  if (scope.name === "Parens" && lastChild?.name === ")") return lastChild.from;
+
+  let end = scope.to;
+  while (end > scope.from && /\s/.test(state.sliceDoc(end - 1, end))) end--;
+  if (end > scope.from && state.sliceDoc(end - 1, end) === ";") end--;
+  return end;
+}
+
+function isQueryParens(state: EditorState, node: SyntaxNode): boolean {
+  for (let child = node.firstChild; child; child = child.nextSibling) {
+    if (child.name === "(" || child.name === "LineComment" || child.name === "BlockComment") continue;
+    if (child.name !== "Keyword") return false;
+    const keyword = state.sliceDoc(child.from, child.to).toUpperCase();
+    return keyword === "SELECT" || keyword === "WITH";
+  }
+  return false;
+}
+
+interface QueryScopeTokens {
+  scope: SyntaxNode;
+  tokens: Array<{ from: number; keyword: "SELECT" | "UNION" }>;
+  hasUnion: boolean;
+}
+
+export function collectUnionBranchFoldRanges(tokens: ReadonlyArray<{ from: number; keyword: "SELECT" | "UNION" }>, scopeEnd: number): FoldRange[] {
+  const ranges: FoldRange[] = [];
+  let branchEnd = scopeEnd;
+
+  for (let index = tokens.length - 1; index >= 0; index--) {
+    const token = tokens[index];
+    if (token.keyword === "UNION") {
+      branchEnd = token.from;
+    } else {
+      ranges.push({ from: token.from, to: branchEnd });
+    }
+  }
+
+  ranges.reverse();
+  return ranges;
+}
+
+function addQueryStructureFoldRanges(state: EditorState, tree: Tree, ranges: Map<number, FoldRange>) {
+  const scopes = new Map<string, QueryScopeTokens>();
+  tree.iterate({
+    enter(node) {
+      if (node.name === "Parens" && isQueryParens(state, node.node)) {
+        addMultilineFoldRange(state, ranges, node.from, queryScopeEnd(state, node.node));
+      }
+      if (node.name !== "Keyword") return;
+      const keyword = state.sliceDoc(node.from, node.to).toUpperCase();
+      if (keyword !== "SELECT" && keyword !== "UNION") return;
+      const scope = queryParent(node.node);
+      if (!scope) return;
+      const key = queryScopeKey(scope);
+      const group = scopes.get(key) ?? { scope, tokens: [], hasUnion: false };
+      group.tokens.push({ from: node.from, keyword });
+      if (keyword === "UNION") group.hasUnion = true;
+      scopes.set(key, group);
+    },
+  });
+
+  for (const { scope, tokens, hasUnion } of scopes.values()) {
+    if (!hasUnion) continue;
+    const branchRanges = collectUnionBranchFoldRanges(tokens, queryScopeEnd(state, scope));
+    for (const range of branchRanges) {
+      addMultilineFoldRange(state, ranges, range.from, range.to);
+    }
+  }
+}
+
 function computeBlockFoldRanges(state: EditorState): Map<number, FoldRange> {
   const tree = syntaxTree(state);
   const cached = rangeCache.get(tree);
@@ -98,7 +187,7 @@ function computeBlockFoldRanges(state: EditorState): Map<number, FoldRange> {
         // from/to formula below would otherwise invert (from > to) since `from` is the end of
         // that shared line while `to` is a position earlier on it.
         if (openerLine.number !== endLine.number) {
-          byOpenerLine.set(openerLine.number, { from: openerLine.to, to: token.from });
+          addMultilineFoldRange(state, byOpenerLine, opener.from, token.from, true);
         }
       }
     } else if (BLOCK_OPENERS.has(text) && (text !== "BEGIN" || !isSqlServerTransactionBegin(state, tokens, i))) {
@@ -106,11 +195,13 @@ function computeBlockFoldRanges(state: EditorState): Map<number, FoldRange> {
     }
   }
 
+  addQueryStructureFoldRanges(state, tree, byOpenerLine);
+
   rangeCache.set(tree, byOpenerLine);
   return byOpenerLine;
 }
 
-/** Folds `BEGIN...END` and `CASE...END` blocks in the SQL editor -- see module doc comment. */
+/** Adds procedural-block, query-parenthesis, and UNION-branch folds to the SQL editor. */
 export const sqlBlockFoldService = foldService.of((state, lineStart) => {
   const range = computeBlockFoldRanges(state).get(state.doc.lineAt(lineStart).number);
   return range ?? null;