Support for the JSON.ARRPOP command

src/commands/cmd_json.cc

@@ -22,6 +22,7 @@
 
 #include "commander.h"
 #include "commands/command_parser.h"
+#include "error_constants.h"
 #include "server/redis_reply.h"
 #include "server/server.h"
 #include "storage/redis_metadata.h"
@@ -222,12 +223,52 @@ class CommandJsonArrLen : public Commander {
   }
 };
 
+class CommandJsonArrPop : public Commander {
+ public:
+  Status Parse(const std::vector<std::string> &args) override {
+    path_ = (args_.size() > 2) ? args_[2] : "$";
+
+    if (args_.size() == 4) {
+      index_ = GET_OR_RET(ParseInt<int64_t>(args_[3], 10));
+    } else if (args_.size() > 4) {
+      return {Status::RedisParseErr, errWrongNumOfArguments};
+    }
+
+    return Status::OK();
+  }
+
+  Status Execute(Server *srv, Connection *conn, std::string *output) override {
+    redis::Json json(srv->storage, conn->GetNamespace());
+
+    std::vector<std::optional<JsonValue>> results;
+
+    auto s = json.ArrPop(args_[1], path_, index_, &results);
+    if (!s.ok()) return {Status::RedisExecErr, s.ToString()};
+
+    *output = redis::MultiLen(results.size());
+    for (const auto &data : results) {
+      if (data.has_value()) {
+        *output += redis::BulkString(GET_OR_RET(data->Print()));
+      } else {
+        *output += redis::NilString();
+      }
+    }
+
+    return Status::OK();
+  }
+
+ private:
+  std::string path_;
+  int64_t index_ = -1;
+};
+
 REDIS_REGISTER_COMMANDS(MakeCmdAttr<CommandJsonSet>("json.set", 4, "write", 1, 1, 1),
                         MakeCmdAttr<CommandJsonGet>("json.get", -2, "read-only", 1, 1, 1),
                         MakeCmdAttr<CommandJsonInfo>("json.info", 2, "read-only", 1, 1, 1),
                         MakeCmdAttr<CommandJsonType>("json.type", -2, "read-only", 1, 1, 1),
                         MakeCmdAttr<CommandJsonArrAppend>("json.arrappend", -4, "write", 1, 1, 1),
                         MakeCmdAttr<CommandJsonClear>("json.clear", -2, "write", 1, 1, 1),
-                        MakeCmdAttr<CommandJsonArrLen>("json.arrlen", -2, "read-only", 1, 1, 1), );
+                        MakeCmdAttr<CommandJsonArrLen>("json.arrlen", -2, "read-only", 1, 1, 1),
+                        MakeCmdAttr<CommandJsonArrPop>("json.arrpop", -2, "write", 1, 1, 1), );
 
 }  // namespace redis

src/types/json.h

@@ -253,6 +253,33 @@ struct JsonValue {
     return Status::OK();
   }
 
+  StatusOr<std::vector<std::optional<JsonValue>>> ArrPop(std::string_view path, int64_t index = -1) {
+    std::vector<std::optional<JsonValue>> popped_values;
+
+    try {
+      jsoncons::jsonpath::json_replace(value, path,
+                                       [&popped_values, index](const std::string & /*path*/, jsoncons::json &val) {
+                                         if (val.is_array() && !val.empty()) {
+                                           auto len = static_cast<int64_t>(val.size());
+                                           auto popped_iter = val.array_range().begin();
+                                           if (index < 0) {
+                                             popped_iter += len - std::min(len, -index);
+                                           } else if (index > 0) {
+                                             popped_iter += std::min(len - 1, index);
+                                           }
+                                           popped_values.emplace_back(*popped_iter);
+                                           val.erase(popped_iter);
+                                         } else {
+                                           popped_values.emplace_back(std::nullopt);
+                                         }
+                                       });
+    } catch (const jsoncons::jsonpath::jsonpath_error &e) {
+      return {Status::NotOK, e.what()};
+    }
+
+    return popped_values;
+  }
+
   JsonValue(const JsonValue &) = default;
   JsonValue(JsonValue &&) = default;
 

src/types/redis_json.cc

@@ -230,4 +230,26 @@ rocksdb::Status Json::ArrLen(const std::string &user_key, const std::string &pat
 
   return rocksdb::Status::OK();
 }
