From 0045120833fceb34e1c446ce442da2b3fa803e56 Mon Sep 17 00:00:00 2001
From: Nathan VanBenschoten <nvanbenschoten@gmail.com>
Date: Thu, 25 Jul 2019 16:25:10 -0400
Subject: [PATCH] workloadccl: inject table statistics during fixtures load

See #39103.

Release note: None
---
 pkg/ccl/workloadccl/cliccl/fixtures.go |  2 +-
 pkg/ccl/workloadccl/fixture.go         | 52 ++++++++++++++++++++++----
 pkg/ccl/workloadccl/fixture_test.go    |  2 +-
 3 files changed, 46 insertions(+), 10 deletions(-)

diff --git a/pkg/ccl/workloadccl/cliccl/fixtures.go b/pkg/ccl/workloadccl/cliccl/fixtures.go
index 452122184917..c769ac3f8711 100644
--- a/pkg/ccl/workloadccl/cliccl/fixtures.go
+++ b/pkg/ccl/workloadccl/cliccl/fixtures.go
@@ -282,7 +282,7 @@ func (l restoreDataLoader) InitialDataLoad(
 ) (int64, error) {
 	log.Infof(ctx, "starting restore of %d tables", len(gen.Tables()))
 	start := timeutil.Now()
-	bytes, err := workloadccl.RestoreFixture(ctx, db, l.fixture, l.database)
+	bytes, err := workloadccl.RestoreFixture(ctx, db, l.fixture, l.database, true /* injectStats */)
 	if err != nil {
 		return 0, errors.Wrap(err, `restoring fixture`)
 	}
diff --git a/pkg/ccl/workloadccl/fixture.go b/pkg/ccl/workloadccl/fixture.go
index 6129f45714a7..6d6c9fff2239 100644
--- a/pkg/ccl/workloadccl/fixture.go
+++ b/pkg/ccl/workloadccl/fixture.go
@@ -358,7 +358,7 @@ func ImportFixture(
 	var bytesAtomic int64
 	g := ctxgroup.WithContext(ctx)
 	tables := gen.Tables()
-	if injectStats && len(tables) > 0 && len(tables[0].Stats) > 0 {
+	if injectStats && tablesHaveStats(tables) {
 		// Turn off automatic stats temporarily so we don't trigger stats creation
 		// after the IMPORT. We will inject stats inside importFixtureTable.
 		// TODO(rytaft): It would be better if the automatic statistics code would
@@ -402,12 +402,7 @@ func importFixtureTable(
 	start := timeutil.Now()
 	var buf bytes.Buffer
 	var params []interface{}
-	var qualifiedTableName string
-	if dbName != `` {
-		qualifiedTableName = fmt.Sprintf(`"%s"."%s"`, dbName, table.Name)
-	} else {
-		qualifiedTableName = fmt.Sprintf(`"%s"`, table.Name)
-	}
+	qualifiedTableName := makeQualifiedTableName(dbName, &table)
 	fmt.Fprintf(&buf, `IMPORT TABLE %s %s CSV DATA (`, qualifiedTableName, table.Schema)
 	// Generate $1,...,$N-1, where N is the number of csv paths.
 	for _, path := range paths {
@@ -448,6 +443,17 @@ func importFixtureTable(
 	return tableBytes, nil
 }
 
+// tablesHaveStats returns whether any of the provided tables have associated
+// table statistics to inject.
+func tablesHaveStats(tables []workload.Table) bool {
+	for _, t := range tables {
+		if len(t.Stats) > 0 {
+			return true
+		}
+	}
+	return false
+}
+
 // disableAutoStats disables automatic stats if they are enabled and returns
 // a function to re-enable them later. If automatic stats are already disabled,
 // disableAutoStats does nothing and returns an empty function.
@@ -494,14 +500,33 @@ func injectStatistics(qualifiedTableName string, table *workload.Table, sqlDB *g
 	return err
 }
 
+// makeQualifiedTableName constructs a qualified table name from the specified
+// database name and table.
+func makeQualifiedTableName(dbName string, table *workload.Table) string {
+	if dbName == "" {
+		return fmt.Sprintf(`"%s"`, table.Name)
+	}
+	return fmt.Sprintf(`"%s"."%s"`, dbName, table.Name)
+}
+
 // RestoreFixture loads a fixture into a CockroachDB cluster. An enterprise
 // license is required to have been set in the cluster.
 func RestoreFixture(
-	ctx context.Context, sqlDB *gosql.DB, fixture Fixture, database string,
+	ctx context.Context, sqlDB *gosql.DB, fixture Fixture, database string, injectStats bool,
 ) (int64, error) {
 	var bytesAtomic int64
 	g := ctxgroup.WithContext(ctx)
 	genName := fixture.Generator.Meta().Name
+	tables := fixture.Generator.Tables()
+	if injectStats && tablesHaveStats(tables) {
+		// Turn off automatic stats temporarily so we don't trigger stats creation
+		// after the RESTORE.
+		// TODO(rytaft): It would be better if the automatic statistics code would
+		// just trigger a no-op if there are new stats available so we wouldn't
+		// have to disable and re-enable automatic stats here.
+		enableFn := disableAutoStats(ctx, sqlDB)
+		defer enableFn()
+	}
 	for _, table := range fixture.Tables {
 		table := table
 		g.GoCtx(func(ctx context.Context) error {
@@ -525,6 +550,17 @@ func RestoreFixture(
 	if err := g.Wait(); err != nil {
 		return 0, err
 	}
+	if injectStats {
+		for i := range tables {
+			t := &tables[i]
+			if len(t.Stats) > 0 {
+				qualifiedTableName := makeQualifiedTableName(genName, t)
+				if err := injectStatistics(qualifiedTableName, t, sqlDB); err != nil {
+					return 0, err
+				}
+			}
+		}
+	}
 	return atomic.LoadInt64(&bytesAtomic), nil
 }
 
diff --git a/pkg/ccl/workloadccl/fixture_test.go b/pkg/ccl/workloadccl/fixture_test.go
index 3413f6342f2b..9dd83a87b911 100644
--- a/pkg/ccl/workloadccl/fixture_test.go
+++ b/pkg/ccl/workloadccl/fixture_test.go
@@ -152,7 +152,7 @@ func TestFixture(t *testing.T) {
 	}
 
 	sqlDB.Exec(t, `CREATE DATABASE test`)
-	if _, err := RestoreFixture(ctx, db, fixture, `test`); err != nil {
+	if _, err := RestoreFixture(ctx, db, fixture, `test`, false); err != nil {
 		t.Fatalf(`%+v`, err)
 	}
 	sqlDB.CheckQueryResults(t,