From dec7a67a3e980f6f5e0d170919da87e1bffe923f Mon Sep 17 00:00:00 2001
From: Takeshi Nakata <takeshi.nakata@mixi.co.jp>
Date: Fri, 23 Jul 2021 23:13:25 +0900
Subject: [PATCH] fix(spanner/spansql): fix DATE and TIMESTAMP parsing. (#4480)

* fix(spanner/spansql): fix DATE, TIMESTAMP parse.

This enables to parse below cases.
- DATE, TIMESTAMP used as identifier.(these words does not reserved
  keyword.)
- DATE, TIMESTAMP used as function.

* fix(spanner/spansql): fix allFuncs index, and add more Date, Timestamp
functions

* fix(spanner/spansql): fix parse DATE, Timestamp literals.

This modified parsing DATE, TIMESTAMP that if next token is string
litelal, then current token treated as DATE, TIMESTAMP literals.
---
 spanner/spansql/keywords.go    | 31 ++++++++++++++++++++++++
 spanner/spansql/parser.go      | 23 ++++++++++++++----
 spanner/spansql/parser_test.go | 44 ++++++++++++++++++++++++++++++++++
 3 files changed, 94 insertions(+), 4 deletions(-)

diff --git a/spanner/spansql/keywords.go b/spanner/spansql/keywords.go
index 9bba9dc4005..927d0ac7763 100644
--- a/spanner/spansql/keywords.go
+++ b/spanner/spansql/keywords.go
@@ -190,4 +190,35 @@ var allFuncs = []string{
 	"ARRAY_REVERSE",
 	"ARRAY_IS_DISTINCT",
 	"SAFE_OFFSET", "SAFE_ORDINAL",
+
+	// Date functions.
+	"CURRENT_DATE",
+	"EXTRACT",
+	"DATE",
+	"DATE_ADD",
+	"DATE_SUB",
+	"DATE_DIFF",
+	"DATE_TRUNC",
+	"DATE_FROM_UNIX_DATE",
+	"FORMAT_DATE",
+	"PARSE_DATE",
+	"UNIX_DATE",
+
+	// Timestamp functions.
+	"CURRENT_TIMESTAMP",
+	"STRING",
+	"TIMESTAMP",
+	"TIMESTAMP_ADD",
+	"TIMESTAMP_SUB",
+	"TIMESTAMP_DIFF",
+	"TIMESTAMP_TRUNC",
+	"FORMAT_TIMESTAMP",
+	"PARSE_TIMESTAMP",
+	"TIMESTAMP_SECONDS",
+	"TIMESTAMP_MILLIS",
+	"TIMESTAMP_MICROS",
+	"UNIX_SECONDS",
+	"UNIX_MILLIS",
+	"UNIX_MICROS",
+	"PENDING_COMMIT_TIMESTAMP",
 }
diff --git a/spanner/spansql/parser.go b/spanner/spansql/parser.go
index ab3220271d5..a575d781069 100644
--- a/spanner/spansql/parser.go
+++ b/spanner/spansql/parser.go
@@ -906,6 +906,17 @@ func (p *parser) sniff(want ...string) bool {
 	return true
 }
 
+// sniffTokenType reports whether the next token type is as specified.
+func (p *parser) sniffTokenType(want tokenType) bool {
+	orig := *p
+	defer func() { *p = orig }()
+
+	if p.next().typ == want {
+		return true
+	}
+	return false
+}
+
 // eat reports whether the next N tokens are as specified,
 // then consumes them.
 func (p *parser) eat(want ...string) bool {
@@ -2714,11 +2725,15 @@ func (p *parser) parseLit() (Expr, *parseError) {
 		p.back()
 		return p.parseArrayLit()
 	case tok.caseEqual("DATE"):
-		p.back()
-		return p.parseDateLit()
+		if p.sniffTokenType(stringToken) {
+			p.back()
+			return p.parseDateLit()
+		}
 	case tok.caseEqual("TIMESTAMP"):
-		p.back()
-		return p.parseTimestampLit()
+		if p.sniffTokenType(stringToken) {
+			p.back()
+			return p.parseTimestampLit()
+		}
 	}
 
 	// TODO: struct literals
diff --git a/spanner/spansql/parser_test.go b/spanner/spansql/parser_test.go
index fd1360c8506..5bbe8276c5d 100644
--- a/spanner/spansql/parser_test.go
+++ b/spanner/spansql/parser_test.go
@@ -83,6 +83,35 @@ func TestParseQuery(t *testing.T) {
 				},
 			},
 		},
