fix(sqlserver): keep IF/ELSE control-flow batches in one range

2026-10-02 02:34:42 +08:00 · 2026-09-22 13:16:56 +08:00
parent f14f4f7390
commit 80f8ed81fb
7 changed files with 293 additions and 7 deletions
@@ -2023,6 +2023,105 @@ WHERE t2.product_name = '12345'
    expect(rangeSqlTexts(ranges)).toEqual([sql]);
    expect(candidateSummaries(candidates)).toEqual([`all:${sql}`]);
  });
+
+  it("keeps a SQL Server IF/ELSE batch whole when it follows another statement", () => {
+    const batch = ["IF NOT EXISTS (SELECT 1 FROM dbo.QRTZ_JOB_DETAILS WHERE job_name = N'x')", "BEGIN", "    SELECT 1;", "END", "ELSE", "BEGIN", "    SELECT 2;", "END"].join("\n");
+    const sql = `SELECT 1;\n${batch}`;
+    const ranges = executableStatementRanges(sql, "sqlserver");
+
+    expect(rangeSqlTexts(ranges)).toEqual(["SELECT 1", batch]);
+  });
+
+  it("keeps a SQL Server IF/ELSE batch whole when its branches hold semicolons", () => {
+    const sql = ["IF NOT EXISTS (SELECT 1 FROM dbo.QRTZ_JOB_DETAILS WHERE job_name = N'x')", "BEGIN", "    SELECT 1;", "END", "ELSE", "BEGIN", "    SELECT 2;", "END"].join("\n");
+    const ranges = executableStatementRanges(sql, "sqlserver");
+    const candidates = buildExecutionCandidates(sql, indexOf(sql, "SELECT 2"), "sqlserver");
+
+    expect(rangeSqlTexts(ranges)).toEqual([sql]);
+    expect(candidateSummaries(candidates)).toEqual([`all:${sql}`]);
+  });
+
+  it("keeps SQL Server IF branches without BEGIN/END blocks whole", () => {
+    const sql = [
+      "IF NOT EXISTS (SELECT 1 FROM ::fn_listextendedproperty(N'MS_Description', N'USER', N'dbo', N'TABLE', N'Categories', N'COLUMN', N'CategoryID'))",
+      "    EXEC sp_addextendedproperty @name=N'MS_Description', @value=N'test'",
+      "ELSE",
+      "    EXEC sp_updateextendedproperty @name=N'MS_Description', @value=N'test'",
+    ].join("\n");
+    const ranges = executableStatementRanges(sql, "sqlserver");
+
+    expect(rangeSqlTexts(ranges)).toEqual([sql]);
+  });
+
+  it("keeps a SQL Server IF/BEGIN/END batch without ELSE whole", () => {
+    const sql = ["IF @x = 1", "BEGIN", "    SELECT 1;", "    SELECT 2;", "END"].join("\n");
+    const ranges = executableStatementRanges(sql, "sqlserver");
+
+    expect(rangeSqlTexts(ranges)).toEqual([sql]);
+  });
+
+  it("keeps SQL Server WHILE batches whole", () => {
+    const sql = ["WHILE @i < 10", "BEGIN", "    SET @i = @i + 1;", "    IF @i = 5 CONTINUE;", "END"].join("\n");
+    const ranges = executableStatementRanges(sql, "sqlserver");
+
+    expect(rangeSqlTexts(ranges)).toEqual([sql]);
+  });
+
+  it("does not merge SQL Server BEGIN TRAN statements with the following batch", () => {
+    const sql = ["BEGIN TRAN;", "UPDATE dbo.T SET x = 1;", "COMMIT;"].join("\n");
+    const ranges = executableStatementRanges(sql, "sqlserver");
+
+    expect(rangeSqlTexts(ranges)).toEqual(["BEGIN TRAN", "UPDATE dbo.T SET x = 1", "COMMIT"]);
+  });
+
+  // The depth carried across fragments decides where a batch ends: fragments
+  // after the closing `END` are independent statements, so every following
+  // statement keeps its own execution icon.
+  it("does not swallow the statement after a SQL Server IF/ELSE batch", () => {
+    const batch = ["IF NOT EXISTS (SELECT 1 FROM dbo.T WHERE n = N'x')", "BEGIN", "    SELECT 1;", "END", "ELSE", "BEGIN", "    SELECT 2;", "END"].join("\n");
+    const ranges = executableStatementRanges(`${batch}\nSELECT 999;`, "sqlserver");
+
+    expect(rangeSqlTexts(ranges)).toEqual([batch, "SELECT 999"]);
+  });
+
+  it("keeps two consecutive SQL Server IF/ELSE batches as two ranges", () => {
+    const first = ["IF NOT EXISTS (SELECT 1 FROM dbo.T WHERE n = N'x')", "BEGIN", "    SELECT 1;", "END", "ELSE", "BEGIN", "    SELECT 2;", "END"].join("\n");
+    const second = ["IF NOT EXISTS (SELECT 1 FROM dbo.T WHERE n = N'y')", "BEGIN", "    SELECT 1;", "END", "ELSE", "BEGIN", "    SELECT 2;", "END"].join("\n");
+    const ranges = executableStatementRanges(`${first}\n${second}`, "sqlserver");
+
+    expect(rangeSqlTexts(ranges)).toEqual([first, second]);
+  });
+
+  it("does not merge a SQL Server batch across a GO separator", () => {
+    const batch = ["IF NOT EXISTS (SELECT 1 FROM dbo.T WHERE n = N'x')", "BEGIN", "    SELECT 1;", "END", "ELSE", "BEGIN", "    SELECT 2;", "END"].join("\n");
+    const insert = "INSERT INTO dbo.T (n) VALUES (N'z')";
+    const ranges = executableStatementRanges(`${batch}\nGO\n${insert};`, "sqlserver");
+
+    expect(rangeSqlTexts(ranges)).toEqual([batch, insert]);
+  });
+
+  it("ends a single-line SQL Server IF/BEGIN/END batch before the next statement", () => {
+    const ranges = executableStatementRanges("IF @x = 1 BEGIN SELECT 1; END\nSELECT 999;", "sqlserver");
+
+    expect(rangeSqlTexts(ranges)).toEqual(["IF @x = 1 BEGIN SELECT 1; END", "SELECT 999"]);
+  });
+
+  // T-SQL semicolons are optional: when the ELSE branch tail carries none, the
+  // branch's own `END` and the next statement share one `;`-fragment, and the
+  // second closure must replace the `ELSE`-continuing first one.
+  it("ends a SQL Server IF/ELSE batch at the ELSE branch END without a semicolon", () => {
+    const batch = ["IF @x = 1", "BEGIN", "    SELECT 1;", "END", "ELSE", "BEGIN", "    SELECT 2", "END"].join("\n");
+    const ranges = executableStatementRanges(`${batch}\nSELECT 999;`, "sqlserver");
+
+    expect(rangeSqlTexts(ranges)).toEqual([batch, "SELECT 999"]);
+  });
+
+  it("keeps the comment after a SQL Server batch out of the next statement", () => {
+    const batch = ["IF @x = 1", "BEGIN", "    SELECT 1;", "END"].join("\n");
+    const ranges = executableStatementRanges(`${batch}\n-- gap\nSELECT 999;`, "sqlserver");
+
+    expect(rangeSqlTexts(ranges)).toEqual([batch, "SELECT 999"]);
+  });
 });

 describe("hasMultipleExecutionTargets", () => {
@@ -668,6 +668,7 @@ export function splitSqlStatementRanges(sql: string, databaseType?: DatabaseType
  // Flush any trailing statement that lacks a terminating semicolon.
  flush();

+  if (databaseType === "sqlserver") return mergeSqlServerControlFlowBatches(sql, statements, databaseType, parameterOptions);
  return statements;
 }

@@ -776,7 +777,7 @@ function splitStatementRangeAtSoftStarts(sql: string, statement: RawStatement, d
  // Routine bodies contain top-level-looking SET/INSERT/SELECT lines that are not independent statements.
  if (isMysqlRoutineBlockDatabase(databaseType) && startsWithMysqlRoutineBlock(statement.sql, parameterOptions)) return [statement];
  // SQL Server control-flow batches use line-oriented BEGIN/EXEC tokens inside one IF/ELSE statement.
-  if (isSqlServerIfElseControlFlowBatch(sql, statement, databaseType, parameterOptions)) return [statement];
+  if (isSqlServerControlFlowBatch(sql, statement, databaseType, parameterOptions)) return [statement];

  const lineStarts = topLevelSoftStatementLineStarts(sql, statement, databaseType, parameterOptions);
  if (lineStarts.length <= 1) return [statement];
@@ -882,11 +883,182 @@ function splitStatementRangeAtSoftStarts(sql: string, statement: RawStatement, d
  return ranges.length > 0 ? ranges : [statement];
 }

-function isSqlServerIfElseControlFlowBatch(sql: string, statement: RawStatement, databaseType?: DatabaseType, parameterOptions?: SqlParameterOptions): boolean {
-  if (databaseType !== "sqlserver" || !startsWithSqlWords(sql, statement.from, ["IF"], databaseType, parameterOptions)) return false;
+// `BEGIN` opens a control-flow block only when it does not start a transaction or a
+// conversation (`BEGIN TRAN`, `BEGIN DISTRIBUTED TRANSACTION`, `BEGIN DIALOG
+// CONVERSATION`); those have no matching `END`.
+const SQLSERVER_NON_BLOCK_BEGIN_KEYWORDS = new Set(["TRANSACTION", "TRAN", "DISTRIBUTED", "DIALOG", "CONVERSATION"]);
+// `END CONVERSATION`/`END DIALOG` close a conversation instead of a BEGIN/CASE block.
+const SQLSERVER_NON_BLOCK_END_KEYWORDS = new Set(["CONVERSATION", "DIALOG"]);

-  const words = topLevelWordsBefore(sql, statement.from, statement.to, 64, databaseType, parameterOptions);
-  return words.includes("ELSE") && words.includes("BEGIN") && words.includes("END");
+interface SqlServerControlFlowScan {
+  /** BEGIN/CASE blocks left open at the end of the scanned fragment. */
+  openBlocks: number;
+  /** Offset just after the token that closed the block carried into this
+   *  fragment, when the carried depth reached zero inside it. `null` when the
+   *  fragment was scanned from depth 0 or leaves the carried block open. */
+  closedAt: number | null;
+  /** First top-level word after the closing token (`""` when nothing follows),
+   *  so the caller can tell `END ELSE ...` (same IF) from `END <next statement>`. */
+  wordAfterClose: string;
+  /** A top-level `BEGIN` that really opens a block was seen. */
+  sawBlockBegin: boolean;
+  /** A top-level `ELSE` belonging to the enclosing `IF` (not to a `CASE`) was seen. */
+  hasControlFlowElse: boolean;
+}
+
+/**
+ * Collect the control-flow facts of a T-SQL fragment: how many BEGIN/CASE blocks
+ * it leaves open and whether it contains an `IF`-level `ELSE`. Tokens inside
+ * parentheses, string literals and comments are ignored, and `CASE ... END`
+ * nesting keeps a `CASE`'s own `ELSE` from looking like an `IF` branch.
+ *
+ * `initialOpenBlocks` is the depth the previous fragment of the same statement
+ * left open. Scanning with the carried depth (instead of every fragment
+ * restarting from zero) is what lets a fragment that begins by closing the
+ * previous block — `END ELSE BEGIN SELECT 2`, or `END` followed by the next
+ * statement — report where the block actually ends.
+ */
+function scanSqlServerControlFlow(sql: string, from: number, to: number, databaseType?: DatabaseType, parameterOptions?: SqlParameterOptions, initialOpenBlocks = 0): SqlServerControlFlowScan {
+  const { words, ends } = topLevelWordsInRange(sql, from, to, databaseType, parameterOptions);
+  let openBlocks = initialOpenBlocks;
+  let closedAt: number | null = null;
+  let wordAfterClose = "";
+  let sawBlockBegin = false;
+  let hasControlFlowElse = false;
+
+  for (let index = 0; index < words.length; index += 1) {
+    const word = words[index];
+    const next = words[index + 1];
+    if (word === "BEGIN") {
+      if (next !== undefined && SQLSERVER_NON_BLOCK_BEGIN_KEYWORDS.has(next)) continue;
+      sawBlockBegin = true;
+      openBlocks += 1;
+      continue;
+    }
+    if (word === "CASE") {
+      openBlocks += 1;
+      continue;
+    }
+    if (word === "END") {
+      if (next !== undefined && SQLSERVER_NON_BLOCK_END_KEYWORDS.has(next)) continue;
+      openBlocks -= 1;
+      // Only the closure of the *carried* block is a statement boundary: a
+      // block opened and closed inside this fragment (`IF ... BEGIN ... END`
+      // in one piece) belongs to the statement that contains it. A closure
+      // followed by `ELSE` only continues the same IF, so a later closure in
+      // the same fragment (the ELSE branch's own `END`, reachable when the
+      // branch tail carries no semicolon) must still be able to replace it;
+      // a closure followed by anything else is final.
+      if (initialOpenBlocks > 0 && openBlocks === 0 && (closedAt === null || wordAfterClose === "ELSE")) {
+        closedAt = ends[index];
+        wordAfterClose = next ?? "";
+      }
+      continue;
+    }
+    if (word === "ELSE" && openBlocks === 0) {
+      hasControlFlowElse = true;
+    }
+  }
+
+  return { openBlocks, closedAt, wordAfterClose, sawBlockBegin, hasControlFlowElse };
+}
+
+/**
+ * True when `statement` is a whole T-SQL control-flow batch: an `IF`/`WHILE`
+ * whose body is a `BEGIN ... END` block or whose branches are plain statements
+ * with an `ELSE`. Such a batch contains no independent statements, so it must
+ * stay one execution range instead of being split at the line-oriented
+ * `BEGIN`/`EXEC`/`END` tokens.
+ */
+function isSqlServerControlFlowBatch(sql: string, statement: RawStatement, databaseType?: DatabaseType, parameterOptions?: SqlParameterOptions): boolean {
+  if (databaseType !== "sqlserver") return false;
+  const firstWord = nextSqlWordToken(sql, statement.from, databaseType, parameterOptions)?.word;
+  if (firstWord !== "IF" && firstWord !== "WHILE") return false;
+
+  const scan = scanSqlServerControlFlow(sql, statement.from, statement.to, databaseType, parameterOptions);
+  return scan.hasControlFlowElse || (scan.sawBlockBegin && scan.openBlocks === 0);
+}
+
+/**
+ * T-SQL does not end a `BEGIN ... END` block at a semicolon, so splitting on
+ * every top-level `;` cuts `IF ... BEGIN ... ; ... END` batches into fragments
+ * whose `BEGIN`/`EXEC`/`END` lines then look like independent statements
+ * (#9336). Re-join the fragments of such a batch so it stays one range.
+ */
+function mergeSqlServerControlFlowBatches(sql: string, statements: RawStatement[], databaseType: DatabaseType, parameterOptions?: SqlParameterOptions): RawStatement[] {
+  const merged: RawStatement[] = [];
+  const pending = [...statements];
+
+  while (pending.length > 0) {
+    const first = pending.shift()!;
+    const firstWord = nextSqlWordToken(sql, first.from, databaseType, parameterOptions)?.word;
+    if (firstWord !== "IF" && firstWord !== "WHILE") {
+      merged.push(first);
+      continue;
+    }
+
+    // Thread the block depth through the fragments instead of summing
+    // independently scanned contributions: `END ELSE BEGIN SELECT 2` closes the
+    // block the previous fragment opened (net 0) and `END` closes it outright
+    // (net -1), so a per-fragment depth restarted at zero would keep the batch
+    // open forever and swallow every following statement.
+    let last = first;
+    let openBlocks = scanSqlServerControlFlow(sql, first.from, first.to, databaseType, parameterOptions).openBlocks;
+    let closedAt: number | null = null;
+    while (openBlocks > 0 && pending.length > 0) {
+      const fragment = pending.shift()!;
+      const scan = scanSqlServerControlFlow(sql, fragment.from, fragment.to, databaseType, parameterOptions, openBlocks);
+      openBlocks = scan.openBlocks;
+      last = fragment;
+      // `END ELSE ...` continues the same IF, so only a closure that is *not*
+      // followed by ELSE ends the batch — and anything the fragment still holds
+      // after that point starts the next statement.
+      if (scan.closedAt !== null && scan.wordAfterClose !== "ELSE") {
+        closedAt = scan.closedAt;
+        break;
+      }
+    }
+
+    const to = closedAt ?? last.to;
+    if (to > first.from) {
+      merged.push({ hitFrom: first.hitFrom, from: first.from, to, sql: sql.slice(first.from, to) });
+    } else {
+      merged.push(first);
+    }
+    if (closedAt !== null) {
+      const remainderFrom = skipSqlTrivia(sql, closedAt, last.to, databaseType, parameterOptions);
+      if (remainderFrom < last.to) {
+        pending.unshift({ hitFrom: closedAt, from: remainderFrom, to: last.to, sql: sql.slice(remainderFrom, last.to) });
+      }
+    }
+  }
+
+  return merged;
+}
+
+/**
+ * First offset in `sql[from, to)` that is neither whitespace nor a comment,
+ * clamped to `to`. The `;`-split path keeps the comment between two statements
+ * out of both of them; the control-flow remainder re-queue must do the same so
+ * a comment after a closing `END` is not glued onto the next statement's range.
+ */
+function skipSqlTrivia(sql: string, from: number, to: number, databaseType?: DatabaseType, parameterOptions?: SqlParameterOptions): number {
+  let index = from;
+  for (;;) {
+    while (index < to && isSqlWhitespace(sql[index])) index += 1;
+    if (index >= to) return index;
+    if (startsLineComment(sql, index, databaseType, parameterOptions)) {
+      while (index < to && sql[index] !== "\n") index += 1;
+      continue;
+    }
+    if (startsBlockComment(sql, index)) {
+      const close = sql.indexOf("*/", index + 2);
+      if (close === -1 || close + 2 > to) return to;
+      index = close + 2;
+      continue;
+    }
+    return index;
+  }
 }

 function topLevelSoftStatementLineStarts(sql: string, statement: RawStatement, databaseType?: DatabaseType, parameterOptions?: SqlParameterOptions): Array<{ hitFrom: number; from: number; keyword: string }> {
@@ -1129,7 +1301,18 @@ function startsWithMysqlCreateTable(sql: string, statementFrom: number): boolean
 }

 function topLevelWordsBefore(sql: string, from: number, to: number, limit: number, databaseType?: DatabaseType, parameterOptions?: SqlParameterOptions): string[] {
+  return topLevelWordsInRange(sql, from, to, databaseType, parameterOptions, limit).words;
+}
+
+/** Top-level (paren-depth 0) keywords of `sql[from, to)`, in order, each paired
+ *  with the offset just after it.
+ *
+ * `tailLimit` keeps only the last N words (callers that only look at the words
+ * right before a position, e.g. `MERGE ... THEN`), so scanning a long fragment
+ * no longer allocates the whole word list. */
+function topLevelWordsInRange(sql: string, from: number, to: number, databaseType?: DatabaseType, parameterOptions?: SqlParameterOptions, tailLimit?: number): { words: string[]; ends: number[] } {
  const words: string[] = [];
+  const ends: number[] = [];
  const backslashEscapes = allowsBackslashStringEscape(databaseType);
  let state: QuoteState | "lineComment" | "blockComment" = "none";
  let dollarTag = "";
@@ -1273,8 +1456,12 @@ function topLevelWordsBefore(sql: string, from: number, to: number, limit: numbe
    if (parenDepth === 0) {
      const match = /^[A-Za-z_][\w$]*/.exec(sql.slice(i));
      if (match) {
+        if (tailLimit !== undefined && words.length >= tailLimit) {
+          words.shift();
+          ends.shift();
+        }
        words.push(match[0].toUpperCase());
-        if (words.length > limit) words.shift();
+        ends.push(i + match[0].length);
        i += match[0].length;
        continue;
      }
@@ -1282,7 +1469,7 @@ function topLevelWordsBefore(sql: string, from: number, to: number, limit: numbe
    i += 1;
  }

-  return words;
+  return { words, ends };
 }

 function nextNonWhitespaceChar(sql: string, pos: number): string | null {