From 57c267e90ee1ecebfb6b3826984a7b956b28cf8d Mon Sep 17 00:00:00 2001
From: Jordan Lewis <jordanthelewis@gmail.com>
Date: Mon, 5 Mar 2018 17:30:37 -0500
Subject: [PATCH] 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 | 37 ++++++++++++++++++++++++++++++++++++
 1 file changed, 37 insertions(+)

diff --git a/pkg/cmd/workload/fixtures.go b/pkg/cmd/workload/fixtures.go
index 1d261fe7c887..1aeff005c236 100644
--- a/pkg/cmd/workload/fixtures.go
+++ b/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
 
@@ -150,6 +154,25 @@ func fixturesList(_ *cobra.Command, _ []string) error {
 	return nil
 }
 
+type filteringGenerator struct {
+	gen    workload.Generator
+	filter map[string]struct{}
+}
+
+func (f filteringGenerator) Meta() workload.Meta {
+	return f.gen.Meta()
+}
+
+func (f filteringGenerator) Tables() []workload.Table {
+	ret := make([]workload.Table, 0)
+	for _, t := range f.gen.Tables() {
+		if _, ok := f.filter[t.Name]; ok {
+			ret = append(ret, t)
+		}
+	}
+	return ret
+}
+
 func fixturesMake(gen workload.Generator, urls []string, dbName string) error {
 	ctx := context.Background()
 	gcs, err := getStorage(ctx)
@@ -162,6 +185,20 @@ 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 = filteringGenerator{
+			gen:    gen,
+			filter: filter,
+		}
+	}
 	fixture, err := workloadccl.MakeFixture(ctx, sqlDB, gcs, config(), gen)
 	if err != nil {
 		return err