Update maintenance service to be compatible with spanner

- Add support to perform schema checks for spanner
- The unit tests currently only support CDB. Will be adding support for running tests against a local spanner db in a separate CL.

Change-Id: Ida0790a6c9b02816cc45025ebeb5c1776ec88bb6
Reviewed-on: https://skia-review.googlesource.com/c/buildbot/+/934278
Reviewed-by: Farid (Mojtaba) Faridzad <[email protected]>
Commit-Queue: Ashwin Verleker <[email protected]>

Author: ashwinpvg

go/sql/schema/exportschema/exportschema.go

@@ -5,7 +5,6 @@ package exportschema
 import (
 	"context"
 	"encoding/json"
-	"flag"
 	"fmt"
 	"io"
 	"math/rand"
@@ -24,16 +23,7 @@ import (
 // Note that this works by requiring a local instance of the CockroachDB
 // emulator to be running, so this is not appropriate to call in a go:generate
 // statement.
-func Main(args []string, tables interface{}, schemaAsString string) error {
-	var out string
-	fs := flag.NewFlagSet("exportschema", flag.ExitOnError)
-	fs.StringVar(&out, "out", "", "Filename of the schema Description to write.")
-
-	err := fs.Parse(args[1:])
-	if err != nil {
-		return skerr.Wrap(err)
-	}
-
+func Main(out string, dbType string, tables interface{}, schemaAsString string) error {
 	if out == "" {
 		return skerr.Fmt("--out flag must be supplied")
 	}
@@ -42,18 +32,23 @@ func Main(args []string, tables interface{}, schemaAsString string) error {
 	rand.Seed(time.Now().UnixNano())
 	databaseName := fmt.Sprintf("%s_%d", "export", rand.Uint64())
 	host := emulators.GetEmulatorHostEnvVar(emulators.CockroachDB)
+	if dbType == schema.SpannerDBType {
+		host = "localhost:5432"
+	}
 	connectionString := fmt.Sprintf("postgresql://root@%s/%s?sslmode=disable", host, databaseName)
 	db, err := pgxpool.Connect(ctx, connectionString)
 	if err != nil {
 		sklog.Fatal(err)
 	}
 
-	// Create the database in cockroachdb.
-	_, err = db.Exec(ctx, fmt.Sprintf(`
-		CREATE DATABASE %s;
-		SET DATABASE = %s;`, databaseName, databaseName))
-	if err != nil {
-		sklog.Fatal(err)
+	if dbType != schema.SpannerDBType {
+		// Create the database in cockroachdb.
+		_, err = db.Exec(ctx, fmt.Sprintf(`
+				CREATE DATABASE %s;
+				SET DATABASE = %s;`, databaseName, databaseName))
+		if err != nil {
+			sklog.Fatal(err)
+		}
 	}
 
 	// Create the tables.
@@ -62,7 +57,7 @@ func Main(args []string, tables interface{}, schemaAsString string) error {
 		sklog.Fatal(err)
 	}
 
-	sch, err := schema.GetDescription(ctx, db, tables)
+	sch, err := schema.GetDescription(ctx, db, tables, dbType)
 	if err != nil {
 		sklog.Fatal(err)
 	}

go/sql/schema/schema.go

@@ -13,7 +13,9 @@ import (
 
 const (
 	// Timeout used on Contexts when making SQL requests.
-	sqlTimeout = time.Minute
+	sqlTimeout      = time.Minute
+	SpannerDBType   = "spanner"
+	CockroachDBType = "cockroachdb"
 )
 
 // TableNames takes in a "table type", that is a table whose fields are slices.
@@ -48,6 +50,17 @@ WHERE
     table_name = $1;
 `
 
+const typesQuerySpanner = `
+SELECT
+    column_name,
+    data_type || ' def:' || COALESCE(column_default, '') || ' nullable:' || COALESCE(is_nullable, '')
+FROM
+    information_schema.columns
+WHERE
+    table_name = $1;
+
+`
+
 // Query to return the index names for each table.
 const indexNameQuery = `
 SELECT DISTINCT
@@ -60,16 +73,31 @@ ORDER BY
 	index_name DESC
 `
 
+const spannerIndexNameQuery = `
+SELECT DISTINCT
+	index_name
+FROM
+	information_schema.indexes
+WHERE
+	table_name = $1
+ORDER BY
+	index_name DESC
+`
+
 // GetDescription returns a Description populated for every table listed in
 // `tables`.
-func GetDescription(ctx context.Context, db pool.Pool, tables interface{}) (*Description, error) {
+func GetDescription(ctx context.Context, db pool.Pool, tables interface{}, databaseType string) (*Description, error) {
 	ctx, cancel := context.WithTimeout(ctx, sqlTimeout)
 	defer cancel()
 	colNameAndType := map[string]string{}
 	indexNames := []string{}
 	for _, tableName := range TableNames(tables) {
 		// Fill in colNameAndType.
-		rows, err := db.Query(ctx, typesQuery, tableName)
+		typeQuery := typesQuery
+		if databaseType == SpannerDBType {
+			typeQuery = typesQuerySpanner
+		}
+		rows, err := db.Query(ctx, typeQuery, tableName)
 		if err != nil {
 			return nil, skerr.Wrap(err)
 		}
@@ -84,7 +112,11 @@ func GetDescription(ctx context.Context, db pool.Pool, tables interface{}) (*Des
 		}
 
 		// Fill in indexNames.
-		rows, err = db.Query(ctx, indexNameQuery, tableName)
+		indexQuery := indexNameQuery
+		if databaseType == SpannerDBType {
+			indexQuery = spannerIndexNameQuery
+		}
+		rows, err = db.Query(ctx, indexQuery, tableName)
 		if err != nil {
 			return nil, skerr.Wrap(err)
 		}

go/sql/schema/schema_test.go

@@ -80,9 +80,9 @@ func setupForTest(t *testing.T) *pgxpool.Pool {
 
 }
 
-func TestGetDescription(t *testing.T) {
+func TestGetDescriptionCDB(t *testing.T) {
 	db := setupForTest(t)
-	desc, err := GetDescription(context.Background(), db, Tables{})
+	desc, err := GetDescription(context.Background(), db, Tables{}, CockroachDBType)
 	require.NoError(t, err)
 
 	expected := Description{

machine/go/machine/store/cdb/cdb.go

@@ -136,7 +136,7 @@ func New(ctx context.Context, db pool.Pool, pools *pools.Pools) (*Store, error)
 		return nil, skerr.Wrap(err)
 	}
 
-	actual, err := schema.GetDescription(ctx, db, Tables{})
+	actual, err := schema.GetDescription(ctx, db, Tables{}, schema.CockroachDBType)
 	if err != nil {
 		return nil, skerr.Wrap(err)
 	}

machine/go/machine/store/cdb/cdb_test.go

@@ -356,7 +356,7 @@ CREATE TABLE IF NOT EXISTS TaskResult (
 `
 
 func getSchema(t *testing.T, db pool.Pool) *schema.Description {
-	ret, err := schema.GetDescription(context.Background(), db, cdb.Tables{})
+	ret, err := schema.GetDescription(context.Background(), db, cdb.Tables{}, schema.CockroachDBType)
 	require.NoError(t, err)
 	require.NotEmpty(t, ret.ColumnNameAndType)
 	return ret

machine/go/machine/store/cdb/exportschema/BUILD.bazel

@@ -9,6 +9,7 @@ go_library(
         "//go/sklog",
         "//go/sql/schema/exportschema:exportschema_lib",
         "//machine/go/machine/store/cdb",
+        "//machine/go/machine/store/cdb/spanner",
     ],
 )
 

machine/go/machine/store/cdb/exportschema/exportschema.go

@@ -2,15 +2,32 @@
 package main
 
 import (
+	"flag"
 	"os"
 
 	"go.skia.org/infra/go/sklog"
 	"go.skia.org/infra/go/sql/schema/exportschema"
 	"go.skia.org/infra/machine/go/machine/store/cdb"
+	"go.skia.org/infra/machine/go/machine/store/cdb/spanner"
 )
 
 func main() {
-	err := exportschema.Main(os.Args, cdb.Tables{}, cdb.Schema)
+	var dbType string
+	var out string
+	fs := flag.NewFlagSet("exportschema", flag.ExitOnError)
+	fs.StringVar(&dbType, "databaseType", "", "Database type for the schema.")
+	fs.StringVar(&out, "out", "", "Filename of the schema Description to write.")
+
+	err := fs.Parse(os.Args[1:])
+	if err != nil {
+		sklog.Fatalf("Error parsing arguments: %v", err)
+	}
+
+	schema := cdb.Schema
+	if dbType == "spanner" {
+		schema = spanner.Schema
+	}
+	err = exportschema.Main(out, dbType, cdb.Tables{}, schema)
 	if err != nil {
 		sklog.Fatal(err)
 	}

perf/Makefile

@@ -9,6 +9,10 @@ generate_sql_schema:
 	cd ./go/sql; go run ./tosql
 	go run ./go/sql/exportschema --out ./go/sql/expectedschema/schema.json
 
+generate_sql_schema_spanner:
+	cd ./go/sql; go run ./tosql
+	go run ./go/sql/exportschema --out ./go/sql/expectedschema/schema_spanner.json --databaseType spanner
+
 run_cockroachdb_emulators:
 	../scripts/run_emulators/run_emulators start
 
@@ -246,11 +250,22 @@ run_frontend:
 run-demo-instance: run-spanner-emulator initdemo ingest_demo run_frontend
 run-demo-instance-db-persist: initdemo run_frontend
 
+run_maintenance:
+	../_bazel_bin/perf/go/perfserver/perfserver_/perfserver maintenance \
+		--local \
+		--config_filename=./configs/demo_spanner.json \
+		--connection_string=postgresql://[email protected]:5432/demo?sslmode=disable \
+		--migrate_regressions \
+		--refresh_query_cache \
+		--tiles_for_query_cache=5
+
+run-demo-maintenance: run-spanner-emulator initdemo ingest_demo run_maintenance
+
 # Before running this target you need to have a local single-instance of
 # cockroachdb running, started like this:
 #
 #  cd /tmp cockroach start-single-node --insecure --listen-addr=127.0.0.1
-run-demo-maintenance:
+run-demo-maintenance-cdb:
 	$(BAZEL) build --config=mayberemote -c dbg //perf/...
 	$(BAZEL) run --config=mayberemote -c dbg //perf/go/initdemo:initdemo
 	../_bazel_bin/perf/go/perfserver/perfserver_/perfserver ingest \

perf/go/builders/builders.go

@@ -117,12 +117,12 @@ func NewCockroachDBFromConfig(ctx context.Context, instanceConfig *config.Instan
 
 	if checkSchema {
 		// Confirm the database has the right schema.
-		expectedSchema, err := expectedschema.Load()
+		expectedSchema, err := expectedschema.Load(instanceConfig.DataStoreConfig.DataStoreType)
 		if err != nil {
 			return nil, skerr.Wrap(err)
 		}
 
-		actual, err := schema.GetDescription(ctx, singletonPool, sql.Tables{})
+		actual, err := schema.GetDescription(ctx, singletonPool, sql.Tables{}, string(instanceConfig.DataStoreConfig.DataStoreType))
 		if err != nil {
 			return nil, skerr.Wrap(err)
 		}

perf/go/maintenance/maintenance.go

@@ -54,7 +54,7 @@ func Start(ctx context.Context, flags config.MaintenanceFlags, instanceConfig *c
 	if err != nil {
 		return skerr.Wrapf(err, "Failed to create CockroachDB instance.")
 	}
-	err = expectedschema.ValidateAndMigrateNewSchema(ctx, db)
+	err = expectedschema.ValidateAndMigrateNewSchema(ctx, db, instanceConfig.DataStoreConfig.DataStoreType)
 	if err != nil {
 		return skerr.Wrapf(err, "Failed to migrate schema.")
 	}