From d31ddffe54f33cd1b5ac9fb39d8a441fca9f28d5 Mon Sep 17 00:00:00 2001
From: Vijay Gupta <vijayrauniyar1818@gmail.com>
Date: Thu, 19 Dec 2024 16:32:04 +0530
Subject: [PATCH] backend:kv query changes with options

---
 flake.lock                                    |  6 +-
 .../src/Kernel/Beam/Functions.hs              | 72 ++++++++++++++++++-
 2 files changed, 74 insertions(+), 4 deletions(-)

diff --git a/flake.lock b/flake.lock
index 4d5dd803..1124c85c 100644
--- a/flake.lock
+++ b/flake.lock
@@ -370,11 +370,11 @@
         "tinylog": "tinylog"
       },
       "locked": {
-        "lastModified": 1734428345,
-        "narHash": "sha256-Nbss1aAIfrgNMinY0AOfaOa8EUXWhAxHrw7Fhyb55d0=",
+        "lastModified": 1734605214,
+        "narHash": "sha256-DPW2S5NRUFA1y8x4j7+GIoadQqt2bebvnRSSbqYSKzw=",
         "owner": "nammayatri",
         "repo": "euler-hs",
-        "rev": "bccffdbc6c27e9ed648a1abb95f288763cff6bcc",
+        "rev": "6e74b58a3e25304e3f476de038f0e72c1d08e5c1",
         "type": "github"
       },
       "original": {
diff --git a/lib/mobility-core/src/Kernel/Beam/Functions.hs b/lib/mobility-core/src/Kernel/Beam/Functions.hs
index 7ef6ef3b..03e79c0d 100644
--- a/lib/mobility-core/src/Kernel/Beam/Functions.hs
+++ b/lib/mobility-core/src/Kernel/Beam/Functions.hs
@@ -30,6 +30,11 @@ module Kernel.Beam.Functions
     findOneWithKVRedis,
     logQueryData,
     findAllFromKvRedis,
+    createWithKVWithOptions,
+    createWithKVSchedulerWithOptions,
+    updateWithKVWithOptions,
+    updateWithKVSchedulerWithOptions,
+    updateOneWithKVWithOptions,
   )
 where
 
@@ -102,7 +107,8 @@ meshConfig =
       kvHardKilled = True,
       cerealEnabled = False,
       tableShardModRange = (0, 128),
-      redisKeyPrefix = ""
+      redisKeyPrefix = "",
+      forceDrainToDB = False
     }
 
 runInReplica :: (L.MonadFlow m, Log m) => m a -> m a
@@ -415,6 +421,18 @@ updateWithKV ::
 updateWithKV setClause whereClause = withUpdatedMeshConfig (Proxy @table) $ \updatedMeshConfig -> do
   updateInternal updatedMeshConfig setClause whereClause
 
+updateWithKVWithOptions ::
+  forall table m r.
+  (BeamTableFlow table m, EsqDBFlow m r) =>
+  [Set Postgres table] ->
+  Where Postgres table ->
+  Maybe Integer ->
+  Bool ->
+  m ()
+updateWithKVWithOptions setClause whereClause ttl forceDrain = withUpdatedMeshConfig (Proxy @table) $ \updatedMeshConfig -> do
+  let updatedMeshConfig' = updatedMeshConfig {redisTtl = fromMaybe (redisTtl updatedMeshConfig) ttl, forceDrainToDB = forceDrain}
+  updateInternal updatedMeshConfig' setClause whereClause
+
 updateWithKVScheduler ::
   forall table m r.
   (BeamTableFlow table m, EsqDBFlow m r) =>
@@ -424,6 +442,18 @@ updateWithKVScheduler ::
 updateWithKVScheduler setClause whereClause = withUpdatedMeshConfig (Proxy @table) $ \updatedMeshConfig -> do
   updateInternal updatedMeshConfig setClause whereClause
 
+updateWithKVSchedulerWithOptions ::
+  forall table m r.
+  (BeamTableFlow table m, EsqDBFlow m r) =>
+  [Set Postgres table] ->
+  Where Postgres table ->
+  Maybe Integer ->
+  Bool ->
+  m ()
+updateWithKVSchedulerWithOptions setClause whereClause ttl forceDrain = withUpdatedMeshConfig (Proxy @table) $ \updatedMeshConfig -> do
+  let updatedMeshConfig' = updatedMeshConfig {redisTtl = fromMaybe (redisTtl updatedMeshConfig) ttl, forceDrainToDB = forceDrain}
+  updateInternal updatedMeshConfig' setClause whereClause
+
 -- updateOne --
 
 updateOneWithKV ::
@@ -435,6 +465,18 @@ updateOneWithKV ::
 updateOneWithKV setClause whereClause = withUpdatedMeshConfig (Proxy @table) $ \updatedMeshConfig -> do
   updateOneInternal updatedMeshConfig setClause whereClause
 
+updateOneWithKVWithOptions ::
+  forall table m r.
+  (BeamTableFlow table m, EsqDBFlow m r) =>
+  [Set Postgres table] ->
+  Where Postgres table ->
+  Maybe Integer ->
+  Bool ->
+  m ()
+updateOneWithKVWithOptions setClause whereClause ttl forceDrain = withUpdatedMeshConfig (Proxy @table) $ \updatedMeshConfig -> do
+  let updatedMeshConfig' = updatedMeshConfig {redisTtl = fromMaybe (redisTtl updatedMeshConfig) ttl, forceDrainToDB = forceDrain}
+  updateOneInternal updatedMeshConfig' setClause whereClause
+
 -- create --
 
 createWithKV ::
@@ -448,6 +490,20 @@ createWithKV ::
 createWithKV a = withUpdatedMeshConfig (Proxy @table) $ \updatedMeshConfig -> do
   createInternal updatedMeshConfig toTType' a
 
+createWithKVWithOptions ::
+  forall table m r a.
+  ( BeamTableFlow table m,
+    EsqDBFlow m r,
+    ToTType' (table Identity) a
+  ) =>
+  a ->
+  Maybe Integer ->
+  Bool ->
+  m ()
+createWithKVWithOptions a ttl forceDrain = withUpdatedMeshConfig (Proxy @table) $ \updatedMeshConfig -> do
+  let updatedMeshConfig' = updatedMeshConfig {redisTtl = fromMaybe (redisTtl updatedMeshConfig) ttl, forceDrainToDB = forceDrain}
+  createInternal updatedMeshConfig' toTType' a
+
 createWithKVScheduler ::
   forall table m r a.
   ( BeamTableFlow table m,
@@ -459,6 +515,20 @@ createWithKVScheduler ::
 createWithKVScheduler a = withUpdatedMeshConfig (Proxy @table) $ \updatedMeshConfig -> do
   createInternal updatedMeshConfig toTType'' a
 
+createWithKVSchedulerWithOptions ::
+  forall table m r a.
+  ( BeamTableFlow table m,
+    EsqDBFlow m r,
+    ToTType'' (table Identity) a
+  ) =>
+  a ->
+  Maybe Integer ->
+  Bool ->
+  m ()
+createWithKVSchedulerWithOptions a ttl forceDrain = withUpdatedMeshConfig (Proxy @table) $ \updatedMeshConfig -> do
+  let updatedMeshConfig' = updatedMeshConfig {redisTtl = fromMaybe (redisTtl updatedMeshConfig) ttl, forceDrainToDB = forceDrain}
+  createInternal updatedMeshConfig' toTType'' a
+
 -- delete --
 
 deleteWithKV ::