Add support for maxRows in the annotation and sent in dynamically

common/src/main/java/com/moparisthebest/jdbc/codegen/JdbcMapper.java

@@ -59,18 +59,14 @@ public interface JdbcMapper extends Closeable {
 		OptionalBool cachePreparedStatement() default OptionalBool.DEFAULT;
 
 		/**
-		 * Maximum array length.
+		 * Maximum rows returned in collections/maps/arrays/etc, < 1 mean no limit
-		 * Optional element.
+		 *
-		 * This element has no effect on the call unless the method return type is an array.
+		 * Use with care because you normally do not want truncated query results
-		 * When used in conjunction with the maxRows element, the size of the array generated
+		 *
-		 * from the result set will be the smaller of maxRows and arrayMaxLength.
+		 * Can also send in dynamically to function as a primitive integer (long/int/short/byte) named one of the options
-		 * <p>
+		 * in JdbcMapper.allowedMaxRowParamNames (default one of maxRows,rowLimit,arrayMaxLength if not set)
-		 * arrayMaxLength's default value is 1024, but may be set to zero to specify that
-		 * there is no size limit for the array generated from the ResultSet.
-		 * Since the generated array is stored in-memory, care should be taken when dealing
-		 * with very large ResultSets when the value of this element is set to zero.
 		 */
-		int arrayMaxLength() default -1;
+		long maxRows() default -1;
 	}
 
 	// these are to annotate parameters for special bind to PreparedStatement behavior

jdbcmapper/src/main/java/com/moparisthebest/jdbc/codegen/CompileTimeResultSetMapper.java

@@ -56,17 +56,17 @@ public class CompileTimeResultSetMapper {
 		return typeMirrorStringNoGenerics(returnType);
 	}
 
-	public void mapToResultType(final Writer w, final String[] keys, final ExecutableElement eeMethod, final int arrayMaxLength, final String cal, final String cleaner) throws IOException, NoSuchMethodException, ClassNotFoundException {
+	public void mapToResultType(final Writer w, final String[] keys, final ExecutableElement eeMethod, final String maxRows, final String cal, final String cleaner) throws IOException, NoSuchMethodException, ClassNotFoundException {
 		//final Method m = fromExecutableElement(eeMethod);
 		//final Class returnType = m.getReturnType();
 		final TypeMirror returnTypeMirror = eeMethod.getReturnType();
 		//final Class returnType = typeMirrorToClass(returnTypeMirror);
 		if (returnTypeMirror.getKind() == TypeKind.ARRAY) {
 			final TypeMirror componentType = ((ArrayType) returnTypeMirror).getComponentType();
-			toArray(w, keys, componentType, arrayMaxLength, cal, cleaner);
+			toArray(w, keys, componentType, maxRows, cal, cleaner);
 		} else if (types.isAssignable(returnTypeMirror, collectionType)) {
 			final List<? extends TypeMirror> typeArguments = ((DeclaredType) returnTypeMirror).getTypeArguments();
-			toCollection(w, keys, returnTypeMirror, typeArguments.get(0), arrayMaxLength, cal, cleaner);
+			toCollection(w, keys, returnTypeMirror, typeArguments.get(0), maxRows, cal, cleaner);
 		} else if (types.isAssignable(returnTypeMirror, mapType)) {
 			final List<? extends TypeMirror> typeArguments = ((DeclaredType) returnTypeMirror).getTypeArguments();
 			//if (types[1] instanceof ParameterizedType) { // for collectionMaps
@@ -78,16 +78,16 @@ public class CompileTimeResultSetMapper {
 						typeArguments.get(0),
 						collectionTypeMirror,
 						componentTypeMirror,
-						arrayMaxLength, cal, cleaner);
+						maxRows, cal, cleaner);
 				return;
 			}
-			toMap(w, keys, returnTypeMirror, typeArguments.get(0), typeArguments.get(1), arrayMaxLength, cal, cleaner);
+			toMap(w, keys, returnTypeMirror, typeArguments.get(0), typeArguments.get(1), maxRows, cal, cleaner);
 		} else if (types.isAssignable(returnTypeMirror, iteratorType)) {
 			final List<? extends TypeMirror> typeArguments = ((DeclaredType) returnTypeMirror).getTypeArguments();
 			if (types.isAssignable(returnTypeMirror, listIteratorType))
-				toListIterator(w, keys, typeArguments.get(0), arrayMaxLength, cal, cleaner);
+				toListIterator(w, keys, typeArguments.get(0), maxRows, cal, cleaner);
 			else
-				toIterator(w, keys, typeArguments.get(0), arrayMaxLength, cal, cleaner);
+				toIterator(w, keys, typeArguments.get(0), maxRows, cal, cleaner);
 		} else {
 			toObject(w, keys, returnTypeMirror, cal, cleaner);
 		}
