Made suggested change.

Signed-off-by: Cody Littley <[email protected]>

common/kvstore/key.go

@@ -7,22 +7,8 @@ var ErrInvalidKey = errors.New("invalid key")
 
 // Key represents a key in a TableStore. Each key is scoped to a specific table.
 type Key interface {
-	// AsString interprets the key as a string and returns it.
-	AsString() string
-	// AsBytes interprets the key as a byte slice and returns it.
-	AsBytes() []byte
-	// AsUint64 interprets the key as a uint64 and returns it.
-	// Returns an error if the data cannot be interpreted as a uint64.
-	AsUint64() (uint64, error)
-	// AsInt64 interprets the key as an int64 and returns it.
-	// Returns an error if the data cannot be interpreted as an int64.
-	AsInt64() (int64, error)
-	// AsUint32 interprets the key as a uint32 and returns it.
-	// Returns an error if the data cannot be interpreted as a uint32.
-	AsUint32() (uint32, error)
-	// AsInt32 interprets the key as an int32 and returns it.
-	// Returns an error if the data cannot be interpreted as an int32.
-	AsInt32() (int32, error)
+	// Bytes returns the key as a byte slice. Does not include internal metadata (i.e. the table).
+	Bytes() []byte
 	// Raw returns the raw byte slice that represents the key. This value
 	// may not be equal to the byte slice that was used to create the key, and
 	// should be treated as an opaque value.
@@ -38,14 +24,4 @@ type KeyBuilder interface {
 	TableName() string
 	// Key creates a key from a byte slice.
 	Key(key []byte) Key
-	// StringKey creates a key from a string. Equivalent to Key([]byte(key)).
-	StringKey(key string) Key
-	// Uint64Key creates a key from a uint64. Resulting key is an 8-byte big-endian representation of the uint64.
-	Uint64Key(key uint64) Key
-	// Int64Key creates a key from an int64. Resulting key is an 8-byte big-endian representation of the int64.
-	Int64Key(key int64) Key
-	// Uint32Key creates a key from a uint32. Resulting key is a 4-byte big-endian representation of the uint32.
-	Uint32Key(key uint32) Key
-	// Int32Key creates a key from an int32. Resulting key is a 4-byte big-endian representation of the int32.
-	Int32Key(key int32) Key
 }

common/kvstore/tablestore/key_test.go

@@ -3,7 +3,6 @@ package tablestore
 import (
 	tu "github.com/Layr-Labs/eigenda/common/testutils"
 	"github.com/stretchr/testify/assert"
-	"math/rand"
 	"testing"
 )
 
@@ -16,27 +15,6 @@ func TestGetName(t *testing.T) {
 	assert.Equal(t, tableName, kb.TableName())
 }
 
-func TestGetBuilder(t *testing.T) {
-	tu.InitializeRandom()
-
-	tableName := tu.RandomString(10)
-
-	kb := newKeyBuilder(tableName, 0)
-	k := kb.StringKey("asdf")
-	assert.Same(t, kb, k.Builder())
-}
-
-func TestStringRoundTrip(t *testing.T) {
-	tu.InitializeRandom()
-
-	tableName := tu.RandomString(10)
-	str := tu.RandomString(10)
-
-	kb := newKeyBuilder(tableName, 0)
-	k := kb.StringKey(str)
-	assert.Equal(t, str, k.AsString())
-}
-
 func TestBytesRoundTrip(t *testing.T) {
 	tu.InitializeRandom()
 
@@ -45,57 +23,5 @@ func TestBytesRoundTrip(t *testing.T) {
 
 	kb := newKeyBuilder(tableName, 0)
 	k := kb.Key(b)
-	assert.Equal(t, b, k.AsBytes())
-}
-
-func TestUint64RoundTrip(t *testing.T) {
-	tu.InitializeRandom()
-
-	tableName := tu.RandomString(10)
-	u := rand.Uint64()
-
-	kb := newKeyBuilder(tableName, 0)
-	k := kb.Uint64Key(u)
-	u2, err := k.AsUint64()
-	assert.NoError(t, err)
-	assert.Equal(t, u, u2)
-}
-
-func TestInt64RoundTrip(t *testing.T) {
-	tu.InitializeRandom()
-
-	tableName := tu.RandomString(10)
-	i := rand.Int63()
-
-	kb := newKeyBuilder(tableName, 0)
-	k := kb.Int64Key(i)
-	u2, err := k.AsInt64()
-	assert.NoError(t, err)
-	assert.Equal(t, i, u2)
-}
-
-func TestUint32RoundTrip(t *testing.T) {
-	tu.InitializeRandom()
-
-	tableName := tu.RandomString(10)
-	u := rand.Uint32()
-
-	kb := newKeyBuilder(tableName, 0)
-	k := kb.Uint32Key(u)
-	u2, err := k.AsUint32()
-	assert.NoError(t, err)
-	assert.Equal(t, u, u2)
-}
-
-func TestInt32RoundTrip(t *testing.T) {
-	tu.InitializeRandom()
-
-	tableName := tu.RandomString(10)
-	i := rand.Int31()
-
-	kb := newKeyBuilder(tableName, 0)
-	k := kb.Int32Key(i)
-	u2, err := k.AsInt32()
-	assert.NoError(t, err)
-	assert.Equal(t, i, u2)
+	assert.Equal(t, b, k.Bytes())
 }

common/kvstore/tablestore/table_store_builder.go

@@ -135,7 +135,7 @@ func validateSchema(
 	base kvstore.Store[[]byte],
 	metadataKeyBuilder kvstore.KeyBuilder) error {
 
-	schemaKey := metadataKeyBuilder.StringKey(metadataSchemaVersionKey)
+	schemaKey := metadataKeyBuilder.Key([]byte(metadataSchemaVersionKey))
 	onDiskSchemaBytes, err := base.Get(schemaKey.Raw())
 
 	if err != nil {
@@ -328,7 +328,7 @@ func handleIncompleteDeletion(
 	metadataKeyBuilder kvstore.KeyBuilder,
 	namespaceKeyBuilder kvstore.KeyBuilder) error {
 
-	deletionKey := metadataKeyBuilder.StringKey(metadataDeletionKey)
+	deletionKey := metadataKeyBuilder.Key([]byte(metadataDeletionKey))
 	deletionValue, err := base.Get(deletionKey.Raw())
 	if errors.Is(err, kvstore.ErrNotFound) {
 		// No deletion in progress, nothing to do.
@@ -389,7 +389,7 @@ func createTable(
 	batch := base.NewBatch()
 
 	var tableID uint32
-	tableIDBytes, err := base.Get(metadataKeyBuilder.StringKey(nextTableIDKey).Raw())
+	tableIDBytes, err := base.Get(metadataKeyBuilder.Key([]byte(nextTableIDKey)).Raw())
 	if err != nil {
 		if errors.Is(err, kvstore.ErrNotFound) {
 			tableIDBytes = make([]byte, 4)
@@ -407,8 +407,10 @@ func createTable(
 	nextTableIDBytes := make([]byte, 4)
 	binary.BigEndian.PutUint32(nextTableIDBytes, nextTableID)
 
-	batch.Put(namespaceKeyBuilder.Uint32Key(tableID).Raw(), []byte(name))
-	batch.Put(metadataKeyBuilder.StringKey(nextTableIDKey).Raw(), nextTableIDBytes)
+	keyBytes := make([]byte, 4)
+	binary.BigEndian.PutUint32(keyBytes, tableID)
+	batch.Put(namespaceKeyBuilder.Key(keyBytes).Raw(), []byte(name))
+	batch.Put(metadataKeyBuilder.Key([]byte(nextTableIDKey)).Raw(), nextTableIDBytes)
 
 	err = batch.Apply()
 	if err != nil {
@@ -429,7 +431,7 @@ func dropTable(
 	// This single atomic operation ensures that the table is deleted completely, even if we crash
 	// in the middle of the operation. When next starting up, if an entry is observed in this location,
 	// then the interrupted deletion can be completed.
-	deletionKey := metadataKeyBuilder.StringKey(metadataDeletionKey)
+	deletionKey := metadataKeyBuilder.Key([]byte(metadataDeletionKey))
 
 	deletionValue := make([]byte, 4+len(name))
 	binary.BigEndian.PutUint32(deletionValue, tableID)
@@ -455,7 +457,9 @@ func dropTable(
 	}
 
 	// All table entries have been deleted. Now delete the table from the namespace table.
-	tableKey := namespaceKeyBuilder.Uint32Key(tableID)
+	keyBytes := make([]byte, 4)
+	binary.BigEndian.PutUint32(keyBytes, tableID)
+	tableKey := namespaceKeyBuilder.Key(keyBytes)
 	err = base.Delete(tableKey.Raw())
 	if err != nil {
 		return fmt.Errorf("error deleting from namespace table: %w", err)

common/kvstore/tablestore/table_store_keys.go

@@ -24,48 +24,11 @@ func (k *key) Raw() []byte {
 	return k.data
 }
 
-// AsString interprets the key as a string and returns it.
-func (k *key) AsString() string {
-	return string(k.data[prefixLength:])
-}
-
-// AsBytes interprets the key as a byte slice and returns it.
-func (k *key) AsBytes() []byte {
+// Bytes interprets the key as a byte slice and returns it.
+func (k *key) Bytes() []byte {
 	return k.data[prefixLength:]
 }
 
-// AsUint64 interprets the key as a uint64 and returns it.
-func (k *key) AsUint64() (uint64, error) {
-	if len(k.data) != 12 {
-		return 0, kvstore.ErrInvalidKey
-	}
-	return binary.BigEndian.Uint64(k.data[prefixLength:]), nil
-}
-
-// AsInt64 interprets the key as an int64 and returns it.
-func (k *key) AsInt64() (int64, error) {
-	if len(k.data) != 12 {
-		return 0, kvstore.ErrInvalidKey
-	}
-	return int64(binary.BigEndian.Uint64(k.data[prefixLength:])), nil
-}
-
-// AsUint32 interprets the key as a uint32 and returns it.
-func (k *key) AsUint32() (uint32, error) {
-	if len(k.data) != 8 {
-		return 0, kvstore.ErrInvalidKey
-	}
-	return binary.BigEndian.Uint32(k.data[prefixLength:]), nil
-}
-
-// AsInt32 interprets the key as an int32 and returns it.
-func (k *key) AsInt32() (int32, error) {
-	if len(k.data) != 8 {
-		return 0, kvstore.ErrInvalidKey
-	}
-	return int32(binary.BigEndian.Uint32(k.data[prefixLength:])), nil
-}
-
 var _ kvstore.KeyBuilder = (*keyBuilder)(nil)
 
 type keyBuilder struct {
@@ -98,52 +61,3 @@ func (k *keyBuilder) Key(data []byte) kvstore.Key {
 		data:       result,
 	}
 }
-
-// StringKey creates a key from a string. Equivalent to Key([]byte(key)).
-func (k *keyBuilder) StringKey(data string) kvstore.Key {
-	return k.Key([]byte(data))
-}
-
-// Uint64Key creates a key from a uint64. Resulting key is an 8-byte big-endian representation of the uint64.
-func (k *keyBuilder) Uint64Key(data uint64) kvstore.Key {
-	result := make([]byte, 12)
-	copy(result, k.prefix)
-	binary.BigEndian.PutUint64(result[prefixLength:], data)
-	return &key{
-		keyBuilder: k,
-		data:       result,
-	}
-}
-
-// Int64Key creates a key from an int64. Resulting key is an 8-byte big-endian representation of the int64.
-func (k *keyBuilder) Int64Key(data int64) kvstore.Key {
-	result := make([]byte, 12)
-	copy(result, k.prefix)
-	binary.BigEndian.PutUint64(result[prefixLength:], uint64(data))
-	return &key{
-		keyBuilder: k,
-		data:       result,
-	}
-}
-
-// Uint32Key creates a key from a uint32. Resulting key is a 4-byte big-endian representation of the uint32.
-func (k *keyBuilder) Uint32Key(data uint32) kvstore.Key {
-	result := make([]byte, 8)
-	copy(result, k.prefix)
-	binary.BigEndian.PutUint32(result[prefixLength:], data)
-	return &key{
-		keyBuilder: k,
-		data:       result,
-	}
-}
-
-// Int32Key creates a key from an int32. Resulting key is a 4-byte big-endian representation of the int32.
-func (k *keyBuilder) Int32Key(data int32) kvstore.Key {
-	result := make([]byte, 8)
-	copy(result, k.prefix)
-	binary.BigEndian.PutUint32(result[prefixLength:], uint32(data))
-	return &key{
-		keyBuilder: k,
-		data:       result,
-	}
-}

common/kvstore/tablestore/table_store_test.go

@@ -168,28 +168,28 @@ func TestUniqueKeySpace(t *testing.T) {
 
 	// Write to the tables
 
-	err = store.Put(kb1.StringKey("key1"), []byte("value1"))
+	err = store.Put(kb1.Key([]byte("key1")), []byte("value1"))
 	assert.NoError(t, err)
-	err = store.Put(kb2.StringKey("key1"), []byte("value2"))
+	err = store.Put(kb2.Key([]byte("key1")), []byte("value2"))
 	assert.NoError(t, err)
 
-	value, err := store.Get(kb1.StringKey("key1"))
+	value, err := store.Get(kb1.Key([]byte("key1")))
 	assert.NoError(t, err)
 	assert.Equal(t, []byte("value1"), value)
 
-	value, err = store.Get(kb2.StringKey("key1"))
+	value, err = store.Get(kb2.Key([]byte("key1")))
 	assert.NoError(t, err)
 	assert.Equal(t, []byte("value2"), value)
 
 	// Delete a key from one table but not the other
 
-	err = store.Delete(kb1.StringKey("key1"))
+	err = store.Delete(kb1.Key([]byte("key1")))
 	assert.NoError(t, err)
 
-	_, err = store.Get(kb1.StringKey("key1"))
+	_, err = store.Get(kb1.Key([]byte("key1")))
 	assert.Equal(t, kvstore.ErrNotFound, err)
 
-	value, err = store.Get(kb2.StringKey("key1"))
+	value, err = store.Get(kb2.Key([]byte("key1")))
 	assert.NoError(t, err)
 	assert.Equal(t, []byte("value2"), value)
 
@@ -781,7 +781,7 @@ func TestIteration(t *testing.T) {
 	it.Release()
 
 	// Iterate over the "qwer" keys from table 1
-	it, err = store.NewIterator(kb1.StringKey("qwer"))
+	it, err = store.NewIterator(kb1.Key([]byte("qwer")))
 	assert.NoError(t, err)
 
 	count = 0
@@ -801,7 +801,7 @@ func TestIteration(t *testing.T) {
 	it.Release()
 
 	// Iterate over the "asdf" keys from table 2
-	it, err = store.NewIterator(kb2.StringKey("asdf"))
+	it, err = store.NewIterator(kb2.Key([]byte("asdf")))
 	assert.NoError(t, err)
 
 	count = 0
@@ -1032,7 +1032,7 @@ func TestRandomOperations(t *testing.T) {
 			}
 
 			delete(expectedData[tableName], k)
-			err = store.Delete(table.StringKey(k))
+			err = store.Delete(table.Key([]byte(k)))
 			assert.NoError(t, err)
 		}
 
@@ -1044,7 +1044,7 @@ func TestRandomOperations(t *testing.T) {
 
 				for k := range tableData {
 					expectedValue := tableData[k]
-					value, err := store.Get(table.StringKey(k))
+					value, err := store.Get(table.Key([]byte(k)))
 					assert.NoError(t, err)
 					assert.Equal(t, expectedValue, value)
 				}