Skip to content

Adding a new returnAndBulkUpsert API in the doc store. #18

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
buchi-busireddy merged 4 commits into from
Dec 17, 2020
Merged

Adding a new returnAndBulkUpsert API in the doc store. #18

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
3169213
Adding a new API to bulkUpsertAndGet the documents in the doc store.
Dec 15, 2020
b36a521
Changing the API to returnAndBulkUpsert() because that's the use case.
Dec 15, 2020
b020782
Improving the assertions on the integration tests.
Dec 16, 2020
37ba364
Renaming the method.
Dec 17, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions document-store/build.gradle.kts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ dependencies {
testImplementation("org.junit.jupiter:junit-jupiter:5.6.2")
testImplementation("org.mockito:mockito-core:2.19.0")
integrationTestImplementation("org.junit.jupiter:junit-jupiter:5.6.2")
integrationTestImplementation("com.github.java-json-tools:json-patch:1.13")
}

tasks.test {
Expand Down
86 changes: 84 additions & 2 deletions ...e/src/integrationTest/java/org/hypertrace/core/documentstore/mongo/MongoDocStoreTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,9 +5,11 @@
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertTrue;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.github.fge.jsonpatch.diff.JsonDiff;
import com.mongodb.BasicDBObject;
import com.mongodb.DBObject;
import com.mongodb.client.FindIterable;
Expand All @@ -22,9 +24,13 @@
import java.time.Instant;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.stream.Collectors;
import org.hypertrace.core.documentstore.Collection;
import org.hypertrace.core.documentstore.Datastore;
import org.hypertrace.core.documentstore.DatastoreProvider;
Expand Down Expand Up @@ -112,8 +118,8 @@ public void testIgnoreCaseLikeQuery() throws IOException {
String persistedDocument = documents.get(0).toJson();
JsonNode jsonNode = OBJECT_MAPPER.reader().readTree(persistedDocument);
Assertions.assertTrue(persistedDocument.contains("Bob"));
Assertions.assertTrue(jsonNode.findValue("createdTime").asLong(0) > now);
Assertions.assertTrue(jsonNode.findValue("lastUpdatedTime").asLong(0) > now);
Assertions.assertTrue(jsonNode.findValue("createdTime").asLong(0) >= now);
Assertions.assertTrue(jsonNode.findValue("lastUpdatedTime").asLong(0) >= now);
}
}

Expand Down Expand Up @@ -394,6 +400,74 @@ public void testBulkUpsert() {
assertEquals(1, collection.count());
}

@Test
public void testReturnAndBulkUpsert() throws IOException {
datastore.createCollection(COLLECTION_NAME, null);
Collection collection = datastore.getCollection(COLLECTION_NAME);
Map<Key, Document> documentMapV1 = Map.of(
new SingleValueKey("default", "testKey1"), createDocument("id", "1", "testKey1", "abc-v1"),
new SingleValueKey("default", "testKey2"), createDocument("id", "2", "testKey2", "xyz-v1")
);

Iterator<Document> iterator = collection.bulkUpsertAndReturnOlderDocuments(documentMapV1);
// Initially there shouldn't be any documents.
Assertions.assertFalse(iterator.hasNext());

// Add more details to the document and bulk upsert again.
Map<Key, Document> documentMapV2 = Map.of(
new SingleValueKey("default", "testKey1"), createDocument("id", "1", "testKey1", "abc-v2"),
new SingleValueKey("default", "testKey2"), createDocument("id", "2", "testKey2", "xyz-v2")
);
iterator = collection.bulkUpsertAndReturnOlderDocuments(documentMapV2);
assertEquals(2, collection.count());
List<Document> documents = new ArrayList<>();
while (iterator.hasNext()) {
documents.add(iterator.next());
}
assertEquals(2, documents.size());

Map<String, JsonNode> expectedDocs = convertToMap(documentMapV1.values(), "id");
Map<String, JsonNode> actualDocs = convertToMap(documents, "id");

// Verify that the documents returned were previous copies.
for (Map.Entry<String, JsonNode> entry: expectedDocs.entrySet()) {
JsonNode expected = entry.getValue();
JsonNode actual = actualDocs.get(entry.getKey());

Assertions.assertNotNull(actual);
JsonNode patch = JsonDiff.asJson(expected, actual);

// Verify that there are only additions and "no" removals in this new node.
Set<String> ops = new HashSet<>();
patch.elements().forEachRemaining(e -> {
if (e.has("op")) {
ops.add(e.get("op").asText());
}
});

Assertions.assertTrue(ops.contains("add"));
Assertions.assertEquals(1, ops.size());
}

// Delete one of the documents and test again.
collection.delete(new SingleValueKey("default", "testKey1"));
assertEquals(1, collection.count());
}

