Added binding for GEOADD and GEORADIUS.

ibc789 · Sep 28, 2015 · 8a05670 · 8a05670
1 parent eef3fd7
commit 8a05670
Show file tree

Hide file tree

Showing 3 changed files with 231 additions and 0 deletions.
diff --git a/command.go b/command.go
@@ -783,3 +783,102 @@ func (cmd *ClusterSlotCmd) parseReply(cn *conn) error {
 	cmd.val = v.([]ClusterSlotInfo)
 	return nil
 }
+
+//------------------------------------------------------------------------------
+
+// Location type for GEO operations in Redis
+type GeoLocation struct {
+	Name string
+	Longitude, Latitude, Distance float64
+	GeoHash int64
+}
+
+type GeoCmd struct {
+	baseCmd
+
+	locations []GeoLocation
+}
+
+// Query type for geo radius
+type GeoRadiusQuery struct {
+	Key string
+	Longitude, Latitude, Radius float64
+	// Unit default to km when nil
+	Unit string
+	WithCoordinates, WithDistance, WithGeoHash bool
+	// Count default to 0 and ignored limit.
+	Count int
+	// Sort default to unsorted, ASC or DESC otherwise
+	Sort string
+}
+
+func NewGeoCmd(args ...interface{}) *GeoCmd {
+	return &GeoCmd{baseCmd: baseCmd{_args: args, _clusterKeyPos: 1}}
+}
+
+func (cmd *GeoCmd) reset() {
+	cmd.locations = nil
+	cmd.err = nil
+}
+
+func (cmd *GeoCmd) Val() ([]GeoLocation) {
+	return cmd.locations
+}
+
+func (cmd *GeoCmd) Result() ([]GeoLocation, error) {
+	return cmd.locations, cmd.err
+}
+
+func (cmd *GeoCmd) String() string {
+	return cmdString(cmd, cmd.locations)
+}
+
+func (cmd *GeoCmd) parseReply(cn *conn) error {
+	vi, err := parseReply(cn, parseSlice)
+	if err != nil {
+		cmd.err = err
+		return cmd.err
+	}
+
+	v := vi.([]interface{})
+
+	if len(v) == 0 {
+		return nil
+	}
+
+	if _, ok := v[0].(string); ok { // Location names only (single level string array)
+		for _, keyi := range v {
+			cmd.locations = append(cmd.locations, GeoLocation{Name: keyi.(string)})
+		}
+	} else { // Full location details (nested arrays)
+		for _, keyi := range v {
+			tmpLocation := GeoLocation{}
+			keyiface := keyi.([]interface{})
+			for _, subKeyi := range keyiface {
+				if strVal, ok := subKeyi.(string); ok {
+					if len(tmpLocation.Name) == 0 {
+						tmpLocation.Name = strVal
+					} else {
+						tmpLocation.Distance, err = strconv.ParseFloat(strVal, 64)
+						if err != nil {
+							return err
+						}
+					}
+				} else if intVal, ok := subKeyi.(int64); ok {
+					tmpLocation.GeoHash = intVal
+				} else if ifcVal, ok := subKeyi.([]interface{}); ok {
+					tmpLocation.Longitude, err = strconv.ParseFloat(ifcVal[0].(string), 64)
+					if err != nil {
+						return err
+					}
+					tmpLocation.Latitude, err = strconv.ParseFloat(ifcVal[1].(string), 64)
+					if err != nil {
+						return err
+					}
+				}
+			}
+			cmd.locations = append(cmd.locations, tmpLocation)
+		}
+	}
+	return nil
+}
diff --git a/commands.go b/commands.go
@@ -1671,3 +1671,75 @@ func (c *commandable) ClusterAddSlotsRange(min, max int) *StatusCmd {
 	}
 	return c.ClusterAddSlots(slots...)
 }
+
+//------------------------------------------------------------------------------
+
+func (c *commandable) GeoAdd(key string, geoLocation ...*GeoLocation) *IntCmd {
+	args := make([]interface{}, 2+3*len(geoLocation))
+	args[0] = "GEOADD"
+	args[1] = key
+	for i, eachLoc := range geoLocation {
+		args[2+3*i] = eachLoc.Longitude
+		args[2+3*i+1] = eachLoc.Latitude
+		args[2+3*i+2] = eachLoc.Name
+	}
+	cmd := NewIntCmd(args...)
+	c.Process(cmd)
+	return cmd
+}
+
+func (c *commandable) GeoRadius(query *GeoRadiusQuery) *GeoCmd {
+	var options, optionsCtr int
+	if query.WithCoordinates {
+		options++
+	}
+	if query.WithDistance {
+		options++
+	}
+	if query.WithGeoHash {
+		options++
+	}
+	if query.Count > 0 {
+		options += 2
+	}
+	if query.Sort != "" {
+		options++
+	}
+
+	args := make([]interface{}, 6 + options)
+	args[0] = "GEORADIUS"
+	args[1] = query.Key
+	args[2] = query.Longitude
+	args[3] = query.Latitude
+	args[4] = query.Radius
+	if query.Unit != "" {
+		args[5] = query.Unit
+	} else {
+		args[5] = "km"
+	}
+	if query.WithCoordinates {
+		args[6+optionsCtr] = "WITHCOORD"
+		optionsCtr++
+	}
+	if query.WithDistance {
+		args[6+optionsCtr] = "WITHDIST"
+		optionsCtr++
+	}
+	if query.WithGeoHash {
+		args[6+optionsCtr] = "WITHHASH"
+		optionsCtr++
+	}
+	if query.Count > 0 {
+		args[6+optionsCtr] = "COUNT"
+		optionsCtr++
+		args[6+optionsCtr] = query.Count
+		optionsCtr++
+	}
+	if query.Sort != "" {
+		args[6+optionsCtr] = query.Sort
+	}
+
+	cmd := NewGeoCmd(args...)
+	c.Process(cmd)
+	return cmd
+}
diff --git a/commands_test.go b/commands_test.go
@@ -2521,6 +2521,66 @@ var _ = Describe("Commands", func() {
 
 	})
 
