From 832f71a91a26847eb435a91ac84b1d2953e2b451 Mon Sep 17 00:00:00 2001
From: Zijian <Shaddoll@users.noreply.github.com>
Date: Thu, 29 Sep 2022 13:12:07 -0700
Subject: [PATCH] Fix Cadence CLI (#5005)

---
 tools/cli/database.go | 7 +++++++
 tools/cli/flags.go    | 1 +
 2 files changed, 8 insertions(+)

diff --git a/tools/cli/database.go b/tools/cli/database.go
index 3d52a1b5d47..e37dec4747a 100644
--- a/tools/cli/database.go
+++ b/tools/cli/database.go
@@ -77,6 +77,10 @@ func getDBFlags() []cli.Flag {
 			Name:  FlagDBRegion,
 			Usage: "persistence region",
 		},
+		cli.IntFlag{
+			Name:  FlagDBShard,
+			Usage: "number of db shards in a sharded SQL database",
+		},
 		cli.StringFlag{
 			Name:  FlagUsername,
 			Usage: "persistence username",
@@ -301,6 +305,9 @@ func overrideSQLDataStore(c *cli.Context, cfg *config.SQL) {
 	if c.IsSet(FlagDBType) || cfg.PluginName == "" {
 		cfg.PluginName = c.String(FlagDBType)
 	}
+	if c.IsSet(FlagDBShard) || cfg.NumShards == 0 {
+		cfg.NumShards = c.Int(FlagDBShard)
+	}
 	if c.IsSet(FlagUsername) || cfg.User == "" {
 		cfg.User = c.String(FlagUsername)
 	}
diff --git a/tools/cli/flags.go b/tools/cli/flags.go
index 32f9a1c792a..e4e866a623c 100644
--- a/tools/cli/flags.go
+++ b/tools/cli/flags.go
@@ -37,6 +37,7 @@ const (
 	FlagDBAddress                         = "db_address"
 	FlagDBPort                            = "db_port"
 	FlagDBRegion                          = "db_region"
+	FlagDBShard                           = "db_shard"
 	FlagHistoryAddressWithAlias           = FlagHistoryAddress + ", had"
 	FlagProtoVersion                      = "protocol_version"
 	FlagDomainID                          = "domain_id"