Skip to content

Commit 94cb3b1

Browse files
cursoragentvenkateshsakamuri-lab
cursoragent
and
venkateshsakamuri-lab
committed
fix: hide out-of-scope schemas in the SQL editor explorer
Filter /objects, /schema, visualization, and table index/stats APIs by the user's allowedSchemas so Chat + Editor users only see schemas their policy permits. Schema object cache is now keyed per username to avoid leaking an admin's full tree. Co-authored-by: Venkat SF <venkatesh.sakamuri@stayflexi.com>
1 parent 07772b4 commit 94cb3b1

6 files changed

Lines changed: 257 additions & 17 deletions

File tree

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

Lines changed: 49 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import com.dbaagent.service.QueryExecutorService;
1010
import com.dbaagent.service.SqlExecutionAuditService;
1111
import com.dbaagent.service.UserDataAccessPolicyException;
12+
import com.dbaagent.service.UserDataAccessPolicyService;
1213
import com.dbaagent.service.SchemaScannerService;
1314
import com.dbaagent.service.VisualizationService;
1415
import com.dbaagent.service.security.AccessControlService;
@@ -36,6 +37,7 @@ public class SchemaController {
3637
private final QueryExecutorService queryExecutorService;
3738
private final AccessControlService accessControlService;
3839
private final SqlExecutionAuditService sqlExecutionAuditService;
40+
private final UserDataAccessPolicyService userDataAccessPolicyService;
3941

4042
@PostMapping("/scan")
4143
public ResponseEntity<Map<String, Object>> scanSchema(@PathVariable String connectionId) {
@@ -48,7 +50,7 @@ public ResponseEntity<Map<String, Object>> scanSchema(@PathVariable String conne
4850
}
4951
accessControlService.assertCanUseChatEditor(connectionId);
5052

51-
SchemaMetadata schema = schemaScannerService.scanSchema(connectionId);
53+
SchemaMetadata schema = scopedSchema(connectionId, schemaScannerService.scanSchema(connectionId));
5254
response.put("success", true);
5355
response.put("schema", schema);
5456
return ResponseEntity.ok(response);
@@ -78,7 +80,7 @@ public ResponseEntity<Map<String, Object>> getSchema(@PathVariable String connec
7880
}
7981
accessControlService.assertCanUseChatEditor(connectionId);
8082

81-
SchemaMetadata schema = schemaScannerService.scanSchema(connectionId);
83+
SchemaMetadata schema = scopedSchema(connectionId, schemaScannerService.scanSchema(connectionId));
8284
response.put("success", true);
8385
response.put("schema", schema);
8486
return ResponseEntity.ok(response);
@@ -108,7 +110,7 @@ public ResponseEntity<Map<String, Object>> getVisualization(@PathVariable String
108110
}
109111
accessControlService.assertCanUseChatEditor(connectionId);
110112

111-
SchemaMetadata schema = schemaScannerService.scanSchema(connectionId);
113+
SchemaMetadata schema = scopedSchema(connectionId, schemaScannerService.scanSchema(connectionId));
112114
ErDiagramData erDiagram = visualizationService.generateErDiagram(schema, connectionId);
113115
DependencyGraphData dependencyGraph = visualizationService.generateDependencyGraph(schema, connectionId);
114116

@@ -143,7 +145,10 @@ public ResponseEntity<Map<String, Object>> getDatabaseObjects(@PathVariable Stri
143145
}
144146
accessControlService.assertCanUseChatEditor(connectionId);
145147

146-
List<DatabaseObject> objects = queryExecutorService.getDatabaseObjects(connectionId);
148+
List<DatabaseObject> objects = scopedObjects(
149+
connectionId,
150+
queryExecutorService.getDatabaseObjects(connectionId)
151+
);
147152
log.info("Successfully fetched {} database objects for connection: {}", objects.size(), connectionId);
148153
response.put("success", true);
149154
response.put("objects", objects);
@@ -273,6 +278,12 @@ public ResponseEntity<Map<String, Object>> getTableIndexes(
273278
return ResponseEntity.status(HttpStatus.NOT_FOUND).body(response);
274279
}
275280
accessControlService.assertCanUseChatEditor(connectionId);
281+
userDataAccessPolicyService.assertTableSchemaAllowed(
282+
connectionId,
283+
accessControlService.getCurrentUsername(),
284+
accessControlService.isCurrentUserAdmin(),
285+
tableName
286+
);
276287

277288
List<TableIndex> indexes = queryExecutorService.getTableIndexes(connectionId, tableName);
278289
response.put("success", true);
@@ -282,6 +293,11 @@ public ResponseEntity<Map<String, Object>> getTableIndexes(
282293
response.put("success", false);
283294
response.put("message", "Failed to fetch table indexes: " + e.getMessage());
284295
return ResponseEntity.status(HttpStatus.INTERNAL_SERVER_ERROR).body(response);
296+
} catch (UserDataAccessPolicyException e) {
297+
response.put("success", false);
298+
response.put("message", e.getMessage());
299+
response.put("errorCode", e.getErrorCode());
300+
return ResponseEntity.status(HttpStatus.FORBIDDEN).body(response);
285301
} catch (ResponseStatusException e) {
286302
response.put("success", false);
287303
response.put("message", e.getReason());
@@ -305,6 +321,12 @@ public ResponseEntity<Map<String, Object>> getTableStats(
305321
return ResponseEntity.status(HttpStatus.NOT_FOUND).body(response);
306322
}
307323
accessControlService.assertCanUseChatEditor(connectionId);
324+
userDataAccessPolicyService.assertTableSchemaAllowed(
325+
connectionId,
326+
accessControlService.getCurrentUsername(),
327+
accessControlService.isCurrentUserAdmin(),
328+
tableName
329+
);
308330

309331
TableStats stats = queryExecutorService.getTableStats(connectionId, tableName);
310332
response.put("success", true);
@@ -314,6 +336,11 @@ public ResponseEntity<Map<String, Object>> getTableStats(
314336
response.put("success", false);
315337
response.put("message", "Failed to fetch table stats: " + e.getMessage());
316338
return ResponseEntity.status(HttpStatus.INTERNAL_SERVER_ERROR).body(response);
339+
} catch (UserDataAccessPolicyException e) {
340+
response.put("success", false);
341+
response.put("message", e.getMessage());
342+
response.put("errorCode", e.getErrorCode());
343+
return ResponseEntity.status(HttpStatus.FORBIDDEN).body(response);
317344
} catch (ResponseStatusException e) {
318345
response.put("success", false);
319346
response.put("message", e.getReason());
@@ -324,4 +351,22 @@ public ResponseEntity<Map<String, Object>> getTableStats(
324351
return ResponseEntity.status(HttpStatus.INTERNAL_SERVER_ERROR).body(response);
325352
}
326353
}
354+
355+
private SchemaMetadata scopedSchema(String connectionId, SchemaMetadata schema) {
356+
return userDataAccessPolicyService.filterSchemaMetadata(
357+
connectionId,
358+
accessControlService.getCurrentUsername(),
359+
accessControlService.isCurrentUserAdmin(),
360+
schema
361+
);
362+
}
363+
364+
private List<DatabaseObject> scopedObjects(String connectionId, List<DatabaseObject> objects) {
365+
return userDataAccessPolicyService.filterDatabaseObjects(
366+
connectionId,
367+
accessControlService.getCurrentUsername(),
368+
accessControlService.isCurrentUserAdmin(),
369+
objects
370+
);
371+
}
327372
}

‎backend/src/main/java/com/dbaagent/service/ConnectionChatAccessPolicyService.java‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -600,10 +600,18 @@ private Set<String> extractAllowedSchemas(String normalized, SchemaMetadata sche
600600
}
601601

602602
private boolean schemaInScope(String schema, Set<String> allowedSchemas) {
603+
return isSchemaInScope(schema, allowedSchemas);
604+
}
605+
606+
public static boolean isSchemaInScope(String schema, Set<String> allowedSchemas) {
603607
if (allowedSchemas == null || allowedSchemas.isEmpty()) {
604608
return true;
605609
}
606-
return allowedSchemas.contains(normalizeName(schema));
610+
String normalized = schema == null ? "" : schema.trim().replace("\"", "").replace("`", "").toLowerCase(Locale.ROOT);
611+
if (normalized.isBlank()) {
612+
normalized = "public";
613+
}
614+
return allowedSchemas.contains(normalized);
607615
}
608616

609617
private SchemaMetadata tryScanSchema(String connectionId) {

‎backend/src/main/java/com/dbaagent/service/UserDataAccessPolicyService.java‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,9 @@
22

33
import com.dbaagent.model.QueryRequest;
44
import com.dbaagent.model.QueryResult;
5+
import com.dbaagent.model.DatabaseObject;
6+
import com.dbaagent.model.SchemaMetadata;
7+
import com.dbaagent.model.TableMetadata;
58
import com.dbaagent.model.SecurityEventOutcome;
69
import com.dbaagent.model.SecurityEventType;
710
import lombok.RequiredArgsConstructor;
@@ -148,6 +151,96 @@ public QueryGuardDecision enforcePreExecution(
148151
return QueryGuardDecision.allow(policy);
149152
}
150153

154+
public List<DatabaseObject> filterDatabaseObjects(
155+
String connectionId,
156+
String username,
157+
boolean actorIsAdmin,
158+
List<DatabaseObject> objects
159+
) {
160+
if (objects == null || objects.isEmpty()) {
161+
return objects;
162+
}
163+
Set<String> allowedSchemas = allowedSchemasForActor(connectionId, username, actorIsAdmin);
164+
if (allowedSchemas.isEmpty()) {
165+
return objects;
166+
}
167+
return objects.stream()
168+
.filter(object -> ConnectionChatAccessPolicyService.isSchemaInScope(object.getSchema(), allowedSchemas))
169+
.toList();
170+
}
171+
172+
public SchemaMetadata filterSchemaMetadata(
173+
String connectionId,
174+
String username,
175+
boolean actorIsAdmin,
176+
SchemaMetadata schema
177+
) {
178+
if (schema == null) {
179+
return null;
180+
}
181+
Set<String> allowedSchemas = allowedSchemasForActor(connectionId, username, actorIsAdmin);
182+
if (allowedSchemas.isEmpty()) {
183+
return schema;
184+
}
185+
186+
SchemaMetadata filtered = new SchemaMetadata();
187+
filtered.setDatabaseName(schema.getDatabaseName());
188+
filtered.setDbType(schema.getDbType());
189+
filtered.setTotalViews(schema.getTotalViews());
190+
filtered.setTotalSizeBytes(schema.getTotalSizeBytes());
191+
List<TableMetadata> tables = schema.getTables() == null ? List.of() : schema.getTables().stream()
192+
.filter(table -> ConnectionChatAccessPolicyService.isSchemaInScope(table.getSchema(), allowedSchemas))
193+
.toList();
194+
filtered.setTables(tables);
195+
filtered.setTotalTables((long) tables.size());
196+
if (schema.getRelationships() != null) {
197+
filtered.setRelationships(schema.getRelationships().stream()
198+
.filter(relationship ->
199+
ConnectionChatAccessPolicyService.isSchemaInScope(schemaFromTableRef(relationship.getFromTable()), allowedSchemas)
200+
&& ConnectionChatAccessPolicyService.isSchemaInScope(schemaFromTableRef(relationship.getToTable()), allowedSchemas))
201+
.toList());
202+
}
203+
return filtered;
204+
}
205+
206+
public void assertTableSchemaAllowed(
207+
String connectionId,
208+
String username,
209+
boolean actorIsAdmin,
210+
String tableRef
211+
) {
212+
Set<String> allowedSchemas = allowedSchemasForActor(connectionId, username, actorIsAdmin);
213+
if (allowedSchemas.isEmpty()) {
214+
return;
215+
}
216+
String schema = schemaFromTableRef(tableRef);
217+
if (!ConnectionChatAccessPolicyService.isSchemaInScope(schema, allowedSchemas)) {
218+
throw new UserDataAccessPolicyException(
219+
"This object is in schema '" + (schema.isBlank() ? "public" : schema)
220+
+ "' which is outside your allowed schema scope.",
221+
"POLICY_SCHEMA_BLOCKED"
222+
);
223+
}
224+
}
225+
226+
private Set<String> allowedSchemasForActor(String connectionId, String username, boolean actorIsAdmin) {
227+
ConnectionChatAccessPolicyService.EffectivePolicy policy =
228+
policyService.resolveEffectivePolicy(connectionId, username, actorIsAdmin);
229+
if (policy == null || policy.allowedSchemas() == null) {
230+
return Set.of();
231+
}
232+
return policy.allowedSchemas();
233+
}
234+
235+
private String schemaFromTableRef(String tableRef) {
236+
if (tableRef == null || tableRef.isBlank()) {
237+
return "";
238+
}
239+
String normalized = normalizeName(tableRef);
240+
int separator = normalized.lastIndexOf('.');
241+
return separator > 0 ? normalized.substring(0, separator) : "";
242+
}
243+
151244
public QueryResult redactResult(
152245
String connectionId,
153246
QueryResult result,

‎backend/src/test/java/com/dbaagent/service/UserDataAccessPolicyServiceTest.java‎

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -169,6 +169,96 @@ void enforcePreExecution_allowsMartsQueriesWhenOtherSchemasHaveProtectedColumns(
169169
).policy().allowedSchemas()).containsExactly("marts");
170170
}
171171

172+
@Test
173+
void filterDatabaseObjects_keepsOnlyAllowedSchemas() {
174+
ConnectionChatAccessPolicyService.EffectivePolicy schemaPolicy = new ConnectionChatAccessPolicyService.EffectivePolicy(
175+
true,
176+
"conn-1",
177+
"analyst",
178+
Set.of(),
179+
Set.of(),
180+
Set.of(),
181+
Set.of("marts"),
182+
true,
183+
true,
184+
"Only schema marts",
185+
List.of(),
186+
List.of()
187+
);
188+
when(policyService.resolveEffectivePolicy("conn-1", "analyst", false)).thenReturn(schemaPolicy);
189+
190+
var marts = new com.dbaagent.model.DatabaseObject("fct_enrollment", "table", "marts", List.of(), 2L, null);
191+
var crm = new com.dbaagent.model.DatabaseObject("customers", "table", "crm", List.of(), 2L, null);
192+
var sales = new com.dbaagent.model.DatabaseObject("orders", "table", "sales", List.of(), 2L, null);
193+
194+
assertThat(service.filterDatabaseObjects("conn-1", "analyst", false, List.of(marts, crm, sales)))
195+
.extracting(com.dbaagent.model.DatabaseObject::getName)
196+
.containsExactly("fct_enrollment");
197+
}
198+
199+
@Test
200+
void filterSchemaMetadata_dropsOutOfScopeTablesAndRelationships() {
201+
ConnectionChatAccessPolicyService.EffectivePolicy schemaPolicy = new ConnectionChatAccessPolicyService.EffectivePolicy(
202+
true,
203+
"conn-1",
204+
"analyst",
205+
Set.of(),
206+
Set.of(),
207+
Set.of(),
208+
Set.of("marts"),
209+
true,
210+
true,
211+
"Only schema marts",
212+
List.of(),
213+
List.of()
214+
);
215+
when(policyService.resolveEffectivePolicy("conn-1", "analyst", false)).thenReturn(schemaPolicy);
216+
217+
var schema = new com.dbaagent.model.SchemaMetadata();
218+
var marts = new com.dbaagent.model.TableMetadata();
219+
marts.setSchema("marts");
220+
marts.setName("fct_enrollment");
221+
var crm = new com.dbaagent.model.TableMetadata();
222+
crm.setSchema("crm");
223+
crm.setName("customers");
224+
schema.setTables(List.of(marts, crm));
225+
var relationship = new com.dbaagent.model.RelationshipMetadata();
226+
relationship.setFromTable("crm.customers");
227+
relationship.setToTable("marts.fct_enrollment");
228+
schema.setRelationships(List.of(relationship));
229+
230+
var filtered = service.filterSchemaMetadata("conn-1", "analyst", false, schema);
231+
assertThat(filtered.getTables()).extracting(com.dbaagent.model.TableMetadata::getName)
232+
.containsExactly("fct_enrollment");
233+
assertThat(filtered.getRelationships()).isEmpty();
234+
}
235+
236+
@Test
237+
void assertTableSchemaAllowed_blocksOutOfScopeTableMetadata() {
238+
ConnectionChatAccessPolicyService.EffectivePolicy schemaPolicy = new ConnectionChatAccessPolicyService.EffectivePolicy(
239+
true,
240+
"conn-1",
241+
"analyst",
242+
Set.of(),
243+
Set.of(),
244+
Set.of(),
245+
Set.of("marts"),
246+
true,
247+
true,
248+
"Only schema marts",
249+
List.of(),
250+
List.of()
251+
);
252+
when(policyService.resolveEffectivePolicy("conn-1", "analyst", false)).thenReturn(schemaPolicy);
253+
254+
UserDataAccessPolicyException exception = assertThrows(
255+
UserDataAccessPolicyException.class,
256+
() -> service.assertTableSchemaAllowed("conn-1", "analyst", false, "crm.customers")
257+
);
258+
assertThat(exception.getErrorCode()).isEqualTo("POLICY_SCHEMA_BLOCKED");
259+
service.assertTableSchemaAllowed("conn-1", "analyst", false, "marts.fct_enrollment");
260+
}
261+
172262
private ConnectionChatAccessPolicyService.EffectivePolicy policy() {
173263
return new ConnectionChatAccessPolicyService.EffectivePolicy(
174264
true,

0 commit comments

Comments
 (0)