// Copyright 2025 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 spannerdriver

import (

"context"

"database/sql"

"database/sql/driver"

"errors"

"fmt"

"log/slog"

"slices"

"sync"

"time"

"cloud.google.com/go/spanner"

adminapi "cloud.google.com/go/spanner/admin/database/apiv1"

adminpb "cloud.google.com/go/spanner/admin/database/apiv1/databasepb"

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

"github.com/googleapis/go-sql-spanner/connectionstate"

"github.com/googleapis/go-sql-spanner/parser"

"google.golang.org/api/iterator"

"google.golang.org/grpc/codes"

"google.golang.org/grpc/status"

)

// SpannerConn is the public interface for the raw Spanner connection for the

// sql driver. This interface can be used with the db.Conn().Raw() method.

type SpannerConn interface {

// StartBatchDDL starts a DDL batch on the connection. After calling this

// method all subsequent DDL statements will be cached locally. Calling

// RunBatch will send all cached DDL statements to Spanner as one batch.

// Use DDL batching to speed up the execution of multiple DDL statements.

// Note that a DDL batch is not atomic. It is possible that some DDL

// statements are executed successfully and some not.

// See https://cloud.google.com/spanner/docs/schema-updates#order_of_execution_of_statements_in_batches

// for more information on how Cloud Spanner handles DDL batches.

StartBatchDDL() error

// StartBatchDML starts a DML batch on the connection. After calling this

// method all subsequent DML statements will be cached locally. Calling

// RunBatch will send all cached DML statements to Spanner as one batch.

// Use DML batching to speed up the execution of multiple DML statements.

// DML batches can be executed both outside of a transaction and during

// a read/write transaction. If a DML batch is executed outside an active

// transaction, the batch will be applied atomically to the database if

// successful and rolled back if one or more of the statements fail.

// If a DML batch is executed as part of a transaction, the error will

// be returned to the application, and the application can decide whether

// to commit or rollback the transaction.

StartBatchDML() error

// RunBatch sends all batched DDL or DML statements to Spanner. This is a

// no-op if no statements have been batched or if there is no active batch.

RunBatch(ctx context.Context) error

// RunDmlBatch sends all batched DML statements to Spanner. This is a

// no-op if no statements have been batched or if there is no active DML batch.

RunDmlBatch(ctx context.Context) (SpannerResult, error)

// AbortBatch aborts the current DDL or DML batch and discards all batched

// statements.

AbortBatch() error

// InDDLBatch returns true if the connection is currently in a DDL batch.

InDDLBatch() bool

// InDMLBatch returns true if the connection is currently in a DML batch.

InDMLBatch() bool

// GetBatchedStatements returns a copy of the statements that are currently

// buffered to be executed as a DML or DDL batch. It returns an empty slice

// if no batch is active, or if there are no statements buffered.

GetBatchedStatements() []spanner.Statement

// AutoBatchDml determines whether DML statements should automatically

// be batched and sent to Spanner when a non-DML statement is encountered.

// The update count that is returned for DML statements that are buffered

// is by default 1. This default can be changed by setting the connection

// variable AutoBatchDmlUpdateCount to a value other than 1.

// This feature is only used in read/write transactions. DML statements

// outside transactions are always executed directly.

AutoBatchDml() bool

SetAutoBatchDml(autoBatch bool) error

// AutoBatchDmlUpdateCount determines the update count that is returned for

// DML statements that are executed when AutoBatchDml is true.

AutoBatchDmlUpdateCount() int64

SetAutoBatchDmlUpdateCount(updateCount int64) error

// AutoBatchDmlUpdateCountVerification enables/disables the verification

// that the update count that was returned for automatically batched DML

// statements was correct.

AutoBatchDmlUpdateCountVerification() bool

SetAutoBatchDmlUpdateCountVerification(verify bool) error

// RetryAbortsInternally returns true if the connection automatically

// retries all aborted transactions.

RetryAbortsInternally() bool

// SetRetryAbortsInternally enables/disables the automatic retry of aborted

// transactions. If disabled, any aborted error from a transaction will be

// propagated to the application.

SetRetryAbortsInternally(retry bool) error

// AutocommitDMLMode returns the current mode that is used for DML

// statements outside a transaction. The default is Transactional.

AutocommitDMLMode() AutocommitDMLMode

// SetAutocommitDMLMode sets the mode to use for DML statements that are

// executed outside transactions. The default is Transactional. Change to

// PartitionedNonAtomic to use Partitioned DML instead of Transactional DML.

// See https://cloud.google.com/spanner/docs/dml-partitioned for more

// information on Partitioned DML.

SetAutocommitDMLMode(mode AutocommitDMLMode) error

// ReadOnlyStaleness returns the current staleness that is used for

// queries in autocommit mode, and for read-only transactions.

ReadOnlyStaleness() spanner.TimestampBound

// SetReadOnlyStaleness sets the staleness to use for queries in autocommit

// mode and for read-only transaction.

SetReadOnlyStaleness(staleness spanner.TimestampBound) error

// IsolationLevel returns the current default isolation level that is

// used for read/write transactions on this connection.

IsolationLevel() sql.IsolationLevel

// SetIsolationLevel sets the default isolation level to use for read/write

// transactions on this connection.

SetIsolationLevel(level sql.IsolationLevel) error

// ReadLockMode returns the current read lock mode that is used for read/write

// transactions on this connection.

ReadLockMode() spannerpb.TransactionOptions_ReadWrite_ReadLockMode

// SetReadLockMode sets the read lock mode to use for read/write transactions

// on this connection.

//

// The read lock mode option controls the locking behavior for read operations and queries within a

// read-write transaction. It works in conjunction with the transaction's isolation level.

//

// PESSIMISTIC: Read locks are acquired immediately on read. This mode only applies to SERIALIZABLE

// isolation. This mode prevents concurrent modifications by locking data throughout the transaction.

// This reduces commit-time aborts due to conflicts but can increase how long transactions wait for

// locks and the overall contention.

//

// OPTIMISTIC: Locks for reads within the transaction are not acquired on read. Instead, the locks

// are acquired on commit to validate that read/queried data has not changed since the transaction

// started. If a conflict is detected, the transaction will fail. This mode only applies to SERIALIZABLE

// isolation. This mode defers locking until commit, which can reduce contention and improve throughput.

// However, be aware that this increases the risk of transaction aborts if there's significant write

// competition on the same data.

//

// READ_LOCK_MODE_UNSPECIFIED: This is the default if no mode is set. The locking behavior depends on

// the isolation level:

//

// REPEATABLE_READ isolation: Locking semantics default to OPTIMISTIC. However, validation checks at

// commit are only performed for queries using SELECT FOR UPDATE, statements with LOCK_SCANNED_RANGES

// hints, and DML statements. Note: It is an error to explicitly set ReadLockMode when the isolation

// level is REPEATABLE_READ.

//

// For all other isolation levels: If the read lock mode is not set, it defaults to PESSIMISTIC locking.

SetReadLockMode(mode spannerpb.TransactionOptions_ReadWrite_ReadLockMode) error

// TransactionTag returns the transaction tag that will be applied to the next

// read/write transaction on this connection. The transaction tag that is set

// on the connection is cleared when a read/write transaction is started.

TransactionTag() string

// SetTransactionTag sets the transaction tag that should be applied to the

// next read/write transaction on this connection. The tag is cleared when a

// read/write transaction is started.

SetTransactionTag(transactionTag string) error

// MaxCommitDelay returns the max commit delay that will be applied to read/write

// transactions on this connection.

MaxCommitDelay() time.Duration

// SetMaxCommitDelay sets the max commit delay that will be applied to read/write

// transactions on this connection.

SetMaxCommitDelay(delay time.Duration) error

// ExcludeTxnFromChangeStreams returns true if the next transaction should be excluded from change streams with the

// DDL option `allow_txn_exclusion=true`.

ExcludeTxnFromChangeStreams() bool

// SetExcludeTxnFromChangeStreams sets whether the next transaction should be excluded from change streams with the

// DDL option `allow_txn_exclusion=true`.

SetExcludeTxnFromChangeStreams(excludeTxnFromChangeStreams bool) error

// DecodeToNativeArrays indicates whether arrays with a Go native type

// should be decoded to those native types instead of the corresponding

// spanner.NullTypeName (e.g. []bool vs []spanner.NullBool).

// See ExecOptions.DecodeToNativeArrays for more information.

DecodeToNativeArrays() bool

// SetDecodeToNativeArrays sets whether arrays with a Go native type

// should be decoded to those native types instead of the corresponding

// spanner.NullTypeName (e.g. []bool vs []spanner.NullBool).

// See ExecOptions.DecodeToNativeArrays for more information.

SetDecodeToNativeArrays(decodeToNativeArrays bool) error

// Apply writes an array of mutations to the database. This method may only be called while the connection

// is outside a transaction. Use BufferWrite to write mutations in a transaction.

// See also spanner.Client#Apply

Apply(ctx context.Context, ms []*spanner.Mutation, opts ...spanner.ApplyOption) (commitTimestamp time.Time, err error)

// BufferWrite writes an array of mutations to the current transaction. This method may only be called while the

// connection is in a read/write transaction. Use Apply to write mutations outside a transaction.

// See also spanner.ReadWriteTransaction#BufferWrite

BufferWrite(ms []*spanner.Mutation) error

// CommitTimestamp returns the commit timestamp of the last implicit or explicit read/write transaction that

// was executed on the connection, or an error if the connection has not executed a read/write transaction

// that committed successfully. The timestamp is in the local timezone.

CommitTimestamp() (commitTimestamp time.Time, err error)

// CommitResponse returns the commit response of the last implicit or explicit read/write transaction that

// was executed on the connection, or an error if the connection has not executed a read/write transaction

// that committed successfully.

CommitResponse() (commitResponse *spanner.CommitResponse, err error)

// UnderlyingClient returns the underlying Spanner client for the database.

// The client cannot be used to access the current transaction or batch on

// this connection. Executing a transaction or batch using the client that is

// returned, does not affect this connection.

// Note that multiple connections share the same Spanner client. Calling

// this function on different connections to the same database, can

// return the same Spanner client.

UnderlyingClient() (client *spanner.Client, err error)

// DetectStatementType returns the type of SQL statement.

DetectStatementType(query string) parser.StatementType

// resetTransactionForRetry resets the current transaction after it has

// been aborted by Spanner. Calling this function on a transaction that

// has not been aborted is not supported and will cause an error to be

// returned.

resetTransactionForRetry(ctx context.Context, errDuringCommit bool) error

// withTransactionCloseFunc sets the close function that should be registered

// on the next transaction on this connection. This method should only be called

// directly before starting a new transaction.

withTransactionCloseFunc(close func())

// setReadWriteTransactionOptions sets the ReadWriteTransactionOptions that should be

// used for the current read/write transaction. This method should be called right

// after starting a new read/write transaction.

setReadWriteTransactionOptions(options *ReadWriteTransactionOptions)

// setReadOnlyTransactionOptions sets the options that should be used

// for the current read-only transaction. This method should be called

// right after starting a new read-only transaction.

setReadOnlyTransactionOptions(options *ReadOnlyTransactionOptions)

// setBatchReadOnlyTransactionOptions sets the options that should be used

// for the current batch read-only transaction. This method should be called

// right after starting a new batch read-only transaction.

setBatchReadOnlyTransactionOptions(options *BatchReadOnlyTransactionOptions)

}

