Copilot commented on code in PR #10546:
URL: https://github.com/apache/gravitino/pull/10546#discussion_r2987835778
##########
server-common/src/test/java/org/apache/gravitino/server/authorization/jcasbin/TestJcasbinAuthorizer.java:
##########
@@ -106,12 +115,33 @@ public class TestJcasbinAuthorizer {
private static MockedStatic<OwnerMetaService> ownerMetaServiceMockedStatic;
+ private static MockedStatic<RoleMetaService> roleMetaServiceMockedStatic;
+
private static JcasbinAuthorizer jcasbinAuthorizer;
private static ObjectMapper objectMapper = new ObjectMapper();
@BeforeAll
public static void setup() throws IOException {
+ // Pre-populate known constant roles so tests don't need to set them up
individually.
+ roleSecurableObjectsMap.put(ALLOW_ROLE_ID,
ImmutableList.of(getAllowSecurableObject()));
+ roleSecurableObjectsMap.put(DENY_ROLE_ID,
ImmutableList.of(getDenySecurableObject()));
+
+ // Mock RoleMetaService.batchListSecurableObjectsForRoles to avoid real DB
access.
+ roleMetaServiceMockedStatic = mockStatic(RoleMetaService.class);
+ roleMetaServiceMockedStatic
+ .when(() -> RoleMetaService.batchListSecurableObjectsForRoles(any()))
+ .thenAnswer(
+ inv -> {
+ List<Long> ids = inv.getArgument(0);
+ com.google.common.collect.ImmutableMap.Builder<Long,
List<SecurableObject>> result =
+ com.google.common.collect.ImmutableMap.builder();
+ for (Long id : ids) {
Review Comment:
Avoid using a fully-qualified class name inside the method body
(`com.google.common.collect.ImmutableMap.Builder`). Please add an import for
`ImmutableMap` and reference it directly to match the repo's Java import
convention.
##########
core/src/main/java/org/apache/gravitino/storage/relational/service/RoleMetaService.java:
##########
@@ -338,6 +339,37 @@ public static List<SecurableObjectPO>
listSecurableObjectsByRoleId(Long roleId)
SecurableObjectMapper.class, mapper ->
mapper.listSecurableObjectsByRoleId(roleId));
}
+ /**
+ * Batch-loads securable objects for multiple roles in a single SQL query
and returns a map from
+ * role ID to the resolved {@link SecurableObject} list. This eliminates the
N+1 query pattern
+ * that occurs when loading securable objects for each role individually.
+ *
+ * @param roleIds the list of role IDs to load
+ * @return a map from role ID to its list of resolved securable objects
+ */
+ @Monitored(
+ metricsSource = GRAVITINO_RELATIONAL_STORE_METRIC_NAME,
+ baseMetricName = "batchListSecurableObjectsForRoles")
+ public static Map<Long, List<SecurableObject>>
batchListSecurableObjectsForRoles(
+ List<Long> roleIds) {
+ if (roleIds.isEmpty()) {
+ return ImmutableMap.of();
+ }
+ List<SecurableObjectPO> allPOs =
+ SessionUtils.getWithoutCommit(
+ SecurableObjectMapper.class, mapper ->
mapper.listSecurableObjectsByRoleIds(roleIds));
+
+ Map<Long, List<SecurableObjectPO>> byRoleId =
+
allPOs.stream().collect(Collectors.groupingBy(SecurableObjectPO::getRoleId));
+
+ ImmutableMap.Builder<Long, List<SecurableObject>> builder =
ImmutableMap.builder();
+ for (Long roleId : roleIds) {
+ List<SecurableObjectPO> pos = byRoleId.getOrDefault(roleId,
Collections.emptyList());
+ builder.put(roleId, buildSecurableObjectsFromPOs(pos));
+ }
+ return builder.build();
+ }
Review Comment:
New behavior (`batchListSecurableObjectsForRoles`) + new mapper/provider
path (`listSecurableObjectsByRoleIds`) is not covered by any unit/integration
test in this repo (no tests reference the new method). Per project guidelines,
please add coverage in `TestRoleMetaService` to validate correctness (multiple
role IDs, missing roles returning empty lists) and that it works across
backends (H2/PostgreSQL).
##########
server-common/src/main/java/org/apache/gravitino/server/authorization/jcasbin/JcasbinAuthorizer.java:
##########
@@ -490,48 +493,56 @@ private void loadRolePrivilege(
() -> {
EntityStore entityStore = GravitinoEnv.getInstance().entityStore();
NameIdentifier userNameIdentifier =
NameIdentifierUtil.ofUser(metalake, username);
- List<RoleEntity> entities;
+ List<RoleEntity> roleStubs;
try {
- entities =
+ roleStubs =
entityStore
.relationOperations()
.listEntitiesByRelation(
SupportsRelationOperations.Type.ROLE_USER_REL,
userNameIdentifier,
Entity.EntityType.USER);
- List<CompletableFuture<Void>> loadRoleFutures = new ArrayList<>();
- for (RoleEntity role : entities) {
- Long roleId = role.id();
- allowEnforcer.addRoleForUser(String.valueOf(userId),
String.valueOf(roleId));
- denyEnforcer.addRoleForUser(String.valueOf(userId),
String.valueOf(roleId));
- if (loadedRoles.getIfPresent(roleId) != null) {
- continue;
- }
- CompletableFuture<Void> loadRoleFuture =
- CompletableFuture.supplyAsync(
- () -> {
- try {
- return entityStore.get(
- NameIdentifierUtil.ofRole(metalake,
role.name()),
- Entity.EntityType.ROLE,
- RoleEntity.class);
- } catch (Exception e) {
- throw new RuntimeException("Failed to load role:
" + role.name(), e);
- }
- },
- executor)
- .thenAcceptAsync(
- roleEntity -> {
- loadPolicyByRoleEntity(roleEntity);
- loadedRoles.put(roleId, true);
- },
- executor);
- loadRoleFutures.add(loadRoleFuture);
- }
- CompletableFuture.allOf(loadRoleFutures.toArray(new
CompletableFuture[0])).join();
} catch (IOException e) {
throw new RuntimeException(e);
}
+
+ // Register user-role associations in enforcers for all roles.
+ for (RoleEntity role : roleStubs) {
+ allowEnforcer.addRoleForUser(String.valueOf(userId),
String.valueOf(role.id()));
+ denyEnforcer.addRoleForUser(String.valueOf(userId),
String.valueOf(role.id()));
+ }
+
+ // Collect stubs for roles whose policies have not yet been loaded
into the enforcer.
+ List<RoleEntity> unloadedRoleStubs =
+ roleStubs.stream()
+ .filter(role -> loadedRoles.getIfPresent(role.id()) == null)
+ .collect(Collectors.toList());
+ if (unloadedRoleStubs.isEmpty()) {
+ return;
+ }
+
+ // Batch-fetch securable objects for all unloaded roles in a single
query,
+ // eliminating the N+1 pattern of per-role entityStore.get() calls.
+ List<Long> unloadedRoleIds =
+
unloadedRoleStubs.stream().map(RoleEntity::id).collect(Collectors.toList());
+ Map<Long, List<SecurableObject>> secObjsByRoleId =
+
RoleMetaService.batchListSecurableObjectsForRoles(unloadedRoleIds);
+
+ for (RoleEntity stub : unloadedRoleStubs) {
+ List<SecurableObject> securableObjects =
+ secObjsByRoleId.getOrDefault(stub.id(),
Collections.emptyList());
+ RoleEntity fullRole =
+ RoleEntity.builder()
+ .withId(stub.id())
+ .withName(stub.name())
+ .withNamespace(stub.namespace())
+ .withProperties(stub.properties())
+ .withAuditInfo(stub.auditInfo())
+ .withSecurableObjects(securableObjects)
+ .build();
+ loadPolicyByRoleEntity(fullRole);
+ loadedRoles.put(stub.id(), true);
+ }
Review Comment:
`loadRolePrivilege()` no longer uses the `executor` thread pool (the
previous `CompletableFuture` path was removed), but `initialize()` still
creates a fixed thread pool (default 100 threads) and `close()` shuts it down.
This now allocates idle threads for no benefit and increases memory/CPU
footprint. Consider removing the `executor` field + thread pool initialization
entirely (and any related test reflection), or reintroduce an async use of it
if still needed.
##########
server-common/src/main/java/org/apache/gravitino/server/authorization/jcasbin/JcasbinAuthorizer.java:
##########
@@ -490,48 +493,56 @@ private void loadRolePrivilege(
() -> {
EntityStore entityStore = GravitinoEnv.getInstance().entityStore();
NameIdentifier userNameIdentifier =
NameIdentifierUtil.ofUser(metalake, username);
- List<RoleEntity> entities;
+ List<RoleEntity> roleStubs;
try {
- entities =
+ roleStubs =
entityStore
.relationOperations()
.listEntitiesByRelation(
SupportsRelationOperations.Type.ROLE_USER_REL,
userNameIdentifier,
Entity.EntityType.USER);
- List<CompletableFuture<Void>> loadRoleFutures = new ArrayList<>();
- for (RoleEntity role : entities) {
- Long roleId = role.id();
- allowEnforcer.addRoleForUser(String.valueOf(userId),
String.valueOf(roleId));
- denyEnforcer.addRoleForUser(String.valueOf(userId),
String.valueOf(roleId));
- if (loadedRoles.getIfPresent(roleId) != null) {
- continue;
- }
- CompletableFuture<Void> loadRoleFuture =
- CompletableFuture.supplyAsync(
- () -> {
- try {
- return entityStore.get(
- NameIdentifierUtil.ofRole(metalake,
role.name()),
- Entity.EntityType.ROLE,
- RoleEntity.class);
- } catch (Exception e) {
- throw new RuntimeException("Failed to load role:
" + role.name(), e);
- }
- },
- executor)
- .thenAcceptAsync(
- roleEntity -> {
- loadPolicyByRoleEntity(roleEntity);
- loadedRoles.put(roleId, true);
- },
- executor);
- loadRoleFutures.add(loadRoleFuture);
- }
- CompletableFuture.allOf(loadRoleFutures.toArray(new
CompletableFuture[0])).join();
} catch (IOException e) {
throw new RuntimeException(e);
}
+
+ // Register user-role associations in enforcers for all roles.
+ for (RoleEntity role : roleStubs) {
+ allowEnforcer.addRoleForUser(String.valueOf(userId),
String.valueOf(role.id()));
+ denyEnforcer.addRoleForUser(String.valueOf(userId),
String.valueOf(role.id()));
+ }
+
+ // Collect stubs for roles whose policies have not yet been loaded
into the enforcer.
+ List<RoleEntity> unloadedRoleStubs =
+ roleStubs.stream()
+ .filter(role -> loadedRoles.getIfPresent(role.id()) == null)
+ .collect(Collectors.toList());
+ if (unloadedRoleStubs.isEmpty()) {
+ return;
+ }
+
+ // Batch-fetch securable objects for all unloaded roles in a single
query,
+ // eliminating the N+1 pattern of per-role entityStore.get() calls.
+ List<Long> unloadedRoleIds =
+
unloadedRoleStubs.stream().map(RoleEntity::id).collect(Collectors.toList());
+ Map<Long, List<SecurableObject>> secObjsByRoleId =
+
RoleMetaService.batchListSecurableObjectsForRoles(unloadedRoleIds);
Review Comment:
The new batch call to
`RoleMetaService.batchListSecurableObjectsForRoles(...)` can now fail as a
single operation, but any exception will be wrapped by
`AuthorizationRequestContext.loadRole()` as `RuntimeException("Failed to load
role: ", e)` without indicating which role(s) were being loaded. Previously the
per-role `entityStore.get(...)` path added role-name context in the thrown
exception. Consider catching failures around the batch call and rethrowing with
actionable context (e.g., include `metalake`, `username`, and the list of
`unloadedRoleIds`).
```suggestion
Map<Long, List<SecurableObject>> secObjsByRoleId;
try {
secObjsByRoleId =
RoleMetaService.batchListSecurableObjectsForRoles(unloadedRoleIds);
} catch (RuntimeException e) {
throw new RuntimeException(
"Failed to batch-load securable objects for roles " +
unloadedRoleIds
+ " of userId " + userId,
e);
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
