From bb75703b17e05223397b693cebe9e7bd9491dbdc Mon Sep 17 00:00:00 2001
From: kukey <wei.kukey@gmail.com>
Date: Sun, 3 Jan 2021 23:13:37 +0800
Subject: [PATCH] GEOADD - add [CH] [NX|XX] options (#8227)

New command flags similar to what SADD already has.

Co-authored-by: huangwei03 <huangwei03@kuaishou.com>
Co-authored-by: Itamar Haber <itamar@redislabs.com>
Co-authored-by: Oran Agra <oran@redislabs.com>
---
 src/geo.c          | 42 ++++++++++++++++++++++++++++--------------
 tests/unit/geo.tcl | 43 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+), 14 deletions(-)

diff --git a/src/geo.c b/src/geo.c
index 417bf897e..e9c8bf5cb 100644
--- a/src/geo.c
+++ b/src/geo.c
@@ -428,31 +428,45 @@ static int sort_gp_desc(const void *a, const void *b) {
  * Commands
  * ==================================================================== */
 
-/* GEOADD key long lat name [long2 lat2 name2 ... longN latN nameN] */
+/* GEOADD key [CH] [NX|XX] long lat name [long2 lat2 name2 ... longN latN nameN] */
 void geoaddCommand(client *c) {
-    /* Check arguments number for sanity. */
-    if ((c->argc - 2) % 3 != 0) {
+    int xx = 0, nx = 0, longidx = 2;
+    int i;
+
+    /* Parse options. At the end 'longidx' is set to the argument position
+     * of the longitude of the first element. */
+    while (longidx < c->argc) {
+        char *opt = c->argv[longidx]->ptr;
+        if (!strcasecmp(opt,"nx")) nx = 1;
+        else if (!strcasecmp(opt,"xx")) xx = 1;
+        else if (!strcasecmp(opt,"ch")) {}
+        else break;
+        longidx++;
+    }
+
+    if ((c->argc - longidx) % 3 || (xx && nx)) {
         /* Need an odd number of arguments if we got this far... */
-        addReplyError(c, "syntax error. Try GEOADD key [x1] [y1] [name1] "
-                         "[x2] [y2] [name2] ... ");
+            addReplyErrorObject(c,shared.syntaxerr);
         return;
     }
 
-    int elements = (c->argc - 2) / 3;
-    int argc = 2+elements*2; /* ZADD key score ele ... */
+    /* Set up the vector for calling ZADD. */
+    int elements = (c->argc - longidx) / 3;
+    int argc = longidx+elements*2; /* ZADD key [CH] [NX|XX] score ele ... */
     robj **argv = zcalloc(argc*sizeof(robj*));
     argv[0] = createRawStringObject("zadd",4);
-    argv[1] = c->argv[1]; /* key */
-    incrRefCount(argv[1]);
+    for (i = 1; i < longidx; i++) {
+        argv[i] = c->argv[i];
+        incrRefCount(argv[i]);
+    }
 
     /* Create the argument vector to call ZADD in order to add all
      * the score,value pairs to the requested zset, where score is actually
      * an encoded version of lat,long. */
-    int i;
     for (i = 0; i < elements; i++) {
         double xy[2];
 
-        if (extractLongLatOrReply(c, (c->argv+2)+(i*3),xy) == C_ERR) {
+        if (extractLongLatOrReply(c, (c->argv+longidx)+(i*3),xy) == C_ERR) {
             for (i = 0; i < argc; i++)
                 if (argv[i]) decrRefCount(argv[i]);
             zfree(argv);
@@ -464,9 +478,9 @@ void geoaddCommand(client *c) {
         geohashEncodeWGS84(xy[0], xy[1], GEO_STEP_MAX, &hash);
         GeoHashFix52Bits bits = geohashAlign52Bits(hash);
         robj *score = createObject(OBJ_STRING, sdsfromlonglong(bits));
-        robj *val = c->argv[2 + i * 3 + 2];
-        argv[2+i*2] = score;
-        argv[3+i*2] = val;
+        robj *val = c->argv[longidx + i * 3 + 2];
+        argv[longidx+i*2] = score;
+        argv[longidx+1+i*2] = val;
         incrRefCount(val);
     }
 
diff --git a/tests/unit/geo.tcl b/tests/unit/geo.tcl
index 119d2bd64..6c1f2ef11 100644
--- a/tests/unit/geo.tcl
+++ b/tests/unit/geo.tcl
@@ -74,6 +74,49 @@ start_server {tags {"geo"}} {
         r geoadd nyc -73.9454966 40.747533 "lic market"
     } {0}
 
+    test {GEOADD update with CH option} {
+        assert_equal 1 [r geoadd nyc CH 40.747533 -73.9454966 "lic market"]
+        lassign [lindex [r geopos nyc "lic market"] 0] x1 y1
+        assert {abs($x1) - 40.747 < 0.001}
+        assert {abs($y1) - 73.945 < 0.001}
+    } {}
+
+    test {GEOADD update with NX option} {
+        assert_equal 0 [r geoadd nyc NX -73.9454966 40.747533 "lic market"]
+        lassign [lindex [r geopos nyc "lic market"] 0] x1 y1
+        assert {abs($x1) - 40.747 < 0.001}
+        assert {abs($y1) - 73.945 < 0.001}
+    } {}
+
+    test {GEOADD update with XX option} {
+        assert_equal 0 [r geoadd nyc XX -83.9454966 40.747533 "lic market"]
+        lassign [lindex [r geopos nyc "lic market"] 0] x1 y1
+        assert {abs($x1) - 83.945 < 0.001}
+        assert {abs($y1) - 40.747 < 0.001}
+    } {}
+
+    test {GEOADD update with CH NX option} {
+        r geoadd nyc CH NX -73.9454966 40.747533 "lic market"
+    } {0}
+
+    test {GEOADD update with CH XX option} {
+        r geoadd nyc CH XX -73.9454966 40.747533 "lic market"
+    } {1}
+
+    test {GEOADD update with XX NX option will return syntax error} {
+        catch {
+            r geoadd nyc xx nx -73.9454966 40.747533 "lic market"
+        } err
+        set err
+    } {ERR*syntax*}
+
+    test {GEOADD update with invalid option} {
+        catch {
+            r geoadd nyc ch xx foo -73.9454966 40.747533 "lic market"
+        } err
+        set err
+    } {ERR*syntax*}
+
     test {GEOADD invalid coordinates} {
         catch {
             r geoadd nyc -73.9454966 40.747533 "lic market" \