diff --git a/src/commands/cmd_server.cc b/src/commands/cmd_server.cc index 442908ffa6c..28c31603c80 100644 --- a/src/commands/cmd_server.cc +++ b/src/commands/cmd_server.cc @@ -840,7 +840,7 @@ class CommandScan : public CommandScanBase { std::vector 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()}; } diff --git a/src/commands/scan_base.h b/src/commands/scan_base.h index da3c81accf6..3a6438c2f09 100644 --- a/src/commands/scan_base.h +++ b/src/commands/scan_base.h @@ -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(iter - RedisTypeNames.begin()); + } else { + return {Status::RedisExecErr, "Invalid type"}; + } } else { return parser.InvalidSyntax(); } @@ -93,6 +99,7 @@ class CommandScanBase : public Commander { std::string cursor_; std::string prefix_; int limit_ = 20; + RedisType type_ = kRedisNone; }; class CommandSubkeyScanBase : public CommandScanBase { diff --git a/src/storage/redis_db.cc b/src/storage/redis_db.cc index fab6562b7c5..a0e06615e76 100644 --- a/src/storage/redis_db.cc +++ b/src/storage/redis_db.cc @@ -310,7 +310,7 @@ rocksdb::Status Database::Keys(const std::string &prefix, std::vector *keys, std::string *end_cursor) { + std::vector *keys, std::string *end_cursor, RedisType type) { end_cursor->clear(); uint64_t cnt = 0; uint16_t slot_start = 0; @@ -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(iter->key(), storage_->IsSlotIdEncoded()); keys->emplace_back(user_key); diff --git a/src/storage/redis_db.h b/src/storage/redis_db.h index 31de41dc668..73a5a6545c0 100644 --- a/src/storage/redis_db.h +++ b/src/storage/redis_db.h @@ -93,7 +93,8 @@ class Database { [[nodiscard]] rocksdb::Status Keys(const std::string &prefix, std::vector *keys = nullptr, KeyNumStats *stats = nullptr); [[nodiscard]] rocksdb::Status Scan(const std::string &cursor, uint64_t limit, const std::string &prefix, - std::vector *keys, std::string *end_cursor = nullptr); + std::vector *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, diff --git a/tests/gocase/unit/scan/scan_test.go b/tests/gocase/unit/scan/scan_test.go index 81a2ff64044..62b91aff9fd 100644 --- a/tests/gocase/unit/scan/scan_test.go +++ b/tests/gocase/unit/scan/scan_test.go @@ -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