fix(editor): ignore SQL Server transaction BEGIN in folding

apps/desktop/src/components/editor/QueryEditor.vue

@@ -115,6 +115,7 @@ import { trimmedSelectionLayer } from "@/lib/editor/codemirrorTrimmedSelectionLa
 import { currentStatementFrameLayer } from "@/lib/editor/codemirrorCurrentStatementFrameLayer";
 import { selectionMatchOccurrences } from "@/lib/editor/codemirrorSelectionMatches";
 import { createInsertValueHintsExtension, requestInsertValueHintsRefresh } from "@/lib/editor/codemirrorInsertValueHints";
+import { sqlBlockFoldService } from "@/lib/editor/codemirrorSqlBlockFolding";
 import { focusEditorView } from "@/lib/editor/queryEditorFocus";
 import { createDbxCodeMirrorSqlDialect, type CodeMirrorSqlDialectName } from "@/lib/editor/codemirrorSqlDialect";
 import { sqlSemanticTableNameSpansForSyntaxTree } from "@/lib/editor/codemirrorSqlSemanticHighlight";
@@ -4915,6 +4916,7 @@ onMounted(async () => {
           return span;
         },
       }),
+      sqlBlockFoldService,
       drawSelection(),
       trimmedSelectionLayer(),
       selectionMatchOccurrences(),

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

