[nifi] branch master updated: NIFI-6554 fix oid parsing for PutMongo processor

[mthomsen]

This is an automated email from the ASF dual-hosted git repository.

mthomsen pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/nifi.git


The following commit(s) were added to refs/heads/master by this push:
     new 9c8f404  NIFI-6554 fix oid parsing for PutMongo processor
9c8f404 is described below

commit 9c8f40415e61ee421e6f12d4e145f07984393f91
Author: andrew <andrewkovale...@gmail.com>
AuthorDate: Wed Aug 14 18:48:12 2019 +0300

    NIFI-6554 fix oid parsing for PutMongo processor
    
    This closes #3650
    
    Signed-off-by: Mike Thomsen <mthom...@apache.org>
---
 .../org/apache/nifi/processors/mongodb/PutMongo.java | 10 ++++++++--
 .../nifi/processors/mongodb/MongoWriteTestBase.java  |  2 ++
 .../apache/nifi/processors/mongodb/PutMongoIT.java   | 20 ++++++++++++++++++++
 3 files changed, 30 insertions(+), 2 deletions(-)

diff --git 
a/nifi-nar-bundles/nifi-mongodb-bundle/nifi-mongodb-processors/src/main/java/org/apache/nifi/processors/mongodb/PutMongo.java
 
b/nifi-nar-bundles/nifi-mongodb-bundle/nifi-mongodb-processors/src/main/java/org/apache/nifi/processors/mongodb/PutMongo.java
index a62620f..e9bf09e 100644
--- 
a/nifi-nar-bundles/nifi-mongodb-bundle/nifi-mongodb-processors/src/main/java/org/apache/nifi/processors/mongodb/PutMongo.java
+++ 
b/nifi-nar-bundles/nifi-mongodb-bundle/nifi-mongodb-processors/src/main/java/org/apache/nifi/processors/mongodb/PutMongo.java
@@ -255,8 +255,14 @@ public class PutMongo extends AbstractMongoProcessor {
 
     private Document parseUpdateKey(String updateKey, Map doc) {
         Document retVal;
-        if (updateKey.equals("_id") && ObjectId.isValid(((String) 
doc.get(updateKey)))) {
-            retVal = new Document("_id", new ObjectId((String) 
doc.get(updateKey)));
+        if (updateKey.equals("_id")) {
+            if (doc.get("_id") instanceof ObjectId) {
+                retVal = new Document("_id", doc.get("_id"));
+            } else if (ObjectId.isValid((String) doc.get("_id"))){
+                retVal = new Document("_id", new ObjectId((String) 
doc.get("_id")));
+            } else {
+                retVal = new Document("_id", doc.get("_id"));
+            }
         } else if (updateKey.contains(",")) {
             String[] parts = updateKey.split(",[\\s]*");
             retVal = new Document();
diff --git 
a/nifi-nar-bundles/nifi-mongodb-bundle/nifi-mongodb-processors/src/test/java/org/apache/nifi/processors/mongodb/MongoWriteTestBase.java
 
b/nifi-nar-bundles/nifi-mongodb-bundle/nifi-mongodb-processors/src/test/java/org/apache/nifi/processors/mongodb/MongoWriteTestBase.java
index 7294ade..814d78f 100644
--- 
a/nifi-nar-bundles/nifi-mongodb-bundle/nifi-mongodb-processors/src/test/java/org/apache/nifi/processors/mongodb/MongoWriteTestBase.java
+++ 
b/nifi-nar-bundles/nifi-mongodb-bundle/nifi-mongodb-processors/src/test/java/org/apache/nifi/processors/mongodb/MongoWriteTestBase.java
@@ -32,6 +32,8 @@ public class MongoWriteTestBase {
     protected static final String COLLECTION_NAME = "test";
     protected String DATABASE_NAME;
 
+    protected static Document oidDocument = Document.parse("{\"_id\": 
{\"$oid\": \"5cd1a7376264b959a71588c1\"}}");
+
     protected static final List<Document> DOCUMENTS = Arrays.asList(
             new Document("_id", "doc_1").append("a", 1).append("b", 
2).append("c", 3),
             new Document("_id", "doc_2").append("a", 1).append("b", 
2).append("c", 4),
diff --git 
a/nifi-nar-bundles/nifi-mongodb-bundle/nifi-mongodb-processors/src/test/java/org/apache/nifi/processors/mongodb/PutMongoIT.java
 
b/nifi-nar-bundles/nifi-mongodb-bundle/nifi-mongodb-processors/src/test/java/org/apache/nifi/processors/mongodb/PutMongoIT.java
index de08a80..cb78f97 100644
--- 
a/nifi-nar-bundles/nifi-mongodb-bundle/nifi-mongodb-processors/src/test/java/org/apache/nifi/processors/mongodb/PutMongoIT.java
+++ 
b/nifi-nar-bundles/nifi-mongodb-bundle/nifi-mongodb-processors/src/test/java/org/apache/nifi/processors/mongodb/PutMongoIT.java
@@ -392,6 +392,26 @@ public class PutMongoIT extends MongoWriteTestBase {
     }
 
     @Test
+    public void testUpsertWithOid() throws Exception {
+        TestRunner runner = init(PutMongo.class);
+        runner.setProperty(PutMongo.UPDATE_QUERY_KEY, "_id");
+        byte[] bytes = documentToByteArray(oidDocument);
+
+        runner.setProperty(PutMongo.MODE, "update");
+        runner.setProperty(PutMongo.UPSERT, "true");
+        runner.enqueue(bytes);
+        runner.run();
+
+        runner.assertAllFlowFilesTransferred(PutMongo.REL_SUCCESS, 1);
+        MockFlowFile out = 
runner.getFlowFilesForRelationship(PutMongo.REL_SUCCESS).get(0);
+        out.assertContentEquals(bytes);
+
+        // verify 1 doc inserted into the collection
+        assertEquals(1, collection.count());
+        assertEquals(oidDocument, collection.find().first());
+    }
+
+    @Test
     public void testUpdate() throws Exception {
         TestRunner runner = init(PutMongo.class);
         runner.setProperty(PutMongo.UPDATE_QUERY_KEY, "_id");