feat(spanner/spannertest): evaluate path expressions (#2931)

This adds support for aliases in SELECT statements, and in particular
table aliases that can then be referenced from other places in the query
evaluation.

Fixes #2463.

spanner/spannertest/README.md

@@ -25,7 +25,6 @@ by ascending esotericism:
 - joins
 - transaction simulation
 - expression type casting, coercion
-- SELECT aliases in FROM clause, ORDER BY
 - subselects
 - set operations (UNION, INTERSECT, EXCEPT)
 - partition support

spanner/spannertest/db.go

@@ -67,7 +67,8 @@ type table struct {
 type colInfo struct {
 	Name     spansql.ID
 	Type     spansql.Type
-	AggIndex int // Index+1 of SELECT list for which this is an aggregate value.
+	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.

spanner/spannertest/db_eval.go

@@ -354,6 +354,8 @@ func (ec evalContext) evalExpr(e spansql.Expr) (interface{}, error) {
 		return nil, fmt.Errorf("TODO: evalExpr(%s %T)", e.SQL(), e)
 	case coercedValue:
 		return e.val, nil
+	case spansql.PathExp:
+		return ec.evalPathExp(e)
 	case spansql.ID:
 		return ec.evalID(e)
 	case spansql.Param:
@@ -449,11 +451,36 @@ func (ec evalContext) evalExpr(e spansql.Expr) (interface{}, error) {
 	}
 }
 
-func (ec evalContext) evalID(id spansql.ID) (interface{}, error) {
-	for i, col := range ec.cols {
-		if col.Name == id {
-			return ec.row.copyDataElem(i), nil
+// resolveColumnIndex turns an ID or PathExp into a table column index.
+func (ec evalContext) resolveColumnIndex(e spansql.Expr) (int, error) {
+	switch e := e.(type) {
+	case spansql.ID:
+		for i, col := range ec.cols {
+			if col.Name == e {
+				return i, nil
+			}
 		}
+	case spansql.PathExp:
+		for i, col := range ec.cols {
+			if pathExpEqual(e, col.Alias) {
+				return i, nil
+			}
+		}
+	}
+	return 0, fmt.Errorf("couldn't resolve [%s] as a table column", e.SQL())
+}
+
+func (ec evalContext) evalPathExp(pe spansql.PathExp) (interface{}, error) {
+	// TODO: support more than only naming an aliased table column.
+	if i, err := ec.resolveColumnIndex(pe); err == nil {
+		return ec.row.copyDataElem(i), nil
+	}
+	return nil, fmt.Errorf("couldn't resolve path expression %s", pe.SQL())
+}
+
+func (ec evalContext) evalID(id spansql.ID) (interface{}, error) {
+	if i, err := ec.resolveColumnIndex(id); err == nil {
+		return ec.row.copyDataElem(i), nil
 	}
 	if e, ok := ec.aliases[id]; ok {
 		// Make a copy of the context without this alias
@@ -681,13 +708,13 @@ func (ec evalContext) colInfo(e spansql.Expr) (colInfo, error) {
 		return colInfo{Type: t}, nil
 	case spansql.LogicalOp, spansql.ComparisonOp, spansql.IsOp:
 		return colInfo{Type: spansql.Type{Base: spansql.Bool}}, nil
-	case spansql.ID:
+	case spansql.PathExp, spansql.ID:
 		// TODO: support more than only naming a table column.
-		for _, col := range ec.cols {
-			if col.Name == e {
-				return col, nil
-			}
+		i, err := ec.resolveColumnIndex(e)
+		if err == nil {
+			return ec.cols[i], nil
 		}
+		// Let errors fall through.
 	case spansql.Param:
 		qp, ok := ec.params[string(e)]
 		if !ok {
@@ -748,6 +775,18 @@ func (ec evalContext) arithColType(ao spansql.ArithOp) (spansql.Type, error) {
 	}
 }
 
+func pathExpEqual(a, b spansql.PathExp) bool {
+	if len(a) != len(b) {
+		return false
+	}
+	for i := range a {
+		if a[i] != b[i] {
+			return false
+		}
+	}
+	return true
+}
+
 func evalLike(str, pat string) bool {
 	/*
 		% matches any number of chars.

spanner/spannertest/db_query.go

@@ -82,9 +82,21 @@ func (ni *nullIter) Next() (row, error) {
 type tableIter struct {
 	t        *table
 	rowIndex int // index of next row to return
+
+	alias spansql.ID // if non-empty, "AS <alias>"
+}
+
+func (ti *tableIter) Cols() []colInfo {
+	var cis []colInfo
+	for _, ci := range ti.t.cols {
+		if ti.alias != "" {
+			ci.Alias = spansql.PathExp{ti.alias, ci.Name}
+		}
+		cis = append(cis, ci)
+	}
+	return cis
 }
 
-func (ti *tableIter) Cols() []colInfo { return ti.t.cols }
 func (ti *tableIter) Next() (row, error) {
 	if ti.rowIndex >= len(ti.t.rows) {
 		return nil, io.EOF
@@ -342,16 +354,16 @@ func (d *database) evalSelect(sel spansql.Select, params queryParams) (ri rowIte
 		if !ok {
 			return nil, fmt.Errorf("selecting with FROM clause of type %T not yet supported", sel.From[0])
 		}
-		// TODO: sft.Alias needs mixing in here.
 		tableName := sft.Table
 		t, err := d.table(tableName)
 		if err != nil {
 			return nil, err
 		}
 		t.mu.Lock()
 		defer t.mu.Unlock()
-		ri = &tableIter{t: t}
-		ec.cols = t.cols
+		ti := &tableIter{t: t, alias: sft.Alias}
+		ri = ti
+		ec.cols = ti.Cols()
 
 		// On the way out, convert the result to a rawIter
 		// so that the table may be safely unlocked.

spanner/spannertest/db_test.go

@@ -548,6 +548,14 @@ func TestTableData(t *testing.T) {
 				{[]interface{}{false, nil, nil, false, true}},
 			},
 		},
+		// SELECT with aliases.
+		{
+			`SELECT s.Name FROM Staff AS s WHERE s.ID = 3 ORDER BY s.Tenure`,
+			nil,
+			[][]interface{}{
+				{"Sam"},
+			},
+		},
 		// Regression test for aggregating no rows; it used to return an empty row.
 		// https://github.com/googleapis/google-cloud-go/issues/2793
 		{