Repository: src-d/go-mysql-server
Pull request #832 (2019-10-01 01:12:27)

Implemented UPDATE
Notice: one of the tests will fail until https://github.com/src-d/go-mysql-server/pull/831 is merged in, as it is required to fix one of the bugs that that test catches.

This is an implementation of UPDATE, which brings with it a new interface. Unlike the other interfaces, this one supplies two rows, the original row, and the updated row. This decision was made so that implementations that rely on primary keys can have access to the old primary key if need be.

~Besides the big change, I've also caught a bug where supplying a non-existent column to ORDER BY will cause the query to still succeed, while ignoring the ORDER BY clause. Testing against MySQL showed that an error is the proper behavior, and this has been corrected. Since the change was small, I added it here rather than making an entirely-new PR for it.~ This was fixed by https://github.com/src-d/go-mysql-server/pull/818.

ajnavarro: @Hydrocharged #831 merged! you can rebase when you wish. Thanks!
Hydrocharged: @ajnavarro It has been rebased!
agarciamontoro: Sorry for the delay on this one, @Hydrocharged. When rebasing, build will fail because of https://github.com/src-d/go-mysql-server/pull/835. You would probably just need to pass a `*sql.Context` wherever it is needed.

Patch:
diff --git a/engine.go b/engine.go
index 4cb90c3..6be2be3 100644
--- a/engine.go
+++ b/engine.go
@@ -120,7 +120,7 @@ func (e *Engine) Query(
 	case *plan.CreateIndex:
 		typ = sql.CreateIndexProcess
 		perm = auth.ReadPerm | auth.WritePerm
-	case *plan.InsertInto, *plan.DeleteFrom, *plan.DropIndex, *plan.UnlockTables, *plan.LockTables:
+	case *plan.InsertInto, *plan.DeleteFrom, *plan.Update, *plan.DropIndex, *plan.UnlockTables, *plan.LockTables:
 		perm = auth.ReadPerm | auth.WritePerm
 	}
 
diff --git a/memory/table.go b/memory/table.go
index a73eced..0889563 100644
--- a/memory/table.go
+++ b/memory/table.go
@@ -290,6 +290,37 @@ func (t *Table) Delete(ctx *sql.Context, row sql.Row) error {
 	return nil
 }
 
