diff --git a/CLAUDE.md b/CLAUDE.md index 652ae89..2cf64a0 100644 --- a/CLAUDE.md +++ b/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 diff --git a/backend/src/main/java/com/dbaagent/controller/SchemaController.java b/backend/src/main/java/com/dbaagent/controller/SchemaController.java index 8bab6c9..72204be 100644 --- a/backend/src/main/java/com/dbaagent/controller/SchemaController.java +++ b/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> scanSchema(@PathVariable String connectionId) { @@ -260,6 +264,63 @@ public ResponseEntity> 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> cancelQuery( + @PathVariable String connectionId, + @PathVariable String executionId) { + Map 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> getTableIndexes( diff --git a/backend/src/main/java/com/dbaagent/model/QueryRequest.java b/backend/src/main/java/com/dbaagent/model/QueryRequest.java index 78a673e..a825075 100644 --- a/backend/src/main/java/com/dbaagent/model/QueryRequest.java +++ b/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; diff --git a/backend/src/main/java/com/dbaagent/model/QueryResult.java b/backend/src/main/java/com/dbaagent/model/QueryResult.java index a815fc7..f521c1e 100644 --- a/backend/src/main/java/com/dbaagent/model/QueryResult.java +++ b/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 columns, + List> rows, + Integer rowCount, + Long totalRowCount, + Boolean isLimited, + Long executionTimeMs, + String query + ) { + this(columns, rows, rowCount, totalRowCount, isLimited, executionTimeMs, query, null); + } } diff --git a/backend/src/main/java/com/dbaagent/provider/api/QueryExecutionProvider.java b/backend/src/main/java/com/dbaagent/provider/api/QueryExecutionProvider.java index c0780e3..9841ce0 100644 --- a/backend/src/main/java/com/dbaagent/provider/api/QueryExecutionProvider.java +++ b/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; + } } diff --git a/backend/src/main/java/com/dbaagent/provider/mysql/MySQLQueryExecutionProvider.java b/backend/src/main/java/com/dbaagent/provider/mysql/MySQLQueryExecutionProvider.java index 7350e86..92d82d5 100644 --- a/backend/src/main/java/com/dbaagent/provider/mysql/MySQLQueryExecutionProvider.java +++ b/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()"; + } } diff --git a/backend/src/main/java/com/dbaagent/provider/postgres/PostgresQueryExecutionProvider.java b/backend/src/main/java/com/dbaagent/provider/postgres/PostgresQueryExecutionProvider.java index ac7e657..a2cf814 100644 --- a/backend/src/main/java/com/dbaagent/provider/postgres/PostgresQueryExecutionProvider.java +++ b/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()"; + } } diff --git a/backend/src/main/java/com/dbaagent/service/ActiveQueryService.java b/backend/src/main/java/com/dbaagent/service/ActiveQueryService.java index 49c2ff4..4c4ea00 100644 --- a/backend/src/main/java/com/dbaagent/service/ActiveQueryService.java +++ b/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)) { diff --git a/backend/src/main/java/com/dbaagent/service/QueryExecutionPolicyService.java b/backend/src/main/java/com/dbaagent/service/QueryExecutionPolicyService.java index 32127d8..d330e6a 100644 --- a/backend/src/main/java/com/dbaagent/service/QueryExecutionPolicyService.java +++ b/backend/src/main/java/com/dbaagent/service/QueryExecutionPolicyService.java @@ -19,7 +19,12 @@ import net.sf.jsqlparser.statement.drop.Drop; import net.sf.jsqlparser.statement.insert.Insert; import net.sf.jsqlparser.statement.merge.Merge; +import net.sf.jsqlparser.schema.Table; +import net.sf.jsqlparser.statement.select.ParenthesedSelect; +import net.sf.jsqlparser.statement.select.PlainSelect; import net.sf.jsqlparser.statement.select.Select; +import net.sf.jsqlparser.statement.select.SetOperationList; +import net.sf.jsqlparser.statement.select.WithItem; import net.sf.jsqlparser.statement.truncate.Truncate; import net.sf.jsqlparser.statement.upsert.Upsert; import net.sf.jsqlparser.statement.update.Update; @@ -51,6 +56,22 @@ public class QueryExecutionPolicyService { Pattern.CASE_INSENSITIVE ); + // Text-level backstops for writes that hide inside a statement which reads as + // a SELECT. Applied to SQL with literals and comments already stripped. + private static final Pattern CTE_WRITE_PATTERN = Pattern.compile( + "\\bAS\\s*\\(\\s*(INSERT|UPDATE|DELETE|MERGE)\\b", + Pattern.CASE_INSENSITIVE + ); + // Anchored on INTO alone, with no leading wildcard. The previous form + // (\bSELECT\b[\s\S]*?\bINTO ...) backtracked quadratically: 224KB of + // repeated "SELECT " took ~44s of CPU in the guard before the query reached + // the database. The caller already knows the statement is a SELECT, so the + // prefix bought nothing. + private static final Pattern SELECT_INTO_PATTERN = Pattern.compile( + "\\bINTO\\s+(?!STRICT\\b)[\"`\\w]", + Pattern.CASE_INSENSITIVE + ); + private final DatabaseProviderRegistry providerRegistry; public PolicyDecision enforce( @@ -163,9 +184,23 @@ private StatementClassification classifyStatement(String statement, QueryExecuti return new StatementClassification("USE", false, false, false, false, true); } + // A pre-parse scan runs first so a statement the parser mis-models — or + // rejects outright — still fails closed. `WITH ... DELETE` parses cleanly + // as a Select, and the keyword fallback below treats anything starting + // with WITH as read-only, so without this a write reaches the database + // classified as a read. + String hiddenWrite = detectHiddenWrite(trimmed); + try { Statement parsed = CCJSqlParserUtil.parse(trimmed); - if (parsed instanceof Select) { + if (parsed instanceof Select select) { + String writeKind = hiddenWrite != null ? hiddenWrite : detectSelectWrite(select); + if (writeKind != null) { + // requiresWhereClause=false: the WHERE guard reads the top-level + // statement, which is a SELECT here. A data-modifying CTE always + // needs explicit confirmation instead. + return new StatementClassification(writeKind, false, true, false, true, false); + } return new StatementClassification("SELECT", true, false, false, false, false); } if (parsed instanceof ExplainStatement) { @@ -220,6 +255,12 @@ private StatementClassification classifyStatement(String statement, QueryExecuti return new StatementClassification(explainWrappedMutation, false, true, requiresWhere, hasWhere, false); } + // The parser failed or produced a type we don't model. `isReadOnlyQuery` + // only looks at the leading keyword, so a hidden write must veto it. + if (hiddenWrite != null) { + return new StatementClassification(hiddenWrite, false, true, false, true, false); + } + String queryType = QueryNormalizer.detectQueryType(trimmed); boolean readOnly = executionProvider.isReadOnlyQuery(trimmed); boolean mutating = !readOnly && !"UNKNOWN".equalsIgnoreCase(queryType); @@ -228,6 +269,133 @@ private StatementClassification classifyStatement(String statement, QueryExecuti return new StatementClassification(queryType, readOnly, mutating, requiresWhere, hasWhere, false); } + /** + * Detects a write hidden inside a statement that parses as a {@link Select}: + * a data-modifying CTE ({@code WITH x AS (DELETE ...) SELECT ...}, which + * PostgreSQL executes for real) or {@code SELECT ... INTO newtable}, which + * creates a table. Returns the write's query type, or null when the select + * really is read-only. + */ + private String detectSelectWrite(Select select) { + String fromCte = detectWriteInWithItems(select.getWithItemsList(), 0); + if (fromCte != null) { + return fromCte; + } + return detectSelectInto(select); + } + + private String detectWriteInWithItems(List> withItems, int depth) { + // CTEs can nest; bound the walk so a pathological query can't spin here. + if (withItems == null || withItems.isEmpty() || depth > 10) { + return null; + } + for (WithItem item : withItems) { + if (item == null) { + continue; + } + if (item.getDelete() != null) return "DELETE (in CTE)"; + if (item.getUpdate() != null) return "UPDATE (in CTE)"; + if (item.getInsert() != null) return "INSERT (in CTE)"; + + // getSelect() throws when the item holds a non-select statement, which + // the getters above already covered. A select here may carry its own + // WITH list, so recurse to catch a write nested one level deeper. + try { + Select inner = item.getSelect(); + if (inner != null) { + String innerWrite = detectSelectWriteQuietly(inner, depth + 1); + if (innerWrite != null) { + return innerWrite; + } + } + } catch (Exception ignored) { + // Not a select-bearing CTE; nothing further to inspect. + } + } + return null; + } + + private String detectSelectWriteQuietly(Select select, int depth) { + if (depth > 10) { + return null; + } + String fromCte = detectWriteInWithItems(select.getWithItemsList(), depth); + return fromCte != null ? fromCte : detectSelectInto(select); + } + + private String detectSelectInto(Select select) { + if (select instanceof PlainSelect plain) { + List into = plain.getIntoTables(); + if (into != null && !into.isEmpty()) { + return "SELECT INTO"; + } + } + if (select instanceof ParenthesedSelect parenthesed) { + Select inner = parenthesed.getSelect(); + return inner == null ? null : detectSelectInto(inner); + } + if (select instanceof SetOperationList setOps && setOps.getSelects() != null) { + for (Select part : setOps.getSelects()) { + String found = detectSelectInto(part); + if (found != null) { + return found; + } + } + } + return null; + } + + /** + * Text-level backstop for the same two shapes, used when the parser is not + * available or disagrees. Quoted literals and comments are stripped first so + * the word DELETE inside a string can't trigger a false positive. + */ + private String detectHiddenWrite(String sql) { + if (sql == null || sql.isBlank()) { + return null; + } + String scrubbed = stripLeadingComments(stripComments(stripQuotedLiterals(sql))); + boolean startsWithWith = scrubbed.regionMatches(true, 0, "WITH", 0, 4); + if (startsWithWith && CTE_WRITE_PATTERN.matcher(scrubbed).find()) { + return "WRITE (in CTE)"; + } + // SELECT_INTO_PATTERN matches a bare INTO target, so it must only be + // consulted for statements that actually read as a SELECT — otherwise + // "INSERT INTO t SELECT ..." would be relabelled SELECT INTO. + if (startsWithWith || scrubbed.regionMatches(true, 0, "SELECT", 0, 6)) { + return SELECT_INTO_PATTERN.matcher(scrubbed).find() ? "SELECT INTO" : null; + } + return null; + } + + // Removes block and line comments with a single linear scan. The regex form + // matched a lazy wildcard between block-comment delimiters and backtracked + // quadratically on an unterminated block comment: 96KB of input cost ~14s + // of CPU. + private String stripComments(String sql) { + StringBuilder out = new StringBuilder(sql.length()); + int i = 0; + int len = sql.length(); + while (i < len) { + char c = sql.charAt(i); + if (c == '/' && i + 1 < len && sql.charAt(i + 1) == '*') { + int end = sql.indexOf("*/", i + 2); + i = end < 0 ? len : end + 2; + out.append(' '); + continue; + } + if (c == '-' && i + 1 < len && sql.charAt(i + 1) == '-') { + int nl = sql.indexOf('\n', i + 2); + i = nl < 0 ? len : nl; + out.append(' '); + continue; + } + out.append(c); + i++; + } + return out.toString(); + } + private boolean isUseStatement(String sql) { return sql != null && USE_ANY_PATTERN.matcher(stripLeadingComments(sql)).find(); } diff --git a/backend/src/main/java/com/dbaagent/service/QueryExecutorService.java b/backend/src/main/java/com/dbaagent/service/QueryExecutorService.java index 0c14cc9..110eaa7 100644 --- a/backend/src/main/java/com/dbaagent/service/QueryExecutorService.java +++ b/backend/src/main/java/com/dbaagent/service/QueryExecutorService.java @@ -29,10 +29,34 @@ public class QueryExecutorService { private final UserDataAccessPolicyService userDataAccessPolicyService; private final com.dbaagent.service.telemetry.TelemetryCounters telemetryCounters; private final KeyColumnAnalysisRepository keyColumnAnalysisRepository; + private final RunningQueryRegistry runningQueryRegistry; @Value("${db.query-timeout-seconds:30}") private int queryTimeoutSeconds; + /** + * Reads the connection's server-side session id via the dialect's own query. + * Never fatal: without it the query simply cannot be cancelled early. + */ + private String resolveSessionPid(Connection connection, String dbType) { + String pidQuery; + try { + pidQuery = providerRegistry.getDialect(dbType).queryExecution().getSessionPidQuery(); + } catch (Exception e) { + return null; + } + if (pidQuery == null) { + return null; + } + try (Statement stmt = connection.createStatement(); + ResultSet rs = stmt.executeQuery(pidQuery)) { + return rs.next() ? rs.getString(1) : null; + } catch (SQLException e) { + log.debug("Could not resolve session pid: {}", e.getMessage()); + return null; + } + } + public List getDatabaseObjects(String connectionId) throws SQLException { // Try to get from cache, but handle cache failures gracefully Cache cache = null; @@ -547,6 +571,23 @@ public QueryResult executeQuery( ); try (Connection connection = connectionService.getConnection(connectionId, connRequest)) { + // Defense in depth: when the caller may not mutate, ask the database + // to enforce that too. Classification is a parser heuristic, so a + // read-only session is what keeps a future gap from becoming data + // loss — PostgreSQL then refuses the write itself. + boolean readOnlySession = + executionContext.mutationMode() == QueryExecutionContext.MutationMode.READ_ONLY_ONLY; + if (readOnlySession) { + try { + connection.setReadOnly(true); + } catch (SQLException e) { + // A driver that refuses the hint must not silently downgrade + // to a writable session. + throw new SQLException( + "Could not open a read-only database session for this connection: " + e.getMessage(), e); + } + } + // Per-query timeout override: use the request's timeout if provided, // otherwise fall back to the server default (db.query-timeout-seconds). int maxAllowedTimeoutSeconds = 600; // 10 minutes ceiling @@ -554,6 +595,14 @@ public QueryResult executeQuery( ? Math.min(queryRequest.getTimeoutSeconds(), maxAllowedTimeoutSeconds) : queryTimeoutSeconds; + // Publish the server-side session id so a client that abandons this + // request can terminate the query rather than leaving it to run out + // its timeout holding a pooled connection. + String sessionPid = resolveSessionPid(connection, connRequest.getDbType()); + result.setSessionPid(sessionPid); + runningQueryRegistry.register( + queryRequest.getExecutionId(), connectionId, sessionPid, executionContext.actorUsername()); + // Execute all preamble statements (SET @var, USE db, etc.) on the same connection // so that session variables are in scope for the final SELECT. for (int si = 0; si < statements.size() - 1; si++) { @@ -573,15 +622,23 @@ public QueryResult executeQuery( stmt.setQueryTimeout(effectiveTimeout); } - // Apply limit if specified + // Row cap. setMaxRows is the enforcement: the driver stops handing + // back rows regardless of the query's shape, so a LIMIT inside a CTE, + // a comment, or a string literal can no longer defeat the cap. We ask + // for one extra row purely to detect truncation. + // + // Appending " LIMIT n" is kept as an optimization for simple SELECTs + // (it lets the database stop early rather than stream rows we drop), + // but correctness no longer depends on that textual check. + boolean limitRequested = queryRequest.getLimit() != null && queryRequest.getLimit() > 0; + int rowLimit = limitRequested ? queryRequest.getLimit() : 0; boolean limitApplied = false; String baseQuery = null; // original query without limit, for COUNT - if (queryRequest.getLimit() != null && queryRequest.getLimit() > 0) { + if (limitRequested) { + stmt.setMaxRows(rowLimit + 1); + String queryTrimmed = finalQuery.trim(); String queryLower = queryTrimmed.toLowerCase(); - // Only add LIMIT to SELECT queries, not to SHOW/DESCRIBE/EXPLAIN commands. - // Use regex to check for actual LIMIT clause (avoids false positives from - // column names, comments, or identifiers that contain the word "limit"). boolean isSelect = queryLower.startsWith("select"); boolean alreadyHasLimit = queryLower.matches("(?s).*\\blimit\\s+\\d+.*"); if (isSelect && !alreadyHasLimit) { @@ -589,7 +646,10 @@ public QueryResult executeQuery( if (baseQuery.endsWith(";")) { baseQuery = baseQuery.substring(0, baseQuery.length() - 1); } - finalQuery = baseQuery + " LIMIT " + queryRequest.getLimit(); + // Ask for one more than we will return, matching setMaxRows, + // so an exactly-at-the-cap result is still detectable as + // truncated rather than looking complete. + finalQuery = baseQuery + " LIMIT " + (rowLimit + 1); limitApplied = true; } } @@ -620,9 +680,16 @@ public QueryResult executeQuery( } result.setColumns(columns); - // Get rows + // Get rows. setMaxRows let one extra row through so we can + // tell "exactly at the cap" from "there were more"; it is + // counted but never returned. List> rows = new ArrayList<>(); + boolean truncated = false; while (rs.next()) { + if (limitRequested && rows.size() >= rowLimit) { + truncated = true; + break; + } List row = new ArrayList<>(); for (int i = 1; i <= columnCount; i++) { Object value = rs.getObject(i); @@ -632,11 +699,14 @@ public QueryResult executeQuery( } result.setRows(rows); result.setRowCount(rows.size()); - - // If a limit was applied, mark result as limited and fetch total row count. - // The COUNT query runs on the same connection so @session variables are in scope. - if (limitApplied && baseQuery != null) { + if (truncated) { result.setIsLimited(true); + } + + // Fetch the true total only when we appended the LIMIT + // ourselves, since that is the one case where baseQuery is + // a valid standalone statement to wrap in COUNT(*). + if (truncated && limitApplied && baseQuery != null) { try (Statement countStmt = connection.createStatement()) { countStmt.setQueryTimeout(effectiveTimeout > 0 ? effectiveTimeout : 60); // Strip trailing ORDER BY before wrapping in COUNT subquery, @@ -661,6 +731,8 @@ public QueryResult executeQuery( result.setRows(Arrays.asList(Arrays.asList(updateCount))); } } + } finally { + runningQueryRegistry.unregister(queryRequest.getExecutionId()); } result.setExecutionTimeMs(System.currentTimeMillis() - startTime); diff --git a/backend/src/main/java/com/dbaagent/service/RunningQueryRegistry.java b/backend/src/main/java/com/dbaagent/service/RunningQueryRegistry.java new file mode 100644 index 0000000..ff1f805 --- /dev/null +++ b/backend/src/main/java/com/dbaagent/service/RunningQueryRegistry.java @@ -0,0 +1,66 @@ +package com.dbaagent.service; + +import lombok.extern.slf4j.Slf4j; +import org.springframework.stereotype.Service; + +import java.time.Duration; +import java.time.Instant; +import java.util.Iterator; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.ConcurrentHashMap; + +/** + * Tracks in-flight queries by a client-supplied execution id so a client that + * gives up can terminate the query it started. Aborting the HTTP request only + * closes the socket; without this the statement keeps running and holds a pooled + * connection until it finishes on its own. + * + *

Entries are short-lived: registered when execution begins and removed in a + * finally block. The sweep on write is a backstop for a process that died between + * those two points. + */ +@Service +@Slf4j +public class RunningQueryRegistry { + + private static final Duration MAX_AGE = Duration.ofHours(1); + private static final int SWEEP_THRESHOLD = 256; + + public record RunningQuery(String connectionId, String sessionPid, String username, Instant startedAt) {} + + private final Map running = new ConcurrentHashMap<>(); + + public void register(String executionId, String connectionId, String sessionPid, String username) { + if (executionId == null || executionId.isBlank() || sessionPid == null || sessionPid.isBlank()) { + return; + } + if (running.size() > SWEEP_THRESHOLD) { + sweepExpired(); + } + running.put(executionId, new RunningQuery(connectionId, sessionPid, username, Instant.now())); + } + + public void unregister(String executionId) { + if (executionId != null && !executionId.isBlank()) { + running.remove(executionId); + } + } + + public Optional find(String executionId) { + if (executionId == null || executionId.isBlank()) { + return Optional.empty(); + } + return Optional.ofNullable(running.get(executionId)); + } + + private void sweepExpired() { + Instant cutoff = Instant.now().minus(MAX_AGE); + Iterator> it = running.entrySet().iterator(); + while (it.hasNext()) { + if (it.next().getValue().startedAt().isBefore(cutoff)) { + it.remove(); + } + } + } +} diff --git a/backend/src/main/java/com/dbaagent/service/UserDataAccessPolicyService.java b/backend/src/main/java/com/dbaagent/service/UserDataAccessPolicyService.java index 6ef3e66..a65a9f6 100644 --- a/backend/src/main/java/com/dbaagent/service/UserDataAccessPolicyService.java +++ b/backend/src/main/java/com/dbaagent/service/UserDataAccessPolicyService.java @@ -216,7 +216,8 @@ public QueryResult redactResult( result.getTotalRowCount(), result.getIsLimited(), result.getExecutionTimeMs(), - result.getQuery() + result.getQuery(), + result.getSessionPid() ); logPolicyEvent(SecurityEventType.CHAT_ACCESS_POLICY_REDACTED, executionContext.actorUsername(), connectionId, "result_redacted", Map.of( "query", truncate(result.getQuery()), diff --git a/backend/src/test/java/com/dbaagent/service/QueryExecutionPolicyServiceTest.java b/backend/src/test/java/com/dbaagent/service/QueryExecutionPolicyServiceTest.java index ef78d7f..566ab0d 100644 --- a/backend/src/test/java/com/dbaagent/service/QueryExecutionPolicyServiceTest.java +++ b/backend/src/test/java/com/dbaagent/service/QueryExecutionPolicyServiceTest.java @@ -4,6 +4,7 @@ import com.dbaagent.provider.DatabaseProviderRegistry; import com.dbaagent.provider.api.DatabaseDialect; import com.dbaagent.provider.api.QueryExecutionProvider; +import com.dbaagent.provider.mysql.MySQLQueryExecutionProvider; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; @@ -23,16 +24,20 @@ class QueryExecutionPolicyServiceTest { @Mock private DatabaseProviderRegistry providerRegistry; @Mock private DatabaseDialect databaseDialect; - @Mock private QueryExecutionProvider queryExecutionProvider; private QueryExecutionPolicyService service; @BeforeEach void setUp() { + // The real provider, not a stub. A stubbed isReadOnlyQuery() that always + // answered "false" made these tests assert the opposite of production: + // MySQLQueryExecutionProvider reports anything starting with WITH as + // read-only, which is how `WITH x AS (DELETE ...) SELECT` reached the + // database classified as a read. + QueryExecutionProvider realProvider = new MySQLQueryExecutionProvider(); when(providerRegistry.getDialect(anyString())).thenReturn(databaseDialect); - when(databaseDialect.queryExecution()).thenReturn(queryExecutionProvider); + when(databaseDialect.queryExecution()).thenReturn(realProvider); lenient().when(providerRegistry.getCanonicalName(anyString())).thenReturn("mysql"); - lenient().when(queryExecutionProvider.isReadOnlyQuery(anyString())).thenReturn(false); service = new QueryExecutionPolicyService(providerRegistry); } @@ -296,4 +301,158 @@ void internalMutation_isStillAllowed() { assertThat(decision.mutating()).isTrue(); assertThat(decision.primaryQueryType()).isEqualTo("TRUNCATE"); } + + // --- Writes hidden inside a statement that reads as a SELECT --------------- + // PostgreSQL executes data-modifying CTEs, and the parser models them as a + // Select. Each of these deleted or rewrote a whole table from a non-admin + // account before the classifier learned to look inside. + + private QueryExecutionPolicyException assertBlockedForViewer(String sql) { + return assertThrows( + QueryExecutionPolicyException.class, + () -> service.enforce( + new QueryRequest(sql, null, null), + QueryExecutionContext.editor("viewer", false, false), + "postgresql" + ) + ); + } + + @Test + void cteDelete_isBlockedForNonAdmin() { + QueryExecutionPolicyException e = + assertBlockedForViewer("WITH x AS (DELETE FROM orders RETURNING *) SELECT * FROM x"); + assertThat(e.getErrorCode()).isEqualTo(QueryExecutionPolicyException.EDITOR_MUTATION_FORBIDDEN); + } + + @Test + void cteUpdate_isBlockedForNonAdmin() { + assertBlockedForViewer("WITH u AS (UPDATE orders SET total = 0 RETURNING *) SELECT * FROM u"); + } + + @Test + void cteInsert_isBlockedForNonAdmin() { + assertBlockedForViewer("WITH i AS (INSERT INTO audit(id) VALUES (1) RETURNING *) SELECT * FROM i"); + } + + @Test + void cteWriteInLaterPosition_isBlockedForNonAdmin() { + assertBlockedForViewer( + "WITH a AS (SELECT 1), b AS (DELETE FROM orders RETURNING *) SELECT * FROM a"); + } + + @Test + void nestedCteWrite_isBlockedForNonAdmin() { + assertBlockedForViewer( + "WITH o AS (WITH i AS (DELETE FROM orders RETURNING *) SELECT * FROM i) SELECT * FROM o"); + } + + @Test + void selectInto_isBlockedForNonAdmin() { + assertBlockedForViewer("SELECT * INTO exfiltrated FROM customers"); + } + + @Test + void cteWrite_requiresConfirmationForAdmin() { + QueryExecutionPolicyException e = assertThrows( + QueryExecutionPolicyException.class, + () -> service.enforce( + new QueryRequest("WITH x AS (DELETE FROM orders RETURNING *) SELECT * FROM x", null, null), + QueryExecutionContext.editor("admin", true, false), + "postgresql" + ) + ); + assertThat(e.getErrorCode()) + .isEqualTo(QueryExecutionPolicyException.EDITOR_MUTATION_CONFIRMATION_REQUIRED); + } + + @Test + void cteWrite_isAllowedForConfirmedAdmin() { + QueryExecutionPolicyService.PolicyDecision decision = service.enforce( + new QueryRequest("WITH x AS (DELETE FROM orders RETURNING *) SELECT * FROM x", null, null), + QueryExecutionContext.editor("admin", true, true), + "postgresql" + ); + assertThat(decision.mutating()).isTrue(); + } + + // --- The guard must not swallow legitimate reads -------------------------- + + @Test + void readOnlyCte_remainsAllowedForNonAdmin() { + QueryExecutionPolicyService.PolicyDecision decision = service.enforce( + new QueryRequest("WITH recent AS (SELECT * FROM orders LIMIT 10) SELECT * FROM recent", null, null), + QueryExecutionContext.editor("viewer", false, false), + "postgresql" + ); + assertThat(decision.mutating()).isFalse(); + assertThat(decision.primaryQueryType()).isEqualTo("SELECT"); + } + + @Test + void writeKeywordInsideStringLiteral_isStillAReadForNonAdmin() { + QueryExecutionPolicyService.PolicyDecision decision = service.enforce( + new QueryRequest("SELECT 'WITH x AS (DELETE FROM t)' AS example", null, null), + QueryExecutionContext.editor("viewer", false, false), + "postgresql" + ); + assertThat(decision.mutating()).isFalse(); + } + + @Test + void writeKeywordInsideComment_isStillAReadForNonAdmin() { + QueryExecutionPolicyService.PolicyDecision decision = service.enforce( + new QueryRequest("SELECT 1 -- WITH x AS (DELETE FROM t)\n", null, null), + QueryExecutionContext.editor("viewer", false, false), + "postgresql" + ); + assertThat(decision.mutating()).isFalse(); + } + + @Test + void hiddenWriteScan_staysLinearOnAdversarialInput() { + // The text backstop used a lazy wildcard (\bSELECT\b[\s\S]*?\bINTO) + // and a regex block-comment strip, both of which backtracked + // quadratically: 224KB of repeated "SELECT " burned ~44s of CPU inside + // the guard, before the query ever reached the database. Any + // authenticated Editor user could stall a request thread with it. + String repeatedSelect = "SELECT " + "SELECT ".repeat(32_000); + String unterminatedBlockComment = "SELECT 1 /*" + "a/*".repeat(32_000); + + for (String hostile : List.of(repeatedSelect, unterminatedBlockComment)) { + long startedAt = System.currentTimeMillis(); + service.enforce( + new QueryRequest(hostile, null, null), + QueryExecutionContext.editor("viewer", false, false), + "postgresql" + ); + long elapsed = System.currentTimeMillis() - startedAt; + assertThat(elapsed) + .as("classification of a %d char statement must not backtrack", hostile.length()) + .isLessThan(5_000L); + } + } + + @Test + void insertIntoSelect_isNotMisreadAsSelectInto() { + // SELECT_INTO_PATTERN matches a bare INTO target now that the SELECT + // prefix is gone, so the caller must gate it on the statement actually + // reading as a SELECT. + QueryExecutionPolicyService.PolicyDecision decision = service.enforce( + new QueryRequest("INSERT INTO archive SELECT * FROM orders", null, null), + QueryExecutionContext.editor("admin", true, true), + "postgresql" + ); + assertThat(decision.primaryQueryType()).isEqualTo("INSERT"); + } + + @Test + void insertIntoSelect_isStillClassifiedAsInsert() { + QueryExecutionPolicyService.PolicyDecision decision = service.enforce( + new QueryRequest("INSERT INTO archive SELECT * FROM orders", null, null), + QueryExecutionContext.editor("admin", true, true), + "postgresql" + ); + assertThat(decision.primaryQueryType()).isEqualTo("INSERT"); + } } diff --git a/backend/src/test/java/com/dbaagent/service/QueryExecutorServiceCounterTest.java b/backend/src/test/java/com/dbaagent/service/QueryExecutorServiceCounterTest.java index 87b5d3e..ed23806 100644 --- a/backend/src/test/java/com/dbaagent/service/QueryExecutorServiceCounterTest.java +++ b/backend/src/test/java/com/dbaagent/service/QueryExecutorServiceCounterTest.java @@ -58,7 +58,8 @@ void setUp() { queryExecutionPolicyService, userDataAccessPolicyService, telemetryCounters, - keyColumnAnalysisRepository + keyColumnAnalysisRepository, + new RunningQueryRegistry() ); } diff --git a/backend/src/test/java/com/dbaagent/service/QueryExecutorServiceTest.java b/backend/src/test/java/com/dbaagent/service/QueryExecutorServiceTest.java index c23d805..a26d479 100644 --- a/backend/src/test/java/com/dbaagent/service/QueryExecutorServiceTest.java +++ b/backend/src/test/java/com/dbaagent/service/QueryExecutorServiceTest.java @@ -56,7 +56,8 @@ void setUp() { queryExecutionPolicyService, userDataAccessPolicyService, telemetryCounters, - keyColumnAnalysisRepository + keyColumnAnalysisRepository, + new RunningQueryRegistry() ); } diff --git a/docker/nginx/default.conf b/docker/nginx/default.conf index e69243f..b4e2f7d 100644 --- a/docker/nginx/default.conf +++ b/docker/nginx/default.conf @@ -1,3 +1,9 @@ +# Ad-hoc SQL execution is the most expensive thing an authenticated user can do: +# each request can occupy a database connection for minutes. Cap the sustained +# rate per client, with a burst so a normal person clicking Run repeatedly is +# never throttled. +limit_req_zone $binary_remote_addr zone=sqlexec:10m rate=30r/m; + server { listen 80; server_name _; @@ -26,6 +32,27 @@ server { add_header Expires "0"; } + # Ad-hoc SQL execution, rate-limited. Declared before the general /api/ rule + # so it wins for this path; everything else is identical to it. + # `nodelay` lets the burst through immediately rather than queueing it, so + # normal interactive use is unaffected and only sustained abuse is rejected. + location ~ ^/api/connections/[^/]+/query$ { + limit_req zone=sqlexec burst=20 nodelay; + limit_req_status 429; + + proxy_pass http://backend:8080$request_uri; + proxy_http_version 1.1; + proxy_set_header Host $host; + proxy_set_header X-Real-IP $remote_addr; + proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for; + proxy_set_header X-Forwarded-Proto $scheme; + proxy_read_timeout 300s; + proxy_send_timeout 300s; + proxy_connect_timeout 10s; + proxy_buffering off; + proxy_cache off; + } + # API proxy to the backend service # Allows the frontend to call /api/... without CORS issues location /api/ { diff --git a/src/components/tabs/Core/SqlRunnerTab.js b/src/components/tabs/Core/SqlRunnerTab.js index 9eedffd..376ea19 100644 --- a/src/components/tabs/Core/SqlRunnerTab.js +++ b/src/components/tabs/Core/SqlRunnerTab.js @@ -37,7 +37,6 @@ import { queryPerformanceAPI, explainAPI, brainAPI, - activeQueryAPI, } from "@/lib/api/client"; import ExplainAnalysisPanel from "./ExplainAnalysisPanel"; import QueryOptimizePanel from "./QueryOptimizePanel"; @@ -51,6 +50,10 @@ import { connectionHasMultipleSchemas, } from "@/lib/schemaNames"; +// Kept below the 300s proxy_read_timeout in docker/nginx/default.conf so a slow +// query fails with a real message rather than an opaque 504. +const QUERY_TIMEOUT_SECONDS = 240; + // Constants for diagram layout const DIAGRAM_NODE_WIDTH = 240; const DIAGRAM_NODE_HEIGHT = 120; @@ -294,6 +297,7 @@ export default function SqlRunnerTab({ connectionId }) { const autocompleteRegisteredRef = useRef(false); const dbObjectsRef = useRef([]); const abortControllerRef = useRef(null); + const executionIdRef = useRef(null); // Note: savedQueriesPanelRef kept for potential future use, but panel UI is now in modal const savedQueriesPanelRef = useRef(null); const hasRowCount = (value) => value !== null && value !== undefined; @@ -1256,19 +1260,27 @@ export default function SqlRunnerTab({ connectionId }) { const abortController = new AbortController(); abortControllerRef.current = abortController; + // Identifies this run so cancelling can terminate it on the database. + const executionId = + globalThis.crypto?.randomUUID?.() ?? + `exec-${Date.now()}-${Math.random().toString(16).slice(2)}`; + executionIdRef.current = executionId; try { - // Use extended timeout (10 min) for SQL editor queries to support long-running queries. - // The server default is 30s which is too short for analytical queries. + // The server default of 30s is too short for analytical queries, but the + // nginx proxy in front of the API gives up at 300s (docker/nginx/default.conf). + // Staying under that means a slow query surfaces as a real error here + // instead of an opaque gateway timeout. const response = await queryAPI.executeQuery( connectionId, trimmedQuery, 1000, - 600, + QUERY_TIMEOUT_SECONDS, abortController.signal, { executionOrigin: "EDITOR", mutationConfirmed, + executionId, }, ); @@ -1344,29 +1356,25 @@ export default function SqlRunnerTab({ connectionId }) { } } finally { abortControllerRef.current = null; + executionIdRef.current = null; setIsRunning(false); } }; const handleStopQuery = () => { + const executionId = executionIdRef.current; if (abortControllerRef.current) { abortControllerRef.current.abort(); } setIsRunning(false); setError(null); - // Best-effort: capture active queries and kill any running ones on the DB side. - // Fire-and-forget — don't block the UI on this. - if (connectionId) { - activeQueryAPI.capture(connectionId) - .then((queries) => { - const running = queries.filter((q) => q.state === "active" || q.state === "running"); - running.forEach((q) => { - if (q.pid) { - activeQueryAPI.kill(connectionId, q.pid).catch(() => {}); - } - }); - }) - .catch(() => {}); + // Aborting above only drops the HTTP response; the statement keeps running + // and holds a pooled connection. Ask the server to terminate this specific + // execution. Previously this killed every active query on the connection, + // which could take out other users' work and DeepSQL's own background jobs. + if (connectionId && executionId) { + executionIdRef.current = null; + queryAPI.cancelQuery(connectionId, executionId).catch(() => {}); } }; diff --git a/src/lib/api/client.js b/src/lib/api/client.js index 9ab49ee..2694008 100644 --- a/src/lib/api/client.js +++ b/src/lib/api/client.js @@ -1386,6 +1386,11 @@ export const queryAPI = { if (requestOptions.mutationConfirmed != null) { payload.mutationConfirmed = requestOptions.mutationConfirmed; } + // Lets the caller cancel this run server-side; aborting the request alone + // leaves the query running on the database. + if (requestOptions.executionId) { + payload.executionId = requestOptions.executionId; + } // Use a longer axios timeout for queries with extended timeouts const axiosTimeout = timeoutSeconds != null @@ -1402,6 +1407,13 @@ export const queryAPI = { return response.data; }, + cancelQuery: async (connectionId, executionId) => { + const response = await apiClient.post( + `/api/connections/${connectionId}/query/${encodeURIComponent(executionId)}/cancel`, + ); + return response.data; + }, + getTableIndexes: async (connectionId, tableName) => { // Encode so schema-qualified ids (`crm.orders`) survive the path segment. const tableId = encodeURIComponent(String(tableName || ""));