support keyspace by sys env or config

Signed-off-by: ystaticy <[email protected]>

.github/workflows/cd.yml

@@ -14,6 +14,7 @@ jobs:
       ECR_REPO: dbaas/dev-tier/tidb
       AWS_ROLE_ARN: ${{ secrets.DBAAS_DEV_AWS_ROLE }}
       AWS_REGION: us-west-2
+      GITHUB_TOKEN: ${{ secrets.GIT_ACTION_BOT }}
     steps:
       - name: Check out Code
         uses: actions/checkout@v2

Dockerfile

@@ -25,6 +25,11 @@ COPY go.sum .
 COPY parser/go.mod parser/go.mod
 COPY parser/go.sum parser/go.sum
 
+ARG GITHUB_TOKEN
+RUN if [ -n "$GITHUB_TOKEN"]; then git config --global url."https://${GITHUB_TOKEN}@github.com/".insteadOf "https://github.com/"; fi
+
+ENV GOPRIVATE=github.com/tidbcloud
+
 RUN GO111MODULE=on go mod download
 
 # Build real binaries

config/config.go

@@ -179,6 +179,7 @@ type Config struct {
 	VersionComment             string                  `toml:"version-comment" json:"version-comment"`
 	TiDBEdition                string                  `toml:"tidb-edition" json:"tidb-edition"`
 	TiDBReleaseVersion         string                  `toml:"tidb-release-version" json:"tidb-release-version"`
+	KeyspaceName               string                  `toml:"keyspace-name" json:"keyspace-name"`
 	Log                        Log                     `toml:"log" json:"log"`
 	Instance                   Instance                `toml:"instance" json:"instance"`
 	Security                   Security                `toml:"security" json:"security"`

config/config.toml.example

@@ -462,3 +462,8 @@ tidb_record_plan_in_slow_log = 1
 
 # The maximum permitted number of simultaneous client connections. When the value is 0, the number of connections is unlimited.
 max_connections = 0
+
+# Set keyspace name in multi-tenant scenario.
+# If don't set keyspace_name or keyspace_name == "", it means the keyspace name is not set, and in single-tenant.
+# All architectures and request should keep the single-tenant design.
+keyspace-name= ""

domain/keyspace.go

@@ -0,0 +1,25 @@
+package domain
+
+import (
+	"os"
+
+	"github.com/pingcap/tidb/config"
+)
+
+const EnvVarKeyspaceName = "KEYSPACE_NAME"
+
+func GetKeyspaceNameBySettings() (keyspaceName string) {
+
+	// The cfg.keyspaceName get higher weights than KEYSPACE_NAME in system env.
+	keyspaceName = config.GetGlobalConfig().KeyspaceName
+	if !IsKeyspaceNameEmpty(keyspaceName) {
+		return keyspaceName
+	}
+
+	keyspaceName = os.Getenv(EnvVarKeyspaceName)
+	return keyspaceName
+}
+
+func IsKeyspaceNameEmpty(keyspaceName string) bool {
+	return keyspaceName == ""
+}

domain/keyspace_test.go

@@ -0,0 +1,80 @@
+// Copyright 2021 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package domain
+
+import (
+	"os"
+	"testing"
+
+	"github.com/pingcap/tidb/config"
+	"github.com/stretchr/testify/suite"
+)
+
+type keyspaceSuite struct {
+	suite.Suite
+}
+
+func TestSetKeyspaceName(t *testing.T) {
+	suite.Run(t, new(keyspaceSuite))
+}
+
+func (k *keyspaceSuite) TearDownTest() {
+	// Clear keyspace setting
+	conf := config.GetGlobalConfig()
+	conf.KeyspaceName = ""
+	config.StoreGlobalConfig(conf)
+
+	os.Unsetenv(EnvVarKeyspaceName)
+}
+
+func (k *keyspaceSuite) TestSetKeyspaceNameInSysEnv() {
+
+	keyspaceName := "test_keyspace_env"
+	os.Setenv(EnvVarKeyspaceName, keyspaceName)
+
+	getKeyspaceName := GetKeyspaceNameBySettings()
+	// Check the keyspaceName which get from GetKeyspaceNameBySettings, equals the keyspaceName which is in system env.
+	k.Equal(keyspaceName, getKeyspaceName)
+	k.Equal(false, IsKeyspaceNameEmpty(getKeyspaceName))
+
+}
+
+func (k *keyspaceSuite) TestSetKeyspaceNameInConf() {
+
+	keyspaceNameInCfg := "test_keyspace_cfg"
+	keyspaceNameInSysenv := "test_keyspace_env"
+	os.Setenv(EnvVarKeyspaceName, keyspaceNameInSysenv)
+
+	// Set KeyspaceName in conf
+	c1 := config.GetGlobalConfig()
+	c1.KeyspaceName = keyspaceNameInCfg
+
+	getKeyspaceName := GetKeyspaceNameBySettings()
+
+	// Check the keyspaceName which get from GetKeyspaceNameBySettings, equals keyspaceNameInCfg which is in conf.
+	// The cfg.keyspaceName get higher weights than KEYSPACE_NAME in system env.
+	k.Equal(keyspaceNameInCfg, getKeyspaceName)
+	k.Equal(false, IsKeyspaceNameEmpty(getKeyspaceName))
+
+}
+
+func (k *keyspaceSuite) TestNoKeyspaceNameSet() {
+
+	getKeyspaceName := GetKeyspaceNameBySettings()
+
+	k.Equal("", getKeyspaceName)
+	k.Equal(true, IsKeyspaceNameEmpty(getKeyspaceName))
+
+}

go.mod

@@ -46,7 +46,7 @@ require (
 	github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c
 	github.com/pingcap/failpoint v0.0.0-20220423142525-ae43b7f4e5c3
 	github.com/pingcap/fn v0.0.0-20200306044125-d5540d389059
-	github.com/pingcap/kvproto v0.0.0-20220711062932-08b02befd813
+	github.com/pingcap/kvproto v0.0.0-20220808072825-3692dfb0dad7
 	github.com/pingcap/log v1.1.0
 	github.com/pingcap/sysutil v0.0.0-20220114020952-ea68d2dbf5b4
 	github.com/pingcap/tidb/parser v0.0.0-20211011031125-9b13dc409c5e
@@ -242,3 +242,9 @@ replace github.com/pingcap/tidb/parser => ./parser
 replace github.com/dgrijalva/jwt-go => github.com/form3tech-oss/jwt-go v3.2.6-0.20210809144907-32ab6a8243d7+incompatible
 
 replace honnef.co/go/tools => honnef.co/go/tools v0.3.2
+
+// Use github.com/tidbcloud/pd-cse/client
+replace github.com/tikv/pd/client => github.com/tidbcloud/pd-cse/client v0.0.0-20220816074230-171c16e01847
+
+// Use github.com/tikv/client-go/v2@api-v2
+replace github.com/tikv/client-go/v2 => github.com/tikv/client-go/v2 v2.0.1-0.20220816031836-dba8cc31ff9f

go.sum

@@ -733,9 +733,8 @@ github.com/pingcap/fn v0.0.0-20200306044125-d5540d389059/go.mod h1:fMRU1BA1y+r89
 github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989 h1:surzm05a8C9dN8dIUmo4Be2+pMRb6f55i+UIYrluu2E=
 github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989/go.mod h1:O17XtbryoCJhkKGbT62+L2OlrniwqiGLSqrmdHCMzZw=
 github.com/pingcap/kvproto v0.0.0-20191211054548-3c6b38ea5107/go.mod h1:WWLmULLO7l8IOcQG+t+ItJ3fEcrL5FxF0Wu+HrMy26w=
-github.com/pingcap/kvproto v0.0.0-20220510035547-0e2f26c0a46a/go.mod h1:OYtxs0786qojVTmkVeufx93xe+jUgm56GUYRIKnmaGI=
-github.com/pingcap/kvproto v0.0.0-20220711062932-08b02befd813 h1:PAXtUVMJnyQQS8t9GzihIFmh6FBXu0JziWbIVknLniA=
-github.com/pingcap/kvproto v0.0.0-20220711062932-08b02befd813/go.mod h1:OYtxs0786qojVTmkVeufx93xe+jUgm56GUYRIKnmaGI=
+github.com/pingcap/kvproto v0.0.0-20220808072825-3692dfb0dad7 h1:ECiYgZszvTI/RHnIDLsHHcgcrdAcxS2PYgMyU6PH43g=
+github.com/pingcap/kvproto v0.0.0-20220808072825-3692dfb0dad7/go.mod h1:OYtxs0786qojVTmkVeufx93xe+jUgm56GUYRIKnmaGI=
 github.com/pingcap/log v0.0.0-20191012051959-b742a5d432e9/go.mod h1:4rbK1p9ILyIfb6hU7OG2CiWSqMXnp3JMbiaVJ6mvoY8=
 github.com/pingcap/log v0.0.0-20200511115504-543df19646ad/go.mod h1:4rbK1p9ILyIfb6hU7OG2CiWSqMXnp3JMbiaVJ6mvoY8=
 github.com/pingcap/log v0.0.0-20210625125904-98ed8e2eb1c7/go.mod h1:8AanEdAHATuRurdGxZXBz0At+9avep+ub7U1AGYLIMM=
@@ -879,10 +878,10 @@ github.com/tdakkota/asciicheck v0.1.1 h1:PKzG7JUTUmVspQTDqtkX9eSiLGossXTybutHwTX
 github.com/tdakkota/asciicheck v0.1.1/go.mod h1:yHp0ai0Z9gUljN3o0xMhYJnH/IcvkdTBOX2fmJ93JEM=
 github.com/tiancaiamao/appdash v0.0.0-20181126055449-889f96f722a2 h1:mbAskLJ0oJfDRtkanvQPiooDH8HvJ2FBh+iKT/OmiQQ=
 github.com/tiancaiamao/appdash v0.0.0-20181126055449-889f96f722a2/go.mod h1:2PfKggNGDuadAa0LElHrByyrz4JPZ9fFx6Gs7nx7ZZU=
-github.com/tikv/client-go/v2 v2.0.1-0.20220725090834-0cdc7c1d0fb9 h1:RW/8xnBK618j7B33mP4TTCVhzj1dgwZv/TT7vHaMpbg=
-github.com/tikv/client-go/v2 v2.0.1-0.20220725090834-0cdc7c1d0fb9/go.mod h1:v3DEt8LS9olI6D6El17pYBWq7B28hw3NnDFTxQHDLpY=
-github.com/tikv/pd/client v0.0.0-20220725055910-7187a7ab72db h1:r1eMh9Rny3hfWuBuxOnbsCRrR4FhthiNxLQ5rAUtaww=
-github.com/tikv/pd/client v0.0.0-20220725055910-7187a7ab72db/go.mod h1:ew8kS0yIcEaSetuuywkTLIUBR+sz3J5XvAYRae11qwc=
+github.com/tidbcloud/pd-cse/client v0.0.0-20220816074230-171c16e01847 h1:X7Y4Wys49aPBDI1RxROubDOVHRE3/lopl/grUK7Cd6U=
+github.com/tidbcloud/pd-cse/client v0.0.0-20220816074230-171c16e01847/go.mod h1:l76QLn0xvJx6k6maY+OmybxJDUnXf5Z4xDsypKRctao=
+github.com/tikv/client-go/v2 v2.0.1-0.20220816031836-dba8cc31ff9f h1:FE+/w1ufHodyeNJFg/4UYrIvoDGG4KN4SmJZ5RBbBcA=
+github.com/tikv/client-go/v2 v2.0.1-0.20220816031836-dba8cc31ff9f/go.mod h1:CSej9vvfDdGZVBkkME5XwQzL6LTONlVa81GUJnZf1gI=
 github.com/tklauser/go-sysconf v0.3.9/go.mod h1:11DU/5sG7UexIrp/O6g35hrWzu0JxlwQ3LSFUzyeuhs=
 github.com/tklauser/go-sysconf v0.3.10 h1:IJ1AZGZRWbY8T5Vfk04D9WOA5WSejdflXxP03OUqALw=
 github.com/tklauser/go-sysconf v0.3.10/go.mod h1:C8XykCvCb+Gn0oNCWPIlcb0RuglQTYaQ2hGm7jmxEFk=

kv/interface_mock_test.go

@@ -235,6 +235,10 @@ func (s *mockStorage) GetMinSafeTS(txnScope string) uint64 {
 	return 0
 }
 
+func (s *mockStorage) GetCodec() tikv.Codec {
+	return nil
+}
+
 // newMockStorage creates a new mockStorage.
 func newMockStorage() Storage {
 	return &mockStorage{}

kv/kv.go

@@ -482,6 +482,8 @@ type Storage interface {
 	GetMinSafeTS(txnScope string) uint64
 	// GetLockWaits return all lock wait information
 	GetLockWaits() ([]*deadlockpb.WaitForEntry, error)
+	// GetCodec gets the codec of the storage.
+	GetCodec() tikv.Codec
 }
 
 // EtcdBackend is used for judging a storage is a real TiKV.

store/driver/tikv_driver.go

@@ -18,6 +18,7 @@ import (
 	"context"
 	"crypto/tls"
 	"fmt"
+	"github.com/pingcap/tidb/domain"
 	"math/rand"
 	"net/url"
 	"strings"
@@ -88,6 +89,7 @@ func WithPDClientConfig(client config.PDClient) Option {
 
 // TiKVDriver implements engine TiKV.
 type TiKVDriver struct {
+	keyspaceName    string
 	pdConfig        config.PDClient
 	security        config.Security
 	tikvConfig      config.TiKVClient
@@ -117,7 +119,7 @@ func (d TiKVDriver) OpenWithOptions(path string, options ...Option) (kv.Storage,
 	mc.Lock()
 	defer mc.Unlock()
 	d.setDefaultAndOptions(options...)
-	etcdAddrs, disableGC, err := config.ParsePath(path)
+	etcdAddrs, disableGC, keyspaceName, err := config.ParsePath(path)
 	if err != nil {
 		return nil, errors.Trace(err)
 	}
@@ -157,8 +159,29 @@ func (d TiKVDriver) OpenWithOptions(path string, options ...Option) (kv.Storage,
 		return nil, errors.Trace(err)
 	}
 
-	pdClient := tikv.CodecPDClient{Client: pdCli}
-	s, err := tikv.NewKVStore(uuid, &pdClient, spkv, tikv.NewRPCClient(tikv.WithSecurity(d.security)))
+	var (
+		pdClient *tikv.CodecPDClient
+	)
+
+	if domain.IsKeyspaceNameEmpty(keyspaceName) {
+		logutil.BgLogger().Info("using API V1.")
+		pdClient = tikv.NewCodecPDClient(tikv.ModeTxn, pdCli)
+	} else {
+		logutil.BgLogger().Info("using API V2.", zap.String("keyspaceName", keyspaceName))
+		pdClient, err = tikv.NewCodecPDClientWithKeyspace(tikv.ModeTxn, pdCli, keyspaceName)
+		if err != nil {
+			return nil, errors.Trace(err)
+		}
+	}
+
+	codec := pdClient.GetCodec()
+
+	rpcClient := tikv.NewRPCClient(
+		tikv.WithSecurity(d.security),
+		tikv.WithCodec(codec),
+	)
+
+	s, err := tikv.NewKVStore(uuid, pdClient, spkv, rpcClient)
 	if err != nil {
 		return nil, errors.Trace(err)
 	}
@@ -178,6 +201,7 @@ func (d TiKVDriver) OpenWithOptions(path string, options ...Option) (kv.Storage,
 		memCache:  kv.NewCacheDB(),
 		enableGC:  !disableGC,
 		coprStore: coprStore,
+		codec:     codec,
 	}
 
 	mc.cache[uuid] = store
@@ -192,6 +216,7 @@ type tikvStore struct {
 	enableGC  bool
 	gcWorker  *gcworker.GCWorker
 	coprStore *copr.Store
+	codec     tikv.Codec
 }
 
 // Name gets the name of the storage engine
@@ -343,3 +368,7 @@ func (s *tikvStore) GetLockWaits() ([]*deadlockpb.WaitForEntry, error) {
 	}
 	return result, nil
 }
+
+func (s *tikvStore) GetCodec() tikv.Codec {
+	return s.codec
+}

store/helper/helper.go

@@ -78,6 +78,7 @@ type Storage interface {
 	Closed() <-chan struct{}
 	GetMinSafeTS(txnScope string) uint64
 	GetLockWaits() ([]*deadlockpb.WaitForEntry, error)
+	GetCodec() tikv.Codec
 }
 
 // Helper is a middleware to get some information from tikv/pd. It can be used for TiDB's http api or mem table.

store/mockstore/mockstorage/storage.go

@@ -117,6 +117,10 @@ func (s *mockStorage) Close() error {
 	return s.KVStore.Close()
 }
 
+func (s *mockStorage) GetCodec() tikv.Codec {
+	return nil
+}
+
 // MockLockWaitSetter is used to set the mocked lock wait information, which helps implementing tests that uses the
 // GetLockWaits function.
 type MockLockWaitSetter interface {

store/mockstore/unistore/pd.go

@@ -17,6 +17,7 @@ package unistore
 import (
 	"context"
 	"errors"
+	"github.com/pingcap/kvproto/pkg/keyspacepb"
 	"math"
 	"sync"
 
@@ -35,6 +36,16 @@ type pdClient struct {
 	globalConfig      map[string]string
 }
 
+func (c *pdClient) LoadKeyspace(ctx context.Context, name string) (*keyspacepb.KeyspaceMeta, error) {
+	//TODO implement me, introduce keyspace
+	panic("unimplemented")
+}
+
+func (c *pdClient) WatchKeyspaces(ctx context.Context) (chan []*keyspacepb.KeyspaceMeta, error) {
+	//TODO implement me, introduce keyspace
+	panic("unimplemented")
+}
+
 func newPDClient(pd *us.MockPD) *pdClient {
 	return &pdClient{
 		MockPD:            pd,

store/mockstore/unistore/tikv/server.go

@@ -52,6 +52,11 @@ type Server struct {
 	stopped       int32
 }
 
+func (svr *Server) GetLockWaitHistory(ctx context.Context, request *kvrpcpb.GetLockWaitHistoryRequest) (*kvrpcpb.GetLockWaitHistoryResponse, error) {
+	//TODO implement me
+	panic("unimplemented")
+}
+
 // NewServer returns a new server.
 func NewServer(rm RegionManager, store *MVCCStore, innerServer InnerServer) *Server {
 	return &Server{

tidb-server/main.go

@@ -203,7 +203,8 @@ func main() {
 	setupBinlogClient()
 	setupMetrics()
 
-	storage, dom := createStoreAndDomain()
+	keyspaceName := domain.GetKeyspaceNameBySettings()
+	storage, dom := createStoreAndDomain(keyspaceName)
 	svr := createServer(storage, dom)
 
 	// Register error API is not thread-safe, the caller MUST NOT register errors after initialization.
@@ -293,9 +294,9 @@ func registerMetrics() {
 	}
 }
 
-func createStoreAndDomain() (kv.Storage, *domain.Domain) {
+func createStoreAndDomain(keyspaceName string) (kv.Storage, *domain.Domain) {
 	cfg := config.GetGlobalConfig()
-	fullPath := fmt.Sprintf("%s://%s", cfg.Store, cfg.Path)
+	fullPath := fmt.Sprintf("%s://%s?keyspaceName=%s", cfg.Store, cfg.Path, keyspaceName)
 	var err error
 	storage, err := kvstore.New(fullPath)
 	terror.MustNil(err)

util/mock/store.go

@@ -80,3 +80,7 @@ func (*Store) GetMinSafeTS(_ string) uint64 {
 func (*Store) GetLockWaits() ([]*deadlockpb.WaitForEntry, error) {
 	return nil, nil
 }
+
+func (*Store) GetCodec() tikv.Codec {
+	return nil
+}