DiceDB#281: added support for DUMP command (DiceDB#633)

TheRanomial · web-flow · commit 96a0202058dc · 2024-10-01T23:14:43.000+05:30
diff --git a/integration_tests/commands/async/dump_test.go b/integration_tests/commands/async/dump_test.go
@@ -0,0 +1,110 @@
+package async
+
+import (
+	"encoding/base64"
+	"testing"
+
+	"github.com/dicedb/dice/testutils"
+	"gotest.tools/v3/assert"
+)
+
+func TestDumpRestore(t *testing.T) {
+	conn := getLocalConnection()
+	defer conn.Close()
+
+	testCases := []struct {
+		name     string
+		commands []string
+		expected []interface{}
+	}{
+		{
+			name: "DUMP and RESTORE string value",
+			commands: []string{
+				"SET mykey hello",
+				"DUMP mykey",
+				"DEL mykey",
+				"RESTORE mykey 2 CQAAAAAFaGVsbG//AEeXk742Rcc=",
+				"GET mykey",
+			},
+			expected: []interface{}{
+				"OK",
+				func(result interface{}) bool {
+					dumped, ok := result.(string)
+					if !ok {
+						return false
+					}
+					decoded, err := base64.StdEncoding.DecodeString(dumped)
+					if err != nil {
+						return false
+					}
+					return len(decoded) > 11 &&
+						decoded[0] == 0x09 && 
+						decoded[1] == 0x00 && 
+						string(decoded[6:11]) == "hello" &&
+						decoded[11] == 0xFF
+				},
+				int64(1),
+				"OK",
+				"hello",
+			},
+		},
+		{
+			name: "DUMP and RESTORE integer value",
+			commands: []string{
+				"SET intkey 42",
+				"DUMP intkey",
+				"DEL intkey",
+				"RESTORE intkey 2 CcAAAAAAAAAAKv9S/ymRDY3rXg==",
+			},
+			expected: []interface{}{
+				"OK",
+				func(result interface{}) bool {
+					dumped, ok := result.(string)
+					if !ok {
+						return false
+					}
+					decoded, err := base64.StdEncoding.DecodeString(dumped)
+					if err != nil {
+						return false
+					}
+					return len(decoded) > 2 &&
+						decoded[0] == 0x09 &&
+						decoded[1] == 0xC0
+				},
+				int64(1),
+				"OK",
+			},
+		},
+		{
+			name: "DUMP non-existent key",
+			commands: []string{
+				"DUMP nonexistentkey",
+			},
+			expected: []interface{}{
+				"ERR nil",
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			FireCommand(conn, "FLUSHALL")
+			for i, cmd := range tc.commands {
+				var result interface{}
+				result = FireCommand(conn, cmd)
+				expected := tc.expected[i]
+
+				switch exp := expected.(type) {
+				case string:
+					assert.DeepEqual(t, exp, result)
+				case []interface{}:
+					assert.Assert(t, testutils.UnorderedEqual(exp, result))
+				case func(interface{}) bool:
+					assert.Assert(t, exp(result), cmd)
+				default:
+					assert.DeepEqual(t, expected, result)
+				}
+			}
+		})
+	}
+}
diff --git a/internal/eval/commands.go b/internal/eval/commands.go
@@ -911,11 +911,28 @@ var (
 		Arity:    3,
 		KeySpecs: KeySpecs{BeginIndex: 1},
 	}