+
+rocksdb::Status Json::ArrPop(const std::string &user_key, const std::string &path, int64_t index,
+                             std::vector<std::optional<JsonValue>> *results) {
+  auto ns_key = AppendNamespacePrefix(user_key);
+
+  LockGuard guard(storage_->GetLockManager(), ns_key);
+
+  JsonMetadata metadata;
+  JsonValue json_val;
+  auto s = read(ns_key, &metadata, &json_val);
+  if (!s.ok()) return s;
+
+  auto pop_res = json_val.ArrPop(path, index);
+  if (!pop_res) return rocksdb::Status::InvalidArgument(pop_res.Msg());
+  *results = *pop_res;
+
+  bool is_write = std::any_of(pop_res->begin(), pop_res->end(),
+                              [](const std::optional<JsonValue> &val) { return val.has_value(); });
+  if (!is_write) return rocksdb::Status::OK();
+
+  return write(ns_key, &metadata, json_val);
+}
 }  // namespace redis

src/types/redis_json.h

@@ -42,6 +42,8 @@ class Json : public Database {
   rocksdb::Status Clear(const std::string &user_key, const std::string &path, size_t *result);
   rocksdb::Status ArrLen(const std::string &user_key, const std::string &path,
                          std::vector<std::optional<uint64_t>> &arr_lens);
+  rocksdb::Status ArrPop(const std::string &user_key, const std::string &path, int64_t index,
+                         std::vector<std::optional<JsonValue>> *results);
 
  private:
   rocksdb::Status write(Slice ns_key, JsonMetadata *metadata, const JsonValue &json_val);

tests/cppunit/types/json_test.cc

@@ -272,3 +272,71 @@ TEST_F(RedisJsonTest, ArrLen) {
   ASSERT_TRUE(json_->ArrLen(key_, "$.not_exists", res).ok());
   ASSERT_TRUE(res.empty());
 }
+
+TEST_F(RedisJsonTest, ArrPop) {
+  std::vector<std::optional<JsonValue>> res;
+
+  // Array
+  ASSERT_TRUE(json_->Set(key_, "$", R"([3,"str",2.1,{},[5,6]])").ok());
+  ASSERT_TRUE(json_->ArrPop(key_, "$", -1, &res).ok());
+  ASSERT_EQ(res.size(), 1);
+  ASSERT_TRUE(res[0].has_value());
+  ASSERT_EQ(res[0]->Dump().GetValue(), "[5,6]");
+  res.clear();
+  ASSERT_TRUE(json_->ArrPop(key_, "$", -2, &res).ok());
+  ASSERT_EQ(res.size(), 1);
+  ASSERT_TRUE(res[0].has_value());
+  ASSERT_EQ(res[0]->Dump().GetValue(), "2.1");
+  res.clear();
+  ASSERT_TRUE(json_->ArrPop(key_, "$", 3, &res).ok());
+  ASSERT_EQ(res.size(), 1);
+  ASSERT_TRUE(res[0].has_value());
+  ASSERT_EQ(res[0]->Dump().GetValue(), "{}");
+  res.clear();
+  ASSERT_TRUE(json_->ArrPop(key_, "$", 1, &res).ok());
+  ASSERT_EQ(res.size(), 1);
+  ASSERT_TRUE(res[0].has_value());
+  ASSERT_EQ(res[0]->Dump().GetValue(), R"("str")");
+  res.clear();
+  ASSERT_TRUE(json_->ArrPop(key_, "$", 0, &res).ok());
+  ASSERT_EQ(res.size(), 1);
+  ASSERT_TRUE(res[0].has_value());
+  ASSERT_EQ(res[0]->Dump().GetValue(), "3");
+  res.clear();
+  ASSERT_TRUE(json_->ArrPop(key_, "$", -1, &res).ok());
+  ASSERT_EQ(res.size(), 1);
+  ASSERT_FALSE(res[0].has_value());
+  res.clear();
+
+  // Non-array
+  ASSERT_TRUE(json_->Set(key_, "$", R"({"o":{"x":1},"s":"str","i":1,"d":2.2})").ok());
+  ASSERT_TRUE(json_->ArrPop(key_, "$.o", 1, &res).ok());
+  ASSERT_EQ(res.size(), 1);
+  ASSERT_FALSE(res[0].has_value());
+  res.clear();
+  ASSERT_TRUE(json_->ArrPop(key_, "$.s", -1, &res).ok());
+  ASSERT_EQ(res.size(), 1);
+  ASSERT_FALSE(res[0].has_value());
+  res.clear();
+  ASSERT_TRUE(json_->ArrPop(key_, "$.i", 0, &res).ok());
+  ASSERT_EQ(res.size(), 1);
+  ASSERT_FALSE(res[0].has_value());
+  res.clear();
+  ASSERT_TRUE(json_->ArrPop(key_, "$.d", 2, &res).ok());
+  ASSERT_EQ(res.size(), 1);
+  ASSERT_FALSE(res[0].has_value());
+  res.clear();
+
+  // Multiple arrays
+  ASSERT_TRUE(json_->Set(key_, "$", R"([[0,1],[3,{"x":2.0}],"str",[4,[5,"6"]]])").ok());
+  ASSERT_TRUE(json_->ArrPop(key_, "$.*", -1, &res).ok());
+  ASSERT_EQ(res.size(), 4);
+  ASSERT_TRUE(res[0].has_value());
+  ASSERT_EQ(res[0]->Dump().GetValue(), R"([5,"6"])");
+  ASSERT_FALSE(res[1].has_value());
+  ASSERT_TRUE(res[2].has_value());
+  ASSERT_EQ(res[2]->Dump().GetValue(), R"({"x":2.0})");
+  ASSERT_TRUE(res[3].has_value());
+  ASSERT_EQ(res[3]->Dump().GetValue(), "1");
+  res.clear();
+}

tests/gocase/unit/type/json/json_test.go

@@ -196,4 +196,28 @@ func TestJson(t *testing.T) {
 		require.EqualValues(t, []uint64{}, lens)
 	})
 
