blob: 4e6b66e88496737ff44e692833cfc011b951c81f [file] [log] [blame]
// Copyright 2019 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// Package spantest implements creation/destruction of a temporary Spanner
// database.
package spantest
import (
"context"
"crypto/rand"
"encoding/binary"
"fmt"
"io/ioutil"
"regexp"
"strings"
"time"
"cloud.google.com/go/spanner"
spandb "cloud.google.com/go/spanner/admin/database/apiv1"
"google.golang.org/api/option"
dbpb "google.golang.org/genproto/googleapis/spanner/admin/database/v1"
"google.golang.org/grpc"
"google.golang.org/grpc/credentials"
"go.chromium.org/luci/auth"
"go.chromium.org/luci/common/errors"
"go.chromium.org/luci/hardcoded/chromeinfra"
)
// TempDBConfig specifies how to create a temporary database.
type TempDBConfig struct {
// InstanceName is the name of Spannner instance where to create the
// temporary database.
// Format: projects/{project}/instances/{instance}.
// Defaults to chromeinfra.TestSpannerInstance.
InstanceName string
// Credentials will be used to authenticate to Spanner.
// If nil, auth.Authenticator with SilentLogin and chrome-infra auth options
// will be used.
// This means that that the user may have to login with luci-auth tool.
Credentials credentials.PerRPCCredentials
// InitScriptPath is a path to a DDL script to initialize the database.
//
// In lieu of a proper DDL parser, it is parsed using regexes.
// Therefore the script MUST:
// - Use `#`` and/or `--`` for comments. No block comments.
// - Separate DDL statements with `;\n`.
//
// If empty, the database is created with no tables.
InitScriptPath string
}
func (cfg *TempDBConfig) credentials(ctx context.Context) (credentials.PerRPCCredentials, error) {
if cfg.Credentials != nil {
return cfg.Credentials, nil
}
opts := chromeinfra.DefaultAuthOptions()
opts.Scopes = spandb.DefaultAuthScopes()
a := auth.NewAuthenticator(ctx, auth.SilentLogin, opts)
if err := a.CheckLoginRequired(); err != nil {
return nil, errors.Annotate(err, "please login with `luci-auth login -scopes %q`", strings.Join(opts.Scopes, " ")).Err()
}
return a.PerRPCCredentials()
}
var ddlStatementSepRe = regexp.MustCompile(`;\s*\n`)
var commentRe = regexp.MustCompile(`(--|#)[^\n]*`)
// readDDLStatements read the file at cfg.InitScriptPath as a sequence of DDL
// statements. If the path is empty, returns (nil, nil).
func (cfg *TempDBConfig) readDDLStatements() ([]string, error) {
if cfg.InitScriptPath == "" {
return nil, nil
}
contents, err := ioutil.ReadFile(cfg.InitScriptPath)
if err != nil {
return nil, err
}
statements := ddlStatementSepRe.Split(string(contents), -1)
ret := statements[:0]
for _, stmt := range statements {
stmt = commentRe.ReplaceAllString(stmt, "")
stmt = strings.TrimSpace(stmt)
if stmt != "" {
ret = append(ret, stmt)
}
}
return ret, nil
}
// adminClient returns a Spanner admin client, it must be closed when done.
func adminClient(ctx context.Context, creds credentials.PerRPCCredentials) (*spandb.DatabaseAdminClient, error) {
return spandb.NewDatabaseAdminClient(ctx,
option.WithGRPCDialOption(grpc.WithPerRPCCredentials(creds)))
}
// TempDB is a temporary Spanner database.
type TempDB struct {
Name string
creds credentials.PerRPCCredentials
}
// Client returns a spanner client connected to the database.
func (db *TempDB) Client(ctx context.Context) (*spanner.Client, error) {
return spanner.NewClient(ctx, db.Name,
option.WithGRPCDialOption(grpc.WithPerRPCCredentials(db.creds)),
)
}
// Drop deletes the database.
func (db *TempDB) Drop(ctx context.Context) error {
client, err := adminClient(ctx, db.creds)
if err != nil {
return err
}
defer client.Close()
return client.DropDatabase(ctx, &dbpb.DropDatabaseRequest{
Database: db.Name,
})
}
var dbNameAlphabetInversedRe = regexp.MustCompile(`[^\w]+`)
// NewTempDB creates a temporary database with a random name.
// The caller is responsible for calling Drop on the returned TempDB to
// cleanup resources after usage.
func NewTempDB(ctx context.Context, cfg TempDBConfig) (*TempDB, error) {
instanceName := cfg.InstanceName
if instanceName == "" {
instanceName = chromeinfra.TestSpannerInstance
}
creds, err := cfg.credentials(ctx)
if err != nil {
return nil, err
}
initStatements, err := cfg.readDDLStatements()
if err != nil {
return nil, errors.Annotate(err, "failed to read %q", cfg.InitScriptPath).Err()
}
client, err := adminClient(ctx, creds)
if err != nil {
return nil, err
}
defer client.Close()
// Generate a random database name.
var random uint32
if err := binary.Read(rand.Reader, binary.LittleEndian, &random); err != nil {
panic(err)
}
dbName := fmt.Sprintf("tmp%s-%d", time.Now().Format("20060102-"), random)
dbName = SanitizeDBName(dbName)
dbOp, err := client.CreateDatabase(ctx, &dbpb.CreateDatabaseRequest{
Parent: instanceName,
CreateStatement: "CREATE DATABASE " + dbName,
ExtraStatements: initStatements,
})
if err != nil {
return nil, errors.Annotate(err, "failed to create database").Err()
}
db, err := dbOp.Wait(ctx)
if err != nil {
return nil, errors.Annotate(err, "failed to create database").Err()
}
return &TempDB{
Name: db.Name,
creds: creds,
}, nil
}
// SanitizeDBName tranforms name to a valid one.
// If name is already valid, returns it without changes.
func SanitizeDBName(name string) string {
name = strings.ToLower(name)
name = dbNameAlphabetInversedRe.ReplaceAllLiteralString(name, "_")
const maxLen = 30
if len(name) > maxLen {
name = name[:maxLen]
}
name = strings.TrimRight(name, "_")
return name
}