Python: adds GEOSEARCHSTORE command (valkey-io#1581)

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))

glide-core/src/protobuf/redis_request.proto

@@ -222,6 +222,7 @@ enum RequestType {
     GeoSearch = 182;
     Watch = 183;
     UnWatch = 184;
+    GeoSearchStore = 185;
     GetEx = 183;
 }
 

glide-core/src/request_type.rs

@@ -192,7 +192,8 @@ pub enum RequestType {
     GeoSearch = 182,
     Watch = 183,
     UnWatch = 184,
-    GetEx = 183,
+    GeoSearchStore = 185,
+    GetEx = 186,
 }
 
 fn get_two_word_command(first: &str, second: &str) -> Cmd {
@@ -388,6 +389,7 @@ impl From<::protobuf::EnumOrUnknown<ProtobufRequestType>> for RequestType {
             ProtobufRequestType::GeoSearch => RequestType::GeoSearch,
             ProtobufRequestType::Watch => RequestType::Watch,
             ProtobufRequestType::UnWatch => RequestType::UnWatch,
+            ProtobufRequestType::GeoSearchStore => RequestType::GeoSearchStore,
             ProtobufRequestType::GetEx => RequestType::GetEx,
         }
     }
@@ -580,6 +582,7 @@ impl RequestType {
             RequestType::GeoSearch => Some(cmd("GEOSEARCH")),
             RequestType::Watch => Some(cmd("WATCH")),
             RequestType::UnWatch => Some(cmd("UNWATCH")),
+            RequestType::GeoSearchStore => Some(cmd("GEOSEARCHSTORE")),
             RequestType::GetEx => Some(cmd("GETEX")),
         }
     }

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,

python/python/glide/async_commands/sorted_set.py

@@ -356,16 +356,17 @@ 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,
     count: Optional[GeoSearchCount] = None,
     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

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,