feat: add transactions

Author: olavloite

conn.go

@@ -59,6 +59,9 @@ type SpannerConn interface {
 	// 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
@@ -202,6 +205,11 @@ type SpannerConn interface {
 	withTempBatchReadOnlyTransactionOptions(options *BatchReadOnlyTransactionOptions)
 }
 
+type SpannerResult interface {
+	driver.Result
+	BatchRowsAffected() ([]int64, error)
+}
+
 var _ SpannerConn = &conn{}
 
 type conn struct {
@@ -261,6 +269,7 @@ type conn struct {
 	// tempBatchReadOnlyTransactionOptions are temporarily set right before a
 	// batch read-only transaction is started on a Spanner connection.
 	tempBatchReadOnlyTransactionOptions *BatchReadOnlyTransactionOptions
+	tempProtoTransactionOptions         *spannerpb.TransactionOptions
 }
 
 func (c *conn) UnderlyingClient() (*spanner.Client, error) {
@@ -444,6 +453,18 @@ func (c *conn) RunBatch(ctx context.Context) error {
 	return err
 }
 
+func (c *conn) RunDmlBatch(ctx context.Context) (SpannerResult, error) {
+	res, err := c.runBatch(ctx)
+	if err != nil {
+		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
@@ -522,7 +543,7 @@ func (c *conn) runDDLBatch(ctx context.Context) (driver.Result, error) {
 	return c.execDDL(ctx, statements...)
 }
 
-func (c *conn) runDMLBatch(ctx context.Context) (driver.Result, error) {
+func (c *conn) runDMLBatch(ctx context.Context) (SpannerResult, error) {
 	statements := c.batch.statements
 	options := c.batch.options
 	options.QueryOptions.LastStatement = true
@@ -541,7 +562,7 @@ func (c *conn) abortBatch() (driver.Result, error) {
 
 func (c *conn) execDDL(ctx context.Context, statements ...spanner.Statement) (driver.Result, error) {
 	if c.batch != nil && c.batch.tp == dml {
-		return nil, spanner.ToSpannerError(status.Error(codes.FailedPrecondition, "This connection has an active DML batch"))
+		return nil, spanner.ToSpannerError(status.Error(codes.FailedPrecondition, "This connection has an active DDL batch"))
 	}
 	if c.batch != nil && c.batch.tp == ddl {
 		c.batch.statements = append(c.batch.statements, statements...)
@@ -567,7 +588,7 @@ func (c *conn) execDDL(ctx context.Context, statements ...spanner.Statement) (dr
 	return driver.ResultNoRows, nil
 }
 
-func (c *conn) execBatchDML(ctx context.Context, statements []spanner.Statement, options ExecOptions) (driver.Result, error) {
+func (c *conn) execBatchDML(ctx context.Context, statements []spanner.Statement, options ExecOptions) (SpannerResult, error) {
 	if len(statements) == 0 {
 		return &result{}, nil
 	}
@@ -586,7 +607,7 @@ func (c *conn) execBatchDML(ctx context.Context, statements []spanner.Statement,
 			return err
 		}, options.TransactionOptions)
 	}
-	return &result{rowsAffected: sum(affected)}, err
+	return &result{rowsAffected: sum(affected), batchUpdateCounts: affected}, err
 }
 
 func sum(affected []int64) int64 {

conn_with_mockserver_test.go

@@ -24,6 +24,7 @@ import (
 	"cloud.google.com/go/spanner"
 	"cloud.google.com/go/spanner/apiv1/spannerpb"
 	"github.com/googleapis/go-sql-spanner/testutil"
+	"google.golang.org/grpc/codes"
 )
 
 func TestBeginTx(t *testing.T) {
@@ -48,6 +49,25 @@ func TestBeginTx(t *testing.T) {
 	}
 }
 
+func TestTwoTransactionsOnOneConn(t *testing.T) {
+	t.Parallel()
+
+	db, _, teardown := setupTestDBConnection(t)
+	defer teardown()
+	ctx := context.Background()
+
+	c, _ := db.Conn(ctx)
+	tx1, err := c.BeginTx(ctx, &sql.TxOptions{})
+	defer tx1.Rollback()
+	if err != nil {
+		t.Fatal(err)
+	}
+	_, err = c.BeginTx(ctx, &sql.TxOptions{})
+	if g, w := spanner.ErrCode(err), codes.FailedPrecondition; g != w {
+		t.Fatalf("BeginTx error code mismatch\n Got: %v\nWant: %v", g, w)
+	}
+}
+
 func TestBeginTxWithIsolationLevel(t *testing.T) {
 	t.Parallel()
 

driver.go

@@ -954,6 +954,10 @@ func BeginReadWriteTransaction(ctx context.Context, db *sql.DB, options ReadWrit
 		// be active when we hit this point.
 		go conn.Close()
 	}
+	return BeginReadWriteTransactionOnConn(ctx, conn, options)
+}
+
+func BeginReadWriteTransactionOnConn(ctx context.Context, conn *sql.Conn, options ReadWriteTransactionOptions) (*sql.Tx, error) {
 	if err := withTempReadWriteTransactionOptions(conn, &options); err != nil {
 		return nil, err
 	}
@@ -1005,6 +1009,10 @@ func BeginReadOnlyTransaction(ctx context.Context, db *sql.DB, options ReadOnlyT
 		// be active when we hit this point.
 		go conn.Close()
 	}
+	return BeginReadOnlyTransactionOnConn(ctx, conn, options)
+}
+
+func BeginReadOnlyTransactionOnConn(ctx context.Context, conn *sql.Conn, options ReadOnlyTransactionOptions) (*sql.Tx, error) {
 	if err := withTempReadOnlyTransactionOptions(conn, &options); err != nil {
 		return nil, err
 	}

driver_with_mockserver_test.go

@@ -4801,9 +4801,8 @@ func setupTestDBConnectionWithParams(t *testing.T, params string) (db *sql.DB, s
 
 func setupTestDBConnectionWithParamsAndDialect(t *testing.T, params string, dialect databasepb.DatabaseDialect) (db *sql.DB, server *testutil.MockedSpannerInMemTestServer, teardown func()) {
 	server, _, serverTeardown := setupMockedTestServerWithDialect(t, dialect)
-	db, err := sql.Open(
-		"spanner",
-		fmt.Sprintf("%s/projects/p/instances/i/databases/d?useplaintext=true;%s", server.Address, params))
+	dsn := fmt.Sprintf("%s/projects/p/instances/i/databases/d?useplaintext=true;%s", server.Address, params)
+	db, err := sql.Open("spanner", dsn)
 	if err != nil {
 		serverTeardown()
 		t.Fatal(err)

spannerlib/backend/db_pool.go

@@ -3,60 +3,35 @@ package backend
 import (
 	"context"
 	"database/sql"
-	"errors"
-	"fmt"
-	"sync"
 
 	spannerdriver "github.com/googleapis/go-sql-spanner"
 )
 
+// Pool is a simple wrapper around sql.DB and contains a pool of connections.
 type Pool struct {
-	Project  string
-	Instance string
-	Database string
-
-	mu      sync.Mutex
-	entries map[string]*sql.DB
-}
-
-func (pool *Pool) Close() (err error) {
-	pool.mu.Lock()
-	defer pool.mu.Unlock()
-	for _, db := range pool.entries {
-		err = errors.Join(err, db.Close())
-	}
-	return err
+	db *sql.DB
 }
 
-func (pool *Pool) Conn(ctx context.Context, project, instance, database string) (*sql.Conn, error) {
-	if project == "" {
-		project = pool.Project
-	}
-	if instance == "" {
-		instance = pool.Instance
-	}
-	if database == "" {
-		database = pool.Database
-	}
-	key := fmt.Sprintf("projects/%s/instances/%s/databases/%s", project, instance, database)
-	pool.mu.Lock()
-	defer pool.mu.Unlock()
-	if db, ok := pool.entries[key]; ok {
-		return db.Conn(ctx)
-	}
-	config := spannerdriver.ConnectorConfig{
-		Project:  project,
-		Instance: instance,
-		Database: database,
+func CreatePool(dsn string) (*Pool, error) {
+	config, err := spannerdriver.ExtractConnectorConfig(dsn)
+	if err != nil {
+		return nil, err
 	}
 	connector, err := spannerdriver.CreateConnector(config)
 	if err != nil {
 		return nil, err
 	}
 	db := sql.OpenDB(connector)
-	if pool.entries == nil {
-		pool.entries = make(map[string]*sql.DB)
+	pool := &Pool{
+		db: db,
 	}
-	pool.entries[key] = db
-	return db.Conn(ctx)
+	return pool, nil
+}
+
+func (pool *Pool) Close() (err error) {
+	return pool.db.Close()
+}
+
+func (pool *Pool) Conn(ctx context.Context) (*sql.Conn, error) {
+	return pool.db.Conn(ctx)
 }