feat: add copy cmd (#5032)

* add copy command
* add tests
This commit is contained in:
Borys 2025-04-30 16:25:21 +03:00 committed by GitHub
parent 10cd22375e
commit 2661fe16b4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 110 additions and 4 deletions

View file

@ -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<std::string_view, std::string_view>();
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)

View file

@ -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);

View file

@ -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