From 8c385659ef313d8fbc9a36833ed1e0f329a58abe Mon Sep 17 00:00:00 2001
From: Tianle Xu <xtl@xtlsoft.top>
Date: Fri, 2 Oct 2020 09:03:05 +0800
Subject: [PATCH] fix(bucket): trim prefix when listing

---
 docs/unikvd-protocol.md             |  2 ++
 drivers/persistent-memory/driver.go | 13 +++++++++++-
 drivers/redis/driver.go             | 32 ++++++++++++++++++++---------
 unikvd/main.go                      | 11 +++++++++-
 utils.go                            | 10 ++++++++-
 5 files changed, 55 insertions(+), 13 deletions(-)

diff --git a/docs/unikvd-protocol.md b/docs/unikvd-protocol.md
index 9d2cc7b..7c88f7d 100644
--- a/docs/unikvd-protocol.md
+++ b/docs/unikvd-protocol.md
@@ -24,6 +24,8 @@ When you POST this address, the body will be the data.
 
 When you DELETE this address, the key will be deleted.
 
+When KEY is set to `_list`, a list of all the keys will be returned in JSON format.
+
 ## System Namespace
 
 When you get `/v1/_system/*`, you are visiting system namespace.
diff --git a/drivers/persistent-memory/driver.go b/drivers/persistent-memory/driver.go
index 4630737..03c6eb7 100644
--- a/drivers/persistent-memory/driver.go
+++ b/drivers/persistent-memory/driver.go
@@ -68,7 +68,18 @@ func (d *Driver) Put(key string, value string) error {
 
 // List lists the keys
 func (d *Driver) List() (interface{}, error) {
-	return reflect.ValueOf(d.data).MapKeys(), nil
+	if d.loadWhenGet {
+		err := d.Load()
+		if err != nil {
+			return nil, err
+		}
+	}
+	keys := reflect.ValueOf(d.data).MapKeys()
+	var rslt []string = make([]string, len(keys))
+	for k, v := range keys {
+		rslt[k] = unikv.TrimPrefix(d.prefix, v.String())
+	}
+	return rslt, nil
 }
 
 // Unset unsets data
diff --git a/drivers/redis/driver.go b/drivers/redis/driver.go
index 42cc179..1500578 100644
--- a/drivers/redis/driver.go
+++ b/drivers/redis/driver.go
@@ -46,10 +46,12 @@ func (d *Driver) Get(key string) (string, error) {
 func (d *Driver) Put(key string, value string) error {
 	key = unikv.ConcatPrefix(d.prefix, key)
 	_, err := d.conn.Do("SET", key, value)
-	if !d.reconnected {
-		d.connect()
-		d.reconnected = true
-		return d.Put(key, value)
+	if err != nil {
+		if !d.reconnected {
+			d.connect()
+			d.reconnected = true
+			return d.Put(key, value)
+		}
 	}
 	if d.reconnected {
 		d.reconnected = false
@@ -61,10 +63,12 @@ func (d *Driver) Put(key string, value string) error {
 func (d *Driver) Unset(key string) error {
 	key = unikv.ConcatPrefix(d.prefix, key)
 	_, err := d.conn.Do("SET", key, "EX", "1")
-	if !d.reconnected {
-		d.connect()
-		d.reconnected = true
-		return d.Unset(key)
+	if err != nil {
+		if !d.reconnected {
+			d.connect()
+			d.reconnected = true
+			return d.Unset(key)
+		}
 	}
 	if d.reconnected {
 		d.reconnected = false
@@ -74,7 +78,7 @@ func (d *Driver) Unset(key string) error {
 
 // List lists the keys
 func (d *Driver) List() (interface{}, error) {
-	data, err := d.conn.Do("KEYS", "*")
+	data, err := d.conn.Do("KEYS", d.prefix+"*")
 	if err != nil {
 		if err == redis.ErrNil {
 			return "", unikv.ErrNotFound
@@ -89,7 +93,15 @@ func (d *Driver) List() (interface{}, error) {
 	if d.reconnected {
 		d.reconnected = false
 	}
-	return data, err
+	dat, ok := data.([]interface{})
+	if !ok {
+		return nil, fmt.Errorf("Invalid response from redis server")
+	}
+	var rslt []string = make([]string, len(dat))
+	for k, v := range dat {
+		rslt[k] = unikv.TrimPrefix(d.prefix, string(v.([]byte)))
+	}
+	return rslt, err
 }
 
 // Close closes driver
diff --git a/unikvd/main.go b/unikvd/main.go
index 917dcf1..28a0810 100644
--- a/unikvd/main.go
+++ b/unikvd/main.go
@@ -19,7 +19,7 @@ var namespaces map[string]*unikv.Namespace
 var enableLog = true
 
 func main() {
-	if len(os.Args) <= 2 {
+	if len(os.Args) < 2 {
 		fmt.Println(`USAGE: unikvd [LISTEN_ADDR] [--disable-log]`)
 		os.Exit(0)
 	}
@@ -171,6 +171,15 @@ func handle(rw http.ResponseWriter, r *http.Request) {
 	}
 	switch r.Method {
 	case "GET":
+		if key == "_list" {
+			rslt, err := bucket.List()
+			if err != nil {
+				resp500(rw, err.Error())
+				return
+			}
+			resp200(rw, rslt)
+			return
+		}
 		rslt, err := bucket.GetString(key)
 		if err == unikv.ErrNotFound {
 			resp404(rw)
diff --git a/utils.go b/utils.go
index b02d1a0..800608a 100644
--- a/utils.go
+++ b/utils.go
@@ -1,6 +1,9 @@
 package unikv
 
-import "encoding/json"
+import (
+	"encoding/json"
+	"strings"
+)
 
 // ConcatPrefix concats two prefixes together
 func ConcatPrefix(prefix string, str string) string {
@@ -14,6 +17,11 @@ func concatPrefix(prefix string, str string) string {
 	return ConcatPrefix(prefix, str)
 }
 
+// TrimPrefix trims the prefix of the given string
+func TrimPrefix(prefix string, str string) string {
+	return strings.TrimPrefix(str, prefix)
+}
+
 type temporaryStringWriter struct {
 	buffer string
 }