|
10 | 10 | import com.dbaagent.service.ClientContext; |
11 | 11 | import com.dbaagent.service.CredentialService; |
12 | 12 | import com.dbaagent.service.ExplainPlanService; |
| 13 | +import com.dbaagent.service.McpTokenService; |
13 | 14 | import com.dbaagent.service.QueryExecutionContext; |
14 | 15 | import com.dbaagent.service.QueryExecutionPolicyException; |
15 | 16 | import com.dbaagent.service.QueryExecutionPolicyService; |
|
19 | 20 | import lombok.Data; |
20 | 21 | import lombok.RequiredArgsConstructor; |
21 | 22 | import lombok.extern.slf4j.Slf4j; |
| 23 | +import org.springframework.http.HttpHeaders; |
22 | 24 | import org.springframework.http.HttpStatus; |
23 | 25 | import org.springframework.http.ResponseEntity; |
24 | 26 | import org.springframework.jdbc.BadSqlGrammarException; |
|
38 | 40 | * 1. `useAnalyze=true` actually runs the underlying statement inside |
39 | 41 | * EXPLAIN ANALYZE — so for any mutating statement, that's a real |
40 | 42 | * database write. We route those through QueryExecutionPolicyService |
41 | | - * with `QueryExecutionContext.editor(...)` so the same role/WHERE/ |
42 | | - * confirmation gates that protect /api/connections/{id}/query |
43 | | - * protect this path too. |
| 43 | + * with `QueryExecutionContext.forSqlSurface(...)` so MCP bearers |
| 44 | + * keep the MCP DROP/TRUNCATE block and Editor JWT callers keep the |
| 45 | + * Editor DROP TABLE block. Role, WHERE, and confirmation gates that |
| 46 | + * protect /api/connections/{id}/query protect this path too. |
44 | 47 | * |
45 | 48 | * 2. Every call — success, blocked, or failed — emits a SecurityEvent so |
46 | 49 | * audit dashboards can see CLI/MCP/Editor traffic with one filter. |
@@ -71,23 +74,27 @@ public ResponseEntity<?> analyzeQuery( |
71 | 74 | ) { |
72 | 75 | ClientContext client = ClientContext.fromRequest(httpRequest); |
73 | 76 | String connectionId = request.getConnectionId(); |
74 | | - QueryRequest auditQueryRequest = buildAuditQueryRequest(request); |
| 77 | + QueryRequest auditQueryRequest = buildAuditQueryRequest(request, httpRequest); |
75 | 78 | ConnectionRequest connectionRequest = null; |
76 | 79 |
|
77 | 80 | try { |
78 | 81 | accessControlService.assertCanUseChatEditor(connectionId); |
79 | 82 | log.info("EXPLAIN analysis requested for connection: {} (useAnalyze={})", connectionId, request.isUseAnalyze()); |
80 | 83 |
|
81 | 84 | // ANALYZE actually executes the SQL. Route the underlying |
82 | | - // statement through the same policy gate the SQL Editor uses so |
83 | | - // a developer can't bypass the mutation guard by sending |
84 | | - // useAnalyze=true with `DELETE FROM users`. |
| 85 | + // statement through the same policy gate /connections/{id}/query |
| 86 | + // uses so a developer can't bypass the mutation guard by sending |
| 87 | + // useAnalyze=true with `DELETE FROM users`, and MCP callers keep |
| 88 | + // the DROP/TRUNCATE block. |
85 | 89 | if (request.isUseAnalyze()) { |
86 | 90 | connectionRequest = credentialService.getDecryptedConnection(connectionId); |
87 | 91 | String dbType = providerRegistry.getCanonicalName(connectionRequest.getDbType()); |
88 | 92 | queryExecutionPolicyService.enforce( |
89 | 93 | auditQueryRequest, |
90 | | - QueryExecutionContext.editor( |
| 94 | + QueryExecutionContext.forSqlSurface( |
| 95 | + McpTokenService.isMcpAuthorizationHeader( |
| 96 | + httpRequest.getHeader(HttpHeaders.AUTHORIZATION) |
| 97 | + ), |
91 | 98 | accessControlService.getCurrentUsername(), |
92 | 99 | accessControlService.isCurrentUserAdmin(), |
93 | 100 | Boolean.TRUE.equals(request.getMutationConfirmed()) |
@@ -161,10 +168,13 @@ public ResponseEntity<?> analyzeQuery( |
161 | 168 | * Carries the user's mutationConfirmed flag through so admin-confirmed |
162 | 169 | * ANALYZE runs aren't stuck on the confirmation gate. |
163 | 170 | */ |
164 | | - private QueryRequest buildAuditQueryRequest(ExplainRequest request) { |
| 171 | + private QueryRequest buildAuditQueryRequest(ExplainRequest request, HttpServletRequest httpRequest) { |
165 | 172 | QueryRequest qr = new QueryRequest(); |
166 | 173 | qr.setQuery(request.getQuery()); |
167 | | - qr.setExecutionOrigin(QueryExecutionOrigin.EDITOR); |
| 174 | + boolean mcpBearer = McpTokenService.isMcpAuthorizationHeader( |
| 175 | + httpRequest.getHeader(HttpHeaders.AUTHORIZATION) |
| 176 | + ); |
| 177 | + qr.setExecutionOrigin(mcpBearer ? QueryExecutionOrigin.MCP : QueryExecutionOrigin.EDITOR); |
168 | 178 | qr.setMutationConfirmed(Boolean.TRUE.equals(request.getMutationConfirmed())); |
169 | 179 | return qr; |
170 | 180 | } |
|
0 commit comments