From 4740424049173cf139b79a79df97878ab5f76f95 Mon Sep 17 00:00:00 2001
From: Qu Chen <quchen@amazon.com>
Date: Fri, 7 Apr 2017 22:31:11 +0000
Subject: [PATCH] Implement getKeys procedure for georadius and
 georadiusbymember commands.

---
 src/db.c     | 38 ++++++++++++++++++++++++++++++++++++++
 src/server.c |  4 ++--
 src/server.h |  1 +
 3 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/src/db.c b/src/db.c
index 6d39bb9b..7d1504d3 100644
--- a/src/db.c
+++ b/src/db.c
@@ -1312,6 +1312,44 @@ int *migrateGetKeys(struct redisCommand *cmd, robj **argv, int argc, int *numkey
     return keys;
 }
 
+/* Helper function to extract keys from following commands:
+ * GEORADIUS key x y radius unit [WITHDIST] [WITHHASH] [WITHCOORD] [ASC|DESC]
+ *                             [COUNT count] [STORE key] [STOREDIST key]
+ * GEORADIUSBYMEMBER key member radius unit ... options ... */
+int *georadiusGetKeys(struct redisCommand *cmd, robj **argv, int argc, int *numkeys) {
+    int i, num, *keys;
+    UNUSED(cmd);
+
+    /* Check for the presence of the stored key in the command */
+    int stored_key = -1;
+    for (i = 5; i < argc; i++) {
+        char *arg = argv[i]->ptr;
+        /* For the case when user specifies both "store" and "storedist" options, the
+         * second key specified would override the first key. This behavior is kept 
+         * the same as in georadiusCommand method.
+         */
+        if ((!strcasecmp(arg, "store") || !strcasecmp(arg, "storedist")) && ((i+1) < argc)) {
+            stored_key = i+1;
+            i++;
+        }
+    }
+    num = 1 + (stored_key == -1 ? 0 : 1);
+
+    /* Keys in the command come from two places:
+     * argv[1] = key,
+     * argv[5...n] = stored key if present
+     */
+    keys = zmalloc(sizeof(int) * num);
+
+    /* Add all key positions to keys[] */
+    keys[0] = 1;
+    if(num > 1) {
+         keys[1] = stored_key;
+    }
+    *numkeys = num; 
+    return keys;
+}
+
 /* Slot to Key API. This is used by Redis Cluster in order to obtain in
  * a fast way a key that belongs to a specified hash slot. This is useful
  * while rehashing the cluster and in other conditions when we need to
diff --git a/src/server.c b/src/server.c
index 9abda682..c08c095c 100644
--- a/src/server.c
+++ b/src/server.c
@@ -290,8 +290,8 @@ struct redisCommand redisCommandTable[] = {
     {"wait",waitCommand,3,"s",0,NULL,0,0,0,0,0},
     {"command",commandCommand,0,"lt",0,NULL,0,0,0,0,0},
     {"geoadd",geoaddCommand,-5,"wm",0,NULL,1,1,1,0,0},
-    {"georadius",georadiusCommand,-6,"w",0,NULL,1,1,1,0,0},
-    {"georadiusbymember",georadiusByMemberCommand,-5,"w",0,NULL,1,1,1,0,0},
+    {"georadius",georadiusCommand,-6,"w",0,georadiusGetKeys,1,1,1,0,0},
+    {"georadiusbymember",georadiusByMemberCommand,-5,"w",0,georadiusGetKeys,1,1,1,0,0},
     {"geohash",geohashCommand,-2,"r",0,NULL,1,1,1,0,0},
     {"geopos",geoposCommand,-2,"r",0,NULL,1,1,1,0,0},
     {"geodist",geodistCommand,-4,"r",0,NULL,1,1,1,0,0},
diff --git a/src/server.h b/src/server.h
index 8403ed5b..18924090 100644
--- a/src/server.h
+++ b/src/server.h
@@ -1730,6 +1730,7 @@ int *zunionInterGetKeys(struct redisCommand *cmd,robj **argv, int argc, int *num
 int *evalGetKeys(struct redisCommand *cmd, robj **argv, int argc, int *numkeys);
 int *sortGetKeys(struct redisCommand *cmd, robj **argv, int argc, int *numkeys);
 int *migrateGetKeys(struct redisCommand *cmd, robj **argv, int argc, int *numkeys);
+int *georadiusGetKeys(struct redisCommand *cmd, robj **argv, int argc, int *numkeys);
 
 /* Cluster */
 void clusterInit(void);