workload: fixture create can make table subset

The new --only-tables flag allows specifying a subset of the tables in a
workload for creation. This is useful for benchmarking the workload
program itself, as well as creating partial backups of very large
datasets.

Release note: None

pkg/cmd/workload/fixtures.go

@@ -70,6 +70,10 @@ var fixturesMakeCSVServerURL = fixturesMakeCmd.PersistentFlags().String(
 	`csv-server`, ``,
 	`Skip saving CSVs to cloud storage, instead get them from a 'csv-server' running at this url`)
 
+var fixturesMakeOnlyTable = fixturesMakeCmd.PersistentFlags().String(
+	`only-tables`, ``,
+	`Only load the tables with the given comma-separated names`)
+
 // gcs-bucket-override and gcs-prefix-override are exposed for testing.
 var gcsBucketOverride, gcsPrefixOverride *string
 
@@ -162,6 +166,17 @@ func fixturesMake(gen workload.Generator, urls []string, dbName string) error {
 	if err != nil {
 		return err
 	}
+	if *fixturesMakeOnlyTable != "" {
+		tableNames := strings.Split(*fixturesMakeOnlyTable, ",")
+		if len(tableNames) == 0 {
+			return errors.New("no table names specified")
+		}
+		filter := make(map[string]struct{}, len(tableNames))
+		for _, tableName := range tableNames {
+			filter[tableName] = struct{}{}
+		}
+		gen = workload.NewFilteringGenerator(gen, filter)
+	}
 	fixture, err := workloadccl.MakeFixture(ctx, sqlDB, gcs, config(), gen)
 	if err != nil {
 		return err

pkg/workload/workload.go

@@ -146,6 +146,32 @@ type QueryLoad struct {
 	ResultHist string
 }
 
+type filteringGenerator struct {
+	gen    Generator
+	filter map[string]struct{}
+}
+
+func NewFilteringGenerator(gen Generator, filter map[string]struct{}) Generator {
+	return filteringGenerator{
+		gen:    gen,
+		filter: filter,
+	}
+}
+
+func (f filteringGenerator) Meta() Meta {
+	return f.gen.Meta()
+}
+
+func (f filteringGenerator) Tables() []Table {
+	ret := make([]Table, 0)
+	for _, t := range f.gen.Tables() {
+		if _, ok := f.filter[t.Name]; ok {
+			ret = append(ret, t)
+		}
+	}
+	return ret
+}
+
 var registered = make(map[string]Meta)
 
 // Register is a hook for init-time registration of Generator implementations.