Skip to content

Commit

Permalink
Python: adds GEOADD command
Browse files Browse the repository at this point in the history
  • Loading branch information
shohamazon committed Apr 10, 2024
1 parent 78cfa33 commit 10b0db3
Show file tree
Hide file tree
Showing 8 changed files with 172 additions and 1 deletion.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down
1 change: 1 addition & 0 deletions glide-core/src/protobuf/redis_request.proto
Original file line number Diff line number Diff line change
Expand Up @@ -143,6 +143,7 @@ enum RequestType {
Blpop = 100;
RPushX = 102;
LPushX = 103;
GeoAdd = 104;
}

message Command {
Expand Down
3 changes: 3 additions & 0 deletions glide-core/src/request_type.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -225,6 +226,7 @@ impl From<::protobuf::EnumOrUnknown<ProtobufRequestType>> for RequestType {
ProtobufRequestType::LPushX => RequestType::LPushX,
ProtobufRequestType::Blpop => RequestType::Blpop,
ProtobufRequestType::Spop => RequestType::Spop,
ProtobufRequestType::GeoAdd => RequestType::GeoAdd,
}
}
}
Expand Down Expand Up @@ -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")),
}
}
}
2 changes: 2 additions & 0 deletions python/python/glide/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

from glide.async_commands.core import (
ConditionalChange,
Coordinate,
ExpireOptions,
ExpirySet,
ExpiryType,
Expand Down Expand Up @@ -56,6 +57,7 @@
"RedisClientConfiguration",
"ScoreBoundary",
"ConditionalChange",
"Coordinate",
"ExpireOptions",
"ExpirySet",
"ExpiryType",
Expand Down
70 changes: 69 additions & 1 deletion python/python/glide/async_commands/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
"""
Expand Down Expand Up @@ -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."""

Expand Down Expand Up @@ -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,
Expand Down
43 changes: 43 additions & 0 deletions python/python/glide/async_commands/transaction.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@

from glide.async_commands.core import (
ConditionalChange,
Coordinate,
ExpireOptions,
ExpirySet,
InfoSection,
Expand Down Expand Up @@ -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,
Expand Down
42 changes: 42 additions & 0 deletions python/python/tests/test_async_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
from glide import ClosingError, RequestError, Script, TimeoutError
from glide.async_commands.core import (
ConditionalChange,
Coordinate,
ExpireOptions,
ExpirySet,
ExpiryType,
Expand Down Expand Up @@ -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):
Expand Down
11 changes: 11 additions & 0 deletions python/python/tests/test_transaction.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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)
Expand Down Expand Up @@ -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


Expand Down

0 comments on commit 10b0db3

Please sign in to comment.