From 196193034a15f84dc3d3c27901990e8be77fca85 Mon Sep 17 00:00:00 2001
From: David Symonds <dsymonds@golang.org>
Date: Thu, 14 Jan 2021 11:16:02 +1100
Subject: [PATCH] feat(spanner/spansql): support DATE and TIMESTAMP literals
 (#3557)

This parses a reasonable subset, and renders in a reasonable format.
Future work will be required to support this in spannertest, as well as
correctly supporting the full range of valid layouts.

Fixes #3548.
---
 spanner/spannertest/README.md  |  2 +-
 spanner/spansql/parser.go      | 88 ++++++++++++++++++++++++++++++++--
 spanner/spansql/parser_test.go |  6 +++
 spanner/spansql/sql.go         | 11 +++++
 spanner/spansql/sql_test.go    | 18 +++++++
 spanner/spansql/types.go       | 15 ++++++
 6 files changed, 135 insertions(+), 5 deletions(-)

diff --git a/spanner/spannertest/README.md b/spanner/spannertest/README.md
index bee71528e70..14caba869b5 100644
--- a/spanner/spannertest/README.md
+++ b/spanner/spannertest/README.md
@@ -22,7 +22,7 @@ by ascending esotericism:
 - more aggregation functions
 - SELECT HAVING
 - case insensitivity
-- alternate literal types (esp. strings)
+- more literal types
 - generated columns
 - expression type casting, coercion
 - multiple joins
diff --git a/spanner/spansql/parser.go b/spanner/spansql/parser.go
index ed93afb4d4e..cbd2ae003aa 100644
--- a/spanner/spansql/parser.go
+++ b/spanner/spansql/parser.go
@@ -47,7 +47,10 @@ import (
 	"os"
 	"strconv"
 	"strings"
+	"time"
 	"unicode/utf8"
+
+	"cloud.google.com/go/civil"
 )
 
 const debug = false
@@ -780,7 +783,7 @@ func (p *parser) advance() {
 	p.cur.err = nil
 	p.cur.line, p.cur.offset = p.line, p.offset
 	p.cur.typ = unknownToken
-	// TODO: struct, date, timestamp literals
+	// TODO: struct literals
 	switch p.s[0] {
 	case ',', ';', '(', ')', '{', '}', '[', ']', '*', '+', '-':
 		// Single character symbol.
@@ -2601,13 +2604,20 @@ func (p *parser) parseLit() (Expr, *parseError) {
 		// case insensitivity for keywords.
 	}
 
-	// Handle array literals.
-	if tok.value == "ARRAY" || tok.value == "[" {
+	// Handle typed literals.
+	switch tok.value {
+	case "ARRAY", "[":
 		p.back()
 		return p.parseArrayLit()
+	case "DATE":
+		p.back()
+		return p.parseDateLit()
+	case "TIMESTAMP":
+		p.back()
+		return p.parseTimestampLit()
 	}
 
-	// TODO: more types of literals (struct, date, timestamp).
+	// TODO: struct literals
 
 	// Try a parameter.
 	// TODO: check character sets.
@@ -2645,6 +2655,76 @@ func (p *parser) parseArrayLit() (Array, *parseError) {
 	return arr, err
 }
 
+// TODO: There should be exported Parse{Date,Timestamp}Literal package-level funcs
+// to support spannertest coercing plain string literals when used in a typed context.
+// Those should wrap parseDateLit and parseTimestampLit below.
+
+func (p *parser) parseDateLit() (DateLiteral, *parseError) {
+	if err := p.expect("DATE"); err != nil {
+		return DateLiteral{}, err
+	}
+	s, err := p.parseStringLit()
+	if err != nil {
+		return DateLiteral{}, err
+	}
+	d, perr := civil.ParseDate(string(s))
+	if perr != nil {
+		return DateLiteral{}, p.errorf("bad date literal %q: %v", s, perr)
+	}
+	// TODO: Enforce valid range.
+	return DateLiteral(d), nil
+}
+
+// TODO: A manual parser is probably better than this.
+// There are a lot of variations that this does not handle.
+var timestampFormats = []string{
+	// 'YYYY-[M]M-[D]D [[H]H:[M]M:[S]S[.DDDDDD] [timezone]]'
+	"2006-01-02",
+	"2006-01-02 15:04:05",
+	"2006-01-02 15:04:05.000000",
+	"2006-01-02 15:04:05 -07:00",
+	"2006-01-02 15:04:05.000000 -07:00",
+}
+
+var defaultLocation = func() *time.Location {
+	// The docs say "America/Los_Angeles" is the default.
+	// Use that if we can load it, but fall back on UTC if we don't have timezone data.
+	loc, err := time.LoadLocation("America/Los_Angeles")
+	if err == nil {
+		return loc
+	}
+	return time.UTC
+}()
+
+func (p *parser) parseTimestampLit() (TimestampLiteral, *parseError) {
+	if err := p.expect("TIMESTAMP"); err != nil {
+		return TimestampLiteral{}, err
+	}
+	s, err := p.parseStringLit()
+	if err != nil {
+		return TimestampLiteral{}, err
+	}
+	for _, format := range timestampFormats {
+		t, err := time.ParseInLocation(format, string(s), defaultLocation)
+		if err == nil {
+			// TODO: Enforce valid range.
+			return TimestampLiteral(t), nil
+		}
+	}
+	return TimestampLiteral{}, p.errorf("invalid timestamp literal %q", s)
+}
+
+func (p *parser) parseStringLit() (StringLiteral, *parseError) {
+	tok := p.next()
+	if tok.err != nil {
+		return "", tok.err
+	}
+	if tok.typ != stringToken {
+		return "", p.errorf("got %q, want string literal", tok.value)
+	}
+	return StringLiteral(tok.string), nil
+}
+
 func (p *parser) parsePathExp() (PathExp, *parseError) {
 	var pe PathExp
 	for {
diff --git a/spanner/spansql/parser_test.go b/spanner/spansql/parser_test.go
index ae5577d5844..f4c425c175f 100644
--- a/spanner/spansql/parser_test.go
+++ b/spanner/spansql/parser_test.go
@@ -21,6 +21,9 @@ import (
 	"math"
 	"reflect"
 	"testing"
+	"time"
+
+	"cloud.google.com/go/civil"
 )
 
 func TestParseQuery(t *testing.T) {
@@ -302,6 +305,9 @@ func TestParseExpr(t *testing.T) {
 		{`RB"""\\//\\//"""`, BytesLiteral("\\\\//\\\\//")},
 		{"RB'''\\\\//\n\\\\//'''", BytesLiteral("\\\\//\n\\\\//")},
 
+		// Date and timestamp literals:
+		{`DATE '2014-09-27'`, DateLiteral(civil.Date{Year: 2014, Month: time.September, Day: 27})},
+
 		// Array literals:
 		// https://cloud.google.com/spanner/docs/lexical#array_literals
 		{`[1, 2, 3]`, Array{IntegerLiteral(1), IntegerLiteral(2), IntegerLiteral(3)}},
diff --git a/spanner/spansql/sql.go b/spanner/spansql/sql.go
index 99160c83d91..99f7b7d1335 100644
--- a/spanner/spansql/sql.go
+++ b/spanner/spansql/sql.go
@@ -27,6 +27,7 @@ import (
 	"fmt"
 	"strconv"
 	"strings"
+	"time"
 )
 
 func buildSQL(x interface{ addSQL(*strings.Builder) }) string {
@@ -585,3 +586,13 @@ func (sl StringLiteral) addSQL(sb *strings.Builder) { fmt.Fprintf(sb, "%q", sl)
 
 func (bl BytesLiteral) SQL() string                { return buildSQL(bl) }
 func (bl BytesLiteral) addSQL(sb *strings.Builder) { fmt.Fprintf(sb, "B%q", bl) }
+
+func (dl DateLiteral) SQL() string { return buildSQL(dl) }
+func (dl DateLiteral) addSQL(sb *strings.Builder) {
+	fmt.Fprintf(sb, "DATE '%04d-%02d-%02d'", dl.Year, dl.Month, dl.Day)
+}
+
+func (tl TimestampLiteral) SQL() string { return buildSQL(tl) }
+func (tl TimestampLiteral) addSQL(sb *strings.Builder) {
+	fmt.Fprintf(sb, "TIMESTAMP '%s'", time.Time(tl).Format("2006-01-02 15:04:05.000000 -07:00"))
+}
diff --git a/spanner/spansql/sql_test.go b/spanner/spansql/sql_test.go
index d50089d63e1..1f9db7db597 100644
--- a/spanner/spansql/sql_test.go
+++ b/spanner/spansql/sql_test.go
@@ -19,6 +19,9 @@ package spansql
 import (
 	"reflect"
 	"testing"
+	"time"
+
+	"cloud.google.com/go/civil"
 )
 
 func boolAddr(b bool) *bool {
@@ -53,6 +56,11 @@ func TestSQL(t *testing.T) {
 		return e, nil
 	}
 
+	latz, err := time.LoadLocation("America/Los_Angeles")
+	if err != nil {
+		t.Fatalf("Loading Los Angeles time zone info: %v", err)
+	}
+
 	line := func(n int) Position { return Position{Line: n} }
 	tests := []struct {
 		data    interface{ SQL() string }
@@ -300,6 +308,16 @@ func TestSQL(t *testing.T) {
 			"SELECT `Desc`",
 			reparseQuery,
 		},
+		{
+			DateLiteral(civil.Date{Year: 2014, Month: time.September, Day: 27}),
+			`DATE '2014-09-27'`,
+			reparseExpr,
+		},
+		{
+			TimestampLiteral(time.Date(2014, time.September, 27, 12, 34, 56, 123456e3, latz)),
+			`TIMESTAMP '2014-09-27 12:34:56.123456 -07:00'`,
+			reparseExpr,
+		},
 	}
 	for _, test := range tests {
 		sql := test.data.SQL()
diff --git a/spanner/spansql/types.go b/spanner/spansql/types.go
index 940df82e4b6..70ea1889ef4 100644
--- a/spanner/spansql/types.go
+++ b/spanner/spansql/types.go
@@ -23,6 +23,9 @@ import (
 	"math"
 	"sort"
 	"strings"
+	"time"
+
+	"cloud.google.com/go/civil"
 )
 
 // TODO: More Position fields throughout; maybe in Query/Select.
@@ -630,6 +633,18 @@ type BytesLiteral string
 
 func (BytesLiteral) isExpr() {}
 
+// DateLiteral represents a date literal.
+// https://cloud.google.com/spanner/docs/lexical#date_literals
+type DateLiteral civil.Date
+
+func (DateLiteral) isExpr() {}
+
+// TimestampLiteral represents a timestamp literal.
+// https://cloud.google.com/spanner/docs/lexical#timestamp_literals
+type TimestampLiteral time.Time
+
+func (TimestampLiteral) isExpr() {}
+
 type StarExpr int
 
 // Star represents a "*" in an expression.