fix(editor): bound CSV export, guard concurrent runs, audit query can… (#76)

Author: notSumit25

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

@@ -275,8 +275,10 @@ public ResponseEntity<Map<String, Object>> executeQuery(
     @PostMapping("/query/{executionId}/cancel")
     public ResponseEntity<Map<String, Object>> cancelQuery(
             @PathVariable String connectionId,
-            @PathVariable String executionId) {
+            @PathVariable String executionId,
+            HttpServletRequest httpRequest) {
         Map<String, Object> response = new HashMap<>();
+        ClientContext client = ClientContext.fromRequest(httpRequest);
         try {
             if (!credentialService.connectionExists(connectionId)) {
                 response.put("success", false);
@@ -287,7 +289,14 @@ public ResponseEntity<Map<String, Object>> cancelQuery(
 
             var running = runningQueryRegistry.find(executionId);
             if (running.isEmpty()) {
-                // Already finished, or never started. Nothing to cancel.
+                // Already finished, or never started. Nothing to cancel. Still
+                // audited: without this, a cancel that misses its target left
+                // no trace at all, deliberate or not.
+                sqlExecutionAuditService.record(SqlExecutionAuditService.AuditRecord.cancelNoOp()
+                    .connectionId(connectionId)
+                    .executionId(executionId)
+                    .httpRequest(httpRequest)
+                    .client(client));
                 response.put("success", true);
                 response.put("cancelled", false);
                 response.put("message", "Query is no longer running");
@@ -301,13 +310,24 @@ public ResponseEntity<Map<String, Object>> cancelQuery(
             String currentUser = accessControlService.getCurrentUsername();
             if (!connectionId.equals(target.connectionId())
                 || (target.username() != null && currentUser != null && !target.username().equals(currentUser))) {
+                sqlExecutionAuditService.record(SqlExecutionAuditService.AuditRecord.blocked(
+                        "cancel requested for an execution id not owned by this caller/connection")
+                    .connectionId(connectionId)
+                    .executionId(executionId)
+                    .httpRequest(httpRequest)
+                    .client(client));
                 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);
+            sqlExecutionAuditService.record(SqlExecutionAuditService.AuditRecord.cancelled(target.sessionPid())
+                .connectionId(connectionId)
+                .executionId(executionId)
+                .httpRequest(httpRequest)
+                .client(client));
             response.put("success", true);
             response.put("cancelled", true);
             response.put("message", "Query cancelled");
@@ -318,6 +338,12 @@ public ResponseEntity<Map<String, Object>> cancelQuery(
             return ResponseEntity.status(e.getStatusCode()).body(response);
         } catch (Exception e) {
             log.warn("Failed to cancel query {} on connection {}: {}", executionId, connectionId, e.getMessage());
+            sqlExecutionAuditService.record(SqlExecutionAuditService.AuditRecord.failed(e.getMessage())
+                .operation("cancel")
+                .connectionId(connectionId)
+                .executionId(executionId)
+                .httpRequest(httpRequest)
+                .client(client));
             response.put("success", false);
             response.put("message", "Failed to cancel query: " + e.getMessage());
             return ResponseEntity.status(HttpStatus.INTERNAL_SERVER_ERROR).body(response);

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

@@ -41,6 +41,7 @@ public enum SecurityEventType {
     EDITOR_QUERY_EXECUTED,
     EDITOR_QUERY_BLOCKED,
     EDITOR_QUERY_FAILED,
+    EDITOR_QUERY_CANCELLED,
     SUSPICIOUS_AUTH_ACTIVITY,
     SMTP_CONFIG_UPDATED,
     SMTP_TEST_SUCCEEDED,

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

@@ -466,8 +466,19 @@ private String stripQuotedLiterals(String sql) {
             .replaceAll("\"([^\"\\\\]|\\\\.)*\"", "\"\"");
     }
 
+    // Word-bounded and comment-stripped: the previous literal " WHERE " match
+    // missed a WHERE preceded by a newline (as in any multi-line formatted
+    // UPDATE/DELETE) and was satisfied by a commented-out "-- WHERE ..." that
+    // never reaches the database. This only runs on the keyword-fallback path —
+    // when JSqlParser succeeds, delete.getWhere()/update.getWhere() are used
+    // instead, which are exact.
+    private static final Pattern WHERE_CLAUSE_PATTERN = Pattern.compile("\\bWHERE\\b", Pattern.CASE_INSENSITIVE);
+
     private boolean containsWhereClause(String sql) {
-        return sql != null && sql.toUpperCase(Locale.ROOT).contains(" WHERE ");
+        if (sql == null) {
+            return false;
+        }
+        return WHERE_CLAUSE_PATTERN.matcher(stripComments(sql)).find();
     }
 
     private boolean looksLikeMultipleStatements(String sql) {

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

@@ -80,6 +80,9 @@ public void record(AuditRecord rec) {
             metadata.put("connectionId", rec.connectionId);
             metadata.put("connectionName", rec.connectionRequest == null ? null : rec.connectionRequest.getConnectionName());
             metadata.put("dbType", rec.connectionRequest == null ? null : rec.connectionRequest.getDbType());
+            if (rec.executionId != null) {
+                metadata.put("executionId", rec.executionId);
+            }
 
             String queryText = rec.queryRequest == null ? null : rec.queryRequest.getQuery();
             metadata.put("queryHash", SecurityHashUtil.sha256Hex(queryText == null ? "" : queryText));
@@ -183,6 +186,7 @@ public static final class AuditRecord {
         private String failureReason;
         private HttpServletRequest httpRequest;
         private ClientContext client;
+        private String executionId;
 
         private AuditRecord() {}
 
@@ -214,6 +218,33 @@ public static AuditRecord failed(String reason) {
             return r;
         }
 
+        /**
+         * A deliberate cancel request against a still-running query, distinct
+         * from {@link #cancelNoOp()}. Without a dedicated event type, the only
+         * trace of a cancel was the killed query's own thread logging
+         * {@code pg_terminate_backend}'s error as an ordinary EDITOR_QUERY_FAILED
+         * — indistinguishable from any other failure, and absent entirely when
+         * the target had already finished before the kill reached it.
+         */
+        public static AuditRecord cancelled(String sessionPid) {
+            AuditRecord r = new AuditRecord();
+            r.eventType = SecurityEventType.EDITOR_QUERY_CANCELLED;
+            r.outcome = SecurityEventOutcome.SUCCESS;
+            r.operation = "cancel";
+            r.failureReason = sessionPid == null ? null : "terminated session pid " + sessionPid;
+            return r;
+        }
+
+        /** Cancel requested for an execution id that was already finished or unknown. */
+        public static AuditRecord cancelNoOp() {
+            AuditRecord r = new AuditRecord();
+            r.eventType = SecurityEventType.EDITOR_QUERY_CANCELLED;
+            r.outcome = SecurityEventOutcome.INFO;
+            r.operation = "cancel";
+            r.failureReason = "query was no longer running";
+            return r;
+        }
+
         // ── fluent setters ────────────────────────────────────────────────
 
         public AuditRecord operation(String op) { this.operation = op; return this; }
@@ -225,6 +256,7 @@ public static AuditRecord failed(String reason) {
         public AuditRecord useAnalyze(Boolean v) { this.useAnalyze = v; return this; }
         public AuditRecord httpRequest(HttpServletRequest r) { this.httpRequest = r; return this; }
         public AuditRecord client(ClientContext c) { this.client = c; return this; }
+        public AuditRecord executionId(String id) { this.executionId = id; return this; }
         public AuditRecord eventType(SecurityEventType t) { this.eventType = t; return this; }
         public AuditRecord outcome(SecurityEventOutcome o) { this.outcome = o; return this; }
     }

backend/src/test/java/com/dbaagent/service/QueryExecutionPolicyServiceTest.java

@@ -112,6 +112,42 @@ void editorConfirmedDeleteWithoutWhere_isBlocked() {
         assertThat(exception.getMessage()).contains("without a WHERE clause");
     }
 
+    // EXPLAIN UPDATE/DELETE is the one reachable path to the keyword-fallback
+    // containsWhereClause check for a real UPDATE/DELETE: JSqlParser only models
+    // EXPLAIN SELECT, so these fall through to detectExplainWrappedMutation
+    // rather than Update.getWhere()/Delete.getWhere(), which every ordinary
+    // (non-EXPLAIN) mutation test above exercises instead.
+    @Test
+    void editorConfirmedExplainUpdateWithMultilineWhere_isAllowed() {
+        QueryExecutionPolicyService.PolicyDecision decision = service.enforce(
+            new QueryRequest(
+                "EXPLAIN UPDATE customers\nSET property_status = 'ACTIVE'\nWHERE customer_id = 9",
+                null,
+                null
+            ),
+            QueryExecutionContext.editor("admin", true, true),
+            "mysql"
+        );
+
+        assertThat(decision.mutating()).isTrue();
+        assertThat(decision.primaryQueryType()).isEqualTo("UPDATE");
+    }
+
+    @Test
+    void editorConfirmedExplainDeleteWithOnlyCommentedOutWhere_isBlocked() {
+        QueryExecutionPolicyException exception = assertThrows(
+            QueryExecutionPolicyException.class,
+            () -> service.enforce(
+                new QueryRequest("EXPLAIN DELETE FROM customers -- WHERE customer_id = 9\n", null, null),
+                QueryExecutionContext.editor("admin", true, true),
+                "mysql"
+            )
+        );
+
+        assertThat(exception.getErrorCode()).isEqualTo(QueryExecutionPolicyException.UNSAFE_MUTATION_BLOCKED);
+        assertThat(exception.getMessage()).contains("without a WHERE clause");
+    }
+
     @Test
     void editorMutation_multiStatementBatchIsBlocked() {
         QueryExecutionPolicyException exception = assertThrows(

backend/src/test/java/com/dbaagent/service/SqlExecutionAuditServiceTest.java

@@ -128,6 +128,46 @@ void record_failed_capturesFailureReason() {
         assertThat(event.metadata()).containsEntry("clientType", "unknown");
     }
 
+    @Test
+    void record_cancelled_logsSuccessWithSessionPidAndExecutionId() {
+        // Before this event type existed, a successful cancel was only visible
+        // as an accidental EDITOR_QUERY_FAILED thrown by the killed query's own
+        // thread when pg_terminate_backend severed it — indistinguishable from
+        // an ordinary query failure, and with no executionId to tie it back to
+        // the run the user actually meant to stop.
+        givenActor("alice@example.com", 1L);
+
+        audit.record(SqlExecutionAuditService.AuditRecord.cancelled("8642")
+            .connectionId("conn-1")
+            .executionId("exec-abc-123")
+            .client(ClientContext.unknown()));
+
+        SecurityEventService.EventRequest event = captureLoggedEvent();
+        assertThat(event.eventType()).isEqualTo(SecurityEventType.EDITOR_QUERY_CANCELLED);
+        assertThat(event.outcome()).isEqualTo(SecurityEventOutcome.SUCCESS);
+        assertThat(event.reason()).contains("8642");
+        assertThat(event.metadata()).containsEntry("operation", "cancel");
+        assertThat(event.metadata()).containsEntry("executionId", "exec-abc-123");
+    }
+
+    @Test
+    void record_cancelNoOp_logsInfoOutcomeNotSilence() {
+        // A cancel that misses its target (already finished, or a guessed/
+        // replayed executionId) must still leave a trace — previously it left
+        // none at all, so "did anyone try to cancel this?" was unanswerable.
+        givenActor("alice@example.com", 1L);
+
+        audit.record(SqlExecutionAuditService.AuditRecord.cancelNoOp()
+            .connectionId("conn-1")
+            .executionId("exec-already-done")
+            .client(ClientContext.unknown()));
+
+        SecurityEventService.EventRequest event = captureLoggedEvent();
+        assertThat(event.eventType()).isEqualTo(SecurityEventType.EDITOR_QUERY_CANCELLED);
+        assertThat(event.outcome()).isEqualTo(SecurityEventOutcome.INFO);
+        assertThat(event.metadata()).containsEntry("executionId", "exec-already-done");
+    }
+
     @Test
     void record_analyzePlan_carriesPlanSignalsInsteadOfRowCount() {
         givenActor("alice@example.com", 1L);

src/components/tabs/Core/SqlRunnerTab.js

@@ -54,6 +54,15 @@ import {
 // query fails with a real message rather than an opaque 504.
 const QUERY_TIMEOUT_SECONDS = 240;
 
+// CSV export re-runs the query with no display cap, so it needs its own bound.
+// Without one, an unbounded SELECT streams every row into a JS array, a CSV
+// string, and a Blob, and the equivalent unbounded read on the backend loads
+// the whole result set into memory before responding — large enough result
+// sets can exhaust the backend heap for every tenant on that instance, not
+// just the exporter. 100k rows is generous for a CSV download while staying
+// well short of that failure mode.
+const EXPORT_ROW_LIMIT = 100000;
+
 // Constants for diagram layout
 const DIAGRAM_NODE_WIDTH = 240;
 const DIAGRAM_NODE_HEIGHT = 120;
@@ -298,6 +307,12 @@ export default function SqlRunnerTab({ connectionId }) {
   const dbObjectsRef = useRef([]);
   const abortControllerRef = useRef(null);
   const executionIdRef = useRef(null);
+  // Guards against a second run starting while one is already in flight (e.g. a
+  // held-down Cmd+Enter): without it, whichever response lands last wins the
+  // results panel regardless of which run the user actually meant to see last,
+  // and the earlier run's abort/cancel handle gets silently discarded.
+  const isRunningRef = useRef(false);
+  const runSeqRef = useRef(0);
   // 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;
@@ -1194,6 +1209,9 @@ export default function SqlRunnerTab({ connectionId }) {
   };
 
   const handleRunQuery = async (queryToRun = null, options = {}) => {
+    if (isRunningRef.current) {
+      return;
+    }
     const mutationConfirmed = options.mutationConfirmed === true;
     let queryText = null;
 
@@ -1252,12 +1270,14 @@ export default function SqlRunnerTab({ connectionId }) {
       return;
     }
 
+    isRunningRef.current = true;
     setIsRunning(true);
     setError(null);
     setResults(null);
     setExplainResults(null); // Clear explain results when running query
     setOptimizeResult(null);
 
+    const seq = ++runSeqRef.current;
     const abortController = new AbortController();
     abortControllerRef.current = abortController;
     // Identifies this run so cancelling can terminate it on the database.
@@ -1284,6 +1304,12 @@ export default function SqlRunnerTab({ connectionId }) {
         },
       );
 
+      if (seq !== runSeqRef.current) {
+        // Superseded by a later run — drop this response rather than let a
+        // stale result overwrite what the user is now looking at.
+        return;
+      }
+
       if (response.success) {
         setPendingMutationConfirmation(null);
         const DISPLAY_LIMIT = 1000;
@@ -1355,9 +1381,12 @@ export default function SqlRunnerTab({ connectionId }) {
         setError(err.message || "Failed to execute query");
       }
     } finally {
-      abortControllerRef.current = null;
-      executionIdRef.current = null;
-      setIsRunning(false);
+      if (seq === runSeqRef.current) {
+        abortControllerRef.current = null;
+        executionIdRef.current = null;
+        isRunningRef.current = false;
+        setIsRunning(false);
+      }
     }
   };
 
@@ -1366,6 +1395,7 @@ export default function SqlRunnerTab({ connectionId }) {
     if (abortControllerRef.current) {
       abortControllerRef.current.abort();
     }
+    isRunningRef.current = false;
     setIsRunning(false);
     setError(null);
     // Aborting above only drops the HTTP response; the statement keeps running
@@ -1406,23 +1436,37 @@ export default function SqlRunnerTab({ connectionId }) {
 
   const handleExportResults = async () => {
     if (!results) return;
-
-    // If the result was limited, always re-fetch the full result set for download.
-    // This ensures the CSV contains all rows, not just the displayed page.
+    // Export re-runs the query, so it must respect the same in-flight guard as
+    // Run. Today `handleRunQuery` clears `results`, which hides the Export
+    // button for the duration of a run and makes this unreachable — but that is
+    // an incidental consequence of unrelated state handling, not a guarantee.
+    // Without this check, any change that keeps results on screen during a run
+    // silently reintroduces two concurrent queries from one tab.
+    if (isRunningRef.current || isExporting) return;
+
+    // If the result was limited, always re-fetch for download so the CSV isn't
+    // just the displayed page — but still capped at EXPORT_ROW_LIMIT, not
+    // unbounded, and on the same timeout budget as Run so it fails with a real
+    // error instead of an opaque 504 from the nginx proxy.
     if (results.isLimited && results.query) {
       setIsExporting(true);
       // Strip trailing semicolon so backends don't reject the re-executed query
       const queryForExport = results.query.trim().replace(/;+$/, "");
+      const exportAbortController = new AbortController();
+      const exportExecutionId =
+        globalThis.crypto?.randomUUID?.() ??
+        `exec-${Date.now()}-${Math.random().toString(16).slice(2)}`;
       try {
         const response = await queryAPI.executeQuery(
           connectionId,
           queryForExport,
-          null,  // no limit — fetch all rows
-          600,
-          null,
+          EXPORT_ROW_LIMIT,
+          QUERY_TIMEOUT_SECONDS,
+          exportAbortController.signal,
           {
             executionOrigin: "EDITOR",
             mutationConfirmed: false,
+            executionId: exportExecutionId,
           },
         );
         if (response.success) {
@@ -2590,9 +2634,9 @@ export default function SqlRunnerTab({ connectionId }) {
                         disabled={isExporting}
                         title={
                           results?.isLimited
-                            ? results?.totalRowCount != null
-                              ? `Download all ${results.totalRowCount.toLocaleString()} rows as CSV`
-                              : "Download full result set as CSV"
+                            ? results?.totalRowCount != null && results.totalRowCount > EXPORT_ROW_LIMIT
+                              ? `Download first ${EXPORT_ROW_LIMIT.toLocaleString()} of ${results.totalRowCount.toLocaleString()} rows as CSV`
+                              : `Download up to ${EXPORT_ROW_LIMIT.toLocaleString()} rows as CSV`
                             : "Export CSV"
                         }
                       >