From c6185cffc7f23741ac4a230aadee74b3def85ced Mon Sep 17 00:00:00 2001
From: Seiichi Uchida <topecongiro@fastmail.com>
Date: Wed, 15 Jun 2022 15:17:48 +0900
Subject: [PATCH] feat(spanner/spansql): add a support for parsing INSERT
 statement (#6148)

* feat(spanner/spansql): add a support for parsing INSERT statement

* Clean up

Co-authored-by: rahul2393 <rahulyadavsep92@gmail.com>
---
 spanner/spansql/parser.go      | 51 +++++++++++++++++++++++++++++++++-
 spanner/spansql/parser_test.go | 42 ++++++++++++++++++++++++++++
 spanner/spansql/sql.go         | 32 +++++++++++++++++++++
 spanner/spansql/types.go       | 24 +++++++++++++++-
 4 files changed, 147 insertions(+), 2 deletions(-)

diff --git a/spanner/spansql/parser.go b/spanner/spansql/parser.go
index 0ebefb1ddd9..2f2d066a0c1 100644
--- a/spanner/spansql/parser.go
+++ b/spanner/spansql/parser.go
@@ -1442,7 +1442,16 @@ func (p *parser) parseDMLStmt() (DMLStmt, *parseError) {
 
 		update_item: path_expression = expression | path_expression = DEFAULT
 
-		TODO: Insert.
+		INSERT [INTO] target_name
+		 (column_name_1 [, ..., column_name_n] )
+		 input
+
+		input:
+		 VALUES (row_1_column_1_expr [, ..., row_1_column_n_expr ] )
+		        [, ..., (row_k_column_1_expr [, ..., row_k_column_n_expr ] ) ]
+		| select_query
+
+		expr: value_expression | DEFAULT
 	*/
 
 	if p.eat("DELETE") {
@@ -1499,6 +1508,46 @@ func (p *parser) parseDMLStmt() (DMLStmt, *parseError) {
 		return u, nil
 	}
 
+	if p.eat("INSERT") {
+		p.eat("INTO") // optional
+		tname, err := p.parseTableOrIndexOrColumnName()
+		if err != nil {
+			return nil, err
+		}
+
+		columns, err := p.parseColumnNameList()
+		if err != nil {
+			return nil, err
+		}
+
+		var input ValuesOrSelect
+		if p.eat("VALUES") {
+			values := make([][]Expr, 0)
+			for {
+				exprs, err := p.parseParenExprList()
+				if err != nil {
+					return nil, err
+				}
+				values = append(values, exprs)
+				if !p.eat(",") {
+					break
+				}
+			}
+			input = Values(values)
+		} else {
+			input, err = p.parseSelect()
+			if err != nil {
+				return nil, err
+			}
+		}
+
+		return &Insert{
+			Table:   tname,
+			Columns: columns,
+			Input:   input,
+		}, nil
+	}
+
 	return nil, p.errorf("unknown DML statement")
 }
 
diff --git a/spanner/spansql/parser_test.go b/spanner/spansql/parser_test.go
index 9b4565b42bd..d42dcf7427b 100644
--- a/spanner/spansql/parser_test.go
+++ b/spanner/spansql/parser_test.go
@@ -300,6 +300,48 @@ func TestParseQuery(t *testing.T) {
 	}
 }
 
+func TestParseDMLStmt(t *testing.T) {
+	tests := []struct {
+		in   string
+		want DMLStmt
+	}{
+		{"INSERT Singers (SingerId, FirstName, LastName) VALUES (1, 'Marc', 'Richards')",
+			&Insert{
+				Table:   "Singers",
+				Columns: []ID{ID("SingerId"), ID("FirstName"), ID("LastName")},
+				Input:   Values{{IntegerLiteral(1), StringLiteral("Marc"), StringLiteral("Richards")}},
+			},
+		},
+		{"INSERT Singers (SingerId, FirstName, LastName) SELECT * FROM UNNEST ([1, 2, 3]) AS data",
+			&Insert{
+				Table:   "Singers",
+				Columns: []ID{ID("SingerId"), ID("FirstName"), ID("LastName")},
+				Input: Select{
+					List: []Expr{Star},
+					From: []SelectFrom{SelectFromUnnest{
+						Expr: Array{
+							IntegerLiteral(1),
+							IntegerLiteral(2),
+							IntegerLiteral(3),
+						},
+						Alias: ID("data"),
+					}},
+				},
+			},
+		},
+	}
+	for _, test := range tests {
+		got, err := ParseDMLStmt(test.in)
+		if err != nil {
+			t.Errorf("ParseDMLStmt(%q): %v", test.in, err)
+			continue
+		}
+		if !reflect.DeepEqual(got, test.want) {
+			t.Errorf("ParseDMLStmt(%q) incorrect.\n got %#v\nwant %#v", test.in, got, test.want)
+		}
+	}
+}
+
 func TestParseExpr(t *testing.T) {
 	tests := []struct {
 		in   string
diff --git a/spanner/spansql/sql.go b/spanner/spansql/sql.go
index 406b5a8648c..32e498d1ac6 100644
--- a/spanner/spansql/sql.go
+++ b/spanner/spansql/sql.go
@@ -260,6 +260,38 @@ func (u *Update) SQL() string {
 	return str
 }
 
+func (i *Insert) SQL() string {
+	str := "INSERT " + i.Table.SQL() + " INTO ("
+	for i, column := range i.Columns {
+		if i > 0 {
+			str += ", "
+		}
+		str += column.SQL()
+	}
+	str += ") "
+	str += i.Input.SQL()
+	return str
+}
+
+func (v Values) SQL() string {
+	str := "VALUES "
+	for j, values := range v {
+		if j > 0 {
+			str += ", "
+		}
+		str += "("
+
+		for k, value := range values {
+			if k > 0 {
+				str += ", "
+			}
+			str += value.SQL()
+		}
+		str += ")"
+	}
+	return str
+}
+
 func (cd ColumnDef) SQL() string {
 	str := cd.Name.SQL() + " " + cd.Type.SQL()
 	if cd.NotNull {
diff --git a/spanner/spansql/types.go b/spanner/spansql/types.go
index 521df50b752..a9b335eee43 100644
--- a/spanner/spansql/types.go
+++ b/spanner/spansql/types.go
@@ -302,7 +302,29 @@ type Delete struct {
 func (d *Delete) String() string { return fmt.Sprintf("%#v", d) }
 func (*Delete) isDMLStmt()       {}
 
-// TODO: Insert.
+// Insert represents an INSERT statement.
+// https://cloud.google.com/spanner/docs/dml-syntax#insert-statement
+type Insert struct {
+	Table   ID
+	Columns []ID
+	Input   ValuesOrSelect
+}
+
+// Values represents one or more lists of expressions passed to an `INSERT` statement.
+type Values [][]Expr
+
+func (v Values) isValuesOrSelect() {}
+func (v Values) String() string    { return fmt.Sprintf("%#v", v) }
+
+type ValuesOrSelect interface {
+	isValuesOrSelect()
+	SQL() string
+}
+
+func (Select) isValuesOrSelect() {}
+
+func (i *Insert) String() string { return fmt.Sprintf("%#v", i) }
+func (*Insert) isDMLStmt()       {}
 
 // Update represents an UPDATE statement.
 // https://cloud.google.com/spanner/docs/dml-syntax#update-statement