Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

testutils: add GrantTenantCapabilities to test interfaces #141490

Merged
merged 1 commit into from
Feb 21, 2025
Merged

testutils: add GrantTenantCapabilities to test interfaces #141490

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 5 additions & 3 deletions pkg/backup/backup_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@ import (
"github.com/cockroachdb/cockroach/pkg/kv/kvserver/protectedts"
"github.com/cockroachdb/cockroach/pkg/kv/kvserver/protectedts/ptutil"
"github.com/cockroachdb/cockroach/pkg/multitenant/mtinfopb"
"github.com/cockroachdb/cockroach/pkg/multitenant/tenantcapabilities"
"github.com/cockroachdb/cockroach/pkg/roachpb"
"github.com/cockroachdb/cockroach/pkg/security/securitytest"
"github.com/cockroachdb/cockroach/pkg/security/username"
Expand Down Expand Up @@ -589,9 +590,10 @@ func TestBackupRestoreAppend(t *testing.T) {
tc, sqlDB, tmpDir, cleanupFn := backupRestoreTestSetupWithParams(t, multiNode, numAccounts, InitManualReplication, params)
defer cleanupFn()

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

// Ensure that each node has at least one leaseholder. (These splits were
Expand Down
1 change: 1 addition & 0 deletions pkg/ccl/changefeedccl/BUILD.bazel
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -261,6 +261,7 @@ go_test(
"//pkg/kv/kvserver/kvserverbase",
"//pkg/kv/kvserver/protectedts",
"//pkg/kv/kvserver/protectedts/ptpb",
"//pkg/multitenant/tenantcapabilities",
"//pkg/roachpb",
"//pkg/scheduledjobs",
"//pkg/scheduledjobs/schedulebase",
Expand Down
8 changes: 7 additions & 1 deletion pkg/ccl/changefeedccl/changefeed_dist_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ import (
"github.com/cockroachdb/cockroach/pkg/jobs"
"github.com/cockroachdb/cockroach/pkg/keys"
"github.com/cockroachdb/cockroach/pkg/kv/kvclient/kvcoord"
"github.com/cockroachdb/cockroach/pkg/multitenant/tenantcapabilities"
"github.com/cockroachdb/cockroach/pkg/roachpb"
"github.com/cockroachdb/cockroach/pkg/sql"
"github.com/cockroachdb/cockroach/pkg/sql/catalog"
Expand Down Expand Up @@ -400,8 +401,13 @@ func newRangeDistributionTester(

systemDB := sqlutils.MakeSQLRunner(tc.SystemLayer(len(tc.Servers) - 1).SQLConn(t))
systemDB.Exec(t, "SET CLUSTER SETTING kv.rangefeed.enabled = true")
if tc.DefaultTenantDeploymentMode().IsExternal() {
tc.GrantTenantCapabilities(
ctx, t, serverutils.TestTenantID(),
map[tenantcapabilities.ID]string{tenantcapabilities.CanAdminRelocateRange: "true"})
}

if tc.StartedDefaultTestTenant() {
systemDB.Exec(t, `ALTER TENANT [$1] GRANT CAPABILITY can_admin_relocate_range=true`, serverutils.TestTenantID().ToUint64())
// Give 1,000,000 upfront tokens to the tenant, and keep the tokens per
// second rate to the default value of 10,000. This helps avoid throttling
// in the tests.
Expand Down
12 changes: 4 additions & 8 deletions pkg/server/authserver/authentication_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,15 +94,11 @@ func TestSSLEnforcement(t *testing.T) {
})
defer srv.Stopper().Stop(ctx)

if srv.TenantController().StartedDefaultTestTenant() {
if srv.DeploymentMode().IsExternal() {
// Enable access to the nodes endpoint for the test tenant.
_, err := srv.SystemLayer().SQLConn(t).Exec(
`ALTER TENANT [$1] GRANT CAPABILITY can_view_node_info=true`, serverutils.TestTenantID().ToUint64())
require.NoError(t, err)

serverutils.WaitForTenantCapabilities(t, srv, serverutils.TestTenantID(), map[tenantcapabilities.ID]string{
tenantcapabilities.CanViewNodeInfo: "true",
}, "")
require.NoError(t, srv.GrantTenantCapabilities(
ctx, serverutils.TestTenantID(),
map[tenantcapabilities.ID]string{tenantcapabilities.CanViewNodeInfo: "true"}))
}

s := srv.ApplicationLayer()
Expand Down
15 changes: 6 additions & 9 deletions pkg/server/storage_api/health_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,6 @@ import (
"github.com/cockroachdb/cockroach/pkg/util/log"
"github.com/cockroachdb/errors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)

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

ctx := context.Background()

// The liveness endpoint needs a special tenant capability.
if tc.Server(0).TenantController().StartedDefaultTestTenant() {
if tc.DefaultTenantDeploymentMode().IsExternal() {
// Enable access to the nodes endpoint for the test tenant.
_, err := tc.SystemLayer(0).SQLConn(t).Exec(
`ALTER TENANT [$1] GRANT CAPABILITY can_view_node_info=true`, serverutils.TestTenantID().ToUint64())
require.NoError(t, err)

tc.WaitForTenantCapabilities(t, serverutils.TestTenantID(), map[tenantcapabilities.ID]string{
tenantcapabilities.CanViewNodeInfo: "true",
})
tc.GrantTenantCapabilities(
ctx, t, serverutils.TestTenantID(),
map[tenantcapabilities.ID]string{tenantcapabilities.CanViewNodeInfo: "true"})
}

ts := tc.Server(0).ApplicationLayer()
Expand Down
17 changes: 5 additions & 12 deletions pkg/server/storage_api/network_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@ package storage_api_test

import (
"context"
"fmt"
"testing"

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

s0 := testCluster.Server(0)

if s0.TenantController().StartedDefaultTestTenant() {
_, err := s0.SystemLayer().SQLConn(t).Exec(
`ALTER TENANT [$1] GRANT CAPABILITY can_debug_process=true`,
serverutils.TestTenantID().ToUint64(),
)
require.NoError(t, err)

serverutils.WaitForTenantCapabilities(t, s0, serverutils.TestTenantID(), map[tenantcapabilities.ID]string{
tenantcapabilities.CanDebugProcess: "true",
}, "")
if s0.DeploymentMode().IsExternal() {
testCluster.GrantTenantCapabilities(
ctx, t, serverutils.TestTenantID(),
map[tenantcapabilities.ID]string{tenantcapabilities.CanDebugProcess: "true"})
}

ts := s0.ApplicationLayer()
Expand Down Expand Up @@ -71,7 +64,7 @@ func TestNetworkConnectivity(t *testing.T) {
return err
}
require.Equal(t, len(resp.Connections), numNodes-1)
fmt.Printf("got status: %s", resp.Connections[s0.StorageLayer().NodeID()].Peers[stoppedNodeID].Status.String())
t.Logf("got status: %s\n", resp.Connections[s0.StorageLayer().NodeID()].Peers[stoppedNodeID].Status.String())
if resp.Connections[s0.StorageLayer().NodeID()].Peers[stoppedNodeID].Status != serverpb.NetworkConnectivityResponse_ERROR {
return errors.New("waiting for connection state to be changed.")
}
Expand Down
12 changes: 4 additions & 8 deletions pkg/server/storage_api/nodes_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -221,15 +221,11 @@ func TestNodesGRPCResponse(t *testing.T) {
})
defer srv.Stopper().Stop(ctx)

if srv.TenantController().StartedDefaultTestTenant() {
if srv.DeploymentMode().IsExternal() {
// Enable access to the nodes endpoint for the test tenant.
_, err := srv.SystemLayer().SQLConn(t).Exec(
`ALTER TENANT [$1] GRANT CAPABILITY can_view_node_info=true`, serverutils.TestTenantID().ToUint64())
require.NoError(t, err)

serverutils.WaitForTenantCapabilities(t, srv, serverutils.TestTenantID(), map[tenantcapabilities.ID]string{
tenantcapabilities.CanViewNodeInfo: "true",
}, "")
require.NoError(t, srv.GrantTenantCapabilities(
ctx, serverutils.TestTenantID(),
map[tenantcapabilities.ID]string{tenantcapabilities.CanViewNodeInfo: "true"}))
}

var request serverpb.NodesRequest
Expand Down
24 changes: 24 additions & 0 deletions pkg/server/testserver.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1265,6 +1265,30 @@ func (t *testTenant) DeploymentMode() serverutils.DeploymentMode {
return t.deploymentMode
}

// GrantTenantCapabilities is part of the serverutils.TenantControlInterface.
func (ts *testServer) GrantTenantCapabilities(
ctx context.Context, tenID roachpb.TenantID, targetCaps map[tenantcapabilities.ID]string,
) error {
conn, err := ts.SQLConnE()
if err != nil {
return err
}

var parts []string
for k, v := range targetCaps {
parts = append(parts, k.String()+"="+v)
}
capabilities := strings.Join(parts, ",")

if _, err = conn.Exec(fmt.Sprintf(`ALTER TENANT [$1] GRANT CAPABILITY %s`, capabilities),
tenID.ToUint64(),
); err != nil {
return err
}

return ts.WaitForTenantCapabilities(ctx, tenID, targetCaps, "")
}

// WaitForTenantCapabilities is part of the serverutils.TenantControlInterface.
func (ts *testServer) WaitForTenantCapabilities(
ctx context.Context,
Expand Down
11 changes: 4 additions & 7 deletions pkg/sql/colflow/draining_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,13 +54,10 @@ func TestDrainingAfterRemoteError(t *testing.T) {
tc := testcluster.StartTestCluster(t, 2 /* nodes */, args)
defer tc.Stopper().Stop(ctx)

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

// Create two tables, one with small values, and another with large rows.
Expand Down
8 changes: 4 additions & 4 deletions pkg/sql/distsql_physical_planner_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ import (
"github.com/cockroachdb/cockroach/pkg/kv"
"github.com/cockroachdb/cockroach/pkg/kv/kvclient/kvcoord"
"github.com/cockroachdb/cockroach/pkg/kv/kvclient/rangecache"
"github.com/cockroachdb/cockroach/pkg/multitenant/tenantcapabilities"
"github.com/cockroachdb/cockroach/pkg/roachpb"
"github.com/cockroachdb/cockroach/pkg/rpc"
"github.com/cockroachdb/cockroach/pkg/server/serverpb"
Expand Down Expand Up @@ -449,10 +450,9 @@ func TestDistSQLUnavailableHosts(t *testing.T) {
}

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

// Connect to node 1 (gateway node)
Expand Down
11 changes: 4 additions & 7 deletions pkg/sql/execstats/traceanalyzer_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,13 +78,10 @@ func TestTraceAnalyzer(t *testing.T) {

const gatewayNode = 0
srv, s := tc.Server(gatewayNode), tc.ApplicationLayer(gatewayNode)
if srv.TenantController().StartedDefaultTestTenant() {
systemSqlDB := srv.SystemLayer().SQLConn(t, serverutils.DBName("system"))
_, err := systemSqlDB.Exec(`ALTER TENANT [$1] GRANT CAPABILITY can_admin_relocate_range=true`, serverutils.TestTenantID().ToUint64())
require.NoError(t, err)
serverutils.WaitForTenantCapabilities(t, srv, serverutils.TestTenantID(), map[tenantcapabilities.ID]string{
tenantcapabilities.CanAdminRelocateRange: "true",
}, "")
if srv.DeploymentMode().IsExternal() {
require.NoError(t, srv.GrantTenantCapabilities(
ctx, serverutils.TestTenantID(),
map[tenantcapabilities.ID]string{tenantcapabilities.CanAdminRelocateRange: "true"}))
}
db := s.SQLConn(t)
sqlDB := sqlutils.MakeSQLRunner(db)
Expand Down
11 changes: 4 additions & 7 deletions pkg/sql/flowinfra/cluster_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -594,13 +594,10 @@ func TestEvalCtxTxnOnRemoteNodes(t *testing.T) {
})
defer tc.Stopper().Stop(ctx)

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

db := tc.ServerConn(0)
Expand Down
11 changes: 4 additions & 7 deletions pkg/sql/importer/exportcsv_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -663,14 +663,11 @@ func TestProcessorEncountersUncertaintyError(t *testing.T) {
ctx := context.Background()
defer tc.Stopper().Stop(ctx)

if tc.StartedDefaultTestTenant() {
systemSqlDB := tc.Server(0).SystemLayer().SQLConn(t, serverutils.DBName("system"))
_, err := systemSqlDB.Exec(`ALTER TENANT [$1] GRANT CAPABILITY can_admin_relocate_range=true`, serverutils.TestTenantID().ToUint64())
require.NoError(t, err)
serverutils.WaitForTenantCapabilities(t, tc.Server(0), serverutils.TestTenantID(), map[tenantcapabilities.ID]string{
tenantcapabilities.CanAdminRelocateRange: "true",
}, "")
if tc.DefaultTenantDeploymentMode().IsExternal() {
tc.GrantTenantCapabilities(ctx, t, serverutils.TestTenantID(),
map[tenantcapabilities.ID]string{tenantcapabilities.CanAdminRelocateRange: "true"})
}

origDB0 := tc.ServerConn(0)

sqlutils.CreateTable(t, origDB0, "t",
Expand Down
8 changes: 1 addition & 7 deletions pkg/sql/logictest/logic.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1728,12 +1728,6 @@ func (t *logicTest) newCluster(

capabilities := toa.capabilities
if len(capabilities) > 0 {
for name, value := range capabilities {
query := fmt.Sprintf("ALTER TENANT [$1] GRANT CAPABILITY %s = $2", name)
if _, err := conn.Exec(query, tenantID.ToUint64(), value); err != nil {
t.Fatal(err)
}
}
capabilityMap := make(map[tenantcapabilities.ID]string, len(capabilities))
for k, v := range capabilities {
capability, ok := tenantcapabilities.FromName(k)
Expand All @@ -1742,7 +1736,7 @@ func (t *logicTest) newCluster(
}
capabilityMap[capability.ID()] = v
}
t.cluster.WaitForTenantCapabilities(t.t(), tenantID, capabilityMap)
t.cluster.GrantTenantCapabilities(context.Background(), t.t(), tenantID, capabilityMap)
}
}

Expand Down
10 changes: 3 additions & 7 deletions pkg/sql/opt/exec/explain/output_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -163,13 +163,9 @@ func TestCPUTimeEndToEnd(t *testing.T) {
ctx := context.Background()
defer tc.Stopper().Stop(ctx)

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

db := sqlutils.MakeSQLRunner(tc.Conns[0])
Expand Down
14 changes: 14 additions & 0 deletions pkg/testutils/serverutils/api.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -141,6 +141,10 @@ const (
ExternalProcess
)

func (d DeploymentMode) IsExternal() bool {
return d == ExternalProcess
}

// ApplicationLayerInterface defines accessors to the application
// layer of a test server. Tests written against this interface are
// effectively agnostic to whether they use a virtual cluster or not.
Expand Down Expand Up @@ -506,6 +510,16 @@ type TenantControlInterface interface {
// if the tenant record exists in KV.
WaitForTenantReadiness(ctx context.Context, tenantID roachpb.TenantID) error

// GrantTenantCapabilities grants a capability to a tenant and waits until the
// in-memory cache reflects the change.
//
// Note: There is no need to call WaitForTenantCapabilities separately.
GrantTenantCapabilities(
context.Context,
roachpb.TenantID,
map[tenantcapabilities.ID]string,
) error

// WaitForTenantCapabilities waits until the in-RAM cache of
// tenant capabilities has been populated for the given tenant ID
// with the expected target capability values.
Expand Down
16 changes: 16 additions & 0 deletions pkg/testutils/serverutils/test_cluster_shim.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -250,6 +250,11 @@ type TestClusterInterface interface {
// cluster.
StorageLayer(idx int) StorageLayerInterface

// DefaultTenantDeploymentMode returns the deployment mode of the default
// server or tenant, which can be single-tenant (system-only),
// shared-process, or external-process.
DefaultTenantDeploymentMode() DeploymentMode

// SplitTable splits a range in the table, creates a replica for the right
// side of the split on TargetNodeIdx, and moves the lease for the right
// side of the split to TargetNodeIdx for each SplitPoint. This forces the
Expand All @@ -259,6 +264,17 @@ type TestClusterInterface interface {
// are indeed distributed as intended.
SplitTable(t TestFataler, desc catalog.TableDescriptor, sps []SplitPoint)

// GrantTenantCapabilities grants a capability to a tenant and waits until all
// servers have the in-memory cache reflects the change.
//
// Note: There is no need to call WaitForTenantCapabilities separately.
GrantTenantCapabilities(
context.Context,
TestFataler,
roachpb.TenantID,
map[tenantcapabilities.ID]string,
)

// WaitForTenantCapabilities waits until all servers have the specified
// tenant capabilities for the specified tenant ID.
// Only boolean capabilities are currently supported as we wait for the
Expand Down
Loading