spanner/spannertest: support DELETE DML statements

This parses them in spansql, and executes them when they appear in an
ExecuteSql RPC.

Change-Id: I3a40c2ed039886f4ff58433b9d765a1ad8c6fb77
Reviewed-on: https://code-review.googlesource.com/c/gocloud/+/51112
Reviewed-by: kokoro <noreply+kokoro@google.com>
Reviewed-by: Knut Olav Løite <koloite@gmail.com>

spanner/spannertest/README.md

@@ -19,7 +19,7 @@ by ascending esotericism:
 - SELECT HAVING
 - arithmetic expressions (operators, parens)
 - transaction simulation
-- DML statements
+- INSERT/UPDATE DML statements
 - case insensitivity
 - alternate literal types (esp. strings)
 - allow_commit_timestamp

spanner/spannertest/db.go

@@ -716,3 +716,41 @@ func (r *keyRange) String() string {
 }
 
 type keyRangeList []*keyRange
+
+// Execute runs a DML statement.
+// It returns the number of affected rows.
+func (d *database) Execute(stmt spansql.DMLStmt, params queryParams) (int, error) { // TODO: return *status.Status instead?
+	switch stmt := stmt.(type) {
+	default:
+		return 0, status.Errorf(codes.Unimplemented, "unhandled DML statement type %T", stmt)
+	case *spansql.Delete:
+		t, err := d.table(stmt.Table)
+		if err != nil {
+			return 0, err
+		}
+
+		t.mu.Lock()
+		defer t.mu.Unlock()
+
+		n := 0
+		for i := 0; i < len(t.rows); {
+			ec := evalContext{
+				table:  t,
+				row:    t.rows[i],
+				params: params,
+			}
+			b, err := ec.evalBoolExpr(stmt.Where)
+			if err != nil {
+				return 0, err
+			}
+			if b {
+				copy(t.rows[i:], t.rows[i+1:])
+				t.rows = t.rows[:len(t.rows)-1]
+				n++
+				continue
+			}
+			i++
+		}
+		return n, nil
+	}
+}

spanner/spannertest/db_test.go

