Repository: cloudspannerecosystem/memefish (go)

Issue:
Support REPLACE_FIELDS
Protocol buffer function `REPLACE_FIELDS` can't be parsed as a normal function, so it is needed to be implemented as an `ast.Expr`.

## Syntax

https://cloud.google.com/spanner/docs/reference/standard-sql/protocol_buffer_functions#replace_fields

```
REPLACE_FIELDS(proto_expression, value AS field_path [, ... ])
```

## ZetaSQL Refs:
https://github.com/google/zetasql/blob/194cd32b5d766d60e3ca442651d792c7fe54ea74/zetasql/parser/bison_parser.y#L9206-L9233
https://github.com/google/zetasql/blob/194cd32b5d766d60e3ca442651d792c7fe54ea74/zetasql/parser/gen_parse_tree.py#L7103-L7135


Pull request:
Implement REPLACE_FIELDS
This PR implements `REPLACE_FIELDS` expression

- close #124 

Patch:
diff --git a/ast/ast.go b/ast/ast.go
index 1de8aac..a51cf9e 100644
--- 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 (ReplaceFieldsExpr) isExpr()     {}
 func (CaseExpr) isExpr()              {}
 func (IfExpr) isExpr()                {}
 func (ParenExpr) isExpr()             {}
@@ -1273,6 +1274,31 @@ type ExtractExpr struct {
 	AtTimeZone *AtTimeZone // optional
 }
 
+// ReplaceFieldsArg is value AS field_path node in ReplaceFieldsExpr.
+//
+//	{{.Expr | sql}} AS {{.Field | sql}}
+type ReplaceFieldsArg struct {
+	// pos = Expr.pos
+	// end = Field.end
+
+	Expr  Expr
+	Field *Path
+}
+
+// ReplaceFieldsExpr is REPLACE_FIELDS call expression node.
+//
+//	REPLACE_FIELDS({{.Expr.| sql}}, {{.Fields | sqlJoin ", "}})
+type ReplaceFieldsExpr struct {
+	// pos = ReplaceFields
+	// end = Rparen + 1
+
+	ReplaceFields token.Pos // position of "REPLACE_FIELDS" keyword
+	Rparen        token.Pos // position of ")"
+
+	Expr   Expr
+	Fields []*ReplaceFieldsArg
+}
+
 // AtTimeZone is AT TIME ZONE clause in EXTRACT call.
 //
 //	AT TIME ZONE {{.Expr | sql}}
diff --git a/ast/pos.go b/ast/pos.go
index dd0cd88..8bcac94 100644
--- a/ast/pos.go
+++ b/ast/pos.go
@@ -478,6 +478,22 @@ func (e *ExtractExpr) End() token.Pos {
 	return posAdd(e.Rparen, 1)
 }
 
+func (r *ReplaceFieldsArg) Pos() token.Pos {
+	return nodePos(wrapNode(r.Expr))
+}
+
+func (r *ReplaceFieldsArg) End() token.Pos {
+	return nodeEnd(wrapNode(r.Field))
+}
+
+func (r *ReplaceFieldsExpr) Pos() token.Pos {
+	return r.ReplaceFields
+}
+
+func (r *ReplaceFieldsExpr) End() token.Pos {
+	return posAdd(r.Rparen, 1)
+}
+
 func (a *AtTimeZone) Pos() token.Pos {
 	return a.At
 }
diff --git a/ast/sql.go b/ast/sql.go
index 12f864d..e7a2ec7 100644
--- a/ast/sql.go
+++ b/ast/sql.go
@@ -549,6 +549,12 @@ func (a *AtTimeZone) SQL() string {
 	return "AT TIME ZONE " + a.Expr.SQL()
 }
 
+func (r *ReplaceFieldsArg) SQL() string { return r.Expr.SQL() + " AS " + r.Field.SQL() }
+
+func (r *ReplaceFieldsExpr) SQL() string {
+	return "REPLACE_FIELDS(" + r.Expr.SQL() + ", " + sqlJoin(r.Fields, ", ") + ")"
+}
+
 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