var _ SpannerConn = &conn{}

type conn struct {

parser *parser.StatementParser

connector *connector

closed bool

client *spanner.Client

adminClient *adminapi.DatabaseAdminClient

connId string

logger *slog.Logger

tx *delegatingTransaction

prevTx *delegatingTransaction

resetForRetry bool

instance string

database string

execSingleQuery func(ctx context.Context, c *spanner.Client, statement spanner.Statement, statementInfo *parser.StatementInfo, bound spanner.TimestampBound, options *ExecOptions) *spanner.RowIterator

execSingleQueryTransactional func(ctx context.Context, c *spanner.Client, statement spanner.Statement, statementInfo *parser.StatementInfo, options *ExecOptions) (rowIterator, *spanner.CommitResponse, error)

execSingleDMLTransactional func(ctx context.Context, c *spanner.Client, statement spanner.Statement, statementInfo *parser.StatementInfo, options *ExecOptions) (*result, *spanner.CommitResponse, error)

execSingleDMLPartitioned func(ctx context.Context, c *spanner.Client, statement spanner.Statement, options *ExecOptions) (int64, error)

// state contains the current ConnectionState for this connection.

state *connectionstate.ConnectionState

// batch is the currently active DDL or DML batch on this connection.

batch *batch

// tempExecOptions can be set by passing it in as an argument to ExecContext or QueryContext

// and are applied only to that statement.

tempExecOptions *ExecOptions

// tempTransactionCloseFunc is set right before a transaction is started, and is set as the

// close function for that transaction.

tempTransactionCloseFunc func()

}

