From 9c8f40415e61ee421e6f12d4e145f07984393f91 Mon Sep 17 00:00:00 2001 From: andrew Date: Wed, 14 Aug 2019 18:48:12 +0300 Subject: [PATCH] NIFI-6554 fix oid parsing for PutMongo processor This closes #3650 Signed-off-by: Mike Thomsen --- .../nifi/processors/mongodb/PutMongo.java | 10 ++++++++-- .../mongodb/MongoWriteTestBase.java | 2 ++ .../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 a62620f8dc..e9bf09e47d 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 7294adeba6..814d78fd19 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 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 de08a80754..cb78f97b4e 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 @@ -391,6 +391,26 @@ public class PutMongoIT extends MongoWriteTestBase { assertEquals(doc, collection.find().first()); } + @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);