diff --git a/src/connectors/__tests__/sqlite.integration.test.ts b/src/connectors/__tests__/sqlite.integration.test.ts index a61ec1d4..1c760bfd 100644 --- a/src/connectors/__tests__/sqlite.integration.test.ts +++ b/src/connectors/__tests__/sqlite.integration.test.ts @@ -348,6 +348,20 @@ describe('SQLite Connector Integration Tests', () => { expect(result.resultSets[0].truncated).toBe(true); }); + it('should cap the row count of LIMIT offset, count without moving the offset', async () => { + const numbers = + 'SELECT n FROM (WITH RECURSIVE c(n) AS (SELECT 1 UNION ALL SELECT n + 1 FROM c WHERE n < 30) SELECT n FROM c)'; + + const capped = await sqliteTest.connector.executeSQL(`${numbers} ORDER BY n LIMIT 10, 5`, { maxRows: 3 }); + expect(capped.resultSets[0].rows.map(row => Number(row.n))).toEqual([11, 12, 13]); + expect(capped.resultSets[0].truncated).toBe(true); + + // A small offset must not be mistaken for a row count within the cap + const smallOffset = await sqliteTest.connector.executeSQL(`${numbers} ORDER BY n LIMIT 2, 10`, { maxRows: 5 }); + expect(smallOffset.resultSets[0].rows.map(row => Number(row.n))).toEqual([3, 4, 5, 6, 7]); + expect(smallOffset.resultSets[0].truncated).toBe(true); + }); + it('should flag truncated when maxRows cuts off rows', async () => { // users has 3+ rows; the cap of 2 provably cuts rows off const result = await sqliteTest.connector.executeSQL( diff --git a/src/utils/__tests__/sql-row-limiter.test.ts b/src/utils/__tests__/sql-row-limiter.test.ts index b86f8202..c9be8724 100644 --- a/src/utils/__tests__/sql-row-limiter.test.ts +++ b/src/utils/__tests__/sql-row-limiter.test.ts @@ -115,6 +115,24 @@ describe("SQLRowLimiter", () => { expect(result).toBe("SELECT * FROM users LIMIT 100"); }); + it("tightens the row count of LIMIT offset, count and keeps the offset", () => { + expect(SQLRowLimiter.applyMaxRows("SELECT * FROM users LIMIT 10, 200", 100, "mysql")).toBe( + "SELECT * FROM users LIMIT 10, 100" + ); + expect(SQLRowLimiter.applyMaxRows("SELECT * FROM users LIMIT 500, 20", 100, "sqlite")).toBe( + "SELECT * FROM users LIMIT 500, 20" + ); + expect(SQLRowLimiter.applyMaxRows("SELECT * FROM users LIMIT ?, 200", 100, "mysql")).toBe( + "SELECT * FROM users LIMIT ?, 100" + ); + }); + + it("wraps LIMIT offset, count when the row count is a parameter", () => { + expect(SQLRowLimiter.applyMaxRows("SELECT * FROM users LIMIT 10, ?", 100, "mysql")).toBe( + "SELECT * FROM (SELECT * FROM users LIMIT 10, ?\n) AS subq LIMIT 100" + ); + }); + it("should handle complex query with parameterized LIMIT", () => { const sql = "SELECT emp_no, first_name, last_name, hire_date FROM employee WHERE first_name ILIKE '%' || $1 || '%' OR last_name ILIKE '%' || $1 || '%' LIMIT $2"; const result = SQLRowLimiter.applyMaxRows(sql, 1000); @@ -431,6 +449,18 @@ describe("SQLRowLimiter", () => { }); }); + it("judges LIMIT offset, count by its row count, not its offset", () => { + // The offset 2 is within the cap, but the query asks for 200 rows + expect(SQLRowLimiter.applyMaxRowsWithTruncationProbe("SELECT * FROM t LIMIT 2, 200", 100, "mysql")).toEqual({ + sql: "SELECT * FROM t LIMIT 2, 101", + probeApplied: true, + }); + expect(SQLRowLimiter.applyMaxRowsWithTruncationProbe("SELECT * FROM t LIMIT 500, 20", 100, "mysql")).toEqual({ + sql: "SELECT * FROM t LIMIT 500, 20", + probeApplied: false, + }); + }); + it("does not probe a data-modifying CTE", () => { const sql = "WITH d AS (DELETE FROM t RETURNING *) SELECT * FROM d"; expect(SQLRowLimiter.applyMaxRowsWithTruncationProbe(sql, 100)).toEqual({ diff --git a/src/utils/sql-row-limiter.ts b/src/utils/sql-row-limiter.ts index 66d8cc17..c474ceed 100644 --- a/src/utils/sql-row-limiter.ts +++ b/src/utils/sql-row-limiter.ts @@ -22,6 +22,8 @@ interface TopLevelClause { length: number; /** Null when the clause holds a parameter placeholder instead of a literal. */ value: number | null; + /** The offset of a MySQL/SQLite `LIMIT offset, count` clause, as written. */ + offset?: string; } /** @@ -119,7 +121,8 @@ export class SQLRowLimiter { // LIMIT found textually, which on a CTE would rewrite the CTE's cap and // leave the statement itself uncapped. const effectiveLimit = Math.min(limit.value, maxRows); - return `${sql.slice(0, limit.index)}LIMIT ${effectiveLimit}${sql.slice(limit.index + limit.length)}`; + const offset = limit.offset !== undefined ? `${limit.offset}, ` : ""; + return `${sql.slice(0, limit.index)}LIMIT ${offset}${effectiveLimit}${sql.slice(limit.index + limit.length)}`; } // Add LIMIT clause to the end of the query @@ -173,21 +176,27 @@ export class SQLRowLimiter { return found; } - /** The statement's own LIMIT clause, literal or parameterized ($1, ?, @p1). */ + /** + * The statement's own LIMIT clause, literal or parameterized ($1, ?, @p1). + * In the MySQL/SQLite `LIMIT offset, count` form the row count is the + * second operand, so that is the value reported and tightened. + */ private static findTopLevelLimit(sql: string, dialect?: ConnectorType): TopLevelClause | null { const match = this.findTopLevelMatch( sql, - /\(|\)|\blimit\s+(?:(\d+)|\$\d+|\?|@p\d+)/gi, + /\(|\)|\blimit\s+(\d+|\$\d+|\?|@p\d+)(?:\s*,\s*(\d+|\$\d+|\?|@p\d+))?/gi, "last", dialect ); if (match === null) { return null; } + const count = match[2] ?? match[1]; return { index: match.index, length: match[0].length, - value: match[1] !== undefined ? parseInt(match[1], 10) : null, + value: /^\d+$/.test(count) ? parseInt(count, 10) : null, + offset: match[2] !== undefined ? match[1] : undefined, }; }