From 023e33dffe48e3a602b7fadb214191ee6c6f2d2d Mon Sep 17 00:00:00 2001 From: monkey92t <golang@88.com> Date: Tue, 27 Jul 2021 18:12:19 +0800 Subject: [PATCH 1/2] add cmd: geosearch, geosearchstore Signed-off-by: monkey92t <golang@88.com> --- command.go | 182 +++++++++++++++++++++++++++++++++++++++++++ commands.go | 35 +++++++++ commands_test.go | 198 +++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 415 insertions(+) diff --git a/command.go b/command.go index 77d0bada4..00c251115 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 GeoSearchWithOptionsQuery 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 geoSearchWithOptionsArgs(q GeoSearchWithOptionsQuery, 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 GeoSearchWithOptionsQuery + val []GeoLocation +} + +var _ Cmder = (*GeoSearchLocationCmd)(nil) + +func NewGeoSearchLocationCmd( + ctx context.Context, opt GeoSearchWithOptionsQuery, 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..800eca584 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 + GeoSearchWithOptions(ctx context.Context, key string, q GeoSearchWithOptionsQuery) *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) GeoSearchWithOptions( + ctx context.Context, key string, q GeoSearchWithOptionsQuery, +) *GeoSearchLocationCmd { + args := make([]interface{}, 0, 16) + args = append(args, "geosearch", key) + args = geoSearchWithOptionsArgs(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..d74b71817 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.GeoSearchWithOptionsQuery{ + GeoSearchQuery: redis.GeoSearchQuery{ + Longitude: 15, + Latitude: 37, + Radius: 200, + RadiusUnit: "km", + Sort: "asc", + }, + WithHash: true, + WithDist: true, + WithCoord: true, + } + val, err := client.GeoSearchWithOptions(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.GeoSearchWithOptions(ctx, "key1", redis.GeoSearchWithOptionsQuery{ + 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() { From 3ca8a3c941565fd7b134481e2b02b310ea416aa2 Mon Sep 17 00:00:00 2001 From: monkey92t <golang@88.com> Date: Wed, 28 Jul 2021 10:27:28 +0800 Subject: [PATCH 2/2] GeoSearchQuery and GeoSearchLocationQuery changed to pointer passing Signed-off-by: monkey92t <golang@88.com> --- command.go | 12 ++++++------ commands.go | 18 +++++++++--------- commands_test.go | 16 ++++++++-------- 3 files changed, 23 insertions(+), 23 deletions(-) diff --git a/command.go b/command.go index 00c251115..6f77258bd 100644 --- a/command.go +++ b/command.go @@ -2662,7 +2662,7 @@ type GeoSearchQuery struct { CountAny bool } -type GeoSearchWithOptionsQuery struct { +type GeoSearchLocationQuery struct { GeoSearchQuery WithCoord bool @@ -2679,8 +2679,8 @@ type GeoSearchStoreQuery struct { StoreDist bool } -func geoSearchWithOptionsArgs(q GeoSearchWithOptionsQuery, args []interface{}) []interface{} { - args = geoSearchArgs(q.GeoSearchQuery, args) +func geoSearchLocationArgs(q *GeoSearchLocationQuery, args []interface{}) []interface{} { + args = geoSearchArgs(&q.GeoSearchQuery, args) if q.WithCoord { args = append(args, "withcoord") @@ -2695,7 +2695,7 @@ func geoSearchWithOptionsArgs(q GeoSearchWithOptionsQuery, args []interface{}) [ return args } -func geoSearchArgs(q GeoSearchQuery, args []interface{}) []interface{} { +func geoSearchArgs(q *GeoSearchQuery, args []interface{}) []interface{} { if q.Member != "" { args = append(args, "frommember", q.Member) } else { @@ -2731,14 +2731,14 @@ func geoSearchArgs(q GeoSearchQuery, args []interface{}) []interface{} { type GeoSearchLocationCmd struct { baseCmd - opt GeoSearchWithOptionsQuery + opt *GeoSearchLocationQuery val []GeoLocation } var _ Cmder = (*GeoSearchLocationCmd)(nil) func NewGeoSearchLocationCmd( - ctx context.Context, opt GeoSearchWithOptionsQuery, args ...interface{}, + ctx context.Context, opt *GeoSearchLocationQuery, args ...interface{}, ) *GeoSearchLocationCmd { return &GeoSearchLocationCmd{ baseCmd: baseCmd{ diff --git a/commands.go b/commands.go index 800eca584..fb1fd82ed 100644 --- a/commands.go +++ b/commands.go @@ -383,9 +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 - GeoSearchWithOptions(ctx context.Context, key string, q GeoSearchWithOptionsQuery) *GeoSearchLocationCmd - GeoSearchStore(ctx context.Context, key, store string, q GeoSearchStoreQuery) *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 } @@ -3350,7 +3350,7 @@ func (c cmdable) GeoRadiusByMemberStore( return cmd } -func (c cmdable) GeoSearch(ctx context.Context, key string, q GeoSearchQuery) *StringSliceCmd { +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) @@ -3359,21 +3359,21 @@ func (c cmdable) GeoSearch(ctx context.Context, key string, q GeoSearchQuery) *S return cmd } -func (c cmdable) GeoSearchWithOptions( - ctx context.Context, key string, q GeoSearchWithOptionsQuery, +func (c cmdable) GeoSearchLocation( + ctx context.Context, key string, q *GeoSearchLocationQuery, ) *GeoSearchLocationCmd { args := make([]interface{}, 0, 16) args = append(args, "geosearch", key) - args = geoSearchWithOptionsArgs(q, args) + 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 { +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) + args = geoSearchArgs(&q.GeoSearchQuery, args) if q.StoreDist { args = append(args, "storedist") } diff --git a/commands_test.go b/commands_test.go index d74b71817..a331f7f98 100644 --- a/commands_test.go +++ b/commands_test.go @@ -5144,7 +5144,7 @@ var _ = Describe("Commands", func() { }) It("should geo search", func() { - q := redis.GeoSearchQuery{ + q := &redis.GeoSearchQuery{ Member: "Catania", BoxWidth: 400, BoxHeight: 100, @@ -5170,7 +5170,7 @@ var _ = Describe("Commands", func() { Expect(err).NotTo(HaveOccurred()) Expect(val).To(Equal([]string{"Palermo"})) - q = redis.GeoSearchQuery{ + q = &redis.GeoSearchQuery{ Member: "Catania", Radius: 100, RadiusUnit: "km", @@ -5195,7 +5195,7 @@ var _ = Describe("Commands", func() { Expect(err).NotTo(HaveOccurred()) Expect(val).To(Equal([]string{"Palermo"})) - q = redis.GeoSearchQuery{ + q = &redis.GeoSearchQuery{ Longitude: 15, Latitude: 37, BoxWidth: 200, @@ -5222,7 +5222,7 @@ var _ = Describe("Commands", func() { Expect(err).NotTo(HaveOccurred()) Expect(val).To(Equal([]string{"Palermo"})) - q = redis.GeoSearchQuery{ + q = &redis.GeoSearchQuery{ Longitude: 15, Latitude: 37, Radius: 100, @@ -5250,7 +5250,7 @@ var _ = Describe("Commands", func() { }) It("should geo search with options", func() { - q := redis.GeoSearchWithOptionsQuery{ + q := &redis.GeoSearchLocationQuery{ GeoSearchQuery: redis.GeoSearchQuery{ Longitude: 15, Latitude: 37, @@ -5262,7 +5262,7 @@ var _ = Describe("Commands", func() { WithDist: true, WithCoord: true, } - val, err := client.GeoSearchWithOptions(ctx, "Sicily", q).Result() + val, err := client.GeoSearchLocation(ctx, "Sicily", q).Result() Expect(err).NotTo(HaveOccurred()) Expect(val).To(Equal([]redis.GeoLocation{ { @@ -5283,7 +5283,7 @@ var _ = Describe("Commands", func() { }) It("should geo search store", func() { - q := redis.GeoSearchStoreQuery{ + q := &redis.GeoSearchStoreQuery{ GeoSearchQuery: redis.GeoSearchQuery{ Longitude: 15, Latitude: 37, @@ -5303,7 +5303,7 @@ var _ = Describe("Commands", func() { Expect(err).NotTo(HaveOccurred()) Expect(val).To(Equal(int64(2))) - loc, err := client.GeoSearchWithOptions(ctx, "key1", redis.GeoSearchWithOptionsQuery{ + loc, err := client.GeoSearchLocation(ctx, "key1", &redis.GeoSearchLocationQuery{ GeoSearchQuery: q.GeoSearchQuery, WithCoord: true, WithDist: true,