feat(spanner/spansql): add support for interval arg of some date/time…

…stamp functions (#6950) * feat(spanner/spansql): add support for interval arg of date/timestamp functions * feat(spanner/spansql): add support for interval arg of GENERATE_DATE_ARRAY Co-authored-by: rahul2393 <irahul@google.com>
googleapis · Nov 3, 2022 · 1ce0f7d · 1ce0f7d
1 parent 5b14658
commit 1ce0f7d
Show file tree

Hide file tree

Showing 5 changed files with 92 additions and 0 deletions.
diff --git a/spanner/spansql/keywords.go b/spanner/spansql/keywords.go
@@ -138,6 +138,13 @@ func init() {
 	funcArgParsers["CAST"] = typedArgParser
 	funcArgParsers["SAFE_CAST"] = typedArgParser
 	funcArgParsers["EXTRACT"] = extractArgParser
+	// Spacial case of INTERVAL arg for DATE_ADD, DATE_SUB, GENERATE_DATE_ARRAY
+	funcArgParsers["DATE_ADD"] = dateIntervalArgParser
+	funcArgParsers["DATE_SUB"] = dateIntervalArgParser
+	funcArgParsers["GENERATE_DATE_ARRAY"] = dateIntervalArgParser
+	// Spacial case of INTERVAL arg for TIMESTAMP_ADD, TIMESTAMP_SUB
+	funcArgParsers["TIMESTAMP_ADD"] = timestampIntervalArgParser
+	funcArgParsers["TIMESTAMP_SUB"] = timestampIntervalArgParser
 }
 
 var allFuncs = []string{

diff --git a/spanner/spansql/parser.go b/spanner/spansql/parser.go
@@ -2818,6 +2818,69 @@ var extractArgParser = func(p *parser) (Expr, *parseError) {
 	}, nil
 }
 
+var intervalArgParser = func(parseDatePart func(*parser) (string, *parseError)) func(*parser) (Expr, *parseError) {
+	return func(p *parser) (Expr, *parseError) {
+		if p.eat("INTERVAL") {
+			expr, err := p.parseExpr()
+			if err != nil {
+				return nil, err
+			}
+			datePart, err := parseDatePart(p)
+			if err != nil {
+				return nil, err
+			}
+			return IntervalExpr{Expr: expr, DatePart: datePart}, nil
+		}
+		return p.parseExpr()
+	}
+}
+
+var dateIntervalDateParts map[string]bool = map[string]bool{
+	"DAY":     true,
+	"WEEK":    true,
+	"MONTH":   true,
+	"QUARTER": true,
+	"YEAR":    true,
+}
+
+func (p *parser) parseDateIntervalDatePart() (string, *parseError) {
+	tok := p.next()
+	if tok.err != nil {
+		return "", tok.err
+	}
+	if dateIntervalDateParts[strings.ToUpper(tok.value)] {
+		return strings.ToUpper(tok.value), nil
+	}
+	return "", p.errorf("got %q, want valid date part names", tok.value)
+}
+
+var timestampIntervalDateParts map[string]bool = map[string]bool{
+	"NANOSECOND":  true,
+	"MICROSECOND": true,
+	"MILLISECOND": true,
+	"SECOND":      true,
+	"MINUTE":      true,
+	"HOUR":        true,
+	"DAY":         true,
+}
+
+func (p *parser) parseTimestampIntervalDatePart() (string, *parseError) {
+	tok := p.next()
+	if tok.err != nil {
+		return "", tok.err
+	}
+	if timestampIntervalDateParts[strings.ToUpper(tok.value)] {
+		return strings.ToUpper(tok.value), nil
+	}
+	return "", p.errorf("got %q, want valid date part names", tok.value)
+}
+
+// Special argument parser for DATE_ADD, DATE_SUB
+var dateIntervalArgParser = intervalArgParser((*parser).parseDateIntervalDatePart)
+
+// Special argument parser for TIMESTAMP_ADD, TIMESTAMP_SUB
+var timestampIntervalArgParser = intervalArgParser((*parser).parseTimestampIntervalDatePart)
+
 /*
 Expressions
 

diff --git a/spanner/spansql/parser_test.go b/spanner/spansql/parser_test.go
@@ -411,6 +411,11 @@ func TestParseExpr(t *testing.T) {
 		{`SAFE_CAST(Bar AS INT64)`, Func{Name: "SAFE_CAST", Args: []Expr{TypedExpr{Expr: ID("Bar"), Type: Type{Base: Int64}}}}},
 		{`EXTRACT(DATE FROM TIMESTAMP AT TIME ZONE "America/Los_Angeles")`, Func{Name: "EXTRACT", Args: []Expr{ExtractExpr{Part: "DATE", Type: Type{Base: Date}, Expr: AtTimeZoneExpr{Expr: ID("TIMESTAMP"), Zone: "America/Los_Angeles", Type: Type{Base: Timestamp}}}}}},
 		{`EXTRACT(DAY FROM DATE)`, Func{Name: "EXTRACT", Args: []Expr{ExtractExpr{Part: "DAY", Expr: ID("DATE"), Type: Type{Base: Int64}}}}},
+		{`DATE_ADD(CURRENT_DATE(), INTERVAL 1 DAY)`, Func{Name: "DATE_ADD", Args: []Expr{Func{Name: "CURRENT_DATE"}, IntervalExpr{Expr: IntegerLiteral(1), DatePart: "DAY"}}}},
+		{`DATE_SUB(CURRENT_DATE(), INTERVAL 1 WEEK)`, Func{Name: "DATE_SUB", Args: []Expr{Func{Name: "CURRENT_DATE"}, IntervalExpr{Expr: IntegerLiteral(1), DatePart: "WEEK"}}}},
+		{`GENERATE_DATE_ARRAY('2022-01-01', CURRENT_DATE(), INTERVAL 1 MONTH)`, Func{Name: "GENERATE_DATE_ARRAY", Args: []Expr{StringLiteral("2022-01-01"), Func{Name: "CURRENT_DATE"}, IntervalExpr{Expr: IntegerLiteral(1), DatePart: "MONTH"}}}},
+		{`TIMESTAMP_ADD(CURRENT_TIMESTAMP(), INTERVAL 1 HOUR)`, Func{Name: "TIMESTAMP_ADD", Args: []Expr{Func{Name: "CURRENT_TIMESTAMP"}, IntervalExpr{Expr: IntegerLiteral(1), DatePart: "HOUR"}}}},
+		{`TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL 1 MINUTE)`, Func{Name: "TIMESTAMP_SUB", Args: []Expr{Func{Name: "CURRENT_TIMESTAMP"}, IntervalExpr{Expr: IntegerLiteral(1), DatePart: "MINUTE"}}}},
 
 		// Conditional expressions
 		{

diff --git a/spanner/spansql/sql.go b/spanner/spansql/sql.go
@@ -694,6 +694,15 @@ func (aze AtTimeZoneExpr) addSQL(sb *strings.Builder) {
 	sb.WriteString(aze.Zone)
 }
 
+func (ie IntervalExpr) SQL() string { return buildSQL(ie) }
+func (ie IntervalExpr) addSQL(sb *strings.Builder) {
+	sb.WriteString("INTERVAL")
+	sb.WriteString(" ")
+	ie.Expr.addSQL(sb)
+	sb.WriteString(" ")
+	sb.WriteString(ie.DatePart)
+}
+
 func idList(l []ID, join string) string {
 	var ss []string
 	for _, s := range l {

diff --git a/spanner/spansql/types.go b/spanner/spansql/types.go
@@ -702,6 +702,14 @@ type AtTimeZoneExpr struct {
 func (AtTimeZoneExpr) isBoolExpr() {} // possibly bool
 func (AtTimeZoneExpr) isExpr()     {}
 
+type IntervalExpr struct {
+	Expr     Expr
+	DatePart string
+}
+
+func (IntervalExpr) isBoolExpr() {} // possibly bool
+func (IntervalExpr) isExpr()     {}
+
 // Paren represents a parenthesised expression.
 type Paren struct {
 	Expr Expr