From c2351b491aad1003529e1cfdaf41b69188bbe46d Mon Sep 17 00:00:00 2001
From: Leandro Forain <forain@gmail.com>
Date: Fri, 8 Jan 2021 16:36:20 -0300
Subject: [PATCH] Add ScanType command to Scan with 'type' option

As of version 6.0 you can use this 'type' option to ask SCAN to only
return objects that match a given type, allowing you to
iterate through the database looking for keys of a specific type.
---
 commands.go      | 17 +++++++++++++++++
 commands_test.go | 12 ++++++++++++
 example_test.go  | 28 ++++++++++++++++++++++++++++
 3 files changed, 57 insertions(+)

diff --git a/commands.go b/commands.go
index 79698bab..422a7c5b 100644
--- a/commands.go
+++ b/commands.go
@@ -141,6 +141,7 @@ type Cmdable interface {
 	BitField(ctx context.Context, key string, args ...interface{}) *IntSliceCmd
 
 	Scan(ctx context.Context, cursor uint64, match string, count int64) *ScanCmd
+	ScanType(ctx context.Context, cursor uint64, match string, count int64, keyType string) *ScanCmd
 	SScan(ctx context.Context, key string, cursor uint64, match string, count int64) *ScanCmd
 	HScan(ctx context.Context, key string, cursor uint64, match string, count int64) *ScanCmd
 	ZScan(ctx context.Context, key string, cursor uint64, match string, count int64) *ScanCmd
@@ -965,6 +966,22 @@ func (c cmdable) Scan(ctx context.Context, cursor uint64, match string, count in
 	return cmd
 }
 
+func (c cmdable) ScanType(ctx context.Context, cursor uint64, match string, count int64, keyType string) *ScanCmd {
+	args := []interface{}{"scan", cursor}
+	if match != "" {
+		args = append(args, "match", match)
+	}
+	if count > 0 {
+		args = append(args, "count", count)
+	}
+	if keyType != "" {
+		args = append(args, "type", keyType)
+	}
+	cmd := NewScanCmd(ctx, c, args...)
+	_ = c(ctx, cmd)
+	return cmd
+}
+
 func (c cmdable) SScan(ctx context.Context, key string, cursor uint64, match string, count int64) *ScanCmd {
 	args := []interface{}{"sscan", key, cursor}
 	if match != "" {
diff --git a/commands_test.go b/commands_test.go
index b5855c6d..707aea1b 100644
--- a/commands_test.go
+++ b/commands_test.go
@@ -770,6 +770,18 @@ var _ = Describe("Commands", func() {
 			Expect(cursor).NotTo(BeZero())
 		})
 
+		It("should ScanType", func() {
+			for i := 0; i < 1000; i++ {
+				set := client.Set(ctx, fmt.Sprintf("key%d", i), "hello", 0)
+				Expect(set.Err()).NotTo(HaveOccurred())
+			}
+
+			keys, cursor, err := client.ScanType(ctx, 0, "", 0, "string").Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(keys).NotTo(BeEmpty())
+			Expect(cursor).NotTo(BeZero())
+		})
+
 		It("should SScan", func() {
 			for i := 0; i < 1000; i++ {
 				sadd := client.SAdd(ctx, "myset", fmt.Sprintf("member%d", i))
diff --git a/example_test.go b/example_test.go
index c4099dc8..161eabf6 100644
--- a/example_test.go
+++ b/example_test.go
@@ -248,6 +248,34 @@ func ExampleClient_Scan() {
 	// Output: found 33 keys
 }
 
+func ExampleClient_ScanType() {
+	rdb.FlushDB(ctx)
+	for i := 0; i < 33; i++ {
+		err := rdb.Set(ctx, fmt.Sprintf("key%d", i), "value", 0).Err()
+		if err != nil {
+			panic(err)
+		}
+	}
+
+	var cursor uint64
+	var n int
+	for {
+		var keys []string
+		var err error
+		keys, cursor, err = rdb.ScanType(ctx, cursor, "key*", 10, "string").Result()
+		if err != nil {
+			panic(err)
+		}
+		n += len(keys)
+		if cursor == 0 {
+			break
+		}
+	}
+
+	fmt.Printf("found %d keys\n", n)
+	// Output: found 33 keys
+}
+
 func ExampleClient_Pipelined() {
 	var incr *redis.IntCmd
 	_, err := rdb.Pipelined(ctx, func(pipe redis.Pipeliner) error {