Repository: tomarrell/lbadd (go)

Issue:
compiler: support UPDATE statement

Pull request:
compiler: add support for UPDATE statements
Add support for `UPDATE` statements without `WITH` clauses.

Closes #138 

Patch:
diff --git a/Makefile b/Makefile
index 9650a8d..91e8eea 100644
--- a/Makefile
+++ b/Makefile
@@ -14,7 +14,7 @@ lint: ## Runs the linters (including internal ones)
 	# internal analysis tools
 	go run ./internal/tool/analysis ./...;
 	# external analysis tools
-	golint ./...;
+	golint -set_exit_status ./...;
 	errcheck ./...;
 	gosec -quiet ./...;
 	staticcheck ./...;
diff --git a/internal/compiler/command/command.go b/internal/compiler/command/command.go
index 3239f3d..1c5ab5c 100644
--- a/internal/compiler/command/command.go
+++ b/internal/compiler/command/command.go
@@ -41,6 +41,22 @@ const (
 	JoinCross
 )
 
+//go:generate stringer -type=UpdateOr
+
+// UpdateOr is the type of update alternative that is specified in an update
+// statement.
+type UpdateOr uint8
+
+// Known UpdateOrs
+const (
+	UpdateOrUnknown UpdateOr = iota
+	UpdateOrRollback
+	UpdateOrAbort
+	UpdateOrReplace
+	UpdateOrFail
+	UpdateOrIgnore
+)
+
 type (
 	// Explain instructs the executor to explain the nested command instead of
 	// executing it.
@@ -52,7 +68,7 @@ type (
 	// List is a marker interface that facilitates creating a type hierarchy for
 	// the command model.
 	List interface {
-		fmt.Stringer
+		Command
 		_list()
 	}
 
@@ -151,6 +167,31 @@ type (
 	// schema defined in this command.
 	DropTrigger drop
 
+	// Update instructs the executor to update all datasets, for which the
+	// filter expression evaluates to true, with the defined updates.
+	Update struct {
+		// UpdateOr is the OR clause in an update statement.
+		UpdateOr UpdateOr
+		// Table is the table on which the update should be performed.
+		Table Table
+		// Updates is a list of updates that must be applied.
+		Updates []UpdateSetter
+		// Filter is the filter expression, that determines, which datasets are
+		// to be updated.
+		Filter Expr
+	}
+
+	// UpdateSetter is an update that can be applied to a value in a dataset.
+	UpdateSetter struct {
+		// Cols are the columns of the dataset, that have to be updated. Because
+		// the columns must be inside the table that this update refers to, and
+		// no alias can be specified according to the grammar, this is just a
+		// string, and not a full blown column.
+		Cols []string
+		// Value is the expression that the columns have to be set to.
+		Value Expr
+	}
+
 	// Column represents a database table column.
 	Column struct {
 		// Table is the table name that this column belongs to. May be empty, as
@@ -293,6 +334,18 @@ func (d DropView) String() string {
 	return fmt.Sprintf("DropView[view=%v,ifexists=%v]()", view, d.IfExists)
 }
 
+func (u Update) String() string {
+	var sets []string
+	for _, set := range u.Updates {
+		sets = append(sets, set.String())
+	}
+	return fmt.Sprintf("Update[or=%v,table=%v,sets=(%v),filter=%v]", u.UpdateOr, u.Table, strings.Join(sets, ","), u.Filter)
+}
+
+func (u UpdateSetter) String() string {
+	return fmt.Sprintf("%v=%v", strings.Join(u.Cols, ","), u.Value)
+}
+
 func (c Column) String() string {
 	if c.Alias == "" {
 		return c.Column.String()
diff --git a/internal/compiler/command/updateor_string.go b/internal/compiler/command/updateor_string.go
new file mode 100644
index 0000000..d1ebb99
--- /dev/null
+++ b/internal/compiler/command/updateor_string.go
@@ -0,0 +1,28 @@
+// Code generated by "stringer -type=UpdateOr"; DO NOT EDIT.
+
+package command
+
+import "strconv"
+
+func _() {
+	// An "invalid array index" compiler error signifies that the constant values have changed.
+	// Re-run the stringer command to generate them again.
+	var x [1]struct{}
+	_ = x[UpdateOrUnknown-0]
+	_ = x[UpdateOrRollback-1]
+	_ = x[UpdateOrAbort-2]
+	_ = x[UpdateOrReplace-3]
+	_ = x[UpdateOrFail-4]
+	_ = x[UpdateOrIgnore-5]
+}
+
+const _UpdateOr_name = "UpdateOrUnknownUpdateOrRollbackUpdateOrAbortUpdateOrReplaceUpdateOrFailUpdateOrIgnore"
+
+var _UpdateOr_index = [...]uint8{0, 15, 31, 44, 59, 71, 85}
+
+func (i UpdateOr) String() string {
+	if i >= UpdateOr(len(_UpdateOr_index)-1) {
+		return "UpdateOr(" + strconv.FormatInt(int64(i), 10) + ")"
+	}
+	return _UpdateOr_name[_UpdateOr_index[i]:_UpdateOr_index[i+1]]
+}
diff --git a/internal/compiler/simple_compiler.go b/internal/compiler/simple_compiler.go
index 94a348f..972fd92 100644
--- a/internal/compiler/simple_compiler.go
+++ b/internal/compiler/simple_compiler.go
@@ -2,6 +2,7 @@ package compiler
 
 import (
 	"fmt"
+	"strings"
 
 	"github.com/tomarrell/lbadd/internal/compiler/command"
 	"github.com/tomarrell/lbadd/internal/compiler/optimization"
@@ -87,11 +88,84 @@ func (c *simpleCompiler) compileInternal(ast *ast.SQLStmt) (command.Command, err
 			return nil, fmt.Errorf("drop view: %w", err)
 		}
 		return cmd, nil
+	case ast.UpdateStmt != nil:
+		cmd, err := c.compileUpdate(ast.UpdateStmt)
+		if err != nil {
+			return nil, fmt.Errorf("update: %w", err)
+		}
+		return cmd, nil
 	}
 	return nil, fmt.Errorf("statement type: %w", ErrUnsupported)
 }
 
-func (c *simpleCompiler) compileDropTable(stmt *ast.DropTableStmt) (command.Command, error) {
+func (c *simpleCompiler) compileUpdate(stmt *ast.UpdateStmt) (command.Update, error) {
+	updateOr := command.UpdateOrIgnore // ignore as default or
+	switch {
+	case stmt.Rollback != nil:
+		updateOr = command.UpdateOrRollback
+	case stmt.Abort != nil:
+		updateOr = command.UpdateOrAbort
+	case stmt.Replace != nil:
+		updateOr = command.UpdateOrReplace
+	case stmt.Fail != nil:
+		updateOr = command.UpdateOrFail
+	case stmt.Ignore != nil:
+		updateOr = command.UpdateOrIgnore
+	}
+
+	qtn, err := c.compileQualifiedTableName(stmt.QualifiedTableName)
+	if err != nil {
+		return command.Update{}, fmt.Errorf("qualified table name: %w", err)
+	}
+
+	var sets []command.UpdateSetter
+	for _, set := range stmt.UpdateSetter {
+		compiledSet, err := c.compileUpdateSetter(set)
+		if err != nil {
+			return command.Update{}, fmt.Errorf("update setter: %w", err)
+		}
+		sets = append(sets, compiledSet)
+	}
+
+	var filter command.Expr
+	if stmt.Where != nil {
+		filterExpr, err := c.compileExpr(stmt.Expr)
+		if err != nil {
+			return command.Update{}, fmt.Errorf("expr: %w", err)
+		}
+		filter = filterExpr
+	} else {
+		filter = command.ConstantBooleanExpr{Value: true}
+	}
+
+	return command.Update{
+		UpdateOr: updateOr,
+		Table:    qtn,
+		Updates:  sets,
+		Filter:   filter,
+	}, nil
+}
+
+func (c *simpleCompiler) compileUpdateSetter(setter *ast.UpdateSetter) (command.UpdateSetter, error) {
+	var cols []string
+	if setter.ColumnName != nil {
+		cols = append(cols, setter.ColumnName.Value())
+	} else {
+		for _, col := range setter.ColumnNameList.ColumnName {
+			cols = append(cols, col.Value())
+		}
+	}
+	expr, err := c.compileExpr(setter.Expr)
+	if err != nil {
+		return command.UpdateSetter{}, fmt.Errorf("expr: %w", err)
+	}
+	return command.UpdateSetter{
+		Cols:  cols,
+		Value: expr,
+	}, nil
+}
+
+func (c *simpleCompiler) compileDropTable(stmt *ast.DropTableStmt) (command.DropTable, error) {
 	cmd := command.DropTable{
 		IfExists: stmt.If != nil,
 		Name:     stmt.TableName.Value(),
@@ -102,7 +176,7 @@ func (c *simpleCompiler) compileDropTable(stmt *ast.DropTableStmt) (command.Comm
 	return cmd, nil
 }
 
-func (c *simpleCompiler) compileDropIndex(stmt *ast.DropIndexStmt) (command.Command, error) {
+func (c *simpleCompiler) compileDropIndex(stmt *ast.DropIndexStmt) (command.DropIndex, error) {
 	cmd := command.DropIndex{
 		IfExists: stmt.If != nil,
 		Name:     stmt.IndexName.Value(),
@@ -113,7 +187,7 @@ func (c *simpleCompiler) compileDropIndex(stmt *ast.DropIndexStmt) (command.Comm
 	return cmd, nil
 }
 
-func (c *simpleCompiler) compileDropTrigger(stmt *ast.DropTriggerStmt) (command.Command, error) {
+func (c *simpleCompiler) compileDropTrigger(stmt *ast.DropTriggerStmt) (command.DropTrigger, error) {
 	cmd := command.DropTrigger{
 		IfExists: stmt.If != nil,
 		Name:     stmt.TriggerName.Value(),
@@ -124,7 +198,7 @@ func (c *simpleCompiler) compileDropTrigger(stmt *ast.DropTriggerStmt) (command.
 	return cmd, nil
 }
 
-func (c *simpleCompiler) compileDropView(stmt *ast.DropViewStmt) (command.Command, error) {
+func (c *simpleCompiler) compileDropView(stmt *ast.DropViewStmt) (command.DropView, error) {
 	cmd := command.DropView{
 		IfExists: stmt.If != nil,
 		Name:     stmt.ViewName.Value(),
@@ -135,16 +209,16 @@ func (c *simpleCompiler) compileDropView(stmt *ast.DropViewStmt) (command.Comman
 	return cmd, nil
 }
 
-func (c *simpleCompiler) compileDelete(stmt *ast.DeleteStmt) (command.Command, error) {
+func (c *simpleCompiler) compileDelete(stmt *ast.DeleteStmt) (command.Delete, error) {
 	if stmt.WithClause != nil {
-		return nil, fmt.Errorf("with: %w", ErrUnsupported)
+		return command.Delete{}, fmt.Errorf("with: %w", ErrUnsupported)
 	}
 
 	var filter command.Expr
 	if stmt.Where != nil {
 		compiled, err := c.compileExpr(stmt.Expr)
 		if err != nil {
-			return nil, fmt.Errorf("expr: %w", err)
+			return command.Delete{}, fmt.Errorf("expr: %w", err)
 		}
 		filter = compiled
 	} else {
@@ -153,7 +227,7 @@ func (c *simpleCompiler) compileDelete(stmt *ast.DeleteStmt) (command.Command, e
 
 	table, err := c.compileQualifiedTableName(stmt.QualifiedTableName)
 	if err != nil {
-		return nil, fmt.Errorf("qualified table name: %w", err)
+		return command.Delete{}, fmt.Errorf("qualified table name: %w", err)
 	}
 	return command.Delete{
 		Table:  table,
@@ -235,14 +309,14 @@ func (c *simpleCompiler) compileSelectCore(core *ast.SelectCore) (command.Comman
 	return c.compileSelectCoreSelect(core)
 }
 
-func (c *simpleCompiler) compileSelectCoreValues(core *ast.SelectCore) (command.Command, error) {
+func (c *simpleCompiler) compileSelectCoreValues(core *ast.SelectCore) (command.Values, error) {
 	var datasets [][]command.Expr
 	for _, parExpr := range core.ParenthesizedExpressions {
 		var values []command.Expr
 		for _, expr := range parExpr.Exprs {
 			compiled, err := c.compileExpr(expr)
 			if err != nil {
-				return nil, fmt.Errorf("expr: %w", err)
+				return command.Values{}, fmt.Errorf("expr: %w", err)
 			}
 			values = append(values, compiled)
 		}
@@ -357,6 +431,9 @@ func (c *simpleCompiler) compileResultColumn(col *ast.ResultColumn) (command.Col
 func (c *simpleCompiler) compileExpr(expr *ast.Expr) (command.Expr, error) {
 	switch {
 	case expr.LiteralValue != nil:
+		if val := strings.ToLower(expr.LiteralValue.Value()); val == "true" || val == "false" {
+			return command.ConstantBooleanExpr{Value: val == "true"}, nil
+		}
 		return command.LiteralExpr{Value: expr.LiteralValue.Value()}, nil
 	case expr.UnaryOperator != nil:
 		val, err := c.compileExpr(expr.Expr1)


Test patch:
diff --git a/internal/compiler/simple_compiler_test.go b/internal/compiler/simple_compiler_test.go
index 602f8ce..71d80a5 100644
--- a/internal/compiler/simple_compiler_test.go
+++ b/internal/compiler/simple_compiler_test.go
@@ -19,6 +19,77 @@ func Test_simpleCompiler_Compile_NoOptimizations(t *testing.T) {
 	t.Run("select", _TestSimpleCompilerCompileSelectNoOptimizations)
 	t.Run("delete", _TestSimpleCompilerCompileDeleteNoOptimizations)
 	t.Run("drop", _TestSimpleCompilerCompileDropNoOptimizations)
+	t.Run("update", _TestSimpleCompilerCompileUpdateNoOptimizations)
+}
+
+func _TestSimpleCompilerCompileUpdateNoOptimizations(t *testing.T) {
+	tests := []testcase{
+		{
+			"simple update",
+			"UPDATE myTable SET myCol = 7",
+			command.Update{
+				UpdateOr: command.UpdateOrIgnore, // default
+				Table: command.SimpleTable{
+					Table: "myTable",
+				},
+				Filter: command.ConstantBooleanExpr{Value: true},
+				Updates: []command.UpdateSetter{
+					{
+						Cols:  []string{"myCol"},
+						Value: command.LiteralExpr{Value: "7"},
+					},
+				},
+			},
+			false,
+		},
+		{
+			"filtered update",
+			"UPDATE myTable SET myCol = 7 WHERE myOtherCol == 9",
+			command.Update{
+				UpdateOr: command.UpdateOrIgnore, // default
+				Table: command.SimpleTable{
+					Table: "myTable",
+				},
+				Filter: command.BinaryExpr{
+					Left:     command.LiteralExpr{Value: "myOtherCol"},
+					Operator: "==",
+					Right:    command.LiteralExpr{Value: "9"},
+				},
+				Updates: []command.UpdateSetter{
+					{
+						Cols:  []string{"myCol"},
+						Value: command.LiteralExpr{Value: "7"},
+					},
+				},
+			},
+			false,
+		},
+		{
+			"filtered update or fail",
+			"UPDATE OR FAIL myTable SET myCol = 7 WHERE myOtherCol == 9",
+			command.Update{
+				UpdateOr: command.UpdateOrFail,
+				Table: command.SimpleTable{
+					Table: "myTable",
+				},
+				Filter: command.BinaryExpr{
+					Left:     command.LiteralExpr{Value: "myOtherCol"},
+					Operator: "==",
+					Right:    command.LiteralExpr{Value: "9"},
+				},
+				Updates: []command.UpdateSetter{
+					{
+						Cols:  []string{"myCol"},
+						Value: command.LiteralExpr{Value: "7"},
+					},
+				},
+			},
+			false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, _TestCompile(tt))
+	}
 }
 
 func _TestSimpleCompilerCompileDropNoOptimizations(t *testing.T) {
@@ -186,9 +257,7 @@ func _TestSimpleCompilerCompileDeleteNoOptimizations(t *testing.T) {
 				Table: command.SimpleTable{
 					Table: "myTable",
 				},
-				Filter: command.ConstantBooleanExpr{
-					Value: true,
-				},
+				Filter: command.ConstantBooleanExpr{Value: true},
 			},
 			false,
 		},
@@ -200,9 +269,7 @@ func _TestSimpleCompilerCompileDeleteNoOptimizations(t *testing.T) {
 					Table:  "myTable",
 					Schema: "mySchema",
 				},
-				Filter: command.ConstantBooleanExpr{
-					Value: true,
-				},
+				Filter: command.ConstantBooleanExpr{Value: true},
 			},
 			false,
 		},
@@ -280,7 +347,7 @@ func _TestSimpleCompilerCompileSelectNoOptimizations(t *testing.T) {
 					},
 				},
 				Input: command.Select{
-					Filter: command.LiteralExpr{Value: "true"},
+					Filter: command.ConstantBooleanExpr{Value: true},
 					Input: command.Scan{
 						Table: command.SimpleTable{
 							Table: "myTable",
@@ -367,7 +434,7 @@ func _TestSimpleCompilerCompileSelectNoOptimizations(t *testing.T) {
 						},
 					},
 					Input: command.Select{
-						Filter: command.LiteralExpr{Value: "true"},
+						Filter: command.ConstantBooleanExpr{Value: true},
 						Input: command.Scan{
 							Table: command.SimpleTable{
 								Table: "myTable",
@@ -388,7 +455,7 @@ func _TestSimpleCompilerCompileSelectNoOptimizations(t *testing.T) {
 					},
 				},
 				Input: command.Select{
-					Filter: command.LiteralExpr{Value: "true"},
+					Filter: command.ConstantBooleanExpr{Value: true},
 					Input: command.Join{
 						Left: command.Scan{
 							Table: command.SimpleTable{
@@ -415,7 +482,7 @@ func _TestSimpleCompilerCompileSelectNoOptimizations(t *testing.T) {
 					},
 				},
 				Input: command.Select{
-					Filter: command.LiteralExpr{Value: "true"},
+					Filter: command.ConstantBooleanExpr{Value: true},
 					Input: command.Join{
 						Left: command.Scan{
 							Table: command.SimpleTable{
@@ -442,7 +509,7 @@ func _TestSimpleCompilerCompileSelectNoOptimizations(t *testing.T) {
 					},
 				},
 				Input: command.Select{
-					Filter: command.LiteralExpr{Value: "true"},
+					Filter: command.ConstantBooleanExpr{Value: true},
 					Input: command.Join{
 						Left: command.Join{
 							Left: command.Scan{