From 9bb2c6811d607f7296165ad4beb99718158702cf Mon Sep 17 00:00:00 2001 From: Sloane Sturzenegger Date: Tue, 9 Mar 2021 08:59:00 -0800 Subject: [PATCH] Support STORE in Georadius commands --- redis/src/main/scala/zio/redis/api/Geo.scala | 155 ++++++++++++++++-- .../main/scala/zio/redis/options/Geo.scala | 18 ++ redis/src/test/scala/zio/redis/GeoSpec.scala | 16 ++ 3 files changed, 172 insertions(+), 17 deletions(-) diff --git a/redis/src/main/scala/zio/redis/api/Geo.scala b/redis/src/main/scala/zio/redis/api/Geo.scala index 58e92a937..2828f06c2 100644 --- a/redis/src/main/scala/zio/redis/api/Geo.scala +++ b/redis/src/main/scala/zio/redis/api/Geo.scala @@ -79,9 +79,6 @@ trait Geo { * @param withHash flag to include raw geohash sorted set score of each member in the result * @param count limit the results to the first N matching items * @param order sort returned items in the given `Order` - * @param store sorted set where the results should be stored - * @param storeDist sorted set where the results populated with their distance from the center as a floating point - * number should be stored * @return chunk of members within the specified are */ final def geoRadius( @@ -93,11 +90,43 @@ trait Geo { withDist: Option[WithDist] = None, withHash: Option[WithHash] = None, count: Option[Count] = None, - order: Option[Order] = None, - store: Option[Store] = None, - storeDist: Option[StoreDist] = None + order: Option[Order] = None ): ZIO[RedisExecutor, RedisError, Chunk[GeoView]] = - GeoRadius.run((key, center, radius, radiusUnit, withCoord, withDist, withHash, count, order, store, storeDist)) + GeoRadius.run((key, center, radius, radiusUnit, withCoord, withDist, withHash, count, order)) + + /** + * Similar to geoRadius, but store the results to the argument passed to store and return the number of + * elements stored. Added as a separate Scala function because it returns a Long instead of the list of results. + * Find the geospatial members of a sorted set which are within the area specified with a *center location* and the + * *maximum distance from the center*. + * + * @param key sorted set of geospatial members + * @param center position + * @param radius distance from the center + * @param store sorted set where the results and/or distances should be stored + * @param radiusUnit Unit of distance ("m", "km", "ft", "mi") + * @param withCoord flag to include the position of each member in the result + * @param withDist flag to include the distance of each member from the center in the result + * @param withHash flag to include raw geohash sorted set score of each member in the result + * @param count limit the results to the first N matching items + * @param order sort returned items in the given `Order` + * @return chunk of members within the specified are + */ + final def geoRadiusStore( + key: String, + center: LongLat, + radius: Double, + radiusUnit: RadiusUnit, + store: StoreOptions, + withCoord: Option[WithCoord] = None, + withDist: Option[WithDist] = None, + withHash: Option[WithHash] = None, + count: Option[Count] = None, + order: Option[Order] = None + ): ZIO[RedisExecutor, RedisError, Long] = + GeoRadiusStore.run( + (key, center, radius, radiusUnit, withCoord, withDist, withHash, count, order, store.store, store.storeDist) + ) /** * Return geospatial members of a sorted set which are within the area specified with an *existing member* in the set @@ -111,8 +140,6 @@ trait Geo { * @param withHash flag to include raw geohash sorted set score of each member in the result * @param count limit the results to the first N matching items * @param order sort returned items in the given `Order` - * @param store sorted set where the results should be stored - * @param storeDist sorted set where the results populated with their distance from the center as a floating point * number should be stored * @return chunk of members within the specified area, or an error if the member is not in the set */ @@ -125,12 +152,44 @@ trait Geo { withDist: Option[WithDist] = None, withHash: Option[WithHash] = None, count: Option[Count] = None, - order: Option[Order] = None, - store: Option[Store] = None, - storeDist: Option[StoreDist] = None + order: Option[Order] = None ): ZIO[RedisExecutor, RedisError, Chunk[GeoView]] = GeoRadiusByMember.run( - (key, member, radius, radiusUnit, withCoord, withDist, withHash, count, order, store, storeDist) + (key, member, radius, radiusUnit, withCoord, withDist, withHash, count, order) + ) + + /** + * Similar to geoRadiusByMember, but store the results to the argument passed to store and return the number of + * elements stored. Added as a separate Scala function because it returns a Long instead of the list of results. + * Find geospatial members of a sorted set which are within the area specified with an *existing member* in the set + * and the *maximum distance from the location of that member*. + * + * @param key sorted set of geospatial members + * @param member member in the set + * @param radius distance from the member + * @param radiusUnit Unit of distance ("m", "km", "ft", "mi") + * @param store sorted set where the results and/or distances should be stored + * @param withCoord flag to include the position of each member in the result + * @param withDist flag to include the distance of each member from the center in the result + * @param withHash flag to include raw geohash sorted set score of each member in the result + * @param count limit the results to the first N matching items + * @param order sort returned items in the given `Order` + * @return chunk of members within the specified area, or an error if the member is not in the set + */ + final def geoRadiusByMemberStore( + key: String, + member: String, + radius: Double, + radiusUnit: RadiusUnit, + store: StoreOptions, + withCoord: Option[WithCoord] = None, + withDist: Option[WithDist] = None, + withHash: Option[WithHash] = None, + count: Option[Count] = None, + order: Option[Order] = None + ): ZIO[RedisExecutor, RedisError, Long] = + GeoRadiusByMemberStore.run( + (key, member, radius, radiusUnit, withCoord, withDist, withHash, count, order, store.store, store.storeDist) ) } @@ -152,6 +211,37 @@ private[redis] object Geo { RedisCommand("GEOPOS", Tuple2(StringInput, NonEmptyList(StringInput)), GeoOutput) final val GeoRadius: RedisCommand[ + ( + String, + LongLat, + Double, + RadiusUnit, + Option[WithCoord], + Option[WithDist], + Option[WithHash], + Option[Count], + Option[Order] + ), + Chunk[GeoView] + ] = + RedisCommand( + "GEORADIUS", + Tuple9( + StringInput, + LongLatInput, + DoubleInput, + RadiusUnitInput, + OptionalInput(WithCoordInput), + OptionalInput(WithDistInput), + OptionalInput(WithHashInput), + OptionalInput(CountInput), + OptionalInput(OrderInput) + ), + GeoRadiusOutput + ) + + // We expect at least one of Option[Store] and Option[StoreDist] to be passed here. + final val GeoRadiusStore: RedisCommand[ ( String, LongLat, @@ -165,7 +255,7 @@ private[redis] object Geo { Option[Store], Option[StoreDist] ), - Chunk[GeoView] + Long ] = RedisCommand( "GEORADIUS", @@ -182,10 +272,41 @@ private[redis] object Geo { OptionalInput(StoreInput), OptionalInput(StoreDistInput) ), - GeoRadiusOutput + LongOutput ) final val GeoRadiusByMember: RedisCommand[ + ( + String, + String, + Double, + RadiusUnit, + Option[WithCoord], + Option[WithDist], + Option[WithHash], + Option[Count], + Option[Order] + ), + Chunk[GeoView] + ] = + RedisCommand( + "GEORADIUSBYMEMBER", + Tuple9( + StringInput, + StringInput, + DoubleInput, + RadiusUnitInput, + OptionalInput(WithCoordInput), + OptionalInput(WithDistInput), + OptionalInput(WithHashInput), + OptionalInput(CountInput), + OptionalInput(OrderInput) + ), + GeoRadiusOutput + ) + + // We expect at least one of Option[Store] and Option[StoreDist] to be passed here. + final val GeoRadiusByMemberStore: RedisCommand[ ( String, String, @@ -199,7 +320,7 @@ private[redis] object Geo { Option[Store], Option[StoreDist] ), - Chunk[GeoView] + Long ] = RedisCommand( "GEORADIUSBYMEMBER", @@ -216,6 +337,6 @@ private[redis] object Geo { OptionalInput(StoreInput), OptionalInput(StoreDistInput) ), - GeoRadiusOutput + LongOutput ) } diff --git a/redis/src/main/scala/zio/redis/options/Geo.scala b/redis/src/main/scala/zio/redis/options/Geo.scala index 1b14559b5..47e724993 100644 --- a/redis/src/main/scala/zio/redis/options/Geo.scala +++ b/redis/src/main/scala/zio/redis/options/Geo.scala @@ -1,6 +1,7 @@ package zio.redis.options trait Geo { + this: Shared => sealed case class LongLat(longitude: Double, latitude: Double) @@ -23,6 +24,23 @@ trait Geo { case object Miles extends RadiusUnit } + sealed trait StoreOptions { + def store: Option[Store] + def storeDist: Option[StoreDist] + } + case class StoreResults(results: Store) extends StoreOptions { + override def store: Option[Store] = Some(results) + override def storeDist: Option[StoreDist] = None + } + case class StoreDistances(distances: StoreDist) extends StoreOptions { + override def store: Option[Store] = None + override def storeDist: Option[StoreDist] = Some(distances) + } + case class StoreBoth(results: Store, distances: StoreDist) extends StoreOptions { + override def store: Option[Store] = Some(results) + override def storeDist: Option[StoreDist] = Some(distances) + } + sealed case class StoreDist(key: String) case object WithCoord { diff --git a/redis/src/test/scala/zio/redis/GeoSpec.scala b/redis/src/test/scala/zio/redis/GeoSpec.scala index 74e489e7c..ef1d7b9fb 100644 --- a/redis/src/test/scala/zio/redis/GeoSpec.scala +++ b/redis/src/test/scala/zio/redis/GeoSpec.scala @@ -45,6 +45,14 @@ trait GeoSpec extends BaseSpec { hasSameElements(Chunk(GeoView(member1, None, None, None), GeoView(member2, None, None, None))) ) }, + testM("storing the result") { + import GeoSpec.Sicily._ + for { + dest <- uuid + _ <- geoAdd(key, member1LongLat -> member1, member2LongLat -> member2) + numStored <- geoRadiusStore(key, member1LongLat, 200d, RadiusUnit.Kilometers, StoreResults(Store(dest))) + } yield assert(numStored)(equalTo(2L)) + }, testM("with coordinates") { import GeoSpec.Sicily._ for { @@ -170,6 +178,14 @@ trait GeoSpec extends BaseSpec { hasSameElements(Chunk(GeoView(member1, None, None, None), GeoView(member2, None, None, None))) ) }, + testM("storing the result") { + import GeoSpec.Sicily._ + for { + dest <- uuid + _ <- geoAdd(key, member1LongLat -> member1, member2LongLat -> member2) + numStored <- geoRadiusByMemberStore(key, member1, 200d, RadiusUnit.Kilometers, StoreResults(Store(dest))) + } yield assert(numStored)(equalTo(2L)) + }, testM("with coordinates") { import GeoSpec.Sicily._ for {