fix: Ensure Index.Walk fetches matching foreign keys only

This commit modifies the behaviour of the indexWalk function to ensure
it parses the key parts and matches the foreign key exactly.

Closes influxdata#20096

kv/index.go

@@ -128,7 +128,9 @@ func (i *Index) sourceBucket(tx Tx) (Bucket, error) {
 	return tx.Bucket(i.SourceBucket())
 }
 
-func indexKey(foreignKey, primaryKey []byte) (newKey []byte) {
+// IndexKey returns a value suitable for use as the key component
+// when storing values in the index.
+func IndexKey(foreignKey, primaryKey []byte) (newKey []byte) {
 	newKey = make([]byte, len(primaryKey)+len(foreignKey)+1)
 	copy(newKey, foreignKey)
 	newKey[len(foreignKey)] = '/'
@@ -157,7 +159,7 @@ func (i *Index) Insert(tx Tx, foreignKey, primaryKey []byte) error {
 		return err
 	}
 
-	return bkt.Put(indexKey(foreignKey, primaryKey), primaryKey)
+	return bkt.Put(IndexKey(foreignKey, primaryKey), primaryKey)
 }
 
 // Delete removes the foreignKey and primaryKey mapping from the underlying index.
@@ -167,7 +169,7 @@ func (i *Index) Delete(tx Tx, foreignKey, primaryKey []byte) error {
 		return err
 	}
 
-	return bkt.Delete(indexKey(foreignKey, primaryKey))
+	return bkt.Delete(IndexKey(foreignKey, primaryKey))
 }
 
 // Walk walks the source bucket using keys found in the index using the provided foreign key
@@ -195,16 +197,20 @@ func (i *Index) Walk(ctx context.Context, tx Tx, foreignKey []byte, visitFn Visi
 		return err
 	}
 
-	return indexWalk(ctx, cursor, sourceBucket, visitFn)
+	return indexWalk(ctx, foreignKey, cursor, sourceBucket, visitFn)
 }
 
 // indexWalk consumes the indexKey and primaryKey pairs in the index bucket and looks up their
 // associated primaryKey's value in the provided source bucket.
 // When an item is located in the source, the provided visit function is called with primary key and associated value.
