yuqi1129 commented on code in PR #13262: URL: https://github.com/apache/gravitino/pull/13262#discussion_r4036337951
########## server-common/src/test/java/org/apache/gravitino/server/authorization/TestPrincipalListQueryCount.java: ########## @@ -0,0 +1,262 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * http://www.apache.org/licenses/LICENSE-2.0 + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.gravitino.server.authorization; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyList; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.when; + +import java.nio.file.Path; +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.Statement; +import java.time.Instant; +import java.util.ArrayList; +import java.util.List; +import java.util.Locale; +import java.util.concurrent.Executor; +import org.apache.commons.lang3.reflect.FieldUtils; +import org.apache.gravitino.Config; +import org.apache.gravitino.Configs; +import org.apache.gravitino.Entity; +import org.apache.gravitino.EntityStore; +import org.apache.gravitino.GravitinoEnv; +import org.apache.gravitino.HasIdentifier; +import org.apache.gravitino.MetadataObject; +import org.apache.gravitino.NameIdentifier; +import org.apache.gravitino.Namespace; +import org.apache.gravitino.SupportsRelationOperations; +import org.apache.gravitino.UserPrincipal; +import org.apache.gravitino.authorization.AuthorizationUtils; +import org.apache.gravitino.authorization.GravitinoAuthorizer; +import org.apache.gravitino.authorization.Privilege; +import org.apache.gravitino.json.JsonUtils; +import org.apache.gravitino.meta.AuditInfo; +import org.apache.gravitino.meta.EntityIdResolver; +import org.apache.gravitino.server.authorization.expression.AuthorizationExpressionConstants; +import org.apache.gravitino.storage.relational.JDBCBackend; +import org.apache.gravitino.storage.relational.RelationalEntityStoreIdResolver; +import org.apache.gravitino.storage.relational.service.EntityIdService; +import org.apache.gravitino.utils.PrincipalUtils; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.function.Executable; +import org.junit.jupiter.api.io.TempDir; +import org.mockito.MockedStatic; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** Counts real database queries in list loading and filtering, independently of list size. */ +class TestPrincipalListQueryCount { + private static final Logger LOG = LoggerFactory.getLogger(TestPrincipalListQueryCount.class); + + @TempDir Path tempDir; + + @Test + void testManagementListsHaveBoundedQueries() throws Exception { + Config config = new Config(false) {}; + config.set( + Configs.ENTITY_RELATIONAL_JDBC_BACKEND_URL, + "jdbc:h2:file:" + tempDir.resolve("metadata") + ";MODE=MYSQL;AUTO_SERVER=FALSE"); + config.set(Configs.ENTITY_RELATIONAL_JDBC_BACKEND_DRIVER, "org.h2.Driver"); + config.set(Configs.ENTITY_RELATIONAL_JDBC_BACKEND_USER, "root"); + config.set(Configs.ENTITY_RELATIONAL_JDBC_BACKEND_PASSWORD, "test"); + config.set(Configs.ENABLE_AUTHORIZATION, true); + EntityIdResolver previousResolver = + (EntityIdResolver) + FieldUtils.readStaticField(EntityIdService.class, "entityIdResolver", true); + Object previousExecutor = + FieldUtils.readStaticField(MetadataAuthzHelper.class, "executor", true); + List<Executable> queryCountAssertions = new ArrayList<>(); + try (MockedStatic<GravitinoEnv> envStatic = mockStatic(GravitinoEnv.class); + MockedStatic<GravitinoAuthorizerProvider> providerStatic = + mockStatic(GravitinoAuthorizerProvider.class); + JDBCBackend backend = new JDBCBackend()) { + GravitinoEnv env = mock(GravitinoEnv.class); + envStatic.when(GravitinoEnv::getInstance).thenReturn(env); + when(env.config()).thenReturn(config); + when(env.cacheEnabled()).thenReturn(true); + EntityStore store = mock(EntityStore.class); + SupportsRelationOperations relations = mock(SupportsRelationOperations.class); + when(env.entityStore()).thenReturn(store); + when(store.relationOperations()).thenReturn(relations); + when(relations.batchListEntitiesByRelation( + eq(SupportsRelationOperations.Type.OWNER_REL), anyList(), any())) + .thenAnswer( + call -> + backend.batchListEntitiesByRelation( + call.getArgument(0), call.getArgument(1), call.getArgument(2))); + GravitinoAuthorizerProvider provider = mock(GravitinoAuthorizerProvider.class); + providerStatic.when(GravitinoAuthorizerProvider::getInstance).thenReturn(provider); + GravitinoAuthorizer authorizer = mock(GravitinoAuthorizer.class); + when(provider.getGravitinoAuthorizer()).thenReturn(authorizer); + // Isolate the storage/list-filter path: the caller has a metalake management grant. + when(authorizer.authorize(any(), any(), any(), any(), any())) + .thenAnswer( + call -> { + MetadataObject object = call.getArgument(2); + Privilege.Name privilege = call.getArgument(3); + return object.type() == MetadataObject.Type.METALAKE + && (privilege == Privilege.Name.MANAGE_USERS + || privilege == Privilege.Name.MANAGE_GROUPS + || privilege == Privilege.Name.MANAGE_GRANTS); + }); + FieldUtils.writeStaticField( + MetadataAuthzHelper.class, "executor", (Executor) Runnable::run, true); + backend.initialize(config); + EntityIdService.initialize(new RelationalEntityStoreIdResolver()); + try (Connection connection = + DriverManager.getConnection( + config.get(Configs.ENTITY_RELATIONAL_JDBC_BACKEND_URL), "root", "test")) { + for (int size : new int[] {1, 1003, 10000}) { + String metalake = "scale" + size; + insertPrincipals(connection, size, metalake); + for (Entity.EntityType type : + List.of(Entity.EntityType.USER, Entity.EntityType.GROUP, Entity.EntityType.ROLE)) { Review Comment: Addressed the missing path in adf6cc00fa: the H2 test now exercises USER/GROUP self-filtering fallback at all three sizes and fails when owner skipping is removed. ROLE is tested for management shortcuts, plus a unit regression verifies owner preloading is retained on its fallback path. The enumeration is intentionally scoped to this PR's principal-list behavior, rather than claiming constant-query filtering for every entity type or future registry entry. Other metadata types have different expressions/preload requirements and need their own fixtures when optimized. The description and test comments now make that boundary explicit. ########## server-common/src/test/java/org/apache/gravitino/server/authorization/TestMetadataAuthzHelper.java: ########## @@ -821,4 +979,33 @@ private static void makeCompletableFutureUseCurrentThread() { throw new RuntimeException(e); } } + + private static NameIdentifier[] principalIdentifiers(Entity.EntityType type, int count) { + return IntStream.range(0, count) + .mapToObj( + i -> + switch (type) { + case USER -> NameIdentifierUtil.ofUser("testMetalake", "user" + i); + case GROUP -> NameIdentifierUtil.ofGroup("testMetalake", "group" + i); + default -> NameIdentifierUtil.ofRole("testMetalake", "role" + i); + }) + .toArray(NameIdentifier[]::new); + } + + private static String principalListExpression(Entity.EntityType type) { Review Comment: Addressed in adf6cc00fa. Extracted PrincipalListTestUtils for the expression, namespace/identifier, and management-privilege mappings. Both tests now share the expression and namespace mapping, and unsupported types fail explicitly rather than silently taking the ROLE branch. -- 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]
