Skip to content

Commit 8b1073d

Browse files
authored
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
1 parent d1b4eae commit 8b1073d

File tree

2 files changed

+780
-729
lines changed

2 files changed

+780
-729
lines changed

probabilistic.go

+52-20
Original file line numberDiff line numberDiff line change
@@ -319,37 +319,69 @@ func (cmd *BFInfoCmd) Result() (BFInfo, error) {
319319
}
320320

321321
func (cmd *BFInfoCmd) readReply(rd *proto.Reader) (err error) {
322-
n, err := rd.ReadMapLen()
322+
result := BFInfo{}
323+
324+
// Create a mapping from key names to pointers of struct fields
325+
respMapping := map[string]*int64{
326+
"Capacity": &result.Capacity,
327+
"CAPACITY": &result.Capacity,
328+
"Size": &result.Size,
329+
"SIZE": &result.Size,
330+
"Number of filters": &result.Filters,
331+
"FILTERS": &result.Filters,
332+
"Number of items inserted": &result.ItemsInserted,
333+
"ITEMS": &result.ItemsInserted,
334+
"Expansion rate": &result.ExpansionRate,
335+
"EXPANSION": &result.ExpansionRate,
336+
}
337+
338+
// Helper function to read and assign a value based on the key
339+
readAndAssignValue := func(key string) error {
340+
fieldPtr, exists := respMapping[key]
341+
if !exists {
342+
return fmt.Errorf("redis: BLOOM.INFO unexpected key %s", key)
343+
}
344+
345+
// Read the integer and assign to the field via pointer dereferencing
346+
val, err := rd.ReadInt()
347+
if err != nil {
348+
return err
349+
}
350+
*fieldPtr = val
351+
return nil
352+
}
353+
354+
readType, err := rd.PeekReplyType()
323355
if err != nil {
324356
return err
325357
}
326358

327-
var key string
328-
var result BFInfo
329-
for f := 0; f < n; f++ {
330-
key, err = rd.ReadString()
359+
if len(cmd.args) > 2 && readType == proto.RespArray {
360+
n, err := rd.ReadArrayLen()
331361
if err != nil {
332362
return err
333363
}
334-
335-
switch key {
336-
case "Capacity":
337-
result.Capacity, err = rd.ReadInt()
338-
case "Size":
339-
result.Size, err = rd.ReadInt()
340-
case "Number of filters":
341-
result.Filters, err = rd.ReadInt()
342-
case "Number of items inserted":
343-
result.ItemsInserted, err = rd.ReadInt()
344-
case "Expansion rate":
345-
result.ExpansionRate, err = rd.ReadInt()
346-
default:
347-
return fmt.Errorf("redis: BLOOM.INFO unexpected key %s", key)
364+
if key, ok := cmd.args[2].(string); ok && n == 1 {
365+
if err := readAndAssignValue(key); err != nil {
366+
return err
367+
}
368+
} else {
369+
return fmt.Errorf("redis: BLOOM.INFO invalid argument key type")
348370
}
349-
371+
} else {
372+
n, err := rd.ReadMapLen()
350373
if err != nil {
351374
return err
352375
}
376+
for i := 0; i < n; i++ {
377+
key, err := rd.ReadString()
378+
if err != nil {
379+
return err
380+
}
381+
if err := readAndAssignValue(key); err != nil {
382+
return err
383+
}
384+
}
353385
}
354386

355387
cmd.val = result

0 commit comments

Comments
 (0)