func (c *conn) UnderlyingClient() (*spanner.Client, error) {

return c.client, nil

}

func (c *conn) DetectStatementType(query string) parser.StatementType {

info := c.parser.DetectStatementType(query)

return info.StatementType

}

func (c *conn) CommitTimestamp() (time.Time, error) {

ts := propertyCommitTimestamp.GetValueOrDefault(c.state)

if ts == nil {

return time.Time{}, spanner.ToSpannerError(status.Error(codes.FailedPrecondition, "this connection has not executed a read/write transaction that committed successfully"))

}

return *ts, nil

}

func (c *conn) CommitResponse() (commitResponse *spanner.CommitResponse, err error) {

resp := propertyCommitResponse.GetValueOrDefault(c.state)

if resp == nil {

return nil, spanner.ToSpannerError(status.Error(codes.FailedPrecondition, "this connection has not executed a read/write transaction that committed successfully"))

}

return resp, nil

}

func (c *conn) clearCommitResponse() {

_ = propertyCommitResponse.SetValue(c.state, nil, connectionstate.ContextUser)

_ = propertyCommitTimestamp.SetValue(c.state, nil, connectionstate.ContextUser)

}

func (c *conn) setCommitResponse(commitResponse *spanner.CommitResponse) {

if commitResponse == nil {

c.clearCommitResponse()

return

}

_ = propertyCommitResponse.SetValue(c.state, commitResponse, connectionstate.ContextUser)

_ = propertyCommitTimestamp.SetValue(c.state, &commitResponse.CommitTs, connectionstate.ContextUser)

}

