Merge pull request #482 from marklogic/feature/metadata-values

MLE-21267 Can now configure metadata values and document properties

Author: rjrudin

marklogic-spark-api/src/main/java/com/marklogic/spark/Options.java

@@ -118,6 +118,17 @@ public abstract class Options {
     public static final String WRITE_COLLECTIONS = "spark.marklogic.write.collections";
     public static final String WRITE_PERMISSIONS = "spark.marklogic.write.permissions";
     public static final String WRITE_JSON_ROOT_NAME = "spark.marklogic.write.jsonRootName";
+
+    /**
+     * @since 2.6.0
+     */
+    public static final String WRITE_METADATA_VALUES_PREFIX = "spark.marklogic.write.metadataValues.";
+
+    /**
+     * @since 2.6.0
+     */
+    public static final String WRITE_DOCUMENT_PROPERTIES_PREFIX = "spark.marklogic.write.documentProperties.";
+
     public static final String WRITE_TEMPORAL_COLLECTION = "spark.marklogic.write.temporalCollection";
     public static final String WRITE_URI_PREFIX = "spark.marklogic.write.uriPrefix";
     public static final String WRITE_URI_REPLACE = "spark.marklogic.write.uriReplace";

marklogic-spark-connector/src/main/java/com/marklogic/spark/writer/DocBuilderFactory.java

@@ -42,6 +42,16 @@ DocBuilderFactory withPermissions(String permissionsString) {
         return this;
     }
 
+    DocBuilderFactory withMetadataValue(String key, String value) {
+        metadata.withMetadataValue(key, value);
+        return this;
+    }
+
+    DocBuilderFactory withDocumentProperty(String name, Object value) {
+        metadata.withProperty(name, value);
+        return this;
+    }
+
     DocBuilderFactory withUriMaker(DocBuilder.UriMaker uriMaker) {
         this.uriMaker = uriMaker;
         return this;

marklogic-spark-connector/src/main/java/com/marklogic/spark/writer/WriteContext.java

@@ -22,6 +22,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.function.BiConsumer;
 import java.util.stream.Stream;
 
 public class WriteContext extends ContextSupport {
@@ -144,9 +145,27 @@ DocBuilder newDocBuilder() {
             configureStandardUriMaker(factory);
         }
 
+        forEachOptionStartingWith(Options.WRITE_METADATA_VALUES_PREFIX, factory::withMetadataValue);
+        forEachOptionStartingWith(Options.WRITE_DOCUMENT_PROPERTIES_PREFIX, factory::withDocumentProperty);
+
         return factory.newDocBuilder();
     }
 
+    /**
+     * Convenience for finding and processing dynamic options that start with a certain prefix.
+     *
+     * @param prefix
+     * @param consumer processes the name (the option minus the prefix) and the option value
+     */
+    private void forEachOptionStartingWith(final String prefix, BiConsumer<String, String> consumer) {
+        getProperties().entrySet().stream()
+            .filter(entry -> entry.getKey().startsWith(prefix))
+            .forEach(entry -> {
+                String name = entry.getKey().substring(prefix.length());
+                consumer.accept(name, entry.getValue());
+            });
+    }
+
     public Format getDocumentFormat() {
         if (hasOption(Options.WRITE_DOCUMENT_TYPE)) {
             String value = getStringOption(Options.WRITE_DOCUMENT_TYPE);

tests/src/test/java/com/marklogic/spark/writer/WriteRowsTest.java

@@ -4,6 +4,7 @@
 package com.marklogic.spark.writer;
 
 import com.fasterxml.jackson.databind.JsonNode;
+import com.marklogic.client.io.DocumentMetadataHandle;
 import com.marklogic.junit5.PermissionsTester;
 import com.marklogic.spark.ConnectorException;
 import com.marklogic.spark.Options;
@@ -139,6 +140,34 @@ void testWithCustomConfig() {
         assertTrue(uri.endsWith(".txt"));
     }
 
+    @Test
+    void metadataValues() {
+        newWriter()
+            .option(Options.WRITE_METADATA_VALUES_PREFIX + "color", "red")
+            .option(Options.WRITE_METADATA_VALUES_PREFIX + "size", "medium")
+            .save();
+
+        String firstUri = verifyTwoHundredDocsWereWritten();
+        DocumentMetadataHandle.DocumentMetadataValues values = readMetadata(firstUri).getMetadataValues();
+        assertEquals(2, values.size());
+        assertEquals("red", values.get("color"));
+        assertEquals("medium", values.get("size"));
+    }
+
+    @Test
+    void documentProperties() {
+        newWriter()
+            .option(Options.WRITE_DOCUMENT_PROPERTIES_PREFIX + "color", "red")
+            .option(Options.WRITE_DOCUMENT_PROPERTIES_PREFIX + "size", "medium")
+            .save();
+
+        String firstUri = verifyTwoHundredDocsWereWritten();
+        DocumentMetadataHandle.DocumentProperties props = readMetadata(firstUri).getProperties();
+        assertEquals(2, props.size());
+        assertEquals("red", props.get("color"));
+        assertEquals("medium", props.get("size"));
+    }
+
     @Test
     void invalidThreadCount() {
         DataFrameWriter writer = newWriter().option(Options.WRITE_THREAD_COUNT, 0);
@@ -257,7 +286,7 @@ private void verifyFailureIsDueToLackOfPermission(SparkException ex) {
         verifyNoDocsWereWritten();
     }
 
-    private void verifyTwoHundredDocsWereWritten() {
+    private String verifyTwoHundredDocsWereWritten() {
         final int expectedCollectionSize = 200;
         String uri = getUrisInCollection(COLLECTION, expectedCollectionSize).get(0);
         assertTrue(uri.startsWith("/test/"), "URI should start with '/test/' due to uriPrefix option: " + uri);
@@ -270,6 +299,7 @@ private void verifyTwoHundredDocsWereWritten() {
         PermissionsTester perms = readDocumentPermissions(uri);
         perms.assertReadPermissionExists("spark-user-role");
         perms.assertUpdatePermissionExists("spark-user-role");
+        return uri;
     }
 
     private void verifyNoDocsWereWritten() {