From 485ef8d7372594ed504add9edbbed7d7c91be5a4 Mon Sep 17 00:00:00 2001
From: Nidhi Nair <nidhi@appsmith.com>
Date: Mon, 14 Oct 2024 16:18:25 +0530
Subject: [PATCH] chore: Convert redis flush to lettuce commands

---
 .../server/migrations/DatabaseChangelog1.java |  2 +-
 .../Migration063CacheBustSpringBoot3_3.java   | 34 ++++++++++++++-----
 2 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/app/server/appsmith-server/src/main/java/com/appsmith/server/migrations/DatabaseChangelog1.java b/app/server/appsmith-server/src/main/java/com/appsmith/server/migrations/DatabaseChangelog1.java
index eab30d2af91..3e5b49a1a08 100644
--- a/app/server/appsmith-server/src/main/java/com/appsmith/server/migrations/DatabaseChangelog1.java
+++ b/app/server/appsmith-server/src/main/java/com/appsmith/server/migrations/DatabaseChangelog1.java
@@ -788,7 +788,7 @@ public void addPluginNameForGoogleSheets(MongoTemplate mongoTemplate) {
         mongoTemplate.save(googleSheetsPlugin);
     }
 
-    protected static void doClearRedisKeys(ReactiveRedisOperations<String, String> reactiveRedisOperations) {
+    public static void doClearRedisKeys(ReactiveRedisOperations<String, String> reactiveRedisOperations) {
         final String script = "for _,k in ipairs(redis.call('keys','spring:session:sessions:*'))"
                 + " do redis.call('del',k) " + "end";
         final Flux<Object> flushdb = reactiveRedisOperations.execute(RedisScript.of(script));
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();
     }
 }