Updated memory tables and tests to use new insert / update / delete interfaces

Signed-off-by: Zach Musgrave <zach@liquidata.co>

Author: zachmu

engine_test.go

@@ -3815,9 +3815,12 @@ func TestGenerators(t *testing.T) {
 func insertRows(t *testing.T, table sql.InsertableTable, rows ...sql.Row) {
 	t.Helper()
 
+	ctx := newCtx()
+	inserter := table.Inserter(ctx)
 	for _, r := range rows {
-		require.NoError(t, table.Insert(newCtx(), r))
+		require.NoError(t, inserter.Insert(ctx, r))
 	}
+	require.NoError(t, inserter.Close(ctx))
 }
 
 var pid uint64

memory/table.go

@@ -29,6 +29,9 @@ type Table struct {
 
 var _ sql.Table = (*Table)(nil)
 var _ sql.InsertableTable = (*Table)(nil)
+var _ sql.UpdatableTable = (*Table)(nil)
+var _ sql.DeletableTable = (*Table)(nil)
+var _ sql.ReplaceableTable = (*Table)(nil)
 var _ sql.FilteredTable = (*Table)(nil)
 var _ sql.ProjectedTable = (*Table)(nil)
 var _ sql.IndexableTable = (*Table)(nil)
@@ -241,30 +244,70 @@ func encodeIndexValue(value *indexValue) ([]byte, error) {
 	return buf.Bytes(), nil
 }
 
-// Insert a new row into the table.
+type tableEditor struct {
+	table *Table
+}
+
+var _ sql.RowReplacer = (*tableEditor)(nil)
+var _ sql.RowUpdater = (*tableEditor)(nil)
+var _ sql.RowInserter = (*tableEditor)(nil)
+var _ sql.RowDeleter = (*tableEditor)(nil)
+
+func (t tableEditor) Close(*sql.Context) error {
+	// TODO: it would be nice to apply all pending updates here at once, rather than directly in the Insert / Update
+	//  / Delete methods.
+	return nil
+}
+
+func (t *Table) Inserter(*sql.Context) sql.RowInserter {
+	return &tableEditor{t}
+}
+
+func (t *Table) Updater(*sql.Context) sql.RowUpdater {
+	return &tableEditor{t}
+}
+
+func (t *Table) Replacer(*sql.Context) sql.RowReplacer {
+	return &tableEditor{t}
+}
+
+func (t *Table) Deleter(*sql.Context) sql.RowDeleter {
+	return &tableEditor{t}
+}
+
+// Convenience method to avoid having to create an inserter in test setup
 func (t *Table) Insert(ctx *sql.Context, row sql.Row) error {
-	if err := checkRow(t.schema, row); err != nil {
+	inserter := t.Inserter(ctx)
+	if err := inserter.Insert(ctx, row); err != nil {
+		return err
+	}
+	return inserter.Close(ctx)
+}
+
+// Insert a new row into the table.
+func (t *tableEditor) Insert(ctx *sql.Context, row sql.Row) error {
+	if err := checkRow(t.table.schema, row); err != nil {
 		return err
 	}
 
-	key := string(t.keys[t.insert])
-	t.insert++
-	if t.insert == len(t.keys) {
-		t.insert = 0
+	key := string(t.table.keys[t.table.insert])
+	t.table.insert++
+	if t.table.insert == len(t.table.keys) {
+		t.table.insert = 0
 	}
 
-	t.partitions[key] = append(t.partitions[key], row)
+	t.table.partitions[key] = append(t.table.partitions[key], row)
 	return nil
 }
 
 // Delete the given row from the table.
-func (t *Table) Delete(ctx *sql.Context, row sql.Row) error {
-	if err := checkRow(t.schema, row); err != nil {
+func (t *tableEditor) Delete(ctx *sql.Context, row sql.Row) error {
+	if err := checkRow(t.table.schema, row); err != nil {
 		return err
 	}
 
 	matches := false
-	for partitionIndex, partition := range t.partitions {
+	for partitionIndex, partition := range t.table.partitions {
 		for partitionRowIndex, partitionRow := range partition {
 			matches = true
 			for rIndex, val := range row {
@@ -274,7 +317,7 @@ func (t *Table) Delete(ctx *sql.Context, row sql.Row) error {
 				}
 			}
 			if matches {
-				t.partitions[partitionIndex] = append(partition[:partitionRowIndex], partition[partitionRowIndex+1:]...)
+				t.table.partitions[partitionIndex] = append(partition[:partitionRowIndex], partition[partitionRowIndex+1:]...)
 				break
 			}
 		}
@@ -290,16 +333,16 @@ func (t *Table) Delete(ctx *sql.Context, row sql.Row) error {
 	return nil
 }
 
-func (t *Table) Update(ctx *sql.Context, oldRow sql.Row, newRow sql.Row) error {
-	if err := checkRow(t.schema, oldRow); err != nil {
+func (t *tableEditor) Update(ctx *sql.Context, oldRow sql.Row, newRow sql.Row) error {
+	if err := checkRow(t.table.schema, oldRow); err != nil {
 		return err
 	}
-	if err := checkRow(t.schema, newRow); err != nil {
+	if err := checkRow(t.table.schema, newRow); err != nil {
 		return err
 	}
 
 	matches := false
-	for partitionIndex, partition := range t.partitions {
+	for partitionIndex, partition := range t.table.partitions {
 		for partitionRowIndex, partitionRow := range partition {
 			matches = true
 			for rIndex, val := range oldRow {
@@ -309,7 +352,7 @@ func (t *Table) Update(ctx *sql.Context, oldRow sql.Row, newRow sql.Row) error {
 				}
 			}
 			if matches {
-				t.partitions[partitionIndex][partitionRowIndex] = newRow
+				t.table.partitions[partitionIndex][partitionRowIndex] = newRow
 				break
 			}
 		}

sql/core.go

@@ -226,7 +226,7 @@ type RowInserter interface {
 type DeletableTable interface {
 	// Deleter returns a RowDeleter for this table. The RowDeleter will get one call to Delete for each row to be deleted,
 	// and will end with a call to Close() to finalize the delete operation.
-	Deleter(*Context, Row) RowDeleter
+	Deleter(*Context) RowDeleter
 }
 
 // RowDeleter is a delete cursor that can delete one or more rows from a table.
@@ -260,7 +260,7 @@ type RowReplacer interface {
 type ReplaceableTable interface {
 	// Replacer returns a RowReplacer for this table. The RowReplacer will have Insert and optionally Delete called once
 	// for each row, followed by a call to Close() when all rows have been processed.
-	Replacer() RowReplacer
+	Replacer(ctx *Context) RowReplacer
 }
 
 // UpdateableTable is a table that can process updates of existing rows via update statements.

sql/plan/delete.go

@@ -76,6 +76,8 @@ func (p *DeleteFrom) Execute(ctx *sql.Context) (int, error) {
 		return 0, err
 	}
 
+	deleter := deletable.Deleter(ctx)
+
 	i := 0
 	for {
 		row, err := iter.Next()
@@ -88,14 +90,18 @@ func (p *DeleteFrom) Execute(ctx *sql.Context) (int, error) {
 			return i, err
 		}
 
-		if err := deletable.Delete(ctx, row); err != nil {
+		if err := deleter.Delete(ctx, row); err != nil {
 			_ = iter.Close()
 			return i, err
 		}
 
 		i++
 	}
 
+	if err := deleter.Close(ctx); err != nil {
+		return 0, err
+	}
+
 	return i, nil
 }
 

sql/plan/insert.go

@@ -132,6 +132,14 @@ func (p *InsertInto) Execute(ctx *sql.Context) (int, error) {
 		return 0, err
 	}
 
+	var inserter sql.RowInserter
+	var replacer sql.RowReplacer
+	if replaceable != nil {
+		replacer = replaceable.Replacer(ctx)
+	} else {
+		inserter = insertable.Inserter(ctx)
+	}
+
 	i := 0
 	for {
 		row, err := iter.Next()
@@ -163,8 +171,8 @@ func (p *InsertInto) Execute(ctx *sql.Context) (int, error) {
 			}
 		}
 
-		if replaceable != nil {
-			if err = replaceable.Delete(ctx, row); err != nil {
+		if replacer != nil {
+			if err = replacer.Delete(ctx, row); err != nil {
 				if err != sql.ErrDeleteRowNotFound {
 					_ = iter.Close()
 					return i, err
@@ -173,19 +181,29 @@ func (p *InsertInto) Execute(ctx *sql.Context) (int, error) {
 				i++
 			}
 
-			if err = replaceable.Insert(ctx, row); err != nil {
+			if err = replacer.Insert(ctx, row); err != nil {
 				_ = iter.Close()
 				return i, err
 			}
 		} else {
-			if err := insertable.Insert(ctx, row); err != nil {
+			if err := inserter.Insert(ctx, row); err != nil {
 				_ = iter.Close()
 				return i, err
 			}
 		}
 		i++
 	}
 
+	if replacer != nil {
+		if err := replacer.Close(ctx); err != nil {
+			return 0, err
+		}
+	} else {
+		if err := inserter.Close(ctx); err != nil {
+			return 0, err
+		}
+	}
+
 	return i, nil
 }
 

sql/plan/update.go

@@ -100,6 +100,8 @@ func (p *Update) Execute(ctx *sql.Context) (int, int, error) {
 		return 0, 0, err
 	}
 
+	updater := updatable.Updater(ctx)
+
 	rowsMatched := 0
 	rowsUpdated := 0
 	for {
@@ -120,7 +122,7 @@ func (p *Update) Execute(ctx *sql.Context) (int, int, error) {
 		}
 		if equals, err := oldRow.Equals(newRow, schema); err == nil {
 			if !equals {
-				err = updatable.Update(ctx, oldRow, newRow)
+				err = updater.Update(ctx, oldRow, newRow)
 				if err != nil {
 					_ = iter.Close()
 					return rowsMatched, rowsUpdated, err
@@ -133,6 +135,10 @@ func (p *Update) Execute(ctx *sql.Context) (int, int, error) {
 		}
 	}
 
+	if err := updater.Close(ctx); err != nil {
+		return 0, 0, err
+	}
+
 	return rowsMatched, rowsUpdated, nil
 }