Support inclusive/exclusive ranges for sorted sets #35

modules/api/src/main/scala/com/github/scoquelin/arugula/commands/RedisSortedSetAsyncCommands.scala

@@ -1,6 +1,7 @@
 package com.github.scoquelin.arugula.commands
 
 
+import scala.collection.immutable.NumericRange.Inclusive
 import scala.concurrent.Future
 import scala.concurrent.duration.FiniteDuration
 
@@ -672,11 +673,106 @@ object RedisSortedSetAsyncCommands {
   /**
    * A range of values
    *
-   * @param start The start value
-   * @param end The end value
+   * @param lower The lower bound
+   * @param upper The upper bound
    * @tparam T The value type
    */
-  final case class ZRange[T](start: T, end: T)
+  final case class ZRange[T](lower: ZRange.Boundary[T], upper: ZRange.Boundary[T])
+
+  object ZRange{
+    /**
+     * A boundary value, used for range queries (upper or lower bounds)
+     * @param value An optional value to use as the boundary. If None, it is unbounded.
+     * @param inclusive Whether the boundary is inclusive. default is false.
+     * @tparam T The value type
+     */
+    case class Boundary[T](value: Option[T] = None, inclusive: Boolean = false)
+
+    object Boundary{
+      /**
+       * Create a new inclusive boundary
+       * @param value The value
+       * @tparam T The value type
+       * @return The boundary
+       */
+      def including[T](value: T): Boundary[T] = Boundary(Some(value), inclusive = true)
+
+      /**
+       * Create a new exclusive boundary
+       * @param value The value
+       * @tparam T The value type
+       * @return The boundary
+       */
+      def excluding[T](value: T): Boundary[T] = Boundary(Some(value))
+
+      /**
+       * Create an unbounded boundary
+       * @tparam T The value type
+       * @return The boundary
+       */
+      def unbounded[T]: Boundary[T] = Boundary[T](None)
+    }
+
+    /**
+     * Create a new range
+     * @param lower The lower bound
+     * @param upper The upper bound
+     * @tparam T The value type
+     * @return The range
+     */
+    def apply[T](lower: T, upper: T): ZRange[T] = ZRange(Boundary.including(lower), Boundary.including(upper))
+
+    /**
+     * Create a new range
+     * @param lower The lower bound
+     * @param upper The upper bound
+     * @tparam T The value type
+     * @return The range
+     */
+    def including[T](lower: T, upper: T): ZRange[T] = ZRange(Boundary.including(lower), Boundary.including(upper))
+
+    /**
+     * Create a new range
+     * @param lower The lower bound
+     * @param upper The upper bound
+     * @tparam T The value type
+     * @return The range
+     */
+    def from[T](lower: T, upper: T, inclusive: Boolean = false): ZRange[T] = {
+      if(inclusive) ZRange(Boundary.including(lower), Boundary.including(upper))
+      else ZRange(Boundary.excluding(lower), Boundary.excluding(upper))
+    }
+
+
+    /**
+     * Create a new range from lower to unbounded
+     * @param lower The start boundary
+     * @tparam T The value type
+     * @return The range
+     */
+    def fromLower[T](lower: T, inclusive: Boolean = false): ZRange[T] = {
+      if(inclusive) ZRange(Boundary.including(lower), Boundary.unbounded[T])
+      else ZRange(Boundary.excluding(lower), Boundary.unbounded[T])
+    }
+
+    /**
+     * Create a new range to upper bound
+     * @param upper The upper boundary
+     * @tparam T The value type
+     * @return The range
+     */
+    def toUpper[T](upper: T, inclusive: Boolean = false): ZRange[T] = {
+      if(inclusive) ZRange(Boundary.unbounded[T], Boundary.including(upper))
+      else ZRange(Boundary.unbounded[T], Boundary.excluding(upper))
+    }
+
+    /**
+     * Create a new unbounded range
+     * @tparam T The value type
+     * @return The range
+     */
+    def unbounded[T]: ZRange[T] = new ZRange[T](Boundary.unbounded, Boundary.unbounded)
+  }
 
   /**
    * A range limit

modules/core/src/main/scala/com/github/scoquelin/arugula/commands/LettuceRedisSortedSetAsyncCommands.scala

@@ -106,7 +106,7 @@ private[arugula] trait LettuceRedisSortedSetAsyncCommands[K, V] extends RedisSor
     delegateRedisClusterCommandAndLift(_.zdiffWithScores(keys: _*)).map(_.asScala.toList.map(scoredValue => ScoreWithValue(scoredValue.getScore, scoredValue.getValue)))
 
   override def zLexCount(key: K, range: ZRange[V]): Future[Long] = {
-    delegateRedisClusterCommandAndLift(_.zlexcount(key, Range.create(range.start, range.end))).map(Long2long)
+    delegateRedisClusterCommandAndLift(_.zlexcount(key, LettuceRedisSortedSetAsyncCommands.toJavaRange(range))).map(Long2long)
   }
 
   override def zMPop(direction: SortOrder, keys: K*): Future[Option[ScoreWithKeyValue[K, V]]] = {
@@ -186,7 +186,7 @@ private[arugula] trait LettuceRedisSortedSetAsyncCommands[K, V] extends RedisSor
       case Some(rangeLimit) => io.lettuce.core.Limit.create(rangeLimit.offset, rangeLimit.count)
       case None => io.lettuce.core.Limit.unlimited()
     }
-    delegateRedisClusterCommandAndLift(_.zrangestorebylex(destination, key, Range.create(range.start, range.end), args)).map(Long2long)
+    delegateRedisClusterCommandAndLift(_.zrangestorebylex(destination, key, LettuceRedisSortedSetAsyncCommands.toJavaRange(range), args)).map(Long2long)
   }
 
   override def zRangeStoreByScore[T: Numeric](destination: K, key: K, range: ZRange[T], limit: Option[RangeLimit] = None): Future[Long] = {
@@ -222,7 +222,7 @@ private[arugula] trait LettuceRedisSortedSetAsyncCommands[K, V] extends RedisSor
       case Some(rangeLimit) => io.lettuce.core.Limit.create(rangeLimit.offset, rangeLimit.count)
       case None => io.lettuce.core.Limit.unlimited()
     }
-    delegateRedisClusterCommandAndLift(_.zrevrangebylex(key, Range.create(range.start, range.end), args)).map(_.asScala.toList)
+    delegateRedisClusterCommandAndLift(_.zrevrangebylex(key, LettuceRedisSortedSetAsyncCommands.toJavaRange(range), args)).map(_.asScala.toList)
   }
 
   override def zRevRangeByScore[T: Numeric](key: K, range: ZRange[T], limit: Option[RangeLimit] = None): Future[List[V]] =
@@ -296,7 +296,7 @@ private[arugula] trait LettuceRedisSortedSetAsyncCommands[K, V] extends RedisSor
       case Some(rangeLimit) => io.lettuce.core.Limit.create(rangeLimit.offset, rangeLimit.count)
       case None => io.lettuce.core.Limit.unlimited()
     }
-    delegateRedisClusterCommandAndLift(_.zrangebylex(key, Range.create(range.start, range.end), args)).map(_.asScala.toList)
+    delegateRedisClusterCommandAndLift(_.zrangebylex(key, LettuceRedisSortedSetAsyncCommands.toJavaRange(range), args)).map(_.asScala.toList)
   }
 
   override def zScan(key: K, cursor: String = InitialCursor, limit: Option[Long] = None, matchPattern: Option[String] = None): Future[ScanResults[List[ScoreWithValue[V]]]] = {
@@ -347,7 +347,7 @@ private[arugula] trait LettuceRedisSortedSetAsyncCommands[K, V] extends RedisSor
     delegateRedisClusterCommandAndLift(_.zrem(key, values: _*)).map(Long2long)
 
   override def zRemRangeByLex(key: K, range: ZRange[V]): Future[Long] =
-    delegateRedisClusterCommandAndLift(_.zremrangebylex(key, Range.create(range.start, range.end))).map(Long2long)
+    delegateRedisClusterCommandAndLift(_.zremrangebylex(key, LettuceRedisSortedSetAsyncCommands.toJavaRange(range))).map(Long2long)
 
   override def zRemRangeByRank(key: K, start: Long, stop: Long): Future[Long] =
     delegateRedisClusterCommandAndLift(_.zremrangebyrank(key, start, stop)).map(Long2long)
@@ -366,7 +366,7 @@ private[arugula] trait LettuceRedisSortedSetAsyncCommands[K, V] extends RedisSor
       case Some(rangeLimit) => io.lettuce.core.Limit.create(rangeLimit.offset, rangeLimit.count)
       case None => io.lettuce.core.Limit.unlimited()
     }
-    delegateRedisClusterCommandAndLift(_.zrevrangestorebylex(destination, key, Range.create(range.start, range.end), args)).map(Long2long)
+    delegateRedisClusterCommandAndLift(_.zrevrangestorebylex(destination, key, LettuceRedisSortedSetAsyncCommands.toJavaRange(range), args)).map(Long2long)
   }
 
   override def zRevRangeStoreByScore[T: Numeric](destination: K,
@@ -409,8 +409,12 @@ private[arugula] trait LettuceRedisSortedSetAsyncCommands[K, V] extends RedisSor
 
 }
 
-private[this] object LettuceRedisSortedSetAsyncCommands{
-  private[commands] def toJavaNumberRange[T: Numeric](range: ZRange[T]): Range[Number] = {
+private[arugula] object LettuceRedisSortedSetAsyncCommands{
+  private[arugula] def toJavaNumberRange[T: Numeric](range: ZRange[T]): io.lettuce.core.Range[java.lang.Number] = {
+    io.lettuce.core.Range.from(toJavaNumberBoundary(range.lower), toJavaNumberBoundary(range.upper))
+  }
+
+  private[commands] def toJavaNumberBoundary[T: Numeric](boundary: RedisSortedSetAsyncCommands.ZRange.Boundary[T]): io.lettuce.core.Range.Boundary[java.lang.Number] = {
     def toJavaNumber(t: T): java.lang.Number = t match {
       case b: Byte => b
       case s: Short => s
@@ -419,7 +423,23 @@ private[this] object LettuceRedisSortedSetAsyncCommands{
       case f: Float => f
       case _ => implicitly[Numeric[T]].toDouble(t)
     }
-    Range.create(toJavaNumber(range.start), toJavaNumber(range.end))
+    boundary.value match {
+      case Some(value) if boundary.inclusive => io.lettuce.core.Range.Boundary.including(toJavaNumber(value))
+      case Some(value) => io.lettuce.core.Range.Boundary.excluding(toJavaNumber(value))
+      case None => io.lettuce.core.Range.Boundary.unbounded[java.lang.Number]()
+    }
+  }
+
+  private [arugula] def toJavaRange[T](range: ZRange[T]): io.lettuce.core.Range[T] = {
+    io.lettuce.core.Range.from[T](toJavaBoundary(range.lower), toJavaBoundary(range.upper))
+  }
+
+  private [commands] def toJavaBoundary[T](boundary: RedisSortedSetAsyncCommands.ZRange.Boundary[T]): io.lettuce.core.Range.Boundary[T] = {
+    boundary.value match {
+      case Some(value) if boundary.inclusive => io.lettuce.core.Range.Boundary.including(value)
+      case Some(value) => io.lettuce.core.Range.Boundary.excluding(value)
+      case None => io.lettuce.core.Range.Boundary.unbounded[T]()
+    }
   }
 
 

modules/tests/it/src/test/scala/com/github/scoquelin/arugula/RedisCommandsIntegrationSpec.scala

@@ -883,7 +883,7 @@ class RedisCommandsIntegrationSpec extends BaseRedisCommandsIntegrationSpec with
           val destination = randomKey("sorted-set-destination") + suffix
           for {
             _ <- client.zAdd(key, ScoreWithValue(1, "a"), ScoreWithValue(2, "b"), ScoreWithValue(3, "c"), ScoreWithValue(4, "d"), ScoreWithValue(5, "e"))
-            lexRange <- client.zRangeByLex(key, ZRange("a", "c"))
+            lexRange <- client.zRangeByLex(key, ZRange(ZRange.Boundary.including("a"), ZRange.Boundary.including("c")))
             _ <- lexRange shouldBe List("a", "b", "c")
             lexRange <- client.zRangeByLex(key, ZRange("a", "c"), Some(RangeLimit(0, 2)))
             _ <- lexRange shouldBe List("a", "b")

modules/tests/test/src/test/scala/com/github/scoquelin/arugula/LettuceRedisSortedSetAsyncCommandsSpec.scala

@@ -3,8 +3,8 @@ package com.github.scoquelin.arugula
 
 import scala.concurrent.duration.DurationInt
 
-import com.github.scoquelin.arugula.commands.RedisSortedSetAsyncCommands
-import com.github.scoquelin.arugula.commands.RedisSortedSetAsyncCommands.{Aggregate, RangeLimit, ScoreWithValue, SortOrder, ZAddOptions, AggregationArgs, ZRange}
+import com.github.scoquelin.arugula.commands.{LettuceRedisSortedSetAsyncCommands, RedisSortedSetAsyncCommands}
+import com.github.scoquelin.arugula.commands.RedisSortedSetAsyncCommands.{Aggregate, AggregationArgs, RangeLimit, ScoreWithValue, SortOrder, ZAddOptions, ZRange}
 import io.lettuce.core.{KeyValue, RedisFuture, ScoredValue, ScoredValueScanCursor}
 import org.mockito.ArgumentCaptor
 import org.mockito.ArgumentMatchers.{any, eq => meq}
@@ -952,4 +952,31 @@ class LettuceRedisSortedSetAsyncCommandsSpec extends wordspec.FixtureAsyncWordSp
       }
     }
   }
+
+  "LettuceRedisSortedSetAsyncCommands.toJavaNumberRange" should {
+    "convert a ZRange with Double.NegativeInfinity and Double.PositiveInfinity to a Range with Double.NEGATIVE_INFINITY and Double.POSITIVE_INFINITY" in { _ =>
+      val range = ZRange(Double.NegativeInfinity, Double.PositiveInfinity)
+      val result = LettuceRedisSortedSetAsyncCommands.toJavaNumberRange(range)
+      result mustBe io.lettuce.core.Range.create(Double.NegativeInfinity, Double.PositiveInfinity)
+    }
+
+    "convert a ZRange with a lower bound to a Range with the lower bound" in { _ =>
+      val range = ZRange.fromLower(1.0)
+      val result = LettuceRedisSortedSetAsyncCommands.toJavaNumberRange(range)
+      result mustBe io.lettuce.core.Range.from(io.lettuce.core.Range.Boundary.excluding(1.0), io.lettuce.core.Range.Boundary.unbounded())
+    }
+
+    "convert a ZRange with an upper bound to a Range with the upper bound" in { _ =>
+      val range = ZRange.toUpper(1.0)
+      val result = LettuceRedisSortedSetAsyncCommands.toJavaNumberRange(range)
+      result mustBe io.lettuce.core.Range.from(io.lettuce.core.Range.Boundary.unbounded(), io.lettuce.core.Range.Boundary.excluding(1.0))
+    }
+
+    "convert a ZRange excluding both bounds to a Range with the lower and upper bounds" in { _ =>
+      val range = ZRange.from(1.0, 2.0)
+      val result = LettuceRedisSortedSetAsyncCommands.toJavaNumberRange(range)
+      result mustBe io.lettuce.core.Range.from(io.lettuce.core.Range.Boundary.excluding(1.0), io.lettuce.core.Range.Boundary.excluding(2.0))
+    }
+  }
+
 }