Skip to content

Commit 622643b

Browse files
christophstroblmp911de
authored andcommitted
DATAMONGO-2130 - Fix Repository count & exists inside transaction.
We now make sure invocations on repository count and exists methods delegate to countDocuments when inside a transaction. Original pull request: #618.
1 parent 51cc55b commit 622643b

File tree

8 files changed

+155
-19
lines changed

8 files changed

+155
-19
lines changed

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/MongoDatabaseUtils.java

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -110,14 +110,34 @@ private static MongoDatabase doGetMongoDatabase(@Nullable String dbName, MongoDb
110110

111111
ClientSession session = doGetSession(factory, sessionSynchronization);
112112

113-
if(session == null) {
113+
if (session == null) {
114114
return StringUtils.hasText(dbName) ? factory.getDb(dbName) : factory.getDb();
115115
}
116116

117117
MongoDbFactory factoryToUse = factory.withSession(session);
118118
return StringUtils.hasText(dbName) ? factoryToUse.getDb(dbName) : factoryToUse.getDb();
119119
}
120120

121+
/**
122+
* Check if the {@link MongoDbFactory} is actually bound to a {@link ClientSession} that has an active transaction, or
123+
* if a {@link TransactionSynchronization} has been registered for the {@link MongoDbFactory resource} and if the
124+
* associated {@link ClientSession} has an {@link ClientSession#hasActiveTransaction() active transaction}.
125+
*
126+
* @param dbFactory the resource to check transactions for. Must not be {@literal null}.
127+
* @return {@literal true} if the factory has an ongoing transaction.
128+
* @since 2.1.3
129+
*/
130+
public static boolean isTransactionActive(MongoDbFactory dbFactory) {
131+
132+
if (dbFactory.isTransactionActive()) {
133+
return true;
134+
}
135+
136+
MongoResourceHolder resourceHolder = (MongoResourceHolder) TransactionSynchronizationManager.getResource(dbFactory);
137+
return resourceHolder != null
138+
&& (resourceHolder.hasSession() && resourceHolder.getSession().hasActiveTransaction());
139+
}
140+
121141
@Nullable
122142
private static ClientSession doGetSession(MongoDbFactory dbFactory, SessionSynchronization sessionSynchronization) {
123143

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/MongoDbFactory.java

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,4 +108,15 @@ default MongoDbFactory withSession(ClientSessionOptions options) {
108108
* @since 2.1
109109
*/
110110
MongoDbFactory withSession(ClientSession session);
111+
112+
/**
113+
* Returns if the given {@link MongoDbFactory} is bound to a {@link ClientSession} that has an
114+
* {@link ClientSession#hasActiveTransaction() active transaction}.
115+
*
116+
* @return {@literal true} if there's an active transaction, {@literal false} otherwise.
117+
* @since 2.1.3
118+
*/
119+
default boolean isTransactionActive() {
120+
return false;
121+
}
111122
}

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoDbFactorySupport.java

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -233,6 +233,15 @@ public MongoDbFactory withSession(ClientSession session) {
233233
return delegate.withSession(session);
234234
}
235235

236+
/*
237+
* (non-Javadoc)
238+
* @see org.springframework.data.mongodb.MongoDbFactory#isTransactionActive()
239+
*/
240+
@Override
241+
public boolean isTransactionActive() {
242+
return session != null && session.hasActiveTransaction();
243+
}
244+
236245
private MongoDatabase proxyMongoDatabase(MongoDatabase database) {
237246
return createProxyInstance(session, database, MongoDatabase.class);
238247
}
@@ -241,7 +250,8 @@ private MongoDatabase proxyDatabase(com.mongodb.session.ClientSession session, M
241250
return createProxyInstance(session, database, MongoDatabase.class);
242251
}
243252

244-
private MongoCollection<?> proxyCollection(com.mongodb.session.ClientSession session, MongoCollection<?> collection) {
253+
private MongoCollection<?> proxyCollection(com.mongodb.session.ClientSession session,
254+
MongoCollection<?> collection) {
245255
return createProxyInstance(session, collection, MongoCollection.class);
246256
}
247257

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoTemplate.java

Lines changed: 19 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1119,7 +1119,16 @@ public long count(Query query, @Nullable Class<?> entityClass, String collection
11191119
Document document = queryMapper.getMappedObject(query.getQueryObject(),
11201120
Optional.ofNullable(entityClass).map(it -> mappingContext.getPersistentEntity(entityClass)));
11211121

1122-
return execute(collectionName, collection -> collection.count(document, options));
1122+
return doCount(collectionName, document, options);
1123+
}
1124+
1125+
protected long doCount(String collectionName, Document filter, CountOptions options) {
1126+
1127+
if (!MongoDatabaseUtils.isTransactionActive(getMongoDbFactory())) {
1128+
return execute(collectionName, collection -> collection.count(filter, options));
1129+
}
1130+
1131+
return execute(collectionName, collection -> collection.countDocuments(filter, options));
11231132
}
11241133

11251134
/*
@@ -2820,21 +2829,23 @@ public FindIterable<Document> doInCollection(MongoCollection<Document> collectio
28202829
}
28212830

28222831
/**
2823-
* Optimized {@link CollectionCallback} that takes an already mappend query and a nullable
2832+
* Optimized {@link CollectionCallback} that takes an already mapped query and a nullable
28242833
* {@link com.mongodb.client.model.Collation} to execute a count query limited to one element.
28252834
*
28262835
* @author Christoph Strobl
28272836
* @since 2.0
28282837
*/
28292838
@RequiredArgsConstructor
2830-
private static class ExistsCallback implements CollectionCallback<Boolean> {
2839+
private class ExistsCallback implements CollectionCallback<Boolean> {
28312840

28322841
private final Document mappedQuery;
28332842
private final com.mongodb.client.model.Collation collation;
28342843

28352844
@Override
28362845
public Boolean doInCollection(MongoCollection<Document> collection) throws MongoException, DataAccessException {
2837-
return collection.count(mappedQuery, new CountOptions().limit(1).collation(collation)) > 0;
2846+
2847+
return doCount(collection.getNamespace().getCollectionName(), mappedQuery,
2848+
new CountOptions().limit(1).collation(collation)) > 0;
28382849
}
28392850
}
28402851

@@ -3343,23 +3354,16 @@ public MongoDatabase getDb() {
33433354

33443355
/*
33453356
* (non-Javadoc)
3346-
* @see org.springframework.data.mongodb.core.MongoTemplate#count(org.springframework.data.mongodb.core.query.Query, java.lang.Class, java.lang.String)
3357+
* @see org.springframework.data.mongodb.core.MongoTemplate#doCount(java.lang.String, org.bson.Document, com.mongodb.client.model.CountOptions)
33473358
*/
33483359
@Override
3349-
@SuppressWarnings("unchecked")
3350-
public long count(Query query, @Nullable Class<?> entityClass, String collectionName) {
3360+
protected long doCount(String collectionName, Document filter, CountOptions options) {
33513361

33523362
if (!session.hasActiveTransaction()) {
3353-
return super.count(query, entityClass, collectionName);
3363+
return super.doCount(collectionName, filter, options);
33543364
}
33553365

3356-
CountOptions options = new CountOptions();
3357-
query.getCollation().map(Collation::toMongoCollation).ifPresent(options::collation);
3358-
3359-
Document document = delegate.queryMapper.getMappedObject(query.getQueryObject(),
3360-
Optional.ofNullable(entityClass).map(it -> delegate.mappingContext.getPersistentEntity(entityClass)));
3361-
3362-
return execute(collectionName, collection -> collection.countDocuments(document, options));
3366+
return execute(collectionName, collection -> collection.countDocuments(filter, options));
33633367
}
33643368
}
33653369
}

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/repository/support/SimpleMongoRepository.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -137,7 +137,7 @@ public boolean existsById(ID id) {
137137
*/
138138
@Override
139139
public long count() {
140-
return mongoOperations.getCollection(entityInformation.getCollectionName()).count();
140+
return mongoOperations.count(new Query(), entityInformation.getCollectionName());
141141
}
142142

143143
/*

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/MongoDatabaseUtilsUnitTests.java

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,39 @@ public void verifyTransactionSynchronizationManagerState() {
7878
assertFalse(TransactionSynchronizationManager.isActualTransactionActive());
7979
}
8080

81+
@Test // DATAMONGO-2130
82+
public void isTransactionActiveShouldDetectTxViaFactory() {
83+
84+
when(dbFactory.isTransactionActive()).thenReturn(true);
85+
86+
assertThat(MongoDatabaseUtils.isTransactionActive(dbFactory)).isTrue();
87+
}
88+
89+
@Test // DATAMONGO-2130
90+
public void isTransactionActiveShouldReturnFalseIfNoTxActive() {
91+
92+
when(dbFactory.isTransactionActive()).thenReturn(false);
93+
94+
assertThat(MongoDatabaseUtils.isTransactionActive(dbFactory)).isFalse();
95+
}
96+
97+
@Test // DATAMONGO-2130
98+
public void isTransactionActiveShouldLookupTxForActiveTransactionSynchronizationViaTxManager() {
99+
100+
when(dbFactory.isTransactionActive()).thenReturn(false);
101+
102+
MongoTransactionManager txManager = new MongoTransactionManager(dbFactory);
103+
TransactionTemplate txTemplate = new TransactionTemplate(txManager);
104+
105+
txTemplate.execute(new TransactionCallbackWithoutResult() {
106+
107+
@Override
108+
protected void doInTransactionWithoutResult(TransactionStatus transactionStatus) {
109+
assertThat(MongoDatabaseUtils.isTransactionActive(dbFactory)).isTrue();
110+
}
111+
});
112+
}
113+
81114
@Test // DATAMONGO-1920
82115
public void shouldNotStartSessionWhenNoTransactionOngoing() {
83116

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateUnitTests.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import static org.springframework.data.mongodb.core.aggregation.Aggregation.*;
2323
import static org.springframework.data.mongodb.test.util.IsBsonObject.*;
2424

25+
import com.mongodb.MongoNamespace;
2526
import lombok.Data;
2627

2728
import java.math.BigInteger;
@@ -135,6 +136,7 @@ public void setUp() {
135136
when(collection.find(any(org.bson.Document.class), any(Class.class))).thenReturn(findIterable);
136137
when(collection.mapReduce(any(), any(), eq(Document.class))).thenReturn(mapReduceIterable);
137138
when(collection.count(any(Bson.class), any(CountOptions.class))).thenReturn(1L);
139+
when(collection.getNamespace()).thenReturn(new MongoNamespace("db.mock-collection"));
138140
when(collection.aggregate(any(List.class), any())).thenReturn(aggregateIterable);
139141
when(collection.withReadPreference(any())).thenReturn(collection);
140142
when(findIterable.projection(any())).thenReturn(findIterable);

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/repository/support/SimpleMongoRepositoryTests.java

Lines changed: 57 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
package org.springframework.data.mongodb.repository.support;
1717

1818
import static org.assertj.core.api.Assertions.*;
19+
import static org.assertj.core.api.Assumptions.*;
1920
import static org.springframework.data.domain.ExampleMatcher.*;
2021

2122
import java.util.ArrayList;
@@ -28,14 +29,16 @@
2829
import java.util.UUID;
2930

3031
import org.junit.Before;
32+
import org.junit.Rule;
3133
import org.junit.Test;
3234
import org.junit.runner.RunWith;
3335
import org.springframework.beans.factory.annotation.Autowired;
3436
import org.springframework.data.domain.Example;
35-
import org.springframework.data.domain.ExampleMatcher.StringMatcher;
3637
import org.springframework.data.domain.Page;
3738
import org.springframework.data.domain.PageRequest;
39+
import org.springframework.data.domain.ExampleMatcher.*;
3840
import org.springframework.data.geo.Point;
41+
import org.springframework.data.mongodb.MongoTransactionManager;
3942
import org.springframework.data.mongodb.core.MongoTemplate;
4043
import org.springframework.data.mongodb.core.geo.GeoJsonPoint;
4144
import org.springframework.data.mongodb.core.mapping.Document;
@@ -44,9 +47,13 @@
4447
import org.springframework.data.mongodb.repository.Person.Sex;
4548
import org.springframework.data.mongodb.repository.User;
4649
import org.springframework.data.mongodb.repository.query.MongoEntityInformation;
50+
import org.springframework.data.mongodb.test.util.MongoVersion;
51+
import org.springframework.data.mongodb.test.util.MongoVersionRule;
52+
import org.springframework.data.mongodb.test.util.ReplicaSet;
4753
import org.springframework.test.context.ContextConfiguration;
4854
import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
4955
import org.springframework.test.util.ReflectionTestUtils;
56+
import org.springframework.transaction.support.TransactionTemplate;
5057

5158
/**
5259
* @author A. B. M. Kowser
@@ -59,6 +66,7 @@
5966
public class SimpleMongoRepositoryTests {
6067

6168
@Autowired private MongoTemplate template;
69+
public @Rule MongoVersionRule mongoVersion = MongoVersionRule.any();
6270

6371
private Person oliver, dave, carter, boyd, stefan, leroi, alicia;
6472
private List<Person> all;
@@ -383,6 +391,54 @@ public void saveAllUsesEntityCollection() {
383391
assertThat(repository.findAll()).containsExactlyInAnyOrder(first, second);
384392
}
385393

394+
@Test // DATAMONGO-2130
395+
@MongoVersion(asOf = "4.0")
396+
public void countShouldBePossibleInTransaction() {
397+
398+
assumeThat(ReplicaSet.required().runsAsReplicaSet()).isTrue();
399+
400+
MongoTransactionManager txmgr = new MongoTransactionManager(template.getMongoDbFactory());
401+
TransactionTemplate tt = new TransactionTemplate(txmgr);
402+
tt.afterPropertiesSet();
403+
404+
long countPreTx = repository.count();
405+
406+
long count = tt.execute(status -> {
407+
408+
Person sample = new Person();
409+
sample.setLastname("Matthews");
410+
411+
repository.save(sample);
412+
413+
return repository.count();
414+
});
415+
416+
assertThat(count).isEqualTo(countPreTx + 1);
417+
}
418+
419+
@Test // DATAMONGO-2130
420+
@MongoVersion(asOf = "4.0")
421+
public void existsShouldBePossibleInTransaction() {
422+
423+
assumeThat(ReplicaSet.required().runsAsReplicaSet()).isTrue();
424+
425+
MongoTransactionManager txmgr = new MongoTransactionManager(template.getMongoDbFactory());
426+
TransactionTemplate tt = new TransactionTemplate(txmgr);
427+
tt.afterPropertiesSet();
428+
429+
boolean exists = tt.execute(status -> {
430+
431+
Person sample = new Person();
432+
sample.setLastname("Matthews");
433+
434+
repository.save(sample);
435+
436+
return repository.existsById(sample.getId());
437+
});
438+
439+
assertThat(exists).isTrue();
440+
}
441+
386442
private void assertThatAllReferencePersonsWereStoredCorrectly(Map<String, Person> references, List<Person> saved) {
387443

388444
for (Person person : saved) {

0 commit comments

Comments
 (0)