+	Describe("Geo add and radius search", func() {
+		It("should add one geo location", func() {
+			geoAdd := client.GeoAdd("Sicily", &redis.GeoLocation{Longitude: 13.361389, Latitude: 38.115556, Name: "Palermo"})
+			Expect(geoAdd.Err()).NotTo(HaveOccurred())
+			Expect(geoAdd.Val()).To(Equal(int64(1)))
+		})
+
+		It("should add multiple geo locations", func() {
+			geoAdd := client.GeoAdd("Sicily", &redis.GeoLocation{Longitude: 13.361389, Latitude: 38.115556, Name: "Palermo"},
+				&redis.GeoLocation{Longitude: 15.087269, Latitude: 37.502669, Name: "Catania"})
+			Expect(geoAdd.Err()).NotTo(HaveOccurred())
+			Expect(geoAdd.Val()).To(Equal(int64(2)))
+		})
+
+		It("should search geo radius", func() {
+			geoAdd := client.GeoAdd("Sicily", &redis.GeoLocation{Longitude: 13.361389, Latitude: 38.115556, Name: "Palermo"},
+				&redis.GeoLocation{Longitude: 15.087269, Latitude: 37.502669, Name: "Catania"})
+			Expect(geoAdd.Err()).NotTo(HaveOccurred())
+			Expect(geoAdd.Val()).To(Equal(int64(2)))
+
+			geoRadius := client.GeoRadius(&redis.GeoRadiusQuery{Key: "Sicily", Longitude: 15, Latitude: 37, Radius: 200})
+			Expect(geoRadius.Err()).NotTo(HaveOccurred())
+			Expect(geoRadius.Val()[0].Name).To(Equal("Palermo"))
+			Expect(geoRadius.Val()[1].Name).To(Equal("Catania"))
+		})
+
+		It("should search geo radius with options", func() {
+			locations := []*redis.GeoLocation{&redis.GeoLocation{Longitude: 13.361389, Latitude: 38.115556, Name: "Palermo"},
+				&redis.GeoLocation{Longitude: 15.087269, Latitude: 37.502669, Name: "Catania"}}
+
+			geoAdd := client.GeoAdd("Sicily", locations...)
+			Expect(geoAdd.Err()).NotTo(HaveOccurred())
+			Expect(geoAdd.Val()).To(Equal(int64(2)))
+
+			geoRadius := client.GeoRadius(&redis.GeoRadiusQuery{Key: "Sicily", Longitude: 15, Latitude: 37, Radius: 200, Unit: "km", WithGeoHash: true, WithCoordinates: true, WithDistance: true, Count: 2, Sort: "ASC"})
+			Expect(geoRadius.Err()).NotTo(HaveOccurred())
+			Expect(geoRadius.Val()[1].Name).To(Equal("Palermo"))
+			Expect(geoRadius.Val()[1].Distance).To(Equal(190.4424))
+			Expect(geoRadius.Val()[1].GeoHash).To(Equal(int64(3479099956230698)))
+			Expect(geoRadius.Val()[1].Longitude).To(Equal(13.361389338970184))
+			Expect(geoRadius.Val()[1].Latitude).To(Equal(38.115556395496299))
+			Expect(geoRadius.Val()[0].Name).To(Equal("Catania"))
+			Expect(geoRadius.Val()[0].Distance).To(Equal(56.4413))
+			Expect(geoRadius.Val()[0].GeoHash).To(Equal(int64(3479447370796909)))
+			Expect(geoRadius.Val()[0].Longitude).To(Equal(15.087267458438873))
+			Expect(geoRadius.Val()[0].Latitude).To(Equal(37.50266842333162))
+		})
+
+		It("should search geo radius with no results", func() {
+			geoAdd := client.GeoAdd("Sicily", &redis.GeoLocation{Longitude: 13.361389, Latitude: 38.115556, Name: "Palermo"},
+				&redis.GeoLocation{Longitude: 15.087269, Latitude: 37.502669, Name: "Catania"})
+			Expect(geoAdd.Err()).NotTo(HaveOccurred())
+			Expect(geoAdd.Val()).To(Equal(int64(2)))
+
+			geoRadius := client.GeoRadius(&redis.GeoRadiusQuery{Key: "Sicily", Longitude: 99, Latitude: 37, Radius: 200, Unit: "km", WithGeoHash: true, WithCoordinates: true, WithDistance: true})
+			Expect(geoRadius.Err()).NotTo(HaveOccurred())
+			Expect(len(geoRadius.Val())).To(Equal(0))
+		})
+	})
+
 	Describe("marshaling/unmarshaling", func() {
 
 		type convTest struct {