Merge #62784

62784: workload/schemachange: add ALTER TABLE ... LOCALITY r=ajwerner a=otan Added LOCALITY for REGIONAL BY TABLE and GLOBAL tables. Held off on REGIONAL BY ROW for now as that requires rejecting everything else in the transaction, which requires more massaging. Refs: #62191 Release note: None Co-authored-by: Oliver Tan <[email protected]>
cockroachdb · Mar 30, 2021 · 8639dd1 · 8639dd1
2 parents 194c7b9 + ffe6533
commit 8639dd1
Show file tree

Hide file tree

Showing 2 changed files with 77 additions and 31 deletions.
diff --git a/pkg/workload/schemachange/operation_generator.go b/pkg/workload/schemachange/operation_generator.go
@@ -99,6 +99,8 @@ const (
 	addRegion                             // ALTER DATABASE <db> ADD REGION <region>
 	addUniqueConstraint                   // ALTER TABLE <table> ADD CONSTRAINT <constraint> UNIQUE (<column>)
 
+	alterTableLocality // ALTER TABLE <table> LOCALITY <locality>
+
 	createIndex    // CREATE INDEX <index> ON <table> <def>
 	createSequence // CREATE SEQUENCE <sequence> <def>
 	createTable    // CREATE TABLE <table> <def>
@@ -145,6 +147,7 @@ var opFuncs = map[opType]func(*operationGenerator, *pgx.Tx) (string, error){
 	addForeignKeyConstraint: (*operationGenerator).addForeignKeyConstraint,
 	addRegion:               (*operationGenerator).addRegion,
 	addUniqueConstraint:     (*operationGenerator).addUniqueConstraint,
+	alterTableLocality:      (*operationGenerator).alterTableLocality,
 	createIndex:             (*operationGenerator).createIndex,
 	createSequence:          (*operationGenerator).createSequence,
 	createTable:             (*operationGenerator).createTable,
@@ -189,6 +192,7 @@ var opWeights = []int{
 	addForeignKeyConstraint: 0,
 	addRegion:               1,
 	addUniqueConstraint:     0,
+	alterTableLocality:      1,
 	createIndex:             1,
 	createSequence:          1,
 	createTable:             1,
@@ -382,6 +386,47 @@ func (og *operationGenerator) addUniqueConstraint(tx *pgx.Tx) (string, error) {
 	return fmt.Sprintf(`ALTER TABLE %s ADD CONSTRAINT %s UNIQUE (%s)`, tableName, constaintName, columnForConstraint.name), nil
 }
 
+func (og *operationGenerator) alterTableLocality(tx *pgx.Tx) (string, error) {
+	tableName, err := og.randTable(tx, og.pctExisting(true), "")
+	if err != nil {
+		return "", err
+	}
+	tableExists, err := tableExists(tx, tableName)
+	if err != nil {
+		return "", err
+	}
+	if !tableExists {
+		og.expectedExecErrors.add(pgcode.UndefinedTable)
+		return fmt.Sprintf(`ALTER TABLE %s SET LOCALITY REGIONAL BY ROW`, tableName), nil
+	}
+
+	databaseRegionNames, err := getDatabaseRegionNames(tx)
+	if err != nil {
+		return "", err
+	}
+	if len(databaseRegionNames) == 0 {
+		og.expectedExecErrors.add(pgcode.InvalidTableDefinition)
+		return fmt.Sprintf(`ALTER TABLE %s SET LOCALITY REGIONAL BY ROW`, tableName), nil
+	}
+
+	localityOptions := []func() string{
+		func() string {
+			return "REGIONAL BY TABLE"
+		},
+		func() string {
+			idx := og.params.rng.Intn(len(databaseRegionNames))
+			regionName := tree.Name(databaseRegionNames[idx])
+			return fmt.Sprintf(`REGIONAL BY TABLE IN %s`, regionName.String())
+		},
+		func() string {
+			return "GLOBAL"
+		},
+		// TODO(#62191): do REGIONAL BY ROW and REGIONAL BY ROW AS <column>
+	}
+	idx := og.params.rng.Intn(len(localityOptions))
+	return fmt.Sprintf(`ALTER TABLE %s SET LOCALITY %s`, tableName, localityOptions[idx]()), nil
+}
+
 func getClusterRegionNames(tx *pgx.Tx) (descpb.RegionNames, error) {
 	return scanRegionNames(tx, "SELECT region FROM [SHOW REGIONS FROM CLUSTER]")
 }

diff --git a/pkg/workload/schemachange/optype_string.go b/pkg/workload/schemachange/optype_string.go