demo: Add location information to default demo localities.

pkg/cli/demo.go

@@ -52,7 +52,8 @@ to avoid pre-loading a dataset.`,
 const defaultGeneratorName = "movr"
 
 var defaultGenerator workload.Generator
-var defaultLocalities = []roachpb.Locality{
+
+var defaultLocalities = demoLocalityList{
 	// Default localities for a 3 node cluster
 	{Tiers: []roachpb.Tier{{Key: "region", Value: "us-east1"}, {Key: "az", Value: "b"}}},
 	{Tiers: []roachpb.Tier{{Key: "region", Value: "us-east1"}, {Key: "az", Value: "c"}}},

pkg/roachpb/metadata.go

@@ -514,3 +514,37 @@ func (l *Locality) Set(value string) error {
 	l.Tiers = tiers
 	return nil
 }
+
+// DefaultLocationInformation is used to populate the system.locations
+// table. The region values here are specific to GCP.
+var DefaultLocationInformation = []struct {
+	Locality  Locality
+	Latitude  string
+	Longitude string
+}{
+	{
+		Locality:  Locality{Tiers: []Tier{{Key: "region", Value: "us-east1"}}},
+		Latitude:  "33.836082",
+		Longitude: "-81.163727",
+	},
+	{
+		Locality:  Locality{Tiers: []Tier{{Key: "region", Value: "us-east4"}}},
+		Latitude:  "37.478397",
+		Longitude: "-76.453077",
+	},
+	{
+		Locality:  Locality{Tiers: []Tier{{Key: "region", Value: "us-central1"}}},
+		Latitude:  "42.032974",
+		Longitude: "-93.581543",
+	},
+	{
+		Locality:  Locality{Tiers: []Tier{{Key: "region", Value: "us-west1"}}},
+		Latitude:  "43.804133",
+		Longitude: "-120.554201",
+	},
+	{
+		Locality:  Locality{Tiers: []Tier{{Key: "region", Value: "europe-west1"}}},
+		Latitude:  "50.44816",
+		Longitude: "3.81886",
+	},
+}

pkg/sqlmigrations/migrations.go

@@ -206,6 +206,11 @@ var backwardCompatibleMigrations = []migrationDescriptor{
 		name:   "propagate the ts purge interval to the new setting names",
 		workFn: retireOldTsPurgeIntervalSettings,
 	},
+	{
+		// Introduced in ? TODO (rohany): what version is this?
+		name:   "update system.locations with default location data",
+		workFn: updateSystemLocationData,
+	},
 }
 
 func staticIDs(ids ...sqlbase.ID) func(ctx context.Context, db db) ([]sqlbase.ID, error) {
@@ -918,3 +923,27 @@ ON CONFLICT (name) DO NOTHING`,
 
 	return nil
 }
+
+func updateSystemLocationData(ctx context.Context, r runner) error {
+	// See if the system.locations table already has data in it.
+	// If so, we don't want to do anything.
+	row, err := r.sqlExecutor.QueryRow(ctx, "update-system-locations",
+		nil, `SELECT count(*) FROM system.locations`)
+	if err != nil {
+		return err
+	}
+	count := int(tree.MustBeDInt(row[0]))
+	if count != 0 {
+		return nil
+	}
+
+	for _, loc := range roachpb.DefaultLocationInformation {
+		stmt := `UPSERT INTO system.locations VALUES ($1, $2, $3, $4)`
+		tier := loc.Locality.Tiers[0]
+		if _, err := r.sqlExecutor.Exec(ctx, "update-system-locations", nil,
+			stmt, tier.Key, tier.Value, loc.Latitude, loc.Longitude); err != nil {
+			return err
+		}
+	}
+	return nil
+}

pkg/sqlmigrations/migrations_test.go

@@ -639,3 +639,32 @@ func TestExpectedInitialRangeCount(t *testing.T) {
 		return nil
 	})
 }
+
+func TestUpdateSystemLocationData(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	ctx := context.Background()
+
+	mt := makeMigrationTest(ctx, t)
+	defer mt.close(ctx)
+
+	migration := mt.pop(t, "update system.locations with default location data")
+	mt.start(t, base.TestServerArgs{})
+
+	// Check that we don't have any data in the system.locations table without the migration.
+	var count int
+	mt.sqlDB.QueryRow(t, `SELECT count(*) FROM system.locations`).Scan(&count)
+	if count != 0 {
+		t.Fatalf("Exected to find 0 rows in system.locations. Found  %d instead", count)
+	}
+
+	// Run the migration to insert locations.
+	if err := mt.runMigration(ctx, migration); err != nil {
+		t.Errorf("expected success, got %q", err)
+	}
+
+	// Check that we have all of the expected locations.
+	mt.sqlDB.QueryRow(t, `SELECT count(*) FROM system.locations`).Scan(&count)
+	if count != len(roachpb.DefaultLocationInformation) {
+		t.Fatalf("Exected to find 0 rows in system.locations. Found  %d instead", count)
+	}
+}