DATAMONGO-2130 - Polishing.

Replace duplicate checks to ClientSession.hasActiveTransaction() with MongoResourceHolder.hasActiveTransaction(). Introduce MongoResourceHolder.getRequiredSession() to avoid nullability warnings.

Original pull request: #618.

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

@@ -134,8 +134,7 @@ public class MongoDatabaseUtils {
 		}
 
 		MongoResourceHolder resourceHolder = (MongoResourceHolder) TransactionSynchronizationManager.getResource(dbFactory);
-		return resourceHolder != null
-				&& (resourceHolder.hasSession() && resourceHolder.getSession().hasActiveTransaction());
+		return resourceHolder != null && resourceHolder.hasActiveTransaction();
 	}
 
 	@Nullable
@@ -160,7 +159,7 @@ public class MongoDatabaseUtils {
 		// init a non native MongoDB transaction by registering a MongoSessionSynchronization
 
 		resourceHolder = new MongoResourceHolder(createClientSession(dbFactory), dbFactory);
-		resourceHolder.getSession().startTransaction();
+		resourceHolder.getRequiredSession().startTransaction();
 
 		TransactionSynchronizationManager
 				.registerSynchronization(new MongoSessionSynchronization(resourceHolder, dbFactory));
@@ -207,8 +206,8 @@ public class MongoDatabaseUtils {
 		@Override
 		protected void processResourceAfterCommit(MongoResourceHolder resourceHolder) {
 
-			if (isTransactionActive(resourceHolder)) {
-				resourceHolder.getSession().commitTransaction();
+			if (resourceHolder.hasActiveTransaction()) {
+				resourceHolder.getRequiredSession().commitTransaction();
 			}
 		}
 
@@ -219,8 +218,8 @@ public class MongoDatabaseUtils {
 		@Override
 		public void afterCompletion(int status) {
 
-			if (status == TransactionSynchronization.STATUS_ROLLED_BACK && isTransactionActive(this.resourceHolder)) {
-				resourceHolder.getSession().abortTransaction();
+			if (status == TransactionSynchronization.STATUS_ROLLED_BACK && this.resourceHolder.hasActiveTransaction()) {
+				resourceHolder.getRequiredSession().abortTransaction();
 			}
 
 			super.afterCompletion(status);
@@ -234,17 +233,8 @@ public class MongoDatabaseUtils {
 		protected void releaseResource(MongoResourceHolder resourceHolder, Object resourceKey) {
 
 			if (resourceHolder.hasActiveSession()) {
-				resourceHolder.getSession().close();
+				resourceHolder.getRequiredSession().close();
 			}
 		}
-
-		private boolean isTransactionActive(MongoResourceHolder resourceHolder) {
-
-			if (!resourceHolder.hasSession()) {
-				return false;
-			}
-
-			return resourceHolder.getSession().hasActiveTransaction();
-		}
 	}
 }

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

@@ -28,6 +28,7 @@ import com.mongodb.client.ClientSession;
  * <strong>Note:</strong> Intended for internal usage only.
  *
  * @author Christoph Strobl
+ * @author Mark Paluch
  * @since 2.1
  * @see MongoTransactionManager
  * @see org.springframework.data.mongodb.core.MongoTemplate
@@ -57,6 +58,22 @@ class MongoResourceHolder extends ResourceHolderSupport {
 		return session;
 	}
 
+	/**
+	 * @return the required associated {@link ClientSession}.
+	 * @throws IllegalStateException if no {@link ClientSession} is associated with this {@link MongoResourceHolder}.
+	 * @since 2.1.3
+	 */
+	ClientSession getRequiredSession() {
+
+		ClientSession session = getSession();
+
+		if (session == null) {
+			throw new IllegalStateException("No session available!");
+		}
+
+		return session;
+	}
+
 	/**
 	 * @return the associated {@link MongoDbFactory}.
 	 */
@@ -101,7 +118,21 @@ class MongoResourceHolder extends ResourceHolderSupport {
 			return false;
 		}
 
-		return hasServerSession() && !getSession().getServerSession().isClosed();
+		return hasServerSession() && !getRequiredSession().getServerSession().isClosed();
+	}
+
+	/**
+	 * @return {@literal true} if the session has an active transaction.
+	 * @since 2.1.3
+	 * @see #hasActiveSession()
+	 */
+	boolean hasActiveTransaction() {
+
+		if (!hasActiveSession()) {
+			return false;
+		}
+
+		return getRequiredSession().hasActiveTransaction();
 	}
 
 	/**
@@ -111,7 +142,7 @@ class MongoResourceHolder extends ResourceHolderSupport {
 	boolean hasServerSession() {
 
 		try {
-			return getSession().getServerSession() != null;
+			return getRequiredSession().getServerSession() != null;
 		} catch (IllegalStateException serverSessionClosed) {
 			// ignore
 		}

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

@@ -1124,11 +1124,11 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware,
 
 	protected long doCount(String collectionName, Document filter, CountOptions options) {
 
-		if (!MongoDatabaseUtils.isTransactionActive(getMongoDbFactory())) {
-			return execute(collectionName, collection -> collection.count(filter, options));
+		if (MongoDatabaseUtils.isTransactionActive(getMongoDbFactory())) {
+			return execute(collectionName, collection -> collection.countDocuments(filter, options));
 		}
 
-		return execute(collectionName, collection -> collection.countDocuments(filter, options));
+		return execute(collectionName, collection -> collection.count(filter, options));
 	}
 
 	/*