fix: apply max_rows to comment-prefixed and CTE queries

src/connectors/__tests__/mysql.integration.test.ts

@@ -391,18 +391,18 @@ describe('MySQL Connector Integration Tests', () => {
       expect(result.resultSets[0].rows[0]).toHaveProperty('total');
     });
 
-    it('should not apply maxRows to CTE queries (WITH clause)', async () => {
-      // Test that maxRows is not applied to CTE queries (WITH clause)
+    it('should apply maxRows to CTE queries (WITH clause)', async () => {
+      // A CTE is the ordinary shape of an analytical query, so leaving it
+      // uncapped left max_rows silently inert for most real queries.
       try {
         const result = await mysqlTest.connector.executeSQL(`
           WITH user_summary AS (
             SELECT name, age FROM users WHERE age IS NOT NULL
           )
           SELECT * FROM user_summary ORDER BY age
         `, { maxRows: 2 });
-
-        // Should return all rows since WITH queries are not limited
-        expect(result.resultSets[0].rows.length).toBeGreaterThan(2);
+
+        expect(result.resultSets[0].rows).toHaveLength(2);
         expect(result.resultSets[0].rows[0]).toHaveProperty('name');
         expect(result.resultSets[0].rows[0]).toHaveProperty('age');
       } catch (error) {

src/connectors/__tests__/postgres.integration.test.ts

@@ -460,21 +460,60 @@ describe('PostgreSQL Connector Integration Tests', () => {
       expect(result.resultSets[0].rows[0]).toHaveProperty('total');
     });
 
-    it('should not apply maxRows to CTE queries (WITH clause)', async () => {
-      // Test that maxRows is not applied to CTE queries (WITH clause)
+    it('should apply maxRows to CTE queries (WITH clause)', async () => {
+      // A CTE is the ordinary shape of an analytical query, so leaving it
+      // uncapped left max_rows silently inert for most real queries.
       const result = await postgresTest.connector.executeSQL(`
         WITH user_summary AS (
           SELECT name, age FROM users WHERE age IS NOT NULL
         )
         SELECT * FROM user_summary ORDER BY age
       `, { maxRows: 2 });
-
-      // Should return all rows since WITH queries are not limited anymore
-      expect(result.resultSets[0].rows.length).toBeGreaterThan(2);
+
+      expect(result.resultSets[0].rows).toHaveLength(2);
       expect(result.resultSets[0].rows[0]).toHaveProperty('name');
       expect(result.resultSets[0].rows[0]).toHaveProperty('age');
     });
 
+    it('should apply maxRows to a query introduced by a comment', async () => {
+      const result = await postgresTest.connector.executeSQL(
+        '-- dbhub attribution tag\nSELECT name FROM users ORDER BY name',
+        { maxRows: 2 }
+      );
+
+      expect(result.resultSets[0].rows).toHaveLength(2);
+    });
+
+    it('should cap the statement itself rather than tightening a CTE\'s own LIMIT', async () => {
+      // The inner LIMIT caps only the CTE; the statement can still return more
+      // rows than that, so it needs a cap of its own.
+      const result = await postgresTest.connector.executeSQL(`
+        WITH first_three AS (
+          SELECT name FROM users ORDER BY name LIMIT 3
+        )
+        SELECT * FROM first_three
+      `, { maxRows: 2 });
+
+      expect(result.resultSets[0].rows).toHaveLength(2);
+    });
+
+    it('should not apply maxRows to a data-modifying CTE', async () => {
+      const result = await postgresTest.connector.executeSQL(`
+        WITH inserted AS (
+          INSERT INTO users (name, email, age)
+          VALUES ('dm1', 'dm1@dm.com', 41), ('dm2', 'dm2@dm.com', 42), ('dm3', 'dm3@dm.com', 43)
+          RETURNING id, name
+        )
+        SELECT * FROM inserted
+      `, { maxRows: 2 });
+
+      // A LIMIT here would cap the rows handed back while all three rows were
+      // still written - a cap that isn't one.
+      expect(result.resultSets[0].rows).toHaveLength(3);
+
+      await postgresTest.connector.executeSQL("DELETE FROM users WHERE email LIKE '%@dm.com'", {});
+    });
+
     it('should handle maxRows in multi-statement execution with transactions', async () => {
       // Test maxRows with multiple statements where some are SELECT
       const result = await postgresTest.connector.executeSQL(`

src/connectors/__tests__/sqlite.integration.test.ts

@@ -379,17 +379,17 @@ describe('SQLite Connector Integration Tests', () => {
       expect(result.resultSets[0].rows[0]).toHaveProperty('total');
     });
 
-    it('should not apply maxRows to CTE queries (WITH clause)', async () => {
-      // Test that maxRows is not applied to CTE queries (WITH clause)
+    it('should apply maxRows to CTE queries (WITH clause)', async () => {
+      // A CTE is the ordinary shape of an analytical query, so leaving it
+      // uncapped left max_rows silently inert for most real queries.
       const result = await sqliteTest.connector.executeSQL(`
         WITH user_summary AS (
           SELECT name, age FROM users WHERE age IS NOT NULL
         )
         SELECT * FROM user_summary ORDER BY age
       `, { maxRows: 2 });
-
-      // Should return all rows since WITH queries are not limited anymore
-      expect(result.resultSets[0].rows.length).toBeGreaterThan(2);
+
+      expect(result.resultSets[0].rows).toHaveLength(2);
       expect(result.resultSets[0].rows[0]).toHaveProperty('name');
       expect(result.resultSets[0].rows[0]).toHaveProperty('age');
     });

src/utils/__tests__/sql-row-limiter.test.ts

@@ -155,7 +155,168 @@ describe("SQLRowLimiter", () => {
     });
   });
 
+  describe("applyMaxRows - leading comments", () => {
+    // A query introduced by a comment (an attribution tag, say) is still a
+    // SELECT. Classifying on the raw text made `max_rows` silently inert for
+    // every one of them.
+    it("caps a query introduced by a -- line comment", () => {
+      const sql = "-- dbhub agent query\nSELECT * FROM users";
+      expect(SQLRowLimiter.applyMaxRows(sql, 100)).toBe(
+        "-- dbhub agent query\nSELECT * FROM users\nLIMIT 100"
+      );
+    });
+
+    it("caps a query introduced by a block comment", () => {
+      const sql = "/* tag: report */ SELECT * FROM users";
+      expect(SQLRowLimiter.applyMaxRows(sql, 100)).toBe(
+        "/* tag: report */ SELECT * FROM users\nLIMIT 100"
+      );
+    });
+
+    it("caps a query introduced by several mixed leading comments", () => {
+      const sql = "-- one\n/* two */\n-- three\nSELECT * FROM users";
+      expect(SQLRowLimiter.applyMaxRows(sql, 100)).toBe(
+        "-- one\n/* two */\n-- three\nSELECT * FROM users\nLIMIT 100"
+      );
+    });
+
+    it("leaves the caller's comment text untouched", () => {
+      // The comment is load-bearing for query attribution, so only the
+      // classification sees the stripped form - the SQL sent to the server
+      // keeps it verbatim.
+      const sql = "/* app=dbhub; user='bob' */\nSELECT * FROM users";
+      expect(SQLRowLimiter.applyMaxRows(sql, 100)).toContain("/* app=dbhub; user='bob' */");
+    });
+
+    it("still leaves a non-SELECT hidden behind a comment alone", () => {
+      const sql = "-- looks harmless\nUPDATE users SET active = true";
+      expect(SQLRowLimiter.applyMaxRows(sql, 100)).toBe(sql);
+    });
+  });
+
+  describe("applyMaxRows - CTEs", () => {
+    it("caps a WITH ... SELECT query", () => {
+      const sql = "WITH recent AS (SELECT * FROM orders) SELECT * FROM recent";
+      expect(SQLRowLimiter.applyMaxRows(sql, 100)).toBe(
+        "WITH recent AS (SELECT * FROM orders) SELECT * FROM recent\nLIMIT 100"
+      );
+    });
+
+    it("appends its own LIMIT instead of tightening a CTE's inner LIMIT", () => {
+      // The CTE's LIMIT caps only the CTE; the statement can still return far
+      // more rows than that (here via the join), so it needs a cap of its own.
+      const sql =
+        "WITH recent AS (SELECT * FROM orders LIMIT 5) SELECT * FROM recent JOIN big ON true";
+      expect(SQLRowLimiter.applyMaxRows(sql, 100)).toBe(
+        "WITH recent AS (SELECT * FROM orders LIMIT 5) SELECT * FROM recent JOIN big ON true\nLIMIT 100"
+      );
+    });
+
+    it("tightens the statement's own LIMIT on a CTE query", () => {
+      const sql = "WITH recent AS (SELECT * FROM orders LIMIT 5) SELECT * FROM recent LIMIT 500";
+      expect(SQLRowLimiter.applyMaxRows(sql, 100)).toBe(
+        "WITH recent AS (SELECT * FROM orders LIMIT 5) SELECT * FROM recent LIMIT 100"
+      );
+    });
+
+    it("wraps a CTE query whose own LIMIT is parameterized", () => {
+      const sql = "WITH recent AS (SELECT * FROM orders) SELECT * FROM recent LIMIT $1";
+      expect(SQLRowLimiter.applyMaxRows(sql, 100)).toBe(
+        "SELECT * FROM (WITH recent AS (SELECT * FROM orders) SELECT * FROM recent LIMIT $1\n) AS subq LIMIT 100"
+      );
+    });
+
+    it.each([
+      ["DELETE", "WITH d AS (DELETE FROM t RETURNING *) SELECT * FROM d"],
+      ["INSERT", "WITH i AS (INSERT INTO t SELECT * FROM s RETURNING *) SELECT * FROM i"],
+      ["UPDATE", "WITH u AS (UPDATE t SET a = 1 RETURNING *) SELECT * FROM u"],
+    ])("does not cap a data-modifying CTE (%s)", (_label, sql) => {
+      // A LIMIT here would cap the rows handed back while the write still runs
+      // in full - a cap that isn't one.
+      expect(SQLRowLimiter.applyMaxRows(sql, 100)).toBe(sql);
+    });
+  });
+
+  describe("applyMaxRows - set operations and nesting", () => {
+    it("caps a parenthesised set operation", () => {
+      const sql = "(SELECT id FROM a) UNION (SELECT id FROM b)";
+      expect(SQLRowLimiter.applyMaxRows(sql, 100)).toBe(
+        "(SELECT id FROM a) UNION (SELECT id FROM b)\nLIMIT 100"
+      );
+    });
+
+    it("keeps appending a trailing LIMIT to a bare UNION ALL, which binds to the whole set operation", () => {
+      const sql = "SELECT id FROM a UNION ALL SELECT id FROM b";
+      expect(SQLRowLimiter.applyMaxRows(sql, 100)).toBe(
+        "SELECT id FROM a UNION ALL SELECT id FROM b\nLIMIT 100"
+      );
+    });
+
+    it("appends its own LIMIT instead of tightening a subquery's LIMIT", () => {
+      const sql = "SELECT * FROM (SELECT * FROM t LIMIT 5) s";
+      expect(SQLRowLimiter.applyMaxRows(sql, 100)).toBe(
+        "SELECT * FROM (SELECT * FROM t LIMIT 5) s\nLIMIT 100"
+      );
+    });
+  });
+
+  describe("clause detection ignores nested clauses", () => {
+    it("does not report a subquery's LIMIT as the statement's own", () => {
+      const sql = "SELECT * FROM (SELECT * FROM t LIMIT 5) s";
+      expect(SQLRowLimiter.hasLimitClause(sql)).toBe(false);
+      expect(SQLRowLimiter.extractLimitValue(sql)).toBe(null);
+    });
+
+    it("does not report a CTE's parameterized LIMIT as the statement's own", () => {
+      const sql = "WITH x AS (SELECT * FROM t LIMIT $1) SELECT * FROM x";
+      expect(SQLRowLimiter.hasParameterizedLimit(sql)).toBe(false);
+    });
+
+    it("does not report a CTE's TOP as the statement's own", () => {
+      const sql = "WITH x AS (SELECT TOP 5 id FROM t) SELECT * FROM x";
+      expect(SQLRowLimiter.hasTopClause(sql)).toBe(false);
+      expect(SQLRowLimiter.extractTopValue(sql)).toBe(null);
+    });
+  });
+
   describe("applyMaxRowsForSQLServer", () => {
+    it("caps a query introduced by a leading comment", () => {
+      const sql = "-- tag\nSELECT * FROM users";
+      expect(SQLRowLimiter.applyMaxRowsForSQLServer(sql, 100)).toBe(
+        "-- tag\nSELECT TOP 100 * FROM users"
+      );
+    });
+
+    it("puts TOP on the statement's own SELECT, not on the CTE's", () => {
+      const sql = "WITH x AS (SELECT id FROM t) SELECT * FROM x";
+      expect(SQLRowLimiter.applyMaxRowsForSQLServer(sql, 100)).toBe(
+        "WITH x AS (SELECT id FROM t) SELECT TOP 100 * FROM x"
+      );
+    });
+
+    it("leaves a CTE's own TOP alone and caps the final SELECT", () => {
+      const sql = "WITH x AS (SELECT TOP 5 id FROM t) SELECT * FROM x";
+      expect(SQLRowLimiter.applyMaxRowsForSQLServer(sql, 100)).toBe(
+        "WITH x AS (SELECT TOP 5 id FROM t) SELECT TOP 100 * FROM x"
+      );
+    });
+
+    it("keeps a leading CTE outside the wrapped subquery for a set operation", () => {
+      // T-SQL has no `SELECT ... FROM (WITH ...) AS subq` form, but a CTE
+      // declared before the SELECT is in scope inside the derived table.
+      const sql = "WITH x AS (SELECT id FROM t) SELECT id FROM x UNION ALL SELECT id FROM y";
+      expect(SQLRowLimiter.applyMaxRowsForSQLServer(sql, 100)).toBe(
+        "WITH x AS (SELECT id FROM t) SELECT TOP 100 * FROM (SELECT id FROM x UNION ALL SELECT id FROM y\n) AS subq"
+      );
+    });
+
+    it("does not cap a data-modifying CTE", () => {
+      const sql = "WITH d AS (DELETE FROM t OUTPUT deleted.*) SELECT * FROM d";
+      expect(SQLRowLimiter.applyMaxRowsForSQLServer(sql, 100)).toBe(sql);
+    });
+  });
+
+  describe("applyMaxRowsForSQLServer - pre-existing behaviour", () => {
     it("should not modify SQL when maxRows is undefined", () => {
       const sql = "SELECT * FROM users";
       expect(SQLRowLimiter.applyMaxRowsForSQLServer(sql, undefined)).toBe(sql);

src/utils/allowed-keywords.ts

@@ -178,6 +178,19 @@ const mutatingPatterns: Record<ConnectorType, RegExp> = {
   sqlserver: mutatingPatternSqlServer,
 };
 
+/**
+ * Whether (comment/string-stripped) SQL contains a data-modifying keyword.
+ * Shared by the read-only classifier's CTE check and the row limiter, so both
+ * agree on what counts as a write hidden inside a WITH. Falls back to the
+ * dialect-independent keyword set when no connector type is given.
+ */
+export function hasMutatingKeyword(
+  strippedSQL: string,
+  connectorType?: ConnectorType | string
+): boolean {
+  return (mutatingPatterns[connectorType as ConnectorType] ?? mutatingPattern).test(strippedSQL);
+}
+
 const selectIntoPattern = /\bselect\b[\s\S]+\binto\b/i;
 
 /**
@@ -251,11 +264,8 @@ function checkReadOnly(cleanedSQL: string, connectorType: ConnectorType | string
   }
 
   // WITH statements can embed DML in CTEs (e.g. WITH cte AS (UPDATE ...))
-  if (firstWord === "with") {
-    const pattern = mutatingPatterns[connectorType as ConnectorType] ?? mutatingPattern;
-    if (pattern.test(cleanedSQL)) {
-      return false;
-    }
+  if (firstWord === "with" && hasMutatingKeyword(cleanedSQL, connectorType)) {
+    return false;
   }
 
   // SQLite PRAGMA: a pragma that sets a value mutates durable or session state and