Skip to content

Commit fbcabe0

Browse files
craig[bot]shubhamdhama
craig[bot]
and
shubhamdhama
committed
Merge #141490
141490: testutils: add `GrantTenantCapability` to test interfaces r=stevendanna,cthumuluru-crdb a=shubhamdhama Many tests require granting a capability when using external-process mode. Since this pattern is repeated in multiple places, it makes sense to provide it as a utility. As part of this cleanup, I've also adjusted when it runs: it now applies only to external-process tenants, as shared-process tenants always have all capabilities. Informs: #138912 Epic: CRDB-38970 Release note: None Co-authored-by: Shubham Dhama <[email protected]>
2 parents b0b0e69 + 2ea4166 commit fbcabe0

18 files changed

+126
-87
lines changed

‎pkg/backup/backup_test.go

+5-3
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,7 @@ import (
6262
"github.com/cockroachdb/cockroach/pkg/kv/kvserver/protectedts"
6363
"github.com/cockroachdb/cockroach/pkg/kv/kvserver/protectedts/ptutil"
6464
"github.com/cockroachdb/cockroach/pkg/multitenant/mtinfopb"
65+
"github.com/cockroachdb/cockroach/pkg/multitenant/tenantcapabilities"
6566
"github.com/cockroachdb/cockroach/pkg/roachpb"
6667
"github.com/cockroachdb/cockroach/pkg/security/securitytest"
6768
"github.com/cockroachdb/cockroach/pkg/security/username"
@@ -589,9 +590,10 @@ func TestBackupRestoreAppend(t *testing.T) {
589590
tc, sqlDB, tmpDir, cleanupFn := backupRestoreTestSetupWithParams(t, multiNode, numAccounts, InitManualReplication, params)
590591
defer cleanupFn()
591592

592-
if !tc.ApplicationLayer(0).Codec().ForSystemTenant() {
593-
systemRunner := sqlutils.MakeSQLRunner(tc.SystemLayer(0).SQLConn(t))
594-
systemRunner.Exec(t, `ALTER TENANT [$1] GRANT CAPABILITY can_admin_relocate_range=true`, serverutils.TestTenantID().ToUint64())
593+
if tc.DefaultTenantDeploymentMode().IsExternal() {
594+
tc.GrantTenantCapabilities(
595+
ctx, t, serverutils.TestTenantID(),
596+
map[tenantcapabilities.ID]string{tenantcapabilities.CanAdminRelocateRange: "true"})
595597
}
596598

597599
// Ensure that each node has at least one leaseholder. (These splits were

‎pkg/ccl/changefeedccl/BUILD.bazel

+1
Original file line numberDiff line numberDiff line change
@@ -263,6 +263,7 @@ go_test(
263263
"//pkg/kv/kvserver/kvserverbase",
264264
"//pkg/kv/kvserver/protectedts",
265265
"//pkg/kv/kvserver/protectedts/ptpb",
266+
"//pkg/multitenant/tenantcapabilities",
266267
"//pkg/roachpb",
267268
"//pkg/scheduledjobs",
268269
"//pkg/scheduledjobs/schedulebase",

‎pkg/ccl/changefeedccl/changefeed_dist_test.go

+7-1
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import (
1919
"github.com/cockroachdb/cockroach/pkg/jobs"
2020
"github.com/cockroachdb/cockroach/pkg/keys"
2121
"github.com/cockroachdb/cockroach/pkg/kv/kvclient/kvcoord"
22+
"github.com/cockroachdb/cockroach/pkg/multitenant/tenantcapabilities"
2223
"github.com/cockroachdb/cockroach/pkg/roachpb"
2324
"github.com/cockroachdb/cockroach/pkg/sql"
2425
"github.com/cockroachdb/cockroach/pkg/sql/catalog"
@@ -400,8 +401,13 @@ func newRangeDistributionTester(
400401

401402
systemDB := sqlutils.MakeSQLRunner(tc.SystemLayer(len(tc.Servers) - 1).SQLConn(t))
402403
systemDB.Exec(t, "SET CLUSTER SETTING kv.rangefeed.enabled = true")
404+
if tc.DefaultTenantDeploymentMode().IsExternal() {
405+
tc.GrantTenantCapabilities(
406+
ctx, t, serverutils.TestTenantID(),
407+
map[tenantcapabilities.ID]string{tenantcapabilities.CanAdminRelocateRange: "true"})
408+
}
409+
403410
if tc.StartedDefaultTestTenant() {
404-
systemDB.Exec(t, `ALTER TENANT [$1] GRANT CAPABILITY can_admin_relocate_range=true`, serverutils.TestTenantID().ToUint64())
405411
// Give 1,000,000 upfront tokens to the tenant, and keep the tokens per
406412
// second rate to the default value of 10,000. This helps avoid throttling
407413
// in the tests.

‎pkg/server/authserver/authentication_test.go

+4-8
Original file line numberDiff line numberDiff line change
@@ -94,15 +94,11 @@ func TestSSLEnforcement(t *testing.T) {
9494
})
9595
defer srv.Stopper().Stop(ctx)
9696

97-
if srv.TenantController().StartedDefaultTestTenant() {
97+
if srv.DeploymentMode().IsExternal() {
9898
// Enable access to the nodes endpoint for the test tenant.
99-
_, err := srv.SystemLayer().SQLConn(t).Exec(
100-
`ALTER TENANT [$1] GRANT CAPABILITY can_view_node_info=true`, serverutils.TestTenantID().ToUint64())
101-
require.NoError(t, err)
102-
103-
serverutils.WaitForTenantCapabilities(t, srv, serverutils.TestTenantID(), map[tenantcapabilities.ID]string{
104-
tenantcapabilities.CanViewNodeInfo: "true",
105-
}, "")
99+
require.NoError(t, srv.GrantTenantCapabilities(
100+
ctx, serverutils.TestTenantID(),
101+
map[tenantcapabilities.ID]string{tenantcapabilities.CanViewNodeInfo: "true"}))
106102
}
107103

108104
s := srv.ApplicationLayer()

‎pkg/server/storage_api/health_test.go

+6-9
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,6 @@ import (
2323
"github.com/cockroachdb/cockroach/pkg/util/log"
2424
"github.com/cockroachdb/errors"
2525
"github.com/stretchr/testify/assert"
26-
"github.com/stretchr/testify/require"
2726
)
2827

2928
func TestHealthAPI(t *testing.T) {
@@ -123,16 +122,14 @@ func TestLivenessAPI(t *testing.T) {
123122
tc := testcluster.StartTestCluster(t, 3, base.TestClusterArgs{})
124123
defer tc.Stopper().Stop(context.Background())
125124

125+
ctx := context.Background()
126+
126127
// The liveness endpoint needs a special tenant capability.
127-
if tc.Server(0).TenantController().StartedDefaultTestTenant() {
128+
if tc.DefaultTenantDeploymentMode().IsExternal() {
128129
// Enable access to the nodes endpoint for the test tenant.
129-
_, err := tc.SystemLayer(0).SQLConn(t).Exec(
130-
`ALTER TENANT [$1] GRANT CAPABILITY can_view_node_info=true`, serverutils.TestTenantID().ToUint64())
131-
require.NoError(t, err)
132-
133-
tc.WaitForTenantCapabilities(t, serverutils.TestTenantID(), map[tenantcapabilities.ID]string{
134-
tenantcapabilities.CanViewNodeInfo: "true",
135-
})
130+
tc.GrantTenantCapabilities(
131+
ctx, t, serverutils.TestTenantID(),
132+
map[tenantcapabilities.ID]string{tenantcapabilities.CanViewNodeInfo: "true"})
136133
}
137134

138135
ts := tc.Server(0).ApplicationLayer()

‎pkg/server/storage_api/network_test.go

+5-12
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@ package storage_api_test
77

88
import (
99
"context"
10-
"fmt"
1110
"testing"
1211

1312
"github.com/cockroachdb/cockroach/pkg/base"
@@ -34,16 +33,10 @@ func TestNetworkConnectivity(t *testing.T) {
3433

3534
s0 := testCluster.Server(0)
3635

37-
if s0.TenantController().StartedDefaultTestTenant() {
38-
_, err := s0.SystemLayer().SQLConn(t).Exec(
39-
`ALTER TENANT [$1] GRANT CAPABILITY can_debug_process=true`,
40-
serverutils.TestTenantID().ToUint64(),
41-
)
42-
require.NoError(t, err)
43-
44-
serverutils.WaitForTenantCapabilities(t, s0, serverutils.TestTenantID(), map[tenantcapabilities.ID]string{
45-
tenantcapabilities.CanDebugProcess: "true",
46-
}, "")
36+
if s0.DeploymentMode().IsExternal() {
37+
testCluster.GrantTenantCapabilities(
38+
ctx, t, serverutils.TestTenantID(),
39+
map[tenantcapabilities.ID]string{tenantcapabilities.CanDebugProcess: "true"})
4740
}
4841

4942
ts := s0.ApplicationLayer()
@@ -71,7 +64,7 @@ func TestNetworkConnectivity(t *testing.T) {
7164
return err
7265
}
7366
require.Equal(t, len(resp.Connections), numNodes-1)
74-
fmt.Printf("got status: %s", resp.Connections[s0.StorageLayer().NodeID()].Peers[stoppedNodeID].Status.String())
67+
t.Logf("got status: %s\n", resp.Connections[s0.StorageLayer().NodeID()].Peers[stoppedNodeID].Status.String())
7568
if resp.Connections[s0.StorageLayer().NodeID()].Peers[stoppedNodeID].Status != serverpb.NetworkConnectivityResponse_ERROR {
7669
return errors.New("waiting for connection state to be changed.")
7770
}

‎pkg/server/storage_api/nodes_test.go

+4-8
Original file line numberDiff line numberDiff line change
@@ -221,15 +221,11 @@ func TestNodesGRPCResponse(t *testing.T) {
221221
})
222222
defer srv.Stopper().Stop(ctx)
223223

224-
if srv.TenantController().StartedDefaultTestTenant() {
224+
if srv.DeploymentMode().IsExternal() {
225225
// Enable access to the nodes endpoint for the test tenant.
226-
_, err := srv.SystemLayer().SQLConn(t).Exec(
227-
`ALTER TENANT [$1] GRANT CAPABILITY can_view_node_info=true`, serverutils.TestTenantID().ToUint64())
228-
require.NoError(t, err)
229-
230-
serverutils.WaitForTenantCapabilities(t, srv, serverutils.TestTenantID(), map[tenantcapabilities.ID]string{
231-
tenantcapabilities.CanViewNodeInfo: "true",
232-
}, "")
226+
require.NoError(t, srv.GrantTenantCapabilities(
227+
ctx, serverutils.TestTenantID(),
228+
map[tenantcapabilities.ID]string{tenantcapabilities.CanViewNodeInfo: "true"}))
233229
}
234230

235231
var request serverpb.NodesRequest

‎pkg/server/testserver.go

+24
Original file line numberDiff line numberDiff line change
@@ -1265,6 +1265,30 @@ func (t *testTenant) DeploymentMode() serverutils.DeploymentMode {
12651265
return t.deploymentMode
12661266
}
12671267

1268+
// GrantTenantCapabilities is part of the serverutils.TenantControlInterface.
1269+
func (ts *testServer) GrantTenantCapabilities(
1270+
ctx context.Context, tenID roachpb.TenantID, targetCaps map[tenantcapabilities.ID]string,
1271+
) error {
1272+
conn, err := ts.SQLConnE()
1273+
if err != nil {
1274+
return err
1275+
}
1276+
1277+
var parts []string
1278+
for k, v := range targetCaps {
1279+
parts = append(parts, k.String()+"="+v)
1280+
}
1281+
capabilities := strings.Join(parts, ",")
1282+
1283+
if _, err = conn.Exec(fmt.Sprintf(`ALTER TENANT [$1] GRANT CAPABILITY %s`, capabilities),
1284+
tenID.ToUint64(),
1285+
); err != nil {
1286+
return err
1287+
}
1288+
1289+
return ts.WaitForTenantCapabilities(ctx, tenID, targetCaps, "")
1290+
}
1291+
12681292
// WaitForTenantCapabilities is part of the serverutils.TenantControlInterface.
12691293
func (ts *testServer) WaitForTenantCapabilities(
12701294
ctx context.Context,

‎pkg/sql/colflow/draining_test.go

+4-7
Original file line numberDiff line numberDiff line change
@@ -54,13 +54,10 @@ func TestDrainingAfterRemoteError(t *testing.T) {
5454
tc := testcluster.StartTestCluster(t, 2 /* nodes */, args)
5555
defer tc.Stopper().Stop(ctx)
5656

57-
if srv := tc.Server(0); srv.TenantController().StartedDefaultTestTenant() {
58-
systemSqlDB := srv.SystemLayer().SQLConn(t, serverutils.DBName("system"))
59-
_, err := systemSqlDB.Exec(`ALTER TENANT [$1] GRANT CAPABILITY can_admin_relocate_range=true`, serverutils.TestTenantID().ToUint64())
60-
require.NoError(t, err)
61-
serverutils.WaitForTenantCapabilities(t, srv, serverutils.TestTenantID(), map[tenantcapabilities.ID]string{
62-
tenantcapabilities.CanAdminRelocateRange: "true",
63-
}, "")
57+
if tc.DefaultTenantDeploymentMode().IsExternal() {
58+
tc.GrantTenantCapabilities(
59+
ctx, t, serverutils.TestTenantID(),
60+
map[tenantcapabilities.ID]string{tenantcapabilities.CanAdminRelocateRange: "true"})
6461
}
6562

6663
// Create two tables, one with small values, and another with large rows.

‎pkg/sql/distsql_physical_planner_test.go

+4-4
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ import (
2323
"github.com/cockroachdb/cockroach/pkg/kv"
2424
"github.com/cockroachdb/cockroach/pkg/kv/kvclient/kvcoord"
2525
"github.com/cockroachdb/cockroach/pkg/kv/kvclient/rangecache"
26+
"github.com/cockroachdb/cockroach/pkg/multitenant/tenantcapabilities"
2627
"github.com/cockroachdb/cockroach/pkg/roachpb"
2728
"github.com/cockroachdb/cockroach/pkg/rpc"
2829
"github.com/cockroachdb/cockroach/pkg/server/serverpb"
@@ -449,10 +450,9 @@ func TestDistSQLUnavailableHosts(t *testing.T) {
449450
}
450451

451452
// Grant capability to run RELOCATE to secondary (test) tenant.
452-
systemDB := sqlutils.MakeSQLRunner(tc.SystemLayer(0).SQLConn(t))
453-
systemDB.Exec(t,
454-
`ALTER TENANT [$1] GRANT CAPABILITY can_admin_relocate_range=true`,
455-
serverutils.TestTenantID().ToUint64())
453+
tc.GrantTenantCapabilities(
454+
ctx, t, serverutils.TestTenantID(),
455+
map[tenantcapabilities.ID]string{tenantcapabilities.CanAdminRelocateRange: "true"})
456456
}
457457

458458
// Connect to node 1 (gateway node)

‎pkg/sql/execstats/traceanalyzer_test.go

+4-7
Original file line numberDiff line numberDiff line change
@@ -78,13 +78,10 @@ func TestTraceAnalyzer(t *testing.T) {
7878

7979
const gatewayNode = 0
8080
srv, s := tc.Server(gatewayNode), tc.ApplicationLayer(gatewayNode)
81-
if srv.TenantController().StartedDefaultTestTenant() {
82-
systemSqlDB := srv.SystemLayer().SQLConn(t, serverutils.DBName("system"))
83-
_, err := systemSqlDB.Exec(`ALTER TENANT [$1] GRANT CAPABILITY can_admin_relocate_range=true`, serverutils.TestTenantID().ToUint64())
84-
require.NoError(t, err)
85-
serverutils.WaitForTenantCapabilities(t, srv, serverutils.TestTenantID(), map[tenantcapabilities.ID]string{
86-
tenantcapabilities.CanAdminRelocateRange: "true",
87-
}, "")
81+
if srv.DeploymentMode().IsExternal() {
82+
require.NoError(t, srv.GrantTenantCapabilities(
83+
ctx, serverutils.TestTenantID(),
84+
map[tenantcapabilities.ID]string{tenantcapabilities.CanAdminRelocateRange: "true"}))
8885
}
8986
db := s.SQLConn(t)
9087
sqlDB := sqlutils.MakeSQLRunner(db)

‎pkg/sql/flowinfra/cluster_test.go

+4-7
Original file line numberDiff line numberDiff line change
@@ -594,13 +594,10 @@ func TestEvalCtxTxnOnRemoteNodes(t *testing.T) {
594594
})
595595
defer tc.Stopper().Stop(ctx)
596596

597-
if srv := tc.Server(0); srv.TenantController().StartedDefaultTestTenant() {
598-
systemSqlDB := srv.SystemLayer().SQLConn(t, serverutils.DBName("system"))
599-
_, err := systemSqlDB.Exec(`ALTER TENANT [$1] GRANT CAPABILITY can_admin_relocate_range=true`, serverutils.TestTenantID().ToUint64())
600-
require.NoError(t, err)
601-
serverutils.WaitForTenantCapabilities(t, srv, serverutils.TestTenantID(), map[tenantcapabilities.ID]string{
602-
tenantcapabilities.CanAdminRelocateRange: "true",
603-
}, "")
597+
if tc.DefaultTenantDeploymentMode().IsExternal() {
598+
tc.GrantTenantCapabilities(
599+
ctx, t, serverutils.TestTenantID(),
600+
map[tenantcapabilities.ID]string{tenantcapabilities.CanAdminRelocateRange: "true"})
604601
}
605602

606603
db := tc.ServerConn(0)

‎pkg/sql/importer/exportcsv_test.go

+4-7
Original file line numberDiff line numberDiff line change
@@ -663,14 +663,11 @@ func TestProcessorEncountersUncertaintyError(t *testing.T) {
663663
ctx := context.Background()
664664
defer tc.Stopper().Stop(ctx)
665665

666-
if tc.StartedDefaultTestTenant() {
667-
systemSqlDB := tc.Server(0).SystemLayer().SQLConn(t, serverutils.DBName("system"))
668-
_, err := systemSqlDB.Exec(`ALTER TENANT [$1] GRANT CAPABILITY can_admin_relocate_range=true`, serverutils.TestTenantID().ToUint64())
669-
require.NoError(t, err)
670-
serverutils.WaitForTenantCapabilities(t, tc.Server(0), serverutils.TestTenantID(), map[tenantcapabilities.ID]string{
671-
tenantcapabilities.CanAdminRelocateRange: "true",
672-
}, "")
666+
if tc.DefaultTenantDeploymentMode().IsExternal() {
667+
tc.GrantTenantCapabilities(ctx, t, serverutils.TestTenantID(),
668+
map[tenantcapabilities.ID]string{tenantcapabilities.CanAdminRelocateRange: "true"})
673669
}
670+
674671
origDB0 := tc.ServerConn(0)
675672

676673
sqlutils.CreateTable(t, origDB0, "t",

‎pkg/sql/logictest/logic.go

+1-7
Original file line numberDiff line numberDiff line change
@@ -1728,12 +1728,6 @@ func (t *logicTest) newCluster(
17281728

17291729
capabilities := toa.capabilities
17301730
if len(capabilities) > 0 {
1731-
for name, value := range capabilities {
1732-
query := fmt.Sprintf("ALTER TENANT [$1] GRANT CAPABILITY %s = $2", name)
1733-
if _, err := conn.Exec(query, tenantID.ToUint64(), value); err != nil {
1734-
t.Fatal(err)
1735-
}
1736-
}
17371731
capabilityMap := make(map[tenantcapabilities.ID]string, len(capabilities))
17381732
for k, v := range capabilities {
17391733
capability, ok := tenantcapabilities.FromName(k)
@@ -1742,7 +1736,7 @@ func (t *logicTest) newCluster(
17421736
}
17431737
capabilityMap[capability.ID()] = v
17441738
}
1745-
t.cluster.WaitForTenantCapabilities(t.t(), tenantID, capabilityMap)
1739+
t.cluster.GrantTenantCapabilities(context.Background(), t.t(), tenantID, capabilityMap)
17461740
}
17471741
}
17481742

‎pkg/sql/opt/exec/explain/output_test.go

+3-7
Original file line numberDiff line numberDiff line change
@@ -163,13 +163,9 @@ func TestCPUTimeEndToEnd(t *testing.T) {
163163
ctx := context.Background()
164164
defer tc.Stopper().Stop(ctx)
165165

166-
if srv := tc.Server(0); srv.TenantController().StartedDefaultTestTenant() {
167-
systemSqlDB := srv.SystemLayer().SQLConn(t, serverutils.DBName("system"))
168-
_, err := systemSqlDB.Exec(`ALTER TENANT [$1] GRANT CAPABILITY can_admin_relocate_range=true`, serverutils.TestTenantID().ToUint64())
169-
require.NoError(t, err)
170-
serverutils.WaitForTenantCapabilities(t, srv, serverutils.TestTenantID(), map[tenantcapabilities.ID]string{
171-
tenantcapabilities.CanAdminRelocateRange: "true",
172-
}, "")
166+
if tc.DefaultTenantDeploymentMode().IsExternal() {
167+
tc.GrantTenantCapabilities(ctx, t, serverutils.TestTenantID(),
168+
map[tenantcapabilities.ID]string{tenantcapabilities.CanAdminRelocateRange: "true"})
173169
}
174170

175171
db := sqlutils.MakeSQLRunner(tc.Conns[0])

‎pkg/testutils/serverutils/api.go

+14
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,10 @@ const (
141141
ExternalProcess
142142
)
143143

144+
func (d DeploymentMode) IsExternal() bool {
145+
return d == ExternalProcess
146+
}
147+
144148
// ApplicationLayerInterface defines accessors to the application
145149
// layer of a test server. Tests written against this interface are
146150
// effectively agnostic to whether they use a virtual cluster or not.
@@ -506,6 +510,16 @@ type TenantControlInterface interface {
506510
// if the tenant record exists in KV.
507511
WaitForTenantReadiness(ctx context.Context, tenantID roachpb.TenantID) error
508512

513+
// GrantTenantCapabilities grants a capability to a tenant and waits until the
514+
// in-memory cache reflects the change.
515+
//
516+
// Note: There is no need to call WaitForTenantCapabilities separately.
517+
GrantTenantCapabilities(
518+
context.Context,
519+
roachpb.TenantID,
520+
map[tenantcapabilities.ID]string,
521+
) error
522+
509523
// WaitForTenantCapabilities waits until the in-RAM cache of
510524
// tenant capabilities has been populated for the given tenant ID
511525
// with the expected target capability values.

‎pkg/testutils/serverutils/test_cluster_shim.go

+16
Original file line numberDiff line numberDiff line change
@@ -250,6 +250,11 @@ type TestClusterInterface interface {
250250
// cluster.
251251
StorageLayer(idx int) StorageLayerInterface
252252

253+
// DefaultTenantDeploymentMode returns the deployment mode of the default
254+
// server or tenant, which can be single-tenant (system-only),
255+
// shared-process, or external-process.
256+
DefaultTenantDeploymentMode() DeploymentMode
257+
253258
// SplitTable splits a range in the table, creates a replica for the right
254259
// side of the split on TargetNodeIdx, and moves the lease for the right
255260
// side of the split to TargetNodeIdx for each SplitPoint. This forces the
@@ -259,6 +264,17 @@ type TestClusterInterface interface {
259264
// are indeed distributed as intended.
260265
SplitTable(t TestFataler, desc catalog.TableDescriptor, sps []SplitPoint)
261266

267+
// GrantTenantCapabilities grants a capability to a tenant and waits until all
268+
// servers have the in-memory cache reflects the change.
269+
//
270+
// Note: There is no need to call WaitForTenantCapabilities separately.
271+
GrantTenantCapabilities(
272+
context.Context,
273+
TestFataler,
274+
roachpb.TenantID,
275+
map[tenantcapabilities.ID]string,
276+
)
277+
262278
// WaitForTenantCapabilities waits until all servers have the specified
263279
// tenant capabilities for the specified tenant ID.
264280
// Only boolean capabilities are currently supported as we wait for the

0 commit comments

Comments
 (0)