Add key-only scan interface

Author: bootjp

store/bolt_store.go

@@ -95,6 +95,34 @@ func (s *boltStore) Scan(ctx context.Context, start []byte, end []byte, limit in
 	return res, errors.WithStack(err)
 }
 
+func (s *boltStore) ScanKeys(ctx context.Context, start []byte, end []byte, limit int) ([][]byte, error) {
+	s.log.InfoContext(ctx, "ScanKeys",
+		slog.String("start", string(start)),
+		slog.String("end", string(end)),
+		slog.Int("limit", limit),
+	)
+
+	var res [][]byte
+
+	err := s.bbolt.View(func(tx *bbolt.Tx) error {
+		b := tx.Bucket(defaultBucket)
+		if b == nil {
+			return nil
+		}
+
+		c := b.Cursor()
+		for k, _ := c.Seek(start); k != nil && (end == nil || bytes.Compare(k, end) < 0); k, _ = c.Next() {
+			res = append(res, k)
+			if len(res) >= limit {
+				break
+			}
+		}
+		return nil
+	})
+
+	return res, errors.WithStack(err)
+}
+
 func (s *boltStore) Put(ctx context.Context, key []byte, value []byte) error {
 	s.log.InfoContext(ctx, "put",
 		slog.String("key", string(key)),

store/bolt_store_test.go

@@ -83,6 +83,45 @@ func TestBoltStore_Scan(t *testing.T) {
 	assert.Equal(t, 100, cnt)
 }
 
+func TestBoltStore_ScanKeys(t *testing.T) {
+	ctx := context.Background()
+	t.Parallel()
+	st := mustStore(NewBoltStore(t.TempDir() + "/bolt.db"))
+
+	for i := 0; i < 999; i++ {
+		keyStr := "prefix " + strconv.Itoa(i) + "foo"
+		key := []byte(keyStr)
+		b := make([]byte, 8)
+		binary.PutVarint(b, int64(i))
+		err := st.Put(ctx, key, b)
+		assert.NoError(t, err)
+	}
+
+	res, err := st.ScanKeys(ctx, []byte("prefix"), []byte("z"), 100)
+	assert.NoError(t, err)
+	assert.Equal(t, 100, len(res))
+
+	sortedKeys := make([][]byte, 999)
+
+	for _, k := range res {
+		str := string(k)
+		i, err := strconv.Atoi(str[7 : len(str)-3])
+		assert.NoError(t, err)
+		sortedKeys[i] = k
+	}
+
+	cnt := 0
+	for i, k := range sortedKeys {
+		if k == nil {
+			continue
+		}
+		cnt++
+		assert.Equal(t, []byte("prefix "+strconv.Itoa(i)+"foo"), k)
+	}
+
+	assert.Equal(t, 100, cnt)
+}
+
 func TestBoltStore_Txn(t *testing.T) {
 	t.Parallel()
 	t.Run("success", func(t *testing.T) {

store/rb_memory_store.go

@@ -126,6 +126,36 @@ func (s *rbMemoryStore) Scan(ctx context.Context, start []byte, end []byte, limi
 	return result, nil
 }
 
+func (s *rbMemoryStore) ScanKeys(ctx context.Context, start []byte, end []byte, limit int) ([][]byte, error) {
+	s.mtx.RLock()
+	defer s.mtx.RUnlock()
+
+	var result [][]byte
+
+	s.tree.Each(func(key interface{}, _ interface{}) {
+		k, ok := key.([]byte)
+		if !ok {
+			return
+		}
+
+		if start != nil && bytes.Compare(k, start) < 0 {
+			return
+		}
+
+		if end != nil && bytes.Compare(k, end) > 0 {
+			return
+		}
+
+		if len(result) >= limit {
+			return
+		}
+
+		result = append(result, k)
+
+	})
+	return result, nil
+}
+
 func (s *rbMemoryStore) Put(ctx context.Context, key []byte, value []byte) error {
 	s.mtx.Lock()
 	defer s.mtx.Unlock()

store/rb_memory_store_test.go

@@ -85,6 +85,45 @@ func TestRbMemoryStore_Scan(t *testing.T) {
 	assert.Equal(t, 100, cnt)
 }
 
+func TestRbMemoryStore_ScanKeys(t *testing.T) {
+	ctx := context.Background()
+	t.Parallel()
+	st := NewRbMemoryStore()
+
+	for i := 0; i < 9999; i++ {
+		keyStr := "prefix " + strconv.Itoa(i) + "foo"
+		key := []byte(keyStr)
+		b := make([]byte, 8)
+		binary.PutVarint(b, int64(i))
+		err := st.Put(ctx, key, b)
+		assert.NoError(t, err)
+	}
+
+	res, err := st.ScanKeys(ctx, []byte("prefix"), []byte("z"), 100)
+	assert.NoError(t, err)
+	assert.Equal(t, 100, len(res))
+
+	sortedKeys := make([][]byte, 9999)
+
+	for _, k := range res {
+		str := string(k)
+		i, err := strconv.Atoi(str[7 : len(str)-3])
+		assert.NoError(t, err)
+		sortedKeys[i] = k
+	}
+
+	cnt := 0
+	for i, k := range sortedKeys {
+		if k == nil {
+			continue
+		}
+		cnt++
+		assert.Equal(t, []byte("prefix "+strconv.Itoa(i)+"foo"), k)
+	}
+
+	assert.Equal(t, 100, cnt)
+}
+
 func TestRbMemoryStore_Txn(t *testing.T) {
 	t.Parallel()
 	t.Run("success", func(t *testing.T) {

store/store.go

@@ -32,6 +32,7 @@ type Store interface {
 type ScanStore interface {
 	Store
 	Scan(ctx context.Context, start []byte, end []byte, limit int) ([]*KVPair, error)
+	ScanKeys(ctx context.Context, start []byte, end []byte, limit int) ([][]byte, error)
 }
 
 type TTLStore interface {
@@ -56,6 +57,7 @@ type Txn interface {
 type ScanTxn interface {
 	Txn
 	Scan(ctx context.Context, start []byte, end []byte, limit int) ([]*KVPair, error)
+	ScanKeys(ctx context.Context, start []byte, end []byte, limit int) ([][]byte, error)
 }
 
 type TTLTxn interface {