chore: add CreatePool and CreateConnection functions for SpannerLib

Creates a spannerlib module and adds functions for CreatePool and CreateConnection.

Author: olavloite

.github/workflows/spanner-lib-tests.yml

@@ -0,0 +1,25 @@
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+permissions:
+  contents: read
+  pull-requests: write
+name: Spanner Lib Tests
+jobs:
+  test:
+    strategy:
+      matrix:
+        go-version: [1.25.x]
+        os: [ubuntu-latest, macos-latest, windows-latest]
+    runs-on: ${{ matrix.os }}
+    steps:
+      - name: Install Go
+        uses: actions/setup-go@v5
+        with:
+          go-version: ${{ matrix.go-version }}
+      - name: Checkout code
+        uses: actions/checkout@v4
+      - name: Run unit tests
+        working-directory: spannerlib
+        run: go test ./... -race -short

spannerlib/.gitignore

@@ -0,0 +1,3 @@
+spannerlib.h
+spannerlib.so
+grpc_server

spannerlib/README.md

@@ -0,0 +1,6 @@
+# Shared Library (Internal)
+
+__This module can receive breaking changes without prior notice.__
+
+This is an internal module that is used to expose the features in the database/sql driver
+to drivers in other programming languages.

spannerlib/api/README.md

@@ -0,0 +1,6 @@
+# Shared Library (Internal API)
+
+__This package can receive breaking changes without prior notice.__
+
+This is the common internal API for the various external APIs for exposing the
+features in the database/sql driver to drivers in other programming languages.

spannerlib/api/connection.go

