feat(server): Implement SMISMEMBER command (#377)

* feat(server): Implement SMISMEMBER command (#361)

Signed-off-by: ATM SALEH <saleh.cse08@gmail.com>
Signed-off-by: Roman Gershman <roman@dragonflydb.io>
Co-authored-by: Roman Gershman <roman@dragonflydb.io>
This commit is contained in:
SALEH 2022-10-13 03:21:27 +01:00 committed by GitHub
parent 4833f93366
commit 45a5f30cdd
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 73 additions and 2 deletions

View file

@ -8,4 +8,5 @@
* **[Ryan Russell](https://github.com/ryanrussell)**
* Docs & Code Readability
* **[Ali-Akber Saifee](https://github.com/alisaifee)**
* **[Elle Y](https://github.com/inohime)**
* **[Elle Y](https://github.com/inohime)**
* **[ATM SALEH](https://github.com/ATM-SALEH)**

View file

@ -32,7 +32,6 @@ using absl::GetFlag;
using ResultStringVec = vector<OpResult<StringVec>>;
using ResultSetView = OpResult<absl::flat_hash_set<std::string_view>>;
using SvArray = vector<std::string_view>;
namespace {
constexpr uint32_t kMaxIntSetEntries = 256;
@ -283,6 +282,15 @@ bool IsInSet(const DbContext& db_context, const SetType& st, string_view member)
}
}
void FindInSet(StringVec& memberships,
const DbContext& db_context, const SetType& st,
const vector<string_view>& members) {
for (const auto& member : members) {
bool status = IsInSet(db_context, st, member);
memberships.emplace_back(to_string(status));
}
}
// Removes arg from result.
void DiffStrSet(const DbContext& db_context, const SetType& st,
absl::flat_hash_set<string>* result) {
@ -1020,6 +1028,37 @@ void SIsMember(CmdArgList args, ConnectionContext* cntx) {
}
}
void SMIsMember(CmdArgList args, ConnectionContext* cntx) {
string_view key = ArgS(args, 1);
vector<string_view> vals(args.size() - 2);
for (size_t i = 2; i < args.size(); ++i) {
vals[i - 2] = ArgS(args, i);
}
StringVec memberships;
memberships.reserve(vals.size());
auto cb = [&](Transaction* t, EngineShard* shard) {
OpResult<PrimeIterator> find_res = shard->db_slice().Find(t->db_context(), key, OBJ_SET);
if (find_res) {
SetType st{find_res.value()->second.RObjPtr(), find_res.value()->second.Encoding()};
FindInSet(memberships, t->db_context(), st, vals);
return OpStatus::OK;
}
return find_res.status();
};
OpResult<void> result = cntx->transaction->ScheduleSingleHop(std::move(cb));
if (result == OpStatus::KEY_NOTFOUND) {
memberships.assign(vals.size(), "0");
return (*cntx)->SendStringArr(memberships);
} else if (result == OpStatus::OK) {
return (*cntx)->SendStringArr(memberships);
}
(*cntx)->SendError(result.status());
}
void SMove(CmdArgList args, ConnectionContext* cntx) {
string_view src = ArgS(args, 1);
string_view dest = ArgS(args, 2);
@ -1464,6 +1503,7 @@ void SetFamily::Register(CommandRegistry* registry) {
<< CI{"SINTERSTORE", CO::WRITE | CO::DENYOOM, -3, 1, -1, 1}.HFUNC(SInterStore)
<< CI{"SMEMBERS", CO::READONLY, 2, 1, 1, 1}.HFUNC(SMembers)
<< CI{"SISMEMBER", CO::FAST | CO::READONLY, 3, 1, 1, 1}.HFUNC(SIsMember)
<< CI{"SMISMEMBER", CO::READONLY, -3, 1, 1, 1}.HFUNC(SMIsMember)
<< CI{"SMOVE", CO::FAST | CO::WRITE, 4, 1, 2, 1}.HFUNC(SMove)
<< CI{"SREM", CO::WRITE | CO::FAST | CO::DENYOOM, -3, 1, 1, 1}.HFUNC(SRem)
<< CI{"SCARD", CO::READONLY | CO::FAST, 2, 1, 1, 1}.HFUNC(SCard)

View file

@ -132,6 +132,36 @@ TEST_F(SetFamilyTest, SPop) {
EXPECT_THAT(resp.GetVec(), IsSubsetOf({"a", "b", "c"}));
}
TEST_F(SetFamilyTest, SMIsMember) {
Run({"sadd", "foo", "a"});
Run({"sadd", "foo", "b"});
auto resp = Run({"smismember", "foo"});
EXPECT_THAT(resp, ErrArg("wrong number of arguments"));
resp = Run({"smismember", "foo1", "a", "b"});
ASSERT_THAT(resp, ArgType(RespExpr::ARRAY));
EXPECT_THAT(resp.GetVec(), ElementsAre("0", "0"));
resp = Run({"smismember", "foo", "a", "c"});
ASSERT_THAT(resp, ArgType(RespExpr::ARRAY));
EXPECT_THAT(resp.GetVec(), ElementsAre("1", "0"));
resp = Run({"smismember", "foo", "a", "b"});
ASSERT_THAT(resp, ArgType(RespExpr::ARRAY));
EXPECT_THAT(resp.GetVec(), ElementsAre("1", "1"));
resp = Run({"smismember", "foo", "d", "e"});
ASSERT_THAT(resp, ArgType(RespExpr::ARRAY));
EXPECT_THAT(resp.GetVec(), ElementsAre("0", "0"));
resp = Run({"smismember", "foo", "b"});
EXPECT_THAT(resp, "1");
resp = Run({"smismember", "foo", "x"});
EXPECT_THAT(resp, "0");
}
TEST_F(SetFamilyTest, Empty) {
auto resp = Run({"smembers", "x"});
ASSERT_THAT(resp, ArrLen(0));