Add support for specified DB scans

tair-opensource · Oct 23, 2023 · 3358ad5 · 3358ad5
1 parent 8d239ae
commit 3358ad5
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 2 deletions.
diff --git a/docs/src/en/reader/scan_reader.md b/docs/src/en/reader/scan_reader.md
@@ -24,6 +24,7 @@ username = ""              # keep empty if not using ACL
 password = ""              # keep empty if no authentication is required
 tls = false
 ksn = false                # set to true to enabled Redis keyspace notifications (KSN) subscription
+dbs = []                   # set you want to scan dbs, if you don't want to scan all
 ```
 
 * `cluster`：源端是否为集群
@@ -35,6 +36,7 @@ ksn = false                # set to true to enabled Redis keyspace notifications
 * `tls`：源端是否开启 TLS/SSL，不需要配置证书因为 RedisShake 没有校验服务器证书
 * `ksn`：开启 `ksn` 参数后 RedisShake 会在 `SCAN` 之前使用 [Redis keyspace notifications](https://redis.io/docs/manual/keyspace-notifications/)
 能力来订阅 Key 的变化。当 Key 发生变化时，RedisShake 会使用 `DUMP` 与 `RESTORE` 命令来从源端读取 Key 的内容，并写入目标端。
+* `dbs`: 源端为非集群模式时，支持指定DB库
 
 ::: warning
 Redis keyspace notifications 不会感知到 `FLUSHALL` 与 `FLUSHDB` 命令，因此在使用 `ksn` 参数时，需要确保源端数据库不会执行这两个命令。

diff --git a/docs/src/zh/reader/scan_reader.md b/docs/src/zh/reader/scan_reader.md
@@ -24,6 +24,7 @@ username = ""              # keep empty if not using ACL
 password = ""              # keep empty if no authentication is required
 tls = false
 ksn = false                # set to true to enabled Redis keyspace notifications (KSN) subscription
+dbs = []                   # set you want to scan dbs, if you don't want to scan all
 ```
 
 * `cluster`：源端是否为集群
@@ -35,6 +36,7 @@ ksn = false                # set to true to enabled Redis keyspace notifications
 * `tls`：源端是否开启 TLS/SSL，不需要配置证书因为 RedisShake 没有校验服务器证书
 * `ksn`：开启 `ksn` 参数后 RedisShake 会在 `SCAN` 之前使用 [Redis keyspace notifications](https://redis.io/docs/manual/keyspace-notifications/)
 能力来订阅 Key 的变化。当 Key 发生变化时，RedisShake 会使用 `DUMP` 与 `RESTORE` 命令来从源端读取 Key 的内容，并写入目标端。
+* `dbs`: 源端为非集群模式时，支持指定DB库
 
 ::: warning
 Redis keyspace notifications 不会感知到 `FLUSHALL` 与 `FLUSHDB` 命令，因此在使用 `ksn` 参数时，需要确保源端数据库不会执行这两个命令。

diff --git a/internal/reader/scan_standalone_reader.go b/internal/reader/scan_standalone_reader.go
@@ -22,6 +22,7 @@ type ScanReaderOptions struct {
 	Password string `mapstructure:"password" default:""`
 	Tls      bool   `mapstructure:"tls" default:"false"`
 	KSN      bool   `mapstructure:"ksn" default:"false"`
+	DBS      []int  `mapstructure:"dbs"`
 }
 
 type dbKey struct {
@@ -52,7 +53,11 @@ func NewScanStandaloneReader(opts *ScanReaderOptions) Reader {
 	if c.IsCluster() { // not use opts.Cluster, because user may use standalone mode to scan a cluster node
 		r.dbs = []int{0}
 	} else {
-		r.dbs = []int{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}
+		if len(opts.DBS) == 0 {
+			r.dbs = []int{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}
+		} else {
+			r.dbs = opts.DBS
+		}
 	}
 	r.opts = opts
 	r.ch = make(chan *entry.Entry, 1024)
@@ -99,7 +104,7 @@ func (r *scanStandaloneReader) subscript() {
 
 func (r *scanStandaloneReader) scan() {
 	c := client.NewRedisClient(r.opts.Address, r.opts.Username, r.opts.Password, r.opts.Tls)
-	for dbId := range r.dbs {
+	for _, dbId := range r.dbs {
 		if dbId != 0 {
 			reply := c.DoWithStringReply("SELECT", strconv.Itoa(dbId))
 			if reply != "OK" {

diff --git a/shake.toml b/shake.toml
@@ -17,6 +17,7 @@ sync_aof = true # set to false if you don't want to sync aof
 # password = ""              # keep empty if no authentication is required
 # ksn = false                # set to true to enabled Redis keyspace notifications (KSN) subscription
 # tls = false
+# dbs = []                   # set you want to scan dbs such as [1,5,7], if you don't want to scan all
 
 # [rdb_reader]
 # filepath = "/tmp/dump.rdb"