func (c *conn) showConnectionVariable(identifier parser.Identifier) (any, bool, error) {

extension, name, err := toExtensionAndName(identifier)

if err != nil {

return nil, false, err

}

return c.state.GetValue(extension, name)

}

func (c *conn) setConnectionVariable(identifier parser.Identifier, value string, local bool, transaction bool, statementScoped bool) error {

if transaction && !local {

// When transaction == true, then local must also be true.

// We should never hit this condition, as this is an indication of a bug in the driver code.

return status.Errorf(codes.FailedPrecondition, "transaction properties must be set as a local value")

}

if statementScoped && local {

return status.Errorf(codes.FailedPrecondition, "cannot specify both statementScoped and local")

}

extension, name, err := toExtensionAndName(identifier)

if err != nil {

return err

}

if statementScoped {

return c.state.SetStatementScopedValue(extension, name, value)

}

if local {

return c.state.SetLocalValue(extension, name, value, transaction)

}

return c.state.SetValue(extension, name, value, connectionstate.ContextUser)

}

func toExtensionAndName(identifier parser.Identifier) (string, string, error) {

var extension string

var name string

if len(identifier.Parts) == 1 {

extension = ""

name = identifier.Parts[0]

} else if len(identifier.Parts) == 2 {

extension = identifier.Parts[0]

name = identifier.Parts[1]

} else {

return "", "", status.Errorf(codes.InvalidArgument, "invalid variable name: %s", identifier)

}

return extension, name, nil

}

func (c *conn) RetryAbortsInternally() bool {

return propertyRetryAbortsInternally.GetValueOrDefault(c.state)

}

func (c *conn) SetRetryAbortsInternally(retry bool) error {

_, err := c.setRetryAbortsInternally(retry)

return err

}

func (c *conn) setRetryAbortsInternally(retry bool) (driver.Result, error) {

if err := propertyRetryAbortsInternally.SetValue(c.state, retry, connectionstate.ContextUser); err != nil {

return nil, err

}

return driver.ResultNoRows, nil

}

func (c *conn) AutocommitDMLMode() AutocommitDMLMode {

return propertyAutocommitDmlMode.GetValueOrDefault(c.state)

}

func (c *conn) SetAutocommitDMLMode(mode AutocommitDMLMode) error {

if mode == Unspecified {

return spanner.ToSpannerError(status.Error(codes.InvalidArgument, "autocommit dml mode cannot be unspecified"))

}

_, err := c.setAutocommitDMLMode(mode)

return err

}

func (c *conn) setAutocommitDMLMode(mode AutocommitDMLMode) (driver.Result, error) {

if err := propertyAutocommitDmlMode.SetValue(c.state, mode, connectionstate.ContextUser); err != nil {

return nil, err

}

return driver.ResultNoRows, nil

}

func (c *conn) ReadOnlyStaleness() spanner.TimestampBound {

return propertyReadOnlyStaleness.GetValueOrDefault(c.state)

}

func (c *conn) SetReadOnlyStaleness(staleness spanner.TimestampBound) error {

_, err := c.setReadOnlyStaleness(staleness)

return err

}

func (c *conn) setReadOnlyStaleness(staleness spanner.TimestampBound) (driver.Result, error) {

if err := propertyReadOnlyStaleness.SetValue(c.state, staleness, connectionstate.ContextUser); err != nil {

return nil, err

}

return driver.ResultNoRows, nil

}

func (c *conn) readOnlyStalenessPointer() *spanner.TimestampBound {

val := propertyReadOnlyStaleness.GetConnectionPropertyValue(c.state)

if val == nil || !val.HasValue() {

return nil

}

staleness, _ := val.GetValue()

timestampBound := staleness.(spanner.TimestampBound)

return &timestampBound

}

func (c *conn) IsolationLevel() sql.IsolationLevel {

return propertyIsolationLevel.GetValueOrDefault(c.state)

}

func (c *conn) SetIsolationLevel(level sql.IsolationLevel) error {

return propertyIsolationLevel.SetValue(c.state, level, connectionstate.ContextUser)

}

func (c *conn) ReadLockMode() spannerpb.TransactionOptions_ReadWrite_ReadLockMode {

return propertyReadLockMode.GetValueOrDefault(c.state)

}