@@ -109,8 +109,8 @@ public class CompileTimeResultSetMapper {
 		clean(w, cleaner).write(";\n\t\t\t} else {\n\t\t\t\treturn null;\n\t\t\t}\n");
 	}
 
-	public void writeCollection(final Writer w, final String[] keys, final String returnTypeString, final String concreteTypeString, final TypeMirror componentTypeMirror, int arrayMaxLength, String cal, final String cleaner) throws IOException, ClassNotFoundException {
+	public void writeCollection(final Writer w, final String[] keys, final String returnTypeString, final String concreteTypeString, final TypeMirror componentTypeMirror, String maxRows, String cal, final String cleaner) throws IOException, ClassNotFoundException {
-		w.write("\t\t\tfinal ");
+		maxRowInit(w, maxRows).write("\t\t\tfinal ");
 		w.write(returnTypeString);
 		w.write(" _colret = new ");
 		w.write(concreteTypeString);
@@ -118,39 +118,40 @@ public class CompileTimeResultSetMapper {
 		w.write("();\n\t\t\twhile(rs.next()) {\n");
 		writeObject(w, keys, componentTypeMirror, cal);
 		w.write("\t\t\t\t_colret.add(");
-		clean(w, cleaner).write(");\n\t\t\t}\n");
+		clean(w, cleaner).write(");\n");
+		maxRowBreak(w, maxRows).write("\t\t\t}\n");
 	}
 
-	public  void toCollection(final Writer w, final String[] keys, final TypeMirror collectionTypeMirror, final TypeMirror componentTypeMirror, int arrayMaxLength, String cal, final String cleaner) throws IOException, ClassNotFoundException {
+	public  void toCollection(final Writer w, final String[] keys, final TypeMirror collectionTypeMirror, final TypeMirror componentTypeMirror, String maxRows, String cal, final String cleaner) throws IOException, ClassNotFoundException {
 		final String collectionType = getConcreteClassCanonicalName(collectionTypeMirror, ArrayList.class);
-		writeCollection(w, keys, collectionTypeMirror.toString(), collectionType, componentTypeMirror, arrayMaxLength, cal, cleaner);
+		writeCollection(w, keys, collectionTypeMirror.toString(), collectionType, componentTypeMirror, maxRows, cal, cleaner);
 		w.write("\t\t\treturn _colret;\n");
 	}
 
-	public void toArray(final Writer w, final String[] keys, final TypeMirror componentTypeMirror, int arrayMaxLength, String cal, final String cleaner) throws IOException, ClassNotFoundException {
+	public void toArray(final Writer w, final String[] keys, final TypeMirror componentTypeMirror, String maxRows, String cal, final String cleaner) throws IOException, ClassNotFoundException {
 		final String returnTypeString = componentTypeMirror.toString();
-		writeCollection(w, keys, "java.util.List<" + returnTypeString + ">", "java.util.ArrayList", componentTypeMirror, arrayMaxLength, cal, cleaner);
+		writeCollection(w, keys, "java.util.List<" + returnTypeString + ">", "java.util.ArrayList", componentTypeMirror, maxRows, cal, cleaner);
 		w.write("\t\t\treturn _colret.toArray(new ");
 		w.write(returnTypeString);
 		w.write("[_colret.size()]);\n");
 	}
 
-	public void toListIterator(final Writer w, final String[] keys, final TypeMirror componentTypeMirror, int arrayMaxLength, String cal, final String cleaner) throws IOException, ClassNotFoundException {
+	public void toListIterator(final Writer w, final String[] keys, final TypeMirror componentTypeMirror, String maxRows, String cal, final String cleaner) throws IOException, ClassNotFoundException {
 		final String returnTypeString = componentTypeMirror.toString();
-		writeCollection(w, keys, "java.util.List<" + returnTypeString + ">", "java.util.ArrayList", componentTypeMirror, arrayMaxLength, cal, cleaner);
+		writeCollection(w, keys, "java.util.List<" + returnTypeString + ">", "java.util.ArrayList", componentTypeMirror, maxRows, cal, cleaner);
 		w.write("\t\t\treturn _colret.listIterator();\n");
 	}
 
-	public void toIterator(final Writer w, final String[] keys, final TypeMirror componentTypeMirror, int arrayMaxLength, String cal, final String cleaner) throws IOException, ClassNotFoundException {
+	public void toIterator(final Writer w, final String[] keys, final TypeMirror componentTypeMirror, String maxRows, String cal, final String cleaner) throws IOException, ClassNotFoundException {
 		final String returnTypeString = componentTypeMirror.toString();
-		writeCollection(w, keys, "java.util.List<" + returnTypeString + ">", "java.util.ArrayList", componentTypeMirror, arrayMaxLength, cal, cleaner);
+		writeCollection(w, keys, "java.util.List<" + returnTypeString + ">", "java.util.ArrayList", componentTypeMirror, maxRows, cal, cleaner);
 		w.write("\t\t\treturn _colret.iterator();\n");
 	}
 
-	public void toMap(final Writer w, final String[] keys, final TypeMirror mapTypeMirror, final TypeMirror mapKeyTypeMirror, final TypeMirror componentTypeMirror, int arrayMaxLength, String cal, final String cleaner) throws IOException, ClassNotFoundException {
+	public void toMap(final Writer w, final String[] keys, final TypeMirror mapTypeMirror, final TypeMirror mapKeyTypeMirror, final TypeMirror componentTypeMirror, String maxRows, String cal, final String cleaner) throws IOException, ClassNotFoundException {
 		final String mapType = getConcreteClassCanonicalName(mapTypeMirror, HashMap.class);
 		final String returnTypeString = mapTypeMirror.toString();
-		w.write("\t\t\tfinal ");
+		maxRowInit(w, maxRows).write("\t\t\tfinal ");
 		w.write(returnTypeString);
 		w.write(" _colret = new ");
 		w.write(mapType);
@@ -163,7 +164,8 @@ public class CompileTimeResultSetMapper {
 		w.write("\t\t\t\t_colret.put(");
 		rm.extractColumnValueString(w, 1, mapKeyTypeMirror);
 		w.write(", ");
-		clean(w, cleaner).write(");\n\t\t\t}\n");
+		clean(w, cleaner).write(");\n");
+		maxRowBreak(w, maxRows).write("\t\t\t}\n");
 		w.write("\t\t\treturn _colret;\n");
 	}
 
@@ -172,12 +174,13 @@ public class CompileTimeResultSetMapper {
 								final TypeMirror mapKeyTypeMirror,
 								final TypeMirror collectionTypeMirror,
 								final TypeMirror componentTypeMirror,
-								int arrayMaxLength, String cal, final String cleaner) throws IOException, ClassNotFoundException {
+								String maxRows, String cal, final String cleaner) throws IOException, ClassNotFoundException {
 		final String mapType = getConcreteClassCanonicalName(mapTypeMirror, HashMap.class);
 		final String collectionType = getConcreteClassCanonicalName(collectionTypeMirror, ArrayList.class);
 		final String returnTypeString = mapTypeMirror.toString();
 		final String collectionTypeString = collectionTypeMirror.toString();
-		w.write("\t\t\tfinal ");
+
+		maxRowInit(w, maxRows).write("\t\t\tfinal ");
 		w.write(returnTypeString);
 		w.write(" _colret = new ");
 		w.write(mapType);
@@ -199,7 +202,30 @@ public class CompileTimeResultSetMapper {
 		w.write(collectionType);
 		w.write(collectionTypeString.substring(collectionTypeString.indexOf('<')));
 		w.write("();\n\t\t\t\t\t_colret.put(_colkey, _collist);\n\t\t\t\t}\n\t\t\t\t_collist.add(");
-		clean(w, cleaner).write(");\n\t\t\t}\n\t\t\treturn _colret;\n");
+		clean(w, cleaner).write(");\n");
+		maxRowBreak(w, maxRows).write("\t\t\t}\n\t\t\treturn _colret;\n");
+	}
+
+	private Writer maxRowInit(final Writer w, final String maxRows) throws IOException {
+		if(maxRows != null)
+			w.write("\t\t\tlong _rowCount = 1;\n"); // todo: use type of dynamic param, or minimum the hard coded number will fit in?
+		return w;
+	}
+
+	private Writer maxRowBreak(final Writer w, final String maxRows) throws IOException {
+		if(maxRows != null) {
+			w.append("\t\t\t\tif(");
+			// annoying and hacky, but if this is a dynamic param name we want to only check rowCount if it's < 1
+			// if it's a hard coded number, it'll be null so we won't even be here
+			try {
+				Long.parseLong(maxRows);
+			} catch (NumberFormatException e) {
+				// dynamic
+				w.append(maxRows).append(" > 0 && ");
+			}
+			w.append("++_rowCount > ").append(maxRows).append(")\n\t\t\t\t\tbreak;\n");
+		}
+		return w;
 	}
 
 	private Writer clean(final Writer w, final String cleaner) throws IOException {

jdbcmapper/src/main/java/com/moparisthebest/jdbc/codegen/JdbcMapperProcessor.java

@@ -23,7 +23,7 @@ import static com.moparisthebest.jdbc.TryClose.tryClose;
  * Created by mopar on 5/24/17.
  */
 @SupportedAnnotationTypes("com.moparisthebest.jdbc.codegen.JdbcMapper.Mapper")
-@SupportedOptions({"jdbcMapper.databaseType", "jdbcMapper.arrayNumberTypeName", "jdbcMapper.arrayStringTypeName"})
+@SupportedOptions({"jdbcMapper.databaseType", "jdbcMapper.arrayNumberTypeName", "jdbcMapper.arrayStringTypeName", "JdbcMapper.allowedMaxRowParamNames"})
 @SupportedSourceVersion(SourceVersion.RELEASE_5)
 public class JdbcMapperProcessor extends AbstractProcessor {
 
@@ -35,6 +35,7 @@ public class JdbcMapperProcessor extends AbstractProcessor {
 	private TypeElement cleanerElement;
 	private JdbcMapper.DatabaseType defaultDatabaseType;
 	private String defaultArrayNumberTypeName, defaultArrayStringTypeName;
+	private Set<String> allowedMaxRowParamNames;
 	private CompileTimeResultSetMapper rsm;
 
 	public JdbcMapperProcessor() {
@@ -76,6 +77,11 @@ public class JdbcMapperProcessor extends AbstractProcessor {
 		if (defaultArrayStringTypeName == null || defaultArrayStringTypeName.isEmpty())
 			defaultArrayStringTypeName = defaultDatabaseType.arrayStringTypeName;
 
+		String allowedMaxRowParamNames = processingEnv.getOptions().get("JdbcMapper.allowedMaxRowParamNames");
+		if (allowedMaxRowParamNames == null || allowedMaxRowParamNames.isEmpty())
+			allowedMaxRowParamNames = "maxRows,rowLimit,arrayMaxLength";
+		this.allowedMaxRowParamNames = new HashSet<String>(Arrays.asList(allowedMaxRowParamNames.split(",")));
+
 		rsm = new CompileTimeResultSetMapper(processingEnv);
 	}
 
@@ -219,7 +225,7 @@ public class JdbcMapperProcessor extends AbstractProcessor {
 							// build query and bind param order
 							final List<VariableElement> bindParams = new ArrayList<VariableElement>();
 							final String sqlStatement;
-							String calendarName = null, cleanerName = null;
+							String calendarName = null, cleanerName = null, maxRowsName = sql.maxRows() < 1 ? null : Long.toString(sql.maxRows());
 							boolean sqlExceptionThrown = false;
 							{
 								// now parameters
@@ -311,6 +317,8 @@ public class JdbcMapperProcessor extends AbstractProcessor {
 													String.format("@JdbcMapper.SQL method has unused parameter '%s' of cleaner type '%s' when cleaner type required is '%s'", unusedParam.getKey(), unusedType, requiredType),
 													unusedParam.getValue());
 										*/
+									} else if(isPrimitiveInteger(unusedType.getKind()) && maxRowsName == null && this.allowedMaxRowParamNames.contains(unusedParam.getKey())) {
+										maxRowsName = unusedParam.getKey();
 									} else
 										processingEnv.getMessager().printMessage(Diagnostic.Kind.ERROR, String.format("@JdbcMapper.SQL method has unused parameter '%s'", unusedParam.getKey()), unusedParam.getValue());
 								}
@@ -361,7 +369,7 @@ public class JdbcMapperProcessor extends AbstractProcessor {
 										processingEnv.getMessager().printMessage(Diagnostic.Kind.ERROR, "@JdbcMapper.SQL sql parsed a wildcard column name which is not supported", methodElement);
 										return false;
 									}
-								rsm.mapToResultType(w, keys, eeMethod, sql.arrayMaxLength(), calendarName, cleanerName);
+								rsm.mapToResultType(w, keys, eeMethod, maxRowsName, calendarName, cleanerName);
 							}
 
 							// if no SQLException is thrown, we have to catch it here and wrap it with RuntimeException...
@@ -668,4 +676,16 @@ public class JdbcMapperProcessor extends AbstractProcessor {
 				methodElement.getSimpleName().toString().equals("close") &&
 				methodElement.getParameters().isEmpty() ? methodElement : null;
 	}
+
+	public static boolean isPrimitiveInteger(final TypeKind kind) {
+		switch(kind) {
+			case BYTE:
+			case SHORT:
+			case INT:
+			case LONG:
+				return true;
+			default:
+				return false;
+		}
+	}
 }

jdbcmapper/src/test/java/com/moparisthebest/jdbc/codegen/PersonDAO.java

@@ -112,4 +112,20 @@ public interface PersonDAO {
 	@JdbcMapper.SQL("SELECT first_name, last_name, birth_date FROM person WHERE person_no = {personNo}")
 	FieldPerson getPersonCleanNumber(long personNo, Cleaner<Number> clean) throws SQLException;
 	*/
+
+	// max row checks
+	@JdbcMapper.SQL("SELECT first_name, last_name, birth_date FROM person WHERE person_no = {personNo}")
+	List<FieldPerson> getPersonDynamicLimit(long personNo, int maxRows) throws SQLException;
+	@JdbcMapper.SQL("SELECT first_name, last_name, birth_date FROM person WHERE person_no = {personNo}")
+	Map<String, FieldPerson> getPersonDynamicLimit(long personNo, short arrayMaxLength) throws SQLException;
+	@JdbcMapper.SQL("SELECT first_name, last_name, birth_date FROM person WHERE person_no = {personNo}")
+	Map<String, List<FieldPerson>> getPersonDynamicLimit(long personNo, long rowLimit) throws SQLException;
+
+
+	@JdbcMapper.SQL(value = "SELECT first_name, last_name, birth_date FROM person WHERE person_no = {personNo}", maxRows = 1)
+	List<FieldPerson> getPersonStaticLimitList(long personNo) throws SQLException;
+	@JdbcMapper.SQL(value = "SELECT first_name, last_name, birth_date FROM person WHERE person_no = {personNo}", maxRows = 1)
+	Map<String, FieldPerson> getPersonStaticLimitMap(long personNo) throws SQLException;
+	@JdbcMapper.SQL(value = "SELECT first_name, last_name, birth_date FROM person WHERE person_no = {personNo}", maxRows = 1)
+	Map<String, List<FieldPerson>> getPersonStaticLimitMapList(long personNo) throws SQLException;
 }