workload: Allow multi-region TPC-C to use IMPORT

Previously, multi-region TPC-C could only leverage insert for data loading.
This commit enables the IMPORT path by adding the pre-create step to the
workload Setup. The commit also works around a problem in IMPORT where
for multi-region databases, if certain table types (REGIONAL BY ROW, REGIONAL
BY TABLE IN) are created during a running IMPORT, the IMPORT will fail due to
the fact that the crdb_internal_region type is modified under the covers (to
install back references). To work around this problem, we pre-create the
tables before running an IMPORT INTO.

Release note (bug fix): Fixes IMPORT in tpcc workload.

Author: ajstorm

pkg/ccl/workloadccl/fixture.go

@@ -401,6 +401,19 @@ func ImportFixture(
 		pathPrefix = `workload://`
 	}
 
+	// Pre-create tables. It's required that we pre-create the tables before we
+	// parallelize the IMPORT because for multi-region setups, the create table
+	// will end up modifying the crdb_internal_region type (to install back
+	// references). If create table is done in parallel with IMPORT, some IMPORT
+	// jobs may fail because the type is being modified concurrently with the
+	// IMPORT. Removing the need to pre-create is being tracked with #70987.
+	for _, table := range tables {
+		err := createFixtureTable(sqlDB, dbName, table)
+		if err != nil {
+			return 0, errors.Wrapf(err, `creating table %s`, table.Name)
+		}
+	}
+
 	for _, t := range tables {
 		table := t
 		paths := csvServerPaths(pathPrefix, gen, table, numNodes*filesPerNode)
@@ -417,6 +430,16 @@ func ImportFixture(
 	return atomic.LoadInt64(&bytesAtomic), nil
 }
 
+func createFixtureTable(sqlDB *gosql.DB, dbName string, table workload.Table) error {
+	qualifiedTableName := makeQualifiedTableName(dbName, &table)
+	createTable := fmt.Sprintf(
+		`CREATE TABLE IF NOT EXISTS %s %s`,
+		qualifiedTableName,
+		table.Schema)
+	_, err := sqlDB.Exec(createTable)
+	return err
+}
+
 func importFixtureTable(
 	ctx context.Context,
 	sqlDB *gosql.DB,
@@ -429,8 +452,9 @@ func importFixtureTable(
 	start := timeutil.Now()
 	var buf bytes.Buffer
 	var params []interface{}
+
 	qualifiedTableName := makeQualifiedTableName(dbName, &table)
-	fmt.Fprintf(&buf, `IMPORT TABLE %s %s CSV DATA (`, qualifiedTableName, table.Schema)
+	fmt.Fprintf(&buf, `IMPORT INTO %s CSV DATA (`, qualifiedTableName)
 	// Generate $1,...,$N-1, where N is the number of csv paths.
 	for _, path := range paths {
 		params = append(params, path)

pkg/ccl/workloadccl/fixture_test.go

@@ -196,9 +196,9 @@ func TestImportFixture(t *testing.T) {
 
 	// Since we did not inject stats, the IMPORT should have triggered
 	// automatic stats collection.
-	sqlDB.CheckQueryResultsRetry(t,
-		`SELECT statistics_name, column_names, row_count, distinct_count, null_count
-           FROM [SHOW STATISTICS FOR TABLE ingest.fx]`,
+	statsQuery := fmt.Sprintf(`SELECT statistics_name, column_names, row_count, distinct_count, null_count
+           FROM [SHOW STATISTICS FOR TABLE ingest.fx] WHERE row_count = %d`, fixtureTestGenRows)
+	sqlDB.CheckQueryResultsRetry(t, statsQuery,
 		[][]string{
 			{"__auto__", "{key}", "10", "10", "0"},
 			{"__auto__", "{value}", "10", "1", "0"},

pkg/workload/tpcc/tpcc.go

@@ -230,6 +230,31 @@ var tpccMeta = workload.Meta{
 	},
 }
 
+func queryDatabaseRegions(db *gosql.DB) (map[string]struct{}, error) {
+	regions := make(map[string]struct{})
+	rows, err := db.Query(`SELECT region FROM [SHOW REGIONS FROM DATABASE]`)
+	if err != nil {
+		return regions, err
+	}
+	defer func() {
+		_ = rows.Close()
+	}()
+	for rows.Next() {
+		if rows.Err() != nil {
+			return regions, err
+		}
+		var region string
+		if err := rows.Scan(&region); err != nil {
+			return regions, err
+		}
+		regions[region] = struct{}{}
+	}
+	if rows.Err() != nil {
+		return regions, err
+	}
+	return regions, nil
+}
+
 // Meta implements the Generator interface.
 func (*tpcc) Meta() workload.Meta { return tpccMeta }
 
@@ -354,27 +379,10 @@ func (w *tpcc) Hooks() workload.Hooks {
 				return nil
 			}
 
-			regions := make(map[string]struct{})
-			rows, err := db.Query(`SELECT region FROM [SHOW REGIONS FROM DATABASE]`)
+			regions, err := queryDatabaseRegions(db)
 			if err != nil {
 				return err
 			}
-			defer func() {
-				_ = rows.Close()
-			}()
-			for rows.Next() {
-				if rows.Err() != nil {
-					return err
-				}
-				var region string
-				if err := rows.Scan(&region); err != nil {
-					return err
-				}
-				regions[region] = struct{}{}
-			}
-			if rows.Err() != nil {
-				return err
-			}
 
 			var dbName string
 			if err := db.QueryRow(`SHOW DATABASE`).Scan(&dbName); err != nil {

pkg/workload/workloadsql/workloadsql.go

@@ -36,6 +36,17 @@ import (
 func Setup(
 	ctx context.Context, db *gosql.DB, gen workload.Generator, l workload.InitialDataLoader,
 ) (int64, error) {
+	var hooks workload.Hooks
+	if h, ok := gen.(workload.Hookser); ok {
+		hooks = h.Hooks()
+	}
+
+	if hooks.PreCreate != nil {
+		if err := hooks.PreCreate(db); err != nil {
+			return 0, errors.Wrapf(err, "Could not pre-create")
+		}
+	}
+
 	bytes, err := l.InitialDataLoad(ctx, db, gen)
 	if err != nil {
 		return 0, err
@@ -48,10 +59,6 @@ func Setup(
 		}
 	}
 
-	var hooks workload.Hooks
-	if h, ok := gen.(workload.Hookser); ok {
-		hooks = h.Hooks()
-	}
 	if hooks.PostLoad != nil {
 		if err := hooks.PostLoad(db); err != nil {
 			return 0, errors.Wrapf(err, "Could not postload")