Add index methods on engine transaction interface

engine/engine.go

@@ -3,6 +3,8 @@ package engine
 import (
 	"errors"
 
+	"github.com/asdine/genji/index"
+
 	"github.com/asdine/genji/table"
 )
 
@@ -21,4 +23,6 @@ type Transaction interface {
 	Commit() error
 	Table(name string) (table.Table, error)
 	CreateTable(name string) (table.Table, error)
+	Index(name string) (index.Index, error)
+	CreateIndex(name string) (index.Index, error)
 }

engine/memory/engine.go

@@ -7,20 +7,23 @@ import (
 	"sync"
 
 	"github.com/asdine/genji/engine"
+	"github.com/asdine/genji/index"
 	"github.com/asdine/genji/table"
 	"modernc.org/b"
 )
 
 type Engine struct {
-	closed bool
+	closed  bool
-	tables map[string]*b.Tree
+	tables  map[string]*b.Tree
+	indexes map[string]*Index
 
 	mu sync.RWMutex
 }
 
 func NewEngine() *Engine {
 	return &Engine{
-		tables: make(map[string]*b.Tree),
+		tables:  make(map[string]*b.Tree),
+		indexes: make(map[string]*Index),
 	}
 }
 
@@ -123,3 +126,31 @@ func (tx *transaction) CreateTable(name string) (table.Table, error) {
 
 	return &tableTx{tx: tx, tree: tr}, nil
 }
+
+func (tx *transaction) Index(name string) (index.Index, error) {
+	idx, ok := tx.ng.indexes[name]
+	if !ok {
+		return nil, errors.New("index not found")
+	}
+
+	return idx, nil
+}
+
+func (tx *transaction) CreateIndex(name string) (index.Index, error) {
+	if !tx.writable {
+		return nil, errors.New("can't create index in read-only transaction")
+	}
+
+	_, err := tx.Index(name)
+	if err == nil {
+		return nil, fmt.Errorf("index '%s' already exists", name)
+	}
+
+	tx.ng.indexes[name] = NewIndex()
+
+	tx.undos = append(tx.undos, func() {
+		delete(tx.ng.indexes, name)
+	})
+
+	return tx.ng.indexes[name], nil
+}

engine/memory/engine_test.go

@@ -0,0 +1,24 @@
+package memory
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestTransaction(t *testing.T) {
+	t.Run("index", func(t *testing.T) {
+		ng := NewEngine()
+		tx, err := ng.Begin(true)
+		require.NoError(t, err)
+		defer tx.Rollback()
+
+		_, err = tx.Index("test")
+		require.Error(t, err)
+		idx1, err := tx.CreateIndex("test")
+		require.NoError(t, err)
+		idx2, err := tx.Index("test")
+		require.NoError(t, err)
+		require.Equal(t, idx1, idx2)
+	})
+}