From 27e00727d869732579f86ebd62f2547f47fb39c7 Mon Sep 17 00:00:00 2001
From: neglect-yp <grandis.hebitombo@gmail.com>
Date: Fri, 15 Jul 2022 00:26:27 +0900
Subject: [PATCH] feat(spanner/spansql): add support for if expressions

---
 spanner/spansql/parser.go      | 35 ++++++++++++++++++++++++++++++++++
 spanner/spansql/parser_test.go |  7 +++++++
 spanner/spansql/sql.go         | 11 +++++++++++
 spanner/spansql/sql_test.go    | 15 +++++++++++++++
 spanner/spansql/types.go       |  9 +++++++++
 5 files changed, 77 insertions(+)

diff --git a/spanner/spansql/parser.go b/spanner/spansql/parser.go
index 2f2d066a0c11..8d12c481c88d 100644
--- a/spanner/spansql/parser.go
+++ b/spanner/spansql/parser.go
@@ -3007,6 +3007,9 @@ func (p *parser) parseLit() (Expr, *parseError) {
 	case tok.caseEqual("CASE"):
 		p.back()
 		return p.parseCaseExpr()
+	case tok.caseEqual("IF"):
+		p.back()
+		return p.parseIfExpr()
 	}
 
 	// Handle typed literals.
@@ -3117,6 +3120,38 @@ func (p *parser) parseWhenClause() (WhenClause, *parseError) {
 	return WhenClause{Cond: cond, Result: result}, nil
 }
 
+func (p *parser) parseIfExpr() (If, *parseError) {
+	if err := p.expect("IF", "("); err != nil {
+		return If{}, err
+	}
+
+	expr, err := p.parseBoolExpr()
+	if err != nil {
+		return If{}, err
+	}
+	if err := p.expect(","); err != nil {
+		return If{}, err
+	}
+
+	trueResult, err := p.parseExpr()
+	if err != nil {
+		return If{}, err
+	}
+	if err := p.expect(","); err != nil {
+		return If{}, err
+	}
+
+	elseResult, err := p.parseExpr()
+	if err != nil {
+		return If{}, err
+	}
+	if err := p.expect(")"); err != nil {
+		return If{}, err
+	}
+
+	return If{Expr: expr, TrueResult: trueResult, ElseResult: elseResult}, nil
+}
+
 func (p *parser) parseArrayLit() (Array, *parseError) {
 	// ARRAY keyword is optional.
 	// TODO: If it is present, consume any <T> after it.
diff --git a/spanner/spansql/parser_test.go b/spanner/spansql/parser_test.go
index 83f827bd7ae3..94342bbd3740 100644
--- a/spanner/spansql/parser_test.go
+++ b/spanner/spansql/parser_test.go
@@ -411,6 +411,13 @@ func TestParseExpr(t *testing.T) {
 				},
 			},
 		},
+		{`IF(A < B, TRUE, FALSE)`,
+			If{
+				Expr:       ComparisonOp{LHS: ID("A"), Op: Lt, RHS: ID("B")},
+				TrueResult: True,
+				ElseResult: False,
+			},
+		},
 
 		// String literal:
 		// Accept double quote and single quote.
diff --git a/spanner/spansql/sql.go b/spanner/spansql/sql.go
index 1f92deee2458..a45cda4bdb32 100644
--- a/spanner/spansql/sql.go
+++ b/spanner/spansql/sql.go
@@ -733,6 +733,17 @@ func (c Case) addSQL(sb *strings.Builder) {
 	sb.WriteString("END")
 }
 
+func (i If) SQL() string { return buildSQL(i) }
+func (i If) addSQL(sb *strings.Builder) {
+	sb.WriteString("IF(")
+	i.Expr.addSQL(sb)
+	sb.WriteString(", ")
+	i.TrueResult.addSQL(sb)
+	sb.WriteString(", ")
+	i.ElseResult.addSQL(sb)
+	sb.WriteString(")")
+}
+
 func (b BoolLiteral) SQL() string { return buildSQL(b) }
 func (b BoolLiteral) addSQL(sb *strings.Builder) {
 	if b {
diff --git a/spanner/spansql/sql_test.go b/spanner/spansql/sql_test.go
index febbbda63daf..6e51263b65b3 100644
--- a/spanner/spansql/sql_test.go
+++ b/spanner/spansql/sql_test.go
@@ -656,6 +656,21 @@ func TestSQL(t *testing.T) {
 			`SELECT CASE WHEN TRUE THEN "X" WHEN FALSE THEN "Y" END`,
 			reparseQuery,
 		},
+		{
+			Query{
+				Select: Select{
+					List: []Expr{
+						If{
+							Expr:       ComparisonOp{LHS: IntegerLiteral(1), Op: Lt, RHS: IntegerLiteral(2)},
+							TrueResult: True,
+							ElseResult: False,
+						},
+					},
+				},
+			},
+			`SELECT IF(1 < 2, TRUE, FALSE)`,
+			reparseQuery,
+		},
 	}
 	for _, test := range tests {
 		sql := test.data.SQL()
diff --git a/spanner/spansql/types.go b/spanner/spansql/types.go
index a9b335eee43f..b9e44df1cc75 100644
--- a/spanner/spansql/types.go
+++ b/spanner/spansql/types.go
@@ -738,6 +738,15 @@ type WhenClause struct {
 	Result Expr
 }
 
+type If struct {
+	Expr       Expr
+	TrueResult Expr
+	ElseResult Expr
+}
+
+func (If) isBoolExpr() {} // possibly bool
+func (If) isExpr()     {}
+
 type BoolLiteral bool
 
 const (