feat(server): implement JSON.GET command (#104)

Signed-off-by: iko1 <me@remotecpp.dev>

src/server/CMakeLists.txt

@@ -8,13 +8,13 @@ cxx_link(dfly_transaction uring_fiber_lib dfly_core strings_lib)
 
 add_library(dragonfly_lib  channel_slice.cc command_registry.cc
             config_flags.cc conn_context.cc debugcmd.cc dflycmd.cc
-            generic_family.cc hset_family.cc
+            generic_family.cc hset_family.cc json_family.cc 
             list_family.cc main_service.cc  rdb_load.cc rdb_save.cc replica.cc
             snapshot.cc script_mgr.cc server_family.cc
             set_family.cc stream_family.cc string_family.cc
             zset_family.cc version.cc)
 
-cxx_link(dragonfly_lib dfly_transaction dfly_facade redis_lib strings_lib html_lib)
+cxx_link(dragonfly_lib dfly_transaction dfly_facade redis_lib strings_lib html_lib TRDP::jsoncons)
 
 add_library(dfly_test_lib test_utils.cc)
 cxx_link(dfly_test_lib dragonfly_lib facade_test gtest_main_ext)
@@ -31,9 +31,10 @@ cxx_test(rdb_test dfly_test_lib DATA testdata/empty.rdb testdata/redis6_small.rd
 cxx_test(zset_family_test dfly_test_lib LABELS DFLY)
 cxx_test(blocking_controller_test dragonfly_lib LABELS DFLY)
 cxx_test(snapshot_test dragonfly_lib LABELS DFLY)
+cxx_test(json_family_test dfly_test_lib LABELS DFLY)
 
 
 add_custom_target(check_dfly WORKING_DIRECTORY .. COMMAND ctest -L DFLY)
-add_dependencies(check_dfly dragonfly_test list_family_test
+add_dependencies(check_dfly dragonfly_test json_family_test list_family_test
                  generic_family_test memcache_parser_test rdb_test
-                 redis_parser_test stream_family_test string_family_test)
+                 redis_parser_test snapshot_test stream_family_test string_family_test)

src/server/json_family.cc

@@ -0,0 +1,140 @@
+// Copyright 2022, Roman Gershman.  All rights reserved.
+// See LICENSE for licensing terms.
+//
+
+#include "server/json_family.h"
+
+extern "C" {
+#include "redis/object.h"
+}
+
+#include "base/logging.h"
+#include "server/command_registry.h"
+#include "server/error.h"
+#include "server/tiered_storage.h"
+#include "server/transaction.h"
+
+#include <jsoncons/json.hpp>
+#include <jsoncons_ext/jsonpath/jsonpath.hpp>
+#include <absl/strings/str_join.h>
+
+namespace dfly {
+
+using namespace std;
+using namespace jsoncons;
+
+using JsonExpression = jsonpath::jsonpath_expression<json>;
+using CI = CommandId;
+
+namespace {
+
+string GetString(EngineShard* shard, const PrimeValue& pv) {
+  string res;
+  if (pv.IsExternal()) {
+    auto* tiered = shard->tiered_storage();
+    auto [offset, size] = pv.GetExternalPtr();
+    res.resize(size);
+
+    error_code ec = tiered->Read(offset, size, res.data());
+    CHECK(!ec) << "TBD: " << ec;
+  } else {
+    pv.GetString(&res);
+  }
+
+  return res;
+}
+
+OpResult<string> OpGet(const OpArgs& op_args, string_view key, vector<pair<string_view, JsonExpression>> expressions) {
+  OpResult<PrimeIterator> it_res = op_args.shard->db_slice().Find(op_args.db_ind, key, OBJ_STRING);
+  if (!it_res.ok())
+    return it_res.status();
+
+  const PrimeValue& pv = it_res.value()->second;
+  string val = GetString(op_args.shard, pv);
+
+  auto callback = [](json_errc ec, const ser_context&) {
+    VLOG(1) << "Error while decode JSON: " << make_error_code(ec).message();
+    return false;
+  };
+
+  error_code ec;
+  json_decoder<json> decoder;
+  basic_json_parser<char> parser(basic_json_decode_options<char>{}, callback);
+
+  parser.update(val);
+  parser.finish_parse(decoder, ec);
+
+  if (!decoder.is_valid()) {
+    return OpStatus::SYNTAX_ERR;
+  }
+
+  if (expressions.size() == 1) {
+    json out = expressions[0].second.evaluate(decoder.get_result());
+    return out.as<string>();
+  }
+
+  json out;
+  json result = decoder.get_result();
+  for (auto& expr: expressions) {
+    json eval = expr.second.evaluate(result);
+    out[expr.first] = eval;
+  }
+
+  return out.as<string>();
+}
+
+} // namespace
+
+void JsonFamily::Get(CmdArgList args, ConnectionContext* cntx) {
+  DCHECK_GE(args.size(), 3U);
+  string_view key = ArgS(args, 1);
+
+  vector<pair<string_view, JsonExpression>> expressions;
+  for (size_t i = 2; i < args.size(); ++i) {
+    string_view path = ArgS(args, i);
+
+    error_code ec;
+    JsonExpression expr = jsonpath::make_expression<json>(path, ec);
+
+    if (ec) {
+      VLOG(1) << "Invalid JSONPath: " << ec.message();
+      (*cntx)->SendError(kWrongTypeErr);
+      return;
+    }
+
+    expressions.emplace_back(path, move(expr));
+  }
+
+  auto cb = [&](Transaction* t, EngineShard* shard) {
+    return OpGet(t->GetOpArgs(shard), key, move(expressions));
+  };
+
+  DVLOG(1) << "Before Get::ScheduleSingleHopT " << key;
+  Transaction* trans = cntx->transaction;
+  OpResult<string> result = trans->ScheduleSingleHopT(move(cb));
+
+  if (result) {
+    DVLOG(1) << "JSON.GET " << trans->DebugId() << ": " << key << " " << result.value();
+    (*cntx)->SendBulkString(*result);
+  } else {
+    switch (result.status()) {
+      case OpStatus::WRONG_TYPE:
+        (*cntx)->SendError(kWrongTypeErr);
+        break;
+      case OpStatus::SYNTAX_ERR:
+        (*cntx)->SendError(kSyntaxErr);
+        break;
+      default:
+        DVLOG(1) << "JSON.GET " << key << " nil";
+        (*cntx)->SendNull();
+    }
+  }
+}
+
+#define HFUNC(x) SetHandler(&JsonFamily::x)
+
+void JsonFamily::Register(CommandRegistry* registry) {
+  *registry << CI{"JSON.GET", CO::READONLY | CO::FAST, -3, 1, 1, 1}.HFUNC(Get);
+}
+
+}  // namespace dfly

src/server/json_family.h

@@ -0,0 +1,25 @@
+// Copyright 2022, Roman Gershman.  All rights reserved.
+// See LICENSE for licensing terms.
+//
+
+#pragma once
+
+#include "server/common.h"
+#include "server/engine_shard_set.h"
+
+namespace dfly {
+
+class ConnectionContext;
+class CommandRegistry;
+using facade::OpResult;
+using facade::OpStatus;
+
+class JsonFamily {
+ public:
+  static void Register(CommandRegistry* registry);
+
+ private:
+  static void Get(CmdArgList args, ConnectionContext* cntx);
+};
+
+}  // namespace dfly

src/server/json_family_test.cc

@@ -0,0 +1,125 @@
+// Copyright 2022, Roman Gershman.  All rights reserved.
+// See LICENSE for licensing terms.
+//
+
+#include "server/json_family.h"
+
+#include "base/gtest.h"
+#include "base/logging.h"
+#include "facade/facade_test.h"
+#include "server/command_registry.h"
+#include "server/test_utils.h"
+
+using namespace testing;
+using namespace std;
+using namespace util;
+
+namespace dfly {
+
+class JsonFamilyTest : public BaseFamilyTest {
+ protected:
+};
+
+TEST_F(JsonFamilyTest, SetGetBasic) {
+  string json = R"(
+    {
+       "store": {
+        "book": [
+         {
+           "category": "Fantasy",
+           "author": "J. K. Rowling",
+           "title": "Harry Potter and the Philosopher's Stone",
+           "isbn": 9780747532743,
+           "price": 5.99
+         }
+       ]
+      }
+    }
+)";
+
+  string xml = R"(
+    <?xml version="1.0" encoding="UTF-8" ?>
+    <store>
+      <book>
+        <category>Fantasy</category>
+        <author>J. K. Rowling</author>
+        <title>Harry Potter and the Philosopher&#x27;s Stone</title>
+        <isbn>9780747532743</isbn>
+        <price>5.99</price>
+      </book>
+    </store>
+)";
+
+  auto resp = Run({"set", "json", json});
+  ASSERT_THAT(resp, "OK");
+
+  resp = Run({"JSON.GET", "json", "$..*"});
+  ASSERT_THAT(resp, ArgType(RespExpr::STRING));
+
+  resp = Run({"JSON.GET", "json", "$..book[0].price"});
+  EXPECT_THAT(resp, ArgType(RespExpr::STRING));
+
+  resp = Run({"JSON.GET", "json", "//*"});
+  EXPECT_THAT(resp, ArgType(RespExpr::ERROR));
+
+  resp = Run({"JSON.GET", "json", "//book[0]"});
+  EXPECT_THAT(resp, ArgType(RespExpr::ERROR));
+
+  resp = Run({"set", "xml", xml});
+  ASSERT_THAT(resp, "OK");
+
+  resp = Run({"JSON.GET", "xml", "$..*"});
+  EXPECT_THAT(resp, ArgType(RespExpr::ERROR));
+}
+
+TEST_F(JsonFamilyTest, SetGetFromPhonebook) {
+    string json = R"(
+    {
+      "firstName":"John",
+      "lastName":"Smith",
+      "age":27,
+      "weight":135.25,
+      "isAlive":true,
+      "address":{
+          "street":"21 2nd Street",
+          "city":"New York",
+          "state":"NY",
+          "zipcode":"10021-3100"
+      },
+      "phoneNumbers":[
+          {
+            "type":"home",
+            "number":"212 555-1234"
+          },
+          {
+            "type":"office",
+            "number":"646 555-4567"
+          }
+      ],
+      "children":[
+
+      ],
+      "spouse":null
+    }
+  )";
+
+  auto resp = Run({"set", "json", json});
+  ASSERT_THAT(resp, "OK");
+
+  resp = Run({"JSON.GET", "json", "$.address.*"});
+  EXPECT_EQ(resp, string_view(R"(["New York","NY","21 2nd Street","10021-3100"])", 46));
+
+  resp = Run({"JSON.GET", "json", "$.firstName", "$.age", "$.lastName"});
+  EXPECT_EQ(resp, string_view(R"({"$.age":[27],"$.firstName":["John"],"$.lastName":["Smith"]})", 60));
+
+  resp = Run({"JSON.GET", "json", "$.spouse.*"});
+  EXPECT_EQ(resp, string_view("[]", 2));
+
+  resp = Run({"JSON.GET", "json", "$.children.*"});
+  EXPECT_EQ(resp, string_view("[]", 2));
+
+  resp = Run({"JSON.GET", "json", "$..phoneNumbers[1].*"});
+  EXPECT_EQ(resp, string_view(R"(["646 555-4567","office"])", 25));
+}
+
+}  // namespace dfly

src/server/main_service.cc

@@ -24,6 +24,7 @@ extern "C" {
 #include "server/error.h"
 #include "server/generic_family.h"
 #include "server/hset_family.h"
+#include "server/json_family.h"
 #include "server/list_family.h"
 #include "server/script_mgr.h"
 #include "server/server_state.h"
@@ -1161,6 +1162,7 @@ void Service::RegisterCommands() {
   SetFamily::Register(&registry_);
   HSetFamily::Register(&registry_);
   ZSetFamily::Register(&registry_);
+  JsonFamily::Register(&registry_);
 
   server_family_.Register(&registry_);