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,45 @@
+package exported
+
+import "C"
+import (
+	"cloud.google.com/go/spanner/apiv1/spannerpb"
+	"context"
+	"sync"
+	"sync/atomic"
+
+	spannerdriver "github.com/googleapis/go-sql-spanner"
+	"spannerlib/backend"
+)
+
+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 {
+	it, err := conn.backend.Conn.QueryContext(context.Background(), statement.Sql, spannerdriver.ExecOptions{DecodeOption: spannerdriver.DecodeOptionProto})
+	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,132 @@
+package exported
+
+import "C"
+import (
+	"cloud.google.com/go/spanner/apiv1/spannerpb"
+	"context"
+	"fmt"
+	"google.golang.org/protobuf/proto"
+	"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, 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(), "appdev-soda-spanner-staging", "knut-test-ycsb", "knut-test-db")
+	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 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)
+}
+
+func Next(poolId, connId, rowsId int64) *Message {
+	res, err := findRows(poolId, connId, rowsId)
+	if err != nil {
+		return errMessage(err)
+	}
+	return res.Next()
+}
+
+func CloseRows(poolId, connId, rowsId int64) *Message {
+	res, err := findRows(poolId, connId, rowsId)
+	if err != nil {
+		return errMessage(err)
+	}
+	return res.Close()
+}
+
+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/rows.go

@@ -0,0 +1,51 @@
+package exported
+
+import (
+	"database/sql"
+
+	"cloud.google.com/go/spanner"
+	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/types/known/structpb"
+)
+
+type rows struct {
+	backend *sql.Rows
+}
+
+func (rows *rows) Close() *Message {
+	err := rows.backend.Close()
+	if err != nil {
+		return errMessage(err)
+	}
+	return &Message{}
+}
+
+func (rows *rows) Next() *Message {
+	ok := rows.backend.Next()
+	if !ok {
+		// An empty message indicates no more rows.
+		return &Message{}
+	}
+	cols, err := rows.backend.Columns()
+	if err != nil {
+		return errMessage(err)
+	}
+	buffer := make([]any, len(cols))
+	for i := range buffer {
+		buffer[i] = &spanner.GenericColumnValue{}
+	}
+	if err := rows.backend.Scan(buffer...); err != nil {
+		return errMessage(err)
+	}
+	row := &structpb.ListValue{
+		Values: make([]*structpb.Value, len(buffer)),
+	}
+	for i := range buffer {
+		row.Values[i] = buffer[i].(*spanner.GenericColumnValue).Value
+	}
+	res, err := proto.Marshal(row)
+	if err != nil {
+		return errMessage(err)
+	}
+	return &Message{Res: res}
+}