From 10b0db3a9383f14cb8a53de163c8b3e288534817 Mon Sep 17 00:00:00 2001 From: Shoham Elias Date: Wed, 10 Apr 2024 09:50:53 +0000 Subject: [PATCH] Python: adds GEOADD command --- CHANGELOG.md | 1 + glide-core/src/protobuf/redis_request.proto | 1 + glide-core/src/request_type.rs | 3 + python/python/glide/__init__.py | 2 + python/python/glide/async_commands/core.py | 70 ++++++++++++++++++- .../glide/async_commands/transaction.py | 43 ++++++++++++ python/python/tests/test_async_client.py | 42 +++++++++++ python/python/tests/test_transaction.py | 11 +++ 8 files changed, 172 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a609b1dd79..6d6148b8ae 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ * Python: Added HKEYS command ([#1228](https://github.com/aws/glide-for-redis/pull/1228)) * Python: Added ZREMRANGEBYSCORE command ([#1151](https://github.com/aws/glide-for-redis/pull/1151)) * Node: Added SPOP, SPOPCOUNT commands. ([#1117](https://github.com/aws/glide-for-redis/pull/1117)) +* Python: Added GEOADD command ([#1258](https://github.com/aws/glide-for-redis/pull/1258)) #### Fixes * Python: Fix typing error "‘type’ object is not subscriptable" ([#1203](https://github.com/aws/glide-for-redis/pull/1203)) diff --git a/glide-core/src/protobuf/redis_request.proto b/glide-core/src/protobuf/redis_request.proto index 08428b92ea..1688e4acc4 100644 --- a/glide-core/src/protobuf/redis_request.proto +++ b/glide-core/src/protobuf/redis_request.proto @@ -143,6 +143,7 @@ enum RequestType { Blpop = 100; RPushX = 102; LPushX = 103; + GeoAdd = 104; } message Command { diff --git a/glide-core/src/request_type.rs b/glide-core/src/request_type.rs index ffed834871..7cb361b635 100644 --- a/glide-core/src/request_type.rs +++ b/glide-core/src/request_type.rs @@ -111,6 +111,7 @@ pub enum RequestType { Blpop = 100, RPushX = 102, LPushX = 103, + GeoAdd = 104, } fn get_two_word_command(first: &str, second: &str) -> Cmd { @@ -225,6 +226,7 @@ impl From<::protobuf::EnumOrUnknown> for RequestType { ProtobufRequestType::LPushX => RequestType::LPushX, ProtobufRequestType::Blpop => RequestType::Blpop, ProtobufRequestType::Spop => RequestType::Spop, + ProtobufRequestType::GeoAdd => RequestType::GeoAdd, } } } @@ -335,6 +337,7 @@ impl RequestType { RequestType::LPushX => Some(cmd("LPUSHX")), RequestType::Blpop => Some(cmd("BLPOP")), RequestType::Spop => Some(cmd("SPOP")), + RequestType::GeoAdd => Some(cmd("GEOADD")), } } } diff --git a/python/python/glide/__init__.py b/python/python/glide/__init__.py index 2750f617fd..2a10f291ea 100644 --- a/python/python/glide/__init__.py +++ b/python/python/glide/__init__.py @@ -2,6 +2,7 @@ from glide.async_commands.core import ( ConditionalChange, + Coordinate, ExpireOptions, ExpirySet, ExpiryType, @@ -56,6 +57,7 @@ "RedisClientConfiguration", "ScoreBoundary", "ConditionalChange", + "Coordinate", "ExpireOptions", "ExpirySet", "ExpiryType", diff --git a/python/python/glide/async_commands/core.py b/python/python/glide/async_commands/core.py index 7bc02edf76..4724021801 100644 --- a/python/python/glide/async_commands/core.py +++ b/python/python/glide/async_commands/core.py @@ -33,7 +33,7 @@ class ConditionalChange(Enum): """ - A condition to the "SET" and "ZADD" commands. + A condition to the `SET`, `ZADD` and `GEOADD` commands. - ONLY_IF_EXISTS - Only update key / elements that already exist. Equivalent to `XX` in the Redis API - ONLY_IF_DOES_NOT_EXIST - Only set key / add elements that does not already exist. Equivalent to `NX` in the Redis API """ @@ -131,6 +131,23 @@ class UpdateOptions(Enum): GREATER_THAN = "GT" +class Coordinate: + def __init__(self, longitude: float, latitude: float): + """ + Represents a geographic position defined by longitude and latitude. + + The exact limits, as specified by EPSG:900913 / EPSG:3785 / OSGEO:41001 are the following: + - Valid longitudes are from -180 to 180 degrees. + - Valid latitudes are from -85.05112878 to 85.05112878 degrees. + + Args: + longitude (float): The longitude coordinate. + latitude (float): The latitude coordinate. + """ + self.longitude = longitude + self.latitude = latitude + + class ExpirySet: """SET option: Represents the expiry type and value to be executed with "SET" command.""" @@ -1386,6 +1403,57 @@ async def type(self, key: str) -> str: """ return cast(str, await self._execute_command(RequestType.Type, [key])) + async def geoadd( + self, + key: str, + members_coordinates: Mapping[str, Coordinate], + existing_options: Optional[ConditionalChange] = None, + changed: bool = False, + ) -> int: + """ + Adds geospatial members with their positions to the specified sorted set stored at `key`. + If a member is already a part of the sorted set, its position is updated. + + See https://redis.io/commands/geoadd for more details. + + Args: + key (str): The key of the sorted set. + members_coordinates (Mapping[str, Coordinate]): A mapping of member names to their corresponding positions. See `Coordinate`. + The command will report an error when the user attempts to index coordinates outside the specified ranges. + existing_options (Optional[ConditionalChange]): Options for handling existing members. + - NX: Only add new elements. + - XX: Only update existing elements. + changed (bool): Modify the return value to return the number of changed elements, instead of the number of new elements added. + + Returns: + int: The number of elements added to the sorted set. + If `changed` is set, returns the number of elements updated in the sorted set. + + Examples: + >>> await geoadd("my_sorted_set", {"Palermo": Coordinate(13.361389, 38.115556), "Catania": Coordinate(15.087269, 37.502669)}) + 2 # Indicates that two elements have been added to the sorted set "my_sorted_set". + >>> await geoadd("my_sorted_set", {"Palermo": Coordinate(14.361389, 38.115556)}, existing_options=ConditionalChange.XX, changed=True) + 1 # Updates the position of an existing member in the sorted set "my_sorted_set". + """ + args = [key] + if existing_options: + args.append(existing_options.value) + + if changed: + args.append("CH") + + members_coordinates_list = [ + coord + for member, position in members_coordinates.items() + for coord in [str(position.longitude), str(position.latitude), member] + ] + args += members_coordinates_list + + return cast( + int, + await self._execute_command(RequestType.GeoAdd, args), + ) + async def zadd( self, key: str, diff --git a/python/python/glide/async_commands/transaction.py b/python/python/glide/async_commands/transaction.py index 233c584c0c..6dbf0655b8 100644 --- a/python/python/glide/async_commands/transaction.py +++ b/python/python/glide/async_commands/transaction.py @@ -5,6 +5,7 @@ from glide.async_commands.core import ( ConditionalChange, + Coordinate, ExpireOptions, ExpirySet, InfoSection, @@ -1068,6 +1069,48 @@ def type(self: TTransaction, key: str) -> TTransaction: """ return self.append_command(RequestType.Type, [key]) + def geoadd( + self: TTransaction, + key: str, + members_coordinates: Mapping[str, Coordinate], + existing_options: Optional[ConditionalChange] = None, + changed: bool = False, + ) -> TTransaction: + """ + Adds geospatial members with their positions to the specified sorted set stored at `key`. + If a member is already a part of the sorted set, its position is updated. + + See https://redis.io/commands/geoadd for more details. + + Args: + key (str): The key of the sorted set. + members_coordinates (Mapping[str, Coordinate]): A mapping of member names to their corresponding positions. See `Coordinate`. + The command will report an error when the user attempts to index coordinates outside the specified ranges. + existing_options (Optional[ConditionalChange]): Options for handling existing members. + - NX: Only add new elements. + - XX: Only update existing elements. + changed (bool): Modify the return value to return the number of changed elements, instead of the number of new elements added. + + Returns: + int: The number of elements added to the sorted set. + If `changed` is set, returns the number of elements updated in the sorted set. + """ + args = [key] + if existing_options: + args.append(existing_options.value) + + if changed: + args.append("CH") + + members_coordinates_list = [ + coord + for member, position in members_coordinates.items() + for coord in [str(position.longitude), str(position.latitude), member] + ] + args += members_coordinates_list + + return self.append_command(RequestType.GeoAdd, 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 7d5131b201..15b59dd41e 100644 --- a/python/python/tests/test_async_client.py +++ b/python/python/tests/test_async_client.py @@ -13,6 +13,7 @@ from glide import ClosingError, RequestError, Script, TimeoutError from glide.async_commands.core import ( ConditionalChange, + Coordinate, ExpireOptions, ExpirySet, ExpiryType, @@ -1140,6 +1141,47 @@ async def test_persist(self, redis_client: TRedisClient): assert await redis_client.expire(key, 10) assert await redis_client.persist(key) + @pytest.mark.parametrize("cluster_mode", [True, False]) + @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3]) + async def test_geoadd(self, redis_client: TRedisClient): + key, key2 = get_random_string(10), get_random_string(10) + members_coordinates = { + "Palmero": Coordinate(13.361389, 38.115556), + "Catania": Coordinate(15.087269, 37.502669), + } + assert await redis_client.geoadd(key, members_coordinates) == 2 + members_coordinates["Catania"].latitude = 39 + assert ( + await redis_client.geoadd( + key, + members_coordinates, + existing_options=ConditionalChange.ONLY_IF_DOES_NOT_EXIST, + ) + == 0 + ) + assert ( + await redis_client.geoadd( + key, + members_coordinates, + existing_options=ConditionalChange.ONLY_IF_EXISTS, + ) + == 0 + ) + members_coordinates["Catania"].latitude = 40 + members_coordinates.update({"Tel-Aviv": Coordinate(32.0853, 34.7818)}) + assert ( + await redis_client.geoadd( + key, + members_coordinates, + changed=True, + ) + == 2 + ) + + assert await redis_client.set(key2, "value") == OK + with pytest.raises(RequestError): + await redis_client.geoadd(key2, members_coordinates) + @pytest.mark.parametrize("cluster_mode", [True, False]) @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3]) async def test_zadd_zaddincr(self, redis_client: TRedisClient): diff --git a/python/python/tests/test_transaction.py b/python/python/tests/test_transaction.py index 2828bdd620..e8d4540d10 100644 --- a/python/python/tests/test_transaction.py +++ b/python/python/tests/test_transaction.py @@ -5,6 +5,7 @@ import pytest from glide import RequestError +from glide.async_commands.core import Coordinate from glide.async_commands.sorted_set import InfBound, RangeByIndex, ScoreBoundary from glide.async_commands.transaction import ( BaseTransaction, @@ -31,6 +32,7 @@ async def transaction_test( key6 = "{{{}}}:{}".format(keyslot, get_random_string(3)) key7 = "{{{}}}:{}".format(keyslot, get_random_string(3)) key8 = "{{{}}}:{}".format(keyslot, get_random_string(3)) + key9 = "{{{}}}:{}".format(keyslot, get_random_string(3)) value = datetime.now(timezone.utc).strftime("%m/%d/%Y, %H:%M:%S") value2 = get_random_string(5) @@ -181,6 +183,15 @@ async def transaction_test( args.append({"four": 4}) transaction.zremrangebyscore(key8, InfBound.NEG_INF, InfBound.POS_INF) args.append(1) + + transaction.geoadd( + key9, + { + "Palermo": Coordinate(13.361389, 38.115556), + "Catania": Coordinate(15.087269, 37.502669), + }, + ) + args.append(2) return args