#74 Added BatchHandler interface

Author: rjrudin

src/main/java/com/marklogic/client/ext/batch/BatchHandler.java

@@ -0,0 +1,15 @@
+package com.marklogic.client.ext.batch;
+
+import com.marklogic.client.DatabaseClient;
+import com.marklogic.client.document.DocumentWriteOperation;
+
+import java.util.List;
+
+/**
+ * Allows for customizing what RestBatchWriter does with each batch of documents.
+ */
+public interface BatchHandler {
+
+	void handleBatch(DatabaseClient client, List<? extends DocumentWriteOperation> items);
+
+}

src/main/java/com/marklogic/client/ext/batch/DefaultBatchHandler.java

@@ -0,0 +1,67 @@
+package com.marklogic.client.ext.batch;
+
+import com.marklogic.client.DatabaseClient;
+import com.marklogic.client.document.DocumentManager;
+import com.marklogic.client.document.DocumentWriteOperation;
+import com.marklogic.client.document.DocumentWriteSet;
+import com.marklogic.client.document.ServerTransform;
+import com.marklogic.client.ext.helper.LoggingObject;
+import com.marklogic.client.io.Format;
+
+import java.util.List;
+
+/**
+ * Default implementation of BatchHandler that uses a DocumentManager to write documents, along with an optional
+ * ServerTransform.
+ */
+public class DefaultBatchHandler extends LoggingObject implements BatchHandler {
+
+	private Format contentFormat;
+	private ServerTransform serverTransform;
+
+	public DefaultBatchHandler() {
+	}
+
+	public DefaultBatchHandler(Format contentFormat, ServerTransform serverTransform) {
+		this.contentFormat = contentFormat;
+		this.serverTransform = serverTransform;
+	}
+
+	@Override
+	public void handleBatch(DatabaseClient client, List<? extends DocumentWriteOperation> items) {
+		DocumentManager<?, ?> mgr = buildDocumentManager(client);
+		if (contentFormat != null) {
+			mgr.setContentFormat(contentFormat);
+		}
+
+		DocumentWriteSet set = mgr.newWriteSet();
+		for (DocumentWriteOperation item : items) {
+			set.add(item);
+		}
+		int count = set.size();
+		if (logger.isDebugEnabled()) {
+			logger.debug("Writing " + count + " documents to MarkLogic");
+		}
+		if (serverTransform != null) {
+			mgr.write(set, serverTransform);
+		} else {
+			mgr.write(set);
+		}
+		if (logger.isInfoEnabled()) {
+			logger.info("Wrote " + count + " documents to MarkLogic");
+		}
+	}
+
+	protected DocumentManager<?, ?> buildDocumentManager(DatabaseClient client) {
+		return client.newDocumentManager();
+	}
+
+
+	public void setServerTransform(ServerTransform serverTransform) {
+		this.serverTransform = serverTransform;
+	}
+
+	public void setContentFormat(Format contentFormat) {
+		this.contentFormat = contentFormat;
+	}
+}

src/main/java/com/marklogic/client/ext/batch/RestBatchWriter.java

@@ -1,9 +1,7 @@
 package com.marklogic.client.ext.batch;
 
 import com.marklogic.client.DatabaseClient;
-import com.marklogic.client.document.DocumentManager;
 import com.marklogic.client.document.DocumentWriteOperation;
-import com.marklogic.client.document.DocumentWriteSet;
 import com.marklogic.client.document.ServerTransform;
 import com.marklogic.client.io.Format;
 
@@ -14,16 +12,19 @@
  * REST API-based implementation, using the Java Client API. By default, this will call release() on each of the
  * DatabaseClient objects that are passed in. Be sure to disable this if you want to keep using those DatabaseClient
  * objects.
+ * <p>
+ * To customize what this does with every batch, you can set a new instance of BatchHandler. This class defaults to using
+ * DefaultBatchHandler; it'll pass its instances of Format and ServerTransform to that class.
  */
 public class RestBatchWriter extends BatchWriterSupport {
 
 	private List<DatabaseClient> databaseClients;
 	private int clientIndex = 0;
 	private boolean releaseDatabaseClients = true;
-	private ServerTransform serverTransform;
 
-	// Specific to Client API 4.0.+
 	private Format contentFormat;
+	private ServerTransform serverTransform;
+	private BatchHandler batchHandler;
 
 	public RestBatchWriter(DatabaseClient databaseClient) {
 		this(databaseClient, true);
@@ -46,6 +47,17 @@ public void write(List<? extends DocumentWriteOperation> items) {
 		executeRunnable(runnable, items);
 	}
 
+	@Override
+	public void initialize() {
+		super.initialize();
+		if (batchHandler == null) {
+			DefaultBatchHandler dbh = new DefaultBatchHandler();
+			dbh.setContentFormat(contentFormat);
+			dbh.setServerTransform(serverTransform);
+			this.batchHandler = dbh;
+		}
+	}
+
 	protected DatabaseClient determineDatabaseClientToUse() {
 		if (clientIndex >= databaseClients.size()) {
 			clientIndex = 0;
@@ -59,42 +71,11 @@ protected Runnable buildRunnable(final DatabaseClient client, final List<? exten
 		return new Runnable() {
 			@Override
 			public void run() {
-				DocumentManager<?, ?> mgr = buildDocumentManager(client);
-				if (contentFormat != null) {
-					mgr.setContentFormat(contentFormat);
-				}
-
-				DocumentWriteSet set = mgr.newWriteSet();
-				for (DocumentWriteOperation item : items) {
-					set.add(item);
-				}
-				int count = set.size();
-				if (logger.isDebugEnabled()) {
-					logger.debug("Writing " + count + " documents to MarkLogic");
-				}
-				if (serverTransform != null) {
-					mgr.write(set, serverTransform);
-				} else {
-					mgr.write(set);
-				}
-				if (logger.isInfoEnabled()) {
-					logger.info("Wrote " + count + " documents to MarkLogic");
-				}
+				batchHandler.handleBatch(client, items);
 			}
 		};
 	}
 
-	/**
-	 * Factored out so it can be overridden for e.g. those already using MarkLogic 9, who may need to set the content
-	 * format on the manager.
-	 *
-	 * @param client
-	 * @return
-	 */
-	protected DocumentManager<?, ?> buildDocumentManager(DatabaseClient client) {
-		return client.newDocumentManager();
-	}
-
 	@Override
 	public void waitForCompletion() {
 		super.waitForCompletion();

src/main/java/com/marklogic/xcc/template/XccTemplate.java

@@ -38,7 +38,7 @@ public XccTemplate(String uri) {
     public XccTemplate(ContentSource contentSource) {
     	this.contentSource = contentSource;
     }
-    
+
     public <T> T execute(XccCallback<T> callback) {
         Session session = contentSource.newSession();
         try {