jerryshao commented on code in PR #9560: URL: https://github.com/apache/gravitino/pull/9560#discussion_r2672260993
########## core/src/test/java/org/apache/gravitino/catalog/TestManagedFunctionOperations.java: ########## @@ -0,0 +1,390 @@ +/* + * 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.catalog; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import java.util.Arrays; +import java.util.HashMap; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; +import org.apache.gravitino.Entity; +import org.apache.gravitino.EntityAlreadyExistsException; +import org.apache.gravitino.EntityStore; +import org.apache.gravitino.NameIdentifier; +import org.apache.gravitino.Namespace; +import org.apache.gravitino.exceptions.FunctionAlreadyExistsException; +import org.apache.gravitino.exceptions.NoSuchEntityException; +import org.apache.gravitino.exceptions.NoSuchFunctionException; +import org.apache.gravitino.function.Function; +import org.apache.gravitino.function.FunctionDefinition; +import org.apache.gravitino.function.FunctionDefinitions; +import org.apache.gravitino.function.FunctionImpl; +import org.apache.gravitino.function.FunctionImpls; +import org.apache.gravitino.function.FunctionParam; +import org.apache.gravitino.function.FunctionParams; +import org.apache.gravitino.function.FunctionType; +import org.apache.gravitino.meta.FunctionEntity; +import org.apache.gravitino.rel.expressions.literals.Literals; +import org.apache.gravitino.rel.types.Types; +import org.apache.gravitino.storage.IdGenerator; +import org.apache.gravitino.storage.RandomIdGenerator; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +public class TestManagedFunctionOperations { + + private static final String METALAKE_NAME = "test_metalake"; + private static final String CATALOG_NAME = "test_catalog"; + private static final String SCHEMA_NAME = "schema1"; + + private final IdGenerator idGenerator = new RandomIdGenerator(); + private final Map<NameIdentifier, FunctionEntity> entityMap = new HashMap<>(); + + private EntityStore store; + private ManagedFunctionOperations functionOperations; + + @BeforeEach + public void setUp() throws Exception { + entityMap.clear(); + store = createMockEntityStore(); + functionOperations = new ManagedFunctionOperations(store, idGenerator); + } + + @Test + public void testRegisterAndListFunctions() { + NameIdentifier func1Ident = getFunctionIdent("func1"); + FunctionParam[] params1 = new FunctionParam[] {FunctionParams.of("a", Types.IntegerType.get())}; + FunctionDefinition[] definitions1 = new FunctionDefinition[] {createSimpleDefinition(params1)}; + + functionOperations.registerFunction( + func1Ident, + "Test function 1", + FunctionType.SCALAR, + true, + Types.StringType.get(), + definitions1); + + NameIdentifier func2Ident = getFunctionIdent("func2"); + FunctionParam[] params2 = + new FunctionParam[] { + FunctionParams.of("x", Types.StringType.get()), + FunctionParams.of("y", Types.StringType.get()) + }; + FunctionDefinition[] definitions2 = new FunctionDefinition[] {createSimpleDefinition(params2)}; + + functionOperations.registerFunction( + func2Ident, + "Test function 2", + FunctionType.SCALAR, + false, + Types.IntegerType.get(), + definitions2); + + // List functions + NameIdentifier[] functionIdents = functionOperations.listFunctions(getFunctionNamespace()); + Assertions.assertEquals(2, functionIdents.length); + Set<String> functionNames = + Arrays.stream(functionIdents).map(NameIdentifier::name).collect(Collectors.toSet()); + + Assertions.assertTrue(functionNames.contains("func1")); + Assertions.assertTrue(functionNames.contains("func2")); + } + + @Test + public void testRegisterAndGetFunction() { + NameIdentifier funcIdent = getFunctionIdent("my_func"); + FunctionParam[] params = + new FunctionParam[] {FunctionParams.of("input", Types.StringType.get())}; + FunctionDefinition[] definitions = new FunctionDefinition[] {createSimpleDefinition(params)}; + + org.apache.gravitino.function.Function newFunc = + functionOperations.registerFunction( + funcIdent, + "My test function", + FunctionType.SCALAR, + true, + Types.IntegerType.get(), + definitions); + + Assertions.assertEquals("my_func", newFunc.name()); + Assertions.assertEquals("My test function", newFunc.comment()); + Assertions.assertEquals(FunctionType.SCALAR, newFunc.functionType()); + Assertions.assertTrue(newFunc.deterministic()); + Assertions.assertEquals(Types.IntegerType.get(), newFunc.returnType()); + Assertions.assertEquals(0, newFunc.version()); + + // Get function (latest version) + Function loadedFunc = functionOperations.getFunction(funcIdent); + Assertions.assertEquals(newFunc.name(), loadedFunc.name()); + Assertions.assertEquals(newFunc.comment(), loadedFunc.comment()); + + // Test register function that already exists + Assertions.assertThrows( + FunctionAlreadyExistsException.class, + () -> + functionOperations.registerFunction( + funcIdent, + "Another function", + FunctionType.SCALAR, + true, + Types.StringType.get(), + definitions)); + + // Test get non-existing function + NameIdentifier nonExistingIdent = getFunctionIdent("non_existing_func"); + Assertions.assertThrows( + NoSuchFunctionException.class, () -> functionOperations.getFunction(nonExistingIdent)); + } + + @Test + public void testRegisterAndDropFunction() { + NameIdentifier funcIdent = getFunctionIdent("func_to_drop"); + FunctionParam[] params = new FunctionParam[] {FunctionParams.of("a", Types.IntegerType.get())}; + FunctionDefinition[] definitions = new FunctionDefinition[] {createSimpleDefinition(params)}; + + functionOperations.registerFunction( + funcIdent, + "Function to drop", + FunctionType.SCALAR, + true, + Types.StringType.get(), + definitions); + + // Drop the function + boolean dropped = functionOperations.dropFunction(funcIdent); + Assertions.assertTrue(dropped); + + // Verify the function is dropped + Assertions.assertThrows( + NoSuchFunctionException.class, () -> functionOperations.getFunction(funcIdent)); + + // Test drop non-existing function + Assertions.assertFalse(functionOperations.dropFunction(funcIdent)); + } + + @Test + public void testRegisterFunctionWithOverlappingDefinitions() { + NameIdentifier funcIdent = getFunctionIdent("func_overlap_register"); + + // Try to register with two definitions that have overlapping arities + FunctionParam[] params1 = + new FunctionParam[] { + FunctionParams.of("a", Types.IntegerType.get()), + FunctionParams.of("b", Types.FloatType.get(), null, Literals.floatLiteral(1.0f)) + }; + FunctionParam[] params2 = + new FunctionParam[] { + FunctionParams.of("a", Types.IntegerType.get()), + FunctionParams.of("c", Types.StringType.get(), null, Literals.stringLiteral("x")) + }; + + FunctionDefinition[] definitions = + new FunctionDefinition[] {createSimpleDefinition(params1), createSimpleDefinition(params2)}; + + Assertions.assertThrows( + IllegalArgumentException.class, + () -> + functionOperations.registerFunction( + funcIdent, + "Test function", + FunctionType.SCALAR, + true, + Types.StringType.get(), + definitions)); + } + + @SuppressWarnings("unchecked") + private EntityStore createMockEntityStore() throws Exception { + EntityStore mockStore = mock(EntityStore.class); + + // Mock put operation + doAnswer( + invocation -> { + FunctionEntity entity = invocation.getArgument(0); + boolean overwrite = invocation.getArgument(1); + NameIdentifier ident = entity.nameIdentifier(); + + if (!overwrite && entityMap.containsKey(ident)) { + throw new EntityAlreadyExistsException("Entity %s already exists", ident); + } + entityMap.put(ident, entity); + return null; + }) + .when(mockStore) + .put(any(FunctionEntity.class), any(Boolean.class)); + + // Mock get operation + when(mockStore.get( + any(NameIdentifier.class), eq(Entity.EntityType.FUNCTION), eq(FunctionEntity.class))) + .thenAnswer( + invocation -> { + NameIdentifier ident = invocation.getArgument(0); + FunctionEntity entity = findEntityByIdent(ident); + if (entity == null) { + throw new NoSuchEntityException("Entity %s does not exist", ident); + } + return entity; + }); + + // Mock delete operation (2 parameters - default method that calls 3-parameter version) + when(mockStore.delete(any(NameIdentifier.class), eq(Entity.EntityType.FUNCTION))) + .thenAnswer( + invocation -> { + NameIdentifier ident = invocation.getArgument(0); + FunctionEntity entity = findEntityByIdent(ident); + if (entity == null) { + return false; + } + entityMap.remove(entity.nameIdentifier()); + return true; + }); + + // Mock list operation + when(mockStore.list( + any(Namespace.class), eq(FunctionEntity.class), eq(Entity.EntityType.FUNCTION))) + .thenAnswer( + invocation -> { + Namespace namespace = invocation.getArgument(0); + return entityMap.values().stream() + .filter(e -> e.namespace().equals(namespace)) + .collect(Collectors.toList()); + }); + + return mockStore; + } + + /** + * Finds an entity by identifier. This method handles both versioned identifiers (used by + * getFunction) and original identifiers (used by alterFunction and dropFunction). + * + * <p>Versioned identifier format: namespace = original_namespace + function_name, name = version + * Original identifier format: namespace = schema_namespace, name = function_name + */ + private FunctionEntity findEntityByIdent(NameIdentifier ident) { + // First, try to find by original identifier (direct match) + FunctionEntity directMatch = entityMap.get(ident); + if (directMatch != null) { + return directMatch; + } + + // If not found, try to interpret as versioned identifier + String[] levels = ident.namespace().levels(); + if (levels.length < 1) { + return null; + } + String functionName = levels[levels.length - 1]; + Namespace originalNamespace = Namespace.of(Arrays.copyOf(levels, levels.length - 1)); + + for (FunctionEntity entity : entityMap.values()) { + if (entity.name().equals(functionName) && entity.namespace().equals(originalNamespace)) { + return entity; + } + } + return null; + } + + private Namespace getFunctionNamespace() { + return Namespace.of(METALAKE_NAME, CATALOG_NAME, SCHEMA_NAME); + } + + private NameIdentifier getFunctionIdent(String functionName) { + return NameIdentifier.of(getFunctionNamespace(), functionName); + } + + private FunctionDefinition createSimpleDefinition(FunctionParam[] params) { + FunctionImpl impl = FunctionImpls.ofJava(FunctionImpl.RuntimeType.SPARK, "com.example.TestUDF"); + return FunctionDefinitions.of(params, new FunctionImpl[] {impl}); + } + + private FunctionDefinition createDefinitionWithImpls( + FunctionParam[] params, FunctionImpl[] impls) { + return FunctionDefinitions.of(params, impls); + } + + @Test + public void testInvalidParameterOrder() { Review Comment: Move all the public test methods above the private methods. -- 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]
