Python: adds GEOSEARCH command

valkey-io · May 28, 2024 · df781a6 · df781a6
1 parent b375b57
commit df781a6
Show file tree

Hide file tree

Showing 10 changed files with 702 additions and 102 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,6 +1,7 @@
 #### Changes
 * Python: Added OBJECT ENCODING command ([#1471](https://github.com/aws/glide-for-redis/pull/1471))
 * Python: Added OBJECT FREQ command ([#1472](https://github.com/aws/glide-for-redis/pull/1472))
+* Python: Added GEOSEARCH command ([#1482](https://github.com/aws/glide-for-redis/pull/1482))
 
 ## 0.4.0 (2024-05-26)
 
@@ -44,7 +45,7 @@
 * Python: Added SDIFFSTORE command ([#1449](https://github.com/aws/glide-for-redis/pull/1449))
 * Python: Added SINTERSTORE command ([#1459](https://github.com/aws/glide-for-redis/pull/1459))
 * Python: Added SMISMEMBER command ([#1461](https://github.com/aws/glide-for-redis/pull/1461))
-* Python: Added SETRANGE command ([#1453](https://github.com/aws/glide-for-redis/pull/1453)
+* Python: Added SETRANGE command ([#1453](https://github.com/aws/glide-for-redis/pull/1453))
 
 #### 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/client/value_conversion.rs b/glide-core/src/client/value_conversion.rs
@@ -26,6 +26,7 @@ pub(crate) enum ExpectedReturnType {
     ArrayOfMemberScorePairs,
     ZMPopReturnType,
     KeyWithMemberAndScore,
+    GeoSearchReturnType(bool, bool, bool),
 }
 
 pub(crate) fn convert_to_expected_type(
@@ -379,6 +380,68 @@ pub(crate) fn convert_to_expected_type(
             )
                 .into()),
         },
+        ExpectedReturnType::GeoSearchReturnType(dist, hash, coord) => match value {
+            Value::Array(array) => {
+                if !dist && !hash && !coord {
+                    return Ok(Value::Array(array));
+                }
+
+                let mut converted_array = Vec::with_capacity(array.len());
+                let num_true = dist as usize + hash as usize + coord as usize;
+                for item in array.iter() {
+                    if let Value::Array(inner_array) = item {
+                        if let Some((name, rest)) = inner_array.split_first() {
+                            if rest.len() != num_true {
+                                return Err((
+                                    ErrorKind::TypeError,
+                                    "Response couldn't be converted to GeoSearchReturn type",
+                                    format!("(response was {:?})", array),
+                                )
+                                    .into());
+                            }
+                            let rest = rest
+                                .iter()
+                                .map(|v| match v {
+                                    Value::Array(coord) => {
+                                        if coord.len() != 2 {
+                                            return Err((
+                                                ErrorKind::TypeError,
+                                                "Inner Array must contain exactly two elements, longitude and latitude",
+                                            )
+                                                .into());
+                                        }
+                                        let longitude = convert_to_expected_type(
+                                            coord[0].clone(),
+                                            Some(ExpectedReturnType::Double),
+                                        )?;
+                                        let latitude = convert_to_expected_type(
+                                            coord[1].clone(),
+                                            Some(ExpectedReturnType::Double),
+                                        )?;
+                                        Ok(Value::Array(vec![longitude, latitude]))
+                                    }
+                                    Value::BulkString(dist) => convert_to_expected_type(
+                                        Value::BulkString(dist.to_vec()),
+                                        Some(ExpectedReturnType::Double),
+                                    ),
+                                    _ => Ok(v).cloned(),
+                                })
+                                .collect::<Result<Vec<Value>, _>>()?;
+
+                            converted_array
+                                .push(Value::Array(vec![name.clone(), Value::Array(rest)]));
+                        }
+                    }
+                }
+                Ok(Value::Array(converted_array))
+            }
+            _ => Err((
+                ErrorKind::TypeError,
+                "Response couldn't be converted to GeoSearchReturnType",
+                format!("(response was {:?})", value),
+            )
+                .into()),
+        },
     }
 }
 
@@ -567,6 +630,11 @@ pub(crate) fn expected_type_for_cmd(cmd: &Cmd) -> Option<ExpectedReturnType> {
             }
         }
         b"LOLWUT" => Some(ExpectedReturnType::Lolwut),
+        b"GEOSEARCH" => Some(ExpectedReturnType::GeoSearchReturnType(
+            cmd.position(b"WITHDIST").is_some(),
+            cmd.position(b"WITHHASH").is_some(),
+            cmd.position(b"WITHCOORD").is_some(),
+        )),
         _ => None,
     }
 }

diff --git a/glide-core/src/protobuf/redis_request.proto b/glide-core/src/protobuf/redis_request.proto
@@ -194,6 +194,7 @@ enum RequestType {
     ExpireTime = 156;
     PExpireTime = 157;
     BLMPop = 158;
+    GeoSearch = 159;
 }
 
 message Command {

diff --git a/glide-core/src/request_type.rs b/glide-core/src/request_type.rs
@@ -164,6 +164,7 @@ pub enum RequestType {
     ExpireTime = 156,
     PExpireTime = 157,
     BLMPop = 158,
+    GeoSearch = 159,
 }
 
 fn get_two_word_command(first: &str, second: &str) -> Cmd {
@@ -331,6 +332,7 @@ impl From<::protobuf::EnumOrUnknown<ProtobufRequestType>> for RequestType {
             ProtobufRequestType::HStrlen => RequestType::HStrlen,
             ProtobufRequestType::ExpireTime => RequestType::ExpireTime,
             ProtobufRequestType::PExpireTime => RequestType::PExpireTime,
+            ProtobufRequestType::GeoSearch => RequestType::GeoSearch,
         }
     }
 }
@@ -494,6 +496,7 @@ impl RequestType {
             RequestType::HStrlen => Some(cmd("HSTRLEN")),
             RequestType::ExpireTime => Some(cmd("EXPIRETIME")),
             RequestType::PExpireTime => Some(cmd("PEXPIRETIME")),
+            RequestType::GeoSearch => Some(cmd("GEOSEARCH")),
         }
     }
 }
diff --git a/python/python/glide/__init__.py b/python/python/glide/__init__.py
@@ -5,8 +5,6 @@
     ExpireOptions,
     ExpirySet,
     ExpiryType,
-    GeospatialData,
-    GeoUnit,
     InfoSection,
     InsertPosition,
     StreamAddOptions,
@@ -18,6 +16,11 @@
 from glide.async_commands.redis_modules import json
 from glide.async_commands.sorted_set import (
     AggregationType,
+    GeoSearchByBox,
+    GeoSearchByRadius,
+    GeoSearchCount,
+    GeospatialData,
+    GeoUnit,
     InfBound,
     LexBoundary,
     Limit,
@@ -89,6 +92,9 @@
     "ExpireOptions",
     "ExpirySet",
     "ExpiryType",
+    "GeoSearchByBox",
+    "GeoSearchByRadius",
+    "GeoSearchCount",
     "GeoUnit",
     "GeospatialData",
     "AggregationType",