diff --git a/CHANGELOG.md b/CHANGELOG.md index 58b4263627..b55b2f1b07 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ * Python: Added OBJECT FREQ command ([#1472](https://github.com/aws/glide-for-redis/pull/1472)) * Python: Added OBJECT IDLETIME command ([#1474](https://github.com/aws/glide-for-redis/pull/1474)) * Python: Added GEOSEARCH command ([#1482](https://github.com/aws/glide-for-redis/pull/1482)) +* Python: Added GEOSEARCHSTORE command ([#1581](https://github.com/aws/glide-for-redis/pull/1581)) * Node: Added RENAMENX command ([#1483](https://github.com/aws/glide-for-redis/pull/1483)) * Python: Added OBJECT REFCOUNT command ([#1485](https://github.com/aws/glide-for-redis/pull/1485)) * Python: Added RENAMENX command ([#1492](https://github.com/aws/glide-for-redis/pull/1492)) diff --git a/glide-core/src/protobuf/redis_request.proto b/glide-core/src/protobuf/redis_request.proto index 8056bf308e..0e38eefe76 100644 --- a/glide-core/src/protobuf/redis_request.proto +++ b/glide-core/src/protobuf/redis_request.proto @@ -219,6 +219,7 @@ enum RequestType { LPos = 180; LCS = 181; GeoSearch = 182; + GeoSearchStore = 183; } message Command { diff --git a/glide-core/src/request_type.rs b/glide-core/src/request_type.rs index 8e417b91f4..6387963e50 100644 --- a/glide-core/src/request_type.rs +++ b/glide-core/src/request_type.rs @@ -189,6 +189,7 @@ pub enum RequestType { LPos = 180, LCS = 181, GeoSearch = 182, + GeoSearchStore = 183, } fn get_two_word_command(first: &str, second: &str) -> Cmd { @@ -381,6 +382,7 @@ impl From<::protobuf::EnumOrUnknown> for RequestType { ProtobufRequestType::LPos => RequestType::LPos, ProtobufRequestType::LCS => RequestType::LCS, ProtobufRequestType::GeoSearch => RequestType::GeoSearch, + ProtobufRequestType::GeoSearchStore => RequestType::GeoSearchStore, } } } @@ -569,6 +571,7 @@ impl RequestType { RequestType::LPos => Some(cmd("LPOS")), RequestType::LCS => Some(cmd("LCS")), RequestType::GeoSearch => Some(cmd("GEOSEARCH")), + RequestType::GeoSearchStore => Some(cmd("GEOSEARCHSTORE")), } } } diff --git a/python/python/glide/async_commands/core.py b/python/python/glide/async_commands/core.py index 4e33247192..fd0ec6720b 100644 --- a/python/python/glide/async_commands/core.py +++ b/python/python/glide/async_commands/core.py @@ -2684,7 +2684,7 @@ async def geosearch( Since: Redis version 6.2.0. """ args = _create_geosearch_args( - key, + [key], search_from, seach_by, order_by, @@ -2699,6 +2699,70 @@ async def geosearch( await self._execute_command(RequestType.GeoSearch, args), ) + async def geosearchstore( + self, + destination: str, + source: str, + search_from: Union[str, GeospatialData], + search_by: Union[GeoSearchByRadius, GeoSearchByBox], + count: Optional[GeoSearchCount] = None, + store_dist: bool = False, + ) -> int: + """ + Stores the results of a geospatial search into a destination sorted set. + + See https://valkey.io/commands/geosearchstore/ for more details. + + Note: + When in cluster mode, both `source` and `destination` must map to the same hash slot. + + Args: + destination (str): TThe key to store the search results. + source (str): The key of the sorted set representing geospatial data to search from. + search_from (Union[str, GeospatialData]): The location to search from. Can be specified either as a member + from the sorted set or as a geospatial data (see `GeospatialData`). + search_by (Union[GeoSearchByRadius, GeoSearchByBox]): The search criteria. + For circular area search, see `GeoSearchByRadius`. + For rectangular area search, see `GeoSearchByBox`. + count (Optional[GeoSearchCount]): Specifies the maximum number of results to store. See `GeoSearchCount`. + If not specified, stores all results. + store_dist (bool): Whether to store the distance from the center as the sorted set score. Defaults to False. + - The distance is from the center of the circle or box, as a floating-point number, in the same unit specified for that shape. + - If set to False, the geohash of the location will be stored as the sorted set score. + + Returns: + int: The number of elements in the resulting sorted set. + + Examples: + >>> await client.geoadd("my_geo_sorted_set", {"Palermo": GeospatialData(13.361389, 38.115556), "Catania": GeospatialData(15.087269, 37.502669)}) + >>> await client.geosearchstore("my_dest_sorted_set", "my_geo_sorted_set", "Catania", GeoSearchByRadius(175, GeoUnit.MILES)) + 2 # Number of elements stored in "my_dest_sorted_set". + >>> await client.zrange_withscores("my_dest_sorted_set", RangeByIndex(0, -1)) + {"Palermo": 3479099956230698.0, "Catania": 3479447370796909.0} # The elements within te search area, with their geohash as score. + >>> await client.geosearchstore("my_dest_sorted_set", "my_geo_sorted_set", GeospatialData(15, 37), GeoSearchByBox(400, 400, GeoUnit.KILOMETERS), store_dist=True) + 2 # Number of elements stored in "my_dest_sorted_set", with distance as score. + >>> await client.zrange_withscores("my_dest_sorted_set", RangeByIndex(0, -1)) + {"Catania": 56.4412578701582, "Palermo": 190.44242984775784} # The elements within te search area, with the distance as score. + + Since: Redis version 6.2.0. + """ + args = _create_geosearch_args( + [destination, source], + search_from, + search_by, + None, + count, + False, + False, + False, + store_dist, + ) + + return cast( + int, + await self._execute_command(RequestType.GeoSearchStore, args), + ) + async def zadd( self, key: str, diff --git a/python/python/glide/async_commands/sorted_set.py b/python/python/glide/async_commands/sorted_set.py index d9f99a42c9..b0c608fde0 100644 --- a/python/python/glide/async_commands/sorted_set.py +++ b/python/python/glide/async_commands/sorted_set.py @@ -356,7 +356,7 @@ def _create_zinter_zunion_cmd_args( def _create_geosearch_args( - key: str, + keys: List[str], search_from: Union[str, GeospatialData], seach_by: Union[GeoSearchByRadius, GeoSearchByBox], order_by: Optional[OrderBy] = None, @@ -364,8 +364,9 @@ def _create_geosearch_args( with_coord: bool = False, with_dist: bool = False, with_hash: bool = False, + store_dist: bool = False, ) -> List[str]: - args = [key] + args = [*keys] if isinstance(search_from, str): args += ["FROMMEMBER", search_from] else: @@ -389,4 +390,7 @@ def _create_geosearch_args( if with_hash: args.append("WITHHASH") + if store_dist: + args.append("STOREDIST") + return args diff --git a/python/python/glide/async_commands/transaction.py b/python/python/glide/async_commands/transaction.py index 743f3e8fe8..658aa9677e 100644 --- a/python/python/glide/async_commands/transaction.py +++ b/python/python/glide/async_commands/transaction.py @@ -1856,7 +1856,7 @@ def geosearch( Since: Redis version 6.2.0. """ args = _create_geosearch_args( - key, + [key], search_from, seach_by, order_by, @@ -1868,6 +1868,53 @@ def geosearch( return self.append_command(RequestType.GeoSearch, args) + def geosearchstore( + self: TTransaction, + destination: str, + source: str, + search_from: Union[str, GeospatialData], + search_by: Union[GeoSearchByRadius, GeoSearchByBox], + count: Optional[GeoSearchCount] = None, + store_dist: bool = False, + ) -> TTransaction: + """ + Stores the results of a geospatial search into a destination sorted set. + + See https://valkey.io/commands/geosearchstore/ for more details. + + Args: + destination (str): TThe key to store the search results. + source (str): The key of the sorted set representing geospatial data to search from. + search_from (Union[str, GeospatialData]): The location to search from. Can be specified either as a member + from the sorted set or as a geospatial data (see `GeospatialData`). + search_by (Union[GeoSearchByRadius, GeoSearchByBox]): The search criteria. + For circular area search, see `GeoSearchByRadius`. + For rectangular area search, see `GeoSearchByBox`. + count (Optional[GeoSearchCount]): Specifies the maximum number of results to store. See `GeoSearchCount`. + If not specified, stores all results. + store_dist (bool): Whether to store the distance from the center as the sorted set score. Defaults to False. + - The distance is from the center of the circle or box, as a floating-point number, in the same unit specified for that shape. + - If set to False, the geohash of the location will be stored as the sorted set score. + + Commands response: + int: The number of elements in the resulting sorted set. + + Since: Redis version 6.2.0. + """ + args = _create_geosearch_args( + [destination, source], + search_from, + search_by, + None, + count, + False, + False, + False, + store_dist, + ) + + return self.append_command(RequestType.GeoSearchStore, args) + def zadd( self: TTransaction, key: str, diff --git a/python/python/tests/test_async_client.py b/python/python/tests/test_async_client.py index 0f94589515..12a79a5abe 100644 --- a/python/python/tests/test_async_client.py +++ b/python/python/tests/test_async_client.py @@ -2189,7 +2189,7 @@ async def test_geosearch_by_radius(self, redis_client: TRedisClient): == members[:2][::-1] ) - # Test search by radius, unit: miles, from a geospatial data, with limited count to 1 + # Test search by radius, unit: miles, from a geospatial data assert ( await redis_client.geosearch( key, @@ -2212,7 +2212,7 @@ async def test_geosearch_by_radius(self, redis_client: TRedisClient): with_dist=True, with_hash=True, ) - == result[:2] + == result ) # Test search by radius, unit: kilometers, from a geospatial data, with limited ANY count to 1 @@ -2291,6 +2291,308 @@ async def test_geosearch_no_result(self, redis_client: TRedisClient): GeoSearchByBox(10, 10, GeoUnit.MILES), ) + @pytest.mark.parametrize("cluster_mode", [True, False]) + @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3]) + async def test_geosearchstore_by_box(self, redis_client: TRedisClient): + key = f"{{testKey}}:{get_random_string(10)}" + destination_key = f"{{testKey}}:{get_random_string(8)}" + members_coordinates = { + "Palermo": GeospatialData(13.361389, 38.115556), + "Catania": GeospatialData(15.087269, 37.502669), + "edge1": GeospatialData(12.758489, 38.788135), + "edge2": GeospatialData(17.241510, 38.788135), + } + result = { + "Catania": [56.4412578701582, 3479447370796909.0], + "Palermo": [190.44242984775784, 3479099956230698.0], + "edge2": [279.7403417843143, 3481342659049484.0], + "edge1": [279.7404521356343, 3479273021651468.0], + } + assert await redis_client.geoadd(key, members_coordinates) == 4 + + # Test storing results of a box search, unit: kilometes, from a geospatial data + assert ( + await redis_client.geosearchstore( + destination_key, + key, + GeospatialData(15, 37), + GeoSearchByBox(400, 400, GeoUnit.KILOMETERS), + ) + ) == 4 # Number of elements stored + + # Verify the stored results + zrange_map = await redis_client.zrange_withscores( + destination_key, RangeByIndex(0, -1) + ) + expected_map = {member: value[1] for member, value in result.items()} + sorted_expected_map = dict(sorted(expected_map.items(), key=lambda x: x[1])) + assert compare_maps(zrange_map, sorted_expected_map) is True + + # Test storing results of a box search, unit: kilometes, from a geospatial data, with distance + assert ( + await redis_client.geosearchstore( + destination_key, + key, + GeospatialData(15, 37), + GeoSearchByBox(400, 400, GeoUnit.KILOMETERS), + store_dist=True, + ) + ) == 4 # Number of elements stored + + # Verify the stored results + zrange_map = await redis_client.zrange_withscores( + destination_key, RangeByIndex(0, -1) + ) + expected_map = {member: value[0] for member, value in result.items()} + sorted_expected_map = dict(sorted(expected_map.items(), key=lambda x: x[1])) + assert compare_maps(zrange_map, sorted_expected_map) is True + + # Test storing results of a box search, unit: kilometes, from a geospatial data, with count + assert ( + await redis_client.geosearchstore( + destination_key, + key, + GeospatialData(15, 37), + GeoSearchByBox(400, 400, GeoUnit.KILOMETERS), + count=GeoSearchCount(1), + ) + ) == 1 # Number of elements stored + + # Verify the stored results + zrange_map = await redis_client.zrange_withscores( + destination_key, RangeByIndex(0, -1) + ) + assert compare_maps(zrange_map, {"Catania": 3479447370796909.0}) is True + + # Test storing results of a box search, unit: meters, from a member, with distance + meters = 400 * 1000 + assert ( + await redis_client.geosearchstore( + destination_key, + key, + "Catania", + GeoSearchByBox(meters, meters, GeoUnit.METERS), + store_dist=True, + ) + ) == 3 # Number of elements stored + + # Verify the stored results with distances + zrange_map = await redis_client.zrange_withscores( + destination_key, RangeByIndex(0, -1) + ) + expected_distances = { + "Catania": 0.0, + "Palermo": 166274.15156960033, + "edge2": 236529.17986494553, + } + assert compare_maps(zrange_map, expected_distances) is True + + # Test search by box, unit: feet, from a member, with limited ANY count to 2, with hash + feet = 400 * 3280.8399 + assert ( + await redis_client.geosearchstore( + destination_key, + key, + "Palermo", + GeoSearchByBox(feet, feet, GeoUnit.FEET), + count=GeoSearchCount(2), + ) + == 2 + ) + + # Verify the stored results + zrange_map = await redis_client.zrange_withscores( + destination_key, RangeByIndex(0, -1) + ) + for member in zrange_map: + assert member in result + + @pytest.mark.parametrize("cluster_mode", [True, False]) + @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3]) + async def test_geosearchstore_by_radius(self, redis_client: TRedisClient): + key = f"{{testKey}}:{get_random_string(10)}" + destination_key = f"{{testKey}}:{get_random_string(8)}" + members_coordinates = { + "Palermo": GeospatialData(13.361389, 38.115556), + "Catania": GeospatialData(15.087269, 37.502669), + "edge1": GeospatialData(12.758489, 38.788135), + "edge2": GeospatialData(17.241510, 38.788135), + } + result = { + "Catania": [56.4412578701582, 3479447370796909.0], + "Palermo": [190.44242984775784, 3479099956230698.0], + } + assert await redis_client.geoadd(key, members_coordinates) == 4 + + # Test storing results of a radius search, unit: feet, from a member + feet = 200 * 3280.8399 + assert ( + await redis_client.geosearchstore( + destination_key, + key, + "Catania", + GeoSearchByRadius(feet, GeoUnit.FEET), + ) + == 2 + ) + + # Verify the stored results + zrange_map = await redis_client.zrange_withscores( + destination_key, RangeByIndex(0, -1) + ) + expected_map = {member: value[1] for member, value in result.items()} + sorted_expected_map = dict(sorted(expected_map.items(), key=lambda x: x[1])) + assert compare_maps(zrange_map, sorted_expected_map) is True + + # Test search by radius, units: meters, from a member + meters = 200 * 1000 + assert ( + await redis_client.geosearchstore( + destination_key, + key, + "Catania", + GeoSearchByRadius(meters, GeoUnit.METERS), + store_dist=True, + ) + == 2 + ) + + # Verify the stored results + zrange_map = await redis_client.zrange_withscores( + destination_key, RangeByIndex(0, -1) + ) + expected_distances = { + "Catania": 0.0, + "Palermo": 166274.15156960033, + } + assert compare_maps(zrange_map, expected_distances) is True + + # Test search by radius, unit: miles, from a geospatial data + assert ( + await redis_client.geosearchstore( + destination_key, + key, + GeospatialData(15, 37), + GeoSearchByRadius(175, GeoUnit.MILES), + ) + == 4 + ) + + # Test storing results of a radius search, unit: kilometers, from a geospatial data, with limited count to 2 + kilometers = 200 + assert ( + await redis_client.geosearchstore( + destination_key, + key, + GeospatialData(15, 37), + GeoSearchByRadius(kilometers, GeoUnit.KILOMETERS), + count=GeoSearchCount(2), + store_dist=True, + ) + == 2 + ) + + # Verify the stored results + zrange_map = await redis_client.zrange_withscores( + destination_key, RangeByIndex(0, -1) + ) + expected_map = {member: value[0] for member, value in result.items()} + sorted_expected_map = dict(sorted(expected_map.items(), key=lambda x: x[1])) + assert compare_maps(zrange_map, sorted_expected_map) is True + + # Test storing results of a radius search, unit: kilometers, from a geospatial data, with limited ANY count to 1 + assert ( + await redis_client.geosearchstore( + destination_key, + key, + GeospatialData(15, 37), + GeoSearchByRadius(kilometers, GeoUnit.KILOMETERS), + count=GeoSearchCount(1, True), + ) + == 1 + ) + + # Verify the stored results + zrange_map = await redis_client.zrange_withscores( + destination_key, RangeByIndex(0, -1) + ) + + for member in zrange_map: + assert member in result + + @pytest.mark.parametrize("cluster_mode", [True, False]) + @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3]) + async def test_geosearchstore_no_result(self, redis_client: TRedisClient): + key = f"{{testKey}}:{get_random_string(10)}" + destination_key = f"{{testKey}}:{get_random_string(8)}" + members_coordinates = { + "Palermo": GeospatialData(13.361389, 38.115556), + "Catania": GeospatialData(15.087269, 37.502669), + "edge1": GeospatialData(12.758489, 38.788135), + "edge2": GeospatialData(17.241510, 38.788135), + } + assert await redis_client.geoadd(key, members_coordinates) == 4 + + # No members within the area + assert ( + await redis_client.geosearchstore( + destination_key, + key, + GeospatialData(15, 37), + GeoSearchByBox(50, 50, GeoUnit.METERS), + ) + == 0 + ) + + assert ( + await redis_client.geosearchstore( + destination_key, + key, + GeospatialData(15, 37), + GeoSearchByRadius(10, GeoUnit.METERS), + ) + == 0 + ) + + # No members in the area (apart from the member we search from itself) + assert ( + await redis_client.geosearchstore( + destination_key, + key, + "Catania", + GeoSearchByBox(10, 10, GeoUnit.KILOMETERS), + ) + == 1 + ) + + assert ( + await redis_client.geosearchstore( + destination_key, + key, + "Catania", + GeoSearchByRadius(10, GeoUnit.METERS), + ) + == 1 + ) + + # Search from non-existing member + with pytest.raises(RequestError): + await redis_client.geosearchstore( + destination_key, + key, + "non_existing_member", + GeoSearchByBox(10, 10, GeoUnit.MILES), + ) + + assert await redis_client.set(key, "foo") == OK + with pytest.raises(RequestError): + await redis_client.geosearchstore( + destination_key, + key, + "Catania", + GeoSearchByBox(10, 10, GeoUnit.MILES), + ) + @pytest.mark.parametrize("cluster_mode", [True, False]) @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3]) async def test_geohash(self, redis_client: TRedisClient): @@ -4491,6 +4793,12 @@ async def test_multi_key_command_returns_cross_slot_error( "abc", "zxy", ListDirection.LEFT, ListDirection.LEFT, 1 ), redis_client.msetnx({"abc": "abc", "zxy": "zyx"}), + redis_client.geosearchstore( + "abc", + "zxy", + GeospatialData(15, 37), + GeoSearchByBox(400, 400, GeoUnit.KILOMETERS), + ), ] if not await check_if_server_version_lt(redis_client, "7.0.0"): diff --git a/python/python/tests/test_transaction.py b/python/python/tests/test_transaction.py index bd547f6410..97787ac2df 100644 --- a/python/python/tests/test_transaction.py +++ b/python/python/tests/test_transaction.py @@ -10,6 +10,7 @@ from glide.async_commands.core import InsertPosition, StreamAddOptions, TrimByMinId from glide.async_commands.sorted_set import ( AggregationType, + GeoSearchByBox, GeoSearchByRadius, GeospatialData, GeoUnit, @@ -373,10 +374,19 @@ async def transaction_test( None, ] ) + transaction.geosearch( key12, "Catania", GeoSearchByRadius(200, GeoUnit.KILOMETERS), OrderBy.ASC ) args.append(["Catania", "Palermo"]) + transaction.geosearchstore( + key12, + key12, + GeospatialData(15, 37), + GeoSearchByBox(400, 400, GeoUnit.KILOMETERS), + store_dist=True, + ) + args.append(2) transaction.xadd(key11, [("foo", "bar")], StreamAddOptions(id="0-1")) args.append("0-1")