Skip to content

fix: apply max_rows to comment-prefixed and CTE queries - #400

Draft
Elrendio wants to merge 1 commit into
bytebase:mainfrom
Elrendio:fix/max-rows-comment-prefixed-and-cte-queries
Draft

fix: apply max_rows to comment-prefixed and CTE queries#400
Elrendio wants to merge 1 commit into
bytebase:mainfrom
Elrendio:fix/max-rows-comment-prefixed-and-cte-queries

Conversation

@Elrendio

Copy link
Copy Markdown
Contributor

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:unit981 passed (958 on main; +23 new).
  • pnpm run test:integration334 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.

`SQLRowLimiter.isSelectQuery` classified a statement with
`sql.trim().toLowerCase().startsWith("select")`, so `max_rows` was
silently a no-op for every query that does not open with a bare
`SELECT`. No LIMIT was appended at all — the cap simply did not exist
for these, and nothing surfaced that:

  -- reporting job\nSELECT * FROM t     -> returned unchanged
  /* tag */ SELECT * FROM t             -> returned unchanged
  WITH x AS (SELECT ...) SELECT * FROM x-> returned unchanged
  (SELECT a) UNION (SELECT b)           -> returned unchanged

Leading comments are common (query tags/attribution) and CTEs are the
normal shape of an analytical query, so in practice the row cap covered
much less than it appeared to.

Classify on the comment/string-blanked text instead, and accept `WITH`
as row-returning. The SQL handed to the server is never rewritten by the
classification, so an attribution comment survives verbatim. 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 keyword heuristic, now shared as
`hasMutatingKeyword`.

Enabling CTEs exposed a second problem: the LIMIT/TOP helpers matched
the first clause found textually, which on a CTE is the *inner* one.
`WITH x AS (SELECT ... LIMIT 5) SELECT * FROM x JOIN big ON true` would
have had the CTE's own cap rewritten while the statement stayed
uncapped. Every helper now scans with parenthesis-depth tracking (the
mechanism already used for `hasSetOperator` and the ORDER BY hoist) and
reasons 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`).

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 a wrap —
T-SQL has no `SELECT ... FROM (WITH ...) AS subq` form.

`SELECT ... UNION ALL SELECT ...` keeps its existing behaviour: it
starts with `select`, gets a trailing LIMIT, and on PostgreSQL that
binds to the whole set operation. Covered by a regression test.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Elrendio added a commit to Elrendio/dbhub that referenced this pull request Aug 10, 2026
for_stockly_main is the integration branch: upstream main plus one merge per open PR,
so the PR branches stay single-commit and reviewable on their own.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Elrendio added a commit to Elrendio/dbhub that referenced this pull request Aug 10, 2026
for_stockly_main is the integration branch: upstream main plus one merge per open PR,
so the PR branches stay single-commit and reviewable on their own.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant