From 9658eab036e7d51ef6db579c0cbcb0235d4ba135 Mon Sep 17 00:00:00 2001
From: Roman Gershman <roman@dragonflydb.io>
Date: Sun, 21 May 2023 20:51:45 +0300
Subject: [PATCH] fix: support XREAD ... STREAMS ... keys derivation (#1250)

---
 src/server/stream_family.cc      |  5 +++--
 src/server/stream_family_test.cc |  2 +-
 src/server/transaction.cc        | 19 +++++++++++++++++++
 3 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/src/server/stream_family.cc b/src/server/stream_family.cc
index 3738fa6a1..4d89da08e 100644
--- a/src/server/stream_family.cc
+++ b/src/server/stream_family.cc
@@ -618,7 +618,7 @@ vector<RecordVec> OpRead(const OpArgs& op_args, const ArgSlice& args, const Read
 
   vector<RecordVec> response(args.size());
   for (size_t i = 0; i < args.size(); ++i) {
-    const string_view key = args[i];
+    string_view key = args[i];
 
     range_opts.start = opts.stream_ids.at(key);
 
@@ -1446,7 +1446,8 @@ void StreamFamily::Register(CommandRegistry* registry) {
             // Therefore the command has format:
             //   XREAD COUNT <count> STREAMS <stream1> <stream2> <id1> <id2>
             // Where the keys are <stream1> and <stream2>.
-            << CI{"XREAD", CO::READONLY | CO::REVERSE_MAPPING, -4, 4, 5, 1}.HFUNC(XRead)
+            << CI{"XREAD", CO::READONLY | CO::REVERSE_MAPPING | CO::VARIADIC_KEYS, -3, 3, 3, 1}
+                   .HFUNC(XRead)
             << CI{"XSETID", CO::WRITE | CO::DENYOOM, 3, 1, 1, 1}.HFUNC(XSetId)
             << CI{"_XGROUP_HELP", CO::NOSCRIPT | CO::HIDDEN, 2, 0, 0, 0}.SetHandler(XGroupHelp);
 }
diff --git a/src/server/stream_family_test.cc b/src/server/stream_family_test.cc
index 795fba6fd..67409892f 100644
--- a/src/server/stream_family_test.cc
+++ b/src/server/stream_family_test.cc
@@ -170,7 +170,7 @@ TEST_F(StreamFamilyTest, XReadInvalidArgs) {
 
   // Missing STREAMS.
   resp = Run({"xread", "count", "5"});
-  EXPECT_THAT(resp, ErrArg("wrong number of arguments for 'xread' command"));
+  EXPECT_THAT(resp, ErrArg("syntax error"));
 
   // Unbalanced list of streams.
   resp = Run({"xread", "count", "invalid", "streams", "s1", "s2", "s3", "0", "0"});
diff --git a/src/server/transaction.cc b/src/server/transaction.cc
index e60abd491..c790f71a4 100644
--- a/src/server/transaction.cc
+++ b/src/server/transaction.cc
@@ -1389,12 +1389,31 @@ OpResult<KeyIndex> DetermineKeys(const CommandId* cid, CmdArgList args) {
   if (cid->opt_mask() & CO::VARIADIC_KEYS) {
     // ZUNION/INTER <num_keys> <key1> [<key2> ...]
     // EVAL <script> <num_keys>
+    // XREAD ... STREAMS ...
     if (args.size() < 2) {
       return OpStatus::SYNTAX_ERR;
     }
 
     string_view name{cid->name()};
 
+    if (name == "XREAD") {
+      for (size_t i = 0; i < args.size(); ++i) {
+        string_view arg = ArgS(args, i);
+        if (absl::EqualsIgnoreCase(arg, "STREAMS")) {
+          size_t left = args.size() - i - 1;
+          if (left < 2 || left % 2 != 0)
+            return OpStatus::SYNTAX_ERR;
+
+          key_index.start = i + 1;
+          key_index.end = key_index.start + (left / 2);
+          key_index.step = 1;
+
+          return key_index;
+        }
+      }
+      return OpStatus::SYNTAX_ERR;
+    }
+
     if (absl::EndsWith(name, "STORE"))
       key_index.bonus = 0;  // Z<xxx>STORE <key> commands