Merge pull request #131 from kyeett/add-georadiusbymember

add command GEORADIUSBYMEMBER

Author: alicebob

cmd_geo.go

@@ -18,6 +18,8 @@ func commandsGeo(m *Miniredis) {
 	m.srv.Register("GEOPOS", m.cmdGeopos)
 	m.srv.Register("GEORADIUS", m.cmdGeoradius)
 	m.srv.Register("GEORADIUS_RO", m.cmdGeoradius)
+	m.srv.Register("GEORADIUSBYMEMBER", m.cmdGeoradiusbymember)
+	m.srv.Register("GEORADIUSBYMEMBER_RO", m.cmdGeoradiusbymember)
 }
 
 // GEOADD
@@ -370,6 +372,200 @@ func (m *Miniredis) cmdGeoradius(c *server.Peer, cmd string, args []string) {
 	})
 }
 
+// GEORADIUSBYMEMBER and GEORADIUSBYMEMBER_RO
+func (m *Miniredis) cmdGeoradiusbymember(c *server.Peer, cmd string, args []string) {
+	if len(args) < 4 {
+		setDirty(c)
+		c.WriteError(errWrongNumber(cmd))
+		return
+	}
+	if !m.handleAuth(c) {
+		return
+	}
+	if m.checkPubsub(c) {
+		return
+	}
+
+	key := args[0]
+	member := args[1]
+
+	radius, err := strconv.ParseFloat(args[2], 64)
+	if err != nil || radius < 0 {
+		setDirty(c)
+		c.WriteError(errWrongNumber(cmd))
+		return
+	}
+	toMeter := parseUnit(args[3])
+	if toMeter == 0 {
+		setDirty(c)
+		c.WriteError(errWrongNumber(cmd))
+		return
+	}
+	args = args[4:]
+
+	var (
+		withDist      = false
+		withCoord     = false
+		direction     = unsorted
+		count         = 0
+		withStore     = false
+		storeKey      = ""
+		withStoredist = false
+		storedistKey  = ""
+	)
+	for len(args) > 0 {
+		arg := args[0]
+		args = args[1:]
+		switch strings.ToUpper(arg) {
+		case "WITHCOORD":
+			withCoord = true
+		case "WITHDIST":
+			withDist = true
+		case "ASC":
+			direction = asc
+		case "DESC":
+			direction = desc
+		case "COUNT":
+			if len(args) == 0 {
+				setDirty(c)
+				c.WriteError("ERR syntax error")
+				return
+			}
+			n, err := strconv.Atoi(args[0])
+			if err != nil {
+				setDirty(c)
+				c.WriteError(msgInvalidInt)
+				return
+			}
+			if n <= 0 {
+				setDirty(c)
+				c.WriteError("ERR COUNT must be > 0")
+				return
+			}
+			args = args[1:]
+			count = n
+		case "STORE":
+			if len(args) == 0 {
+				setDirty(c)
+				c.WriteError("ERR syntax error")
+				return
+			}
+			withStore = true
+			storeKey = args[0]
+			args = args[1:]
+		case "STOREDIST":
+			if len(args) == 0 {
+				setDirty(c)
+				c.WriteError("ERR syntax error")
+				return
+			}
+			withStoredist = true
+			storedistKey = args[0]
+			args = args[1:]
+		default:
+			setDirty(c)
+			c.WriteError("ERR syntax error")
+			return
+		}
+	}
+
+	if strings.ToUpper(cmd) == "GEORADIUSBYMEMBER_RO" && (withStore || withStoredist) {
+		setDirty(c)
+		c.WriteError("ERR syntax error")
+		return
+	}
+
+	withTx(m, c, func(c *server.Peer, ctx *connCtx) {
+		if (withStore || withStoredist) && (withDist || withCoord) {
+			c.WriteError("ERR STORE option in GEORADIUS is not compatible with WITHDIST, WITHHASH and WITHCOORDS options")
+			return
+		}
+
+		db := m.db(ctx.selectedDB)
+		if !db.exists(key) {
+			c.WriteNull()
+			return
+		}
+
+		if db.t(key) != "zset" {
+			c.WriteError(ErrWrongType.Error())
+			return
+		}
+
+		// get position of member
+		if !db.ssetExists(key, member) {
+			c.WriteError("ERR could not decode requested zset member")
+			return
+		}
+		score := db.ssetScore(key, member)
+		longitude, latitude := fromGeohash(uint64(score))
+
+		members := db.ssetElements(key)
+		matches := withinRadius(members, longitude, latitude, radius*toMeter)
+
+		// deal with ASC/DESC
+		if direction != unsorted {
+			sort.Slice(matches, func(i, j int) bool {
+				if direction == desc {
+					return matches[i].Distance > matches[j].Distance
+				}
+				return matches[i].Distance < matches[j].Distance
+			})
+		}
+
+		// deal with COUNT
+		if count > 0 && len(matches) > count {
+			matches = matches[:count]
+		}
+
+		// deal with "STORE x"
+		if withStore {
+			db.del(storeKey, true)
+			for _, member := range matches {
+				db.ssetAdd(storeKey, member.Score, member.Name)
+			}
+			c.WriteInt(len(matches))
+			return
+		}
+
+		// deal with "STOREDIST x"
+		if withStoredist {
+			db.del(storedistKey, true)
+			for _, member := range matches {
+				db.ssetAdd(storedistKey, member.Distance/toMeter, member.Name)
+			}
+			c.WriteInt(len(matches))
+			return
+		}
+
+		c.WriteLen(len(matches))
+		for _, member := range matches {
+			if !withDist && !withCoord {
+				c.WriteBulk(member.Name)
+				continue
+			}
+
+			len := 1
+			if withDist {
+				len++
+			}
+			if withCoord {
+				len++
+			}
+			c.WriteLen(len)
+			c.WriteBulk(member.Name)
+			if withDist {
+				c.WriteBulk(fmt.Sprintf("%.4f", member.Distance/toMeter))
+			}
+			if withCoord {
+				c.WriteLen(2)
+				c.WriteBulk(fmt.Sprintf("%f", member.Longitude))
+				c.WriteBulk(fmt.Sprintf("%f", member.Latitude))
+			}
+		}
+	})
+}
+
 func withinRadius(members []ssElem, longitude, latitude, radius float64) []geoDistance {
 	matches := []geoDistance{}
 	for _, el := range members {

cmd_geo_test.go

@@ -289,3 +289,160 @@ func TestGeodist(t *testing.T) {
 		mustFail(t, err, "WRONGTYPE Operation against a key holding the wrong kind of value")
 	})
 }
