diff --git a/apps/desktop/src/lib/__tests__/sql/sqlStatementRanges.spec.ts b/apps/desktop/src/lib/__tests__/sql/sqlStatementRanges.spec.ts index 7e37e186f6..7676684314 100644 --- a/apps/desktop/src/lib/__tests__/sql/sqlStatementRanges.spec.ts +++ b/apps/desktop/src/lib/__tests__/sql/sqlStatementRanges.spec.ts @@ -407,6 +407,92 @@ COMMENT = '测试';`; expect(range?.sql.trim()).toBe("EXPLAIN\nSELECT * FROM users"); }); + it("keeps issue #3567 EXPLAIN options and CTE target as one statement", () => { + const sql = "explain (analyze,buffers)\nwith tmp as(select* from test.tt)\nselect * from tmp;"; + + expect(statementRangeAtCursor(sql, indexOf(sql, "explain"), "postgres")?.sql.trim()).toBe(sql.slice(0, -1)); + expect(statementRangeAtCursor(sql, indexOf(sql, "select * from tmp"), "postgres")?.sql.trim()).toBe(sql.slice(0, -1)); + expect(rangeSqlTexts(executableStatementRanges(sql, "postgres"))).toEqual([sql.slice(0, -1)]); + expect(rangeSqlTexts(executableStatementRanges(sql))).toEqual([sql.slice(0, -1)]); + }); + + it("keeps EXPLAIN ANALYZE with a CTE main query as one statement", () => { + const sql = "explain analyze\nwith tmp as (select 1)\nselect * from tmp;"; + + expect(statementRangeAtCursor(sql, indexOf(sql, "explain"), "postgres")?.sql.trim()).toBe(sql.slice(0, -1)); + expect(rangeSqlTexts(executableStatementRanges(sql, "postgres"))).toEqual([sql.slice(0, -1)]); + }); + + it("keeps a plain EXPLAIN CTE target as one statement without merging later queries", () => { + const sql = "EXPLAIN\nWITH tmp AS (SELECT 1)\nSELECT * FROM tmp\nSELECT * FROM logs;"; + const expected = "EXPLAIN\nWITH tmp AS (SELECT 1)\nSELECT * FROM tmp"; + + expect(statementRangeAtCursor(sql, indexOf(sql, "EXPLAIN"))?.sql.trim()).toBe(expected); + expect(rangeSqlTexts(executableStatementRanges(sql))).toEqual([expected, "SELECT * FROM logs"]); + }); + + it("does not merge a query after an inline EXPLAIN options CTE statement", () => { + const sql = "EXPLAIN (ANALYZE) WITH tmp AS (SELECT 1)\nSELECT * FROM tmp\nSELECT 2;"; + const expected = "EXPLAIN (ANALYZE) WITH tmp AS (SELECT 1)\nSELECT * FROM tmp"; + + expect(statementRangeAtCursor(sql, indexOf(sql, "EXPLAIN"), "postgres")?.sql.trim()).toBe(expected); + expect(rangeSqlTexts(executableStatementRanges(sql, "postgres"))).toEqual([expected, "SELECT 2"]); + }); + + it("keeps EXPLAIN options CTE UPDATE assignments as one statement", () => { + const sql = "EXPLAIN (ANALYZE)\nWITH tmp AS (SELECT 1)\nUPDATE t\nSET x = 1;"; + + expect(statementRangeAtCursor(sql, indexOf(sql, "SET"), "postgres")?.sql.trim()).toBe(sql.slice(0, -1)); + expect(rangeSqlTexts(executableStatementRanges(sql, "postgres"))).toEqual([sql.slice(0, -1)]); + }); + + it("keeps CTE INSERT ... SELECT with the WITH statement", () => { + const sql = "WITH tmp AS (SELECT 1)\nINSERT INTO t (id)\nSELECT * FROM tmp;"; + + expect(statementRangeAtCursor(sql, indexOf(sql, "INSERT"))?.sql.trim()).toBe(sql.slice(0, -1)); + expect(rangeSqlTexts(executableStatementRanges(sql))).toEqual([sql.slice(0, -1)]); + }); + + it("does not merge a query after a CTE INSERT ... SELECT statement", () => { + const sql = "WITH tmp AS (SELECT 1)\nINSERT INTO t (id)\nSELECT * FROM tmp\nSELECT 2;"; + const expected = "WITH tmp AS (SELECT 1)\nINSERT INTO t (id)\nSELECT * FROM tmp"; + + expect(statementRangeAtCursor(sql, indexOf(sql, "INSERT"))?.sql.trim()).toBe(expected); + expect(rangeSqlTexts(executableStatementRanges(sql))).toEqual([expected, "SELECT 2"]); + }); + + it("does not merge a query after an INSERT with a CTE source query", () => { + const sql = "INSERT INTO t (id)\nWITH tmp AS (SELECT 1)\nSELECT * FROM tmp\nSELECT 2;"; + const expected = "INSERT INTO t (id)\nWITH tmp AS (SELECT 1)\nSELECT * FROM tmp"; + + expect(statementRangeAtCursor(sql, indexOf(sql, "INSERT"))?.sql.trim()).toBe(expected); + expect(rangeSqlTexts(executableStatementRanges(sql))).toEqual([expected, "SELECT 2"]); + }); + + it("skips block comments inside EXPLAIN options when resolving the target", () => { + const sql = "EXPLAIN (ANALYZE /* ) */) UPDATE t\nSET x = 1;"; + + expect(statementRangeAtCursor(sql, indexOf(sql, "SET"), "postgres")?.sql.trim()).toBe(sql.slice(0, -1)); + expect(rangeSqlTexts(executableStatementRanges(sql, "postgres"))).toEqual([sql.slice(0, -1)]); + }); + + it("recovers later statements from an unclosed EXPLAIN option list", () => { + const sql = "EXPLAIN (ANALYZE\nSELECT 1\nSELECT 2;"; + const explainSql = "EXPLAIN (ANALYZE\nSELECT 1"; + + expect(statementRangeAtCursor(sql, indexOf(sql, "EXPLAIN"), "postgres")?.sql.trim()).toBe(explainSql); + expect(statementRangeAtCursor(sql, indexOf(sql, "SELECT 2"), "postgres")?.sql.trim()).toBe("SELECT 2"); + expect(rangeSqlTexts(executableStatementRanges(sql, "postgres"))).toEqual([explainSql, "SELECT 2"]); + }); + + it("does not treat a DESCRIBE subquery paren as EXPLAIN options", () => { + const sql = "DESCRIBE (SELECT 1)\nSELECT 2;"; + + expect(statementRangeAtCursor(sql, indexOf(sql, "DESCRIBE"), "clickhouse")?.sql.trim()).toBe("DESCRIBE (SELECT 1)"); + expect(rangeSqlTexts(executableStatementRanges(sql, "clickhouse"))).toEqual(["DESCRIBE (SELECT 1)", "SELECT 2"]); + expect(rangeSqlTexts(executableStatementRanges(sql))).toEqual(["DESCRIBE (SELECT 1)", "SELECT 2"]); + }); + it("keeps MySQL DESC UPDATE joins as one statement", () => { const sql = "desc update test_orders a\njoin test_users b\non a.id=b.id \nset a.name = '张三'\nwhere b.id > 10;"; expect(statementRangeAtCursor(sql, indexOf(sql, "desc"), "mysql")?.sql.trim()).toBe(sql.slice(0, -1)); diff --git a/apps/desktop/src/lib/sql/sqlStatementRanges.ts b/apps/desktop/src/lib/sql/sqlStatementRanges.ts index 75aa3a9a19..20e0f5c1b9 100644 --- a/apps/desktop/src/lib/sql/sqlStatementRanges.ts +++ b/apps/desktop/src/lib/sql/sqlStatementRanges.ts @@ -487,8 +487,12 @@ function splitStatementRangeAtSoftStarts(sql: string, statement: RawStatement, d for (const lineStart of lineStarts) { if (lineStart.from <= statement.from) continue; - if (currentKeyword === "WITH" && !consumedWithMainStatement && WITH_MAIN_STATEMENT_KEYWORDS.has(lineStart.keyword)) { + if (currentBodyKeyword === "WITH" && !consumedWithMainStatement && WITH_MAIN_STATEMENT_KEYWORDS.has(lineStart.keyword)) { consumedWithMainStatement = true; + // The CTE main statement also satisfies a pending EXPLAIN target, and its + // own body rules (e.g. UPDATE ... SET) must take over from here. + consumedExplainStatement = true; + currentBodyKeyword = lineStart.keyword; continue; } @@ -511,6 +515,10 @@ function splitStatementRangeAtSoftStarts(sql: string, statement: RawStatement, d } if (currentBodyKeyword === "INSERT" && INSERT_BODY_KEYWORDS.has(lineStart.keyword)) { + // Hand over to the source query's own rules so only its continuations + // (CTE main statement, set operations) stay attached to the INSERT. + currentBodyKeyword = lineStart.keyword; + if (lineStart.keyword === "WITH") consumedWithMainStatement = false; continue; } @@ -553,6 +561,10 @@ function splitStatementRangeAtSoftStarts(sql: string, statement: RawStatement, d function topLevelSoftStatementLineStarts(sql: string, statement: RawStatement, databaseType?: DatabaseType): Array<{ hitFrom: number; from: number; keyword: string }> { const starts: Array<{ hitFrom: number; from: number; keyword: string }> = []; const len = statement.to; + const explainOptionsStart = explainOptionsParenAt(sql, statement.from); + // Recover soft statement boundaries while the user is still typing an + // EXPLAIN option list; otherwise its unmatched opener hides every later line. + const unclosedExplainOptionsStart = explainOptionsStart !== null && skipBalancedParens(sql, explainOptionsStart) === null ? explainOptionsStart : null; let state: QuoteState | "lineComment" | "blockComment" = "none"; let dollarTag = ""; let parenDepth = 0; @@ -702,7 +714,7 @@ function topLevelSoftStatementLineStarts(sql: string, statement: RawStatement, d continue; } } - if (ch === "(") { + if (ch === "(" && i !== unclosedExplainOptionsStart) { parenDepth += 1; } else if (ch === ")" && parenDepth > 0) { parenDepth -= 1; @@ -926,6 +938,15 @@ function isExplainLikeKeyword(keyword: string | null): boolean { return keyword === "EXPLAIN" || keyword === "DESCRIBE" || keyword === "DESC"; } +function explainOptionsParenAt(sql: string, pos: number): number | null { + const prefixMatch = /^[A-Za-z_][\w$]*/.exec(sql.slice(pos)); + if (prefixMatch?.[0]?.toUpperCase() !== "EXPLAIN") return null; + + let i = pos + prefixMatch[0].length; + while (i < sql.length && isSqlWhitespace(sql[i])) i += 1; + return sql[i] === "(" ? i : null; +} + function explainLikeTargetKeywordAt(sql: string, pos: number): string | null { const prefixMatch = /^[A-Za-z_][\w$]*/.exec(sql.slice(pos)); const prefix = prefixMatch?.[0]?.toUpperCase(); @@ -933,11 +954,93 @@ function explainLikeTargetKeywordAt(sql: string, pos: number): string | null { let i = pos + (prefixMatch?.[0].length ?? 0); while (i < sql.length && isSqlWhitespace(sql[i])) i += 1; + // Parenthesized EXPLAIN options (e.g. Postgres `EXPLAIN (ANALYZE, BUFFERS) ...`) + // sit between the keyword and its target statement. DESC/DESCRIBE take no + // options — a paren there is a subquery (ClickHouse `DESCRIBE (SELECT ...)`). + if (prefix === "EXPLAIN" && sql[i] === "(") { + const optionsEnd = skipBalancedParens(sql, i); + if (optionsEnd === null) return null; + i = optionsEnd; + while (i < sql.length && isSqlWhitespace(sql[i])) i += 1; + } const targetMatch = /^[A-Za-z_][\w$]*/.exec(sql.slice(i)); const targetKeyword = targetMatch?.[0]?.toUpperCase(); return targetKeyword && EXPLAIN_STATEMENT_KEYWORDS.has(targetKeyword) ? targetKeyword : null; } +function skipBalancedParens(sql: string, pos: number): number | null { + let state: "none" | "single" | "double" | "lineComment" | "blockComment" = "none"; + let depth = 0; + let i = pos; + + while (i < sql.length) { + const ch = sql[i]; + const next = sql[i + 1] ?? ""; + + if (state === "lineComment") { + if (ch === "\n") state = "none"; + i += 1; + continue; + } + if (state === "blockComment") { + if (ch === "*" && next === "/") { + state = "none"; + i += 2; + continue; + } + i += 1; + continue; + } + if (state === "single") { + if (ch === "'" && next === "'") { + i += 2; + continue; + } + if (ch === "'") state = "none"; + i += 1; + continue; + } + if (state === "double") { + if (ch === '"' && next === '"') { + i += 2; + continue; + } + if (ch === '"') state = "none"; + i += 1; + continue; + } + + if (ch === "-" && next === "-") { + state = "lineComment"; + i += 2; + continue; + } + if (ch === "/" && next === "*") { + state = "blockComment"; + i += 2; + continue; + } + if (ch === "'") { + state = "single"; + i += 1; + continue; + } + if (ch === '"') { + state = "double"; + i += 1; + continue; + } + if (ch === "(") depth += 1; + if (ch === ")") { + depth -= 1; + if (depth === 0) return i + 1; + } + i += 1; + } + + return null; +} + function startsLineComment(sql: string, pos: number): boolean { return (sql[pos] === "-" && sql[pos + 1] === "-") || sql[pos] === "#"; }