Add key prefix matching to KEYS command (#237)

Related to #234 and !236.

This is the implementation that was requested in the original issue. I updated KEYS command to be redis-valid and implemented prefix search. There is also a rather interesting test, I could you use some feedback here.

I noticed that it might not be possible to reduce the complexity of the KEYS command. Because even if you use Scan, you will have to store the counter of all found keys before you do WriteBulk of the actual keys.

@prologic here is what you probably had in mind:

```
s.db.Scan([]byte(prefix), func(key []byte) error {
	conn.WriteBulk(key)
	return nil
})
```

But there is no way to call `conn.WriteArray(n)` with the number of keys until you iterate through all of them, hence the second loop over found keys.

Co-authored-by: Ivan Elfimov <ielfimov@gmail.com>
Co-authored-by: James Mills <james@mills.io>
Reviewed-on: https://git.mills.io/prologic/bitcask/pulls/237
Reviewed-by: James Mills <james@mills.io>
Co-authored-by: biozz <biozz@noreply@mills.io>
Co-committed-by: biozz <biozz@noreply@mills.io>
This commit is contained in:
biozz 2021-09-20 10:35:27 +00:00 committed by James Mills
parent 2279245b8c
commit 21a824e13e
3 changed files with 137 additions and 3 deletions

@ -17,3 +17,4 @@ Yash Chandra <yashschandra@gmail.com>
Yury Fedorov orlangure
o2gy84 <o2gy84@gmail.com>
garsue <labs.garsue@gmail.com>
biozz <ielfimov@gmail.com>

@ -84,10 +84,41 @@ func (s *server) handleGet(cmd redcon.Command, conn redcon.Conn) {
}
func (s *server) handleKeys(cmd redcon.Command, conn redcon.Conn) {
conn.WriteArray(s.db.Len())
for key := range s.db.Keys() {
conn.WriteBulk(key)
if len(cmd.Args) != 2 {
conn.WriteError("ERR wrong number of arguments for '" + string(cmd.Args[0]) + "' command")
return
}
pattern := string(cmd.Args[1])
// Fast-track condition for improved speed
if pattern == "*" {
conn.WriteArray(s.db.Len())
for key := range s.db.Keys() {
conn.WriteBulk(key)
}
return
}
// Prefix handling
if strings.Count(pattern, "*") == 1 && strings.HasSuffix(pattern, "*") {
prefix := strings.ReplaceAll(pattern, "*", "")
count := 0
keys := make([][]byte, 0)
s.db.Scan([]byte(prefix), func(key []byte) error {
keys = append(keys, key)
count++
return nil
})
conn.WriteArray(count)
for _, key := range keys {
conn.WriteBulk(key)
}
return
}
// No results means empty array
conn.WriteArray(0)
}
func (s *server) handleExists(cmd redcon.Command, conn redcon.Conn) {

102
cmd/bitcaskd/server_test.go Normal file

@ -0,0 +1,102 @@
package main
import (
"net"
"strconv"
"testing"
"github.com/tidwall/redcon"
)
func TestHandleKeys(t *testing.T) {
s, err := newServer(":61234", "./test.db")
if err != nil {
t.Fatalf("Unable to create server: %v", err)
}
s.db.Put([]byte("foo"), []byte("bar"))
testCases := []TestCase{
{
Command: redcon.Command{
Raw: []byte("KEYS *"),
Args: [][]byte{[]byte("KEYS"), []byte("*")},
},
Expected: "1,foo",
},
{
Command: redcon.Command{
Raw: []byte("KEYS fo*"),
Args: [][]byte{[]byte("KEYS"), []byte("fo*")},
},
Expected: "1,foo",
},
{
Command: redcon.Command{
Raw: []byte("KEYS ba*"),
Args: [][]byte{[]byte("KEYS"), []byte("ba*")},
},
Expected: "0",
},
{
Command: redcon.Command{
Raw: []byte("KEYS *oo"),
Args: [][]byte{[]byte("KEYS"), []byte("*oo")},
},
Expected: "0",
},
}
for _, testCase := range testCases {
conn := DummyConn{}
s.handleKeys(testCase.Command, &conn)
if testCase.Expected != conn.Result {
t.Fatalf("s.handleKeys failed: expected '%s', got '%s'", testCase.Expected, conn.Result)
}
}
}
type TestCase struct {
Command redcon.Command
Expected string
}
type DummyConn struct {
Result string
}
func (dc *DummyConn) RemoteAddr() string {
return ""
}
func (dc *DummyConn) Close() error {
return nil
}
func (dc *DummyConn) WriteError(msg string) {}
func (dc *DummyConn) WriteString(str string) {}
func (dc *DummyConn) WriteBulk(bulk []byte) {
dc.Result += "," + string(bulk)
}
func (dc *DummyConn) WriteBulkString(bulk string) {}
func (dc *DummyConn) WriteInt(num int) {}
func (dc *DummyConn) WriteInt64(num int64) {}
func (dc *DummyConn) WriteUint64(num uint64) {}
func (dc *DummyConn) WriteArray(count int) {
dc.Result = strconv.Itoa(count)
}
func (dc *DummyConn) WriteNull() {}
func (dc *DummyConn) WriteRaw(data []byte) {}
func (dc *DummyConn) WriteAny(any interface{}) {}
func (dc *DummyConn) Context() interface{} {
return nil
}
func (dc *DummyConn) SetContext(v interface{}) {}
func (dc *DummyConn) SetReadBuffer(bytes int) {}
func (dc *DummyConn) Detach() redcon.DetachedConn {
return nil
}
func (dc *DummyConn) ReadPipeline() []redcon.Command {
return nil
}
func (dc *DummyConn) PeekPipeline() []redcon.Command {
return nil
}
func (dc *DummyConn) NetConn() net.Conn {
return nil
}