From 79acc0efe53ff02dde7db2c2a406355d5f94009a Mon Sep 17 00:00:00 2001
From: Alex Roitman <aroitman@housecanary.com>
Date: Fri, 3 Nov 2017 11:50:03 -0700
Subject: [PATCH] Allow setting multiple fields in a single fset command. Add
 xx flag to fset.

---
 controller/collection/collection.go | 17 +++++++
 controller/crud.go                  | 76 +++++++++++++++++------------
 tests/keys_test.go                  |  7 ++-
 3 files changed, 68 insertions(+), 32 deletions(-)

diff --git a/controller/collection/collection.go b/controller/collection/collection.go
index c3886c3e..55a955b0 100644
--- a/controller/collection/collection.go
+++ b/controller/collection/collection.go
@@ -233,6 +233,23 @@ func (c *Collection) SetField(id, field string, value float64) (obj geojson.Obje
 	return item.object, c.getFieldValues(id), updated, true
 }
 
+// SetFields is similar to SetField, just setting multiple fields at once
+func (c *Collection) SetFields(id string, in_fields []string, in_values []float64) (
+	obj geojson.Object, fields []float64, updated bool, ok bool,
+) {
+	i := c.items.Get(&itemT{id: id})
+	if i == nil {
+		ok = false
+		return
+	}
+	item := i.(*itemT)
+	for idx, field := range in_fields {
+		_upd := c.setField(item, field, in_values[idx])
+		updated = updated || _upd
+	}
+	return item.object, c.getFieldValues(id), updated, true
+}
+
 func (c *Collection) setField(item *itemT, field string, value float64) (updated bool) {
 	idx, ok := c.fieldMap[field]
 	if !ok {
diff --git a/controller/crud.go b/controller/crud.go
index cabde995..7139082e 100644
--- a/controller/crud.go
+++ b/controller/crud.go
@@ -780,8 +780,9 @@ notok:
 	return
 }
 
-func (c *Controller) parseFSetArgs(vs []resp.Value) (d commandDetailsT, err error) {
-	var svalue string
+func (c *Controller) parseFSetArgs(vs []resp.Value) (
+	d commandDetailsT, fields []string, values []float64, xx bool, err error,
+) {
 	var ok bool
 	if vs, d.key, ok = tokenval(vs); !ok || d.key == "" {
 		err = errInvalidNumberOfArguments
@@ -791,26 +792,33 @@ func (c *Controller) parseFSetArgs(vs []resp.Value) (d commandDetailsT, err erro
 		err = errInvalidNumberOfArguments
 		return
 	}
-	if vs, d.field, ok = tokenval(vs); !ok || d.field == "" {
-		err = errInvalidNumberOfArguments
-		return
-	}
-	if isReservedFieldName(d.field) {
-		err = errInvalidNumberOfArguments
-		return
-	}
-	if vs, svalue, ok = tokenval(vs); !ok || svalue == "" {
-		err = errInvalidNumberOfArguments
-		return
-	}
-	if len(vs) != 0 {
-		err = errInvalidNumberOfArguments
-		return
-	}
-	d.value, err = strconv.ParseFloat(svalue, 64)
-	if err != nil {
-		err = errInvalidArgument(svalue)
-		return
+	for len(vs) > 0 {
+		var name string
+		if vs, name, ok = tokenval(vs); !ok || name == "" {
+			err = errInvalidNumberOfArguments
+			return
+		}
+		if lc(name, "xx") {
+			xx = true
+			continue
+		}
+		if isReservedFieldName(name) {
+			err = errInvalidArgument(name)
+			return
+		}
+		var svalue string
+		var value float64
+		if vs, svalue, ok = tokenval(vs); !ok || svalue == "" {
+			err = errInvalidNumberOfArguments
+			return
+		}
+		value, err = strconv.ParseFloat(svalue, 64)
+		if err != nil {
+			err = errInvalidArgument(svalue)
+			return
+		}
+		fields = append(fields, name)
+		values = append(values, value)
 	}
 	return
 }
@@ -818,24 +826,30 @@ func (c *Controller) parseFSetArgs(vs []resp.Value) (d commandDetailsT, err erro
 func (c *Controller) cmdFset(msg *server.Message) (res resp.Value, d commandDetailsT, err error) {
 	start := time.Now()
 	vs := msg.Values[1:]
-	d, err = c.parseFSetArgs(vs)
+	var fields []string
+	var values []float64
+	var xx bool
+	d, fields, values, xx, err = c.parseFSetArgs(vs)
+
 	col := c.getCol(d.key)
 	if col == nil {
 		err = errKeyNotFound
 		return
 	}
 	var ok bool
-	d.obj, d.fields, d.updated, ok = col.SetField(d.id, d.field, d.value)
-	if !ok {
+	d.obj, d.fields, d.updated, ok = col.SetFields(d.id, fields, values)
+	if !(ok || xx) {
 		err = errIDNotFound
 		return
 	}
-	d.command = "fset"
-	d.timestamp = time.Now()
-	fmap := col.FieldMap()
-	d.fmap = make(map[string]int)
-	for key, idx := range fmap {
-		d.fmap[key] = idx
+	if ok {
+		d.command = "fset"
+		d.timestamp = time.Now()
+		fmap := col.FieldMap()
+		d.fmap = make(map[string]int)
+		for key, idx := range fmap {
+			d.fmap[key] = idx
+		}
 	}
 
 	switch msg.OutputType {
diff --git a/tests/keys_test.go b/tests/keys_test.go
index d9d24f5d..b62fa621 100644
--- a/tests/keys_test.go
+++ b/tests/keys_test.go
@@ -80,9 +80,14 @@ func keys_FSET_test(mc *mockServer) error {
 		{"GET", "mykey", "myid", "WITHFIELDS", "HASH", 7}, {"[9my5xp7]"},
 		{"FSET", "mykey", "myid", "f1", 105.6}, {1},
 		{"GET", "mykey", "myid", "WITHFIELDS", "HASH", 7}, {"[9my5xp7 [f1 105.6]]"},
+		{"FSET", "mykey", "myid", "f1", 1.1, "f2", 2.2}, {1},
+		{"GET", "mykey", "myid", "WITHFIELDS", "HASH", 7}, {"[9my5xp7 [f1 1.1 f2 2.2]]"},
 		{"FSET", "mykey", "myid", "f1", 0}, {1},
+		{"GET", "mykey", "myid", "WITHFIELDS", "HASH", 7}, {"[9my5xp7 [f2 2.2]]"},
+		{"FSET", "mykey", "myid", "f2", 0}, {1},
 		{"GET", "mykey", "myid", "WITHFIELDS", "HASH", 7}, {"[9my5xp7]"},
-		{"FSET", "mykey", "myid", "f1", 0}, {0},
+		{"FSET", "mykey", "myid2", "xx", "f1", 1.1, "f2", 2.2}, {0},
+		{"GET", "mykey", "myid2"}, {nil},
 		{"DEL", "mykey", "myid"}, {"1"},
 		{"GET", "mykey", "myid"}, {nil},
 	})