func (c *conn) SetReadLockMode(mode spannerpb.TransactionOptions_ReadWrite_ReadLockMode) error {

return propertyReadLockMode.SetValue(c.state, mode, connectionstate.ContextUser)

}

func (c *conn) MaxCommitDelay() time.Duration {

return propertyMaxCommitDelay.GetValueOrDefault(c.state)

}

func (c *conn) maxCommitDelayPointer() *time.Duration {

val := propertyMaxCommitDelay.GetConnectionPropertyValue(c.state)

if val == nil || !val.HasValue() {

return nil

}

maxCommitDelay, _ := val.GetValue()

duration := maxCommitDelay.(time.Duration)

return &duration

}

func (c *conn) SetMaxCommitDelay(delay time.Duration) error {

return propertyMaxCommitDelay.SetValue(c.state, delay, connectionstate.ContextUser)

}

func (c *conn) ExcludeTxnFromChangeStreams() bool {

return propertyExcludeTxnFromChangeStreams.GetValueOrDefault(c.state)

}

func (c *conn) SetExcludeTxnFromChangeStreams(excludeTxnFromChangeStreams bool) error {

return propertyExcludeTxnFromChangeStreams.SetValue(c.state, excludeTxnFromChangeStreams, connectionstate.ContextUser)

}

func (c *conn) DecodeToNativeArrays() bool {

return propertyDecodeToNativeArrays.GetValueOrDefault(c.state)

}

func (c *conn) SetDecodeToNativeArrays(decodeToNativeArrays bool) error {

return propertyDecodeToNativeArrays.SetValue(c.state, decodeToNativeArrays, connectionstate.ContextUser)

}

func (c *conn) TransactionTag() string {

return propertyTransactionTag.GetValueOrDefault(c.state)

}

func (c *conn) SetTransactionTag(transactionTag string) error {

return propertyTransactionTag.SetValue(c.state, transactionTag, connectionstate.ContextUser)

}

func (c *conn) StatementTag() string {

return propertyStatementTag.GetValueOrDefault(c.state)

}

func (c *conn) SetStatementTag(statementTag string) error {

return propertyStatementTag.SetValue(c.state, statementTag, connectionstate.ContextUser)

}

func (c *conn) AutoBatchDml() bool {

return propertyAutoBatchDml.GetValueOrDefault(c.state)

}

func (c *conn) SetAutoBatchDml(autoBatch bool) error {

return propertyAutoBatchDml.SetValue(c.state, autoBatch, connectionstate.ContextUser)

}

func (c *conn) AutoBatchDmlUpdateCount() int64 {

return propertyAutoBatchDmlUpdateCount.GetValueOrDefault(c.state)

}

func (c *conn) SetAutoBatchDmlUpdateCount(updateCount int64) error {

return propertyAutoBatchDmlUpdateCount.SetValue(c.state, updateCount, connectionstate.ContextUser)

}

func (c *conn) AutoBatchDmlUpdateCountVerification() bool {

return propertyAutoBatchDmlUpdateCountVerification.GetValueOrDefault(c.state)

}

func (c *conn) SetAutoBatchDmlUpdateCountVerification(verify bool) error {

return propertyAutoBatchDmlUpdateCountVerification.SetValue(c.state, verify, connectionstate.ContextUser)

}

func (c *conn) StartBatchDDL() error {

_, err := c.startBatchDDL()

return err

}

func (c *conn) StartBatchDML() error {

_, err := c.startBatchDML( /* automatic = */ false)

return err

}

func (c *conn) RunBatch(ctx context.Context) error {

_, err := c.runBatch(ctx)

return err

}

func (c *conn) RunDmlBatch(ctx context.Context) (SpannerResult, error) {

res, err := c.runBatch(ctx)

if err != nil {

if spannerRes, ok := res.(SpannerResult); ok {

return spannerRes, err

}

return nil, err

}

spannerRes, ok := res.(SpannerResult)

if !ok {

return nil, spanner.ToSpannerError(status.Errorf(codes.FailedPrecondition, "not a DML batch"))

}

return spannerRes, nil

}

func (c *conn) AbortBatch() error {

_, err := c.abortBatch()

return err

}

func (c *conn) InDDLBatch() bool {

return c.batch != nil && c.batch.tp == parser.BatchTypeDdl

}

func (c *conn) InDMLBatch() bool {

return (c.batch != nil && c.batch.tp == parser.BatchTypeDml) || (c.inTransaction() && c.tx.IsInBatch())

}

func (c *conn) GetBatchedStatements() []spanner.Statement {

if c.batch == nil || c.batch.statements == nil {

return []spanner.Statement{}

}

return slices.Clone(c.batch.statements)

}

func (c *conn) inBatch() bool {

return c.InDDLBatch() || c.InDMLBatch()

}