private Map<String, JsonNode> convertToMap(java.util.Collection<Document> docs, String key) {
return docs.stream()
.map(d -> {
try {
return OBJECT_MAPPER.reader().readTree(d.toJson());
} catch (JsonProcessingException e) {
e.printStackTrace();
}
return null;
})
.filter(Objects::nonNull)
.collect(Collectors.toMap(d -> d.get(key).asText(), d -> d));
}

@Test
public void testLike() {
MongoClient mongoClient = MongoClients.create("mongodb://localhost:27017");
Expand Down Expand Up @@ -431,6 +505,14 @@ public void testLike() {
assertEquals(1, results.size());
}

private Document createDocument(String ...keys) {
ObjectNode objectNode = OBJECT_MAPPER.createObjectNode();
for (int i = 0; i < keys.length - 1; i++) {
objectNode.put(keys[i], keys[i + 1]);
}
return new JSONDocument(objectNode);
}

private Document createDocument(String key, String value) {
ObjectNode objectNode = OBJECT_MAPPER.createObjectNode();
objectNode.put(key, value);
Expand Down
61 changes: 53 additions & 8 deletions ...integrationTest/java/org/hypertrace/core/documentstore/postgres/PostgresDocStoreTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -99,7 +99,7 @@ public void testUpsert() throws IOException {
query.setFilter(Filter.eq(ID, "default:testKey"));
Iterator<Document> results = collection.search(query);
List<Document> documents = new ArrayList<>();
for (; results.hasNext(); ) {
while (results.hasNext()) {
documents.add(results.next());
}
Assertions.assertFalse(documents.isEmpty());
Expand All @@ -119,7 +119,7 @@ public void testUpsertAndReturn() throws IOException {
}

@Test
public void testBulkUpsert() throws IOException {
public void testBulkUpsert() {
Collection collection = datastore.getCollection(COLLECTION_NAME);
Map<Key, Document> bulkMap = new HashMap<>();
bulkMap.put(new SingleValueKey("default", "testKey1"), createDocument("name", "Bob"));
Expand Down Expand Up @@ -153,6 +153,51 @@ public void testBulkUpsert() throws IOException {
}
}

@Test
public void testBulkUpsertAndReturn() throws IOException {
Collection collection = datastore.getCollection(COLLECTION_NAME);
Map<Key, Document> bulkMap = new HashMap<>();
bulkMap.put(new SingleValueKey("default", "testKey1"), createDocument("name", "Bob"));
bulkMap.put(new SingleValueKey("default", "testKey2"), createDocument("name", "Alice"));
bulkMap.put(new SingleValueKey("default", "testKey3"), createDocument("name", "Alice"));
bulkMap.put(new SingleValueKey("default", "testKey4"), createDocument("name", "Bob"));
bulkMap.put(new SingleValueKey("default", "testKey5"), createDocument("name", "Alice"));
bulkMap.put(
new SingleValueKey("default", "testKey6"), createDocument("email", "[email protected]"));

Iterator<Document> iterator = collection.bulkUpsertAndReturnOlderDocuments(bulkMap);
// Initially there shouldn't be any documents.
Assertions.assertFalse(iterator.hasNext());

// The operation should be idempotent, so go ahead and try again.
iterator = collection.bulkUpsertAndReturnOlderDocuments(bulkMap);
List<Document> documents = new ArrayList<>();
while (iterator.hasNext()) {
documents.add(iterator.next());
}
Assertions.assertEquals(6, documents.size());

{
// empty query returns all the documents
Query query = new Query();
Assertions.assertEquals(6, collection.total(query));
}

{
Query query = new Query();
query.setFilter(Filter.eq("name", "Bob"));
Assertions.assertEquals(2, collection.total(query));
}

{
// limit should not affect the total
Query query = new Query();
query.setFilter(Filter.eq("name", "Bob"));
query.setLimit(1);
Assertions.assertEquals(2, collection.total(query));
}
}

@Test
public void testSubDocumentUpdate() throws IOException {
Collection collection = datastore.getCollection(COLLECTION_NAME);
Expand All @@ -169,7 +214,7 @@ public void testSubDocumentUpdate() throws IOException {
query.setFilter(Filter.eq(ID, "default:testKey"));
Iterator<Document> results = collection.search(query);
List<Document> documents = new ArrayList<>();
for (; results.hasNext(); ) {
while (results.hasNext()) {
documents.add(results.next());
}
Assertions.assertFalse(documents.isEmpty());
Expand Down Expand Up @@ -229,7 +274,7 @@ public void testDeleteAll() throws IOException {
}

@Test
public void testDrop() throws IOException {
public void testDrop() {
Collection collection = datastore.getCollection(COLLECTION_NAME);

Assertions.assertTrue(datastore.listCollections().contains("postgres." + COLLECTION_NAME));
Expand All @@ -249,7 +294,7 @@ public void testIgnoreCaseLikeQuery() throws IOException {
query.setFilter(new Filter(Filter.Op.LIKE, "name", searchValue));
Iterator<Document> results = collection.search(query);
List<Document> documents = new ArrayList<>();
for (; results.hasNext(); ) {
while (results.hasNext()) {
documents.add(results.next());
}
Assertions.assertFalse(documents.isEmpty());
Expand All @@ -271,7 +316,7 @@ public void testSearch() throws IOException {
query.setFilter(new Filter(Filter.Op.EQ, DOCUMENT_ID, key.toString()));
Iterator<Document> results = collection.search(query);
List<Document> documents = new ArrayList<>();
for (; results.hasNext(); ) {
while (results.hasNext()) {
documents.add(results.next());
}
Assertions.assertEquals(documents.size(), 1);
Expand All @@ -291,7 +336,7 @@ public void testSearchForNestedKey() throws IOException {
.setFilter(new Filter(Filter.Op.EQ, "attributes.span_id.value.string", "6449f1f720c93a67"));
Iterator<Document> results = collection.search(query);
List<Document> documents = new ArrayList<>();
for (; results.hasNext(); ) {
while (results.hasNext()) {
documents.add(results.next());
}
Assertions.assertEquals(documents.size(), 1);
Expand Down Expand Up @@ -347,7 +392,7 @@ public void testOffsetLimitAndOrderBY() throws IOException {

Iterator<Document> results = collection.search(query);
List<Document> documents = new ArrayList<>();
for (; results.hasNext(); ) {
while (results.hasNext()) {
documents.add(results.next());
}

Expand Down
7 changes: 7 additions & 0 deletions document-store/src/main/java/org/hypertrace/core/documentstore/Collection.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,13 @@ public interface Collection {
*/
boolean bulkUpsert(Map<Key, Document> documents);

/**
* Method to bulkUpsert the given documents and return the previous copies of those documents.
* This helps the clients to see how the documents were prior to upserting them and do that
* in one less round trip.
*/
Iterator<Document> bulkUpsertAndReturnOlderDocuments(Map<Key, Document> documents) throws IOException;

/**
* Drops a collections
*/
Expand Down
64 changes: 50 additions & 14 deletions document-store/src/main/java/org/hypertrace/core/documentstore/mongo/MongoCollection.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.function.Function;
import java.util.stream.Collectors;
import net.jodah.failsafe.Failsafe;
Expand Down Expand Up @@ -386,28 +387,58 @@ public long total(Query query) {
@Override
public boolean bulkUpsert(Map<Key, Document> documents) {
try {
List<UpdateOneModel<BasicDBObject>> bulkCollection = new ArrayList<>();
for (Entry<Key, Document> entry : documents.entrySet()) {
Key key = entry.getKey();
// insert or overwrite
bulkCollection.add(new UpdateOneModel<>(
this.selectionCriteriaForKey(key),
prepareUpsert(key, entry.getValue()),
new UpdateOptions().upsert(true)));
}

BulkWriteResult result = Failsafe.with(bulkWriteRetryPolicy)
.get(() -> collection.bulkWrite(bulkCollection, new BulkWriteOptions().ordered(false)));
BulkWriteResult result = bulkUpsertImpl(documents);
LOGGER.debug(result.toString());

return true;

} catch (IOException | MongoServerException e) {
LOGGER.error("Error during bulk upsert for documents:{}", documents, e);
return false;
}
}

private BulkWriteResult bulkUpsertImpl(Map<Key, Document> documents) throws JsonProcessingException {
List<UpdateOneModel<BasicDBObject>> bulkCollection = new ArrayList<>();
for (Entry<Key, Document> entry : documents.entrySet()) {
Key key = entry.getKey();
// insert or overwrite
bulkCollection.add(new UpdateOneModel<>(
this.selectionCriteriaForKey(key),
prepareUpsert(key, entry.getValue()),
new UpdateOptions().upsert(true)));
}

return Failsafe.with(bulkWriteRetryPolicy)
.get(() -> collection.bulkWrite(bulkCollection, new BulkWriteOptions().ordered(false)));
}

@Override
public Iterator<Document> bulkUpsertAndReturnOlderDocuments(Map<Key, Document> documents) throws IOException {
try {
// First get all the documents for the given keys.
FindIterable<BasicDBObject> cursor = collection.find(selectionCriteriaForKeys(documents.keySet()));
final MongoCursor<BasicDBObject> mongoCursor = cursor.cursor();

// Now go ahead and do the bulk upsert.
BulkWriteResult result = bulkUpsertImpl(documents);
LOGGER.debug(result.toString());
Copy link
Contributor

@jcchavezs jcchavezs Dec 17, 2020 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am not sure about the performance of this. Not a Javaer here but it seams whether logger debug is enabled or not we still turn it into string? cc @kotharironak

Copy link
Contributor

@aaron-steinfeld aaron-steinfeld Dec 17, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

👍 we should always be letting the logger doing the stringification for us so we don't have to eat this cost unless the message is needed. That means wrapping it in an if or IMO, more graceful to do LOGGER.debug("{}", result);

Copy link
Contributor Author

@buchi-busireddy buchi-busireddy Dec 18, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixing in a new PR. Thanks!

Copy link
Contributor Author

@buchi-busireddy buchi-busireddy Jan 3, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

#25


return new Iterator<>() {
@Override
public boolean hasNext() {
return mongoCursor.hasNext();
}

@Override
public Document next() {
return MongoCollection.this.dbObjectToDocument(mongoCursor.next());
}
};
} catch (JsonProcessingException e) {
LOGGER.error("Error during bulk upsert for documents:{}", documents, e);
Copy link
Contributor

@jcchavezs jcchavezs Dec 17, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is more like an in general comment. I am usually in favour of either log the error and handle it or bubble up the exception but not both of them because they usually flood the logs. Also, do we want to print the full set of documents in logs? How about privacy concerns and also efficient usage of the log storage? I don't thing dumping the failing documents in the logs is actionable either.

Copy link
Contributor Author

@buchi-busireddy buchi-busireddy Jan 3, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Raised #25

throw new IOException("Error during bulk upsert.");
Copy link
Contributor

@jcchavezs jcchavezs Dec 17, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not passing the previous exception makes us loosing all the context on this error. Is there any reason for not doing it?

Copy link
Contributor Author

@buchi-busireddy buchi-busireddy Dec 18, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was thinking the API should mask the implementation specific exception details but this is actually a library so I'll fix it.

Copy link
Contributor Author

@buchi-busireddy buchi-busireddy Jan 3, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Raised #25

}
}

@Override
public void drop() {
collection.drop();
Expand All @@ -417,6 +448,11 @@ private BasicDBObject selectionCriteriaForKey(Key key) {
return new BasicDBObject(ID_KEY, key.toString());
}

private BasicDBObject selectionCriteriaForKeys(Set<Key> keys) {
Copy link
Contributor

@jcchavezs jcchavezs Dec 17, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can't this be private? A simple inspection tells me yes it does but I am not 100% sure. Tho ID_KEY is static.

Copy link
Contributor Author

@buchi-busireddy buchi-busireddy Dec 18, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think you meant to ask about static right? Fixing it.

Copy link
Contributor

@jcchavezs jcchavezs Dec 18, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yeah I meant static, sorry.

Copy link
Contributor Author

@buchi-busireddy buchi-busireddy Jan 3, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Raised #25

return new BasicDBObject(Map.of(ID_KEY, new BasicDBObject("$in",
keys.stream().map(Key::toString).collect(Collectors.toList()))));
}

private Document dbObjectToDocument(BasicDBObject dbObject) {
try {
// Hack: Remove the _id field since it's an unrecognized field for Proto layer.
Expand Down
Loading