spanner/spannertest, spanner/spansql: support functions, COUNT(*)

The implementation in spannertest only supports COUNT(*) in a SELECT; other functions and aggregations are not yet supported. In particular, other aggregations will require some more restructuring of evalSelect. Updates #1181. Change-Id: Ifc859647cda57895f6940219379a91091654b4e9 Reviewed-on: https://code-review.googlesource.com/c/gocloud/+/43590 Reviewed-by: kokoro <[email protected]> Reviewed-by: Knut Olav Løite <[email protected]>
googleapis · Aug 7, 2019 · 1982229 · 1982229
1 parent 21afd04
commit 1982229
Show file tree

Hide file tree

Showing 6 changed files with 137 additions and 5 deletions.
diff --git a/spanner/spannertest/db_eval.go b/spanner/spannertest/db_eval.go
@@ -34,7 +34,7 @@ type evalContext struct {
 	params queryParams
 }
 
-func (d *database) evalSelect(sel spansql.Select, params queryParams, aux []spansql.Expr) (*resultIter, error) {
+func (d *database) evalSelect(sel spansql.Select, params queryParams, aux []spansql.Expr) (ri *resultIter, evalErr error) {
 	// TODO: weave this in below.
 	if len(sel.From) == 0 && sel.Where == nil {
 		// Simple expressions.
@@ -70,6 +70,24 @@ func (d *database) evalSelect(sel spansql.Select, params queryParams, aux []span
 		return nil, err
 	}
 
+	ri = &resultIter{}
+
+	// Handle COUNT(*) specially.
+	// TODO: Handle aggregation more generally.
+	if len(sel.List) == 1 && isCountStar(sel.List[0]) {
+		// Replace the `COUNT(*)` with `1`, then aggregate on the way out.
+		sel.List[0] = spansql.IntegerLiteral(1)
+		defer func() {
+			if evalErr != nil {
+				return
+			}
+			count := int64(len(ri.rows))
+			ri.rows = []resultRow{
+				{data: []interface{}{count}},
+			}
+		}()
+	}
+
 	// TODO: Support table sampling.
 
 	t.mu.Lock()
@@ -79,7 +97,6 @@ func (d *database) evalSelect(sel spansql.Select, params queryParams, aux []span
 		params: params,
 	}
 
-	ri := &resultIter{}
 	for _, e := range sel.List {
 		ci, err := ec.colInfo(e)
 		if err != nil {
@@ -420,3 +437,14 @@ func evalLike(str, pat string) bool {
 	}
 	return match
 }
+
+func isCountStar(e spansql.Expr) bool {
+	f, ok := e.(spansql.Func)
+	if !ok {
+		return false
+	}
+	if f.Name != "COUNT" || len(f.Args) != 1 {
+		return false
+	}
+	return f.Args[0] == spansql.Star
+}
diff --git a/spanner/spannertest/db_test.go b/spanner/spannertest/db_test.go
@@ -229,6 +229,13 @@ func TestTableData(t *testing.T) {
 				{"Sam", 1.75},
 			},
 		},
+		{
+			`SELECT COUNT(*) FROM Staff WHERE Name < "T"`,
+			nil,
+			[][]interface{}{
+				{int64(4)},
+			},
+		},
 	}
 	for _, test := range tests {
 		q, err := spansql.ParseQuery(test.q)

diff --git a/spanner/spansql/parser.go b/spanner/spansql/parser.go
@@ -410,7 +410,7 @@ func (p *parser) advance() {
 	// TODO: backtick (`) for quoted identifiers.
 	// TODO: array, struct, date, timestamp literals
 	switch p.s[0] {
-	case ',', ';', '(', ')':
+	case ',', ';', '(', ')', '*':
 		// Single character symbol.
 		p.cur.value, p.s = p.s[:1], p.s[1:]
 		return
@@ -1178,6 +1178,39 @@ func (p *parser) parseLimitCount() (Limit, error) {
 	return nil, p.errorf("got %q, want literal or parameter", tok.value)
 }
 
+func (p *parser) parseExprList() ([]Expr, error) {
+	if err := p.expect("("); err != nil {
+		return nil, err
+	}
+	var list []Expr
+	for {
+		if err := p.expect(")"); err == nil {
+			break
+		}
+		p.back()
+
+		e, err := p.parseExpr()
+		if err != nil {
+			return nil, err
+		}
+		list = append(list, e)
+
+		// ")" or "," should be next.
+		tok := p.next()
+		if tok.err != nil {
+			return nil, err
+		}
+		if tok.value == ")" {
+			break
+		} else if tok.value == "," {
+			continue
+		} else {
+			return nil, p.errorf(`got %q, want ")" or ","`, tok.value)
+		}
+	}
+	return list, nil
+}
+
 /*
 Expressions
 
@@ -1414,7 +1447,25 @@ func (p *parser) parseArithOp() (Expr, error) {
 		return Paren{Expr: e}, nil
 	}
 
-	return p.parseLit()
+	lit, err := p.parseLit()
+	if err != nil {
+		return nil, err
+	}
+
+	// If the literal was an identifier, and there's an open paren next,
+	// this is a function invocation.
+	if id, ok := lit.(ID); ok && p.sniff("(") {
+		list, err := p.parseExprList()
+		if err != nil {
+			return nil, err
+		}
+		return Func{
+			Name: string(id),
+			Args: list,
+		}, nil
+	}
+
+	return lit, nil
 }
 
 func (p *parser) parseLit() (Expr, error) {
@@ -1432,7 +1483,7 @@ func (p *parser) parseLit() (Expr, error) {
 		return StringLiteral(tok.string), nil
 	}
 
-	// Handle some reserved keywords that become specific values.
+	// Handle some reserved keywords and special tokens that become specific values.
 	// TODO: Handle the other 92 keywords.
 	switch tok.value {
 	case "TRUE":
@@ -1441,6 +1492,8 @@ func (p *parser) parseLit() (Expr, error) {
 		return False, nil
 	case "NULL":
 		return Null, nil
+	case "*":
+		return Star, nil
 	}
 
 	// TODO: more types of literals (array, struct, date, timestamp).

diff --git a/spanner/spansql/parser_test.go b/spanner/spansql/parser_test.go
@@ -55,6 +55,19 @@ func TestParseQuery(t *testing.T) {
 				Limit: Param("limit"),
 			},
 		},
+		{`SELECT COUNT(*) FROM Packages`,
+			Query{
+				Select: Select{
+					List: []Expr{
+						Func{
+							Name: "COUNT",
+							Args: []Expr{Star},
+						},
+					},
+					From: []SelectFrom{{Table: "Packages"}},
+				},
+			},
+		},
 	}
 	for _, test := range tests {
 		got, err := ParseQuery(test.in)

diff --git a/spanner/spansql/sql.go b/spanner/spansql/sql.go
@@ -243,6 +243,18 @@ func (io IsOp) SQL() string {
 	return str
 }
 
+func (f Func) SQL() string {
+	str := f.Name + "("
+	for i, e := range f.Args {
+		if i > 0 {
+			str += ", "
+		}
+		str += e.SQL()
+	}
+	str += ")"
+	return str
+}
+
 func (p Paren) SQL() string { return "(" + p.Expr.SQL() + ")" }
 
 func (id ID) SQL() string   { return string(id) }
@@ -256,6 +268,7 @@ func (b BoolLiteral) SQL() string {
 }
 
 func (n NullLiteral) SQL() string { return "NULL" }
+func (StarExpr) SQL() string      { return "*" }
 
 func (il IntegerLiteral) SQL() string { return strconv.Itoa(int(il)) }
 func (fl FloatLiteral) SQL() string   { return strconv.FormatFloat(float64(fl), 'g', -1, 64) }

diff --git a/spanner/spansql/types.go b/spanner/spansql/types.go
@@ -250,6 +250,17 @@ type IsExpr interface {
 	SQL() string
 }
 
+// Func represents a function call.
+type Func struct {
+	Name string
+	Args []Expr
+
+	// TODO: various functions permit as-expressions, which might warrant different types in here.
+}
+
+func (Func) isBoolExpr() {} // possibly bool
+func (Func) isExpr()     {}
+
 // Paren represents a parenthesised expression.
 type Paren struct {
 	Expr Expr
@@ -308,6 +319,13 @@ type StringLiteral string
 
 func (StringLiteral) isExpr() {}
 
+type StarExpr int
+
+// Star represents a "*" in an expression.
+const Star = StarExpr(0)
+
+func (StarExpr) isExpr() {}
+
 // DDL
 // https://cloud.google.com/spanner/docs/data-definition-language#ddl_syntax