Skip to content

Commit

Permalink
sql: update TestMultiTenantAdminFunction to not assume a default cap …
Browse files Browse the repository at this point in the history 
…state

Release note: None

Co-authored-by: Evan Wall <[email protected]>
  • Loading branch information
@knz @ecwall
knz and ecwall committed Mar 14, 2023
1 parent 43f7cde commit bacfb0f
Show file tree
Hide file tree
Showing 3 changed files with 52 additions and 43 deletions.
77 changes: 44 additions & 33 deletions pkg/sql/multitenant_admin_function_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -238,9 +238,14 @@ type testCase struct {
queryClusterSetting *settings.BoolSetting
// Used for tests that have a capability prereq
// (eq SPLIT AT is required for UNSPLIT AT).
setupCapability tenantcapabilities.CapabilityID
setupCapability capValue
// Capability required for secondary tenant query.
queryCapability tenantcapabilities.CapabilityID
queryCapability capValue
}

type capValue struct {
cap tenantcapabilities.CapabilityID
value string
}

func (tc testCase) runTest(
Expand Down Expand Up @@ -303,32 +308,30 @@ func (tc testCase) runTest(
var waitForTenantCapabilitiesFns []func()
setCapabilities := func(
tenantID roachpb.TenantID,
capabilityIDs ...tenantcapabilities.CapabilityID,
capVals ...capValue,
) {
// Filter out empty capabilities.
var caps []tenantcapabilities.CapabilityID
for _, capabilityID := range capabilityIDs {
if capabilityID == 0 {
var caps []capValue
for _, capVal := range capVals {
if capVal.cap == 0 {
// This can happen if e.g. setupCapability / queryCapability
// are unset in a test.
continue
}
caps = append(caps, capabilityID)
caps = append(caps, capVal)
}
capabilityIDs = caps
if len(capabilityIDs) > 0 {
var builder strings.Builder
for i, capabilityID := range capabilityIDs {
if i > 0 {
builder.WriteString(", ")
}
builder.WriteString(capabilityID.String())
capVals = caps
if len(capVals) > 0 {
expected := map[tenantcapabilities.CapabilityID]string{}
for _, capVal := range capVals {
query := fmt.Sprintf("ALTER TENANT [$1] GRANT CAPABILITY %s = %s", capVal.cap.String(), capVal.value)
_, err := systemDB.ExecContext(ctx, query, tenantID.ToUint64())
require.NoError(t, err, query)
expected[capVal.cap] = capVal.value
}
query := fmt.Sprintf("ALTER TENANT [$1] GRANT CAPABILITY %s", builder.String())
_, err := systemDB.ExecContext(ctx, query, tenantID.ToUint64())
require.NoError(t, err, query)

waitForTenantCapabilitiesFns = append(waitForTenantCapabilitiesFns, func() {
testCluster.WaitForTenantCapabilities(t, tenantID, capabilityIDs...)
testCluster.WaitForTenantCapabilities(t, tenantID, expected)
})
}
}
Expand Down Expand Up @@ -403,6 +406,10 @@ func (te tenantExpected) validate(t *testing.T, rows *gosql.Rows, err error, mes
}
}

func bcap(cap tenantcapabilities.CapabilityID, val bool) capValue {
return capValue{cap: cap, value: fmt.Sprint(val)}
}

// TestMultiTenantAdminFunction tests the "simple" admin functions that do not require complex setup.
func TestMultiTenantAdminFunction(t *testing.T) {
defer leaktest.AfterTest(t)()
Expand Down Expand Up @@ -464,7 +471,8 @@ func TestMultiTenantAdminFunction(t *testing.T) {
errorMessage: `does not have capability "can_admin_split"`,
},
queryClusterSetting: sql.SecondaryTenantSplitAtEnabled,
queryCapability: tenantcapabilities.CanAdminSplit,
setupCapability: bcap(tenantcapabilities.CanAdminSplit, false),
queryCapability: bcap(tenantcapabilities.CanAdminSplit, true),
},
{
desc: "ALTER INDEX x SPLIT AT",
Expand All @@ -483,7 +491,8 @@ func TestMultiTenantAdminFunction(t *testing.T) {
errorMessage: `does not have capability "can_admin_split"`,
},
queryClusterSetting: sql.SecondaryTenantSplitAtEnabled,
queryCapability: tenantcapabilities.CanAdminSplit,
setupCapability: bcap(tenantcapabilities.CanAdminSplit, false),
queryCapability: bcap(tenantcapabilities.CanAdminSplit, true),
},
{
desc: "ALTER TABLE x UNSPLIT AT",
Expand All @@ -499,8 +508,8 @@ func TestMultiTenantAdminFunction(t *testing.T) {
errorMessage: `does not have capability "can_admin_unsplit"`,
},
setupClusterSetting: sql.SecondaryTenantSplitAtEnabled,
setupCapability: tenantcapabilities.CanAdminSplit,
queryCapability: tenantcapabilities.CanAdminUnsplit,
setupCapability: bcap(tenantcapabilities.CanAdminSplit, true),
queryCapability: bcap(tenantcapabilities.CanAdminUnsplit, true),
},
{
desc: "ALTER INDEX x UNSPLIT AT",
Expand All @@ -519,8 +528,8 @@ func TestMultiTenantAdminFunction(t *testing.T) {
errorMessage: `does not have capability "can_admin_unsplit"`,
},
setupClusterSetting: sql.SecondaryTenantSplitAtEnabled,
setupCapability: tenantcapabilities.CanAdminSplit,
queryCapability: tenantcapabilities.CanAdminUnsplit,
setupCapability: bcap(tenantcapabilities.CanAdminSplit, true),
queryCapability: bcap(tenantcapabilities.CanAdminUnsplit, true),
},
{
desc: "ALTER TABLE x UNSPLIT ALL",
Expand All @@ -536,8 +545,8 @@ func TestMultiTenantAdminFunction(t *testing.T) {
errorMessage: `does not have capability "can_admin_unsplit"`,
},
setupClusterSetting: sql.SecondaryTenantSplitAtEnabled,
setupCapability: tenantcapabilities.CanAdminSplit,
queryCapability: tenantcapabilities.CanAdminUnsplit,
setupCapability: bcap(tenantcapabilities.CanAdminSplit, true),
queryCapability: bcap(tenantcapabilities.CanAdminUnsplit, true),
},
{
desc: "ALTER INDEX x UNSPLIT ALL",
Expand All @@ -556,8 +565,8 @@ func TestMultiTenantAdminFunction(t *testing.T) {
errorMessage: `does not have capability "can_admin_unsplit"`,
},
setupClusterSetting: sql.SecondaryTenantSplitAtEnabled,
setupCapability: tenantcapabilities.CanAdminSplit,
queryCapability: tenantcapabilities.CanAdminUnsplit,
setupCapability: bcap(tenantcapabilities.CanAdminSplit, true),
queryCapability: bcap(tenantcapabilities.CanAdminUnsplit, true),
},
{
desc: "ALTER TABLE x SCATTER",
Expand All @@ -575,7 +584,8 @@ func TestMultiTenantAdminFunction(t *testing.T) {
errorMessage: `does not have capability "can_admin_scatter"`,
},
queryClusterSetting: sql.SecondaryTenantScatterEnabled,
queryCapability: tenantcapabilities.CanAdminScatter,
setupCapability: bcap(tenantcapabilities.CanAdminScatter, false),
queryCapability: bcap(tenantcapabilities.CanAdminScatter, true),
},
{
desc: "ALTER INDEX x SCATTER",
Expand All @@ -594,7 +604,8 @@ func TestMultiTenantAdminFunction(t *testing.T) {
errorMessage: `does not have capability "can_admin_scatter"`,
},
queryClusterSetting: sql.SecondaryTenantScatterEnabled,
queryCapability: tenantcapabilities.CanAdminScatter,
setupCapability: bcap(tenantcapabilities.CanAdminScatter, false),
queryCapability: bcap(tenantcapabilities.CanAdminScatter, true),
},
}

Expand Down Expand Up @@ -648,8 +659,8 @@ func TestTruncateTable(t *testing.T) {
errorMessage: `does not have capability "can_admin_scatter"`,
},
setupClusterSetting: sql.SecondaryTenantSplitAtEnabled,
setupCapability: tenantcapabilities.CanAdminSplit,
queryCapability: tenantcapabilities.CanAdminScatter,
setupCapability: bcap(tenantcapabilities.CanAdminSplit, true),
queryCapability: bcap(tenantcapabilities.CanAdminScatter, true),
}
tc.runTest(
t,
Expand Down
2 changes: 1 addition & 1 deletion pkg/testutils/serverutils/test_cluster_shim.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -236,7 +236,7 @@ type TestClusterInterface interface {
// tenant capabilities for the specified tenant ID.
// Only boolean capabilities are currently supported as we wait for the
// specified capabilities to have a "true" value.
WaitForTenantCapabilities(*testing.T, roachpb.TenantID, ...tenantcapabilities.CapabilityID)
WaitForTenantCapabilities(*testing.T, roachpb.TenantID, map[tenantcapabilities.CapabilityID]string)
}

// SplitPoint describes a split point that is passed to SplitTable.
Expand Down
16 changes: 7 additions & 9 deletions pkg/testutils/testcluster/testcluster.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1778,28 +1778,26 @@ func (tc *TestCluster) SplitTable(

// WaitForTenantCapabilities implements TestClusterInterface.
func (tc *TestCluster) WaitForTenantCapabilities(
t *testing.T, tenID roachpb.TenantID, capIDs ...tenantcapabilities.CapabilityID,
t *testing.T, tenID roachpb.TenantID, targetCaps map[tenantcapabilities.CapabilityID]string,
) {
for i, ts := range tc.Servers {
testutils.SucceedsSoon(t, func() error {
if tenID.IsSystem() {
return nil
}

if len(capIDs) > 0 {
if len(targetCaps) > 0 {
missingCapabilityError := func(capID tenantcapabilities.CapabilityID) error {
return errors.Newf("server=%d tenant %s does not have capability %q", i, tenID, capID)
return errors.Newf("server=%d tenant %s cap %q not at expected value", i, tenID, capID)
}
capabilities, found := ts.Server.TenantCapabilitiesReader().GetCapabilities(tenID)
if !found {
return missingCapabilityError(capIDs[0])
return errors.Newf("capabilities not ready for tenant %s", tenID)
}

for _, capID := range capIDs {
if capID.CapabilityType() != tenantcapabilities.Bool {
return errors.AssertionFailedf("WaitForTenantCapabilities only supports boolean capabilities")
}
if !capabilities.GetBool(capID) {
for capID, expectedValue := range targetCaps {
curVal := capabilities.Cap(capID).Get().String()
if curVal != expectedValue {
return missingCapabilityError(capID)
}
}
Expand Down

0 comments on commit bacfb0f

Please sign in to comment.