Merge pull request #167 from marklogic/feature/uri-fix

Fixing bug with URI generation

Author: rjrudin

build.gradle

@@ -56,7 +56,7 @@ dependencies {
     exclude module: 'jackson-dataformat-csv'
   }
 
-  testImplementation "ch.qos.logback:logback-classic:1.3.5"
+  testImplementation "ch.qos.logback:logback-classic:1.3.14"
   testImplementation "org.slf4j:jcl-over-slf4j:1.7.36"
   testImplementation "org.skyscreamer:jsonassert:1.5.1"
 }

src/main/java/com/marklogic/spark/writer/MarkLogicWrite.java

@@ -38,8 +38,16 @@ public class MarkLogicWrite implements BatchWrite, StreamingWrite {
     private WriteContext writeContext;
 
     // Used solely for testing. Will never be populated in a real world scenario.
-    public static Consumer<Integer> successCountConsumer;
-    public static Consumer<Integer> failureCountConsumer;
+    private static Consumer<Integer> successCountConsumer;
+    private static Consumer<Integer> failureCountConsumer;
+
+    public static void setSuccessCountConsumer(Consumer<Integer> consumer) {
+        successCountConsumer = consumer;
+    }
+
+    public static void setFailureCountConsumer(Consumer<Integer> consumer) {
+        failureCountConsumer = consumer;
+    }
 
     MarkLogicWrite(WriteContext writeContext) {
         this.writeContext = writeContext;

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

@@ -23,6 +23,7 @@
 import com.marklogic.spark.ContextSupport;
 import com.marklogic.spark.Options;
 import com.marklogic.spark.Util;
+import com.marklogic.spark.reader.document.DocumentRowSchema;
 import org.apache.spark.sql.types.StructType;
 
 import java.util.Arrays;
@@ -104,7 +105,7 @@ DocBuilder newDocBuilder() {
             String uriSuffix = null;
             if (hasOption(Options.WRITE_URI_SUFFIX)) {
                 uriSuffix = getProperties().get(Options.WRITE_URI_SUFFIX);
-            } else if (!isUsingFileSchema()) {
+            } else if (!isUsingFileSchema() && !DocumentRowSchema.SCHEMA.equals(this.schema)) {
                 uriSuffix = ".json";
             }
             factory.withUriMaker(new StandardUriMaker(

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

@@ -180,8 +180,8 @@ void invalidPermissionsConfig() {
     void dontAbortOnFailure() {
         AtomicInteger successCount = new AtomicInteger();
         AtomicInteger failureCount = new AtomicInteger();
-        MarkLogicWrite.successCountConsumer = count -> successCount.set(count);
-        MarkLogicWrite.failureCountConsumer = count -> failureCount.set(count);
+        MarkLogicWrite.setSuccessCountConsumer(count -> successCount.set(count));
+        MarkLogicWrite.setFailureCountConsumer(count -> failureCount.set(count));
 
         newWriterWithDefaultConfig("temporal-data-with-invalid-rows.csv", 1)
             .option(Options.WRITE_TEMPORAL_COLLECTION, TEMPORAL_COLLECTION)

src/test/java/com/marklogic/spark/writer/customcode/ProcessWithCustomCodeTest.java

@@ -127,8 +127,8 @@ void abortOnFailure() {
     void dontAbortOnFailure() {
         AtomicInteger successCount = new AtomicInteger();
         AtomicInteger failureCount = new AtomicInteger();
-        MarkLogicWrite.successCountConsumer = count -> successCount.set(count);
-        MarkLogicWrite.failureCountConsumer = count -> failureCount.set(count);
+        MarkLogicWrite.setSuccessCountConsumer(count -> successCount.set(count));
+        MarkLogicWrite.setFailureCountConsumer(count -> failureCount.set(count));
 
         // The lack of an error here indicates that the job did not abort. The connector is expected to have logged
         // each error instead.

src/test/java/com/marklogic/spark/writer/document/WriteDocumentRowsToMarkLogicTest.java

@@ -55,6 +55,25 @@ void retainMetadataFromCopiedDocuments() {
         });
     }
 
+    /**
+     * This ensures that the default suffix of ".json" isn't applied when the incoming row is a "document row" and
+     * thus has an initial URI.
+     */
+    @Test
+    void uriPrefix() {
+        readTheTwoTestDocuments()
+            .write().format(CONNECTOR_IDENTIFIER)
+            .option(Options.CLIENT_URI, makeClientUri())
+            .option(Options.WRITE_URI_PREFIX, "/backup")
+            .option(Options.WRITE_COLLECTIONS, "backup-docs")
+            .mode(SaveMode.Append)
+            .save();
+
+        assertCollectionSize("backup-docs", 2);
+        assertInCollections("/backup/test/1.xml", "backup-docs");
+        assertInCollections("/backup/test/2.xml", "backup-docs");
+    }
+
     @Test
     void overrideMetadataFromCopiedDocuments() {
         readTheTwoTestDocuments()