#88 AssetFileLoader now supports a batch size

Author: rjrudin

src/main/java/com/marklogic/client/ext/file/GenericFileLoader.java

@@ -15,13 +15,18 @@
  * Generic implementation of FileLoader. Delegates to a DocumentFileReader for reading from a set of file paths, and
  * delegates to a BatchWriter for writing to MarkLogic (where that BatchWriter could use XCC, the REST API, or the
  * Data Movement SDK in ML9).
+ *
+ * The batchSize property defaults to null, which means all files are written in one call via the BatchWriter. Setting
+ * this means that the List of DocumentFile objects read from the DocumentFileReader will be written in batches, each
+ * the size of the batchSize property, except for the final one that may be less than this size.
  */
 public class GenericFileLoader extends LoggingObject implements FileLoader {
 
 	private DocumentFileReader documentFileReader;
 	private BatchWriter batchWriter;
 	private boolean waitForCompletion = true;
 	private boolean logFileUris = true;
+	private Integer batchSize;
 
 	// These are passed on to the DefaultDocumentFileReader that is created if one isn't set
 	private List<FileFilter> fileFilters;
@@ -69,20 +74,54 @@ public List<DocumentFile> loadFiles(String... paths) {
 
 		List<DocumentFile> documentFiles = documentFileReader.readDocumentFiles(paths);
 		if (documentFiles != null && !documentFiles.isEmpty()) {
+			writeBatchOfDocuments(documentFiles, 0);
+			if (waitForCompletion) {
+				batchWriter.waitForCompletion();
+			}
+		}
+		return documentFiles;
+	}
+
+	/**
+	 * If batchSize is not set, then this method will load all the documents in one call to the BatchWriter. Otherwise,
+	 * this will divide up the list of documentFiles into batches matching the value of batchSize, with the last batch
+	 * possibly being less than batchSize.
+	 * 
+	 * @param documentFiles
+	 * @param startPosition
+	 */
+	protected void writeBatchOfDocuments(List<DocumentFile> documentFiles, final int startPosition) {
+		final int documentFilesSize = documentFiles.size();
+		if (startPosition >= documentFilesSize) {
+			return;
+		}
+
+		if (batchSize != null && batchSize < 1) {
+			batchSize = null;
+		}
+
+		// The "end" param to subList below is exclusive, so the highest valid value is the list size
+		int endPosition = batchSize == null ? documentFilesSize : startPosition + batchSize;
+		if (endPosition > documentFilesSize) {
+			endPosition = documentFilesSize;
+		}
+
+		List<DocumentFile> batch = documentFiles.subList(startPosition, endPosition);
+		if (!batch.isEmpty()) {
 			if (logger.isInfoEnabled()) {
-				logger.info(format("Writing %d files", documentFiles.size()));
-				if (logFileUris ) {
-					for (DocumentFile df : documentFiles) {
+				logger.info(format("Writing %d files", batch.size()));
+				if (logFileUris) {
+					for (DocumentFile df : batch) {
 						logger.info("Writing: " + df.getUri());
 					}
 				}
 			}
-			batchWriter.write(documentFiles);
-			if (waitForCompletion) {
-				batchWriter.waitForCompletion();
-			}
+			batchWriter.write(batch);
+		}
+
+		if (endPosition < documentFilesSize) {
+			writeBatchOfDocuments(documentFiles, endPosition);
 		}
-		return documentFiles;
 	}
 
 	/**
@@ -228,4 +267,8 @@ public List<FileFilter> getFileFilters() {
 	public List<DocumentFileProcessor> getDocumentFileProcessors() {
 		return documentFileProcessors;
 	}
+
+	public void setBatchSize(Integer batchSize) {
+		this.batchSize = batchSize;
+	}
 }

src/test/java/com/marklogic/client/ext/modulesloader/impl/LoadModulesTest.java

@@ -85,6 +85,21 @@ public void replaceTokens() {
 		assertTrue(transformText.contains("xdmp:log(\"hello-world\")"));
 	}
 
+	@Test
+	public void customBatchSize() {
+		initializeModulesLoaderWithAssetBatchSize(2);
+		verifyModuleCountWithPattern(".*/ext/.*", "Should load every file", 7);
+	}
+
+	/**
+	 * Just ignoring an invalid batch size, which is anything less than 1.
+	 */
+	@Test
+	public void invalidBatchSize() {
+		initializeModulesLoaderWithAssetBatchSize(-1);
+		verifyModuleCountWithPattern(".*/ext/.*", "Should load every file", 7);
+	}
+
 	@Test
 	public void withFilenamePattern() {
 		verifyModuleCountWithPattern(".*options.*(xml)", "Should only load the single XML options file", 1);
@@ -160,4 +175,12 @@ private void assertModuleExists(String uri) {
 			modulesClient.newServerEval().xquery(String.format("fn:doc-available('%s')", uri)).evalAs(String.class)
 		);
 	}
+
+	private void initializeModulesLoaderWithAssetBatchSize(int batchSize) {
+		AssetFileLoader assetFileLoader = new AssetFileLoader(modulesClient);
+		assetFileLoader.setBatchSize(batchSize);
+		modulesLoader = new DefaultModulesLoader(assetFileLoader);
+		modulesLoader.setModulesManager(null);
+	}
+
 }