From b086761e8b64fcc21242d71726f7dabdbb224902 Mon Sep 17 00:00:00 2001
From: Nidhi <nidhi@appsmith.com>
Date: Mon, 14 Oct 2024 17:11:01 +0530
Subject: [PATCH] chore: Use Lettuce client to run clear keys across Redis
 nodes (#36862)

---
 .../Migration063CacheBustSpringBoot3_3.java   | 34 ++++++++++++++-----
 1 file changed, 26 insertions(+), 8 deletions(-)

diff --git a/app/server/appsmith-server/src/main/java/com/appsmith/server/migrations/db/ce/Migration063CacheBustSpringBoot3_3.java b/app/server/appsmith-server/src/main/java/com/appsmith/server/migrations/db/ce/Migration063CacheBustSpringBoot3_3.java
index 7ca5077baba..217feb8abd1 100644
--- a/app/server/appsmith-server/src/main/java/com/appsmith/server/migrations/db/ce/Migration063CacheBustSpringBoot3_3.java
+++ b/app/server/appsmith-server/src/main/java/com/appsmith/server/migrations/db/ce/Migration063CacheBustSpringBoot3_3.java
@@ -5,7 +5,12 @@
 import io.mongock.api.annotations.RollbackExecution;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.beans.factory.annotation.Qualifier;
-import org.springframework.data.redis.core.ReactiveRedisTemplate;
+import org.springframework.data.redis.core.ReactiveRedisOperations;
+import org.springframework.data.redis.core.ScanOptions;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+
+import java.nio.ByteBuffer;
 
 @Slf4j
 @ChangeUnit(order = "063", id = "reset_session_oauth2_spring_3_3")
@@ -16,12 +21,25 @@ public void rollbackExecution() {}
 
     @Execution
     public void execute(
-            @Qualifier("reactiveRedisTemplate") final ReactiveRedisTemplate<String, Object> reactiveRedisTemplate) {
-        reactiveRedisTemplate
-                .getConnectionFactory()
-                .getReactiveConnection()
-                .serverCommands()
-                .flushDb()
-                .block();
+            @Qualifier("reactiveRedisOperations") ReactiveRedisOperations<String, Object> reactiveRedisOperations) {
+        scanForKeysAcrossCluster(reactiveRedisOperations, "*").block();
+    }
+
+    private Mono<Void> scanForKeysAcrossCluster(
+            ReactiveRedisOperations<String, Object> reactiveRedisOperations, String pattern) {
+        return reactiveRedisOperations
+                .execute(connection -> {
+                    Flux<ByteBuffer> scanFlux = connection
+                            .keyCommands()
+                            .scan(ScanOptions.scanOptions()
+                                    .match(pattern)
+                                    .count(1000)
+                                    .build());
+                    return scanFlux.flatMap(scannedKey -> {
+                                return connection.keyCommands().del(scannedKey);
+                            })
+                            .then();
+                })
+                .then();
     }
 }