diff --git a/command.go b/command.go index 77d0bada4..6f77258bd 100644 --- a/command.go +++ b/command.go @@ -2637,6 +2637,188 @@ func newGeoLocationParser(q *GeoRadiusQuery) proto.MultiBulkParse { //------------------------------------------------------------------------------ +// GeoSearchQuery is used for GEOSearch/GEOSearchStore command query. +type GeoSearchQuery struct { + Member string + + // Latitude and Longitude when using FromLonLat option. + Longitude float64 + Latitude float64 + + // Distance and unit when using ByRadius option. + // Can use m, km, ft, or mi. Default is km. + Radius float64 + RadiusUnit string + + // Height, width and unit when using ByBox option. + // Can be m, km, ft, or mi. Default is km. + BoxWidth float64 + BoxHeight float64 + BoxUnit string + + // Can be ASC or DESC. Default is no sort order. + Sort string + Count int + CountAny bool +} + +type GeoSearchLocationQuery struct { + GeoSearchQuery + + WithCoord bool + WithDist bool + WithHash bool +} + +type GeoSearchStoreQuery struct { + GeoSearchQuery + + // When using the StoreDist option, the command stores the items in a + // sorted set populated with their distance from the center of the circle or box, + // as a floating-point number, in the same unit specified for that shape. + StoreDist bool +} + +func geoSearchLocationArgs(q *GeoSearchLocationQuery, args []interface{}) []interface{} { + args = geoSearchArgs(&q.GeoSearchQuery, args) + + if q.WithCoord { + args = append(args, "withcoord") + } + if q.WithDist { + args = append(args, "withdist") + } + if q.WithHash { + args = append(args, "withhash") + } + + return args +} + +func geoSearchArgs(q *GeoSearchQuery, args []interface{}) []interface{} { + if q.Member != "" { + args = append(args, "frommember", q.Member) + } else { + args = append(args, "fromlonlat", q.Longitude, q.Latitude) + } + + if q.Radius > 0 { + if q.RadiusUnit == "" { + q.RadiusUnit = "km" + } + args = append(args, "byradius", q.Radius, q.RadiusUnit) + } else { + if q.BoxUnit == "" { + q.BoxUnit = "km" + } + args = append(args, "bybox", q.BoxWidth, q.BoxHeight, q.BoxUnit) + } + + if q.Sort != "" { + args = append(args, q.Sort) + } + + if q.Count > 0 { + args = append(args, "count", q.Count) + if q.CountAny { + args = append(args, "any") + } + } + + return args +} + +type GeoSearchLocationCmd struct { + baseCmd + + opt *GeoSearchLocationQuery + val []GeoLocation +} + +var _ Cmder = (*GeoSearchLocationCmd)(nil) + +func NewGeoSearchLocationCmd( + ctx context.Context, opt *GeoSearchLocationQuery, args ...interface{}, +) *GeoSearchLocationCmd { + return &GeoSearchLocationCmd{ + baseCmd: baseCmd{ + ctx: ctx, + args: args, + }, + opt: opt, + } +} + +func (cmd *GeoSearchLocationCmd) Val() []GeoLocation { + return cmd.val +} + +func (cmd *GeoSearchLocationCmd) Result() ([]GeoLocation, error) { + return cmd.val, cmd.err +} + +func (cmd *GeoSearchLocationCmd) String() string { + return cmdString(cmd, cmd.val) +} + +func (cmd *GeoSearchLocationCmd) readReply(rd *proto.Reader) error { + n, err := rd.ReadArrayLen() + if err != nil { + return err + } + + cmd.val = make([]GeoLocation, n) + for i := 0; i < n; i++ { + _, err = rd.ReadArrayLen() + if err != nil { + return err + } + + var loc GeoLocation + + loc.Name, err = rd.ReadString() + if err != nil { + return err + } + if cmd.opt.WithDist { + loc.Dist, err = rd.ReadFloatReply() + if err != nil { + return err + } + } + if cmd.opt.WithHash { + loc.GeoHash, err = rd.ReadIntReply() + if err != nil { + return err + } + } + if cmd.opt.WithCoord { + nn, err := rd.ReadArrayLen() + if err != nil { + return err + } + if nn != 2 { + return fmt.Errorf("got %d coordinates, expected 2", nn) + } + + loc.Longitude, err = rd.ReadFloatReply() + if err != nil { + return err + } + loc.Latitude, err = rd.ReadFloatReply() + if err != nil { + return err + } + } + + cmd.val[i] = loc + } + + return nil +} + +//------------------------------------------------------------------------------ + type GeoPos struct { Longitude, Latitude float64 } diff --git a/commands.go b/commands.go index 8fd039703..fb1fd82ed 100644 --- a/commands.go +++ b/commands.go @@ -383,6 +383,9 @@ type Cmdable interface { GeoRadiusStore(ctx context.Context, key string, longitude, latitude float64, query *GeoRadiusQuery) *IntCmd GeoRadiusByMember(ctx context.Context, key, member string, query *GeoRadiusQuery) *GeoLocationCmd GeoRadiusByMemberStore(ctx context.Context, key, member string, query *GeoRadiusQuery) *IntCmd + GeoSearch(ctx context.Context, key string, q *GeoSearchQuery) *StringSliceCmd + GeoSearchLocation(ctx context.Context, key string, q *GeoSearchLocationQuery) *GeoSearchLocationCmd + GeoSearchStore(ctx context.Context, key, store string, q *GeoSearchStoreQuery) *IntCmd GeoDist(ctx context.Context, key string, member1, member2, unit string) *FloatCmd GeoHash(ctx context.Context, key string, members ...string) *StringSliceCmd } @@ -3347,6 +3350,38 @@ func (c cmdable) GeoRadiusByMemberStore( return cmd } +func (c cmdable) GeoSearch(ctx context.Context, key string, q *GeoSearchQuery) *StringSliceCmd { + args := make([]interface{}, 0, 13) + args = append(args, "geosearch", key) + args = geoSearchArgs(q, args) + cmd := NewStringSliceCmd(ctx, args...) + _ = c(ctx, cmd) + return cmd +} + +func (c cmdable) GeoSearchLocation( + ctx context.Context, key string, q *GeoSearchLocationQuery, +) *GeoSearchLocationCmd { + args := make([]interface{}, 0, 16) + args = append(args, "geosearch", key) + args = geoSearchLocationArgs(q, args) + cmd := NewGeoSearchLocationCmd(ctx, q, args...) + _ = c(ctx, cmd) + return cmd +} + +func (c cmdable) GeoSearchStore(ctx context.Context, key, store string, q *GeoSearchStoreQuery) *IntCmd { + args := make([]interface{}, 0, 15) + args = append(args, "geosearchstore", store, key) + args = geoSearchArgs(&q.GeoSearchQuery, args) + if q.StoreDist { + args = append(args, "storedist") + } + cmd := NewIntCmd(ctx, args...) + _ = c(ctx, cmd) + return cmd +} + func (c cmdable) GeoDist( ctx context.Context, key string, member1, member2, unit string, ) *FloatCmd { diff --git a/commands_test.go b/commands_test.go index 736add761..a331f7f98 100644 --- a/commands_test.go +++ b/commands_test.go @@ -5142,6 +5142,204 @@ var _ = Describe("Commands", func() { nil, })) }) + + It("should geo search", func() { + q := &redis.GeoSearchQuery{ + Member: "Catania", + BoxWidth: 400, + BoxHeight: 100, + BoxUnit: "km", + Sort: "asc", + } + val, err := client.GeoSearch(ctx, "Sicily", q).Result() + Expect(err).NotTo(HaveOccurred()) + Expect(val).To(Equal([]string{"Catania"})) + + q.BoxHeight = 400 + val, err = client.GeoSearch(ctx, "Sicily", q).Result() + Expect(err).NotTo(HaveOccurred()) + Expect(val).To(Equal([]string{"Catania", "Palermo"})) + + q.Count = 1 + val, err = client.GeoSearch(ctx, "Sicily", q).Result() + Expect(err).NotTo(HaveOccurred()) + Expect(val).To(Equal([]string{"Catania"})) + + q.CountAny = true + val, err = client.GeoSearch(ctx, "Sicily", q).Result() + Expect(err).NotTo(HaveOccurred()) + Expect(val).To(Equal([]string{"Palermo"})) + + q = &redis.GeoSearchQuery{ + Member: "Catania", + Radius: 100, + RadiusUnit: "km", + Sort: "asc", + } + val, err = client.GeoSearch(ctx, "Sicily", q).Result() + Expect(err).NotTo(HaveOccurred()) + Expect(val).To(Equal([]string{"Catania"})) + + q.Radius = 400 + val, err = client.GeoSearch(ctx, "Sicily", q).Result() + Expect(err).NotTo(HaveOccurred()) + Expect(val).To(Equal([]string{"Catania", "Palermo"})) + + q.Count = 1 + val, err = client.GeoSearch(ctx, "Sicily", q).Result() + Expect(err).NotTo(HaveOccurred()) + Expect(val).To(Equal([]string{"Catania"})) + + q.CountAny = true + val, err = client.GeoSearch(ctx, "Sicily", q).Result() + Expect(err).NotTo(HaveOccurred()) + Expect(val).To(Equal([]string{"Palermo"})) + + q = &redis.GeoSearchQuery{ + Longitude: 15, + Latitude: 37, + BoxWidth: 200, + BoxHeight: 200, + BoxUnit: "km", + Sort: "asc", + } + val, err = client.GeoSearch(ctx, "Sicily", q).Result() + Expect(err).NotTo(HaveOccurred()) + Expect(val).To(Equal([]string{"Catania"})) + + q.BoxWidth, q.BoxHeight = 400, 400 + val, err = client.GeoSearch(ctx, "Sicily", q).Result() + Expect(err).NotTo(HaveOccurred()) + Expect(val).To(Equal([]string{"Catania", "Palermo"})) + + q.Count = 1 + val, err = client.GeoSearch(ctx, "Sicily", q).Result() + Expect(err).NotTo(HaveOccurred()) + Expect(val).To(Equal([]string{"Catania"})) + + q.CountAny = true + val, err = client.GeoSearch(ctx, "Sicily", q).Result() + Expect(err).NotTo(HaveOccurred()) + Expect(val).To(Equal([]string{"Palermo"})) + + q = &redis.GeoSearchQuery{ + Longitude: 15, + Latitude: 37, + Radius: 100, + RadiusUnit: "km", + Sort: "asc", + } + val, err = client.GeoSearch(ctx, "Sicily", q).Result() + Expect(err).NotTo(HaveOccurred()) + Expect(val).To(Equal([]string{"Catania"})) + + q.Radius = 200 + val, err = client.GeoSearch(ctx, "Sicily", q).Result() + Expect(err).NotTo(HaveOccurred()) + Expect(val).To(Equal([]string{"Catania", "Palermo"})) + + q.Count = 1 + val, err = client.GeoSearch(ctx, "Sicily", q).Result() + Expect(err).NotTo(HaveOccurred()) + Expect(val).To(Equal([]string{"Catania"})) + + q.CountAny = true + val, err = client.GeoSearch(ctx, "Sicily", q).Result() + Expect(err).NotTo(HaveOccurred()) + Expect(val).To(Equal([]string{"Palermo"})) + }) + + It("should geo search with options", func() { + q := &redis.GeoSearchLocationQuery{ + GeoSearchQuery: redis.GeoSearchQuery{ + Longitude: 15, + Latitude: 37, + Radius: 200, + RadiusUnit: "km", + Sort: "asc", + }, + WithHash: true, + WithDist: true, + WithCoord: true, + } + val, err := client.GeoSearchLocation(ctx, "Sicily", q).Result() + Expect(err).NotTo(HaveOccurred()) + Expect(val).To(Equal([]redis.GeoLocation{ + { + Name: "Catania", + Longitude: 15.08726745843887329, + Latitude: 37.50266842333162032, + Dist: 56.4413, + GeoHash: 3479447370796909, + }, + { + Name: "Palermo", + Longitude: 13.36138933897018433, + Latitude: 38.11555639549629859, + Dist: 190.4424, + GeoHash: 3479099956230698, + }, + })) + }) + + It("should geo search store", func() { + q := &redis.GeoSearchStoreQuery{ + GeoSearchQuery: redis.GeoSearchQuery{ + Longitude: 15, + Latitude: 37, + Radius: 200, + RadiusUnit: "km", + Sort: "asc", + }, + StoreDist: false, + } + + val, err := client.GeoSearchStore(ctx, "Sicily", "key1", q).Result() + Expect(err).NotTo(HaveOccurred()) + Expect(val).To(Equal(int64(2))) + + q.StoreDist = true + val, err = client.GeoSearchStore(ctx, "Sicily", "key2", q).Result() + Expect(err).NotTo(HaveOccurred()) + Expect(val).To(Equal(int64(2))) + + loc, err := client.GeoSearchLocation(ctx, "key1", &redis.GeoSearchLocationQuery{ + GeoSearchQuery: q.GeoSearchQuery, + WithCoord: true, + WithDist: true, + WithHash: true, + }).Result() + Expect(err).NotTo(HaveOccurred()) + Expect(loc).To(Equal([]redis.GeoLocation{ + { + Name: "Catania", + Longitude: 15.08726745843887329, + Latitude: 37.50266842333162032, + Dist: 56.4413, + GeoHash: 3479447370796909, + }, + { + Name: "Palermo", + Longitude: 13.36138933897018433, + Latitude: 38.11555639549629859, + Dist: 190.4424, + GeoHash: 3479099956230698, + }, + })) + + v, err := client.ZRangeWithScores(ctx, "key2", 0, -1).Result() + Expect(err).NotTo(HaveOccurred()) + Expect(v).To(Equal([]redis.Z{ + { + Score: 56.441257870158204, + Member: "Catania", + }, + { + Score: 190.44242984775784, + Member: "Palermo", + }, + })) + }) }) Describe("marshaling/unmarshaling", func() {