Repository: genjidb/genji
Pull request #176 (2020-09-01 18:37:26)

Transaction statements
This PR adds support for transaction statements.

- `BEGIN [TRANSACTION] [READ ONLY | READ WRITE]`
- `ROLLBACK [TRANSACTION]`
- `COMMIT [TRANSACTION]`

When called, these transactions are attached to the database as a global transaction.


Patch:
diff --git a/database/database.go b/database/database.go
index c959a6b7..2d271a65 100644
--- a/database/database.go
+++ b/database/database.go
@@ -2,6 +2,8 @@
 package database
 
 import (
+	"errors"
+	"sync"
 	"sync/atomic"
 
 	"github.com/genjidb/genji/engine"
@@ -18,6 +20,13 @@ type Database struct {
 	// It starts at 0 at database startup and is
 	// incremented atomically every time Begin is called.
 	lastTransactionID int64
+
+	// If this is non-nil, the user is running an explicit transaction
+	// using the BEGIN statement.
+	// Only one attached transaction can be run at a time and any calls to DB.Begin()
+	// will cause an error until that transaction is rolled back or commited.
+	attachedTransaction *Transaction
+	attachedTxMu        sync.Mutex
 }
 
 // New initializes the DB using the given engine.
@@ -74,7 +83,28 @@ func (db *Database) Close() error {
 // Begin starts a new transaction.
 // The returned transaction must be closed either by calling Rollback or Commit.
 func (db *Database) Begin(writable bool) (*Transaction, error) {
-	ntx, err := db.ng.Begin(writable)
+	return db.BeginTx(&TxOptions{
+		ReadOnly: !writable,
+	})
+}
+
+// BeginTx starts a new transaction with the given options.
+// If opts is empty, it will use the default options.
+// The returned transaction must be closed either by calling Rollback or Commit.
+// If the Global option is passed, it opens a database level transaction.
+func (db *Database) BeginTx(opts *TxOptions) (*Transaction, error) {
+	if opts == nil {
+		opts = new(TxOptions)
+	}
+
+	db.attachedTxMu.Lock()
+	defer db.attachedTxMu.Unlock()
+
+	if db.attachedTransaction != nil {
+		return nil, errors.New("cannot open a transaction within a transaction")
+	}
+
+	ntx, err := db.ng.Begin(!opts.ReadOnly)
 	if err != nil {
 		return nil, err
 	}
@@ -83,7 +113,7 @@ func (db *Database) Begin(writable bool) (*Transaction, error) {
 		id:             atomic.AddInt64(&db.lastTransactionID, 1),
 		db:             db,
 		Tx:             ntx,
-		writable:       writable,
+		writable:       !opts.ReadOnly,
 		tableInfoStore: db.tableInfoStore,
 	}
 
@@ -92,5 +122,29 @@ func (db *Database) Begin(writable bool) (*Transaction, error) {
 		return nil, err
 	}
 
+	if opts.Attached {
+		db.attachedTransaction = &tx
+	}
+
 	return &tx, nil
 }
+
+// TxOptions are passed to Begin to configure transactions.
+type TxOptions struct {
+	// Open a read-only transaction.
+	ReadOnly bool
+	// Set the transaction as global at the database level.
+	// Any queries run by the database will use that transaction until it is
+	// rolled back or commited.
+	Attached bool
+}
+
+// GetAttachedTx returns the transaction attached to the database. It returns nil if there is no
+// such transaction.
+// The returned transaction is not thread safe.
+func (db *Database) GetAttachedTx() *Transaction {
+	db.attachedTxMu.Lock()
+	defer db.attachedTxMu.Unlock()
+
+	return db.attachedTransaction
+}
diff --git a/database/transaction.go b/database/transaction.go
index d721d4b0..e8d325fd 100644
--- a/database/transaction.go
+++ b/database/transaction.go
@@ -33,18 +33,45 @@ type Transaction struct {
 
 // Rollback the transaction. Can be used safely after commit.
 func (tx *Transaction) Rollback() error {
+	tx.db.attachedTxMu.Lock()
+	defer tx.db.attachedTxMu.Unlock()
+
 	if tx.writable {
 		tx.tableInfoStore.rollback(tx)
 	}
-	return tx.Tx.Rollback()
+
+	err := tx.Tx.Rollback()
+	if err != nil {
+		return err
+	}
+
+	if tx.db.attachedTransaction != nil {
+		tx.db.attachedTransaction = nil
+	}
+
+	return nil
 }
 
 // Commit the transaction.
 func (tx *Transaction) Commit() error {
+	tx.db.attachedTxMu.Lock()
+	defer tx.db.attachedTxMu.Unlock()
+
 	if tx.writable {
 		tx.tableInfoStore.commit(tx)
 	}
-	return tx.Tx.Commit()
+
+	err := tx.Tx.Commit()
+	if err != nil {
+		return err
+	}
+
+	if tx.db.attachedTransaction != nil {
+		tx.db.attachedTransaction = nil
+	}
+
+	return nil
+
 }
 
 // Writable indicates if the transaction is writable or not.
diff --git a/sql/parser/parser.go b/sql/parser/parser.go
index ad975d08..3f616643 100644
--- a/sql/parser/parser.go
+++ b/sql/parser/parser.go
@@ -64,6 +64,10 @@ func (p *Parser) ParseStatement() (query.Statement, error) {
 	switch tok {
 	case scanner.ALTER:
 		return p.parseAlterStatement()
+	case scanner.BEGIN:
+		return p.parseBeginStatement()
+	case scanner.COMMIT:
+		return p.parseCommitStatement()
 	case scanner.SELECT:
 		return p.parseSelectStatement()
 	case scanner.DELETE:
@@ -80,10 +84,12 @@ func (p *Parser) ParseStatement() (query.Statement, error) {
 		return p.parseExplainStatement()
 	case scanner.REINDEX:
 		return p.parseReIndexStatement()
+	case scanner.ROLLBACK:
+		return p.parseRollbackStatement()
 	}
 
 	return nil, newParseError(scanner.Tokstr(tok, lit), []string{
-		"SELECT", "DELETE", "UPDATE", "INSERT", "CREATE", "DROP",
+		"ALTER", "BEGIN", "COMMIT", "SELECT", "DELETE", "UPDATE", "INSERT", "CREATE", "DROP", "EXPLAIN", "REINDEX", "ROLLBACK",
 	}, pos)
 }
 
diff --git a/sql/parser/transaction.go b/sql/parser/transaction.go
new file mode 100644
index 00000000..2bd5f7bc
--- /dev/null
+++ b/sql/parser/transaction.go
@@ -0,0 +1,57 @@
+package parser
+
+import (
+	"github.com/genjidb/genji/sql/query"
+	"github.com/genjidb/genji/sql/scanner"
+)
+
+// parseBeginStatement parses a BEGIN statement.
+// This function assumes the BEGIN token has already been consumed.
+func (p *Parser) parseBeginStatement() (query.Statement, error) {
+	// parse optional TRANSCACTION token
+	if tok, _, _ := p.ScanIgnoreWhitespace(); tok != scanner.TRANSACTION {
+		p.Unscan()
+	}
+
+	// parse optional READ token
+	if tok, _, _ := p.ScanIgnoreWhitespace(); tok != scanner.READ {
+		p.Unscan()
+		return query.BeginStmt{Writable: true}, nil
+	}
+
+	// parse ONLY token
+	if tok, _, _ := p.ScanIgnoreWhitespace(); tok == scanner.ONLY {
+		return query.BeginStmt{Writable: false}, nil
+	}
+
+	p.Unscan()
+
+	// parse WRITE token
+	if tok, pos, lit := p.ScanIgnoreWhitespace(); tok != scanner.WRITE {
+		return query.BeginStmt{}, newParseError(scanner.Tokstr(tok, lit), []string{"ONLY", "WRITE"}, pos)
+	}
+
+	return query.BeginStmt{Writable: true}, nil
+}
+
+// parseRollbackStatement parses a ROLLBACK statement.
+// This function assumes the ROLLBACK token has already been consumed.
+func (p *Parser) parseRollbackStatement() (query.Statement, error) {
+	// parse optional TRANSCACTION token
+	if tok, _, _ := p.ScanIgnoreWhitespace(); tok != scanner.TRANSACTION {
+		p.Unscan()
+	}
+
+	return query.RollbackStmt{}, nil
+}
+
+// parseCommitStatement parses a COMMIT statement.
+// This function assumes the COMMIT token has already been consumed.
+func (p *Parser) parseCommitStatement() (query.Statement, error) {
+	// parse optional TRANSCACTION token
+	if tok, _, _ := p.ScanIgnoreWhitespace(); tok != scanner.TRANSACTION {
+		p.Unscan()
+	}
+
+	return query.CommitStmt{}, nil
+}
diff --git a/sql/query/query.go b/sql/query/query.go
index 204bd474..9b679097 100644
--- a/sql/query/query.go
+++ b/sql/query/query.go
@@ -16,47 +16,76 @@ var ErrResultClosed = errors.New("result already closed")
 // Results are returned as streams.
 type Query struct {
 	Statements []Statement
+	tx         *database.Transaction
+	autoCommit bool
 }
 
 // Run executes all the statements in their own transaction and returns the last result.
 func (q Query) Run(db *database.Database, args []expr.Param) (*Result, error) {
 	var res Result
-	var tx *database.Transaction
 	var err error
 
-	for _, stmt := range q.Statements {
+	q.tx = db.GetAttachedTx()
+	if q.tx == nil {
+		q.autoCommit = true
+	}
+
+	type queryAlterer interface {
+		alterQuery(db *database.Database, q *Query) error
+	}
+
+	for i, stmt := range q.Statements {
+		if qa, ok := stmt.(queryAlterer); ok {
+			err = qa.alterQuery(db, &q)
+			if err != nil {
+				if tx := db.GetAttachedTx(); tx != nil {
+					tx.Rollback()
+				}
+				return nil, err
+			}
+
+			continue
+		}
+
+		if q.tx == nil {
+			q.tx, err = db.Begin(!stmt.IsReadOnly())
+			if err != nil {
+				return nil, err
+			}
+		}
+
+		res, err = stmt.Run(q.tx, args)
+		if err != nil {
+			if q.autoCommit {
+				q.tx.Rollback()
+			}
+
+			return nil, err
+		}
+
 		// it there is an opened transaction but there are still statements
 		// to be executed, close the current transaction.
-		if tx != nil {
-			if tx.Writable() {
-				err := tx.Commit()
+		if q.tx != nil && q.autoCommit && i+1 < len(q.Statements) {
+			if q.tx.Writable() {
+				err := q.tx.Commit()
 				if err != nil {
 					return nil, err
 				}
 			} else {
-				err := tx.Rollback()
+				err := q.tx.Rollback()
 				if err != nil {
 					return nil, err
 				}
 			}
-		}
-
-		// start a new transaction for every statement
-		tx, err = db.Begin(!stmt.IsReadOnly())
-		if err != nil {
-			return nil, err
-		}
-
-		res, err = stmt.Run(tx, args)
-		if err != nil {
-			tx.Rollback()
-			return nil, err
+			q.tx = nil
 		}
 	}
 
-	// the returned result will now own the transaction.
-	// its Close method is expected to be called.
-	res.Tx = tx
+	if q.autoCommit {
+		// the returned result will now own the transaction.
+		// its Close method is expected to be called.
+		res.Tx = q.tx
+	}
 
 	return &res, nil
 }
diff --git a/sql/query/transaction.go b/sql/query/transaction.go
new file mode 100644
index 00000000..da95d62c
--- /dev/null
+++ b/sql/query/transaction.go
@@ -0,0 +1,83 @@
+package query
+
+import (
+	"errors"
+
+	"github.com/genjidb/genji/database"
+	"github.com/genjidb/genji/sql/query/expr"
+)
+
+// BeginStmt is a statement that creates a new transaction.
+type BeginStmt struct {
+	Writable bool
+}
+
+func (stmt BeginStmt) alterQuery(db *database.Database, q *Query) error {
+	if q.tx != nil {
+		return errors.New("cannot begin a transaction within a transaction")
+	}
+
+	var err error
+	q.tx, err = db.BeginTx(&database.TxOptions{
+		ReadOnly: !stmt.Writable,
+		Attached: true,
+	})
+	q.autoCommit = false
+	return err
+}
+
+func (stmt BeginStmt) IsReadOnly() bool {
+	return !stmt.Writable
+}
+
+func (stmt BeginStmt) Run(tx *database.Transaction, args []expr.Param) (Result, error) {
+	return Result{}, errors.New("cannot begin a transaction within a transaction")
+}
+
+// RollbackStmt is a statement that rollbacks the current active transaction.
+type RollbackStmt struct{}
+
+func (stmt RollbackStmt) alterQuery(db *database.Database, q *Query) error {
+	if q.tx == nil || q.autoCommit == true {
+		return errors.New("cannot rollback with no active transaction")
+	}
+
+	err := q.tx.Rollback()
+	if err != nil {
+		return err
+	}
+	q.tx = nil
+	return nil
+}
+
+func (stmt RollbackStmt) IsReadOnly() bool {
+	return false
+}
+
+func (stmt RollbackStmt) Run(tx *database.Transaction, args []expr.Param) (Result, error) {
+	return Result{}, errors.New("cannot rollback with no active transaction")
+}
+
+// CommitStmt is a statement that commits the current active transaction.
+type CommitStmt struct{}
+
+func (stmt CommitStmt) alterQuery(db *database.Database, q *Query) error {
+	if q.tx == nil || q.autoCommit == true {
+		return errors.New("cannot commit with no active transaction")
+	}
+
+	err := q.tx.Commit()
+	if err != nil {
+		return err
+	}
+	q.tx = nil
+	return nil
+}
+
+func (stmt CommitStmt) IsReadOnly() bool {
+	return false
+}
+
+func (stmt CommitStmt) Run(tx *database.Transaction, args []expr.Param) (Result, error) {
+	return Result{}, errors.New("cannot commit with no active transaction")
+}
diff --git a/sql/scanner/token.go b/sql/scanner/token.go
index 0cb6093f..b49713a9 100644
--- a/sql/scanner/token.go
+++ b/sql/scanner/token.go
@@ -76,8 +76,10 @@ const (
 	ALTER
 	AS
 	ASC
+	BEGIN
 	BY
 	CAST
+	COMMIT
 	CREATE
 	DELETE
 	DESC
@@ -94,19 +96,24 @@ const (
 	NOT
 	OFFSET
 	ON
+	ONLY
 	ORDER
 	PRIMARY
+	READ
 	REINDEX
 	RENAME
+	ROLLBACK
 	SELECT
 	SET
 	TABLE
 	TO
+	TRANSACTION
 	UNIQUE
 	UNSET
 	UPDATE
 	VALUES
 	WHERE
+	WRITE
 
 	TYPEARRAY
 	TYPEBLOB
@@ -173,40 +180,47 @@ var tokens = [...]string{
 	SEMICOLON:   ";",
 	DOT:         ".",
 
-	ALTER:   "ALTER",
-	AS:      "AS",
-	ASC:     "ASC",
-	BY:      "BY",
-	CREATE:  "CREATE",
-	CAST:    "CAST",
-	DELETE:  "DELETE",
-	DESC:    "DESC",
-	DROP:    "DROP",
-	EXISTS:  "EXISTS",
-	EXPLAIN: "EXPLAIN",
-	KEY:     "KEY",
-	FROM:    "FROM",
-	IF:      "IF",
-	INDEX:   "INDEX",
-	INSERT:  "INSERT",
-	INTO:    "INTO",
-	LIMIT:   "LIMIT",
-	NOT:     "NOT",
-	OFFSET:  "OFFSET",
-	ON:      "ON",
-	ORDER:   "ORDER",
-	PRIMARY: "PRIMARY",
-	REINDEX: "REINDEX",
-	RENAME:  "RENAME",
-	SELECT:  "SELECT",
-	SET:     "SET",
-	TABLE:   "TABLE",
-	TO:      "TO",
-	UNIQUE:  "UNIQUE",
-	UNSET:   "UNSET",
-	UPDATE:  "UPDATE",
-	VALUES:  "VALUES",
-	WHERE:   "WHERE",
+	ALTER:       "ALTER",
+	AS:          "AS",
+	ASC:         "ASC",
+	BEGIN:       "BEGIN",
+	COMMIT:      "COMMIT",
+	BY:          "BY",
+	CREATE:      "CREATE",
+	CAST:        "CAST",
+	DELETE:      "DELETE",
+	DESC:        "DESC",
+	DROP:        "DROP",
+	EXISTS:      "EXISTS",
+	EXPLAIN:     "EXPLAIN",
+	KEY:         "KEY",
+	FROM:        "FROM",
+	IF:          "IF",
+	INDEX:       "INDEX",
+	INSERT:      "INSERT",
+	INTO:        "INTO",
+	LIMIT:       "LIMIT",
+	NOT:         "NOT",
+	OFFSET:      "OFFSET",
+	ON:          "ON",
+	ONLY:        "ONLY",
+	ORDER:       "ORDER",
+	PRIMARY:     "PRIMARY",
+	READ:        "READ",
+	REINDEX:     "REINDEX",
+	RENAME:      "RENAME",
+	ROLLBACK:    "ROLLBACK",
+	SELECT:      "SELECT",
+	SET:         "SET",
+	TABLE:       "TABLE",
+	TO:          "TO",
+	TRANSACTION: "TRANSACTION",
+	UNIQUE:      "UNIQUE",
+	UNSET:       "UNSET",
+	UPDATE:      "UPDATE",
+	VALUES:      "VALUES",
+	WHERE:       "WHERE",
+	WRITE:       "WRITE",
 
 	TYPEARRAY:    "ARRAY",
 	TYPEBLOB:     "BLOB",


Test patch:
diff --git a/sql/parser/transaction_test.go b/sql/parser/transaction_test.go
new file mode 100644
index 00000000..3c247572
--- /dev/null
+++ b/sql/parser/transaction_test.go
@@ -0,0 +1,40 @@
+package parser
+
+import (
+	"testing"
+
+	"github.com/genjidb/genji/sql/query"
+	"github.com/stretchr/testify/require"
+)
+
+func TestParserTransactions(t *testing.T) {
+	tests := []struct {
+		s        string
+		expected query.Statement
+		errored  bool
+	}{
+		{"BEGIN", query.BeginStmt{Writable: true}, false},
+		{"BEGIN TRANSACTION", query.BeginStmt{Writable: true}, false},
+		{"BEGIN READ ONLY", query.BeginStmt{Writable: false}, false},
+		{"BEGIN READ WRITE", query.BeginStmt{Writable: true}, false},
+		{"BEGIN READ", query.BeginStmt{}, true},
+		{"BEGIN WRITE", query.BeginStmt{}, true},
+		{"ROLLBACK", query.RollbackStmt{}, false},
+		{"ROLLBACK TRANSACTION", query.RollbackStmt{}, false},
+		{"COMMIT", query.CommitStmt{}, false},
+		{"COMMIT TRANSACTION", query.CommitStmt{}, false},
+	}
+
+	for _, test := range tests {
+		t.Run(test.s, func(t *testing.T) {
+			q, err := ParseQuery(test.s)
+			if test.errored {
+				require.Error(t, err)
+				return
+			}
+			require.NoError(t, err)
+			require.Len(t, q.Statements, 1)
+			require.EqualValues(t, test.expected, q.Statements[0])
+		})
+	}
+}
diff --git a/sql/query/transaction_test.go b/sql/query/transaction_test.go
new file mode 100644
index 00000000..153bd569
--- /dev/null
+++ b/sql/query/transaction_test.go
@@ -0,0 +1,48 @@
+package query_test
+
+import (
+	"testing"
+
+	"github.com/genjidb/genji"
+	"github.com/stretchr/testify/require"
+)
+
+func TestTransactionRun(t *testing.T) {
+	tests := []struct {
+		name    string
+		queries []string
+		fails   bool
+	}{
+		{"Same exec/ Basic", []string{`BEGIN`}, false},
+		{"Same exec/ Nested transaction", []string{`BEGIN;BEGIN`}, true},
+		{"Same exec/ Begin then commit", []string{`BEGIN;COMMIT`}, false},
+		{"Same exec/ Begin then rollback", []string{`BEGIN;ROLLBACK`}, false},
+		{"Same exec/ Begin, select, then rollback", []string{`BEGIN;SELECT 1;ROLLBACK`}, false},
+		{"Multiple execs/ Begin then rollback", []string{`BEGIN`, `ROLLBACK`}, false},
+		{"Multiple execs/ Begin then commit", []string{`BEGIN`, `COMMIT`}, false},
+		{"Multiple execs/ Double", []string{`BEGIN`, `COMMIT`, `BEGIN`, `COMMIT`}, false},
+		{"Multiple execs/ Begin then begin", []string{`BEGIN`, `BEGIN`}, true},
+		{"Multiple execs/ Nested", []string{`BEGIN`, `BEGIN`, `COMMIT`, `COMMIT`}, true},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			db, err := genji.Open(":memory:")
+			require.NoError(t, err)
+			defer db.Close()
+			defer db.Exec("ROLLBACK")
+
+			for _, q := range test.queries {
+				err = db.Exec(q)
+				if err != nil {
+					break
+				}
+			}
+			if test.fails {
+				require.Error(t, err)
+				return
+			}
+			require.NoError(t, err)
+		})
+	}
+}
diff --git a/sql/scanner/scanner_test.go b/sql/scanner/scanner_test.go
index 1de3d983..b86c1cb7 100644
--- a/sql/scanner/scanner_test.go
+++ b/sql/scanner/scanner_test.go
@@ -131,7 +131,9 @@ func TestScanner_Scan(t *testing.T) {
 		{s: `AS`, tok: scanner.AS, raw: `AS`},
 		{s: `ASC`, tok: scanner.ASC, raw: `ASC`},
 		{s: `BY`, tok: scanner.BY, raw: `BY`},
+		{s: `BEGIN`, tok: scanner.BEGIN, raw: `BEGIN`},
 		{s: `CAST`, tok: scanner.CAST, raw: `CAST`},
+		{s: `COMMIT`, tok: scanner.COMMIT, raw: `COMMIT`},
 		{s: `CREATE`, tok: scanner.CREATE, raw: `CREATE`},
 		{s: `EXPLAIN`, tok: scanner.EXPLAIN, raw: `EXPLAIN`},
 		{s: `DELETE`, tok: scanner.DELETE, raw: `DELETE`},
@@ -141,18 +143,24 @@ func TestScanner_Scan(t *testing.T) {
 		{s: `INSERT`, tok: scanner.INSERT, raw: `INSERT`},
 		{s: `INTO`, tok: scanner.INTO, raw: `INTO`},
 		{s: `LIMIT`, tok: scanner.LIMIT, raw: `LIMIT`},
+		{s: `ONLY`, tok: scanner.ONLY, raw: `ONLY`},
 		{s: `OFFSET`, tok: scanner.OFFSET, raw: `OFFSET`},
 		{s: `ORDER`, tok: scanner.ORDER, raw: `ORDER`},
 		{s: `PRIMARY`, tok: scanner.PRIMARY, raw: `PRIMARY`},
+		{s: `READ`, tok: scanner.READ, raw: `READ`},
 		{s: `REINDEX`, tok: scanner.REINDEX, raw: `REINDEX`},
 		{s: `RENAME`, tok: scanner.RENAME, raw: `RENAME`},
+		{s: `ROLLBACK`, tok: scanner.ROLLBACK, raw: `ROLLBACK`},
 		{s: `SELECT`, tok: scanner.SELECT, raw: `SELECT`},
 		{s: `SET`, tok: scanner.SET, raw: `SET`},
+		{s: `TABLE`, tok: scanner.TABLE, raw: `TABLE`},
 		{s: `TO`, tok: scanner.TO, raw: `TO`},
+		{s: `TRANSACTION`, tok: scanner.TRANSACTION, raw: `TRANSACTION`},
 		{s: `UPDATE`, tok: scanner.UPDATE, raw: `UPDATE`},
 		{s: `UNSET`, tok: scanner.UNSET, raw: `UNSET`},
 		{s: `VALUES`, tok: scanner.VALUES, raw: `VALUES`},
 		{s: `WHERE`, tok: scanner.WHERE, raw: `WHERE`},
+		{s: `WRITE`, tok: scanner.WRITE, raw: `WRITE`},
 		{s: `seLECT`, tok: scanner.SELECT, raw: `seLECT`}, // case insensitive
 
 		// types