From 7d2d83ee1cce58d4014d5570bc599bcef1ed9c22 Mon Sep 17 00:00:00 2001
From: David Symonds <dsymonds@golang.org>
Date: Tue, 20 Oct 2020 12:10:07 +1100
Subject: [PATCH] feat(spanner/spannertest): make SELECT list aliases visible
 to ORDER BY (#3054)

As noted in #3043.
---
 spanner/spannertest/db_query.go         | 15 ++++++++-------
 spanner/spannertest/integration_test.go | 11 +++++++++++
 2 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/spanner/spannertest/db_query.go b/spanner/spannertest/db_query.go
index dc207e8b613..b9bbc01b0d0 100644
--- a/spanner/spannertest/db_query.go
+++ b/spanner/spannertest/db_query.go
@@ -461,18 +461,19 @@ func (d *database) evalSelect(sel spansql.Select, qc *queryContext) (si *selIter
 		}
 	}
 
+	// Load aliases visible to any future iterators,
+	// including GROUP BY and ORDER BY. These are not visible to the WHERE clause.
+	ec.aliases = make(map[spansql.ID]spansql.Expr)
+	for i, alias := range sel.ListAliases {
+		ec.aliases[alias] = sel.List[i]
+	}
+	// TODO: Add aliases for "1", "2", etc.
+
 	// Apply GROUP BY.
 	// This only reorders rows to group rows together;
 	// aggregation happens next.
 	var rowGroups [][2]int // Sequence of half-open intervals of row numbers.
 	if len(sel.GroupBy) > 0 {
-		// Load aliases visible to this GROUP BY.
-		ec.aliases = make(map[spansql.ID]spansql.Expr)
-		for i, alias := range sel.ListAliases {
-			ec.aliases[alias] = sel.List[i]
-		}
-		// TODO: Add aliases for "1", "2", etc.
-
 		raw, err := toRawIter(ri)
 		if err != nil {
 			return nil, err
diff --git a/spanner/spannertest/integration_test.go b/spanner/spannertest/integration_test.go
index 349e7f6053b..6ee0d6f29f4 100644
--- a/spanner/spannertest/integration_test.go
+++ b/spanner/spannertest/integration_test.go
@@ -871,12 +871,23 @@ func TestIntegration_ReadsAndQueries(t *testing.T) {
 		},
 		// SELECT with aliases.
 		{
+			// Aliased table.
 			`SELECT s.Name FROM Staff AS s WHERE s.ID = 3 ORDER BY s.Tenure`,
 			nil,
 			[][]interface{}{
 				{"Sam"},
 			},
 		},
+		{
+			// Aliased expression.
+			`SELECT Name AS nom FROM Staff WHERE ID < 4 ORDER BY nom`,
+			nil,
+			[][]interface{}{
+				{"Daniel"},
+				{"Jack"},
+				{"Sam"},
+			},
+		},
 		// Joins.
 		{
 			`SELECT * FROM JoinA INNER JOIN JoinB ON JoinA.w = JoinB.y ORDER BY w, x, y, z`,