Merge #133794

133794: roachtest: add db and table flag to run-operation. r=vidit-bhat a=shailendra-patel Operations like `add_column`/`add_index` select database and table randomly to run the operation. This change allows us to specify the database name and table via roachtest flag. Epic: none Release note: None Co-authored-by: Shailendra Patel <[email protected]>
cockroachdb · Nov 1, 2024 · f9918d8 · f9918d8
2 parents 482ba56 + 2c6a937
commit f9918d8
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 1 deletion.
diff --git a/pkg/cmd/roachtest/operations/utils.go b/pkg/cmd/roachtest/operations/utils.go
@@ -15,6 +15,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/cluster"
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/operation"
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/option"
+	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/roachtestflags"
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/roachtestutil"
 	"github.com/cockroachdb/cockroach/pkg/util/randutil"
 	"github.com/cockroachdb/cockroach/pkg/util/retry"
@@ -24,11 +25,16 @@ import (
 // may not be mutable and should be excluded by most operations.
 var systemDBs = []string{"system", "information_schema", "crdb_internal", "defaultdb", "postgres"}
 
-// pickRandomDB picks a random DB that isn't one of `excludeDBs` on the
+// pickRandomDB returns roachtestflags.DBName if not empty.
+// Otherwise, picks a random DB that isn't one of `excludeDBs` on the
 // target cluster connected to by `conn`.
 func pickRandomDB(
 	ctx context.Context, o operation.Operation, conn *gosql.DB, excludeDBs []string,
 ) string {
+	if roachtestflags.DBName != "" {
+		return roachtestflags.DBName
+	}
+
 	rng, _ := randutil.NewPseudoRand()
 
 	// Pick a random table.
@@ -60,9 +66,15 @@ func pickRandomDB(
 	return dbNames[rng.Intn(len(dbNames))]
 }
 
+// pickRandomTable returns roachtestflags.TableName if not empty.
+// Otherwise, picks a random table from given database.
 func pickRandomTable(
 	ctx context.Context, o operation.Operation, conn *gosql.DB, dbName string,
 ) string {
+	if roachtestflags.TableName != "" {
+		return roachtestflags.TableName
+	}
+
 	rng, _ := randutil.NewPseudoRand()
 
 	// Pick a random table.

diff --git a/pkg/cmd/roachtest/roachtestflags/flags.go b/pkg/cmd/roachtest/roachtestflags/flags.go
@@ -343,6 +343,19 @@ var (
 		Usage: `A comma-separated list of tags to attach to telemetry data (e.g., key1:val1,key2:val2).`,
 	})
 
+	DBName string = ""
+	_             = registerRunOpsFlag(&DBName, FlagInfo{
+		Name: "db",
+		Usage: "Specify the name of the database to run the operation against (e.g., tpcc, tpch). If the given database " +
+			"does not exist, the operation fails. If not provided, a random database is selected, excluding system-created databases",
+	})
+
+	TableName string = ""
+	_                = registerRunOpsFlag(&TableName, FlagInfo{
+		Name: "db-table",
+		Usage: "Specifies the name of the database table to run the operation against, using the database provided in the --db flag. " +
+			"If the table does not exist, the operation fails. If not provided, a random table is selected.",
+	})
 	SideEyeApiToken string = ""
 	_                      = registerRunFlag(&SideEyeApiToken, FlagInfo{
 		Name: "side-eye-token",