demo: Add a flag to automatically apply geo-partitioned replicas.

Work for cockroachdb#39945.

Release note (cli change): Add an option for cockroach demo to
automatically apply the geo-partitioned replicas topology to the movr
dataset.

pkg/cli/cliflags/flags.go

@@ -918,6 +918,16 @@ to us-east1 and availability zone to 3.
 `,
 	}
 
+	DemoGeoPartitionedReplicas = FlagInfo{
+		Name: "geo-partitioned-replicas",
+		Description: `
+Create a 9 node cluster preloaded with the Movr dataset, and automatically apply
+the geo-partitioned replicas topology across the regions us-east1, us-west1, and
+europe-west1. More information about the geo-partitioned replicas topology can be found
+at this URL: https://www.cockroachlabs.com/docs/v19.1/topology-geo-partitioned-replicas.html
+		`,
+	}
+
 	UseEmptyDatabase = FlagInfo{
 		Name: "empty",
 		Description: `

pkg/cli/context.go

@@ -147,6 +147,7 @@ func initCLIDefaults() {
 	demoCtx.useEmptyDatabase = false
 	demoCtx.runWorkload = false
 	demoCtx.localities = nil
+	demoCtx.geoPartitionedReplicas = false
 
 	initPreFlagsDefaults()
 
@@ -340,4 +341,6 @@ var demoCtx struct {
 	useEmptyDatabase bool
 	runWorkload      bool
 	localities       demoLocalityList
+	// Use the geopartitioned replicas topology
+	geoPartitionedReplicas bool
 }

pkg/cli/demo.go

@@ -124,6 +124,13 @@ func setupTransientServers(
 		return "", "", cleanup, errors.Errorf("must have a positive number of nodes")
 	}
 
+	// This demo only works on a 9 node cluster, so set the node count as such.
+	// Ignore input user localities.
+	if demoCtx.geoPartitionedReplicas {
+		demoCtx.nodes = 9
+		demoCtx.localities = nil
+	}
+
 	// The user specified some localities for their nodes.
 	if len(demoCtx.localities) != 0 {
 		// Error out of localities don't line up with requested node
@@ -255,6 +262,17 @@ func setupTransientServers(
 			return ``, ``, cleanup, err
 		}
 
+		if demoCtx.geoPartitionedReplicas && gen.Meta().Name == "movr" {
+			db, err := gosql.Open("postgres", urlStr)
+			if err != nil {
+				return ``, ``, cleanup, err
+			}
+			defer db.Close()
+			if err := setupGeoPartitionedReplicas(db); err != nil {
+				return ``, ``, cleanup, err
+			}
+		}
+
 		if demoCtx.runWorkload {
 			if err := runWorkload(ctx, gen, urlStr, stopper); err != nil {
 				return ``, ``, cleanup, err
@@ -265,6 +283,108 @@ func setupTransientServers(
 	return urlStr, s.AdminURL(), cleanup, nil
 }
 
+func setupGeoPartitionedReplicas(db *gosql.DB) error {
+	// Create us-west, us-east and europe-west partitions.
+	q := `
+		ALTER TABLE users PARTITION BY LIST (city) ( 
+			PARTITION us_west VALUES IN ('seattle', 'san francisco', 'los angeles'),
+			PARTITION us_east VALUES IN ('new york', 'boston', 'washington dc'),
+			PARTITION europe_west VALUES IN ('amsterdam', 'paris', 'rome') 
+		);
+		ALTER TABLE vehicles PARTITION BY LIST (city) (
+			PARTITION us_west VALUES IN ('seattle', 'san francisco', 'los angeles'),
+			PARTITION us_east VALUES IN ('new york', 'boston', 'washington dc'),
+			PARTITION europe_west VALUES IN ('amsterdam', 'paris', 'rome') 
+		);
+		ALTER INDEX vehicles_auto_index_fk_city_ref_users PARTITION BY LIST (city) ( 
+			PARTITION us_west VALUES IN ('seattle', 'san francisco', 'los angeles'),
+			PARTITION us_east VALUES IN ('new york', 'boston', 'washington dc'),
+			PARTITION europe_west VALUES IN ('amsterdam', 'paris', 'rome') 
+		);
+		ALTER TABLE rides PARTITION BY LIST (city) (
+			PARTITION us_west VALUES IN ('seattle', 'san francisco', 'los angeles'),
+			PARTITION us_east VALUES IN ('new york', 'boston', 'washington dc'),
+			PARTITION europe_west VALUES IN ('amsterdam', 'paris', 'rome') 
+		);
+		ALTER INDEX rides_auto_index_fk_city_ref_users PARTITION BY LIST (city) (
+			PARTITION us_west VALUES IN ('seattle', 'san francisco', 'los angeles'),
+			PARTITION us_east VALUES IN ('new york', 'boston', 'washington dc'),
+			PARTITION europe_west VALUES IN ('amsterdam', 'paris', 'rome') 
+		);
+		ALTER INDEX rides_auto_index_fk_vehicle_city_ref_vehicles PARTITION BY LIST (vehicle_city) (
+			PARTITION us_west VALUES IN ('seattle', 'san francisco', 'los angeles'),
+			PARTITION us_east VALUES IN ('new york', 'boston', 'washington dc'),
+			PARTITION europe_west VALUES IN ('amsterdam', 'paris', 'rome') 
+		);
+		ALTER TABLE user_promo_codes PARTITION BY LIST (city) (
+			PARTITION us_west VALUES IN ('seattle', 'san francisco', 'los angeles'),
+			PARTITION us_east VALUES IN ('new york', 'boston', 'washington dc'),
+			PARTITION europe_west VALUES IN ('amsterdam', 'paris', 'rome') 
+		);
+		ALTER TABLE vehicle_location_histories PARTITION BY LIST (city) (
+			PARTITION us_west VALUES IN ('seattle', 'san francisco', 'los angeles'),
+			PARTITION us_east VALUES IN ('new york', 'boston', 'washington dc'),
+			PARTITION europe_west VALUES IN ('amsterdam', 'paris', 'rome') 
+		);
+	`
+	if _, err := db.Exec(q); err != nil {
+		return err
+	}
+
+	// Alter the partitions to place replicas in the appropriate zones.
+	q = `
+		ALTER PARTITION us_west OF INDEX users@* CONFIGURE ZONE USING CONSTRAINTS='["+region=us-west1"]';
+		ALTER PARTITION us_east OF INDEX users@* CONFIGURE ZONE USING CONSTRAINTS='["+region=us-east1"]';
+		ALTER PARTITION europe_west OF INDEX users@* CONFIGURE ZONE USING CONSTRAINTS='["+region=europe-west1"]';
+		
+		ALTER PARTITION us_west OF INDEX vehicles@* CONFIGURE ZONE USING CONSTRAINTS='["+region=us-west1"]';
+		ALTER PARTITION us_east OF INDEX vehicles@* CONFIGURE ZONE USING CONSTRAINTS='["+region=us-east1"]';
+		ALTER PARTITION europe_west OF INDEX vehicles@* CONFIGURE ZONE USING CONSTRAINTS='["+region=europe-west1"]';
+		
+		ALTER PARTITION us_west OF INDEX rides@* CONFIGURE ZONE USING CONSTRAINTS='["+region=us-west1"]';
+		ALTER PARTITION us_east OF INDEX rides@* CONFIGURE ZONE USING CONSTRAINTS='["+region=us-east1"]';
+		ALTER PARTITION europe_west OF INDEX rides@* CONFIGURE ZONE USING CONSTRAINTS='["+region=europe-west1"]';
+		
+		ALTER PARTITION us_west OF INDEX user_promo_codes@* CONFIGURE ZONE USING CONSTRAINTS='["+region=us-west1"]';
+		ALTER PARTITION us_east OF INDEX user_promo_codes@* CONFIGURE ZONE USING CONSTRAINTS='["+region=us-east1"]';
+		ALTER PARTITION europe_west OF INDEX user_promo_codes@* CONFIGURE ZONE USING CONSTRAINTS='["+region=europe-west1"]';
+		
+		ALTER PARTITION us_west OF INDEX vehicle_location_histories@* CONFIGURE ZONE USING CONSTRAINTS='["+region=us-west1"]';
+		ALTER PARTITION us_east OF INDEX vehicle_location_histories@* CONFIGURE ZONE USING CONSTRAINTS='["+region=us-east1"]';
+		ALTER PARTITION europe_west OF INDEX vehicle_location_histories@* CONFIGURE ZONE USING CONSTRAINTS='["+region=europe-west1"]';
+	`
+	if _, err := db.Exec(q); err != nil {
+		return err
+	}
+
+	// Create some duplicate indexes for the promo_codes table.
+	q = `
+		CREATE INDEX promo_codes_idx_us_west ON promo_codes (code) STORING (description, creation_time, expiration_time, rules);
+		CREATE INDEX promo_codes_idx_europe_west ON promo_codes (code) STORING (description, creation_time, expiration_time, rules);
+	`
+	if _, err := db.Exec(q); err != nil {
+		return err
+	}
+
+	// Apply configurations to the index for fast reads.
+	q = `
+		ALTER TABLE promo_codes CONFIGURE ZONE USING num_replicas = 3,
+			constraints = '{"+region=us-east1": 1}',
+			lease_preferences = '[[+region=us-east1]]';
+		ALTER INDEX promo_codes@promo_codes_idx_us_west CONFIGURE ZONE USING 
+			constraints = '{"+region=us-west1": 1}',
+			lease_preferences = '[[+region=us-west1]]'; 
+		ALTER INDEX promo_codes@promo_codes_idx_europe_west CONFIGURE ZONE USING 
+			constraints = '{"+region=europe-west1": 1}', 
+			lease_preferences = '[[+region=europe-west1]]';
+	`
+	if _, err := db.Exec(q); err != nil {
+		return err
+	}
+
+	return nil
+}
+
 func runWorkload(
 	ctx context.Context, gen workload.Generator, dbURL string, stopper *stop.Stopper,
 ) error {
@@ -321,6 +441,11 @@ func runDemo(cmd *cobra.Command, gen workload.Generator) error {
 		return errors.New("cannot run a workload against an empty database")
 	}
 
+	// Make sure that the user didn't request to have a topology and an empty database.
+	if demoCtx.geoPartitionedReplicas && demoCtx.useEmptyDatabase {
+		return errors.New("cannot setup geo-partitioned replicas topology on an empty database")
+	}
+
 	connURL, adminURL, cleanup, err := setupTransientServers(cmd, gen)
 	defer cleanup()
 	if err != nil {

pkg/cli/flags.go

@@ -588,6 +588,7 @@ func init() {
 	// so we use the regular flag set.
 	BoolFlag(demoCmd.Flags(), &demoCtx.useEmptyDatabase, cliflags.UseEmptyDatabase, false)
 	BoolFlag(demoCmd.Flags(), &demoCtx.runWorkload, cliflags.RunDemoWorkload, false)
+	BoolFlag(demoCmd.Flags(), &demoCtx.geoPartitionedReplicas, cliflags.DemoGeoPartitionedReplicas, false)
 	VarFlag(demoFlags, &demoCtx.localities, cliflags.DemoNodeLocality)
 
 	// sqlfmt command.

pkg/workload/movr/movr.go

@@ -112,9 +112,10 @@ var cities = []struct {
 	{city: "seattle", locality: "us_west"},
 	{city: "san francisco", locality: "us_west"},
 	{city: "los angeles", locality: "us_west"},
-	{city: "chicago", locality: "us_central"},
-	{city: "detroit", locality: "us_central"},
-	{city: "minneapolis", locality: "us_central"},
+	// The demo setup we have is a 9 node, 3 region cluster.
+	//{city: "chicago", locality: "us_central"},
+	//{city: "detroit", locality: "us_central"},
+	//{city: "minneapolis", locality: "us_central"},
 	{city: "amsterdam", locality: "eu_west"},
 	{city: "paris", locality: "eu_west"},
 	{city: "rome", locality: "eu_west"},