Repository: genjidb/genji
Pull request #385 (2021-04-12 17:48:18)

Add INSERT ... SELECT
This adds support for `INSERT ... SELECT`. The only limitation compared to other databases is that inserting and selecting from the same table is prohibited.


Patch:
diff --git a/document/document.go b/document/document.go
index 491f8cde..76d096e5 100644
--- a/document/document.go
+++ b/document/document.go
@@ -106,6 +106,11 @@ func (fb *FieldBuffer) UnmarshalJSON(data []byte) error {
 	})
 }
 
+func (fb *FieldBuffer) String() string {
+	s, _ := fb.MarshalJSON()
+	return string(s)
+}
+
 type fieldValue struct {
 	Field string
 	Value Value
diff --git a/sql/parser/insert.go b/sql/parser/insert.go
index 9f3b27cb..e6bb182b 100644
--- a/sql/parser/insert.go
+++ b/sql/parser/insert.go
@@ -1,6 +1,8 @@
 package parser
 
 import (
+	"errors"
+
 	"github.com/genjidb/genji/expr"
 	"github.com/genjidb/genji/planner"
 	"github.com/genjidb/genji/sql/scanner"
@@ -28,18 +30,30 @@ func (p *Parser) parseInsertStatement() (*planner.Statement, error) {
 	}
 
 	// Parse path list: (a, b, c)
-	fields, err := p.parseFieldList()
+	cfg.Fields, err = p.parseFieldList()
 	if err != nil {
 		return nil, err
 	}
 
-	// Parse VALUES (v1, v2, v3)
-	cfg.Values, err = p.parseValues(fields)
-	if err != nil {
-		return nil, err
+	// Check if VALUES or SELECT token exists.
+	tok, pos, lit := p.ScanIgnoreWhitespace()
+	switch tok {
+	case scanner.VALUES:
+		// Parse VALUES (v1, v2, v3)
+		cfg.Values, err = p.parseValues(cfg.Fields)
+		if err != nil {
+			return nil, err
+		}
+	case scanner.SELECT:
+		cfg.SelectStmt, err = p.parseSelectStatement()
+		if err != nil {
+			return nil, err
+		}
+	default:
+		return nil, newParseError(scanner.Tokstr(tok, lit), []string{"VALUES", "SELECT"}, pos)
 	}
 
-	return cfg.ToStream(), nil
+	return cfg.ToStream()
 }
 
 // parseFieldList parses a list of fields in the form: (path, path, ...), if exists.
@@ -68,11 +82,6 @@ func (p *Parser) parseFieldList() ([]string, error) {
 
 // parseValues parses the "VALUES" clause of the query, if it exists.
 func (p *Parser) parseValues(fields []string) ([]expr.Expr, error) {
-	// Check if the VALUES token exists.
-	if tok, pos, lit := p.ScanIgnoreWhitespace(); tok != scanner.VALUES {
-		return nil, newParseError(scanner.Tokstr(tok, lit), []string{"VALUES"}, pos)
-	}
-
 	if len(fields) > 0 {
 		return p.parseDocumentsWithFields(fields)
 	}
@@ -182,17 +191,35 @@ func (p *Parser) parseParamOrDocument() (expr.Expr, error) {
 
 // insertConfig holds INSERT configuration.
 type insertConfig struct {
-	TableName string
-	Values    []expr.Expr
+	TableName  string
+	Values     []expr.Expr
+	Fields     []string
+	SelectStmt *planner.Statement
 }
 
-func (cfg *insertConfig) ToStream() *planner.Statement {
-	s := stream.New(stream.Expressions(cfg.Values...))
+func (cfg *insertConfig) ToStream() (*planner.Statement, error) {
+	var s *stream.Stream
+	if cfg.Values != nil {
+		s = stream.New(stream.Expressions(cfg.Values...))
 
-	s = s.Pipe(stream.TableInsert(cfg.TableName))
+		s = s.Pipe(stream.TableInsert(cfg.TableName))
+	} else {
+		s = cfg.SelectStmt.Stream
+
+		// ensure we are not reading and writing to the same table.
+		if s.First().(*stream.SeqScanOperator).TableName == cfg.TableName {
+			return nil, errors.New("cannot read and write to the same table")
+		}
+
+		if len(cfg.Fields) > 0 {
+			s = s.Pipe(stream.IterRename(cfg.Fields...))
+		}
+
+		s = s.Pipe(stream.TableInsert(cfg.TableName))
+	}
 
 	return &planner.Statement{
 		Stream:   s,
 		ReadOnly: false,
-	}
+	}, nil
 }
diff --git a/sql/parser/select.go b/sql/parser/select.go
index 74d02f3c..ac78c5a7 100644
--- a/sql/parser/select.go
+++ b/sql/parser/select.go
@@ -5,7 +5,6 @@ import (
 
 	"github.com/genjidb/genji/expr"
 	"github.com/genjidb/genji/planner"
-	"github.com/genjidb/genji/query"
 	"github.com/genjidb/genji/sql/scanner"
 	"github.com/genjidb/genji/stream"
 	"github.com/genjidb/genji/stringutil"
@@ -13,7 +12,7 @@ import (
 
 // parseSelectStatement parses a select string and returns a Statement AST object.
 // This function assumes the SELECT token has already been consumed.
-func (p *Parser) parseSelectStatement() (query.Statement, error) {
+func (p *Parser) parseSelectStatement() (*planner.Statement, error) {
 	var cfg selectConfig
 	var err error
 
diff --git a/stream/operator.go b/stream/operator.go
index 0fea4d50..0e3dcab7 100644
--- a/stream/operator.go
+++ b/stream/operator.go
@@ -5,6 +5,7 @@ import (
 	"container/heap"
 	"container/list"
 	"errors"
+	"strings"
 
 	"github.com/genjidb/genji/database"
 	"github.com/genjidb/genji/document"
@@ -698,3 +699,70 @@ func (op *UnsetOperator) Iterate(in *expr.Environment, f func(out *expr.Environm
 func (op *UnsetOperator) String() string {
 	return stringutil.Sprintf("unset(%s)", op.Field)
 }
+
+// An IterRenameOperator iterates over all fields of the incoming document in order and renames them.
+type IterRenameOperator struct {
+	baseOperator
+	FieldNames []string
+}
+
+// IterRename iterates over all fields of the incoming document in order and renames them.
+// If the number of fields of the incoming document doesn't match the number of expected fields,
+// it returns an error.
+func IterRename(fieldNames ...string) *IterRenameOperator {
+	return &IterRenameOperator{
+		FieldNames: fieldNames,
+	}
+}
+
+// Iterate implements the Operator interface.
+func (op *IterRenameOperator) Iterate(in *expr.Environment, f func(out *expr.Environment) error) error {
+	var fb document.FieldBuffer
+	var newEnv expr.Environment
+
+	return op.Prev.Iterate(in, func(out *expr.Environment) error {
+		fb.Reset()
+
+		d, ok := out.GetDocument()
+		if !ok {
+			return errors.New("missing document")
+		}
+
+		var i int
+		err := d.Iterate(func(field string, value document.Value) error {
+			// if there are too many fields in the incoming document
+			if i >= len(op.FieldNames) {
+				n, err := document.Length(d)
+				if err != nil {
+					return err
+				}
+				return stringutil.Errorf("%d values for %d fields", n, len(op.FieldNames))
+			}
+
+			fb.Add(op.FieldNames[i], value)
+			i++
+			return nil
+		})
+		if err != nil {
+			return err
+		}
+
+		// if there are too few fields in the incoming document
+		if i < len(op.FieldNames) {
+			n, err := document.Length(d)
+			if err != nil {
+				return err
+			}
+			return stringutil.Errorf("%d values for %d fields", n, len(op.FieldNames))
+		}
+
+		newEnv.Outer = out
+		newEnv.SetDocument(&fb)
+
+		return f(&newEnv)
+	})
+}
+
+func (op *IterRenameOperator) String() string {
+	return stringutil.Sprintf("iterRename(%s)", strings.Join(op.FieldNames, ", "))
+}


Test patch:
diff --git a/query/insert_test.go b/query/insert_test.go
index 08c3036b..036b036d 100644
--- a/query/insert_test.go
+++ b/query/insert_test.go
@@ -276,3 +276,54 @@ func TestInsertStmt(t *testing.T) {
 		}
 	})
 }
+
+func TestInsertSelect(t *testing.T) {
+	tests := []struct {
+		name     string
+		query    string
+		fails    bool
+		expected string
+		params   []interface{}
+	}{
+		{"Same table", `INSERT INTO foo SELECT * FROM foo`, true, ``, nil},
+		{"No fields / No projection", `INSERT INTO foo SELECT * FROM bar`, false, `[{"pk()":1, "a":1, "b":10}]`, nil},
+		{"No fields / Projection", `INSERT INTO foo SELECT a FROM bar`, false, `[{"pk()":1, "a":1}]`, nil},
+		{"With fields / No Projection", `INSERT INTO foo (a, b) SELECT * FROM bar`, false, `[{"pk()":1, "a":1, "b":10}]`, nil},
+		{"With fields / Projection", `INSERT INTO foo (c, d) SELECT a, b FROM bar`, false, `[{"pk()":1, "c":1, "d":10}]`, nil},
+		{"Too many fields / No Projection", `INSERT INTO foo (c) SELECT * FROM bar`, true, ``, nil},
+		{"Too many fields / Projection", `INSERT INTO foo (c, d) SELECT a, b, c FROM bar`, true, ``, nil},
+		{"Too few fields / No Projection", `INSERT INTO foo (c, d, e) SELECT * FROM bar`, true, ``, nil},
+		{"Too few fields / Projection", `INSERT INTO foo (c, d) SELECT a FROM bar`, true, ``, nil},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			db, err := genji.Open(":memory:")
+			require.NoError(t, err)
+			defer db.Close()
+
+			err = db.Exec(`
+				CREATE TABLE foo;
+				CREATE TABLE bar;
+				INSERT INTO bar (a, b) VALUES (1, 10)
+			`)
+			require.NoError(t, err)
+
+			err = db.Exec(test.query, test.params...)
+			if test.fails {
+				require.Error(t, err)
+				return
+			}
+			require.NoError(t, err)
+
+			st, err := db.Query("SELECT pk(), * FROM foo")
+			require.NoError(t, err)
+			defer st.Close()
+
+			var buf bytes.Buffer
+			err = document.IteratorToJSONArray(&buf, st)
+			require.NoError(t, err)
+			require.JSONEq(t, test.expected, buf.String())
+		})
+	}
+}
diff --git a/sql/parser/insert_test.go b/sql/parser/insert_test.go
index 6e720c72..27a480b5 100644
--- a/sql/parser/insert_test.go
+++ b/sql/parser/insert_test.go
@@ -76,6 +76,36 @@ func TestParserInsert(t *testing.T) {
 			nil, true},
 		{"Values / Without fields / Wrong values", "INSERT INTO test VALUES {a: 1}, ('e', 'f')",
 			nil, true},
+		{"Select / same table", "INSERT INTO test SELECT * FROM test",
+			nil, true},
+		{"Select / Without fields", "INSERT INTO test SELECT * FROM foo",
+			stream.New(stream.SeqScan("foo")).
+				Pipe(stream.Project(expr.Wildcard{})).
+				Pipe(stream.TableInsert("test")),
+			false},
+		{"Select / Without fields / With projection", "INSERT INTO test SELECT a, b FROM foo",
+			stream.New(stream.SeqScan("foo")).
+				Pipe(stream.Project(parseNamedExpr(t, "a"), parseNamedExpr(t, "b"))).
+				Pipe(stream.TableInsert("test")),
+			false},
+		{"Select / With fields", "INSERT INTO test (a, b) SELECT * FROM foo",
+			stream.New(stream.SeqScan("foo")).
+				Pipe(stream.Project(expr.Wildcard{})).
+				Pipe(stream.IterRename("a", "b")).
+				Pipe(stream.TableInsert("test")),
+			false},
+		{"Select / With fields / With projection", "INSERT INTO test (a, b) SELECT a, b FROM foo",
+			stream.New(stream.SeqScan("foo")).
+				Pipe(stream.Project(parseNamedExpr(t, "a"), parseNamedExpr(t, "b"))).
+				Pipe(stream.IterRename("a", "b")).
+				Pipe(stream.TableInsert("test")),
+			false},
+		{"Select / With fields / With projection / different fields", "INSERT INTO test (a, b) SELECT c, d FROM foo",
+			stream.New(stream.SeqScan("foo")).
+				Pipe(stream.Project(parseNamedExpr(t, "c"), parseNamedExpr(t, "d"))).
+				Pipe(stream.IterRename("a", "b")).
+				Pipe(stream.TableInsert("test")),
+			false},
 	}
 
 	for _, test := range tests {
@@ -89,7 +119,7 @@ func TestParserInsert(t *testing.T) {
 			require.Len(t, q.Statements, 1)
 			stmt := q.Statements[0].(*planner.Statement)
 			require.False(t, stmt.ReadOnly)
-			require.EqualValues(t, test.expected, stmt.Stream)
+			require.EqualValues(t, test.expected.String(), stmt.Stream.String())
 		})
 	}
 }
diff --git a/stream/operator_test.go b/stream/operator_test.go
index 5726ac0b..3c6aacdd 100644
--- a/stream/operator_test.go
+++ b/stream/operator_test.go
@@ -709,3 +709,52 @@ func TestUnset(t *testing.T) {
 		require.Equal(t, stream.Unset("a").String(), "unset(a)")
 	})
 }
+
+func TestIterRename(t *testing.T) {
+	tests := []struct {
+		fieldNames []string
+		in, out    []document.Document
+		fails      bool
+	}{
+		{
+			[]string{"c", "d"},
+			testutil.MakeDocuments(t, `{"a": 10, "b": 20}`),
+			testutil.MakeDocuments(t, `{"c": 10, "d": 20}`),
+			false,
+		},
+		{
+			[]string{"c", "d", "e"},
+			testutil.MakeDocuments(t, `{"a": 10, "b": 20}`),
+			nil,
+			true,
+		},
+		{
+			[]string{"c"},
+			testutil.MakeDocuments(t, `{"a": 10, "b": 20}`),
+			nil,
+			true,
+		},
+	}
+
+	for _, test := range tests {
+		s := stream.New(stream.Documents(test.in...)).Pipe(stream.IterRename(test.fieldNames...))
+		t.Run(s.String(), func(t *testing.T) {
+			i := 0
+			err := s.Iterate(nil, func(out *expr.Environment) error {
+				d, _ := out.GetDocument()
+				require.Equal(t, test.out[i], d)
+				i++
+				return nil
+			})
+			if test.fails {
+				require.Error(t, err)
+			} else {
+				require.NoError(t, err)
+			}
+		})
+	}
+
+	t.Run("String", func(t *testing.T) {
+		require.Equal(t, stream.IterRename("a", "b", "c").String(), "iterRename(a, b, c)")
+	})
+}