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

apps/desktop/src/lib/sql/sqlStatementRanges.ts

@@ -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 {