@@ -0,0 +1,169 @@
+import * as langSql from "@codemirror/lang-sql";
+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";
+
+function stateFor(doc: string, dialectName: "mysql" | "sqlserver" = "mysql"): EditorState {
+  const state = EditorState.create({
+    doc,
+    extensions: [langSql.sql({ dialect: createDbxCodeMirrorSqlDialect(langSql, dialectName) }), sqlBlockFoldService],
+  });
+  ensureSyntaxTree(state, doc.length, 5_000);
+  return state;
+}
+
+function foldedTextAtLine(state: EditorState, lineNumber: number): string | null {
+  const line = state.doc.line(lineNumber);
+  const range = foldable(state, line.from, line.to);
+  return range ? state.sliceDoc(range.from, range.to) : null;
+}
+
+describe("sqlBlockFoldService", () => {
+  it("folds a BEGIN...END block, from the end of the BEGIN line to the start of END", () => {
+    const sql = "CREATE PROCEDURE p()\nBEGIN\n  SELECT 1;\nEND";
+    const state = stateFor(sql);
+
+    expect(foldedTextAtLine(state, 2)).toBe("\n  SELECT 1;\n");
+  });
+
+  it("does not fold IF...END IF (issue #6574's original complaint) -- no node exists for it, and folds nothing else either", () => {
+    const sql = "BEGIN\n  IF 1 = 1 THEN\n    SELECT 1;\n  END IF;\nEND";
+    const state = stateFor(sql);
+
+    // Line 2 ("IF 1 = 1 THEN") is not opened by this service, so no BEGIN/CASE-block fold starts there.
+    expect(foldedTextAtLine(state, 2)).toBeNull();
+  });
+
+  it("regression: END IF inside a BEGIN block must not be mistaken for BEGIN's own closer", () => {
+    // Before the fix, the first `END` (of "END IF") popped the BEGIN entry off the stack early,
+    // so the BEGIN fold ended at "END IF" instead of the real outer END, and the final bare
+    // `END` was left unmatched.
+    const sql = "BEGIN\n  IF 1 = 1 THEN\n    SELECT 1;\n  END IF;\nEND";
+    const state = stateFor(sql);
+
+    const folded = foldedTextAtLine(state, 1);
+    expect(folded).not.toBeNull();
+    expect(folded).toContain("END IF;"); // the real outer END must swallow the inner "END IF;" too
+    expect(state.sliceDoc(0, foldable(state, state.doc.line(1).from, state.doc.line(1).to)!.to)).toBe("BEGIN\n  IF 1 = 1 THEN\n    SELECT 1;\n  END IF;\n");
+  });
+
+  it("folds nested BEGIN...END blocks independently", () => {
+    const sql = "BEGIN\n  BEGIN\n    SELECT 1;\n  END;\n  SELECT 2;\nEND";
+    const state = stateFor(sql);
+
+    expect(foldedTextAtLine(state, 2)).toBe("\n    SELECT 1;\n  ");
+    expect(foldedTextAtLine(state, 1)).toBe("\n  BEGIN\n    SELECT 1;\n  END;\n  SELECT 2;\n");
+  });
+
+  it("folds CASE...END (expression form, bare END)", () => {
+    const sql = "SELECT CASE\n  WHEN a THEN 1\n  ELSE 2\nEND FROM t";
+    const state = stateFor(sql);
+
+    expect(foldedTextAtLine(state, 1)).toBe("\n  WHEN a THEN 1\n  ELSE 2\n");
+  });
+
+  it("folds CASE...END CASE (procedural statement form) and does not treat the trailing CASE as a new opener", () => {
+    const sql = "BEGIN\n  CASE v\n    WHEN 1 THEN SELECT 1;\n  END CASE;\nEND";
+    const state = stateFor(sql);
+
+    expect(foldedTextAtLine(state, 2)).toBe("\n    WHEN 1 THEN SELECT 1;\n  ");
+  });
+
+  it("regression: END TRY / END CATCH must not be treated as closing an enclosing BEGIN", () => {
+    const sql = "BEGIN\n  BEGIN TRY\n    SELECT 1;\n  END TRY\n  BEGIN CATCH\n    SELECT 2;\n  END CATCH;\nEND";
+    const state = stateFor(sql, "sqlserver");
+
+    const folded = foldedTextAtLine(state, 1);
+    expect(folded).not.toBeNull();
+    expect(folded).toContain("END CATCH;"); // outer BEGIN must extend all the way to the real outer END
+  });
+
+  it.each(["BEGIN TRAN", "begin transaction", "BeGiN DiStRiBuTeD TrAnSaCtIoN"])("does not treat SQL Server %s as a block opener inside BEGIN...END", (transactionBegin) => {
+    const sql = `BEGIN
+  ${transactionBegin};
+  SELECT CASE
+    WHEN 1 = 1 THEN 1
+    ELSE 0
+  END;
+  COMMIT TRAN;
+END`;
+    const state = stateFor(sql, "sqlserver");
+
+    expect(foldedTextAtLine(state, 2)).toBeNull();
+    expect(foldedTextAtLine(state, 3)).toBe("\n    WHEN 1 = 1 THEN 1\n    ELSE 0\n  ");
+    const folded = foldedTextAtLine(state, 1);
+    expect(folded).not.toBeNull();
+    expect(folded).toContain("COMMIT TRAN;");
+  });
+
+  it("keeps a similarly prefixed BEGIN word as an ordinary block opener", () => {
+    const sql = "BEGIN TRANS\n  SELECT CASE\n    WHEN 1 = 1 THEN 1\n  END;\nEND";
+    const state = stateFor(sql, "sqlserver");
+
+    expect(foldedTextAtLine(state, 1)).toContain("SELECT CASE");
+    expect(foldedTextAtLine(state, 2)).toBe("\n    WHEN 1 = 1 THEN 1\n  ");
+  });
+
+  it("ignores keywords inside string/identifier/comment text", () => {
+    const sql = "BEGIN\n  SELECT 'BEGIN END CASE', \"END\";\n  /* BEGIN nested comment END */\nEND";
+    const state = stateFor(sql);
+
+    expect(foldedTextAtLine(state, 1)).toBe("\n  SELECT 'BEGIN END CASE', \"END\";\n  /* BEGIN nested comment END */\n");
+  });
+
+  it("returns null for a line with no block opener", () => {
+    const sql = "SELECT 1;\nSELECT 2;";
+    const state = stateFor(sql);
+
+    expect(foldedTextAtLine(state, 1)).toBeNull();
+    expect(foldedTextAtLine(state, 2)).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.
+    const sql = "PROCEDURE p() BEGIN END";
+    const state = stateFor(sql);
+
+    expect(foldedTextAtLine(state, 1)).toBeNull();
+  });
+
+  it("regression: a comment between END and CASE/TRY/CATCH must not break continuation detection", () => {
+    // Before the fix, the gap check required pure whitespace, so a comment there made `END`
+    // look like a bare closer; the literal `CASE`/`TRY`/`CATCH` word right after was then
+    // pushed as an unmatched new opener, desyncing every fold pairing after it.
+    const sql = "BEGIN\n  CASE v\n    WHEN 1 THEN SELECT 1;\n  END /* comment */ CASE;\nEND";
+    const state = stateFor(sql);
+
+    expect(foldedTextAtLine(state, 2)).toBe("\n    WHEN 1 THEN SELECT 1;\n  ");
+    const folded = foldedTextAtLine(state, 1);
+    expect(folded).not.toBeNull();
+    expect(folded).toContain("END /* comment */ CASE;");
+  });
+
+  it("still folds correctly after a dialect reconfigure with the document unchanged", () => {
+    // QueryEditor.vue's databaseType/dialect watcher reconfigures the language extension in a
+    // Compartment, via a dispatch with no doc `changes`, whenever the user switches a tab's
+    // connection dialect without editing text -- exercise that same sequence here.
+    const languageComp = new Compartment();
+    const sql = "BEGIN\n  BEGIN TRY\n    SELECT 1;\n  END TRY\nEND";
+    let state = EditorState.create({
+      doc: sql,
+      extensions: [languageComp.of(langSql.sql({ dialect: createDbxCodeMirrorSqlDialect(langSql, "mysql") })), sqlBlockFoldService],
+    });
+    ensureSyntaxTree(state, sql.length, 5_000);
+    expect(foldedTextAtLine(state, 1)).toContain("END TRY");
+
+    // Reconfigure only -- no `changes`, so `state.doc` (the `Text` instance) stays identical.
+    const before = state.doc;
+    state = state.update({
+      effects: languageComp.reconfigure(langSql.sql({ dialect: createDbxCodeMirrorSqlDialect(langSql, "sqlserver") })),
+    }).state;
+    expect(state.doc).toBe(before);
+    ensureSyntaxTree(state, sql.length, 5_000);
+
+    expect(foldedTextAtLine(state, 1)).toContain("END TRY");
+  });
+});

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

