Implement WITH expression (#215)

* Implement WITH expression * Update testdata * Fix doc comment * Fix lookaheadWithExprVar * Fix parseWithExpr() * Update testdata * Do make gen
cloudspannerecosystem · Dec 18, 2024 · 438cf9d · 438cf9d
1 parent f155be3
commit 438cf9d
Show file tree

Hide file tree

Showing 7 changed files with 321 additions and 1 deletion.
diff --git a/ast/ast.go b/ast/ast.go
@@ -172,6 +172,7 @@ func (CallExpr) isExpr()              {}
 func (CountStarExpr) isExpr()         {}
 func (CastExpr) isExpr()              {}
 func (ExtractExpr) isExpr()           {}
+func (WithExpr) isExpr()              {}
 func (ReplaceFieldsExpr) isExpr()     {}
 func (CaseExpr) isExpr()              {}
 func (IfExpr) isExpr()                {}
@@ -1329,6 +1330,31 @@ type AtTimeZone struct {
 	Expr Expr
 }
 
+// WithExprVar is "name AS expr" node in WITH expression.
+//
+//	{{.Name | sql}} AS {{.Expr | sql}}
+type WithExprVar struct {
+	// pos = Name.pos
+	// end = Expr.end
+
+	Name *Ident
+	Expr Expr
+}
+
+// WithExpr is WITH expression node.
+//
+//	WITH({{.Vars | sqlJoin ", "}}, {{.Expr | sql}})
+type WithExpr struct {
+	// pos = With
+	// end = Rparen + 1
+
+	With   token.Pos // position of "WITH" keyword
+	Rparen token.Pos // position of ")"
+
+	Vars []*WithExprVar // len(Vars) > 0
+	Expr Expr
+}
+
 // CastExpr is CAST/SAFE_CAST call expression node.
 //
 //	{{if .Safe}}SAFE_{{end}}CAST({{.Expr | sql}} AS {{.Type | sql}})

diff --git a/ast/pos.go b/ast/pos.go
diff --git a/ast/sql.go b/ast/sql.go
@@ -80,7 +80,7 @@ func exprPrec(e Expr) prec {
 	case *CallExpr, *CountStarExpr, *CastExpr, *ExtractExpr, *CaseExpr, *IfExpr, *ParenExpr, *ScalarSubQuery,
 		*ArraySubQuery, *ExistsSubQuery, *Param, *Ident, *Path, *ArrayLiteral, *TupleStructLiteral, *TypedStructLiteral,
 		*TypelessStructLiteral, *NullLiteral, *BoolLiteral, *IntLiteral, *FloatLiteral, *StringLiteral, *BytesLiteral,
-		*DateLiteral, *TimestampLiteral, *NumericLiteral, *JSONLiteral:
+		*DateLiteral, *TimestampLiteral, *NumericLiteral, *JSONLiteral, *WithExpr:
 		return precLit
 	case *IndexExpr, *SelectorExpr:
 		return precSelector
@@ -565,6 +565,12 @@ func (r *ReplaceFieldsExpr) SQL() string {
 	return "REPLACE_FIELDS(" + r.Expr.SQL() + ", " + sqlJoin(r.Fields, ", ") + ")"
 }
 
+func (n *WithExprVar) SQL() string { return n.Name.SQL() + " AS " + n.Expr.SQL() }
+
+func (w *WithExpr) SQL() string {
+	return "WITH(" + sqlJoin(w.Vars, ", ") + ", " + w.Expr.SQL() + ")"
+}
+
 func (c *CastExpr) SQL() string {
 	return strOpt(c.Safe, "SAFE_") + "CAST(" + c.Expr.SQL() + " AS " + c.Type.SQL() + ")"
 }

diff --git a/parser.go b/parser.go
@@ -1451,6 +1451,8 @@ func (p *Parser) parseLit() ast.Expr {
 		return p.parseExistsSubQuery()
 	case "EXTRACT":
 		return p.parseExtractExpr()
+	case "WITH":
+		return p.parseWithExpr()
 	case "ARRAY":
 		return p.parseArrayLiteralOrSubQuery()
 	case "STRUCT":
@@ -1924,6 +1926,47 @@ func (p *Parser) tryParseAtTimeZone() *ast.AtTimeZone {
 	}
 }
 
+func (p *Parser) parseWithExprVar() *ast.WithExprVar {
+	name := p.parseIdent()
+	p.expect("AS")
+	expr := p.parseExpr()
+
+	return &ast.WithExprVar{
+		Expr: expr,
+		Name: name,
+	}
+}
+
+func (p *Parser) lookaheadWithExprVar() bool {
+	lexer := p.Lexer.Clone()
+	defer func() {
+		p.Lexer = lexer
+	}()
+
+	p.parseIdent()
+	return p.Token.Kind == "AS"
+}
+
+func (p *Parser) parseWithExpr() *ast.WithExpr {
+	with := p.expect("WITH").Pos
+	p.expect("(")
+
+	var vars []*ast.WithExprVar
+	for p.lookaheadWithExprVar() {
+		vars = append(vars, p.parseWithExprVar())
+		p.expect(",")
+	}
+
+	expr := p.parseExpr()
+	rparen := p.expect(")").Pos
+	return &ast.WithExpr{
+		With:   with,
+		Rparen: rparen,
+		Vars:   vars,
+		Expr:   expr,
+	}
+}
+
 func (p *Parser) parseParenExpr() ast.Expr {
 	paren := p.Token
 

diff --git a/testdata/input/query/select_with.sql b/testdata/input/query/select_with.sql
@@ -0,0 +1,5 @@
+-- https://cloud.google.com/spanner/docs/reference/standard-sql/operators#with_expression
+SELECT WITH(a AS '123',       -- a is '123'
+    b AS CONCAT(a, '456'),    -- b is '123456'
+    c AS '789',               -- c is '789'
+    CONCAT(b, c)) AS result   -- b + c is '123456789'
diff --git a/testdata/result/query/select_with.sql.txt b/testdata/result/query/select_with.sql.txt
@@ -0,0 +1,112 @@
+--- select_with.sql
+-- https://cloud.google.com/spanner/docs/reference/standard-sql/operators#with_expression
+SELECT WITH(a AS '123',       -- a is '123'
+    b AS CONCAT(a, '456'),    -- b is '123456'
+    c AS '789',               -- c is '789'
+    CONCAT(b, c)) AS result   -- b + c is '123456789'
+--- AST
+&ast.QueryStatement{
+  Query: &ast.Select{
+    Select:  90,
+    Results: []ast.SelectItem{
+      &ast.Alias{
+        Expr: &ast.WithExpr{
+          With:   97,
+          Rparen: 241,
+          Vars:   []*ast.WithExprVar{
+            &ast.WithExprVar{
+              Name: &ast.Ident{
+                NamePos: 102,
+                NameEnd: 103,
+                Name:    "a",
+              },
+              Expr: &ast.StringLiteral{
+                ValuePos: 107,
+                ValueEnd: 112,
+                Value:    "123",
+              },
+            },
+            &ast.WithExprVar{
+              Name: &ast.Ident{
+                NamePos: 138,
+                NameEnd: 139,
+                Name:    "b",
+              },
+              Expr: &ast.CallExpr{
+                Rparen: 158,
+                Func:   &ast.Ident{
+                  NamePos: 143,
+                  NameEnd: 149,
+                  Name:    "CONCAT",
+                },
+                Args: []ast.Arg{
+                  &ast.ExprArg{
+                    Expr: &ast.Ident{
+                      NamePos: 150,
+                      NameEnd: 151,
+                      Name:    "a",
+                    },
+                  },
+                  &ast.ExprArg{
+                    Expr: &ast.StringLiteral{
+                      ValuePos: 153,
+                      ValueEnd: 158,
+                      Value:    "456",
+                    },
+                  },
+                },
+              },
+            },
+            &ast.WithExprVar{
+              Name: &ast.Ident{
+                NamePos: 185,
+                NameEnd: 186,
+                Name:    "c",
+              },
+              Expr: &ast.StringLiteral{
+                ValuePos: 190,
+                ValueEnd: 195,
+                Value:    "789",
+              },
+            },
+          },
+          Expr: &ast.CallExpr{
+            Rparen: 240,
+            Func:   &ast.Ident{
+              NamePos: 229,
+              NameEnd: 235,
+              Name:    "CONCAT",
+            },
+            Args: []ast.Arg{
+              &ast.ExprArg{
+                Expr: &ast.Ident{
+                  NamePos: 236,
+                  NameEnd: 237,
+                  Name:    "b",
+                },
+              },
+              &ast.ExprArg{
+                Expr: &ast.Ident{
+                  NamePos: 239,
+                  NameEnd: 240,
+                  Name:    "c",
+                },
+              },
+            },
+          },
+        },
+        As: &ast.AsAlias{
+          As:    243,
+          Alias: &ast.Ident{
+            NamePos: 246,
+            NameEnd: 252,
+            Name:    "result",
+          },
+        },
+      },
+    },
+  },
+}
+
+--- SQL
+SELECT WITH(a AS "123", b AS CONCAT(a, "456"), c AS "789", CONCAT(b, c)) AS result
diff --git a/testdata/result/statement/select_with.sql.txt b/testdata/result/statement/select_with.sql.txt
@@ -0,0 +1,112 @@
+--- select_with.sql
+-- https://cloud.google.com/spanner/docs/reference/standard-sql/operators#with_expression
+SELECT WITH(a AS '123',       -- a is '123'
+    b AS CONCAT(a, '456'),    -- b is '123456'
+    c AS '789',               -- c is '789'
+    CONCAT(b, c)) AS result   -- b + c is '123456789'
+--- AST
+&ast.QueryStatement{
+  Query: &ast.Select{
+    Select:  90,
+    Results: []ast.SelectItem{
+      &ast.Alias{
+        Expr: &ast.WithExpr{
+          With:   97,
+          Rparen: 241,
+          Vars:   []*ast.WithExprVar{
+            &ast.WithExprVar{
+              Name: &ast.Ident{
+                NamePos: 102,
+                NameEnd: 103,
+                Name:    "a",
+              },
+              Expr: &ast.StringLiteral{
+                ValuePos: 107,
+                ValueEnd: 112,
+                Value:    "123",
+              },
+            },
+            &ast.WithExprVar{
+              Name: &ast.Ident{
+                NamePos: 138,
+                NameEnd: 139,
+                Name:    "b",
+              },
+              Expr: &ast.CallExpr{
+                Rparen: 158,
+                Func:   &ast.Ident{
+                  NamePos: 143,
+                  NameEnd: 149,
+                  Name:    "CONCAT",
+                },
+                Args: []ast.Arg{
+                  &ast.ExprArg{
+                    Expr: &ast.Ident{
+                      NamePos: 150,
+                      NameEnd: 151,
+                      Name:    "a",
+                    },
+                  },
+                  &ast.ExprArg{
+                    Expr: &ast.StringLiteral{
+                      ValuePos: 153,
+                      ValueEnd: 158,
+                      Value:    "456",
+                    },
+                  },
+                },
+              },
+            },
+            &ast.WithExprVar{
+              Name: &ast.Ident{
+                NamePos: 185,
+                NameEnd: 186,
+                Name:    "c",
+              },
+              Expr: &ast.StringLiteral{
+                ValuePos: 190,
+                ValueEnd: 195,
+                Value:    "789",
+              },
+            },
+          },
+          Expr: &ast.CallExpr{
+            Rparen: 240,
+            Func:   &ast.Ident{
+              NamePos: 229,
+              NameEnd: 235,
+              Name:    "CONCAT",
+            },
+            Args: []ast.Arg{
+              &ast.ExprArg{
+                Expr: &ast.Ident{
+                  NamePos: 236,
+                  NameEnd: 237,
+                  Name:    "b",
+                },
+              },
+              &ast.ExprArg{
+                Expr: &ast.Ident{
+                  NamePos: 239,
+                  NameEnd: 240,
+                  Name:    "c",
+                },
+              },
+            },
+          },
+        },
+        As: &ast.AsAlias{
+          As:    243,
+          Alias: &ast.Ident{
+            NamePos: 246,
+            NameEnd: 252,
+            Name:    "result",
+          },
+        },
+      },
+    },
+  },
+}
+
+--- SQL
+SELECT WITH(a AS "123", b AS CONCAT(a, "456"), c AS "789", CONCAT(b, c)) AS result