chore(spanner/spansql): restructure FROM, TABLESAMPLE (#2888)

* spanner/spansql: restructure how TABLESAMPLE fits in to AST

This is to make room for other types of FROM clauses.

Updates #2850.

* spanner/spansql: restructure SelectFrom

This is to make room for other kinds of FROM clauses, such as JOIN.

Updates #2850.

spanner/spannertest/db_query.go

@@ -335,10 +335,15 @@ func (d *database) evalSelect(sel spansql.Select, params queryParams) (ri rowIte
 	// First stage is to identify the data source.
 	// If there's a FROM then that names a table to use.
 	if len(sel.From) > 1 {
-		return nil, fmt.Errorf("selecting from more than one table not yet supported")
+		return nil, fmt.Errorf("selecting with more than one FROM clause not yet supported")
 	}
 	if len(sel.From) == 1 {
-		tableName := sel.From[0].Table
+		sft, ok := sel.From[0].(spansql.SelectFromTable)
+		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

spanner/spansql/parser.go

@@ -1686,25 +1686,37 @@ func (p *parser) parseSelect() (Select, *parseError) {
 	sel.List, sel.ListAliases = list, aliases
 
 	if p.eat("FROM") {
+		padTS := func() {
+			for len(sel.TableSamples) < len(sel.From) {
+				sel.TableSamples = append(sel.TableSamples, nil)
+			}
+		}
+
 		for {
 			from, err := p.parseSelectFrom()
 			if err != nil {
 				return Select{}, err
 			}
+			sel.From = append(sel.From, from)
+
 			if p.sniff("TABLESAMPLE") {
 				ts, err := p.parseTableSample()
 				if err != nil {
 					return Select{}, err
 				}
-				from.TableSample = &ts
+				padTS()
+				sel.TableSamples[len(sel.TableSamples)-1] = &ts
 			}
-			sel.From = append(sel.From, from)
 
 			if p.eat(",") {
 				continue
 			}
 			break
 		}
+
+		if sel.TableSamples != nil {
+			padTS()
+		}
 	}
 
 	if p.eat("WHERE") {
@@ -1770,15 +1782,15 @@ func (p *parser) parseSelectFrom() (SelectFrom, *parseError) {
 	// TODO: support more than a single table name.
 	tname, err := p.parseTableOrIndexOrColumnName()
 	if err != nil {
-		return SelectFrom{}, err
+		return nil, err
 	}
-	sf := SelectFrom{Table: tname}
+	sf := SelectFromTable{Table: tname}
 
 	// TODO: The "AS" keyword is optional.
 	if p.eat("AS") {
 		alias, err := p.parseAlias()
 		if err != nil {
-			return SelectFrom{}, err
+			return nil, err
 		}
 		sf.Alias = alias
 	}

spanner/spansql/parser_test.go

@@ -33,7 +33,7 @@ func TestParseQuery(t *testing.T) {
 			Query{
 				Select: Select{
 					List: []Expr{ID("Alias")},
-					From: []SelectFrom{{
+					From: []SelectFrom{SelectFromTable{
 						Table: "Characters",
 					}},
 					Where: LogicalOp{
@@ -68,15 +68,15 @@ func TestParseQuery(t *testing.T) {
 							Args: []Expr{Star},
 						},
 					},
-					From: []SelectFrom{{Table: "Packages"}},
+					From: []SelectFrom{SelectFromTable{Table: "Packages"}},
 				},
 			},
 		},
 		{`SELECT * FROM Packages`,
 			Query{
 				Select: Select{
 					List: []Expr{Star},
-					From: []SelectFrom{{Table: "Packages"}},
+					From: []SelectFrom{SelectFromTable{Table: "Packages"}},
 				},
 			},
 		},
@@ -88,7 +88,7 @@ func TestParseQuery(t *testing.T) {
 						ID("FirstName"),
 						ID("LastName"),
 					},
-					From:        []SelectFrom{{Table: "PlayerStats"}},
+					From:        []SelectFrom{SelectFromTable{Table: "PlayerStats"}},
 					GroupBy:     []Expr{ID("FirstName"), ID("LastName")},
 					ListAliases: []string{"total_points", "", "surname"},
 				},
@@ -103,7 +103,7 @@ func TestParseQuery(t *testing.T) {
 					List: []Expr{
 						Func{Name: "COUNT", Args: []Expr{Star}},
 					},
-					From: []SelectFrom{{Table: "Lists", Alias: "l"}},
+					From: []SelectFrom{SelectFromTable{Table: "Lists", Alias: "l"}},
 					Where: ComparisonOp{
 						Op:  Eq,
 						LHS: ID("l_user_id"),

spanner/spansql/sql.go

@@ -262,7 +262,7 @@ func (sel Select) SQL() string {
 			if i > 0 {
 				str += ", "
 			}
-			str += ID(f.Table).SQL()
+			str += f.SQL()
 		}
 	}
 	if sel.Where != nil {
@@ -280,6 +280,14 @@ func (sel Select) SQL() string {
 	return str
 }
 
+func (sft SelectFromTable) SQL() string {
+	str := ID(sft.Table).SQL()
+	if sft.Alias != "" {
+		str += " AS " + ID(sft.Alias).SQL()
+	}
+	return str
+}
+
 func (o Order) SQL() string {
 	str := o.Expr.SQL()
 	if o.Desc {

spanner/spansql/sql_test.go

@@ -236,7 +236,7 @@ func TestSQL(t *testing.T) {
 			Query{
 				Select: Select{
 					List: []Expr{ID("A"), ID("B")},
-					From: []SelectFrom{{Table: "Table"}},
+					From: []SelectFrom{SelectFromTable{Table: "Table"}},
 					Where: LogicalOp{
 						LHS: ComparisonOp{
 							LHS: ID("C"),

spanner/spansql/types.go

@@ -291,20 +291,34 @@ type Select struct {
 	GroupBy  []Expr
 	// TODO: Having
 
+	// When the FROM clause has TABLESAMPLE operators,
+	// TableSamples will be populated 1:1 with From;
+	// FROM clauses without will have a nil value.
+	TableSamples []*TableSample
+
 	// If the SELECT list has explicit aliases ("AS alias"),
 	// ListAliases will be populated 1:1 with List;
 	// aliases that are present will be non-empty.
 	ListAliases []string
 }
 
-type SelectFrom struct {
-	// This only supports a FROM clause directly from a table.
+// SelectFrom represents the FROM clause of a SELECT.
+// https://cloud.google.com/spanner/docs/query-syntax#from_clause
+type SelectFrom interface {
+	isSelectFrom()
+	SQL() string
+}
+
+// SelectFromTable is a SelectFrom that specifies a table to read from.
+type SelectFromTable struct {
 	Table string
 	Alias string // empty if not aliased
-
-	TableSample *TableSample // TODO: This isn't part of from_item; move elsewhere.
 }
 
+func (SelectFromTable) isSelectFrom() {}
+
+// TODO: SelectFromJoin, SelectFromSubquery, etc.
+
 type Order struct {
 	Expr Expr
 	Desc bool