NIFI-6554 fix oid parsing for PutMongo processor

This closes #3650

Signed-off-by: Mike Thomsen <mthomsen@apache.org>
This commit is contained in:
andrew 2019-08-14 18:48:12 +03:00 committed by Mike Thomsen
parent 5f696482b0
commit 9c8f40415e
No known key found for this signature in database
GPG Key ID: 88511C3D4CAD246F
3 changed files with 30 additions and 2 deletions

View File

@ -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();

View File

@ -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),

View File

@ -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);