TEIID-3901 accounting for user defined aggregates

teiid · Jan 12, 2016 · e0b8da6 · e0b8da6
1 parent 113327b
commit e0b8da6
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 4 deletions.
diff --git a/api/src/main/java/org/teiid/metadata/MetadataFactory.java b/api/src/main/java/org/teiid/metadata/MetadataFactory.java
@@ -25,6 +25,7 @@
 import java.io.Reader;
 import java.io.Serializable;
 import java.lang.reflect.Method;
+import java.lang.reflect.Modifier;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -35,10 +36,12 @@
 import java.util.TreeMap;
 
 import org.teiid.CommandContext;
+import org.teiid.UserDefinedAggregate;
 import org.teiid.adminapi.Model;
 import org.teiid.adminapi.impl.DataPolicyMetadata.PermissionMetaData;
 import org.teiid.adminapi.impl.ModelMetaData;
 import org.teiid.connector.DataPlugin;
+import org.teiid.core.TeiidRuntimeException;
 import org.teiid.core.types.DataTypeManager;
 import org.teiid.core.util.StringUtil;
 import org.teiid.metadata.FunctionMethod.PushDown;
@@ -526,6 +529,20 @@ public FunctionMethod addFunction(String name, Method method) {
 
 	public static FunctionMethod createFunctionFromMethod(String name, Method method) {
 		Class<?> returnTypeClass = method.getReturnType();
+		AggregateAttributes aa = null;
+		//handle user defined aggregates
+		if ((method.getModifiers() & Modifier.STATIC) == 0 && UserDefinedAggregate.class.isAssignableFrom(method.getDeclaringClass())) {
+			aa = new AggregateAttributes();
+			Method m;
+			try {
+				m = method.getDeclaringClass().getMethod("getResult", CommandContext.class); //$NON-NLS-1$
+			} catch (NoSuchMethodException e) {
+				throw new TeiidRuntimeException(e);
+			} catch (SecurityException e) {
+				throw new TeiidRuntimeException(e);
+			}
+			returnTypeClass = m.getReturnType();
+		}
 		if (returnTypeClass.isPrimitive()) {
 			returnTypeClass = TypeFacility.convertPrimitiveToObject(returnTypeClass);
 		}
@@ -549,6 +566,7 @@ public static FunctionMethod createFunctionFromMethod(String name, Method method
 			paramTypes = Arrays.copyOfRange(paramTypes, 1, paramTypes.length);
 		}
 		FunctionMethod func = FunctionMethod.createFunctionMethod(name, null, null, returnType, paramTypes);
+		func.setAggregateAttributes(aa);
 		func.setInvocationMethod(method.getName());
 		func.setPushdown(PushDown.CANNOT_PUSHDOWN);
 		func.setMethod(method);

diff --git a/.../teiid/metadata/TestMetatdataFactory.java → ...g/teiid/metadata/TestMetadataFactory.java b/.../teiid/metadata/TestMetatdataFactory.java → ...g/teiid/metadata/TestMetadataFactory.java
@@ -27,11 +27,13 @@
 import java.util.Collections;
 
 import org.junit.Test;
+import org.teiid.CommandContext;
+import org.teiid.UserDefinedAggregate;
 import org.teiid.adminapi.impl.ModelMetaData;
 import org.teiid.metadata.AbstractMetadataRecord.DataModifiable;
 
 @SuppressWarnings("nls")
-public class TestMetatdataFactory {
+public class TestMetadataFactory {
 
 	@Test public void testSchemaProperties() {
 		ModelMetaData mmd = new ModelMetaData();
@@ -45,13 +47,19 @@ public class TestMetatdataFactory {
 	}
 
 	@Test public void testCreateFunction() throws NoSuchMethodException, SecurityException {
-		FunctionMethod fm = MetadataFactory.createFunctionFromMethod("x", TestMetatdataFactory.class.getMethod("someFunction"));
+		FunctionMethod fm = MetadataFactory.createFunctionFromMethod("x", TestMetadataFactory.class.getMethod("someFunction"));
 		assertEquals(Boolean.class, fm.getOutputParameter().getJavaType());
 
-		fm = MetadataFactory.createFunctionFromMethod("x", TestMetatdataFactory.class.getMethod("someArrayFunction"));
+		fm = MetadataFactory.createFunctionFromMethod("x", TestMetadataFactory.class.getMethod("someArrayFunction"));
 		assertEquals(String[].class, fm.getOutputParameter().getJavaType());
 	}
 
+	@Test public void testCreateAggregateFunction() throws NoSuchMethodException, SecurityException {
+		FunctionMethod fm = MetadataFactory.createFunctionFromMethod("x", MyUDAF.class.getMethod("addInput", String.class));
+		assertEquals(Boolean.class, fm.getOutputParameter().getJavaType());
+		assertNotNull(fm.getAggregateAttributes());
+	}
+
 	public static boolean someFunction() {
 		return true;
 	}
@@ -60,4 +68,20 @@ public static String[] someArrayFunction() {
 		return null;
 	}
 
+	public static class MyUDAF implements UserDefinedAggregate<Boolean> {
+		@Override
+		public Boolean getResult(CommandContext commandContext) {
+			return null;
+		}
+
+		@Override
+		public void reset() {
+
+		}
+
+		public void addInput(String val) {
+
+		}
+	}
+
 }
diff --git a/engine/src/main/java/org/teiid/query/function/FunctionTree.java b/engine/src/main/java/org/teiid/query/function/FunctionTree.java
@@ -362,7 +362,7 @@ private FunctionDescriptor createFunctionDescriptor(
 
         FunctionDescriptor result = new FunctionDescriptor(method, types, outputType, invocationMethod, requiresContext,
                 source.getClassLoader());
-        if (method.getAggregateAttributes() != null && (method.getPushdown() == PushDown.CAN_PUSHDOWN || method.getPushdown() == PushDown.CANNOT_PUSHDOWN)) {
+        if (validateClass && method.getAggregateAttributes() != null && (method.getPushdown() == PushDown.CAN_PUSHDOWN || method.getPushdown() == PushDown.CANNOT_PUSHDOWN)) {
         	result.newInstance();
         }
         result.setHasWrappedArgs(hasWrappedArg);