func (c *conn) startBatchDDL() (driver.Result, error) {

if c.batch != nil {

return nil, spanner.ToSpannerError(status.Errorf(codes.FailedPrecondition, "This connection already has an active batch."))

}

if c.inTransaction() {

return nil, spanner.ToSpannerError(status.Errorf(codes.FailedPrecondition, "This connection has an active transaction. DDL batches in transactions are not supported."))

}

c.logger.Debug("started ddl batch")

c.batch = &batch{tp: parser.BatchTypeDdl}

return driver.ResultNoRows, nil

}

func (c *conn) startBatchDML(automatic bool) (driver.Result, error) {

execOptions, err := c.options( /* reset = */ true)

if err != nil {

return nil, err

}

if c.inTransaction() {

return c.tx.StartBatchDML(execOptions.QueryOptions, automatic)

}

if c.batch != nil {

return nil, spanner.ToSpannerError(status.Errorf(codes.FailedPrecondition, "This connection already has an active batch."))

}

c.logger.Debug("starting dml batch outside transaction")

c.batch = &batch{tp: parser.BatchTypeDml, options: execOptions}

return driver.ResultNoRows, nil

}

func (c *conn) runBatch(ctx context.Context) (driver.Result, error) {

if c.inTransaction() {

return c.tx.RunBatch(ctx)

}

if c.batch == nil {

return nil, spanner.ToSpannerError(status.Errorf(codes.FailedPrecondition, "This connection does not have an active batch"))

}

switch c.batch.tp {

case parser.BatchTypeDdl:

return c.runDDLBatch(ctx)

case parser.BatchTypeDml:

return c.runDMLBatch(ctx)

default:

return nil, spanner.ToSpannerError(status.Errorf(codes.InvalidArgument, "Unknown batch type: %d", c.batch.tp))

}

}

func (c *conn) runDDLBatch(ctx context.Context) (driver.Result, error) {

statements := c.batch.statements

c.batch = nil

return c.execDDL(ctx, statements...)

}

func (c *conn) runDMLBatch(ctx context.Context) (SpannerResult, error) {

if c.inTransaction() {

return c.tx.RunDmlBatch(ctx)

}

statements := c.batch.statements

options := c.batch.options

options.QueryOptions.LastStatement = true

c.batch = nil

return c.execBatchDML(ctx, statements, options)

}

func (c *conn) abortBatch() (driver.Result, error) {

if c.inTransaction() {

return c.tx.AbortBatch()

}

c.batch = nil

return driver.ResultNoRows, nil

}

func (c *conn) execDDL(ctx context.Context, statements ...spanner.Statement) (driver.Result, error) {

if c.batch != nil && c.batch.tp == parser.BatchTypeDml {

return nil, spanner.ToSpannerError(status.Error(codes.FailedPrecondition, "This connection has an active DML batch"))

}

if c.batch != nil && c.batch.tp == parser.BatchTypeDdl {

c.batch.statements = append(c.batch.statements, statements...)

return driver.ResultNoRows, nil

}

if len(statements) > 0 {

ddlStatements := make([]string, len(statements))

for i, s := range statements {

ddlStatements[i] = s.SQL

}

if c.parser.IsCreateDatabaseStatement(ddlStatements[0]) {

op, err := c.adminClient.CreateDatabase(ctx, &adminpb.CreateDatabaseRequest{

CreateStatement: ddlStatements[0],

DatabaseDialect: c.parser.Dialect,

Parent: c.instance,

ExtraStatements: ddlStatements[1:],

})

if err != nil {

return nil, err

}

if _, err := op.Wait(ctx); err != nil {

return nil, err

}

return driver.ResultNoRows, nil

} else if c.parser.IsDropDatabaseStatement(ddlStatements[0]) {

if len(ddlStatements) > 1 {

return nil, spanner.ToSpannerError(status.Error(codes.InvalidArgument, "DROP DATABASE cannot be used in a batch with other statements"))

}

stmt := &parser.ParsedDropDatabaseStatement{}

if err := stmt.Parse(c.parser, ddlStatements[0]); err != nil {

return nil, err

}

return (&executableDropDatabaseStatement{stmt}).execContext(ctx, c, nil)

}

op, err := c.adminClient.UpdateDatabaseDdl(ctx, &adminpb.UpdateDatabaseDdlRequest{

Database: c.database,

Statements: ddlStatements,

})

if err != nil {

return nil, err

}

if err := op.Wait(ctx); err != nil {

if len(statements) > 1 {

be := &BatchError{

Err: err,

BatchUpdateCounts: []int64{},

}

metadata, err := op.Metadata()

if err != nil {

c.logger.WarnContext(ctx, fmt.Sprintf("Error getting metadata for UpdateDatabaseDdl: %v", err))

} else if metadata != nil {

for _, ts := range metadata.CommitTimestamps {

if ts != nil {

be.BatchUpdateCounts = append(be.BatchUpdateCounts, int64(-1))

} else {

break

}

}

}

return nil, be

}

return nil, err

}

}

return driver.ResultNoRows, nil

}

