fix(editor): close SQL guard bypasses in the query editor (#63)

A CHAT_EDITOR (non-admin) user could delete or overwrite every row in
any table by wrapping the write in a CTE. `WITH x AS (DELETE FROM t
RETURNING *) SELECT * FROM x` parses as a Select, so classification
returned read-only and returned before the admin check ever ran.
PostgreSQL executes data-modifying CTEs for real. Verified end to end: 3
rows -> 0, success:true, no confirmation prompt, logged as an ordinary
EDITOR_QUERY_EXECUTED / SUCCESS.

- classifyStatement now walks the parse tree for data-modifying CTEs and
SELECT ... INTO (detectSelectWrite), with a text backstop
(detectHiddenWrite) so an unparseable variant fails closed instead of
reaching isReadOnlyQuery, which reports anything starting with WITH as
safe.
- READ_ONLY_ONLY contexts now open read-only JDBC sessions, so the
database refuses the write even when classification is wrong.
Classification is a parser heuristic; this is what keeps the next parser
gap from being data loss. HikariCP resets the flag on return to the
pool, verified, so it cannot leak into an admin's later write.
- Row caps are enforced with setMaxRows instead of a `\blimit\s+\d+`
text match that hit inside comments, string literals and subqueries. An
inner LIMIT returned 200k rows against a 1,000 cap, into an unbounded
ArrayList and an unvirtualized table.
- Cancel terminates the query instead of only aborting the HTTP request,
which left the statement holding one of the pool's 10 connections. The
client sends an executionId, RunningQueryRegistry maps it to the backend
session pid, and the new cancel endpoint kills exactly that session,
scoped to the connection and the user who started it. The previous UI
behavior killed *every* active query on the connection, including other
users' work.
- Fix pg_terminate_backend binding: setLong sent bigint, so PostgreSQL
found no matching overload and every kill failed, including the Active
Queries screen's own button.
- Editor timeout 600s -> 240s, under nginx's 300s proxy_read_timeout, so
a slow query reports a real error rather than an opaque 504 while still
running.
- Rate-limit /api/connections/*/query (30r/m + burst 20, 429 on reject).

QueryExecutionPolicyServiceTest stubbed isReadOnlyQuery to always return
false — the opposite of what the shipped providers do for WITH — so it
asserted behavior no deployment had, and withInsert_isTreatedAsMutation
passed *because* of the stub. It now uses a real
MySQLQueryExecutionProvider, plus 12 regression tests covering each
bypass and the reads that must keep working.

---------

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Co-authored-by: Krishna Sasank Talasila <sasanktk@gmail.com>

Author: 3 people

CLAUDE.md

@@ -267,6 +267,60 @@ broken. Assert the *outcome*, never the attempt:
   schemas have a `comment` table. Assert `SELECT * FROM comment` is allowed *and*
   that `WITH x AS (DELETE …) SELECT …` / `WITH x AS (…) DELETE …` still are not.
 
+### SQL Editor Guard Rules
+
+The Editor (`EditorSection` → `SqlRunnerTab` → `POST /connections/{id}/query` →
+`QueryExecutionPolicyService` → `QueryExecutorService`) is the only surface where a
+user submits arbitrary SQL. Everything below was a live bug, verified by executing
+it against a real database — not a theoretical hardening pass.
+
+- **Never classify SQL by its leading keyword alone.** `WITH x AS (DELETE FROM t
+  RETURNING *) SELECT * FROM x` parses as a `Select` and *every* leading-keyword
+  check calls it read-only — including `isReadOnlyQuery`, which reports anything
+  starting with `WITH` as safe. PostgreSQL executes data-modifying CTEs for real,
+  so a **non-admin** wiped whole tables through the Editor with no confirmation
+  prompt, logged as an ordinary `EDITOR_QUERY_EXECUTED / SUCCESS`. `SELECT … INTO
+  newtab` is the same class of bug (it is DDL). `classifyStatement` now inspects
+  the parse tree (`detectSelectWrite`) **and** runs a text backstop
+  (`detectHiddenWrite`) so an unparseable variant fails closed instead of falling
+  through to the keyword path.
+- **Read-only contexts open read-only JDBC sessions.** `QueryExecutorService` calls
+  `connection.setReadOnly(true)` whenever `mutationMode() == READ_ONLY_ONLY`, so
+  PostgreSQL refuses the write itself even if classification is wrong. Classification
+  is a parser heuristic; this is what keeps the *next* parser gap from being data
+  loss. A driver that rejects the hint raises rather than silently continuing
+  writable. HikariCP resets the flag on return to the pool (verified), so it cannot
+  leak into an admin's later write.
+- **Row caps are enforced with `setMaxRows`, not by appending `LIMIT n`.** The old
+  check skipped its own LIMIT whenever the regex `\blimit\s+\d+` matched anywhere —
+  including inside a comment, a string literal, or a subquery. `WITH a AS (SELECT …
+  LIMIT 100) SELECT * FROM a` is ordinary analyst SQL and returned **200k rows**
+  against a 1,000 cap, straight into an unbounded `ArrayList` and then an
+  unvirtualized table. The SQL `LIMIT` is still appended for simple SELECTs, but
+  only as an optimization — correctness no longer depends on that text match.
+- **Cancel must terminate the query, not just the HTTP request.** `abortController
+  .abort()` only closes the socket; the statement runs on holding one of the pool's
+  10 connections for up to its timeout. The client now sends an `executionId`,
+  `RunningQueryRegistry` maps it to the backend session pid (via the dialect's
+  `getSessionPidQuery()`), and `POST /connections/{id}/query/{executionId}/cancel`
+  terminates exactly that session. The previous UI behavior was worse than nothing:
+  it killed **every** active query on the connection, including other users' work.
+  The cancel endpoint is scoped to the connection *and* the user who started the
+  run, so an execution id is not a kill primitive for someone else's query.
+- **Keep the client timeout under the proxy's.** `docker/nginx/default.conf` gives
+  up at `proxy_read_timeout 300s`; the Editor used to ask for 600s, so a 6-minute
+  query returned an opaque 504 while still running. `QUERY_TIMEOUT_SECONDS = 240`
+  in `SqlRunnerTab.js` — change both together or not at all.
+- **`/api/connections/*/query` is rate-limited in nginx** (`limit_req zone=sqlexec`,
+  30r/m + burst 20, `429` on reject). It is the most expensive authenticated call
+  in the product.
+- **Test the policy against the real providers.** `QueryExecutionPolicyServiceTest`
+  used to stub `isReadOnlyQuery` to always return `false` — the exact opposite of
+  what the shipped providers do for `WITH`. It asserted behavior no deployment had,
+  and `withInsert_isTreatedAsMutation` passed *because* of the stub. It now
+  constructs a real `MySQLQueryExecutionProvider`. Do not reintroduce a stubbed
+  dialect here; the mock is what let the blocker ship.
+
 ### Data Model Rules
 
 - **`mcp_tokens.user_id` is a non-null FK with no cascade.** Deleting a user who holds

backend/src/main/java/com/dbaagent/controller/SchemaController.java

@@ -6,7 +6,9 @@
 import com.dbaagent.service.CredentialService;
 import com.dbaagent.service.QueryExecutionContext;
 import com.dbaagent.service.QueryExecutionPolicyException;
+import com.dbaagent.service.ActiveQueryService;
 import com.dbaagent.service.QueryExecutorService;
+import com.dbaagent.service.RunningQueryRegistry;
 import com.dbaagent.service.SqlExecutionAuditService;
 import com.dbaagent.service.UserDataAccessPolicyException;
 import com.dbaagent.service.SchemaScannerService;
@@ -36,6 +38,8 @@ public class SchemaController {
     private final QueryExecutorService queryExecutorService;
     private final AccessControlService accessControlService;
     private final SqlExecutionAuditService sqlExecutionAuditService;
+    private final RunningQueryRegistry runningQueryRegistry;
+    private final ActiveQueryService activeQueryService;
 
     @PostMapping("/scan")
     public ResponseEntity<Map<String, Object>> scanSchema(@PathVariable String connectionId) {
@@ -260,6 +264,63 @@ public ResponseEntity<Map<String, Object>> executeQuery(
         }
     }
 
+    /**
+     * Terminates a query this caller started but abandoned. Aborting the HTTP
+     * request only closes the socket — the statement keeps running and holds a
+     * pooled connection until it finishes, so the client must ask for it to stop.
+     */
+    @PostMapping("/query/{executionId}/cancel")
+    public ResponseEntity<Map<String, Object>> cancelQuery(
+            @PathVariable String connectionId,
+            @PathVariable String executionId) {
+        Map<String, Object> response = new HashMap<>();
+        try {
+            if (!credentialService.connectionExists(connectionId)) {
+                response.put("success", false);
+                response.put("message", "Connection not found");
+                return ResponseEntity.status(HttpStatus.NOT_FOUND).body(response);
+            }
+            accessControlService.assertCanUseChatEditor(connectionId);
+
+            var running = runningQueryRegistry.find(executionId);
+            if (running.isEmpty()) {
+                // Already finished, or never started. Nothing to cancel.
+                response.put("success", true);
+                response.put("cancelled", false);
+                response.put("message", "Query is no longer running");
+                return ResponseEntity.ok(response);
+            }
+
+            RunningQueryRegistry.RunningQuery target = running.get();
+            // The execution id is a bearer token for a kill: scope it to this
+            // connection and to the user who started it, so one caller cannot
+            // terminate another's query by guessing an id.
+            String currentUser = accessControlService.getCurrentUsername();
+            if (!connectionId.equals(target.connectionId())
+                || (target.username() != null && currentUser != null && !target.username().equals(currentUser))) {
+                response.put("success", false);
+                response.put("message", "Query not found for this connection");
+                return ResponseEntity.status(HttpStatus.NOT_FOUND).body(response);
+            }
+
+            activeQueryService.killQuery(connectionId, target.sessionPid());
+            runningQueryRegistry.unregister(executionId);
+            response.put("success", true);
+            response.put("cancelled", true);
+            response.put("message", "Query cancelled");
+            return ResponseEntity.ok(response);
+        } catch (ResponseStatusException e) {
+            response.put("success", false);
+            response.put("message", e.getReason());
+            return ResponseEntity.status(e.getStatusCode()).body(response);
+        } catch (Exception e) {
+            log.warn("Failed to cancel query {} on connection {}: {}", executionId, connectionId, e.getMessage());
+            response.put("success", false);
+            response.put("message", "Failed to cancel query: " + e.getMessage());
+            return ResponseEntity.status(HttpStatus.INTERNAL_SERVER_ERROR).body(response);
+        }
+    }
+
     // `{tableName:.+}` keeps schema-qualified ids (`crm.orders`) as one segment.
     @GetMapping("/tables/{tableName:.+}/indexes")
     public ResponseEntity<Map<String, Object>> getTableIndexes(

backend/src/main/java/com/dbaagent/model/QueryRequest.java

@@ -11,6 +11,9 @@ public class QueryRequest {
     private Integer timeoutSeconds; // Optional per-query timeout override (null = use server default)
     private QueryExecutionOrigin executionOrigin = QueryExecutionOrigin.INTERNAL;
     private Boolean mutationConfirmed = Boolean.FALSE;
+    // Client-generated id for this run, used to cancel the query if the caller
+    // gives up before it finishes.
+    private String executionId;
 
     public QueryRequest(String query, Integer limit, Integer timeoutSeconds) {
         this.query = query;

backend/src/main/java/com/dbaagent/model/QueryResult.java

@@ -17,4 +17,23 @@ public class QueryResult {
     private Boolean isLimited;   // True when a row limit was applied to this query
     private Long executionTimeMs;
     private String query;
+    // Server-side session id this query ran on, so a client that gives up can ask
+    // the backend to terminate it instead of leaving it holding a connection.
+    private String sessionPid;
+
+    /**
+     * Kept so adding {@code sessionPid} did not break every positional caller.
+     * Prefer the setter for new code — this class is a response DTO that grows.
+     */
+    public QueryResult(
+        List<String> columns,
+        List<List<Object>> rows,
+        Integer rowCount,
+        Long totalRowCount,
+        Boolean isLimited,
+        Long executionTimeMs,
+        String query
+    ) {
+        this(columns, rows, rowCount, totalRowCount, isLimited, executionTimeMs, query, null);
+    }
 }

backend/src/main/java/com/dbaagent/provider/api/QueryExecutionProvider.java

@@ -51,4 +51,16 @@ public interface QueryExecutionProvider {
      * @return The query signature
      */
     String generateQuerySignature(String query);
+
+    /**
+     * SQL returning the server-side session id of the current connection, in the
+     * form accepted by this dialect's kill statement. Callers use it to cancel a
+     * still-running query after the client has gone away.
+     *
+     * @return a single-column, single-row query, or null when the dialect has no
+     *         session identifier we can act on
+     */
+    default String getSessionPidQuery() {
+        return null;
+    }
 }

backend/src/main/java/com/dbaagent/provider/mysql/MySQLQueryExecutionProvider.java

@@ -123,4 +123,9 @@ public String generateQuerySignature(String query) {
             return String.valueOf(normalized.hashCode());
         }
     }
+
+    @Override
+    public String getSessionPidQuery() {
+        return "SELECT CONNECTION_ID()";
+    }
 }

backend/src/main/java/com/dbaagent/provider/postgres/PostgresQueryExecutionProvider.java

@@ -127,4 +127,9 @@ public String generateQuerySignature(String query) {
             return String.valueOf(normalized.hashCode());
         }
     }
+
+    @Override
+    public String getSessionPidQuery() {
+        return "SELECT pg_backend_pid()";
+    }
 }

backend/src/main/java/com/dbaagent/service/ActiveQueryService.java

@@ -203,8 +203,15 @@ public void killQuery(String connectionId, String pid) {
                 String dbType = providerRegistry.getCanonicalName(connRequest.getDbType());
 
                 if ("postgres".equals(dbType)) {
+                    // pg_terminate_backend takes an integer; binding a long makes
+                    // the driver send bigint and PostgreSQL then finds no matching
+                    // overload ("function pg_terminate_backend(bigint) does not
+                    // exist"), so every kill failed.
+                    if (backendPid > Integer.MAX_VALUE || backendPid < Integer.MIN_VALUE) {
+                        throw new IllegalArgumentException("Not a valid PostgreSQL backend pid: " + backendPid);
+                    }
                     try (PreparedStatement ps = conn.prepareStatement("SELECT pg_terminate_backend(?)")) {
-                        ps.setLong(1, backendPid);
+                        ps.setInt(1, (int) backendPid);
                         ps.execute();
                     }
                 } else if ("mysql".equals(dbType)) {