From aaf7818427ac4497c1d747ed45e92ba20b09db3e Mon Sep 17 00:00:00 2001
From: Petromir Dzhunev <petromir@gmail.com>
Date: Fri, 10 Sep 2021 17:54:20 +0300
Subject: [PATCH] Extract connection creation into
 ReactiveRedisTemplate.getConnection() method.

A new method, called getConnection, is introduced to allow overriding of connection creation.

Closes #2145
Original pull request: #2162.
---
 .../redis/core/ReactiveRedisTemplate.java     | 24 ++++++++++++-------
 1 file changed, 16 insertions(+), 8 deletions(-)

diff --git a/src/main/java/org/springframework/data/redis/core/ReactiveRedisTemplate.java b/src/main/java/org/springframework/data/redis/core/ReactiveRedisTemplate.java
index 82627e864..e4788e914 100644
--- a/src/main/java/org/springframework/data/redis/core/ReactiveRedisTemplate.java
+++ b/src/main/java/org/springframework/data/redis/core/ReactiveRedisTemplate.java
@@ -232,14 +232,7 @@ public class ReactiveRedisTemplate<K, V> implements ReactiveRedisOperations<K, V
 
 		Assert.notNull(action, "Callback object must not be null");
 
-		return Flux.usingWhen(Mono.fromSupplier(() -> {
-
-			ReactiveRedisConnectionFactory factory = getConnectionFactory();
-			ReactiveRedisConnection conn = factory.getReactiveConnection();
-			ReactiveRedisConnection connToUse = preProcessConnection(conn, false);
-
-			return (exposeConnection ? connToUse : createRedisConnectionProxy(connToUse));
-		}), conn -> {
+		return Flux.usingWhen(getConnection(exposeConnection), conn -> {
 			Publisher<T> result = action.doInRedis(conn);
 
 			return postProcessResult(result, conn, false);
@@ -247,6 +240,21 @@ public class ReactiveRedisTemplate<K, V> implements ReactiveRedisOperations<K, V
 		}, ReactiveRedisConnection::closeLater);
 	}
 
+	/**
+	 * Creates a Mono which generates a new connection. The successors of {@link ReactiveRedisTemplate} might override
+	 * the default behaviour.
+	 *
+	 * @param exposeConnection whether to enforce exposure of the native Redis Connection to callback code
+	 * return a {@link Mono} wrapping the {@link ReactiveRedisConnection}.
+	 */
+	protected Mono<ReactiveRedisConnection> getConnection(boolean exposeConnection) {
+		ReactiveRedisConnectionFactory factory = getConnectionFactory();
+		ReactiveRedisConnection conn = factory.getReactiveConnection();
+		ReactiveRedisConnection connToUse = preProcessConnection(conn, false);
+
+		return Mono.just(exposeConnection ? connToUse : createRedisConnectionProxy(connToUse));
+	}
+
 	/*
 	 * (non-Javadoc)
 	 * @see org.springframework.data.redis.core.ReactiveRedisOperations#convertAndSend(java.lang.String, java.lang.Object)