Support build & parse RawKV keys (#13)

Signed-off-by: Ping Yu <yuping@pingcap.com>

README.md

@@ -64,6 +64,15 @@ $ ./mok --keyspace-id 255 --table-id 43 --row-value 81934
 built key: 780000FF74800000FF000000002B5F7280FF0000000001400E00FE
 ```
 
+Build a RawKV key for a given raw key under given keyspace
+```
+$ ./mok --keyspace-id 0 --key-mode rawkv --raw-key test
+built key: 7200000074657374FF0000000000000000F7
+
+$ ./mok --keyspace-id 1 --key-mode rawkv --raw-key ''
+built key: 7200000100000000FB
+```
+
 ## TODO
 
 - [x] build keys

main.go

@@ -14,23 +14,35 @@ import (
 
 var keyFormat = flag.String("format", "proto", "output format (go/hex/base64/proto)")
 var NullSpaceID = int64(0xffffffff)
+var keyMode = flag.String("key-mode", "txnkv", "key mode (txnkv/rawkv)")
 var keyspaceID = flag.Int64("keyspace-id", NullSpaceID, "keyspace ID")
 var tableID = flag.Int64("table-id", 0, "table ID")
 var indexID = flag.Int64("index-id", 0, "index ID")
 var rowValue = flag.String("row-value", "", "row value")
 var indexValue = flag.String("index-value", "", "index value")
+var rawKey = flag.String("raw-key", "", "raw key (rawkv only)")
+var rawKeyFormat = flag.String("raw-key-format", "str", "input format (str/hex, rawkv only)")
 
-func getKeyPrefix(keyspaceID int64) ([]byte, error) {
+func getKeyPrefix(keyModeStr string, keyspaceID int64) (*KeyMode, []byte, error) {
 	if keyspaceID == NullSpaceID {
-		return []byte{'t'}, nil
+		return nil, []byte{'t'}, nil
 	}
 	if keyspaceID > 0xffffff {
-		return nil, fmt.Errorf("invalid keyspace value: %d", keyspaceID)
+		return nil, nil, fmt.Errorf("invalid keyspace value: %d", keyspaceID)
 	}
+
+	keyMode := FromStringToKeyMode(keyModeStr)
+	if keyMode == nil {
+		return nil, nil, fmt.Errorf("invalid key mode: %s", keyModeStr)
+	}
+
 	var prefix [4]byte
 	binary.BigEndian.PutUint32(prefix[:], uint32(keyspaceID))
-	prefix[0] = 'x'
-	return append(prefix[:], 't'), nil
+	prefix[0] = byte(*keyMode)
+	if *keyMode == KeyModeRaw {
+		return keyMode, prefix[:], nil
+	}
+	return keyMode, append(prefix[:], 't'), nil
 }
 
 func main() {
@@ -40,46 +52,64 @@ func main() {
 		n := N("key", []byte(flag.Arg(0)))
 		n.Expand().Print()
 	} else if flag.NArg() == 0 { // Build a key with given flags.
-		key, err := getKeyPrefix(*keyspaceID)
+		keyMode, key, err := getKeyPrefix(*keyMode, *keyspaceID)
 		if err != nil {
 			fmt.Println(err.Error())
 			os.Exit(1)
 		}
 
-		key = codec.EncodeInt(key, *tableID)
-		if *tableID == 0 {
-			fmt.Println("table ID shouldn't be 0")
-			os.Exit(1)
-		}
-
-		if *indexID == 0 {
-			if *rowValue != "" {
-				key = append(key, []byte("_r")...)
-				rowValueInt, err := strconv.ParseInt(*rowValue, 10, 64)
-				if err != nil {
-					fmt.Printf("invalid row value: %s\n", *rowValue)
-					os.Exit(1)
-				}
-				key = codec.EncodeInt(key, rowValueInt)
+		if keyMode != nil && *keyMode == KeyModeRaw {
+			key, err = buildRawKVKey(key, *rawKey, *rawKeyFormat)
+			if err != nil {
+				fmt.Println(err.Error())
+				os.Exit(1)
 			}
 		} else {
-			key = append(key, []byte("_i")...)
-			key = codec.EncodeInt(key, *indexID)
-			if *indexValue != "" {
-				indexValueInt, err := strconv.ParseInt(*indexValue, 10, 64)
-				if err != nil {
-					fmt.Printf("invalid index value: %s\n", *indexValue)
-					os.Exit(1)
+			key = codec.EncodeInt(key, *tableID)
+			if *tableID == 0 {
+				fmt.Println("table ID shouldn't be 0")
+				os.Exit(1)
+			}
+
+			if *indexID == 0 {
+				if *rowValue != "" {
+					key = append(key, []byte("_r")...)
+					rowValueInt, err := strconv.ParseInt(*rowValue, 10, 64)
+					if err != nil {
+						fmt.Printf("invalid row value: %s\n", *rowValue)
+						os.Exit(1)
+					}
+					key = codec.EncodeInt(key, rowValueInt)
+				}
+			} else {
+				key = append(key, []byte("_i")...)
+				key = codec.EncodeInt(key, *indexID)
+				if *indexValue != "" {
+					indexValueInt, err := strconv.ParseInt(*indexValue, 10, 64)
+					if err != nil {
+						fmt.Printf("invalid index value: %s\n", *indexValue)
+						os.Exit(1)
+					}
 					key = codec.EncodeInt(key, indexValueInt)
 				}
 			}
+
+			key = codec.EncodeBytes([]byte{}, key)
 		}
 
-		key = codec.EncodeBytes([]byte{}, key)
 		fmt.Printf("built key: %s\n", strings.ToUpper(hex.EncodeToString(key)))
 	} else {
 		fmt.Println("usage:\nmok {flags} [key]")
 		flag.PrintDefaults()
 		os.Exit(1)
 	}
 }
+
+func buildRawKVKey(key []byte, rawKey string, format string) ([]byte, error) {
+	parsedRawKey, err := ParseRawKey(rawKey, format)
+	if err != nil {
+		return nil, err
+	}
+	key = append(key, parsedRawKey...)
+	return codec.EncodeBytes([]byte{}, key), nil
+}

main_test.go

@@ -0,0 +1,136 @@
+package main
+
+import (
+	"encoding/hex"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestBuildRawKVKey(t *testing.T) {
+	assert := require.New(t)
+
+	keyMode, prefix, err := getKeyPrefix("rawkv", 0)
+	assert.NoError(err)
+	assert.Equal(KeyModeRaw, *keyMode)
+
+	tests := []struct {
+		rawKey   string
+		format   string
+		expected string
+		err      bool
+	}{
+		{
+			rawKey:   "",
+			format:   "str",
+			expected: "7200000000000000fb",
+			err:      false,
+		},
+
+		{
+			rawKey:   "test",
+			format:   "str",
+			expected: "7200000074657374ff0000000000000000f7",
+			err:      false,
+		},
+		{
+			rawKey:   "testtest",
+			format:   "str",
+			expected: "7200000074657374ff7465737400000000fb",
+			err:      false,
+		},
+		{
+			rawKey:   "74657374",
+			format:   "hex",
+			expected: "7200000074657374ff0000000000000000f7",
+			err:      false,
+		},
+		{
+			rawKey:   "test",
+			format:   "invalid",
+			expected: "",
+			err:      true,
+		},
+	}
+
+	for _, tt := range tests {
+		actualKey, err := buildRawKVKey(prefix, tt.rawKey, tt.format)
+
+		if tt.err && err == nil {
+			t.Errorf("expected error, but got nil")
+		}
+
+		if !tt.err && err != nil {
+			t.Errorf("expected nil error, but got %v", err)
+		}
+
+		if hex.EncodeToString(actualKey) != tt.expected {
+			t.Errorf("expected key %s, but got %s", tt.expected, hex.EncodeToString(actualKey))
+		}
+	}
+}
+
+func TestParseRawKVKey(t *testing.T) {
+	// ❯ ./mok 7200000074657374ff0000000000000000f7
+	// "7200000074657374ff0000000000000000f7"
+	// └─## decode hex key
+	//   └─"r\000\000\000test\377\000\000\000\000\000\000\000\000\367"
+	//     ├─## decode mvcc key
+	//     │ └─"r\000\000\000test"
+	//     │   └─## decode keyspace
+	//     │     ├─key mode: rawkv
+	//     │     ├─keyspace: 0
+	//     │     └─"test"
+	//     └─## decode keyspace
+	//       ├─key mode: rawkv
+	//       ├─keyspace: 0
+	//       └─"test\377\000\000\000\000\000\000\000\000\367"
+	assert := require.New(t)
+
+	tests := []struct {
+		encoded  string
+		expected string
+	}{
+		{
+			encoded:  "7200000000000000fb",
+			expected: "",
+		},
+		{
+			encoded:  "7200000074657374ff0000000000000000f7",
+			expected: "test",
+		},
+		{
+			encoded:  "7200000074657374ff7465737400000000fb",
+			expected: "testtest",
+		},
+	}
+
+	for _, tt := range tests {
+		n := N("key", []byte(tt.encoded)).Expand()
+		assert.Len(n.variants, 1) // decode hex key
+		assert.Equal("decode hex key", n.variants[0].method)
+		assert.Len(n.variants[0].children, 1)
+
+		n = n.variants[0].children[0]
+		assert.Len(n.variants, 2) // with and without mvcc
+
+		variant := n.variants[0]
+		assert.Equal("decode mvcc key", variant.method)
+		assert.Len(variant.children, 1)
+
+		n = variant.children[0]
+		assert.Len(n.variants, 1)
+		variant = n.variants[0]
+		assert.Equal("decode keyspace", variant.method)
+		assert.Len(variant.children, 3)
+
+		assert.Equal("key_mode", variant.children[0].typ)
+		assert.Equal("r", string(variant.children[0].val))
+
+		assert.Equal("keyspace_id", variant.children[1].typ)
+		assert.Equal([]byte{0, 0, 0}, variant.children[1].val)
+
+		assert.Equal("raw_key", variant.children[2].typ)
+		assert.Equal(tt.expected, string(variant.children[2].val))
+	}
+}

mok.go

@@ -27,7 +27,7 @@ func N(t string, v []byte) *Node {
 
 func (n *Node) String() string {
 	switch n.typ {
-	case "key", "index_values":
+	case "key", "raw_key", "index_values":
 		switch *keyFormat {
 		case "hex":
 			return `"` + strings.ToUpper(hex.EncodeToString(n.val)) + `"`
@@ -38,6 +38,8 @@ func (n *Node) String() string {
 		default:
 			return fmt.Sprintf("%q", n.val)
 		}
+	case "key_mode":
+		return fmt.Sprintf("key mode: %s", KeyMode(n.val[0]))
 	case "keyspace_id":
 		tmp := []byte{'\x00'}
 		t := append(tmp, n.val...)

rules.go

@@ -81,10 +81,14 @@ func DecodeRocksDBKey(n *Node) *Variant {
 }
 
 func DecodeKeyspace(n *Node) *Variant {
-	if n.typ == "key" && n.val[0] == 'x' && len(n.val) >= 4 {
+	if n.typ == "key" && IsValidKeyMode(n.val[0]) && len(n.val) >= 4 {
+		keyType := "key"
+		if IsRawKeyMode(n.val[0]) {
+			keyType = "raw_key"
+		}
 		return &Variant{
 			method:   "decode keyspace",
-			children: []*Node{N("keyspace_id", n.val[1:4]), N("key", n.val[4:])},
+			children: []*Node{N("key_mode", n.val[0:1]), N("keyspace_id", n.val[1:4]), N(keyType, n.val[4:])},
 		}
 	}
 	return nil

util.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"bytes"
+	"encoding/hex"
 	"fmt"
 	"io"
 	"strings"
@@ -76,3 +77,52 @@ func GetTimeFromTS(ts uint64) time.Time {
 	ms := int64(ts >> 18)
 	return time.Unix(ms/1e3, (ms%1e3)*1e6)
 }
+
+type KeyMode byte
+
+const (
+	KeyModeTxn KeyMode = 'x'
+	KeyModeRaw KeyMode = 'r'
+)
+
+func IsValidKeyMode(b byte) bool {
+	return b == byte(KeyModeTxn) || b == byte(KeyModeRaw)
+}
+
+func IsRawKeyMode(b byte) bool {
+	return b == byte(KeyModeRaw)
+}
+
+func (k KeyMode) String() string {
+	switch k {
+	case KeyModeTxn:
+		return "txnkv"
+	case KeyModeRaw:
+		return "rawkv"
+	default:
+		return "other"
+	}
+}
+
+func FromStringToKeyMode(s string) *KeyMode {
+	var keyMode KeyMode
+	switch s {
+	case "txnkv":
+		keyMode = KeyModeTxn
+	case "rawkv":
+		keyMode = KeyModeRaw
+	default:
+	}
+	return &keyMode
+}
+
+func ParseRawKey(s string, format string) ([]byte, error) {
+	switch format {
+	case "hex":
+		return hex.DecodeString(s)
+	case "str": // for `s` with all characters printable.
+		return []byte(s), nil
+	default:
+		return nil, fmt.Errorf("invalid raw key format: %s", format)
+	}
+}