Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 54 additions & 0 deletions CLAUDE.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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(
Expand Down
3 changes: 3 additions & 0 deletions backend/src/main/java/com/dbaagent/model/QueryRequest.java
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
19 changes: 19 additions & 0 deletions backend/src/main/java/com/dbaagent/model/QueryResult.java
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -123,4 +123,9 @@ public String generateQuerySignature(String query) {
return String.valueOf(normalized.hashCode());
}
}

@Override
public String getSessionPidQuery() {
return "SELECT CONNECTION_ID()";
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -127,4 +127,9 @@ public String generateQuerySignature(String query) {
return String.valueOf(normalized.hashCode());
}
}

@Override
public String getSessionPidQuery() {
return "SELECT pg_backend_pid()";
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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)) {
Expand Down
Loading
Loading