From 4b7533416449f4caeb64794fe2732ce04896aa23 Mon Sep 17 00:00:00 2001 From: zipg Date: Fri, 21 Aug 2026 18:13:55 +0800 Subject: [PATCH] fix(editor): avoid quadratic SQL folding scans --- .../codemirrorSqlBlockFolding.spec.ts | 97 ++++++++++++++++- .../lib/editor/codemirrorSqlBlockFolding.ts | 101 +++++++++++++++++- 2 files changed, 191 insertions(+), 7 deletions(-) diff --git a/apps/desktop/src/lib/editor/__tests__/codemirrorSqlBlockFolding.spec.ts b/apps/desktop/src/lib/editor/__tests__/codemirrorSqlBlockFolding.spec.ts index 6059412440..4593b65a6c 100644 --- a/apps/desktop/src/lib/editor/__tests__/codemirrorSqlBlockFolding.spec.ts +++ b/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. diff --git a/apps/desktop/src/lib/editor/codemirrorSqlBlockFolding.ts b/apps/desktop/src/lib/editor/codemirrorSqlBlockFolding.ts index 174a086116..8ee0f00d60 100644 --- a/apps/desktop/src/lib/editor/codemirrorSqlBlockFolding.ts +++ b/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>(); +function addMultilineFoldRange(state: EditorState, ranges: Map, 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) { + const scopes = new Map(); + 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 { const tree = syntaxTree(state); const cached = rangeCache.get(tree); @@ -98,7 +187,7 @@ function computeBlockFoldRanges(state: EditorState): Map { // 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 { } } + 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;