| // 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 |
| } |