Merge pull request #2 from wwu/main

- add parameter to support keyspace, which is supported in Tikv API V2
ken8203 · May 27, 2024 · 1afed46 · 1afed46
2 parents 5f3b5c6 + e4eeb1d
commit 1afed46
Show file tree

Hide file tree

Showing 6 changed files with 12 additions and 8 deletions.
diff --git a/README.md b/README.md
@@ -27,6 +27,7 @@ Flags:
       --help                 Help for tikv-cli
   -h, --host string          PD host address (default "localhost")
   -m, --mode string          Client mode. raw/txn (default "txn")
+  -k, --keyspace string      Keyspace name. (default empty)
   -p, --port string          PD port (default "2379")
 
 Use "tikv-cli [command] --help" for more information about a command.

diff --git a/cmd/client.go b/cmd/client.go
@@ -18,7 +18,7 @@ func newClient() (client.Client, error) {
 		v = client.APIVersion2
 	}
 
-	c, err := client.New([]string{addr(Host, Port)}, client.Mode(Mode), v)
+	c, err := client.New([]string{addr(Host, Port)}, client.Mode(Mode), v, KeySpace)
 	if err != nil {
 		return nil, err
 	}

diff --git a/cmd/cmd.go b/cmd/cmd.go
@@ -16,6 +16,8 @@ var (
 	Port string
 	// Mode is the client mode: raw/txn
 	Mode string
+	// Keyspace name for tikv V2 storage
+	KeySpace string
 	// APIVersion is the API version: v1/v1ttl/v2
 	APIVersion string
 	// Debug determines whether to enable logging in tikv/client-go.
@@ -52,6 +54,7 @@ func init() {
 	rootCmd.PersistentFlags().StringVarP(&Host, "host", "h", "localhost", "PD host address")
 	rootCmd.PersistentFlags().StringVarP(&Port, "port", "p", "2379", "PD port")
 	rootCmd.PersistentFlags().StringVarP(&Mode, "mode", "m", "txn", "Client mode. raw/txn")
+	rootCmd.PersistentFlags().StringVarP(&KeySpace, "keyspace", "k", "", "Tikv keyspace, default is empty")
 	rootCmd.PersistentFlags().StringVarP(&APIVersion, "api-version", "a", "v2", "API version. v1/v1ttl/v2")
 	rootCmd.PersistentFlags().Bool("help", false, "Help for tikv-cli")
 	rootCmd.PersistentFlags().BoolVar(&Debug, "debug", false, "Debug determines whether to enable logging in tikv/client-go")

diff --git a/internal/client/client.go b/internal/client/client.go
@@ -31,12 +31,12 @@ type Client interface {
 	Close(ctx context.Context) error
 }
 
-func New(addrs []string, mode Mode, apiVersion APIVersion) (Client, error) {
+func New(addrs []string, mode Mode, apiVersion APIVersion, keySpace string) (Client, error) {
 	switch mode {
 	case ModeTxn:
-		return newTxnClient(addrs, kvrpcpb.APIVersion(apiVersion))
+		return newTxnClient(addrs, kvrpcpb.APIVersion(apiVersion), keySpace)
 	case ModeRaw:
-		return newRawClient(addrs, kvrpcpb.APIVersion(apiVersion))
+		return newRawClient(addrs, kvrpcpb.APIVersion(apiVersion), keySpace)
 	default:
 		return nil, fmt.Errorf("invalid mode: %s", mode)
 	}

diff --git a/internal/client/raw.go b/internal/client/raw.go
@@ -13,8 +13,8 @@ type rawClient struct {
 
 var _ Client = (*rawClient)(nil)
 
-func newRawClient(addrs []string, apiVersion kvrpcpb.APIVersion) (*rawClient, error) {
-	client, err := rawkv.NewClientWithOpts(context.Background(), addrs, rawkv.WithAPIVersion(apiVersion))
+func newRawClient(addrs []string, apiVersion kvrpcpb.APIVersion, keySpace string) (*rawClient, error) {
+	client, err := rawkv.NewClientWithOpts(context.Background(), addrs, rawkv.WithAPIVersion(apiVersion), rawkv.WithKeyspace(keySpace))
 	if err != nil {
 		return nil, err
 	}

diff --git a/internal/client/txn.go b/internal/client/txn.go
@@ -14,8 +14,8 @@ type txnClient struct {
 
 var _ Client = (*txnClient)(nil)
 
-func newTxnClient(addrs []string, apiVersion kvrpcpb.APIVersion) (*txnClient, error) {
-	client, err := txnkv.NewClient(addrs, txnkv.WithAPIVersion(apiVersion))
+func newTxnClient(addrs []string, apiVersion kvrpcpb.APIVersion, keySpace string) (*txnClient, error) {
+	client, err := txnkv.NewClient(addrs, txnkv.WithAPIVersion(apiVersion), txnkv.WithKeyspace(keySpace))
 	if err != nil {
 		return nil, err
 	}