@@ -206,6 +206,40 @@ func TestTableData(t *testing.T) {
 	if err != nil {
 		t.Fatalf("Deleting key range: %v", err)
 	}
+	// Re-add the data and delete with DML.
+	err = db.Insert("Staff", []string{"Name", "ID"}, []*structpb.ListValue{
+		listV(stringV("01"), stringV("1")),
+		listV(stringV("03"), stringV("3")),
+		listV(stringV("06"), stringV("6")),
+	})
+	if err != nil {
+		t.Fatalf("Inserting data: %v", err)
+	}
+	n, err := db.Execute(&spansql.Delete{
+		Table: "Staff",
+		Where: spansql.LogicalOp{
+			LHS: spansql.ComparisonOp{
+				LHS: spansql.ID("Name"),
+				Op:  spansql.Ge,
+				RHS: spansql.Param("min"),
+			},
+			Op: spansql.And,
+			RHS: spansql.ComparisonOp{
+				LHS: spansql.ID("Name"),
+				Op:  spansql.Lt,
+				RHS: spansql.Param("max"),
+			},
+		},
+	}, queryParams{
+		"min": "01",
+		"max": "07",
+	})
+	if err != nil {
+		t.Fatalf("Deleting with DML: %v", err)
+	}
+	if n != 3 {
+		t.Errorf("Deleting with DML affected %d rows, want 3", n)
+	}
 
 	// Add a BYTES column, and populate it with some data.
 	st = db.ApplyDDL(&spansql.AlterTable{

spanner/spannertest/inmem.go

@@ -449,7 +449,39 @@ func (s *server) readTx(ctx context.Context, session string, tsel *spannerpb.Tra
 }
 
 func (s *server) ExecuteSql(ctx context.Context, req *spannerpb.ExecuteSqlRequest) (*spannerpb.ResultSet, error) {
-	return nil, status.Errorf(codes.Unimplemented, "ExecuteSql not implemented yet")
+	// Assume this is probably a DML statement. Queries tend to use ExecuteStreamingSql.
+	// TODO: Expand this to support more things.
+
+	obj, ok := req.Transaction.Selector.(*spannerpb.TransactionSelector_Id)
+	if !ok {
+		return nil, fmt.Errorf("unsupported transaction type %T", req.Transaction.Selector)
+	}
+	tid := string(obj.Id)
+	_ = tid // TODO: lookup an existing transaction by ID.
+
+	stmt, err := spansql.ParseDMLStmt(req.Sql)
+	if err != nil {
+		return nil, status.Errorf(codes.InvalidArgument, "bad DML: %v", err)
+	}
+	params, err := parseQueryParams(req.GetParams())
+	if err != nil {
+		return nil, err
+	}
+
+	s.logf("Executing: %s", stmt.SQL())
+	if len(params) > 0 {
+		s.logf("        ▹ %v", params)
+	}
+
+	n, err := s.db.Execute(stmt, params)
+	if err != nil {
+		return nil, err
+	}
+	return &spannerpb.ResultSet{
+		Stats: &spannerpb.ResultSetStats{
+			RowCount: &spannerpb.ResultSetStats_RowCountExact{int64(n)},
+		},
+	}, nil
 }
 
 func (s *server) ExecuteStreamingSql(req *spannerpb.ExecuteSqlRequest, stream spannerpb.Spanner_ExecuteStreamingSqlServer) error {
@@ -465,19 +497,9 @@ func (s *server) ExecuteStreamingSql(req *spannerpb.ExecuteSqlRequest, stream sp
 		return status.Errorf(codes.InvalidArgument, "bad query: %v", err)
 	}
 
-	params := make(queryParams)
-	for k, v := range req.GetParams().GetFields() {
-		switch v := v.Kind.(type) {
-		default:
-			return fmt.Errorf("unsupported well-known type value kind %T", v)
-		case *structpb.Value_NullValue:
-			params[k] = nil
-		case *structpb.Value_NumberValue:
-			params[k] = v.NumberValue
-		case *structpb.Value_StringValue:
-			params[k] = v.StringValue
-		}
-
+	params, err := parseQueryParams(req.GetParams())
+	if err != nil {
+		return err
 	}
 
 	s.logf("Querying: %s", q.SQL())
@@ -681,6 +703,23 @@ func (s *server) Rollback(ctx context.Context, req *spannerpb.RollbackRequest) (
 
 // TODO: PartitionQuery, PartitionRead
 
+func parseQueryParams(p *structpb.Struct) (queryParams, error) {
+	params := make(queryParams)
+	for k, v := range p.GetFields() {
+		switch v := v.Kind.(type) {
+		default:
+			return nil, fmt.Errorf("unsupported well-known type value kind %T", v)
+		case *structpb.Value_NullValue:
+			params[k] = nil
+		case *structpb.Value_NumberValue:
+			params[k] = v.NumberValue
+		case *structpb.Value_StringValue:
+			params[k] = v.StringValue
+		}
+	}
+	return params, nil
+}
+
 func spannerTypeFromType(typ spansql.Type) (*spannerpb.Type, error) {
 	var code spannerpb.TypeCode
 	switch typ.Base {

spanner/spansql/parser.go

@@ -112,6 +112,19 @@ func ParseDDLStmt(s string) (DDLStmt, error) {
 	return stmt, nil
 }
 
+// ParseDMLStmt parses a single DML statement.
+func ParseDMLStmt(s string) (DMLStmt, error) {
+	p := newParser("-", s)
+	stmt, err := p.parseDMLStmt()
+	if err != nil {
+		return nil, err
+	}
+	if p.Rem() != "" {
+		return nil, fmt.Errorf("unexpected trailing contents %q", p.Rem())
+	}
+	return stmt, nil
+}
+
 // ParseQuery parses a query string.
 func ParseQuery(s string) (Query, error) {
 	p := newParser("-", s)
@@ -1064,6 +1077,39 @@ func (p *parser) parseAlterTable() (*AlterTable, *parseError) {
 	// TODO: "ALTER"
 }
 
+func (p *parser) parseDMLStmt() (DMLStmt, *parseError) {
+	debugf("parseDMLStmt: %v", p)
+
+	/*
+		DELETE [FROM] target_name [[AS] alias]
+		WHERE condition
+
+		TODO: Insert, Update.
+	*/
+
+	if p.eat("DELETE") {
+		p.eat("FROM") // optional
+		tname, err := p.parseTableOrIndexOrColumnName()
+		if err != nil {
+			return nil, err
+		}
+		// TODO: parse alias.
+		if err := p.expect("WHERE"); err != nil {
+			return nil, err
+		}
+		where, err := p.parseBoolExpr()
+		if err != nil {
+			return nil, err
+		}
+		return &Delete{
+			Table: tname,
+			Where: where,
+		}, nil
+	}
+
+	return nil, p.errorf("unknown DML statement")
+}
+
 func (p *parser) parseColumnDef() (ColumnDef, *parseError) {
 	debugf("parseColumnDef: %v", p)
 

spanner/spansql/sql.go

@@ -106,6 +106,10 @@ func (od OnDelete) SQL() string {
 
 // TODO func (ac AlterColumn) SQL() string { }
 
+func (d *Delete) SQL() string {
+	return "DELETE FROM " + d.Table + " WHERE " + d.Where.SQL()
+}
+
 func (cd ColumnDef) SQL() string {
 	str := cd.Name + " " + cd.Type.SQL()
 	if cd.NotNull {

spanner/spansql/sql_test.go

@@ -34,6 +34,13 @@ func TestSQL(t *testing.T) {
 		ddl.clearOffset()
 		return ddl, nil
 	}
+	reparseDML := func(s string) (interface{}, error) {
+		dml, err := ParseDMLStmt(s)
+		if err != nil {
+			return nil, err
+		}
+		return dml, nil
+	}
 	reparseQuery := func(s string) (interface{}, error) {
 		q, err := ParseQuery(s)
 		return q, err
@@ -179,6 +186,18 @@ func TestSQL(t *testing.T) {
 			"ALTER TABLE Ta SET ON DELETE CASCADE",
 			reparseDDL,
 		},
+		{
+			&Delete{
+				Table: "Ta",
+				Where: ComparisonOp{
+					LHS: ID("C"),
+					Op:  Gt,
+					RHS: IntegerLiteral(2),
+				},
+			},
+			"DELETE FROM Ta WHERE C > 2",
+			reparseDML,
+		},
 		{
 			Query{
 				Select: Select{

spanner/spansql/types.go

@@ -136,6 +136,20 @@ type AlterColumn struct {
 }
 */
 
+// Delete represents a DELETE statement.
+// https://cloud.google.com/spanner/docs/dml-syntax#delete-statement
+type Delete struct {
+	Table string
+	Where BoolExpr
+
+	// TODO: Alias
+}
+
+func (d *Delete) String() string { return fmt.Sprintf("%#v", d) }
+func (*Delete) isDMLStmt()       {}
+
+// TODO: Insert, Update.
+
 // ColumnDef represents a column definition as part of a CREATE TABLE
 // or ALTER TABLE statement.
 type ColumnDef struct {
@@ -407,6 +421,12 @@ type DDLStmt interface {
 	Node
 }
 
+// DMLStmt is satisfied by a type that is a DML statement.
+type DMLStmt interface {
+	isDMLStmt()
+	SQL() string
+}
+
 // Node is implemented by concrete types in this package that represent things
 // appearing in a DDL file.
 type Node interface {