diff --git a/apps/desktop/src/lib/__tests__/sql/sqlStatementRanges.spec.ts b/apps/desktop/src/lib/__tests__/sql/sqlStatementRanges.spec.ts index 80e27eeebf..e9cfdea9f6 100644 --- a/apps/desktop/src/lib/__tests__/sql/sqlStatementRanges.spec.ts +++ b/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; + 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"]); diff --git a/apps/desktop/src/lib/sql/sqlStatementRanges.ts b/apps/desktop/src/lib/sql/sqlStatementRanges.ts index 0074830147..59ef2c28e3 100644 --- a/apps/desktop/src/lib/sql/sqlStatementRanges.ts +++ b/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); - const carriageReturn = sql.indexOf("\r", pos); - if (newline === -1) return carriageReturn === -1 ? sql.length : carriageReturn; - if (carriageReturn === -1) return newline; - return Math.min(newline, carriageReturn); + for (let cursor = pos; cursor < sql.length; cursor += 1) { + if (sql[cursor] === "\n" || sql[cursor] === "\r") return cursor; + } + return sql.length; } function nextLineStart(sql: string, lineEnd: number): number { diff --git a/apps/desktop/src/stores/__tests__/queryStore.multiStatementError.spec.ts b/apps/desktop/src/stores/__tests__/queryStore.multiStatementError.spec.ts index 63eae634a7..d4093274bc 100644 --- a/apps/desktop/src/stores/__tests__/queryStore.multiStatementError.spec.ts +++ b/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 }, diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index cf2a0ac66d..8d07636562 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -369,6 +369,18 @@ function preservedResultIndex(results: QueryResult[], currentIndex: number | und return currentIndex; } +function findSourceDocumentStatement(statements: ReturnType, 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 = - sourceOffset === undefined - ? undefined - : documentStatements.find((candidate) => { - const sourceFrom = sourceOffset + statement.from; - return candidate.from === sourceFrom || (sourceFrom >= candidate.from && sourceFrom <= candidate.to); - }); + const documentStatement = sourceOffset === undefined ? undefined : findSourceDocumentStatement(documentStatements, sourceOffset + statement.from); 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;