+	t.Run("JSON.ARRPOP basics", func(t *testing.T) {
+		require.NoError(t, rdb.Do(ctx, "JSON.SET", "a", "$", `[3,"str",2.1,{},[5,6]]`).Err())
+		require.EqualValues(t, []interface{}{"[5,6]"}, rdb.Do(ctx, "JSON.ARRPOP", "a").Val())
+		require.EqualValues(t, []interface{}{"2.1"}, rdb.Do(ctx, "JSON.ARRPOP", "a", "$", "-2").Val())
+		require.EqualValues(t, []interface{}{"{}"}, rdb.Do(ctx, "JSON.ARRPOP", "a", "$", "3").Val())
+		require.EqualValues(t, []interface{}{`"str"`}, rdb.Do(ctx, "JSON.ARRPOP", "a", "$", "1").Val())
+		require.EqualValues(t, []interface{}{"3"}, rdb.Do(ctx, "JSON.ARRPOP", "a", "$", "0").Val())
+		require.EqualValues(t, []interface{}{nil}, rdb.Do(ctx, "JSON.ARRPOP", "a", "$").Val())
+
+		require.NoError(t, rdb.Do(ctx, "JSON.SET", "o", "$", `{"o":{"x":1},"s":"str","i":1,"d":2.2}`).Err())
+		require.EqualValues(t, []interface{}{nil}, rdb.Do(ctx, "JSON.ARRPOP", "o", "$.o", 1).Val())
+		require.EqualValues(t, []interface{}{nil}, rdb.Do(ctx, "JSON.ARRPOP", "o", "$.s", -1).Val())
+		require.EqualValues(t, []interface{}{nil}, rdb.Do(ctx, "JSON.ARRPOP", "o", "$.i", 0).Val())
+		require.EqualValues(t, []interface{}{nil}, rdb.Do(ctx, "JSON.ARRPOP", "o", "$.d", 2).Val())
+
+		require.NoError(t, rdb.Do(ctx, "JSON.SET", "a", "$", `[[0,1],[3,{"x":2.0}],"str",[4,[5,"6"]]]`).Err())
+		require.EqualValues(t, []interface{}{`[5,"6"]`, nil, `{"x":2.0}`, "1"}, rdb.Do(ctx, "JSON.ARRPOP", "a", "$.*").Val())
+
+		require.ErrorContains(t, rdb.Do(ctx, "JSON.ARRPOP", "a", "$", "str").Err(), "not started as an integer")
+		require.ErrorContains(t, rdb.Do(ctx, "JSON.ARRPOP", "a", "$", "2str").Err(), "encounter non-integer characters")
+		require.ErrorContains(t, rdb.Do(ctx, "JSON.ARRPOP", "a", "$", "-1str").Err(), "encounter non-integer characters")
+		require.ErrorContains(t, rdb.Do(ctx, "JSON.ARRPOP", "a", "$", "0", "1").Err(), "wrong number of arguments")
+	})
+
 }