index 0655fc5..ad6bf95 100644
--- a/parser.go
+++ b/parser.go
@@ -1469,6 +1469,8 @@ func (p *Parser) parseLit() ast.Expr {
 		switch {
 		case id.IsKeywordLike("SAFE_CAST"):
 			return p.parseCastExpr()
+		case id.IsKeywordLike("REPLACE_FIELDS"):
+			return p.parseReplaceFieldsExpr()
 		}
 		p.nextToken()
 		switch p.Token.Kind {
@@ -1754,6 +1756,33 @@ func (p *Parser) parseIfExpr() *ast.IfExpr {
 	}
 }
 
+func (p *Parser) parseReplaceFieldsArg() *ast.ReplaceFieldsArg {
+	expr := p.parseExpr()
+	p.expect("AS")
+	field := p.parsePath()
+
+	return &ast.ReplaceFieldsArg{
+		Expr:  expr,
+		Field: field,
+	}
+}
+
+func (p *Parser) parseReplaceFieldsExpr() *ast.ReplaceFieldsExpr {
+	replaceFields := p.expectKeywordLike("REPLACE_FIELDS").Pos
+	p.expect("(")
+	expr := p.parseExpr()
+	p.expect(",")
+	fields := parseCommaSeparatedList(p, p.parseReplaceFieldsArg)
+	rparen := p.expect(")").Pos
+
+	return &ast.ReplaceFieldsExpr{
+		ReplaceFields: replaceFields,
+		Rparen:        rparen,
+		Expr:          expr,
+		Fields:        fields,
+	}
+}
+
 func (p *Parser) parseCastExpr() *ast.CastExpr {
 	if p.Token.Kind != "CAST" && !p.Token.IsKeywordLike("SAFE_CAST") {
 		panic(p.errorfAtToken(&p.Token, `expected CAST keyword or SAFE_CAST pseudo keyword, but: %v`, p.Token.Kind))


Test patch:
diff --git a/testdata/input/expr/protocol_buffer_functions_replace_fields.sql b/testdata/input/expr/protocol_buffer_functions_replace_fields.sql
new file mode 100644
index 0000000..64caec0
--- /dev/null
+++ b/testdata/input/expr/protocol_buffer_functions_replace_fields.sql
@@ -0,0 +1,6 @@
+REPLACE_FIELDS(
+  NEW Book(
+    "The Hummingbird" AS title,
+    NEW BookDetails(10 AS chapters) AS details),
+  "The Hummingbird II" AS title,
+  11 AS details.chapters)
\ No newline at end of file
diff --git a/testdata/result/expr/protocol_buffer_functions_replace_fields.sql.txt b/testdata/result/expr/protocol_buffer_functions_replace_fields.sql.txt
new file mode 100644
index 0000000..5587e16
--- /dev/null
+++ b/testdata/result/expr/protocol_buffer_functions_replace_fields.sql.txt
@@ -0,0 +1,126 @@
+--- protocol_buffer_functions_replace_fields.sql
+REPLACE_FIELDS(
+  NEW Book(
+    "The Hummingbird" AS title,
+    NEW BookDetails(10 AS chapters) AS details),
+  "The Hummingbird II" AS title,
+  11 AS details.chapters)
+--- AST
+&ast.ReplaceFieldsExpr{
+  ReplaceFields: 0,
+  Rparen:        166,
+  Expr:          &ast.NewConstructor{
+    New:  18,
+    Type: &ast.NamedType{
+      Path: []*ast.Ident{
+        &ast.Ident{
+          NamePos: 22,
+          NameEnd: 26,
+          Name:    "Book",
+        },
+      },
+    },
+    Args: []ast.NewConstructorArg{
+      &ast.Alias{
+        Expr: &ast.StringLiteral{
+          ValuePos: 32,
+          ValueEnd: 49,
+          Value:    "The Hummingbird",
+        },
+        As: &ast.AsAlias{
+          As:    50,
+          Alias: &ast.Ident{
+            NamePos: 53,
+            NameEnd: 58,
+            Name:    "title",
+          },
+        },
+      },
+      &ast.Alias{
+        Expr: &ast.NewConstructor{
+          New:  64,
+          Type: &ast.NamedType{
+            Path: []*ast.Ident{
+              &ast.Ident{
+                NamePos: 68,
+                NameEnd: 79,
+                Name:    "BookDetails",
+              },
+            },
+          },
+          Args: []ast.NewConstructorArg{
+            &ast.Alias{
+              Expr: &ast.IntLiteral{
+                ValuePos: 80,
+                ValueEnd: 82,
+                Base:     10,
+                Value:    "10",
+              },
+              As: &ast.AsAlias{
+                As:    83,
+                Alias: &ast.Ident{
+                  NamePos: 86,
+                  NameEnd: 94,
+                  Name:    "chapters",
+                },
+              },
+            },
+          },
+          Rparen: 94,
+        },
+        As: &ast.AsAlias{
+          As:    96,
+          Alias: &ast.Ident{
+            NamePos: 99,
+            NameEnd: 106,
+            Name:    "details",
+          },
+        },
+      },
+    },
+    Rparen: 106,
+  },
+  Fields: []*ast.ReplaceFieldsArg{
+    &ast.ReplaceFieldsArg{
+      Expr: &ast.StringLiteral{
+        ValuePos: 111,
+        ValueEnd: 131,
+        Value:    "The Hummingbird II",
+      },
+      Field: &ast.Path{
+        Idents: []*ast.Ident{
+          &ast.Ident{
+            NamePos: 135,
+            NameEnd: 140,
+            Name:    "title",
+          },
+        },
+      },
+    },
+    &ast.ReplaceFieldsArg{
+      Expr: &ast.IntLiteral{
+        ValuePos: 144,
+        ValueEnd: 146,
+        Base:     10,
+        Value:    "11",
+      },
+      Field: &ast.Path{
+        Idents: []*ast.Ident{
+          &ast.Ident{
+            NamePos: 150,
+            NameEnd: 157,
+            Name:    "details",
+          },
+          &ast.Ident{
+            NamePos: 158,
+            NameEnd: 166,
+            Name:    "chapters",
+          },
+        },
+      },
+    },
+  },
+}
+
+--- SQL
+REPLACE_FIELDS(NEW Book("The Hummingbird" AS title, NEW BookDetails(10 AS chapters) AS details), "The Hummingbird II" AS title, 11 AS details.chapters)