chore: create shared lib for connections

Author: olavloite

spannerlib/backend/connection.go

@@ -0,0 +1,13 @@
+package backend
+
+import (
+	"database/sql"
+)
+
+type SpannerConnection struct {
+	Conn *sql.Conn
+}
+
+func (conn *SpannerConnection) Close() error {
+	return conn.Conn.Close()
+}

spannerlib/backend/db_pool.go

@@ -0,0 +1,62 @@
+package backend
+
+import (
+	"context"
+	"database/sql"
+	"errors"
+	"fmt"
+	"sync"
+
+	spannerdriver "github.com/googleapis/go-sql-spanner"
+)
+
+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
+}
+
+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,
+	}
+	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.entries[key] = db
+	return db.Conn(ctx)
+}

spannerlib/exported/connection.go

@@ -0,0 +1,89 @@
+package exported
+
+import "C"
+import (
+	"context"
+	"database/sql"
+	"fmt"
+	"sync"
+	"sync/atomic"
+
+	"cloud.google.com/go/spanner"
+	"cloud.google.com/go/spanner/apiv1/spannerpb"
+	spannerdriver "github.com/googleapis/go-sql-spanner"
+	"google.golang.org/protobuf/proto"
+	"spannerlib/backend"
+)
+
+func CloseConnection(poolId, connId int64) *Message {
+	conn, err := findConnection(poolId, connId)
+	if err != nil {
+		return errMessage(err)
+	}
+	return conn.close()
+}
+
+func Execute(poolId, connId int64, statementBytes []byte) *Message {
+	statement := spannerpb.ExecuteBatchDmlRequest_Statement{}
+	if err := proto.Unmarshal(statementBytes, &statement); err != nil {
+		return errMessage(err)
+	}
+	fmt.Printf("Statement: %v\n", statement.Sql)
+	fmt.Printf("Params: %v\n", statement.Params)
+	conn, err := findConnection(poolId, connId)
+	if err != nil {
+		return errMessage(err)
+	}
+	return conn.Execute(&statement)
+}
+
+type Connection struct {
+	results    *sync.Map
+	resultsIdx atomic.Int64
+
+	backend *backend.SpannerConnection
+}
+
+func (conn *Connection) close() *Message {
+	conn.results.Range(func(key, value interface{}) bool {
+		res := value.(*rows)
+		res.Close()
+		return true
+	})
+	err := conn.backend.Close()
+	if err != nil {
+		return errMessage(err)
+	}
+	return &Message{}
+}
+
+func (conn *Connection) Execute(statement *spannerpb.ExecuteBatchDmlRequest_Statement) *Message {
+	paramsLen := 1
+	if statement.Params != nil {
+		paramsLen = 1 + len(statement.Params.Fields)
+	}
+	params := make([]any, paramsLen)
+	params = append(params, spannerdriver.ExecOptions{DecodeOption: spannerdriver.DecodeOptionProto})
+	if statement.Params != nil {
+		if statement.ParamTypes == nil {
+			statement.ParamTypes = make(map[string]*spannerpb.Type)
+		}
+		for param, value := range statement.Params.Fields {
+			genericValue := spanner.GenericColumnValue{
+				Value: value,
+				Type:  statement.ParamTypes[param],
+			}
+			params = append(params, sql.Named(param, genericValue))
+		}
+	}
+	it, err := conn.backend.Conn.QueryContext(context.Background(), statement.Sql, params...)
+	if err != nil {
+		return errMessage(err)
+	}
+	id := conn.resultsIdx.Add(1)
+	res := &rows{
+		backend: it,
+	}
+	conn.results.Store(id, res)
+	return idMessage(id)
+}

spannerlib/exported/message.go

@@ -0,0 +1,55 @@
+package exported
+
+import (
+	"unsafe"
+
+	"github.com/google/uuid"
+	"google.golang.org/grpc/status"
+)
+
+type Message struct {
+	Code     int32
+	ObjectId int64
+	Res      []byte
+}
+
+func (m *Message) Length() int32 {
+	if m.Res == nil {
+		return 0
+	}
+	return int32(len(m.Res))
+}
+
+func (m *Message) ResPointer() unsafe.Pointer {
+	if m.Res == nil {
+		return nil
+	}
+	return unsafe.Pointer(&(m.Res[0]))
+}
+
+func generateId() (string, error) {
+	id, err := uuid.NewRandom()
+	if err != nil {
+		return "", err
+	}
+	return id.String(), nil
+}
+
+type BaseMsg struct {
+	Allowed bool
+}
+
+type OtherMsg struct {
+	BaseMsg
+	Foo string
+}
+
+func idMessage(id int64) *Message {
+	return &Message{ObjectId: id}
+}
+
+func errMessage(err error) *Message {
+	errCode := status.Code(err)
+	b := []byte(err.Error())
+	return &Message{Code: int32(errCode), Res: b}
+}

