Skip to content

Commit 56f440c

Browse files
rafissrafiss
committed
roachtest: add operations to GRANT roles and SHOW TABLES
These are being added since we had received a bug report of the SHOW TABLES command not working if role memberships have been added. We have fixed the underlying bug, but adding the operations here can still improve our test coverage. Release note: None
1 parent 7a28a7b commit 56f440c

File tree

4 files changed

+198
-0
lines changed

4 files changed

+198
-0
lines changed

‎pkg/cmd/roachtest/operations/BUILD.bazel‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ go_library(
1313
"debug_zip.go",
1414
"disk_stall.go",
1515
"grant_revoke_all.go",
16+
"grant_role.go",
1617
"license_throttle.go",
1718
"manual_compaction.go",
1819
"network_partition.go",
@@ -21,6 +22,7 @@ go_library(
2122
"register.go",
2223
"resize.go",
2324
"session_vars.go",
25+
"show_tables.go",
2426
"sql_objects.go",
2527
],
2628
importpath = "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/operations",

‎pkg/cmd/roachtest/operations/grant_role.go‎

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
// Copyright 2024 The Cockroach Authors.
2+
//
3+
// Use of this software is governed by the CockroachDB Software License
4+
// included in the /LICENSE file.
5+
6+
package operations
7+
8+
import (
9+
"context"
10+
"fmt"
11+
"time"
12+
13+
"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/cluster"
14+
"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/operation"
15+
"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/option"
16+
"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/registry"
17+
"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/roachtestflags"
18+
"github.com/cockroachdb/cockroach/pkg/util/randutil"
19+
)
20+
21+
// This is a schema change operation that does the following:
22+
// Operation Run:
23+
// 1. Create two roles with randomly generated names.
24+
// 2. GRANT one role to the other role.
25+
// Operation Cleanup:
26+
// 1. REVOKE the granted role.
27+
// 2. Drop both roles.
28+
29+
type revokeRole struct {
30+
grantedRole, granteeRole string
31+
waitDuration time.Duration
32+
}
33+
34+
func (cl *revokeRole) Cleanup(ctx context.Context, o operation.Operation, c cluster.Cluster) {
35+
o.Status(fmt.Sprintf("Scheduling cleanup to happen after %s", cl.waitDuration))
36+
37+
// Start a goroutine to handle the wait and cleanup. Since this runs in the
38+
// background, we also create a new context so that the background goroutine
39+
// isn't aborted by the parent context.
40+
go func() {
41+
newCtx := context.Background()
42+
if deadline, ok := ctx.Deadline(); ok {
43+
var cancel context.CancelFunc
44+
newCtx, cancel = context.WithDeadline(newCtx, deadline.Add(cl.waitDuration))
45+
defer cancel()
46+
}
47+
ctx = newCtx
48+
49+
// Wait for the specified duration before performing cleanup.
50+
time.Sleep(cl.waitDuration)
51+
o.Status(fmt.Sprintf("Wait time of %s elapsed, performing cleanup", cl.waitDuration))
52+
53+
conn := c.Conn(ctx, o.L(), 1, option.VirtualClusterName(roachtestflags.VirtualCluster))
54+
defer func() { _ = conn.Close() }()
55+
56+
o.Status(fmt.Sprintf("Revoking role %s from role %s", cl.grantedRole, cl.granteeRole))
57+
_, err := conn.ExecContext(ctx, fmt.Sprintf("REVOKE %s FROM %s", cl.grantedRole, cl.granteeRole))
58+
if err != nil {
59+
o.Fatal(err)
60+
}
61+
o.Status(fmt.Sprintf("Revoked role %s from role %s", cl.grantedRole, cl.granteeRole))
62+
63+
// Drop both roles
64+
for _, role := range []string{cl.granteeRole, cl.grantedRole} {
65+
o.Status(fmt.Sprintf("Dropping role %s", role))
66+
_, err = conn.ExecContext(ctx, fmt.Sprintf("DROP ROLE %s", role))
67+
if err != nil {
68+
o.Fatal(err)
69+
}
70+
o.Status(fmt.Sprintf("Dropped role %s", role))
71+
}
72+
}()
73+
}
74+
75+
func runGrantRole(
76+
ctx context.Context, o operation.Operation, c cluster.Cluster,
77+
) registry.OperationCleanup {
78+
conn := c.Conn(ctx, o.L(), 1, option.VirtualClusterName(roachtestflags.VirtualCluster))
79+
defer func() { _ = conn.Close() }()
80+
81+
rng, _ := randutil.NewTestRand()
82+
// Create random role names with a prefix to avoid conflicts.
83+
grantedRole := fmt.Sprintf("roachprod_ops_granted_role_%s", randutil.RandString(rng, 10, randutil.PrintableKeyAlphabet))
84+
granteeRole := fmt.Sprintf("roachprod_ops_grantee_role_%s", randutil.RandString(rng, 10, randutil.PrintableKeyAlphabet))
85+
86+
// Create both roles.
87+
for _, role := range []string{grantedRole, granteeRole} {
88+
o.Status(fmt.Sprintf("Creating role %s", role))
89+
_, err := conn.ExecContext(ctx, fmt.Sprintf("CREATE ROLE %s", role))
90+
if err != nil {
91+
o.Fatal(err)
92+
}
93+
o.Status(fmt.Sprintf("Created role %s", role))
94+
}
95+
96+
// Grant one role to the other.
97+
o.Status(fmt.Sprintf("Granting role %s to role %s", grantedRole, granteeRole))
98+
_, err := conn.ExecContext(ctx, fmt.Sprintf("GRANT %s TO %s", grantedRole, granteeRole))
99+
if err != nil {
100+
o.Fatal(err)
101+
}
102+
o.Status(fmt.Sprintf("Granted role %s to role %s", grantedRole, granteeRole))
103+
104+
// Return the cleanup struct with wait duration.
105+
waitDuration := 1 * time.Hour
106+
return &revokeRole{
107+
grantedRole: grantedRole,
108+
granteeRole: granteeRole,
109+
waitDuration: waitDuration,
110+
}
111+
}
112+
113+
func registerGrantRole(r registry.Registry) {
114+
r.AddOperation(registry.OperationSpec{
115+
Name: "grant-role",
116+
Owner: registry.OwnerSQLFoundations,
117+
Timeout: 10 * time.Minute,
118+
CompatibleClouds: registry.AllClouds,
119+
CanRunConcurrently: registry.OperationCanRunConcurrently,
120+
Dependencies: []registry.OperationDependency{registry.OperationRequiresNodes},
121+
Run: runGrantRole,
122+
})
123+
}

‎pkg/cmd/roachtest/operations/register.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@ func RegisterOperations(r registry.Registry) {
1616
registerAddDatabase(r)
1717
registerAddIndex(r)
1818
registerAddRLSPolicy(r)
19+
registerShowTables(r)
20+
registerGrantRole(r)
1921
registerGrantRevoke(r)
2022
registerNetworkPartition(r)
2123
registerDiskStall(r)

‎pkg/cmd/roachtest/operations/show_tables.go‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
// Copyright 2025 The Cockroach Authors.
2+
//
3+
// Use of this software is governed by the CockroachDB Software License
4+
// included in the /LICENSE file.
5+
6+
package operations
7+
8+
import (
9+
"context"
10+
"fmt"
11+
"time"
12+
13+
"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/cluster"
14+
"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/operation"
15+
"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/operations/helpers"
16+
"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/option"
17+
"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/registry"
18+
"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/roachtestflags"
19+
)
20+
21+
// This operation runs a SHOW TABLES command on a randomly selected database.
22+
// It does not require cleanup as it only reads data without making changes.
23+
24+
func runShowTables(
25+
ctx context.Context, o operation.Operation, c cluster.Cluster,
26+
) registry.OperationCleanup {
27+
conn := c.Conn(ctx, o.L(), 1, option.VirtualClusterName(roachtestflags.VirtualCluster))
28+
defer func() { _ = conn.Close() }()
29+
30+
// Pick a random database to run SHOW TABLES on
31+
dbName := helpers.PickRandomDB(ctx, o, conn, helpers.SystemDBs)
32+
33+
o.Status(fmt.Sprintf("Running SHOW TABLES command on database %s", dbName))
34+
35+
// Execute the SHOW TABLES command
36+
rows, err := conn.QueryContext(ctx, fmt.Sprintf("SHOW TABLES FROM %s", dbName))
37+
if err != nil {
38+
o.Fatal(err)
39+
}
40+
defer rows.Close()
41+
42+
// Process and log the results
43+
tableCount := 0
44+
for rows.Next() {
45+
var tableName string
46+
if err := rows.Scan(&tableName); err != nil {
47+
o.Fatal(err)
48+
}
49+
tableCount++
50+
}
51+
if err := rows.Err(); err != nil {
52+
o.Fatal(err)
53+
}
54+
55+
o.Status(fmt.Sprintf("SHOW TABLES completed successfully, found %d tables in database %s", tableCount, dbName))
56+
57+
// No cleanup needed since this is a read-only operation
58+
return nil
59+
}
60+
61+
func registerShowTables(r registry.Registry) {
62+
r.AddOperation(registry.OperationSpec{
63+
Name: "show-tables",
64+
Owner: registry.OwnerSQLFoundations,
65+
Timeout: 5 * time.Minute,
66+
CompatibleClouds: registry.AllClouds,
67+
CanRunConcurrently: registry.OperationCanRunConcurrently,
68+
Dependencies: []registry.OperationDependency{registry.OperationRequiresNodes},
69+
Run: runShowTables,
70+
})
71+
}

0 commit comments

Comments
 (0)