diff --git a/client/client.go b/client/client.go index 82af1abd4740..0ae3fd21df05 100644 --- a/client/client.go +++ b/client/client.go @@ -308,8 +308,6 @@ func NewClientWithContext(ctx context.Context, svrAddrs []string, security Secur // NewClientWithKeyspace creates a client with context and the specified keyspace id. func NewClientWithKeyspace(ctx context.Context, keyspaceID uint32, svrAddrs []string, security SecurityOption, opts ...ClientOption) (Client, error) { - log.Info("[pd] create pd client with endpoints and keyspace", zap.Strings("pd-address", svrAddrs), zap.Uint32("keyspace-id", keyspaceID)) - tlsCfg := &tlsutil.TLSConfig{ CAPath: security.CAPath, CertPath: security.CertPath, @@ -345,6 +343,49 @@ func NewClientWithKeyspace(ctx context.Context, keyspaceID uint32, svrAddrs []st return c, nil } +// NewClientWithKeyspaceName creates a client with context and the specified keyspace name. +func NewClientWithKeyspaceName(ctx context.Context, keyspace string, svrAddrs []string, security SecurityOption, opts ...ClientOption) (Client, error) { + log.Info("[pd] create pd client with endpoints and keyspace", zap.Strings("pd-address", svrAddrs), zap.String("keyspace", keyspace)) + + tlsCfg := &tlsutil.TLSConfig{ + CAPath: security.CAPath, + CertPath: security.CertPath, + KeyPath: security.KeyPath, + + SSLCABytes: security.SSLCABytes, + SSLCertBytes: security.SSLCertBytes, + SSLKEYBytes: security.SSLKEYBytes, + } + + clientCtx, clientCancel := context.WithCancel(ctx) + c := &client{ + updateTokenConnectionCh: make(chan struct{}, 1), + ctx: clientCtx, + cancel: clientCancel, + svrUrls: addrsToUrls(svrAddrs), + tlsCfg: tlsCfg, + option: newOption(), + } + + // Inject the client options. + for _, opt := range opts { + opt(c) + } + + c.pdSvcDiscovery = newPDServiceDiscovery(clientCtx, clientCancel, &c.wg, c.setServiceMode, c.svrUrls, c.tlsCfg, c.option) + if err := c.setup(); err != nil { + c.cancel() + return nil, err + } + keyspaceMeta, err := c.LoadKeyspace(context.TODO(), keyspace) + // Here we ignore ENTRY_NOT_FOUND error and it will set the keyspaceID to 0. + if err != nil && !strings.Contains(err.Error(), "ENTRY_NOT_FOUND") { + return nil, err + } + c.keyspaceID = keyspaceMeta.GetId() + return c, nil +} + func (c *client) setup() error { // Init the client base. if err := c.pdSvcDiscovery.Init(); err != nil { diff --git a/tests/integrations/mcs/testutil.go b/tests/integrations/mcs/testutil.go index da23d2fe7f25..f9c47aa56f03 100644 --- a/tests/integrations/mcs/testutil.go +++ b/tests/integrations/mcs/testutil.go @@ -27,7 +27,6 @@ import ( bs "github.com/tikv/pd/pkg/basicserver" rm "github.com/tikv/pd/pkg/mcs/resource_manager/server" tso "github.com/tikv/pd/pkg/mcs/tso/server" - "github.com/tikv/pd/pkg/mcs/utils" "github.com/tikv/pd/pkg/utils/logutil" "github.com/tikv/pd/pkg/utils/testutil" ) @@ -51,14 +50,7 @@ func InitLogger(cfg *tso.Config) (err error) { // SetupClientWithKeyspace creates a TSO client for test. func SetupClientWithKeyspace(ctx context.Context, re *require.Assertions, endpoints []string, opts ...pd.ClientOption) pd.Client { - cli, err := pd.NewClientWithKeyspace(ctx, utils.DefaultKeyspaceID, endpoints, pd.SecurityOption{}, opts...) - re.NoError(err) - return cli -} - -// SetupClient creates a TSO client for test. -func SetupClient(ctx context.Context, re *require.Assertions, endpoints []string, opts ...pd.ClientOption) pd.Client { - cli, err := pd.NewClientWithContext(ctx, endpoints, pd.SecurityOption{}, opts...) + cli, err := pd.NewClientWithKeyspaceName(ctx, "", endpoints, pd.SecurityOption{}, opts...) re.NoError(err) return cli } diff --git a/tests/integrations/mcs/tso/server_test.go b/tests/integrations/mcs/tso/server_test.go index 64c59685d51b..0ff65575cc22 100644 --- a/tests/integrations/mcs/tso/server_test.go +++ b/tests/integrations/mcs/tso/server_test.go @@ -173,6 +173,7 @@ func checkTSOPath(re *require.Assertions, isAPIServiceMode bool) { re.NoError(err) leaderName := cluster.WaitLeader() pdLeader := cluster.GetServer(leaderName) + re.NoError(pdLeader.BootstrapCluster()) backendEndpoints := pdLeader.GetAddr() client := pdLeader.GetEtcdClient() if isAPIServiceMode {