#44 New BatchWriter interface for batched writes to ML9

Author: Rob Rudin

src/main/java/com/marklogic/client/batch/BatchWriter.java

@@ -0,0 +1,29 @@
+package com.marklogic.client.batch;
+
+import com.marklogic.client.document.DocumentWriteOperation;
+
+import java.util.List;
+
+/**
+ * Interface for writing batches of documents to pre-ML9 clusters (DMSDK can be used with ML9+).
+ */
+public interface BatchWriter {
+
+	/**
+	 * Give the writer a chance to perform any initialization it requires before it starts writing documents.
+	 */
+	void initialize();
+
+	/**
+	 * Write the given list of documents, as defined by the Java Client DocumentWriteOperation interface.
+	 *
+	 * @param items
+	 */
+	void write(List<? extends DocumentWriteOperation> items);
+
+	/**
+	 * Assuming that the writer is using a multi-threaded approach, call this to wait for the writer to finish
+	 * performing all of its writes.
+	 */
+	void waitForCompletion();
+}

src/main/java/com/marklogic/client/batch/BatchWriterSupport.java

@@ -0,0 +1,83 @@
+package com.marklogic.client.batch;
+
+import com.marklogic.client.helper.LoggingObject;
+import org.springframework.core.task.AsyncTaskExecutor;
+import org.springframework.core.task.SyncTaskExecutor;
+import org.springframework.core.task.TaskExecutor;
+import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * Support class for BatchWriter implementations that uses Spring's TaskExecutor interface for parallelizing writes to
+ * MarkLogic.
+ */
+public abstract class BatchWriterSupport extends LoggingObject implements BatchWriter {
+
+	private TaskExecutor taskExecutor;
+	private int threadCount = 16;
+
+	/**
+	 * Seems necessary to keep track of each Future instance so that we can properly wait for each one to finish.
+	 * Spring's TaskExecutor library doesn't seem to provide a better way of doing this.
+	 */
+	private List<Future<?>> futures = new ArrayList<>();
+
+	@Override
+	public void initialize() {
+		if (threadCount > 1) {
+			if (logger.isInfoEnabled()) {
+				logger.info("Initializing thread pool with a count of " + threadCount);
+			}
+			ThreadPoolTaskExecutor tpte = new ThreadPoolTaskExecutor();
+			tpte.setCorePoolSize(threadCount);
+			tpte.afterPropertiesSet();
+			this.taskExecutor = tpte;
+		} else {
+			if (logger.isInfoEnabled()) {
+				logger.info("Thread count is 1, so using a synchronous TaskExecutor");
+			}
+			this.taskExecutor = new SyncTaskExecutor();
+		}
+	}
+
+	@Override
+	public void waitForCompletion() {
+		int size = futures.size();
+		if (logger.isDebugEnabled()) {
+			logger.debug("Waiting for threads to finish document processing; futures count: " + size);
+		}
+
+		for (int i = 0; i < size; i++) {
+			Future<?> f = futures.get(i);
+			if (f.isDone()) {
+				continue;
+			}
+			try {
+				// Wait up to 1 hour for a write to ML to finish (should never happen)
+				f.get(1, TimeUnit.HOURS);
+			} catch (Exception ex) {
+				logger.warn("Unable to wait for last set of documents to be processed: " + ex.getMessage(), ex);
+			}
+		}
+	}
+
+	protected void execute(Runnable runnable) {
+		if (taskExecutor instanceof AsyncTaskExecutor) {
+			futures.add(((AsyncTaskExecutor) taskExecutor).submit(runnable));
+		} else {
+			taskExecutor.execute(runnable);
+		}
+	}
+
+	public void setTaskExecutor(TaskExecutor taskExecutor) {
+		this.taskExecutor = taskExecutor;
+	}
+
+	public void setThreadCount(int threadCount) {
+		this.threadCount = threadCount;
+	}
+}

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

@@ -0,0 +1,46 @@
+package com.marklogic.client.batch;
+
+import com.marklogic.client.DatabaseClient;
+import com.marklogic.client.document.DocumentWriteOperation;
+import com.marklogic.client.document.DocumentWriteSet;
+import com.marklogic.client.document.GenericDocumentManager;
+
+import java.util.List;
+
+public class RestBatchWriter extends BatchWriterSupport {
+
+	private List<DatabaseClient> databaseClients;
+	private int clientIndex = 0;
+
+	public RestBatchWriter(List<DatabaseClient> databaseClients) {
+		this.databaseClients = databaseClients;
+	}
+
+	@Override
+	public void write(final List<? extends DocumentWriteOperation> items) {
+		if (clientIndex >= databaseClients.size()) {
+			clientIndex = 0;
+		}
+		final DatabaseClient client = databaseClients.get(clientIndex);
+		clientIndex++;
+
+		execute(new Runnable() {
+			@Override
+			public void run() {
+				GenericDocumentManager mgr = client.newDocumentManager();
+				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");
+				}
+				mgr.write(set);
+				if (logger.isInfoEnabled()) {
+					logger.info("Wrote " + count + " documents to MarkLogic");
+				}
+			}
+		});
+	}
+}

