Repository: genjidb/genji (go)

Issue:
Add support for context.Context in engines
PR #221 added support for passing context to high-level APIs. We should also pass them to the engine implementations. E.g. it’d be possible to pass context.Context to `tikvengine` (though currently there is no such engine, but the API maps well to the Genji engine definitions).

See also: #206


Pull request description:
Add cancellation in engines
This PR fixes #224 by adding support for cancellation to engines.

Patch:
diff --git a/db.go b/db.go
index 5af7e692..02072edb 100644
--- a/db.go
+++ b/db.go
@@ -13,13 +13,14 @@ import (
 type DB struct {
 	DB *database.Database
 
-	Context context.Context
+	ctx context.Context
 }
 
+// WithContext creates a new database handle using the given context for every operation.
 func (db *DB) WithContext(ctx context.Context) *DB {
 	return &DB{
-		DB:      db.DB,
-		Context: ctx,
+		DB:  db.DB,
+		ctx: ctx,
 	}
 }
 
@@ -31,7 +32,7 @@ func (db *DB) Close() error {
 // Begin starts a new transaction.
 // The returned transaction must be closed either by calling Rollback or Commit.
 func (db *DB) Begin(writable bool) (*Tx, error) {
-	tx, err := db.DB.BeginTx(db.Context, &database.TxOptions{
+	tx, err := db.DB.BeginTx(db.ctx, &database.TxOptions{
 		ReadOnly: !writable,
 	})
 	if err != nil {
@@ -88,7 +89,7 @@ func (db *DB) Query(q string, args ...interface{}) (*query.Result, error) {
 		return nil, err
 	}
 
-	return pq.Run(db.Context, db.DB, argsToParams(args))
+	return pq.Run(db.ctx, db.DB, argsToParams(args))
 }
 
 // QueryDocument runs the query and returns the first document.
diff --git a/document/iterator.go b/document/iterator.go
index 98a7153e..cd1f3107 100644
--- a/document/iterator.go
+++ b/document/iterator.go
@@ -2,6 +2,7 @@ package document
 
 import (
 	"bufio"
+	"bytes"
 	"errors"
 	"io"
 )
@@ -282,11 +283,14 @@ func (s Stream) GroupBy(groupFn func(d Document) (Value, error)) Stream {
 // Aggregate builds a list of aggregators for each group of documents and passes each document of the stream to them.
 func (s Stream) Aggregate(aggregatorBuilders ...AggregatorBuilder) Stream {
 	return NewStream(IteratorFunc(func(fn func(d Document) error) error {
-		aggregates := make(map[Value][]Aggregator)
-		var groups []Value
+		aggregates := make(map[string][]Aggregator)
+		var groupKeys []string
 
 		nullValue := NewNullValue()
 
+		var b bytes.Buffer
+
+		enc := NewValueEncoder(&b)
 		err := s.Iterate(func(d Document) error {
 			group := nullValue
 
@@ -294,17 +298,26 @@ func (s Stream) Aggregate(aggregatorBuilders ...AggregatorBuilder) Stream {
 				group = gd.group
 			}
 
-			aggs, ok := aggregates[group]
+			b.Reset()
+
+			err := enc.Encode(group)
+			if err != nil {
+				return err
+			}
+
+			// using b.Bytes() and cast as a string is much faster than using b.String()
+			// because of compiler optimizations.
+			aggs, ok := aggregates[string(b.Bytes())]
 			if !ok {
-				groups = append(groups, group)
+				groupKey := b.String()
+				groupKeys = append(groupKeys, groupKey)
 				aggs = make([]Aggregator, len(aggregatorBuilders))
 				for i, builder := range aggregatorBuilders {
 					aggs[i] = builder.NewAggregator(group)
 				}
-				aggregates[group] = aggs
+				aggregates[groupKey] = aggs
 			}
 
-			var err error
 			for _, agg := range aggs {
 				err = agg.Add(d)
 				if err != nil {
@@ -318,9 +331,9 @@ func (s Stream) Aggregate(aggregatorBuilders ...AggregatorBuilder) Stream {
 			return err
 		}
 
-		for _, group := range groups {
+		for _, groupKey := range groupKeys {
 			fb := NewFieldBuffer()
-			aggs := aggregates[group]
+			aggs := aggregates[groupKey]
 			for _, agg := range aggs {
 				err = agg.Aggregate(fb)
 				if err != nil {
diff --git a/engine/badgerengine/engine.go b/engine/badgerengine/engine.go
index f4321f8a..f0611caa 100644
--- a/engine/badgerengine/engine.go
+++ b/engine/badgerengine/engine.go
@@ -34,9 +34,16 @@ func NewEngine(opt badger.Options) (*Engine, error) {
 
 // Begin creates a transaction using Badger's transaction API.
 func (e *Engine) Begin(ctx context.Context, opts engine.TxOptions) (engine.Transaction, error) {
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	default:
+	}
+
 	tx := e.DB.NewTransaction(opts.Writable)
 
 	return &Transaction{
+		ctx:      ctx,
 		ng:       e,
 		tx:       tx,
 		writable: opts.Writable,
@@ -50,6 +57,7 @@ func (e *Engine) Close() error {
 
 // A Transaction uses Badger's transactions.
 type Transaction struct {
+	ctx       context.Context
 	ng        *Engine
 	tx        *badger.Txn
 	writable  bool
@@ -61,11 +69,23 @@ func (t *Transaction) Rollback() error {
 	t.tx.Discard()
 
 	t.discarded = true
+	select {
+	case <-t.ctx.Done():
+		return t.ctx.Err()
+	default:
+	}
 	return nil
 }
 
 // Commit the transaction.
 func (t *Transaction) Commit() error {
+	select {
+	case <-t.ctx.Done():
+		_ = t.Rollback()
+		return t.ctx.Err()
+	default:
+	}
+
 	if t.discarded {
 		return badger.ErrDiscardedTxn
 	}
@@ -99,6 +119,12 @@ func buildStorePrefixKey(name []byte) []byte {
 
 // GetStore returns a store by name.
 func (t *Transaction) GetStore(name []byte) (engine.Store, error) {
+	select {
+	case <-t.ctx.Done():
+		return nil, t.ctx.Err()
+	default:
+	}
+
 	key := buildStoreKey(name)
 
 	_, err := t.tx.Get(key)
@@ -113,6 +139,7 @@ func (t *Transaction) GetStore(name []byte) (engine.Store, error) {
 	pkey := buildStorePrefixKey(name)
 
 	return &Store{
+		ctx:      t.ctx,
 		ng:       t.ng,
 		tx:       t.tx,
 		prefix:   pkey,
@@ -124,6 +151,12 @@ func (t *Transaction) GetStore(name []byte) (engine.Store, error) {
 // CreateStore creates a store.
 // If the store already exists, returns engine.ErrStoreAlreadyExists.
 func (t *Transaction) CreateStore(name []byte) error {
+	select {
+	case <-t.ctx.Done():
+		return t.ctx.Err()
+	default:
+	}
+
 	if !t.writable {
 		return engine.ErrTransactionReadOnly
 	}
@@ -142,6 +175,12 @@ func (t *Transaction) CreateStore(name []byte) error {
 
 // DropStore deletes the store and all its keys.
 func (t *Transaction) DropStore(name []byte) error {
+	select {
+	case <-t.ctx.Done():
+		return t.ctx.Err()
+	default:
+	}
+
 	if !t.writable {
 		return engine.ErrTransactionReadOnly
 	}
diff --git a/engine/badgerengine/store.go b/engine/badgerengine/store.go
index e3d6f2ea..bc29a3bb 100644
--- a/engine/badgerengine/store.go
+++ b/engine/badgerengine/store.go
@@ -2,6 +2,7 @@ package badgerengine
 
 import (
 	"bytes"
+	"context"
 	"errors"
 
 	"github.com/dgraph-io/badger/v2"
@@ -10,6 +11,7 @@ import (
 
 // A Store is an implementation of the engine.Store interface.
 type Store struct {
+	ctx      context.Context
 	ng       *Engine
 	tx       *badger.Txn
 	prefix   []byte
@@ -34,6 +36,12 @@ func buildKey(prefix, k []byte) []byte {
 
 // Put stores a key value pair. If it already exists, it overrides it.
 func (s *Store) Put(k, v []byte) error {
+	select {
+	case <-s.ctx.Done():
+		return s.ctx.Err()
+	default:
+	}
+
 	if !s.writable {
 		return engine.ErrTransactionReadOnly
 	}
@@ -47,6 +55,12 @@ func (s *Store) Put(k, v []byte) error {
 
 // Get returns a value associated with the given key. If not found, returns engine.ErrKeyNotFound.
 func (s *Store) Get(k []byte) ([]byte, error) {
+	select {
+	case <-s.ctx.Done():
+		return nil, s.ctx.Err()
+	default:
+	}
+
 	it, err := s.tx.Get(buildKey(s.prefix, k))
 	if err != nil {
 		if err == badger.ErrKeyNotFound {
@@ -61,6 +75,12 @@ func (s *Store) Get(k []byte) ([]byte, error) {
 
 // Delete a record by key. If not found, returns engine.ErrKeyNotFound.
 func (s *Store) Delete(k []byte) error {
+	select {
+	case <-s.ctx.Done():
+		return s.ctx.Err()
+	default:
+	}
+
 	if !s.writable {
 		return engine.ErrTransactionReadOnly
 	}
@@ -80,6 +100,12 @@ func (s *Store) Delete(k []byte) error {
 
 // Truncate deletes all the records of the store.
 func (s *Store) Truncate() error {
+	select {
+	case <-s.ctx.Done():
+		return s.ctx.Err()
+	default:
+	}
+
 	if !s.writable {
 		return engine.ErrTransactionReadOnly
 	}
@@ -105,6 +131,12 @@ func (s *Store) Truncate() error {
 
 // NextSequence returns a monotonically increasing integer.
 func (s *Store) NextSequence() (uint64, error) {
+	select {
+	case <-s.ctx.Done():
+		return 0, s.ctx.Err()
+	default:
+	}
+
 	if !s.writable {
 		return 0, engine.ErrTransactionReadOnly
 	}
@@ -138,6 +170,7 @@ func (s *Store) Iterator(opts engine.IteratorOptions) engine.Iterator {
 	it := s.tx.NewIterator(opt)
 
 	return &iterator{
+		ctx:         s.ctx,
 		storePrefix: s.prefix,
 		prefix:      prefix,
 		it:          it,
@@ -147,14 +180,23 @@ func (s *Store) Iterator(opts engine.IteratorOptions) engine.Iterator {
 }
 
 type iterator struct {
+	ctx         context.Context
 	prefix      []byte
 	storePrefix []byte
 	it          *badger.Iterator
 	reverse     bool
 	item        badgerItem
+	err         error
 }
 
 func (it *iterator) Seek(pivot []byte) {
+	select {
+	case <-it.ctx.Done():
+		it.err = it.ctx.Err()
+		return
+	default:
+	}
+
 	var seek []byte
 
 	if !it.reverse {
@@ -175,7 +217,7 @@ func (it *iterator) Seek(pivot []byte) {
 }
 
 func (it *iterator) Valid() bool {
-	return it.it.ValidForPrefix(it.prefix)
+	return it.it.ValidForPrefix(it.prefix) && it.err == nil
 }
 
 func (it *iterator) Next() {
@@ -183,7 +225,7 @@ func (it *iterator) Next() {
 }
 
 func (it *iterator) Err() error {
-	return nil
+	return it.err
 }
 
 func (it *iterator) Item() engine.Item {
diff --git a/engine/boltengine/engine.go b/engine/boltengine/engine.go
index ad013dce..606168ae 100644
--- a/engine/boltengine/engine.go
+++ b/engine/boltengine/engine.go
@@ -32,12 +32,19 @@ func NewEngine(path string, mode os.FileMode, opts *bolt.Options) (*Engine, erro
 
 // Begin creates a transaction using Bolt's transaction API.
 func (e *Engine) Begin(ctx context.Context, opts engine.TxOptions) (engine.Transaction, error) {
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	default:
+	}
+
 	tx, err := e.DB.Begin(opts.Writable)
 	if err != nil {
 		return nil, err
 	}
 
 	return &Transaction{
+		ctx:      ctx,
 		tx:       tx,
 		writable: opts.Writable,
 	}, nil
@@ -50,6 +57,7 @@ func (e *Engine) Close() error {
 
 // A Transaction uses Bolt's transactions.
 type Transaction struct {
+	ctx      context.Context
 	tx       *bolt.Tx
 	writable bool
 }
@@ -61,16 +69,35 @@ func (t *Transaction) Rollback() error {
 		return err
 	}
 
+	select {
+	case <-t.ctx.Done():
+		return t.ctx.Err()
+	default:
+	}
+
 	return nil
 }
 
 // Commit the transaction.
 func (t *Transaction) Commit() error {
+	select {
+	case <-t.ctx.Done():
+		_ = t.Rollback()
+		return t.ctx.Err()
+	default:
+	}
+
 	return t.tx.Commit()
 }
 
 // GetStore returns a store by name. The store uses a Bolt bucket.
 func (t *Transaction) GetStore(name []byte) (engine.Store, error) {
+	select {
+	case <-t.ctx.Done():
+		return nil, t.ctx.Err()
+	default:
+	}
+
 	b := t.tx.Bucket(name)
 	if b == nil {
 		return nil, engine.ErrStoreNotFound
@@ -80,12 +107,19 @@ func (t *Transaction) GetStore(name []byte) (engine.Store, error) {
 		bucket: b,
 		tx:     t.tx,
 		name:   name,
+		ctx:    t.ctx,
 	}, nil
 }
 
 // CreateStore creates a bolt bucket and returns a store.
 // If the store already exists, returns engine.ErrStoreAlreadyExists.
 func (t *Transaction) CreateStore(name []byte) error {
+	select {
+	case <-t.ctx.Done():
+		return t.ctx.Err()
+	default:
+	}
+
 	if !t.writable {
 		return engine.ErrTransactionReadOnly
 	}
@@ -100,6 +134,12 @@ func (t *Transaction) CreateStore(name []byte) error {
 
 // DropStore deletes the underlying bucket.
 func (t *Transaction) DropStore(name []byte) error {
+	select {
+	case <-t.ctx.Done():
+		return t.ctx.Err()
+	default:
+	}
+
 	if !t.writable {
 		return engine.ErrTransactionReadOnly
 	}
diff --git a/engine/boltengine/store.go b/engine/boltengine/store.go
index 2ed80498..653cfbf6 100644
--- a/engine/boltengine/store.go
+++ b/engine/boltengine/store.go
@@ -2,6 +2,7 @@ package boltengine
 
 import (
 	"bytes"
+	"context"
 
 	"github.com/genjidb/genji/engine"
 	bolt "go.etcd.io/bbolt"
@@ -12,10 +13,17 @@ type Store struct {
 	bucket *bolt.Bucket
 	tx     *bolt.Tx
 	name   []byte
+	ctx    context.Context
 }
 
 // Put stores a key value pair. If it already exists, it overrides it.
 func (s *Store) Put(k, v []byte) error {
+	select {
+	case <-s.ctx.Done():
+		return s.ctx.Err()
+	default:
+	}
+
 	if !s.bucket.Writable() {
 		return engine.ErrTransactionReadOnly
 	}
@@ -25,6 +33,12 @@ func (s *Store) Put(k, v []byte) error {
 
 // Get returns a value associated with the given key. If not found, returns engine.ErrKeyNotFound.
 func (s *Store) Get(k []byte) ([]byte, error) {
+	select {
+	case <-s.ctx.Done():
+		return nil, s.ctx.Err()
+	default:
+	}
+
 	v := s.bucket.Get(k)
 	if v == nil {
 		return nil, engine.ErrKeyNotFound
@@ -35,6 +49,12 @@ func (s *Store) Get(k []byte) ([]byte, error) {
 
 // Delete a record by key. If not found, returns table.ErrDocumentNotFound.
 func (s *Store) Delete(k []byte) error {
+	select {
+	case <-s.ctx.Done():
+		return s.ctx.Err()
+	default:
+	}
+
 	if !s.bucket.Writable() {
 		return engine.ErrTransactionReadOnly
 	}
@@ -49,6 +69,12 @@ func (s *Store) Delete(k []byte) error {
 
 // Truncate deletes all the records of the store.
 func (s *Store) Truncate() error {
+	select {
+	case <-s.ctx.Done():
+		return s.ctx.Err()
+	default:
+	}
+
 	if !s.bucket.Writable() {
 		return engine.ErrTransactionReadOnly
 	}
@@ -64,6 +90,12 @@ func (s *Store) Truncate() error {
 
 // NextSequence returns a monotonically increasing integer.
 func (s *Store) NextSequence() (uint64, error) {
+	select {
+	case <-s.ctx.Done():
+		return 0, s.ctx.Err()
+	default:
+	}
+
 	if !s.bucket.Writable() {
 		return 0, engine.ErrTransactionReadOnly
 	}
@@ -76,6 +108,7 @@ func (s *Store) Iterator(opts engine.IteratorOptions) engine.Iterator {
 	return &iterator{
 		c:       s.bucket.Cursor(),
 		reverse: opts.Reverse,
+		ctx:     s.ctx,
 	}
 }
 
@@ -83,9 +116,18 @@ type iterator struct {
 	c       *bolt.Cursor
 	reverse bool
 	item    boltItem
+	err     error
+	ctx     context.Context
 }
 
 func (it *iterator) Seek(pivot []byte) {
+	select {
+	case <-it.ctx.Done():
+		it.err = it.ctx.Err()
+		return
+	default:
+	}
+
 	if !it.reverse {
 		it.item.k, it.item.v = it.c.Seek(pivot)
 		return
@@ -105,7 +147,7 @@ func (it *iterator) Seek(pivot []byte) {
 }
 
 func (it *iterator) Valid() bool {
-	return it.item.k != nil
+	return it.item.k != nil && it.err == nil
 }
 
 func (it *iterator) Next() {
@@ -117,7 +159,7 @@ func (it *iterator) Next() {
 }
 
 func (it *iterator) Err() error {
-	return nil
+	return it.err
 }
 
 func (it *iterator) Item() engine.Item {
diff --git a/engine/memoryengine/engine.go b/engine/memoryengine/engine.go
index 2b01152d..8356c630 100644
--- a/engine/memoryengine/engine.go
+++ b/engine/memoryengine/engine.go
@@ -34,6 +34,12 @@ func NewEngine() *Engine {
 
 // Begin creates a transaction.
 func (ng *Engine) Begin(ctx context.Context, opts engine.TxOptions) (engine.Transaction, error) {
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	default:
+	}
+
 	if opts.Writable {
 		ng.mu.Lock()
 	} else {
@@ -44,7 +50,7 @@ func (ng *Engine) Begin(ctx context.Context, opts engine.TxOptions) (engine.Tran
 		return nil, errors.New("engine closed")
 	}
 
-	return &transaction{ng: ng, writable: opts.Writable}, nil
+	return &transaction{ctx: ctx, ng: ng, writable: opts.Writable}, nil
 }
 
 // Close the engine.
@@ -61,6 +67,7 @@ func (ng *Engine) Close() error {
 
 // This implements the engine.Transaction type.
 type transaction struct {
+	ctx        context.Context
 	ng         *Engine
 	writable   bool
 	onRollback []func() // called during a rollback
@@ -91,6 +98,12 @@ func (tx *transaction) Rollback() error {
 		tx.ng.mu.RUnlock()
 	}
 
+	select {
+	case <-tx.ctx.Done():
+		return tx.ctx.Err()
+	default:
+	}
+
 	return nil
 }
 
@@ -109,6 +122,12 @@ func (tx *transaction) Commit() error {
 
 	tx.wg.Wait()
 
+	select {
+	case <-tx.ctx.Done():
+		return tx.Rollback()
+	default:
+	}
+
 	tx.terminated = true
 
 	for _, fn := range tx.onCommit {
@@ -121,6 +140,12 @@ func (tx *transaction) Commit() error {
 }
 
 func (tx *transaction) GetStore(name []byte) (engine.Store, error) {
+	select {
+	case <-tx.ctx.Done():
+		return nil, tx.ctx.Err()
+	default:
+	}
+
 	tr, ok := tx.ng.stores[string(name)]
 	if !ok {
 		return nil, engine.ErrStoreNotFound
@@ -130,6 +155,12 @@ func (tx *transaction) GetStore(name []byte) (engine.Store, error) {
 }
 
 func (tx *transaction) CreateStore(name []byte) error {
+	select {
+	case <-tx.ctx.Done():
+		return tx.ctx.Err()
+	default:
+	}
+
 	if !tx.writable {
 		return engine.ErrTransactionReadOnly
 	}
@@ -152,6 +183,12 @@ func (tx *transaction) CreateStore(name []byte) error {
 }
 
 func (tx *transaction) DropStore(name []byte) error {
+	select {
+	case <-tx.ctx.Done():
+		return tx.ctx.Err()
+	default:
+	}
+
 	if !tx.writable {
 		return engine.ErrTransactionReadOnly
 	}
diff --git a/engine/memoryengine/store.go b/engine/memoryengine/store.go
index 1fde8f6a..d47957f7 100644
--- a/engine/memoryengine/store.go
+++ b/engine/memoryengine/store.go
@@ -43,6 +43,12 @@ type storeTx struct {
 }
 
 func (s *storeTx) Put(k, v []byte) error {
+	select {
+	case <-s.tx.ctx.Done():
+		return s.tx.ctx.Err()
+	default:
+	}
+
 	if !s.tx.writable {
 		return engine.ErrTransactionReadOnly
 	}
@@ -82,6 +88,12 @@ func (s *storeTx) Put(k, v []byte) error {
 }
 
 func (s *storeTx) Get(k []byte) ([]byte, error) {
+	select {
+	case <-s.tx.ctx.Done():
+		return nil, s.tx.ctx.Err()
+	default:
+	}
+
 	it := s.tr.Get(&item{k: k})
 
 	if it == nil {
@@ -105,6 +117,12 @@ func (s *storeTx) Get(k []byte) ([]byte, error) {
 // which causes iterators to behave incorrectly when looping
 // and deleting at the same time.
 func (s *storeTx) Delete(k []byte) error {
+	select {
+	case <-s.tx.ctx.Done():
+		return s.tx.ctx.Err()
+	default:
+	}
+
 	if !s.tx.writable {
 		return engine.ErrTransactionReadOnly
 	}
@@ -145,6 +163,12 @@ func (s *storeTx) Delete(k []byte) error {
 // one. The current tree will be garbage collected
 // once the transaction is commited.
 func (s *storeTx) Truncate() error {
+	select {
+	case <-s.tx.ctx.Done():
+		return s.tx.ctx.Err()
+	default:
+	}
+
 	if !s.tx.writable {
 		return engine.ErrTransactionReadOnly
 	}
@@ -162,6 +186,12 @@ func (s *storeTx) Truncate() error {
 
 // NextSequence returns a monotonically increasing integer.
 func (s *storeTx) NextSequence() (uint64, error) {
+	select {
+	case <-s.tx.ctx.Done():
+		return 0, s.tx.ctx.Err()
+	default:
+	}
+
 	if !s.tx.writable {
 		return 0, engine.ErrTransactionReadOnly
 	}
@@ -192,6 +222,7 @@ type iterator struct {
 	closed  chan struct{} // closed by the goroutine when it's shutdown
 	ctx     context.Context
 	cancel  func()
+	err     error
 }
 
 func (it *iterator) Seek(pivot []byte) {
@@ -204,7 +235,7 @@ func (it *iterator) Seek(pivot []byte) {
 		it.closed = make(chan struct{})
 	}
 
-	it.ctx, it.cancel = context.WithCancel(context.Background())
+	it.ctx, it.cancel = context.WithCancel(it.tx.ctx)
 
 	it.runIterator(pivot)
 
@@ -259,16 +290,20 @@ func (it *iterator) runIterator(pivot []byte) {
 }
 
 func (it *iterator) Valid() bool {
-	return it.item != nil
+	return it.item != nil && it.err == nil
 }
 
 // Read the next item from the goroutine
 func (it *iterator) Next() {
-	it.item = <-it.ch
+	select {
+	case it.item = <-it.ch:
+	case <-it.tx.ctx.Done():
+		it.err = it.tx.ctx.Err()
+	}
 }
 
 func (it *iterator) Err() error {
-	return nil
+	return it.err
 }
 
 func (it *iterator) Item() engine.Item {
diff --git a/new.go b/new.go
index d691c284..a80be632 100644
--- a/new.go
+++ b/new.go
@@ -18,7 +18,7 @@ func New(ctx context.Context, ng engine.Engine) (*DB, error) {
 	}
 
 	return &DB{
-		DB:      db,
-		Context: context.Background(),
+		DB:  db,
+		ctx: context.Background(),
 	}, nil
 }
diff --git a/new_wasm.go b/new_wasm.go
index a8b43b24..44c05aff 100644
--- a/new_wasm.go
+++ b/new_wasm.go
@@ -18,7 +18,7 @@ func New(ctx context.Context, ng engine.Engine) (*DB, error) {
 	}
 
 	return &DB{
-		DB:      db,
-		Context: context.Background(),
+		DB:  db,
+		ctx: context.Background(),
 	}, nil
 }
diff --git a/sql/planner/tree.go b/sql/planner/tree.go
index dbe1eb11..d45c9d83 100644
--- a/sql/planner/tree.go
+++ b/sql/planner/tree.go
@@ -42,7 +42,7 @@ const (
 	// Unset is an operation that removes a value at a given path from every document of a stream
 	Unset
 	// Group is an operation that groups documents based on a given path.
-	_
+	Group
 	// Dedup is an operation that removes duplicate documents from a stream
 	Dedup
 )
@@ -444,9 +444,9 @@ func (n *unsetNode) String() string {
 type GroupingNode struct {
 	node
 
-	Expr   expr.Expr
 	Tx     *database.Transaction
 	Params []expr.Param
+	Expr   expr.Expr
 }
 
 var _ operationNode = (*GroupingNode)(nil)
@@ -455,7 +455,7 @@ var _ operationNode = (*GroupingNode)(nil)
 func NewGroupingNode(n Node, e expr.Expr) Node {
 	return &GroupingNode{
 		node: node{
-			op:   Projection,
+			op:   Group,
 			left: n,
 		},
 		Expr: e,
@@ -469,6 +469,8 @@ func (n *GroupingNode) Bind(tx *database.Transaction, params []expr.Param) (err
 	return
 }
 
+// toStream uses the GroupBy stream operation. It evaluates Expr for every field and returns
+// the result.
 func (n *GroupingNode) toStream(st document.Stream) (document.Stream, error) {
 	return st.GroupBy(func(d document.Document) (document.Value, error) {
 		return n.Expr.Eval(expr.EvalStack{


Test patch:
diff --git a/engine/enginetest/testing.go b/engine/enginetest/testing.go
index edc71cde..c21c38a9 100644
--- a/engine/enginetest/testing.go
+++ b/engine/enginetest/testing.go
@@ -31,7 +31,7 @@ func TestSuite(t *testing.T, builder Builder) {
 	}{
 		{"Engine", TestEngine},
 		{"Transaction/Commit-Rollback", TestTransactionCommitRollback},
-		{"Transaction/Store", TestTransactionStore},
+		{"Transaction/GetStore", TestTransactionGetStore},
 		{"Transaction/CreateStore", TestTransactionCreateStore},
 		{"Transaction/DropStore", TestTransactionDropStore},
 		{"Store/Iterator", TestStoreIterator},
@@ -95,6 +95,35 @@ func TestTransactionCommitRollback(t *testing.T, builder Builder) {
 		require.Error(t, err)
 	})
 
+	t.Run("Commit after context canceled should fail", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(context.Background())
+		tx, err := ng.Begin(ctx, engine.TxOptions{
+			Writable: true,
+		})
+		require.NoError(t, err)
+
+		require.NoError(t, tx.CreateStore([]byte("test")))
+		st, err := tx.GetStore([]byte("test"))
+		require.NoError(t, err)
+		err = st.Put([]byte("a"), []byte("b"))
+		require.NoError(t, err)
+
+		cancel()
+
+		err = tx.Commit()
+		require.Error(t, err)
+
+		// ensure data has not been persisted
+		tx, err = ng.Begin(context.Background(), engine.TxOptions{
+			Writable: false,
+		})
+		require.NoError(t, err)
+		defer tx.Rollback()
+
+		st, err = tx.GetStore([]byte("test"))
+		require.Error(t, err)
+	})
+
 	t.Run("Rollback after commit should not fail", func(t *testing.T) {
 		tx, err := ng.Begin(context.Background(), engine.TxOptions{
 			Writable: true,
@@ -137,6 +166,19 @@ func TestTransactionCommitRollback(t *testing.T, builder Builder) {
 		require.NoError(t, err)
 	})
 
+	t.Run("Rollback after context canceled should return context.Canceled", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(context.Background())
+		tx, err := ng.Begin(ctx, engine.TxOptions{
+			Writable: true,
+		})
+		require.NoError(t, err)
+
+		cancel()
+
+		err = tx.Rollback()
+		require.Equal(t, context.Canceled, err)
+	})
+
 	t.Run("Read-Only write attempts", func(t *testing.T) {
 		tx, err := ng.Begin(context.Background(), engine.TxOptions{
 			Writable: true,
@@ -389,10 +431,29 @@ func TestTransactionCreateStore(t *testing.T, builder Builder) {
 		err = tx.CreateStore([]byte("store"))
 		require.Equal(t, engine.ErrStoreAlreadyExists, err)
 	})
+
+	t.Run("Should fail if context canceled", func(t *testing.T) {
+		ng, cleanup := builder()
+		defer cleanup()
+		defer func() {
+			require.NoError(t, ng.Close())
+		}()
+
+		ctx, cancel := context.WithCancel(context.Background())
+		tx, err := ng.Begin(ctx, engine.TxOptions{
+			Writable: true,
+		})
+		require.NoError(t, err)
+		defer tx.Rollback()
+
+		cancel()
+		err = tx.CreateStore([]byte("store"))
+		require.Equal(t, context.Canceled, err)
+	})
 }
 
-// TestTransactionStore verifies Store behaviour.
-func TestTransactionStore(t *testing.T, builder Builder) {
+// TestTransactionGetStore verifies GetStore behaviour.
+func TestTransactionGetStore(t *testing.T, builder Builder) {
 	t.Run("Should fail if store not found", func(t *testing.T) {
 		ng, cleanup := builder()
 		defer cleanup()
@@ -451,6 +512,30 @@ func TestTransactionStore(t *testing.T, builder Builder) {
 		_, err = stb.Get([]byte("foo"))
 		require.Equal(t, engine.ErrKeyNotFound, err)
 	})
+
+	t.Run("Should fail if context canceled", func(t *testing.T) {
+		ng, cleanup := builder()
+		defer cleanup()
+		defer func() {
+			require.NoError(t, ng.Close())
+		}()
+
+		ctx, cancel := context.WithCancel(context.Background())
+		tx, err := ng.Begin(ctx, engine.TxOptions{
+			Writable: true,
+		})
+		require.NoError(t, err)
+		defer tx.Rollback()
+
+		// create two stores
+		err = tx.CreateStore([]byte("store"))
+		require.NoError(t, err)
+
+		cancel()
+
+		_, err = tx.GetStore([]byte("store"))
+		require.Equal(t, context.Canceled, err)
+	})
 }
 
 // TestTransactionDropStore verifies DropStore behaviour.
@@ -494,11 +579,39 @@ func TestTransactionDropStore(t *testing.T, builder Builder) {
 		err = tx.DropStore([]byte("store"))
 		require.Equal(t, engine.ErrStoreNotFound, err)
 	})
+
+	t.Run("Should fail if context canceled", func(t *testing.T) {
+		ng, cleanup := builder()
+		defer cleanup()
+		defer func() {
+			require.NoError(t, ng.Close())
+		}()
+
+		ctx, cancel := context.WithCancel(context.Background())
+		tx, err := ng.Begin(ctx, engine.TxOptions{
+			Writable: true,
+		})
+		require.NoError(t, err)
+		defer tx.Rollback()
+
+		// create two stores
+		err = tx.CreateStore([]byte("store"))
+		require.NoError(t, err)
+
+		cancel()
+
+		err = tx.DropStore([]byte("store"))
+		require.Equal(t, context.Canceled, err)
+	})
 }
 
 func storeBuilder(t testing.TB, builder Builder) (engine.Store, func()) {
+	return storeBuilderWithContext(context.Background(), t, builder)
+}
+
+func storeBuilderWithContext(ctx context.Context, t testing.TB, builder Builder) (engine.Store, func()) {
 	ng, cleanup := builder()
-	tx, err := ng.Begin(context.Background(), engine.TxOptions{
+	tx, err := ng.Begin(ctx, engine.TxOptions{
 		Writable: true,
 	})
 	require.NoError(t, err)
@@ -540,6 +653,31 @@ func TestStoreIterator(t *testing.T, builder Builder) {
 		})
 	})
 
+	t.Run("Should stop the iteration if context canceled", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+
+		st, cleanup := storeBuilderWithContext(ctx, t, builder)
+		defer cleanup()
+
+		for i := 1; i <= 10; i++ {
+			err := st.Put([]byte{uint8(i)}, []byte{uint8(i + 20)})
+			require.NoError(t, err)
+		}
+
+		it := st.Iterator(engine.IteratorOptions{})
+		defer it.Close()
+
+		cancel()
+
+		var i int
+		for it.Seek(nil); it.Valid(); it.Next() {
+			i++
+		}
+		require.Equal(t, context.Canceled, it.Err())
+		require.Zero(t, i)
+	})
+
 	t.Run("With no pivot, should iterate over all documents in order", func(t *testing.T) {
 		st, cleanup := storeBuilder(t, builder)
 		defer cleanup()
@@ -771,6 +909,18 @@ func TestStorePut(t *testing.T, builder Builder) {
 		err = st.Put([]byte("foo"), []byte(""))
 		require.NoError(t, err)
 	})
+
+	t.Run("Should fail if context canceled", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+
+		st, cleanup := storeBuilderWithContext(ctx, t, builder)
+		defer cleanup()
+
+		cancel()
+		err := st.Put([]byte("foo"), []byte("FOO"))
+		require.Equal(t, context.Canceled, err)
+	})
 }
 
 // TestStoreGet verifies Get behaviour.
@@ -801,6 +951,21 @@ func TestStoreGet(t *testing.T, builder Builder) {
 		require.NoError(t, err)
 		require.Equal(t, []byte("BAR"), v)
 	})
+
+	t.Run("Should fail if context canceled", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+
+		st, cleanup := storeBuilderWithContext(ctx, t, builder)
+		defer cleanup()
+
+		err := st.Put([]byte("foo"), []byte("FOO"))
+		require.NoError(t, err)
+
+		cancel()
+		_, err = st.Get([]byte("foo"))
+		require.Equal(t, context.Canceled, err)
+	})
 }
 
 // TestStoreDelete verifies Delete behaviour.
@@ -839,6 +1004,21 @@ func TestStoreDelete(t *testing.T, builder Builder) {
 		require.NoError(t, err)
 		require.Equal(t, []byte("BAR"), v)
 	})
+
+	t.Run("Should fail if context canceled", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+
+		st, cleanup := storeBuilderWithContext(ctx, t, builder)
+		defer cleanup()
+
+		err := st.Put([]byte("foo"), []byte("FOO"))
+		require.NoError(t, err)
+
+		cancel()
+		err = st.Delete([]byte("foo"))
+		require.Equal(t, context.Canceled, err)
+	})
 }
 
 // TestStoreTruncate verifies Truncate behaviour.
@@ -869,6 +1049,21 @@ func TestStoreTruncate(t *testing.T, builder Builder) {
 		require.NoError(t, it.Err())
 		require.False(t, it.Valid())
 	})
+
+	t.Run("Should fail if context canceled", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+
+		st, cleanup := storeBuilderWithContext(ctx, t, builder)
+		defer cleanup()
+
+		err := st.Put([]byte("foo"), []byte("FOO"))
+		require.NoError(t, err)
+
+		cancel()
+		err = st.Truncate()
+		require.Equal(t, context.Canceled, err)
+	})
 }
 
 // TestStoreNextSequence verifies NextSequence behaviour.
@@ -949,6 +1144,18 @@ func TestStoreNextSequence(t *testing.T, builder Builder) {
 		require.NoError(t, err)
 		require.Equal(t, s1+1, s2)
 	})
+
+	t.Run("Should fail if context canceled", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+
+		st, cleanup := storeBuilderWithContext(ctx, t, builder)
+		defer cleanup()
+
+		cancel()
+		_, err := st.NextSequence()
+		require.Equal(t, context.Canceled, err)
+	})
 }
 
 // TestQueries test simple queries against the engine.
diff --git a/sql/query/select_test.go b/sql/query/select_test.go
index 160817a6..5cfb7f35 100644
--- a/sql/query/select_test.go
+++ b/sql/query/select_test.go
@@ -3,6 +3,7 @@ package query_test
 import (
 	"bytes"
 	"database/sql"
+	"encoding/json"
 	"strconv"
 	"testing"
 
@@ -221,6 +222,24 @@ func TestSelectStmt(t *testing.T) {
 		require.NoError(t, err)
 		require.JSONEq(t, `[{"foo": true},{"foo": 1}, {"foo": 2},{"foo": "hello"}]`, buf.String())
 	})
+
+	// https://github.com/genjidb/genji/issues/208
+	t.Run("group by with arrays", func(t *testing.T) {
+		db, err := genji.Open(":memory:")
+		require.NoError(t, err)
+		defer db.Close()
+
+		err = db.Exec("CREATE TABLE test; INSERT INTO test (a) VALUES ([1, 2, 3]);")
+		require.NoError(t, err)
+
+		d, err := db.QueryDocument("SELECT MAX(a) from test GROUP BY a")
+		require.NoError(t, err)
+
+		enc, err := json.Marshal(d)
+		require.NoError(t, err)
+
+		require.JSONEq(t, `{"MAX(a)": [1, 2, 3]}`, string(enc))
+	})
 }
 
 func TestDistinct(t *testing.T) {
@@ -297,63 +316,3 @@ func TestDistinct(t *testing.T) {
 		})
 	}
 }
-
-// func TestSelectIndexedDoubleAndIntegers(t *testing.T) {
-// 	tests := []struct {
-// 		a      interface{}
-// 		query  string
-// 		isTrue bool
-// 		param  interface{}
-// 	}{
-// 		{math.MaxInt64, "SELECT true FROM test WHERE a = ?", true, math.MaxInt64},
-// 		{math.MaxInt64, "SELECT true FROM test WHERE a > ?", false, math.MaxInt64},
-// 		{math.MaxInt64, "SELECT true FROM test WHERE a > ?", true, math.MaxInt64 - 1},
-// 		{math.MaxInt64, "SELECT true FROM test WHERE a >= ?", true, math.MaxInt64},
-// 		{math.MaxInt64, "SELECT true FROM test WHERE a >= 9223372036854775807.0", true, nil},
-// 		{math.MaxInt64, "SELECT true FROM test WHERE a > 9223372036854775800.0", false, nil},
-// 		{math.MaxInt64, "SELECT true FROM test WHERE a < 9223372036854775800.0", false, nil},
-// 		{math.MaxInt64, "SELECT true FROM test WHERE a = 9223372036854775800.0", true, nil},
-// 	}
-
-// 	for _, indexed := range []bool{false, true} {
-// 		for _, test := range tests {
-// 			t.Run(fmt.Sprintf("indexed: %v / %q", indexed, test.query), func(t *testing.T) {
-// 				fmt.Println("---")
-// 				db, err := genji.Open(":memory:")
-// 				require.NoError(t, err)
-// 				defer db.Close()
-
-// 				tx, err := db.Begin(true)
-// 				require.NoError(t, err)
-// 				defer tx.Rollback()
-
-// 				err = tx.Exec("CREATE TABLE test")
-// 				require.NoError(t, err)
-
-// 				if indexed {
-// 					err = tx.Exec("CREATE INDEX idx_foo_a ON test(a)")
-// 					require.NoError(t, err)
-// 				}
-
-// 				// a = math.MaxInt64
-// 				err = tx.Exec("INSERT INTO test (a) VALUES (?)", test.a)
-// 				require.NoError(t, err)
-
-// 				var params []interface{}
-// 				if test.param != nil {
-// 					params = append(params, test.param)
-// 				}
-// 				doc, err := tx.QueryDocument(test.query, params...)
-// 				if test.isTrue {
-// 					require.NoError(t, err)
-// 					var got bool
-// 					err = document.Scan(doc, &got)
-// 					require.NoError(t, err)
-// 					require.True(t, got)
-// 				} else {
-// 					require.Equal(t, database.ErrDocumentNotFound, err)
-// 				}
-// 			})
-// 		}
-// 	}
-// }