From efc5f39da72f59db126386df19d2f13b7bba5ebb Mon Sep 17 00:00:00 2001 From: Shoham Elias Date: Sun, 16 Jun 2024 12:07:34 +0000 Subject: [PATCH] Python: adds GEOSEARCHSTORE command --- CHANGELOG.md | 1 + glide-core/src/protobuf/redis_request.proto | 1 + glide-core/src/request_type.rs | 3 + python/python/glide/async_commands/core.py | 68 +++- .../python/glide/async_commands/sorted_set.py | 8 +- .../glide/async_commands/transaction.py | 53 ++- python/python/tests/test_async_client.py | 318 +++++++++++++++++- python/python/tests/test_transaction.py | 10 + 8 files changed, 456 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6d0bb11e3a..d62796a2bf 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 bfd8f15cf2..2d68131274 100644 --- a/glide-core/src/protobuf/redis_request.proto +++ b/glide-core/src/protobuf/redis_request.proto @@ -222,6 +222,7 @@ enum RequestType { GeoSearch = 182; Watch = 183; UnWatch = 184; + GeoSearchStore = 185; } message Command { diff --git a/glide-core/src/request_type.rs b/glide-core/src/request_type.rs index 02155947e6..947236672e 100644 --- a/glide-core/src/request_type.rs +++ b/glide-core/src/request_type.rs @@ -192,6 +192,7 @@ pub enum RequestType { GeoSearch = 182, Watch = 183, UnWatch = 184, + GeoSearchStore = 185, } fn get_two_word_command(first: &str, second: &str) -> Cmd { @@ -387,6 +388,7 @@ impl From<::protobuf::EnumOrUnknown> for RequestType { ProtobufRequestType::GeoSearch => RequestType::GeoSearch, ProtobufRequestType::Watch => RequestType::Watch, ProtobufRequestType::UnWatch => RequestType::UnWatch, + ProtobufRequestType::GeoSearchStore => RequestType::GeoSearchStore, } } } @@ -578,6 +580,7 @@ impl RequestType { RequestType::GeoSearch => Some(cmd("GEOSEARCH")), RequestType::Watch => Some(cmd("WATCH")), RequestType::UnWatch => Some(cmd("UNWATCH")), + RequestType::GeoSearchStore => Some(cmd("GEOSEARCHSTORE")), } } } diff --git a/python/python/glide/async_commands/core.py b/python/python/glide/async_commands/core.py index 5d2473f493..0ebc864673 100644 --- a/python/python/glide/async_commands/core.py +++ b/python/python/glide/async_commands/core.py @@ -2708,7 +2708,7 @@ async def geosearch( Since: Redis version 6.2.0. """ args = _create_geosearch_args( - key, + [key], search_from, seach_by, order_by, @@ -2723,6 +2723,72 @@ 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: + """ + Searches for members in a sorted set stored at `key` representing geospatial data within a circular or rectangular area and stores the result in `destination`. + If `destination` already exists, it is overwritten. Otherwise, a new sorted set will be created. + + To get the result directly, see `geosearch`. + + Note: + When in cluster mode, both `source` and `destination` must map to the same hash slot. + + Args: + destination (str): The 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): Determines what is stored as the sorted set score. Defaults to False. + - If set to False, the geohash of the location will be stored as the sorted set score. + - If set to True, the distance from the center of the shape (circle or box) will be stored as the sorted set score. + The distance is represented as a floating-point number in the same unit specified for that shape. + + Returns: + int: The number of elements in the resulting sorted set stored at `destination`. + + 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..274501ba7f 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 383c3487d3..cabd250157 100644 --- a/python/python/glide/async_commands/transaction.py +++ b/python/python/glide/async_commands/transaction.py @@ -1872,7 +1872,7 @@ def geosearch( Since: Redis version 6.2.0. """ args = _create_geosearch_args( - key, + [key], search_from, seach_by, order_by, @@ -1884,6 +1884,57 @@ 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: + """ + Searches for members in a sorted set stored at `key` representing geospatial data within a circular or rectangular area and stores the result in `destination`. + If `destination` already exists, it is overwritten. Otherwise, a new sorted set will be created. + + To get the result directly, see `geosearch`. + + See https://valkey.io/commands/geosearch/ for more details. + + Args: + destination (str): The 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): Determines what is stored as the sorted set score. Defaults to False. + - If set to False, the geohash of the location will be stored as the sorted set score. + - If set to True, the distance from the center of the shape (circle or box) will be stored as the sorted set score. + The distance is represented as a floating-point number in the same unit specified for that shape. + + Commands response: + int: The number of elements in the resulting sorted set stored at `destination`.s + + 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 eae0d7a325..337967cd70 100644 --- a/python/python/tests/test_async_client.py +++ b/python/python/tests/test_async_client.py @@ -2205,7 +2205,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, @@ -2228,7 +2228,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 @@ -2307,6 +2307,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): @@ -4598,6 +4900,18 @@ async def test_multi_key_command_returns_cross_slot_error( redis_client.msetnx({"abc": "abc", "zxy": "zyx"}), ] + if not await check_if_server_version_lt(redis_client, "6.2.0"): + promises.extend( + [ + 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"): promises.extend( [ diff --git a/python/python/tests/test_transaction.py b/python/python/tests/test_transaction.py index f576371906..f933ad8d87 100644 --- a/python/python/tests/test_transaction.py +++ b/python/python/tests/test_transaction.py @@ -11,6 +11,7 @@ from glide.async_commands.core import InsertPosition, StreamAddOptions, TrimByMinId from glide.async_commands.sorted_set import ( AggregationType, + GeoSearchByBox, GeoSearchByRadius, GeospatialData, GeoUnit, @@ -390,10 +391,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")