+func (t *Table) Update(ctx *sql.Context, oldRow sql.Row, newRow sql.Row) error {
+	if err := checkRow(t.schema, oldRow); err != nil {
+		return err
+	}
+	if err := checkRow(t.schema, newRow); err != nil {
+		return err
+	}
+
+	matches := false
+	for partitionIndex, partition := range t.partitions {
+		for partitionRowIndex, partitionRow := range partition {
+			matches = true
+			for rIndex, val := range oldRow {
+				if val != partitionRow[rIndex] {
+					matches = false
+					break
+				}
+			}
+			if matches {
+				t.partitions[partitionIndex][partitionRowIndex] = newRow
+				break
+			}
+		}
+		if matches {
+			break
+		}
+	}
+
+	return nil
+}
+
 func checkRow(schema sql.Schema, row sql.Row) error {
 	if len(row) != len(schema) {
 		return sql.ErrUnexpectedRowLength.New(len(schema), len(row))
diff --git a/sql/analyzer/pushdown.go b/sql/analyzer/pushdown.go
index 34312b1..07049c4 100644
--- a/sql/analyzer/pushdown.go
+++ b/sql/analyzer/pushdown.go
@@ -20,7 +20,7 @@ func pushdown(ctx *sql.Context, a *Analyzer, n sql.Node) (sql.Node, error) {
 
 	// don't do pushdown on certain queries
 	switch n.(type) {
-	case *plan.InsertInto, *plan.DeleteFrom, *plan.CreateIndex:
+	case *plan.InsertInto, *plan.DeleteFrom, *plan.Update, *plan.CreateIndex:
 		return n, nil
 	}
 
diff --git a/sql/core.go b/sql/core.go
index e91ce0b..7a88750 100644
--- a/sql/core.go
+++ b/sql/core.go
@@ -217,6 +217,12 @@ type Replacer interface {
 	Inserter
 }
 
+// Updater allows rows to be updated.
+type Updater interface {
+	// Update the given row. Provides both the old and new rows.
+	Update(ctx *Context, old Row, new Row) error
+}
+
 // Database represents the database.
 type Database interface {
 	Nameable
diff --git a/sql/expression/set.go b/sql/expression/set.go
new file mode 100644
index 0000000..d18bde3
--- /dev/null
+++ b/sql/expression/set.go
@@ -0,0 +1,61 @@
+package expression
+
+import (
+	"fmt"
+	"github.com/src-d/go-mysql-server/sql"
+	"gopkg.in/src-d/go-errors.v1"
+)
+
+var errCannotSetField = errors.NewKind("Expected GetField expression on left but got %T")
+
+// SetField updates the value of a field from a row.
+type SetField struct {
+	BinaryExpression
+}
+
+// NewSetField creates a new SetField expression.
+func NewSetField(colName, expr sql.Expression) sql.Expression {
+	return &SetField{BinaryExpression{Left: colName, Right: expr}}
+}
+
+func (s *SetField) String() string {
+	return fmt.Sprintf("SETFIELD %s = %s", s.Left, s.Right)
+}
+
+// Type implements the Expression interface.
+func (s *SetField) Type() sql.Type {
+	return s.Left.Type()
+}
+
+// Eval implements the Expression interface.
+// Returns a copy of the given row with an updated value.
+func (s *SetField) Eval(ctx *sql.Context, row sql.Row) (interface{}, error) {
+	getField, ok := s.Left.(*GetField)
+	if !ok {
+		return nil, errCannotSetField.New(s.Left)
+	}
+	if getField.fieldIndex < 0 || getField.fieldIndex >= len(row) {
+		return nil, ErrIndexOutOfBounds.New(getField.fieldIndex, len(row))
+	}
+	val, err := s.Right.Eval(ctx, row)
+	if err != nil {
+		return nil, err
+	}
+	if val != nil {
+		val, err = getField.fieldType.Convert(val)
+		if err != nil {
+			return nil, err
+		}
+	}
+	updatedRow := row.Copy()
+	updatedRow[getField.fieldIndex] = val
+	return updatedRow, nil
+}
+
+// WithChildren implements the Expression interface.
+func (s *SetField) WithChildren(children ...sql.Expression) (sql.Expression, error) {
+	if len(children) != 2 {
+		return nil, sql.ErrInvalidChildrenNumber.New(s, len(children), 2)
+	}
+	return NewSetField(children[0], children[1]), nil
+}
\ No newline at end of file
diff --git a/sql/parse/parse.go b/sql/parse/parse.go
index e77819f..5b6b9ca 100644
--- a/sql/parse/parse.go
+++ b/sql/parse/parse.go
@@ -153,6 +153,8 @@ func convert(ctx *sql.Context, stmt sqlparser.Statement, query string) (sql.Node
 		return plan.NewRollback(), nil
 	case *sqlparser.Delete:
 		return convertDelete(ctx, n)
+	case *sqlparser.Update:
+		return convertUpdate(ctx, n)
 	}
 }
 
@@ -429,6 +431,49 @@ func convertDelete(ctx *sql.Context, d *sqlparser.Delete) (sql.Node, error) {
 	return plan.NewDeleteFrom(node), nil
 }
 
+func convertUpdate(ctx *sql.Context, d *sqlparser.Update) (sql.Node, error) {
+	node, err := tableExprsToTable(ctx, d.TableExprs)
+	if err != nil {
+		return nil, err
+	}
+
+	updateExprs, err := updateExprsToExpressions(ctx, d.Exprs)
+	if err != nil {
+		return nil, err
+	}
+
+	if d.Where != nil {
+		node, err = whereToFilter(ctx, d.Where, node)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	if len(d.OrderBy) != 0 {
+		node, err = orderByToSort(ctx, d.OrderBy, node)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	// Limit must wrap offset, and not vice-versa, so that skipped rows don't count toward the returned row count.
+	if d.Limit != nil && d.Limit.Offset != nil {
+		node, err = offsetToOffset(ctx, d.Limit.Offset, node)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	if d.Limit != nil {
+		node, err = limitToLimit(ctx, d.Limit.Rowcount, node)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	return plan.NewUpdate(node, updateExprs), nil
+}
+
 func columnDefinitionToSchema(colDef []*sqlparser.ColumnDefinition) (sql.Schema, error) {
 	var schema sql.Schema
 	for _, cd := range colDef {
@@ -1241,6 +1286,22 @@ func intervalExprToExpression(ctx *sql.Context, e *sqlparser.IntervalExpr) (sql.
 	return expression.NewInterval(expr, e.Unit), nil
 }
 
+func updateExprsToExpressions(ctx *sql.Context, e sqlparser.UpdateExprs) ([]sql.Expression, error) {
+	res := make([]sql.Expression, len(e))
+	for i, updateExpr := range e {
+		colName, err := exprToExpression(ctx, updateExpr.Name)
+		if err != nil {
+			return nil, err
+		}
+		innerExpr, err := exprToExpression(ctx, updateExpr.Expr)
+		if err != nil {
+			return nil, err
+		}
+		res[i] = expression.NewSetField(colName, innerExpr)
+	}
+	return res, nil
+}
+
 func removeComments(s string) string {
 	r := bufio.NewReader(strings.NewReader(s))
 	var result []rune
diff --git a/sql/plan/update.go b/sql/plan/update.go
new file mode 100644
index 0000000..a29fb73
--- /dev/null
+++ b/sql/plan/update.go
@@ -0,0 +1,189 @@
+package plan
+
+import (
+	"github.com/src-d/go-mysql-server/sql"
+	"gopkg.in/src-d/go-errors.v1"
+	"io"
+)
+
+var ErrUpdateNotSupported = errors.NewKind("table doesn't support UPDATE")
+var ErrUpdateUnexpectedSetResult = errors.NewKind("attempted to set field but expression returned %T")
+
+// Update is a node for updating rows on tables.
+type Update struct {
+	sql.Node
+	UpdateExprs []sql.Expression
+}
+
+// NewUpdate creates an Update node.
+func NewUpdate(n sql.Node, updateExprs []sql.Expression) *Update {
+	return &Update{n, updateExprs}
+}
+
+// Expressions implements the Expressioner interface.
+func (p *Update) Expressions() []sql.Expression {
+	return p.UpdateExprs
+}
+
+// Schema implements the Node interface.
+func (p *Update) Schema() sql.Schema {
+	return sql.Schema{
+		{
+			Name:     "matched",
+			Type:     sql.Int64,
+			Default:  int64(0),
+			Nullable: false,
+		},
+		{
+			Name:     "updated",
+			Type:     sql.Int64,
+			Default:  int64(0),
+			Nullable: false,
+		},
+	}
+}
+
+// Resolved implements the Resolvable interface.
+func (p *Update) Resolved() bool {
+	if !p.Node.Resolved() {
+		return false
+	}
+	for _, updateExpr := range p.UpdateExprs {
+		if !updateExpr.Resolved() {
+			return false
+		}
+	}
+	return true
+}
+
+func (p *Update) Children() []sql.Node {
+	return []sql.Node{p.Node}
+}
+
+func getUpdatable(node sql.Node) (sql.Updater, error) {
+	switch node := node.(type) {
+	case sql.Updater:
+		return node, nil
+	case *ResolvedTable:
+		return getUpdatableTable(node.Table)
+	}
+	for _, child := range node.Children() {
+		updater, _ := getUpdatable(child)
+		if updater != nil {
+			return updater, nil
+		}
+	}
+	return nil, ErrUpdateNotSupported.New()
+}
+
+func getUpdatableTable(t sql.Table) (sql.Updater, error) {
+	switch t := t.(type) {
+	case sql.Updater:
+		return t, nil
+	case sql.TableWrapper:
+		return getUpdatableTable(t.Underlying())
+	default:
+		return nil, ErrUpdateNotSupported.New()
+	}
+}
+
+// Execute inserts the rows in the database.
+func (p *Update) Execute(ctx *sql.Context) (int, int, error) {
+	updatable, err := getUpdatable(p.Node)
+	if err != nil {
+		return 0, 0, err
+	}
+	schema := p.Node.Schema()
+
+	iter, err := p.Node.RowIter(ctx)
+	if err != nil {
+		return 0, 0, err
+	}
+
+	rowsMatched := 0
+	rowsUpdated := 0
+	for {
+		oldRow, err := iter.Next()
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			_ = iter.Close()
+			return rowsMatched, rowsUpdated, err
+		}
+		rowsMatched++
+
+		newRow, err := p.applyUpdates(ctx, oldRow)
+		if err != nil {
+			_ = iter.Close()
+			return rowsMatched, rowsUpdated, err
+		}
+		if equals, err := oldRow.Equals(newRow, schema); err == nil {
+			if !equals {
+				err = updatable.Update(ctx, oldRow, newRow)
+				if err != nil {
+					_ = iter.Close()
+					return rowsMatched, rowsUpdated, err
+				}
+				rowsUpdated++
+			}
+		} else {
+			_ = iter.Close()
+			return rowsMatched, rowsUpdated, err
+		}
+	}
+
+	return rowsMatched, rowsUpdated, nil
+}
+
+// RowIter implements the Node interface.
+func (p *Update) RowIter(ctx *sql.Context) (sql.RowIter, error) {
+	matched, updated, err := p.Execute(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	return sql.RowsToRowIter(sql.NewRow(int64(matched), int64(updated))), nil
+}
+
+// WithChildren implements the Node interface.
+func (p *Update) WithChildren(children ...sql.Node) (sql.Node, error) {
+	if len(children) != 1 {
+		return nil, sql.ErrInvalidChildrenNumber.New(p, len(children), 1)
+	}
+	return NewUpdate(children[0], p.UpdateExprs), nil
+}
+
+// WithExpressions implements the Expressioner interface.
+func (p *Update) WithExpressions(newExprs ...sql.Expression) (sql.Node, error) {
+	if len(newExprs) != len(p.UpdateExprs) {
+		return nil, sql.ErrInvalidChildrenNumber.New(p, len(p.UpdateExprs), 1)
+	}
+	return NewUpdate(p.Node, newExprs), nil
+}
+
+func (p Update) String() string {
+	pr := sql.NewTreePrinter()
+	_ = pr.WriteNode("Update")
+	_ = pr.WriteChildren(p.Node.String())
+	for _, updateExpr := range p.UpdateExprs {
+		_ = pr.WriteChildren(updateExpr.String())
+	}
+	return pr.String()
+}
+
+func (p *Update) applyUpdates(ctx *sql.Context, row sql.Row) (sql.Row, error) {
+	var ok bool
+	prev := row
+	for _, updateExpr := range p.UpdateExprs {
+		val, err := updateExpr.Eval(ctx, prev)
+		if err != nil {
+			return nil, err
+		}
+		prev, ok = val.(sql.Row)
+		if !ok {
+			return nil, ErrUpdateUnexpectedSetResult.New(val)
+		}
+	}
+	return prev, nil
+}


Test patch:
diff --git a/engine_test.go b/engine_test.go
index 536e020..e1274e2 100644
--- a/engine_test.go
+++ b/engine_test.go
@@ -18,6 +18,7 @@ import (
 	"github.com/src-d/go-mysql-server/sql/plan"
 	"github.com/src-d/go-mysql-server/test"
 
+
 	"github.com/stretchr/testify/require"
 )
 
@@ -2245,6 +2246,142 @@ func TestReplaceIntoErrors(t *testing.T) {
 	}
 }
 
+func TestUpdate(t *testing.T) {
+	var updates = []struct {
+		updateQuery    string
+		expectedUpdate []sql.Row
+		selectQuery     string
+		expectedSelect  []sql.Row
+	}{
+		{
+			"UPDATE mytable SET s = 'updated';",
+			[]sql.Row{{int64(3), int64(3)}},
+			"SELECT * FROM mytable;",
+			[]sql.Row{{int64(1), "updated"}, {int64(2), "updated"}, {int64(3), "updated"}},
+		},
+		{
+			"UPDATE mytable SET s = 'updated' WHERE i > 9999;",
+			[]sql.Row{{int64(0), int64(0)}},
+			"SELECT * FROM mytable;",
+			[]sql.Row{{int64(1), "first row"}, {int64(2), "second row"}, {int64(3), "third row"}},
+		},
+		{
+			"UPDATE mytable SET s = 'updated' WHERE i = 1;",
+			[]sql.Row{{int64(1), int64(1)}},
+			"SELECT * FROM mytable;",
+			[]sql.Row{{int64(1), "updated"}, {int64(2), "second row"}, {int64(3), "third row"}},
+		},
+		{
+			"UPDATE mytable SET s = 'updated' WHERE i <> 9999;",
+			[]sql.Row{{int64(3), int64(3)}},
+			"SELECT * FROM mytable;",
+			[]sql.Row{{int64(1), "updated"},{int64(2), "updated"},{int64(3), "updated"}},
+		},
+		{
+			"UPDATE floattable SET f32 = f32 + f32, f64 = f32 * f64 WHERE i = 2;",
+			[]sql.Row{{int64(1), int64(1)}},
+			"SELECT * FROM floattable WHERE i = 2;",
+			[]sql.Row{{int64(2), float32(3.0), float64(4.5)}},
+		},
+		{
+			"UPDATE floattable SET f32 = 5, f32 = 4 WHERE i = 1;",
+			[]sql.Row{{int64(1), int64(1)}},
+			"SELECT f32 FROM floattable WHERE i = 1;",
+			[]sql.Row{{float32(4.0)}},
+		},
+		{
+			"UPDATE mytable SET s = 'first row' WHERE i = 1;",
+			[]sql.Row{{int64(1), int64(0)}},
+			"SELECT * FROM mytable;",
+			[]sql.Row{{int64(1), "first row"}, {int64(2), "second row"}, {int64(3), "third row"}},
+		},
+		{
+			"UPDATE niltable SET b = NULL WHERE f IS NULL;",
+			[]sql.Row{{int64(2), int64(1)}},
+			"SELECT * FROM niltable WHERE f IS NULL;",
+			[]sql.Row{{int64(4), nil, nil}, {nil, nil, nil}},
+		},
+		{
+			"UPDATE mytable SET s = 'updated' ORDER BY i ASC LIMIT 2;",
+			[]sql.Row{{int64(2), int64(2)}},
+			"SELECT * FROM mytable;",
+			[]sql.Row{{int64(1), "updated"}, {int64(2), "updated"}, {int64(3), "third row"}},
+		},
+		{
+			"UPDATE mytable SET s = 'updated' ORDER BY i DESC LIMIT 2;",
+			[]sql.Row{{int64(2), int64(2)}},
+			"SELECT * FROM mytable;",
+			[]sql.Row{{int64(1), "first row"}, {int64(2), "updated"}, {int64(3), "updated"}},
+		},
+		{
+			"UPDATE mytable SET s = 'updated' ORDER BY i LIMIT 1 OFFSET 1;",
+			[]sql.Row{{int64(1), int64(1)}},
+			"SELECT * FROM mytable;",
+			[]sql.Row{{int64(1), "first row"}, {int64(2), "updated"}, {int64(3), "third row"}},
+		},
+		{
+			"UPDATE mytable SET s = 'updated';",
+			[]sql.Row{{int64(3), int64(3)}},
+			"SELECT * FROM mytable;",
+			[]sql.Row{{int64(1), "updated"}, {int64(2), "updated"}, {int64(3), "updated"}},
+		},
+	}
+
+	for _, update := range updates {
+		e := newEngine(t)
+		ctx := newCtx()
+		testQueryWithContext(ctx, t, e, update.updateQuery, update.expectedUpdate)
+		testQueryWithContext(ctx, t, e, update.selectQuery, update.expectedSelect)
+	}
+}
+
+func TestUpdateErrors(t *testing.T) {
+	var expectedFailures = []struct {
+		name  string
+		query string
+	}{
+		{
+			"invalid table",
+			"UPDATE doesnotexist SET i = 0;",
+		},
+		{
+			"invalid column set",
+			"UPDATE mytable SET z = 0;",
+		},
+		{
+			"invalid column set value",
+			"UPDATE mytable SET i = z;",
+		},
+		{
+			"invalid column where",
+			"UPDATE mytable SET s = 'hi' WHERE z = 1;",
+		},
+		{
+			"invalid column order by",
+			"UPDATE mytable SET s = 'hi' ORDER BY z;",
+		},
+		{
+			"negative limit",
+			"UPDATE mytable SET s = 'hi' LIMIT -1;",
+		},
+		{
+			"negative offset",
+			"UPDATE mytable SET s = 'hi' LIMIT 1 OFFSET -1;",
+		},
+		{
+			"set null on non-nullable",
+			"UPDATE mytable SET s = NULL;",
+		},
+	}
+
+	for _, expectedFailure := range expectedFailures {
+		t.Run(expectedFailure.name, func(t *testing.T) {
+			_, _, err := newEngine(t).Query(newCtx(), expectedFailure.query)
+			require.Error(t, err)
+		})
+	}
+}
+
 const testNumPartitions = 5
 
 func TestAmbiguousColumnResolution(t *testing.T) {
@@ -2670,12 +2807,12 @@ func newEngineWithParallelism(t *testing.T, parallelism int) *sqle.Engine {
 
 	insertRows(
 		t, floatTable,
-		sql.NewRow(1, float32(1.0), float64(1.0)),
-		sql.NewRow(2, float32(1.5), float64(1.5)),
-		sql.NewRow(3, float32(2.0), float64(2.0)),
-		sql.NewRow(4, float32(2.5), float64(2.5)),
-		sql.NewRow(-1, float32(-1.0), float64(-1.0)),
-		sql.NewRow(-2, float32(-1.5), float64(-1.5)),
+		sql.NewRow(int64(1), float32(1.0), float64(1.0)),
+		sql.NewRow(int64(2), float32(1.5), float64(1.5)),
+		sql.NewRow(int64(3), float32(2.0), float64(2.0)),
+		sql.NewRow(int64(4), float32(2.5), float64(2.5)),
+		sql.NewRow(int64(-1), float32(-1.0), float64(-1.0)),
+		sql.NewRow(int64(-2), float32(-1.5), float64(-1.5)),
 	)
 
 	nilTable := memory.NewPartitionedTable("niltable", sql.Schema{