src/main/java/com/marklogic/client/batch/XccBatchWriter.java

@@ -0,0 +1,67 @@
+package com.marklogic.client.batch;
+
+import com.marklogic.client.document.DocumentWriteOperation;
+import com.marklogic.client.xcc.DefaultDocumentWriteOperationAdapter;
+import com.marklogic.client.xcc.DocumentWriteOperationAdapter;
+import com.marklogic.xcc.Content;
+import com.marklogic.xcc.ContentSource;
+import com.marklogic.xcc.Session;
+import com.marklogic.xcc.exceptions.RequestException;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * XCC implementation for batched writes. Most important thing here is we depend on an instance of
+ * DocumentWriteOperationAdapter to adapt a DocumentWriteOperation instance into a Content instance.
+ */
+public class XccBatchWriter extends BatchWriterSupport {
+
+	private List<ContentSource> contentSources;
+	private int contentSourceIndex = 0;
+	private DocumentWriteOperationAdapter documentWriteOperationAdapter;
+
+	public XccBatchWriter(List<ContentSource> contentSources) {
+		this.contentSources = contentSources;
+		this.documentWriteOperationAdapter = new DefaultDocumentWriteOperationAdapter();
+	}
+
+	@Override
+	public void write(final List<? extends DocumentWriteOperation> items) {
+		if (contentSourceIndex >= contentSources.size()) {
+			contentSourceIndex = 0;
+		}
+
+		final ContentSource contentSource = contentSources.get(contentSourceIndex);
+		contentSourceIndex++;
+
+		execute(new Runnable() {
+			@Override
+			public void run() {
+				Session session = contentSource.newSession();
+				int count = items.size();
+				Content[] array = new Content[count];
+				for (int i = 0; i < count; i++) {
+					array[i] = documentWriteOperationAdapter.adapt(items.get(i));
+				}
+				if (logger.isDebugEnabled()) {
+					logger.debug("Writing " + count + " documents to MarkLogic");
+				}
+				try {
+					session.insertContent(array);
+					if (logger.isInfoEnabled()) {
+						logger.info("Wrote " + count + " documents to MarkLogic");
+					}
+				} catch (RequestException e) {
+					throw new RuntimeException("Unable to insert content: " + e.getMessage(), e);
+				} finally {
+					session.close();
+				}
+			}
+		});
+	}
+
+	public void setDocumentWriteOperationAdapter(DocumentWriteOperationAdapter documentWriteOperationAdapter) {
+		this.documentWriteOperationAdapter = documentWriteOperationAdapter;
+	}
+}

src/main/java/com/marklogic/client/xcc/DefaultDocumentWriteOperationAdapter.java

