Proxy MongoCluster for passing in ClientSession if needed

Author: christophstrobl

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

@@ -54,6 +54,7 @@ public class SessionAwareMethodInterceptor<D, C> implements MethodInterceptor {
 	private final ClientSessionOperator databaseDecorator;
 	private final Object target;
 	private final Class<?> targetType;
+	private final Class<?> clientType;
 	private final Class<?> collectionType;
 	private final Class<?> databaseType;
 	private final Class<? extends ClientSession> sessionType;
@@ -71,7 +72,7 @@ public class SessionAwareMethodInterceptor<D, C> implements MethodInterceptor {
 	 *          {@code MongoCollection}.
 	 * @param <T> target object type.
 	 */
-	public <T> SessionAwareMethodInterceptor(ClientSession session, T target, Class<? extends ClientSession> sessionType,
+	public <T> SessionAwareMethodInterceptor(ClientSession session, T target, Class<?> clientType, Class<? extends ClientSession> sessionType,
 			Class<D> databaseType, ClientSessionOperator<D> databaseDecorator, Class<C> collectionType,
 			ClientSessionOperator<C> collectionDecorator) {
 
@@ -85,15 +86,24 @@ public <T> SessionAwareMethodInterceptor(ClientSession session, T target, Class<
 
 		this.session = session;
 		this.target = target;
+		this.clientType = ClassUtils.getUserClass(clientType);
 		this.databaseType = ClassUtils.getUserClass(databaseType);
 		this.collectionType = ClassUtils.getUserClass(collectionType);
 		this.collectionDecorator = collectionDecorator;
 		this.databaseDecorator = databaseDecorator;
 
-		this.targetType = ClassUtils.isAssignable(databaseType, target.getClass()) ? databaseType : collectionType;
+		this.targetType = targetType(target.getClass());
 		this.sessionType = sessionType;
 	}
 
+	Class<?> targetType(@Nullable Class<?> targetType) {
+
+		if(ClassUtils.isAssignable(clientType, targetType)) {
+			return clientType;
+		}
+		return ClassUtils.isAssignable(databaseType, target.getClass()) ? databaseType : collectionType;
+	}
+
 	@Override
 	public @Nullable Object invoke(MethodInvocation methodInvocation) throws Throwable {
 

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

@@ -15,9 +15,7 @@
  */
 package org.springframework.data.mongodb.core;
 
-import com.mongodb.client.MongoCluster;
 import org.jspecify.annotations.Nullable;
-
 import org.springframework.aop.framework.ProxyFactory;
 import org.springframework.dao.DataAccessException;
 import org.springframework.dao.support.PersistenceExceptionTranslator;
@@ -29,6 +27,7 @@
 import com.mongodb.ClientSessionOptions;
 import com.mongodb.WriteConcern;
 import com.mongodb.client.ClientSession;
+import com.mongodb.client.MongoCluster;
 import com.mongodb.client.MongoCollection;
 import com.mongodb.client.MongoDatabase;
 
@@ -126,7 +125,6 @@ public MongoDatabase getMongoDatabase(String dbName) throws DataAccessException
 	 */
 	protected abstract MongoDatabase doGetMongoDatabase(String dbName);
 
-
 	public void destroy() throws Exception {
 		if (mongoInstanceCreated) {
 			closeClient();
@@ -170,9 +168,7 @@ record ClientSessionBoundMongoDbFactory(ClientSession session,
 
 		@Override
 		public MongoCluster getCluster() {
-
-			// TODO: we need to proxy the cluster and methods that accept a client session
-			return delegate.getCluster();
+			return proxyMongoCluster(delegate.getCluster());
 		}
 
 		@Override
@@ -209,6 +205,10 @@ private MongoDatabase proxyMongoDatabase(MongoDatabase database) {
 			return createProxyInstance(session, database, MongoDatabase.class);
 		}
 
+		private MongoCluster proxyMongoCluster(MongoCluster cluster) {
+			return createProxyInstance(session, cluster, MongoCluster.class);
+		}
+
 		private MongoDatabase proxyDatabase(com.mongodb.session.ClientSession session, MongoDatabase database) {
 			return createProxyInstance(session, database, MongoDatabase.class);
 		}
@@ -225,7 +225,7 @@ private <T> T createProxyInstance(com.mongodb.session.ClientSession session, T t
 			factory.setInterfaces(targetType);
 			factory.setOpaque(true);
 
-			factory.addAdvice(new SessionAwareMethodInterceptor<>(session, target, ClientSession.class, MongoDatabase.class,
+			factory.addAdvice(new SessionAwareMethodInterceptor<>(session, target, MongoCluster.class, ClientSession.class, MongoDatabase.class,
 					this::proxyDatabase, MongoCollection.class, this::proxyCollection));
 
 			return targetType.cast(factory.getProxy(target.getClass().getClassLoader()));

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

@@ -15,6 +15,7 @@
  */
 package org.springframework.data.mongodb.core;
 
+import com.mongodb.reactivestreams.client.MongoCluster;
 import reactor.core.publisher.Mono;
 
 import org.bson.codecs.configuration.CodecRegistry;
@@ -226,7 +227,7 @@ private <T> T createProxyInstance(com.mongodb.session.ClientSession session, T t
 			factory.setInterfaces(targetType);
 			factory.setOpaque(true);
 
-			factory.addAdvice(new SessionAwareMethodInterceptor<>(session, target, ClientSession.class, MongoDatabase.class,
+			factory.addAdvice(new SessionAwareMethodInterceptor<>(session, target, MongoCluster.class, ClientSession.class, MongoDatabase.class,
 					this::proxyDatabase, MongoCollection.class, this::proxyCollection));
 
 			return targetType.cast(factory.getProxy(target.getClass().getClassLoader()));

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

@@ -22,6 +22,7 @@
 import java.lang.reflect.Method;
 import java.lang.reflect.Proxy;
 
+import com.mongodb.client.MongoCluster;
 import org.bson.Document;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -175,7 +176,7 @@ private <T> T createProxyInstance(com.mongodb.session.ClientSession session, T t
 		factory.setInterfaces(targetType);
 		factory.setOpaque(true);
 
-		factory.addAdvice(new SessionAwareMethodInterceptor<>(session, target, ClientSession.class, MongoDatabase.class,
+		factory.addAdvice(new SessionAwareMethodInterceptor<>(session, target, MongoCluster.class, ClientSession.class, MongoDatabase.class,
 				this::proxyDatabase, MongoCollection.class, this::proxyCollection));
 
 		return targetType.cast(factory.getProxy());

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

@@ -20,6 +20,7 @@
 
 import java.lang.reflect.InvocationHandler;
 import java.lang.reflect.Proxy;
+import java.util.List;
 
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
@@ -103,6 +104,18 @@ void cascadedWithSessionUsesRootFactory() {
 		assertThat(singletonTarget).isSameAs(database);
 	}
 
+	@Test // GH-5087
+	void passesOnClientSessionWhenInvokingMethodsOnMongoCluster() {
+
+		MongoDatabaseFactory factory = MongoDatabaseFactory.create(mongo, "foo");
+		MongoDatabaseFactory wrapped = factory.withSession(clientSession);
+
+		assertThat(wrapped.getCluster()).isInstanceOf(Proxy.class);
+		wrapped.getCluster().bulkWrite(List.of());
+
+		verify(mongo).bulkWrite(eq(clientSession), any());
+	}
+
 	private void rejectsDatabaseName(String databaseName) {
 		assertThatThrownBy(() -> new SimpleMongoClientDatabaseFactory(mongo, databaseName))
 				.isInstanceOf(IllegalArgumentException.class);