Skip to content

Commit a7fd620

Browse files
YongwuHeYongwuHe
authored
feat: mongo key holder (#469)
1 parent e439eb0 commit a7fd620

File tree

4 files changed

+120
-14
lines changed

4 files changed

+120
-14
lines changed

‎arex-instrumentation/database/arex-database-common/src/main/java/io/arex/inst/database/common/DatabaseExtractor.java

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -103,13 +103,13 @@ public MockResult replay(String serializer) {
103103
replayResult = Serializer.deserialize(replayMocker.getTargetResponse().getBody(),
104104
replayMocker.getTargetResponse().getType(), serializer);
105105
}
106-
107-
if (replayResult != null) {
108-
// restore keyHolder
109-
setKeyHolder(replayMocker.getTargetResponse().attributeAsString(KEY_HOLDER));
110-
setPage(replayMocker.getTargetResponse().attributeAsString(PAGE_NAME));
111-
setKeyHolderName(replayMocker.getTargetResponse().attributeAsString(KEY_HOLDER_NAME));
112-
}
106+
}
107+
// compatible with methods whose return type is void but need to restore the keyHolder. ex: mongo insert(version < 4.0.1)
108+
if (replayMocker != null) {
109+
// restore keyHolder
110+
setKeyHolder(replayMocker.getTargetResponse().attributeAsString(KEY_HOLDER));
111+
setPage(replayMocker.getTargetResponse().attributeAsString(PAGE_NAME));
112+
setKeyHolderName(replayMocker.getTargetResponse().attributeAsString(KEY_HOLDER_NAME));
113113
}
114114

115115
return MockResult.success(ignoreMockResult, replayResult);

‎arex-instrumentation/database/arex-database-mongo/src/main/java/io/arex/inst/database/mongo/MongoHelper.java

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,20 @@
11
package io.arex.inst.database.mongo;
22

3+
import com.google.common.collect.Sets;
34
import com.mongodb.MongoNamespace;
45

56
import io.arex.agent.bootstrap.model.MockResult;
7+
import io.arex.agent.bootstrap.util.StringUtil;
68
import io.arex.inst.database.common.DatabaseExtractor;
79
import io.arex.inst.database.mongo.wrapper.ResultWrapper;
810
import io.arex.inst.runtime.serializer.Serializer;
11+
import io.arex.inst.runtime.util.TypeUtil;
912
import org.bson.Document;
1013

14+
import java.util.ArrayList;
15+
import java.util.List;
16+
import java.util.Set;
17+
1118
import static io.arex.inst.runtime.model.ArexConstants.GSON_REQUEST_SERIALIZER;
1219
import static io.arex.inst.runtime.model.ArexConstants.GSON_SERIALIZER;
1320

