fix(sql): avoid repeated full scans in large script execution

Refs: #3335

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

@@ -1,4 +1,4 @@
-import { describe, expect, it } from "vitest";
+import { describe, expect, it, vi } from "vitest";
 import { executionCandidateForMode } from "@/lib/sql/sqlExecutionTarget";
 import { buildExecutionCandidates, currentExecutableStatementRange, executableStatementRanges, fullSqlRange, hasMultipleExecutionTargets, splitSqlStatementRanges, statementRangeAtCursor, stripMysqlClientDisplayCommand, supportsExecutionTargetPicker } from "@/lib/sql/sqlStatementRanges";
 
@@ -422,6 +422,48 @@ describe("stripMysqlClientDisplayCommand", () => {
 });
 
 describe("splitSqlStatementRanges", () => {
+  it.each([
+    { databaseType: "postgres" as const, lineEnding: "\n" },
+    { databaseType: "postgres" as const, lineEnding: "\r" },
+    { databaseType: "sqlserver" as const, lineEnding: "\n" },
+    { databaseType: "sqlserver" as const, lineEnding: "\r" },
+    { databaseType: "mysql" as const, lineEnding: "\n" },
+    { databaseType: "mysql" as const, lineEnding: "\r" },
+  ])("does not rescan the remaining script for an absent line ending: %j", ({ databaseType, lineEnding }) => {
+    const statements = Array.from({ length: 500 }, (_, index) => `${databaseType === "mysql" ? "/* trace */ " : ""}SELECT ${index};`);
+    const sql = statements.join(lineEnding);
+    const originalIndexOf = String.prototype.indexOf;
+    let absentLineEndingSearches = 0;
+    const indexOfSpy = vi.spyOn(String.prototype, "indexOf").mockImplementation(function (this: string, searchString: string, position?: number) {
+      const found = originalIndexOf.call(this, searchString, position);
+      if (String(this) === sql && (searchString === "\r" || searchString === "\n") && found === -1) absentLineEndingSearches += 1;
+      return found;
+    });
+    let ranges: ReturnType<typeof splitSqlStatementRanges>;
+    try {
+      ranges = splitSqlStatementRanges(sql, databaseType);
+    } finally {
+      indexOfSpy.mockRestore();
+    }
+
+    expect(absentLineEndingSearches).toBeLessThanOrEqual(1);
+    expect(ranges.map((range) => range.sql)).toEqual(statements.map((statement) => statement.slice(0, -1)));
+    for (const range of ranges) expect(sql.slice(range.from, range.to)).toBe(range.sql);
+  });
+
+  it.each(["\n", "\r", "\r\n"])("preserves standalone slash, GO, and DELIMITER lines with %j", (lineEnding) => {
+    const slashSql = ["SELECT '/ literal';", "\t /  ", "SELECT 2;"].join(lineEnding);
+    expect(rangeSqlTexts(splitSqlStatementRanges(slashSql, "postgres"))).toEqual(["SELECT '/ literal'", "SELECT 2"]);
+    expect(rangeSqlTexts(splitSqlStatementRanges(slashSql, "oracle"))).toEqual(["SELECT '/ literal'", "SELECT 2"]);
+    expect(rangeSqlTexts(splitSqlStatementRanges(["SELECT 1", "  GO  ", "SELECT 2"].join(lineEnding), "sqlserver"))).toEqual(["SELECT 1", "SELECT 2"]);
+    expect(rangeSqlTexts(splitSqlStatementRanges(["DELIMITER //", "SELECT ';'//", "DELIMITER ;", "SELECT 2;"].join(lineEnding), "mysql"))).toEqual(["SELECT ';'", "SELECT 2"]);
+  });
+
+  it("preserves mixed line endings and a final delimiter without a newline", () => {
+    const sql = "SELECT 1;\r\n /\rSELECT 2;\n\t/\r\nSELECT 3;\r/";
+    expect(rangeSqlTexts(splitSqlStatementRanges(sql, "postgres"))).toEqual(["SELECT 1", "SELECT 2", "SELECT 3"]);
+  });
+
   it("splits multiple top-level statements", () => {
     const sql = "SELECT 1;\nSELECT 2;\nSELECT 3;";
     expect(rangeSqlTexts(splitSqlStatementRanges(sql))).toEqual(["SELECT 1", "SELECT 2", "SELECT 3"]);

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

@@ -385,14 +385,6 @@ export function splitSqlStatementRanges(sql: string, databaseType?: DatabaseType
   };
 
   const isWhitespace = (ch: string) => ch === " " || ch === "\t" || ch === "\r" || ch === "\n";
-  const nextLineBreak = (from: number) => {
-    const carriageReturn = sql.indexOf("\r", from);
-    const lineFeed = sql.indexOf("\n", from);
-    if (carriageReturn === -1) return lineFeed === -1 ? len : lineFeed;
-    if (lineFeed === -1) return carriageReturn;
-    return Math.min(carriageReturn, lineFeed);
-  };
-
   const markContent = (pos: number) => {
     if (statementStart === -1) {
       // TDSQL accepts arbitrary leading block directives on the SQL line. The
@@ -554,7 +546,7 @@ export function splitSqlStatementRanges(sql: string, databaseType?: DatabaseType
         if (databaseType === "mysql") {
           if (pendingMysqlDirectiveStart === -1 || i > pendingMysqlDirectiveLineEnd) {
             pendingMysqlDirectiveStart = i;
-            pendingMysqlDirectiveLineEnd = nextLineBreak(i);
+            pendingMysqlDirectiveLineEnd = findLineEnd(sql, i);
           }
         } else {
           pendingMysqlDirectiveStart = -1;
@@ -2506,11 +2498,10 @@ function mysqlDelimiterDirectiveCursorRange(sql: string, pos: number, databaseTy
 }
 
 function findLineEnd(sql: string, pos: number): number {
-  const newline = sql.indexOf("\n", pos);
+  for (let cursor = pos; cursor < sql.length; cursor += 1) {
-  const carriageReturn = sql.indexOf("\r", pos);
+    if (sql[cursor] === "\n" || sql[cursor] === "\r") return cursor;
-  if (newline === -1) return carriageReturn === -1 ? sql.length : carriageReturn;
+  }
-  if (carriageReturn === -1) return newline;
+  return sql.length;
-  return Math.min(newline, carriageReturn);
 }
 
 function nextLineStart(sql: string, lineEnd: number): number {

apps/desktop/src/stores/__tests__/queryStore.multiStatementError.spec.ts

@@ -883,6 +883,54 @@ describe("queryStore multi-statement errors", () => {
     expect(store.tabs.find((item) => item.id === tabId)?.results?.map((result) => result.sourceLabel)).toEqual(["Orders", "Users"]);
   });
 
+  it("does not repeatedly scan the full document when naming a selected large batch", async () => {
+    const statementCount = 500;
+    const statements = Array.from({ length: statementCount }, (_, index) => `-- Name: Result ${index}\nSELECT ${index};`);
+    const sql = statements.join("\n");
+    mocks.executeMulti.mockResolvedValue(Array.from({ length: statementCount }, (_, index) => ({ columns: ["value"], rows: [[index]], affected_rows: 0, execution_time_ms: 1, statement_index: index })));
+    const { useQueryStore } = await import("@/stores/queryStore");
+    const store = useQueryStore();
+    const tabId = store.createTab("mysql-1", "app", "Query");
+    store.updateSql(tabId, sql);
+    const originalFind = Array.prototype.find;
+    let statementVisits = 0;
+    const findSpy = vi.spyOn(Array.prototype, "find").mockImplementation(function (this: any[], predicate, thisArg) {
+      const isStatementList = this.length === statementCount && typeof this[0]?.hitFrom === "number";
+      return originalFind.call(this, (value, index, values) => {
+        if (isStatementList) statementVisits += 1;
+        return predicate.call(thisArg, value, index, values);
+      });
+    });
+    try {
+      await store.executeTabSql(tabId, sql, { sourceOffset: 0 });
+    } finally {
+      findSpy.mockRestore();
+    }
+
+    expect(statementVisits).toBeLessThan(statementCount * 20);
+    expect(store.tabs.find((tab) => tab.id === tabId)?.results?.map((result) => result.sourceLabel)).toEqual(statements.map((_, index) => `Result ${index}`));
+  });
+
+  it("preserves document names for out-of-order results and a selection starting inside a statement", async () => {
+    mocks.executeMulti.mockResolvedValue([
+      { columns: ["value"], rows: [[2]], affected_rows: 0, execution_time_ms: 1, statement_index: 1 },
+      { columns: ["value"], rows: [[1]], affected_rows: 0, execution_time_ms: 1, statement_index: 0 },
+    ]);
+    const { useQueryStore } = await import("@/stores/queryStore");
+    const store = useQueryStore();
+    const tabId = store.createTab("mysql-1", "app", "Query");
+    const sql = "-- Name: Users\nEXPLAIN SELECT * FROM users;\n-- Name: Orders\nSELECT * FROM orders;";
+    const sourceOffset = sql.indexOf("SELECT");
+    store.updateSql(tabId, sql);
+
+    await store.executeTabSql(tabId, sql.slice(sourceOffset), { sourceOffset });
+
+    expect(store.tabs.find((tab) => tab.id === tabId)?.results).toMatchObject([
+      { sourceLabel: "Orders", sourceStatement: "SELECT * FROM orders", sourceFrom: sql.lastIndexOf("SELECT") },
+      { sourceLabel: "Users", sourceStatement: "SELECT * FROM users", sourceFrom: sourceOffset },
+    ]);
+  });
+
   it("does not promote an unmarked Error alias without type metadata as a batch failure", async () => {
     mocks.executeMulti.mockResolvedValue([
       { columns: ["value"], rows: [[1]], affected_rows: 0, execution_time_ms: 1 },

apps/desktop/src/stores/queryStore.ts

@@ -369,6 +369,18 @@ function preservedResultIndex(results: QueryResult[], currentIndex: number | und
   return currentIndex;
 }
 
+function findSourceDocumentStatement(statements: ReturnType<typeof splitSqlStatementRanges>, sourceFrom: number) {
+  let lower = 0;
+  let upper = statements.length;
+  while (lower < upper) {
+    const middle = Math.floor((lower + upper) / 2);
+    if (statements[middle].to < sourceFrom) lower = middle + 1;
+    else upper = middle;
+  }
+  const statement = statements[lower];
+  return statement && statement.from <= sourceFrom ? statement : undefined;
+}
+
 function annotateQueryResultSources(results: QueryResult[], sql: string, database: string | undefined, databaseType?: DatabaseType, sourceOffset?: number, parameterOptions?: SqlParameterOptions, executedSql?: string, sourceDocumentSql?: string): { results: QueryResult[]; useDatabase?: string } {
   const statements = splitSqlStatementRanges(sql, databaseType, parameterOptions);
   // The backend positions errors against the SQL it actually received. When the
@@ -401,13 +413,7 @@ function annotateQueryResultSources(results: QueryResult[], sql: string, databas
         });
       }
     }
-    const documentStatement =
+    const documentStatement = sourceOffset === undefined ? undefined : findSourceDocumentStatement(documentStatements, sourceOffset + statement.from);
-      sourceOffset === undefined
-        ? undefined
-        : documentStatements.find((candidate) => {
-            const sourceFrom = sourceOffset + statement.from;
-            return candidate.from === sourceFrom || (sourceFrom >= candidate.from && sourceFrom <= candidate.to);
-          });
     const preamble = documentStatement ? sourceDocumentSql!.slice(documentStatement.hitFrom, documentStatement.from) : sql.slice(statement.hitFrom, statement.from);
     const customName = queryResultNameFromPreamble(preamble, { databaseType });
     if (customName) result.sourceLabel = customName;