Skip to content

Commit

Permalink
Support TYPE flag in the SCAN command (#2255)
Browse files Browse the repository at this point in the history
  • Loading branch information
@jackyyyyyssss
jackyyyyyssss committed Apr 21, 2024
1 parent 3dc16f3 commit eb94fbb
Show file tree
Hide file tree
Showing 5 changed files with 81 additions and 4 deletions.
2 changes: 1 addition & 1 deletion src/commands/cmd_server.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -840,7 +840,7 @@ class CommandScan : public CommandScanBase {

std::vector<std::string> keys;
std::string end_key;
auto s = redis_db.Scan(key_name, limit_, prefix_, &keys, &end_key);
auto s = redis_db.Scan(key_name, limit_, prefix_, &keys, &end_key, type_);
if (!s.ok()) {
return {Status::RedisExecErr, s.ToString()};
}
Expand Down
9 changes: 8 additions & 1 deletion src/commands/scan_base.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,13 @@ class CommandScanBase : public Commander {
return {Status::RedisParseErr, "limit should be a positive integer"};
}
} else if (IsScan && parser.EatEqICase("type")) {
return {Status::RedisParseErr, "TYPE flag is currently not supported"};
std::string type_str = GET_OR_RET(parser.TakeStr());
if (auto iter = std::find(RedisTypeNames.begin(), RedisTypeNames.end(), type_str);
iter != RedisTypeNames.end()) {
type_ = static_cast<RedisType>(iter - RedisTypeNames.begin());
} else {
return {Status::RedisExecErr, "Invalid type"};
}
} else {
return parser.InvalidSyntax();
}
Expand Down Expand Up @@ -93,6 +99,7 @@ class CommandScanBase : public Commander {
std::string cursor_;
std::string prefix_;
int limit_ = 20;
RedisType type_ = kRedisNone;
};

class CommandSubkeyScanBase : public CommandScanBase {
Expand Down
4 changes: 3 additions & 1 deletion src/storage/redis_db.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -310,7 +310,7 @@ rocksdb::Status Database::Keys(const std::string &prefix, std::vector<std::strin
}

rocksdb::Status Database::Scan(const std::string &cursor, uint64_t limit, const std::string &prefix,
std::vector<std::string> *keys, std::string *end_cursor) {
std::vector<std::string> *keys, std::string *end_cursor, RedisType type) {
end_cursor->clear();
uint64_t cnt = 0;
uint16_t slot_start = 0;
Expand Down Expand Up @@ -355,6 +355,8 @@ rocksdb::Status Database::Scan(const std::string &cursor, uint64_t limit, const
auto s = metadata.Decode(iter->value());
if (!s.ok()) continue;

if (type != kRedisNone && type != metadata.Type()) continue;

if (metadata.Expired()) continue;
std::tie(std::ignore, user_key) = ExtractNamespaceKey<std::string>(iter->key(), storage_->IsSlotIdEncoded());
keys->emplace_back(user_key);
Expand Down
3 changes: 2 additions & 1 deletion src/storage/redis_db.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -93,7 +93,8 @@ class Database {
[[nodiscard]] rocksdb::Status Keys(const std::string &prefix, std::vector<std::string> *keys = nullptr,
KeyNumStats *stats = nullptr);
[[nodiscard]] rocksdb::Status Scan(const std::string &cursor, uint64_t limit, const std::string &prefix,
std::vector<std::string> *keys, std::string *end_cursor = nullptr);
std::vector<std::string> *keys, std::string *end_cursor = nullptr,
RedisType type = kRedisNone);
[[nodiscard]] rocksdb::Status RandomKey(const std::string &cursor, std::string *key);
std::string AppendNamespacePrefix(const Slice &user_key);
[[nodiscard]] rocksdb::Status FindKeyRangeWithPrefix(const std::string &prefix, const std::string &prefix_end,
Expand Down
67 changes: 67 additions & 0 deletions tests/gocase/unit/scan/scan_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -301,6 +301,73 @@ func ScanTest(t *testing.T, rdb *redis.Client, ctx context.Context) {
util.ErrorRegexp(t, rdb.Do(ctx, "SCAN", "0", "count", "1", "match", "a*", "hello").Err(), ".*syntax error.*")
util.ErrorRegexp(t, rdb.Do(ctx, "SCAN", "0", "count", "1", "match", "a*", "hello", "hi").Err(), ".*syntax error.*")
})

t.Run("SCAN with type args ", func(t *testing.T) {
//string type
require.NoError(t, rdb.Set(ctx, "stringtype1", "fee1", 0).Err())
require.NoError(t, rdb.Set(ctx, "stringtype2", "fee1", 0).Err())
require.NoError(t, rdb.Set(ctx, "stringtype3", "fee1", 0).Err())
require.Equal(t, []string{"stringtype1", "stringtype2", "stringtype3"}, scanAll(t, rdb, "match", "stringtype*", "type", "string"))
require.Equal(t, []string{"stringtype1", "stringtype2", "stringtype3"}, scanAll(t, rdb, "match", "stringtype*", "count", "3", "type", "string"))
//hash type
require.NoError(t, rdb.HSet(ctx, "hashtype1", "key1", "val1", "key2", "val2").Err())
require.NoError(t, rdb.HSet(ctx, "hashtype2", "key1", "val1", "key2", "val2").Err())
require.NoError(t, rdb.HSet(ctx, "hashtype3", "key1", "val1", "key2", "val2").Err())
require.Equal(t, []string{"hashtype1", "hashtype2", "hashtype3"}, scanAll(t, rdb, "match", "hashtype*", "type", "hash"))
require.Equal(t, []string{"hashtype1", "hashtype2", "hashtype3"}, scanAll(t, rdb, "match", "hashtype*", "count", "3", "type", "hash"))
//list type
require.NoError(t, rdb.RPush(ctx, "listtype1", "1").Err())
require.NoError(t, rdb.RPush(ctx, "listtype2", "2").Err())
require.NoError(t, rdb.RPush(ctx, "listtype3", "3").Err())
require.Equal(t, []string{"listtype1", "listtype2", "listtype3"}, scanAll(t, rdb, "match", "listtype*", "type", "list"))
require.Equal(t, []string{"listtype1", "listtype2", "listtype3"}, scanAll(t, rdb, "match", "listtype*", "count", "3", "type", "list"))
//set type
require.NoError(t, rdb.SAdd(ctx, "settype1", "1").Err())
require.NoError(t, rdb.SAdd(ctx, "settype2", "1").Err())
require.NoError(t, rdb.SAdd(ctx, "settype3", "1").Err())
require.Equal(t, []string{"settype1", "settype2", "settype3"}, scanAll(t, rdb, "match", "settype*", "type", "set"))
require.Equal(t, []string{"settype1", "settype2", "settype3"}, scanAll(t, rdb, "match", "settype*", "count", "3", "type", "set"))
//zet type
members := []redis.Z{
{Score: 1, Member: "1"},
{Score: 2, Member: "2"},
{Score: 3, Member: "3"},
{Score: 10, Member: "4"},
}
require.NoError(t, rdb.ZAdd(ctx, "zsettype1", members...).Err())
require.NoError(t, rdb.ZAdd(ctx, "zsettype2", members...).Err())
require.NoError(t, rdb.ZAdd(ctx, "zsettype3", members...).Err())
require.Equal(t, []string{"zsettype1", "zsettype2", "zsettype3"}, scanAll(t, rdb, "match", "zsettype*", "type", "zset"))
require.Equal(t, []string{"zsettype1", "zsettype2", "zsettype3"}, scanAll(t, rdb, "match", "zsettype*", "count", "3", "type", "zset"))
//bitmap type
require.NoError(t, rdb.SetBit(ctx, "bitmaptype1", 0, 0).Err())
require.NoError(t, rdb.SetBit(ctx, "bitmaptype2", 0, 0).Err())
require.NoError(t, rdb.SetBit(ctx, "bitmaptype3", 0, 0).Err())
require.Equal(t, []string{"bitmaptype1", "bitmaptype2", "bitmaptype3"}, scanAll(t, rdb, "match", "bitmaptype*", "type", "bitmap"))
require.Equal(t, []string{"bitmaptype1", "bitmaptype2", "bitmaptype3"}, scanAll(t, rdb, "match", "bitmaptype*", "count", "3", "type", "bitmap"))
//stream type
require.NoError(t, rdb.XAdd(ctx, &redis.XAddArgs{Stream: "streamtype1", Values: []string{"item", "1", "value", "a"}}).Err())
require.NoError(t, rdb.XAdd(ctx, &redis.XAddArgs{Stream: "streamtype2", Values: []string{"item", "1", "value", "a"}}).Err())
require.NoError(t, rdb.XAdd(ctx, &redis.XAddArgs{Stream: "streamtype3", Values: []string{"item", "1", "value", "a"}}).Err())
require.Equal(t, []string{"streamtype1", "streamtype2", "streamtype3"}, scanAll(t, rdb, "match", "streamtype*", "type", "stream"))
require.Equal(t, []string{"streamtype1", "streamtype2", "streamtype3"}, scanAll(t, rdb, "match", "streamtype*", "count", "3", "type", "stream"))
//MBbloom type
require.NoError(t, rdb.Do(ctx, "bf.reserve", "MBbloomtype1", "0.02", "1000").Err())
require.NoError(t, rdb.Do(ctx, "bf.reserve", "MBbloomtype2", "0.02", "1000").Err())
require.NoError(t, rdb.Do(ctx, "bf.reserve", "MBbloomtype3", "0.02", "1000").Err())
require.Equal(t, []string{"MBbloomtype1", "MBbloomtype2", "MBbloomtype3"}, scanAll(t, rdb, "match", "MBbloomtype*", "type", "MBbloom--"))
require.Equal(t, []string{"MBbloomtype1", "MBbloomtype2", "MBbloomtype3"}, scanAll(t, rdb, "match", "MBbloomtype*", "count", "3", "type", "MBbloom--"))
//ReJSON-RL type
require.NoError(t, rdb.Do(ctx, "JSON.SET", "ReJSONtype1", "$", ` {"x":1, "y":2} `).Err())
require.NoError(t, rdb.Do(ctx, "JSON.SET", "ReJSONtype2", "$", ` {"x":1, "y":2} `).Err())
require.NoError(t, rdb.Do(ctx, "JSON.SET", "ReJSONtype3", "$", ` {"x":1, "y":2} `).Err())
require.Equal(t, []string{"ReJSONtype1", "ReJSONtype2", "ReJSONtype3"}, scanAll(t, rdb, "match", "ReJSONtype*", "type", "ReJSON-RL"))
require.Equal(t, []string{"ReJSONtype1", "ReJSONtype2", "ReJSONtype3"}, scanAll(t, rdb, "match", "ReJSONtype*", "count", "3", "type", "ReJSON-RL"))
//invalid type
util.ErrorRegexp(t, rdb.Do(ctx, "SCAN", "0", "count", "1", "match", "a*", "type", "hi").Err(), "Invalid type")

})

}

// SCAN of Kvrocks returns _cursor instead of cursor. Thus, redis.Client Scan can fail with
Expand Down

0 comments on commit eb94fbb

Please sign in to comment.