@@ -0,0 +1,62 @@
+// 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 api
+
+import (
+	"context"
+	"database/sql"
+	"sync"
+	"sync/atomic"
+)
+
+// CloseConnection looks up the connection with the given poolId and connId and closes it.
+func CloseConnection(ctx context.Context, poolId, connId int64) error {
+	pool, err := findPool(poolId)
+	if err != nil {
+		return err
+	}
+	c, ok := pool.connections.LoadAndDelete(connId)
+	if !ok {
+		// Closing an unknown connection or a connection that has previously been closed is a no-op.
+		return nil
+	}
+	conn := c.(*Connection)
+	return conn.close(ctx)
+}
+
+type Connection struct {
+	// results contains the open query results for this connection.
+	results    *sync.Map
+	resultsIdx atomic.Int64
+
+	// backend is the database/sql connection of this connection.
+	backend *sql.Conn
+}
+
+func (conn *Connection) close(ctx context.Context) error {
+	conn.closeResults(ctx)
+	err := conn.backend.Close()
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func (conn *Connection) closeResults(ctx context.Context) {
+	conn.results.Range(func(key, value interface{}) bool {
+		// TODO: Implement
+		return true
+	})
+}

spannerlib/api/connection_test.go

@@ -0,0 +1,145 @@
+// 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 api
+
+import (
+	"context"
+	"fmt"
+	"reflect"
+	"testing"
+
+	"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 TestCreateAndCloseConnection(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	server, teardown := setupMockServer(t)
+	defer teardown()
+	dsn := fmt.Sprintf("%s/projects/p/instances/i/databases/d?useplaintext=true", server.Address)
+
+	poolId, err := CreatePool(ctx, dsn)
+	if err != nil {
+		t.Fatalf("CreatePool returned unexpected error: %v", err)
+	}
+	connId, err := CreateConnection(ctx, poolId)
+	if err != nil {
+		t.Fatalf("CreateConnection returned unexpected error: %v", err)
+	}
+	if connId == 0 {
+		t.Fatal("CreateConnection returned unexpected zero id")
+	}
+	p, ok := pools.Load(poolId)
+	if !ok {
+		t.Fatal("pool not found in map")
+	}
+	pool := p.(*Pool)
+	if _, ok := pool.connections.Load(connId); !ok {
+		t.Fatal("connection not in map")
+	}
+
+	requests := server.TestSpanner.DrainRequestsFromServer()
+	createSessionRequests := testutil.RequestsOfType(requests, reflect.TypeOf(&spannerpb.CreateSessionRequest{}))
+	if g, w := len(createSessionRequests), 1; g != w {
+		t.Fatalf("num CreateSession requests mismatch\n Got: %d\nWant: %d", g, w)
+	}
+
+	if err := CloseConnection(ctx, poolId, connId); err != nil {
+		t.Fatalf("CloseConnection returned unexpected error: %v", err)
+	}
+	if _, ok := pool.connections.Load(connId); ok {
+		t.Fatal("connection still in map")
+	}
+
+	if err := ClosePool(ctx, poolId); err != nil {
+		t.Fatalf("ClosePool returned unexpected error: %v", err)
+	}
+}
+
+func TestCreateConnectionWithUnknownPool(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	// Try to create a connection for an unknown pool.
+	_, err := CreateConnection(ctx, -1)
+	if g, w := spanner.ErrCode(err), codes.NotFound; g != w {
+		t.Fatalf("error code mismatch\n Got: %d\nWant: %d", g, w)
+	}
+}
+
+func TestCreateTwoConnections(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	server, teardown := setupMockServer(t)
+	defer teardown()
+	dsn := fmt.Sprintf("%s/projects/p/instances/i/databases/d?useplaintext=true", server.Address)
+
+	poolId, err := CreatePool(ctx, dsn)
+	if err != nil {
+		t.Fatalf("CreatePool returned unexpected error: %v", err)
+	}
+
+	for range 2 {
+		connId, err := CreateConnection(ctx, poolId)
+		if err != nil {
+			t.Fatalf("CreateConnection returned unexpected error: %v", err)
+		}
+		//goland:noinspection GoDeferInLoop
+		defer func() { _ = CloseConnection(ctx, poolId, connId) }()
+	}
+
+	// Two connections in one pool should only create one multiplexed session.
+	requests := server.TestSpanner.DrainRequestsFromServer()
+	createSessionRequests := testutil.RequestsOfType(requests, reflect.TypeOf(&spannerpb.CreateSessionRequest{}))
+	if g, w := len(createSessionRequests), 1; g != w {
+		t.Fatalf("num CreateSession requests mismatch\n Got: %d\nWant: %d", g, w)
+	}
+
+	if err := ClosePool(ctx, poolId); err != nil {
+		t.Fatalf("ClosePool returned unexpected error: %v", err)
+	}
+}
+
+func TestCloseConnectionTwice(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	server, teardown := setupMockServer(t)
+	defer teardown()
+	dsn := fmt.Sprintf("%s/projects/p/instances/i/databases/d?useplaintext=true", server.Address)
+
+	poolId, err := CreatePool(ctx, dsn)
+	if err != nil {
+		t.Fatalf("CreatePool returned unexpected error: %v", err)
+	}
+	connId, err := CreateConnection(ctx, poolId)
+	if err != nil {
+		t.Fatalf("CreateConnection returned unexpected error: %v", err)
+	}
+
+	for range 2 {
+		if err := CloseConnection(ctx, poolId, connId); err != nil {
+			t.Fatalf("CloseConnection returned unexpected error: %v", err)
+		}
+	}
+	if err := ClosePool(ctx, poolId); err != nil {
+		t.Fatalf("ClosePool returned unexpected error: %v", err)
+	}
+}

spannerlib/api/pool.go

@@ -0,0 +1,133 @@
+// 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 api
+
+import (
+	"context"
+	"database/sql"
+	"sync"
+	"sync/atomic"
+
+	spannerdriver "github.com/googleapis/go-sql-spanner"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+// pools contains all open pools in this runtime.
+var pools = sync.Map{}
+var poolsIdx = atomic.Int64{}
+
+// Pool is the equivalent of a sql.DB. It contains a pool of connections to the same database.
+// All connections in a pool share the same Spanner client.
+type Pool struct {
+	db             *sql.DB
+	connections    *sync.Map
+	connectionsIdx atomic.Int64
+}
+
+// CreatePool creates a new Pool and stores it in the global map of pool.
+// The connectionString must be in the form
+// [host:port]/project/<project>/instances/<instance>/databases/<database>[?option1=value1[;option2=value2...]]
+//
+// Creating a pool is a relatively expensive operation, as each pool has its own Spanner client.
+func CreatePool(ctx context.Context, connectionString string) (int64, error) {
+	config, err := spannerdriver.ExtractConnectorConfig(connectionString)
+	if err != nil {
+		return 0, err
+	}
+	connector, err := spannerdriver.CreateConnector(config)
+	if err != nil {
+		return 0, err
+	}
+	db := sql.OpenDB(connector)
+	// Create a connection to force an error if the connection string is invalid.
+	conn, err := db.Conn(ctx)
+	if err != nil {
+		return 0, err
+	}
+	_ = conn.Close()
+
+	id := poolsIdx.Add(1)
+	pool := &Pool{
+		db:          db,
+		connections: &sync.Map{},
+	}
+	pools.Store(id, pool)
+	return id, nil
+}
+
+// ClosePool closes the pool with the given id. All open connections in the pool are also closed, and the pool cannot
+// be used to create any new connections.
+func ClosePool(ctx context.Context, id int64) error {
+	p, ok := pools.LoadAndDelete(id)
+	if !ok {
+		// Closing an unknown pool or a pool that has previously been closed is a no-op.
+		return nil
+	}
+	pool := p.(*Pool)
+	pool.connections.Range(func(key, value interface{}) bool {
+		conn := value.(*Connection)
+		_ = conn.close(ctx)
+		return true
+	})
+	if err := pool.db.Close(); err != nil {
+		return err
+	}
+	return nil
+}
+
+// CreateConnection creates a new connection in the given pool. This is a relatively cheap operation, as a connection
+// does not have its own physical connection to Spanner. Instead, all connections in the same pool share the same
+// underlying Spanner client, which again contains a gRPC channel pool.
+func CreateConnection(ctx context.Context, poolId int64) (int64, error) {
+	pool, err := findPool(poolId)
+	if err != nil {
+		return 0, err
+	}
+	sqlConn, err := pool.db.Conn(ctx)
+	if err != nil {
+		return 0, err
+	}
+	id := poolsIdx.Add(1)
+	conn := &Connection{
+		backend: sqlConn,
+		results: &sync.Map{},
+	}
+	pool.connections.Store(id, conn)
+
+	return id, nil
+}
+
+func findPool(poolId int64) (*Pool, error) {
+	p, ok := pools.Load(poolId)
+	if !ok {
+		return nil, status.Errorf(codes.NotFound, "pool %v not found", poolId)
+	}
+	pool := p.(*Pool)
+	return pool, nil
+}
+
+func findConnection(poolId, connId int64) (*Connection, error) {
+	pool, err := findPool(poolId)
+	if err != nil {
+		return nil, err
+	}
+	c, ok := pool.connections.Load(connId)
+	if !ok {
+		return nil, status.Errorf(codes.NotFound, "connection %v not found", connId)
+	}
+	conn := c.(*Connection)
+	return conn, nil
+}