Merge pull request #1119 from go-redis/fix/geo-radious-store

Add writing version of GeoRadius commands

command.go

@@ -1580,6 +1580,13 @@ type GeoLocationCmd struct {
 var _ Cmder = (*GeoLocationCmd)(nil)
 
 func NewGeoLocationCmd(q *GeoRadiusQuery, args ...interface{}) *GeoLocationCmd {
+	return &GeoLocationCmd{
+		baseCmd: baseCmd{_args: geoLocationArgs(q, args...)},
+		q:       q,
+	}
+}
+
+func geoLocationArgs(q *GeoRadiusQuery, args ...interface{}) []interface{} {
 	args = append(args, q.Radius)
 	if q.Unit != "" {
 		args = append(args, q.Unit)
@@ -1609,10 +1616,7 @@ func NewGeoLocationCmd(q *GeoRadiusQuery, args ...interface{}) *GeoLocationCmd {
 		args = append(args, "storedist")
 		args = append(args, q.StoreDist)
 	}
-	return &GeoLocationCmd{
+	return args
-		baseCmd: baseCmd{_args: args},
-		q:       q,
-	}
 }
 
 func (cmd *GeoLocationCmd) Val() []GeoLocation {

commands.go

@@ -282,9 +282,9 @@ type Cmdable interface {
 	GeoAdd(key string, geoLocation ...*GeoLocation) *IntCmd
 	GeoPos(key string, members ...string) *GeoPosCmd
 	GeoRadius(key string, longitude, latitude float64, query *GeoRadiusQuery) *GeoLocationCmd
-	GeoRadiusRO(key string, longitude, latitude float64, query *GeoRadiusQuery) *GeoLocationCmd
+	GeoRadiusStore(key string, longitude, latitude float64, query *GeoRadiusQuery) *IntCmd
 	GeoRadiusByMember(key, member string, query *GeoRadiusQuery) *GeoLocationCmd
-	GeoRadiusByMemberRO(key, member string, query *GeoRadiusQuery) *GeoLocationCmd
+	GeoRadiusByMemberStore(key, member string, query *GeoRadiusQuery) *IntCmd
 	GeoDist(key string, member1, member2, unit string) *FloatCmd
 	GeoHash(key string, members ...string) *StringSliceCmd
 	ReadOnly() *StatusCmd
@@ -2513,26 +2513,48 @@ func (c cmdable) GeoAdd(key string, geoLocation ...*GeoLocation) *IntCmd {
 	return cmd
 }
 
+// GeoRadius is a read-only GEORADIUS_RO command.
 func (c cmdable) GeoRadius(key string, longitude, latitude float64, query *GeoRadiusQuery) *GeoLocationCmd {
-	cmd := NewGeoLocationCmd(query, "georadius", key, longitude, latitude)
-	_ = c(cmd)
-	return cmd
-}
-
-func (c cmdable) GeoRadiusRO(key string, longitude, latitude float64, query *GeoRadiusQuery) *GeoLocationCmd {
 	cmd := NewGeoLocationCmd(query, "georadius_ro", key, longitude, latitude)
+	if query.Store != "" || query.StoreDist != "" {
+		cmd.setErr(errors.New("GeoRadius does not support Store or StoreDist"))
+		return cmd
+	}
 	_ = c(cmd)
 	return cmd
 }
 
+// GeoRadiusStore is a writing GEORADIUS command.
+func (c cmdable) GeoRadiusStore(key string, longitude, latitude float64, query *GeoRadiusQuery) *IntCmd {
+	args := geoLocationArgs(query, "georadius", key, longitude, latitude)
+	cmd := NewIntCmd(args...)
+	if query.Store == "" && query.StoreDist == "" {
+		cmd.setErr(errors.New("GeoRadiusStore requires Store or StoreDist"))
+		return cmd
+	}
+	_ = c(cmd)
+	return cmd
+}
+
+// GeoRadius is a read-only GEORADIUSBYMEMBER_RO command.
 func (c cmdable) GeoRadiusByMember(key, member string, query *GeoRadiusQuery) *GeoLocationCmd {
-	cmd := NewGeoLocationCmd(query, "georadiusbymember", key, member)
+	cmd := NewGeoLocationCmd(query, "georadiusbymember_ro", key, member)
+	if query.Store != "" || query.StoreDist != "" {
+		cmd.setErr(errors.New("GeoRadiusByMember does not support Store or StoreDist"))
+		return cmd
+	}
 	_ = c(cmd)
 	return cmd
 }
 
-func (c cmdable) GeoRadiusByMemberRO(key, member string, query *GeoRadiusQuery) *GeoLocationCmd {
+// GeoRadiusByMemberStore is a writing GEORADIUSBYMEMBER command.
-	cmd := NewGeoLocationCmd(query, "georadiusbymember_ro", key, member)
+func (c cmdable) GeoRadiusByMemberStore(key, member string, query *GeoRadiusQuery) *IntCmd {
+	args := geoLocationArgs(query, "georadiusbymember", key, member)
+	cmd := NewIntCmd(args...)
+	if query.Store == "" && query.StoreDist == "" {
+		cmd.setErr(errors.New("GeoRadiusByMemberStore requires Store or StoreDist"))
+		return cmd
+	}
 	_ = c(cmd)
 	return cmd
 }

commands_test.go

@@ -3733,6 +3733,46 @@ var _ = Describe("Commands", func() {
 			Expect(res[1].Name).To(Equal("Catania"))
 		})
 
+		It("should geo radius and store the result", func() {
+			n, err := client.GeoRadiusStore("Sicily", 15, 37, &redis.GeoRadiusQuery{
+				Radius: 200,
+				Store:  "result",
+			}).Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(n).To(Equal(int64(2)))
+
+			res, err := client.ZRangeWithScores("result", 0, -1).Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(res).To(ContainElement(redis.Z{
+				Score:  3.479099956230698e+15,
+				Member: "Palermo",
+			}))
+			Expect(res).To(ContainElement(redis.Z{
+				Score:  3.479447370796909e+15,
+				Member: "Catania",
+			}))
+		})
+
+		It("should geo radius and store dist", func() {
+			n, err := client.GeoRadiusStore("Sicily", 15, 37, &redis.GeoRadiusQuery{
+				Radius:    200,
+				StoreDist: "result",
+			}).Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(n).To(Equal(int64(2)))
+
+			res, err := client.ZRangeWithScores("result", 0, -1).Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(res).To(ContainElement(redis.Z{
+				Score:  190.44242984775784,
+				Member: "Palermo",
+			}))
+			Expect(res).To(ContainElement(redis.Z{
+				Score:  56.4412578701582,
+				Member: "Catania",
+			}))
+		})
+
 		It("should search geo radius with options", func() {
 			res, err := client.GeoRadius("Sicily", 15, 37, &redis.GeoRadiusQuery{
 				Radius:      200,

internal/pool/pool.go

@@ -310,6 +310,12 @@ func (p *ConnPool) popIdle() *Conn {
 }
 
 func (p *ConnPool) Put(cn *Conn) {
+	if cn.rd.Buffered() > 0 {
+		internal.Logger.Printf("Conn has unread data")
+		p.Remove(cn, BadConnError{})
+		return
+	}
+
 	if !cn.pooled {
 		p.Remove(cn, nil)
 		return

internal/proto/reader.go

@@ -41,6 +41,10 @@ func NewReader(rd io.Reader) *Reader {
 	}
 }
 
+func (r *Reader) Buffered() int {
+	return r.rd.Buffered()
+}
+
 func (r *Reader) Reset(rd io.Reader) {
 	r.rd.Reset(rd)
 }