Merge pull request #107 from NoRedInk/fxn-more-sorted-set

Add zrangeByScoreWithScores
NoRedInk · May 14, 2024 · b75698e · b75698e
2 parents 5e0c40b + ed5b9d2
commit b75698e
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 0 deletions.
diff --git a/nri-redis/src/Redis/Handler.hs b/nri-redis/src/Redis/Handler.hs
@@ -313,6 +313,13 @@ doRawQuery query =
         (Prelude.fromIntegral stop)
         |> PreparedQuery
         |> map Ok
+    Internal.ZrangeByScoreWithScores key start stop ->
+      Database.Redis.zrangebyscoreWithscores
+        (toB key)
+        start
+        stop
+        |> PreparedQuery
+        |> map Ok
     Internal.Zrank key member ->
       Database.Redis.zrank (toB key) member
         |> PreparedQuery

diff --git a/nri-redis/src/Redis/Internal.hs b/nri-redis/src/Redis/Internal.hs
@@ -109,6 +109,7 @@ cmds query'' =
     Smembers key -> [unwords ["SMEMBERS", key]]
     Zadd key vals -> [unwords ("ZADD" : key : List.concatMap (\(_, val) -> ["*****", Text.fromFloat val]) (Dict.toList vals))]
     Zrange key start stop -> [unwords ["ZRANGE", key, Text.fromInt start, Text.fromInt stop]]
+    ZrangeByScoreWithScores key start stop -> [unwords ["ZRANGE", key, "BYSCORE", Text.fromFloat start, Text.fromFloat stop, "WITHSCORES"]]
     Zrank key _ -> [unwords ["ZRANK", key, "*****"]]
     Zrevrank key _ -> [unwords ["ZREVRANK", key, "*****"]]
     Pure _ -> []
@@ -164,6 +165,7 @@ data Query a where
   Smembers :: Text -> Query (List ByteString)
   Zadd :: Text -> Dict.Dict ByteString Float -> Query Int
   Zrange :: Text -> Int -> Int -> Query [ByteString]
+  ZrangeByScoreWithScores :: Text -> Float -> Float -> Query [(ByteString, Float)]
   Zrank :: Text -> ByteString -> Query (Maybe Int)
   Zrevrank :: Text -> ByteString -> Query (Maybe Int)
   -- The constructors below are not Redis-related, but support using functions
@@ -302,6 +304,7 @@ mapKeys fn query' =
     Smembers key -> Task.map Smembers (fn key)
     Zadd key vals -> Task.map (\newKey -> Zadd newKey vals) (fn key)
     Zrange key start stop -> Task.map (\newKey -> Zrange newKey start stop) (fn key)
+    ZrangeByScoreWithScores key start stop -> Task.map (\newKey -> ZrangeByScoreWithScores newKey start stop) (fn key)
     Zrank key member -> Task.map (\newKey -> Zrank newKey member) (fn key)
     Zrevrank key member -> Task.map (\newKey -> Zrevrank newKey member) (fn key)
     Pure x -> Task.succeed (Pure x)
@@ -343,6 +346,7 @@ mapReturnedKeys fn query' =
     Smembers key -> Smembers key
     Zadd key vals -> Zadd key vals
     Zrange key start stop -> Zrange key start stop
+    ZrangeByScoreWithScores key start stop -> ZrangeByScoreWithScores key start stop
     Zrank key member -> Zrank key member
     Zrevrank key member -> Zrevrank key member
     Pure x -> Pure x
@@ -399,6 +403,7 @@ keysTouchedByQuery query' =
     Smembers key -> Set.singleton key
     Zadd key _ -> Set.singleton key
     Zrange key _ _ -> Set.singleton key
+    ZrangeByScoreWithScores key _ _ -> Set.singleton key
     Zrank key _ -> Set.singleton key
     Zrevrank key _ -> Set.singleton key
     WithResult _ q -> keysTouchedByQuery q

diff --git a/nri-redis/src/Redis/SortedSet.hs b/nri-redis/src/Redis/SortedSet.hs
@@ -26,6 +26,7 @@ module Redis.SortedSet
     ping,
     zadd,
     zrange,
+    zrangeByScoreWithScores,
     zrank,
     zrevrank,
 
@@ -97,6 +98,9 @@ data Api key a = Api
     --
     -- https://redis.io/commands/zrange
     zrange :: key -> Int -> Int -> Internal.Query (List a),
+    -- | Like `zrange`, but with the bounds being scores rather than offsets,
+    -- and with the result including the scores for each returned result.
+    zrangeByScoreWithScores :: key -> Float -> Float -> Internal.Query [(a, Float)],
     -- | Returns the rank of member in the sorted set stored at key, with the
     -- scores ordered from low to high. The rank (or index) is 0-based, which
     -- means that the member with the lowest score has rank 0.
@@ -149,6 +153,14 @@ makeApi Codec.Codec {Codec.codecEncoder, Codec.codecDecoder} toKey =
       zrange = \key start stop ->
         Internal.Zrange (toKey key) start stop
           |> Internal.WithResult (Prelude.traverse codecDecoder),
+      zrangeByScoreWithScores = \key start stop ->
+        Internal.ZrangeByScoreWithScores (toKey key) start stop
+          |> Internal.WithResult
+            ( Prelude.traverse
+                ( \(v, score) ->
+                    codecDecoder v |> Result.map (\val -> (val, score))
+                )
+            ),
       zrank = \key member -> Internal.Zrank (toKey key) (codecEncoder member),
       zrevrank = \key member -> Internal.Zrevrank (toKey key) (codecEncoder member)
     }