feat(spanner/spansql): support table_hint_expr at from_clause on query_statement (#4457)

* feat(spanner/spansql): support table_hint_expr at from_clause on
query_statement

This fixes parse error when query statement includes table hint expr.
This add function parseHints and use in parsing table hints and join
hints.

* feat(spanner/spansql): modify to enable multiple table hint keys.

* feat(spanner/spansql): use ++ increment for lint

Co-authored-by: Hengfeng Li <hengfeng@google.com>

Author: Takeshi Nakata

spanner/spansql/parser.go

@@ -1885,7 +1885,7 @@ func (p *parser) parseQuery() (Query, *parseError) {
 			[ LIMIT count [ OFFSET skip_rows ] ]
 	*/
 
-	// TODO: hints, sub-selects, etc.
+	// TODO: sub-selects, etc.
 
 	if err := p.expect("SELECT"); err != nil {
 		return Query{}, err
@@ -2111,6 +2111,13 @@ func (p *parser) parseSelectFrom() (SelectFrom, *parseError) {
 		return nil, err
 	}
 	sf := SelectFromTable{Table: tname}
+	if p.eat("@") {
+		hints, err := p.parseHints(map[string]string{})
+		if err != nil {
+			return nil, err
+		}
+		sf.Hints = hints
+	}
 
 	// TODO: The "AS" keyword is optional.
 	if p.eat("AS") {
@@ -2159,46 +2166,20 @@ func (p *parser) parseSelectFrom() (SelectFrom, *parseError) {
 		Type: jt,
 		LHS:  sf,
 	}
-	setHint := func(k, v string) {
-		if sfj.Hints == nil {
-			sfj.Hints = make(map[string]string)
-		}
-		sfj.Hints[k] = v
-	}
+	var hints map[string]string
 	if hashJoin {
-		setHint("JOIN_METHOD", "HASH_JOIN")
+		hints = map[string]string{}
+		hints["JOIN_METHOD"] = "HASH_JOIN"
 	}
 
 	if p.eat("@") {
-		if err := p.expect("{"); err != nil {
-			return nil, err
-		}
-		for {
-			if p.sniff("}") {
-				break
-			}
-			tok := p.next()
-			if tok.err != nil {
-				return nil, tok.err
-			}
-			k := tok.value
-			if err := p.expect("="); err != nil {
-				return nil, err
-			}
-			tok = p.next()
-			if tok.err != nil {
-				return nil, tok.err
-			}
-			v := tok.value
-			setHint(k, v)
-			if !p.eat(",") {
-				break
-			}
-		}
-		if err := p.expect("}"); err != nil {
+		h, err := p.parseHints(hints)
+		if err != nil {
 			return nil, err
 		}
+		hints = h
 	}
+	sfj.Hints = hints
 
 	sfj.RHS, err = p.parseSelectFrom()
 	if err != nil {
@@ -2889,6 +2870,41 @@ func (p *parser) parseAlias() (ID, *parseError) {
 	return p.parseTableOrIndexOrColumnName()
 }
 
+func (p *parser) parseHints(hints map[string]string) (map[string]string, *parseError) {
+	if hints == nil {
+		hints = map[string]string{}
+	}
+	if err := p.expect("{"); err != nil {
+		return nil, err
+	}
+	for {
+		if p.sniff("}") {
+			break
+		}
+		tok := p.next()
+		if tok.err != nil {
+			return nil, tok.err
+		}
+		k := tok.value
+		if err := p.expect("="); err != nil {
+			return nil, err
+		}
+		tok = p.next()
+		if tok.err != nil {
+			return nil, tok.err
+		}
+		v := tok.value
+		hints[k] = v
+		if !p.eat(",") {
+			break
+		}
+	}
+	if err := p.expect("}"); err != nil {
+		return nil, err
+	}
+	return hints, nil
+}
+
 func (p *parser) parseTableOrIndexOrColumnName() (ID, *parseError) {
 	/*
 		table_name and column_name and index_name:

spanner/spansql/parser_test.go

@@ -114,6 +114,34 @@ func TestParseQuery(t *testing.T) {
 				},
 			},
 		},
+		// with single table hint
+		{`SELECT * FROM Packages@{FORCE_INDEX=PackagesIdx} WHERE package_idx=@packageIdx`,
+			Query{
+				Select: Select{
+					List: []Expr{Star},
+					From: []SelectFrom{SelectFromTable{Table: "Packages", Hints: map[string]string{"FORCE_INDEX": "PackagesIdx"}}},
+					Where: ComparisonOp{
+						Op:  Eq,
+						LHS: ID("package_idx"),
+						RHS: Param("packageIdx"),
+					},
+				},
+			},
+		},
+		// with multiple table hints
+		{`SELECT * FROM Packages@{ FORCE_INDEX=PackagesIdx, GROUPBY_SCAN_OPTIMIZATION=TRUE } WHERE package_idx=@packageIdx`,
+			Query{
+				Select: Select{
+					List: []Expr{Star},
+					From: []SelectFrom{SelectFromTable{Table: "Packages", Hints: map[string]string{"FORCE_INDEX": "PackagesIdx", "GROUPBY_SCAN_OPTIMIZATION": "TRUE"}}},
+					Where: ComparisonOp{
+						Op:  Eq,
+						LHS: ID("package_idx"),
+						RHS: Param("packageIdx"),
+					},
+				},
+			},
+		},
 		{`SELECT * FROM A INNER JOIN B ON A.w = B.y`,
 			Query{
 				Select: Select{

spanner/spansql/sql.go

@@ -365,6 +365,18 @@ func (sel Select) addSQL(sb *strings.Builder) {
 
 func (sft SelectFromTable) SQL() string {
 	str := sft.Table.SQL()
+	if len(sft.Hints) > 0 {
+		str += "@{"
+		kvs := make([]string, len(sft.Hints))
+		i := 0
+		for k, v := range sft.Hints {
+			kvs[i] = fmt.Sprintf("%s=%s", k, v)
+			i++
+		}
+		str += strings.Join(kvs, ",")
+		str += "}"
+	}
+
 	if sft.Alias != "" {
 		str += " AS " + sft.Alias.SQL()
 	}

spanner/spansql/sql_test.go

@@ -309,6 +309,42 @@ func TestSQL(t *testing.T) {
 			`SELECT A, B AS banana FROM Table WHERE C < "whelp" AND D IS NOT NULL ORDER BY OCol DESC LIMIT 1000`,
 			reparseQuery,
 		},
+		{
+			Query{
+				Select: Select{
+					List: []Expr{ID("A")},
+					From: []SelectFrom{SelectFromTable{
+						Table: "Table",
+						Hints: map[string]string{"FORCE_INDEX": "Idx"},
+					}},
+					Where: ComparisonOp{
+						LHS: ID("B"),
+						Op:  Eq,
+						RHS: Param("b"),
+					},
+				},
+			},
+			`SELECT A FROM Table@{FORCE_INDEX=Idx} WHERE B = @b`,
+			reparseQuery,
+		},
+		{
+			Query{
+				Select: Select{
+					List: []Expr{ID("A")},
+					From: []SelectFrom{SelectFromTable{
+						Table: "Table",
+						Hints: map[string]string{"FORCE_INDEX": "Idx", "GROUPBY_SCAN_OPTIMIZATION": "TRUE"},
+					}},
+					Where: ComparisonOp{
+						LHS: ID("B"),
+						Op:  Eq,
+						RHS: Param("b"),
+					},
+				},
+			},
+			`SELECT A FROM Table@{FORCE_INDEX=Idx,GROUPBY_SCAN_OPTIMIZATION=TRUE} WHERE B = @b`,
+			reparseQuery,
+		},
 		{
 			Query{
 				Select: Select{

spanner/spansql/types.go

@@ -394,6 +394,7 @@ type SelectFrom interface {
 type SelectFromTable struct {
 	Table ID
 	Alias ID // empty if not aliased
+	Hints map[string]string
 }
 
 func (SelectFromTable) isSelectFrom() {}