/*

Copyright 2017 Google LLC

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

https://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 spanner

import (

"context"

"fmt"

"io"

"log"

"os"

"regexp"

"strconv"

"strings"

"time"

"cloud.google.com/go/internal/trace"

sppb "cloud.google.com/go/spanner/apiv1/spannerpb"

"github.com/GoogleCloudPlatform/grpc-gcp-go/grpcgcp"

grpcgcppb "github.com/GoogleCloudPlatform/grpc-gcp-go/grpcgcp/grpc_gcp"

"github.com/googleapis/gax-go/v2"

"go.opentelemetry.io/otel/attribute"

"go.opentelemetry.io/otel/metric"

"go.opentelemetry.io/otel/metric/noop"

otrace "go.opentelemetry.io/otel/trace"

"google.golang.org/api/iterator"

"google.golang.org/api/option"

"google.golang.org/api/option/internaloption"

gtransport "google.golang.org/api/transport/grpc"

"google.golang.org/grpc"

"google.golang.org/grpc/codes"

"google.golang.org/grpc/credentials/insecure"

"google.golang.org/grpc/encoding/gzip"

"google.golang.org/grpc/keepalive"

"google.golang.org/grpc/metadata"

"google.golang.org/grpc/peer"

"google.golang.org/grpc/status"

vkit "cloud.google.com/go/spanner/apiv1"

"cloud.google.com/go/spanner/internal"

)

const (

// resourcePrefixHeader is the name of the metadata header used to indicate

// the resource being operated on.

resourcePrefixHeader = "google-cloud-resource-prefix"

// routeToLeaderHeader is the name of the metadata header if RW/PDML

// requests need to route to leader.

routeToLeaderHeader = "x-goog-spanner-route-to-leader"

requestsCompressionHeader = "x-response-encoding"

// endToEndTracingHeader is the name of the metadata header if client

// has opted-in for the creation of trace spans on the Spanner layer.

endToEndTracingHeader = "x-goog-spanner-end-to-end-tracing"

// afeMetricHeader is the name of the metadata header if client

// has opted-in for the receiving Spanner API Frontend server timing metrics.

afeMetricHeader = "x-goog-spanner-enable-afe-server-timing"

// numChannels is the default value for NumChannels of client.

numChannels = 4

// Server timing header constants

serverTimingHeaderKey = "server-timing"

gfeTimingHeader = "gfet4t7"

afeTimingHeader = "afe"

// MinSessions for Experimental Host connection

experimentalHostMinSessions = 0

)

const (

// Scope is the scope for Cloud Spanner Data API.

Scope = "https://www.googleapis.com/auth/spanner.data"

// AdminScope is the scope for Cloud Spanner Admin APIs.

AdminScope = "https://www.googleapis.com/auth/spanner.admin"

)

var (

validDBPattern = regexp.MustCompile("^projects/(?P<project>[^/]+)/instances/(?P<instance>[^/]+)/databases/(?P<database>[^/]+)$")

serverTimingPattern = regexp.MustCompile(`([a-zA-Z0-9_-]+);\s*dur=(\d*\.?\d+)`)

)

func validDatabaseName(db string) error {

if matched := validDBPattern.MatchString(db); !matched {

return fmt.Errorf("database name %q should conform to pattern %q",

db, validDBPattern.String())

}

return nil

}

func parseDatabaseName(db string) (project, instance, database string, err error) {

matches := validDBPattern.FindStringSubmatch(db)

if len(matches) == 0 {

return "", "", "", fmt.Errorf("failed to parse database name from %q according to pattern %q",

db, validDBPattern.String())

}

return matches[1], matches[2], matches[3], nil

}

// Client is a client for reading and writing data to a Cloud Spanner database.

// A client is safe to use concurrently, except for its Close method.

type Client struct {

sc *sessionClient

sm *sessionManager

logger *log.Logger

qo QueryOptions

ro ReadOptions

ao []ApplyOption

txo TransactionOptions

bwo BatchWriteOptions

ct *commonTags

disableRouteToLeader bool

dro *sppb.DirectedReadOptions

otConfig *openTelemetryConfig

metricsTracerFactory *builtinMetricsTracerFactory

clientContext *sppb.RequestOptions_ClientContext

locationRouter *locationRouter

}

// DatabaseName returns the full name of a database, e.g.,

// "projects/spanner-cloud-test/instances/foo/databases/foodb".

func (c *Client) DatabaseName() string {

return c.sc.database

}

// ClientID returns the id of the Client. This is not recommended for customer applications and used internally for testing.

func (c *Client) ClientID() string {

return c.sc.id

}

func createGCPMultiEndpoint(cfg *grpcgcp.GCPMultiEndpointOptions, config ClientConfig, opts ...option.ClientOption) (*grpcgcp.GCPMultiEndpoint, error) {

if cfg.GRPCgcpConfig == nil {

cfg.GRPCgcpConfig = &grpcgcppb.ApiConfig{}

}

if len(cfg.GRPCgcpConfig.Method) == 0 {

cfg.GRPCgcpConfig.Method = []*grpcgcppb.MethodConfig{

{

Name: []string{"/google.spanner.v1.Spanner/CreateSession"},

Affinity: &grpcgcppb.AffinityConfig{

Command: grpcgcppb.AffinityConfig_BIND,

AffinityKey: "name",

},

},

{

Name: []string{"/google.spanner.v1.Spanner/BatchCreateSessions"},

Affinity: &grpcgcppb.AffinityConfig{

Command: grpcgcppb.AffinityConfig_BIND,

AffinityKey: "session.name",

},

},

{

Name: []string{"/google.spanner.v1.Spanner/DeleteSession"},

Affinity: &grpcgcppb.AffinityConfig{

Command: grpcgcppb.AffinityConfig_UNBIND,

AffinityKey: "name",

},

},

{

Name: []string{"/google.spanner.v1.Spanner/GetSession"},

Affinity: &grpcgcppb.AffinityConfig{

Command: grpcgcppb.AffinityConfig_BOUND,

AffinityKey: "name",

},

},

{

Name: []string{

"/google.spanner.v1.Spanner/BeginTransaction",

"/google.spanner.v1.Spanner/Commit",

"/google.spanner.v1.Spanner/ExecuteBatchDml",

"/google.spanner.v1.Spanner/ExecuteSql",

"/google.spanner.v1.Spanner/ExecuteStreamingSql",

"/google.spanner.v1.Spanner/PartitionQuery",

"/google.spanner.v1.Spanner/PartitionRead",

"/google.spanner.v1.Spanner/Read",

"/google.spanner.v1.Spanner/Rollback",

"/google.spanner.v1.Spanner/StreamingRead",

},

Affinity: &grpcgcppb.AffinityConfig{

Command: grpcgcppb.AffinityConfig_BOUND,

AffinityKey: "session",

},

},

}

}

// Append emulator options if SPANNER_EMULATOR_HOST has been set.

if emulatorAddr := os.Getenv("SPANNER_EMULATOR_HOST"); emulatorAddr != "" {

schemeRemoved := regexp.MustCompile("^(https://|https://|passthrough:///)").ReplaceAllString(emulatorAddr, "")

emulatorOpts := []option.ClientOption{

option.WithEndpoint("passthrough:///" + schemeRemoved),

option.WithGRPCDialOption(grpc.WithTransportCredentials(insecure.NewCredentials())),

option.WithoutAuthentication(),

internaloption.SkipDialSettingsValidation(),

}

opts = append(opts, emulatorOpts...)

// Replace all endpoints with emulator target.

for _, meo := range cfg.MultiEndpoints {

meo.Endpoints = []string{emulatorAddr}

}

}

// Set the number of channels to the default value if not specified.

if cfg.GRPCgcpConfig.GetChannelPool() == nil || cfg.GRPCgcpConfig.GetChannelPool().GetMaxSize() == 0 {

cfg.GRPCgcpConfig.ChannelPool = &grpcgcppb.ChannelPoolConfig{

MinSize: numChannels,

MaxSize: numChannels,

}

}

// Set MinSize equal to MaxSize to create all the channels beforehand.

cfg.GRPCgcpConfig.ChannelPool.MinSize = cfg.GRPCgcpConfig.ChannelPool.GetMaxSize()

cfg.GRPCgcpConfig.ChannelPool.BindPickStrategy = grpcgcppb.ChannelPoolConfig_ROUND_ROBIN

cfg.DialFunc = func(ctx context.Context, target string, dopts ...grpc.DialOption) (*grpc.ClientConn, error) {

copts := opts

for _, do := range dopts {

copts = append(copts, option.WithGRPCDialOption(do))

}

allOpts := allClientOpts(1, config.Compression, copts...)

// Overwrite endpoint and pool config.

allOpts = append(allOpts,

option.WithEndpoint(target),

option.WithGRPCConnectionPool(1),

option.WithGRPCConn(nil),

)

return gtransport.Dial(ctx, allOpts...)

}

gme, err := grpcgcp.NewGCPMultiEndpoint(cfg)

return gme, err

}

// To use GCPMultiEndpoint in gtransport.Dial (via gtransport.WithConnPool option)

// we implement gtransport.ConnPool interface using this wrapper.

type gmeWrapper struct {

*grpcgcp.GCPMultiEndpoint

}

// Make sure gmeWrapper implements ConnPool interface.

var _ gtransport.ConnPool = (*gmeWrapper)(nil)

func (gw *gmeWrapper) Conn() *grpc.ClientConn {

// GCPMultiEndpoint does not expose any ClientConn.

// This is safe because Cloud Spanner client doesn't use this function and instead

// makes calls directly using Invoke and NewStream from the grpc.ClientConnInterface

// which GCPMultiEndpoint implements.

return nil

}

func (gw *gmeWrapper) Num() int {

return int(gw.GCPMultiEndpoint.GCPConfig().GetChannelPool().GetMaxSize())

}

// ClientConfig has configurations for the client.

type ClientConfig struct {

// NumChannels is the number of gRPC channels.

// If zero, a reasonable default is used based on the execution environment.

//

// Deprecated: The Spanner client now uses a pool of gRPC connections. Use

// option.WithGRPCConnectionPool(numConns) instead to specify the number of

// connections the client should use. The client will default to a

// reasonable default if this option is not specified.

NumChannels int

// SessionPoolConfig is the configuration for session pool.

SessionPoolConfig

// SessionLabels for the sessions created by this client.

// See https://cloud.google.com/spanner/docs/reference/rpc/google.spanner.v1#session

// for more info.

SessionLabels map[string]string

// QueryOptions is the configuration for executing a sql query.

QueryOptions QueryOptions

// ReadOptions is the configuration for reading rows from a database

ReadOptions ReadOptions

// ApplyOptions is the configuration for applying

ApplyOptions []ApplyOption

// TransactionOptions is the configuration for a transaction.

TransactionOptions TransactionOptions

// BatchWriteOptions is the configuration for a BatchWrite request.

BatchWriteOptions BatchWriteOptions

// CallOptions is the configuration for providing custom retry settings that

// override the default values.

CallOptions *vkit.CallOptions

// UserAgent is the prefix to the user agent header. This is used to supply information

// such as application name or partner tool.

//

// Internal Use Only: This field is for internal tracking purpose only,

// setting the value for this config is not required.

//

// Recommended format: ``application-or-tool-ID/major.minor.version``.

UserAgent string

// DatabaseRole specifies the role to be assumed for all operations on the

// database by this client.

DatabaseRole string

// DisableRouteToLeader specifies if all the requests of type read-write and PDML

// need to be routed to the leader region.

//

// Default: false

DisableRouteToLeader bool

// Logger is the logger to use for this client. If it is nil, all logging

// will be directed to the standard logger.

Logger *log.Logger

//

// Sets the compression to use for all gRPC calls. The compressor must be a valid name.

// This will enable compression both from the client to the

// server and from the server to the client.

//

// Supported values are:

// gzip: Enable gzip compression

// identity: Disable compression

//

// Default: identity

Compression string

// BatchTimeout specifies the timeout for a batch of sessions managed sessionClient.

BatchTimeout time.Duration

// ClientConfig options used to set the DirectedReadOptions for all ReadRequests

// and ExecuteSqlRequests for the Client which indicate which replicas or regions

// should be used for non-transactional reads or queries.

DirectedReadOptions *sppb.DirectedReadOptions

OpenTelemetryMeterProvider metric.MeterProvider

// EnableEndToEndTracing indicates whether end to end tracing is enabled or not. If

// it is enabled, trace spans will be created at Spanner layer. Enabling end to end

// tracing requires OpenTelemetry to be set up. Simply enabling this option won't

// generate traces at Spanner layer.

//

// Default: false

EnableEndToEndTracing bool

// DisableNativeMetrics indicates whether native metrics should be disabled or not.

// If true, native metrics will not be emitted.

//

// Default: false

DisableNativeMetrics bool

// Default: false

IsExperimentalHost bool

// ClientContext is the default context for all requests made by the client.

ClientContext *sppb.RequestOptions_ClientContext

}

type openTelemetryConfig struct {

enabled bool

meterProvider metric.MeterProvider

commonTraceStartOptions []otrace.SpanStartOption

attributeMap []attribute.KeyValue

attributeMapWithMultiplexed []attribute.KeyValue

otMetricRegistration metric.Registration

openSessionCount metric.Int64ObservableGauge

getSessionTimeoutsCount metric.Int64Counter

acquiredSessionsCount metric.Int64Counter

releasedSessionsCount metric.Int64Counter

gfeLatency metric.Int64Histogram

gfeHeaderMissingCount metric.Int64Counter

}

func contextWithOutgoingMetadata(ctx context.Context, md metadata.MD, disableRouteToLeader bool) context.Context {

existing, ok := metadata.FromOutgoingContext(ctx)

if ok {

// Make sure that we only send one resource header.

existing.Delete(resourcePrefixHeader)

md = metadata.Join(existing, md)

}

if !disableRouteToLeader {

md = metadata.Join(md, metadata.Pairs(routeToLeaderHeader, "true"))

}

return metadata.NewOutgoingContext(ctx, md)

}

// NewClient creates a client to a database. A valid database name has the

// form projects/PROJECT_ID/instances/INSTANCE_ID/databases/DATABASE_ID. It uses

// a default configuration.

func NewClient(ctx context.Context, database string, opts ...option.ClientOption) (*Client, error) {

return NewClientWithConfig(ctx, database, ClientConfig{SessionPoolConfig: DefaultSessionPoolConfig, DisableRouteToLeader: false}, opts...)

}

// NewClientWithConfig creates a client to a database. A valid database name has

// the form projects/PROJECT_ID/instances/INSTANCE_ID/databases/DATABASE_ID.

func NewClientWithConfig(ctx context.Context, database string, config ClientConfig, opts ...option.ClientOption) (c *Client, err error) {

return newClientWithConfig(ctx, database, config, nil, opts...)

}

type fallbackWrapper struct {

*grpcgcp.GCPFallback

primaryConn gtransport.ConnPool

fallbackConn gtransport.ConnPool

}

// Conn returns nil because GCPFallback hides the underlying ClientConn.

// The Spanner client handles this by using the interface methods (Invoke/NewStream).

func (fw *fallbackWrapper) Conn() *grpc.ClientConn {

return nil

}

func (fw *fallbackWrapper) Num() int {

return fw.primaryConn.Num()

}

func (fw *fallbackWrapper) Close() error {

fw.GCPFallback.Close()

err1 := fw.primaryConn.Close()

err2 := fw.fallbackConn.Close()

if err1 != nil {

return err1

}

return err2

}

func newClientWithConfig(ctx context.Context, database string, config ClientConfig, gme *grpcgcp.GCPMultiEndpoint, opts ...option.ClientOption) (c *Client, err error) {

// Validate database path.

if err := validDatabaseName(database); err != nil {

return nil, err

}

ctx, _ = startSpan(ctx, "NewClient")

defer func() { endSpan(ctx, err) }()

// Explicitly disable some gRPC experiments as they are not stable yet.

gRPCPickFirstEnvVarName := "GRPC_EXPERIMENTAL_ENABLE_NEW_PICK_FIRST"

if os.Getenv(gRPCPickFirstEnvVarName) == "" {

err := os.Setenv(gRPCPickFirstEnvVarName, "false")

if err != nil {

logf(config.Logger, "Error overriding GRPC_EXPERIMENTAL_ENABLE_NEW_PICK_FIRST to false: %v. Ignoring.", err)

}

}

// Append emulator options if SPANNER_EMULATOR_HOST has been set.

if emulatorAddr := os.Getenv("SPANNER_EMULATOR_HOST"); emulatorAddr != "" {

schemeRemoved := regexp.MustCompile("^(https://|https://|passthrough:///)").ReplaceAllString(emulatorAddr, "")

emulatorOpts := []option.ClientOption{

option.WithEndpoint("passthrough:///" + schemeRemoved),

option.WithGRPCDialOption(grpc.WithTransportCredentials(insecure.NewCredentials())),

option.WithoutAuthentication(),

internaloption.SkipDialSettingsValidation(),

}

opts = append(emulatorOpts, opts...)

}

// Prepare gRPC channels.

hasNumChannelsConfig := config.NumChannels > 0

if config.NumChannels == 0 {

config.NumChannels = numChannels

}

var metricsProvider metric.MeterProvider

if emulatorAddr := os.Getenv("SPANNER_EMULATOR_HOST"); emulatorAddr != "" {

// Do not emit native metrics when emulator is being used

metricsProvider = noop.NewMeterProvider()

}

// Check if native metrics are disabled via env.

if disableNativeMetrics, _ := strconv.ParseBool(os.Getenv("SPANNER_DISABLE_BUILTIN_METRICS")); disableNativeMetrics {

config.DisableNativeMetrics = true

}

if config.DisableNativeMetrics {

// Do not emit native metrics when DisableNativeMetrics is set

metricsProvider = noop.NewMeterProvider()

}

isAFEBuiltInMetricEnabled := strings.EqualFold("false", os.Getenv("SPANNER_DISABLE_AFE_SERVER_TIMING"))

isGRPCBuiltInMetricsEnabled := strings.EqualFold("false", os.Getenv("SPANNER_DISABLE_DIRECT_ACCESS_GRPC_BUILTIN_METRICS"))

// enable the AFE/GRPC built-in metrics if direct-path is enabled

isDirectPathEnabled, _ := strconv.ParseBool(os.Getenv("GOOGLE_SPANNER_ENABLE_DIRECT_ACCESS"))

if isDirectPathEnabled {

isAFEBuiltInMetricEnabled = true

isGRPCBuiltInMetricsEnabled = true

}

// disable the AFE/GRPC built-in metrics if the env var is explicitly set

if ok, _ := strconv.ParseBool(os.Getenv("SPANNER_DISABLE_AFE_SERVER_TIMING")); ok {

isAFEBuiltInMetricEnabled = false

}

if ok, _ := strconv.ParseBool(os.Getenv("SPANNER_DISABLE_DIRECT_ACCESS_GRPC_BUILTIN_METRICS")); ok {

isGRPCBuiltInMetricsEnabled = false

}

metricsTracerFactory, err := newBuiltinMetricsTracerFactory(ctx, database, config.Compression, isAFEBuiltInMetricEnabled, isGRPCBuiltInMetricsEnabled, metricsProvider, opts...)

if err != nil {

return nil, err

}

if len(metricsTracerFactory.clientOpts) > 0 {

opts = append(opts, metricsTracerFactory.clientOpts...)

}

var pool gtransport.ConnPool

var endpointClientOpts []option.ClientOption

if gme != nil {

// Use GCPMultiEndpoint if provided.

pool = &gmeWrapper{gme}

endpointClientOpts = append(endpointClientOpts, opts...)

} else if isFallbackEnabled, _ := strconv.ParseBool(os.Getenv("GOOGLE_SPANNER_ENABLE_GCP_FALLBACK")); isFallbackEnabled && isDirectPathEnabled {

var primaryConn gtransport.ConnPool

var fallbackConn gtransport.ConnPool

reqIDInjector := new(requestIDHeaderInjector)

opts = append(opts,

option.WithGRPCDialOption(grpc.WithChainStreamInterceptor(reqIDInjector.interceptStream)),

option.WithGRPCDialOption(grpc.WithChainUnaryInterceptor(reqIDInjector.interceptUnary)),

)

allOpts := allClientOpts(config.NumChannels, config.Compression, opts...)

endpointClientOpts = append(endpointClientOpts, allOpts...)

primaryConn, err = gtransport.DialPool(ctx, allOpts...)

if err != nil {

return nil, err

}

fallbackConnOpts := append(allOpts, internaloption.EnableDirectPath(false))

fallbackConn, err = gtransport.DialPool(ctx, fallbackConnOpts...)

if err != nil {

primaryConn.Close()

return nil, err

}

if hasNumChannelsConfig && ((primaryConn.Num() != config.NumChannels) || (fallbackConn.Num() != config.NumChannels)) {

primaryConn.Close()

fallbackConn.Close()

return nil, spannerErrorf(codes.InvalidArgument, "Connection pool mismatch: NumChannels=%v, primaryConn.Num()=%v, fallbackConn.Num()=%v", config.NumChannels, primaryConn.Num(), fallbackConn.Num())

}

fbOpts := grpcgcp.NewGCPFallbackOptions()

fbOpts.EnableFallback = true

fbOpts.ErrorRateThreshold = 1

fbOpts.MinFailedCalls = 1

fbOpts.MeterProvider = config.OpenTelemetryMeterProvider

gcpFallback, err := grpcgcp.NewGCPFallback(ctx, primaryConn, fallbackConn, fbOpts)

if err != nil {

primaryConn.Close()

fallbackConn.Close()

return nil, err

}

pool = &fallbackWrapper{gcpFallback, primaryConn, fallbackConn}

} else {

// Create gtransport ConnPool as usual if MultiEndpoint is not used.

// gRPC options.

// Add a unaryClientInterceptor and streamClientInterceptor.

reqIDInjector := new(requestIDHeaderInjector)

opts = append(opts,

option.WithGRPCDialOption(grpc.WithChainStreamInterceptor(reqIDInjector.interceptStream)),

option.WithGRPCDialOption(grpc.WithChainUnaryInterceptor(reqIDInjector.interceptUnary)),

)

allOpts := allClientOpts(config.NumChannels, config.Compression, opts...)

endpointClientOpts = append(endpointClientOpts, allOpts...)

pool, err = gtransport.DialPool(ctx, allOpts...)

if err != nil {

return nil, err

}

if hasNumChannelsConfig && pool.Num() != config.NumChannels {

pool.Close()

return nil, spannerErrorf(codes.InvalidArgument, "Connection pool mismatch: NumChannels=%v, WithGRPCConnectionPool=%v. Only set one of these options, or set both to the same value.", config.NumChannels, pool.Num())

}

}

// TODO(loite): Remove as the original map cannot be changed by the user

// anyways, and the client library is also not changing it.

// Make a copy of labels.

sessionLabels := make(map[string]string)

for k, v := range config.SessionLabels {

sessionLabels[k] = v

}

// Default configs for session pool.

if config.MaxOpened == 0 {

config.MaxOpened = uint64(pool.Num() * 100)

}

if config.MaxBurst == 0 {

config.MaxBurst = DefaultSessionPoolConfig.MaxBurst

}

if config.BatchTimeout == 0 {

config.BatchTimeout = time.Minute

}

md := metadata.Pairs(resourcePrefixHeader, database)

if config.Compression == gzip.Name {

md.Append(requestsCompressionHeader, gzip.Name)

}

// Append end to end tracing header if SPANNER_ENABLE_END_TO_END_TRACING

// environment variable has been set or client has passed the opt-in

// option in ClientConfig.

endToEndTracingEnvironmentVariable := os.Getenv("SPANNER_ENABLE_END_TO_END_TRACING")

if config.EnableEndToEndTracing || endToEndTracingEnvironmentVariable == "true" {

md.Append(endToEndTracingHeader, "true")

}

if isAFEBuiltInMetricEnabled {

md.Append(afeMetricHeader, "true")

}

// Multiplexed sessions are always enabled as the session pool has been removed.

// Create a session client.

sc := newSessionClient(pool, database, config.UserAgent, sessionLabels, config.DatabaseRole, config.DisableRouteToLeader, md, config.BatchTimeout, config.Logger, config.CallOptions)

// Create an OpenTelemetry configuration

otConfig, err := createOpenTelemetryConfig(ctx, config.OpenTelemetryMeterProvider, config.Logger, sc.id, database)

if err != nil {

// The error returned here will be due to database name parsing

return nil, err

}

// To prevent data race in unit tests (ex: TestClient_SessionNotFound)

sc.mu.Lock()

sc.otConfig = otConfig

sc.metricsTracerFactory = metricsTracerFactory

sc.mu.Unlock()

var locationRouter *locationRouter

if isExperimentalLocationAPIEnabled() {

sc.baseClientOpts = endpointClientOpts

if conn := pool.Conn(); conn != nil {

sc.endpointAuthority = normalizeAuthorityTarget(conn.Target())

}

epCache := newEndpointClientCache(sc.createEndpointClient)

locationRouter = newLocationRouter(epCache)

}

// Create a session manager.

sp, err := newSessionManager(sc, config.SessionPoolConfig)

if err != nil {

sc.close()

return nil, err

}

if locationRouter != nil {

sp.locationRouter = locationRouter

}

if enableLogClientOptions() {

projectID, _, _, _ := parseDatabaseName(database)

logf(config.Logger, `

-----Client Options--------

Project ID: %v

Num of gRPC channels: %v

LAR enabled: %v

Directpath enabled: %v

End To End Tracing enabled: %v

Built-in metrics enabled: %v

gRPC metrics enabled: %v

Multiplexed session enabled: true

-----------------------------`,

projectID, config.NumChannels, !config.DisableRouteToLeader, isDirectPathEnabled,

config.EnableEndToEndTracing, !config.DisableNativeMetrics, isGRPCBuiltInMetricsEnabled)

}

c = &Client{

sc: sc,

sm: sp,

logger: config.Logger,

qo: getQueryOptions(config.QueryOptions),

ro: config.ReadOptions,

ao: config.ApplyOptions,

txo: config.TransactionOptions,

bwo: config.BatchWriteOptions,

ct: getCommonTags(sc),

disableRouteToLeader: config.DisableRouteToLeader,

dro: config.DirectedReadOptions,

otConfig: otConfig,

metricsTracerFactory: metricsTracerFactory,

clientContext: config.ClientContext,

locationRouter: locationRouter,

}

return c, nil

}

func normalizeAuthorityTarget(target string) string {

if idx := strings.Index(target, ":///"); idx >= 0 {

return strings.TrimSuffix(target[idx+4:], "/")

}

if idx := strings.Index(target, "://"); idx >= 0 {

return strings.TrimSuffix(target[idx+3:], "/")

}

return strings.TrimSuffix(target, "/")

}

// NewMultiEndpointClient is the same as NewMultiEndpointClientWithConfig with

// the default client configuration.

//

// A valid database name has the

// form projects/PROJECT_ID/instances/INSTANCE_ID/databases/DATABASE_ID.

func NewMultiEndpointClient(ctx context.Context, database string, gmeCfg *grpcgcp.GCPMultiEndpointOptions, opts ...option.ClientOption) (*Client, *grpcgcp.GCPMultiEndpoint, error) {

return NewMultiEndpointClientWithConfig(ctx, database, ClientConfig{SessionPoolConfig: DefaultSessionPoolConfig, DisableRouteToLeader: false}, gmeCfg, opts...)

}

// NewMultiEndpointClientWithConfig creates a client to a database using GCPMultiEndpoint.

//

// The purposes of GCPMultiEndpoint are:

//

// - Fallback to an alternative endpoint (host:port) when the original

// endpoint is completely unavailable.

// - Be able to route a Cloud Spanner call to a specific group of endpoints.

// - Be able to reconfigure endpoints in runtime.

//

// The GRPCgcpConfig and DialFunc in the GCPMultiEndpointOptions are optional

// and will be configured automatically.

//

// For GCPMultiEndpoint the number of channels is configured via MaxSize of the

// ChannelPool config in the GRPCgcpConfig.

//

// The GCPMultiEndpoint returned can be used to update the endpoints in runtime.

//

// A valid database name has the

// form projects/PROJECT_ID/instances/INSTANCE_ID/databases/DATABASE_ID.

func NewMultiEndpointClientWithConfig(ctx context.Context, database string, config ClientConfig, gmeCfg *grpcgcp.GCPMultiEndpointOptions, opts ...option.ClientOption) (c *Client, gme *grpcgcp.GCPMultiEndpoint, err error) {

gme, err = createGCPMultiEndpoint(gmeCfg, config, opts...)

if err != nil {

return nil, nil, err

}

// Align number of channels.

config.NumChannels = int(gme.GCPConfig().GetChannelPool().GetMaxSize())

c, err = newClientWithConfig(ctx, database, config, gme, opts...)

if err != nil {

return nil, nil, err

}

return

}

// Combines the default options from the generated client, the default options

// of the hand-written client and the user options to one list of options.

// Precedence: userOpts > clientDefaultOpts > generatedDefaultOpts

func allClientOpts(numChannels int, compression string, userOpts ...option.ClientOption) []option.ClientOption {

generatedDefaultOpts := vkit.DefaultClientOptions()

clientDefaultOpts := []option.ClientOption{

option.WithGRPCConnectionPool(numChannels),

option.WithUserAgent(fmt.Sprintf("spanner-go/v%s", internal.Version)),

option.WithGRPCDialOption(grpc.WithChainUnaryInterceptor(addNativeMetricsInterceptor()...)),

option.WithGRPCDialOption(grpc.WithChainStreamInterceptor(addStreamNativeMetricsInterceptor()...)),

option.WithGRPCDialOption(grpc.WithKeepaliveParams(keepalive.ClientParameters{Time: 120 * time.Second})),

}

if enableDirectPathXds, _ := strconv.ParseBool(os.Getenv("GOOGLE_SPANNER_ENABLE_DIRECT_ACCESS")); enableDirectPathXds {

clientDefaultOpts = append(clientDefaultOpts, internaloption.AllowNonDefaultServiceAccount(true))

clientDefaultOpts = append(clientDefaultOpts, internaloption.EnableDirectPath(true), internaloption.EnableDirectPathXds())

if disableBoundToken, _ := strconv.ParseBool(os.Getenv("GOOGLE_SPANNER_DISABLE_DIRECT_ACCESS_BOUND_TOKEN")); !disableBoundToken {

clientDefaultOpts = append(clientDefaultOpts, internaloption.AllowHardBoundTokens("ALTS"))

}

}

if compression == "gzip" {

userOpts = append(userOpts, option.WithGRPCDialOption(grpc.WithDefaultCallOptions(

grpc.UseCompressor(gzip.Name))))

}

allDefaultOpts := append(generatedDefaultOpts, clientDefaultOpts...)

return append(allDefaultOpts, userOpts...)

}

// metricsInterceptor is a gRPC unary client interceptor that records metrics for unary RPCs.

func metricsInterceptor() grpc.UnaryClientInterceptor {

return func(

ctx context.Context,

method string,

req interface{},

reply interface{},

cc *grpc.ClientConn,

invoker grpc.UnaryInvoker,

opts ...grpc.CallOption,

) error {

mt, ok := ctx.Value(metricsTracerKey).(*builtinMetricsTracer)

if !ok {

return invoker(ctx, method, req, reply, cc, opts...)

}

mt.method = method

mt.currOp.incrementAttemptCount()

mt.currOp.currAttempt = &attemptTracer{}

mt.currOp.currAttempt.setStartTime(time.Now())

if strings.HasPrefix(cc.Target(), "google-c2p") {

mt.currOp.setDirectPathEnabled(true)

}

var md metadata.MD

peerInfo := &peer.Peer{}

opts = append(opts, grpc.Header(&md), grpc.Peer(peerInfo))

err := invoker(ctx, method, req, reply, cc, opts...)

statusCode, _ := status.FromError(err)

mt.currOp.currAttempt.setStatus(statusCode.Code().String())

mt.currOp.currAttempt.setDirectPathUsed(peer.NewContext(ctx, peerInfo))

latencies := parseServerTimingHeader(md)

span := otrace.SpanFromContext(ctx)

setGFEAndAFESpanAttributes(span, latencies)

mt.currOp.currAttempt.setServerTimingMetrics(latencies)

recordAttemptCompletion(mt)

return err

}

}

// metricsInterceptor is a gRPC stream client interceptor that records metrics for stream RPCs.

func metricsStreamInterceptor() grpc.StreamClientInterceptor {

return func(

ctx context.Context,

desc *grpc.StreamDesc,

cc *grpc.ClientConn,

method string,

streamer grpc.Streamer,

opts ...grpc.CallOption,

) (grpc.ClientStream, error) {

s, err := streamer(ctx, desc, cc, method, opts...)

if err != nil {

return nil, err

}

mt, ok := ctx.Value(metricsTracerKey).(*builtinMetricsTracer)

if ok && mt != nil {

mt.method = method

if strings.HasPrefix(cc.Target(), "google-c2p") {

mt.currOp.setDirectPathEnabled(true)

}

}

return s, nil

}

}

// AddNativeMetricsInterceptor intercepts unary requests and records metrics for them.

func addNativeMetricsInterceptor() []grpc.UnaryClientInterceptor {

unaryInterceptors := []grpc.UnaryClientInterceptor{}

unaryInterceptors = append(unaryInterceptors, metricsInterceptor())

return unaryInterceptors

}

// AddStreamNativeMetricsInterceptor intercepts stream requests and records metrics for them.

func addStreamNativeMetricsInterceptor() []grpc.StreamClientInterceptor {

streamInterceptors := []grpc.StreamClientInterceptor{}

streamInterceptors = append(streamInterceptors, metricsStreamInterceptor())

return streamInterceptors

}

// getQueryOptions returns the query options overwritten by the environment

// variables if exist. The input parameter is the query options set by users

// via application-level configuration. If the environment variables are set,

// this will return the overwritten query options.

func getQueryOptions(opts QueryOptions) QueryOptions {

if opts.Options == nil {

opts.Options = &sppb.ExecuteSqlRequest_QueryOptions{}

}

opv := os.Getenv("SPANNER_OPTIMIZER_VERSION")

if opv != "" {

opts.Options.OptimizerVersion = opv

}

opsp := os.Getenv("SPANNER_OPTIMIZER_STATISTICS_PACKAGE")

if opsp != "" {

opts.Options.OptimizerStatisticsPackage = opsp

}

return opts

}

// Close closes the client.

func (c *Client) Close() {

if c.metricsTracerFactory != nil {

c.metricsTracerFactory.shutdown(context.Background())

}

if c.sm != nil {

ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)

defer cancel()

c.sm.close(ctx)

}

if c.locationRouter != nil {

c.locationRouter.Close()

}

c.sc.close()

}

// Single provides a read-only snapshot transaction optimized for the case

// where only a single read or query is needed. This is more efficient than

// using ReadOnlyTransaction() for a single read or query.

//

// Single will use a strong TimestampBound by default. Use

// ReadOnlyTransaction.WithTimestampBound to specify a different

// TimestampBound. A non-strong bound can be used to reduce latency, or

// "time-travel" to prior versions of the database, see the documentation of

// TimestampBound for details.

func (c *Client) Single() *ReadOnlyTransaction {

t := &ReadOnlyTransaction{singleUse: true}

t.txReadOnly.sm = c.sm

t.txReadOnly.txReadEnv = t

t.txReadOnly.qo = c.qo

t.txReadOnly.ro = c.ro

t.txReadOnly.disableRouteToLeader = true

t.txReadOnly.qo.DirectedReadOptions = c.dro

t.txReadOnly.ro.DirectedReadOptions = c.dro

t.txReadOnly.ro.LockHint = sppb.ReadRequest_LOCK_HINT_UNSPECIFIED

t.txReadOnly.clientContext = c.clientContext

t.ct = c.ct

t.otConfig = c.otConfig

return t

}

// ReadOnlyTransaction returns a ReadOnlyTransaction that can be used for

// multiple reads from the database. You must call Close() when the

// ReadOnlyTransaction is no longer needed to release resources on the server.

//

// ReadOnlyTransaction will use a strong TimestampBound by default. Use

// ReadOnlyTransaction.WithTimestampBound to specify a different

// TimestampBound. A non-strong bound can be used to reduce latency, or

// "time-travel" to prior versions of the database, see the documentation of

// TimestampBound for details.

func (c *Client) ReadOnlyTransaction() *ReadOnlyTransaction {

t := &ReadOnlyTransaction{

singleUse: false,

txReadyOrClosed: make(chan struct{}),

}

t.txReadOnly.sm = c.sm

t.txReadOnly.txReadEnv = t

t.txReadOnly.qo = c.qo

t.txReadOnly.ro = c.ro

t.txReadOnly.disableRouteToLeader = true

t.txReadOnly.qo.DirectedReadOptions = c.dro

t.txReadOnly.ro.DirectedReadOptions = c.dro

t.txReadOnly.ro.LockHint = sppb.ReadRequest_LOCK_HINT_UNSPECIFIED

t.txReadOnly.clientContext = c.clientContext

t.ct = c.ct

t.otConfig = c.otConfig

return t

}

// BatchReadOnlyTransaction returns a BatchReadOnlyTransaction that can be used

// for partitioned reads or queries from a snapshot of the database. This is

// useful in batch processing pipelines where one wants to divide the work of

// reading from the database across multiple machines.

//

// Note: This transaction does not use the underlying session pool but creates a

// new session each time, and the session is reused across clients.

//

// You should call Close() after the txn is no longer needed on local

// client, and call Cleanup() when the txn is finished for all clients, to free

// the session.

func (c *Client) BatchReadOnlyTransaction(ctx context.Context, tb TimestampBound) (*BatchReadOnlyTransaction, error) {

var (

tx transactionID

rts time.Time

sh *sessionHandle

err error

)

// Always use multiplexed sessions for batch read operations.

sh, err = c.sm.takeMultiplexed(ctx)

if err != nil {

return nil, err

}

// Begin transaction.

res, err := sh.getClient().BeginTransaction(contextWithOutgoingMetadata(ctx, sh.getMetadata(), true), &sppb.BeginTransactionRequest{

Session: sh.getID(),

Options: &sppb.TransactionOptions{

Mode: &sppb.TransactionOptions_ReadOnly_{

ReadOnly: buildTransactionOptionsReadOnly(tb, true),

},

},

RequestOptions: createRequestOptions(sppb.RequestOptions_PRIORITY_UNSPECIFIED, "", "", c.clientContext),

})

if err != nil {

return nil, ToSpannerError(err)

}

tx = res.Id

if res.ReadTimestamp != nil {

rts = time.Unix(res.ReadTimestamp.Seconds, int64(res.ReadTimestamp.Nanos))

}

t := &BatchReadOnlyTransaction{

ReadOnlyTransaction: ReadOnlyTransaction{

tx: tx,

txReadyOrClosed: make(chan struct{}),

state: txActive,

rts: rts,

isLongRunningTransaction: true,

},

ID: BatchReadOnlyTransactionID{