feat(spanner/spannertest): Support generated columns (#4742)

* feat(spanner/spannertest): Support generated columns

* Add support for LOWER function

* Remove checking of dependent columns and improve testing

* docs: fix a couple of comments

* Make CI happy with casing

* Don't panic if columns are nil in generated columns using funcs

* No longer returning code.InvalidArgument when calling db_eval.evalFunc
  but that seems consistent with the rest of the module, no tests failed
  so I assume its not behaviour relied upon, we could check for the
  error type instead

* Add ErrorIs and make functions public to allow spannertest to use

* fix: add errorIs to spannertest package

* fix: remove unrelated changes

* fix: remove specific error checking and return null from evaluations

* fix: restore to original errors

* Prettify comments

* test: add error msg + delete test

Co-authored-by: Knut Olav Løite <koloite@gmail.com>

spanner/spannertest/README.md

@@ -22,7 +22,9 @@ by ascending esotericism:
 - more aggregation functions
 - SELECT HAVING
 - more literal types
-- generated columns
+- expressions that return null for generated columns
+- generated columns referencing other generated columns
+- checking dependencies on a generated column before deleting a column
 - expression type casting, coercion
 - multiple joins
 - subselects

spanner/spannertest/db.go

@@ -66,11 +66,12 @@ type table struct {
 
 // colInfo represents information about a column in a table or result set.
 type colInfo struct {
-	Name     spansql.ID
-	Type     spansql.Type
-	NotNull  bool            // only set for table columns
-	AggIndex int             // Index+1 of SELECT list for which this is an aggregate value.
-	Alias    spansql.PathExp // an alternate name for this column (result sets only)
+	Name      spansql.ID
+	Type      spansql.Type
+	Generated spansql.Expr
+	NotNull   bool            // only set for table columns
+	AggIndex  int             // Index+1 of SELECT list for which this is an aggregate value.
+	Alias     spansql.PathExp // an alternate name for this column (result sets only)
 }
 
 // commitTimestampSentinel is a sentinel value for TIMESTAMP fields with allow_commit_timestamp=true.
@@ -395,6 +396,9 @@ func (d *database) writeValues(tx *transaction, tbl spansql.ID, cols []spansql.I
 		for j, v := range vs.Values {
 			i := colIndexes[j]
 
+			if t.cols[i].Generated != nil {
+				return status.Error(codes.InvalidArgument, "values can't be written to a generated column")
+			}
 			x, err := valForType(v, t.cols[i].Type)
 			if err != nil {
 				return err
@@ -414,6 +418,33 @@ func (d *database) writeValues(tx *transaction, tbl spansql.ID, cols []spansql.I
 		if err := f(t, colIndexes, r); err != nil {
 			return err
 		}
+
+		// Get row again after potential update merge to ensure we compute
+		// generated columns with fresh data.
+		pk := r[:t.pkCols]
+		rowNum, found := t.rowForPK(pk)
+		// This should never fail as the row was just inserted.
+		if !found {
+			return status.Error(codes.Internal, "row failed to be inserted")
+		}
+		row := t.rows[rowNum]
+		ec := evalContext{
+			cols: t.cols,
+			row:  row,
+		}
+
+		// TODO: We would need to do a topological sort on dependencies
+		// (i.e. what other columns the expression references) to ensure we
+		// can handle generated columns which reference other generated columns
+		for i, col := range t.cols {
+			if col.Generated != nil {
+				res, err := ec.evalExpr(col.Generated)
+				if err != nil {
+					return err
+				}
+				row[i] = res
+			}
+		}
 	}
 
 	return nil
@@ -634,6 +665,10 @@ func (t *table) addColumn(cd spansql.ColumnDef, newTable bool) *status.Status {
 			// TODO: what happens in this case?
 			return status.Newf(codes.Unimplemented, "can't add NOT NULL columns to non-empty tables yet")
 		}
+		if cd.Generated != nil {
+			// TODO: should backfill the data to maintain behaviour with real spanner
+			return status.Newf(codes.Unimplemented, "can't add generated columns to non-empty tables yet")
+		}
 		for i := range t.rows {
 			t.rows[i] = append(t.rows[i], nil)
 		}
@@ -643,6 +678,11 @@ func (t *table) addColumn(cd spansql.ColumnDef, newTable bool) *status.Status {
 		Name:    cd.Name,
 		Type:    cd.Type,
 		NotNull: cd.NotNull,
+		// TODO: We should figure out what columns the Generator expression
+		// relies on and check it is valid at this time currently it will
+		// fail when writing data instead as it is the first time we
+		// evaluate the expression.
+		Generated: cd.Generated,
 	})
 	t.colIndex[cd.Name] = len(t.cols) - 1
 	if !newTable {

spanner/spannertest/db_eval.go

@@ -237,12 +237,16 @@ func (ec evalContext) evalBoolExpr(be spansql.BoolExpr) (*bool, error) {
 }
 
 func (ec evalContext) evalArithOp(e spansql.ArithOp) (interface{}, error) {
+	// TODO: Better NULL handling
 	switch e.Op {
 	case spansql.Neg:
 		rhs, err := ec.evalExpr(e.RHS)
 		if err != nil {
 			return nil, err
 		}
+		if rhs == nil {
+			return nil, nil
+		}
 		switch rhs := rhs.(type) {
 		case float64:
 			return -rhs, nil
@@ -255,6 +259,9 @@ func (ec evalContext) evalArithOp(e spansql.ArithOp) (interface{}, error) {
 		if err != nil {
 			return nil, err
 		}
+		if rhs == nil {
+			return nil, nil
+		}
 		switch rhs := rhs.(type) {
 		case int64:
 			return ^rhs, nil
@@ -285,10 +292,16 @@ func (ec evalContext) evalArithOp(e spansql.ArithOp) (interface{}, error) {
 		if err != nil {
 			return nil, err
 		}
+		if lhs == nil {
+			return nil, nil
+		}
 		rhs, err := ec.evalExpr(e.RHS)
 		if err != nil {
 			return nil, err
 		}
+		if rhs == nil {
+			return nil, nil
+		}
 		i1, ok1 := lhs.(int64)
 		i2, ok2 := rhs.(int64)
 		if ok1 && ok2 {
@@ -322,10 +335,16 @@ func (ec evalContext) evalArithOp(e spansql.ArithOp) (interface{}, error) {
 		if err != nil {
 			return nil, err
 		}
+		if lhs == nil {
+			return nil, nil
+		}
 		rhs, err := ec.evalExpr(e.RHS)
 		if err != nil {
 			return nil, err
 		}
+		if rhs == nil {
+			return nil, nil
+		}
 		i1, ok1 := lhs.(int64)
 		i2, ok2 := rhs.(int64)
 		if ok1 && ok2 {

spanner/spannertest/db_test.go

@@ -22,10 +22,12 @@ import (
 	"fmt"
 	"io"
 	"reflect"
+	"strings"
 	"sync"
 	"testing"
 
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 
 	structpb "github.com/golang/protobuf/ptypes/struct"
 
@@ -361,6 +363,90 @@ func TestConcurrentReadInsert(t *testing.T) {
 	}
 }
 
+func TestGeneratedColumn(t *testing.T) {
+	sql := `CREATE TABLE Songwriters (
+		Id INT64 NOT NULL,
+		Name STRING(20),
+		Age INT64,
+		Over18 BOOL AS (Age > 18) STORED,
+	) PRIMARY KEY (Id);`
+	var db database
+
+	ddl, err := spansql.ParseDDL("filename", sql)
+	if err != nil {
+		t.Fatalf("%s: Bad DDL", err)
+	}
+	for _, stmt := range ddl.List {
+		if st := db.ApplyDDL(stmt); st.Code() != codes.OK {
+			t.Fatalf("ApplyDDL failed: %v", st)
+		}
+	}
+
+	addColSQL := `ALTER TABLE Songwriters ADD COLUMN CanonicalName STRING(20) AS (LOWER(Name)) STORED;`
+	ddl, err = spansql.ParseDDL("filename", addColSQL)
+	if err != nil {
+		t.Fatalf("%s: Bad DDL", err)
+	}
+	if st := db.ApplyDDL(ddl.List[0]); st.Code() != codes.OK {
+		t.Fatalf("Should have been able to add a generated column to empty table\n status: %v", st)
+	}
+
+	tx := db.NewTransaction()
+	err = db.Insert(tx, "Songwriters",
+		[]spansql.ID{"Id", "Over18"},
+		[]*structpb.ListValue{
+			listV(stringV("3"), boolV(true)),
+		})
+	if err == nil || status.Code(err) != codes.InvalidArgument {
+		t.Fatal("Should have failed to insert to generated column")
+	}
+
+	err = db.Insert(tx, "Songwriters",
+		[]spansql.ID{"Id"},
+		[]*structpb.ListValue{
+			listV(stringV("1")),
+		})
+	if err != nil {
+		t.Fatalf("Should have succeeded to insert to with no dependent columns: %v", err)
+	}
+
+	name := "Famous Writer"
+	err = db.Insert(tx, "Songwriters",
+		[]spansql.ID{"Id", "Name", "Age"},
+		[]*structpb.ListValue{
+			listV(stringV("3"), stringV(name), stringV("40")),
+		})
+	if err != nil {
+		t.Fatalf("Should have succeeded to insert to without generated column: %v", err)
+	}
+
+	var kr keyRangeList
+	iter, err := db.Read("Songwriters", []spansql.ID{"Id", "CanonicalName", "Over18"},
+		[]*structpb.ListValue{
+			listV(stringV("3")),
+		}, kr, 0)
+	if err != nil {
+		t.Fatalf("failed to read: %v", err)
+	}
+	rows := slurp(t, iter)
+	if rows[0][1].(string) != strings.ToLower(name) {
+		t.Fatalf("Generated value for CanonicalName mismatch\n Got: %v\n Want: %v", rows[0][1].(string), strings.ToLower(name))
+	}
+	if !rows[0][2].(bool) {
+		t.Fatalf("Generated value for Over18 mismatch\n Got: %v\n Want: true", rows[0][2].(bool))
+	}
+
+	addColSQL = `ALTER TABLE Songwriters ADD COLUMN Under18 BOOL AS (Age < 18) STORED;`
+	ddl, err = spansql.ParseDDL("filename", addColSQL)
+	if err != nil {
+		t.Fatalf("%s: Bad DDL", err)
+	}
+	if st := db.ApplyDDL(ddl.List[0]); st.Code() == codes.OK {
+		t.Fatalf("Should have failed to add a generated column to non-empty table\n status: %v", st)
+	}
+
+}
+
 func slurp(t *testing.T, ri rowIter) (all [][]interface{}) {
 	t.Helper()
 	for {

spanner/spannertest/funcs.go

@@ -52,6 +52,20 @@ var functions = map[string]function{
 			return strings.HasPrefix(s, prefix), spansql.Type{Base: spansql.Bool}, nil
 		},
 	},
+	"LOWER": {
+		Eval: func(values []interface{}) (interface{}, spansql.Type, error) {
+			if len(values) != 1 {
+				return nil, spansql.Type{}, status.Error(codes.InvalidArgument, "No matching signature for function LOWER for the given argument types")
+			}
+			if values[0] == nil {
+				return nil, spansql.Type{Base: spansql.String}, nil
+			}
+			if _, ok := values[0].(string); !ok {
+				return nil, spansql.Type{}, status.Error(codes.InvalidArgument, "No matching signature for function LOWER for the given argument types")
+			}
+			return strings.ToLower(values[0].(string)), spansql.Type{Base: spansql.String}, nil
+		},
+	},
 }
 
 type aggregateFunc struct {