fix: apply max_rows to comment-prefixed and CTE queries

[Elrendio]

Hey 👋 This looks like a bug but if it's wanted behavior, feel free to drop 😊

SQLRowLimiter.isSelectQuery decided whether to cap a statement with:

static isSelectQuery(sql: string): boolean {
  const trimmed = sql.trim().toLowerCase();
  return trimmed.startsWith('select');
}

so anything not opening with a bare SELECT got no LIMIT appended at all. Not a wrong limit —
no limit, and nothing anywhere says so. Verified by executing the shipped v1.2.0 module
(applyMaxRows(sql, 100)):

input	v1.2.0 output
-- reporting job\nSELECT * FROM t	returned unchanged
/* tag: report */ SELECT * FROM t	returned unchanged
WITH x AS (SELECT * FROM orders) SELECT * FROM x	returned unchanged
(SELECT id FROM a) UNION (SELECT id FROM b)	returned unchanged
SELECT * FROM (SELECT * FROM t LIMIT 5) s	returned unchanged
SELECT * FROM t	SELECT * FROM t\nLIMIT 100 ✅

This is not an exotic corner. Leading comments are how queries get attributed (-- app=…,
/* trace-id */), and a CTE is the ordinary shape of an analytical query — which is exactly the
kind of query a row cap exists to contain. In practice max_rows covered far less than it appeared
to. We hit this running DBHub against ~24 production PostgreSQL read replicas with agents driving
execute_sql.

The existing suite pinned the gap rather than catching it: three integration tests were named
should not apply maxRows to CTE queries (WITH clause). This PR inverts them.

To be clear about what those tests were pinning: git log -S shows the startsWith('select') guard
and those three tests both arrive in the same commit — 4fa886e feat: support max-row limit, the
commit that introduced max_rows. The CTE exclusion was never a separate, deliberate carve-out; it
is a consequence of that one-line classifier, and the tests recorded whatever it happened to do.
(One of the assertions is commented "not limited anymore", which reads as though a decision had been
reversed — the history shows there was no earlier behaviour to reverse.) If the exclusion is
wanted, it deserves to be an explicit rule rather than a side effect of the guard, and this PR
should be redirected to documenting it instead.

Fix

Classify on the comment/string-blanked text, and treat a leading WITH as row-returning. The SQL
sent to the server is never rewritten by the classification, so an attribution comment survives
verbatim — only the classifier sees the blanked form.

A data-modifying CTE (WITH x AS (DELETE … RETURNING *) SELECT …) is deliberately still not
limited: a LIMIT would cap the rows handed back while the write ran in full, which reads as a cap
that isn't one. That check reuses the read-only classifier's own keyword heuristic, extracted as
hasMutatingKeyword, so both layers agree on what counts as a write hidden inside a WITH. A false
positive there only means the statement keeps today's behaviour.

Enabling CTEs exposed a second bug, which is why the diff is larger than "add with to a list".
Every LIMIT/TOP helper matched the first clause found textually, which on a CTE is the inner one:

WITH recent AS (SELECT * FROM orders LIMIT 5)
SELECT * FROM recent JOIN big ON true

Naively enabling CTEs would have rewritten the CTE's own LIMIT 5 and left the statement — the part
that can return millions of rows — uncapped. So the helpers now scan with parenthesis-depth tracking
(the mechanism already in this file for hasSetOperator and the ORDER BY hoist) and reason only
about the statement's own clause; a nested LIMIT/TOP caps only its branch, so the statement still
gets a cap appended. This also fixes the same pre-existing hole for plain subqueries
(SELECT * FROM (SELECT * FROM t LIMIT 5) s, last row of the table above).

On SQL Server, TOP now lands on the statement's own SELECT rather than the first one textually
(for a CTE, the final SELECT), and a leading CTE is kept outside the derived table when a set
operation forces the #387 wrap — T-SQL has no SELECT … FROM (WITH …) AS subq form.

Explicitly preserved: SELECT … UNION ALL SELECT … starts with select, still gets a trailing
LIMIT, and on PostgreSQL that binds to the whole set operation. Pinned by a regression test.

Testing

• pnpm run test:unit — 981 passed (958 on main; +23 new).
• pnpm run test:integration — 334 passed (331 on main; +3) across PostgreSQL, MySQL, MariaDB,
  SQL Server and SQLite containers.
• pnpm run build:backend — clean.

New coverage: leading --//* *//mixed comments; CTE; CTE with an inner LIMIT; CTE with its own
LIMIT; CTE with a parameterized LIMIT; data-modifying CTE (all three of INSERT/UPDATE/DELETE);
parenthesised set operation; nested-subquery LIMIT; UNION ALL regression; and SQL Server CTE/TOP
cases.

Notes for reviewers

• npx tsc --noEmit reports 132 pre-existing errors on main; this branch reports the same 132,
  none in the changed files.
• Found alongside a separate PostgreSQL cancellation bug, sent as its own PR — the two are
  independent and share no files but one integration test file.

[Copilot]

Pull request overview

This PR fixes max_rows enforcement gaps by correctly classifying comment-prefixed and CTE (WITH) queries as row-returning, and by ensuring LIMIT/TOP detection targets the statement’s own top-level clauses (not nested CTE/subquery clauses). It also updates integration tests to reflect the corrected behavior across connectors, including SQL Server set-operation handling related to #387.

Changes:

• Improve row-returning classification to handle leading comments, leading parentheses, and CTEs while avoiding data-modifying CTEs.
• Rework LIMIT/TOP/ORDER BY detection to ignore nested clauses via top-level parenthesis-depth scanning, and adjust SQL Server wrapping behavior for set operations + CTE prefixes.
• Update/expand unit + integration tests to cover the corrected max_rows behavior for comments, CTEs, nesting, and SQL Server cases.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
src/utils/sql-row-limiter.ts	Refactors classification and top-level clause detection; updates LIMIT/TOP application logic including SQL Server set-op wrapping.
src/utils/allowed-keywords.ts	Extracts hasMutatingKeyword to share mutating detection logic across components.
src/utils/tests/sql-row-limiter.test.ts	Adds extensive unit coverage for comment-prefixed queries, CTEs, nesting, and SQL Server TOP behavior.
src/connectors/tests/sqlite.integration.test.ts	Updates SQLite integration test to expect maxRows applied to CTE queries.
src/connectors/tests/postgres.integration.test.ts	Updates Postgres integration test expectations and adds integration coverage for comment-prefixed and CTE-specific cases.
src/connectors/tests/mysql.integration.test.ts	Updates MySQL integration test to expect maxRows applied to CTE queries.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Elrendio]

Hey @tianzhou 👋

I saw you closed this PR, can I fix the issue raised and re-open or do you not want this?

Thanks 🙏