From c9c33b476bffdac092269648b7f3fe9d7babe6c3 Mon Sep 17 00:00:00 2001
From: iko1 <me@remotecpp.dev>
Date: Sun, 13 Nov 2022 09:23:02 +0200
Subject: [PATCH] feat(server): implement json.resp command (#104) (#482)

Signed-off-by: iko1 <me@remotecpp.dev>
---
 src/server/json_family.cc      | 73 ++++++++++++++++++++++++++++++++++
 src/server/json_family.h       |  1 +
 src/server/json_family_test.cc | 24 +++++++++--
 3 files changed, 95 insertions(+), 3 deletions(-)

diff --git a/src/server/json_family.cc b/src/server/json_family.cc
index 84a29a959..ff30ef40f 100644
--- a/src/server/json_family.cc
+++ b/src/server/json_family.cc
@@ -314,6 +314,34 @@ size_t CountJsonFields(const json& j) {
   return res;
 }
 
+void SendJsonValue(ConnectionContext* cntx, const json& j) {
+  if (j.is_double()) {
+    (*cntx)->SendDouble(j.as_double());
+  } else if (j.is_number()) {
+    (*cntx)->SendLong(j.as_integer<long>());
+  } else if (j.is_bool()) {
+    (*cntx)->SendSimpleString(j.as_bool() ? "true" : "false");
+  } else if (j.is_null()) {
+    (*cntx)->SendNull();
+  } else if (j.is_string()) {
+    (*cntx)->SendSimpleString(j.as_string_view());
+  } else if (j.is_object()) {
+    (*cntx)->StartArray(j.size() + 1);
+    (*cntx)->SendSimpleString("{");
+    for (const auto& item : j.object_range()) {
+      (*cntx)->StartArray(2);
+      (*cntx)->SendSimpleString(item.key());
+      SendJsonValue(cntx, item.value());
+    }
+  } else if (j.is_array()) {
+    (*cntx)->StartArray(j.size() + 1);
+    (*cntx)->SendSimpleString("[");
+    for (const auto& item : j.array_range()) {
+      SendJsonValue(cntx, item);
+    }
+  }
+}
+
 OpResult<string> OpGet(const OpArgs& op_args, string_view key,
                        vector<pair<string_view, JsonExpression>> expressions) {
   OpResult<json> result = GetJson(op_args, key);
@@ -950,8 +978,52 @@ OpResult<vector<OptSizeT>> OpFields(const OpArgs& op_args, string_view key,
   return vec;
 }
 
+// Returns json vector that represents the result of the json query.
+OpResult<vector<json>> OpResp(const OpArgs& op_args, string_view key, JsonExpression expression) {
+  OpResult<json> result = GetJson(op_args, key);
+  if (!result) {
+    return result.status();
+  }
+
+  vector<json> vec;
+  auto cb = [&vec](const string_view& path, const json& val) { vec.emplace_back(val); };
+
+  expression.evaluate(*result, cb);
+  return vec;
+}
+
 }  // namespace
 
+void JsonFamily::Resp(CmdArgList args, ConnectionContext* cntx) {
+  string_view key = ArgS(args, 1);
+  string_view path = ArgS(args, 2);
+
+  error_code ec;
+  JsonExpression expression = jsonpath::make_expression<json>(path, ec);
+
+  if (ec) {
+    VLOG(1) << "Invalid JSONPath syntax: " << ec.message();
+    (*cntx)->SendError(kSyntaxErr);
+    return;
+  }
+
+  auto cb = [&](Transaction* t, EngineShard* shard) {
+    return OpResp(t->GetOpArgs(shard), key, move(expression));
+  };
+
+  Transaction* trans = cntx->transaction;
+  OpResult<vector<json>> result = trans->ScheduleSingleHopT(move(cb));
+
+  if (result) {
+    (*cntx)->StartArray(result->size());
+    for (const auto& it : *result) {
+      SendJsonValue(cntx, it);
+    }
+  } else {
+    (*cntx)->SendError(result.status());
+  }
+}
+
 void JsonFamily::Debug(CmdArgList args, ConnectionContext* cntx) {
   function<decltype(OpFields)> func;
   string_view command = ArgS(args, 1);
@@ -1578,6 +1650,7 @@ void JsonFamily::Register(CommandRegistry* registry) {
       ArrAppend);
   *registry << CI{"JSON.ARRINDEX", CO::READONLY | CO::FAST, -4, 1, 1, 1}.HFUNC(ArrIndex);
   *registry << CI{"JSON.DEBUG", CO::READONLY | CO::FAST, -2, 1, 1, 1}.HFUNC(Debug);
+  *registry << CI{"JSON.RESP", CO::READONLY | CO::FAST, 3, 1, 1, 1}.HFUNC(Resp);
 }
 
 }  // namespace dfly
diff --git a/src/server/json_family.h b/src/server/json_family.h
index 99ad8288b..9a1a9f5c2 100644
--- a/src/server/json_family.h
+++ b/src/server/json_family.h
@@ -39,6 +39,7 @@ class JsonFamily {
   static void ArrAppend(CmdArgList args, ConnectionContext* cntx);
   static void ArrIndex(CmdArgList args, ConnectionContext* cntx);
   static void Debug(CmdArgList args, ConnectionContext* cntx);
+  static void Resp(CmdArgList args, ConnectionContext* cntx);
 };
 
 }  // namespace dfly
diff --git a/src/server/json_family_test.cc b/src/server/json_family_test.cc
index 26c8afbb5..04276e6f9 100644
--- a/src/server/json_family_test.cc
+++ b/src/server/json_family_test.cc
@@ -72,8 +72,7 @@ TEST_F(JsonFamilyTest, SetGetBasic) {
   EXPECT_THAT(resp, ArgType(RespExpr::ERROR));
 }
 
-TEST_F(JsonFamilyTest, SetGetFromPhonebook) {
-  string json = R"(
+static const string PhonebookJson = R"(
     {
       "firstName":"John",
       "lastName":"Smith",
@@ -103,7 +102,8 @@ TEST_F(JsonFamilyTest, SetGetFromPhonebook) {
     }
   )";
 
-  auto resp = Run({"set", "json", json});
+TEST_F(JsonFamilyTest, SetGetFromPhonebook) {
+  auto resp = Run({"set", "json", PhonebookJson});
   ASSERT_THAT(resp, "OK");
 
   resp = Run({"JSON.GET", "json", "$.address.*"});
@@ -892,4 +892,22 @@ TEST_F(JsonFamilyTest, DebugFields) {
   EXPECT_THAT(resp, IntArg(16));
 }
 
+TEST_F(JsonFamilyTest, Resp) {
+  auto resp = Run({"set", "json", PhonebookJson});
+  ASSERT_THAT(resp, "OK");
+
+  resp = Run({"JSON.RESP", "json", "$.address.*"});
+  ASSERT_EQ(RespExpr::ARRAY, resp.type);
+  EXPECT_THAT(resp.GetVec(), ElementsAre("New York", "NY", "21 2nd Street", "10021-3100"));
+
+  resp = Run({"JSON.RESP", "json", "$.isAlive"});
+  EXPECT_THAT(resp, "true");
+
+  resp = Run({"JSON.RESP", "json", "$.age"});
+  EXPECT_THAT(resp, IntArg(27));
+
+  resp = Run({"JSON.RESP", "json", "$.weight"});
+  EXPECT_THAT(resp, "135.25");
+}
+
 }  // namespace dfly