Fix count and update in MongoDb stores

* The `into()` query can't infer the type for entity if we provide `Object.class`
* The `updateFirst()` does not support sort queries any more - replace with `findAndModify()`
* Add `getMessageGroupCount()` into tests

**Cherry-pick to 5.3.x & 5.2.x**

Author: artembilan

spring-integration-mongodb/src/main/java/org/springframework/integration/mongodb/store/ConfigurableMongoDbMessageStore.java

@@ -20,12 +20,14 @@
 import java.util.Collection;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
 import java.util.UUID;
 import java.util.stream.Collectors;
 import java.util.stream.StreamSupport;
 
 import org.springframework.data.domain.Sort;
 import org.springframework.data.mongodb.MongoDatabaseFactory;
+import org.springframework.data.mongodb.core.FindAndModifyOptions;
 import org.springframework.data.mongodb.core.MongoTemplate;
 import org.springframework.data.mongodb.core.convert.MappingMongoConverter;
 import org.springframework.data.mongodb.core.query.Criteria;
@@ -247,9 +249,8 @@ public int getMessageCountForAllMessageGroups() {
 	@ManagedAttribute
 	public int getMessageGroupCount() {
 		Query query = Query.query(Criteria.where(MessageDocumentFields.GROUP_ID).exists(true));
-		return getMongoTemplate().getCollection(this.collectionName)
-				.distinct(MessageDocumentFields.GROUP_ID, query.getQueryObject(), Object.class)
-				.into(new ArrayList<>())
+		return getMongoTemplate()
+				.findDistinct(query, MessageDocumentFields.GROUP_ID, this.collectionName, Object.class)
 				.size();
 	}
 
@@ -278,7 +279,9 @@ public Collection<Message<?>> getMessagesForGroup(Object groupId) {
 	}
 
 	private void updateGroup(Object groupId, Update update) {
-		getMongoTemplate().updateFirst(groupOrderQuery(groupId), update, this.collectionName);
+		getMongoTemplate()
+				.findAndModify(groupOrderQuery(groupId), update, FindAndModifyOptions.none(), Map.class,
+						this.collectionName);
 	}
 
 	private static Update lastModifiedUpdate() {

spring-integration-mongodb/src/main/java/org/springframework/integration/mongodb/store/MongoDbMessageStore.java

@@ -428,20 +428,15 @@ public Collection<Message<?>> getMessagesForGroup(Object groupId) {
 	@Override
 	@ManagedAttribute
 	public int getMessageCountForAllMessageGroups() {
-		Query query = Query.query(Criteria.where(MessageDocumentFields.MESSAGE_ID).exists(true)
-				.and(MessageDocumentFields.GROUP_ID).exists(true));
-		long count = this.template.count(query, this.collectionName);
-		Assert.isTrue(count <= Integer.MAX_VALUE, "Message count is out of Integer's range");
-		return (int) count;
+		Query query = Query.query(Criteria.where(GROUP_ID_KEY).exists(true));
+		return (int) this.template.count(query, this.collectionName);
 	}
 
 	@Override
 	@ManagedAttribute
 	public int getMessageGroupCount() {
-		Query query = Query.query(Criteria.where(MessageDocumentFields.GROUP_ID).exists(true));
-		return this.template.getCollection(this.collectionName)
-				.distinct(MessageDocumentFields.GROUP_ID, query.getQueryObject(), Object.class)
-				.into(new ArrayList<>())
+		Query query = Query.query(Criteria.where(GROUP_ID_KEY).exists(true));
+		return this.template.findDistinct(query, GROUP_ID_KEY, this.collectionName, Object.class)
 				.size();
 	}
 
@@ -471,7 +466,7 @@ private static Query whereGroupIdIs(Object groupId) {
 
 	private void updateGroup(Object groupId, Update update) {
 		Query query = whereGroupIdIs(groupId).with(Sort.by(Sort.Direction.DESC, GROUP_UPDATE_TIMESTAMP_KEY, SEQUENCE));
-		this.template.updateFirst(query, update, this.collectionName);
+		this.template.findAndModify(query, update, FindAndModifyOptions.none(), Map.class, this.collectionName);
 	}
 
 	private int getNextId() {

spring-integration-mongodb/src/test/java/org/springframework/integration/mongodb/store/AbstractMongoDbMessageGroupStoreTests.java

@@ -95,6 +95,7 @@ public void testMessageGroupWithAddedMessagePrimitiveGroupId() {
 		assertThat(messageA.getHeaders().getId()).isEqualTo(retrievedMessage.getHeaders().getId());
 		// ensure that 'message_group' header that is only used internally is not propagated
 		assertThat(retrievedMessage.getHeaders().get("message_group")).isNull();
+		assertThat(store.getMessageGroupCount()).isEqualTo(1);
 	}
 
 	@Test