Support Probabilistic commands with RESP 2 protocol (#3176)

* Support bloom resp 2

* Support Resp2 for BF.Info

* simplify BFInfoCmd field assignment using map-based key-to-field references
This commit is contained in:
ofekshenawa 2024-11-13 11:15:19 +02:00 committed by GitHub
parent d1b4eaed41
commit 8b1073d2d6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 780 additions and 729 deletions

View File

@ -319,38 +319,70 @@ func (cmd *BFInfoCmd) Result() (BFInfo, error) {
}
func (cmd *BFInfoCmd) readReply(rd *proto.Reader) (err error) {
result := BFInfo{}
// Create a mapping from key names to pointers of struct fields
respMapping := map[string]*int64{
"Capacity": &result.Capacity,
"CAPACITY": &result.Capacity,
"Size": &result.Size,
"SIZE": &result.Size,
"Number of filters": &result.Filters,
"FILTERS": &result.Filters,
"Number of items inserted": &result.ItemsInserted,
"ITEMS": &result.ItemsInserted,
"Expansion rate": &result.ExpansionRate,
"EXPANSION": &result.ExpansionRate,
}
// Helper function to read and assign a value based on the key
readAndAssignValue := func(key string) error {
fieldPtr, exists := respMapping[key]
if !exists {
return fmt.Errorf("redis: BLOOM.INFO unexpected key %s", key)
}
// Read the integer and assign to the field via pointer dereferencing
val, err := rd.ReadInt()
if err != nil {
return err
}
*fieldPtr = val
return nil
}
readType, err := rd.PeekReplyType()
if err != nil {
return err
}
if len(cmd.args) > 2 && readType == proto.RespArray {
n, err := rd.ReadArrayLen()
if err != nil {
return err
}
if key, ok := cmd.args[2].(string); ok && n == 1 {
if err := readAndAssignValue(key); err != nil {
return err
}
} else {
return fmt.Errorf("redis: BLOOM.INFO invalid argument key type")
}
} else {
n, err := rd.ReadMapLen()
if err != nil {
return err
}
var key string
var result BFInfo
for f := 0; f < n; f++ {
key, err = rd.ReadString()
for i := 0; i < n; i++ {
key, err := rd.ReadString()
if err != nil {
return err
}
switch key {
case "Capacity":
result.Capacity, err = rd.ReadInt()
case "Size":
result.Size, err = rd.ReadInt()
case "Number of filters":
result.Filters, err = rd.ReadInt()
case "Number of items inserted":
result.ItemsInserted, err = rd.ReadInt()
case "Expansion rate":
result.ExpansionRate, err = rd.ReadInt()
default:
return fmt.Errorf("redis: BLOOM.INFO unexpected key %s", key)
}
if err != nil {
if err := readAndAssignValue(key); err != nil {
return err
}
}
}
cmd.val = result
return nil

View File

@ -13,15 +13,32 @@ import (
var _ = Describe("Probabilistic commands", Label("probabilistic"), func() {
ctx := context.TODO()
setupRedisClient := func(protocolVersion int) *redis.Client {
return redis.NewClient(&redis.Options{
Addr: "localhost:6379",
DB: 0,
Protocol: protocolVersion,
})
}
protocols := []int{2, 3}
for _, protocol := range protocols {
protocol := protocol // capture loop variable for each context
Context(fmt.Sprintf("with protocol version %d", protocol), func() {
var client *redis.Client
BeforeEach(func() {
client = redis.NewClient(&redis.Options{Addr: ":6379"})
Expect(client.FlushDB(ctx).Err()).NotTo(HaveOccurred())
client = setupRedisClient(protocol)
Expect(client.FlushAll(ctx).Err()).NotTo(HaveOccurred())
})
AfterEach(func() {
Expect(client.Close()).NotTo(HaveOccurred())
if client != nil {
client.FlushDB(ctx)
client.Close()
}
})
Describe("bloom", Label("bloom"), func() {
@ -117,7 +134,7 @@ var _ = Describe("Probabilistic commands", Label("probabilistic"), func() {
NoCreate: true,
}
resultInsert, err := client.BFInsert(ctx, "testbf1", options, "item1").Result()
_, err := client.BFInsert(ctx, "testbf1", options, "item1").Result()
Expect(err).To(HaveOccurred())
Expect(err).To(MatchError("ERR not found"))
@ -129,7 +146,7 @@ var _ = Describe("Probabilistic commands", Label("probabilistic"), func() {
NoCreate: false,
}
resultInsert, err = client.BFInsert(ctx, "testbf1", options, "item1").Result()
resultInsert, err := client.BFInsert(ctx, "testbf1", options, "item1").Result()
Expect(err).NotTo(HaveOccurred())
Expect(len(resultInsert)).To(BeEquivalentTo(1))
@ -396,7 +413,7 @@ var _ = Describe("Probabilistic commands", Label("probabilistic"), func() {
NoCreate: true,
}
result, err := client.CFInsertNX(ctx, "testcf1", args, "item1", "item2", "item2").Result()
_, err := client.CFInsertNX(ctx, "testcf1", args, "item1", "item2", "item2").Result()
Expect(err).To(HaveOccurred())
args = &redis.CFInsertOptions{
@ -404,7 +421,7 @@ var _ = Describe("Probabilistic commands", Label("probabilistic"), func() {
NoCreate: false,
}
result, err = client.CFInsertNX(ctx, "testcf2", args, "item1", "item2", "item2").Result()
result, err := client.CFInsertNX(ctx, "testcf2", args, "item1", "item2", "item2").Result()
Expect(err).NotTo(HaveOccurred())
Expect(len(result)).To(BeEquivalentTo(3))
Expect(result[0]).To(BeEquivalentTo(int64(1)))
@ -730,4 +747,6 @@ var _ = Describe("Probabilistic commands", Label("probabilistic"), func() {
Expect(max).To(BeEquivalentTo(float64(140)))
})
})
})
}
})