-func indexWalk(ctx context.Context, indexCursor ForwardCursor, sourceBucket Bucket, visit VisitFunc) (err error) {
+func indexWalk(ctx context.Context, foreignKey []byte, indexCursor ForwardCursor, sourceBucket Bucket, visit VisitFunc) (err error) {
 	var keys [][]byte
 	for ik, pk := indexCursor.Next(); ik != nil; ik, pk = indexCursor.Next() {
-		keys = append(keys, pk)
+		if fk, _, err := indexKeyParts(ik); err != nil {
+			return err
+		} else if string(fk) == string(foreignKey) {
+			keys = append(keys, pk)
+		}
 	}
 
 	if err := indexCursor.Err(); err != nil {

kv/index_migration.go

@@ -128,7 +128,7 @@ func (i *IndexMigration) Populate(ctx context.Context, store Store) (n int, err
 
 	for fk, fkm := range diff.MissingFromIndex {
 		for pk := range fkm {
-			batch = append(batch, [2][]byte{indexKey([]byte(fk), []byte(pk)), []byte(pk)})
+			batch = append(batch, [2][]byte{IndexKey([]byte(fk), []byte(pk)), []byte(pk)})
 
 			if len(batch) >= i.operationBatchSize {
 				if err := flush(batch); err != nil {
@@ -183,7 +183,7 @@ func (i *IndexMigration) remove(ctx context.Context, store Store, mappings map[s
 
 	for fk, fkm := range mappings {
 		for pk := range fkm {
-			batch = append(batch, indexKey([]byte(fk), []byte(pk)))
+			batch = append(batch, IndexKey([]byte(fk), []byte(pk)))
 
 			if len(batch) >= i.operationBatchSize {
 				if err := flush(batch); err != nil {

kv/index_test.go

@@ -7,9 +7,14 @@ import (
 	"os"
 	"testing"
 
+	"github.com/golang/mock/gomock"
 	"github.com/influxdata/influxdb/v2/bolt"
 	"github.com/influxdata/influxdb/v2/inmem"
+	"github.com/influxdata/influxdb/v2/kv"
+	"github.com/influxdata/influxdb/v2/kv/mock"
 	influxdbtesting "github.com/influxdata/influxdb/v2/testing"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"go.uber.org/zap/zaptest"
 )
 
@@ -33,6 +38,102 @@ func Test_Bolt_Index(t *testing.T) {
 	influxdbtesting.TestIndex(t, s)
 }
 
+func TestIndex_Walk(t *testing.T) {
+	t.Run("only selects exact keys", func(t *testing.T) {
+		ctrl := gomock.NewController(t)
+		t.Cleanup(ctrl.Finish)
+
+		type keyValue struct{ key, val []byte }
+		makeIndexKV := func(fk, pk string) keyValue {
+			return keyValue{
+				key: kv.IndexKey([]byte(fk), []byte(pk)),
+				val: []byte(pk),
+			}
+		}
+
+		makeKV := func(key, val string) keyValue {
+			return keyValue{[]byte(key), []byte(val)}
+		}
+
+		var (
+			sourceBucket = []byte("source")
+			indexBucket  = []byte("index")
+			foreignKey   = []byte("jenkins")
+			idxkeyvals   = []keyValue{
+				makeIndexKV("jenkins-aws", "pk1"),
+				makeIndexKV("jenkins-aws", "pk2"),
+				makeIndexKV("jenkins-aws", "pk3"),
+				makeIndexKV("jenkins", "pk4"),
+				makeIndexKV("jenkins", "pk5"),
+			}
+			srckeyvals = []struct{ key, val []byte }{
+				makeKV("pk4", "val4"),
+				makeKV("pk5", "val5"),
+			}
+		)
+
+		mapping := kv.NewIndexMapping(sourceBucket, indexBucket, func(data []byte) ([]byte, error) {
+			return nil, nil
+		})
+
+		tx := mock.NewMockTx(ctrl)
+
+		src := mock.NewMockBucket(ctrl)
+		src.EXPECT().
+			GetBatch(srckeyvals[0].key, srckeyvals[1].key).
+			Return([][]byte{srckeyvals[0].val, srckeyvals[1].val}, nil)
+
+		tx.EXPECT().
+			Bucket(sourceBucket).
+			Return(src, nil)
+
+		idx := mock.NewMockBucket(ctrl)
+		tx.EXPECT().
+			Bucket(indexBucket).
+			Return(idx, nil)
+
+		cur := mock.NewMockForwardCursor(ctrl)
+
+		i := 0
+		cur.EXPECT().
+			Next().
+			DoAndReturn(func() ([]byte, []byte) {
+				var k, v []byte
+				if i < len(idxkeyvals) {
+					elem := idxkeyvals[i]
+					i++
+					k, v = elem.key, elem.val
+				}
+
+				return k, v
+			}).
+			Times(len(idxkeyvals) + 1)
+		cur.EXPECT().
+			Err().
+			Return(nil)
+		cur.EXPECT().
+			Close().
+			Return(nil)
+		idx.EXPECT().
+			ForwardCursor(foreignKey, gomock.Any()).
+			Return(cur, nil)
+
+		ctx := context.Background()
+		index := kv.NewIndex(mapping, kv.WithIndexReadPathEnabled)
+
+		j := 0
+		err := index.Walk(ctx, tx, foreignKey, func(k, v []byte) (bool, error) {
+			require.Less(t, j, len(srckeyvals))
+			assert.Equal(t, srckeyvals[j].key, k)
+			assert.Equal(t, srckeyvals[j].val, v)
+			j++
+			return true, nil
+		})
+
+		assert.NoError(t, err)
+	})
+}
+
 func Benchmark_Inmem_Index_Walk(b *testing.B) {
 	influxdbtesting.BenchmarkIndexWalk(b, inmem.NewKVStore(), 1000, 200)
 }