@@ -0,0 +1,117 @@
+import { foldService, syntaxTree } from "@codemirror/language";
+import type { EditorState } from "@codemirror/state";
+
+// `@lezer/common` is only a transitive dependency here (see sqlSyntaxTreeWindow.ts's comment on
+// the same pattern), so derive the node types structurally instead of importing them.
+type SyntaxNode = ReturnType<ReturnType<typeof syntaxTree>["resolve"]>;
+type Tree = ReturnType<typeof syntaxTree>;
+
+interface FoldRange {
+  from: number;
+  to: number;
+}
+
+// `@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.
+//
+// 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
+// trailing word. `IF`/`WHILE`/`LOOP` are deliberately excluded as openers: MySQL's
+// `IF(cond, a, b)` is a builtin scalar function, so treating bare `IF` as a block opener would
+// misparse every such call as an unterminated block and desync all fold ranges after it. T-SQL
+// (the dialect in #6574) doesn't need `IF` support anyway -- its `IF`/`ELSE` bodies are
+// themselves wrapped in `BEGIN...END`, which this already folds.
+const BLOCK_OPENERS = new Set(["BEGIN", "CASE"]);
+// Continuation words whose `END <word>` pairing closes something this service DID push an opener
+// for, so the `END` must pop the stack: `END CASE` closes a procedural `CASE` (the CASE-expression
+// form closes with plain `END` instead); `END TRY`/`END CATCH` close T-SQL's `BEGIN TRY`/
+// `BEGIN CATCH` -- the leading `BEGIN` there already pushed a generic entry (see BLOCK_OPENERS;
+// the following `TRY`/`CATCH` word itself isn't inspected on push), so its closer must be
+// recognized too or that entry is never popped and desyncs every fold range after it.
+const TRACKED_END_CONTINUATIONS = new Set(["CASE", "TRY", "CATCH"]);
+// `END IF`/`END WHILE`/`END LOOP` close constructs with no `BEGIN`/`CASE` prefix at all -- bare
+// `IF`/`WHILE`/`LOOP` are never pushed (see BLOCK_OPENERS comment), so these `END`s must NOT pop
+// the stack; nothing on it belongs to them.
+const UNTRACKED_END_CONTINUATIONS = new Set(["IF", "WHILE", "LOOP"]);
+
+// A comment is legal SQL wherever whitespace is, so "END /* note */ CASE" must still read as a
+// continuation; this matches a run of pure whitespace and/or `--`/`/* */` comments so real code
+// between the two words (e.g. a `;` starting an unrelated statement) still fails the check.
+const GAP_IS_WHITESPACE_OR_COMMENT = /^(?:\s|--[^\n]*|\/\*[\s\S]*?\*\/)*$/;
+
+// T-SQL transaction openers do not have a matching `END`, so they must not consume the closer
+// of an enclosing procedural `BEGIN...END` block. Match complete keyword tokens only: `TRAN` is
+// SQL Server's documented abbreviation, while similar identifiers such as `TRANS` remain blocks.
+const SQLSERVER_TRANSACTION_BEGIN_WORDS = new Set(["TRAN", "TRANSACTION"]);
+
+function isSqlServerTransactionBegin(state: EditorState, tokens: SyntaxNode[], index: number): boolean {
+  const next = tokens[index + 1];
+  if (!next || !GAP_IS_WHITESPACE_OR_COMMENT.test(state.sliceDoc(tokens[index].to, next.from))) return false;
+
+  const nextText = state.sliceDoc(next.from, next.to).toUpperCase();
+  if (SQLSERVER_TRANSACTION_BEGIN_WORDS.has(nextText)) return true;
+  if (nextText !== "DISTRIBUTED") return false;
+
+  const transaction = tokens[index + 2];
+  return Boolean(transaction && GAP_IS_WHITESPACE_OR_COMMENT.test(state.sliceDoc(next.to, transaction.from)) && state.sliceDoc(transaction.from, transaction.to).toUpperCase() === "TRANSACTION");
+}
+
+// Keyed by the `Tree` instance (not `Text`): the syntax tree is also invalidated when the SQL
+// dialect is reconfigured with the document unchanged (QueryEditor.vue's databaseType/dialect
+// watcher), and when Lezer's incremental parser finishes covering more of a large document in the
+// background -- both produce a new `Tree` without a new `Text`, so keying on `Text` alone would
+// keep serving ranges computed from a stale or partial parse.
+const rangeCache = new WeakMap<Tree, Map<number, FoldRange>>();
+
+function computeBlockFoldRanges(state: EditorState): Map<number, FoldRange> {
+  const tree = syntaxTree(state);
+  const cached = rangeCache.get(tree);
+  if (cached) return cached;
+
+  const byOpenerLine = new Map<number, FoldRange>();
+  const tokens: SyntaxNode[] = [];
+  tree.iterate({
+    enter(node) {
+      if (node.name === "Keyword") tokens.push(node.node);
+    },
+  });
+
+  const stack: SyntaxNode[] = [];
+  for (let i = 0; i < tokens.length; i++) {
+    const token = tokens[i];
+    const text = state.sliceDoc(token.from, token.to).toUpperCase();
+    if (text === "END") {
+      const next = tokens[i + 1];
+      const continuation = next && GAP_IS_WHITESPACE_OR_COMMENT.test(state.sliceDoc(token.to, next.from)) ? state.sliceDoc(next.from, next.to).toUpperCase() : null;
+      if (continuation && UNTRACKED_END_CONTINUATIONS.has(continuation)) {
+        i++; // "END IF"/"END WHILE"/... -- closes something we don't track; skip past it untouched
+        continue;
+      }
+      if (continuation && TRACKED_END_CONTINUATIONS.has(continuation)) i++; // "END CASE" -- consume the word, not a new opener
+      const opener = stack.pop();
+      if (opener) {
+        const openerLine = state.doc.lineAt(opener.to);
+        const endLine = state.doc.lineAt(token.from);
+        // Same source line as its opener (e.g. "BEGIN END"): nothing to collapse, and the
+        // 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 });
+        }
+      }
+    } else if (BLOCK_OPENERS.has(text) && (text !== "BEGIN" || !isSqlServerTransactionBegin(state, tokens, i))) {
+      stack.push(token);
+    }
+  }
+
+  rangeCache.set(tree, byOpenerLine);
+  return byOpenerLine;
+}
+
+/** Folds `BEGIN...END` and `CASE...END` blocks in the SQL editor -- see module doc comment. */
+export const sqlBlockFoldService = foldService.of((state, lineStart) => {
+  const range = computeBlockFoldRanges(state).get(state.doc.lineAt(lineStart).number);
+  return range ?? null;
+});