@@ -0,0 +1,114 @@
+package com.marklogic.client.xcc;
+
+import com.marklogic.client.document.DocumentWriteOperation;
+import com.marklogic.client.io.*;
+import com.marklogic.client.io.marker.AbstractWriteHandle;
+import com.marklogic.client.io.marker.DocumentMetadataWriteHandle;
+import com.marklogic.xcc.*;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.util.HashSet;
+import java.util.Set;
+
+public class DefaultDocumentWriteOperationAdapter implements DocumentWriteOperationAdapter {
+
+	private final static Logger logger = LoggerFactory.getLogger(DefaultDocumentWriteOperationAdapter.class);
+
+	@Override
+	public Content adapt(DocumentWriteOperation operation) {
+		String uri = operation.getUri();
+		ContentCreateOptions options = adaptMetadata(operation.getMetadata());
+		AbstractWriteHandle handle = operation.getContent();
+		if (handle instanceof StringHandle) {
+			return ContentFactory.newContent(uri, ((StringHandle) handle).get(), options);
+		} else if (handle instanceof FileHandle) {
+			return ContentFactory.newContent(uri, ((FileHandle) handle).get(), options);
+		} else if (handle instanceof BytesHandle) {
+			return ContentFactory.newContent(uri, ((BytesHandle) handle).get(), options);
+		} else if (handle instanceof InputStreamHandle) {
+			try {
+				return ContentFactory.newContent(uri, ((InputStreamHandle) handle).get(), options);
+			} catch (IOException e) {
+				throw new RuntimeException("Unable to read content input stream: " + e.getMessage(), e);
+			}
+		} else if (handle instanceof DOMHandle) {
+			return ContentFactory.newContent(uri, ((DOMHandle) handle).get(), options);
+		} else throw new IllegalArgumentException("No support yet for content class: " + handle.getClass().getName());
+	}
+
+	/**
+	 * TODO Only adapts collections, quality, format, and permissions so far.
+	 *
+	 * @param handle
+	 * @return
+	 */
+	protected ContentCreateOptions adaptMetadata(DocumentMetadataWriteHandle handle) {
+		ContentCreateOptions options = new ContentCreateOptions();
+		if (handle instanceof DocumentMetadataHandle) {
+			DocumentMetadataHandle metadata = (DocumentMetadataHandle) handle;
+			options.setQuality(metadata.getQuality());
+			options.setCollections(metadata.getCollections().toArray(new String[]{}));
+			adaptPermissions(options, metadata);
+			adaptFormat(options, metadata);
+		} else {
+			logger.warn("Only supports DocumentMetadataHandle; unsupported metadata class: " + handle.getClass().getName());
+		}
+		return options;
+	}
+
+	/**
+	 * The REST API Format class has a "getDefaultMimetype" method on it, but there doesn't appear to be anything
+	 * useful to do with that for XCC. So we just do a simple translation from Format to DocumentFormat.
+	 *
+	 * @param options
+	 * @param metadata
+	 */
+	protected void adaptFormat(ContentCreateOptions options, DocumentMetadataHandle metadata) {
+		Format format = metadata.getFormat();
+		DocumentFormat xccFormat = null;
+		if (format != null) {
+			if (Format.BINARY.equals(format)) {
+				xccFormat = DocumentFormat.BINARY;
+			} else if (Format.JSON.equals(format)) {
+				xccFormat = DocumentFormat.JSON;
+			} else if (Format.TEXT.equals(format)) {
+				xccFormat = DocumentFormat.TEXT;
+			} else if (Format.XML.equals(format)) {
+				xccFormat = DocumentFormat.XML;
+			} else if (Format.UNKNOWN.equals(format)) {
+				xccFormat = DocumentFormat.NONE;
+			} else if (logger.isDebugEnabled()) {
+				logger.debug("Unsupported format, can't adapt to an XCC DocumentFormat; " + format.toString());
+			}
+		}
+		if (xccFormat != null) {
+			if (logger.isDebugEnabled()) {
+				logger.debug("Adapted REST format " + format + " to XCC format: " + xccFormat.toString());
+			}
+			options.setFormat(xccFormat);
+		}
+	}
+
+	protected void adaptPermissions(ContentCreateOptions options, DocumentMetadataHandle metadata) {
+		Set<ContentPermission> contentPermissions = new HashSet<>();
+		DocumentMetadataHandle.DocumentPermissions permissions = metadata.getPermissions();
+		for (String role : permissions.keySet()) {
+			for (DocumentMetadataHandle.Capability capability : permissions.get(role)) {
+				ContentCapability contentCapability;
+				if (DocumentMetadataHandle.Capability.EXECUTE.equals(capability)) {
+					contentCapability = ContentCapability.EXECUTE;
+				} else if (DocumentMetadataHandle.Capability.INSERT.equals(capability)) {
+					contentCapability = ContentCapability.INSERT;
+				} else if (DocumentMetadataHandle.Capability.READ.equals(capability)) {
+					contentCapability = ContentCapability.READ;
+				} else if (DocumentMetadataHandle.Capability.UPDATE.equals(capability)) {
+					contentCapability = ContentCapability.UPDATE;
+				} else throw new IllegalArgumentException("Unrecognized permission capability: " + capability);
+				contentPermissions.add(new ContentPermission(contentCapability, role));
+			}
+		}
+		options.setPermissions(contentPermissions.toArray(new ContentPermission[]{}));
+	}
+}

src/main/java/com/marklogic/client/xcc/DocumentWriteOperationAdapter.java

@@ -0,0 +1,15 @@
+package com.marklogic.client.xcc;
+
+import com.marklogic.client.document.DocumentWriteOperation;
+import com.marklogic.xcc.Content;
+
+/**
+ * Interface for adapting a REST DocumentWriteOperation instance into an XCC Content instance. Intended to support
+ * objects that prefer using the REST API to insert a document, where the document to insert is defined via a
+ * DocumentWriteOperation instance, but can then shift to XCC if needed.
+ */
+public interface DocumentWriteOperationAdapter {
+
+	Content adapt(DocumentWriteOperation operation);
+
+}