From 2661fe16b4c408a40a7011fd4e2e5ebb75e7a28a Mon Sep 17 00:00:00 2001 From: Borys Date: Wed, 30 Apr 2025 16:25:21 +0300 Subject: [PATCH] feat: add copy cmd (#5032) * add copy command * add tests --- src/server/generic_family.cc | 31 ++++++++++-- src/server/generic_family.h | 1 + src/server/generic_family_test.cc | 82 +++++++++++++++++++++++++++++++ 3 files changed, 110 insertions(+), 4 deletions(-) diff --git a/src/server/generic_family.cc b/src/server/generic_family.cc index fb7e9677b..866833f3a 100644 --- a/src/server/generic_family.cc +++ b/src/server/generic_family.cc @@ -300,12 +300,14 @@ OpStatus OpPersist(const OpArgs& op_args, string_view key); class Renamer { public: - Renamer(Transaction* t, std::string_view src_key, std::string_view dest_key, unsigned shard_count) + Renamer(Transaction* t, std::string_view src_key, std::string_view dest_key, unsigned shard_count, + bool do_copy = false) : transaction_(t), src_key_(src_key), dest_key_(dest_key), src_sid_(Shard(src_key, shard_count)), - dest_sid_(Shard(dest_key, shard_count)) { + dest_sid_(Shard(dest_key, shard_count)), + do_copy_(do_copy) { } ErrorReply Rename(bool destination_should_not_exist); @@ -337,6 +339,7 @@ class Renamer { bool src_found_ = false; bool dest_found_ = false; + bool do_copy_ = false; SerializedValue serialized_value_; }; @@ -366,7 +369,7 @@ ErrorReply Renamer::Rename(bool destination_should_not_exist) { void Renamer::FetchData() { auto cb = [this](Transaction* t, EngineShard* shard) { auto args = t->GetShardArgs(shard->shard_id()); - DCHECK_EQ(1u, args.Size()); + DCHECK(1 == args.Size() || do_copy_); const ShardId shard_id = shard->shard_id(); @@ -388,7 +391,7 @@ void Renamer::FinalizeRename() { auto cb = [this](Transaction* t, EngineShard* shard) { const ShardId shard_id = shard->shard_id(); - if (shard_id == src_sid_) { + if (!do_copy_ && shard_id == src_sid_) { return DelSrc(t, shard); } @@ -1654,6 +1657,24 @@ void GenericFamily::RenameNx(CmdArgList args, const CommandContext& cmd_cntx) { } } +void GenericFamily::Copy(CmdArgList args, const CommandContext& cmd_cntx) { + CmdArgParser parser(args); + auto [k1, k2] = parser.Next(); + bool replace = parser.Check("REPLACE"); + + if (!parser.Finalize()) { + return cmd_cntx.rb->SendError(parser.Error()->MakeReply()); + } + + if (k1 == k2) { + cmd_cntx.rb->SendError("source and destination objects are the same"); + return; + } + + Renamer renamer(cmd_cntx.tx, k1, k2, shard_set->size(), true); + cmd_cntx.rb->SendError(renamer.Rename(!replace)); +} + void GenericFamily::ExpireTime(CmdArgList args, const CommandContext& cmd_cntx) { ExpireTimeGeneric(args, TimeUnit::SEC, cmd_cntx.tx, cmd_cntx.rb); } @@ -1869,6 +1890,7 @@ constexpr uint32_t kKeys = KEYSPACE | READ | SLOW | DANGEROUS; constexpr uint32_t kPExpireAt = KEYSPACE | WRITE | FAST; constexpr uint32_t kPExpire = KEYSPACE | WRITE | FAST; constexpr uint32_t kRename = KEYSPACE | WRITE | SLOW; +constexpr uint32_t kCopy = KEYSPACE | WRITE | SLOW; constexpr uint32_t kRenamNX = KEYSPACE | WRITE | FAST; constexpr uint32_t kSelect = FAST | CONNECTION; constexpr uint32_t kScan = KEYSPACE | READ | SLOW; @@ -1914,6 +1936,7 @@ void GenericFamily::Register(CommandRegistry* registry) { << CI{"FIELDEXPIRE", CO::WRITE | CO::FAST | CO::DENYOOM, -4, 1, 1, acl::kFieldExpire}.HFUNC( FieldExpire) << CI{"RENAME", CO::WRITE | CO::NO_AUTOJOURNAL, 3, 1, 2, acl::kRename}.HFUNC(Rename) + << CI{"COPY", CO::WRITE | CO::NO_AUTOJOURNAL, -3, 1, 2, acl::kCopy}.HFUNC(Copy) << CI{"RENAMENX", CO::WRITE | CO::NO_AUTOJOURNAL, 3, 1, 2, acl::kRenamNX}.HFUNC(RenameNx) << CI{"SELECT", kSelectOpts, 2, 0, 0, acl::kSelect}.HFUNC(Select) << CI{"SCAN", CO::READONLY | CO::FAST | CO::LOADING, -2, 0, 0, acl::kScan}.HFUNC(Scan) diff --git a/src/server/generic_family.h b/src/server/generic_family.h index 5876ad7ff..2812de074 100644 --- a/src/server/generic_family.h +++ b/src/server/generic_family.h @@ -43,6 +43,7 @@ class GenericFamily { static void Rename(CmdArgList args, const CommandContext& cmd_cntx); static void RenameNx(CmdArgList args, const CommandContext& cmd_cntx); + static void Copy(CmdArgList args, const CommandContext& cmd_cntx); static void ExpireTime(CmdArgList args, const CommandContext& cmd_cntx); static void PExpireTime(CmdArgList args, const CommandContext& cmd_cntx); static void Ttl(CmdArgList args, const CommandContext& cmd_cntx); diff --git a/src/server/generic_family_test.cc b/src/server/generic_family_test.cc index ed6587795..af5ea8f5c 100644 --- a/src/server/generic_family_test.cc +++ b/src/server/generic_family_test.cc @@ -939,4 +939,86 @@ TEST_F(GenericFamilyTest, Unlink) { EXPECT_THAT(resp, IntArg(2)); } +TEST_F(GenericFamilyTest, Copy) { + RespExpr resp; + string b_val(32, 'b'); + string x_val(32, 'x'); + + resp = Run({"mset", "x", x_val, "b", b_val}); + ASSERT_EQ(resp, "OK"); + ASSERT_EQ(2, last_cmd_dbg_info_.shards_count); + + resp = Run({"COPY", "z", "b"}); + ASSERT_THAT(resp, ErrArg("no such key")); + + resp = Run({"COPY", "b", "c"}); + ASSERT_EQ(resp, "OK"); + ASSERT_EQ(b_val, Run({"get", "c"})); + + resp = Run({"COPY", "x", "b", "REPLACE"}); + ASSERT_EQ(resp, "OK"); + + ASSERT_EQ(x_val, Run({"get", "x"})); + ASSERT_EQ(x_val, Run({"get", "b"})); + EXPECT_EQ(CheckedInt({"exists", "x", "b"}), 2); + + const char* keys[2] = {"b", "x"}; + auto ren_fb = pp_->at(0)->LaunchFiber([&] { + for (size_t i = 0; i < 200; ++i) { + int j = i % 2; + auto resp = Run({"COPY", keys[j], keys[1 - j], "REPLACE"}); + ASSERT_EQ(resp, "OK"); + } + }); + + auto exist_fb = pp_->at(2)->LaunchFiber([&] { + for (size_t i = 0; i < 300; ++i) { + int64_t resp = CheckedInt({"exists", "x", "b"}); + ASSERT_EQ(2, resp); + } + }); + + exist_fb.Join(); + ren_fb.Join(); +} + +TEST_F(GenericFamilyTest, CopyNonString) { + EXPECT_EQ(1, CheckedInt({"lpush", "x", "elem"})); + auto resp = Run({"COPY", "x", "b"}); + ASSERT_EQ(resp, "OK"); + ASSERT_EQ(2, last_cmd_dbg_info_.shards_count); + + EXPECT_EQ(1, CheckedInt({"del", "x"})); + EXPECT_EQ(1, CheckedInt({"del", "b"})); +} + +TEST_F(GenericFamilyTest, CopyBinary) { + const char kKey1[] = "\x01\x02\x03\x04"; + const char kKey2[] = "\x05\x06\x07\x08"; + + Run({"set", kKey1, "bar"}); + Run({"COPY", kKey1, kKey2}); + EXPECT_EQ(Run({"get", kKey1}), "bar"); + EXPECT_EQ(Run({"get", kKey2}), "bar"); +} + +TEST_F(GenericFamilyTest, CopyTTL) { + Run({"setex", "k1", "10", "bar"}); + + ASSERT_EQ(Run({"COPY", "k1", "k2"}), "OK"); + EXPECT_THAT(Run({"ttl", "k2"}), 10); +} + +TEST_F(GenericFamilyTest, CopySameName) { + ASSERT_THAT(Run({"COPY", "k1", "k1"}), ErrArg("source and destination objects are the same")); + + ASSERT_EQ(Run({"set", "k1", "v"}), "OK"); + ASSERT_THAT(Run({"COPY", "k1", "k1"}), ErrArg("source and destination objects are the same")); +} + +TEST_F(GenericFamilyTest, CopyToDB) { + // we don't support DB arg for now + ASSERT_THAT(Run({"COPY", "k1", "k1", "DB", "SOME_DB"}), ErrArg("syntax error")); +} + } // namespace dfly