func (c *conn) execBatchDML(ctx context.Context, statements []spanner.Statement, options *ExecOptions) (SpannerResult, error) {

if len(statements) == 0 {

return &result{}, nil

}

var affected []int64

var err error

if c.inTransaction() && c.tx.contextTransaction != nil {

tx, ok := c.tx.contextTransaction.(*readWriteTransaction)

if !ok {

return nil, status.Errorf(codes.FailedPrecondition, "connection is in a transaction that is not a read/write transaction")

}

affected, err = tx.rwTx.BatchUpdateWithOptions(ctx, statements, options.QueryOptions)

} else {

_, err = c.client.ReadWriteTransactionWithOptions(ctx, func(ctx context.Context, transaction *spanner.ReadWriteTransaction) error {

affected, err = transaction.BatchUpdateWithOptions(ctx, statements, options.QueryOptions)

return err

}, options.TransactionOptions)

}

res := &result{rowsAffected: sum(affected), batchUpdateCounts: affected}

ba := toBatchError(res, err)

return res, ba

}

func sum(affected []int64) int64 {

sum := int64(0)

for _, c := range affected {

sum += c

}

return sum

}

// WriteMutations is not part of the public API of the database/sql driver.

// It is exported for internal reasons, and may receive breaking changes without prior notice.

//

// WriteMutations writes mutations using this connection. The mutations are either buffered in the current transaction,

// or written directly to Spanner using a new read/write transaction if the connection does not have a transaction.

//

// The function returns an error if the connection currently has a read-only transaction.

//

// The returned CommitResponse is nil if the connection currently has a transaction, as the mutations will only be

// applied to Spanner when the transaction commits.

func (c *conn) WriteMutations(ctx context.Context, ms []*spanner.Mutation) (*spanner.CommitResponse, error) {

if c.inTransaction() {

return nil, c.BufferWrite(ms)

}

ts, err := c.Apply(ctx, ms)

if err != nil {

return nil, err

}

return &spanner.CommitResponse{CommitTs: ts}, nil

}

func (c *conn) Apply(ctx context.Context, ms []*spanner.Mutation, opts ...spanner.ApplyOption) (commitTimestamp time.Time, err error) {

if c.inTransaction() {

return time.Time{}, spanner.ToSpannerError(

status.Error(

codes.FailedPrecondition,

"Apply may not be called while the connection is in a transaction. Use BufferWrite to write mutations in a transaction."))

}

return c.client.Apply(ctx, ms, opts...)

}

func (c *conn) BufferWrite(ms []*spanner.Mutation) error {

if !c.inTransaction() {

return spanner.ToSpannerError(

status.Error(

codes.FailedPrecondition,

"BufferWrite may not be called while the connection is not in a transaction. Use Apply to write mutations outside a transaction."))

}

return c.tx.BufferWrite(ms)

}

// Ping implements the driver.Pinger interface.

// returns ErrBadConn if the connection is no longer valid.

func (c *conn) Ping(ctx context.Context) error {

if c.closed {

return driver.ErrBadConn

}

rows, err := c.QueryContext(ctx, "SELECT 1", []driver.NamedValue{})

if err != nil {

return driver.ErrBadConn

}

defer func() { _ = rows.Close() }()

values := make([]driver.Value, 1)

if err := rows.Next(values); err != nil {

return driver.ErrBadConn

}

if values[0] != int64(1) {

return driver.ErrBadConn

}

return nil

}

// ResetSession implements the driver.SessionResetter interface.

// returns ErrBadConn if the connection is no longer valid.

func (c *conn) ResetSession(_ context.Context) error {

if c.closed {

return driver.ErrBadConn

}

if c.inTransaction() {

if err := c.tx.Rollback(); err != nil {

return driver.ErrBadConn

}

}

c.batch = nil

_ = c.state.Reset(connectionstate.ContextUser)

c.tempExecOptions = nil

return nil

}

// IsValid implements the driver.Validator interface.

func (c *conn) IsValid() bool {

// Mark the connection as invalid if it is still in a read/write transaction.

// This ensures that the database/sql package discards the connection and releases

// any locks this connection might hold.

return !c.closed && !c.inReadWriteTransaction()

}