@@ -16,6 +23,12 @@
1623
* ex: FindOperation
1724
*/
1825
public class MongoHelper {
26+
private static final String ID_FIELD = "_id";
27+
private static final Set<String> INSERT_METHODS = Sets.newHashSet("executeInsertOne", "executeInsertMany");
28+
/**
29+
* used to separate keyHolder and keyHolderType
30+
*/
31+
private static final char KEYHOLDER_TYPE_SEPARATOR = ';';
1932
private MongoHelper() {
2033
}
2134
static {
@@ -27,12 +40,17 @@ public static MockResult replay(String methodName, MongoNamespace namespace, Obj
2740
final DatabaseExtractor extractor = new DatabaseExtractor(dbName, methodName, parameter, methodName);
2841
final MockResult mockResult = extractor.replay(GSON_SERIALIZER);
2942

43+
if (INSERT_METHODS.contains(methodName)) {
44+
restoreKeyHolder(extractor.getKeyHolder(), filter);
45+
}
46+
3047
if (mockResult == null || mockResult.getThrowable() != null) {
3148
return mockResult;
3249
}
3350

3451
return ResultWrapper.unwrap(mockResult);
3552
}
53+
3654
public static void record(String methodName, MongoNamespace namespace, Object filter, Object result, Throwable throwable) {
3755
String dbName = namespace.getFullName();
3856
String parameter = Serializer.serialize(filter, GSON_REQUEST_SERIALIZER);
@@ -41,7 +59,62 @@ public static void record(String methodName, MongoNamespace namespace, Object fi
4159
extractor.recordDb(throwable, GSON_SERIALIZER);
4260
return;
4361
}
62+
63+
if (INSERT_METHODS.contains(methodName)) {
64+
saveKeyHolder(extractor, filter);
65+
}
66+
4467
Object actualResult = ResultWrapper.wrap(result);
4568
extractor.recordDb(actualResult, GSON_SERIALIZER);
4669
}
70+
71+
72+
private static void restoreKeyHolder(String keyHolder, Object filter) {
73+
if (StringUtil.isEmpty(keyHolder)) {
74+
return;
75+
}
76+
String[] keyHolderArray = StringUtil.split(keyHolder, KEYHOLDER_TYPE_SEPARATOR);
77+
if (keyHolderArray.length != 2) {
78+
return;
79+
}
80+
Object insertId = Serializer.deserialize(keyHolderArray[0], keyHolderArray[1], GSON_SERIALIZER);
81+
restoreInsertId(filter, insertId);
82+
}
83+
84+
private static void restoreInsertId(Object filter, Object insertId) {
85+
if (insertId == null) {
86+
return;
87+
}
88+
if (filter instanceof Document) {
89+
((Document) filter).put(ID_FIELD, insertId);
90+
} else if (filter instanceof List) {
91+
List<?> filterList = (List<?>) filter;
92+
for (int i = 0; i < filterList.size(); i++) {
93+
restoreInsertId(filterList.get(i), ((List<?>) insertId).get(i));
94+
}
95+
}
96+
}
97+
98+
private static void saveKeyHolder(DatabaseExtractor extractor, Object filter) {
99+
Object insertId = getInsertId(filter);
100+
if (insertId == null) {
101+
return;
102+
}
103+
extractor.setKeyHolder(Serializer.serialize(insertId, GSON_SERIALIZER) + KEYHOLDER_TYPE_SEPARATOR + TypeUtil.getName(insertId));
104+
}
105+
106+
private static Object getInsertId(Object filter) {
107+
if (filter instanceof Document) {
108+
Document document = (Document) filter;
109+
return document.get(ID_FIELD);
110+
}
111+
if (filter instanceof List) {
112+
List<Object> insertIds = new ArrayList<>(((List<?>) filter).size());
113+
for (Object o : (List<?>) filter) {
114+
insertIds.add(getInsertId(o));
115+
}
116+
return insertIds;
117+
}
118+
return null;
119+
}
47120
}

‎arex-instrumentation/database/arex-database-mongo/src/main/java/io/arex/inst/database/mongo/WriteOperationInstrumentation.java

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -74,9 +74,6 @@ public static boolean onEnter(@Advice.FieldValue("namespace") MongoNamespace nam
7474
RepeatedCollectManager.enter();
7575
}
7676
if (ContextManager.needReplay()) {
77-
if (void.class.isAssignableFrom(method.getReturnType())) {
78-
return true;
79-
}
8077
mockResult = MongoHelper.replay(methodName, namespace, filter);
8178
return mockResult != null && mockResult.notIgnoreMockResult();
8279
}
@@ -100,7 +97,7 @@ public static void onExit(@Advice.Origin Method method,
10097
return;
10198
}
10299

103-
if (ContextManager.needRecord() && RepeatedCollectManager.exitAndValidate() && !void.class.isAssignableFrom(method.getReturnType())) {
100+
if (ContextManager.needRecord() && RepeatedCollectManager.exitAndValidate()) {
104101
MongoHelper.record(methodName, namespace, filter, result, throwable);
105102
}
106103
}

‎arex-instrumentation/database/arex-database-mongo/src/test/java/io/arex/inst/database/mongo/MongoHelperTest.java

Lines changed: 39 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package io.arex.inst.database.mongo;
22

33
import static org.junit.jupiter.api.Assertions.*;
4+
import static org.mockito.ArgumentMatchers.anyString;
45

56
import com.mongodb.MongoNamespace;
67
import io.arex.agent.bootstrap.model.MockResult;
@@ -10,6 +11,12 @@
1011
import io.arex.inst.runtime.serializer.Serializer.Builder;
1112
import io.arex.inst.runtime.serializer.StringSerializable;
1213
import java.lang.reflect.Field;
14+
import java.util.ArrayList;
15+
import java.util.List;
16+
17+
import io.arex.inst.runtime.util.TypeUtil;
18+
import org.bson.Document;
19+
import org.bson.types.ObjectId;
1320
import org.junit.jupiter.api.AfterAll;
1421
import org.junit.jupiter.api.BeforeAll;
1522
import org.junit.jupiter.api.Test;
@@ -18,11 +25,10 @@
1825
import org.mockito.Mockito;
1926

2027
class MongoHelperTest {
21-
28+
static StringSerializable serializable = Mockito.mock(StringSerializable.class);
2229

2330
@BeforeAll
2431
static void setUp() {
25-
final StringSerializable serializable = Mockito.mock(StringSerializable.class);
2632
final Builder builder = Serializer.builder(serializable);
2733
builder.addSerializer("gson", serializable);
2834
builder.build();
@@ -34,6 +40,7 @@ static void tearDown() throws Exception {
3440
final Field instance = Serializer.class.getDeclaredField("INSTANCE");
3541
instance.setAccessible(true);
3642
instance.set(null, null);
43+
serializable = null;
3744
}
3845

3946
@Test
@@ -59,6 +66,24 @@ void replay() {
5966
})) {
6067
assertNull(MongoHelper.replay("test", Mockito.mock(MongoNamespace.class), null));
6168
}
69+
70+
// restore keyHolder
71+
try(final MockedConstruction<DatabaseExtractor> construction = Mockito.mockConstruction(DatabaseExtractor.class, (mock, context) ->{
72+
Mockito.when(mock.replay("gson")).thenReturn(null);
73+
String keyHolder = "test;java.lang.String";
74+
Mockito.when(mock.getKeyHolder()).thenReturn(keyHolder);
75+
})) {
76+
List<ObjectId> insertIds = new ArrayList<>();
77+
ObjectId objectId = new ObjectId();
78+
insertIds.add(objectId);
79+
List<Document> documentList = new ArrayList<>();
80+
documentList.add(new Document());
81+
Mockito.when(serializable.deserialize("test", TypeUtil.forName("java.lang.String"))).thenReturn(insertIds);
82+
assertNull(MongoHelper.replay("executeInsertMany", Mockito.mock(MongoNamespace.class), documentList));
83+
assertEquals(objectId, documentList.get(0).get("_id"));
84+
} catch (Throwable e) {
85+
e.printStackTrace();
86+
}
6287
}
6388

6489
@Test
@@ -69,7 +94,18 @@ void record() {
6994
final RuntimeException runtimeException = new RuntimeException();
7095
MongoHelper.record("test", Mockito.mock(MongoNamespace.class), null, "test", runtimeException);
7196
Mockito.verify(construction.constructed().get(1), Mockito.times(1)).recordDb(runtimeException, "gson");
97+
// insert save keyHolder
98+
List<Document> documentList = new ArrayList<>();
99+
Document document = new Document();
100+
ObjectId objectId = new ObjectId();
101+
document.put("_id", objectId);
102+
documentList.add(document);
103+
MongoHelper.record("executeInsertMany", Mockito.mock(MongoNamespace.class), documentList, "test", null);
104+
Mockito.verify(construction.constructed().get(2), Mockito.times(1)).setKeyHolder(anyString());
105+
// no Document
106+
MongoHelper.record("executeInsertMany", Mockito.mock(MongoNamespace.class), "notDocument", "test", null);
107+
Mockito.verify(construction.constructed().get(3), Mockito.never()).setKeyHolder(anyString());
72108
}
73109

74110
}
75-
}
111+
}

0 commit comments

Comments
 (0)