+
+// Test GEOADD / GEORADIUSBYMEMBER / GEORADIUSBYMEMBER_RO
+func TestGeobymember(t *testing.T) {
+	s, err := Run()
+	ok(t, err)
+	defer s.Close()
+	c, err := redis.Dial("tcp", s.Addr())
+	ok(t, err)
+	defer c.Close()
+
+	_, err = c.Do("GEOADD", "Sicily", 13.361389, 38.115556, "Palermo")
+	ok(t, err)
+	_, err = c.Do("GEOADD", "Sicily", 15.087269, 37.502669, "Catania")
+	ok(t, err)
+
+	t.Run("WITHDIST WITHCOORD", func(t *testing.T) {
+		res, err := redis.Values(c.Do("GEORADIUSBYMEMBER", "Sicily", "Palermo", 200, "km", "WITHDIST", "WITHCOORD"))
+		ok(t, err)
+		equals(t, 2, len(res))
+		for _, loc := range res {
+			item := loc.([]interface{})
+			var (
+				name, _  = redis.String(item[0], nil)
+				coord, _ = redis.Float64s(item[2], nil)
+			)
+			if name != "Catania" && name != "Palermo" {
+				t.Errorf("unexpected name %q", name)
+			}
+
+			equals(t, 2, len(coord))
+			if coord[0] == 0.00 || coord[1] == 0.00 {
+				t.Errorf("latitude/longitude shouldn't be empty")
+			}
+		}
+	})
+
+	t.Run("WITHCOORD", func(t *testing.T) {
+		res, err := redis.Values(c.Do("GEORADIUSBYMEMBER", "Sicily", "Palermo", 200, "km", "WITHCOORD"))
+		ok(t, err)
+		equals(t, 2, len(res))
+		for _, loc := range res {
+			item := loc.([]interface{})
+			var (
+				name, _  = redis.String(item[0], nil)
+				coord, _ = redis.Float64s(item[1], nil)
+			)
+			equals(t, 2, len(item))
+			if name != "Catania" && name != "Palermo" {
+				t.Errorf("unexpected name %q", name)
+			}
+			equals(t, 2, len(coord))
+			if coord[0] == 0.00 || coord[1] == 0.00 {
+				t.Errorf("latitude/longitude shouldn't be empty")
+			}
+		}
+	})
+
+	t.Run("WITHDIST", func(t *testing.T) {
+		// in km
+		res, err := redis.Values(c.Do("GEORADIUSBYMEMBER", "Sicily", "Palermo", 200, "km", "WITHDIST"))
+		ok(t, err)
+		equals(t, 2, len(res))
+		var (
+			name1, name2 string
+			dist1, dist2 float64
+		)
+		leftover, err := redis.Scan(res[0].([]interface{}), &name1, &dist1)
+		ok(t, err)
+		equals(t, 0, len(leftover))
+		equals(t, "Palermo", name1)
+		equals(t, 0.0, dist1) // in km
+		_, err = redis.Scan(res[1].([]interface{}), &name2, &dist2)
+		ok(t, err)
+		equals(t, "Catania", name2)
+		equals(t, 166.2742, dist2)
+
+		// in meter
+		res, err = redis.Values(c.Do("GEORADIUSBYMEMBER", "Sicily", "Catania", 200000, "m", "WITHDIST"))
+		ok(t, err)
+		equals(t, 2, len(res))
+		distance, err := redis.Float64(res[0].([]interface{})[1], nil)
+		ok(t, err)
+		equals(t, 166274.1514, distance) // in meter
+	})
+
+	t.Run("ASC DESC", func(t *testing.T) {
+		asc, err := redis.Strings(c.Do("GEORADIUSBYMEMBER", "Sicily", "Palermo", 200, "km", "ASC"))
+		ok(t, err)
+		equals(t, []string{"Palermo", "Catania"}, asc)
+
+		asc2, err := redis.Strings(c.Do("GEORADIUSBYMEMBER", "Sicily", "Catania", 200, "km", "ASC"))
+		ok(t, err)
+		equals(t, []string{"Catania", "Palermo"}, asc2)
+
+		desc, err := redis.Strings(c.Do("GEORADIUSBYMEMBER", "Sicily", "Palermo", 200, "km", "DESC"))
+		ok(t, err)
+		equals(t, []string{"Catania", "Palermo"}, desc)
+	})
+
+	t.Run("COUNT", func(t *testing.T) {
+		count1, err := redis.Strings(c.Do("GEORADIUSBYMEMBER", "Sicily", "Palermo", 200, "km", "ASC", "COUNT", 1))
+		ok(t, err)
+		equals(t, []string{"Palermo"}, count1)
+
+		count99, err := redis.Strings(c.Do("GEORADIUSBYMEMBER", "Sicily", "Palermo", 200, "km", "ASC", "COUNT", 99))
+		ok(t, err)
+		equals(t, []string{"Palermo", "Catania"}, count99)
+
+		_, err = c.Do("GEORADIUSBYMEMBER", "Sicily", "Palermo", 200, "km", "COUNT")
+		mustFail(t, err, "ERR syntax error")
+
+		_, err = c.Do("GEORADIUSBYMEMBER", "Sicily", "Palermo", 200, "km", "COUNT", "notanumber")
+		mustFail(t, err, msgInvalidInt)
+
+		_, err = c.Do("GEORADIUSBYMEMBER", "Sicily", "Palermo", 200, "km", "COUNT", -12)
+		mustFail(t, err, "ERR COUNT must be > 0")
+	})
+
+	t.Run("no args", func(t *testing.T) {
+		res, err := redis.Strings(c.Do("GEORADIUSBYMEMBER", "Sicily", "Palermo", 200, "km"))
+		ok(t, err)
+		equals(t, 2, len(res))
+		equals(t, []string{"Palermo", "Catania"}, res)
+
+		// Wrong map key
+		n, err := c.Do("GEORADIUSBYMEMBER", "Capri", "Palermo", 200, "km")
+		ok(t, err)
+		equals(t, nil, n)
+
+		// Missing member
+		res, err = redis.Strings(c.Do("GEORADIUSBYMEMBER", "Sicily", "nosuch", 200, "km"))
+		mustFail(t, err, "ERR could not decode requested zset member")
+		equals(t, 0, len(res))
+
+		// Unsupported/unknown distance unit
+		res, err = redis.Strings(c.Do("GEORADIUSBYMEMBER", "Sicily", "Palermo", 200, "mm"))
+		mustFail(t, err, "ERR wrong number of arguments for 'georadiusbymember' command")
+		equals(t, 0, len(res))
+
+		// Wrong parameter type
+		res, err = redis.Strings(c.Do("GEORADIUSBYMEMBER", "Sicily", "abc", "def", "ghi", "m"))
+		mustFail(t, err, "ERR wrong number of arguments for 'georadiusbymember' command")
+		equals(t, 0, len(res))
+	})
+
+	t.Run("GEORADIUSBYMEMBER_RO", func(t *testing.T) {
+		asc, err := redis.Strings(c.Do("GEORADIUSBYMEMBER_RO", "Sicily", "Palermo", 200, "km", "ASC"))
+		ok(t, err)
+		equals(t, []string{"Palermo", "Catania"}, asc)
+
+		_, err = c.Do("GEORADIUSBYMEMBER_RO", "Sicily", "Palermo", 200, "km", "STORE", "foo")
+		mustFail(t, err, "ERR syntax error")
+
+		_, err = c.Do("GEORADIUSBYMEMBER_RO", "Sicily", "Palermo", 200, "km", "STOREDIST", "foo")
+		mustFail(t, err, "ERR syntax error")
+	})
+}