spannerlib/exported/pool.go

@@ -0,0 +1,92 @@
+package exported
+
+import "C"
+import (
+	"context"
+	"fmt"
+	"sync"
+	"sync/atomic"
+
+	"spannerlib/backend"
+)
+
+var pools = sync.Map{}
+var poolsIdx = atomic.Int64{}
+
+type Pool struct {
+	backend        *backend.Pool
+	connections    *sync.Map
+	connectionsIdx atomic.Int64
+}
+
+func CreatePool() *Message {
+	backendPool := &backend.Pool{}
+	id := poolsIdx.Add(1)
+	pool := &Pool{
+		backend:     backendPool,
+		connections: &sync.Map{},
+	}
+	pools.Store(id, pool)
+	return idMessage(id)
+}
+
+func ClosePool(id int64) *Message {
+	p, ok := pools.LoadAndDelete(id)
+	if !ok {
+		return errMessage(fmt.Errorf("pool %v not found", id))
+	}
+	pool := p.(*Pool)
+	pool.connections.Range(func(key, value interface{}) bool {
+		conn := value.(*Connection)
+		conn.close()
+		return true
+	})
+	return &Message{}
+}
+
+func CreateConnection(poolId int64, project, instance, database string) *Message {
+	p, ok := pools.Load(poolId)
+	if !ok {
+		return errMessage(fmt.Errorf("pool %v not found", poolId))
+	}
+	pool := p.(*Pool)
+	sqlConn, err := pool.backend.Conn(context.Background(), project, instance, database)
+	if err != nil {
+		return errMessage(err)
+	}
+	id := poolsIdx.Add(1)
+	conn := &Connection{
+		backend: &backend.SpannerConnection{Conn: sqlConn},
+		results: &sync.Map{},
+	}
+	pool.connections.Store(id, conn)
+
+	return idMessage(id)
+}
+
+func findConnection(poolId, connId int64) (*Connection, error) {
+	p, ok := pools.Load(poolId)
+	if !ok {
+		return nil, fmt.Errorf("pool %v not found", poolId)
+	}
+	pool := p.(*Pool)
+	c, ok := pool.connections.Load(connId)
+	if !ok {
+		return nil, fmt.Errorf("connection %v not found", connId)
+	}
+	conn := c.(*Connection)
+	return conn, nil
+}
+
+func findRows(poolId, connId, rowsId int64) (*rows, error) {
+	conn, err := findConnection(poolId, connId)
+	if err != nil {
+		return nil, err
+	}
+	r, ok := conn.results.Load(rowsId)
+	if !ok {
+		return nil, fmt.Errorf("rows %v not found", rowsId)
+	}
+	res := r.(*rows)
+	return res, nil
+}

spannerlib/exported/pool_test.go

@@ -0,0 +1,34 @@
+package exported
+
+import (
+	"cloud.google.com/go/spanner/apiv1/spannerpb"
+	"fmt"
+	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/types/known/structpb"
+	"testing"
+)
+
+func TestExecute(t *testing.T) {
+	pool := CreatePool()
+	conn := CreateConnection(pool.ObjectId, "appdev-soda-spanner-staging", "knut-test-ycsb", "knut-test-db")
+	stmt := spannerpb.ExecuteBatchDmlRequest_Statement{
+		Sql: "select * from all_types where col_varchar=$1 /*and col_bigint=@id*/ limit 10",
+		Params: &structpb.Struct{
+			Fields: map[string]*structpb.Value{"p1": {Kind: &structpb.Value_StringValue{StringValue: "61763b0e7feb3ea8fc9e734a6700f6a4"}}},
+		},
+	}
+	stmtBytes, _ := proto.Marshal(&stmt)
+	rows := Execute(pool.ObjectId, conn.ObjectId, stmtBytes)
+	for {
+		row := Next(pool.ObjectId, conn.ObjectId, rows.ObjectId)
+		rowValue := structpb.ListValue{}
+		_ = proto.Unmarshal(row.Res, &rowValue)
+		if row.Length() == 0 {
+			break
+		}
+		fmt.Printf("row: %v\n", rowValue.Values)
+	}
+	CloseRows(pool.ObjectId, conn.ObjectId, rows.ObjectId)
+	CloseConnection(pool.ObjectId, conn.ObjectId)
+	ClosePool(pool.ObjectId)
+}