re
/
redis
forked from mirror/redis
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

add cmd: geosearch, geosearchstore (#1836) 

* add cmd: geosearch, geosearchstore

Signed-off-by: monkey92t <golang@88.com>

* GeoSearchQuery and GeoSearchLocationQuery changed to pointer passing

Signed-off-by: monkey92t <golang@88.com>
This commit is contained in:
monkey92t 2021-07-28 14:38:14 +08:00 committed by GitHub
parent 437184bded
commit 7dad93efa2
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 415 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

182
command.go
View File

@ -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
}

35
commands.go
View File

@ -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 {

198
commands_test.go
View File

@ -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() {