blob: e3ec0fbd686711d8f7c055d20abe7cef49c678f3 [file] [log] [blame]
// Copyright 2020 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
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.
// Package span implements a server module for communicating with Cloud Spanner.
package span
import (
// ModuleName can be used to refer to this module when declaring dependencies.
var ModuleName = module.RegisterName("")
// ModuleOptions contain configuration of the Cloud Spanner server module.
type ModuleOptions struct {
SpannerEndpoint string // the Spanner endpoint to connect to
SpannerDatabase string // identifier of Cloud Spanner database to connect to
// Register registers the command line flags.
func (o *ModuleOptions) Register(f *flag.FlagSet) {
"The Spanner endpoint to connect to. "+
"The default is defined by the Cloud Spanner library, "+
"but usually it is",
"Identifier of the Cloud Spanner database to connect to. A valid database "+
"name has the form projects/PROJECT_ID/instances/INSTANCE_ID/databases/DATABASE_ID. Required.",
// NewModule returns a server module that sets up a Spanner client connected
// to some single Cloud Spanner database.
// Client's functionality is exposed via Single(ctx), ReadOnlyTransaction(ctx),
// ReadWriteTransaction(ctx), etc.
// The underlying *spanner.Client is intentionally not exposed to make sure
// all callers use the functions mentioned above since they generally add
// additional functionality on top of the raw Spanner client that other LUCI
// packages assume to be present. Using the Spanner client directly may violate
// such assumptions leading to undefined behavior when multiple packages are
// used together.
func NewModule(opts *ModuleOptions) module.Module {
if opts == nil {
opts = &ModuleOptions{}
return &spannerModule{opts: opts}
// NewModuleFromFlags is a variant of NewModule that initializes options through
// command line flags.
// Calling this function registers flags in flag.CommandLine. They are usually
// parsed in server.Main(...).
func NewModuleFromFlags() module.Module {
opts := &ModuleOptions{}
return NewModule(opts)
// spannerModule implements module.Module.
type spannerModule struct {
opts *ModuleOptions
// Name is part of module.Module interface.
func (*spannerModule) Name() module.Name {
return ModuleName
// Dependencies is part of module.Module interface.
func (*spannerModule) Dependencies() []module.Dependency {
return nil
// Initialize is part of module.Module interface.
func (m *spannerModule) Initialize(ctx context.Context, host module.Host, opts module.HostOptions) (context.Context, error) {
switch {
case m.opts.SpannerDatabase == "":
return nil, errors.New("Cloud Spanner database name is required")
case !isValidDB(m.opts.SpannerDatabase):
return nil, errors.New("Cloud Spanner database name must have form `projects/.../instances/.../databases/...`")
// Credentials with Cloud scope.
creds, err := auth.GetPerRPCCredentials(ctx, auth.AsSelf, auth.WithScopes(auth.CloudOAuthScopes...))
if err != nil {
return nil, errors.Annotate(err, "failed to get PerRPCCredentials").Err()
// Initialize the client.
options := []option.ClientOption{option.WithGRPCDialOption(grpc.WithPerRPCCredentials(creds))}
if m.opts.SpannerEndpoint != "" {
options = append(options, option.WithEndpoint(m.opts.SpannerEndpoint))
client, err := spanner.NewClientWithConfig(ctx,
SessionPoolConfig: spanner.SessionPoolConfig{
TrackSessionHandles: !opts.Prod,
if err != nil {
return nil, errors.Annotate(err, "failed to instantiate Cloud Spanner client").Err()
ctx = UseClient(ctx, client)
// Close the client when exiting gracefully.
host.RegisterCleanup(func(ctx context.Context) { client.Close() })
// Run a "select 1" query to verify the database exists and we can access it
// before we actually serve any requests.
if err := pingDB(ctx); err != nil {
return nil, errors.Annotate(err, "failed to ping Cloud Spanner database").Err()
return ctx, nil
func isValidDB(name string) bool {
chunks := strings.Split(name, "/")
if len(chunks) != 6 {
return false
for _, ch := range chunks {
if ch == "" {
return false
return chunks[0] == "projects" && chunks[2] == "instances" && chunks[4] == "databases"
func pingDB(ctx context.Context) error {
ctx, done := context.WithTimeout(Single(ctx), 30*time.Second)
defer done()
return Query(ctx, spanner.NewStatement("SELECT 1;")).Do(func(*spanner.Row) error { return nil })