func (c *conn) CheckNamedValue(value *driver.NamedValue) error {

if value == nil {

return nil

}

if execOptions, ok := value.Value.(ExecOptions); ok {

c.tempExecOptions = &execOptions

return driver.ErrRemoveArgument

}

if execOptions, ok := value.Value.(*ExecOptions); ok {

c.tempExecOptions = execOptions

return driver.ErrRemoveArgument

}

if checkIsValidType(value.Value) {

return nil

}

// Convert directly if the value implements driver.Valuer. Although this

// is also done by the default converter in the sql driver, that conversion

// requires the value that is returned by Value() to be one of the base

// types that is supported by database/sql. By doing this check here first,

// we also support driver.Valuer types that return a Spanner-supported type,

// such as for example []string.

if valuer, ok := value.Value.(driver.Valuer); ok {

if v, err := callValuerValue(valuer); err == nil {

if checkIsValidType(v) {

value.Value = v

return nil

}

}

}

// Convert the value using the default sql driver. This uses driver.Valuer,

// if implemented, and falls back to reflection. If the converted value is

// a supported spanner type, use it. Otherwise, ignore any errors and

// continue checking other supported spanner specific types.

if v, err := driver.DefaultParameterConverter.ConvertValue(value.Value); err == nil {

if checkIsValidType(v) {

value.Value = v

return nil

}

}

// google-cloud-go/spanner knows how to deal with these

if isStructOrArrayOfStructValue(value.Value) || isAnArrayOfProtoColumn(value.Value) {

return nil

}

return spanner.ToSpannerError(status.Errorf(codes.InvalidArgument, "unsupported value type: %T", value.Value))

}

func (c *conn) Prepare(query string) (driver.Stmt, error) {

return c.PrepareContext(context.Background(), query)

}

func (c *conn) PrepareContext(_ context.Context, query string) (driver.Stmt, error) {

execOptions, err := c.options( /* reset = */ true)

if err != nil {

return nil, err

}

parsedSQL, args, err := c.parser.ParseParameters(query)

if err != nil {

return nil, err

}

return &stmt{conn: c, query: parsedSQL, numArgs: len(args), execOptions: execOptions}, nil

}

// Adds any statement or transaction timeout to the given context. The deadline of the returned

// context will be the earliest of:

// 1. Any existing deadline on the input context.

// 2. Any existing transaction deadline.

// 3. A deadline calculated from the current time + the value of statement_timeout.

func (c *conn) addStatementAndTransactionTimeout(ctx context.Context) (context.Context, context.CancelFunc, error) {

var statementDeadline time.Time

var transactionDeadline time.Time

var deadline time.Time

var hasStatementDeadline bool

var hasTransactionDeadline bool

// Check if the connection has a value for statement_timeout.

statementTimeout := propertyStatementTimeout.GetValueOrDefault(c.state)

if statementTimeout != time.Duration(0) {

hasStatementDeadline = true

statementDeadline = time.Now().Add(statementTimeout)

}

// Check if the current transaction has a deadline.

transactionDeadline, hasTransactionDeadline, err := c.transactionDeadline()

if err != nil {

return nil, nil, err

}

// If there is no statement_timeout and no current transaction deadline,

// then can just use the input context as-is.

if !hasStatementDeadline && !hasTransactionDeadline {

return ctx, func() {}, nil

}

// If there is both a transaction and a statement deadline, then we use the earliest

// of those two.

if hasTransactionDeadline && hasStatementDeadline {

if statementDeadline.Before(transactionDeadline) {

deadline = statementDeadline

} else {

deadline = transactionDeadline

}

} else if hasStatementDeadline {

deadline = statementDeadline

} else {

deadline = transactionDeadline

}

// context.WithDeadline automatically selects the earliest deadline of

// the existing deadline on the context and the given deadline.

newCtx, cancel := context.WithDeadline(ctx, deadline)

return newCtx, cancel, nil

}

// transactionDeadline returns the deadline of the current transaction

// on the connection. This also activates the transaction if it is not

// yet activated.

func (c *conn) transactionDeadline() (time.Time, bool, error) {

if c.tx == nil {

return time.Time{}, false, nil

}

if err := c.tx.ensureActivated(); err != nil {

return time.Time{}, false, err

}

deadline, hasDeadline := c.tx.deadline()

return deadline, hasDeadline, nil

}

func (c *conn) applyStatementScopedValues(execOptions *ExecOptions) (cleanup func(), returnedErr error) {

if execOptions == nil {

return func() {}, nil

}

defer func() {

if returnedErr != nil {

// Clear any statement values that might have been set if we return an error, as that also means that we

// are not returning a cleanup function.

c.state.ClearStatementScopedValues()

}

}()

values := execOptions.PropertyValues

for _, value := range values {

if err := c.setConnectionVariable(value.Identifier, value.Value /*local=*/, false /*transaction=*/, false /*statementScoped=*/, true); err != nil {

return func() {}, err

}

}

return c.state.ClearStatementScopedValues, nil

}

func (c *conn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {

// Remove the ExecOptions once all statements in the SQL string have been executed.

defer func() { c.tempExecOptions = nil }()