+	dumpkeyCMmdMeta=DiceCmdMeta{
+		Name:	 "DUMP",
+		Info:	`Serialize the value stored at key in a Redis-specific format and return it to the user.
+				The returned value can be synthesized back into a Redis key using the RESTORE command.`,
+		Eval:   evalDUMP,
+		Arity: 	1,
+		KeySpecs:   KeySpecs{BeginIndex: 1},
+	}
+	restorekeyCmdMeta=DiceCmdMeta{
+		Name:	"RESTORE",
+		Info:  `Serialize the value stored at key in a Redis-specific format and return it to the user.
+				The returned value can be synthesized back into a Redis key using the RESTORE command.`,
+		Eval: evalRestore,
+		Arity:	2,
+		KeySpecs: KeySpecs{BeginIndex: 1},
+	}
 	typeCmdMeta = DiceCmdMeta{
 		Name:     "TYPE",
 		Info:     `Returns the string representation of the type of the value stored at key. The different types that can be returned are: string, list, set, zset, hash and stream.`,
 		Eval:     evalTYPE,
 		Arity:    1,
+
 		KeySpecs: KeySpecs{BeginIndex: 1},
 	}
 	incrbyCmdMeta = DiceCmdMeta{
@@ -999,6 +1016,8 @@ func init() {
 	DiceCmds["PING"] = pingCmdMeta
 	DiceCmds["ECHO"] = echoCmdMeta
 	DiceCmds["AUTH"] = authCmdMeta
+	DiceCmds["DUMP"]=dumpkeyCMmdMeta
+	DiceCmds["RESTORE"]=restorekeyCmdMeta
 	DiceCmds["SET"] = setCmdMeta
 	DiceCmds["GET"] = getCmdMeta
 	DiceCmds["MSET"] = msetCmdMeta
diff --git a/internal/eval/dump_restore.go b/internal/eval/dump_restore.go
@@ -0,0 +1,158 @@
+package eval
+
+import (
+	"bytes"
+	"encoding/base64"
+	"encoding/binary"
+	"errors"
+	"hash/crc64"
+	"strconv"
+
+	"github.com/dicedb/dice/internal/clientio"
+	diceerrors "github.com/dicedb/dice/internal/errors"
+	"github.com/dicedb/dice/internal/object"
+	dstore "github.com/dicedb/dice/internal/store"
+)
+
+func evalDUMP(args []string, store *dstore.Store) []byte {
+    if len(args) < 1 {
+        return diceerrors.NewErrArity("DUMP")
+    }
+    key := args[0]
+    obj := store.Get(key)
+    if obj == nil {
+        return diceerrors.NewErrWithFormattedMessage("nil")
+    }
+	
+    serializedValue, err := rdbSerialize(obj)
+    if err != nil {
+        return diceerrors.NewErrWithMessage("serialization failed")
+    }
+	encodedResult := base64.StdEncoding.EncodeToString(serializedValue)
+    return clientio.Encode(encodedResult, false)
+}
+
+func evalRestore(args []string, store *dstore.Store) []byte {
+	if len(args) < 3 {
+		return diceerrors.NewErrArity("RESTORE")
+	}
+
+	key := args[0]
+	ttlStr:=args[1]
+	ttl, _ := strconv.ParseInt(ttlStr, 10, 64)
+	
+	encodedValue := args[2]
+	serializedData, err := base64.StdEncoding.DecodeString(encodedValue)
+	if err != nil {
+		return diceerrors.NewErrWithMessage("failed to decode base64 value")
+	}
+
+	obj, err := rdbDeserialize(serializedData)
+	if err != nil {
+		return diceerrors.NewErrWithMessage("deserialization failed: " + err.Error())
+	}
+
+	newobj:=store.NewObj(obj.Value,ttl,obj.TypeEncoding,obj.TypeEncoding)
+	var keepttl=true
+
+	if(ttl>0){
+		store.Put(key, newobj, dstore.WithKeepTTL(keepttl))
+	}else{
+		store.Put(key,obj)
+	}
+	
+	return clientio.RespOK
+}
+
+func rdbDeserialize(data []byte) (*object.Obj, error) {
+	if len(data) < 3 {
+		return nil, errors.New("insufficient data for deserialization")
+	}
+	objType := data[1] 
+	switch objType {
+	case 0x00:
+		return readString(data[2:])
+	case 0xC0: // Integer type
+		return readInt(data[2:])
+	default:
+		return nil, errors.New("unsupported object type")
+	}
+}
+
+func readString(data []byte) (*object.Obj, error) {
+	buf := bytes.NewReader(data)
+	var strLen uint32
+	if err := binary.Read(buf, binary.BigEndian, &strLen); err != nil {
+		return nil, err
+	}
+
+	strBytes := make([]byte, strLen)
+	if _, err := buf.Read(strBytes); err != nil {
+		return nil, err
+	}
+
+	return &object.Obj{TypeEncoding: object.ObjTypeString, Value: string(strBytes)}, nil
+}
+
+func readInt(data []byte) (*object.Obj, error) {
+	var intVal int64
+	if err := binary.Read(bytes.NewReader(data), binary.BigEndian, &intVal); err != nil {
+		return nil, err
+	}
+
+	return &object.Obj{TypeEncoding: object.ObjTypeInt, Value: intVal}, nil
+}
+
+func rdbSerialize(obj *object.Obj) ([]byte, error) {
+    var buf bytes.Buffer
+    buf.WriteByte(0x09)
+
+    switch object.GetType(obj.TypeEncoding) {
+    case object.ObjTypeString:
+        str, ok := obj.Value.(string)
+        if !ok {
+            return nil, errors.New("invalid string value")
+        }
+        buf.WriteByte(0x00) 
+        if err := writeString(&buf, str); err != nil {
+            return nil, err
+        }
+
+    case object.ObjTypeInt:
+        intVal, ok := obj.Value.(int64)
+        if !ok {
+            return nil, errors.New("invalid integer value")
+        }
+        buf.WriteByte(0xC0)
+        writeInt(&buf, intVal);
+
+    default:
+        return nil, errors.New("unsupported object type")
+    }
+
+    buf.WriteByte(0xFF) // End marker
+
+    return appendChecksum(buf.Bytes()), nil
+}
+
+func writeString(buf *bytes.Buffer, str string) error {
+    strLen := uint32(len(str))
+    if err := binary.Write(buf, binary.BigEndian, strLen); err != nil {
+        return err
+    }
+    buf.WriteString(str)
+    return nil
+}
+
+func writeInt(buf *bytes.Buffer, intVal int64){
+    tempBuf := make([]byte, 8)
+    binary.BigEndian.PutUint64(tempBuf, uint64(intVal))
+    buf.Write(tempBuf)
+}
+
+func appendChecksum(data []byte) []byte {
+    checksum := crc64.Checksum(data, crc64.MakeTable(crc64.ECMA))
+    checksumBuf := make([]byte, 8)
+    binary.BigEndian.PutUint64(checksumBuf, checksum)
+    return append(data, checksumBuf...)
+}
diff --git a/internal/eval/eval.go b/internal/eval/eval.go
@@ -2,10 +2,14 @@ package eval
 
 import (
 	"bytes"
+
 	"crypto/rand"
+
 	"errors"
 	"fmt"
+
 	"log/slog"
+
 	"math"
 	"math/big"
 	"math/bits"
@@ -4147,6 +4151,8 @@ func evalJSONNUMINCRBY(args []string, store *dstore.Store) []byte {
 	return clientio.Encode(resultString, false)
 }
 
+
+
 // evalJSONOBJKEYS retrieves the keys of a JSON object stored at path specified.
 // It takes two arguments: the key where the JSON document is stored, and an optional JSON path.
 // It returns a list of keys from the object at the specified path or an error if the path is invalid.
@@ -4249,6 +4255,7 @@ func evalTYPE(args []string, store *dstore.Store) []byte {
 	return clientio.Encode(typeStr, true)
 }
 
+
 // evalGETRANGE returns the substring of the string value stored at key, determined by the offsets start and end
 // The offsets are zero-based and can be negative values to index from the end of the string
 //
@@ -4680,3 +4687,4 @@ func evalHINCRBYFLOAT(args []string, store *dstore.Store) []byte {
 
 	return clientio.Encode(numkey, false)
 }
+
diff --git a/internal/eval/eval_test.go b/internal/eval/eval_test.go
