fix(json_family): fix JSON.GET commmand for JSON legacy mode (#3261)

* fix(json_family): fix JSON.GET commmand for JSON legacy mode

fixes dragonflydb#1523

Signed-off-by: Stepan Bagritsevich <bagr.stepan@gmail.com>

* refactor(json_family): address comments

Signed-off-by: Stepan Bagritsevich <bagr.stepan@gmail.com>

---------

Signed-off-by: Stepan Bagritsevich <bagr.stepan@gmail.com>

src/server/detail/wrapped_json_path.h

@@ -0,0 +1,128 @@
+// Copyright 2022, DragonflyDB authors.  All rights reserved.
+// See LICENSE for licensing terms.
+//
+
+#pragma once
+
+#include <string_view>
+#include <utility>
+#include <variant>
+
+#include "core/json/json_object.h"
+#include "core/json/path.h"
+#include "core/string_or_view.h"
+
+namespace dfly {
+
+using JsonExpression = jsoncons::jsonpath::jsonpath_expression<JsonType>;
+
+template <typename T>
+using JsonPathEvaluateCallback = absl::FunctionRef<T(std::string_view, const JsonType&)>;
+
+template <typename T> class JsonCallbackResult {
+ public:
+  using JsonV1Result = std::optional<T>;
+  using JsonV2Result = std::vector<T>;
+
+  explicit JsonCallbackResult(bool legacy_mode_is_enabled)
+      : legacy_mode_is_enabled_(legacy_mode_is_enabled) {
+    if (!legacy_mode_is_enabled_) {
+      result_ = JsonV2Result{};
+    }
+  }
+
+  void AddValue(T&& value) {
+    if (IsV1()) {
+      AsV1().emplace(std::forward<T>(value));
+    } else {
+      AsV2().emplace_back(std::forward<T>(value));
+    }
+  }
+
+  bool IsV1() const {
+    return legacy_mode_is_enabled_;
+  }
+
+  JsonV1Result& AsV1() {
+    return std::get<JsonV1Result>(result_);
+  }
+
+  JsonV2Result& AsV2() {
+    return std::get<JsonV2Result>(result_);
+  }
+
+ private:
+  std::variant<JsonV1Result, JsonV2Result> result_;
+  bool legacy_mode_is_enabled_;
+};
+
+class WrappedJsonPath {
+ public:
+  static constexpr std::string_view kV1PathRootElement = ".";
+  static constexpr std::string_view kV2PathRootElement = "$";
+
+  WrappedJsonPath(json::Path json_path, StringOrView path, bool is_legacy_mode_path)
+      : parsed_path_(std::move(json_path)),
+        path_(std::move(path)),
+        is_legacy_mode_path_(is_legacy_mode_path) {
+  }
+
+  WrappedJsonPath(JsonExpression expression, StringOrView path, bool is_legacy_mode_path)
+      : parsed_path_(std::move(expression)),
+        path_(std::move(path)),
+        is_legacy_mode_path_(is_legacy_mode_path) {
+  }
+
+  template <typename T>
+  JsonCallbackResult<T> Evaluate(const JsonType* json_entry, JsonPathEvaluateCallback<T> cb) const {
+    return Evaluate(json_entry, cb, IsLegacyModePath());
+  }
+
+  template <typename T>
+  JsonCallbackResult<T> Evaluate(const JsonType* json_entry, JsonPathEvaluateCallback<T> cb,
+                                 bool legacy_mode_is_enabled) const {
+    JsonCallbackResult<T> eval_result{legacy_mode_is_enabled};
+
+    auto eval_callback = [&cb, &eval_result](std::string_view path, const JsonType& val) {
+      eval_result.AddValue(cb(path, val));
+    };
+
+    if (HoldsJsonPath()) {
+      const auto& json_path = AsJsonPath();
+      json::EvaluatePath(
+          json_path, *json_entry,
+          [&eval_callback](std::optional<std::string_view> key, const JsonType& val) {
+            eval_callback(key ? *key : std::string_view{}, val);
+          });
+    } else {
+      const auto& json_expression = AsJsonExpression();
+      json_expression.evaluate(*json_entry, eval_callback);
+    }
+
+    return eval_result;
+  }
+
+  bool IsLegacyModePath() const {
+    return is_legacy_mode_path_;
+  }
+
+ private:
+  bool HoldsJsonPath() const {
+    return std::holds_alternative<json::Path>(parsed_path_);
+  }
+
+  const json::Path& AsJsonPath() const {
+    return std::get<json::Path>(parsed_path_);
+  }
+
+  const JsonExpression& AsJsonExpression() const {
+    return std::get<JsonExpression>(parsed_path_);
+  }
+
+ private:
+  std::variant<json::Path, JsonExpression> parsed_path_;
+  StringOrView path_;
+  bool is_legacy_mode_path_;
+};
+
+}  // namespace dfly

src/server/error.h

@@ -44,6 +44,31 @@ using facade::kWrongTypeErr;
 
 #endif  // RETURN_ON_BAD_STATUS
 
+#ifndef RETURN_UNEXPECTED
+
+#define RETURN_UNEXPECTED(x)       \
+  do {                             \
+    if (!(x)) {                    \
+      return (x).get_unexpected(); \
+    }                              \
+  } while (0)
+
+#endif  // RETURN_UNEXPECTED
+
+#ifndef GET_OR_SEND_UNEXPECTED
+
+#define GET_OR_SEND_UNEXPECTED(expr)     \
+  ({                                     \
+    auto expr_res = (expr);              \
+    if (!expr_res) {                     \
+      cntx->SendError(expr_res.error()); \
+      return;                            \
+    }                                    \
+    std::move(expr_res).value();         \
+  })
+
+#endif  // GET_OR_SEND_UNEXPECTED
+
 namespace rdb {
 
 enum errc {

src/server/json_family.cc

@@ -24,6 +24,7 @@
 #include "server/acl/acl_commands_def.h"
 #include "server/command_registry.h"
 #include "server/common.h"
+#include "server/detail/wrapped_json_path.h"
 #include "server/error.h"
 #include "server/journal/journal.h"
 #include "server/search/doc_index.h"
@@ -55,6 +56,56 @@ static const char DefaultJsonPath[] = "$";
 
 namespace {
 
+namespace json_parser {
+
+template <typename T> using ParseResult = io::Result<T, std::string>;
+
+ParseResult<JsonExpression> ParseJsonPathAsExpression(std::string_view path) {
+  std::error_code ec;
+  JsonExpression res = MakeJsonPathExpr(path, ec);
+  if (ec)
+    return nonstd::make_unexpected(kSyntaxErr);
+  return res;
+}
+
+ParseResult<WrappedJsonPath> ParseJsonPath(StringOrView path, bool is_legacy_mode_path) {
+  if (absl::GetFlag(FLAGS_jsonpathv2)) {
+    auto path_result = json::ParsePath(path.view());
+    RETURN_UNEXPECTED(path_result);
+    return WrappedJsonPath{std::move(path_result).value(), std::move(path), is_legacy_mode_path};
+  }
+
+  auto expr_result = ParseJsonPathAsExpression(path.view());
+  RETURN_UNEXPECTED(expr_result);
+  return WrappedJsonPath{std::move(expr_result).value(), std::move(path), is_legacy_mode_path};
+}
+
+ParseResult<WrappedJsonPath> ParseJsonPathV1(std::string_view path) {
+  if (path == WrappedJsonPath::kV1PathRootElement) {
+    return ParseJsonPath(StringOrView::FromView(WrappedJsonPath::kV2PathRootElement), true);
+  }
+
+  std::string v2_path = absl::StrCat(
+      WrappedJsonPath::kV2PathRootElement, path.front() != '.' && path.front() != '[' ? "." : "",
+      path);  // Convert to V2 path; TODO(path.front() != all kinds of symbols)
+  return ParseJsonPath(StringOrView::FromString(std::move(v2_path)), true);
+}
+
+ParseResult<WrappedJsonPath> ParseJsonPathV2(std::string_view path) {
+  return ParseJsonPath(StringOrView::FromView(path), false);
+}
+
+bool IsJsonPathV2(std::string_view path) {
+  return path.front() == '$';
+}
+
+ParseResult<WrappedJsonPath> ParseJsonPath(std::string_view path) {
+  DCHECK(!path.empty());
+  return IsJsonPathV2(path) ? ParseJsonPathV2(path) : ParseJsonPathV1(path);
+}
+
+}  // namespace json_parser
+
 using JsonPathV2 = variant<json::Path, JsonExpression>;
 using ExprCallback = absl::FunctionRef<void(string_view, const JsonType&)>;
 
@@ -454,34 +505,36 @@ void SendJsonValue(RedisReplyBuilder* rb, const JsonType& j) {
   }
 }
 
-OpResult<string> OpJsonGet(const OpArgs& op_args, string_view key,
+bool LegacyModeIsEnabled(const std::vector<std::pair<std::string_view, WrappedJsonPath>>& paths) {
-                           const vector<pair<string_view, optional<JsonPathV2>>>& expressions,
+  return std::all_of(paths.begin(), paths.end(),
-                           bool should_format, const OptString& indent, const OptString& new_line,
+                     [](auto& parsed_path) { return parsed_path.second.IsLegacyModePath(); });
-                           const OptString& space) {
-  OpResult<JsonType*> result = GetJson(op_args, key);
-  if (!result) {
-    return result.status();
 }
 
+OpResult<std::string> OpJsonGet(const OpArgs& op_args, string_view key,
+                                const vector<pair<string_view, WrappedJsonPath>>& paths,
+                                const std::optional<std::string>& indent,
+                                const std::optional<std::string>& new_line,
+                                const std::optional<std::string>& space) {
+  OpResult<JsonType*> result = GetJson(op_args, key);
+  RETURN_ON_BAD_STATUS(result);
+
   const JsonType& json_entry = *(result.value());
-  if (expressions.empty()) {
+  if (paths.empty()) {
     // this implicitly means that we're using $ which
     // means we just brings all values
     return json_entry.to_string();
   }
 
   json_options options;
-  if (should_format) {
   options.spaces_around_comma(spaces_option::no_spaces)
       .spaces_around_colon(spaces_option::no_spaces)
       .object_array_line_splits(line_split_kind::multi_line)
-        .indent_size(1)
+      .indent_size(0)
       .new_line_chars("");
 
   if (indent) {
+    options.indent_size(1);
     options.indent_chars(*indent);
-    } else {
-      options.indent_size(0);
   }
 
   if (new_line) {
@@ -491,31 +544,42 @@ OpResult<string> OpJsonGet(const OpArgs& op_args, string_view key,
   if (space) {
     options.after_key_chars(*space);
   }
+
+  const bool legacy_mode_is_enabled = LegacyModeIsEnabled(paths);
+
+  auto cb = [](std::string_view, const JsonType& val) { return val; };
+
+  auto eval_wrapped = [&json_entry, &cb, legacy_mode_is_enabled](
+                          const WrappedJsonPath& json_path) -> std::optional<JsonType> {
+    auto eval_result = json_path.Evaluate<JsonType>(&json_entry, cb, legacy_mode_is_enabled);
+
+    DCHECK(legacy_mode_is_enabled == eval_result.IsV1());
+
+    if (eval_result.IsV1()) {
+      return eval_result.AsV1();
     }
 
-  auto eval_wrapped = [&json_entry](const optional<JsonPathV2>& expr) -> JsonType {
+    return JsonType{eval_result.AsV2()};
-    return expr ? visit([&](auto& arg) { return Evaluate(arg, json_entry); }, *expr) : json_entry;
   };
 
-  JsonType out{json_object_arg};  // see https://github.com/danielaparker/jsoncons/issues/482
+  JsonType out{
-  if (expressions.size() == 1) {
+      jsoncons::json_object_arg};  // see https://github.com/danielaparker/jsoncons/issues/482
-    out = eval_wrapped(expressions[0].second);
+  if (paths.size() == 1) {
+    auto eval_result = eval_wrapped(paths[0].second);
+    out = std::move(eval_result).value();  // TODO(Print not existing path to the user)
   } else {
-    for (const auto& [expr_str, expr] : expressions) {
+    for (const auto& [path_str, path] : paths) {
-      out[expr_str] = eval_wrapped(expr);
+      auto eval_result = eval_wrapped(path);
+      out[path_str] = std::move(eval_result).value();  // TODO(Print not existing path to the user)
     }
   }
 
-  if (should_format) {
+  jsoncons::json_printable jp(out, options, jsoncons::indenting::indent);
-    json_printable jp(out, options, indenting::indent);
   std::stringstream ss;
   jp.dump(ss);
   return ss.str();
 }
 
-  return out.as<string>();
-}
-
 OpResult<vector<string>> OpType(const OpArgs& op_args, string_view key, JsonPathV2 expression) {
   OpResult<JsonType*> result = GetJson(op_args, key);
   if (!result) {
@@ -2067,8 +2131,7 @@ void JsonFamily::Get(CmdArgList args, ConnectionContext* cntx) {
   OptString new_line;
   OptString space;
 
-  // '.' corresponds to the legacy, non-array format and is passed as nullopt.
+  vector<pair<string_view, WrappedJsonPath>> paths;
-  vector<pair<string_view, optional<JsonPathV2>>> expressions;
 
   while (parser.HasNext()) {
     if (parser.Check("SPACE").IgnoreCase().ExpectTail(1)) {
@@ -2084,23 +2147,17 @@ void JsonFamily::Get(CmdArgList args, ConnectionContext* cntx) {
       continue;
     }
 
-    optional<JsonPathV2> expr;
+    string_view path_str = parser.Next();
-    string_view expr_str = parser.Next();
+    WrappedJsonPath json_path = GET_OR_SEND_UNEXPECTED(json_parser::ParseJsonPath(path_str));
 
-    if (expr_str != ".") {
+    paths.emplace_back(path_str, std::move(json_path));
-      JsonPathV2 expression = PARSE_PATHV2(expr_str);
-      expr.emplace(std::move(expression));
-    }
-
-    expressions.emplace_back(expr_str, std::move(expr));
   }
 
   if (auto err = parser.Error(); err)
     return cntx->SendError(err->MakeReply());
 
-  bool should_format = (indent || new_line || space);
   auto cb = [&](Transaction* t, EngineShard* shard) {
-    return OpJsonGet(t->GetOpArgs(shard), key, expressions, should_format, indent, new_line, space);
+    return OpJsonGet(t->GetOpArgs(shard), key, paths, indent, new_line, space);
   };
 
   Transaction* trans = cntx->transaction;

src/server/json_family_test.cc

@@ -70,10 +70,10 @@ TEST_F(JsonFamilyTest, SetGetBasic) {
   EXPECT_THAT(resp, ArgType(RespExpr::ERROR));
 
   resp = Run({"JSON.GET", "json", "store.book[0].category"});
-  EXPECT_EQ(resp, "[\"Fantasy\"]");
+  EXPECT_EQ(resp, "\"Fantasy\"");
 
   resp = Run({"JSON.GET", "json", ".store.book[0].category"});
-  EXPECT_EQ(resp, "[\"Fantasy\"]");
+  EXPECT_EQ(resp, "\"Fantasy\"");
 
   resp = Run({"SET", "xml", xml});
   ASSERT_THAT(resp, "OK");
@@ -82,6 +82,34 @@ TEST_F(JsonFamilyTest, SetGetBasic) {
   EXPECT_THAT(resp, ArgType(RespExpr::ERROR));
 }
 
+TEST_F(JsonFamilyTest, GetLegacy) {
+  string json = R"({"name":"Leonard Cohen","lastSeen":1478476800,"loggedOut": true})";
+
+  auto resp = Run({"JSON.SET", "json", "$", json});
+  ASSERT_THAT(resp, "OK");
+
+  resp = Run({"JSON.GET", "json", "."});  // V1 Response
+  ASSERT_THAT(resp, "{\"lastSeen\":1478476800,\"loggedOut\":true,\"name\":\"Leonard Cohen\"}");
+
+  resp = Run({"JSON.GET", "json", "$"});  // V2 Response
+  ASSERT_THAT(resp, "[{\"lastSeen\":1478476800,\"loggedOut\":true,\"name\":\"Leonard Cohen\"}]");
+
+  resp = Run({"JSON.GET", "json", ".name"});  // V1 Response
+  ASSERT_THAT(resp, "\"Leonard Cohen\"");
+
+  resp = Run({"JSON.GET", "json", "$.name"});  // V2 Response
+  ASSERT_THAT(resp, "[\"Leonard Cohen\"]");
+
+  resp = Run({"JSON.GET", "json", ".name", "$.lastSeen"});  // V2 Response
+  ASSERT_THAT(resp, "{\"$.lastSeen\":[1478476800],\".name\":[\"Leonard Cohen\"]}");
+
+  resp = Run({"JSON.GET", "json", ".name", ".lastSeen"});  // V1 Response
+  ASSERT_THAT(resp, "{\".lastSeen\":1478476800,\".name\":\"Leonard Cohen\"}");
+
+  resp = Run({"JSON.GET", "json", "$.name", "$.lastSeen"});  // V2 Response
+  ASSERT_THAT(resp, "{\"$.lastSeen\":[1478476800],\"$.name\":[\"Leonard Cohen\"]}");
+}
+
 static const string PhonebookJson = R"(
     {
       "firstName":"John",