+		{`SELECT date, timestamp as timestamp FROM Packages WHERE date = DATE '2014-09-27' AND timestamp = TIMESTAMP '2014-09-27 12:30:00'`,
+			Query{
+				Select: Select{
+					List: []Expr{ID("date"), ID("timestamp")},
+					From: []SelectFrom{SelectFromTable{Table: "Packages"}},
+					Where: LogicalOp{
+						Op: And,
+						LHS: ComparisonOp{
+							Op:  Eq,
+							LHS: ID("date"),
+							RHS: DateLiteral{Year: 2014, Month: 9, Day: 27},
+						},
+						RHS: ComparisonOp{
+							Op:  Eq,
+							LHS: ID("timestamp"),
+							RHS: TimestampLiteral(timef(t, "2006-01-02 15:04:05", "2014-09-27 12:30:00")),
+						},
+					},
+					ListAliases: []ID{"", "timestamp"},
+				},
+			},
+		},
+		{`SELECT UNIX_DATE(DATE "2008-12-25")`,
+			Query{
+				Select: Select{
+					List: []Expr{Func{Name: "UNIX_DATE", Args: []Expr{DateLiteral{Year: 2008, Month: 12, Day: 25}}}},
+				},
+			},
+		},
 		{`SELECT SUM(PointsScored) AS total_points, FirstName, LastName AS surname FROM PlayerStats GROUP BY FirstName, LastName`,
 			Query{
 				Select: Select{
@@ -336,7 +365,14 @@ func TestParseExpr(t *testing.T) {
 
 		// Date and timestamp literals:
 		{`DATE '2014-09-27'`, DateLiteral(civil.Date{Year: 2014, Month: time.September, Day: 27})},
+		{`TIMESTAMP '2014-09-27 12:30:00'`, TimestampLiteral(timef(t, "2006-01-02 15:04:05", "2014-09-27 12:30:00"))},
 
+		// date and timestamp funclit
+		{`DATE('2014-09-27')`, Func{Name: "DATE", Args: []Expr{StringLiteral("2014-09-27")}}},
+		{`TIMESTAMP('2014-09-27 12:30:00')`, Func{Name: "TIMESTAMP", Args: []Expr{StringLiteral("2014-09-27 12:30:00")}}},
+		// date and timestamp identifier
+		{`DATE = '2014-09-27'`, ComparisonOp{LHS: ID("DATE"), Op: Eq, RHS: StringLiteral("2014-09-27")}},
+		{`TIMESTAMP = '2014-09-27 12:30:00'`, ComparisonOp{LHS: ID("TIMESTAMP"), Op: Eq, RHS: StringLiteral("2014-09-27 12:30:00")}},
 		// Array literals:
 		// https://cloud.google.com/spanner/docs/lexical#array_literals
 		{`[1, 2, 3]`, Array{IntegerLiteral(1), IntegerLiteral(2), IntegerLiteral(3)}},
@@ -796,3 +832,11 @@ func TestParseFailures(t *testing.T) {
 		}
 	}
 }
+
+func timef(t *testing.T, format, s string) time.Time {
+	ti, err := time.ParseInLocation(format, string(s), defaultLocation)
+	if err != nil {
+		t.Errorf("parsing %s [%s] time.ParseInLocation failed.", s, format)
+	}
+	return ti
+}