From 868a9fd03e337da9d09c0cee2f8dca31680669b4 Mon Sep 17 00:00:00 2001 From: tinswzy Date: Mon, 2 Dec 2024 22:43:19 +0800 Subject: [PATCH] refine interfaces for kv operation Signed-off-by: tinswzy --- cmd/milvus/mck.go | 22 +- internal/kv/etcd/embed_etcd_config_test.go | 3 +- internal/kv/etcd/embed_etcd_kv.go | 112 ++--- internal/kv/etcd/embed_etcd_kv_test.go | 177 ++++---- internal/kv/etcd/embed_etcd_restart_test.go | 9 +- internal/kv/etcd/etcd_kv.go | 110 ++--- internal/kv/etcd/etcd_kv_test.go | 153 +++---- internal/kv/mem/mem_kv.go | 49 +- internal/kv/mem/mem_kv_test.go | 67 +-- internal/kv/mock_snapshot_kv.go | 39 +- internal/kv/mock_snapshot_kv_test.go | 37 +- internal/kv/mocks/meta_kv.go | 359 ++++++++------- internal/kv/mocks/snapshot_kv.go | 152 ++++--- internal/kv/mocks/txn_kv.go | 309 +++++++------ internal/kv/mocks/watch_kv.go | 419 +++++++++--------- internal/kv/tikv/txn_tikv.go | 48 +- internal/kv/tikv/txn_tikv_test.go | 172 +++---- internal/metastore/kv/datacoord/kv_catalog.go | 132 +++--- .../metastore/kv/datacoord/kv_catalog_test.go | 204 ++++----- .../metastore/kv/querycoord/kv_catalog.go | 40 +- .../kv/querycoord/kv_catalog_test.go | 4 +- internal/metastore/kv/rootcoord/kv_catalog.go | 166 +++---- .../metastore/kv/rootcoord/kv_catalog_test.go | 417 ++++++++--------- .../metastore/kv/rootcoord/suffix_snapshot.go | 85 ++-- .../kv/rootcoord/suffix_snapshot_test.go | 111 ++--- .../metastore/kv/streamingcoord/kv_catalog.go | 4 +- .../kv/streamingcoord/kv_catalog_test.go | 12 +- .../metastore/kv/streamingnode/kv_catalog.go | 6 +- .../kv/streamingnode/kv_catalog_test.go | 6 +- pkg/kv/kv.go | 50 ++- pkg/kv/rocksdb/rocksdb_kv.go | 41 +- pkg/kv/rocksdb/rocksdb_kv_test.go | 153 +++---- pkg/mocks/mock_kv/mock_MetaKv.go | 359 ++++++++------- pkg/mq/mqimpl/rocksmq/server/rocksmq_impl.go | 25 +- .../rocksmq/server/rocksmq_retention.go | 7 +- .../rocksmq/server/rocksmq_retention_test.go | 41 +- 36 files changed, 2107 insertions(+), 1993 deletions(-) diff --git a/cmd/milvus/mck.go b/cmd/milvus/mck.go index 6581356c252b6..b22b28114c856 100644 --- a/cmd/milvus/mck.go +++ b/cmd/milvus/mck.go @@ -107,7 +107,7 @@ func (c *mck) execute(args []string, flags *flag.FlagSet) { func (c *mck) run() { c.connectMinio() - _, values, err := c.metaKV.LoadWithPrefix(segmentPrefix) + _, values, err := c.metaKV.LoadWithPrefix(context.TODO(), segmentPrefix) if err != nil { log.Fatal("failed to list the segment info", zap.String("key", segmentPrefix), zap.Error(err)) } @@ -121,7 +121,7 @@ func (c *mck) run() { c.segmentIDMap[info.ID] = info } - _, values, err = c.metaKV.LoadWithPrefix(collectionPrefix) + _, values, err = c.metaKV.LoadWithPrefix(context.TODO(), collectionPrefix) if err != nil { log.Fatal("failed to list the collection info", zap.String("key", collectionPrefix), zap.Error(err)) } @@ -150,13 +150,13 @@ func (c *mck) run() { } log.Info("partition ids", zap.Int64s("ids", ids)) - keys, values, err := c.metaKV.LoadWithPrefix(triggerTaskPrefix) + keys, values, err := c.metaKV.LoadWithPrefix(context.TODO(), triggerTaskPrefix) if err != nil { log.Fatal("failed to list the trigger task info", zap.Error(err)) } c.extractTask(triggerTaskPrefix, keys, values) - keys, values, err = c.metaKV.LoadWithPrefix(activeTaskPrefix) + keys, values, err = c.metaKV.LoadWithPrefix(context.TODO(), activeTaskPrefix) if err != nil { log.Fatal("failed to list the active task info", zap.Error(err)) } @@ -259,7 +259,7 @@ func getConfigValue(a string, b string, name string) string { } func (c *mck) cleanTrash() { - keys, _, err := c.metaKV.LoadWithPrefix(MckTrash) + keys, _, err := c.metaKV.LoadWithPrefix(context.TODO(), MckTrash) if err != nil { log.Error("failed to load backup info", zap.Error(err)) return @@ -273,7 +273,7 @@ func (c *mck) cleanTrash() { deleteAll := "" fmt.Scanln(&deleteAll) if deleteAll == "Y" { - err = c.metaKV.RemoveWithPrefix(MckTrash) + err = c.metaKV.RemoveWithPrefix(context.TODO(), MckTrash) if err != nil { log.Error("failed to remove backup infos", zap.String("key", MckTrash), zap.Error(err)) return @@ -395,31 +395,31 @@ func (c *mck) extractTask(prefix string, keys []string, values []string) { func (c *mck) removeTask(invalidTask int64) bool { taskType := c.taskNameMap[invalidTask] key := c.taskKeyMap[invalidTask] - err := c.metaKV.Save(getTrashKey(taskType, key), c.allTaskInfo[key]) + err := c.metaKV.Save(context.TODO(), getTrashKey(taskType, key), c.allTaskInfo[key]) if err != nil { log.Warn("failed to backup task", zap.String("key", getTrashKey(taskType, key)), zap.Int64("task_id", invalidTask), zap.Error(err)) return false } fmt.Printf("Back up task successfully, back path: %s\n", getTrashKey(taskType, key)) - err = c.metaKV.Remove(key) + err = c.metaKV.Remove(context.TODO(), key) if err != nil { log.Warn("failed to remove task", zap.Int64("task_id", invalidTask), zap.Error(err)) return false } key = fmt.Sprintf("%s/%d", taskInfoPrefix, invalidTask) - taskInfo, err := c.metaKV.Load(key) + taskInfo, err := c.metaKV.Load(context.TODO(), key) if err != nil { log.Warn("failed to load task info", zap.Int64("task_id", invalidTask), zap.Error(err)) return false } - err = c.metaKV.Save(getTrashKey(taskType, key), taskInfo) + err = c.metaKV.Save(context.TODO(), getTrashKey(taskType, key), taskInfo) if err != nil { log.Warn("failed to backup task info", zap.Int64("task_id", invalidTask), zap.Error(err)) return false } fmt.Printf("Back up task info successfully, back path: %s\n", getTrashKey(taskType, key)) - err = c.metaKV.Remove(key) + err = c.metaKV.Remove(context.TODO(), key) if err != nil { log.Warn("failed to remove task info", zap.Int64("task_id", invalidTask), zap.Error(err)) } diff --git a/internal/kv/etcd/embed_etcd_config_test.go b/internal/kv/etcd/embed_etcd_config_test.go index 1a121f6a09bf9..72c9ee156df3a 100644 --- a/internal/kv/etcd/embed_etcd_config_test.go +++ b/internal/kv/etcd/embed_etcd_config_test.go @@ -17,6 +17,7 @@ package etcdkv_test import ( + "context" "os" "testing" @@ -49,7 +50,7 @@ func TestEtcdConfigLoad(te *testing.T) { require.NoError(t, err) defer metaKv.Close() - defer metaKv.RemoveWithPrefix("") + defer metaKv.RemoveWithPrefix(context.TODO(), "") kv := metaKv.(*embed_etcd_kv.EmbedEtcdKV) assert.Equal(t, kv.GetConfig().SnapshotCount, uint64(1000)) diff --git a/internal/kv/etcd/embed_etcd_kv.go b/internal/kv/etcd/embed_etcd_kv.go index 754fb81e052ae..1c982d1f77053 100644 --- a/internal/kv/etcd/embed_etcd_kv.go +++ b/internal/kv/etcd/embed_etcd_kv.go @@ -123,9 +123,9 @@ func (kv *EmbedEtcdKV) GetPath(key string) string { return path.Join(kv.rootPath, key) } -func (kv *EmbedEtcdKV) WalkWithPrefix(prefix string, paginationSize int, fn func([]byte, []byte) error) error { +func (kv *EmbedEtcdKV) WalkWithPrefix(ctx context.Context, prefix string, paginationSize int, fn func([]byte, []byte) error) error { prefix = path.Join(kv.rootPath, prefix) - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() batch := int64(paginationSize) @@ -158,10 +158,10 @@ func (kv *EmbedEtcdKV) WalkWithPrefix(prefix string, paginationSize int, fn func } // LoadWithPrefix returns all the keys and values with the given key prefix -func (kv *EmbedEtcdKV) LoadWithPrefix(key string) ([]string, []string, error) { +func (kv *EmbedEtcdKV) LoadWithPrefix(ctx context.Context, key string) ([]string, []string, error) { key = path.Join(kv.rootPath, key) log.Debug("LoadWithPrefix ", zap.String("prefix", key)) - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() resp, err := kv.client.Get(ctx, key, clientv3.WithPrefix(), clientv3.WithSort(clientv3.SortByKey, clientv3.SortAscend)) @@ -178,10 +178,10 @@ func (kv *EmbedEtcdKV) LoadWithPrefix(key string) ([]string, []string, error) { return keys, values, nil } -func (kv *EmbedEtcdKV) Has(key string) (bool, error) { +func (kv *EmbedEtcdKV) Has(ctx context.Context, key string) (bool, error) { key = path.Join(kv.rootPath, key) log.Debug("Has", zap.String("key", key)) - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() resp, err := kv.client.Get(ctx, key, clientv3.WithCountOnly()) if err != nil { @@ -190,11 +190,11 @@ func (kv *EmbedEtcdKV) Has(key string) (bool, error) { return resp.Count != 0, nil } -func (kv *EmbedEtcdKV) HasPrefix(prefix string) (bool, error) { +func (kv *EmbedEtcdKV) HasPrefix(ctx context.Context, prefix string) (bool, error) { prefix = path.Join(kv.rootPath, prefix) log.Debug("HasPrefix", zap.String("prefix", prefix)) - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() resp, err := kv.client.Get(ctx, prefix, clientv3.WithPrefix(), clientv3.WithCountOnly(), clientv3.WithLimit(1)) @@ -206,10 +206,10 @@ func (kv *EmbedEtcdKV) HasPrefix(prefix string) (bool, error) { } // LoadBytesWithPrefix returns all the keys and values with the given key prefix -func (kv *EmbedEtcdKV) LoadBytesWithPrefix(key string) ([]string, [][]byte, error) { +func (kv *EmbedEtcdKV) LoadBytesWithPrefix(ctx context.Context, key string) ([]string, [][]byte, error) { key = path.Join(kv.rootPath, key) log.Debug("LoadBytesWithPrefix ", zap.String("prefix", key)) - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() resp, err := kv.client.Get(ctx, key, clientv3.WithPrefix(), clientv3.WithSort(clientv3.SortByKey, clientv3.SortAscend)) @@ -226,10 +226,10 @@ func (kv *EmbedEtcdKV) LoadBytesWithPrefix(key string) ([]string, [][]byte, erro } // LoadBytesWithPrefix2 returns all the keys and values with versions by the given key prefix -func (kv *EmbedEtcdKV) LoadBytesWithPrefix2(key string) ([]string, [][]byte, []int64, error) { +func (kv *EmbedEtcdKV) LoadBytesWithPrefix2(ctx context.Context, key string) ([]string, [][]byte, []int64, error) { key = path.Join(kv.rootPath, key) log.Debug("LoadBytesWithPrefix2 ", zap.String("prefix", key)) - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() resp, err := kv.client.Get(ctx, key, clientv3.WithPrefix(), clientv3.WithSort(clientv3.SortByKey, clientv3.SortAscend)) @@ -248,9 +248,9 @@ func (kv *EmbedEtcdKV) LoadBytesWithPrefix2(key string) ([]string, [][]byte, []i } // Load returns value of the given key -func (kv *EmbedEtcdKV) Load(key string) (string, error) { +func (kv *EmbedEtcdKV) Load(ctx context.Context, key string) (string, error) { key = path.Join(kv.rootPath, key) - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() resp, err := kv.client.Get(ctx, key) if err != nil { @@ -264,9 +264,9 @@ func (kv *EmbedEtcdKV) Load(key string) (string, error) { } // LoadBytes returns value of the given key -func (kv *EmbedEtcdKV) LoadBytes(key string) ([]byte, error) { +func (kv *EmbedEtcdKV) LoadBytes(ctx context.Context, key string) ([]byte, error) { key = path.Join(kv.rootPath, key) - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() resp, err := kv.client.Get(ctx, key) if err != nil { @@ -280,13 +280,13 @@ func (kv *EmbedEtcdKV) LoadBytes(key string) ([]byte, error) { } // MultiLoad returns values of a set of keys -func (kv *EmbedEtcdKV) MultiLoad(keys []string) ([]string, error) { +func (kv *EmbedEtcdKV) MultiLoad(ctx context.Context, keys []string) ([]string, error) { ops := make([]clientv3.Op, 0, len(keys)) for _, keyLoad := range keys { ops = append(ops, clientv3.OpGet(path.Join(kv.rootPath, keyLoad))) } - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() resp, err := kv.client.Txn(ctx).If().Then(ops...).Commit() if err != nil { @@ -316,13 +316,13 @@ func (kv *EmbedEtcdKV) MultiLoad(keys []string) ([]string, error) { } // MultiLoadBytes returns values of a set of keys -func (kv *EmbedEtcdKV) MultiLoadBytes(keys []string) ([][]byte, error) { +func (kv *EmbedEtcdKV) MultiLoadBytes(ctx context.Context, keys []string) ([][]byte, error) { ops := make([]clientv3.Op, 0, len(keys)) for _, keyLoad := range keys { ops = append(ops, clientv3.OpGet(path.Join(kv.rootPath, keyLoad))) } - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() resp, err := kv.client.Txn(ctx).If().Then(ops...).Commit() if err != nil { @@ -352,10 +352,10 @@ func (kv *EmbedEtcdKV) MultiLoadBytes(keys []string) ([][]byte, error) { } // LoadBytesWithRevision returns keys, values and revision with given key prefix. -func (kv *EmbedEtcdKV) LoadBytesWithRevision(key string) ([]string, [][]byte, int64, error) { +func (kv *EmbedEtcdKV) LoadBytesWithRevision(ctx context.Context, key string) ([]string, [][]byte, int64, error) { key = path.Join(kv.rootPath, key) log.Debug("LoadBytesWithRevision ", zap.String("prefix", key)) - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() resp, err := kv.client.Get(ctx, key, clientv3.WithPrefix(), clientv3.WithSort(clientv3.SortByKey, clientv3.SortAscend)) @@ -372,40 +372,40 @@ func (kv *EmbedEtcdKV) LoadBytesWithRevision(key string) ([]string, [][]byte, in } // Save saves the key-value pair. -func (kv *EmbedEtcdKV) Save(key, value string) error { +func (kv *EmbedEtcdKV) Save(ctx context.Context, key, value string) error { key = path.Join(kv.rootPath, key) - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() _, err := kv.client.Put(ctx, key, value) return err } // SaveBytes saves the key-value pair. -func (kv *EmbedEtcdKV) SaveBytes(key string, value []byte) error { +func (kv *EmbedEtcdKV) SaveBytes(ctx context.Context, key string, value []byte) error { key = path.Join(kv.rootPath, key) - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() _, err := kv.client.Put(ctx, key, string(value)) return err } // SaveBytesWithLease is a function to put value in etcd with etcd lease options. -func (kv *EmbedEtcdKV) SaveBytesWithLease(key string, value []byte, id clientv3.LeaseID) error { +func (kv *EmbedEtcdKV) SaveBytesWithLease(ctx context.Context, key string, value []byte, id clientv3.LeaseID) error { key = path.Join(kv.rootPath, key) - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() _, err := kv.client.Put(ctx, key, string(value), clientv3.WithLease(id)) return err } // MultiSave saves the key-value pairs in a transaction. -func (kv *EmbedEtcdKV) MultiSave(kvs map[string]string) error { +func (kv *EmbedEtcdKV) MultiSave(ctx context.Context, kvs map[string]string) error { ops := make([]clientv3.Op, 0, len(kvs)) for key, value := range kvs { ops = append(ops, clientv3.OpPut(path.Join(kv.rootPath, key), value)) } - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() _, err := kv.client.Txn(ctx).If().Then(ops...).Commit() @@ -413,13 +413,13 @@ func (kv *EmbedEtcdKV) MultiSave(kvs map[string]string) error { } // MultiSaveBytes saves the key-value pairs in a transaction. -func (kv *EmbedEtcdKV) MultiSaveBytes(kvs map[string][]byte) error { +func (kv *EmbedEtcdKV) MultiSaveBytes(ctx context.Context, kvs map[string][]byte) error { ops := make([]clientv3.Op, 0, len(kvs)) for key, value := range kvs { ops = append(ops, clientv3.OpPut(path.Join(kv.rootPath, key), string(value))) } - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() _, err := kv.client.Txn(ctx).If().Then(ops...).Commit() @@ -427,9 +427,9 @@ func (kv *EmbedEtcdKV) MultiSaveBytes(kvs map[string][]byte) error { } // RemoveWithPrefix removes the keys with given prefix. -func (kv *EmbedEtcdKV) RemoveWithPrefix(prefix string) error { +func (kv *EmbedEtcdKV) RemoveWithPrefix(ctx context.Context, prefix string) error { key := path.Join(kv.rootPath, prefix) - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() _, err := kv.client.Delete(ctx, key, clientv3.WithPrefix()) @@ -437,9 +437,9 @@ func (kv *EmbedEtcdKV) RemoveWithPrefix(prefix string) error { } // Remove removes the key. -func (kv *EmbedEtcdKV) Remove(key string) error { +func (kv *EmbedEtcdKV) Remove(ctx context.Context, key string) error { key = path.Join(kv.rootPath, key) - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() _, err := kv.client.Delete(ctx, key) @@ -447,13 +447,13 @@ func (kv *EmbedEtcdKV) Remove(key string) error { } // MultiRemove removes the keys in a transaction. -func (kv *EmbedEtcdKV) MultiRemove(keys []string) error { +func (kv *EmbedEtcdKV) MultiRemove(ctx context.Context, keys []string) error { ops := make([]clientv3.Op, 0, len(keys)) for _, key := range keys { ops = append(ops, clientv3.OpDelete(path.Join(kv.rootPath, key))) } - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() _, err := kv.client.Txn(ctx).If().Then(ops...).Commit() @@ -461,7 +461,7 @@ func (kv *EmbedEtcdKV) MultiRemove(keys []string) error { } // MultiSaveAndRemove saves the key-value pairs and removes the keys in a transaction. -func (kv *EmbedEtcdKV) MultiSaveAndRemove(saves map[string]string, removals []string, preds ...predicates.Predicate) error { +func (kv *EmbedEtcdKV) MultiSaveAndRemove(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error { cmps, err := parsePredicates(kv.rootPath, preds...) if err != nil { return err @@ -476,7 +476,7 @@ func (kv *EmbedEtcdKV) MultiSaveAndRemove(saves map[string]string, removals []st ops = append(ops, clientv3.OpDelete(path.Join(kv.rootPath, keyDelete))) } - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() resp, err := kv.client.Txn(ctx).If(cmps...).Then(ops...).Commit() @@ -491,7 +491,7 @@ func (kv *EmbedEtcdKV) MultiSaveAndRemove(saves map[string]string, removals []st } // MultiSaveBytesAndRemove saves the key-value pairs and removes the keys in a transaction. -func (kv *EmbedEtcdKV) MultiSaveBytesAndRemove(saves map[string][]byte, removals []string) error { +func (kv *EmbedEtcdKV) MultiSaveBytesAndRemove(ctx context.Context, saves map[string][]byte, removals []string) error { ops := make([]clientv3.Op, 0, len(saves)+len(removals)) for key, value := range saves { ops = append(ops, clientv3.OpPut(path.Join(kv.rootPath, key), string(value))) @@ -501,33 +501,33 @@ func (kv *EmbedEtcdKV) MultiSaveBytesAndRemove(saves map[string][]byte, removals ops = append(ops, clientv3.OpDelete(path.Join(kv.rootPath, keyDelete))) } - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() _, err := kv.client.Txn(ctx).If().Then(ops...).Commit() return err } -func (kv *EmbedEtcdKV) Watch(key string) clientv3.WatchChan { +func (kv *EmbedEtcdKV) Watch(ctx context.Context, key string) clientv3.WatchChan { key = path.Join(kv.rootPath, key) - rch := kv.client.Watch(context.Background(), key, clientv3.WithCreatedNotify()) + rch := kv.client.Watch(ctx, key, clientv3.WithCreatedNotify()) return rch } -func (kv *EmbedEtcdKV) WatchWithPrefix(key string) clientv3.WatchChan { +func (kv *EmbedEtcdKV) WatchWithPrefix(ctx context.Context, key string) clientv3.WatchChan { key = path.Join(kv.rootPath, key) - rch := kv.client.Watch(context.Background(), key, clientv3.WithPrefix(), clientv3.WithCreatedNotify()) + rch := kv.client.Watch(ctx, key, clientv3.WithPrefix(), clientv3.WithCreatedNotify()) return rch } -func (kv *EmbedEtcdKV) WatchWithRevision(key string, revision int64) clientv3.WatchChan { +func (kv *EmbedEtcdKV) WatchWithRevision(ctx context.Context, key string, revision int64) clientv3.WatchChan { key = path.Join(kv.rootPath, key) - rch := kv.client.Watch(context.Background(), key, clientv3.WithPrefix(), clientv3.WithPrevKV(), clientv3.WithRev(revision)) + rch := kv.client.Watch(ctx, key, clientv3.WithPrefix(), clientv3.WithPrevKV(), clientv3.WithRev(revision)) return rch } // MultiSaveAndRemoveWithPrefix saves kv in @saves and removes the keys with given prefix in @removals. -func (kv *EmbedEtcdKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals []string, preds ...predicates.Predicate) error { +func (kv *EmbedEtcdKV) MultiSaveAndRemoveWithPrefix(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error { cmps, err := parsePredicates(kv.rootPath, preds...) if err != nil { return err @@ -542,7 +542,7 @@ func (kv *EmbedEtcdKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, rem ops = append(ops, clientv3.OpDelete(path.Join(kv.rootPath, keyDelete), clientv3.WithPrefix())) } - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() resp, err := kv.client.Txn(ctx).If(cmps...).Then(ops...).Commit() @@ -557,7 +557,7 @@ func (kv *EmbedEtcdKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, rem } // MultiSaveBytesAndRemoveWithPrefix saves kv in @saves and removes the keys with given prefix in @removals. -func (kv *EmbedEtcdKV) MultiSaveBytesAndRemoveWithPrefix(saves map[string][]byte, removals []string) error { +func (kv *EmbedEtcdKV) MultiSaveBytesAndRemoveWithPrefix(ctx context.Context, saves map[string][]byte, removals []string) error { ops := make([]clientv3.Op, 0, len(saves)+len(removals)) for key, value := range saves { ops = append(ops, clientv3.OpPut(path.Join(kv.rootPath, key), string(value))) @@ -567,7 +567,7 @@ func (kv *EmbedEtcdKV) MultiSaveBytesAndRemoveWithPrefix(saves map[string][]byte ops = append(ops, clientv3.OpDelete(path.Join(kv.rootPath, keyDelete), clientv3.WithPrefix())) } - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() _, err := kv.client.Txn(ctx).If().Then(ops...).Commit() @@ -576,8 +576,8 @@ func (kv *EmbedEtcdKV) MultiSaveBytesAndRemoveWithPrefix(saves map[string][]byte // CompareVersionAndSwap compares the existing key-value's version with version, and if // they are equal, the target is stored in etcd. -func (kv *EmbedEtcdKV) CompareVersionAndSwap(key string, version int64, target string) (bool, error) { - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) +func (kv *EmbedEtcdKV) CompareVersionAndSwap(ctx context.Context, key string, version int64, target string) (bool, error) { + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() resp, err := kv.client.Txn(ctx).If( clientv3.Compare( @@ -593,8 +593,8 @@ func (kv *EmbedEtcdKV) CompareVersionAndSwap(key string, version int64, target s // CompareVersionAndSwapBytes compares the existing key-value's version with version, and if // they are equal, the target is stored in etcd. -func (kv *EmbedEtcdKV) CompareVersionAndSwapBytes(key string, version int64, target []byte, opts ...clientv3.OpOption) (bool, error) { - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) +func (kv *EmbedEtcdKV) CompareVersionAndSwapBytes(ctx context.Context, key string, version int64, target []byte, opts ...clientv3.OpOption) (bool, error) { + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() resp, err := kv.client.Txn(ctx).If( clientv3.Compare( diff --git a/internal/kv/etcd/embed_etcd_kv_test.go b/internal/kv/etcd/embed_etcd_kv_test.go index 7e01b5b5516b9..f18351417e4bc 100644 --- a/internal/kv/etcd/embed_etcd_kv_test.go +++ b/internal/kv/etcd/embed_etcd_kv_test.go @@ -17,6 +17,7 @@ package etcdkv_test import ( + "context" "fmt" "path" "sort" @@ -55,7 +56,7 @@ func TestEmbedEtcd(te *testing.T) { require.NoError(t, err) defer metaKv.Close() - defer metaKv.RemoveWithPrefix("") + defer metaKv.RemoveWithPrefix(context.TODO(), "") saveAndLoadTests := []struct { key string @@ -69,11 +70,11 @@ func TestEmbedEtcd(te *testing.T) { for i, test := range saveAndLoadTests { if i < 4 { - err = metaKv.Save(test.key, test.value) + err = metaKv.Save(context.TODO(), test.key, test.value) assert.NoError(t, err) } - val, err := metaKv.Load(test.key) + val, err := metaKv.Load(context.TODO(), test.key) assert.NoError(t, err) assert.Equal(t, test.value, val) } @@ -87,7 +88,7 @@ func TestEmbedEtcd(te *testing.T) { } for _, test := range invalidLoadTests { - val, err := metaKv.Load(test.invalidKey) + val, err := metaKv.Load(context.TODO(), test.invalidKey) assert.Error(t, err) assert.Zero(t, val) } @@ -124,7 +125,7 @@ func TestEmbedEtcd(te *testing.T) { } for _, test := range loadPrefixTests { - actualKeys, actualValues, err := metaKv.LoadWithPrefix(test.prefix) + actualKeys, actualValues, err := metaKv.LoadWithPrefix(context.TODO(), test.prefix) assert.ElementsMatch(t, test.expectedKeys, actualKeys) assert.ElementsMatch(t, test.expectedValues, actualValues) assert.Equal(t, test.expectedError, err) @@ -141,15 +142,15 @@ func TestEmbedEtcd(te *testing.T) { } for _, test := range removeTests { - err = metaKv.Remove(test.validKey) + err = metaKv.Remove(context.TODO(), test.validKey) assert.NoError(t, err) - _, err = metaKv.Load(test.validKey) + _, err = metaKv.Load(context.TODO(), test.validKey) assert.Error(t, err) - err = metaKv.Remove(test.validKey) + err = metaKv.Remove(context.TODO(), test.validKey) assert.NoError(t, err) - err = metaKv.Remove(test.invalidKey) + err = metaKv.Remove(context.TODO(), test.invalidKey) assert.NoError(t, err) } }) @@ -163,7 +164,7 @@ func TestEmbedEtcd(te *testing.T) { require.NoError(t, err) defer metaKv.Close() - defer metaKv.RemoveWithPrefix("") + defer metaKv.RemoveWithPrefix(context.TODO(), "") saveAndLoadTests := []struct { key string @@ -177,11 +178,11 @@ func TestEmbedEtcd(te *testing.T) { for i, test := range saveAndLoadTests { if i < 4 { - err = metaKv.SaveBytes(test.key, test.value) + err = metaKv.SaveBytes(context.TODO(), test.key, test.value) assert.NoError(t, err) } - val, err := metaKv.LoadBytes(test.key) + val, err := metaKv.LoadBytes(context.TODO(), test.key) assert.NoError(t, err) assert.Equal(t, test.value, val) } @@ -195,7 +196,7 @@ func TestEmbedEtcd(te *testing.T) { } for _, test := range invalidLoadTests { - val, err := metaKv.LoadBytes(test.invalidKey) + val, err := metaKv.LoadBytes(context.TODO(), test.invalidKey) assert.Error(t, err) assert.Zero(t, val) } @@ -232,12 +233,12 @@ func TestEmbedEtcd(te *testing.T) { } for _, test := range loadPrefixTests { - actualKeys, actualValues, err := metaKv.LoadBytesWithPrefix(test.prefix) + actualKeys, actualValues, err := metaKv.LoadBytesWithPrefix(context.TODO(), test.prefix) assert.ElementsMatch(t, test.expectedKeys, actualKeys) assert.ElementsMatch(t, test.expectedValues, actualValues) assert.Equal(t, test.expectedError, err) - actualKeys, actualValues, versions, err := metaKv.LoadBytesWithPrefix2(test.prefix) + actualKeys, actualValues, versions, err := metaKv.LoadBytesWithPrefix2(context.TODO(), test.prefix) assert.ElementsMatch(t, test.expectedKeys, actualKeys) assert.ElementsMatch(t, test.expectedValues, actualValues) assert.NotZero(t, versions) @@ -255,15 +256,15 @@ func TestEmbedEtcd(te *testing.T) { } for _, test := range removeTests { - err = metaKv.Remove(test.validKey) + err = metaKv.Remove(context.TODO(), test.validKey) assert.NoError(t, err) - _, err = metaKv.Load(test.validKey) + _, err = metaKv.Load(context.TODO(), test.validKey) assert.Error(t, err) - err = metaKv.Remove(test.validKey) + err = metaKv.Remove(context.TODO(), test.validKey) assert.NoError(t, err) - err = metaKv.Remove(test.invalidKey) + err = metaKv.Remove(context.TODO(), test.invalidKey) assert.NoError(t, err) } }) @@ -275,7 +276,7 @@ func TestEmbedEtcd(te *testing.T) { assert.NoError(t, err) defer metaKv.Close() - defer metaKv.RemoveWithPrefix("") + defer metaKv.RemoveWithPrefix(context.TODO(), "") prepareKV := []struct { inKey string @@ -289,7 +290,7 @@ func TestEmbedEtcd(te *testing.T) { } for _, test := range prepareKV { - err = metaKv.SaveBytes(test.inKey, test.inValue) + err = metaKv.SaveBytes(context.TODO(), test.inKey, test.inValue) require.NoError(t, err) } @@ -305,7 +306,7 @@ func TestEmbedEtcd(te *testing.T) { } for _, test := range loadWithRevisionTests { - keys, values, revision, err := metaKv.LoadBytesWithRevision(test.inKey) + keys, values, revision, err := metaKv.LoadBytesWithRevision(context.TODO(), test.inKey) assert.NoError(t, err) assert.Equal(t, test.expectedKeyNo, len(keys)) assert.ElementsMatch(t, test.expectedValues, values) @@ -319,7 +320,7 @@ func TestEmbedEtcd(te *testing.T) { assert.NoError(t, err) defer metaKv.Close() - defer metaKv.RemoveWithPrefix("") + defer metaKv.RemoveWithPrefix(context.TODO(), "") multiSaveTests := map[string]string{ "key_1": "value_1", @@ -330,10 +331,10 @@ func TestEmbedEtcd(te *testing.T) { "_": "other", } - err = metaKv.MultiSave(multiSaveTests) + err = metaKv.MultiSave(context.TODO(), multiSaveTests) assert.NoError(t, err) for k, v := range multiSaveTests { - actualV, err := metaKv.Load(k) + actualV, err := metaKv.Load(context.TODO(), k) assert.NoError(t, err) assert.Equal(t, v, actualV) } @@ -349,7 +350,7 @@ func TestEmbedEtcd(te *testing.T) { } for _, test := range multiLoadTests { - vs, err := metaKv.MultiLoad(test.inputKeys) + vs, err := metaKv.MultiLoad(context.TODO(), test.inputKeys) assert.NoError(t, err) assert.Equal(t, test.expectedValues, vs) } @@ -365,7 +366,7 @@ func TestEmbedEtcd(te *testing.T) { } for _, test := range invalidMultiLoad { - vs, err := metaKv.MultiLoad(test.invalidKeys) + vs, err := metaKv.MultiLoad(context.TODO(), test.invalidKeys) assert.Error(t, err) assert.Equal(t, test.expectedValues, vs) } @@ -376,10 +377,10 @@ func TestEmbedEtcd(te *testing.T) { } for _, k := range removeWithPrefixTests { - err = metaKv.RemoveWithPrefix(k) + err = metaKv.RemoveWithPrefix(context.TODO(), k) assert.NoError(t, err) - ks, vs, err := metaKv.LoadWithPrefix(k) + ks, vs, err := metaKv.LoadWithPrefix(context.TODO(), k) assert.Empty(t, ks) assert.Empty(t, vs) assert.NoError(t, err) @@ -392,10 +393,10 @@ func TestEmbedEtcd(te *testing.T) { "_", } - err = metaKv.MultiRemove(multiRemoveTests) + err = metaKv.MultiRemove(context.TODO(), multiRemoveTests) assert.NoError(t, err) - ks, vs, err := metaKv.LoadWithPrefix("") + ks, vs, err := metaKv.LoadWithPrefix(context.TODO(), "") assert.NoError(t, err) assert.Empty(t, ks) assert.Empty(t, vs) @@ -412,11 +413,11 @@ func TestEmbedEtcd(te *testing.T) { {make(map[string]string), []string{"multikey_2"}}, } for _, test := range multiSaveAndRemoveTests { - err = metaKv.MultiSaveAndRemove(test.multiSaves, test.multiRemoves) + err = metaKv.MultiSaveAndRemove(context.TODO(), test.multiSaves, test.multiRemoves) assert.NoError(t, err) } - ks, vs, err = metaKv.LoadWithPrefix("") + ks, vs, err = metaKv.LoadWithPrefix(context.TODO(), "") assert.NoError(t, err) assert.Empty(t, ks) assert.Empty(t, vs) @@ -429,7 +430,7 @@ func TestEmbedEtcd(te *testing.T) { assert.NoError(t, err) defer metaKv.Close() - defer metaKv.RemoveWithPrefix("") + defer metaKv.RemoveWithPrefix(context.TODO(), "") multiSaveTests := map[string][]byte{ "key_1": []byte("value_1"), @@ -440,10 +441,10 @@ func TestEmbedEtcd(te *testing.T) { "_": []byte("other"), } - err = metaKv.MultiSaveBytes(multiSaveTests) + err = metaKv.MultiSaveBytes(context.TODO(), multiSaveTests) assert.NoError(t, err) for k, v := range multiSaveTests { - actualV, err := metaKv.LoadBytes(k) + actualV, err := metaKv.LoadBytes(context.TODO(), k) assert.NoError(t, err) assert.Equal(t, v, actualV) } @@ -459,7 +460,7 @@ func TestEmbedEtcd(te *testing.T) { } for _, test := range multiLoadTests { - vs, err := metaKv.MultiLoadBytes(test.inputKeys) + vs, err := metaKv.MultiLoadBytes(context.TODO(), test.inputKeys) assert.NoError(t, err) assert.Equal(t, test.expectedValues, vs) } @@ -475,7 +476,7 @@ func TestEmbedEtcd(te *testing.T) { } for _, test := range invalidMultiLoad { - vs, err := metaKv.MultiLoadBytes(test.invalidKeys) + vs, err := metaKv.MultiLoadBytes(context.TODO(), test.invalidKeys) assert.Error(t, err) assert.Equal(t, test.expectedValues, vs) } @@ -486,10 +487,10 @@ func TestEmbedEtcd(te *testing.T) { } for _, k := range removeWithPrefixTests { - err = metaKv.RemoveWithPrefix(k) + err = metaKv.RemoveWithPrefix(context.TODO(), k) assert.NoError(t, err) - ks, vs, err := metaKv.LoadBytesWithPrefix(k) + ks, vs, err := metaKv.LoadBytesWithPrefix(context.TODO(), k) assert.Empty(t, ks) assert.Empty(t, vs) assert.NoError(t, err) @@ -502,10 +503,10 @@ func TestEmbedEtcd(te *testing.T) { "_", } - err = metaKv.MultiRemove(multiRemoveTests) + err = metaKv.MultiRemove(context.TODO(), multiRemoveTests) assert.NoError(t, err) - ks, vs, err := metaKv.LoadBytesWithPrefix("") + ks, vs, err := metaKv.LoadBytesWithPrefix(context.TODO(), "") assert.NoError(t, err) assert.Empty(t, ks) assert.Empty(t, vs) @@ -522,11 +523,11 @@ func TestEmbedEtcd(te *testing.T) { {map[string][]byte{}, []string{"multikey_2"}}, } for _, test := range multiSaveAndRemoveTests { - err = metaKv.MultiSaveBytesAndRemove(test.multiSaves, test.multiRemoves) + err = metaKv.MultiSaveBytesAndRemove(context.TODO(), test.multiSaves, test.multiRemoves) assert.NoError(t, err) } - ks, vs, err = metaKv.LoadBytesWithPrefix("") + ks, vs, err = metaKv.LoadBytesWithPrefix(context.TODO(), "") assert.NoError(t, err) assert.Empty(t, ks) assert.Empty(t, vs) @@ -538,7 +539,7 @@ func TestEmbedEtcd(te *testing.T) { require.NoError(t, err) defer metaKv.Close() - defer metaKv.RemoveWithPrefix("") + defer metaKv.RemoveWithPrefix(context.TODO(), "") prepareTests := map[string]string{ "x/abc/1": "1", @@ -550,7 +551,7 @@ func TestEmbedEtcd(te *testing.T) { } // MultiSaveAndRemoveWithPrefix - err = metaKv.MultiSave(prepareTests) + err = metaKv.MultiSave(context.TODO(), prepareTests) require.NoError(t, err) multiSaveAndRemoveWithPrefixTests := []struct { multiSave map[string]string @@ -569,14 +570,14 @@ func TestEmbedEtcd(te *testing.T) { } for _, test := range multiSaveAndRemoveWithPrefixTests { - k, _, err := metaKv.LoadWithPrefix(test.loadPrefix) + k, _, err := metaKv.LoadWithPrefix(context.TODO(), test.loadPrefix) assert.NoError(t, err) assert.Equal(t, test.lengthBeforeRemove, len(k)) - err = metaKv.MultiSaveAndRemoveWithPrefix(test.multiSave, test.prefix) + err = metaKv.MultiSaveAndRemoveWithPrefix(context.TODO(), test.multiSave, test.prefix) assert.NoError(t, err) - k, _, err = metaKv.LoadWithPrefix(test.loadPrefix) + k, _, err = metaKv.LoadWithPrefix(context.TODO(), test.loadPrefix) assert.NoError(t, err) assert.Equal(t, test.lengthAfterRemove, len(k)) } @@ -589,7 +590,7 @@ func TestEmbedEtcd(te *testing.T) { require.NoError(t, err) defer metaKv.Close() - defer metaKv.RemoveWithPrefix("") + defer metaKv.RemoveWithPrefix(context.TODO(), "") prepareTests := map[string][]byte{ "x/abc/1": []byte("1"), @@ -600,13 +601,13 @@ func TestEmbedEtcd(te *testing.T) { "x/den/2": []byte("200"), } - k, v, err := metaKv.LoadBytesWithPrefix("/") + k, v, err := metaKv.LoadBytesWithPrefix(context.TODO(), "/") assert.NoError(t, err) assert.Zero(t, len(k)) assert.Zero(t, len(v)) // MultiSaveAndRemoveWithPrefix - err = metaKv.MultiSaveBytes(prepareTests) + err = metaKv.MultiSaveBytes(context.TODO(), prepareTests) require.NoError(t, err) multiSaveAndRemoveWithPrefixTests := []struct { multiSave map[string][]byte @@ -624,14 +625,14 @@ func TestEmbedEtcd(te *testing.T) { } for _, test := range multiSaveAndRemoveWithPrefixTests { - k, _, err = metaKv.LoadBytesWithPrefix(test.loadPrefix) + k, _, err = metaKv.LoadBytesWithPrefix(context.TODO(), test.loadPrefix) assert.NoError(t, err) assert.Equal(t, test.lengthBeforeRemove, len(k)) - err = metaKv.MultiSaveBytesAndRemoveWithPrefix(test.multiSave, test.prefix) + err = metaKv.MultiSaveBytesAndRemoveWithPrefix(context.TODO(), test.multiSave, test.prefix) assert.NoError(t, err) - k, _, err = metaKv.LoadBytesWithPrefix(test.loadPrefix) + k, _, err = metaKv.LoadBytesWithPrefix(context.TODO(), test.loadPrefix) assert.NoError(t, err) assert.Equal(t, test.lengthAfterRemove, len(k)) } @@ -643,13 +644,13 @@ func TestEmbedEtcd(te *testing.T) { assert.NoError(t, err) defer watchKv.Close() - defer watchKv.RemoveWithPrefix("") + defer watchKv.RemoveWithPrefix(context.TODO(), "") - ch := watchKv.Watch("x") + ch := watchKv.Watch(context.TODO(), "x") resp := <-ch assert.True(t, resp.Created) - ch = watchKv.WatchWithPrefix("x") + ch = watchKv.WatchWithPrefix(context.TODO(), "x") resp = <-ch assert.True(t, resp.Created) }) @@ -661,7 +662,7 @@ func TestEmbedEtcd(te *testing.T) { assert.NoError(t, err) defer metaKv.Close() - defer metaKv.RemoveWithPrefix("") + defer metaKv.RemoveWithPrefix(context.TODO(), "") revisionTests := []struct { inKey string @@ -674,13 +675,13 @@ func TestEmbedEtcd(te *testing.T) { } for _, test := range revisionTests { - err = metaKv.SaveBytes(test.inKey, test.fistValue) + err = metaKv.SaveBytes(context.TODO(), test.inKey, test.fistValue) require.NoError(t, err) - _, _, revision, _ := metaKv.LoadBytesWithRevision(test.inKey) - ch := metaKv.WatchWithRevision(test.inKey, revision+1) + _, _, revision, _ := metaKv.LoadBytesWithRevision(context.TODO(), test.inKey) + ch := metaKv.WatchWithRevision(context.TODO(), test.inKey, revision+1) - err = metaKv.SaveBytes(test.inKey, test.secondValue) + err = metaKv.SaveBytes(context.TODO(), test.inKey, test.secondValue) require.NoError(t, err) resp := <-ch @@ -689,15 +690,15 @@ func TestEmbedEtcd(te *testing.T) { assert.Equal(t, revision+1, resp.Header.Revision) } - success, err := metaKv.CompareVersionAndSwapBytes("a/b/c", 0, []byte("1")) + success, err := metaKv.CompareVersionAndSwapBytes(context.TODO(), "a/b/c", 0, []byte("1")) assert.NoError(t, err) assert.True(t, success) - value, err := metaKv.LoadBytes("a/b/c") + value, err := metaKv.LoadBytes(context.TODO(), "a/b/c") assert.NoError(t, err) assert.Equal(t, value, []byte("1")) - success, err = metaKv.CompareVersionAndSwapBytes("a/b/c", 0, []byte("1")) + success, err = metaKv.CompareVersionAndSwapBytes(context.TODO(), "a/b/c", 0, []byte("1")) assert.NoError(t, err) assert.False(t, success) }) @@ -708,7 +709,7 @@ func TestEmbedEtcd(te *testing.T) { assert.NoError(t, err) defer metaKv.Close() - defer metaKv.RemoveWithPrefix("") + defer metaKv.RemoveWithPrefix(context.TODO(), "") kvs := map[string]string{ "A/100": "v1", @@ -718,23 +719,23 @@ func TestEmbedEtcd(te *testing.T) { "B/100": "v5", } - err = metaKv.MultiSave(kvs) + err = metaKv.MultiSave(context.TODO(), kvs) assert.NoError(t, err) for k, v := range kvs { - actualV, err := metaKv.Load(k) + actualV, err := metaKv.Load(context.TODO(), k) assert.NoError(t, err) assert.Equal(t, v, actualV) } t.Run("apply function error ", func(t *testing.T) { - err = metaKv.WalkWithPrefix("A", 5, func(key []byte, value []byte) error { + err = metaKv.WalkWithPrefix(context.TODO(), "A", 5, func(key []byte, value []byte) error { return errors.New("error") }) assert.Error(t, err) }) t.Run("get with non-exist prefix ", func(t *testing.T) { - err = metaKv.WalkWithPrefix("non-exist-prefix", 5, func(key []byte, value []byte) error { + err = metaKv.WalkWithPrefix(context.TODO(), "non-exist-prefix", 5, func(key []byte, value []byte) error { return nil }) assert.NoError(t, err) @@ -755,7 +756,7 @@ func TestEmbedEtcd(te *testing.T) { ret := make(map[string]string) actualSortedKey := make([]string, 0) - err = metaKv.WalkWithPrefix("A", pagination, func(key []byte, value []byte) error { + err = metaKv.WalkWithPrefix(context.TODO(), "A", pagination, func(key []byte, value []byte) error { k := string(key) k = k[len(rootPath)+1:] ret[k] = string(value) @@ -783,23 +784,23 @@ func TestEmbedEtcd(te *testing.T) { assert.NoError(t, err) defer kv.Close() - defer kv.RemoveWithPrefix("") + defer kv.RemoveWithPrefix(context.TODO(), "") - has, err := kv.Has("key1") + has, err := kv.Has(context.TODO(), "key1") assert.NoError(t, err) assert.False(t, has) - err = kv.Save("key1", "value1") + err = kv.Save(context.TODO(), "key1", "value1") assert.NoError(t, err) - has, err = kv.Has("key1") + has, err = kv.Has(context.TODO(), "key1") assert.NoError(t, err) assert.True(t, has) - err = kv.Remove("key1") + err = kv.Remove(context.TODO(), "key1") assert.NoError(t, err) - has, err = kv.Has("key1") + has, err = kv.Has(context.TODO(), "key1") assert.NoError(t, err) assert.False(t, has) }) @@ -810,23 +811,23 @@ func TestEmbedEtcd(te *testing.T) { assert.NoError(t, err) defer kv.Close() - defer kv.RemoveWithPrefix("") + defer kv.RemoveWithPrefix(context.TODO(), "") - has, err := kv.HasPrefix("key") + has, err := kv.HasPrefix(context.TODO(), "key") assert.NoError(t, err) assert.False(t, has) - err = kv.Save("key1", "value1") + err = kv.Save(context.TODO(), "key1", "value1") assert.NoError(t, err) - has, err = kv.HasPrefix("key") + has, err = kv.HasPrefix(context.TODO(), "key") assert.NoError(t, err) assert.True(t, has) - err = kv.Remove("key1") + err = kv.Remove(context.TODO(), "key1") assert.NoError(t, err) - has, err = kv.HasPrefix("key") + has, err = kv.HasPrefix(context.TODO(), "key") assert.NoError(t, err) assert.False(t, has) }) @@ -865,7 +866,7 @@ func (s *EmbedEtcdKVSuite) SetupTest() { func (s *EmbedEtcdKVSuite) TearDownTest() { if s.kv != nil { - s.kv.RemoveWithPrefix("") + s.kv.RemoveWithPrefix(context.TODO(), "") s.kv.Close() s.kv = nil } @@ -879,7 +880,7 @@ func (s *EmbedEtcdKVSuite) TestTxnWithPredicates() { "lease2": "2", } - err := etcdKV.MultiSave(prepareKV) + err := etcdKV.MultiSave(context.TODO(), prepareKV) s.Require().NoError(err) badPredicate := predicates.NewMockPredicate(s.T()) @@ -899,13 +900,13 @@ func (s *EmbedEtcdKVSuite) TestTxnWithPredicates() { for _, test := range multiSaveAndRemovePredTests { s.Run(test.tag, func() { - err := etcdKV.MultiSaveAndRemove(test.multiSave, nil, test.preds...) + err := etcdKV.MultiSaveAndRemove(context.TODO(), test.multiSave, nil, test.preds...) if test.expectSuccess { s.NoError(err) } else { s.Error(err) } - err = etcdKV.MultiSaveAndRemoveWithPrefix(test.multiSave, nil, test.preds...) + err = etcdKV.MultiSaveAndRemoveWithPrefix(context.TODO(), test.multiSave, nil, test.preds...) if test.expectSuccess { s.NoError(err) } else { diff --git a/internal/kv/etcd/embed_etcd_restart_test.go b/internal/kv/etcd/embed_etcd_restart_test.go index 97fa8600a33d1..a93db26e7efe1 100644 --- a/internal/kv/etcd/embed_etcd_restart_test.go +++ b/internal/kv/etcd/embed_etcd_restart_test.go @@ -17,6 +17,7 @@ package etcdkv_test import ( + "context" "os" "testing" @@ -49,7 +50,7 @@ func TestEtcdRestartLoad(te *testing.T) { require.NoError(t, err) defer metaKv.Close() - defer metaKv.RemoveWithPrefix("") + defer metaKv.RemoveWithPrefix(context.TODO(), "") saveAndLoadTests := []struct { key string @@ -64,14 +65,14 @@ func TestEtcdRestartLoad(te *testing.T) { // save some data for i, test := range saveAndLoadTests { if i < 4 { - err = metaKv.Save(test.key, test.value) + err = metaKv.Save(context.TODO(), test.key, test.value) assert.NoError(t, err) } } // check test result for _, test := range saveAndLoadTests { - val, err := metaKv.Load(test.key) + val, err := metaKv.Load(context.TODO(), test.key) assert.NoError(t, err) assert.Equal(t, test.value, val) } @@ -83,7 +84,7 @@ func TestEtcdRestartLoad(te *testing.T) { metaKv, _ = embed_etcd_kv.NewMetaKvFactory(rootPath, ¶m.EtcdCfg) for _, test := range saveAndLoadTests { - val, err := metaKv.Load(test.key) + val, err := metaKv.Load(context.TODO(), test.key) assert.NoError(t, err) assert.Equal(t, test.value, val) } diff --git a/internal/kv/etcd/etcd_kv.go b/internal/kv/etcd/etcd_kv.go index 5003d0fae4b1f..f8fd41f5a7561 100644 --- a/internal/kv/etcd/etcd_kv.go +++ b/internal/kv/etcd/etcd_kv.go @@ -20,6 +20,7 @@ import ( "context" "encoding/binary" "fmt" + "github.com/milvus-io/milvus/pkg/kv" "path" "time" @@ -38,6 +39,9 @@ const ( defaultRequestTimeout = 10 * time.Second ) +// implementation assertion +var _ kv.WatchKV = (*etcdKV)(nil) + // etcdKV implements TxnKV interface, it supports to process multiple kvs in a transaction. type etcdKV struct { client *clientv3.Client @@ -71,7 +75,7 @@ func (kv *etcdKV) GetPath(key string) string { return path.Join(kv.rootPath, key) } -func (kv *etcdKV) WalkWithPrefix(prefix string, paginationSize int, fn func([]byte, []byte) error) error { +func (kv *etcdKV) WalkWithPrefix(ctx context.Context, prefix string, paginationSize int, fn func([]byte, []byte) error) error { start := time.Now() prefix = path.Join(kv.rootPath, prefix) @@ -84,7 +88,7 @@ func (kv *etcdKV) WalkWithPrefix(prefix string, paginationSize int, fn func([]by key := prefix for { - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() resp, err := kv.getEtcdMeta(ctx, key, opts...) if err != nil { @@ -109,10 +113,10 @@ func (kv *etcdKV) WalkWithPrefix(prefix string, paginationSize int, fn func([]by } // LoadWithPrefix returns all the keys and values with the given key prefix. -func (kv *etcdKV) LoadWithPrefix(key string) ([]string, []string, error) { +func (kv *etcdKV) LoadWithPrefix(ctx context.Context, key string) ([]string, []string, error) { start := time.Now() key = path.Join(kv.rootPath, key) - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() resp, err := kv.getEtcdMeta(ctx, key, clientv3.WithPrefix(), clientv3.WithSort(clientv3.SortByKey, clientv3.SortAscend)) @@ -129,10 +133,10 @@ func (kv *etcdKV) LoadWithPrefix(key string) ([]string, []string, error) { return keys, values, nil } -func (kv *etcdKV) Has(key string) (bool, error) { +func (kv *etcdKV) Has(ctx context.Context, key string) (bool, error) { start := time.Now() key = path.Join(kv.rootPath, key) - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() resp, err := kv.getEtcdMeta(ctx, key, clientv3.WithCountOnly()) @@ -144,10 +148,10 @@ func (kv *etcdKV) Has(key string) (bool, error) { return resp.Count != 0, nil } -func (kv *etcdKV) HasPrefix(prefix string) (bool, error) { +func (kv *etcdKV) HasPrefix(ctx context.Context, prefix string) (bool, error) { start := time.Now() prefix = path.Join(kv.rootPath, prefix) - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() resp, err := kv.getEtcdMeta(ctx, prefix, clientv3.WithPrefix(), clientv3.WithLimit(1), clientv3.WithCountOnly()) @@ -160,10 +164,10 @@ func (kv *etcdKV) HasPrefix(prefix string) (bool, error) { } // LoadBytesWithPrefix returns all the keys and values with the given key prefix. -func (kv *etcdKV) LoadBytesWithPrefix(key string) ([]string, [][]byte, error) { +func (kv *etcdKV) LoadBytesWithPrefix(ctx context.Context, key string) ([]string, [][]byte, error) { start := time.Now() key = path.Join(kv.rootPath, key) - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() resp, err := kv.getEtcdMeta(ctx, key, clientv3.WithPrefix(), clientv3.WithSort(clientv3.SortByKey, clientv3.SortAscend)) @@ -181,10 +185,10 @@ func (kv *etcdKV) LoadBytesWithPrefix(key string) ([]string, [][]byte, error) { } // LoadBytesWithPrefix2 returns all the keys,values and key versions with the given key prefix. -func (kv *etcdKV) LoadBytesWithPrefix2(key string) ([]string, [][]byte, []int64, error) { +func (kv *etcdKV) LoadBytesWithPrefix2(ctx context.Context, key string) ([]string, [][]byte, []int64, error) { start := time.Now() key = path.Join(kv.rootPath, key) - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() resp, err := kv.getEtcdMeta(ctx, key, clientv3.WithPrefix(), clientv3.WithSort(clientv3.SortByKey, clientv3.SortAscend)) @@ -204,10 +208,10 @@ func (kv *etcdKV) LoadBytesWithPrefix2(key string) ([]string, [][]byte, []int64, } // Load returns value of the key. -func (kv *etcdKV) Load(key string) (string, error) { +func (kv *etcdKV) Load(ctx context.Context, key string) (string, error) { start := time.Now() key = path.Join(kv.rootPath, key) - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() resp, err := kv.getEtcdMeta(ctx, key) if err != nil { @@ -221,10 +225,10 @@ func (kv *etcdKV) Load(key string) (string, error) { } // LoadBytes returns value of the key. -func (kv *etcdKV) LoadBytes(key string) ([]byte, error) { +func (kv *etcdKV) LoadBytes(ctx context.Context, key string) ([]byte, error) { start := time.Now() key = path.Join(kv.rootPath, key) - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() resp, err := kv.getEtcdMeta(ctx, key) if err != nil { @@ -238,14 +242,14 @@ func (kv *etcdKV) LoadBytes(key string) ([]byte, error) { } // MultiLoad gets the values of the keys in a transaction. -func (kv *etcdKV) MultiLoad(keys []string) ([]string, error) { +func (kv *etcdKV) MultiLoad(ctx context.Context, keys []string) ([]string, error) { start := time.Now() ops := make([]clientv3.Op, 0, len(keys)) for _, keyLoad := range keys { ops = append(ops, clientv3.OpGet(path.Join(kv.rootPath, keyLoad))) } - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() resp, err := kv.executeTxn(kv.getTxnWithCmp(ctx), ops...) if err != nil { @@ -273,14 +277,14 @@ func (kv *etcdKV) MultiLoad(keys []string) ([]string, error) { } // MultiLoadBytes gets the values of the keys in a transaction. -func (kv *etcdKV) MultiLoadBytes(keys []string) ([][]byte, error) { +func (kv *etcdKV) MultiLoadBytes(ctx context.Context, keys []string) ([][]byte, error) { start := time.Now() ops := make([]clientv3.Op, 0, len(keys)) for _, keyLoad := range keys { ops = append(ops, clientv3.OpGet(path.Join(kv.rootPath, keyLoad))) } - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() resp, err := kv.executeTxn(kv.getTxnWithCmp(ctx), ops...) if err != nil { @@ -308,10 +312,10 @@ func (kv *etcdKV) MultiLoadBytes(keys []string) ([][]byte, error) { } // LoadBytesWithRevision returns keys, values and revision with given key prefix. -func (kv *etcdKV) LoadBytesWithRevision(key string) ([]string, [][]byte, int64, error) { +func (kv *etcdKV) LoadBytesWithRevision(ctx context.Context, key string) ([]string, [][]byte, int64, error) { start := time.Now() key = path.Join(kv.rootPath, key) - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() resp, err := kv.getEtcdMeta(ctx, key, clientv3.WithPrefix(), clientv3.WithSort(clientv3.SortByKey, clientv3.SortAscend)) @@ -329,10 +333,10 @@ func (kv *etcdKV) LoadBytesWithRevision(key string) ([]string, [][]byte, int64, } // Save saves the key-value pair. -func (kv *etcdKV) Save(key, value string) error { +func (kv *etcdKV) Save(ctx context.Context, key, value string) error { start := time.Now() key = path.Join(kv.rootPath, key) - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() CheckValueSizeAndWarn(key, value) _, err := kv.putEtcdMeta(ctx, key, value) @@ -341,10 +345,10 @@ func (kv *etcdKV) Save(key, value string) error { } // SaveBytes saves the key-value pair. -func (kv *etcdKV) SaveBytes(key string, value []byte) error { +func (kv *etcdKV) SaveBytes(ctx context.Context, key string, value []byte) error { start := time.Now() key = path.Join(kv.rootPath, key) - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() CheckValueSizeAndWarn(key, value) _, err := kv.putEtcdMeta(ctx, key, string(value)) @@ -353,10 +357,10 @@ func (kv *etcdKV) SaveBytes(key string, value []byte) error { } // SaveBytesWithLease is a function to put value in etcd with etcd lease options. -func (kv *etcdKV) SaveBytesWithLease(key string, value []byte, id clientv3.LeaseID) error { +func (kv *etcdKV) SaveBytesWithLease(ctx context.Context, key string, value []byte, id clientv3.LeaseID) error { start := time.Now() key = path.Join(kv.rootPath, key) - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() CheckValueSizeAndWarn(key, value) _, err := kv.putEtcdMeta(ctx, key, string(value), clientv3.WithLease(id)) @@ -365,7 +369,7 @@ func (kv *etcdKV) SaveBytesWithLease(key string, value []byte, id clientv3.Lease } // MultiSave saves the key-value pairs in a transaction. -func (kv *etcdKV) MultiSave(kvs map[string]string) error { +func (kv *etcdKV) MultiSave(ctx context.Context, kvs map[string]string) error { start := time.Now() ops := make([]clientv3.Op, 0, len(kvs)) var keys []string @@ -374,7 +378,7 @@ func (kv *etcdKV) MultiSave(kvs map[string]string) error { ops = append(ops, clientv3.OpPut(path.Join(kv.rootPath, key), value)) } - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() CheckTnxStringValueSizeAndWarn(kvs) @@ -387,7 +391,7 @@ func (kv *etcdKV) MultiSave(kvs map[string]string) error { } // MultiSaveBytes saves the key-value pairs in a transaction. -func (kv *etcdKV) MultiSaveBytes(kvs map[string][]byte) error { +func (kv *etcdKV) MultiSaveBytes(ctx context.Context, kvs map[string][]byte) error { start := time.Now() ops := make([]clientv3.Op, 0, len(kvs)) var keys []string @@ -396,7 +400,7 @@ func (kv *etcdKV) MultiSaveBytes(kvs map[string][]byte) error { ops = append(ops, clientv3.OpPut(path.Join(kv.rootPath, key), string(value))) } - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() CheckTnxBytesValueSizeAndWarn(kvs) @@ -409,10 +413,10 @@ func (kv *etcdKV) MultiSaveBytes(kvs map[string][]byte) error { } // RemoveWithPrefix removes the keys with given prefix. -func (kv *etcdKV) RemoveWithPrefix(prefix string) error { +func (kv *etcdKV) RemoveWithPrefix(ctx context.Context, prefix string) error { start := time.Now() key := path.Join(kv.rootPath, prefix) - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() _, err := kv.removeEtcdMeta(ctx, key, clientv3.WithPrefix()) @@ -421,10 +425,10 @@ func (kv *etcdKV) RemoveWithPrefix(prefix string) error { } // Remove removes the key. -func (kv *etcdKV) Remove(key string) error { +func (kv *etcdKV) Remove(ctx context.Context, key string) error { start := time.Now() key = path.Join(kv.rootPath, key) - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() _, err := kv.removeEtcdMeta(ctx, key) @@ -433,14 +437,14 @@ func (kv *etcdKV) Remove(key string) error { } // MultiRemove removes the keys in a transaction. -func (kv *etcdKV) MultiRemove(keys []string) error { +func (kv *etcdKV) MultiRemove(ctx context.Context, keys []string) error { start := time.Now() ops := make([]clientv3.Op, 0, len(keys)) for _, key := range keys { ops = append(ops, clientv3.OpDelete(path.Join(kv.rootPath, key))) } - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() _, err := kv.executeTxn(kv.getTxnWithCmp(ctx), ops...) @@ -452,7 +456,7 @@ func (kv *etcdKV) MultiRemove(keys []string) error { } // MultiSaveAndRemove saves the key-value pairs and removes the keys in a transaction. -func (kv *etcdKV) MultiSaveAndRemove(saves map[string]string, removals []string, preds ...predicates.Predicate) error { +func (kv *etcdKV) MultiSaveAndRemove(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error { cmps, err := parsePredicates(kv.rootPath, preds...) if err != nil { return err @@ -470,7 +474,7 @@ func (kv *etcdKV) MultiSaveAndRemove(saves map[string]string, removals []string, ops = append(ops, clientv3.OpDelete(path.Join(kv.rootPath, keyDelete))) } - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() resp, err := kv.executeTxn(kv.getTxnWithCmp(ctx, cmps...), ops...) @@ -492,7 +496,7 @@ func (kv *etcdKV) MultiSaveAndRemove(saves map[string]string, removals []string, } // MultiSaveBytesAndRemove saves the key-value pairs and removes the keys in a transaction. -func (kv *etcdKV) MultiSaveBytesAndRemove(saves map[string][]byte, removals []string) error { +func (kv *etcdKV) MultiSaveBytesAndRemove(ctx context.Context, saves map[string][]byte, removals []string) error { start := time.Now() ops := make([]clientv3.Op, 0, len(saves)+len(removals)) var keys []string @@ -505,7 +509,7 @@ func (kv *etcdKV) MultiSaveBytesAndRemove(saves map[string][]byte, removals []st ops = append(ops, clientv3.OpDelete(path.Join(kv.rootPath, keyDelete))) } - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() _, err := kv.executeTxn(kv.getTxnWithCmp(ctx), ops...) @@ -522,7 +526,7 @@ func (kv *etcdKV) MultiSaveBytesAndRemove(saves map[string][]byte, removals []st } // Watch starts watching a key, returns a watch channel. -func (kv *etcdKV) Watch(key string) clientv3.WatchChan { +func (kv *etcdKV) Watch(ctx context.Context, key string) clientv3.WatchChan { start := time.Now() key = path.Join(kv.rootPath, key) rch := kv.client.Watch(context.Background(), key, clientv3.WithCreatedNotify()) @@ -531,7 +535,7 @@ func (kv *etcdKV) Watch(key string) clientv3.WatchChan { } // WatchWithPrefix starts watching a key with prefix, returns a watch channel. -func (kv *etcdKV) WatchWithPrefix(key string) clientv3.WatchChan { +func (kv *etcdKV) WatchWithPrefix(ctx context.Context, key string) clientv3.WatchChan { start := time.Now() key = path.Join(kv.rootPath, key) rch := kv.client.Watch(context.Background(), key, clientv3.WithPrefix(), clientv3.WithCreatedNotify()) @@ -540,7 +544,7 @@ func (kv *etcdKV) WatchWithPrefix(key string) clientv3.WatchChan { } // WatchWithRevision starts watching a key with revision, returns a watch channel. -func (kv *etcdKV) WatchWithRevision(key string, revision int64) clientv3.WatchChan { +func (kv *etcdKV) WatchWithRevision(ctx context.Context, key string, revision int64) clientv3.WatchChan { start := time.Now() key = path.Join(kv.rootPath, key) rch := kv.client.Watch(context.Background(), key, clientv3.WithPrefix(), clientv3.WithPrevKV(), clientv3.WithRev(revision)) @@ -549,7 +553,7 @@ func (kv *etcdKV) WatchWithRevision(key string, revision int64) clientv3.WatchCh } // MultiSaveAndRemoveWithPrefix saves kv in @saves and removes the keys with given prefix in @removals. -func (kv *etcdKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals []string, preds ...predicates.Predicate) error { +func (kv *etcdKV) MultiSaveAndRemoveWithPrefix(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error { cmps, err := parsePredicates(kv.rootPath, preds...) if err != nil { return err @@ -567,7 +571,7 @@ func (kv *etcdKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals ops = append(ops, clientv3.OpDelete(path.Join(kv.rootPath, keyDelete), clientv3.WithPrefix())) } - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() resp, err := kv.executeTxn(kv.getTxnWithCmp(ctx, cmps...), ops...) @@ -588,7 +592,7 @@ func (kv *etcdKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals } // MultiSaveBytesAndRemoveWithPrefix saves kv in @saves and removes the keys with given prefix in @removals. -func (kv *etcdKV) MultiSaveBytesAndRemoveWithPrefix(saves map[string][]byte, removals []string) error { +func (kv *etcdKV) MultiSaveBytesAndRemoveWithPrefix(ctx context.Context, saves map[string][]byte, removals []string) error { start := time.Now() ops := make([]clientv3.Op, 0, len(saves)) var keys []string @@ -601,7 +605,7 @@ func (kv *etcdKV) MultiSaveBytesAndRemoveWithPrefix(saves map[string][]byte, rem ops = append(ops, clientv3.OpDelete(path.Join(kv.rootPath, keyDelete), clientv3.WithPrefix())) } - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() _, err := kv.executeTxn(kv.getTxnWithCmp(ctx), ops...) @@ -619,9 +623,9 @@ func (kv *etcdKV) MultiSaveBytesAndRemoveWithPrefix(saves map[string][]byte, rem // CompareVersionAndSwap compares the existing key-value's version with version, and if // they are equal, the target is stored in etcd. -func (kv *etcdKV) CompareVersionAndSwap(key string, source int64, target string) (bool, error) { +func (kv *etcdKV) CompareVersionAndSwap(ctx context.Context, key string, source int64, target string) (bool, error) { start := time.Now() - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() resp, err := kv.executeTxn(kv.getTxnWithCmp(ctx, clientv3.Compare(clientv3.Version(path.Join(kv.rootPath, key)), "=", source)), @@ -635,9 +639,9 @@ func (kv *etcdKV) CompareVersionAndSwap(key string, source int64, target string) // CompareVersionAndSwapBytes compares the existing key-value's version with version, and if // they are equal, the target is stored in etcd. -func (kv *etcdKV) CompareVersionAndSwapBytes(key string, source int64, target []byte, opts ...clientv3.OpOption) (bool, error) { +func (kv *etcdKV) CompareVersionAndSwapBytes(ctx context.Context, key string, source int64, target []byte, opts ...clientv3.OpOption) (bool, error) { start := time.Now() - ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() resp, err := kv.executeTxn(kv.getTxnWithCmp(ctx, clientv3.Compare(clientv3.Version(path.Join(kv.rootPath, key)), "=", source)), diff --git a/internal/kv/etcd/etcd_kv_test.go b/internal/kv/etcd/etcd_kv_test.go index 93e42d95f347a..123a0a209ce13 100644 --- a/internal/kv/etcd/etcd_kv_test.go +++ b/internal/kv/etcd/etcd_kv_test.go @@ -17,6 +17,7 @@ package etcdkv import ( + "context" "fmt" "os" "path" @@ -79,7 +80,7 @@ func (s *EtcdKVSuite) SetupTest() { } func (s *EtcdKVSuite) TearDownTest() { - s.etcdKV.RemoveWithPrefix("") + s.etcdKV.RemoveWithPrefix(context.TODO(), "") s.etcdKV.Close() } @@ -97,11 +98,11 @@ func (s *EtcdKVSuite) TestSaveLoad() { for i, test := range saveAndLoadTests { if i < 4 { - err := etcdKV.Save(test.key, test.value) + err := etcdKV.Save(context.TODO(), test.key, test.value) s.Require().NoError(err) } - val, err := etcdKV.Load(test.key) + val, err := etcdKV.Load(context.TODO(), test.key) s.Require().NoError(err) s.Equal(test.value, val) } @@ -115,7 +116,7 @@ func (s *EtcdKVSuite) TestSaveLoad() { } for _, test := range invalidLoadTests { - val, err := etcdKV.Load(test.invalidKey) + val, err := etcdKV.Load(context.TODO(), test.invalidKey) s.Error(err) s.Zero(val) } @@ -152,7 +153,7 @@ func (s *EtcdKVSuite) TestSaveLoad() { } for _, test := range loadPrefixTests { - actualKeys, actualValues, err := etcdKV.LoadWithPrefix(test.prefix) + actualKeys, actualValues, err := etcdKV.LoadWithPrefix(context.TODO(), test.prefix) s.ElementsMatch(test.expectedKeys, actualKeys) s.ElementsMatch(test.expectedValues, actualValues) s.Equal(test.expectedError, err) @@ -169,15 +170,15 @@ func (s *EtcdKVSuite) TestSaveLoad() { } for _, test := range removeTests { - err := etcdKV.Remove(test.validKey) + err := etcdKV.Remove(context.TODO(), test.validKey) s.NoError(err) - _, err = etcdKV.Load(test.validKey) + _, err = etcdKV.Load(context.TODO(), test.validKey) s.Error(err) - err = etcdKV.Remove(test.validKey) + err = etcdKV.Remove(context.TODO(), test.validKey) s.NoError(err) - err = etcdKV.Remove(test.invalidKey) + err = etcdKV.Remove(context.TODO(), test.invalidKey) s.NoError(err) } } @@ -197,11 +198,11 @@ func (s *EtcdKVSuite) TestSaveAndLoadBytes() { for i, test := range saveAndLoadTests { if i < 4 { - err := etcdKV.SaveBytes(test.key, []byte(test.value)) + err := etcdKV.SaveBytes(context.TODO(), test.key, []byte(test.value)) s.Require().NoError(err) } - val, err := etcdKV.LoadBytes(test.key) + val, err := etcdKV.LoadBytes(context.TODO(), test.key) s.NoError(err) s.Equal(test.value, string(val)) } @@ -215,7 +216,7 @@ func (s *EtcdKVSuite) TestSaveAndLoadBytes() { } for _, test := range invalidLoadTests { - val, err := etcdKV.LoadBytes(test.invalidKey) + val, err := etcdKV.LoadBytes(context.TODO(), test.invalidKey) s.Error(err) s.Zero(string(val)) } @@ -252,7 +253,7 @@ func (s *EtcdKVSuite) TestSaveAndLoadBytes() { } for _, test := range loadPrefixTests { - actualKeys, actualValues, err := etcdKV.LoadBytesWithPrefix(test.prefix) + actualKeys, actualValues, err := etcdKV.LoadBytesWithPrefix(context.TODO(), test.prefix) actualStringValues := make([]string, len(actualValues)) for i := range actualValues { actualStringValues[i] = string(actualValues[i]) @@ -261,7 +262,7 @@ func (s *EtcdKVSuite) TestSaveAndLoadBytes() { s.ElementsMatch(test.expectedValues, actualStringValues) s.Equal(test.expectedError, err) - actualKeys, actualValues, versions, err := etcdKV.LoadBytesWithPrefix2(test.prefix) + actualKeys, actualValues, versions, err := etcdKV.LoadBytesWithPrefix2(context.TODO(), test.prefix) actualStringValues = make([]string, len(actualValues)) for i := range actualValues { actualStringValues[i] = string(actualValues[i]) @@ -283,15 +284,15 @@ func (s *EtcdKVSuite) TestSaveAndLoadBytes() { } for _, test := range removeTests { - err := etcdKV.Remove(test.validKey) + err := etcdKV.Remove(context.TODO(), test.validKey) s.NoError(err) - _, err = etcdKV.Load(test.validKey) + _, err = etcdKV.Load(context.TODO(), test.validKey) s.Error(err) - err = etcdKV.Remove(test.validKey) + err = etcdKV.Remove(context.TODO(), test.validKey) s.NoError(err) - err = etcdKV.Remove(test.invalidKey) + err = etcdKV.Remove(context.TODO(), test.invalidKey) s.NoError(err) } } @@ -311,7 +312,7 @@ func (s *EtcdKVSuite) TestLoadBytesWithRevision() { } for _, test := range prepareKV { - err := etcdKV.SaveBytes(test.inKey, []byte(test.inValue)) + err := etcdKV.SaveBytes(context.TODO(), test.inKey, []byte(test.inValue)) s.NoError(err) } @@ -327,7 +328,7 @@ func (s *EtcdKVSuite) TestLoadBytesWithRevision() { } for _, test := range loadWithRevisionTests { - keys, values, revision, err := etcdKV.LoadBytesWithRevision(test.inKey) + keys, values, revision, err := etcdKV.LoadBytesWithRevision(context.TODO(), test.inKey) s.NoError(err) s.Equal(test.expectedKeyNo, len(keys)) stringValues := make([]string, len(values)) @@ -350,10 +351,10 @@ func (s *EtcdKVSuite) TestMultiSaveAndMultiLoad() { "_": "other", } - err := etcdKV.MultiSave(multiSaveTests) + err := etcdKV.MultiSave(context.TODO(), multiSaveTests) s.Require().NoError(err) for k, v := range multiSaveTests { - actualV, err := etcdKV.Load(k) + actualV, err := etcdKV.Load(context.TODO(), k) s.NoError(err) s.Equal(v, actualV) } @@ -369,7 +370,7 @@ func (s *EtcdKVSuite) TestMultiSaveAndMultiLoad() { } for _, test := range multiLoadTests { - vs, err := etcdKV.MultiLoad(test.inputKeys) + vs, err := etcdKV.MultiLoad(context.TODO(), test.inputKeys) s.NoError(err) s.Equal(test.expectedValues, vs) } @@ -385,7 +386,7 @@ func (s *EtcdKVSuite) TestMultiSaveAndMultiLoad() { } for _, test := range invalidMultiLoad { - vs, err := etcdKV.MultiLoad(test.invalidKeys) + vs, err := etcdKV.MultiLoad(context.TODO(), test.invalidKeys) s.Error(err) s.Equal(test.expectedValues, vs) } @@ -396,10 +397,10 @@ func (s *EtcdKVSuite) TestMultiSaveAndMultiLoad() { } for _, k := range removeWithPrefixTests { - err = etcdKV.RemoveWithPrefix(k) + err = etcdKV.RemoveWithPrefix(context.TODO(), k) s.NoError(err) - ks, vs, err := etcdKV.LoadWithPrefix(k) + ks, vs, err := etcdKV.LoadWithPrefix(context.TODO(), k) s.Empty(ks) s.Empty(vs) s.NoError(err) @@ -412,10 +413,10 @@ func (s *EtcdKVSuite) TestMultiSaveAndMultiLoad() { "_", } - err = etcdKV.MultiRemove(multiRemoveTests) + err = etcdKV.MultiRemove(context.TODO(), multiRemoveTests) s.NoError(err) - ks, vs, err := etcdKV.LoadWithPrefix("") + ks, vs, err := etcdKV.LoadWithPrefix(context.TODO(), "") s.NoError(err) s.Empty(ks) s.Empty(vs) @@ -432,11 +433,11 @@ func (s *EtcdKVSuite) TestMultiSaveAndMultiLoad() { {make(map[string]string), []string{"multikey_2"}}, } for _, test := range multiSaveAndRemoveTests { - err = etcdKV.MultiSaveAndRemove(test.multiSaves, test.multiRemoves) + err = etcdKV.MultiSaveAndRemove(context.TODO(), test.multiSaves, test.multiRemoves) s.NoError(err) } - ks, vs, err = etcdKV.LoadWithPrefix("") + ks, vs, err = etcdKV.LoadWithPrefix(context.TODO(), "") s.NoError(err) s.Empty(ks) s.Empty(vs) @@ -458,10 +459,10 @@ func (s *EtcdKVSuite) TestMultiSaveBytesAndMultiLoadBytes() { multiSaveBytesTests[k] = []byte(v) } - err := etcdKV.MultiSaveBytes(multiSaveBytesTests) + err := etcdKV.MultiSaveBytes(context.TODO(), multiSaveBytesTests) s.Require().NoError(err) for k, v := range multiSaveTests { - actualV, err := etcdKV.LoadBytes(k) + actualV, err := etcdKV.LoadBytes(context.TODO(), k) s.NoError(err) s.Equal(v, string(actualV)) } @@ -477,7 +478,7 @@ func (s *EtcdKVSuite) TestMultiSaveBytesAndMultiLoadBytes() { } for _, test := range multiLoadTests { - vs, err := etcdKV.MultiLoadBytes(test.inputKeys) + vs, err := etcdKV.MultiLoadBytes(context.TODO(), test.inputKeys) stringVs := make([]string, len(vs)) for i := range vs { stringVs[i] = string(vs[i]) @@ -497,7 +498,7 @@ func (s *EtcdKVSuite) TestMultiSaveBytesAndMultiLoadBytes() { } for _, test := range invalidMultiLoad { - vs, err := etcdKV.MultiLoadBytes(test.invalidKeys) + vs, err := etcdKV.MultiLoadBytes(context.TODO(), test.invalidKeys) stringVs := make([]string, len(vs)) for i := range vs { stringVs[i] = string(vs[i]) @@ -512,10 +513,10 @@ func (s *EtcdKVSuite) TestMultiSaveBytesAndMultiLoadBytes() { } for _, k := range removeWithPrefixTests { - err = etcdKV.RemoveWithPrefix(k) + err = etcdKV.RemoveWithPrefix(context.TODO(), k) s.NoError(err) - ks, vs, err := etcdKV.LoadBytesWithPrefix(k) + ks, vs, err := etcdKV.LoadBytesWithPrefix(context.TODO(), k) s.Empty(ks) s.Empty(vs) s.NoError(err) @@ -528,10 +529,10 @@ func (s *EtcdKVSuite) TestMultiSaveBytesAndMultiLoadBytes() { "_", } - err = etcdKV.MultiRemove(multiRemoveTests) + err = etcdKV.MultiRemove(context.TODO(), multiRemoveTests) s.NoError(err) - ks, vs, err := etcdKV.LoadBytesWithPrefix("") + ks, vs, err := etcdKV.LoadBytesWithPrefix(context.TODO(), "") s.NoError(err) s.Empty(ks) s.Empty(vs) @@ -549,11 +550,11 @@ func (s *EtcdKVSuite) TestMultiSaveBytesAndMultiLoadBytes() { } for _, test := range multiSaveAndRemoveTests { - err = etcdKV.MultiSaveBytesAndRemove(test.multiSaves, test.multiRemoves) + err = etcdKV.MultiSaveBytesAndRemove(context.TODO(), test.multiSaves, test.multiRemoves) s.NoError(err) } - ks, vs, err = etcdKV.LoadBytesWithPrefix("") + ks, vs, err = etcdKV.LoadBytesWithPrefix(context.TODO(), "") s.NoError(err) s.Empty(ks) s.Empty(vs) @@ -567,7 +568,7 @@ func (s *EtcdKVSuite) TestTxnWithPredicates() { "lease2": "2", } - err := etcdKV.MultiSave(prepareKV) + err := etcdKV.MultiSave(context.TODO(), prepareKV) s.Require().NoError(err) badPredicate := predicates.NewMockPredicate(s.T()) @@ -587,13 +588,13 @@ func (s *EtcdKVSuite) TestTxnWithPredicates() { for _, test := range multiSaveAndRemovePredTests { s.Run(test.tag, func() { - err := etcdKV.MultiSaveAndRemove(test.multiSave, nil, test.preds...) + err := etcdKV.MultiSaveAndRemove(context.TODO(), test.multiSave, nil, test.preds...) if test.expectSuccess { s.NoError(err) } else { s.Error(err) } - err = etcdKV.MultiSaveAndRemoveWithPrefix(test.multiSave, nil, test.preds...) + err = etcdKV.MultiSaveAndRemoveWithPrefix(context.TODO(), test.multiSave, nil, test.preds...) if test.expectSuccess { s.NoError(err) } else { @@ -616,7 +617,7 @@ func (s *EtcdKVSuite) TestMultiSaveAndRemoveWithPrefix() { } // MultiSaveAndRemoveWithPrefix - err := etcdKV.MultiSave(prepareTests) + err := etcdKV.MultiSave(context.TODO(), prepareTests) s.Require().NoError(err) multiSaveAndRemoveWithPrefixTests := []struct { multiSave map[string]string @@ -634,14 +635,14 @@ func (s *EtcdKVSuite) TestMultiSaveAndRemoveWithPrefix() { } for _, test := range multiSaveAndRemoveWithPrefixTests { - k, _, err := etcdKV.LoadWithPrefix(test.loadPrefix) + k, _, err := etcdKV.LoadWithPrefix(context.TODO(), test.loadPrefix) s.NoError(err) s.Equal(test.lengthBeforeRemove, len(k)) - err = etcdKV.MultiSaveAndRemoveWithPrefix(test.multiSave, test.prefix) + err = etcdKV.MultiSaveAndRemoveWithPrefix(context.TODO(), test.multiSave, test.prefix) s.NoError(err) - k, _, err = etcdKV.LoadWithPrefix(test.loadPrefix) + k, _, err = etcdKV.LoadWithPrefix(context.TODO(), test.loadPrefix) s.NoError(err) s.Equal(test.lengthAfterRemove, len(k)) } @@ -650,11 +651,11 @@ func (s *EtcdKVSuite) TestMultiSaveAndRemoveWithPrefix() { func (s *EtcdKVSuite) TestWatch() { etcdKV := s.etcdKV - ch := etcdKV.Watch("x") + ch := etcdKV.Watch(context.TODO(), "x") resp := <-ch s.True(resp.Created) - ch = etcdKV.WatchWithPrefix("x") + ch = etcdKV.WatchWithPrefix(context.TODO(), "x") resp = <-ch s.True(resp.Created) } @@ -673,13 +674,13 @@ func (s *EtcdKVSuite) TestRevisionBytes() { } for _, test := range revisionTests { - err := etcdKV.SaveBytes(test.inKey, test.fistValue) + err := etcdKV.SaveBytes(context.TODO(), test.inKey, test.fistValue) s.Require().NoError(err) - _, _, revision, _ := etcdKV.LoadBytesWithRevision(test.inKey) - ch := etcdKV.WatchWithRevision(test.inKey, revision+1) + _, _, revision, _ := etcdKV.LoadBytesWithRevision(context.TODO(), test.inKey) + ch := etcdKV.WatchWithRevision(context.TODO(), test.inKey, revision+1) - err = etcdKV.SaveBytes(test.inKey, test.secondValue) + err = etcdKV.SaveBytes(context.TODO(), test.inKey, test.secondValue) s.Require().NoError(err) resp := <-ch @@ -688,15 +689,15 @@ func (s *EtcdKVSuite) TestRevisionBytes() { s.Equal(revision+1, resp.Header.Revision) } - success, err := etcdKV.CompareVersionAndSwapBytes("a/b/c", 0, []byte("1")) + success, err := etcdKV.CompareVersionAndSwapBytes(context.TODO(), "a/b/c", 0, []byte("1")) s.NoError(err) s.True(success) - value, err := etcdKV.LoadBytes("a/b/c") + value, err := etcdKV.LoadBytes(context.TODO(), "a/b/c") s.NoError(err) s.Equal(string(value), "1") - success, err = etcdKV.CompareVersionAndSwapBytes("a/b/c", 0, []byte("1")) + success, err = etcdKV.CompareVersionAndSwapBytes(context.TODO(), "a/b/c", 0, []byte("1")) s.NoError(err) s.False(success) } @@ -721,7 +722,7 @@ func Test_WalkWithPagination(t *testing.T) { etcdKV := NewEtcdKV(etcdCli, rootPath) defer etcdKV.Close() - defer etcdKV.RemoveWithPrefix("") + defer etcdKV.RemoveWithPrefix(context.TODO(), "") kvs := map[string]string{ "A/100": "v1", @@ -731,23 +732,23 @@ func Test_WalkWithPagination(t *testing.T) { "B/100": "v5", } - err = etcdKV.MultiSave(kvs) + err = etcdKV.MultiSave(context.TODO(), kvs) assert.NoError(t, err) for k, v := range kvs { - actualV, err := etcdKV.Load(k) + actualV, err := etcdKV.Load(context.TODO(), k) assert.NoError(t, err) assert.Equal(t, v, actualV) } t.Run("apply function error ", func(t *testing.T) { - err = etcdKV.WalkWithPrefix("A", 5, func(key []byte, value []byte) error { + err = etcdKV.WalkWithPrefix(context.TODO(), "A", 5, func(key []byte, value []byte) error { return errors.New("error") }) assert.Error(t, err) }) t.Run("get with non-exist prefix ", func(t *testing.T) { - err = etcdKV.WalkWithPrefix("non-exist-prefix", 5, func(key []byte, value []byte) error { + err = etcdKV.WalkWithPrefix(context.TODO(), "non-exist-prefix", 5, func(key []byte, value []byte) error { return nil }) assert.NoError(t, err) @@ -768,7 +769,7 @@ func Test_WalkWithPagination(t *testing.T) { ret := make(map[string]string) actualSortedKey := make([]string, 0) - err = etcdKV.WalkWithPrefix("A", pagination, func(key []byte, value []byte) error { + err = etcdKV.WalkWithPrefix(context.TODO(), "A", pagination, func(key []byte, value []byte) error { k := string(key) k = k[len(rootPath)+1:] ret[k] = string(value) @@ -844,27 +845,27 @@ func TestHas(t *testing.T) { assert.NoError(t, err) rootPath := "/etcd/test/root/has" kv := NewEtcdKV(etcdCli, rootPath) - err = kv.RemoveWithPrefix("") + err = kv.RemoveWithPrefix(context.TODO(), "") require.NoError(t, err) defer kv.Close() - defer kv.RemoveWithPrefix("") + defer kv.RemoveWithPrefix(context.TODO(), "") - has, err := kv.Has("key1") + has, err := kv.Has(context.TODO(), "key1") assert.NoError(t, err) assert.False(t, has) - err = kv.Save("key1", "value1") + err = kv.Save(context.TODO(), "key1", "value1") assert.NoError(t, err) - has, err = kv.Has("key1") + has, err = kv.Has(context.TODO(), "key1") assert.NoError(t, err) assert.True(t, has) - err = kv.Remove("key1") + err = kv.Remove(context.TODO(), "key1") assert.NoError(t, err) - has, err = kv.Has("key1") + has, err = kv.Has(context.TODO(), "key1") assert.NoError(t, err) assert.False(t, has) } @@ -882,27 +883,27 @@ func TestHasPrefix(t *testing.T) { assert.NoError(t, err) rootPath := "/etcd/test/root/hasprefix" kv := NewEtcdKV(etcdCli, rootPath) - err = kv.RemoveWithPrefix("") + err = kv.RemoveWithPrefix(context.TODO(), "") require.NoError(t, err) defer kv.Close() - defer kv.RemoveWithPrefix("") + defer kv.RemoveWithPrefix(context.TODO(), "") - has, err := kv.HasPrefix("key") + has, err := kv.HasPrefix(context.TODO(), "key") assert.NoError(t, err) assert.False(t, has) - err = kv.Save("key1", "value1") + err = kv.Save(context.TODO(), "key1", "value1") assert.NoError(t, err) - has, err = kv.HasPrefix("key") + has, err = kv.HasPrefix(context.TODO(), "key") assert.NoError(t, err) assert.True(t, has) - err = kv.Remove("key1") + err = kv.Remove(context.TODO(), "key1") assert.NoError(t, err) - has, err = kv.HasPrefix("key") + has, err = kv.HasPrefix(context.TODO(), "key") assert.NoError(t, err) assert.False(t, has) } diff --git a/internal/kv/mem/mem_kv.go b/internal/kv/mem/mem_kv.go index 0b95569ca74d7..7b8381c252f91 100644 --- a/internal/kv/mem/mem_kv.go +++ b/internal/kv/mem/mem_kv.go @@ -17,15 +17,20 @@ package memkv import ( + "context" "strings" "sync" "github.com/google/btree" + "github.com/milvus-io/milvus/pkg/kv" "github.com/milvus-io/milvus/pkg/kv/predicates" "github.com/milvus-io/milvus/pkg/util/merr" ) +// implementation assertion +var _ kv.TxnKV = (*MemoryKV)(nil) + // MemoryKV implements BaseKv interface and relies on underling btree.BTree. // As its name implies, all data is stored in memory. type MemoryKV struct { @@ -78,7 +83,7 @@ func (s memoryKVItem) Less(than btree.Item) bool { } // Load loads an object with @key. -func (kv *MemoryKV) Load(key string) (string, error) { +func (kv *MemoryKV) Load(ctx context.Context, key string) (string, error) { kv.RLock() defer kv.RUnlock() item := kv.tree.Get(memoryKVItem{key: key}) @@ -89,7 +94,7 @@ func (kv *MemoryKV) Load(key string) (string, error) { } // LoadBytes loads an object with @key. -func (kv *MemoryKV) LoadBytes(key string) ([]byte, error) { +func (kv *MemoryKV) LoadBytes(ctx context.Context, key string) ([]byte, error) { kv.RLock() defer kv.RUnlock() item := kv.tree.Get(memoryKVItem{key: key}) @@ -100,7 +105,7 @@ func (kv *MemoryKV) LoadBytes(key string) ([]byte, error) { } // Get return value if key exists, or return empty string -func (kv *MemoryKV) Get(key string) string { +func (kv *MemoryKV) Get(ctx context.Context, key string) string { kv.RLock() defer kv.RUnlock() item := kv.tree.Get(memoryKVItem{key: key}) @@ -111,7 +116,7 @@ func (kv *MemoryKV) Get(key string) string { } // LoadBytesWithDefault loads an object with @key. If the object does not exist, @defaultValue will be returned. -func (kv *MemoryKV) LoadBytesWithDefault(key string, defaultValue []byte) []byte { +func (kv *MemoryKV) LoadBytesWithDefault(ctx context.Context, key string, defaultValue []byte) []byte { kv.RLock() defer kv.RUnlock() item := kv.tree.Get(memoryKVItem{key: key}) @@ -122,7 +127,7 @@ func (kv *MemoryKV) LoadBytesWithDefault(key string, defaultValue []byte) []byte } // LoadBytesRange loads objects with range @startKey to @endKey with @limit number of objects. -func (kv *MemoryKV) LoadBytesRange(key, endKey string, limit int) ([]string, [][]byte, error) { +func (kv *MemoryKV) LoadBytesRange(ctx context.Context, key, endKey string, limit int) ([]string, [][]byte, error) { kv.RLock() defer kv.RUnlock() keys := make([]string, 0, limit) @@ -139,7 +144,7 @@ func (kv *MemoryKV) LoadBytesRange(key, endKey string, limit int) ([]string, [][ } // Save object with @key to btree. Object value is @value. -func (kv *MemoryKV) Save(key, value string) error { +func (kv *MemoryKV) Save(ctx context.Context, key, value string) error { kv.Lock() defer kv.Unlock() kv.tree.ReplaceOrInsert(memoryKVItem{key, StringValue(value)}) @@ -147,7 +152,7 @@ func (kv *MemoryKV) Save(key, value string) error { } // SaveBytes object with @key to btree. Object value is @value. -func (kv *MemoryKV) SaveBytes(key string, value []byte) error { +func (kv *MemoryKV) SaveBytes(ctx context.Context, key string, value []byte) error { kv.Lock() defer kv.Unlock() kv.tree.ReplaceOrInsert(memoryKVItem{key, ByteSliceValue(value)}) @@ -155,7 +160,7 @@ func (kv *MemoryKV) SaveBytes(key string, value []byte) error { } // Remove deletes an object with @key. -func (kv *MemoryKV) Remove(key string) error { +func (kv *MemoryKV) Remove(ctx context.Context, key string) error { kv.Lock() defer kv.Unlock() @@ -164,7 +169,7 @@ func (kv *MemoryKV) Remove(key string) error { } // MultiLoad loads objects with multi @keys. -func (kv *MemoryKV) MultiLoad(keys []string) ([]string, error) { +func (kv *MemoryKV) MultiLoad(ctx context.Context, keys []string) ([]string, error) { kv.RLock() defer kv.RUnlock() result := make([]string, 0, len(keys)) @@ -176,7 +181,7 @@ func (kv *MemoryKV) MultiLoad(keys []string) ([]string, error) { } // MultiLoadBytes loads objects with multi @keys. -func (kv *MemoryKV) MultiLoadBytes(keys []string) ([][]byte, error) { +func (kv *MemoryKV) MultiLoadBytes(ctx context.Context, keys []string) ([][]byte, error) { kv.RLock() defer kv.RUnlock() result := make([][]byte, 0, len(keys)) @@ -188,7 +193,7 @@ func (kv *MemoryKV) MultiLoadBytes(keys []string) ([][]byte, error) { } // MultiSave saves given key-value pairs in MemoryKV atomicly. -func (kv *MemoryKV) MultiSave(kvs map[string]string) error { +func (kv *MemoryKV) MultiSave(ctx context.Context, kvs map[string]string) error { kv.Lock() defer kv.Unlock() for key, value := range kvs { @@ -198,7 +203,7 @@ func (kv *MemoryKV) MultiSave(kvs map[string]string) error { } // MultiSaveBytes saves given key-value pairs in MemoryKV atomicly. -func (kv *MemoryKV) MultiSaveBytes(kvs map[string][]byte) error { +func (kv *MemoryKV) MultiSaveBytes(ctx context.Context, kvs map[string][]byte) error { kv.Lock() defer kv.Unlock() for key, value := range kvs { @@ -208,7 +213,7 @@ func (kv *MemoryKV) MultiSaveBytes(kvs map[string][]byte) error { } // MultiRemove removes given @keys in MemoryKV atomicly. -func (kv *MemoryKV) MultiRemove(keys []string) error { +func (kv *MemoryKV) MultiRemove(ctx context.Context, keys []string) error { kv.Lock() defer kv.Unlock() for _, key := range keys { @@ -218,7 +223,7 @@ func (kv *MemoryKV) MultiRemove(keys []string) error { } // MultiSaveAndRemove saves and removes given key-value pairs in MemoryKV atomicly. -func (kv *MemoryKV) MultiSaveAndRemove(saves map[string]string, removals []string, preds ...predicates.Predicate) error { +func (kv *MemoryKV) MultiSaveAndRemove(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error { if len(preds) > 0 { return merr.WrapErrServiceUnavailable("predicates not supported") } @@ -234,7 +239,7 @@ func (kv *MemoryKV) MultiSaveAndRemove(saves map[string]string, removals []strin } // MultiSaveBytesAndRemove saves and removes given key-value pairs in MemoryKV atomicly. -func (kv *MemoryKV) MultiSaveBytesAndRemove(saves map[string][]byte, removals []string) error { +func (kv *MemoryKV) MultiSaveBytesAndRemove(ctx context.Context, saves map[string][]byte, removals []string) error { kv.Lock() defer kv.Unlock() for key, value := range saves { @@ -247,7 +252,7 @@ func (kv *MemoryKV) MultiSaveBytesAndRemove(saves map[string][]byte, removals [] } // LoadWithPrefix returns all keys & values with given prefix. -func (kv *MemoryKV) LoadWithPrefix(key string) ([]string, []string, error) { +func (kv *MemoryKV) LoadWithPrefix(ctx context.Context, key string) ([]string, []string, error) { kv.Lock() defer kv.Unlock() @@ -265,7 +270,7 @@ func (kv *MemoryKV) LoadWithPrefix(key string) ([]string, []string, error) { } // LoadBytesWithPrefix returns all keys & values with given prefix. -func (kv *MemoryKV) LoadBytesWithPrefix(key string) ([]string, [][]byte, error) { +func (kv *MemoryKV) LoadBytesWithPrefix(ctx context.Context, key string) ([]string, [][]byte, error) { kv.Lock() defer kv.Unlock() @@ -287,7 +292,7 @@ func (kv *MemoryKV) Close() { } // MultiSaveAndRemoveWithPrefix saves key-value pairs in @saves, & remove key with prefix in @removals in MemoryKV atomically. -func (kv *MemoryKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals []string, preds ...predicates.Predicate) error { +func (kv *MemoryKV) MultiSaveAndRemoveWithPrefix(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error { if len(preds) > 0 { return merr.WrapErrServiceUnavailable("predicates not supported") } @@ -314,7 +319,7 @@ func (kv *MemoryKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, remova } // MultiSaveBytesAndRemoveWithPrefix saves key-value pairs in @saves, & remove key with prefix in @removals in MemoryKV atomically. -func (kv *MemoryKV) MultiSaveBytesAndRemoveWithPrefix(saves map[string][]byte, removals []string) error { +func (kv *MemoryKV) MultiSaveBytesAndRemoveWithPrefix(ctx context.Context, saves map[string][]byte, removals []string) error { kv.Lock() defer kv.Unlock() @@ -338,7 +343,7 @@ func (kv *MemoryKV) MultiSaveBytesAndRemoveWithPrefix(saves map[string][]byte, r } // RemoveWithPrefix remove key of given prefix in MemoryKV atomicly. -func (kv *MemoryKV) RemoveWithPrefix(key string) error { +func (kv *MemoryKV) RemoveWithPrefix(ctx context.Context, key string) error { kv.Lock() defer kv.Unlock() @@ -356,13 +361,13 @@ func (kv *MemoryKV) RemoveWithPrefix(key string) error { return nil } -func (kv *MemoryKV) Has(key string) (bool, error) { +func (kv *MemoryKV) Has(ctx context.Context, key string) (bool, error) { kv.Lock() defer kv.Unlock() return kv.tree.Has(memoryKVItem{key: key}), nil } -func (kv *MemoryKV) HasPrefix(prefix string) (bool, error) { +func (kv *MemoryKV) HasPrefix(ctx context.Context, prefix string) (bool, error) { kv.Lock() defer kv.Unlock() diff --git a/internal/kv/mem/mem_kv_test.go b/internal/kv/mem/mem_kv_test.go index 7fbb66bc13e28..4d81294c5e5ae 100644 --- a/internal/kv/mem/mem_kv_test.go +++ b/internal/kv/mem/mem_kv_test.go @@ -17,6 +17,7 @@ package memkv import ( + "context" "testing" "github.com/stretchr/testify/assert" @@ -30,15 +31,15 @@ func TestMemoryKV_SaveAndLoadBytes(t *testing.T) { key := "key" value := []byte("value") - err := mem.SaveBytes(key, value) + err := mem.SaveBytes(context.TODO(), key, value) assert.NoError(t, err) - _value, err := mem.LoadBytes(key) + _value, err := mem.LoadBytes(context.TODO(), key) assert.NoError(t, err) assert.Equal(t, value, _value) noKey := "no_key" - _value, err = mem.LoadBytes(noKey) + _value, err = mem.LoadBytes(context.TODO(), noKey) assert.Error(t, err) assert.Empty(t, _value) } @@ -56,16 +57,16 @@ func TestMemoryKV_LoadBytesRange(t *testing.T) { mem := NewMemoryKV() for _, kv := range saveAndLoadBytesTests { - err := mem.SaveBytes(kv.key, kv.value) + err := mem.SaveBytes(context.TODO(), kv.key, kv.value) assert.NoError(t, err) } - keys, values, err := mem.LoadBytesRange("test1", "test1/b", 0) + keys, values, err := mem.LoadBytesRange(context.TODO(), "test1", "test1/b", 0) assert.Equal(t, len(keys), 2) assert.Equal(t, len(values), 2) assert.NoError(t, err) - keys, values, err = mem.LoadBytesRange("test1", "test1/a", 2) + keys, values, err = mem.LoadBytesRange(context.TODO(), "test1", "test1/a", 2) assert.Equal(t, len(keys), 1) assert.Equal(t, len(values), 1) assert.NoError(t, err) @@ -76,15 +77,15 @@ func TestMemoryKV_LoadBytesWithDefault(t *testing.T) { key := "key" value := []byte("value") - err := mem.SaveBytes(key, value) + err := mem.SaveBytes(context.TODO(), key, value) assert.NoError(t, err) _default := []byte("default") - _value := mem.LoadBytesWithDefault(key, _default) + _value := mem.LoadBytesWithDefault(context.TODO(), key, _default) assert.Equal(t, value, _value) noKey := "no_key" - _value = mem.LoadBytesWithDefault(noKey, _default) + _value = mem.LoadBytesWithDefault(context.TODO(), noKey, _default) assert.Equal(t, _value, _default) } @@ -101,21 +102,21 @@ func TestMemoryKV_LoadBytesWithPrefix(t *testing.T) { mem := NewMemoryKV() for _, kv := range saveAndLoadBytesTests { - err := mem.SaveBytes(kv.key, kv.value) + err := mem.SaveBytes(context.TODO(), kv.key, kv.value) assert.NoError(t, err) } - keys, values, err := mem.LoadBytesWithPrefix("test1") + keys, values, err := mem.LoadBytesWithPrefix(context.TODO(), "test1") assert.Equal(t, len(keys), 3) assert.Equal(t, len(values), 3) assert.NoError(t, err) - keys, values, err = mem.LoadBytesWithPrefix("test") + keys, values, err = mem.LoadBytesWithPrefix(context.TODO(), "test") assert.Equal(t, len(keys), 4) assert.Equal(t, len(values), 4) assert.NoError(t, err) - keys, values, err = mem.LoadBytesWithPrefix("a") + keys, values, err = mem.LoadBytesWithPrefix(context.TODO(), "a") assert.Equal(t, len(keys), 0) assert.Equal(t, len(values), 0) assert.NoError(t, err) @@ -137,14 +138,14 @@ func TestMemoryKV_MultiSaveBytes(t *testing.T) { } mem := NewMemoryKV() - err := mem.MultiSaveBytes(saveAndLoadBytesTests) + err := mem.MultiSaveBytes(context.TODO(), saveAndLoadBytesTests) assert.NoError(t, err) - _values, err := mem.MultiLoadBytes(keys) + _values, err := mem.MultiLoadBytes(context.TODO(), keys) assert.Equal(t, values, _values) assert.NoError(t, err) - _values, err = mem.MultiLoadBytes([]string{}) + _values, err = mem.MultiLoadBytes(context.TODO(), []string{}) assert.Empty(t, _values) assert.NoError(t, err) } @@ -165,14 +166,14 @@ func TestMemoryKV_MultiSaveBytesAndRemove(t *testing.T) { } mem := NewMemoryKV() - err := mem.MultiSaveBytesAndRemove(saveAndLoadBytesTests, []string{keys[0]}) + err := mem.MultiSaveBytesAndRemove(context.TODO(), saveAndLoadBytesTests, []string{keys[0]}) assert.NoError(t, err) - _value, err := mem.LoadBytes(keys[0]) + _value, err := mem.LoadBytes(context.TODO(), keys[0]) assert.Empty(t, _value) assert.Error(t, err) - _values, err := mem.MultiLoadBytes(keys[1:]) + _values, err := mem.MultiLoadBytes(context.TODO(), keys[1:]) assert.Equal(t, values[1:], _values) assert.NoError(t, err) } @@ -193,10 +194,10 @@ func TestMemoryKV_MultiSaveBytesAndRemoveWithPrefix(t *testing.T) { } mem := NewMemoryKV() - err := mem.MultiSaveBytesAndRemoveWithPrefix(saveAndLoadBytesTests, []string{"test1"}) + err := mem.MultiSaveBytesAndRemoveWithPrefix(context.TODO(), saveAndLoadBytesTests, []string{"test1"}) assert.NoError(t, err) - _keys, _values, err := mem.LoadBytesWithPrefix("test") + _keys, _values, err := mem.LoadBytesWithPrefix(context.TODO(), "test") assert.ElementsMatch(t, keys, _keys) assert.ElementsMatch(t, values, _values) assert.NoError(t, err) @@ -205,21 +206,21 @@ func TestMemoryKV_MultiSaveBytesAndRemoveWithPrefix(t *testing.T) { func TestHas(t *testing.T) { kv := NewMemoryKV() - has, err := kv.Has("key1") + has, err := kv.Has(context.TODO(), "key1") assert.NoError(t, err) assert.False(t, has) - err = kv.Save("key1", "value1") + err = kv.Save(context.TODO(), "key1", "value1") assert.NoError(t, err) - has, err = kv.Has("key1") + has, err = kv.Has(context.TODO(), "key1") assert.NoError(t, err) assert.True(t, has) - err = kv.Remove("key1") + err = kv.Remove(context.TODO(), "key1") assert.NoError(t, err) - has, err = kv.Has("key1") + has, err = kv.Has(context.TODO(), "key1") assert.NoError(t, err) assert.False(t, has) } @@ -227,21 +228,21 @@ func TestHas(t *testing.T) { func TestHasPrefix(t *testing.T) { kv := NewMemoryKV() - has, err := kv.HasPrefix("key") + has, err := kv.HasPrefix(context.TODO(), "key") assert.NoError(t, err) assert.False(t, has) - err = kv.Save("key1", "value1") + err = kv.Save(context.TODO(), "key1", "value1") assert.NoError(t, err) - has, err = kv.HasPrefix("key") + has, err = kv.HasPrefix(context.TODO(), "key") assert.NoError(t, err) assert.True(t, has) - err = kv.Remove("key1") + err = kv.Remove(context.TODO(), "key1") assert.NoError(t, err) - has, err = kv.HasPrefix("key") + has, err = kv.HasPrefix(context.TODO(), "key") assert.NoError(t, err) assert.False(t, has) } @@ -250,11 +251,11 @@ func TestPredicates(t *testing.T) { kv := NewMemoryKV() // predicates not supported for mem kv for now - err := kv.MultiSaveAndRemove(map[string]string{}, []string{}, predicates.ValueEqual("a", "b")) + err := kv.MultiSaveAndRemove(context.TODO(), map[string]string{}, []string{}, predicates.ValueEqual("a", "b")) assert.Error(t, err) assert.ErrorIs(t, err, merr.ErrServiceUnavailable) - err = kv.MultiSaveAndRemoveWithPrefix(map[string]string{}, []string{}, predicates.ValueEqual("a", "b")) + err = kv.MultiSaveAndRemoveWithPrefix(context.TODO(), map[string]string{}, []string{}, predicates.ValueEqual("a", "b")) assert.Error(t, err) assert.ErrorIs(t, err, merr.ErrServiceUnavailable) } diff --git a/internal/kv/mock_snapshot_kv.go b/internal/kv/mock_snapshot_kv.go index 4c5ac49a72083..fda86ef107b87 100644 --- a/internal/kv/mock_snapshot_kv.go +++ b/internal/kv/mock_snapshot_kv.go @@ -1,62 +1,65 @@ package kv import ( + "context" "github.com/milvus-io/milvus/pkg/kv" "github.com/milvus-io/milvus/pkg/util/typeutil" ) +var _ kv.SnapShotKV = &mockSnapshotKV{} + type mockSnapshotKV struct { kv.SnapShotKV - SaveFunc func(key string, value string, ts typeutil.Timestamp) error - LoadFunc func(key string, ts typeutil.Timestamp) (string, error) - MultiSaveFunc func(kvs map[string]string, ts typeutil.Timestamp) error - LoadWithPrefixFunc func(key string, ts typeutil.Timestamp) ([]string, []string, error) - MultiSaveAndRemoveWithPrefixFunc func(saves map[string]string, removals []string, ts typeutil.Timestamp) error - MultiSaveAndRemoveFunc func(saves map[string]string, removals []string, ts typeutil.Timestamp) error + SaveFunc func(ctx context.Context, key string, value string, ts typeutil.Timestamp) error + LoadFunc func(ctx context.Context, key string, ts typeutil.Timestamp) (string, error) + MultiSaveFunc func(ctx context.Context, kvs map[string]string, ts typeutil.Timestamp) error + LoadWithPrefixFunc func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) + MultiSaveAndRemoveWithPrefixFunc func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error + MultiSaveAndRemoveFunc func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error } func NewMockSnapshotKV() *mockSnapshotKV { return &mockSnapshotKV{} } -func (m mockSnapshotKV) Save(key string, value string, ts typeutil.Timestamp) error { +func (m mockSnapshotKV) Save(ctx context.Context, key string, value string, ts typeutil.Timestamp) error { if m.SaveFunc != nil { - return m.SaveFunc(key, value, ts) + return m.SaveFunc(ctx, key, value, ts) } return nil } -func (m mockSnapshotKV) Load(key string, ts typeutil.Timestamp) (string, error) { +func (m mockSnapshotKV) Load(ctx context.Context, key string, ts typeutil.Timestamp) (string, error) { if m.LoadFunc != nil { - return m.LoadFunc(key, ts) + return m.LoadFunc(ctx, key, ts) } return "", nil } -func (m mockSnapshotKV) MultiSave(kvs map[string]string, ts typeutil.Timestamp) error { +func (m mockSnapshotKV) MultiSave(ctx context.Context, kvs map[string]string, ts typeutil.Timestamp) error { if m.MultiSaveFunc != nil { - return m.MultiSaveFunc(kvs, ts) + return m.MultiSaveFunc(ctx, kvs, ts) } return nil } -func (m mockSnapshotKV) LoadWithPrefix(key string, ts typeutil.Timestamp) ([]string, []string, error) { +func (m mockSnapshotKV) LoadWithPrefix(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) { if m.LoadWithPrefixFunc != nil { - return m.LoadWithPrefixFunc(key, ts) + return m.LoadWithPrefixFunc(ctx, key, ts) } return nil, nil, nil } -func (m mockSnapshotKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals []string, ts typeutil.Timestamp) error { +func (m mockSnapshotKV) MultiSaveAndRemoveWithPrefix(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error { if m.MultiSaveAndRemoveWithPrefixFunc != nil { - return m.MultiSaveAndRemoveWithPrefixFunc(saves, removals, ts) + return m.MultiSaveAndRemoveWithPrefixFunc(ctx, saves, removals, ts) } return nil } -func (m mockSnapshotKV) MultiSaveAndRemove(saves map[string]string, removals []string, ts typeutil.Timestamp) error { +func (m mockSnapshotKV) MultiSaveAndRemove(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error { if m.MultiSaveAndRemoveFunc != nil { - return m.MultiSaveAndRemoveFunc(saves, removals, ts) + return m.MultiSaveAndRemoveFunc(ctx, saves, removals, ts) } return nil } diff --git a/internal/kv/mock_snapshot_kv_test.go b/internal/kv/mock_snapshot_kv_test.go index 0b2df70f9173f..19fabf54b0540 100644 --- a/internal/kv/mock_snapshot_kv_test.go +++ b/internal/kv/mock_snapshot_kv_test.go @@ -1,6 +1,7 @@ package kv import ( + "context" "testing" "github.com/stretchr/testify/assert" @@ -11,15 +12,15 @@ import ( func Test_mockSnapshotKV_Save(t *testing.T) { t.Run("func not set", func(t *testing.T) { snapshot := NewMockSnapshotKV() - err := snapshot.Save("k", "v", 0) + err := snapshot.Save(context.TODO(), "k", "v", 0) assert.NoError(t, err) }) t.Run("func set", func(t *testing.T) { snapshot := NewMockSnapshotKV() - snapshot.SaveFunc = func(key string, value string, ts typeutil.Timestamp) error { + snapshot.SaveFunc = func(ctx context.Context, key string, value string, ts typeutil.Timestamp) error { return nil } - err := snapshot.Save("k", "v", 0) + err := snapshot.Save(context.TODO(), "k", "v", 0) assert.NoError(t, err) }) } @@ -27,15 +28,15 @@ func Test_mockSnapshotKV_Save(t *testing.T) { func Test_mockSnapshotKV_Load(t *testing.T) { t.Run("func not set", func(t *testing.T) { snapshot := NewMockSnapshotKV() - _, err := snapshot.Load("k", 0) + _, err := snapshot.Load(context.TODO(), "k", 0) assert.NoError(t, err) }) t.Run("func set", func(t *testing.T) { snapshot := NewMockSnapshotKV() - snapshot.LoadFunc = func(key string, ts typeutil.Timestamp) (string, error) { + snapshot.LoadFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) (string, error) { return "", nil } - _, err := snapshot.Load("k", 0) + _, err := snapshot.Load(context.TODO(), "k", 0) assert.NoError(t, err) }) } @@ -43,15 +44,15 @@ func Test_mockSnapshotKV_Load(t *testing.T) { func Test_mockSnapshotKV_MultiSave(t *testing.T) { t.Run("func not set", func(t *testing.T) { snapshot := NewMockSnapshotKV() - err := snapshot.MultiSave(nil, 0) + err := snapshot.MultiSave(context.TODO(), nil, 0) assert.NoError(t, err) }) t.Run("func set", func(t *testing.T) { snapshot := NewMockSnapshotKV() - snapshot.MultiSaveFunc = func(kvs map[string]string, ts typeutil.Timestamp) error { + snapshot.MultiSaveFunc = func(ctx context.Context, kvs map[string]string, ts typeutil.Timestamp) error { return nil } - err := snapshot.MultiSave(nil, 0) + err := snapshot.MultiSave(context.TODO(), nil, 0) assert.NoError(t, err) }) } @@ -59,15 +60,15 @@ func Test_mockSnapshotKV_MultiSave(t *testing.T) { func Test_mockSnapshotKV_LoadWithPrefix(t *testing.T) { t.Run("func not set", func(t *testing.T) { snapshot := NewMockSnapshotKV() - _, _, err := snapshot.LoadWithPrefix("prefix", 0) + _, _, err := snapshot.LoadWithPrefix(context.TODO(), "prefix", 0) assert.NoError(t, err) }) t.Run("func set", func(t *testing.T) { snapshot := NewMockSnapshotKV() - snapshot.LoadWithPrefixFunc = func(key string, ts typeutil.Timestamp) ([]string, []string, error) { + snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) { return nil, nil, nil } - _, _, err := snapshot.LoadWithPrefix("prefix", 0) + _, _, err := snapshot.LoadWithPrefix(context.TODO(), "prefix", 0) assert.NoError(t, err) }) } @@ -75,15 +76,15 @@ func Test_mockSnapshotKV_LoadWithPrefix(t *testing.T) { func Test_mockSnapshotKV_MultiSaveAndRemoveWithPrefix(t *testing.T) { t.Run("func not set", func(t *testing.T) { snapshot := NewMockSnapshotKV() - err := snapshot.MultiSaveAndRemoveWithPrefix(nil, nil, 0) + err := snapshot.MultiSaveAndRemoveWithPrefix(context.TODO(), nil, nil, 0) assert.NoError(t, err) }) t.Run("func set", func(t *testing.T) { snapshot := NewMockSnapshotKV() - snapshot.MultiSaveAndRemoveWithPrefixFunc = func(saves map[string]string, removals []string, ts typeutil.Timestamp) error { + snapshot.MultiSaveAndRemoveWithPrefixFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error { return nil } - err := snapshot.MultiSaveAndRemoveWithPrefix(nil, nil, 0) + err := snapshot.MultiSaveAndRemoveWithPrefix(context.TODO(), nil, nil, 0) assert.NoError(t, err) }) } @@ -91,15 +92,15 @@ func Test_mockSnapshotKV_MultiSaveAndRemoveWithPrefix(t *testing.T) { func Test_mockSnapshotKV_MultiSaveAndRemove(t *testing.T) { t.Run("func not set", func(t *testing.T) { snapshot := NewMockSnapshotKV() - err := snapshot.MultiSaveAndRemove(nil, nil, 0) + err := snapshot.MultiSaveAndRemove(context.TODO(), nil, nil, 0) assert.NoError(t, err) }) t.Run("func set", func(t *testing.T) { snapshot := NewMockSnapshotKV() - snapshot.MultiSaveAndRemoveWithPrefixFunc = func(saves map[string]string, removals []string, ts typeutil.Timestamp) error { + snapshot.MultiSaveAndRemoveWithPrefixFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error { return nil } - err := snapshot.MultiSaveAndRemove(nil, nil, 0) + err := snapshot.MultiSaveAndRemove(context.TODO(), nil, nil, 0) assert.NoError(t, err) }) } diff --git a/internal/kv/mocks/meta_kv.go b/internal/kv/mocks/meta_kv.go index 85094ae478440..82813bf3cb6df 100644 --- a/internal/kv/mocks/meta_kv.go +++ b/internal/kv/mocks/meta_kv.go @@ -3,8 +3,11 @@ package mocks import ( - predicates "github.com/milvus-io/milvus/pkg/kv/predicates" + context "context" + mock "github.com/stretchr/testify/mock" + + predicates "github.com/milvus-io/milvus/pkg/kv/predicates" ) // MetaKv is an autogenerated mock type for the MetaKv type @@ -52,9 +55,9 @@ func (_c *MetaKv_Close_Call) RunAndReturn(run func()) *MetaKv_Close_Call { return _c } -// CompareVersionAndSwap provides a mock function with given fields: key, version, target -func (_m *MetaKv) CompareVersionAndSwap(key string, version int64, target string) (bool, error) { - ret := _m.Called(key, version, target) +// CompareVersionAndSwap provides a mock function with given fields: ctx, key, version, target +func (_m *MetaKv) CompareVersionAndSwap(ctx context.Context, key string, version int64, target string) (bool, error) { + ret := _m.Called(ctx, key, version, target) if len(ret) == 0 { panic("no return value specified for CompareVersionAndSwap") @@ -62,17 +65,17 @@ func (_m *MetaKv) CompareVersionAndSwap(key string, version int64, target string var r0 bool var r1 error - if rf, ok := ret.Get(0).(func(string, int64, string) (bool, error)); ok { - return rf(key, version, target) + if rf, ok := ret.Get(0).(func(context.Context, string, int64, string) (bool, error)); ok { + return rf(ctx, key, version, target) } - if rf, ok := ret.Get(0).(func(string, int64, string) bool); ok { - r0 = rf(key, version, target) + if rf, ok := ret.Get(0).(func(context.Context, string, int64, string) bool); ok { + r0 = rf(ctx, key, version, target) } else { r0 = ret.Get(0).(bool) } - if rf, ok := ret.Get(1).(func(string, int64, string) error); ok { - r1 = rf(key, version, target) + if rf, ok := ret.Get(1).(func(context.Context, string, int64, string) error); ok { + r1 = rf(ctx, key, version, target) } else { r1 = ret.Error(1) } @@ -86,16 +89,17 @@ type MetaKv_CompareVersionAndSwap_Call struct { } // CompareVersionAndSwap is a helper method to define mock.On call +// - ctx context.Context // - key string // - version int64 // - target string -func (_e *MetaKv_Expecter) CompareVersionAndSwap(key interface{}, version interface{}, target interface{}) *MetaKv_CompareVersionAndSwap_Call { - return &MetaKv_CompareVersionAndSwap_Call{Call: _e.mock.On("CompareVersionAndSwap", key, version, target)} +func (_e *MetaKv_Expecter) CompareVersionAndSwap(ctx interface{}, key interface{}, version interface{}, target interface{}) *MetaKv_CompareVersionAndSwap_Call { + return &MetaKv_CompareVersionAndSwap_Call{Call: _e.mock.On("CompareVersionAndSwap", ctx, key, version, target)} } -func (_c *MetaKv_CompareVersionAndSwap_Call) Run(run func(key string, version int64, target string)) *MetaKv_CompareVersionAndSwap_Call { +func (_c *MetaKv_CompareVersionAndSwap_Call) Run(run func(ctx context.Context, key string, version int64, target string)) *MetaKv_CompareVersionAndSwap_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string), args[1].(int64), args[2].(string)) + run(args[0].(context.Context), args[1].(string), args[2].(int64), args[3].(string)) }) return _c } @@ -105,7 +109,7 @@ func (_c *MetaKv_CompareVersionAndSwap_Call) Return(_a0 bool, _a1 error) *MetaKv return _c } -func (_c *MetaKv_CompareVersionAndSwap_Call) RunAndReturn(run func(string, int64, string) (bool, error)) *MetaKv_CompareVersionAndSwap_Call { +func (_c *MetaKv_CompareVersionAndSwap_Call) RunAndReturn(run func(context.Context, string, int64, string) (bool, error)) *MetaKv_CompareVersionAndSwap_Call { _c.Call.Return(run) return _c } @@ -156,9 +160,9 @@ func (_c *MetaKv_GetPath_Call) RunAndReturn(run func(string) string) *MetaKv_Get return _c } -// Has provides a mock function with given fields: key -func (_m *MetaKv) Has(key string) (bool, error) { - ret := _m.Called(key) +// Has provides a mock function with given fields: ctx, key +func (_m *MetaKv) Has(ctx context.Context, key string) (bool, error) { + ret := _m.Called(ctx, key) if len(ret) == 0 { panic("no return value specified for Has") @@ -166,17 +170,17 @@ func (_m *MetaKv) Has(key string) (bool, error) { var r0 bool var r1 error - if rf, ok := ret.Get(0).(func(string) (bool, error)); ok { - return rf(key) + if rf, ok := ret.Get(0).(func(context.Context, string) (bool, error)); ok { + return rf(ctx, key) } - if rf, ok := ret.Get(0).(func(string) bool); ok { - r0 = rf(key) + if rf, ok := ret.Get(0).(func(context.Context, string) bool); ok { + r0 = rf(ctx, key) } else { r0 = ret.Get(0).(bool) } - if rf, ok := ret.Get(1).(func(string) error); ok { - r1 = rf(key) + if rf, ok := ret.Get(1).(func(context.Context, string) error); ok { + r1 = rf(ctx, key) } else { r1 = ret.Error(1) } @@ -190,14 +194,15 @@ type MetaKv_Has_Call struct { } // Has is a helper method to define mock.On call +// - ctx context.Context // - key string -func (_e *MetaKv_Expecter) Has(key interface{}) *MetaKv_Has_Call { - return &MetaKv_Has_Call{Call: _e.mock.On("Has", key)} +func (_e *MetaKv_Expecter) Has(ctx interface{}, key interface{}) *MetaKv_Has_Call { + return &MetaKv_Has_Call{Call: _e.mock.On("Has", ctx, key)} } -func (_c *MetaKv_Has_Call) Run(run func(key string)) *MetaKv_Has_Call { +func (_c *MetaKv_Has_Call) Run(run func(ctx context.Context, key string)) *MetaKv_Has_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string)) + run(args[0].(context.Context), args[1].(string)) }) return _c } @@ -207,14 +212,14 @@ func (_c *MetaKv_Has_Call) Return(_a0 bool, _a1 error) *MetaKv_Has_Call { return _c } -func (_c *MetaKv_Has_Call) RunAndReturn(run func(string) (bool, error)) *MetaKv_Has_Call { +func (_c *MetaKv_Has_Call) RunAndReturn(run func(context.Context, string) (bool, error)) *MetaKv_Has_Call { _c.Call.Return(run) return _c } -// HasPrefix provides a mock function with given fields: prefix -func (_m *MetaKv) HasPrefix(prefix string) (bool, error) { - ret := _m.Called(prefix) +// HasPrefix provides a mock function with given fields: ctx, prefix +func (_m *MetaKv) HasPrefix(ctx context.Context, prefix string) (bool, error) { + ret := _m.Called(ctx, prefix) if len(ret) == 0 { panic("no return value specified for HasPrefix") @@ -222,17 +227,17 @@ func (_m *MetaKv) HasPrefix(prefix string) (bool, error) { var r0 bool var r1 error - if rf, ok := ret.Get(0).(func(string) (bool, error)); ok { - return rf(prefix) + if rf, ok := ret.Get(0).(func(context.Context, string) (bool, error)); ok { + return rf(ctx, prefix) } - if rf, ok := ret.Get(0).(func(string) bool); ok { - r0 = rf(prefix) + if rf, ok := ret.Get(0).(func(context.Context, string) bool); ok { + r0 = rf(ctx, prefix) } else { r0 = ret.Get(0).(bool) } - if rf, ok := ret.Get(1).(func(string) error); ok { - r1 = rf(prefix) + if rf, ok := ret.Get(1).(func(context.Context, string) error); ok { + r1 = rf(ctx, prefix) } else { r1 = ret.Error(1) } @@ -246,14 +251,15 @@ type MetaKv_HasPrefix_Call struct { } // HasPrefix is a helper method to define mock.On call +// - ctx context.Context // - prefix string -func (_e *MetaKv_Expecter) HasPrefix(prefix interface{}) *MetaKv_HasPrefix_Call { - return &MetaKv_HasPrefix_Call{Call: _e.mock.On("HasPrefix", prefix)} +func (_e *MetaKv_Expecter) HasPrefix(ctx interface{}, prefix interface{}) *MetaKv_HasPrefix_Call { + return &MetaKv_HasPrefix_Call{Call: _e.mock.On("HasPrefix", ctx, prefix)} } -func (_c *MetaKv_HasPrefix_Call) Run(run func(prefix string)) *MetaKv_HasPrefix_Call { +func (_c *MetaKv_HasPrefix_Call) Run(run func(ctx context.Context, prefix string)) *MetaKv_HasPrefix_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string)) + run(args[0].(context.Context), args[1].(string)) }) return _c } @@ -263,14 +269,14 @@ func (_c *MetaKv_HasPrefix_Call) Return(_a0 bool, _a1 error) *MetaKv_HasPrefix_C return _c } -func (_c *MetaKv_HasPrefix_Call) RunAndReturn(run func(string) (bool, error)) *MetaKv_HasPrefix_Call { +func (_c *MetaKv_HasPrefix_Call) RunAndReturn(run func(context.Context, string) (bool, error)) *MetaKv_HasPrefix_Call { _c.Call.Return(run) return _c } -// Load provides a mock function with given fields: key -func (_m *MetaKv) Load(key string) (string, error) { - ret := _m.Called(key) +// Load provides a mock function with given fields: ctx, key +func (_m *MetaKv) Load(ctx context.Context, key string) (string, error) { + ret := _m.Called(ctx, key) if len(ret) == 0 { panic("no return value specified for Load") @@ -278,17 +284,17 @@ func (_m *MetaKv) Load(key string) (string, error) { var r0 string var r1 error - if rf, ok := ret.Get(0).(func(string) (string, error)); ok { - return rf(key) + if rf, ok := ret.Get(0).(func(context.Context, string) (string, error)); ok { + return rf(ctx, key) } - if rf, ok := ret.Get(0).(func(string) string); ok { - r0 = rf(key) + if rf, ok := ret.Get(0).(func(context.Context, string) string); ok { + r0 = rf(ctx, key) } else { r0 = ret.Get(0).(string) } - if rf, ok := ret.Get(1).(func(string) error); ok { - r1 = rf(key) + if rf, ok := ret.Get(1).(func(context.Context, string) error); ok { + r1 = rf(ctx, key) } else { r1 = ret.Error(1) } @@ -302,14 +308,15 @@ type MetaKv_Load_Call struct { } // Load is a helper method to define mock.On call +// - ctx context.Context // - key string -func (_e *MetaKv_Expecter) Load(key interface{}) *MetaKv_Load_Call { - return &MetaKv_Load_Call{Call: _e.mock.On("Load", key)} +func (_e *MetaKv_Expecter) Load(ctx interface{}, key interface{}) *MetaKv_Load_Call { + return &MetaKv_Load_Call{Call: _e.mock.On("Load", ctx, key)} } -func (_c *MetaKv_Load_Call) Run(run func(key string)) *MetaKv_Load_Call { +func (_c *MetaKv_Load_Call) Run(run func(ctx context.Context, key string)) *MetaKv_Load_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string)) + run(args[0].(context.Context), args[1].(string)) }) return _c } @@ -319,14 +326,14 @@ func (_c *MetaKv_Load_Call) Return(_a0 string, _a1 error) *MetaKv_Load_Call { return _c } -func (_c *MetaKv_Load_Call) RunAndReturn(run func(string) (string, error)) *MetaKv_Load_Call { +func (_c *MetaKv_Load_Call) RunAndReturn(run func(context.Context, string) (string, error)) *MetaKv_Load_Call { _c.Call.Return(run) return _c } -// LoadWithPrefix provides a mock function with given fields: key -func (_m *MetaKv) LoadWithPrefix(key string) ([]string, []string, error) { - ret := _m.Called(key) +// LoadWithPrefix provides a mock function with given fields: ctx, key +func (_m *MetaKv) LoadWithPrefix(ctx context.Context, key string) ([]string, []string, error) { + ret := _m.Called(ctx, key) if len(ret) == 0 { panic("no return value specified for LoadWithPrefix") @@ -335,27 +342,27 @@ func (_m *MetaKv) LoadWithPrefix(key string) ([]string, []string, error) { var r0 []string var r1 []string var r2 error - if rf, ok := ret.Get(0).(func(string) ([]string, []string, error)); ok { - return rf(key) + if rf, ok := ret.Get(0).(func(context.Context, string) ([]string, []string, error)); ok { + return rf(ctx, key) } - if rf, ok := ret.Get(0).(func(string) []string); ok { - r0 = rf(key) + if rf, ok := ret.Get(0).(func(context.Context, string) []string); ok { + r0 = rf(ctx, key) } else { if ret.Get(0) != nil { r0 = ret.Get(0).([]string) } } - if rf, ok := ret.Get(1).(func(string) []string); ok { - r1 = rf(key) + if rf, ok := ret.Get(1).(func(context.Context, string) []string); ok { + r1 = rf(ctx, key) } else { if ret.Get(1) != nil { r1 = ret.Get(1).([]string) } } - if rf, ok := ret.Get(2).(func(string) error); ok { - r2 = rf(key) + if rf, ok := ret.Get(2).(func(context.Context, string) error); ok { + r2 = rf(ctx, key) } else { r2 = ret.Error(2) } @@ -369,14 +376,15 @@ type MetaKv_LoadWithPrefix_Call struct { } // LoadWithPrefix is a helper method to define mock.On call +// - ctx context.Context // - key string -func (_e *MetaKv_Expecter) LoadWithPrefix(key interface{}) *MetaKv_LoadWithPrefix_Call { - return &MetaKv_LoadWithPrefix_Call{Call: _e.mock.On("LoadWithPrefix", key)} +func (_e *MetaKv_Expecter) LoadWithPrefix(ctx interface{}, key interface{}) *MetaKv_LoadWithPrefix_Call { + return &MetaKv_LoadWithPrefix_Call{Call: _e.mock.On("LoadWithPrefix", ctx, key)} } -func (_c *MetaKv_LoadWithPrefix_Call) Run(run func(key string)) *MetaKv_LoadWithPrefix_Call { +func (_c *MetaKv_LoadWithPrefix_Call) Run(run func(ctx context.Context, key string)) *MetaKv_LoadWithPrefix_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string)) + run(args[0].(context.Context), args[1].(string)) }) return _c } @@ -386,14 +394,14 @@ func (_c *MetaKv_LoadWithPrefix_Call) Return(_a0 []string, _a1 []string, _a2 err return _c } -func (_c *MetaKv_LoadWithPrefix_Call) RunAndReturn(run func(string) ([]string, []string, error)) *MetaKv_LoadWithPrefix_Call { +func (_c *MetaKv_LoadWithPrefix_Call) RunAndReturn(run func(context.Context, string) ([]string, []string, error)) *MetaKv_LoadWithPrefix_Call { _c.Call.Return(run) return _c } -// MultiLoad provides a mock function with given fields: keys -func (_m *MetaKv) MultiLoad(keys []string) ([]string, error) { - ret := _m.Called(keys) +// MultiLoad provides a mock function with given fields: ctx, keys +func (_m *MetaKv) MultiLoad(ctx context.Context, keys []string) ([]string, error) { + ret := _m.Called(ctx, keys) if len(ret) == 0 { panic("no return value specified for MultiLoad") @@ -401,19 +409,19 @@ func (_m *MetaKv) MultiLoad(keys []string) ([]string, error) { var r0 []string var r1 error - if rf, ok := ret.Get(0).(func([]string) ([]string, error)); ok { - return rf(keys) + if rf, ok := ret.Get(0).(func(context.Context, []string) ([]string, error)); ok { + return rf(ctx, keys) } - if rf, ok := ret.Get(0).(func([]string) []string); ok { - r0 = rf(keys) + if rf, ok := ret.Get(0).(func(context.Context, []string) []string); ok { + r0 = rf(ctx, keys) } else { if ret.Get(0) != nil { r0 = ret.Get(0).([]string) } } - if rf, ok := ret.Get(1).(func([]string) error); ok { - r1 = rf(keys) + if rf, ok := ret.Get(1).(func(context.Context, []string) error); ok { + r1 = rf(ctx, keys) } else { r1 = ret.Error(1) } @@ -427,14 +435,15 @@ type MetaKv_MultiLoad_Call struct { } // MultiLoad is a helper method to define mock.On call +// - ctx context.Context // - keys []string -func (_e *MetaKv_Expecter) MultiLoad(keys interface{}) *MetaKv_MultiLoad_Call { - return &MetaKv_MultiLoad_Call{Call: _e.mock.On("MultiLoad", keys)} +func (_e *MetaKv_Expecter) MultiLoad(ctx interface{}, keys interface{}) *MetaKv_MultiLoad_Call { + return &MetaKv_MultiLoad_Call{Call: _e.mock.On("MultiLoad", ctx, keys)} } -func (_c *MetaKv_MultiLoad_Call) Run(run func(keys []string)) *MetaKv_MultiLoad_Call { +func (_c *MetaKv_MultiLoad_Call) Run(run func(ctx context.Context, keys []string)) *MetaKv_MultiLoad_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].([]string)) + run(args[0].(context.Context), args[1].([]string)) }) return _c } @@ -444,22 +453,22 @@ func (_c *MetaKv_MultiLoad_Call) Return(_a0 []string, _a1 error) *MetaKv_MultiLo return _c } -func (_c *MetaKv_MultiLoad_Call) RunAndReturn(run func([]string) ([]string, error)) *MetaKv_MultiLoad_Call { +func (_c *MetaKv_MultiLoad_Call) RunAndReturn(run func(context.Context, []string) ([]string, error)) *MetaKv_MultiLoad_Call { _c.Call.Return(run) return _c } -// MultiRemove provides a mock function with given fields: keys -func (_m *MetaKv) MultiRemove(keys []string) error { - ret := _m.Called(keys) +// MultiRemove provides a mock function with given fields: ctx, keys +func (_m *MetaKv) MultiRemove(ctx context.Context, keys []string) error { + ret := _m.Called(ctx, keys) if len(ret) == 0 { panic("no return value specified for MultiRemove") } var r0 error - if rf, ok := ret.Get(0).(func([]string) error); ok { - r0 = rf(keys) + if rf, ok := ret.Get(0).(func(context.Context, []string) error); ok { + r0 = rf(ctx, keys) } else { r0 = ret.Error(0) } @@ -473,14 +482,15 @@ type MetaKv_MultiRemove_Call struct { } // MultiRemove is a helper method to define mock.On call +// - ctx context.Context // - keys []string -func (_e *MetaKv_Expecter) MultiRemove(keys interface{}) *MetaKv_MultiRemove_Call { - return &MetaKv_MultiRemove_Call{Call: _e.mock.On("MultiRemove", keys)} +func (_e *MetaKv_Expecter) MultiRemove(ctx interface{}, keys interface{}) *MetaKv_MultiRemove_Call { + return &MetaKv_MultiRemove_Call{Call: _e.mock.On("MultiRemove", ctx, keys)} } -func (_c *MetaKv_MultiRemove_Call) Run(run func(keys []string)) *MetaKv_MultiRemove_Call { +func (_c *MetaKv_MultiRemove_Call) Run(run func(ctx context.Context, keys []string)) *MetaKv_MultiRemove_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].([]string)) + run(args[0].(context.Context), args[1].([]string)) }) return _c } @@ -490,22 +500,22 @@ func (_c *MetaKv_MultiRemove_Call) Return(_a0 error) *MetaKv_MultiRemove_Call { return _c } -func (_c *MetaKv_MultiRemove_Call) RunAndReturn(run func([]string) error) *MetaKv_MultiRemove_Call { +func (_c *MetaKv_MultiRemove_Call) RunAndReturn(run func(context.Context, []string) error) *MetaKv_MultiRemove_Call { _c.Call.Return(run) return _c } -// MultiSave provides a mock function with given fields: kvs -func (_m *MetaKv) MultiSave(kvs map[string]string) error { - ret := _m.Called(kvs) +// MultiSave provides a mock function with given fields: ctx, kvs +func (_m *MetaKv) MultiSave(ctx context.Context, kvs map[string]string) error { + ret := _m.Called(ctx, kvs) if len(ret) == 0 { panic("no return value specified for MultiSave") } var r0 error - if rf, ok := ret.Get(0).(func(map[string]string) error); ok { - r0 = rf(kvs) + if rf, ok := ret.Get(0).(func(context.Context, map[string]string) error); ok { + r0 = rf(ctx, kvs) } else { r0 = ret.Error(0) } @@ -519,14 +529,15 @@ type MetaKv_MultiSave_Call struct { } // MultiSave is a helper method to define mock.On call +// - ctx context.Context // - kvs map[string]string -func (_e *MetaKv_Expecter) MultiSave(kvs interface{}) *MetaKv_MultiSave_Call { - return &MetaKv_MultiSave_Call{Call: _e.mock.On("MultiSave", kvs)} +func (_e *MetaKv_Expecter) MultiSave(ctx interface{}, kvs interface{}) *MetaKv_MultiSave_Call { + return &MetaKv_MultiSave_Call{Call: _e.mock.On("MultiSave", ctx, kvs)} } -func (_c *MetaKv_MultiSave_Call) Run(run func(kvs map[string]string)) *MetaKv_MultiSave_Call { +func (_c *MetaKv_MultiSave_Call) Run(run func(ctx context.Context, kvs map[string]string)) *MetaKv_MultiSave_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(map[string]string)) + run(args[0].(context.Context), args[1].(map[string]string)) }) return _c } @@ -536,19 +547,19 @@ func (_c *MetaKv_MultiSave_Call) Return(_a0 error) *MetaKv_MultiSave_Call { return _c } -func (_c *MetaKv_MultiSave_Call) RunAndReturn(run func(map[string]string) error) *MetaKv_MultiSave_Call { +func (_c *MetaKv_MultiSave_Call) RunAndReturn(run func(context.Context, map[string]string) error) *MetaKv_MultiSave_Call { _c.Call.Return(run) return _c } -// MultiSaveAndRemove provides a mock function with given fields: saves, removals, preds -func (_m *MetaKv) MultiSaveAndRemove(saves map[string]string, removals []string, preds ...predicates.Predicate) error { +// MultiSaveAndRemove provides a mock function with given fields: ctx, saves, removals, preds +func (_m *MetaKv) MultiSaveAndRemove(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error { _va := make([]interface{}, len(preds)) for _i := range preds { _va[_i] = preds[_i] } var _ca []interface{} - _ca = append(_ca, saves, removals) + _ca = append(_ca, ctx, saves, removals) _ca = append(_ca, _va...) ret := _m.Called(_ca...) @@ -557,8 +568,8 @@ func (_m *MetaKv) MultiSaveAndRemove(saves map[string]string, removals []string, } var r0 error - if rf, ok := ret.Get(0).(func(map[string]string, []string, ...predicates.Predicate) error); ok { - r0 = rf(saves, removals, preds...) + if rf, ok := ret.Get(0).(func(context.Context, map[string]string, []string, ...predicates.Predicate) error); ok { + r0 = rf(ctx, saves, removals, preds...) } else { r0 = ret.Error(0) } @@ -572,23 +583,24 @@ type MetaKv_MultiSaveAndRemove_Call struct { } // MultiSaveAndRemove is a helper method to define mock.On call +// - ctx context.Context // - saves map[string]string // - removals []string // - preds ...predicates.Predicate -func (_e *MetaKv_Expecter) MultiSaveAndRemove(saves interface{}, removals interface{}, preds ...interface{}) *MetaKv_MultiSaveAndRemove_Call { +func (_e *MetaKv_Expecter) MultiSaveAndRemove(ctx interface{}, saves interface{}, removals interface{}, preds ...interface{}) *MetaKv_MultiSaveAndRemove_Call { return &MetaKv_MultiSaveAndRemove_Call{Call: _e.mock.On("MultiSaveAndRemove", - append([]interface{}{saves, removals}, preds...)...)} + append([]interface{}{ctx, saves, removals}, preds...)...)} } -func (_c *MetaKv_MultiSaveAndRemove_Call) Run(run func(saves map[string]string, removals []string, preds ...predicates.Predicate)) *MetaKv_MultiSaveAndRemove_Call { +func (_c *MetaKv_MultiSaveAndRemove_Call) Run(run func(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate)) *MetaKv_MultiSaveAndRemove_Call { _c.Call.Run(func(args mock.Arguments) { - variadicArgs := make([]predicates.Predicate, len(args)-2) - for i, a := range args[2:] { + variadicArgs := make([]predicates.Predicate, len(args)-3) + for i, a := range args[3:] { if a != nil { variadicArgs[i] = a.(predicates.Predicate) } } - run(args[0].(map[string]string), args[1].([]string), variadicArgs...) + run(args[0].(context.Context), args[1].(map[string]string), args[2].([]string), variadicArgs...) }) return _c } @@ -598,19 +610,19 @@ func (_c *MetaKv_MultiSaveAndRemove_Call) Return(_a0 error) *MetaKv_MultiSaveAnd return _c } -func (_c *MetaKv_MultiSaveAndRemove_Call) RunAndReturn(run func(map[string]string, []string, ...predicates.Predicate) error) *MetaKv_MultiSaveAndRemove_Call { +func (_c *MetaKv_MultiSaveAndRemove_Call) RunAndReturn(run func(context.Context, map[string]string, []string, ...predicates.Predicate) error) *MetaKv_MultiSaveAndRemove_Call { _c.Call.Return(run) return _c } -// MultiSaveAndRemoveWithPrefix provides a mock function with given fields: saves, removals, preds -func (_m *MetaKv) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals []string, preds ...predicates.Predicate) error { +// MultiSaveAndRemoveWithPrefix provides a mock function with given fields: ctx, saves, removals, preds +func (_m *MetaKv) MultiSaveAndRemoveWithPrefix(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error { _va := make([]interface{}, len(preds)) for _i := range preds { _va[_i] = preds[_i] } var _ca []interface{} - _ca = append(_ca, saves, removals) + _ca = append(_ca, ctx, saves, removals) _ca = append(_ca, _va...) ret := _m.Called(_ca...) @@ -619,8 +631,8 @@ func (_m *MetaKv) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals } var r0 error - if rf, ok := ret.Get(0).(func(map[string]string, []string, ...predicates.Predicate) error); ok { - r0 = rf(saves, removals, preds...) + if rf, ok := ret.Get(0).(func(context.Context, map[string]string, []string, ...predicates.Predicate) error); ok { + r0 = rf(ctx, saves, removals, preds...) } else { r0 = ret.Error(0) } @@ -634,23 +646,24 @@ type MetaKv_MultiSaveAndRemoveWithPrefix_Call struct { } // MultiSaveAndRemoveWithPrefix is a helper method to define mock.On call +// - ctx context.Context // - saves map[string]string // - removals []string // - preds ...predicates.Predicate -func (_e *MetaKv_Expecter) MultiSaveAndRemoveWithPrefix(saves interface{}, removals interface{}, preds ...interface{}) *MetaKv_MultiSaveAndRemoveWithPrefix_Call { +func (_e *MetaKv_Expecter) MultiSaveAndRemoveWithPrefix(ctx interface{}, saves interface{}, removals interface{}, preds ...interface{}) *MetaKv_MultiSaveAndRemoveWithPrefix_Call { return &MetaKv_MultiSaveAndRemoveWithPrefix_Call{Call: _e.mock.On("MultiSaveAndRemoveWithPrefix", - append([]interface{}{saves, removals}, preds...)...)} + append([]interface{}{ctx, saves, removals}, preds...)...)} } -func (_c *MetaKv_MultiSaveAndRemoveWithPrefix_Call) Run(run func(saves map[string]string, removals []string, preds ...predicates.Predicate)) *MetaKv_MultiSaveAndRemoveWithPrefix_Call { +func (_c *MetaKv_MultiSaveAndRemoveWithPrefix_Call) Run(run func(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate)) *MetaKv_MultiSaveAndRemoveWithPrefix_Call { _c.Call.Run(func(args mock.Arguments) { - variadicArgs := make([]predicates.Predicate, len(args)-2) - for i, a := range args[2:] { + variadicArgs := make([]predicates.Predicate, len(args)-3) + for i, a := range args[3:] { if a != nil { variadicArgs[i] = a.(predicates.Predicate) } } - run(args[0].(map[string]string), args[1].([]string), variadicArgs...) + run(args[0].(context.Context), args[1].(map[string]string), args[2].([]string), variadicArgs...) }) return _c } @@ -660,22 +673,22 @@ func (_c *MetaKv_MultiSaveAndRemoveWithPrefix_Call) Return(_a0 error) *MetaKv_Mu return _c } -func (_c *MetaKv_MultiSaveAndRemoveWithPrefix_Call) RunAndReturn(run func(map[string]string, []string, ...predicates.Predicate) error) *MetaKv_MultiSaveAndRemoveWithPrefix_Call { +func (_c *MetaKv_MultiSaveAndRemoveWithPrefix_Call) RunAndReturn(run func(context.Context, map[string]string, []string, ...predicates.Predicate) error) *MetaKv_MultiSaveAndRemoveWithPrefix_Call { _c.Call.Return(run) return _c } -// Remove provides a mock function with given fields: key -func (_m *MetaKv) Remove(key string) error { - ret := _m.Called(key) +// Remove provides a mock function with given fields: ctx, key +func (_m *MetaKv) Remove(ctx context.Context, key string) error { + ret := _m.Called(ctx, key) if len(ret) == 0 { panic("no return value specified for Remove") } var r0 error - if rf, ok := ret.Get(0).(func(string) error); ok { - r0 = rf(key) + if rf, ok := ret.Get(0).(func(context.Context, string) error); ok { + r0 = rf(ctx, key) } else { r0 = ret.Error(0) } @@ -689,14 +702,15 @@ type MetaKv_Remove_Call struct { } // Remove is a helper method to define mock.On call +// - ctx context.Context // - key string -func (_e *MetaKv_Expecter) Remove(key interface{}) *MetaKv_Remove_Call { - return &MetaKv_Remove_Call{Call: _e.mock.On("Remove", key)} +func (_e *MetaKv_Expecter) Remove(ctx interface{}, key interface{}) *MetaKv_Remove_Call { + return &MetaKv_Remove_Call{Call: _e.mock.On("Remove", ctx, key)} } -func (_c *MetaKv_Remove_Call) Run(run func(key string)) *MetaKv_Remove_Call { +func (_c *MetaKv_Remove_Call) Run(run func(ctx context.Context, key string)) *MetaKv_Remove_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string)) + run(args[0].(context.Context), args[1].(string)) }) return _c } @@ -706,22 +720,22 @@ func (_c *MetaKv_Remove_Call) Return(_a0 error) *MetaKv_Remove_Call { return _c } -func (_c *MetaKv_Remove_Call) RunAndReturn(run func(string) error) *MetaKv_Remove_Call { +func (_c *MetaKv_Remove_Call) RunAndReturn(run func(context.Context, string) error) *MetaKv_Remove_Call { _c.Call.Return(run) return _c } -// RemoveWithPrefix provides a mock function with given fields: key -func (_m *MetaKv) RemoveWithPrefix(key string) error { - ret := _m.Called(key) +// RemoveWithPrefix provides a mock function with given fields: ctx, key +func (_m *MetaKv) RemoveWithPrefix(ctx context.Context, key string) error { + ret := _m.Called(ctx, key) if len(ret) == 0 { panic("no return value specified for RemoveWithPrefix") } var r0 error - if rf, ok := ret.Get(0).(func(string) error); ok { - r0 = rf(key) + if rf, ok := ret.Get(0).(func(context.Context, string) error); ok { + r0 = rf(ctx, key) } else { r0 = ret.Error(0) } @@ -735,14 +749,15 @@ type MetaKv_RemoveWithPrefix_Call struct { } // RemoveWithPrefix is a helper method to define mock.On call +// - ctx context.Context // - key string -func (_e *MetaKv_Expecter) RemoveWithPrefix(key interface{}) *MetaKv_RemoveWithPrefix_Call { - return &MetaKv_RemoveWithPrefix_Call{Call: _e.mock.On("RemoveWithPrefix", key)} +func (_e *MetaKv_Expecter) RemoveWithPrefix(ctx interface{}, key interface{}) *MetaKv_RemoveWithPrefix_Call { + return &MetaKv_RemoveWithPrefix_Call{Call: _e.mock.On("RemoveWithPrefix", ctx, key)} } -func (_c *MetaKv_RemoveWithPrefix_Call) Run(run func(key string)) *MetaKv_RemoveWithPrefix_Call { +func (_c *MetaKv_RemoveWithPrefix_Call) Run(run func(ctx context.Context, key string)) *MetaKv_RemoveWithPrefix_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string)) + run(args[0].(context.Context), args[1].(string)) }) return _c } @@ -752,22 +767,22 @@ func (_c *MetaKv_RemoveWithPrefix_Call) Return(_a0 error) *MetaKv_RemoveWithPref return _c } -func (_c *MetaKv_RemoveWithPrefix_Call) RunAndReturn(run func(string) error) *MetaKv_RemoveWithPrefix_Call { +func (_c *MetaKv_RemoveWithPrefix_Call) RunAndReturn(run func(context.Context, string) error) *MetaKv_RemoveWithPrefix_Call { _c.Call.Return(run) return _c } -// Save provides a mock function with given fields: key, value -func (_m *MetaKv) Save(key string, value string) error { - ret := _m.Called(key, value) +// Save provides a mock function with given fields: ctx, key, value +func (_m *MetaKv) Save(ctx context.Context, key string, value string) error { + ret := _m.Called(ctx, key, value) if len(ret) == 0 { panic("no return value specified for Save") } var r0 error - if rf, ok := ret.Get(0).(func(string, string) error); ok { - r0 = rf(key, value) + if rf, ok := ret.Get(0).(func(context.Context, string, string) error); ok { + r0 = rf(ctx, key, value) } else { r0 = ret.Error(0) } @@ -781,15 +796,16 @@ type MetaKv_Save_Call struct { } // Save is a helper method to define mock.On call +// - ctx context.Context // - key string // - value string -func (_e *MetaKv_Expecter) Save(key interface{}, value interface{}) *MetaKv_Save_Call { - return &MetaKv_Save_Call{Call: _e.mock.On("Save", key, value)} +func (_e *MetaKv_Expecter) Save(ctx interface{}, key interface{}, value interface{}) *MetaKv_Save_Call { + return &MetaKv_Save_Call{Call: _e.mock.On("Save", ctx, key, value)} } -func (_c *MetaKv_Save_Call) Run(run func(key string, value string)) *MetaKv_Save_Call { +func (_c *MetaKv_Save_Call) Run(run func(ctx context.Context, key string, value string)) *MetaKv_Save_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string), args[1].(string)) + run(args[0].(context.Context), args[1].(string), args[2].(string)) }) return _c } @@ -799,22 +815,22 @@ func (_c *MetaKv_Save_Call) Return(_a0 error) *MetaKv_Save_Call { return _c } -func (_c *MetaKv_Save_Call) RunAndReturn(run func(string, string) error) *MetaKv_Save_Call { +func (_c *MetaKv_Save_Call) RunAndReturn(run func(context.Context, string, string) error) *MetaKv_Save_Call { _c.Call.Return(run) return _c } -// WalkWithPrefix provides a mock function with given fields: prefix, paginationSize, fn -func (_m *MetaKv) WalkWithPrefix(prefix string, paginationSize int, fn func([]byte, []byte) error) error { - ret := _m.Called(prefix, paginationSize, fn) +// WalkWithPrefix provides a mock function with given fields: ctx, prefix, paginationSize, fn +func (_m *MetaKv) WalkWithPrefix(ctx context.Context, prefix string, paginationSize int, fn func([]byte, []byte) error) error { + ret := _m.Called(ctx, prefix, paginationSize, fn) if len(ret) == 0 { panic("no return value specified for WalkWithPrefix") } var r0 error - if rf, ok := ret.Get(0).(func(string, int, func([]byte, []byte) error) error); ok { - r0 = rf(prefix, paginationSize, fn) + if rf, ok := ret.Get(0).(func(context.Context, string, int, func([]byte, []byte) error) error); ok { + r0 = rf(ctx, prefix, paginationSize, fn) } else { r0 = ret.Error(0) } @@ -828,16 +844,17 @@ type MetaKv_WalkWithPrefix_Call struct { } // WalkWithPrefix is a helper method to define mock.On call +// - ctx context.Context // - prefix string // - paginationSize int // - fn func([]byte , []byte) error -func (_e *MetaKv_Expecter) WalkWithPrefix(prefix interface{}, paginationSize interface{}, fn interface{}) *MetaKv_WalkWithPrefix_Call { - return &MetaKv_WalkWithPrefix_Call{Call: _e.mock.On("WalkWithPrefix", prefix, paginationSize, fn)} +func (_e *MetaKv_Expecter) WalkWithPrefix(ctx interface{}, prefix interface{}, paginationSize interface{}, fn interface{}) *MetaKv_WalkWithPrefix_Call { + return &MetaKv_WalkWithPrefix_Call{Call: _e.mock.On("WalkWithPrefix", ctx, prefix, paginationSize, fn)} } -func (_c *MetaKv_WalkWithPrefix_Call) Run(run func(prefix string, paginationSize int, fn func([]byte, []byte) error)) *MetaKv_WalkWithPrefix_Call { +func (_c *MetaKv_WalkWithPrefix_Call) Run(run func(ctx context.Context, prefix string, paginationSize int, fn func([]byte, []byte) error)) *MetaKv_WalkWithPrefix_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string), args[1].(int), args[2].(func([]byte, []byte) error)) + run(args[0].(context.Context), args[1].(string), args[2].(int), args[3].(func([]byte, []byte) error)) }) return _c } @@ -847,7 +864,7 @@ func (_c *MetaKv_WalkWithPrefix_Call) Return(_a0 error) *MetaKv_WalkWithPrefix_C return _c } -func (_c *MetaKv_WalkWithPrefix_Call) RunAndReturn(run func(string, int, func([]byte, []byte) error) error) *MetaKv_WalkWithPrefix_Call { +func (_c *MetaKv_WalkWithPrefix_Call) RunAndReturn(run func(context.Context, string, int, func([]byte, []byte) error) error) *MetaKv_WalkWithPrefix_Call { _c.Call.Return(run) return _c } diff --git a/internal/kv/mocks/snapshot_kv.go b/internal/kv/mocks/snapshot_kv.go index 42bc5d7ad9ed9..220385636b8d7 100644 --- a/internal/kv/mocks/snapshot_kv.go +++ b/internal/kv/mocks/snapshot_kv.go @@ -2,7 +2,11 @@ package mocks -import mock "github.com/stretchr/testify/mock" +import ( + context "context" + + mock "github.com/stretchr/testify/mock" +) // SnapShotKV is an autogenerated mock type for the SnapShotKV type type SnapShotKV struct { @@ -17,9 +21,9 @@ func (_m *SnapShotKV) EXPECT() *SnapShotKV_Expecter { return &SnapShotKV_Expecter{mock: &_m.Mock} } -// Load provides a mock function with given fields: key, ts -func (_m *SnapShotKV) Load(key string, ts uint64) (string, error) { - ret := _m.Called(key, ts) +// Load provides a mock function with given fields: ctx, key, ts +func (_m *SnapShotKV) Load(ctx context.Context, key string, ts uint64) (string, error) { + ret := _m.Called(ctx, key, ts) if len(ret) == 0 { panic("no return value specified for Load") @@ -27,17 +31,17 @@ func (_m *SnapShotKV) Load(key string, ts uint64) (string, error) { var r0 string var r1 error - if rf, ok := ret.Get(0).(func(string, uint64) (string, error)); ok { - return rf(key, ts) + if rf, ok := ret.Get(0).(func(context.Context, string, uint64) (string, error)); ok { + return rf(ctx, key, ts) } - if rf, ok := ret.Get(0).(func(string, uint64) string); ok { - r0 = rf(key, ts) + if rf, ok := ret.Get(0).(func(context.Context, string, uint64) string); ok { + r0 = rf(ctx, key, ts) } else { r0 = ret.Get(0).(string) } - if rf, ok := ret.Get(1).(func(string, uint64) error); ok { - r1 = rf(key, ts) + if rf, ok := ret.Get(1).(func(context.Context, string, uint64) error); ok { + r1 = rf(ctx, key, ts) } else { r1 = ret.Error(1) } @@ -51,15 +55,16 @@ type SnapShotKV_Load_Call struct { } // Load is a helper method to define mock.On call +// - ctx context.Context // - key string // - ts uint64 -func (_e *SnapShotKV_Expecter) Load(key interface{}, ts interface{}) *SnapShotKV_Load_Call { - return &SnapShotKV_Load_Call{Call: _e.mock.On("Load", key, ts)} +func (_e *SnapShotKV_Expecter) Load(ctx interface{}, key interface{}, ts interface{}) *SnapShotKV_Load_Call { + return &SnapShotKV_Load_Call{Call: _e.mock.On("Load", ctx, key, ts)} } -func (_c *SnapShotKV_Load_Call) Run(run func(key string, ts uint64)) *SnapShotKV_Load_Call { +func (_c *SnapShotKV_Load_Call) Run(run func(ctx context.Context, key string, ts uint64)) *SnapShotKV_Load_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string), args[1].(uint64)) + run(args[0].(context.Context), args[1].(string), args[2].(uint64)) }) return _c } @@ -69,14 +74,14 @@ func (_c *SnapShotKV_Load_Call) Return(_a0 string, _a1 error) *SnapShotKV_Load_C return _c } -func (_c *SnapShotKV_Load_Call) RunAndReturn(run func(string, uint64) (string, error)) *SnapShotKV_Load_Call { +func (_c *SnapShotKV_Load_Call) RunAndReturn(run func(context.Context, string, uint64) (string, error)) *SnapShotKV_Load_Call { _c.Call.Return(run) return _c } -// LoadWithPrefix provides a mock function with given fields: key, ts -func (_m *SnapShotKV) LoadWithPrefix(key string, ts uint64) ([]string, []string, error) { - ret := _m.Called(key, ts) +// LoadWithPrefix provides a mock function with given fields: ctx, key, ts +func (_m *SnapShotKV) LoadWithPrefix(ctx context.Context, key string, ts uint64) ([]string, []string, error) { + ret := _m.Called(ctx, key, ts) if len(ret) == 0 { panic("no return value specified for LoadWithPrefix") @@ -85,27 +90,27 @@ func (_m *SnapShotKV) LoadWithPrefix(key string, ts uint64) ([]string, []string, var r0 []string var r1 []string var r2 error - if rf, ok := ret.Get(0).(func(string, uint64) ([]string, []string, error)); ok { - return rf(key, ts) + if rf, ok := ret.Get(0).(func(context.Context, string, uint64) ([]string, []string, error)); ok { + return rf(ctx, key, ts) } - if rf, ok := ret.Get(0).(func(string, uint64) []string); ok { - r0 = rf(key, ts) + if rf, ok := ret.Get(0).(func(context.Context, string, uint64) []string); ok { + r0 = rf(ctx, key, ts) } else { if ret.Get(0) != nil { r0 = ret.Get(0).([]string) } } - if rf, ok := ret.Get(1).(func(string, uint64) []string); ok { - r1 = rf(key, ts) + if rf, ok := ret.Get(1).(func(context.Context, string, uint64) []string); ok { + r1 = rf(ctx, key, ts) } else { if ret.Get(1) != nil { r1 = ret.Get(1).([]string) } } - if rf, ok := ret.Get(2).(func(string, uint64) error); ok { - r2 = rf(key, ts) + if rf, ok := ret.Get(2).(func(context.Context, string, uint64) error); ok { + r2 = rf(ctx, key, ts) } else { r2 = ret.Error(2) } @@ -119,15 +124,16 @@ type SnapShotKV_LoadWithPrefix_Call struct { } // LoadWithPrefix is a helper method to define mock.On call +// - ctx context.Context // - key string // - ts uint64 -func (_e *SnapShotKV_Expecter) LoadWithPrefix(key interface{}, ts interface{}) *SnapShotKV_LoadWithPrefix_Call { - return &SnapShotKV_LoadWithPrefix_Call{Call: _e.mock.On("LoadWithPrefix", key, ts)} +func (_e *SnapShotKV_Expecter) LoadWithPrefix(ctx interface{}, key interface{}, ts interface{}) *SnapShotKV_LoadWithPrefix_Call { + return &SnapShotKV_LoadWithPrefix_Call{Call: _e.mock.On("LoadWithPrefix", ctx, key, ts)} } -func (_c *SnapShotKV_LoadWithPrefix_Call) Run(run func(key string, ts uint64)) *SnapShotKV_LoadWithPrefix_Call { +func (_c *SnapShotKV_LoadWithPrefix_Call) Run(run func(ctx context.Context, key string, ts uint64)) *SnapShotKV_LoadWithPrefix_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string), args[1].(uint64)) + run(args[0].(context.Context), args[1].(string), args[2].(uint64)) }) return _c } @@ -137,22 +143,22 @@ func (_c *SnapShotKV_LoadWithPrefix_Call) Return(_a0 []string, _a1 []string, _a2 return _c } -func (_c *SnapShotKV_LoadWithPrefix_Call) RunAndReturn(run func(string, uint64) ([]string, []string, error)) *SnapShotKV_LoadWithPrefix_Call { +func (_c *SnapShotKV_LoadWithPrefix_Call) RunAndReturn(run func(context.Context, string, uint64) ([]string, []string, error)) *SnapShotKV_LoadWithPrefix_Call { _c.Call.Return(run) return _c } -// MultiSave provides a mock function with given fields: kvs, ts -func (_m *SnapShotKV) MultiSave(kvs map[string]string, ts uint64) error { - ret := _m.Called(kvs, ts) +// MultiSave provides a mock function with given fields: ctx, kvs, ts +func (_m *SnapShotKV) MultiSave(ctx context.Context, kvs map[string]string, ts uint64) error { + ret := _m.Called(ctx, kvs, ts) if len(ret) == 0 { panic("no return value specified for MultiSave") } var r0 error - if rf, ok := ret.Get(0).(func(map[string]string, uint64) error); ok { - r0 = rf(kvs, ts) + if rf, ok := ret.Get(0).(func(context.Context, map[string]string, uint64) error); ok { + r0 = rf(ctx, kvs, ts) } else { r0 = ret.Error(0) } @@ -166,15 +172,16 @@ type SnapShotKV_MultiSave_Call struct { } // MultiSave is a helper method to define mock.On call +// - ctx context.Context // - kvs map[string]string // - ts uint64 -func (_e *SnapShotKV_Expecter) MultiSave(kvs interface{}, ts interface{}) *SnapShotKV_MultiSave_Call { - return &SnapShotKV_MultiSave_Call{Call: _e.mock.On("MultiSave", kvs, ts)} +func (_e *SnapShotKV_Expecter) MultiSave(ctx interface{}, kvs interface{}, ts interface{}) *SnapShotKV_MultiSave_Call { + return &SnapShotKV_MultiSave_Call{Call: _e.mock.On("MultiSave", ctx, kvs, ts)} } -func (_c *SnapShotKV_MultiSave_Call) Run(run func(kvs map[string]string, ts uint64)) *SnapShotKV_MultiSave_Call { +func (_c *SnapShotKV_MultiSave_Call) Run(run func(ctx context.Context, kvs map[string]string, ts uint64)) *SnapShotKV_MultiSave_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(map[string]string), args[1].(uint64)) + run(args[0].(context.Context), args[1].(map[string]string), args[2].(uint64)) }) return _c } @@ -184,22 +191,22 @@ func (_c *SnapShotKV_MultiSave_Call) Return(_a0 error) *SnapShotKV_MultiSave_Cal return _c } -func (_c *SnapShotKV_MultiSave_Call) RunAndReturn(run func(map[string]string, uint64) error) *SnapShotKV_MultiSave_Call { +func (_c *SnapShotKV_MultiSave_Call) RunAndReturn(run func(context.Context, map[string]string, uint64) error) *SnapShotKV_MultiSave_Call { _c.Call.Return(run) return _c } -// MultiSaveAndRemove provides a mock function with given fields: saves, removals, ts -func (_m *SnapShotKV) MultiSaveAndRemove(saves map[string]string, removals []string, ts uint64) error { - ret := _m.Called(saves, removals, ts) +// MultiSaveAndRemove provides a mock function with given fields: ctx, saves, removals, ts +func (_m *SnapShotKV) MultiSaveAndRemove(ctx context.Context, saves map[string]string, removals []string, ts uint64) error { + ret := _m.Called(ctx, saves, removals, ts) if len(ret) == 0 { panic("no return value specified for MultiSaveAndRemove") } var r0 error - if rf, ok := ret.Get(0).(func(map[string]string, []string, uint64) error); ok { - r0 = rf(saves, removals, ts) + if rf, ok := ret.Get(0).(func(context.Context, map[string]string, []string, uint64) error); ok { + r0 = rf(ctx, saves, removals, ts) } else { r0 = ret.Error(0) } @@ -213,16 +220,17 @@ type SnapShotKV_MultiSaveAndRemove_Call struct { } // MultiSaveAndRemove is a helper method to define mock.On call +// - ctx context.Context // - saves map[string]string // - removals []string // - ts uint64 -func (_e *SnapShotKV_Expecter) MultiSaveAndRemove(saves interface{}, removals interface{}, ts interface{}) *SnapShotKV_MultiSaveAndRemove_Call { - return &SnapShotKV_MultiSaveAndRemove_Call{Call: _e.mock.On("MultiSaveAndRemove", saves, removals, ts)} +func (_e *SnapShotKV_Expecter) MultiSaveAndRemove(ctx interface{}, saves interface{}, removals interface{}, ts interface{}) *SnapShotKV_MultiSaveAndRemove_Call { + return &SnapShotKV_MultiSaveAndRemove_Call{Call: _e.mock.On("MultiSaveAndRemove", ctx, saves, removals, ts)} } -func (_c *SnapShotKV_MultiSaveAndRemove_Call) Run(run func(saves map[string]string, removals []string, ts uint64)) *SnapShotKV_MultiSaveAndRemove_Call { +func (_c *SnapShotKV_MultiSaveAndRemove_Call) Run(run func(ctx context.Context, saves map[string]string, removals []string, ts uint64)) *SnapShotKV_MultiSaveAndRemove_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(map[string]string), args[1].([]string), args[2].(uint64)) + run(args[0].(context.Context), args[1].(map[string]string), args[2].([]string), args[3].(uint64)) }) return _c } @@ -232,22 +240,22 @@ func (_c *SnapShotKV_MultiSaveAndRemove_Call) Return(_a0 error) *SnapShotKV_Mult return _c } -func (_c *SnapShotKV_MultiSaveAndRemove_Call) RunAndReturn(run func(map[string]string, []string, uint64) error) *SnapShotKV_MultiSaveAndRemove_Call { +func (_c *SnapShotKV_MultiSaveAndRemove_Call) RunAndReturn(run func(context.Context, map[string]string, []string, uint64) error) *SnapShotKV_MultiSaveAndRemove_Call { _c.Call.Return(run) return _c } -// MultiSaveAndRemoveWithPrefix provides a mock function with given fields: saves, removals, ts -func (_m *SnapShotKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals []string, ts uint64) error { - ret := _m.Called(saves, removals, ts) +// MultiSaveAndRemoveWithPrefix provides a mock function with given fields: ctx, saves, removals, ts +func (_m *SnapShotKV) MultiSaveAndRemoveWithPrefix(ctx context.Context, saves map[string]string, removals []string, ts uint64) error { + ret := _m.Called(ctx, saves, removals, ts) if len(ret) == 0 { panic("no return value specified for MultiSaveAndRemoveWithPrefix") } var r0 error - if rf, ok := ret.Get(0).(func(map[string]string, []string, uint64) error); ok { - r0 = rf(saves, removals, ts) + if rf, ok := ret.Get(0).(func(context.Context, map[string]string, []string, uint64) error); ok { + r0 = rf(ctx, saves, removals, ts) } else { r0 = ret.Error(0) } @@ -261,16 +269,17 @@ type SnapShotKV_MultiSaveAndRemoveWithPrefix_Call struct { } // MultiSaveAndRemoveWithPrefix is a helper method to define mock.On call +// - ctx context.Context // - saves map[string]string // - removals []string // - ts uint64 -func (_e *SnapShotKV_Expecter) MultiSaveAndRemoveWithPrefix(saves interface{}, removals interface{}, ts interface{}) *SnapShotKV_MultiSaveAndRemoveWithPrefix_Call { - return &SnapShotKV_MultiSaveAndRemoveWithPrefix_Call{Call: _e.mock.On("MultiSaveAndRemoveWithPrefix", saves, removals, ts)} +func (_e *SnapShotKV_Expecter) MultiSaveAndRemoveWithPrefix(ctx interface{}, saves interface{}, removals interface{}, ts interface{}) *SnapShotKV_MultiSaveAndRemoveWithPrefix_Call { + return &SnapShotKV_MultiSaveAndRemoveWithPrefix_Call{Call: _e.mock.On("MultiSaveAndRemoveWithPrefix", ctx, saves, removals, ts)} } -func (_c *SnapShotKV_MultiSaveAndRemoveWithPrefix_Call) Run(run func(saves map[string]string, removals []string, ts uint64)) *SnapShotKV_MultiSaveAndRemoveWithPrefix_Call { +func (_c *SnapShotKV_MultiSaveAndRemoveWithPrefix_Call) Run(run func(ctx context.Context, saves map[string]string, removals []string, ts uint64)) *SnapShotKV_MultiSaveAndRemoveWithPrefix_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(map[string]string), args[1].([]string), args[2].(uint64)) + run(args[0].(context.Context), args[1].(map[string]string), args[2].([]string), args[3].(uint64)) }) return _c } @@ -280,22 +289,22 @@ func (_c *SnapShotKV_MultiSaveAndRemoveWithPrefix_Call) Return(_a0 error) *SnapS return _c } -func (_c *SnapShotKV_MultiSaveAndRemoveWithPrefix_Call) RunAndReturn(run func(map[string]string, []string, uint64) error) *SnapShotKV_MultiSaveAndRemoveWithPrefix_Call { +func (_c *SnapShotKV_MultiSaveAndRemoveWithPrefix_Call) RunAndReturn(run func(context.Context, map[string]string, []string, uint64) error) *SnapShotKV_MultiSaveAndRemoveWithPrefix_Call { _c.Call.Return(run) return _c } -// Save provides a mock function with given fields: key, value, ts -func (_m *SnapShotKV) Save(key string, value string, ts uint64) error { - ret := _m.Called(key, value, ts) +// Save provides a mock function with given fields: ctx, key, value, ts +func (_m *SnapShotKV) Save(ctx context.Context, key string, value string, ts uint64) error { + ret := _m.Called(ctx, key, value, ts) if len(ret) == 0 { panic("no return value specified for Save") } var r0 error - if rf, ok := ret.Get(0).(func(string, string, uint64) error); ok { - r0 = rf(key, value, ts) + if rf, ok := ret.Get(0).(func(context.Context, string, string, uint64) error); ok { + r0 = rf(ctx, key, value, ts) } else { r0 = ret.Error(0) } @@ -309,16 +318,17 @@ type SnapShotKV_Save_Call struct { } // Save is a helper method to define mock.On call +// - ctx context.Context // - key string // - value string // - ts uint64 -func (_e *SnapShotKV_Expecter) Save(key interface{}, value interface{}, ts interface{}) *SnapShotKV_Save_Call { - return &SnapShotKV_Save_Call{Call: _e.mock.On("Save", key, value, ts)} +func (_e *SnapShotKV_Expecter) Save(ctx interface{}, key interface{}, value interface{}, ts interface{}) *SnapShotKV_Save_Call { + return &SnapShotKV_Save_Call{Call: _e.mock.On("Save", ctx, key, value, ts)} } -func (_c *SnapShotKV_Save_Call) Run(run func(key string, value string, ts uint64)) *SnapShotKV_Save_Call { +func (_c *SnapShotKV_Save_Call) Run(run func(ctx context.Context, key string, value string, ts uint64)) *SnapShotKV_Save_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string), args[1].(string), args[2].(uint64)) + run(args[0].(context.Context), args[1].(string), args[2].(string), args[3].(uint64)) }) return _c } @@ -328,7 +338,7 @@ func (_c *SnapShotKV_Save_Call) Return(_a0 error) *SnapShotKV_Save_Call { return _c } -func (_c *SnapShotKV_Save_Call) RunAndReturn(run func(string, string, uint64) error) *SnapShotKV_Save_Call { +func (_c *SnapShotKV_Save_Call) RunAndReturn(run func(context.Context, string, string, uint64) error) *SnapShotKV_Save_Call { _c.Call.Return(run) return _c } diff --git a/internal/kv/mocks/txn_kv.go b/internal/kv/mocks/txn_kv.go index 76811bbfb7a9f..7a149148fdc7e 100644 --- a/internal/kv/mocks/txn_kv.go +++ b/internal/kv/mocks/txn_kv.go @@ -3,8 +3,11 @@ package mocks import ( - predicates "github.com/milvus-io/milvus/pkg/kv/predicates" + context "context" + mock "github.com/stretchr/testify/mock" + + predicates "github.com/milvus-io/milvus/pkg/kv/predicates" ) // TxnKV is an autogenerated mock type for the TxnKV type @@ -52,9 +55,9 @@ func (_c *TxnKV_Close_Call) RunAndReturn(run func()) *TxnKV_Close_Call { return _c } -// Has provides a mock function with given fields: key -func (_m *TxnKV) Has(key string) (bool, error) { - ret := _m.Called(key) +// Has provides a mock function with given fields: ctx, key +func (_m *TxnKV) Has(ctx context.Context, key string) (bool, error) { + ret := _m.Called(ctx, key) if len(ret) == 0 { panic("no return value specified for Has") @@ -62,17 +65,17 @@ func (_m *TxnKV) Has(key string) (bool, error) { var r0 bool var r1 error - if rf, ok := ret.Get(0).(func(string) (bool, error)); ok { - return rf(key) + if rf, ok := ret.Get(0).(func(context.Context, string) (bool, error)); ok { + return rf(ctx, key) } - if rf, ok := ret.Get(0).(func(string) bool); ok { - r0 = rf(key) + if rf, ok := ret.Get(0).(func(context.Context, string) bool); ok { + r0 = rf(ctx, key) } else { r0 = ret.Get(0).(bool) } - if rf, ok := ret.Get(1).(func(string) error); ok { - r1 = rf(key) + if rf, ok := ret.Get(1).(func(context.Context, string) error); ok { + r1 = rf(ctx, key) } else { r1 = ret.Error(1) } @@ -86,14 +89,15 @@ type TxnKV_Has_Call struct { } // Has is a helper method to define mock.On call +// - ctx context.Context // - key string -func (_e *TxnKV_Expecter) Has(key interface{}) *TxnKV_Has_Call { - return &TxnKV_Has_Call{Call: _e.mock.On("Has", key)} +func (_e *TxnKV_Expecter) Has(ctx interface{}, key interface{}) *TxnKV_Has_Call { + return &TxnKV_Has_Call{Call: _e.mock.On("Has", ctx, key)} } -func (_c *TxnKV_Has_Call) Run(run func(key string)) *TxnKV_Has_Call { +func (_c *TxnKV_Has_Call) Run(run func(ctx context.Context, key string)) *TxnKV_Has_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string)) + run(args[0].(context.Context), args[1].(string)) }) return _c } @@ -103,14 +107,14 @@ func (_c *TxnKV_Has_Call) Return(_a0 bool, _a1 error) *TxnKV_Has_Call { return _c } -func (_c *TxnKV_Has_Call) RunAndReturn(run func(string) (bool, error)) *TxnKV_Has_Call { +func (_c *TxnKV_Has_Call) RunAndReturn(run func(context.Context, string) (bool, error)) *TxnKV_Has_Call { _c.Call.Return(run) return _c } -// HasPrefix provides a mock function with given fields: prefix -func (_m *TxnKV) HasPrefix(prefix string) (bool, error) { - ret := _m.Called(prefix) +// HasPrefix provides a mock function with given fields: ctx, prefix +func (_m *TxnKV) HasPrefix(ctx context.Context, prefix string) (bool, error) { + ret := _m.Called(ctx, prefix) if len(ret) == 0 { panic("no return value specified for HasPrefix") @@ -118,17 +122,17 @@ func (_m *TxnKV) HasPrefix(prefix string) (bool, error) { var r0 bool var r1 error - if rf, ok := ret.Get(0).(func(string) (bool, error)); ok { - return rf(prefix) + if rf, ok := ret.Get(0).(func(context.Context, string) (bool, error)); ok { + return rf(ctx, prefix) } - if rf, ok := ret.Get(0).(func(string) bool); ok { - r0 = rf(prefix) + if rf, ok := ret.Get(0).(func(context.Context, string) bool); ok { + r0 = rf(ctx, prefix) } else { r0 = ret.Get(0).(bool) } - if rf, ok := ret.Get(1).(func(string) error); ok { - r1 = rf(prefix) + if rf, ok := ret.Get(1).(func(context.Context, string) error); ok { + r1 = rf(ctx, prefix) } else { r1 = ret.Error(1) } @@ -142,14 +146,15 @@ type TxnKV_HasPrefix_Call struct { } // HasPrefix is a helper method to define mock.On call +// - ctx context.Context // - prefix string -func (_e *TxnKV_Expecter) HasPrefix(prefix interface{}) *TxnKV_HasPrefix_Call { - return &TxnKV_HasPrefix_Call{Call: _e.mock.On("HasPrefix", prefix)} +func (_e *TxnKV_Expecter) HasPrefix(ctx interface{}, prefix interface{}) *TxnKV_HasPrefix_Call { + return &TxnKV_HasPrefix_Call{Call: _e.mock.On("HasPrefix", ctx, prefix)} } -func (_c *TxnKV_HasPrefix_Call) Run(run func(prefix string)) *TxnKV_HasPrefix_Call { +func (_c *TxnKV_HasPrefix_Call) Run(run func(ctx context.Context, prefix string)) *TxnKV_HasPrefix_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string)) + run(args[0].(context.Context), args[1].(string)) }) return _c } @@ -159,14 +164,14 @@ func (_c *TxnKV_HasPrefix_Call) Return(_a0 bool, _a1 error) *TxnKV_HasPrefix_Cal return _c } -func (_c *TxnKV_HasPrefix_Call) RunAndReturn(run func(string) (bool, error)) *TxnKV_HasPrefix_Call { +func (_c *TxnKV_HasPrefix_Call) RunAndReturn(run func(context.Context, string) (bool, error)) *TxnKV_HasPrefix_Call { _c.Call.Return(run) return _c } -// Load provides a mock function with given fields: key -func (_m *TxnKV) Load(key string) (string, error) { - ret := _m.Called(key) +// Load provides a mock function with given fields: ctx, key +func (_m *TxnKV) Load(ctx context.Context, key string) (string, error) { + ret := _m.Called(ctx, key) if len(ret) == 0 { panic("no return value specified for Load") @@ -174,17 +179,17 @@ func (_m *TxnKV) Load(key string) (string, error) { var r0 string var r1 error - if rf, ok := ret.Get(0).(func(string) (string, error)); ok { - return rf(key) + if rf, ok := ret.Get(0).(func(context.Context, string) (string, error)); ok { + return rf(ctx, key) } - if rf, ok := ret.Get(0).(func(string) string); ok { - r0 = rf(key) + if rf, ok := ret.Get(0).(func(context.Context, string) string); ok { + r0 = rf(ctx, key) } else { r0 = ret.Get(0).(string) } - if rf, ok := ret.Get(1).(func(string) error); ok { - r1 = rf(key) + if rf, ok := ret.Get(1).(func(context.Context, string) error); ok { + r1 = rf(ctx, key) } else { r1 = ret.Error(1) } @@ -198,14 +203,15 @@ type TxnKV_Load_Call struct { } // Load is a helper method to define mock.On call +// - ctx context.Context // - key string -func (_e *TxnKV_Expecter) Load(key interface{}) *TxnKV_Load_Call { - return &TxnKV_Load_Call{Call: _e.mock.On("Load", key)} +func (_e *TxnKV_Expecter) Load(ctx interface{}, key interface{}) *TxnKV_Load_Call { + return &TxnKV_Load_Call{Call: _e.mock.On("Load", ctx, key)} } -func (_c *TxnKV_Load_Call) Run(run func(key string)) *TxnKV_Load_Call { +func (_c *TxnKV_Load_Call) Run(run func(ctx context.Context, key string)) *TxnKV_Load_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string)) + run(args[0].(context.Context), args[1].(string)) }) return _c } @@ -215,14 +221,14 @@ func (_c *TxnKV_Load_Call) Return(_a0 string, _a1 error) *TxnKV_Load_Call { return _c } -func (_c *TxnKV_Load_Call) RunAndReturn(run func(string) (string, error)) *TxnKV_Load_Call { +func (_c *TxnKV_Load_Call) RunAndReturn(run func(context.Context, string) (string, error)) *TxnKV_Load_Call { _c.Call.Return(run) return _c } -// LoadWithPrefix provides a mock function with given fields: key -func (_m *TxnKV) LoadWithPrefix(key string) ([]string, []string, error) { - ret := _m.Called(key) +// LoadWithPrefix provides a mock function with given fields: ctx, key +func (_m *TxnKV) LoadWithPrefix(ctx context.Context, key string) ([]string, []string, error) { + ret := _m.Called(ctx, key) if len(ret) == 0 { panic("no return value specified for LoadWithPrefix") @@ -231,27 +237,27 @@ func (_m *TxnKV) LoadWithPrefix(key string) ([]string, []string, error) { var r0 []string var r1 []string var r2 error - if rf, ok := ret.Get(0).(func(string) ([]string, []string, error)); ok { - return rf(key) + if rf, ok := ret.Get(0).(func(context.Context, string) ([]string, []string, error)); ok { + return rf(ctx, key) } - if rf, ok := ret.Get(0).(func(string) []string); ok { - r0 = rf(key) + if rf, ok := ret.Get(0).(func(context.Context, string) []string); ok { + r0 = rf(ctx, key) } else { if ret.Get(0) != nil { r0 = ret.Get(0).([]string) } } - if rf, ok := ret.Get(1).(func(string) []string); ok { - r1 = rf(key) + if rf, ok := ret.Get(1).(func(context.Context, string) []string); ok { + r1 = rf(ctx, key) } else { if ret.Get(1) != nil { r1 = ret.Get(1).([]string) } } - if rf, ok := ret.Get(2).(func(string) error); ok { - r2 = rf(key) + if rf, ok := ret.Get(2).(func(context.Context, string) error); ok { + r2 = rf(ctx, key) } else { r2 = ret.Error(2) } @@ -265,14 +271,15 @@ type TxnKV_LoadWithPrefix_Call struct { } // LoadWithPrefix is a helper method to define mock.On call +// - ctx context.Context // - key string -func (_e *TxnKV_Expecter) LoadWithPrefix(key interface{}) *TxnKV_LoadWithPrefix_Call { - return &TxnKV_LoadWithPrefix_Call{Call: _e.mock.On("LoadWithPrefix", key)} +func (_e *TxnKV_Expecter) LoadWithPrefix(ctx interface{}, key interface{}) *TxnKV_LoadWithPrefix_Call { + return &TxnKV_LoadWithPrefix_Call{Call: _e.mock.On("LoadWithPrefix", ctx, key)} } -func (_c *TxnKV_LoadWithPrefix_Call) Run(run func(key string)) *TxnKV_LoadWithPrefix_Call { +func (_c *TxnKV_LoadWithPrefix_Call) Run(run func(ctx context.Context, key string)) *TxnKV_LoadWithPrefix_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string)) + run(args[0].(context.Context), args[1].(string)) }) return _c } @@ -282,14 +289,14 @@ func (_c *TxnKV_LoadWithPrefix_Call) Return(_a0 []string, _a1 []string, _a2 erro return _c } -func (_c *TxnKV_LoadWithPrefix_Call) RunAndReturn(run func(string) ([]string, []string, error)) *TxnKV_LoadWithPrefix_Call { +func (_c *TxnKV_LoadWithPrefix_Call) RunAndReturn(run func(context.Context, string) ([]string, []string, error)) *TxnKV_LoadWithPrefix_Call { _c.Call.Return(run) return _c } -// MultiLoad provides a mock function with given fields: keys -func (_m *TxnKV) MultiLoad(keys []string) ([]string, error) { - ret := _m.Called(keys) +// MultiLoad provides a mock function with given fields: ctx, keys +func (_m *TxnKV) MultiLoad(ctx context.Context, keys []string) ([]string, error) { + ret := _m.Called(ctx, keys) if len(ret) == 0 { panic("no return value specified for MultiLoad") @@ -297,19 +304,19 @@ func (_m *TxnKV) MultiLoad(keys []string) ([]string, error) { var r0 []string var r1 error - if rf, ok := ret.Get(0).(func([]string) ([]string, error)); ok { - return rf(keys) + if rf, ok := ret.Get(0).(func(context.Context, []string) ([]string, error)); ok { + return rf(ctx, keys) } - if rf, ok := ret.Get(0).(func([]string) []string); ok { - r0 = rf(keys) + if rf, ok := ret.Get(0).(func(context.Context, []string) []string); ok { + r0 = rf(ctx, keys) } else { if ret.Get(0) != nil { r0 = ret.Get(0).([]string) } } - if rf, ok := ret.Get(1).(func([]string) error); ok { - r1 = rf(keys) + if rf, ok := ret.Get(1).(func(context.Context, []string) error); ok { + r1 = rf(ctx, keys) } else { r1 = ret.Error(1) } @@ -323,14 +330,15 @@ type TxnKV_MultiLoad_Call struct { } // MultiLoad is a helper method to define mock.On call +// - ctx context.Context // - keys []string -func (_e *TxnKV_Expecter) MultiLoad(keys interface{}) *TxnKV_MultiLoad_Call { - return &TxnKV_MultiLoad_Call{Call: _e.mock.On("MultiLoad", keys)} +func (_e *TxnKV_Expecter) MultiLoad(ctx interface{}, keys interface{}) *TxnKV_MultiLoad_Call { + return &TxnKV_MultiLoad_Call{Call: _e.mock.On("MultiLoad", ctx, keys)} } -func (_c *TxnKV_MultiLoad_Call) Run(run func(keys []string)) *TxnKV_MultiLoad_Call { +func (_c *TxnKV_MultiLoad_Call) Run(run func(ctx context.Context, keys []string)) *TxnKV_MultiLoad_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].([]string)) + run(args[0].(context.Context), args[1].([]string)) }) return _c } @@ -340,22 +348,22 @@ func (_c *TxnKV_MultiLoad_Call) Return(_a0 []string, _a1 error) *TxnKV_MultiLoad return _c } -func (_c *TxnKV_MultiLoad_Call) RunAndReturn(run func([]string) ([]string, error)) *TxnKV_MultiLoad_Call { +func (_c *TxnKV_MultiLoad_Call) RunAndReturn(run func(context.Context, []string) ([]string, error)) *TxnKV_MultiLoad_Call { _c.Call.Return(run) return _c } -// MultiRemove provides a mock function with given fields: keys -func (_m *TxnKV) MultiRemove(keys []string) error { - ret := _m.Called(keys) +// MultiRemove provides a mock function with given fields: ctx, keys +func (_m *TxnKV) MultiRemove(ctx context.Context, keys []string) error { + ret := _m.Called(ctx, keys) if len(ret) == 0 { panic("no return value specified for MultiRemove") } var r0 error - if rf, ok := ret.Get(0).(func([]string) error); ok { - r0 = rf(keys) + if rf, ok := ret.Get(0).(func(context.Context, []string) error); ok { + r0 = rf(ctx, keys) } else { r0 = ret.Error(0) } @@ -369,14 +377,15 @@ type TxnKV_MultiRemove_Call struct { } // MultiRemove is a helper method to define mock.On call +// - ctx context.Context // - keys []string -func (_e *TxnKV_Expecter) MultiRemove(keys interface{}) *TxnKV_MultiRemove_Call { - return &TxnKV_MultiRemove_Call{Call: _e.mock.On("MultiRemove", keys)} +func (_e *TxnKV_Expecter) MultiRemove(ctx interface{}, keys interface{}) *TxnKV_MultiRemove_Call { + return &TxnKV_MultiRemove_Call{Call: _e.mock.On("MultiRemove", ctx, keys)} } -func (_c *TxnKV_MultiRemove_Call) Run(run func(keys []string)) *TxnKV_MultiRemove_Call { +func (_c *TxnKV_MultiRemove_Call) Run(run func(ctx context.Context, keys []string)) *TxnKV_MultiRemove_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].([]string)) + run(args[0].(context.Context), args[1].([]string)) }) return _c } @@ -386,22 +395,22 @@ func (_c *TxnKV_MultiRemove_Call) Return(_a0 error) *TxnKV_MultiRemove_Call { return _c } -func (_c *TxnKV_MultiRemove_Call) RunAndReturn(run func([]string) error) *TxnKV_MultiRemove_Call { +func (_c *TxnKV_MultiRemove_Call) RunAndReturn(run func(context.Context, []string) error) *TxnKV_MultiRemove_Call { _c.Call.Return(run) return _c } -// MultiSave provides a mock function with given fields: kvs -func (_m *TxnKV) MultiSave(kvs map[string]string) error { - ret := _m.Called(kvs) +// MultiSave provides a mock function with given fields: ctx, kvs +func (_m *TxnKV) MultiSave(ctx context.Context, kvs map[string]string) error { + ret := _m.Called(ctx, kvs) if len(ret) == 0 { panic("no return value specified for MultiSave") } var r0 error - if rf, ok := ret.Get(0).(func(map[string]string) error); ok { - r0 = rf(kvs) + if rf, ok := ret.Get(0).(func(context.Context, map[string]string) error); ok { + r0 = rf(ctx, kvs) } else { r0 = ret.Error(0) } @@ -415,14 +424,15 @@ type TxnKV_MultiSave_Call struct { } // MultiSave is a helper method to define mock.On call +// - ctx context.Context // - kvs map[string]string -func (_e *TxnKV_Expecter) MultiSave(kvs interface{}) *TxnKV_MultiSave_Call { - return &TxnKV_MultiSave_Call{Call: _e.mock.On("MultiSave", kvs)} +func (_e *TxnKV_Expecter) MultiSave(ctx interface{}, kvs interface{}) *TxnKV_MultiSave_Call { + return &TxnKV_MultiSave_Call{Call: _e.mock.On("MultiSave", ctx, kvs)} } -func (_c *TxnKV_MultiSave_Call) Run(run func(kvs map[string]string)) *TxnKV_MultiSave_Call { +func (_c *TxnKV_MultiSave_Call) Run(run func(ctx context.Context, kvs map[string]string)) *TxnKV_MultiSave_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(map[string]string)) + run(args[0].(context.Context), args[1].(map[string]string)) }) return _c } @@ -432,19 +442,19 @@ func (_c *TxnKV_MultiSave_Call) Return(_a0 error) *TxnKV_MultiSave_Call { return _c } -func (_c *TxnKV_MultiSave_Call) RunAndReturn(run func(map[string]string) error) *TxnKV_MultiSave_Call { +func (_c *TxnKV_MultiSave_Call) RunAndReturn(run func(context.Context, map[string]string) error) *TxnKV_MultiSave_Call { _c.Call.Return(run) return _c } -// MultiSaveAndRemove provides a mock function with given fields: saves, removals, preds -func (_m *TxnKV) MultiSaveAndRemove(saves map[string]string, removals []string, preds ...predicates.Predicate) error { +// MultiSaveAndRemove provides a mock function with given fields: ctx, saves, removals, preds +func (_m *TxnKV) MultiSaveAndRemove(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error { _va := make([]interface{}, len(preds)) for _i := range preds { _va[_i] = preds[_i] } var _ca []interface{} - _ca = append(_ca, saves, removals) + _ca = append(_ca, ctx, saves, removals) _ca = append(_ca, _va...) ret := _m.Called(_ca...) @@ -453,8 +463,8 @@ func (_m *TxnKV) MultiSaveAndRemove(saves map[string]string, removals []string, } var r0 error - if rf, ok := ret.Get(0).(func(map[string]string, []string, ...predicates.Predicate) error); ok { - r0 = rf(saves, removals, preds...) + if rf, ok := ret.Get(0).(func(context.Context, map[string]string, []string, ...predicates.Predicate) error); ok { + r0 = rf(ctx, saves, removals, preds...) } else { r0 = ret.Error(0) } @@ -468,23 +478,24 @@ type TxnKV_MultiSaveAndRemove_Call struct { } // MultiSaveAndRemove is a helper method to define mock.On call +// - ctx context.Context // - saves map[string]string // - removals []string // - preds ...predicates.Predicate -func (_e *TxnKV_Expecter) MultiSaveAndRemove(saves interface{}, removals interface{}, preds ...interface{}) *TxnKV_MultiSaveAndRemove_Call { +func (_e *TxnKV_Expecter) MultiSaveAndRemove(ctx interface{}, saves interface{}, removals interface{}, preds ...interface{}) *TxnKV_MultiSaveAndRemove_Call { return &TxnKV_MultiSaveAndRemove_Call{Call: _e.mock.On("MultiSaveAndRemove", - append([]interface{}{saves, removals}, preds...)...)} + append([]interface{}{ctx, saves, removals}, preds...)...)} } -func (_c *TxnKV_MultiSaveAndRemove_Call) Run(run func(saves map[string]string, removals []string, preds ...predicates.Predicate)) *TxnKV_MultiSaveAndRemove_Call { +func (_c *TxnKV_MultiSaveAndRemove_Call) Run(run func(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate)) *TxnKV_MultiSaveAndRemove_Call { _c.Call.Run(func(args mock.Arguments) { - variadicArgs := make([]predicates.Predicate, len(args)-2) - for i, a := range args[2:] { + variadicArgs := make([]predicates.Predicate, len(args)-3) + for i, a := range args[3:] { if a != nil { variadicArgs[i] = a.(predicates.Predicate) } } - run(args[0].(map[string]string), args[1].([]string), variadicArgs...) + run(args[0].(context.Context), args[1].(map[string]string), args[2].([]string), variadicArgs...) }) return _c } @@ -494,19 +505,19 @@ func (_c *TxnKV_MultiSaveAndRemove_Call) Return(_a0 error) *TxnKV_MultiSaveAndRe return _c } -func (_c *TxnKV_MultiSaveAndRemove_Call) RunAndReturn(run func(map[string]string, []string, ...predicates.Predicate) error) *TxnKV_MultiSaveAndRemove_Call { +func (_c *TxnKV_MultiSaveAndRemove_Call) RunAndReturn(run func(context.Context, map[string]string, []string, ...predicates.Predicate) error) *TxnKV_MultiSaveAndRemove_Call { _c.Call.Return(run) return _c } -// MultiSaveAndRemoveWithPrefix provides a mock function with given fields: saves, removals, preds -func (_m *TxnKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals []string, preds ...predicates.Predicate) error { +// MultiSaveAndRemoveWithPrefix provides a mock function with given fields: ctx, saves, removals, preds +func (_m *TxnKV) MultiSaveAndRemoveWithPrefix(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error { _va := make([]interface{}, len(preds)) for _i := range preds { _va[_i] = preds[_i] } var _ca []interface{} - _ca = append(_ca, saves, removals) + _ca = append(_ca, ctx, saves, removals) _ca = append(_ca, _va...) ret := _m.Called(_ca...) @@ -515,8 +526,8 @@ func (_m *TxnKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals } var r0 error - if rf, ok := ret.Get(0).(func(map[string]string, []string, ...predicates.Predicate) error); ok { - r0 = rf(saves, removals, preds...) + if rf, ok := ret.Get(0).(func(context.Context, map[string]string, []string, ...predicates.Predicate) error); ok { + r0 = rf(ctx, saves, removals, preds...) } else { r0 = ret.Error(0) } @@ -530,23 +541,24 @@ type TxnKV_MultiSaveAndRemoveWithPrefix_Call struct { } // MultiSaveAndRemoveWithPrefix is a helper method to define mock.On call +// - ctx context.Context // - saves map[string]string // - removals []string // - preds ...predicates.Predicate -func (_e *TxnKV_Expecter) MultiSaveAndRemoveWithPrefix(saves interface{}, removals interface{}, preds ...interface{}) *TxnKV_MultiSaveAndRemoveWithPrefix_Call { +func (_e *TxnKV_Expecter) MultiSaveAndRemoveWithPrefix(ctx interface{}, saves interface{}, removals interface{}, preds ...interface{}) *TxnKV_MultiSaveAndRemoveWithPrefix_Call { return &TxnKV_MultiSaveAndRemoveWithPrefix_Call{Call: _e.mock.On("MultiSaveAndRemoveWithPrefix", - append([]interface{}{saves, removals}, preds...)...)} + append([]interface{}{ctx, saves, removals}, preds...)...)} } -func (_c *TxnKV_MultiSaveAndRemoveWithPrefix_Call) Run(run func(saves map[string]string, removals []string, preds ...predicates.Predicate)) *TxnKV_MultiSaveAndRemoveWithPrefix_Call { +func (_c *TxnKV_MultiSaveAndRemoveWithPrefix_Call) Run(run func(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate)) *TxnKV_MultiSaveAndRemoveWithPrefix_Call { _c.Call.Run(func(args mock.Arguments) { - variadicArgs := make([]predicates.Predicate, len(args)-2) - for i, a := range args[2:] { + variadicArgs := make([]predicates.Predicate, len(args)-3) + for i, a := range args[3:] { if a != nil { variadicArgs[i] = a.(predicates.Predicate) } } - run(args[0].(map[string]string), args[1].([]string), variadicArgs...) + run(args[0].(context.Context), args[1].(map[string]string), args[2].([]string), variadicArgs...) }) return _c } @@ -556,22 +568,22 @@ func (_c *TxnKV_MultiSaveAndRemoveWithPrefix_Call) Return(_a0 error) *TxnKV_Mult return _c } -func (_c *TxnKV_MultiSaveAndRemoveWithPrefix_Call) RunAndReturn(run func(map[string]string, []string, ...predicates.Predicate) error) *TxnKV_MultiSaveAndRemoveWithPrefix_Call { +func (_c *TxnKV_MultiSaveAndRemoveWithPrefix_Call) RunAndReturn(run func(context.Context, map[string]string, []string, ...predicates.Predicate) error) *TxnKV_MultiSaveAndRemoveWithPrefix_Call { _c.Call.Return(run) return _c } -// Remove provides a mock function with given fields: key -func (_m *TxnKV) Remove(key string) error { - ret := _m.Called(key) +// Remove provides a mock function with given fields: ctx, key +func (_m *TxnKV) Remove(ctx context.Context, key string) error { + ret := _m.Called(ctx, key) if len(ret) == 0 { panic("no return value specified for Remove") } var r0 error - if rf, ok := ret.Get(0).(func(string) error); ok { - r0 = rf(key) + if rf, ok := ret.Get(0).(func(context.Context, string) error); ok { + r0 = rf(ctx, key) } else { r0 = ret.Error(0) } @@ -585,14 +597,15 @@ type TxnKV_Remove_Call struct { } // Remove is a helper method to define mock.On call +// - ctx context.Context // - key string -func (_e *TxnKV_Expecter) Remove(key interface{}) *TxnKV_Remove_Call { - return &TxnKV_Remove_Call{Call: _e.mock.On("Remove", key)} +func (_e *TxnKV_Expecter) Remove(ctx interface{}, key interface{}) *TxnKV_Remove_Call { + return &TxnKV_Remove_Call{Call: _e.mock.On("Remove", ctx, key)} } -func (_c *TxnKV_Remove_Call) Run(run func(key string)) *TxnKV_Remove_Call { +func (_c *TxnKV_Remove_Call) Run(run func(ctx context.Context, key string)) *TxnKV_Remove_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string)) + run(args[0].(context.Context), args[1].(string)) }) return _c } @@ -602,22 +615,22 @@ func (_c *TxnKV_Remove_Call) Return(_a0 error) *TxnKV_Remove_Call { return _c } -func (_c *TxnKV_Remove_Call) RunAndReturn(run func(string) error) *TxnKV_Remove_Call { +func (_c *TxnKV_Remove_Call) RunAndReturn(run func(context.Context, string) error) *TxnKV_Remove_Call { _c.Call.Return(run) return _c } -// RemoveWithPrefix provides a mock function with given fields: key -func (_m *TxnKV) RemoveWithPrefix(key string) error { - ret := _m.Called(key) +// RemoveWithPrefix provides a mock function with given fields: ctx, key +func (_m *TxnKV) RemoveWithPrefix(ctx context.Context, key string) error { + ret := _m.Called(ctx, key) if len(ret) == 0 { panic("no return value specified for RemoveWithPrefix") } var r0 error - if rf, ok := ret.Get(0).(func(string) error); ok { - r0 = rf(key) + if rf, ok := ret.Get(0).(func(context.Context, string) error); ok { + r0 = rf(ctx, key) } else { r0 = ret.Error(0) } @@ -631,14 +644,15 @@ type TxnKV_RemoveWithPrefix_Call struct { } // RemoveWithPrefix is a helper method to define mock.On call +// - ctx context.Context // - key string -func (_e *TxnKV_Expecter) RemoveWithPrefix(key interface{}) *TxnKV_RemoveWithPrefix_Call { - return &TxnKV_RemoveWithPrefix_Call{Call: _e.mock.On("RemoveWithPrefix", key)} +func (_e *TxnKV_Expecter) RemoveWithPrefix(ctx interface{}, key interface{}) *TxnKV_RemoveWithPrefix_Call { + return &TxnKV_RemoveWithPrefix_Call{Call: _e.mock.On("RemoveWithPrefix", ctx, key)} } -func (_c *TxnKV_RemoveWithPrefix_Call) Run(run func(key string)) *TxnKV_RemoveWithPrefix_Call { +func (_c *TxnKV_RemoveWithPrefix_Call) Run(run func(ctx context.Context, key string)) *TxnKV_RemoveWithPrefix_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string)) + run(args[0].(context.Context), args[1].(string)) }) return _c } @@ -648,22 +662,22 @@ func (_c *TxnKV_RemoveWithPrefix_Call) Return(_a0 error) *TxnKV_RemoveWithPrefix return _c } -func (_c *TxnKV_RemoveWithPrefix_Call) RunAndReturn(run func(string) error) *TxnKV_RemoveWithPrefix_Call { +func (_c *TxnKV_RemoveWithPrefix_Call) RunAndReturn(run func(context.Context, string) error) *TxnKV_RemoveWithPrefix_Call { _c.Call.Return(run) return _c } -// Save provides a mock function with given fields: key, value -func (_m *TxnKV) Save(key string, value string) error { - ret := _m.Called(key, value) +// Save provides a mock function with given fields: ctx, key, value +func (_m *TxnKV) Save(ctx context.Context, key string, value string) error { + ret := _m.Called(ctx, key, value) if len(ret) == 0 { panic("no return value specified for Save") } var r0 error - if rf, ok := ret.Get(0).(func(string, string) error); ok { - r0 = rf(key, value) + if rf, ok := ret.Get(0).(func(context.Context, string, string) error); ok { + r0 = rf(ctx, key, value) } else { r0 = ret.Error(0) } @@ -677,15 +691,16 @@ type TxnKV_Save_Call struct { } // Save is a helper method to define mock.On call +// - ctx context.Context // - key string // - value string -func (_e *TxnKV_Expecter) Save(key interface{}, value interface{}) *TxnKV_Save_Call { - return &TxnKV_Save_Call{Call: _e.mock.On("Save", key, value)} +func (_e *TxnKV_Expecter) Save(ctx interface{}, key interface{}, value interface{}) *TxnKV_Save_Call { + return &TxnKV_Save_Call{Call: _e.mock.On("Save", ctx, key, value)} } -func (_c *TxnKV_Save_Call) Run(run func(key string, value string)) *TxnKV_Save_Call { +func (_c *TxnKV_Save_Call) Run(run func(ctx context.Context, key string, value string)) *TxnKV_Save_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string), args[1].(string)) + run(args[0].(context.Context), args[1].(string), args[2].(string)) }) return _c } @@ -695,7 +710,7 @@ func (_c *TxnKV_Save_Call) Return(_a0 error) *TxnKV_Save_Call { return _c } -func (_c *TxnKV_Save_Call) RunAndReturn(run func(string, string) error) *TxnKV_Save_Call { +func (_c *TxnKV_Save_Call) RunAndReturn(run func(context.Context, string, string) error) *TxnKV_Save_Call { _c.Call.Return(run) return _c } diff --git a/internal/kv/mocks/watch_kv.go b/internal/kv/mocks/watch_kv.go index af3328ef5dbe2..60a38a5f97bb0 100644 --- a/internal/kv/mocks/watch_kv.go +++ b/internal/kv/mocks/watch_kv.go @@ -3,6 +3,8 @@ package mocks import ( + context "context" + clientv3 "go.etcd.io/etcd/client/v3" mock "github.com/stretchr/testify/mock" @@ -55,9 +57,9 @@ func (_c *WatchKV_Close_Call) RunAndReturn(run func()) *WatchKV_Close_Call { return _c } -// CompareVersionAndSwap provides a mock function with given fields: key, version, target -func (_m *WatchKV) CompareVersionAndSwap(key string, version int64, target string) (bool, error) { - ret := _m.Called(key, version, target) +// CompareVersionAndSwap provides a mock function with given fields: ctx, key, version, target +func (_m *WatchKV) CompareVersionAndSwap(ctx context.Context, key string, version int64, target string) (bool, error) { + ret := _m.Called(ctx, key, version, target) if len(ret) == 0 { panic("no return value specified for CompareVersionAndSwap") @@ -65,17 +67,17 @@ func (_m *WatchKV) CompareVersionAndSwap(key string, version int64, target strin var r0 bool var r1 error - if rf, ok := ret.Get(0).(func(string, int64, string) (bool, error)); ok { - return rf(key, version, target) + if rf, ok := ret.Get(0).(func(context.Context, string, int64, string) (bool, error)); ok { + return rf(ctx, key, version, target) } - if rf, ok := ret.Get(0).(func(string, int64, string) bool); ok { - r0 = rf(key, version, target) + if rf, ok := ret.Get(0).(func(context.Context, string, int64, string) bool); ok { + r0 = rf(ctx, key, version, target) } else { r0 = ret.Get(0).(bool) } - if rf, ok := ret.Get(1).(func(string, int64, string) error); ok { - r1 = rf(key, version, target) + if rf, ok := ret.Get(1).(func(context.Context, string, int64, string) error); ok { + r1 = rf(ctx, key, version, target) } else { r1 = ret.Error(1) } @@ -89,16 +91,17 @@ type WatchKV_CompareVersionAndSwap_Call struct { } // CompareVersionAndSwap is a helper method to define mock.On call +// - ctx context.Context // - key string // - version int64 // - target string -func (_e *WatchKV_Expecter) CompareVersionAndSwap(key interface{}, version interface{}, target interface{}) *WatchKV_CompareVersionAndSwap_Call { - return &WatchKV_CompareVersionAndSwap_Call{Call: _e.mock.On("CompareVersionAndSwap", key, version, target)} +func (_e *WatchKV_Expecter) CompareVersionAndSwap(ctx interface{}, key interface{}, version interface{}, target interface{}) *WatchKV_CompareVersionAndSwap_Call { + return &WatchKV_CompareVersionAndSwap_Call{Call: _e.mock.On("CompareVersionAndSwap", ctx, key, version, target)} } -func (_c *WatchKV_CompareVersionAndSwap_Call) Run(run func(key string, version int64, target string)) *WatchKV_CompareVersionAndSwap_Call { +func (_c *WatchKV_CompareVersionAndSwap_Call) Run(run func(ctx context.Context, key string, version int64, target string)) *WatchKV_CompareVersionAndSwap_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string), args[1].(int64), args[2].(string)) + run(args[0].(context.Context), args[1].(string), args[2].(int64), args[3].(string)) }) return _c } @@ -108,7 +111,7 @@ func (_c *WatchKV_CompareVersionAndSwap_Call) Return(_a0 bool, _a1 error) *Watch return _c } -func (_c *WatchKV_CompareVersionAndSwap_Call) RunAndReturn(run func(string, int64, string) (bool, error)) *WatchKV_CompareVersionAndSwap_Call { +func (_c *WatchKV_CompareVersionAndSwap_Call) RunAndReturn(run func(context.Context, string, int64, string) (bool, error)) *WatchKV_CompareVersionAndSwap_Call { _c.Call.Return(run) return _c } @@ -159,9 +162,9 @@ func (_c *WatchKV_GetPath_Call) RunAndReturn(run func(string) string) *WatchKV_G return _c } -// Has provides a mock function with given fields: key -func (_m *WatchKV) Has(key string) (bool, error) { - ret := _m.Called(key) +// Has provides a mock function with given fields: ctx, key +func (_m *WatchKV) Has(ctx context.Context, key string) (bool, error) { + ret := _m.Called(ctx, key) if len(ret) == 0 { panic("no return value specified for Has") @@ -169,17 +172,17 @@ func (_m *WatchKV) Has(key string) (bool, error) { var r0 bool var r1 error - if rf, ok := ret.Get(0).(func(string) (bool, error)); ok { - return rf(key) + if rf, ok := ret.Get(0).(func(context.Context, string) (bool, error)); ok { + return rf(ctx, key) } - if rf, ok := ret.Get(0).(func(string) bool); ok { - r0 = rf(key) + if rf, ok := ret.Get(0).(func(context.Context, string) bool); ok { + r0 = rf(ctx, key) } else { r0 = ret.Get(0).(bool) } - if rf, ok := ret.Get(1).(func(string) error); ok { - r1 = rf(key) + if rf, ok := ret.Get(1).(func(context.Context, string) error); ok { + r1 = rf(ctx, key) } else { r1 = ret.Error(1) } @@ -193,14 +196,15 @@ type WatchKV_Has_Call struct { } // Has is a helper method to define mock.On call +// - ctx context.Context // - key string -func (_e *WatchKV_Expecter) Has(key interface{}) *WatchKV_Has_Call { - return &WatchKV_Has_Call{Call: _e.mock.On("Has", key)} +func (_e *WatchKV_Expecter) Has(ctx interface{}, key interface{}) *WatchKV_Has_Call { + return &WatchKV_Has_Call{Call: _e.mock.On("Has", ctx, key)} } -func (_c *WatchKV_Has_Call) Run(run func(key string)) *WatchKV_Has_Call { +func (_c *WatchKV_Has_Call) Run(run func(ctx context.Context, key string)) *WatchKV_Has_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string)) + run(args[0].(context.Context), args[1].(string)) }) return _c } @@ -210,14 +214,14 @@ func (_c *WatchKV_Has_Call) Return(_a0 bool, _a1 error) *WatchKV_Has_Call { return _c } -func (_c *WatchKV_Has_Call) RunAndReturn(run func(string) (bool, error)) *WatchKV_Has_Call { +func (_c *WatchKV_Has_Call) RunAndReturn(run func(context.Context, string) (bool, error)) *WatchKV_Has_Call { _c.Call.Return(run) return _c } -// HasPrefix provides a mock function with given fields: prefix -func (_m *WatchKV) HasPrefix(prefix string) (bool, error) { - ret := _m.Called(prefix) +// HasPrefix provides a mock function with given fields: ctx, prefix +func (_m *WatchKV) HasPrefix(ctx context.Context, prefix string) (bool, error) { + ret := _m.Called(ctx, prefix) if len(ret) == 0 { panic("no return value specified for HasPrefix") @@ -225,17 +229,17 @@ func (_m *WatchKV) HasPrefix(prefix string) (bool, error) { var r0 bool var r1 error - if rf, ok := ret.Get(0).(func(string) (bool, error)); ok { - return rf(prefix) + if rf, ok := ret.Get(0).(func(context.Context, string) (bool, error)); ok { + return rf(ctx, prefix) } - if rf, ok := ret.Get(0).(func(string) bool); ok { - r0 = rf(prefix) + if rf, ok := ret.Get(0).(func(context.Context, string) bool); ok { + r0 = rf(ctx, prefix) } else { r0 = ret.Get(0).(bool) } - if rf, ok := ret.Get(1).(func(string) error); ok { - r1 = rf(prefix) + if rf, ok := ret.Get(1).(func(context.Context, string) error); ok { + r1 = rf(ctx, prefix) } else { r1 = ret.Error(1) } @@ -249,14 +253,15 @@ type WatchKV_HasPrefix_Call struct { } // HasPrefix is a helper method to define mock.On call +// - ctx context.Context // - prefix string -func (_e *WatchKV_Expecter) HasPrefix(prefix interface{}) *WatchKV_HasPrefix_Call { - return &WatchKV_HasPrefix_Call{Call: _e.mock.On("HasPrefix", prefix)} +func (_e *WatchKV_Expecter) HasPrefix(ctx interface{}, prefix interface{}) *WatchKV_HasPrefix_Call { + return &WatchKV_HasPrefix_Call{Call: _e.mock.On("HasPrefix", ctx, prefix)} } -func (_c *WatchKV_HasPrefix_Call) Run(run func(prefix string)) *WatchKV_HasPrefix_Call { +func (_c *WatchKV_HasPrefix_Call) Run(run func(ctx context.Context, prefix string)) *WatchKV_HasPrefix_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string)) + run(args[0].(context.Context), args[1].(string)) }) return _c } @@ -266,14 +271,14 @@ func (_c *WatchKV_HasPrefix_Call) Return(_a0 bool, _a1 error) *WatchKV_HasPrefix return _c } -func (_c *WatchKV_HasPrefix_Call) RunAndReturn(run func(string) (bool, error)) *WatchKV_HasPrefix_Call { +func (_c *WatchKV_HasPrefix_Call) RunAndReturn(run func(context.Context, string) (bool, error)) *WatchKV_HasPrefix_Call { _c.Call.Return(run) return _c } -// Load provides a mock function with given fields: key -func (_m *WatchKV) Load(key string) (string, error) { - ret := _m.Called(key) +// Load provides a mock function with given fields: ctx, key +func (_m *WatchKV) Load(ctx context.Context, key string) (string, error) { + ret := _m.Called(ctx, key) if len(ret) == 0 { panic("no return value specified for Load") @@ -281,17 +286,17 @@ func (_m *WatchKV) Load(key string) (string, error) { var r0 string var r1 error - if rf, ok := ret.Get(0).(func(string) (string, error)); ok { - return rf(key) + if rf, ok := ret.Get(0).(func(context.Context, string) (string, error)); ok { + return rf(ctx, key) } - if rf, ok := ret.Get(0).(func(string) string); ok { - r0 = rf(key) + if rf, ok := ret.Get(0).(func(context.Context, string) string); ok { + r0 = rf(ctx, key) } else { r0 = ret.Get(0).(string) } - if rf, ok := ret.Get(1).(func(string) error); ok { - r1 = rf(key) + if rf, ok := ret.Get(1).(func(context.Context, string) error); ok { + r1 = rf(ctx, key) } else { r1 = ret.Error(1) } @@ -305,14 +310,15 @@ type WatchKV_Load_Call struct { } // Load is a helper method to define mock.On call +// - ctx context.Context // - key string -func (_e *WatchKV_Expecter) Load(key interface{}) *WatchKV_Load_Call { - return &WatchKV_Load_Call{Call: _e.mock.On("Load", key)} +func (_e *WatchKV_Expecter) Load(ctx interface{}, key interface{}) *WatchKV_Load_Call { + return &WatchKV_Load_Call{Call: _e.mock.On("Load", ctx, key)} } -func (_c *WatchKV_Load_Call) Run(run func(key string)) *WatchKV_Load_Call { +func (_c *WatchKV_Load_Call) Run(run func(ctx context.Context, key string)) *WatchKV_Load_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string)) + run(args[0].(context.Context), args[1].(string)) }) return _c } @@ -322,14 +328,14 @@ func (_c *WatchKV_Load_Call) Return(_a0 string, _a1 error) *WatchKV_Load_Call { return _c } -func (_c *WatchKV_Load_Call) RunAndReturn(run func(string) (string, error)) *WatchKV_Load_Call { +func (_c *WatchKV_Load_Call) RunAndReturn(run func(context.Context, string) (string, error)) *WatchKV_Load_Call { _c.Call.Return(run) return _c } -// LoadWithPrefix provides a mock function with given fields: key -func (_m *WatchKV) LoadWithPrefix(key string) ([]string, []string, error) { - ret := _m.Called(key) +// LoadWithPrefix provides a mock function with given fields: ctx, key +func (_m *WatchKV) LoadWithPrefix(ctx context.Context, key string) ([]string, []string, error) { + ret := _m.Called(ctx, key) if len(ret) == 0 { panic("no return value specified for LoadWithPrefix") @@ -338,27 +344,27 @@ func (_m *WatchKV) LoadWithPrefix(key string) ([]string, []string, error) { var r0 []string var r1 []string var r2 error - if rf, ok := ret.Get(0).(func(string) ([]string, []string, error)); ok { - return rf(key) + if rf, ok := ret.Get(0).(func(context.Context, string) ([]string, []string, error)); ok { + return rf(ctx, key) } - if rf, ok := ret.Get(0).(func(string) []string); ok { - r0 = rf(key) + if rf, ok := ret.Get(0).(func(context.Context, string) []string); ok { + r0 = rf(ctx, key) } else { if ret.Get(0) != nil { r0 = ret.Get(0).([]string) } } - if rf, ok := ret.Get(1).(func(string) []string); ok { - r1 = rf(key) + if rf, ok := ret.Get(1).(func(context.Context, string) []string); ok { + r1 = rf(ctx, key) } else { if ret.Get(1) != nil { r1 = ret.Get(1).([]string) } } - if rf, ok := ret.Get(2).(func(string) error); ok { - r2 = rf(key) + if rf, ok := ret.Get(2).(func(context.Context, string) error); ok { + r2 = rf(ctx, key) } else { r2 = ret.Error(2) } @@ -372,14 +378,15 @@ type WatchKV_LoadWithPrefix_Call struct { } // LoadWithPrefix is a helper method to define mock.On call +// - ctx context.Context // - key string -func (_e *WatchKV_Expecter) LoadWithPrefix(key interface{}) *WatchKV_LoadWithPrefix_Call { - return &WatchKV_LoadWithPrefix_Call{Call: _e.mock.On("LoadWithPrefix", key)} +func (_e *WatchKV_Expecter) LoadWithPrefix(ctx interface{}, key interface{}) *WatchKV_LoadWithPrefix_Call { + return &WatchKV_LoadWithPrefix_Call{Call: _e.mock.On("LoadWithPrefix", ctx, key)} } -func (_c *WatchKV_LoadWithPrefix_Call) Run(run func(key string)) *WatchKV_LoadWithPrefix_Call { +func (_c *WatchKV_LoadWithPrefix_Call) Run(run func(ctx context.Context, key string)) *WatchKV_LoadWithPrefix_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string)) + run(args[0].(context.Context), args[1].(string)) }) return _c } @@ -389,14 +396,14 @@ func (_c *WatchKV_LoadWithPrefix_Call) Return(_a0 []string, _a1 []string, _a2 er return _c } -func (_c *WatchKV_LoadWithPrefix_Call) RunAndReturn(run func(string) ([]string, []string, error)) *WatchKV_LoadWithPrefix_Call { +func (_c *WatchKV_LoadWithPrefix_Call) RunAndReturn(run func(context.Context, string) ([]string, []string, error)) *WatchKV_LoadWithPrefix_Call { _c.Call.Return(run) return _c } -// MultiLoad provides a mock function with given fields: keys -func (_m *WatchKV) MultiLoad(keys []string) ([]string, error) { - ret := _m.Called(keys) +// MultiLoad provides a mock function with given fields: ctx, keys +func (_m *WatchKV) MultiLoad(ctx context.Context, keys []string) ([]string, error) { + ret := _m.Called(ctx, keys) if len(ret) == 0 { panic("no return value specified for MultiLoad") @@ -404,19 +411,19 @@ func (_m *WatchKV) MultiLoad(keys []string) ([]string, error) { var r0 []string var r1 error - if rf, ok := ret.Get(0).(func([]string) ([]string, error)); ok { - return rf(keys) + if rf, ok := ret.Get(0).(func(context.Context, []string) ([]string, error)); ok { + return rf(ctx, keys) } - if rf, ok := ret.Get(0).(func([]string) []string); ok { - r0 = rf(keys) + if rf, ok := ret.Get(0).(func(context.Context, []string) []string); ok { + r0 = rf(ctx, keys) } else { if ret.Get(0) != nil { r0 = ret.Get(0).([]string) } } - if rf, ok := ret.Get(1).(func([]string) error); ok { - r1 = rf(keys) + if rf, ok := ret.Get(1).(func(context.Context, []string) error); ok { + r1 = rf(ctx, keys) } else { r1 = ret.Error(1) } @@ -430,14 +437,15 @@ type WatchKV_MultiLoad_Call struct { } // MultiLoad is a helper method to define mock.On call +// - ctx context.Context // - keys []string -func (_e *WatchKV_Expecter) MultiLoad(keys interface{}) *WatchKV_MultiLoad_Call { - return &WatchKV_MultiLoad_Call{Call: _e.mock.On("MultiLoad", keys)} +func (_e *WatchKV_Expecter) MultiLoad(ctx interface{}, keys interface{}) *WatchKV_MultiLoad_Call { + return &WatchKV_MultiLoad_Call{Call: _e.mock.On("MultiLoad", ctx, keys)} } -func (_c *WatchKV_MultiLoad_Call) Run(run func(keys []string)) *WatchKV_MultiLoad_Call { +func (_c *WatchKV_MultiLoad_Call) Run(run func(ctx context.Context, keys []string)) *WatchKV_MultiLoad_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].([]string)) + run(args[0].(context.Context), args[1].([]string)) }) return _c } @@ -447,22 +455,22 @@ func (_c *WatchKV_MultiLoad_Call) Return(_a0 []string, _a1 error) *WatchKV_Multi return _c } -func (_c *WatchKV_MultiLoad_Call) RunAndReturn(run func([]string) ([]string, error)) *WatchKV_MultiLoad_Call { +func (_c *WatchKV_MultiLoad_Call) RunAndReturn(run func(context.Context, []string) ([]string, error)) *WatchKV_MultiLoad_Call { _c.Call.Return(run) return _c } -// MultiRemove provides a mock function with given fields: keys -func (_m *WatchKV) MultiRemove(keys []string) error { - ret := _m.Called(keys) +// MultiRemove provides a mock function with given fields: ctx, keys +func (_m *WatchKV) MultiRemove(ctx context.Context, keys []string) error { + ret := _m.Called(ctx, keys) if len(ret) == 0 { panic("no return value specified for MultiRemove") } var r0 error - if rf, ok := ret.Get(0).(func([]string) error); ok { - r0 = rf(keys) + if rf, ok := ret.Get(0).(func(context.Context, []string) error); ok { + r0 = rf(ctx, keys) } else { r0 = ret.Error(0) } @@ -476,14 +484,15 @@ type WatchKV_MultiRemove_Call struct { } // MultiRemove is a helper method to define mock.On call +// - ctx context.Context // - keys []string -func (_e *WatchKV_Expecter) MultiRemove(keys interface{}) *WatchKV_MultiRemove_Call { - return &WatchKV_MultiRemove_Call{Call: _e.mock.On("MultiRemove", keys)} +func (_e *WatchKV_Expecter) MultiRemove(ctx interface{}, keys interface{}) *WatchKV_MultiRemove_Call { + return &WatchKV_MultiRemove_Call{Call: _e.mock.On("MultiRemove", ctx, keys)} } -func (_c *WatchKV_MultiRemove_Call) Run(run func(keys []string)) *WatchKV_MultiRemove_Call { +func (_c *WatchKV_MultiRemove_Call) Run(run func(ctx context.Context, keys []string)) *WatchKV_MultiRemove_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].([]string)) + run(args[0].(context.Context), args[1].([]string)) }) return _c } @@ -493,22 +502,22 @@ func (_c *WatchKV_MultiRemove_Call) Return(_a0 error) *WatchKV_MultiRemove_Call return _c } -func (_c *WatchKV_MultiRemove_Call) RunAndReturn(run func([]string) error) *WatchKV_MultiRemove_Call { +func (_c *WatchKV_MultiRemove_Call) RunAndReturn(run func(context.Context, []string) error) *WatchKV_MultiRemove_Call { _c.Call.Return(run) return _c } -// MultiSave provides a mock function with given fields: kvs -func (_m *WatchKV) MultiSave(kvs map[string]string) error { - ret := _m.Called(kvs) +// MultiSave provides a mock function with given fields: ctx, kvs +func (_m *WatchKV) MultiSave(ctx context.Context, kvs map[string]string) error { + ret := _m.Called(ctx, kvs) if len(ret) == 0 { panic("no return value specified for MultiSave") } var r0 error - if rf, ok := ret.Get(0).(func(map[string]string) error); ok { - r0 = rf(kvs) + if rf, ok := ret.Get(0).(func(context.Context, map[string]string) error); ok { + r0 = rf(ctx, kvs) } else { r0 = ret.Error(0) } @@ -522,14 +531,15 @@ type WatchKV_MultiSave_Call struct { } // MultiSave is a helper method to define mock.On call +// - ctx context.Context // - kvs map[string]string -func (_e *WatchKV_Expecter) MultiSave(kvs interface{}) *WatchKV_MultiSave_Call { - return &WatchKV_MultiSave_Call{Call: _e.mock.On("MultiSave", kvs)} +func (_e *WatchKV_Expecter) MultiSave(ctx interface{}, kvs interface{}) *WatchKV_MultiSave_Call { + return &WatchKV_MultiSave_Call{Call: _e.mock.On("MultiSave", ctx, kvs)} } -func (_c *WatchKV_MultiSave_Call) Run(run func(kvs map[string]string)) *WatchKV_MultiSave_Call { +func (_c *WatchKV_MultiSave_Call) Run(run func(ctx context.Context, kvs map[string]string)) *WatchKV_MultiSave_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(map[string]string)) + run(args[0].(context.Context), args[1].(map[string]string)) }) return _c } @@ -539,19 +549,19 @@ func (_c *WatchKV_MultiSave_Call) Return(_a0 error) *WatchKV_MultiSave_Call { return _c } -func (_c *WatchKV_MultiSave_Call) RunAndReturn(run func(map[string]string) error) *WatchKV_MultiSave_Call { +func (_c *WatchKV_MultiSave_Call) RunAndReturn(run func(context.Context, map[string]string) error) *WatchKV_MultiSave_Call { _c.Call.Return(run) return _c } -// MultiSaveAndRemove provides a mock function with given fields: saves, removals, preds -func (_m *WatchKV) MultiSaveAndRemove(saves map[string]string, removals []string, preds ...predicates.Predicate) error { +// MultiSaveAndRemove provides a mock function with given fields: ctx, saves, removals, preds +func (_m *WatchKV) MultiSaveAndRemove(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error { _va := make([]interface{}, len(preds)) for _i := range preds { _va[_i] = preds[_i] } var _ca []interface{} - _ca = append(_ca, saves, removals) + _ca = append(_ca, ctx, saves, removals) _ca = append(_ca, _va...) ret := _m.Called(_ca...) @@ -560,8 +570,8 @@ func (_m *WatchKV) MultiSaveAndRemove(saves map[string]string, removals []string } var r0 error - if rf, ok := ret.Get(0).(func(map[string]string, []string, ...predicates.Predicate) error); ok { - r0 = rf(saves, removals, preds...) + if rf, ok := ret.Get(0).(func(context.Context, map[string]string, []string, ...predicates.Predicate) error); ok { + r0 = rf(ctx, saves, removals, preds...) } else { r0 = ret.Error(0) } @@ -575,23 +585,24 @@ type WatchKV_MultiSaveAndRemove_Call struct { } // MultiSaveAndRemove is a helper method to define mock.On call +// - ctx context.Context // - saves map[string]string // - removals []string // - preds ...predicates.Predicate -func (_e *WatchKV_Expecter) MultiSaveAndRemove(saves interface{}, removals interface{}, preds ...interface{}) *WatchKV_MultiSaveAndRemove_Call { +func (_e *WatchKV_Expecter) MultiSaveAndRemove(ctx interface{}, saves interface{}, removals interface{}, preds ...interface{}) *WatchKV_MultiSaveAndRemove_Call { return &WatchKV_MultiSaveAndRemove_Call{Call: _e.mock.On("MultiSaveAndRemove", - append([]interface{}{saves, removals}, preds...)...)} + append([]interface{}{ctx, saves, removals}, preds...)...)} } -func (_c *WatchKV_MultiSaveAndRemove_Call) Run(run func(saves map[string]string, removals []string, preds ...predicates.Predicate)) *WatchKV_MultiSaveAndRemove_Call { +func (_c *WatchKV_MultiSaveAndRemove_Call) Run(run func(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate)) *WatchKV_MultiSaveAndRemove_Call { _c.Call.Run(func(args mock.Arguments) { - variadicArgs := make([]predicates.Predicate, len(args)-2) - for i, a := range args[2:] { + variadicArgs := make([]predicates.Predicate, len(args)-3) + for i, a := range args[3:] { if a != nil { variadicArgs[i] = a.(predicates.Predicate) } } - run(args[0].(map[string]string), args[1].([]string), variadicArgs...) + run(args[0].(context.Context), args[1].(map[string]string), args[2].([]string), variadicArgs...) }) return _c } @@ -601,19 +612,19 @@ func (_c *WatchKV_MultiSaveAndRemove_Call) Return(_a0 error) *WatchKV_MultiSaveA return _c } -func (_c *WatchKV_MultiSaveAndRemove_Call) RunAndReturn(run func(map[string]string, []string, ...predicates.Predicate) error) *WatchKV_MultiSaveAndRemove_Call { +func (_c *WatchKV_MultiSaveAndRemove_Call) RunAndReturn(run func(context.Context, map[string]string, []string, ...predicates.Predicate) error) *WatchKV_MultiSaveAndRemove_Call { _c.Call.Return(run) return _c } -// MultiSaveAndRemoveWithPrefix provides a mock function with given fields: saves, removals, preds -func (_m *WatchKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals []string, preds ...predicates.Predicate) error { +// MultiSaveAndRemoveWithPrefix provides a mock function with given fields: ctx, saves, removals, preds +func (_m *WatchKV) MultiSaveAndRemoveWithPrefix(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error { _va := make([]interface{}, len(preds)) for _i := range preds { _va[_i] = preds[_i] } var _ca []interface{} - _ca = append(_ca, saves, removals) + _ca = append(_ca, ctx, saves, removals) _ca = append(_ca, _va...) ret := _m.Called(_ca...) @@ -622,8 +633,8 @@ func (_m *WatchKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, removal } var r0 error - if rf, ok := ret.Get(0).(func(map[string]string, []string, ...predicates.Predicate) error); ok { - r0 = rf(saves, removals, preds...) + if rf, ok := ret.Get(0).(func(context.Context, map[string]string, []string, ...predicates.Predicate) error); ok { + r0 = rf(ctx, saves, removals, preds...) } else { r0 = ret.Error(0) } @@ -637,23 +648,24 @@ type WatchKV_MultiSaveAndRemoveWithPrefix_Call struct { } // MultiSaveAndRemoveWithPrefix is a helper method to define mock.On call +// - ctx context.Context // - saves map[string]string // - removals []string // - preds ...predicates.Predicate -func (_e *WatchKV_Expecter) MultiSaveAndRemoveWithPrefix(saves interface{}, removals interface{}, preds ...interface{}) *WatchKV_MultiSaveAndRemoveWithPrefix_Call { +func (_e *WatchKV_Expecter) MultiSaveAndRemoveWithPrefix(ctx interface{}, saves interface{}, removals interface{}, preds ...interface{}) *WatchKV_MultiSaveAndRemoveWithPrefix_Call { return &WatchKV_MultiSaveAndRemoveWithPrefix_Call{Call: _e.mock.On("MultiSaveAndRemoveWithPrefix", - append([]interface{}{saves, removals}, preds...)...)} + append([]interface{}{ctx, saves, removals}, preds...)...)} } -func (_c *WatchKV_MultiSaveAndRemoveWithPrefix_Call) Run(run func(saves map[string]string, removals []string, preds ...predicates.Predicate)) *WatchKV_MultiSaveAndRemoveWithPrefix_Call { +func (_c *WatchKV_MultiSaveAndRemoveWithPrefix_Call) Run(run func(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate)) *WatchKV_MultiSaveAndRemoveWithPrefix_Call { _c.Call.Run(func(args mock.Arguments) { - variadicArgs := make([]predicates.Predicate, len(args)-2) - for i, a := range args[2:] { + variadicArgs := make([]predicates.Predicate, len(args)-3) + for i, a := range args[3:] { if a != nil { variadicArgs[i] = a.(predicates.Predicate) } } - run(args[0].(map[string]string), args[1].([]string), variadicArgs...) + run(args[0].(context.Context), args[1].(map[string]string), args[2].([]string), variadicArgs...) }) return _c } @@ -663,22 +675,22 @@ func (_c *WatchKV_MultiSaveAndRemoveWithPrefix_Call) Return(_a0 error) *WatchKV_ return _c } -func (_c *WatchKV_MultiSaveAndRemoveWithPrefix_Call) RunAndReturn(run func(map[string]string, []string, ...predicates.Predicate) error) *WatchKV_MultiSaveAndRemoveWithPrefix_Call { +func (_c *WatchKV_MultiSaveAndRemoveWithPrefix_Call) RunAndReturn(run func(context.Context, map[string]string, []string, ...predicates.Predicate) error) *WatchKV_MultiSaveAndRemoveWithPrefix_Call { _c.Call.Return(run) return _c } -// Remove provides a mock function with given fields: key -func (_m *WatchKV) Remove(key string) error { - ret := _m.Called(key) +// Remove provides a mock function with given fields: ctx, key +func (_m *WatchKV) Remove(ctx context.Context, key string) error { + ret := _m.Called(ctx, key) if len(ret) == 0 { panic("no return value specified for Remove") } var r0 error - if rf, ok := ret.Get(0).(func(string) error); ok { - r0 = rf(key) + if rf, ok := ret.Get(0).(func(context.Context, string) error); ok { + r0 = rf(ctx, key) } else { r0 = ret.Error(0) } @@ -692,14 +704,15 @@ type WatchKV_Remove_Call struct { } // Remove is a helper method to define mock.On call +// - ctx context.Context // - key string -func (_e *WatchKV_Expecter) Remove(key interface{}) *WatchKV_Remove_Call { - return &WatchKV_Remove_Call{Call: _e.mock.On("Remove", key)} +func (_e *WatchKV_Expecter) Remove(ctx interface{}, key interface{}) *WatchKV_Remove_Call { + return &WatchKV_Remove_Call{Call: _e.mock.On("Remove", ctx, key)} } -func (_c *WatchKV_Remove_Call) Run(run func(key string)) *WatchKV_Remove_Call { +func (_c *WatchKV_Remove_Call) Run(run func(ctx context.Context, key string)) *WatchKV_Remove_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string)) + run(args[0].(context.Context), args[1].(string)) }) return _c } @@ -709,22 +722,22 @@ func (_c *WatchKV_Remove_Call) Return(_a0 error) *WatchKV_Remove_Call { return _c } -func (_c *WatchKV_Remove_Call) RunAndReturn(run func(string) error) *WatchKV_Remove_Call { +func (_c *WatchKV_Remove_Call) RunAndReturn(run func(context.Context, string) error) *WatchKV_Remove_Call { _c.Call.Return(run) return _c } -// RemoveWithPrefix provides a mock function with given fields: key -func (_m *WatchKV) RemoveWithPrefix(key string) error { - ret := _m.Called(key) +// RemoveWithPrefix provides a mock function with given fields: ctx, key +func (_m *WatchKV) RemoveWithPrefix(ctx context.Context, key string) error { + ret := _m.Called(ctx, key) if len(ret) == 0 { panic("no return value specified for RemoveWithPrefix") } var r0 error - if rf, ok := ret.Get(0).(func(string) error); ok { - r0 = rf(key) + if rf, ok := ret.Get(0).(func(context.Context, string) error); ok { + r0 = rf(ctx, key) } else { r0 = ret.Error(0) } @@ -738,14 +751,15 @@ type WatchKV_RemoveWithPrefix_Call struct { } // RemoveWithPrefix is a helper method to define mock.On call +// - ctx context.Context // - key string -func (_e *WatchKV_Expecter) RemoveWithPrefix(key interface{}) *WatchKV_RemoveWithPrefix_Call { - return &WatchKV_RemoveWithPrefix_Call{Call: _e.mock.On("RemoveWithPrefix", key)} +func (_e *WatchKV_Expecter) RemoveWithPrefix(ctx interface{}, key interface{}) *WatchKV_RemoveWithPrefix_Call { + return &WatchKV_RemoveWithPrefix_Call{Call: _e.mock.On("RemoveWithPrefix", ctx, key)} } -func (_c *WatchKV_RemoveWithPrefix_Call) Run(run func(key string)) *WatchKV_RemoveWithPrefix_Call { +func (_c *WatchKV_RemoveWithPrefix_Call) Run(run func(ctx context.Context, key string)) *WatchKV_RemoveWithPrefix_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string)) + run(args[0].(context.Context), args[1].(string)) }) return _c } @@ -755,22 +769,22 @@ func (_c *WatchKV_RemoveWithPrefix_Call) Return(_a0 error) *WatchKV_RemoveWithPr return _c } -func (_c *WatchKV_RemoveWithPrefix_Call) RunAndReturn(run func(string) error) *WatchKV_RemoveWithPrefix_Call { +func (_c *WatchKV_RemoveWithPrefix_Call) RunAndReturn(run func(context.Context, string) error) *WatchKV_RemoveWithPrefix_Call { _c.Call.Return(run) return _c } -// Save provides a mock function with given fields: key, value -func (_m *WatchKV) Save(key string, value string) error { - ret := _m.Called(key, value) +// Save provides a mock function with given fields: ctx, key, value +func (_m *WatchKV) Save(ctx context.Context, key string, value string) error { + ret := _m.Called(ctx, key, value) if len(ret) == 0 { panic("no return value specified for Save") } var r0 error - if rf, ok := ret.Get(0).(func(string, string) error); ok { - r0 = rf(key, value) + if rf, ok := ret.Get(0).(func(context.Context, string, string) error); ok { + r0 = rf(ctx, key, value) } else { r0 = ret.Error(0) } @@ -784,15 +798,16 @@ type WatchKV_Save_Call struct { } // Save is a helper method to define mock.On call +// - ctx context.Context // - key string // - value string -func (_e *WatchKV_Expecter) Save(key interface{}, value interface{}) *WatchKV_Save_Call { - return &WatchKV_Save_Call{Call: _e.mock.On("Save", key, value)} +func (_e *WatchKV_Expecter) Save(ctx interface{}, key interface{}, value interface{}) *WatchKV_Save_Call { + return &WatchKV_Save_Call{Call: _e.mock.On("Save", ctx, key, value)} } -func (_c *WatchKV_Save_Call) Run(run func(key string, value string)) *WatchKV_Save_Call { +func (_c *WatchKV_Save_Call) Run(run func(ctx context.Context, key string, value string)) *WatchKV_Save_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string), args[1].(string)) + run(args[0].(context.Context), args[1].(string), args[2].(string)) }) return _c } @@ -802,22 +817,22 @@ func (_c *WatchKV_Save_Call) Return(_a0 error) *WatchKV_Save_Call { return _c } -func (_c *WatchKV_Save_Call) RunAndReturn(run func(string, string) error) *WatchKV_Save_Call { +func (_c *WatchKV_Save_Call) RunAndReturn(run func(context.Context, string, string) error) *WatchKV_Save_Call { _c.Call.Return(run) return _c } -// WalkWithPrefix provides a mock function with given fields: prefix, paginationSize, fn -func (_m *WatchKV) WalkWithPrefix(prefix string, paginationSize int, fn func([]byte, []byte) error) error { - ret := _m.Called(prefix, paginationSize, fn) +// WalkWithPrefix provides a mock function with given fields: ctx, prefix, paginationSize, fn +func (_m *WatchKV) WalkWithPrefix(ctx context.Context, prefix string, paginationSize int, fn func([]byte, []byte) error) error { + ret := _m.Called(ctx, prefix, paginationSize, fn) if len(ret) == 0 { panic("no return value specified for WalkWithPrefix") } var r0 error - if rf, ok := ret.Get(0).(func(string, int, func([]byte, []byte) error) error); ok { - r0 = rf(prefix, paginationSize, fn) + if rf, ok := ret.Get(0).(func(context.Context, string, int, func([]byte, []byte) error) error); ok { + r0 = rf(ctx, prefix, paginationSize, fn) } else { r0 = ret.Error(0) } @@ -831,16 +846,17 @@ type WatchKV_WalkWithPrefix_Call struct { } // WalkWithPrefix is a helper method to define mock.On call +// - ctx context.Context // - prefix string // - paginationSize int // - fn func([]byte , []byte) error -func (_e *WatchKV_Expecter) WalkWithPrefix(prefix interface{}, paginationSize interface{}, fn interface{}) *WatchKV_WalkWithPrefix_Call { - return &WatchKV_WalkWithPrefix_Call{Call: _e.mock.On("WalkWithPrefix", prefix, paginationSize, fn)} +func (_e *WatchKV_Expecter) WalkWithPrefix(ctx interface{}, prefix interface{}, paginationSize interface{}, fn interface{}) *WatchKV_WalkWithPrefix_Call { + return &WatchKV_WalkWithPrefix_Call{Call: _e.mock.On("WalkWithPrefix", ctx, prefix, paginationSize, fn)} } -func (_c *WatchKV_WalkWithPrefix_Call) Run(run func(prefix string, paginationSize int, fn func([]byte, []byte) error)) *WatchKV_WalkWithPrefix_Call { +func (_c *WatchKV_WalkWithPrefix_Call) Run(run func(ctx context.Context, prefix string, paginationSize int, fn func([]byte, []byte) error)) *WatchKV_WalkWithPrefix_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string), args[1].(int), args[2].(func([]byte, []byte) error)) + run(args[0].(context.Context), args[1].(string), args[2].(int), args[3].(func([]byte, []byte) error)) }) return _c } @@ -850,22 +866,22 @@ func (_c *WatchKV_WalkWithPrefix_Call) Return(_a0 error) *WatchKV_WalkWithPrefix return _c } -func (_c *WatchKV_WalkWithPrefix_Call) RunAndReturn(run func(string, int, func([]byte, []byte) error) error) *WatchKV_WalkWithPrefix_Call { +func (_c *WatchKV_WalkWithPrefix_Call) RunAndReturn(run func(context.Context, string, int, func([]byte, []byte) error) error) *WatchKV_WalkWithPrefix_Call { _c.Call.Return(run) return _c } -// Watch provides a mock function with given fields: key -func (_m *WatchKV) Watch(key string) clientv3.WatchChan { - ret := _m.Called(key) +// Watch provides a mock function with given fields: ctx, key +func (_m *WatchKV) Watch(ctx context.Context, key string) clientv3.WatchChan { + ret := _m.Called(ctx, key) if len(ret) == 0 { panic("no return value specified for Watch") } var r0 clientv3.WatchChan - if rf, ok := ret.Get(0).(func(string) clientv3.WatchChan); ok { - r0 = rf(key) + if rf, ok := ret.Get(0).(func(context.Context, string) clientv3.WatchChan); ok { + r0 = rf(ctx, key) } else { if ret.Get(0) != nil { r0 = ret.Get(0).(clientv3.WatchChan) @@ -881,14 +897,15 @@ type WatchKV_Watch_Call struct { } // Watch is a helper method to define mock.On call +// - ctx context.Context // - key string -func (_e *WatchKV_Expecter) Watch(key interface{}) *WatchKV_Watch_Call { - return &WatchKV_Watch_Call{Call: _e.mock.On("Watch", key)} +func (_e *WatchKV_Expecter) Watch(ctx interface{}, key interface{}) *WatchKV_Watch_Call { + return &WatchKV_Watch_Call{Call: _e.mock.On("Watch", ctx, key)} } -func (_c *WatchKV_Watch_Call) Run(run func(key string)) *WatchKV_Watch_Call { +func (_c *WatchKV_Watch_Call) Run(run func(ctx context.Context, key string)) *WatchKV_Watch_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string)) + run(args[0].(context.Context), args[1].(string)) }) return _c } @@ -898,22 +915,22 @@ func (_c *WatchKV_Watch_Call) Return(_a0 clientv3.WatchChan) *WatchKV_Watch_Call return _c } -func (_c *WatchKV_Watch_Call) RunAndReturn(run func(string) clientv3.WatchChan) *WatchKV_Watch_Call { +func (_c *WatchKV_Watch_Call) RunAndReturn(run func(context.Context, string) clientv3.WatchChan) *WatchKV_Watch_Call { _c.Call.Return(run) return _c } -// WatchWithPrefix provides a mock function with given fields: key -func (_m *WatchKV) WatchWithPrefix(key string) clientv3.WatchChan { - ret := _m.Called(key) +// WatchWithPrefix provides a mock function with given fields: ctx, key +func (_m *WatchKV) WatchWithPrefix(ctx context.Context, key string) clientv3.WatchChan { + ret := _m.Called(ctx, key) if len(ret) == 0 { panic("no return value specified for WatchWithPrefix") } var r0 clientv3.WatchChan - if rf, ok := ret.Get(0).(func(string) clientv3.WatchChan); ok { - r0 = rf(key) + if rf, ok := ret.Get(0).(func(context.Context, string) clientv3.WatchChan); ok { + r0 = rf(ctx, key) } else { if ret.Get(0) != nil { r0 = ret.Get(0).(clientv3.WatchChan) @@ -929,14 +946,15 @@ type WatchKV_WatchWithPrefix_Call struct { } // WatchWithPrefix is a helper method to define mock.On call +// - ctx context.Context // - key string -func (_e *WatchKV_Expecter) WatchWithPrefix(key interface{}) *WatchKV_WatchWithPrefix_Call { - return &WatchKV_WatchWithPrefix_Call{Call: _e.mock.On("WatchWithPrefix", key)} +func (_e *WatchKV_Expecter) WatchWithPrefix(ctx interface{}, key interface{}) *WatchKV_WatchWithPrefix_Call { + return &WatchKV_WatchWithPrefix_Call{Call: _e.mock.On("WatchWithPrefix", ctx, key)} } -func (_c *WatchKV_WatchWithPrefix_Call) Run(run func(key string)) *WatchKV_WatchWithPrefix_Call { +func (_c *WatchKV_WatchWithPrefix_Call) Run(run func(ctx context.Context, key string)) *WatchKV_WatchWithPrefix_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string)) + run(args[0].(context.Context), args[1].(string)) }) return _c } @@ -946,22 +964,22 @@ func (_c *WatchKV_WatchWithPrefix_Call) Return(_a0 clientv3.WatchChan) *WatchKV_ return _c } -func (_c *WatchKV_WatchWithPrefix_Call) RunAndReturn(run func(string) clientv3.WatchChan) *WatchKV_WatchWithPrefix_Call { +func (_c *WatchKV_WatchWithPrefix_Call) RunAndReturn(run func(context.Context, string) clientv3.WatchChan) *WatchKV_WatchWithPrefix_Call { _c.Call.Return(run) return _c } -// WatchWithRevision provides a mock function with given fields: key, revision -func (_m *WatchKV) WatchWithRevision(key string, revision int64) clientv3.WatchChan { - ret := _m.Called(key, revision) +// WatchWithRevision provides a mock function with given fields: ctx, key, revision +func (_m *WatchKV) WatchWithRevision(ctx context.Context, key string, revision int64) clientv3.WatchChan { + ret := _m.Called(ctx, key, revision) if len(ret) == 0 { panic("no return value specified for WatchWithRevision") } var r0 clientv3.WatchChan - if rf, ok := ret.Get(0).(func(string, int64) clientv3.WatchChan); ok { - r0 = rf(key, revision) + if rf, ok := ret.Get(0).(func(context.Context, string, int64) clientv3.WatchChan); ok { + r0 = rf(ctx, key, revision) } else { if ret.Get(0) != nil { r0 = ret.Get(0).(clientv3.WatchChan) @@ -977,15 +995,16 @@ type WatchKV_WatchWithRevision_Call struct { } // WatchWithRevision is a helper method to define mock.On call +// - ctx context.Context // - key string // - revision int64 -func (_e *WatchKV_Expecter) WatchWithRevision(key interface{}, revision interface{}) *WatchKV_WatchWithRevision_Call { - return &WatchKV_WatchWithRevision_Call{Call: _e.mock.On("WatchWithRevision", key, revision)} +func (_e *WatchKV_Expecter) WatchWithRevision(ctx interface{}, key interface{}, revision interface{}) *WatchKV_WatchWithRevision_Call { + return &WatchKV_WatchWithRevision_Call{Call: _e.mock.On("WatchWithRevision", ctx, key, revision)} } -func (_c *WatchKV_WatchWithRevision_Call) Run(run func(key string, revision int64)) *WatchKV_WatchWithRevision_Call { +func (_c *WatchKV_WatchWithRevision_Call) Run(run func(ctx context.Context, key string, revision int64)) *WatchKV_WatchWithRevision_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string), args[1].(int64)) + run(args[0].(context.Context), args[1].(string), args[2].(int64)) }) return _c } @@ -995,7 +1014,7 @@ func (_c *WatchKV_WatchWithRevision_Call) Return(_a0 clientv3.WatchChan) *WatchK return _c } -func (_c *WatchKV_WatchWithRevision_Call) RunAndReturn(run func(string, int64) clientv3.WatchChan) *WatchKV_WatchWithRevision_Call { +func (_c *WatchKV_WatchWithRevision_Call) RunAndReturn(run func(context.Context, string, int64) clientv3.WatchChan) *WatchKV_WatchWithRevision_Call { _c.Call.Return(run) return _c } diff --git a/internal/kv/tikv/txn_tikv.go b/internal/kv/tikv/txn_tikv.go index f7fcd7d444a7e..47af16f46995a 100644 --- a/internal/kv/tikv/txn_tikv.go +++ b/internal/kv/tikv/txn_tikv.go @@ -138,10 +138,10 @@ func logWarnOnFailure(err *error, msg string, fields ...zap.Field) { } // Has returns if a key exists. -func (kv *txnTiKV) Has(key string) (bool, error) { +func (kv *txnTiKV) Has(ctx context.Context, key string) (bool, error) { start := time.Now() key = path.Join(kv.rootPath, key) - ctx, cancel := context.WithTimeout(context.Background(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() var loggingErr error @@ -167,7 +167,7 @@ func rollbackOnFailure(err *error, txn *transaction.KVTxn) { } // HasPrefix returns if a key prefix exists. -func (kv *txnTiKV) HasPrefix(prefix string) (bool, error) { +func (kv *txnTiKV) HasPrefix(ctx context.Context, prefix string) (bool, error) { start := time.Now() prefix = path.Join(kv.rootPath, prefix) @@ -197,10 +197,10 @@ func (kv *txnTiKV) HasPrefix(prefix string) (bool, error) { } // Load returns value of the key. -func (kv *txnTiKV) Load(key string) (string, error) { +func (kv *txnTiKV) Load(ctx context.Context, key string) (string, error) { start := time.Now() key = path.Join(kv.rootPath, key) - ctx, cancel := context.WithTimeout(context.Background(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() var loggingErr error @@ -229,9 +229,9 @@ func batchConvertFromString(prefix string, keys []string) [][]byte { } // MultiLoad gets the values of input keys in a transaction. -func (kv *txnTiKV) MultiLoad(keys []string) ([]string, error) { +func (kv *txnTiKV) MultiLoad(ctx context.Context, keys []string) ([]string, error) { start := time.Now() - ctx, cancel := context.WithTimeout(context.Background(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() var loggingErr error @@ -271,7 +271,7 @@ func (kv *txnTiKV) MultiLoad(keys []string) ([]string, error) { } // LoadWithPrefix returns all the keys and values for the given key prefix. -func (kv *txnTiKV) LoadWithPrefix(prefix string) ([]string, []string, error) { +func (kv *txnTiKV) LoadWithPrefix(ctx context.Context, prefix string) ([]string, []string, error) { start := time.Now() prefix = path.Join(kv.rootPath, prefix) @@ -311,9 +311,9 @@ func (kv *txnTiKV) LoadWithPrefix(prefix string) ([]string, []string, error) { } // Save saves the input key-value pair. -func (kv *txnTiKV) Save(key, value string) error { +func (kv *txnTiKV) Save(ctx context.Context, key, value string) error { key = path.Join(kv.rootPath, key) - ctx, cancel := context.WithTimeout(context.Background(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() var loggingErr error @@ -324,9 +324,9 @@ func (kv *txnTiKV) Save(key, value string) error { } // MultiSave saves the input key-value pairs in transaction manner. -func (kv *txnTiKV) MultiSave(kvs map[string]string) error { +func (kv *txnTiKV) MultiSave(ctx context.Context, kvs map[string]string) error { start := time.Now() - ctx, cancel := context.WithTimeout(context.Background(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() var loggingErr error @@ -366,9 +366,9 @@ func (kv *txnTiKV) MultiSave(kvs map[string]string) error { } // Remove removes the input key. -func (kv *txnTiKV) Remove(key string) error { +func (kv *txnTiKV) Remove(ctx context.Context, key string) error { key = path.Join(kv.rootPath, key) - ctx, cancel := context.WithTimeout(context.Background(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() var loggingErr error @@ -379,9 +379,9 @@ func (kv *txnTiKV) Remove(key string) error { } // MultiRemove removes the input keys in transaction manner. -func (kv *txnTiKV) MultiRemove(keys []string) error { +func (kv *txnTiKV) MultiRemove(ctx context.Context, keys []string) error { start := time.Now() - ctx, cancel := context.WithTimeout(context.Background(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() var loggingErr error @@ -415,10 +415,10 @@ func (kv *txnTiKV) MultiRemove(keys []string) error { } // RemoveWithPrefix removes the keys for the given prefix. -func (kv *txnTiKV) RemoveWithPrefix(prefix string) error { +func (kv *txnTiKV) RemoveWithPrefix(ctx context.Context, prefix string) error { start := time.Now() prefix = path.Join(kv.rootPath, prefix) - ctx, cancel := context.WithTimeout(context.Background(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() var loggingErr error @@ -436,9 +436,9 @@ func (kv *txnTiKV) RemoveWithPrefix(prefix string) error { } // MultiSaveAndRemove saves the key-value pairs and removes the keys in a transaction. -func (kv *txnTiKV) MultiSaveAndRemove(saves map[string]string, removals []string, preds ...predicates.Predicate) error { +func (kv *txnTiKV) MultiSaveAndRemove(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error { start := time.Now() - ctx, cancel := context.WithTimeout(context.Background(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() var loggingErr error @@ -499,9 +499,9 @@ func (kv *txnTiKV) MultiSaveAndRemove(saves map[string]string, removals []string } // MultiSaveAndRemoveWithPrefix saves kv in @saves and removes the keys with given prefix in @removals. -func (kv *txnTiKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals []string, preds ...predicates.Predicate) error { +func (kv *txnTiKV) MultiSaveAndRemoveWithPrefix(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error { start := time.Now() - ctx, cancel := context.WithTimeout(context.Background(), kv.requestTimeout) + ctx, cancel := context.WithTimeout(ctx, kv.requestTimeout) defer cancel() var loggingErr error @@ -585,7 +585,7 @@ func (kv *txnTiKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, removal } // WalkWithPrefix visits each kv with input prefix and apply given fn to it. -func (kv *txnTiKV) WalkWithPrefix(prefix string, paginationSize int, fn func([]byte, []byte) error) error { +func (kv *txnTiKV) WalkWithPrefix(ctx context.Context, prefix string, paginationSize int, fn func([]byte, []byte) error) error { start := time.Now() prefix = path.Join(kv.rootPath, prefix) @@ -746,7 +746,7 @@ func (kv *txnTiKV) removeTiKVMeta(ctx context.Context, key string) error { return err } -func (kv *txnTiKV) CompareVersionAndSwap(key string, version int64, target string) (bool, error) { +func (kv *txnTiKV) CompareVersionAndSwap(ctx context.Context, key string, version int64, target string) (bool, error) { err := fmt.Errorf("Unimplemented! CompareVersionAndSwap is under deprecation") logWarnOnFailure(&err, "Unimplemented") return false, err diff --git a/internal/kv/tikv/txn_tikv_test.go b/internal/kv/tikv/txn_tikv_test.go index eeff2338ffb9f..fd62f4ffc1249 100644 --- a/internal/kv/tikv/txn_tikv_test.go +++ b/internal/kv/tikv/txn_tikv_test.go @@ -37,11 +37,11 @@ func TestTiKVLoad(te *testing.T) { te.Run("kv SaveAndLoad", func(t *testing.T) { rootPath := "/tikv/test/root/saveandload" kv := NewTiKV(txnClient, rootPath) - err := kv.RemoveWithPrefix("") + err := kv.RemoveWithPrefix(context.TODO(), "") require.NoError(t, err) defer kv.Close() - defer kv.RemoveWithPrefix("") + defer kv.RemoveWithPrefix(context.TODO(), "") saveAndLoadTests := []struct { key string @@ -55,11 +55,11 @@ func TestTiKVLoad(te *testing.T) { for i, test := range saveAndLoadTests { if i < 4 { - err = kv.Save(test.key, test.value) + err = kv.Save(context.TODO(), test.key, test.value) assert.NoError(t, err) } - val, err := kv.Load(test.key) + val, err := kv.Load(context.TODO(), test.key) assert.NoError(t, err) assert.Equal(t, test.value, val) } @@ -73,7 +73,7 @@ func TestTiKVLoad(te *testing.T) { } for _, test := range invalidLoadTests { - val, err := kv.Load(test.invalidKey) + val, err := kv.Load(context.TODO(), test.invalidKey) assert.Error(t, err) assert.Zero(t, val) } @@ -110,7 +110,7 @@ func TestTiKVLoad(te *testing.T) { } for _, test := range loadPrefixTests { - actualKeys, actualValues, err := kv.LoadWithPrefix(test.prefix) + actualKeys, actualValues, err := kv.LoadWithPrefix(context.TODO(), test.prefix) assert.ElementsMatch(t, test.expectedKeys, actualKeys) assert.ElementsMatch(t, test.expectedValues, actualValues) assert.Equal(t, test.expectedError, err) @@ -127,15 +127,15 @@ func TestTiKVLoad(te *testing.T) { } for _, test := range removeTests { - err = kv.Remove(test.validKey) + err = kv.Remove(context.TODO(), test.validKey) assert.NoError(t, err) - _, err = kv.Load(test.validKey) + _, err = kv.Load(context.TODO(), test.validKey) assert.Error(t, err) - err = kv.Remove(test.validKey) + err = kv.Remove(context.TODO(), test.validKey) assert.NoError(t, err) - err = kv.Remove(test.invalidKey) + err = kv.Remove(context.TODO(), test.invalidKey) assert.NoError(t, err) } }) @@ -145,7 +145,7 @@ func TestTiKVLoad(te *testing.T) { kv := NewTiKV(txnClient, rootPath) defer kv.Close() - defer kv.RemoveWithPrefix("") + defer kv.RemoveWithPrefix(context.TODO(), "") multiSaveTests := map[string]string{ "key_1": "value_1", @@ -156,10 +156,10 @@ func TestTiKVLoad(te *testing.T) { "_": "other", } - err := kv.MultiSave(multiSaveTests) + err := kv.MultiSave(context.TODO(), multiSaveTests) assert.NoError(t, err) for k, v := range multiSaveTests { - actualV, err := kv.Load(k) + actualV, err := kv.Load(context.TODO(), k) assert.NoError(t, err) assert.Equal(t, v, actualV) } @@ -175,7 +175,7 @@ func TestTiKVLoad(te *testing.T) { } for _, test := range multiLoadTests { - vs, err := kv.MultiLoad(test.inputKeys) + vs, err := kv.MultiLoad(context.TODO(), test.inputKeys) assert.NoError(t, err) assert.Equal(t, test.expectedValues, vs) } @@ -191,7 +191,7 @@ func TestTiKVLoad(te *testing.T) { } for _, test := range invalidMultiLoad { - vs, err := kv.MultiLoad(test.invalidKeys) + vs, err := kv.MultiLoad(context.TODO(), test.invalidKeys) assert.Error(t, err) assert.Equal(t, test.expectedValues, vs) } @@ -202,10 +202,10 @@ func TestTiKVLoad(te *testing.T) { } for _, k := range removeWithPrefixTests { - err = kv.RemoveWithPrefix(k) + err = kv.RemoveWithPrefix(context.TODO(), k) assert.NoError(t, err) - ks, vs, err := kv.LoadWithPrefix(k) + ks, vs, err := kv.LoadWithPrefix(context.TODO(), k) assert.Empty(t, ks) assert.Empty(t, vs) assert.NoError(t, err) @@ -218,10 +218,10 @@ func TestTiKVLoad(te *testing.T) { "_", } - err = kv.MultiRemove(multiRemoveTests) + err = kv.MultiRemove(context.TODO(), multiRemoveTests) assert.NoError(t, err) - ks, vs, err := kv.LoadWithPrefix("") + ks, vs, err := kv.LoadWithPrefix(context.TODO(), "") assert.NoError(t, err) assert.Empty(t, ks) assert.Empty(t, vs) @@ -238,11 +238,11 @@ func TestTiKVLoad(te *testing.T) { {make(map[string]string), []string{"multikey_2"}}, } for _, test := range multiSaveAndRemoveTests { - err = kv.MultiSaveAndRemove(test.multiSaves, test.multiRemoves) + err = kv.MultiSaveAndRemove(context.TODO(), test.multiSaves, test.multiRemoves) assert.NoError(t, err) } - ks, vs, err = kv.LoadWithPrefix("") + ks, vs, err = kv.LoadWithPrefix(context.TODO(), "") assert.NoError(t, err) assert.Empty(t, ks) assert.Empty(t, vs) @@ -252,7 +252,7 @@ func TestTiKVLoad(te *testing.T) { rootPath := "/tikv/test/root/multi_remove_with_prefix" kv := NewTiKV(txnClient, rootPath) defer kv.Close() - defer kv.RemoveWithPrefix("") + defer kv.RemoveWithPrefix(context.TODO(), "") prepareTests := map[string]string{ "x/abc/1": "1", @@ -264,7 +264,7 @@ func TestTiKVLoad(te *testing.T) { } // MultiSaveAndRemoveWithPrefix - err := kv.MultiSave(prepareTests) + err := kv.MultiSave(context.TODO(), prepareTests) require.NoError(t, err) multiSaveAndRemoveWithPrefixTests := []struct { multiSave map[string]string @@ -282,14 +282,14 @@ func TestTiKVLoad(te *testing.T) { } for _, test := range multiSaveAndRemoveWithPrefixTests { - k, _, err := kv.LoadWithPrefix(test.loadPrefix) + k, _, err := kv.LoadWithPrefix(context.TODO(), test.loadPrefix) assert.NoError(t, err) assert.Equal(t, test.lengthBeforeRemove, len(k)) - err = kv.MultiSaveAndRemoveWithPrefix(test.multiSave, test.prefix) + err = kv.MultiSaveAndRemoveWithPrefix(context.TODO(), test.multiSave, test.prefix) assert.NoError(t, err) - k, _, err = kv.LoadWithPrefix(test.loadPrefix) + k, _, err = kv.LoadWithPrefix(context.TODO(), test.loadPrefix) assert.NoError(t, err) assert.Equal(t, test.lengthAfterRemove, len(k)) } @@ -306,17 +306,17 @@ func TestTiKVLoad(te *testing.T) { defer func() { beginTxn = tiTxnBegin }() - err := kv.Save("key1", "v1") + err := kv.Save(context.TODO(), "key1", "v1") assert.Error(t, err) - err = kv.MultiSave(map[string]string{"A/100": "v1"}) + err = kv.MultiSave(context.TODO(), map[string]string{"A/100": "v1"}) assert.Error(t, err) - err = kv.Remove("key1") + err = kv.Remove(context.TODO(), "key1") assert.Error(t, err) - err = kv.MultiRemove([]string{"key_1", "key_2"}) + err = kv.MultiRemove(context.TODO(), []string{"key_1", "key_2"}) assert.Error(t, err) - err = kv.MultiSaveAndRemove(map[string]string{"key_1": "value_1"}, []string{}) + err = kv.MultiSaveAndRemove(context.TODO(), map[string]string{"key_1": "value_1"}, []string{}) assert.Error(t, err) - err = kv.MultiSaveAndRemoveWithPrefix(map[string]string{"y/c": "vvv"}, []string{"/"}) + err = kv.MultiSaveAndRemoveWithPrefix(context.TODO(), map[string]string{"y/c": "vvv"}, []string{"/"}) assert.Error(t, err) }) @@ -332,17 +332,17 @@ func TestTiKVLoad(te *testing.T) { commitTxn = tiTxnCommit }() var err error - err = kv.Save("key1", "v1") + err = kv.Save(context.TODO(), "key1", "v1") assert.Error(t, err) - err = kv.MultiSave(map[string]string{"A/100": "v1"}) + err = kv.MultiSave(context.TODO(), map[string]string{"A/100": "v1"}) assert.Error(t, err) - err = kv.Remove("key1") + err = kv.Remove(context.TODO(), "key1") assert.Error(t, err) - err = kv.MultiRemove([]string{"key_1", "key_2"}) + err = kv.MultiRemove(context.TODO(), []string{"key_1", "key_2"}) assert.Error(t, err) - err = kv.MultiSaveAndRemove(map[string]string{"key_1": "value_1"}, []string{}) + err = kv.MultiSaveAndRemove(context.TODO(), map[string]string{"key_1": "value_1"}, []string{}) assert.Error(t, err) - err = kv.MultiSaveAndRemoveWithPrefix(map[string]string{"y/c": "vvv"}, []string{"/"}) + err = kv.MultiSaveAndRemoveWithPrefix(context.TODO(), map[string]string{"y/c": "vvv"}, []string{"/"}) assert.Error(t, err) }) } @@ -352,7 +352,7 @@ func TestWalkWithPagination(t *testing.T) { kv := NewTiKV(txnClient, rootPath) defer kv.Close() - defer kv.RemoveWithPrefix("") + defer kv.RemoveWithPrefix(context.TODO(), "") kvs := map[string]string{ "A/100": "v1", @@ -362,23 +362,23 @@ func TestWalkWithPagination(t *testing.T) { "B/100": "v5", } - err := kv.MultiSave(kvs) + err := kv.MultiSave(context.TODO(), kvs) assert.NoError(t, err) for k, v := range kvs { - actualV, err := kv.Load(k) + actualV, err := kv.Load(context.TODO(), k) assert.NoError(t, err) assert.Equal(t, v, actualV) } t.Run("apply function error ", func(t *testing.T) { - err = kv.WalkWithPrefix("A", 5, func(key []byte, value []byte) error { + err = kv.WalkWithPrefix(context.TODO(), "A", 5, func(key []byte, value []byte) error { return errors.New("error") }) assert.Error(t, err) }) t.Run("get with non-exist prefix ", func(t *testing.T) { - err = kv.WalkWithPrefix("non-exist-prefix", 5, func(key []byte, value []byte) error { + err = kv.WalkWithPrefix(context.TODO(), "non-exist-prefix", 5, func(key []byte, value []byte) error { return nil }) assert.NoError(t, err) @@ -399,7 +399,7 @@ func TestWalkWithPagination(t *testing.T) { ret := make(map[string]string) actualKeys := make([]string, 0) - err = kv.WalkWithPrefix("A", pagination, func(key []byte, value []byte) error { + err = kv.WalkWithPrefix(context.TODO(), "A", pagination, func(key []byte, value []byte) error { k := string(key) k = k[len(rootPath)+1:] ret[k] = string(value) @@ -434,30 +434,30 @@ func TestElapse(t *testing.T) { func TestHas(t *testing.T) { rootPath := "/tikv/test/root/pagination" kv := NewTiKV(txnClient, rootPath) - err := kv.RemoveWithPrefix("") + err := kv.RemoveWithPrefix(context.TODO(), "") require.NoError(t, err) defer kv.Close() - defer kv.RemoveWithPrefix("") + defer kv.RemoveWithPrefix(context.TODO(), "") - has, err := kv.Has("key1") + has, err := kv.Has(context.TODO(), "key1") assert.NoError(t, err) assert.False(t, has) - err = kv.Save("key1", "value1") + err = kv.Save(context.TODO(), "key1", "value1") assert.NoError(t, err) - err = kv.Save("key1", EmptyValueString) + err = kv.Save(context.TODO(), "key1", EmptyValueString) assert.Error(t, err) - has, err = kv.Has("key1") + has, err = kv.Has(context.TODO(), "key1") assert.NoError(t, err) assert.True(t, has) - err = kv.Remove("key1") + err = kv.Remove(context.TODO(), "key1") assert.NoError(t, err) - has, err = kv.Has("key1") + has, err = kv.Has(context.TODO(), "key1") assert.NoError(t, err) assert.False(t, has) } @@ -465,27 +465,27 @@ func TestHas(t *testing.T) { func TestHasPrefix(t *testing.T) { rootPath := "/etcd/test/root/hasprefix" kv := NewTiKV(txnClient, rootPath) - err := kv.RemoveWithPrefix("") + err := kv.RemoveWithPrefix(context.TODO(), "") require.NoError(t, err) defer kv.Close() - defer kv.RemoveWithPrefix("") + defer kv.RemoveWithPrefix(context.TODO(), "") - has, err := kv.HasPrefix("key") + has, err := kv.HasPrefix(context.TODO(), "key") assert.NoError(t, err) assert.False(t, has) - err = kv.Save("key1", "value1") + err = kv.Save(context.TODO(), "key1", "value1") assert.NoError(t, err) - has, err = kv.HasPrefix("key") + has, err = kv.HasPrefix(context.TODO(), "key") assert.NoError(t, err) assert.True(t, has) - err = kv.Remove("key1") + err = kv.Remove(context.TODO(), "key1") assert.NoError(t, err) - has, err = kv.HasPrefix("key") + has, err = kv.HasPrefix(context.TODO(), "key") assert.NoError(t, err) assert.False(t, has) } @@ -493,32 +493,32 @@ func TestHasPrefix(t *testing.T) { func TestEmptyKey(t *testing.T) { rootPath := "/etcd/test/root/loadempty" kv := NewTiKV(txnClient, rootPath) - err := kv.RemoveWithPrefix("") + err := kv.RemoveWithPrefix(context.TODO(), "") require.NoError(t, err) defer kv.Close() - defer kv.RemoveWithPrefix("") + defer kv.RemoveWithPrefix(context.TODO(), "") - has, err := kv.HasPrefix("key") + has, err := kv.HasPrefix(context.TODO(), "key") assert.NoError(t, err) assert.False(t, has) - err = kv.Save("key", "") + err = kv.Save(context.TODO(), "key", "") assert.NoError(t, err) - has, err = kv.HasPrefix("key") + has, err = kv.HasPrefix(context.TODO(), "key") assert.NoError(t, err) assert.True(t, has) - val, err := kv.Load("key") + val, err := kv.Load(context.TODO(), "key") assert.NoError(t, err) assert.Equal(t, val, "") - _, vals, err := kv.LoadWithPrefix("key") + _, vals, err := kv.LoadWithPrefix(context.TODO(), "key") assert.NoError(t, err) assert.Equal(t, vals[0], "") - vals, err = kv.MultiLoad([]string{"key"}) + vals, err = kv.MultiLoad(context.TODO(), []string{"key"}) assert.NoError(t, err) assert.Equal(t, vals[0], "") @@ -528,34 +528,34 @@ func TestEmptyKey(t *testing.T) { return nil } - err = kv.WalkWithPrefix("", 1, nothing) + err = kv.WalkWithPrefix(context.TODO(), "", 1, nothing) assert.NoError(t, err) assert.Equal(t, res, "") multiSaveTests := map[string]string{ "key1": "", } - err = kv.MultiSave(multiSaveTests) + err = kv.MultiSave(context.TODO(), multiSaveTests) assert.NoError(t, err) - val, err = kv.Load("key1") + val, err = kv.Load(context.TODO(), "key1") assert.NoError(t, err) assert.Equal(t, val, "") multiSaveTests = map[string]string{ "key2": "", } - err = kv.MultiSaveAndRemove(multiSaveTests, nil) + err = kv.MultiSaveAndRemove(context.TODO(), multiSaveTests, nil) assert.NoError(t, err) - val, err = kv.Load("key2") + val, err = kv.Load(context.TODO(), "key2") assert.NoError(t, err) assert.Equal(t, val, "") multiSaveTests = map[string]string{ "key3": "", } - err = kv.MultiSaveAndRemoveWithPrefix(multiSaveTests, nil) + err = kv.MultiSaveAndRemoveWithPrefix(context.TODO(), multiSaveTests, nil) assert.NoError(t, err) - val, err = kv.Load("key3") + val, err = kv.Load(context.TODO(), "key3") assert.NoError(t, err) assert.Equal(t, val, "") } @@ -563,11 +563,11 @@ func TestEmptyKey(t *testing.T) { func TestScanSize(t *testing.T) { scanSize := SnapshotScanSize kv := NewTiKV(txnClient, "/") - err := kv.RemoveWithPrefix("") + err := kv.RemoveWithPrefix(context.TODO(), "") require.NoError(t, err) defer kv.Close() - defer kv.RemoveWithPrefix("") + defer kv.RemoveWithPrefix(context.TODO(), "") // Test total > scansize keyMap := map[string]string{} @@ -576,32 +576,32 @@ func TestScanSize(t *testing.T) { keyMap[a] = a } - err = kv.MultiSave(keyMap) + err = kv.MultiSave(context.TODO(), keyMap) assert.NoError(t, err) - keys, _, err := kv.LoadWithPrefix("") + keys, _, err := kv.LoadWithPrefix(context.TODO(), "") assert.NoError(t, err) assert.Equal(t, len(keys), scanSize+100) - err = kv.RemoveWithPrefix("") + err = kv.RemoveWithPrefix(context.TODO(), "") require.NoError(t, err) } func TestTiKVUnimplemented(t *testing.T) { kv := NewTiKV(txnClient, "/") - err := kv.RemoveWithPrefix("") + err := kv.RemoveWithPrefix(context.TODO(), "") require.NoError(t, err) defer kv.Close() - defer kv.RemoveWithPrefix("") + defer kv.RemoveWithPrefix(context.TODO(), "") - _, err = kv.CompareVersionAndSwap("k", 1, "target") + _, err = kv.CompareVersionAndSwap(context.TODO(), "k", 1, "target") assert.Error(t, err) } func TestTxnWithPredicates(t *testing.T) { kv := NewTiKV(txnClient, "/") - err := kv.RemoveWithPrefix("") + err := kv.RemoveWithPrefix(context.TODO(), "") require.NoError(t, err) prepareKV := map[string]string{ @@ -609,7 +609,7 @@ func TestTxnWithPredicates(t *testing.T) { "lease2": "2", } - err = kv.MultiSave(prepareKV) + err = kv.MultiSave(context.TODO(), prepareKV) require.NoError(t, err) multiSaveAndRemovePredTests := []struct { @@ -624,14 +624,14 @@ func TestTxnWithPredicates(t *testing.T) { for _, test := range multiSaveAndRemovePredTests { t.Run(test.tag, func(t *testing.T) { - err := kv.MultiSaveAndRemove(test.multiSave, nil, test.preds...) + err := kv.MultiSaveAndRemove(context.TODO(), test.multiSave, nil, test.preds...) t.Log(err) if test.expectSuccess { assert.NoError(t, err) } else { assert.Error(t, err) } - err = kv.MultiSaveAndRemoveWithPrefix(test.multiSave, nil, test.preds...) + err = kv.MultiSaveAndRemoveWithPrefix(context.TODO(), test.multiSave, nil, test.preds...) if test.expectSuccess { assert.NoError(t, err) } else { diff --git a/internal/metastore/kv/datacoord/kv_catalog.go b/internal/metastore/kv/datacoord/kv_catalog.go index 2233029ee30f6..7bea3819433ca 100644 --- a/internal/metastore/kv/datacoord/kv_catalog.go +++ b/internal/metastore/kv/datacoord/kv_catalog.go @@ -69,7 +69,7 @@ func (kc *Catalog) ListSegments(ctx context.Context) ([]*datapb.SegmentInfo, err executeFn := func(binlogType storage.BinlogType, result map[typeutil.UniqueID][]*datapb.FieldBinlog) { group.Go(func() error { - ret, err := kc.listBinlogs(binlogType) + ret, err := kc.listBinlogs(ctx, binlogType) if err != nil { return err } @@ -85,7 +85,7 @@ func (kc *Catalog) ListSegments(ctx context.Context) ([]*datapb.SegmentInfo, err executeFn(storage.StatsBinlog, statsLogs) executeFn(storage.BM25Binlog, bm25Logs) group.Go(func() error { - ret, err := kc.listSegments() + ret, err := kc.listSegments(ctx) if err != nil { return err } @@ -105,7 +105,7 @@ func (kc *Catalog) ListSegments(ctx context.Context) ([]*datapb.SegmentInfo, err return segments, nil } -func (kc *Catalog) listSegments() ([]*datapb.SegmentInfo, error) { +func (kc *Catalog) listSegments(ctx context.Context) ([]*datapb.SegmentInfo, error) { segments := make([]*datapb.SegmentInfo, 0) applyFn := func(key []byte, value []byte) error { @@ -130,7 +130,7 @@ func (kc *Catalog) listSegments() ([]*datapb.SegmentInfo, error) { return nil } - err := kc.MetaKv.WalkWithPrefix(SegmentPrefix+"/", paginationSize, applyFn) + err := kc.MetaKv.WalkWithPrefix(ctx, SegmentPrefix+"/", paginationSize, applyFn) if err != nil { return nil, err } @@ -163,7 +163,7 @@ func (kc *Catalog) parseBinlogKey(key string, prefixIdx int) (int64, int64, int6 return collectionID, partitionID, segmentID, nil } -func (kc *Catalog) listBinlogs(binlogType storage.BinlogType) (map[typeutil.UniqueID][]*datapb.FieldBinlog, error) { +func (kc *Catalog) listBinlogs(ctx context.Context, binlogType storage.BinlogType) (map[typeutil.UniqueID][]*datapb.FieldBinlog, error) { ret := make(map[typeutil.UniqueID][]*datapb.FieldBinlog) var err error @@ -215,7 +215,7 @@ func (kc *Catalog) listBinlogs(binlogType storage.BinlogType) (map[typeutil.Uniq return nil } - err = kc.MetaKv.WalkWithPrefix(logPathPrefix, paginationSize, applyFn) + err = kc.MetaKv.WalkWithPrefix(ctx, logPathPrefix, paginationSize, applyFn) if err != nil { return nil, err } @@ -263,13 +263,13 @@ func (kc *Catalog) AddSegment(ctx context.Context, segment *datapb.SegmentInfo) if err != nil { return err } - return kc.MetaKv.MultiSave(kvs) + return kc.MetaKv.MultiSave(ctx, kvs) } // LoadFromSegmentPath loads segment info from persistent storage by given segment path. // # TESTING ONLY # -func (kc *Catalog) LoadFromSegmentPath(colID, partID, segID typeutil.UniqueID) (*datapb.SegmentInfo, error) { - v, err := kc.MetaKv.Load(buildSegmentPath(colID, partID, segID)) +func (kc *Catalog) LoadFromSegmentPath(ctx context.Context, colID, partID, segID typeutil.UniqueID) (*datapb.SegmentInfo, error) { + v, err := kc.MetaKv.Load(ctx, buildSegmentPath(colID, partID, segID)) if err != nil { log.Error("(testing only) failed to load segment info by segment path") return nil, err @@ -301,7 +301,7 @@ func (kc *Catalog) AlterSegments(ctx context.Context, segments []*datapb.Segment } if segment.GetState() == commonpb.SegmentState_Dropped { - binlogs, err := kc.handleDroppedSegment(segment) + binlogs, err := kc.handleDroppedSegment(ctx, segment) if err != nil { return err } @@ -327,12 +327,12 @@ func (kc *Catalog) AlterSegments(ctx context.Context, segments []*datapb.Segment maps.Copy(kvs, binlogKvs) } - return kc.SaveByBatch(kvs) + return kc.SaveByBatch(ctx, kvs) } -func (kc *Catalog) handleDroppedSegment(segment *datapb.SegmentInfo) (kvs map[string]string, err error) { +func (kc *Catalog) handleDroppedSegment(ctx context.Context, segment *datapb.SegmentInfo) (kvs map[string]string, err error) { var has bool - has, err = kc.hasBinlogPrefix(segment) + has, err = kc.hasBinlogPrefix(ctx, segment) if err != nil { return } @@ -346,9 +346,9 @@ func (kc *Catalog) handleDroppedSegment(segment *datapb.SegmentInfo) (kvs map[st return } -func (kc *Catalog) SaveByBatch(kvs map[string]string) error { +func (kc *Catalog) SaveByBatch(ctx context.Context, kvs map[string]string) error { saveFn := func(partialKvs map[string]string) error { - return kc.MetaKv.MultiSave(partialKvs) + return kc.MetaKv.MultiSave(ctx, partialKvs) } err := etcd.SaveByBatchWithLimit(kvs, util.MaxEtcdTxnNum, saveFn) if err != nil { @@ -358,22 +358,22 @@ func (kc *Catalog) SaveByBatch(kvs map[string]string) error { return nil } -func (kc *Catalog) hasBinlogPrefix(segment *datapb.SegmentInfo) (bool, error) { +func (kc *Catalog) hasBinlogPrefix(ctx context.Context, segment *datapb.SegmentInfo) (bool, error) { collectionID, partitionID, segmentID := segment.GetCollectionID(), segment.GetPartitionID(), segment.GetID() prefix := buildFieldBinlogPathPrefix(collectionID, partitionID, segmentID) - hasBinlogPrefix, err := kc.MetaKv.HasPrefix(prefix) + hasBinlogPrefix, err := kc.MetaKv.HasPrefix(ctx, prefix) if err != nil { return false, err } prefix = buildFieldDeltalogPathPrefix(collectionID, partitionID, segmentID) - hasDeltaPrefix, err := kc.MetaKv.HasPrefix(prefix) + hasDeltaPrefix, err := kc.MetaKv.HasPrefix(ctx, prefix) if err != nil { return false, err } prefix = buildFieldStatslogPathPrefix(collectionID, partitionID, segmentID) - hasStatsPrefix, err := kc.MetaKv.HasPrefix(prefix) + hasStatsPrefix, err := kc.MetaKv.HasPrefix(ctx, prefix) if err != nil { return false, err } @@ -400,7 +400,7 @@ func (kc *Catalog) SaveDroppedSegmentsInBatch(ctx context.Context, segments []*d } saveFn := func(partialKvs map[string]string) error { - return kc.MetaKv.MultiSave(partialKvs) + return kc.MetaKv.MultiSave(ctx, partialKvs) } if err := etcd.SaveByBatchWithLimit(kvs, util.MaxEtcdTxnNum, saveFn); err != nil { return err @@ -417,7 +417,7 @@ func (kc *Catalog) DropSegment(ctx context.Context, segment *datapb.SegmentInfo) bm25logPrefix := fmt.Sprintf("%s/%d/%d/%d", SegmentBM25logPathPrefix, segment.GetCollectionID(), segment.GetPartitionID(), segment.GetID()) keys := []string{segKey, binlogPreix, deltalogPreix, statelogPreix, bm25logPrefix} - if err := kc.MetaKv.MultiSaveAndRemoveWithPrefix(nil, keys); err != nil { + if err := kc.MetaKv.MultiSaveAndRemoveWithPrefix(ctx, nil, keys); err != nil { return err } @@ -426,7 +426,7 @@ func (kc *Catalog) DropSegment(ctx context.Context, segment *datapb.SegmentInfo) func (kc *Catalog) MarkChannelAdded(ctx context.Context, channel string) error { key := buildChannelRemovePath(channel) - err := kc.MetaKv.Save(key, NonRemoveFlagTomestone) + err := kc.MetaKv.Save(ctx, key, NonRemoveFlagTomestone) if err != nil { log.Error("failed to mark channel added", zap.String("channel", channel), zap.Error(err)) return err @@ -437,7 +437,7 @@ func (kc *Catalog) MarkChannelAdded(ctx context.Context, channel string) error { func (kc *Catalog) MarkChannelDeleted(ctx context.Context, channel string) error { key := buildChannelRemovePath(channel) - err := kc.MetaKv.Save(key, RemoveFlagTomestone) + err := kc.MetaKv.Save(ctx, key, RemoveFlagTomestone) if err != nil { log.Error("Failed to mark channel dropped", zap.String("channel", channel), zap.Error(err)) return err @@ -448,7 +448,7 @@ func (kc *Catalog) MarkChannelDeleted(ctx context.Context, channel string) error func (kc *Catalog) ShouldDropChannel(ctx context.Context, channel string) bool { key := buildChannelRemovePath(channel) - v, err := kc.MetaKv.Load(key) + v, err := kc.MetaKv.Load(ctx, key) if err != nil || v != RemoveFlagTomestone { return false } @@ -457,7 +457,7 @@ func (kc *Catalog) ShouldDropChannel(ctx context.Context, channel string) bool { func (kc *Catalog) ChannelExists(ctx context.Context, channel string) bool { key := buildChannelRemovePath(channel) - v, err := kc.MetaKv.Load(key) + v, err := kc.MetaKv.Load(ctx, key) return err == nil && v == NonRemoveFlagTomestone } @@ -465,7 +465,7 @@ func (kc *Catalog) ChannelExists(ctx context.Context, channel string) bool { func (kc *Catalog) DropChannel(ctx context.Context, channel string) error { key := buildChannelRemovePath(channel) log.Info("removing channel remove path", zap.String("channel", channel)) - return kc.MetaKv.Remove(key) + return kc.MetaKv.Remove(ctx, key) } func (kc *Catalog) ListChannelCheckpoint(ctx context.Context) (map[string]*msgpb.MsgPosition, error) { @@ -483,7 +483,7 @@ func (kc *Catalog) ListChannelCheckpoint(ctx context.Context) (map[string]*msgpb return nil } - err := kc.MetaKv.WalkWithPrefix(ChannelCheckpointPrefix, paginationSize, applyFn) + err := kc.MetaKv.WalkWithPrefix(ctx, ChannelCheckpointPrefix, paginationSize, applyFn) if err != nil { return nil, err } @@ -497,7 +497,7 @@ func (kc *Catalog) SaveChannelCheckpoint(ctx context.Context, vChannel string, p if err != nil { return err } - return kc.MetaKv.Save(k, string(v)) + return kc.MetaKv.Save(ctx, k, string(v)) } func (kc *Catalog) SaveChannelCheckpoints(ctx context.Context, positions []*msgpb.MsgPosition) error { @@ -510,15 +510,15 @@ func (kc *Catalog) SaveChannelCheckpoints(ctx context.Context, positions []*msgp } kvs[k] = string(v) } - return kc.SaveByBatch(kvs) + return kc.SaveByBatch(ctx, kvs) } func (kc *Catalog) DropChannelCheckpoint(ctx context.Context, vChannel string) error { k := buildChannelCPKey(vChannel) - return kc.MetaKv.Remove(k) + return kc.MetaKv.Remove(ctx, k) } -func (kc *Catalog) getBinlogsWithPrefix(binlogType storage.BinlogType, collectionID, partitionID, +func (kc *Catalog) getBinlogsWithPrefix(ctx context.Context, binlogType storage.BinlogType, collectionID, partitionID, segmentID typeutil.UniqueID, ) ([]string, []string, error) { var binlogPrefix string @@ -532,7 +532,7 @@ func (kc *Catalog) getBinlogsWithPrefix(binlogType storage.BinlogType, collectio default: return nil, nil, fmt.Errorf("invalid binlog type: %d", binlogType) } - keys, values, err := kc.MetaKv.LoadWithPrefix(binlogPrefix) + keys, values, err := kc.MetaKv.LoadWithPrefix(ctx, binlogPrefix) if err != nil { return nil, nil, err } @@ -547,7 +547,7 @@ func (kc *Catalog) CreateIndex(ctx context.Context, index *model.Index) error { return err } - err = kc.MetaKv.Save(key, string(value)) + err = kc.MetaKv.Save(ctx, key, string(value)) if err != nil { return err } @@ -568,7 +568,7 @@ func (kc *Catalog) ListIndexes(ctx context.Context) ([]*model.Index, error) { return nil } - err := kc.MetaKv.WalkWithPrefix(util.FieldIndexPrefix, paginationSize, applyFn) + err := kc.MetaKv.WalkWithPrefix(ctx, util.FieldIndexPrefix, paginationSize, applyFn) if err != nil { return nil, err } @@ -588,7 +588,7 @@ func (kc *Catalog) AlterIndexes(ctx context.Context, indexes []*model.Index) err kvs[key] = string(value) // TODO when we have better txn kv we should make this as a transaction if len(kvs) >= 64 { - err = kc.MetaKv.MultiSave(kvs) + err = kc.MetaKv.MultiSave(ctx, kvs) if err != nil { return err } @@ -596,7 +596,7 @@ func (kc *Catalog) AlterIndexes(ctx context.Context, indexes []*model.Index) err } } if len(kvs) != 0 { - return kc.MetaKv.MultiSave(kvs) + return kc.MetaKv.MultiSave(ctx, kvs) } return nil } @@ -604,7 +604,7 @@ func (kc *Catalog) AlterIndexes(ctx context.Context, indexes []*model.Index) err func (kc *Catalog) DropIndex(ctx context.Context, collID typeutil.UniqueID, dropIdxID typeutil.UniqueID) error { key := BuildIndexKey(collID, dropIdxID) - err := kc.MetaKv.Remove(key) + err := kc.MetaKv.Remove(ctx, key) if err != nil { log.Error("drop collection index meta fail", zap.Int64("collectionID", collID), zap.Int64("indexID", dropIdxID), zap.Error(err)) @@ -621,7 +621,7 @@ func (kc *Catalog) CreateSegmentIndex(ctx context.Context, segIdx *model.Segment if err != nil { return err } - err = kc.MetaKv.Save(key, string(value)) + err = kc.MetaKv.Save(ctx, key, string(value)) if err != nil { log.Error("failed to save segment index meta in etcd", zap.Int64("buildID", segIdx.BuildID), zap.Int64("segmentID", segIdx.SegmentID), zap.Error(err)) @@ -644,7 +644,7 @@ func (kc *Catalog) ListSegmentIndexes(ctx context.Context) ([]*model.SegmentInde return nil } - err := kc.MetaKv.WalkWithPrefix(util.SegmentIndexPrefix, paginationSize, applyFn) + err := kc.MetaKv.WalkWithPrefix(ctx, util.SegmentIndexPrefix, paginationSize, applyFn) if err != nil { return nil, err } @@ -662,13 +662,13 @@ func (kc *Catalog) AlterSegmentIndexes(ctx context.Context, segIdxes []*model.Se } kvs[key] = string(value) } - return kc.MetaKv.MultiSave(kvs) + return kc.MetaKv.MultiSave(ctx, kvs) } func (kc *Catalog) DropSegmentIndex(ctx context.Context, collID, partID, segID, buildID typeutil.UniqueID) error { key := BuildSegmentIndexKey(collID, partID, segID, buildID) - err := kc.MetaKv.Remove(key) + err := kc.MetaKv.Remove(ctx, key) if err != nil { log.Error("drop segment index meta fail", zap.Int64("buildID", buildID), zap.Error(err)) return err @@ -683,7 +683,7 @@ func (kc *Catalog) SaveImportJob(ctx context.Context, job *datapb.ImportJob) err if err != nil { return err } - return kc.MetaKv.Save(key, string(value)) + return kc.MetaKv.Save(ctx, key, string(value)) } func (kc *Catalog) ListImportJobs(ctx context.Context) ([]*datapb.ImportJob, error) { @@ -698,7 +698,7 @@ func (kc *Catalog) ListImportJobs(ctx context.Context) ([]*datapb.ImportJob, err return nil } - err := kc.MetaKv.WalkWithPrefix(ImportJobPrefix, paginationSize, applyFn) + err := kc.MetaKv.WalkWithPrefix(ctx, ImportJobPrefix, paginationSize, applyFn) if err != nil { return nil, err } @@ -707,7 +707,7 @@ func (kc *Catalog) ListImportJobs(ctx context.Context) ([]*datapb.ImportJob, err func (kc *Catalog) DropImportJob(ctx context.Context, jobID int64) error { key := buildImportJobKey(jobID) - return kc.MetaKv.Remove(key) + return kc.MetaKv.Remove(ctx, key) } func (kc *Catalog) SavePreImportTask(ctx context.Context, task *datapb.PreImportTask) error { @@ -716,7 +716,7 @@ func (kc *Catalog) SavePreImportTask(ctx context.Context, task *datapb.PreImport if err != nil { return err } - return kc.MetaKv.Save(key, string(value)) + return kc.MetaKv.Save(ctx, key, string(value)) } func (kc *Catalog) ListPreImportTasks(ctx context.Context) ([]*datapb.PreImportTask, error) { @@ -732,7 +732,7 @@ func (kc *Catalog) ListPreImportTasks(ctx context.Context) ([]*datapb.PreImportT return nil } - err := kc.MetaKv.WalkWithPrefix(PreImportTaskPrefix, paginationSize, applyFn) + err := kc.MetaKv.WalkWithPrefix(ctx, PreImportTaskPrefix, paginationSize, applyFn) if err != nil { return nil, err } @@ -741,7 +741,7 @@ func (kc *Catalog) ListPreImportTasks(ctx context.Context) ([]*datapb.PreImportT func (kc *Catalog) DropPreImportTask(ctx context.Context, taskID int64) error { key := buildPreImportTaskKey(taskID) - return kc.MetaKv.Remove(key) + return kc.MetaKv.Remove(ctx, key) } func (kc *Catalog) SaveImportTask(ctx context.Context, task *datapb.ImportTaskV2) error { @@ -750,7 +750,7 @@ func (kc *Catalog) SaveImportTask(ctx context.Context, task *datapb.ImportTaskV2 if err != nil { return err } - return kc.MetaKv.Save(key, string(value)) + return kc.MetaKv.Save(ctx, key, string(value)) } func (kc *Catalog) ListImportTasks(ctx context.Context) ([]*datapb.ImportTaskV2, error) { @@ -766,7 +766,7 @@ func (kc *Catalog) ListImportTasks(ctx context.Context) ([]*datapb.ImportTaskV2, return nil } - err := kc.MetaKv.WalkWithPrefix(ImportTaskPrefix, paginationSize, applyFn) + err := kc.MetaKv.WalkWithPrefix(ctx, ImportTaskPrefix, paginationSize, applyFn) if err != nil { return nil, err } @@ -775,7 +775,7 @@ func (kc *Catalog) ListImportTasks(ctx context.Context) ([]*datapb.ImportTaskV2, func (kc *Catalog) DropImportTask(ctx context.Context, taskID int64) error { key := buildImportTaskKey(taskID) - return kc.MetaKv.Remove(key) + return kc.MetaKv.Remove(ctx, key) } // GcConfirm returns true if related collection/partition is not found. @@ -785,7 +785,7 @@ func (kc *Catalog) GcConfirm(ctx context.Context, collectionID, partitionID type if partitionID != common.AllPartitionsID { prefix = buildPartitionPrefix(collectionID, partitionID) } - keys, values, err := kc.MetaKv.LoadWithPrefix(prefix) + keys, values, err := kc.MetaKv.LoadWithPrefix(ctx, prefix) if err != nil { // error case can be regarded as not finished. return false @@ -806,7 +806,7 @@ func (kc *Catalog) ListCompactionTask(ctx context.Context) ([]*datapb.Compaction return nil } - err := kc.MetaKv.WalkWithPrefix(CompactionTaskPrefix, paginationSize, applyFn) + err := kc.MetaKv.WalkWithPrefix(ctx, CompactionTaskPrefix, paginationSize, applyFn) if err != nil { return nil, err } @@ -824,12 +824,12 @@ func (kc *Catalog) SaveCompactionTask(ctx context.Context, coll *datapb.Compacti } kvs := make(map[string]string) kvs[k] = v - return kc.SaveByBatch(kvs) + return kc.SaveByBatch(ctx, kvs) } func (kc *Catalog) DropCompactionTask(ctx context.Context, task *datapb.CompactionTask) error { key := buildCompactionTaskPath(task) - return kc.MetaKv.Remove(key) + return kc.MetaKv.Remove(ctx, key) } func (kc *Catalog) ListAnalyzeTasks(ctx context.Context) ([]*indexpb.AnalyzeTask, error) { @@ -845,7 +845,7 @@ func (kc *Catalog) ListAnalyzeTasks(ctx context.Context) ([]*indexpb.AnalyzeTask return nil } - err := kc.MetaKv.WalkWithPrefix(AnalyzeTaskPrefix, paginationSize, applyFn) + err := kc.MetaKv.WalkWithPrefix(ctx, AnalyzeTaskPrefix, paginationSize, applyFn) if err != nil { return nil, err } @@ -860,7 +860,7 @@ func (kc *Catalog) SaveAnalyzeTask(ctx context.Context, task *indexpb.AnalyzeTas return err } - err = kc.MetaKv.Save(key, string(value)) + err = kc.MetaKv.Save(ctx, key, string(value)) if err != nil { return err } @@ -869,7 +869,7 @@ func (kc *Catalog) SaveAnalyzeTask(ctx context.Context, task *indexpb.AnalyzeTas func (kc *Catalog) DropAnalyzeTask(ctx context.Context, taskID typeutil.UniqueID) error { key := buildAnalyzeTaskKey(taskID) - return kc.MetaKv.Remove(key) + return kc.MetaKv.Remove(ctx, key) } func (kc *Catalog) ListPartitionStatsInfos(ctx context.Context) ([]*datapb.PartitionStatsInfo, error) { @@ -885,7 +885,7 @@ func (kc *Catalog) ListPartitionStatsInfos(ctx context.Context) ([]*datapb.Parti return nil } - err := kc.MetaKv.WalkWithPrefix(PartitionStatsInfoPrefix, paginationSize, applyFn) + err := kc.MetaKv.WalkWithPrefix(ctx, PartitionStatsInfoPrefix, paginationSize, applyFn) if err != nil { return nil, err } @@ -903,23 +903,23 @@ func (kc *Catalog) SavePartitionStatsInfo(ctx context.Context, coll *datapb.Part } kvs := make(map[string]string) kvs[k] = v - return kc.SaveByBatch(kvs) + return kc.SaveByBatch(ctx, kvs) } func (kc *Catalog) DropPartitionStatsInfo(ctx context.Context, info *datapb.PartitionStatsInfo) error { key := buildPartitionStatsInfoPath(info) - return kc.MetaKv.Remove(key) + return kc.MetaKv.Remove(ctx, key) } func (kc *Catalog) SaveCurrentPartitionStatsVersion(ctx context.Context, collID, partID int64, vChannel string, currentVersion int64) error { key := buildCurrentPartitionStatsVersionPath(collID, partID, vChannel) value := strconv.FormatInt(currentVersion, 10) - return kc.MetaKv.Save(key, value) + return kc.MetaKv.Save(ctx, key, value) } func (kc *Catalog) GetCurrentPartitionStatsVersion(ctx context.Context, collID, partID int64, vChannel string) (int64, error) { key := buildCurrentPartitionStatsVersionPath(collID, partID, vChannel) - valueStr, err := kc.MetaKv.Load(key) + valueStr, err := kc.MetaKv.Load(ctx, key) if err != nil { if errors.Is(err, merr.ErrIoKeyNotFound) { return 0, nil @@ -932,7 +932,7 @@ func (kc *Catalog) GetCurrentPartitionStatsVersion(ctx context.Context, collID, func (kc *Catalog) DropCurrentPartitionStatsVersion(ctx context.Context, collID, partID int64, vChannel string) error { key := buildCurrentPartitionStatsVersionPath(collID, partID, vChannel) - return kc.MetaKv.Remove(key) + return kc.MetaKv.Remove(ctx, key) } func (kc *Catalog) ListStatsTasks(ctx context.Context) ([]*indexpb.StatsTask, error) { @@ -948,7 +948,7 @@ func (kc *Catalog) ListStatsTasks(ctx context.Context) ([]*indexpb.StatsTask, er return nil } - err := kc.MetaKv.WalkWithPrefix(StatsTaskPrefix, paginationSize, applyFn) + err := kc.MetaKv.WalkWithPrefix(ctx, StatsTaskPrefix, paginationSize, applyFn) if err != nil { return nil, err } @@ -962,7 +962,7 @@ func (kc *Catalog) SaveStatsTask(ctx context.Context, task *indexpb.StatsTask) e return err } - err = kc.MetaKv.Save(key, string(value)) + err = kc.MetaKv.Save(ctx, key, string(value)) if err != nil { return err } @@ -971,5 +971,5 @@ func (kc *Catalog) SaveStatsTask(ctx context.Context, task *indexpb.StatsTask) e func (kc *Catalog) DropStatsTask(ctx context.Context, taskID typeutil.UniqueID) error { key := buildStatsTaskKey(taskID) - return kc.MetaKv.Remove(key) + return kc.MetaKv.Remove(ctx, key) } diff --git a/internal/metastore/kv/datacoord/kv_catalog_test.go b/internal/metastore/kv/datacoord/kv_catalog_test.go index baeb8bedd385c..992aeb0830aed 100644 --- a/internal/metastore/kv/datacoord/kv_catalog_test.go +++ b/internal/metastore/kv/datacoord/kv_catalog_test.go @@ -148,7 +148,7 @@ var ( func Test_ListSegments(t *testing.T) { t.Run("load failed", func(t *testing.T) { metakv := mocks.NewMetaKv(t) - metakv.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return(errors.New("error")) + metakv.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(errors.New("error")) catalog := NewCatalog(metakv, rootPath, "") ret, err := catalog.ListSegments(context.TODO()) @@ -190,7 +190,7 @@ func Test_ListSegments(t *testing.T) { assert.NoError(t, err) metakv := mocks.NewMetaKv(t) - metakv.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(s string, i int, f func([]byte, []byte) error) error { + metakv.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, s string, i int, f func([]byte, []byte) error) error { if strings.HasPrefix(k5, s) { return f([]byte(k5), segBytes) } @@ -207,7 +207,7 @@ func Test_ListSegments(t *testing.T) { t.Run("test compatibility with stats task", func(t *testing.T) { metakv := mocks.NewMetaKv(t) - metakv.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(s string, i int, f func([]byte, []byte) error) error { + metakv.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, s string, i int, f func([]byte, []byte) error) error { if strings.HasPrefix(k9, path.Join(s)) { return f([]byte(k9), nil) } @@ -225,7 +225,7 @@ func Test_ListSegments(t *testing.T) { var savedKvs map[string]string metakv := mocks.NewMetaKv(t) - metakv.EXPECT().MultiSave(mock.Anything).RunAndReturn(func(m map[string]string) error { + metakv.EXPECT().MultiSave(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, m map[string]string) error { savedKvs = m return nil }) @@ -234,7 +234,7 @@ func Test_ListSegments(t *testing.T) { err := catalog.AddSegment(context.TODO(), segment1) assert.NoError(t, err) - metakv.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(s string, i int, f func([]byte, []byte) error) error { + metakv.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, s string, i int, f func([]byte, []byte) error) error { if strings.HasPrefix(k5, s) { return f([]byte(k5), []byte(savedKvs[k5])) } @@ -267,7 +267,7 @@ func Test_ListSegments(t *testing.T) { func Test_AddSegments(t *testing.T) { t.Run("generate binlog kvs failed", func(t *testing.T) { metakv := mocks.NewMetaKv(t) - metakv.EXPECT().MultiSave(mock.Anything).Return(errors.New("error")).Maybe() + metakv.EXPECT().MultiSave(mock.Anything, mock.Anything).Return(errors.New("error")).Maybe() catalog := NewCatalog(metakv, rootPath, "") @@ -277,7 +277,7 @@ func Test_AddSegments(t *testing.T) { t.Run("save error", func(t *testing.T) { metakv := mocks.NewMetaKv(t) - metakv.EXPECT().MultiSave(mock.Anything).Return(errors.New("error")) + metakv.EXPECT().MultiSave(mock.Anything, mock.Anything).Return(errors.New("error")) catalog := NewCatalog(metakv, rootPath, "") err := catalog.AddSegment(context.TODO(), segment1) @@ -287,11 +287,11 @@ func Test_AddSegments(t *testing.T) { t.Run("save successfully", func(t *testing.T) { savedKvs := make(map[string]string) metakv := mocks.NewMetaKv(t) - metakv.EXPECT().MultiSave(mock.Anything).RunAndReturn(func(m map[string]string) error { + metakv.EXPECT().MultiSave(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, m map[string]string) error { savedKvs = m return nil }) - metakv.EXPECT().Load(mock.Anything).RunAndReturn(func(s string) (string, error) { + metakv.EXPECT().Load(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, s string) (string, error) { if v, ok := savedKvs[s]; ok { return v, nil } @@ -301,7 +301,7 @@ func Test_AddSegments(t *testing.T) { catalog := NewCatalog(metakv, rootPath, "") err := catalog.AddSegment(context.TODO(), segment1) assert.NoError(t, err) - adjustedSeg, err := catalog.LoadFromSegmentPath(segment1.CollectionID, segment1.PartitionID, segment1.ID) + adjustedSeg, err := catalog.LoadFromSegmentPath(context.TODO(), segment1.CollectionID, segment1.PartitionID, segment1.ID) assert.NoError(t, err) // Check that num of rows is corrected from 100 to 5. assert.Equal(t, int64(100), segment1.GetNumOfRows()) @@ -364,7 +364,7 @@ func Test_AddSegments(t *testing.T) { func Test_AlterSegments(t *testing.T) { t.Run("generate binlog kvs failed", func(t *testing.T) { metakv := mocks.NewMetaKv(t) - metakv.EXPECT().MultiSave(mock.Anything).Return(errors.New("error")).Maybe() + metakv.EXPECT().MultiSave(mock.Anything, mock.Anything).Return(errors.New("error")).Maybe() catalog := NewCatalog(metakv, rootPath, "") err := catalog.AlterSegments(context.TODO(), []*datapb.SegmentInfo{invalidSegment}, metastore.BinlogsIncrement{ @@ -375,7 +375,7 @@ func Test_AlterSegments(t *testing.T) { t.Run("save error", func(t *testing.T) { metakv := mocks.NewMetaKv(t) - metakv.EXPECT().MultiSave(mock.Anything).Return(errors.New("error")) + metakv.EXPECT().MultiSave(mock.Anything, mock.Anything).Return(errors.New("error")) catalog := NewCatalog(metakv, rootPath, "") err := catalog.AlterSegments(context.TODO(), []*datapb.SegmentInfo{segment1}) @@ -386,7 +386,7 @@ func Test_AlterSegments(t *testing.T) { t.Run("save successfully", func(t *testing.T) { var savedKvs map[string]string metakv := mocks.NewMetaKv(t) - metakv.EXPECT().MultiSave(mock.Anything).RunAndReturn(func(m map[string]string) error { + metakv.EXPECT().MultiSave(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, m map[string]string) error { savedKvs = m return nil }) @@ -410,12 +410,12 @@ func Test_AlterSegments(t *testing.T) { savedKvs := make(map[string]string) opGroupCount := 0 metakv := mocks.NewMetaKv(t) - metakv.EXPECT().MultiSave(mock.Anything).RunAndReturn(func(m map[string]string) error { + metakv.EXPECT().MultiSave(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, m map[string]string) error { maps.Copy(savedKvs, m) opGroupCount++ return nil }) - metakv.EXPECT().Load(mock.Anything).RunAndReturn(func(s string) (string, error) { + metakv.EXPECT().Load(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, s string) (string, error) { if v, ok := savedKvs[s]; ok { return v, nil } @@ -458,7 +458,7 @@ func Test_AlterSegments(t *testing.T) { assert.Equal(t, 255+3, len(savedKvs)) assert.Equal(t, 3, opGroupCount) - adjustedSeg, err := catalog.LoadFromSegmentPath(segmentXL.CollectionID, segmentXL.PartitionID, segmentXL.ID) + adjustedSeg, err := catalog.LoadFromSegmentPath(context.TODO(), segmentXL.CollectionID, segmentXL.PartitionID, segmentXL.ID) assert.NoError(t, err) // Check that num of rows is corrected from 100 to 1275. assert.Equal(t, int64(100), segmentXL.GetNumOfRows()) @@ -507,7 +507,7 @@ func Test_AlterSegments(t *testing.T) { func Test_DropSegment(t *testing.T) { t.Run("remove failed", func(t *testing.T) { metakv := mocks.NewMetaKv(t) - metakv.EXPECT().MultiSaveAndRemoveWithPrefix(mock.Anything, mock.Anything).Return(errors.New("error")) + metakv.EXPECT().MultiSaveAndRemoveWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return(errors.New("error")) catalog := NewCatalog(metakv, rootPath, "") err := catalog.DropSegment(context.TODO(), segment1) @@ -517,7 +517,7 @@ func Test_DropSegment(t *testing.T) { t.Run("remove successfully", func(t *testing.T) { removedKvs := make(map[string]struct{}, 0) metakv := mocks.NewMetaKv(t) - metakv.EXPECT().MultiSaveAndRemoveWithPrefix(mock.Anything, mock.Anything).RunAndReturn(func(m map[string]string, s []string, p ...predicates.Predicate) error { + metakv.EXPECT().MultiSaveAndRemoveWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, m map[string]string, s []string, p ...predicates.Predicate) error { for _, key := range s { removedKvs[key] = struct{}{} } @@ -543,7 +543,7 @@ func Test_DropSegment(t *testing.T) { func Test_SaveDroppedSegmentsInBatch_SaveError(t *testing.T) { txn := mocks.NewMetaKv(t) - txn.EXPECT().MultiSave(mock.Anything).Return(errors.New("mock error")) + txn.EXPECT().MultiSave(mock.Anything, mock.Anything).Return(errors.New("mock error")) catalog := NewCatalog(txn, rootPath, "") segments := []*datapb.SegmentInfo{ @@ -563,8 +563,8 @@ func Test_SaveDroppedSegmentsInBatch_MultiSave(t *testing.T) { ) txn := mocks.NewMetaKv(t) txn.EXPECT(). - MultiSave(mock.Anything). - Run(func(kvs map[string]string) { + MultiSave(mock.Anything, mock.Anything). + Run(func(ctx context.Context, kvs map[string]string) { count++ kvSize += len(kvs) }). @@ -632,12 +632,12 @@ func TestChannelCP(t *testing.T) { t.Run("ListChannelCheckpoint", func(t *testing.T) { txn := mocks.NewMetaKv(t) - txn.EXPECT().Save(mock.Anything, mock.Anything).Return(nil) + txn.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(nil) catalog := NewCatalog(txn, rootPath, "") err := catalog.SaveChannelCheckpoint(context.TODO(), mockVChannel, pos) assert.NoError(t, err) - txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ string, _ int, f func([]byte, []byte) error) error { + txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, _ string, _ int, f func([]byte, []byte) error) error { return f([]byte(k), v) }) res, err := catalog.ListChannelCheckpoint(context.TODO()) @@ -648,14 +648,14 @@ func TestChannelCP(t *testing.T) { t.Run("ListChannelCheckpoint failed", func(t *testing.T) { txn := mocks.NewMetaKv(t) catalog := NewCatalog(txn, rootPath, "") - txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return(errors.New("mock error")) + txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(errors.New("mock error")) _, err = catalog.ListChannelCheckpoint(context.TODO()) assert.Error(t, err) }) t.Run("SaveChannelCheckpoint", func(t *testing.T) { txn := mocks.NewMetaKv(t) - txn.EXPECT().Save(mock.Anything, mock.Anything).Return(nil) + txn.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(nil) catalog := NewCatalog(txn, rootPath, "") err := catalog.SaveChannelCheckpoint(context.TODO(), mockVChannel, pos) assert.NoError(t, err) @@ -664,14 +664,14 @@ func TestChannelCP(t *testing.T) { t.Run("SaveChannelCheckpoint failed", func(t *testing.T) { txn := mocks.NewMetaKv(t) catalog := NewCatalog(txn, rootPath, "") - txn.EXPECT().Save(mock.Anything, mock.Anything).Return(errors.New("mock error")) + txn.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(errors.New("mock error")) err = catalog.SaveChannelCheckpoint(context.TODO(), mockVChannel, &msgpb.MsgPosition{}) assert.Error(t, err) }) t.Run("SaveChannelCheckpoints", func(t *testing.T) { txn := mocks.NewMetaKv(t) - txn.EXPECT().MultiSave(mock.Anything).Return(nil) + txn.EXPECT().MultiSave(mock.Anything, mock.Anything).Return(nil) catalog := NewCatalog(txn, rootPath, "") err := catalog.SaveChannelCheckpoints(context.TODO(), []*msgpb.MsgPosition{pos}) assert.NoError(t, err) @@ -680,20 +680,20 @@ func TestChannelCP(t *testing.T) { t.Run("SaveChannelCheckpoints failed", func(t *testing.T) { txn := mocks.NewMetaKv(t) catalog := NewCatalog(txn, rootPath, "") - txn.EXPECT().MultiSave(mock.Anything).Return(errors.New("mock error")) + txn.EXPECT().MultiSave(mock.Anything, mock.Anything).Return(errors.New("mock error")) err = catalog.SaveChannelCheckpoints(context.TODO(), []*msgpb.MsgPosition{pos}) assert.Error(t, err) }) t.Run("DropChannelCheckpoint", func(t *testing.T) { txn := mocks.NewMetaKv(t) - txn.EXPECT().Save(mock.Anything, mock.Anything).Return(nil) + txn.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(nil) catalog := NewCatalog(txn, rootPath, "") err := catalog.SaveChannelCheckpoint(context.TODO(), mockVChannel, pos) assert.NoError(t, err) - txn.EXPECT().Remove(mock.Anything).Return(nil) - txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return(nil) + txn.EXPECT().Remove(mock.Anything, mock.Anything).Return(nil) + txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil) err = catalog.DropChannelCheckpoint(context.TODO(), mockVChannel) assert.NoError(t, err) res, err := catalog.ListChannelCheckpoint(context.TODO()) @@ -704,7 +704,7 @@ func TestChannelCP(t *testing.T) { t.Run("DropChannelCheckpoint failed", func(t *testing.T) { txn := mocks.NewMetaKv(t) catalog := NewCatalog(txn, rootPath, "") - txn.EXPECT().Remove(mock.Anything).Return(errors.New("mock error")) + txn.EXPECT().Remove(mock.Anything, mock.Anything).Return(errors.New("mock error")) err = catalog.DropChannelCheckpoint(context.TODO(), mockVChannel) assert.Error(t, err) }) @@ -713,7 +713,7 @@ func TestChannelCP(t *testing.T) { func Test_MarkChannelDeleted_SaveError(t *testing.T) { txn := mocks.NewMetaKv(t) txn.EXPECT(). - Save(mock.Anything, mock.Anything). + Save(mock.Anything, mock.Anything, mock.Anything). Return(errors.New("mock error")) catalog := NewCatalog(txn, rootPath, "") @@ -724,7 +724,7 @@ func Test_MarkChannelDeleted_SaveError(t *testing.T) { func Test_MarkChannelAdded_SaveError(t *testing.T) { txn := mocks.NewMetaKv(t) txn.EXPECT(). - Save(mock.Anything, mock.Anything). + Save(mock.Anything, mock.Anything, mock.Anything). Return(errors.New("mock error")) catalog := NewCatalog(txn, rootPath, "") @@ -735,7 +735,7 @@ func Test_MarkChannelAdded_SaveError(t *testing.T) { func Test_ChannelExists_SaveError(t *testing.T) { txn := mocks.NewMetaKv(t) txn.EXPECT(). - Load(mock.Anything). + Load(mock.Anything, mock.Anything). Return("", errors.New("mock error")) catalog := NewCatalog(txn, rootPath, "") @@ -854,7 +854,7 @@ func verifySavedKvsForDroppedSegment(t *testing.T, savedKvs map[string]string) { func TestCatalog_CreateIndex(t *testing.T) { t.Run("success", func(t *testing.T) { metakv := mocks.NewMetaKv(t) - metakv.EXPECT().Save(mock.Anything, mock.Anything).Return(nil) + metakv.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(nil) catalog := &Catalog{ MetaKv: metakv, @@ -866,7 +866,7 @@ func TestCatalog_CreateIndex(t *testing.T) { t.Run("failed", func(t *testing.T) { metakv := mocks.NewMetaKv(t) - metakv.EXPECT().Save(mock.Anything, mock.Anything).Return(errors.New("error")) + metakv.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(errors.New("error")) catalog := &Catalog{ MetaKv: metakv, @@ -880,7 +880,7 @@ func TestCatalog_CreateIndex(t *testing.T) { func TestCatalog_ListIndexes(t *testing.T) { t.Run("success", func(t *testing.T) { metakv := mocks.NewMetaKv(t) - metakv.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ string, _ int, f func([]byte, []byte) error) error { + metakv.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, _ string, _ int, f func([]byte, []byte) error) error { i := &indexpb.FieldIndex{ IndexInfo: &indexpb.IndexInfo{ CollectionID: 0, @@ -908,7 +908,7 @@ func TestCatalog_ListIndexes(t *testing.T) { t.Run("failed", func(t *testing.T) { txn := mocks.NewMetaKv(t) - txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return(errors.New("error")) + txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(errors.New("error")) catalog := &Catalog{ MetaKv: txn, } @@ -918,7 +918,7 @@ func TestCatalog_ListIndexes(t *testing.T) { t.Run("unmarshal failed", func(t *testing.T) { txn := mocks.NewMetaKv(t) - txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ string, _ int, f func([]byte, []byte) error) error { + txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, _ string, _ int, f func([]byte, []byte) error) error { return f([]byte("1"), []byte("invalid")) }) @@ -943,7 +943,7 @@ func TestCatalog_AlterIndexes(t *testing.T) { } metakv := mocks.NewMetaKv(t) - metakv.EXPECT().MultiSave(mock.Anything).Return(nil) + metakv.EXPECT().MultiSave(mock.Anything, mock.Anything).Return(nil) catalog := &Catalog{ MetaKv: metakv, } @@ -970,14 +970,14 @@ func TestCatalog_AlterSegmentMultiIndexes(t *testing.T) { t.Run("batchAdd", func(t *testing.T) { metakv := mocks.NewMetaKv(t) - metakv.EXPECT().MultiSave(mock.Anything).RunAndReturn(func(m map[string]string) error { + metakv.EXPECT().MultiSave(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, m map[string]string) error { if len(m) > 64 { return errors.New("fail") } return nil }) - metakv.EXPECT().MultiSave(mock.Anything).Return(nil) + metakv.EXPECT().MultiSave(mock.Anything, mock.Anything).Return(nil) catalog := &Catalog{ MetaKv: metakv, } @@ -990,7 +990,7 @@ func TestCatalog_AlterSegmentMultiIndexes(t *testing.T) { func TestCatalog_DropIndex(t *testing.T) { t.Run("success", func(t *testing.T) { metakv := mocks.NewMetaKv(t) - metakv.EXPECT().Remove(mock.Anything).Return(nil) + metakv.EXPECT().Remove(mock.Anything, mock.Anything).Return(nil) catalog := &Catalog{ MetaKv: metakv, } @@ -1001,7 +1001,7 @@ func TestCatalog_DropIndex(t *testing.T) { t.Run("failed", func(t *testing.T) { metakv := mocks.NewMetaKv(t) - metakv.EXPECT().Remove(mock.Anything).Return(errors.New("error")) + metakv.EXPECT().Remove(mock.Anything, mock.Anything).Return(errors.New("error")) catalog := &Catalog{ MetaKv: metakv, } @@ -1031,7 +1031,7 @@ func TestCatalog_CreateSegmentIndex(t *testing.T) { t.Run("success", func(t *testing.T) { metakv := mocks.NewMetaKv(t) - metakv.EXPECT().Save(mock.Anything, mock.Anything).Return(nil) + metakv.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(nil) catalog := &Catalog{ MetaKv: metakv, } @@ -1042,7 +1042,7 @@ func TestCatalog_CreateSegmentIndex(t *testing.T) { t.Run("failed", func(t *testing.T) { metakv := mocks.NewMetaKv(t) - metakv.EXPECT().Save(mock.Anything, mock.Anything).Return(errors.New("error")) + metakv.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(errors.New("error")) catalog := &Catalog{ MetaKv: metakv, } @@ -1074,7 +1074,7 @@ func TestCatalog_ListSegmentIndexes(t *testing.T) { assert.NoError(t, err) metakv := mocks.NewMetaKv(t) - metakv.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ string, _ int, f func([]byte, []byte) error) error { + metakv.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, _ string, _ int, f func([]byte, []byte) error) error { return f([]byte("key"), v) }) catalog := &Catalog{ @@ -1088,7 +1088,7 @@ func TestCatalog_ListSegmentIndexes(t *testing.T) { t.Run("failed", func(t *testing.T) { metakv := mocks.NewMetaKv(t) - metakv.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return(errors.New("error")) + metakv.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(errors.New("error")) catalog := &Catalog{ MetaKv: metakv, } @@ -1099,7 +1099,7 @@ func TestCatalog_ListSegmentIndexes(t *testing.T) { t.Run("unmarshal failed", func(t *testing.T) { metakv := mocks.NewMetaKv(t) - metakv.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ string, _ int, f func([]byte, []byte) error) error { + metakv.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, _ string, _ int, f func([]byte, []byte) error) error { return f([]byte("key"), []byte("invalid")) }) catalog := &Catalog{ @@ -1131,7 +1131,7 @@ func TestCatalog_AlterSegmentIndexes(t *testing.T) { t.Run("add", func(t *testing.T) { metakv := mocks.NewMetaKv(t) - metakv.EXPECT().MultiSave(mock.Anything).Return(nil) + metakv.EXPECT().MultiSave(mock.Anything, mock.Anything).Return(nil) catalog := &Catalog{ MetaKv: metakv, } @@ -1144,7 +1144,7 @@ func TestCatalog_AlterSegmentIndexes(t *testing.T) { func TestCatalog_DropSegmentIndex(t *testing.T) { t.Run("success", func(t *testing.T) { metakv := mocks.NewMetaKv(t) - metakv.EXPECT().Remove(mock.Anything).Return(nil) + metakv.EXPECT().Remove(mock.Anything, mock.Anything).Return(nil) catalog := &Catalog{ MetaKv: metakv, } @@ -1155,7 +1155,7 @@ func TestCatalog_DropSegmentIndex(t *testing.T) { t.Run("fail", func(t *testing.T) { metakv := mocks.NewMetaKv(t) - metakv.EXPECT().Remove(mock.Anything).Return(errors.New("error")) + metakv.EXPECT().Remove(mock.Anything, mock.Anything).Return(errors.New("error")) catalog := &Catalog{ MetaKv: metakv, } @@ -1336,12 +1336,14 @@ func TestCatalog_GcConfirm(t *testing.T) { kc.MetaKv = txn txn.On("LoadWithPrefix", + mock.Anything, mock.AnythingOfType("string")). Return(nil, nil, errors.New("error mock LoadWithPrefix")). Once() assert.False(t, kc.GcConfirm(context.TODO(), 100, 10000)) txn.On("LoadWithPrefix", + mock.Anything, mock.AnythingOfType("string")). Return(nil, nil, nil) assert.True(t, kc.GcConfirm(context.TODO(), 100, 10000)) @@ -1365,7 +1367,7 @@ func TestCatalog_Import(t *testing.T) { t.Run("SaveImportJob", func(t *testing.T) { txn := mocks.NewMetaKv(t) - txn.EXPECT().Save(mock.Anything, mock.Anything).Return(nil) + txn.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(nil) kc.MetaKv = txn err := kc.SaveImportJob(context.TODO(), job) assert.NoError(t, err) @@ -1374,7 +1376,7 @@ func TestCatalog_Import(t *testing.T) { assert.NoError(t, err) txn = mocks.NewMetaKv(t) - txn.EXPECT().Save(mock.Anything, mock.Anything).Return(mockErr) + txn.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(mockErr) kc.MetaKv = txn err = kc.SaveImportJob(context.TODO(), job) assert.Error(t, err) @@ -1384,7 +1386,7 @@ func TestCatalog_Import(t *testing.T) { txn := mocks.NewMetaKv(t) value, err := proto.Marshal(job) assert.NoError(t, err) - txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ string, _ int, f func([]byte, []byte) error) error { + txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, _ string, _ int, f func([]byte, []byte) error) error { return f(nil, value) }) kc.MetaKv = txn @@ -1393,7 +1395,7 @@ func TestCatalog_Import(t *testing.T) { assert.Equal(t, 1, len(jobs)) txn = mocks.NewMetaKv(t) - txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ string, _ int, f func([]byte, []byte) error) error { + txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, _ string, _ int, f func([]byte, []byte) error) error { return f(nil, []byte("@#%#^#")) }) kc.MetaKv = txn @@ -1401,7 +1403,7 @@ func TestCatalog_Import(t *testing.T) { assert.Error(t, err) txn = mocks.NewMetaKv(t) - txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return(mockErr) + txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(mockErr) kc.MetaKv = txn _, err = kc.ListImportJobs(context.TODO()) assert.Error(t, err) @@ -1409,13 +1411,13 @@ func TestCatalog_Import(t *testing.T) { t.Run("DropImportJob", func(t *testing.T) { txn := mocks.NewMetaKv(t) - txn.EXPECT().Remove(mock.Anything).Return(nil) + txn.EXPECT().Remove(mock.Anything, mock.Anything).Return(nil) kc.MetaKv = txn err := kc.DropImportJob(context.TODO(), job.GetJobID()) assert.NoError(t, err) txn = mocks.NewMetaKv(t) - txn.EXPECT().Remove(mock.Anything).Return(mockErr) + txn.EXPECT().Remove(mock.Anything, mock.Anything).Return(mockErr) kc.MetaKv = txn err = kc.DropImportJob(context.TODO(), job.GetJobID()) assert.Error(t, err) @@ -1423,7 +1425,7 @@ func TestCatalog_Import(t *testing.T) { t.Run("SavePreImportTask", func(t *testing.T) { txn := mocks.NewMetaKv(t) - txn.EXPECT().Save(mock.Anything, mock.Anything).Return(nil) + txn.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(nil) kc.MetaKv = txn err := kc.SavePreImportTask(context.TODO(), pit) assert.NoError(t, err) @@ -1432,7 +1434,7 @@ func TestCatalog_Import(t *testing.T) { assert.NoError(t, err) txn = mocks.NewMetaKv(t) - txn.EXPECT().Save(mock.Anything, mock.Anything).Return(mockErr) + txn.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(mockErr) kc.MetaKv = txn err = kc.SavePreImportTask(context.TODO(), pit) assert.Error(t, err) @@ -1442,7 +1444,7 @@ func TestCatalog_Import(t *testing.T) { txn := mocks.NewMetaKv(t) value, err := proto.Marshal(pit) assert.NoError(t, err) - txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ string, _ int, f func([]byte, []byte) error) error { + txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, _ string, _ int, f func([]byte, []byte) error) error { return f(nil, value) }) kc.MetaKv = txn @@ -1451,7 +1453,7 @@ func TestCatalog_Import(t *testing.T) { assert.Equal(t, 1, len(tasks)) txn = mocks.NewMetaKv(t) - txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ string, _ int, f func([]byte, []byte) error) error { + txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, _ string, _ int, f func([]byte, []byte) error) error { return f(nil, []byte("@#%#^#")) }) kc.MetaKv = txn @@ -1459,7 +1461,7 @@ func TestCatalog_Import(t *testing.T) { assert.Error(t, err) txn = mocks.NewMetaKv(t) - txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return(mockErr) + txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(mockErr) kc.MetaKv = txn _, err = kc.ListPreImportTasks(context.TODO()) assert.Error(t, err) @@ -1467,13 +1469,13 @@ func TestCatalog_Import(t *testing.T) { t.Run("DropPreImportTask", func(t *testing.T) { txn := mocks.NewMetaKv(t) - txn.EXPECT().Remove(mock.Anything).Return(nil) + txn.EXPECT().Remove(mock.Anything, mock.Anything).Return(nil) kc.MetaKv = txn err := kc.DropPreImportTask(context.TODO(), pit.GetTaskID()) assert.NoError(t, err) txn = mocks.NewMetaKv(t) - txn.EXPECT().Remove(mock.Anything).Return(mockErr) + txn.EXPECT().Remove(mock.Anything, mock.Anything).Return(mockErr) kc.MetaKv = txn err = kc.DropPreImportTask(context.TODO(), pit.GetTaskID()) assert.Error(t, err) @@ -1481,7 +1483,7 @@ func TestCatalog_Import(t *testing.T) { t.Run("SaveImportTask", func(t *testing.T) { txn := mocks.NewMetaKv(t) - txn.EXPECT().Save(mock.Anything, mock.Anything).Return(nil) + txn.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(nil) kc.MetaKv = txn err := kc.SaveImportTask(context.TODO(), it) assert.NoError(t, err) @@ -1490,7 +1492,7 @@ func TestCatalog_Import(t *testing.T) { assert.NoError(t, err) txn = mocks.NewMetaKv(t) - txn.EXPECT().Save(mock.Anything, mock.Anything).Return(mockErr) + txn.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(mockErr) kc.MetaKv = txn err = kc.SaveImportTask(context.TODO(), it) assert.Error(t, err) @@ -1500,7 +1502,7 @@ func TestCatalog_Import(t *testing.T) { txn := mocks.NewMetaKv(t) value, err := proto.Marshal(it) assert.NoError(t, err) - txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ string, _ int, f func([]byte, []byte) error) error { + txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, _ string, _ int, f func([]byte, []byte) error) error { return f(nil, value) }) kc.MetaKv = txn @@ -1509,7 +1511,7 @@ func TestCatalog_Import(t *testing.T) { assert.Equal(t, 1, len(tasks)) txn = mocks.NewMetaKv(t) - txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ string, _ int, f func([]byte, []byte) error) error { + txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, _ string, _ int, f func([]byte, []byte) error) error { return f(nil, []byte("@#%#^#")) }) kc.MetaKv = txn @@ -1517,7 +1519,7 @@ func TestCatalog_Import(t *testing.T) { assert.Error(t, err) txn = mocks.NewMetaKv(t) - txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return(mockErr) + txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(mockErr) kc.MetaKv = txn _, err = kc.ListImportTasks(context.TODO()) assert.Error(t, err) @@ -1525,13 +1527,13 @@ func TestCatalog_Import(t *testing.T) { t.Run("DropImportTask", func(t *testing.T) { txn := mocks.NewMetaKv(t) - txn.EXPECT().Remove(mock.Anything).Return(nil) + txn.EXPECT().Remove(mock.Anything, mock.Anything).Return(nil) kc.MetaKv = txn err := kc.DropImportTask(context.TODO(), it.GetTaskID()) assert.NoError(t, err) txn = mocks.NewMetaKv(t) - txn.EXPECT().Remove(mock.Anything).Return(mockErr) + txn.EXPECT().Remove(mock.Anything, mock.Anything).Return(mockErr) kc.MetaKv = txn err = kc.DropImportTask(context.TODO(), it.GetTaskID()) assert.Error(t, err) @@ -1544,7 +1546,7 @@ func TestCatalog_AnalyzeTask(t *testing.T) { t.Run("ListAnalyzeTasks", func(t *testing.T) { txn := mocks.NewMetaKv(t) - txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return(mockErr) + txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(mockErr) kc.MetaKv = txn tasks, err := kc.ListAnalyzeTasks(context.Background()) @@ -1570,7 +1572,7 @@ func TestCatalog_AnalyzeTask(t *testing.T) { assert.NoError(t, err) txn = mocks.NewMetaKv(t) - txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ string, _ int, f func([]byte, []byte) error) error { + txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, _ string, _ int, f func([]byte, []byte) error) error { return f([]byte("key1"), value) }) kc.MetaKv = txn @@ -1580,7 +1582,7 @@ func TestCatalog_AnalyzeTask(t *testing.T) { assert.Equal(t, 1, len(tasks)) txn = mocks.NewMetaKv(t) - txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ string, _ int, f func([]byte, []byte) error) error { + txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, _ string, _ int, f func([]byte, []byte) error) error { return f([]byte("key1"), []byte("1234")) }) kc.MetaKv = txn @@ -1608,14 +1610,14 @@ func TestCatalog_AnalyzeTask(t *testing.T) { } txn := mocks.NewMetaKv(t) - txn.EXPECT().Save(mock.Anything, mock.Anything).Return(nil) + txn.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(nil) kc.MetaKv = txn err := kc.SaveAnalyzeTask(context.Background(), task) assert.NoError(t, err) txn = mocks.NewMetaKv(t) - txn.EXPECT().Save(mock.Anything, mock.Anything).Return(mockErr) + txn.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(mockErr) kc.MetaKv = txn err = kc.SaveAnalyzeTask(context.Background(), task) @@ -1624,14 +1626,14 @@ func TestCatalog_AnalyzeTask(t *testing.T) { t.Run("DropAnalyzeTask", func(t *testing.T) { txn := mocks.NewMetaKv(t) - txn.EXPECT().Remove(mock.Anything).Return(nil) + txn.EXPECT().Remove(mock.Anything, mock.Anything).Return(nil) kc.MetaKv = txn err := kc.DropAnalyzeTask(context.Background(), 1) assert.NoError(t, err) txn = mocks.NewMetaKv(t) - txn.EXPECT().Remove(mock.Anything).Return(mockErr) + txn.EXPECT().Remove(mock.Anything, mock.Anything).Return(mockErr) kc.MetaKv = txn err = kc.DropAnalyzeTask(context.Background(), 1) @@ -1645,7 +1647,7 @@ func Test_PartitionStatsInfo(t *testing.T) { t.Run("ListPartitionStatsInfo", func(t *testing.T) { txn := mocks.NewMetaKv(t) - txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return(mockErr) + txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(mockErr) kc.MetaKv = txn infos, err := kc.ListPartitionStatsInfos(context.Background()) @@ -1665,7 +1667,7 @@ func Test_PartitionStatsInfo(t *testing.T) { assert.NoError(t, err) txn = mocks.NewMetaKv(t) - txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ string, _ int, f func([]byte, []byte) error) error { + txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, _ string, _ int, f func([]byte, []byte) error) error { return f([]byte("key1"), value) }) kc.MetaKv = txn @@ -1675,7 +1677,7 @@ func Test_PartitionStatsInfo(t *testing.T) { assert.Equal(t, 1, len(infos)) txn = mocks.NewMetaKv(t) - txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ string, _ int, f func([]byte, []byte) error) error { + txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, _ string, _ int, f func([]byte, []byte) error) error { return f([]byte("key1"), []byte("1234")) }) kc.MetaKv = txn @@ -1687,7 +1689,7 @@ func Test_PartitionStatsInfo(t *testing.T) { t.Run("SavePartitionStatsInfo", func(t *testing.T) { txn := mocks.NewMetaKv(t) - txn.EXPECT().MultiSave(mock.Anything).Return(mockErr) + txn.EXPECT().MultiSave(mock.Anything, mock.Anything).Return(mockErr) kc.MetaKv = txn info := &datapb.PartitionStatsInfo{ @@ -1704,7 +1706,7 @@ func Test_PartitionStatsInfo(t *testing.T) { assert.Error(t, err) txn = mocks.NewMetaKv(t) - txn.EXPECT().MultiSave(mock.Anything).Return(nil) + txn.EXPECT().MultiSave(mock.Anything, mock.Anything).Return(nil) kc.MetaKv = txn err = kc.SavePartitionStatsInfo(context.Background(), info) @@ -1713,7 +1715,7 @@ func Test_PartitionStatsInfo(t *testing.T) { t.Run("DropPartitionStatsInfo", func(t *testing.T) { txn := mocks.NewMetaKv(t) - txn.EXPECT().Remove(mock.Anything).Return(mockErr) + txn.EXPECT().Remove(mock.Anything, mock.Anything).Return(mockErr) kc.MetaKv = txn info := &datapb.PartitionStatsInfo{ @@ -1730,7 +1732,7 @@ func Test_PartitionStatsInfo(t *testing.T) { assert.Error(t, err) txn = mocks.NewMetaKv(t) - txn.EXPECT().Remove(mock.Anything).Return(nil) + txn.EXPECT().Remove(mock.Anything, mock.Anything).Return(nil) kc.MetaKv = txn err = kc.DropPartitionStatsInfo(context.Background(), info) @@ -1748,14 +1750,14 @@ func Test_CurrentPartitionStatsVersion(t *testing.T) { t.Run("SaveCurrentPartitionStatsVersion", func(t *testing.T) { txn := mocks.NewMetaKv(t) - txn.EXPECT().Save(mock.Anything, mock.Anything).Return(mockErr) + txn.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(mockErr) kc.MetaKv = txn err := kc.SaveCurrentPartitionStatsVersion(context.Background(), collID, partID, vChannel, currentVersion) assert.Error(t, err) txn = mocks.NewMetaKv(t) - txn.EXPECT().Save(mock.Anything, mock.Anything).Return(nil) + txn.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(nil) kc.MetaKv = txn err = kc.SaveCurrentPartitionStatsVersion(context.Background(), collID, partID, vChannel, currentVersion) @@ -1764,7 +1766,7 @@ func Test_CurrentPartitionStatsVersion(t *testing.T) { t.Run("GetCurrentPartitionStatsVersion", func(t *testing.T) { txn := mocks.NewMetaKv(t) - txn.EXPECT().Load(mock.Anything).Return("", mockErr) + txn.EXPECT().Load(mock.Anything, mock.Anything).Return("", mockErr) kc.MetaKv = txn version, err := kc.GetCurrentPartitionStatsVersion(context.Background(), collID, partID, vChannel) @@ -1772,7 +1774,7 @@ func Test_CurrentPartitionStatsVersion(t *testing.T) { assert.Equal(t, int64(0), version) txn = mocks.NewMetaKv(t) - txn.EXPECT().Load(mock.Anything).Return("1", nil) + txn.EXPECT().Load(mock.Anything, mock.Anything).Return("1", nil) kc.MetaKv = txn version, err = kc.GetCurrentPartitionStatsVersion(context.Background(), collID, partID, vChannel) @@ -1782,14 +1784,14 @@ func Test_CurrentPartitionStatsVersion(t *testing.T) { t.Run("DropCurrentPartitionStatsVersion", func(t *testing.T) { txn := mocks.NewMetaKv(t) - txn.EXPECT().Remove(mock.Anything).Return(mockErr) + txn.EXPECT().Remove(mock.Anything, mock.Anything).Return(mockErr) kc.MetaKv = txn err := kc.DropCurrentPartitionStatsVersion(context.Background(), collID, partID, vChannel) assert.Error(t, err) txn = mocks.NewMetaKv(t) - txn.EXPECT().Remove(mock.Anything).Return(nil) + txn.EXPECT().Remove(mock.Anything, mock.Anything).Return(nil) kc.MetaKv = txn err = kc.DropCurrentPartitionStatsVersion(context.Background(), collID, partID, vChannel) @@ -1803,7 +1805,7 @@ func Test_StatsTasks(t *testing.T) { t.Run("ListStatsTasks", func(t *testing.T) { txn := mocks.NewMetaKv(t) - txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return(mockErr) + txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(mockErr) kc.MetaKv = txn tasks, err := kc.ListStatsTasks(context.Background()) @@ -1825,7 +1827,7 @@ func Test_StatsTasks(t *testing.T) { assert.NoError(t, err) txn = mocks.NewMetaKv(t) - txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ string, _ int, f func([]byte, []byte) error) error { + txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, _ string, _ int, f func([]byte, []byte) error) error { return f([]byte("key1"), value) }) kc.MetaKv = txn @@ -1835,7 +1837,7 @@ func Test_StatsTasks(t *testing.T) { assert.Equal(t, 1, len(tasks)) txn = mocks.NewMetaKv(t) - txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ string, _ int, f func([]byte, []byte) error) error { + txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, _ string, _ int, f func([]byte, []byte) error) error { return f([]byte("key1"), []byte("1234")) }) kc.MetaKv = txn @@ -1847,7 +1849,7 @@ func Test_StatsTasks(t *testing.T) { t.Run("SaveStatsTask", func(t *testing.T) { txn := mocks.NewMetaKv(t) - txn.EXPECT().Save(mock.Anything, mock.Anything).Return(mockErr) + txn.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(mockErr) kc.MetaKv = txn task := &indexpb.StatsTask{ @@ -1866,7 +1868,7 @@ func Test_StatsTasks(t *testing.T) { assert.Error(t, err) txn = mocks.NewMetaKv(t) - txn.EXPECT().Save(mock.Anything, mock.Anything).Return(nil) + txn.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(nil) kc.MetaKv = txn err = kc.SaveStatsTask(context.Background(), task) @@ -1875,14 +1877,14 @@ func Test_StatsTasks(t *testing.T) { t.Run("DropStatsTask", func(t *testing.T) { txn := mocks.NewMetaKv(t) - txn.EXPECT().Remove(mock.Anything).Return(mockErr) + txn.EXPECT().Remove(mock.Anything, mock.Anything).Return(mockErr) kc.MetaKv = txn err := kc.DropStatsTask(context.Background(), 1) assert.Error(t, err) txn = mocks.NewMetaKv(t) - txn.EXPECT().Remove(mock.Anything).Return(nil) + txn.EXPECT().Remove(mock.Anything, mock.Anything).Return(nil) kc.MetaKv = txn err = kc.DropStatsTask(context.Background(), 1) diff --git a/internal/metastore/kv/querycoord/kv_catalog.go b/internal/metastore/kv/querycoord/kv_catalog.go index f121462e47848..fd0c1a8476c1b 100644 --- a/internal/metastore/kv/querycoord/kv_catalog.go +++ b/internal/metastore/kv/querycoord/kv_catalog.go @@ -51,7 +51,7 @@ func (s Catalog) SaveCollection(ctx context.Context, collection *querypb.Collect if err != nil { return err } - err = s.cli.Save(k, string(v)) + err = s.cli.Save(ctx, k, string(v)) if err != nil { return err } @@ -65,7 +65,7 @@ func (s Catalog) SavePartition(ctx context.Context, info ...*querypb.PartitionLo if err != nil { return err } - err = s.cli.Save(k, string(v)) + err = s.cli.Save(ctx, k, string(v)) if err != nil { return err } @@ -83,7 +83,7 @@ func (s Catalog) SaveReplica(ctx context.Context, replicas ...*querypb.Replica) } kvs[key] = string(value) } - return s.cli.MultiSave(kvs) + return s.cli.MultiSave(ctx, kvs) } func (s Catalog) SaveResourceGroup(ctx context.Context, rgs ...*querypb.ResourceGroup) error { @@ -98,12 +98,12 @@ func (s Catalog) SaveResourceGroup(ctx context.Context, rgs ...*querypb.Resource ret[key] = string(value) } - return s.cli.MultiSave(ret) + return s.cli.MultiSave(ctx, ret) } func (s Catalog) RemoveResourceGroup(ctx context.Context, rgName string) error { key := encodeResourceGroupKey(rgName) - return s.cli.Remove(key) + return s.cli.Remove(ctx, key) } func (s Catalog) GetCollections(ctx context.Context) ([]*querypb.CollectionLoadInfo, error) { @@ -117,7 +117,7 @@ func (s Catalog) GetCollections(ctx context.Context) ([]*querypb.CollectionLoadI return nil } - err := s.cli.WalkWithPrefix(CollectionLoadInfoPrefix, paginationSize, applyFn) + err := s.cli.WalkWithPrefix(ctx, CollectionLoadInfoPrefix, paginationSize, applyFn) if err != nil { return nil, err } @@ -136,7 +136,7 @@ func (s Catalog) GetPartitions(ctx context.Context) (map[int64][]*querypb.Partit return nil } - err := s.cli.WalkWithPrefix(PartitionLoadInfoPrefix, paginationSize, applyFn) + err := s.cli.WalkWithPrefix(ctx, PartitionLoadInfoPrefix, paginationSize, applyFn) if err != nil { return nil, err } @@ -155,7 +155,7 @@ func (s Catalog) GetReplicas(ctx context.Context) ([]*querypb.Replica, error) { return nil } - err := s.cli.WalkWithPrefix(ReplicaPrefix, paginationSize, applyFn) + err := s.cli.WalkWithPrefix(ctx, ReplicaPrefix, paginationSize, applyFn) if err != nil { return nil, err } @@ -170,7 +170,7 @@ func (s Catalog) GetReplicas(ctx context.Context) ([]*querypb.Replica, error) { } func (s Catalog) getReplicasFromV1(ctx context.Context) ([]*querypb.Replica, error) { - _, replicaValues, err := s.cli.LoadWithPrefix(ReplicaMetaPrefixV1) + _, replicaValues, err := s.cli.LoadWithPrefix(ctx, ReplicaMetaPrefixV1) if err != nil { return nil, err } @@ -193,7 +193,7 @@ func (s Catalog) getReplicasFromV1(ctx context.Context) ([]*querypb.Replica, err } func (s Catalog) GetResourceGroups(ctx context.Context) ([]*querypb.ResourceGroup, error) { - _, rgs, err := s.cli.LoadWithPrefix(ResourceGroupPrefix) + _, rgs, err := s.cli.LoadWithPrefix(ctx, ResourceGroupPrefix) if err != nil { return nil, err } @@ -214,12 +214,12 @@ func (s Catalog) GetResourceGroups(ctx context.Context) ([]*querypb.ResourceGrou func (s Catalog) ReleaseCollection(ctx context.Context, collection int64) error { // remove collection and obtained partitions collectionKey := EncodeCollectionLoadInfoKey(collection) - err := s.cli.Remove(collectionKey) + err := s.cli.Remove(ctx, collectionKey) if err != nil { return err } partitionsPrefix := fmt.Sprintf("%s/%d", PartitionLoadInfoPrefix, collection) - return s.cli.RemoveWithPrefix(partitionsPrefix) + return s.cli.RemoveWithPrefix(ctx, partitionsPrefix) } func (s Catalog) ReleasePartition(ctx context.Context, collection int64, partitions ...int64) error { @@ -233,7 +233,7 @@ func (s Catalog) ReleasePartition(ctx context.Context, collection int64, partiti if endIndex > len(partitions) { endIndex = len(partitions) } - err := s.cli.MultiRemove(keys[index:endIndex]) + err := s.cli.MultiRemove(ctx, keys[index:endIndex]) if err != nil { return err } @@ -241,12 +241,12 @@ func (s Catalog) ReleasePartition(ctx context.Context, collection int64, partiti } return nil } - return s.cli.MultiRemove(keys) + return s.cli.MultiRemove(ctx, keys) } func (s Catalog) ReleaseReplicas(ctx context.Context, collectionID int64) error { key := encodeCollectionReplicaKey(collectionID) - return s.cli.RemoveWithPrefix(key) + return s.cli.RemoveWithPrefix(ctx, key) } func (s Catalog) ReleaseReplica(ctx context.Context, collection int64, replicas ...int64) error { @@ -260,7 +260,7 @@ func (s Catalog) ReleaseReplica(ctx context.Context, collection int64, replicas if endIndex > len(replicas) { endIndex = len(replicas) } - err := s.cli.MultiRemove(keys[index:endIndex]) + err := s.cli.MultiRemove(ctx, keys[index:endIndex]) if err != nil { return err } @@ -268,7 +268,7 @@ func (s Catalog) ReleaseReplica(ctx context.Context, collection int64, replicas } return nil } - return s.cli.MultiRemove(keys) + return s.cli.MultiRemove(ctx, keys) } func (s Catalog) SaveCollectionTargets(ctx context.Context, targets ...*querypb.CollectionTarget) error { @@ -285,7 +285,7 @@ func (s Catalog) SaveCollectionTargets(ctx context.Context, targets ...*querypb. } // to reduce the target size, we do compress before write to etcd - err := s.cli.MultiSave(kvs) + err := s.cli.MultiSave(ctx, kvs) if err != nil { return err } @@ -294,7 +294,7 @@ func (s Catalog) SaveCollectionTargets(ctx context.Context, targets ...*querypb. func (s Catalog) RemoveCollectionTarget(ctx context.Context, collectionID int64) error { k := encodeCollectionTargetKey(collectionID) - return s.cli.Remove(k) + return s.cli.Remove(ctx, k) } func (s Catalog) GetCollectionTargets(ctx context.Context) (map[int64]*querypb.CollectionTarget, error) { @@ -312,7 +312,7 @@ func (s Catalog) GetCollectionTargets(ctx context.Context) (map[int64]*querypb.C return nil } - err := s.cli.WalkWithPrefix(CollectionTargetPrefix, paginationSize, applyFn) + err := s.cli.WalkWithPrefix(ctx, CollectionTargetPrefix, paginationSize, applyFn) if err != nil { return nil, err } diff --git a/internal/metastore/kv/querycoord/kv_catalog_test.go b/internal/metastore/kv/querycoord/kv_catalog_test.go index 638ee61b524be..1c119e838f611 100644 --- a/internal/metastore/kv/querycoord/kv_catalog_test.go +++ b/internal/metastore/kv/querycoord/kv_catalog_test.go @@ -238,8 +238,8 @@ func (suite *CatalogTestSuite) TestCollectionTarget() { // test access meta store failed mockStore := mocks.NewMetaKv(suite.T()) mockErr := errors.New("failed to access etcd") - mockStore.EXPECT().MultiSave(mock.Anything).Return(mockErr) - mockStore.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return(mockErr) + mockStore.EXPECT().MultiSave(mock.Anything, mock.Anything).Return(mockErr) + mockStore.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(mockErr) suite.catalog.cli = mockStore err = suite.catalog.SaveCollectionTargets(ctx, &querypb.CollectionTarget{}) diff --git a/internal/metastore/kv/rootcoord/kv_catalog.go b/internal/metastore/kv/rootcoord/kv_catalog.go index 7bc1e4e170496..e2f4c1bb704da 100644 --- a/internal/metastore/kv/rootcoord/kv_catalog.go +++ b/internal/metastore/kv/rootcoord/kv_catalog.go @@ -96,16 +96,16 @@ func BuildAliasPrefixWithDB(dbID int64) string { // since SnapshotKV may save both snapshot key and the original key if the original key is newest // MaxEtcdTxnNum need to divided by 2 -func batchMultiSaveAndRemove(snapshot kv.SnapShotKV, limit int, saves map[string]string, removals []string, ts typeutil.Timestamp) error { +func batchMultiSaveAndRemove(ctx context.Context, snapshot kv.SnapShotKV, limit int, saves map[string]string, removals []string, ts typeutil.Timestamp) error { saveFn := func(partialKvs map[string]string) error { - return snapshot.MultiSave(partialKvs, ts) + return snapshot.MultiSave(ctx, partialKvs, ts) } if err := etcd.SaveByBatchWithLimit(saves, limit, saveFn); err != nil { return err } removeFn := func(partialKeys []string) error { - return snapshot.MultiSaveAndRemove(nil, partialKeys, ts) + return snapshot.MultiSaveAndRemove(ctx, nil, partialKeys, ts) } return etcd.RemoveByBatchWithLimit(removals, limit, removeFn) } @@ -117,7 +117,7 @@ func (kc *Catalog) CreateDatabase(ctx context.Context, db *model.Database, ts ty if err != nil { return err } - return kc.Snapshot.Save(key, string(v), ts) + return kc.Snapshot.Save(ctx, key, string(v), ts) } func (kc *Catalog) AlterDatabase(ctx context.Context, newColl *model.Database, ts typeutil.Timestamp) error { @@ -127,16 +127,16 @@ func (kc *Catalog) AlterDatabase(ctx context.Context, newColl *model.Database, t if err != nil { return err } - return kc.Snapshot.Save(key, string(v), ts) + return kc.Snapshot.Save(ctx, key, string(v), ts) } func (kc *Catalog) DropDatabase(ctx context.Context, dbID int64, ts typeutil.Timestamp) error { key := BuildDatabaseKey(dbID) - return kc.Snapshot.MultiSaveAndRemove(nil, []string{key}, ts) + return kc.Snapshot.MultiSaveAndRemove(ctx, nil, []string{key}, ts) } func (kc *Catalog) ListDatabases(ctx context.Context, ts typeutil.Timestamp) ([]*model.Database, error) { - _, vals, err := kc.Snapshot.LoadWithPrefix(DBInfoMetaPrefix, ts) + _, vals, err := kc.Snapshot.LoadWithPrefix(ctx, DBInfoMetaPrefix, ts) if err != nil { return nil, err } @@ -171,7 +171,7 @@ func (kc *Catalog) CreateCollection(ctx context.Context, coll *model.Collection, // If we succeeded to save collection but failed to save other related keys, the garbage meta can be removed // outside and the collection won't be seen by any others (since it's of creating state). // However, if we save other keys first, there is no chance to remove the intermediate meta. - if err := kc.Snapshot.Save(k1, string(v1), ts); err != nil { + if err := kc.Snapshot.Save(ctx, k1, string(v1), ts); err != nil { return err } @@ -216,13 +216,13 @@ func (kc *Catalog) CreateCollection(ctx context.Context, coll *model.Collection, // since SnapshotKV may save both snapshot key and the original key if the original key is newest // MaxEtcdTxnNum need to divided by 2 return etcd.SaveByBatchWithLimit(kvs, util.MaxEtcdTxnNum/2, func(partialKvs map[string]string) error { - return kc.Snapshot.MultiSave(partialKvs, ts) + return kc.Snapshot.MultiSave(ctx, partialKvs, ts) }) } func (kc *Catalog) loadCollectionFromDb(ctx context.Context, dbID int64, collectionID typeutil.UniqueID, ts typeutil.Timestamp) (*pb.CollectionInfo, error) { collKey := BuildCollectionKey(dbID, collectionID) - collVal, err := kc.Snapshot.Load(collKey, ts) + collVal, err := kc.Snapshot.Load(ctx, collKey, ts) if err != nil { return nil, merr.WrapErrCollectionNotFound(collectionID, err.Error()) } @@ -280,7 +280,7 @@ func (kc *Catalog) CreatePartition(ctx context.Context, dbID int64, partition *m if err != nil { return err } - return kc.Snapshot.Save(k, string(v), ts) + return kc.Snapshot.Save(ctx, k, string(v), ts) } if partitionExistByID(collMeta, partition.PartitionID) { @@ -302,7 +302,7 @@ func (kc *Catalog) CreatePartition(ctx context.Context, dbID int64, partition *m if err != nil { return err } - return kc.Snapshot.Save(k, string(v), ts) + return kc.Snapshot.Save(ctx, k, string(v), ts) } func (kc *Catalog) CreateAlias(ctx context.Context, alias *model.Alias, ts typeutil.Timestamp) error { @@ -315,7 +315,7 @@ func (kc *Catalog) CreateAlias(ctx context.Context, alias *model.Alias, ts typeu return err } kvs := map[string]string{k: string(v)} - return kc.Snapshot.MultiSaveAndRemove(kvs, []string{oldKBefore210, oldKeyWithoutDb}, ts) + return kc.Snapshot.MultiSaveAndRemove(ctx, kvs, []string{oldKBefore210, oldKeyWithoutDb}, ts) } func (kc *Catalog) CreateCredential(ctx context.Context, credential *model.Credential) error { @@ -326,7 +326,7 @@ func (kc *Catalog) CreateCredential(ctx context.Context, credential *model.Crede return err } - err = kc.Txn.Save(k, string(v)) + err = kc.Txn.Save(ctx, k, string(v)) if err != nil { log.Error("create credential persist meta fail", zap.String("key", k), zap.Error(err)) return err @@ -341,7 +341,7 @@ func (kc *Catalog) AlterCredential(ctx context.Context, credential *model.Creden func (kc *Catalog) listPartitionsAfter210(ctx context.Context, collectionID typeutil.UniqueID, ts typeutil.Timestamp) ([]*model.Partition, error) { prefix := BuildPartitionPrefix(collectionID) - _, values, err := kc.Snapshot.LoadWithPrefix(prefix, ts) + _, values, err := kc.Snapshot.LoadWithPrefix(ctx, prefix, ts) if err != nil { return nil, err } @@ -358,7 +358,7 @@ func (kc *Catalog) listPartitionsAfter210(ctx context.Context, collectionID type } func (kc *Catalog) batchListPartitionsAfter210(ctx context.Context, ts typeutil.Timestamp) (map[int64][]*model.Partition, error) { - _, values, err := kc.Snapshot.LoadWithPrefix(PartitionMetaPrefix, ts) + _, values, err := kc.Snapshot.LoadWithPrefix(ctx, PartitionMetaPrefix, ts) if err != nil { return nil, err } @@ -385,7 +385,7 @@ func fieldVersionAfter210(collMeta *pb.CollectionInfo) bool { func (kc *Catalog) listFieldsAfter210(ctx context.Context, collectionID typeutil.UniqueID, ts typeutil.Timestamp) ([]*model.Field, error) { prefix := BuildFieldPrefix(collectionID) - _, values, err := kc.Snapshot.LoadWithPrefix(prefix, ts) + _, values, err := kc.Snapshot.LoadWithPrefix(ctx, prefix, ts) if err != nil { return nil, err } @@ -402,7 +402,7 @@ func (kc *Catalog) listFieldsAfter210(ctx context.Context, collectionID typeutil } func (kc *Catalog) batchListFieldsAfter210(ctx context.Context, ts typeutil.Timestamp) (map[int64][]*model.Field, error) { - keys, values, err := kc.Snapshot.LoadWithPrefix(FieldMetaPrefix, ts) + keys, values, err := kc.Snapshot.LoadWithPrefix(ctx, FieldMetaPrefix, ts) if err != nil { return nil, err } @@ -427,9 +427,9 @@ func (kc *Catalog) batchListFieldsAfter210(ctx context.Context, ts typeutil.Time return ret, nil } -func (kc *Catalog) listFunctions(collectionID typeutil.UniqueID, ts typeutil.Timestamp) ([]*model.Function, error) { +func (kc *Catalog) listFunctions(ctx context.Context, collectionID typeutil.UniqueID, ts typeutil.Timestamp) ([]*model.Function, error) { prefix := BuildFunctionPrefix(collectionID) - _, values, err := kc.Snapshot.LoadWithPrefix(prefix, ts) + _, values, err := kc.Snapshot.LoadWithPrefix(ctx, prefix, ts) if err != nil { return nil, err } @@ -445,8 +445,8 @@ func (kc *Catalog) listFunctions(collectionID typeutil.UniqueID, ts typeutil.Tim return functions, nil } -func (kc *Catalog) batchListFunctions(ts typeutil.Timestamp) (map[int64][]*model.Function, error) { - keys, values, err := kc.Snapshot.LoadWithPrefix(FunctionMetaPrefix, ts) +func (kc *Catalog) batchListFunctions(ctx context.Context, ts typeutil.Timestamp) (map[int64][]*model.Function, error) { + keys, values, err := kc.Snapshot.LoadWithPrefix(ctx, FunctionMetaPrefix, ts) if err != nil { return nil, err } @@ -490,7 +490,7 @@ func (kc *Catalog) appendPartitionAndFieldsInfo(ctx context.Context, collMeta *p } collection.Fields = fields - functions, err := kc.listFunctions(collection.CollectionID, ts) + functions, err := kc.listFunctions(ctx, collection.CollectionID, ts) if err != nil { return nil, err } @@ -520,7 +520,7 @@ func (kc *Catalog) batchAppendPartitionAndFieldsInfo(ctx context.Context, collMe return nil, err } - functionMetaMap, err = kc.batchListFunctions(ts) + functionMetaMap, err = kc.batchListFunctions(ctx, ts) if err != nil { return nil, err } @@ -558,7 +558,7 @@ func (kc *Catalog) CollectionExists(ctx context.Context, dbID int64, collectionI func (kc *Catalog) GetCredential(ctx context.Context, username string) (*model.Credential, error) { k := fmt.Sprintf("%s/%s", CredentialPrefix, username) - v, err := kc.Txn.Load(k) + v, err := kc.Txn.Load(ctx, k) if err != nil { if errors.Is(err, merr.ErrIoKeyNotFound) { log.Debug("not found the user", zap.String("key", k)) @@ -612,15 +612,15 @@ func (kc *Catalog) DropCollection(ctx context.Context, collectionInfo *model.Col // However, if we remove collection first, we cannot remove other metas. // since SnapshotKV may save both snapshot key and the original key if the original key is newest // MaxEtcdTxnNum need to divided by 2 - if err := batchMultiSaveAndRemove(kc.Snapshot, util.MaxEtcdTxnNum/2, nil, delMetakeysSnap, ts); err != nil { + if err := batchMultiSaveAndRemove(ctx, kc.Snapshot, util.MaxEtcdTxnNum/2, nil, delMetakeysSnap, ts); err != nil { return err } // if we found collection dropping, we should try removing related resources. - return kc.Snapshot.MultiSaveAndRemove(nil, collectionKeys, ts) + return kc.Snapshot.MultiSaveAndRemove(ctx, nil, collectionKeys, ts) } -func (kc *Catalog) alterModifyCollection(oldColl *model.Collection, newColl *model.Collection, ts typeutil.Timestamp) error { +func (kc *Catalog) alterModifyCollection(ctx context.Context, oldColl *model.Collection, newColl *model.Collection, ts typeutil.Timestamp) error { if oldColl.TenantID != newColl.TenantID || oldColl.CollectionID != newColl.CollectionID { return fmt.Errorf("altering tenant id or collection id is forbidden") } @@ -646,19 +646,19 @@ func (kc *Catalog) alterModifyCollection(oldColl *model.Collection, newColl *mod } saves := map[string]string{newKey: string(value)} if oldKey == newKey { - return kc.Snapshot.Save(newKey, string(value), ts) + return kc.Snapshot.Save(ctx, newKey, string(value), ts) } - return kc.Snapshot.MultiSaveAndRemove(saves, []string{oldKey}, ts) + return kc.Snapshot.MultiSaveAndRemove(ctx, saves, []string{oldKey}, ts) } func (kc *Catalog) AlterCollection(ctx context.Context, oldColl *model.Collection, newColl *model.Collection, alterType metastore.AlterType, ts typeutil.Timestamp) error { if alterType == metastore.MODIFY { - return kc.alterModifyCollection(oldColl, newColl, ts) + return kc.alterModifyCollection(ctx, oldColl, newColl, ts) } return fmt.Errorf("altering collection doesn't support %s", alterType.String()) } -func (kc *Catalog) alterModifyPartition(oldPart *model.Partition, newPart *model.Partition, ts typeutil.Timestamp) error { +func (kc *Catalog) alterModifyPartition(ctx context.Context, oldPart *model.Partition, newPart *model.Partition, ts typeutil.Timestamp) error { if oldPart.CollectionID != newPart.CollectionID || oldPart.PartitionID != newPart.PartitionID { return fmt.Errorf("altering collection id or partition id is forbidden") } @@ -672,12 +672,12 @@ func (kc *Catalog) alterModifyPartition(oldPart *model.Partition, newPart *model if err != nil { return err } - return kc.Snapshot.Save(key, string(value), ts) + return kc.Snapshot.Save(ctx, key, string(value), ts) } func (kc *Catalog) AlterPartition(ctx context.Context, dbID int64, oldPart *model.Partition, newPart *model.Partition, alterType metastore.AlterType, ts typeutil.Timestamp) error { if alterType == metastore.MODIFY { - return kc.alterModifyPartition(oldPart, newPart, ts) + return kc.alterModifyPartition(ctx, oldPart, newPart, ts) } return fmt.Errorf("altering partition doesn't support %s", alterType.String()) } @@ -716,7 +716,7 @@ func (kc *Catalog) DropPartition(ctx context.Context, dbID int64, collectionID t if partitionVersionAfter210(collMeta) { k := BuildPartitionKey(collectionID, partitionID) - return kc.Snapshot.MultiSaveAndRemove(nil, []string{k}, ts) + return kc.Snapshot.MultiSaveAndRemove(ctx, nil, []string{k}, ts) } k := BuildCollectionKey(util.NonDBID, collectionID) @@ -725,7 +725,7 @@ func (kc *Catalog) DropPartition(ctx context.Context, dbID int64, collectionID t if err != nil { return err } - return kc.Snapshot.Save(k, string(v), ts) + return kc.Snapshot.Save(ctx, k, string(v), ts) } func (kc *Catalog) DropCredential(ctx context.Context, username string) error { @@ -745,7 +745,7 @@ func (kc *Catalog) DropCredential(ctx context.Context, username string) error { } } } - err = kc.Txn.MultiRemove(deleteKeys) + err = kc.Txn.MultiRemove(ctx, deleteKeys) if err != nil { log.Warn("fail to drop credential", zap.String("key", k), zap.Error(err)) return err @@ -758,12 +758,12 @@ func (kc *Catalog) DropAlias(ctx context.Context, dbID int64, alias string, ts t oldKBefore210 := BuildAliasKey210(alias) oldKeyWithoutDb := BuildAliasKey(alias) k := BuildAliasKeyWithDB(dbID, alias) - return kc.Snapshot.MultiSaveAndRemove(nil, []string{k, oldKeyWithoutDb, oldKBefore210}, ts) + return kc.Snapshot.MultiSaveAndRemove(ctx, nil, []string{k, oldKeyWithoutDb, oldKBefore210}, ts) } func (kc *Catalog) GetCollectionByName(ctx context.Context, dbID int64, collectionName string, ts typeutil.Timestamp) (*model.Collection, error) { prefix := getDatabasePrefix(dbID) - _, vals, err := kc.Snapshot.LoadWithPrefix(prefix, ts) + _, vals, err := kc.Snapshot.LoadWithPrefix(ctx, prefix, ts) if err != nil { log.Warn("get collection meta fail", zap.String("collectionName", collectionName), zap.Error(err)) return nil, err @@ -787,7 +787,7 @@ func (kc *Catalog) GetCollectionByName(ctx context.Context, dbID int64, collecti func (kc *Catalog) ListCollections(ctx context.Context, dbID int64, ts typeutil.Timestamp) ([]*model.Collection, error) { prefix := getDatabasePrefix(dbID) - _, vals, err := kc.Snapshot.LoadWithPrefix(prefix, ts) + _, vals, err := kc.Snapshot.LoadWithPrefix(ctx, prefix, ts) if err != nil { log.Error("get collections meta fail", zap.String("prefix", prefix), @@ -825,19 +825,19 @@ func (kc *Catalog) ListCollections(ctx context.Context, dbID int64, ts typeutil. // all collections in default database should be marked with dbID 1. // this method also update dbid in meta store when dbid is 0 // see also: https://github.com/milvus-io/milvus/issues/33608 -func (kv *Catalog) fixDefaultDBIDConsistency(_ context.Context, collMeta *pb.CollectionInfo, ts typeutil.Timestamp) { +func (kv *Catalog) fixDefaultDBIDConsistency(ctx context.Context, collMeta *pb.CollectionInfo, ts typeutil.Timestamp) { if collMeta.DbId == util.NonDBID { coll := model.UnmarshalCollectionModel(collMeta) cloned := coll.Clone() cloned.DBID = util.DefaultDBID - kv.alterModifyCollection(coll, cloned, ts) + kv.alterModifyCollection(ctx, coll, cloned, ts) collMeta.DbId = util.DefaultDBID } } func (kc *Catalog) listAliasesBefore210(ctx context.Context, ts typeutil.Timestamp) ([]*model.Alias, error) { - _, values, err := kc.Snapshot.LoadWithPrefix(CollectionAliasMetaPrefix210, ts) + _, values, err := kc.Snapshot.LoadWithPrefix(ctx, CollectionAliasMetaPrefix210, ts) if err != nil { return nil, err } @@ -861,7 +861,7 @@ func (kc *Catalog) listAliasesBefore210(ctx context.Context, ts typeutil.Timesta func (kc *Catalog) listAliasesAfter210WithDb(ctx context.Context, dbID int64, ts typeutil.Timestamp) ([]*model.Alias, error) { prefix := BuildAliasPrefixWithDB(dbID) - _, values, err := kc.Snapshot.LoadWithPrefix(prefix, ts) + _, values, err := kc.Snapshot.LoadWithPrefix(ctx, prefix, ts) if err != nil { return nil, err } @@ -917,7 +917,7 @@ func (kc *Catalog) ListCredentials(ctx context.Context) ([]string, error) { } func (kc *Catalog) ListCredentialsWithPasswd(ctx context.Context) (map[string]string, error) { - keys, values, err := kc.Txn.LoadWithPrefix(CredentialPrefix) + keys, values, err := kc.Txn.LoadWithPrefix(ctx, CredentialPrefix) if err != nil { log.Error("list all credential usernames fail", zap.String("prefix", CredentialPrefix), zap.Error(err)) return nil, err @@ -942,32 +942,32 @@ func (kc *Catalog) ListCredentialsWithPasswd(ctx context.Context) (map[string]st return users, nil } -func (kc *Catalog) save(k string) error { +func (kc *Catalog) save(ctx context.Context, k string) error { var err error - if _, err = kc.Txn.Load(k); err != nil && !errors.Is(err, merr.ErrIoKeyNotFound) { + if _, err = kc.Txn.Load(ctx, k); err != nil && !errors.Is(err, merr.ErrIoKeyNotFound) { return err } if err == nil { log.Debug("the key has existed", zap.String("key", k)) return common.NewIgnorableError(fmt.Errorf("the key[%s] has existed", k)) } - return kc.Txn.Save(k, "") + return kc.Txn.Save(ctx, k, "") } -func (kc *Catalog) remove(k string) error { +func (kc *Catalog) remove(ctx context.Context, k string) error { var err error - if _, err = kc.Txn.Load(k); err != nil && !errors.Is(err, merr.ErrIoKeyNotFound) { + if _, err = kc.Txn.Load(ctx, k); err != nil && !errors.Is(err, merr.ErrIoKeyNotFound) { return err } if err != nil && errors.Is(err, merr.ErrIoKeyNotFound) { return common.NewIgnorableError(fmt.Errorf("the key[%s] isn't existed", k)) } - return kc.Txn.Remove(k) + return kc.Txn.Remove(ctx, k) } func (kc *Catalog) CreateRole(ctx context.Context, tenant string, entity *milvuspb.RoleEntity) error { k := funcutil.HandleTenantForEtcdKey(RolePrefix, tenant, entity.Name) - err := kc.save(k) + err := kc.save(ctx, k) if err != nil && !common.IsIgnorableError(err) { log.Warn("fail to save the role", zap.String("key", k), zap.Error(err)) } @@ -993,7 +993,7 @@ func (kc *Catalog) DropRole(ctx context.Context, tenant string, roleName string) } } - err = kc.Txn.MultiRemove(deleteKeys) + err = kc.Txn.MultiRemove(ctx, deleteKeys) if err != nil { log.Warn("fail to drop role", zap.String("key", k), zap.Error(err)) return err @@ -1005,12 +1005,12 @@ func (kc *Catalog) AlterUserRole(ctx context.Context, tenant string, userEntity k := funcutil.HandleTenantForEtcdKey(RoleMappingPrefix, tenant, fmt.Sprintf("%s/%s", userEntity.Name, roleEntity.Name)) var err error if operateType == milvuspb.OperateUserRoleType_AddUserToRole { - err = kc.save(k) + err = kc.save(ctx, k) if err != nil { log.Error("fail to save the user-role", zap.String("key", k), zap.Error(err)) } } else if operateType == milvuspb.OperateUserRoleType_RemoveUserFromRole { - err = kc.remove(k) + err = kc.remove(ctx, k) if err != nil { log.Error("fail to remove the user-role", zap.String("key", k), zap.Error(err)) } @@ -1026,7 +1026,7 @@ func (kc *Catalog) ListRole(ctx context.Context, tenant string, entity *milvuspb roleToUsers := make(map[string][]string) if includeUserInfo { roleMappingKey := funcutil.HandleTenantForEtcdKey(RoleMappingPrefix, tenant, "") - keys, _, err := kc.Txn.LoadWithPrefix(roleMappingKey) + keys, _, err := kc.Txn.LoadWithPrefix(ctx, roleMappingKey) if err != nil { log.Error("fail to load role mappings", zap.String("key", roleMappingKey), zap.Error(err)) return results, err @@ -1057,7 +1057,7 @@ func (kc *Catalog) ListRole(ctx context.Context, tenant string, entity *milvuspb if entity == nil { roleKey := funcutil.HandleTenantForEtcdKey(RolePrefix, tenant, "") - keys, _, err := kc.Txn.LoadWithPrefix(roleKey) + keys, _, err := kc.Txn.LoadWithPrefix(ctx, roleKey) if err != nil { log.Error("fail to load roles", zap.String("key", roleKey), zap.Error(err)) return results, err @@ -1075,7 +1075,7 @@ func (kc *Catalog) ListRole(ctx context.Context, tenant string, entity *milvuspb return results, fmt.Errorf("role name in the role entity is empty") } roleKey := funcutil.HandleTenantForEtcdKey(RolePrefix, tenant, entity.Name) - _, err := kc.Txn.Load(roleKey) + _, err := kc.Txn.Load(ctx, roleKey) if err != nil { log.Warn("fail to load a role", zap.String("key", roleKey), zap.Error(err)) return results, err @@ -1086,10 +1086,10 @@ func (kc *Catalog) ListRole(ctx context.Context, tenant string, entity *milvuspb return results, nil } -func (kc *Catalog) getRolesByUsername(tenant string, username string) ([]string, error) { +func (kc *Catalog) getRolesByUsername(ctx context.Context, tenant string, username string) ([]string, error) { var roles []string k := funcutil.HandleTenantForEtcdKey(RoleMappingPrefix, tenant, username) - keys, _, err := kc.Txn.LoadWithPrefix(k) + keys, _, err := kc.Txn.LoadWithPrefix(ctx, k) if err != nil { log.Error("fail to load role mappings by the username", zap.String("key", k), zap.Error(err)) return roles, err @@ -1106,12 +1106,12 @@ func (kc *Catalog) getRolesByUsername(tenant string, username string) ([]string, } // getUserResult get the user result by the username. And never return the error because the error means the user isn't added to a role. -func (kc *Catalog) getUserResult(tenant string, username string, includeRoleInfo bool) (*milvuspb.UserResult, error) { +func (kc *Catalog) getUserResult(ctx context.Context, tenant string, username string, includeRoleInfo bool) (*milvuspb.UserResult, error) { result := &milvuspb.UserResult{User: &milvuspb.UserEntity{Name: username}} if !includeRoleInfo { return result, nil } - roleNames, err := kc.getRolesByUsername(tenant, username) + roleNames, err := kc.getRolesByUsername(ctx, tenant, username) if err != nil { log.Warn("fail to get roles by the username", zap.Error(err)) return result, err @@ -1132,7 +1132,7 @@ func (kc *Catalog) ListUser(ctx context.Context, tenant string, entity *milvuspb ) appendUserResult := func(username string) error { - result, err := kc.getUserResult(tenant, username, includeRoleInfo) + result, err := kc.getUserResult(ctx, tenant, username, includeRoleInfo) if err != nil { return err } @@ -1175,13 +1175,13 @@ func (kc *Catalog) AlterGrant(ctx context.Context, tenant string, entity *milvus // Compatible with logic without db if entity.DbName == util.DefaultDBName { - v, err = kc.Txn.Load(funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, fmt.Sprintf("%s/%s/%s", entity.Role.Name, entity.Object.Name, entity.ObjectName))) + v, err = kc.Txn.Load(ctx, funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, fmt.Sprintf("%s/%s/%s", entity.Role.Name, entity.Object.Name, entity.ObjectName))) if err == nil { idStr = v } } if idStr == "" { - if v, err = kc.Txn.Load(k); err == nil { + if v, err = kc.Txn.Load(ctx, k); err == nil { idStr = v } else { log.Warn("fail to load grant privilege entity", zap.String("key", k), zap.Any("type", operateType), zap.Error(err)) @@ -1196,7 +1196,7 @@ func (kc *Catalog) AlterGrant(ctx context.Context, tenant string, entity *milvus } idStr = crypto.MD5(k) - err = kc.Txn.Save(k, idStr) + err = kc.Txn.Save(ctx, k, idStr) if err != nil { log.Error("fail to allocate id when altering the grant", zap.Error(err)) return err @@ -1204,7 +1204,7 @@ func (kc *Catalog) AlterGrant(ctx context.Context, tenant string, entity *milvus } } k = funcutil.HandleTenantForEtcdKey(GranteeIDPrefix, tenant, fmt.Sprintf("%s/%s", idStr, privilegeName)) - _, err = kc.Txn.Load(k) + _, err = kc.Txn.Load(ctx, k) if err != nil { log.Warn("fail to load the grantee id", zap.String("key", k), zap.Error(err)) if !errors.Is(err, merr.ErrIoKeyNotFound) { @@ -1216,7 +1216,7 @@ func (kc *Catalog) AlterGrant(ctx context.Context, tenant string, entity *milvus return common.NewIgnorableError(fmt.Errorf("the grantee-id[%s] isn't existed", k)) } if funcutil.IsGrant(operateType) { - if err = kc.Txn.Save(k, entity.Grantor.User.Name); err != nil { + if err = kc.Txn.Save(ctx, k, entity.Grantor.User.Name); err != nil { log.Error("fail to save the grantee id", zap.String("key", k), zap.Error(err)) } return err @@ -1224,7 +1224,7 @@ func (kc *Catalog) AlterGrant(ctx context.Context, tenant string, entity *milvus return nil } if funcutil.IsRevoke(operateType) { - if err = kc.Txn.Remove(k); err != nil { + if err = kc.Txn.Remove(ctx, k); err != nil { log.Error("fail to remove the grantee id", zap.String("key", k), zap.Error(err)) return err } @@ -1244,7 +1244,7 @@ func (kc *Catalog) ListGrant(ctx context.Context, tenant string, entity *milvusp return nil } granteeIDKey := funcutil.HandleTenantForEtcdKey(GranteeIDPrefix, tenant, v) - keys, values, err := kc.Txn.LoadWithPrefix(granteeIDKey) + keys, values, err := kc.Txn.LoadWithPrefix(ctx, granteeIDKey) if err != nil { log.Error("fail to load the grantee ids", zap.String("key", granteeIDKey), zap.Error(err)) return err @@ -1276,7 +1276,7 @@ func (kc *Catalog) ListGrant(ctx context.Context, tenant string, entity *milvusp if !funcutil.IsEmptyString(entity.ObjectName) && entity.Object != nil && !funcutil.IsEmptyString(entity.Object.Name) { if entity.DbName == util.DefaultDBName { granteeKey = funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, fmt.Sprintf("%s/%s/%s", entity.Role.Name, entity.Object.Name, entity.ObjectName)) - v, err := kc.Txn.Load(granteeKey) + v, err := kc.Txn.Load(ctx, granteeKey) if err == nil { err = appendGrantEntity(v, entity.Object.Name, entity.ObjectName) if err == nil { @@ -1287,14 +1287,14 @@ func (kc *Catalog) ListGrant(ctx context.Context, tenant string, entity *milvusp if entity.DbName != util.AnyWord { granteeKey = funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, fmt.Sprintf("%s/%s/%s", entity.Role.Name, entity.Object.Name, funcutil.CombineObjectName(util.AnyWord, entity.ObjectName))) - v, err := kc.Txn.Load(granteeKey) + v, err := kc.Txn.Load(ctx, granteeKey) if err == nil { _ = appendGrantEntity(v, entity.Object.Name, funcutil.CombineObjectName(util.AnyWord, entity.ObjectName)) } } granteeKey = funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, fmt.Sprintf("%s/%s/%s", entity.Role.Name, entity.Object.Name, funcutil.CombineObjectName(entity.DbName, entity.ObjectName))) - v, err := kc.Txn.Load(granteeKey) + v, err := kc.Txn.Load(ctx, granteeKey) if err != nil { log.Error("fail to load the grant privilege entity", zap.String("key", granteeKey), zap.Error(err)) return entities, err @@ -1305,7 +1305,7 @@ func (kc *Catalog) ListGrant(ctx context.Context, tenant string, entity *milvusp } } else { granteeKey = funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, entity.Role.Name) - keys, values, err := kc.Txn.LoadWithPrefix(granteeKey) + keys, values, err := kc.Txn.LoadWithPrefix(ctx, granteeKey) if err != nil { log.Error("fail to load grant privilege entities", zap.String("key", granteeKey), zap.Error(err)) return entities, err @@ -1336,7 +1336,7 @@ func (kc *Catalog) DeleteGrant(ctx context.Context, tenant string, role *milvusp removeKeys = append(removeKeys, k) // the values are the grantee id list - _, values, err := kc.Txn.LoadWithPrefix(k) + _, values, err := kc.Txn.LoadWithPrefix(ctx, k) if err != nil { log.Warn("fail to load grant privilege entities", zap.String("key", k), zap.Error(err)) return err @@ -1346,7 +1346,7 @@ func (kc *Catalog) DeleteGrant(ctx context.Context, tenant string, role *milvusp removeKeys = append(removeKeys, granteeIDKey) } - if err = kc.Txn.MultiSaveAndRemoveWithPrefix(nil, removeKeys); err != nil { + if err = kc.Txn.MultiSaveAndRemoveWithPrefix(ctx, nil, removeKeys); err != nil { log.Error("fail to remove with the prefix", zap.String("key", k), zap.Error(err)) } return err @@ -1355,7 +1355,7 @@ func (kc *Catalog) DeleteGrant(ctx context.Context, tenant string, role *milvusp func (kc *Catalog) ListPolicy(ctx context.Context, tenant string) ([]string, error) { var grantInfoStrs []string granteeKey := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, "") - keys, values, err := kc.Txn.LoadWithPrefix(granteeKey) + keys, values, err := kc.Txn.LoadWithPrefix(ctx, granteeKey) if err != nil { log.Error("fail to load all grant privilege entities", zap.String("key", granteeKey), zap.Error(err)) return []string{}, err @@ -1368,7 +1368,7 @@ func (kc *Catalog) ListPolicy(ctx context.Context, tenant string) ([]string, err continue } granteeIDKey := funcutil.HandleTenantForEtcdKey(GranteeIDPrefix, tenant, values[i]) - idKeys, _, err := kc.Txn.LoadWithPrefix(granteeIDKey) + idKeys, _, err := kc.Txn.LoadWithPrefix(ctx, granteeIDKey) if err != nil { log.Error("fail to load the grantee ids", zap.String("key", granteeIDKey), zap.Error(err)) return []string{}, err @@ -1390,7 +1390,7 @@ func (kc *Catalog) ListPolicy(ctx context.Context, tenant string) ([]string, err func (kc *Catalog) ListUserRole(ctx context.Context, tenant string) ([]string, error) { var userRoles []string k := funcutil.HandleTenantForEtcdKey(RoleMappingPrefix, tenant, "") - keys, _, err := kc.Txn.LoadWithPrefix(k) + keys, _, err := kc.Txn.LoadWithPrefix(ctx, k) if err != nil { log.Error("fail to load all user-role mappings", zap.String("key", k), zap.Error(err)) return []string{}, err @@ -1613,7 +1613,7 @@ func (kc *Catalog) RestoreRBAC(ctx context.Context, tenant string, meta *milvusp func (kc *Catalog) GetPrivilegeGroup(ctx context.Context, groupName string) (*milvuspb.PrivilegeGroupInfo, error) { k := BuildPrivilegeGroupkey(groupName) - val, err := kc.Txn.Load(k) + val, err := kc.Txn.Load(ctx, k) if err != nil { if errors.Is(err, merr.ErrIoKeyNotFound) { return nil, fmt.Errorf("privilege group [%s] does not exist", groupName) @@ -1632,7 +1632,7 @@ func (kc *Catalog) GetPrivilegeGroup(ctx context.Context, groupName string) (*mi func (kc *Catalog) DropPrivilegeGroup(ctx context.Context, groupName string) error { k := BuildPrivilegeGroupkey(groupName) - err := kc.Txn.Remove(k) + err := kc.Txn.Remove(ctx, k) if err != nil { log.Warn("fail to drop privilege group", zap.String("key", k), zap.Error(err)) return err @@ -1651,7 +1651,7 @@ func (kc *Catalog) SavePrivilegeGroup(ctx context.Context, data *milvuspb.Privil log.Error("failed to marshal privilege group info", zap.Error(err)) return err } - if err = kc.Txn.Save(k, string(v)); err != nil { + if err = kc.Txn.Save(ctx, k, string(v)); err != nil { log.Warn("fail to put privilege group", zap.String("key", k), zap.Error(err)) return err } @@ -1659,7 +1659,7 @@ func (kc *Catalog) SavePrivilegeGroup(ctx context.Context, data *milvuspb.Privil } func (kc *Catalog) ListPrivilegeGroups(ctx context.Context) ([]*milvuspb.PrivilegeGroupInfo, error) { - _, vals, err := kc.Txn.LoadWithPrefix(PrivilegeGroupPrefix) + _, vals, err := kc.Txn.LoadWithPrefix(ctx, PrivilegeGroupPrefix) if err != nil { log.Error("failed to list privilege groups", zap.String("prefix", PrivilegeGroupPrefix), zap.Error(err)) return nil, err diff --git a/internal/metastore/kv/rootcoord/kv_catalog_test.go b/internal/metastore/kv/rootcoord/kv_catalog_test.go index c0dcc6368b880..9cf8dfcfcfd38 100644 --- a/internal/metastore/kv/rootcoord/kv_catalog_test.go +++ b/internal/metastore/kv/rootcoord/kv_catalog_test.go @@ -102,7 +102,7 @@ func TestCatalog_ListCollections(t *testing.T) { t.Run("load collection with prefix fail", func(t *testing.T) { kv := mocks.NewSnapShotKV(t) ts := uint64(1) - kv.On("LoadWithPrefix", CollectionMetaPrefix, ts). + kv.On("LoadWithPrefix", mock.Anything, CollectionMetaPrefix, ts). Return(nil, nil, targetErr) kc := Catalog{Snapshot: kv} @@ -117,9 +117,9 @@ func TestCatalog_ListCollections(t *testing.T) { bColl, err := proto.Marshal(coll2) assert.NoError(t, err) - kv.On("LoadWithPrefix", CollectionMetaPrefix, ts). + kv.On("LoadWithPrefix", mock.Anything, CollectionMetaPrefix, ts). Return([]string{"key"}, []string{string(bColl)}, nil) - kv.On("LoadWithPrefix", mock.MatchedBy( + kv.On("LoadWithPrefix", mock.Anything, mock.MatchedBy( func(prefix string) bool { return strings.HasPrefix(prefix, PartitionMetaPrefix) }), ts). @@ -137,20 +137,20 @@ func TestCatalog_ListCollections(t *testing.T) { bColl, err := proto.Marshal(coll2) assert.NoError(t, err) - kv.On("LoadWithPrefix", CollectionMetaPrefix, ts). + kv.On("LoadWithPrefix", mock.Anything, CollectionMetaPrefix, ts). Return([]string{"key"}, []string{string(bColl)}, nil) partitionMeta := &pb.PartitionInfo{} pm, err := proto.Marshal(partitionMeta) assert.NoError(t, err) - kv.On("LoadWithPrefix", mock.MatchedBy( + kv.On("LoadWithPrefix", mock.Anything, mock.MatchedBy( func(prefix string) bool { return strings.HasPrefix(prefix, PartitionMetaPrefix) }), ts). Return([]string{"key"}, []string{string(pm)}, nil) - kv.On("LoadWithPrefix", mock.MatchedBy( + kv.On("LoadWithPrefix", mock.Anything, mock.MatchedBy( func(prefix string) bool { return strings.HasPrefix(prefix, FieldMetaPrefix) }), ts). @@ -168,9 +168,9 @@ func TestCatalog_ListCollections(t *testing.T) { bColl, err := proto.Marshal(coll1) assert.NoError(t, err) - kv.On("LoadWithPrefix", CollectionMetaPrefix, ts). + kv.On("LoadWithPrefix", mock.Anything, CollectionMetaPrefix, ts). Return([]string{"key"}, []string{string(bColl)}, nil) - kv.On("MultiSaveAndRemove", mock.Anything, mock.Anything, ts).Return(nil) + kv.On("MultiSaveAndRemove", mock.Anything, mock.Anything, mock.Anything, ts).Return(nil) kc := Catalog{Snapshot: kv} ret, err := kc.ListCollections(ctx, util.NonDBID, ts) @@ -186,14 +186,14 @@ func TestCatalog_ListCollections(t *testing.T) { bColl, err := proto.Marshal(coll2) assert.NoError(t, err) - kv.On("LoadWithPrefix", BuildDatabasePrefixWithDBID(testDb), ts). + kv.On("LoadWithPrefix", mock.Anything, BuildDatabasePrefixWithDBID(testDb), ts). Return([]string{"key"}, []string{string(bColl)}, nil) partitionMeta := &pb.PartitionInfo{} pm, err := proto.Marshal(partitionMeta) assert.NoError(t, err) - kv.On("LoadWithPrefix", mock.MatchedBy( + kv.On("LoadWithPrefix", mock.Anything, mock.MatchedBy( func(prefix string) bool { return strings.HasPrefix(prefix, PartitionMetaPrefix) }), ts). @@ -203,7 +203,7 @@ func TestCatalog_ListCollections(t *testing.T) { fm, err := proto.Marshal(fieldMeta) assert.NoError(t, err) - kv.On("LoadWithPrefix", mock.MatchedBy( + kv.On("LoadWithPrefix", mock.Anything, mock.MatchedBy( func(prefix string) bool { return strings.HasPrefix(prefix, FieldMetaPrefix) }), ts). @@ -212,7 +212,7 @@ func TestCatalog_ListCollections(t *testing.T) { functionMeta := &schemapb.FunctionSchema{} fcm, err := proto.Marshal(functionMeta) assert.NoError(t, err) - kv.On("LoadWithPrefix", mock.MatchedBy( + kv.On("LoadWithPrefix", mock.Anything, mock.MatchedBy( func(prefix string) bool { return strings.HasPrefix(prefix, FunctionMetaPrefix) }), ts). @@ -236,14 +236,14 @@ func TestCatalog_ListCollections(t *testing.T) { aColl, err := proto.Marshal(coll3) assert.NoError(t, err) - kv.On("LoadWithPrefix", CollectionMetaPrefix, ts). + kv.On("LoadWithPrefix", mock.Anything, CollectionMetaPrefix, ts). Return([]string{"key", "key2"}, []string{string(bColl), string(aColl)}, nil) partitionMeta := &pb.PartitionInfo{} pm, err := proto.Marshal(partitionMeta) assert.NoError(t, err) - kv.On("LoadWithPrefix", mock.MatchedBy( + kv.On("LoadWithPrefix", mock.Anything, mock.MatchedBy( func(prefix string) bool { return strings.HasPrefix(prefix, PartitionMetaPrefix) }), ts). @@ -253,7 +253,7 @@ func TestCatalog_ListCollections(t *testing.T) { fm, err := proto.Marshal(fieldMeta) assert.NoError(t, err) - kv.On("LoadWithPrefix", mock.MatchedBy( + kv.On("LoadWithPrefix", mock.Anything, mock.MatchedBy( func(prefix string) bool { return strings.HasPrefix(prefix, FieldMetaPrefix) }), ts). @@ -262,13 +262,13 @@ func TestCatalog_ListCollections(t *testing.T) { functionMeta := &schemapb.FunctionSchema{} fcm, err := proto.Marshal(functionMeta) assert.NoError(t, err) - kv.On("LoadWithPrefix", mock.MatchedBy( + kv.On("LoadWithPrefix", mock.Anything, mock.MatchedBy( func(prefix string) bool { return strings.HasPrefix(prefix, FunctionMetaPrefix) }), ts). Return([]string{"rootcoord/functions/1/1"}, []string{string(fcm)}, nil) - kv.On("MultiSaveAndRemove", mock.Anything, mock.Anything, ts).Return(nil) + kv.On("MultiSaveAndRemove", mock.Anything, mock.Anything, mock.Anything, ts).Return(nil) kc := Catalog{Snapshot: kv} ret, err := kc.ListCollections(ctx, util.NonDBID, ts) @@ -284,7 +284,7 @@ func TestCatalog_loadCollection(t *testing.T) { t.Run("load failed", func(t *testing.T) { ctx := context.Background() kv := mocks.NewSnapShotKV(t) - kv.EXPECT().Load(mock.Anything, mock.Anything).Return("", errors.New("mock")) + kv.EXPECT().Load(mock.Anything, mock.Anything, mock.Anything).Return("", errors.New("mock")) kc := Catalog{Snapshot: kv} _, err := kc.loadCollection(ctx, testDb, 1, 0) assert.Error(t, err) @@ -293,7 +293,7 @@ func TestCatalog_loadCollection(t *testing.T) { t.Run("load, not collection info", func(t *testing.T) { ctx := context.Background() kv := mocks.NewSnapShotKV(t) - kv.EXPECT().Load(mock.Anything, mock.Anything).Return("not in pb format", nil) + kv.EXPECT().Load(mock.Anything, mock.Anything, mock.Anything).Return("not in pb format", nil) kc := Catalog{Snapshot: kv} _, err := kc.loadCollection(ctx, testDb, 1, 0) assert.Error(t, err) @@ -305,7 +305,7 @@ func TestCatalog_loadCollection(t *testing.T) { value, err := proto.Marshal(coll) assert.NoError(t, err) kv := mocks.NewSnapShotKV(t) - kv.EXPECT().Load(mock.Anything, mock.Anything).Return(string(value), nil) + kv.EXPECT().Load(mock.Anything, mock.Anything, mock.Anything).Return(string(value), nil) kc := Catalog{Snapshot: kv} got, err := kc.loadCollection(ctx, util.DefaultDBID, 1, 0) assert.NoError(t, err) @@ -322,8 +322,8 @@ func TestCatalog_loadCollection(t *testing.T) { value, err := proto.Marshal(coll) assert.NoError(t, err) kv := mocks.NewSnapShotKV(t) - kv.EXPECT().Load(mock.Anything, mock.Anything).Return(string(value), nil) - kv.EXPECT().MultiSaveAndRemove(mock.Anything, mock.Anything, mock.Anything).Return(nil) + kv.EXPECT().Load(mock.Anything, mock.Anything, mock.Anything).Return(string(value), nil) + kv.EXPECT().MultiSaveAndRemove(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil) kc := Catalog{Snapshot: kv} got, err := kc.loadCollection(ctx, util.NonDBID, 1, 0) assert.NoError(t, err) @@ -380,12 +380,12 @@ func TestCatalog_GetCollectionByID(t *testing.T) { ss := mocks.NewSnapShotKV(t) c := Catalog{Snapshot: ss} - ss.EXPECT().Load(mock.Anything, mock.Anything).Return("", errors.New("load error")).Twice() + ss.EXPECT().Load(mock.Anything, mock.Anything, mock.Anything).Return("", errors.New("load error")).Twice() coll, err := c.GetCollectionByID(ctx, 0, 1, 1) assert.Error(t, err) assert.Nil(t, coll) - ss.EXPECT().Load(mock.Anything, mock.Anything).RunAndReturn(func(key string, ts uint64) (string, error) { + ss.EXPECT().Load(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, key string, ts uint64) (string, error) { collByte, err := proto.Marshal(&pb.CollectionInfo{ ID: 1, Schema: &schemapb.CollectionSchema{ @@ -400,7 +400,7 @@ func TestCatalog_GetCollectionByID(t *testing.T) { require.NoError(t, err) return string(collByte), nil }).Once() - ss.EXPECT().MultiSaveAndRemove(mock.Anything, mock.Anything, mock.Anything).Return(nil) + ss.EXPECT().MultiSaveAndRemove(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil) coll, err = c.GetCollectionByID(ctx, 0, 10000, 1) assert.NoError(t, err) @@ -412,7 +412,7 @@ func TestCatalog_CreatePartitionV2(t *testing.T) { t.Run("collection not exist", func(t *testing.T) { ctx := context.Background() snapshot := kv.NewMockSnapshotKV() - snapshot.LoadFunc = func(key string, ts typeutil.Timestamp) (string, error) { + snapshot.LoadFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) (string, error) { return "", errors.New("mock") } kc := Catalog{Snapshot: snapshot} @@ -430,10 +430,10 @@ func TestCatalog_CreatePartitionV2(t *testing.T) { assert.NoError(t, err) snapshot := kv.NewMockSnapshotKV() - snapshot.LoadFunc = func(key string, ts typeutil.Timestamp) (string, error) { + snapshot.LoadFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) (string, error) { return string(value), nil } - snapshot.SaveFunc = func(key string, value string, ts typeutil.Timestamp) error { + snapshot.SaveFunc = func(ctx context.Context, key string, value string, ts typeutil.Timestamp) error { return errors.New("mock") } @@ -442,7 +442,7 @@ func TestCatalog_CreatePartitionV2(t *testing.T) { err = kc.CreatePartition(ctx, 0, &model.Partition{}, 0) assert.Error(t, err) - snapshot.SaveFunc = func(key string, value string, ts typeutil.Timestamp) error { + snapshot.SaveFunc = func(ctx context.Context, key string, value string, ts typeutil.Timestamp) error { return nil } err = kc.CreatePartition(ctx, 0, &model.Partition{}, 0) @@ -458,7 +458,7 @@ func TestCatalog_CreatePartitionV2(t *testing.T) { assert.NoError(t, err) snapshot := kv.NewMockSnapshotKV() - snapshot.LoadFunc = func(key string, ts typeutil.Timestamp) (string, error) { + snapshot.LoadFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) (string, error) { return string(value), nil } @@ -477,7 +477,7 @@ func TestCatalog_CreatePartitionV2(t *testing.T) { assert.NoError(t, err) snapshot := kv.NewMockSnapshotKV() - snapshot.LoadFunc = func(key string, ts typeutil.Timestamp) (string, error) { + snapshot.LoadFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) (string, error) { return string(value), nil } @@ -500,10 +500,10 @@ func TestCatalog_CreatePartitionV2(t *testing.T) { assert.NoError(t, err) snapshot := kv.NewMockSnapshotKV() - snapshot.LoadFunc = func(key string, ts typeutil.Timestamp) (string, error) { + snapshot.LoadFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) (string, error) { return string(value), nil } - snapshot.SaveFunc = func(key string, value string, ts typeutil.Timestamp) error { + snapshot.SaveFunc = func(ctx context.Context, key string, value string, ts typeutil.Timestamp) error { return errors.New("mock") } @@ -512,7 +512,7 @@ func TestCatalog_CreatePartitionV2(t *testing.T) { err = kc.CreatePartition(ctx, 0, &model.Partition{}, 0) assert.Error(t, err) - snapshot.SaveFunc = func(key string, value string, ts typeutil.Timestamp) error { + snapshot.SaveFunc = func(ctx context.Context, key string, value string, ts typeutil.Timestamp) error { return nil } err = kc.CreatePartition(ctx, 0, &model.Partition{}, 0) @@ -524,7 +524,7 @@ func TestCatalog_CreateAliasV2(t *testing.T) { ctx := context.Background() snapshot := kv.NewMockSnapshotKV() - snapshot.MultiSaveAndRemoveFunc = func(saves map[string]string, removals []string, ts typeutil.Timestamp) error { + snapshot.MultiSaveAndRemoveFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error { return errors.New("mock") } @@ -533,7 +533,7 @@ func TestCatalog_CreateAliasV2(t *testing.T) { err := kc.CreateAlias(ctx, &model.Alias{}, 0) assert.Error(t, err) - snapshot.MultiSaveAndRemoveFunc = func(saves map[string]string, removals []string, ts typeutil.Timestamp) error { + snapshot.MultiSaveAndRemoveFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error { return nil } err = kc.CreateAlias(ctx, &model.Alias{}, 0) @@ -545,7 +545,7 @@ func TestCatalog_listPartitionsAfter210(t *testing.T) { ctx := context.Background() snapshot := kv.NewMockSnapshotKV() - snapshot.LoadWithPrefixFunc = func(key string, ts typeutil.Timestamp) ([]string, []string, error) { + snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) { return nil, nil, errors.New("mock") } @@ -559,7 +559,7 @@ func TestCatalog_listPartitionsAfter210(t *testing.T) { ctx := context.Background() snapshot := kv.NewMockSnapshotKV() - snapshot.LoadWithPrefixFunc = func(key string, ts typeutil.Timestamp) ([]string, []string, error) { + snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) { return []string{"key"}, []string{"not in pb format"}, nil } @@ -577,7 +577,7 @@ func TestCatalog_listPartitionsAfter210(t *testing.T) { assert.NoError(t, err) snapshot := kv.NewMockSnapshotKV() - snapshot.LoadWithPrefixFunc = func(key string, ts typeutil.Timestamp) ([]string, []string, error) { + snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) { return []string{"key"}, []string{string(value)}, nil } @@ -603,7 +603,7 @@ func TestCatalog_listFieldsAfter210(t *testing.T) { ctx := context.Background() snapshot := kv.NewMockSnapshotKV() - snapshot.LoadWithPrefixFunc = func(key string, ts typeutil.Timestamp) ([]string, []string, error) { + snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) { return nil, nil, errors.New("mock") } @@ -617,7 +617,7 @@ func TestCatalog_listFieldsAfter210(t *testing.T) { ctx := context.Background() snapshot := kv.NewMockSnapshotKV() - snapshot.LoadWithPrefixFunc = func(key string, ts typeutil.Timestamp) ([]string, []string, error) { + snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) { return []string{"key"}, []string{"not in pb format"}, nil } @@ -635,7 +635,7 @@ func TestCatalog_listFieldsAfter210(t *testing.T) { assert.NoError(t, err) snapshot := kv.NewMockSnapshotKV() - snapshot.LoadWithPrefixFunc = func(key string, ts typeutil.Timestamp) ([]string, []string, error) { + snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) { return []string{"key"}, []string{string(value)}, nil } @@ -652,7 +652,7 @@ func TestCatalog_AlterAliasV2(t *testing.T) { ctx := context.Background() snapshot := kv.NewMockSnapshotKV() - snapshot.MultiSaveAndRemoveFunc = func(saves map[string]string, removals []string, ts typeutil.Timestamp) error { + snapshot.MultiSaveAndRemoveFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error { return errors.New("mock") } @@ -661,7 +661,7 @@ func TestCatalog_AlterAliasV2(t *testing.T) { err := kc.AlterAlias(ctx, &model.Alias{}, 0) assert.Error(t, err) - snapshot.MultiSaveAndRemoveFunc = func(saves map[string]string, removals []string, ts typeutil.Timestamp) error { + snapshot.MultiSaveAndRemoveFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error { return nil } err = kc.AlterAlias(ctx, &model.Alias{}, 0) @@ -703,7 +703,7 @@ func TestCatalog_DropPartitionV2(t *testing.T) { snapshot := mocks.NewSnapShotKV(t) snapshot.On("Load", - mock.Anything, mock.Anything).Return("not in codec format", nil) + mock.Anything, mock.Anything, mock.Anything).Return("not in codec format", nil) kc := Catalog{Snapshot: snapshot} @@ -716,7 +716,7 @@ func TestCatalog_DropPartitionV2(t *testing.T) { snapshot := mocks.NewSnapShotKV(t) snapshot.On("Load", - mock.Anything, mock.Anything).Return("", merr.WrapErrIoKeyNotFound("partition")) + mock.Anything, mock.Anything, mock.Anything).Return("", merr.WrapErrIoKeyNotFound("partition")) kc := Catalog{Snapshot: snapshot} @@ -732,10 +732,10 @@ func TestCatalog_DropPartitionV2(t *testing.T) { assert.NoError(t, err) snapshot := kv.NewMockSnapshotKV() - snapshot.LoadFunc = func(key string, ts typeutil.Timestamp) (string, error) { + snapshot.LoadFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) (string, error) { return string(value), nil } - snapshot.MultiSaveAndRemoveFunc = func(saves map[string]string, removals []string, ts typeutil.Timestamp) error { + snapshot.MultiSaveAndRemoveFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error { return errors.New("mock") } @@ -744,7 +744,7 @@ func TestCatalog_DropPartitionV2(t *testing.T) { err = kc.DropPartition(ctx, 0, 100, 101, 0) assert.Error(t, err) - snapshot.MultiSaveAndRemoveFunc = func(saves map[string]string, removals []string, ts typeutil.Timestamp) error { + snapshot.MultiSaveAndRemoveFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error { return nil } err = kc.DropPartition(ctx, 0, 100, 101, 0) @@ -764,10 +764,10 @@ func TestCatalog_DropPartitionV2(t *testing.T) { assert.NoError(t, err) snapshot := kv.NewMockSnapshotKV() - snapshot.LoadFunc = func(key string, ts typeutil.Timestamp) (string, error) { + snapshot.LoadFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) (string, error) { return string(value), nil } - snapshot.SaveFunc = func(key string, value string, ts typeutil.Timestamp) error { + snapshot.SaveFunc = func(ctx context.Context, key string, value string, ts typeutil.Timestamp) error { return errors.New("mock") } @@ -776,7 +776,7 @@ func TestCatalog_DropPartitionV2(t *testing.T) { err = kc.DropPartition(ctx, 0, 100, 101, 0) assert.Error(t, err) - snapshot.SaveFunc = func(key string, value string, ts typeutil.Timestamp) error { + snapshot.SaveFunc = func(ctx context.Context, key string, value string, ts typeutil.Timestamp) error { return nil } err = kc.DropPartition(ctx, 0, 100, 102, 0) @@ -788,7 +788,7 @@ func TestCatalog_DropAliasV2(t *testing.T) { ctx := context.Background() snapshot := kv.NewMockSnapshotKV() - snapshot.MultiSaveAndRemoveFunc = func(saves map[string]string, removals []string, ts typeutil.Timestamp) error { + snapshot.MultiSaveAndRemoveFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error { return errors.New("mock") } @@ -797,7 +797,7 @@ func TestCatalog_DropAliasV2(t *testing.T) { err := kc.DropAlias(ctx, testDb, "alias", 0) assert.Error(t, err) - snapshot.MultiSaveAndRemoveFunc = func(saves map[string]string, removals []string, ts typeutil.Timestamp) error { + snapshot.MultiSaveAndRemoveFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error { return nil } err = kc.DropAlias(ctx, testDb, "alias", 0) @@ -809,7 +809,7 @@ func TestCatalog_listAliasesBefore210(t *testing.T) { ctx := context.Background() snapshot := kv.NewMockSnapshotKV() - snapshot.LoadWithPrefixFunc = func(key string, ts typeutil.Timestamp) ([]string, []string, error) { + snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) { return nil, nil, errors.New("mock") } @@ -823,7 +823,7 @@ func TestCatalog_listAliasesBefore210(t *testing.T) { ctx := context.Background() snapshot := kv.NewMockSnapshotKV() - snapshot.LoadWithPrefixFunc = func(key string, ts typeutil.Timestamp) ([]string, []string, error) { + snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) { return []string{"key"}, []string{"not in pb format"}, nil } @@ -841,7 +841,7 @@ func TestCatalog_listAliasesBefore210(t *testing.T) { assert.NoError(t, err) snapshot := kv.NewMockSnapshotKV() - snapshot.LoadWithPrefixFunc = func(key string, ts typeutil.Timestamp) ([]string, []string, error) { + snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) { return []string{"key"}, []string{string(value)}, nil } @@ -859,7 +859,7 @@ func TestCatalog_listAliasesAfter210(t *testing.T) { ctx := context.Background() snapshot := kv.NewMockSnapshotKV() - snapshot.LoadWithPrefixFunc = func(key string, ts typeutil.Timestamp) ([]string, []string, error) { + snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) { return nil, nil, errors.New("mock") } @@ -873,7 +873,7 @@ func TestCatalog_listAliasesAfter210(t *testing.T) { ctx := context.Background() snapshot := kv.NewMockSnapshotKV() - snapshot.LoadWithPrefixFunc = func(key string, ts typeutil.Timestamp) ([]string, []string, error) { + snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) { return []string{"key"}, []string{"not in pb format"}, nil } @@ -891,7 +891,7 @@ func TestCatalog_listAliasesAfter210(t *testing.T) { assert.NoError(t, err) snapshot := kv.NewMockSnapshotKV() - snapshot.LoadWithPrefixFunc = func(key string, ts typeutil.Timestamp) ([]string, []string, error) { + snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) { return []string{"key"}, []string{string(value)}, nil } @@ -909,7 +909,7 @@ func TestCatalog_ListAliasesV2(t *testing.T) { ctx := context.Background() snapshot := kv.NewMockSnapshotKV() - snapshot.LoadWithPrefixFunc = func(key string, ts typeutil.Timestamp) ([]string, []string, error) { + snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) { return []string{"key"}, []string{"not in pb format"}, nil } @@ -927,7 +927,7 @@ func TestCatalog_ListAliasesV2(t *testing.T) { assert.NoError(t, err) snapshot := kv.NewMockSnapshotKV() - snapshot.LoadWithPrefixFunc = func(key string, ts typeutil.Timestamp) ([]string, []string, error) { + snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) { if key == AliasMetaPrefix { return nil, nil, errors.New("mock") } @@ -955,7 +955,7 @@ func TestCatalog_ListAliasesV2(t *testing.T) { assert.NoError(t, err) snapshot := kv.NewMockSnapshotKV() - snapshot.LoadWithPrefixFunc = func(key string, ts typeutil.Timestamp) ([]string, []string, error) { + snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) { dbStr := fmt.Sprintf("%d", testDb) if strings.Contains(key, dbStr) && strings.Contains(key, Aliases) { return []string{"key1"}, []string{string(value2)}, nil @@ -975,33 +975,33 @@ func TestCatalog_ListAliasesV2(t *testing.T) { func Test_batchMultiSaveAndRemove(t *testing.T) { t.Run("failed to save", func(t *testing.T) { snapshot := kv.NewMockSnapshotKV() - snapshot.MultiSaveFunc = func(kvs map[string]string, ts typeutil.Timestamp) error { + snapshot.MultiSaveFunc = func(ctx context.Context, kvs map[string]string, ts typeutil.Timestamp) error { return errors.New("error mock MultiSave") } saves := map[string]string{"k": "v"} - err := batchMultiSaveAndRemove(snapshot, util.MaxEtcdTxnNum/2, saves, []string{}, 0) + err := batchMultiSaveAndRemove(context.TODO(), snapshot, util.MaxEtcdTxnNum/2, saves, []string{}, 0) assert.Error(t, err) }) t.Run("failed to remove", func(t *testing.T) { snapshot := kv.NewMockSnapshotKV() - snapshot.MultiSaveFunc = func(kvs map[string]string, ts typeutil.Timestamp) error { + snapshot.MultiSaveFunc = func(ctx context.Context, kvs map[string]string, ts typeutil.Timestamp) error { return nil } - snapshot.MultiSaveAndRemoveFunc = func(saves map[string]string, removals []string, ts typeutil.Timestamp) error { + snapshot.MultiSaveAndRemoveFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error { return errors.New("error mock MultiSaveAndRemove") } saves := map[string]string{"k": "v"} removals := []string{"prefix1", "prefix2"} - err := batchMultiSaveAndRemove(snapshot, util.MaxEtcdTxnNum/2, saves, removals, 0) + err := batchMultiSaveAndRemove(context.TODO(), snapshot, util.MaxEtcdTxnNum/2, saves, removals, 0) assert.Error(t, err) }) t.Run("normal case", func(t *testing.T) { snapshot := kv.NewMockSnapshotKV() - snapshot.MultiSaveFunc = func(kvs map[string]string, ts typeutil.Timestamp) error { + snapshot.MultiSaveFunc = func(ctx context.Context, kvs map[string]string, ts typeutil.Timestamp) error { log.Info("multi save", zap.Any("len", len(kvs)), zap.Any("saves", kvs)) return nil } - snapshot.MultiSaveAndRemoveFunc = func(saves map[string]string, removals []string, ts typeutil.Timestamp) error { + snapshot.MultiSaveAndRemoveFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error { log.Info("multi save and remove with prefix", zap.Any("len of saves", len(saves)), zap.Any("len of removals", len(removals)), zap.Any("saves", saves), zap.Any("removals", removals)) return nil @@ -1013,7 +1013,7 @@ func Test_batchMultiSaveAndRemove(t *testing.T) { saves[fmt.Sprintf("k%d", i)] = fmt.Sprintf("v%d", i) removals = append(removals, fmt.Sprintf("k%d", i)) } - err := batchMultiSaveAndRemove(snapshot, util.MaxEtcdTxnNum/2, saves, removals, 0) + err := batchMultiSaveAndRemove(context.TODO(), snapshot, util.MaxEtcdTxnNum/2, saves, removals, 0) assert.NoError(t, err) }) } @@ -1036,7 +1036,7 @@ func TestCatalog_AlterCollection(t *testing.T) { t.Run("modify", func(t *testing.T) { snapshot := kv.NewMockSnapshotKV() kvs := map[string]string{} - snapshot.SaveFunc = func(key string, value string, ts typeutil.Timestamp) error { + snapshot.SaveFunc = func(ctx context.Context, key string, value string, ts typeutil.Timestamp) error { kvs[key] = value return nil } @@ -1070,7 +1070,7 @@ func TestCatalog_AlterCollection(t *testing.T) { t.Run("modify db name", func(t *testing.T) { var collectionID int64 = 1 snapshot := kv.NewMockSnapshotKV() - snapshot.MultiSaveAndRemoveFunc = func(saves map[string]string, removals []string, ts typeutil.Timestamp) error { + snapshot.MultiSaveAndRemoveFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error { assert.ElementsMatch(t, []string{BuildCollectionKey(0, collectionID)}, removals) assert.Equal(t, len(saves), 1) assert.Contains(t, maps.Keys(saves), BuildCollectionKey(1, collectionID)) @@ -1104,7 +1104,7 @@ func TestCatalog_AlterPartition(t *testing.T) { t.Run("modify", func(t *testing.T) { snapshot := kv.NewMockSnapshotKV() kvs := map[string]string{} - snapshot.SaveFunc = func(key string, value string, ts typeutil.Timestamp) error { + snapshot.SaveFunc = func(ctx context.Context, key string, value string, ts typeutil.Timestamp) error { kvs[key] = value return nil } @@ -1151,6 +1151,7 @@ func withMockSave(saveErr error) mockSnapshotOpt { return func(ss *mocks.SnapShotKV) { ss.On( "Save", + mock.Anything, mock.AnythingOfType("string"), mock.AnythingOfType("string"), mock.AnythingOfType("uint64")). @@ -1162,6 +1163,7 @@ func withMockMultiSave(multiSaveErr error) mockSnapshotOpt { return func(ss *mocks.SnapShotKV) { ss.On( "MultiSave", + mock.Anything, mock.AnythingOfType("map[string]string"), mock.AnythingOfType("uint64")). Return(multiSaveErr) @@ -1172,6 +1174,7 @@ func withMockMultiSaveAndRemoveWithPrefix(err error) mockSnapshotOpt { return func(ss *mocks.SnapShotKV) { ss.On( "MultiSaveAndRemoveWithPrefix", + mock.Anything, mock.AnythingOfType("map[string]string"), mock.AnythingOfType("[]string"), mock.AnythingOfType("uint64")). @@ -1183,6 +1186,7 @@ func withMockMultiSaveAndRemove(err error) mockSnapshotOpt { return func(ss *mocks.SnapShotKV) { ss.On( "MultiSaveAndRemove", + mock.Anything, mock.AnythingOfType("map[string]string"), mock.AnythingOfType("[]string"), mock.AnythingOfType("uint64")). @@ -1296,19 +1300,21 @@ func TestCatalog_DropCollection(t *testing.T) { removeCollectionCalled := false mockSnapshot.On( "MultiSaveAndRemove", + mock.Anything, mock.AnythingOfType("map[string]string"), mock.AnythingOfType("[]string"), mock.AnythingOfType("uint64")). - Return(func(map[string]string, []string, typeutil.Timestamp) error { + Return(func(context.Context, map[string]string, []string, typeutil.Timestamp) error { removeOtherCalled = true return nil }).Once() mockSnapshot.On( "MultiSaveAndRemove", + mock.Anything, mock.AnythingOfType("map[string]string"), mock.AnythingOfType("[]string"), mock.AnythingOfType("uint64")). - Return(func(map[string]string, []string, typeutil.Timestamp) error { + Return(func(context.Context, map[string]string, []string, typeutil.Timestamp) error { removeCollectionCalled = true return errors.New("error mock MultiSaveAndRemove") }).Once() @@ -1399,10 +1405,10 @@ func TestRBAC_Credential(t *testing.T) { marshalFailkey = fmt.Sprintf("%s/%s", CredentialPrefix, marshalFailName) ) - kvmock.EXPECT().Load(loadFailKey).Return("", errors.New("Mock invalid load")) - kvmock.EXPECT().Load(marshalFailkey).Return("random", nil) - kvmock.EXPECT().Load(mock.Anything).Call.Return( - func(key string) string { + kvmock.EXPECT().Load(mock.Anything, loadFailKey).Return("", errors.New("Mock invalid load")) + kvmock.EXPECT().Load(mock.Anything, marshalFailkey).Return("random", nil) + kvmock.EXPECT().Load(mock.Anything, mock.Anything).Call.Return( + func(ctx context.Context, key string) string { v, err := json.Marshal(&internalpb.CredentialInfo{EncryptedPassword: key}) require.NoError(t, err) return string(v) @@ -1447,8 +1453,8 @@ func TestRBAC_Credential(t *testing.T) { ) mockFailPath := fmt.Sprintf("%s/%s", CredentialPrefix, invalidName) - kvmock.EXPECT().Save(mockFailPath, mock.Anything).Return(errors.New("Mock invalid save")) - kvmock.EXPECT().Save(mock.Anything, mock.Anything).Return(nil) + kvmock.EXPECT().Save(mock.Anything, mockFailPath, mock.Anything).Return(errors.New("Mock invalid save")) + kvmock.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(nil) tests := []struct { description string @@ -1502,26 +1508,27 @@ func TestRBAC_Credential(t *testing.T) { getFailName = "get-fail" ) - kvmock.EXPECT().MultiRemove([]string{fmt.Sprintf("%s/%s", CredentialPrefix, dropFailName)}).Return(errors.New("Mock drop fail")) + kvmock.EXPECT().MultiRemove(mock.Anything, []string{fmt.Sprintf("%s/%s", CredentialPrefix, dropFailName)}).Return(errors.New("Mock drop fail")) kvmock.EXPECT().MultiRemove( + mock.Anything, []string{ fmt.Sprintf("%s/%s", CredentialPrefix, validName), validUserRoleKeyPrefix + "/role1", validUserRoleKeyPrefix + "/role2", }, ).Return(nil) - kvmock.EXPECT().MultiRemove(mock.Anything).Return(errors.New("Mock invalid multi remove")) + kvmock.EXPECT().MultiRemove(mock.Anything, mock.Anything).Return(errors.New("Mock invalid multi remove")) - kvmock.EXPECT().Load(fmt.Sprintf("%s/%s", CredentialPrefix, getFailName)).Return("", errors.New("Mock invalid load")) - kvmock.EXPECT().Load(fmt.Sprintf("%s/%s", CredentialPrefix, validName)).Return(getUserInfoMetaString(validName), nil) - kvmock.EXPECT().Load(fmt.Sprintf("%s/%s", CredentialPrefix, dropFailName)).Return(getUserInfoMetaString(dropFailName), nil) + kvmock.EXPECT().Load(mock.Anything, fmt.Sprintf("%s/%s", CredentialPrefix, getFailName)).Return("", errors.New("Mock invalid load")) + kvmock.EXPECT().Load(mock.Anything, fmt.Sprintf("%s/%s", CredentialPrefix, validName)).Return(getUserInfoMetaString(validName), nil) + kvmock.EXPECT().Load(mock.Anything, fmt.Sprintf("%s/%s", CredentialPrefix, dropFailName)).Return(getUserInfoMetaString(dropFailName), nil) - kvmock.EXPECT().LoadWithPrefix(validUserRoleKeyPrefix).Return( + kvmock.EXPECT().LoadWithPrefix(mock.Anything, validUserRoleKeyPrefix).Return( []string{validUserRoleKeyPrefix + "/role1", validUserRoleKeyPrefix + "/role2"}, []string{"", ""}, nil, ) - kvmock.EXPECT().LoadWithPrefix(dropUserRoleKeyPrefix).Return([]string{}, []string{}, nil) + kvmock.EXPECT().LoadWithPrefix(mock.Anything, dropUserRoleKeyPrefix).Return([]string{}, []string{}, nil) tests := []struct { description string @@ -1557,8 +1564,8 @@ func TestRBAC_Credential(t *testing.T) { // Return valid keys if count==0 // return error if count!=0 - kvmock.EXPECT().LoadWithPrefix(mock.Anything).Call.Return( - func(key string) []string { + kvmock.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Call.Return( + func(ctx context.Context, key string) []string { cmu.RLock() defer cmu.RUnlock() if count == 0 { @@ -1571,7 +1578,7 @@ func TestRBAC_Credential(t *testing.T) { } return nil }, - func(key string) []string { + func(ctx context.Context, key string) []string { cmu.RLock() defer cmu.RUnlock() passwd, _ := json.Marshal(&model.Credential{EncryptedPassword: crypto.Base64Encode("passwd")}) @@ -1585,7 +1592,7 @@ func TestRBAC_Credential(t *testing.T) { } return nil }, - func(key string) error { + func(ctx context.Context, key string) error { cmu.RLock() defer cmu.RUnlock() if count == 0 { @@ -1638,10 +1645,10 @@ func TestRBAC_Role(t *testing.T) { otherError = fmt.Errorf("mock load error") ) - kvmock.EXPECT().Load(notExistKey).Return("", merr.WrapErrIoKeyNotFound(notExistKey)).Once() - kvmock.EXPECT().Load(errorKey).Return("", otherError).Once() - kvmock.EXPECT().Load(mock.Anything).Return("", nil).Once() - kvmock.EXPECT().Remove(mock.Anything).Call.Return(nil).Once() + kvmock.EXPECT().Load(mock.Anything, notExistKey).Return("", merr.WrapErrIoKeyNotFound(notExistKey)).Once() + kvmock.EXPECT().Load(mock.Anything, errorKey).Return("", otherError).Once() + kvmock.EXPECT().Load(mock.Anything, mock.Anything).Return("", nil).Once() + kvmock.EXPECT().Remove(mock.Anything, mock.Anything).Call.Return(nil).Once() tests := []struct { description string @@ -1657,7 +1664,7 @@ func TestRBAC_Role(t *testing.T) { } for _, test := range tests { t.Run(test.description, func(t *testing.T) { - err := c.remove(test.key) + err := c.remove(context.TODO(), test.key) if test.isValid { assert.NoError(t, err) } else { @@ -1682,10 +1689,10 @@ func TestRBAC_Role(t *testing.T) { otherError = fmt.Errorf("mock load error") ) - kvmock.EXPECT().Load(notExistKey).Return("", merr.WrapErrIoKeyNotFound(notExistKey)).Once() - kvmock.EXPECT().Load(errorKey).Return("", otherError).Once() - kvmock.EXPECT().Load(mock.Anything).Return("", nil).Once() - kvmock.EXPECT().Save(mock.Anything, mock.Anything).Call.Return(nil).Once() + kvmock.EXPECT().Load(mock.Anything, notExistKey).Return("", merr.WrapErrIoKeyNotFound(notExistKey)).Once() + kvmock.EXPECT().Load(mock.Anything, errorKey).Return("", otherError).Once() + kvmock.EXPECT().Load(mock.Anything, mock.Anything).Return("", nil).Once() + kvmock.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Call.Return(nil).Once() tests := []struct { description string @@ -1703,7 +1710,7 @@ func TestRBAC_Role(t *testing.T) { for _, test := range tests { t.Run(test.description, func(t *testing.T) { - err := c.save(test.key) + err := c.save(context.TODO(), test.key) if test.isValid { assert.NoError(t, err) } else { @@ -1729,10 +1736,10 @@ func TestRBAC_Role(t *testing.T) { otherError = fmt.Errorf("mock load error") ) - kvmock.EXPECT().Load(notExistPath).Return("", merr.WrapErrIoKeyNotFound(notExistName)).Once() - kvmock.EXPECT().Load(errorPath).Return("", otherError).Once() - kvmock.EXPECT().Load(mock.Anything).Return("", nil).Once() - kvmock.EXPECT().Save(mock.Anything, mock.Anything).Call.Return(nil).Once() + kvmock.EXPECT().Load(mock.Anything, notExistPath).Return("", merr.WrapErrIoKeyNotFound(notExistName)).Once() + kvmock.EXPECT().Load(mock.Anything, errorPath).Return("", otherError).Once() + kvmock.EXPECT().Load(mock.Anything, mock.Anything).Return("", nil).Once() + kvmock.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Call.Return(nil).Once() tests := []struct { description string @@ -1770,27 +1777,27 @@ func TestRBAC_Role(t *testing.T) { getFailName = "get-fail" ) - kvmock.EXPECT().MultiRemove([]string{funcutil.HandleTenantForEtcdKey(RolePrefix, tenant, errorName)}).Return(errors.New("remove error")) - kvmock.EXPECT().MultiRemove([]string{ + kvmock.EXPECT().MultiRemove(mock.Anything, []string{funcutil.HandleTenantForEtcdKey(RolePrefix, tenant, errorName)}).Return(errors.New("remove error")) + kvmock.EXPECT().MultiRemove(mock.Anything, []string{ funcutil.HandleTenantForEtcdKey(RolePrefix, tenant, validName), funcutil.HandleTenantForEtcdKey(RoleMappingPrefix, tenant, fmt.Sprintf("%s/%s", "user1", validName)), funcutil.HandleTenantForEtcdKey(RoleMappingPrefix, tenant, fmt.Sprintf("%s/%s", "user2", validName)), }).Return(nil) - kvmock.EXPECT().MultiRemove(mock.Anything).Return(errors.New("mock multi remove error")) + kvmock.EXPECT().MultiRemove(mock.Anything, mock.Anything).Return(errors.New("mock multi remove error")) getRoleMappingKey := func(username, rolename string) string { return funcutil.HandleTenantForEtcdKey(RoleMappingPrefix, tenant, fmt.Sprintf("%s/%s", username, rolename)) } - kvmock.EXPECT().LoadWithPrefix(funcutil.HandleTenantForEtcdKey(RoleMappingPrefix, tenant, "")).Return( + kvmock.EXPECT().LoadWithPrefix(mock.Anything, funcutil.HandleTenantForEtcdKey(RoleMappingPrefix, tenant, "")).Return( []string{getRoleMappingKey("user1", validName), getRoleMappingKey("user2", validName), getRoleMappingKey("user3", "role3")}, []string{}, nil, ) - kvmock.EXPECT().Load(funcutil.HandleTenantForEtcdKey(RolePrefix, tenant, getFailName)).Return("", errors.New("mock load error")) - kvmock.EXPECT().Load(funcutil.HandleTenantForEtcdKey(RolePrefix, tenant, validName)).Return("", nil) - kvmock.EXPECT().Load(funcutil.HandleTenantForEtcdKey(RolePrefix, tenant, errorName)).Return("", nil) + kvmock.EXPECT().Load(mock.Anything, funcutil.HandleTenantForEtcdKey(RolePrefix, tenant, getFailName)).Return("", errors.New("mock load error")) + kvmock.EXPECT().Load(mock.Anything, funcutil.HandleTenantForEtcdKey(RolePrefix, tenant, validName)).Return("", nil) + kvmock.EXPECT().Load(mock.Anything, funcutil.HandleTenantForEtcdKey(RolePrefix, tenant, errorName)).Return("", nil) tests := []struct { description string @@ -1829,20 +1836,20 @@ func TestRBAC_Role(t *testing.T) { errorRoleRemove = "error-role-remove" errorRoleRemovepath = funcutil.HandleTenantForEtcdKey(RoleMappingPrefix, tenant, fmt.Sprintf("%s/%s", user, errorRoleRemove)) ) - kvmock.EXPECT().Save(mock.Anything, mock.Anything).Return(nil) - kvmock.EXPECT().Remove(mock.Anything).Return(nil) + kvmock.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(nil) + kvmock.EXPECT().Remove(mock.Anything, mock.Anything).Return(nil) // Catalog.save() returns error - kvmock.EXPECT().Load(errorRoleSavepath).Return("", nil) + kvmock.EXPECT().Load(mock.Anything, errorRoleSavepath).Return("", nil) // Catalog.save() returns nil - kvmock.EXPECT().Load(noErrorRoleSavepath).Return("", merr.WrapErrIoKeyNotFound(noErrorRoleSavepath)) + kvmock.EXPECT().Load(mock.Anything, noErrorRoleSavepath).Return("", merr.WrapErrIoKeyNotFound(noErrorRoleSavepath)) // Catalog.remove() returns error - kvmock.EXPECT().Load(errorRoleRemovepath).Return("", errors.New("not exists")) + kvmock.EXPECT().Load(mock.Anything, errorRoleRemovepath).Return("", errors.New("not exists")) // Catalog.remove() returns nil - kvmock.EXPECT().Load(mock.Anything).Return("", nil) + kvmock.EXPECT().Load(mock.Anything, mock.Anything).Return("", nil) tests := []struct { description string @@ -1885,13 +1892,13 @@ func TestRBAC_Role(t *testing.T) { errorLoadPath = funcutil.HandleTenantForEtcdKey(RolePrefix, tenant, errorLoad) ) - kvmock.EXPECT().Load(errorLoadPath).Call.Return("", errors.New("mock load error")) - kvmock.EXPECT().Load(mock.Anything).Call.Return("", nil) + kvmock.EXPECT().Load(mock.Anything, errorLoadPath).Call.Return("", errors.New("mock load error")) + kvmock.EXPECT().Load(mock.Anything, mock.Anything).Call.Return("", nil) // Return valid keys if loadWithPrefixReturn == True // return error if loadWithPrefixReturn == False - kvmock.EXPECT().LoadWithPrefix(mock.Anything).Call.Return( - func(key string) []string { + kvmock.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Call.Return( + func(ctx context.Context, key string) []string { if loadWithPrefixReturn.Load() { return []string{ fmt.Sprintf("%s/%s/%s", RoleMappingPrefix, tenant, "user1/role1"), @@ -1903,7 +1910,7 @@ func TestRBAC_Role(t *testing.T) { return nil }, nil, - func(key string) error { + func(ctx context.Context, key string) error { if loadWithPrefixReturn.Load() { return nil } @@ -1961,8 +1968,8 @@ func TestRBAC_Role(t *testing.T) { // Return valid keys if loadWithPrefixReturn == True // return error if loadWithPrefixReturn == False // Mocking the return of kv_catalog.go:L699 - kvmock.EXPECT().LoadWithPrefix(mock.Anything).Call.Return( - func(key string) []string { + kvmock.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Call.Return( + func(ctx context.Context, key string) []string { if loadWithPrefixReturn.Load() { return []string{ fmt.Sprintf("%s/%s/%s", RolePrefix, tenant, "role1"), @@ -1974,7 +1981,7 @@ func TestRBAC_Role(t *testing.T) { return nil }, nil, - func(key string) error { + func(ctx context.Context, key string) error { if loadWithPrefixReturn.Load() { return nil } @@ -2018,13 +2025,13 @@ func TestRBAC_Role(t *testing.T) { invalidUserKey = funcutil.HandleTenantForEtcdKey(RoleMappingPrefix, tenant, invalidUser) ) // returns error for invalidUserKey - kvmock.EXPECT().LoadWithPrefix(invalidUserKey).Call.Return( + kvmock.EXPECT().LoadWithPrefix(mock.Anything, invalidUserKey).Call.Return( nil, nil, errors.New("Mock load with prefix wrong")) // Returns keys for RoleMappingPrefix/tenant/user1 user1Key := fmt.Sprintf("%s/%s/%s", RoleMappingPrefix, tenant, "user1") - kvmock.EXPECT().LoadWithPrefix(user1Key).Call.Return( - func(key string) []string { + kvmock.EXPECT().LoadWithPrefix(mock.Anything, user1Key).Call.Return( + func(ctx context.Context, key string) []string { return []string{ fmt.Sprintf("%s/%s/%s", RoleMappingPrefix, tenant, "user1/role1"), fmt.Sprintf("%s/%s/%s", RoleMappingPrefix, tenant, "user1/role2"), @@ -2034,8 +2041,8 @@ func TestRBAC_Role(t *testing.T) { // Returns keys for CredentialPrefix var loadCredentialPrefixReturn atomic.Bool - kvmock.EXPECT().LoadWithPrefix(CredentialPrefix).Call.Return( - func(key string) []string { + kvmock.EXPECT().LoadWithPrefix(mock.Anything, CredentialPrefix).Call.Return( + func(ctx context.Context, key string) []string { if loadCredentialPrefixReturn.Load() { return []string{ fmt.Sprintf("%s/%s/%s", CredentialPrefix, UserSubPrefix, "user1"), @@ -2043,14 +2050,14 @@ func TestRBAC_Role(t *testing.T) { } return nil }, - func(key string) []string { + func(ctx context.Context, key string) []string { if loadCredentialPrefixReturn.Load() { passwd, _ := json.Marshal(&model.Credential{EncryptedPassword: crypto.Base64Encode("passwd")}) return []string{string(passwd)} } return nil }, - func(key string) error { + func(ctx context.Context, key string) error { if loadCredentialPrefixReturn.Load() { return nil } @@ -2074,7 +2081,7 @@ func TestRBAC_Role(t *testing.T) { for _, test := range tests { t.Run(test.description, func(t *testing.T) { - res, err := c.getUserResult(tenant, test.user, test.includeRoleInfo) + res, err := c.getUserResult(context.TODO(), tenant, test.user, test.includeRoleInfo) assert.Equal(t, test.user, res.GetUser().GetName()) @@ -2101,9 +2108,9 @@ func TestRBAC_Role(t *testing.T) { invalidUserLoadKey = fmt.Sprintf("%s/%s", CredentialPrefix, invalidUserLoad) ) // Returns error for invalidUserLoadKey - kvmock.EXPECT().Load(invalidUserLoadKey).Call.Return("", errors.New("Mock load wrong")) - kvmock.EXPECT().Load(mock.Anything).Call.Return( - func(key string) string { + kvmock.EXPECT().Load(mock.Anything, invalidUserLoadKey).Call.Return("", errors.New("Mock load wrong")) + kvmock.EXPECT().Load(mock.Anything, mock.Anything).Call.Return( + func(ctx context.Context, key string) string { v, err := json.Marshal(&internalpb.CredentialInfo{EncryptedPassword: key}) require.NoError(t, err) return string(v) @@ -2160,8 +2167,8 @@ func TestRBAC_Role(t *testing.T) { // Return valid keys if loadWithPrefixReturn == True // return error if loadWithPrefixReturn == False // Mocking the return of kv_catalog.go:ListUserRole:L982 - kvmock.EXPECT().LoadWithPrefix(mock.Anything).Call.Return( - func(key string) []string { + kvmock.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Call.Return( + func(ctx context.Context, key string) []string { if loadWithPrefixReturn.Load() { return []string{ fmt.Sprintf("%s/%s/%s", RoleMappingPrefix, tenant, "user1/role1"), @@ -2173,7 +2180,7 @@ func TestRBAC_Role(t *testing.T) { return nil }, nil, - func(key string) error { + func(ctx context.Context, key string) error { if loadWithPrefixReturn.Load() { return nil } @@ -2247,35 +2254,35 @@ func TestRBAC_Grant(t *testing.T) { errorSaveRoleKeyWithDb := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, fmt.Sprintf("%s/%s/%s", errorSaveRole, object, funcutil.CombineObjectName(util.DefaultDBName, objName))) // Mock return in kv_catalog.go:AlterGrant:L815 - kvmock.EXPECT().Load(validRoleKey).Call. - Return(func(key string) string { return validRoleValue }, nil) + kvmock.EXPECT().Load(mock.Anything, validRoleKey).Call. + Return(func(ctx context.Context, key string) string { return validRoleValue }, nil) - kvmock.EXPECT().Load(invalidRoleKey).Call. - Return("", func(key string) error { + kvmock.EXPECT().Load(mock.Anything, invalidRoleKey).Call. + Return("", func(ctx context.Context, key string) error { return fmt.Errorf("mock load error, key=%s", key) }) - kvmock.EXPECT().Load(invalidRoleKeyWithDb).Call. - Return("", func(key string) error { + kvmock.EXPECT().Load(mock.Anything, invalidRoleKeyWithDb).Call. + Return("", func(ctx context.Context, key string) error { return fmt.Errorf("mock load error, key=%s", key) }) - kvmock.EXPECT().Load(keyNotExistRoleKey).Call. - Return("", func(key string) error { + kvmock.EXPECT().Load(mock.Anything, keyNotExistRoleKey).Call. + Return("", func(ctx context.Context, key string) error { return merr.WrapErrIoKeyNotFound(key) }) - kvmock.EXPECT().Load(keyNotExistRoleKeyWithDb).Call. - Return("", func(key string) error { + kvmock.EXPECT().Load(mock.Anything, keyNotExistRoleKeyWithDb).Call. + Return("", func(ctx context.Context, key string) error { return merr.WrapErrIoKeyNotFound(key) }) - kvmock.EXPECT().Load(errorSaveRoleKey).Call. - Return("", func(key string) error { + kvmock.EXPECT().Load(mock.Anything, errorSaveRoleKey).Call. + Return("", func(ctx context.Context, key string) error { return merr.WrapErrIoKeyNotFound(key) }) - kvmock.EXPECT().Load(errorSaveRoleKeyWithDb).Call. - Return("", func(key string) error { + kvmock.EXPECT().Load(mock.Anything, errorSaveRoleKeyWithDb).Call. + Return("", func(ctx context.Context, key string) error { return merr.WrapErrIoKeyNotFound(key) }) - kvmock.EXPECT().Save(keyNotExistRoleKeyWithDb, mock.Anything).Return(nil) - kvmock.EXPECT().Save(errorSaveRoleKeyWithDb, mock.Anything).Return(errors.New("mock save error role")) + kvmock.EXPECT().Save(mock.Anything, keyNotExistRoleKeyWithDb, mock.Anything).Return(nil) + kvmock.EXPECT().Save(mock.Anything, errorSaveRoleKeyWithDb, mock.Anything).Return(errors.New("mock save error role")) validPrivilegeKey := funcutil.HandleTenantForEtcdKey(GranteeIDPrefix, tenant, fmt.Sprintf("%s/%s", validRoleValue, validPrivilege)) invalidPrivilegeKey := funcutil.HandleTenantForEtcdKey(GranteeIDPrefix, tenant, fmt.Sprintf("%s/%s", validRoleValue, invalidPrivilege)) @@ -2283,24 +2290,24 @@ func TestRBAC_Grant(t *testing.T) { keyNotExistPrivilegeKey2WithDb := funcutil.HandleTenantForEtcdKey(GranteeIDPrefix, tenant, fmt.Sprintf("%s/%s", keyNotExistRoleValueWithDb, keyNotExistPrivilege2)) // Mock return in kv_catalog.go:AlterGrant:L838 - kvmock.EXPECT().Load(validPrivilegeKey).Call.Return("", nil) - kvmock.EXPECT().Load(invalidPrivilegeKey).Call. - Return("", func(key string) error { + kvmock.EXPECT().Load(mock.Anything, validPrivilegeKey).Call.Return("", nil) + kvmock.EXPECT().Load(mock.Anything, invalidPrivilegeKey).Call. + Return("", func(ctx context.Context, key string) error { return fmt.Errorf("mock load error, key=%s", key) }) - kvmock.EXPECT().Load(keyNotExistPrivilegeKey).Call. - Return("", func(key string) error { + kvmock.EXPECT().Load(mock.Anything, keyNotExistPrivilegeKey).Call. + Return("", func(ctx context.Context, key string) error { return merr.WrapErrIoKeyNotFound(key) }) - kvmock.EXPECT().Load(keyNotExistPrivilegeKey2WithDb).Call. - Return("", func(key string) error { + kvmock.EXPECT().Load(mock.Anything, keyNotExistPrivilegeKey2WithDb).Call. + Return("", func(ctx context.Context, key string) error { return merr.WrapErrIoKeyNotFound(key) }) - kvmock.EXPECT().Load(mock.Anything).Call.Return("", nil) + kvmock.EXPECT().Load(mock.Anything, mock.Anything).Call.Return("", nil) t.Run("test Grant", func(t *testing.T) { - kvmock.EXPECT().Save(mock.Anything, validUser).Return(nil) - kvmock.EXPECT().Save(mock.Anything, invalidUser).Return(errors.New("mock save invalid user")) + kvmock.EXPECT().Save(mock.Anything, mock.Anything, validUser).Return(nil) + kvmock.EXPECT().Save(mock.Anything, mock.Anything, invalidUser).Return(errors.New("mock save invalid user")) tests := []struct { isValid bool @@ -2359,9 +2366,9 @@ func TestRBAC_Grant(t *testing.T) { invalidPrivilegeRemove := "p-remove" invalidPrivilegeRemoveKey := funcutil.HandleTenantForEtcdKey(GranteeIDPrefix, tenant, fmt.Sprintf("%s/%s", validRoleValue, invalidPrivilegeRemove)) - kvmock.EXPECT().Load(invalidPrivilegeRemoveKey).Call.Return("", nil) - kvmock.EXPECT().Remove(invalidPrivilegeRemoveKey).Return(errors.New("mock remove error")) - kvmock.EXPECT().Remove(mock.Anything).Return(nil) + kvmock.EXPECT().Load(mock.Anything, invalidPrivilegeRemoveKey).Call.Return("", nil) + kvmock.EXPECT().Remove(mock.Anything, invalidPrivilegeRemoveKey).Return(errors.New("mock remove error")) + kvmock.EXPECT().Remove(mock.Anything, mock.Anything).Return(nil) tests := []struct { isValid bool @@ -2426,10 +2433,10 @@ func TestRBAC_Grant(t *testing.T) { granteePrefix = funcutil.HandleTenantForEtcdKey(GranteeIDPrefix, tenant, granteeID+"/") ) - kvmock.EXPECT().LoadWithPrefix(loadErrorRolePrefix).Call.Return(nil, nil, errors.New("mock loadWithPrefix error")) - kvmock.EXPECT().LoadWithPrefix(mock.Anything).Call.Return(nil, []string{granteeID}, nil) - kvmock.EXPECT().MultiSaveAndRemoveWithPrefix(mock.Anything, []string{errorRolePrefix, granteePrefix}, mock.Anything).Call.Return(errors.New("mock removeWithPrefix error")) - kvmock.EXPECT().MultiSaveAndRemoveWithPrefix(mock.Anything, mock.Anything, mock.Anything).Call.Return(nil) + kvmock.EXPECT().LoadWithPrefix(mock.Anything, loadErrorRolePrefix).Call.Return(nil, nil, errors.New("mock loadWithPrefix error")) + kvmock.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Call.Return(nil, []string{granteeID}, nil) + kvmock.EXPECT().MultiSaveAndRemoveWithPrefix(mock.Anything, mock.Anything, []string{errorRolePrefix, granteePrefix}, mock.Anything).Call.Return(errors.New("mock removeWithPrefix error")) + kvmock.EXPECT().MultiSaveAndRemoveWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Call.Return(nil) tests := []struct { isValid bool @@ -2462,24 +2469,24 @@ func TestRBAC_Grant(t *testing.T) { // Mock Load in kv_catalog.go:L901 validGranteeKey := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, fmt.Sprintf("%s/%s/%s", "role1", "obj1", "obj_name1")) - kvmock.EXPECT().Load(validGranteeKey).Call. - Return(func(key string) string { return crypto.MD5(key) }, nil) + kvmock.EXPECT().Load(mock.Anything, validGranteeKey).Call. + Return(func(ctx context.Context, key string) string { return crypto.MD5(key) }, nil) validGranteeKey2 := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, fmt.Sprintf("%s/%s/%s", "role1", "obj2", "foo.obj_name2")) - kvmock.EXPECT().Load(validGranteeKey2).Call. - Return(func(key string) string { return crypto.MD5(key) }, nil) + kvmock.EXPECT().Load(mock.Anything, validGranteeKey2).Call. + Return(func(ctx context.Context, key string) string { return crypto.MD5(key) }, nil) validGranteeKey3 := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, fmt.Sprintf("%s/%s/%s", "role1", "obj3", "*.obj_name3")) - kvmock.EXPECT().Load(validGranteeKey3).Call. - Return(func(key string) string { return crypto.MD5(key) }, nil) + kvmock.EXPECT().Load(mock.Anything, validGranteeKey3).Call. + Return(func(ctx context.Context, key string) string { return crypto.MD5(key) }, nil) - kvmock.EXPECT().Load(mock.Anything).Call. + kvmock.EXPECT().Load(mock.Anything, mock.Anything).Call. Return("", errors.New("mock Load error")) invalidRoleKey := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, invalidRole) - kvmock.EXPECT().LoadWithPrefix(invalidRoleKey).Call.Return(nil, nil, errors.New("mock loadWithPrefix error")) - kvmock.EXPECT().LoadWithPrefix(mock.Anything).Call.Return( - func(key string) []string { + kvmock.EXPECT().LoadWithPrefix(mock.Anything, invalidRoleKey).Call.Return(nil, nil, errors.New("mock loadWithPrefix error")) + kvmock.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Call.Return( + func(ctx context.Context, key string) []string { // Mock kv_catalog.go:ListGrant:L871 if strings.Contains(key, GranteeIDPrefix) { return []string{ @@ -2495,7 +2502,7 @@ func TestRBAC_Grant(t *testing.T) { fmt.Sprintf("%s/%s", key, "obj3/*.obj_name3"), } }, - func(key string) []string { + func(ctx context.Context, key string) []string { if strings.Contains(key, GranteeIDPrefix) { return []string{"user1", "user2"} } @@ -2582,8 +2589,8 @@ func TestRBAC_Grant(t *testing.T) { secondLoadWithPrefixReturn atomic.Bool ) - kvmock.EXPECT().LoadWithPrefix(mock.Anything).Call.Return( - func(key string) []string { + kvmock.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Call.Return( + func(ctx context.Context, key string) []string { contains := strings.Contains(key, GranteeIDPrefix) if contains { if secondLoadWithPrefixReturn.Load() { @@ -2606,7 +2613,7 @@ func TestRBAC_Grant(t *testing.T) { return nil }, - func(key string) []string { + func(ctx context.Context, key string) []string { if firstLoadWithPrefixReturn.Load() { return []string{ crypto.MD5(fmt.Sprintf("%s/%s", key, "obj1/obj_name1")), @@ -2616,7 +2623,7 @@ func TestRBAC_Grant(t *testing.T) { return nil }, - func(key string) error { + func(ctx context.Context, key string) error { contains := strings.Contains(key, GranteeIDPrefix) if contains { if secondLoadWithPrefixReturn.Load() { @@ -2680,7 +2687,7 @@ func TestRBAC_Backup(t *testing.T) { Params.EtcdCfg.EtcdTLSMinVersion.GetValue()) rootPath := "/test/rbac" metaKV := etcdkv.NewEtcdKV(etcdCli, rootPath) - defer metaKV.RemoveWithPrefix("") + defer metaKV.RemoveWithPrefix(context.TODO(), "") defer metaKV.Close() c := &Catalog{Txn: metaKV} @@ -2733,7 +2740,7 @@ func TestRBAC_Restore(t *testing.T) { Params.EtcdCfg.EtcdTLSMinVersion.GetValue()) rootPath := "/test/rbac" metaKV := etcdkv.NewEtcdKV(etcdCli, rootPath) - defer metaKV.RemoveWithPrefix("") + defer metaKV.RemoveWithPrefix(context.TODO(), "") defer metaKV.Close() c := &Catalog{Txn: metaKV} @@ -2901,8 +2908,8 @@ func TestRBAC_PrivilegeGroup(t *testing.T) { kvmock = mocks.NewTxnKV(t) c = &Catalog{Txn: kvmock} ) - kvmock.EXPECT().Load(key1).Return(string(v1), nil) - kvmock.EXPECT().Load(key2).Return("", merr.ErrIoKeyNotFound) + kvmock.EXPECT().Load(mock.Anything, key1).Return(string(v1), nil) + kvmock.EXPECT().Load(mock.Anything, key2).Return("", merr.ErrIoKeyNotFound) tests := []struct { description string @@ -2932,8 +2939,8 @@ func TestRBAC_PrivilegeGroup(t *testing.T) { c = &Catalog{Txn: kvmock} ) - kvmock.EXPECT().Remove(key1).Return(nil) - kvmock.EXPECT().Remove(key2).Return(errors.New("Mock remove failure")) + kvmock.EXPECT().Remove(mock.Anything, key1).Return(nil) + kvmock.EXPECT().Remove(mock.Anything, key2).Return(errors.New("Mock remove failure")) tests := []struct { description string @@ -2962,8 +2969,8 @@ func TestRBAC_PrivilegeGroup(t *testing.T) { c = &Catalog{Txn: kvmock} ) - kvmock.EXPECT().Save(key1, mock.Anything).Return(nil) - kvmock.EXPECT().Save(key2, mock.Anything).Return(nil) + kvmock.EXPECT().Save(mock.Anything, key1, mock.Anything).Return(nil) + kvmock.EXPECT().Save(mock.Anything, key2, mock.Anything).Return(nil) tests := []struct { description string @@ -2992,7 +2999,7 @@ func TestRBAC_PrivilegeGroup(t *testing.T) { c = &Catalog{Txn: kvmock} ) - kvmock.EXPECT().LoadWithPrefix(PrivilegeGroupPrefix).Return( + kvmock.EXPECT().LoadWithPrefix(mock.Anything, PrivilegeGroupPrefix).Return( []string{key1, key2}, []string{string(v1), string(v2)}, nil, @@ -3022,7 +3029,7 @@ func TestCatalog_AlterDatabase(t *testing.T) { c := &Catalog{Snapshot: kvmock} db := model.NewDatabase(1, "db", pb.DatabaseState_DatabaseCreated, nil) - kvmock.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(nil) + kvmock.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil) ctx := context.Background() // test alter database success @@ -3039,7 +3046,7 @@ func TestCatalog_AlterDatabase(t *testing.T) { // test alter database fail mockErr := errors.New("access kv store error") kvmock.ExpectedCalls = nil - kvmock.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(mockErr) + kvmock.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(mockErr) err = c.AlterDatabase(ctx, newDB, typeutil.ZeroTimestamp) assert.ErrorIs(t, err, mockErr) } @@ -3047,11 +3054,11 @@ func TestCatalog_AlterDatabase(t *testing.T) { func TestCatalog_listFunctionError(t *testing.T) { mockSnapshot := newMockSnapshot(t) kc := &Catalog{Snapshot: mockSnapshot} - mockSnapshot.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Return(nil, nil, fmt.Errorf("mock error")) - _, err := kc.listFunctions(1, 1) + mockSnapshot.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return(nil, nil, fmt.Errorf("mock error")) + _, err := kc.listFunctions(context.TODO(), 1, 1) assert.Error(t, err) - mockSnapshot.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Return([]string{"test-key"}, []string{"invalid bytes"}, nil) - _, err = kc.listFunctions(1, 1) + mockSnapshot.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return([]string{"test-key"}, []string{"invalid bytes"}, nil) + _, err = kc.listFunctions(context.TODO(), 1, 1) assert.Error(t, err) } diff --git a/internal/metastore/kv/rootcoord/suffix_snapshot.go b/internal/metastore/kv/rootcoord/suffix_snapshot.go index 5a625d1347450..0457f631fcace 100644 --- a/internal/metastore/kv/rootcoord/suffix_snapshot.go +++ b/internal/metastore/kv/rootcoord/suffix_snapshot.go @@ -18,6 +18,7 @@ package rootcoord import ( "bytes" + "context" "fmt" "path" "sort" @@ -119,7 +120,7 @@ func NewSuffixSnapshot(metaKV kv.MetaKv, sep, root, snapshot string) (*SuffixSna rootLen: rootLen, closeGC: make(chan struct{}, 1), } - go ss.startBackgroundGC() + go ss.startBackgroundGC(context.TODO()) return ss, nil } @@ -196,10 +197,10 @@ func (ss *SuffixSnapshot) isTSOfKey(key string, groupKey string) (typeutil.Times // checkKeyTS checks provided key's latest ts is before provided ts // lock is needed -func (ss *SuffixSnapshot) checkKeyTS(key string, ts typeutil.Timestamp) (bool, error) { +func (ss *SuffixSnapshot) checkKeyTS(ctx context.Context, key string, ts typeutil.Timestamp) (bool, error) { latest, has := ss.lastestTS[key] if !has { - err := ss.loadLatestTS(key) + err := ss.loadLatestTS(ctx, key) if err != nil { return false, err } @@ -209,9 +210,9 @@ func (ss *SuffixSnapshot) checkKeyTS(key string, ts typeutil.Timestamp) (bool, e } // loadLatestTS load the loatest ts for specified key -func (ss *SuffixSnapshot) loadLatestTS(key string) error { +func (ss *SuffixSnapshot) loadLatestTS(ctx context.Context, key string) error { prefix := ss.composeSnapshotPrefix(key) - keys, _, err := ss.MetaKv.LoadWithPrefix(prefix) + keys, _, err := ss.MetaKv.LoadWithPrefix(ctx, prefix) if err != nil { log.Warn("SuffixSnapshot MetaKv LoadWithPrefix failed", zap.String("key", key), zap.Error(err)) @@ -273,11 +274,11 @@ func binarySearchRecords(records []tsv, ts typeutil.Timestamp) (string, bool) { // if ts == 0, SuffixSnapshot works as a MetaKv // otherwise, SuffixSnapshot will store a ts-key as "key[sep]ts"-value pair in snapshot path // and for acceleration store original key-value if ts is the latest -func (ss *SuffixSnapshot) Save(key string, value string, ts typeutil.Timestamp) error { +func (ss *SuffixSnapshot) Save(ctx context.Context, key string, value string, ts typeutil.Timestamp) error { // if ts == 0, act like MetaKv // will not update lastestTs since ts not not valid if ts == 0 { - return ss.MetaKv.Save(key, value) + return ss.MetaKv.Save(ctx, key, value) } ss.Lock() @@ -287,12 +288,12 @@ func (ss *SuffixSnapshot) Save(key string, value string, ts typeutil.Timestamp) // provided key value is latest // stores both tsKey and original key - after, err := ss.checkKeyTS(key, ts) + after, err := ss.checkKeyTS(ctx, key, ts) if err != nil { return err } if after { - err := ss.MetaKv.MultiSave(map[string]string{ + err := ss.MetaKv.MultiSave(ctx, map[string]string{ key: value, tsKey: value, }) @@ -304,14 +305,14 @@ func (ss *SuffixSnapshot) Save(key string, value string, ts typeutil.Timestamp) } // modifying history key, just save tskey-value - return ss.MetaKv.Save(tsKey, value) + return ss.MetaKv.Save(ctx, tsKey, value) } -func (ss *SuffixSnapshot) Load(key string, ts typeutil.Timestamp) (string, error) { +func (ss *SuffixSnapshot) Load(ctx context.Context, key string, ts typeutil.Timestamp) (string, error) { // if ts == 0 or typeutil.MaxTimestamp, load latest by definition // and with acceleration logic, just do load key will do if ts == 0 || ts == typeutil.MaxTimestamp { - value, err := ss.MetaKv.Load(key) + value, err := ss.MetaKv.Load(ctx, key) if ss.isTombstone(value) { return "", errors.New("no value found") } @@ -319,7 +320,7 @@ func (ss *SuffixSnapshot) Load(key string, ts typeutil.Timestamp) (string, error } ss.Lock() - after, err := ss.checkKeyTS(key, ts) + after, err := ss.checkKeyTS(ctx, key, ts) ss.Unlock() ss.RLock() @@ -329,7 +330,7 @@ func (ss *SuffixSnapshot) Load(key string, ts typeutil.Timestamp) (string, error return "", err } if after { - value, err := ss.MetaKv.Load(key) + value, err := ss.MetaKv.Load(ctx, key) if ss.isTombstone(value) { return "", errors.New("no value found") } @@ -338,7 +339,7 @@ func (ss *SuffixSnapshot) Load(key string, ts typeutil.Timestamp) (string, error // before ts, do time travel // 1. load all tsKey with key/ prefix - keys, values, err := ss.MetaKv.LoadWithPrefix(ss.composeSnapshotPrefix(key)) + keys, values, err := ss.MetaKv.LoadWithPrefix(ctx, ss.composeSnapshotPrefix(key)) if err != nil { log.Warn("prefixSnapshot MetaKv LoadWithPrefix failed", zap.String("key", key), zap.Error(err)) return "", err @@ -380,23 +381,23 @@ func (ss *SuffixSnapshot) Load(key string, ts typeutil.Timestamp) (string, error // MultiSave save multiple kvs // if ts == 0, act like MetaKv // each key-value will be treated using same logic like Save -func (ss *SuffixSnapshot) MultiSave(kvs map[string]string, ts typeutil.Timestamp) error { +func (ss *SuffixSnapshot) MultiSave(ctx context.Context, kvs map[string]string, ts typeutil.Timestamp) error { // if ts == 0, act like MetaKv if ts == 0 { - return ss.MetaKv.MultiSave(kvs) + return ss.MetaKv.MultiSave(ctx, kvs) } ss.Lock() defer ss.Unlock() var err error // process each key, checks whether is the latest - execute, updateList, err := ss.generateSaveExecute(kvs, ts) + execute, updateList, err := ss.generateSaveExecute(ctx, kvs, ts) if err != nil { return err } // multi save execute map; if succeeds, update ts in the update list - err = ss.MetaKv.MultiSave(execute) + err = ss.MetaKv.MultiSave(ctx, execute) if err == nil { for _, key := range updateList { ss.lastestTS[key] = ts @@ -407,7 +408,7 @@ func (ss *SuffixSnapshot) MultiSave(kvs map[string]string, ts typeutil.Timestamp // generateSaveExecute examine each key is the after the corresponding latest // returns calculated execute map and update ts list -func (ss *SuffixSnapshot) generateSaveExecute(kvs map[string]string, ts typeutil.Timestamp) (map[string]string, []string, error) { +func (ss *SuffixSnapshot) generateSaveExecute(ctx context.Context, kvs map[string]string, ts typeutil.Timestamp) (map[string]string, []string, error) { var after bool var err error execute := make(map[string]string) @@ -416,7 +417,7 @@ func (ss *SuffixSnapshot) generateSaveExecute(kvs map[string]string, ts typeutil tsKey := ss.composeTSKey(key, ts) // provided key value is latest // stores both tsKey and original key - after, err = ss.checkKeyTS(key, ts) + after, err = ss.checkKeyTS(ctx, key, ts) if err != nil { return nil, nil, err } @@ -432,10 +433,10 @@ func (ss *SuffixSnapshot) generateSaveExecute(kvs map[string]string, ts typeutil } // LoadWithPrefix load keys with provided prefix and returns value in the ts -func (ss *SuffixSnapshot) LoadWithPrefix(key string, ts typeutil.Timestamp) ([]string, []string, error) { +func (ss *SuffixSnapshot) LoadWithPrefix(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) { // ts 0 case shall be treated as fetch latest/current value if ts == 0 || ts == typeutil.MaxTimestamp { - keys, values, err := ss.MetaKv.LoadWithPrefix(key) + keys, values, err := ss.MetaKv.LoadWithPrefix(ctx, key) fks := keys[:0] // make([]string, 0, len(keys)) fvs := values[:0] // make([]string, 0, len(values)) // hide rootPrefix from return value @@ -469,7 +470,7 @@ func (ss *SuffixSnapshot) LoadWithPrefix(key string, ts typeutil.Timestamp) ([]s resultValues = append(resultValues, value) } - err := ss.MetaKv.WalkWithPrefix(prefix, PaginationSize, func(k []byte, v []byte) error { + err := ss.MetaKv.WalkWithPrefix(ctx, prefix, PaginationSize, func(k []byte, v []byte) error { sKey := string(k) sValue := string(v) @@ -506,24 +507,24 @@ func (ss *SuffixSnapshot) LoadWithPrefix(key string, ts typeutil.Timestamp) ([]s // MultiSaveAndRemove save muiltple kvs and remove as well // if ts == 0, act like MetaKv // each key-value will be treated in same logic like Save -func (ss *SuffixSnapshot) MultiSaveAndRemove(saves map[string]string, removals []string, ts typeutil.Timestamp) error { +func (ss *SuffixSnapshot) MultiSaveAndRemove(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error { // if ts == 0, act like MetaKv if ts == 0 { - return ss.MetaKv.MultiSaveAndRemove(saves, removals) + return ss.MetaKv.MultiSaveAndRemove(ctx, saves, removals) } ss.Lock() defer ss.Unlock() var err error // process each key, checks whether is the latest - execute, updateList, err := ss.generateSaveExecute(saves, ts) + execute, updateList, err := ss.generateSaveExecute(ctx, saves, ts) if err != nil { return err } // load each removal, change execution to adding tombstones for _, removal := range removals { - value, err := ss.MetaKv.Load(removal) + value, err := ss.MetaKv.Load(ctx, removal) if err != nil { log.Warn("SuffixSnapshot MetaKv Load failed", zap.String("key", removal), zap.Error(err)) if errors.Is(err, merr.ErrIoKeyNotFound) { @@ -541,7 +542,7 @@ func (ss *SuffixSnapshot) MultiSaveAndRemove(saves map[string]string, removals [ } // multi save execute map; if succeeds, update ts in the update list - err = ss.MetaKv.MultiSave(execute) + err = ss.MetaKv.MultiSave(ctx, execute) if err == nil { for _, key := range updateList { ss.lastestTS[key] = ts @@ -553,24 +554,24 @@ func (ss *SuffixSnapshot) MultiSaveAndRemove(saves map[string]string, removals [ // MultiSaveAndRemoveWithPrefix save muiltple kvs and remove as well // if ts == 0, act like MetaKv // each key-value will be treated in same logic like Save -func (ss *SuffixSnapshot) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals []string, ts typeutil.Timestamp) error { +func (ss *SuffixSnapshot) MultiSaveAndRemoveWithPrefix(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error { // if ts == 0, act like MetaKv if ts == 0 { - return ss.MetaKv.MultiSaveAndRemoveWithPrefix(saves, removals) + return ss.MetaKv.MultiSaveAndRemoveWithPrefix(ctx, saves, removals) } ss.Lock() defer ss.Unlock() var err error // process each key, checks whether is the latest - execute, updateList, err := ss.generateSaveExecute(saves, ts) + execute, updateList, err := ss.generateSaveExecute(ctx, saves, ts) if err != nil { return err } // load each removal, change execution to adding tombstones for _, removal := range removals { - keys, values, err := ss.MetaKv.LoadWithPrefix(removal) + keys, values, err := ss.MetaKv.LoadWithPrefix(ctx, removal) if err != nil { log.Warn("SuffixSnapshot MetaKv LoadwithPrefix failed", zap.String("key", removal), zap.Error(err)) return err @@ -589,7 +590,7 @@ func (ss *SuffixSnapshot) MultiSaveAndRemoveWithPrefix(saves map[string]string, } // multi save execute map; if succeeds, update ts in the update list - err = ss.MetaKv.MultiSave(execute) + err = ss.MetaKv.MultiSave(ctx, execute) if err == nil { for _, key := range updateList { ss.lastestTS[key] = ts @@ -603,7 +604,7 @@ func (ss *SuffixSnapshot) Close() { } // startBackgroundGC the data will clean up if key ts!=0 and expired -func (ss *SuffixSnapshot) startBackgroundGC() { +func (ss *SuffixSnapshot) startBackgroundGC(ctx context.Context) { log.Debug("suffix snapshot GC goroutine start!") ticker := time.NewTicker(60 * time.Minute) defer ticker.Stop() @@ -614,7 +615,7 @@ func (ss *SuffixSnapshot) startBackgroundGC() { log.Warn("quit suffix snapshot GC goroutine!") return case now := <-ticker.C: - err := ss.removeExpiredKvs(now) + err := ss.removeExpiredKvs(ctx, now) if err != nil { log.Warn("remove expired data fail during GC", zap.Error(err)) } @@ -635,7 +636,7 @@ func (ss *SuffixSnapshot) getOriginalKey(snapshotKey string) (string, error) { return prefix[ss.snapshotLen:], nil } -func (ss *SuffixSnapshot) batchRemoveExpiredKvs(keyGroup []string, originalKey string, includeOriginalKey bool) error { +func (ss *SuffixSnapshot) batchRemoveExpiredKvs(ctx context.Context, keyGroup []string, originalKey string, includeOriginalKey bool) error { if includeOriginalKey { keyGroup = append(keyGroup, originalKey) } @@ -643,7 +644,7 @@ func (ss *SuffixSnapshot) batchRemoveExpiredKvs(keyGroup []string, originalKey s // to protect txn finished with ascend order, reverse the latest kv with tombstone to tail of array sort.Strings(keyGroup) removeFn := func(partialKeys []string) error { - return ss.MetaKv.MultiRemove(partialKeys) + return ss.MetaKv.MultiRemove(ctx, partialKeys) } return etcd.RemoveByBatchWithLimit(keyGroup, util.MaxEtcdTxnNum, removeFn) } @@ -651,7 +652,7 @@ func (ss *SuffixSnapshot) batchRemoveExpiredKvs(keyGroup []string, originalKey s // removeExpiredKvs removes expired key-value pairs from the snapshot // It walks through all keys with the snapshot prefix, groups them by original key, // and removes expired versions or all versions if the original key has been deleted -func (ss *SuffixSnapshot) removeExpiredKvs(now time.Time) error { +func (ss *SuffixSnapshot) removeExpiredKvs(ctx context.Context, now time.Time) error { ttlTime := paramtable.Get().ServiceParam.MetaStoreCfg.SnapshotTTLSeconds.GetAsDuration(time.Second) reserveTime := paramtable.Get().ServiceParam.MetaStoreCfg.SnapshotReserveTimeSeconds.GetAsDuration(time.Second) @@ -667,7 +668,7 @@ func (ss *SuffixSnapshot) removeExpiredKvs(now time.Time) error { } if ss.isTombstone(latestOriginValue) { // If deleted, remove all versions including the original key - return ss.batchRemoveExpiredKvs(candidateExpiredKeys, curOriginalKey, totalVersions == len(candidateExpiredKeys)) + return ss.batchRemoveExpiredKvs(ctx, candidateExpiredKeys, curOriginalKey, totalVersions == len(candidateExpiredKeys)) } // If not deleted, check for expired versions @@ -680,13 +681,13 @@ func (ss *SuffixSnapshot) removeExpiredKvs(now time.Time) error { } } if len(expiredKeys) > 0 { - return ss.batchRemoveExpiredKvs(expiredKeys, curOriginalKey, false) + return ss.batchRemoveExpiredKvs(ctx, expiredKeys, curOriginalKey, false) } return nil } // Walk through all keys with the snapshot prefix - err := ss.MetaKv.WalkWithPrefix(ss.snapshotPrefix, PaginationSize, func(k []byte, v []byte) error { + err := ss.MetaKv.WalkWithPrefix(ctx, ss.snapshotPrefix, PaginationSize, func(k []byte, v []byte) error { key := ss.hideRootPrefix(string(k)) ts, ok := ss.isTSKey(key) if !ok { diff --git a/internal/metastore/kv/rootcoord/suffix_snapshot_test.go b/internal/metastore/kv/rootcoord/suffix_snapshot_test.go index 326716919ff85..ec33cc364d508 100644 --- a/internal/metastore/kv/rootcoord/suffix_snapshot_test.go +++ b/internal/metastore/kv/rootcoord/suffix_snapshot_test.go @@ -17,6 +17,7 @@ package rootcoord import ( + "context" "fmt" "math/rand" "os" @@ -304,27 +305,27 @@ func Test_SuffixSnapshotLoad(t *testing.T) { for i := 0; i < 20; i++ { vtso = typeutil.Timestamp(100 + i*5) ts := ftso() - err = ss.Save("key", fmt.Sprintf("value-%d", i), ts) + err = ss.Save(context.TODO(), "key", fmt.Sprintf("value-%d", i), ts) assert.NoError(t, err) assert.Equal(t, vtso, ts) } for i := 0; i < 20; i++ { - val, err := ss.Load("key", typeutil.Timestamp(100+i*5+2)) + val, err := ss.Load(context.TODO(), "key", typeutil.Timestamp(100+i*5+2)) t.Log("ts:", typeutil.Timestamp(100+i*5+2), i, val) assert.NoError(t, err) assert.Equal(t, fmt.Sprintf("value-%d", i), val) } - val, err := ss.Load("key", 0) + val, err := ss.Load(context.TODO(), "key", 0) assert.NoError(t, err) assert.Equal(t, "value-19", val) for i := 0; i < 20; i++ { - val, err := ss.Load("key", typeutil.Timestamp(100+i*5+2)) + val, err := ss.Load(context.TODO(), "key", typeutil.Timestamp(100+i*5+2)) assert.NoError(t, err) assert.Equal(t, val, fmt.Sprintf("value-%d", i)) } - ss.RemoveWithPrefix("") + ss.RemoveWithPrefix(context.TODO(), "") } func Test_SuffixSnapshotMultiSave(t *testing.T) { @@ -360,12 +361,12 @@ func Test_SuffixSnapshotMultiSave(t *testing.T) { saves := map[string]string{"k1": fmt.Sprintf("v1-%d", i), "k2": fmt.Sprintf("v2-%d", i)} vtso = typeutil.Timestamp(100 + i*5) ts := ftso() - err = ss.MultiSave(saves, ts) + err = ss.MultiSave(context.TODO(), saves, ts) assert.NoError(t, err) assert.Equal(t, vtso, ts) } for i := 0; i < 20; i++ { - keys, vals, err := ss.LoadWithPrefix("k", typeutil.Timestamp(100+i*5+2)) + keys, vals, err := ss.LoadWithPrefix(context.TODO(), "k", typeutil.Timestamp(100+i*5+2)) t.Log(i, keys, vals) assert.NoError(t, err) assert.Equal(t, len(keys), len(vals)) @@ -375,7 +376,7 @@ func Test_SuffixSnapshotMultiSave(t *testing.T) { assert.Equal(t, vals[0], fmt.Sprintf("v1-%d", i)) assert.Equal(t, vals[1], fmt.Sprintf("v2-%d", i)) } - keys, vals, err := ss.LoadWithPrefix("k", 0) + keys, vals, err := ss.LoadWithPrefix(context.TODO(), "k", 0) assert.NoError(t, err) assert.Equal(t, len(keys), len(vals)) assert.Equal(t, len(keys), 2) @@ -385,7 +386,7 @@ func Test_SuffixSnapshotMultiSave(t *testing.T) { assert.Equal(t, vals[1], "v2-19") for i := 0; i < 20; i++ { - keys, vals, err := ss.LoadWithPrefix("k", typeutil.Timestamp(100+i*5+2)) + keys, vals, err := ss.LoadWithPrefix(context.TODO(), "k", typeutil.Timestamp(100+i*5+2)) assert.NoError(t, err) assert.Equal(t, len(keys), len(vals)) assert.Equal(t, len(keys), 2) @@ -393,9 +394,9 @@ func Test_SuffixSnapshotMultiSave(t *testing.T) { assert.ElementsMatch(t, vals, []string{fmt.Sprintf("v1-%d", i), fmt.Sprintf("v2-%d", i)}) } // mix non ts k-v - err = ss.Save("kextra", "extra-value", 0) + err = ss.Save(context.TODO(), "kextra", "extra-value", 0) assert.NoError(t, err) - keys, vals, err = ss.LoadWithPrefix("k", typeutil.Timestamp(300)) + keys, vals, err = ss.LoadWithPrefix(context.TODO(), "k", typeutil.Timestamp(300)) assert.NoError(t, err) assert.Equal(t, len(keys), len(vals)) assert.Equal(t, len(keys), 2) @@ -403,7 +404,7 @@ func Test_SuffixSnapshotMultiSave(t *testing.T) { assert.ElementsMatch(t, vals, []string{"v1-19", "v2-19"}) // clean up - ss.RemoveWithPrefix("") + ss.RemoveWithPrefix(context.TODO(), "") } func Test_SuffixSnapshotRemoveExpiredKvs(t *testing.T) { @@ -433,12 +434,12 @@ func Test_SuffixSnapshotRemoveExpiredKvs(t *testing.T) { defer ss.Close() saveFn := func(key, value string, ts typeutil.Timestamp) { - err = ss.Save(key, value, ts) + err = ss.Save(context.TODO(), key, value, ts) assert.NoError(t, err) } multiSaveFn := func(kvs map[string]string, ts typeutil.Timestamp) { - err = ss.MultiSave(kvs, ts) + err = ss.MultiSave(context.TODO(), kvs, ts) assert.NoError(t, err) } @@ -477,7 +478,7 @@ func Test_SuffixSnapshotRemoveExpiredKvs(t *testing.T) { countPrefix := func(prefix string) int { cnt := 0 - err := etcdkv.WalkWithPrefix("", 10, func(key []byte, value []byte) error { + err := etcdkv.WalkWithPrefix(context.TODO(), "", 10, func(key []byte, value []byte) error { cnt++ return nil }) @@ -495,14 +496,14 @@ func Test_SuffixSnapshotRemoveExpiredKvs(t *testing.T) { cnt := countPrefix(prefix) assert.Equal(t, keyCnt*keyVersion+keyCnt, cnt) - err = ss.removeExpiredKvs(now) + err = ss.removeExpiredKvs(context.TODO(), now) assert.NoError(t, err) cnt = countPrefix(prefix) assert.Equal(t, keyCnt*keyVersion+keyCnt-(expiredKCnt*keyVersion+expiredKCnt), cnt) // clean all data - err := etcdkv.RemoveWithPrefix("") + err := etcdkv.RemoveWithPrefix(context.TODO(), "") assert.NoError(t, err) }) @@ -536,14 +537,14 @@ func Test_SuffixSnapshotRemoveExpiredKvs(t *testing.T) { assert.Equal(t, 12, cnt) // err = ss.removeExpiredKvs(now, time.Duration(50)*time.Millisecond) - err = ss.removeExpiredKvs(now) + err = ss.removeExpiredKvs(context.TODO(), now) assert.NoError(t, err) cnt = countPrefix(prefix) assert.Equal(t, 4, cnt) // clean all data - err := etcdkv.RemoveWithPrefix("") + err := etcdkv.RemoveWithPrefix(context.TODO(), "") assert.NoError(t, err) }) @@ -577,31 +578,31 @@ func Test_SuffixSnapshotRemoveExpiredKvs(t *testing.T) { assert.Equal(t, 12, cnt) // err = ss.removeExpiredKvs(now, time.Duration(50)*time.Millisecond) - err = ss.removeExpiredKvs(now) + err = ss.removeExpiredKvs(context.TODO(), now) assert.NoError(t, err) cnt = countPrefix(prefix) assert.Equal(t, 1, cnt) // clean all data - err := etcdkv.RemoveWithPrefix("") + err := etcdkv.RemoveWithPrefix(context.TODO(), "") assert.NoError(t, err) }) t.Run("parse ts fail", func(t *testing.T) { prefix := fmt.Sprintf("prefix%d", rand.Int()) key := fmt.Sprintf("%s-%s", prefix, "ts_error-ts") - err = etcdkv.Save(ss.composeSnapshotPrefix(key), "") + err = etcdkv.Save(context.TODO(), ss.composeSnapshotPrefix(key), "") assert.NoError(t, err) - err = ss.removeExpiredKvs(now) + err = ss.removeExpiredKvs(context.TODO(), now) assert.NoError(t, err) cnt := countPrefix(prefix) assert.Equal(t, 1, cnt) // clean all data - err := etcdkv.RemoveWithPrefix("") + err := etcdkv.RemoveWithPrefix(context.TODO(), "") assert.NoError(t, err) }) @@ -610,14 +611,14 @@ func Test_SuffixSnapshotRemoveExpiredKvs(t *testing.T) { rootPath := "root/" kv := mocks.NewMetaKv(t) kv.EXPECT(). - WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything). + WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything). Return(errors.New("error")) ss, err := NewSuffixSnapshot(kv, sep, rootPath, snapshotPrefix) assert.NotNil(t, ss) assert.NoError(t, err) - err = ss.removeExpiredKvs(time.Now()) + err = ss.removeExpiredKvs(context.TODO(), time.Now()) assert.Error(t, err) }) } @@ -656,7 +657,7 @@ func Test_SuffixSnapshotMultiSaveAndRemoveWithPrefix(t *testing.T) { for i := 0; i < 20; i++ { vtso = typeutil.Timestamp(100 + i*5) ts := ftso() - err = ss.Save(fmt.Sprintf("kd-%04d", i), fmt.Sprintf("value-%d", i), ts) + err = ss.Save(context.TODO(), fmt.Sprintf("kd-%04d", i), fmt.Sprintf("value-%d", i), ts) assert.NoError(t, err) assert.Equal(t, vtso, ts) } @@ -665,53 +666,53 @@ func Test_SuffixSnapshotMultiSaveAndRemoveWithPrefix(t *testing.T) { dm := []string{fmt.Sprintf("kd-%04d", i-20)} vtso = typeutil.Timestamp(100 + i*5) ts := ftso() - err = ss.MultiSaveAndRemoveWithPrefix(sm, dm, ts) + err = ss.MultiSaveAndRemoveWithPrefix(context.TODO(), sm, dm, ts) assert.NoError(t, err) assert.Equal(t, vtso, ts) } for i := 0; i < 20; i++ { - val, err := ss.Load(fmt.Sprintf("kd-%04d", i), typeutil.Timestamp(100+i*5+2)) + val, err := ss.Load(context.TODO(), fmt.Sprintf("kd-%04d", i), typeutil.Timestamp(100+i*5+2)) assert.NoError(t, err) assert.Equal(t, fmt.Sprintf("value-%d", i), val) - _, vals, err := ss.LoadWithPrefix("kd-", typeutil.Timestamp(100+i*5+2)) + _, vals, err := ss.LoadWithPrefix(context.TODO(), "kd-", typeutil.Timestamp(100+i*5+2)) assert.NoError(t, err) assert.Equal(t, i+1, len(vals)) } for i := 20; i < 40; i++ { - val, err := ss.Load("ks", typeutil.Timestamp(100+i*5+2)) + val, err := ss.Load(context.TODO(), "ks", typeutil.Timestamp(100+i*5+2)) assert.NoError(t, err) assert.Equal(t, fmt.Sprintf("value-%d", i), val) - _, vals, err := ss.LoadWithPrefix("kd-", typeutil.Timestamp(100+i*5+2)) + _, vals, err := ss.LoadWithPrefix(context.TODO(), "kd-", typeutil.Timestamp(100+i*5+2)) assert.NoError(t, err) assert.Equal(t, 39-i, len(vals)) } for i := 0; i < 20; i++ { - val, err := ss.Load(fmt.Sprintf("kd-%04d", i), typeutil.Timestamp(100+i*5+2)) + val, err := ss.Load(context.TODO(), fmt.Sprintf("kd-%04d", i), typeutil.Timestamp(100+i*5+2)) assert.NoError(t, err) assert.Equal(t, fmt.Sprintf("value-%d", i), val) - _, vals, err := ss.LoadWithPrefix("kd-", typeutil.Timestamp(100+i*5+2)) + _, vals, err := ss.LoadWithPrefix(context.TODO(), "kd-", typeutil.Timestamp(100+i*5+2)) assert.NoError(t, err) assert.Equal(t, i+1, len(vals)) } for i := 20; i < 40; i++ { - val, err := ss.Load("ks", typeutil.Timestamp(100+i*5+2)) + val, err := ss.Load(context.TODO(), "ks", typeutil.Timestamp(100+i*5+2)) assert.NoError(t, err) assert.Equal(t, fmt.Sprintf("value-%d", i), val) - _, vals, err := ss.LoadWithPrefix("kd-", typeutil.Timestamp(100+i*5+2)) + _, vals, err := ss.LoadWithPrefix(context.TODO(), "kd-", typeutil.Timestamp(100+i*5+2)) assert.NoError(t, err) assert.Equal(t, 39-i, len(vals)) } // try to load - _, err = ss.Load("kd-0000", 500) + _, err = ss.Load(context.TODO(), "kd-0000", 500) assert.Error(t, err) - _, err = ss.Load("kd-0000", 0) + _, err = ss.Load(context.TODO(), "kd-0000", 0) assert.Error(t, err) - _, err = ss.Load("kd-0000", 1) + _, err = ss.Load(context.TODO(), "kd-0000", 1) assert.Error(t, err) // cleanup - ss.MultiSaveAndRemoveWithPrefix(map[string]string{}, []string{""}, 0) + ss.MultiSaveAndRemoveWithPrefix(context.TODO(), map[string]string{}, []string{""}, 0) } func Test_SuffixSnapshotMultiSaveAndRemove(t *testing.T) { @@ -748,7 +749,7 @@ func Test_SuffixSnapshotMultiSaveAndRemove(t *testing.T) { for i := 0; i < 20; i++ { vtso = typeutil.Timestamp(100 + i*5) ts := ftso() - err = ss.Save(fmt.Sprintf("kd-%04d", i), fmt.Sprintf("value-%d", i), ts) + err = ss.Save(context.TODO(), fmt.Sprintf("kd-%04d", i), fmt.Sprintf("value-%d", i), ts) assert.NoError(t, err) assert.Equal(t, vtso, ts) } @@ -757,37 +758,37 @@ func Test_SuffixSnapshotMultiSaveAndRemove(t *testing.T) { dm := []string{fmt.Sprintf("kd-%04d", i-20)} vtso = typeutil.Timestamp(100 + i*5) ts := ftso() - err = ss.MultiSaveAndRemove(sm, dm, ts) + err = ss.MultiSaveAndRemove(context.TODO(), sm, dm, ts) assert.NoError(t, err) assert.Equal(t, vtso, ts) } for i := 0; i < 20; i++ { - val, err := ss.Load(fmt.Sprintf("kd-%04d", i), typeutil.Timestamp(100+i*5+2)) + val, err := ss.Load(context.TODO(), fmt.Sprintf("kd-%04d", i), typeutil.Timestamp(100+i*5+2)) assert.NoError(t, err) assert.Equal(t, fmt.Sprintf("value-%d", i), val) - _, vals, err := ss.LoadWithPrefix("kd-", typeutil.Timestamp(100+i*5+2)) + _, vals, err := ss.LoadWithPrefix(context.TODO(), "kd-", typeutil.Timestamp(100+i*5+2)) assert.NoError(t, err) assert.Equal(t, i+1, len(vals)) } for i := 20; i < 40; i++ { - val, err := ss.Load("ks", typeutil.Timestamp(100+i*5+2)) + val, err := ss.Load(context.TODO(), "ks", typeutil.Timestamp(100+i*5+2)) assert.NoError(t, err) assert.Equal(t, fmt.Sprintf("value-%d", i), val) - _, vals, err := ss.LoadWithPrefix("kd-", typeutil.Timestamp(100+i*5+2)) + _, vals, err := ss.LoadWithPrefix(context.TODO(), "kd-", typeutil.Timestamp(100+i*5+2)) assert.NoError(t, err) assert.Equal(t, 39-i, len(vals)) } // try to load - _, err = ss.Load("kd-0000", 500) + _, err = ss.Load(context.TODO(), "kd-0000", 500) assert.Error(t, err) - _, err = ss.Load("kd-0000", 0) + _, err = ss.Load(context.TODO(), "kd-0000", 0) assert.Error(t, err) - _, err = ss.Load("kd-0000", 1) + _, err = ss.Load(context.TODO(), "kd-0000", 1) assert.Error(t, err) // cleanup - ss.MultiSaveAndRemoveWithPrefix(map[string]string{}, []string{""}, 0) + ss.MultiSaveAndRemoveWithPrefix(context.TODO(), map[string]string{}, []string{""}, 0) } func TestSuffixSnapshot_LoadWithPrefix(t *testing.T) { @@ -819,16 +820,16 @@ func TestSuffixSnapshot_LoadWithPrefix(t *testing.T) { t.Run("parse ts fail", func(t *testing.T) { prefix := fmt.Sprintf("prefix%d", rand.Int()) key := fmt.Sprintf("%s-%s", prefix, "ts_error-ts") - err = etcdkv.Save(ss.composeSnapshotPrefix(key), "") + err = etcdkv.Save(context.TODO(), ss.composeSnapshotPrefix(key), "") assert.NoError(t, err) - keys, values, err := ss.LoadWithPrefix(prefix, 100) + keys, values, err := ss.LoadWithPrefix(context.TODO(), prefix, 100) assert.NoError(t, err) assert.Equal(t, 0, len(keys)) assert.Equal(t, 0, len(values)) // clean all data - err = etcdkv.RemoveWithPrefix("") + err = etcdkv.RemoveWithPrefix(context.TODO(), "") assert.NoError(t, err) }) @@ -837,14 +838,14 @@ func TestSuffixSnapshot_LoadWithPrefix(t *testing.T) { rootPath := "root/" kv := mocks.NewMetaKv(t) kv.EXPECT(). - WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything). + WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything). Return(errors.New("error")) ss, err := NewSuffixSnapshot(kv, sep, rootPath, snapshotPrefix) assert.NotNil(t, ss) assert.NoError(t, err) - keys, values, err := ss.LoadWithPrefix("t", 100) + keys, values, err := ss.LoadWithPrefix(context.TODO(), "t", 100) assert.Error(t, err) assert.Nil(t, keys) assert.Nil(t, values) diff --git a/internal/metastore/kv/streamingcoord/kv_catalog.go b/internal/metastore/kv/streamingcoord/kv_catalog.go index 390818a15c552..d3d804052e026 100644 --- a/internal/metastore/kv/streamingcoord/kv_catalog.go +++ b/internal/metastore/kv/streamingcoord/kv_catalog.go @@ -27,7 +27,7 @@ type catalog struct { // ListPChannels returns all pchannels func (c *catalog) ListPChannel(ctx context.Context) ([]*streamingpb.PChannelMeta, error) { - keys, values, err := c.metaKV.LoadWithPrefix(PChannelMeta) + keys, values, err := c.metaKV.LoadWithPrefix(ctx, PChannelMeta) if err != nil { return nil, err } @@ -56,7 +56,7 @@ func (c *catalog) SavePChannels(ctx context.Context, infos []*streamingpb.PChann kvs[key] = string(v) } return etcd.SaveByBatchWithLimit(kvs, util.MaxEtcdTxnNum, func(partialKvs map[string]string) error { - return c.metaKV.MultiSave(partialKvs) + return c.metaKV.MultiSave(ctx, partialKvs) }) } diff --git a/internal/metastore/kv/streamingcoord/kv_catalog_test.go b/internal/metastore/kv/streamingcoord/kv_catalog_test.go index 2544b9c852965..227ad0469bca3 100644 --- a/internal/metastore/kv/streamingcoord/kv_catalog_test.go +++ b/internal/metastore/kv/streamingcoord/kv_catalog_test.go @@ -16,7 +16,7 @@ func TestCatalog(t *testing.T) { kv := mock_kv.NewMockMetaKv(t) kvStorage := make(map[string]string) - kv.EXPECT().LoadWithPrefix(mock.Anything).RunAndReturn(func(s string) ([]string, []string, error) { + kv.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, s string) ([]string, []string, error) { keys := make([]string, 0, len(kvStorage)) vals := make([]string, 0, len(kvStorage)) for k, v := range kvStorage { @@ -25,7 +25,7 @@ func TestCatalog(t *testing.T) { } return keys, vals, nil }) - kv.EXPECT().MultiSave(mock.Anything).RunAndReturn(func(kvs map[string]string) error { + kv.EXPECT().MultiSave(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, kvs map[string]string) error { for k, v := range kvs { kvStorage[k] = v } @@ -54,13 +54,13 @@ func TestCatalog(t *testing.T) { assert.Len(t, metas, 2) // error path. - kv.EXPECT().LoadWithPrefix(mock.Anything).Unset() - kv.EXPECT().LoadWithPrefix(mock.Anything).Return(nil, nil, errors.New("load error")) + kv.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Unset() + kv.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Return(nil, nil, errors.New("load error")) metas, err = catalog.ListPChannel(context.Background()) assert.Error(t, err) assert.Nil(t, metas) - kv.EXPECT().MultiSave(mock.Anything).Unset() - kv.EXPECT().MultiSave(mock.Anything).Return(errors.New("save error")) + kv.EXPECT().MultiSave(mock.Anything, mock.Anything).Unset() + kv.EXPECT().MultiSave(mock.Anything, mock.Anything).Return(errors.New("save error")) assert.Error(t, err) } diff --git a/internal/metastore/kv/streamingnode/kv_catalog.go b/internal/metastore/kv/streamingnode/kv_catalog.go index a7b4ac803f521..2d279ab6ea6ca 100644 --- a/internal/metastore/kv/streamingnode/kv_catalog.go +++ b/internal/metastore/kv/streamingnode/kv_catalog.go @@ -29,7 +29,7 @@ type catalog struct { func (c *catalog) ListSegmentAssignment(ctx context.Context, pChannelName string) ([]*streamingpb.SegmentAssignmentMeta, error) { prefix := buildSegmentAssignmentMetaPath(pChannelName) - keys, values, err := c.metaKV.LoadWithPrefix(prefix) + keys, values, err := c.metaKV.LoadWithPrefix(ctx, prefix) if err != nil { return nil, err } @@ -66,7 +66,7 @@ func (c *catalog) SaveSegmentAssignments(ctx context.Context, pChannelName strin if len(removes) > 0 { if err := etcd.RemoveByBatchWithLimit(removes, util.MaxEtcdTxnNum, func(partialRemoves []string) error { - return c.metaKV.MultiRemove(partialRemoves) + return c.metaKV.MultiRemove(ctx, partialRemoves) }); err != nil { return err } @@ -74,7 +74,7 @@ func (c *catalog) SaveSegmentAssignments(ctx context.Context, pChannelName strin if len(kvs) > 0 { return etcd.SaveByBatchWithLimit(kvs, util.MaxEtcdTxnNum, func(partialKvs map[string]string) error { - return c.metaKV.MultiSave(partialKvs) + return c.metaKV.MultiSave(ctx, partialKvs) }) } return nil diff --git a/internal/metastore/kv/streamingnode/kv_catalog_test.go b/internal/metastore/kv/streamingnode/kv_catalog_test.go index 32aeadae0c788..35e8769664f57 100644 --- a/internal/metastore/kv/streamingnode/kv_catalog_test.go +++ b/internal/metastore/kv/streamingnode/kv_catalog_test.go @@ -19,15 +19,15 @@ func TestCatalog(t *testing.T) { vs, err := proto.Marshal(&v) assert.NoError(t, err) - kv.EXPECT().LoadWithPrefix(mock.Anything).Return([]string{k}, []string{string(vs)}, nil) + kv.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Return([]string{k}, []string{string(vs)}, nil) catalog := NewCataLog(kv) ctx := context.Background() metas, err := catalog.ListSegmentAssignment(ctx, "p1") assert.Len(t, metas, 1) assert.NoError(t, err) - kv.EXPECT().MultiRemove(mock.Anything).Return(nil) - kv.EXPECT().MultiSave(mock.Anything).Return(nil) + kv.EXPECT().MultiRemove(mock.Anything, mock.Anything).Return(nil) + kv.EXPECT().MultiSave(mock.Anything, mock.Anything).Return(nil) err = catalog.SaveSegmentAssignments(ctx, "p1", []*streamingpb.SegmentAssignmentMeta{ { diff --git a/pkg/kv/kv.go b/pkg/kv/kv.go index c78aa0a960f7a..2386506b000f9 100644 --- a/pkg/kv/kv.go +++ b/pkg/kv/kv.go @@ -17,6 +17,8 @@ package kv import ( + "context" + clientv3 "go.etcd.io/etcd/client/v3" "github.com/milvus-io/milvus/pkg/kv/predicates" @@ -40,16 +42,16 @@ func NewCompareFailedError(err error) error { // BaseKV contains base operations of kv. Include save, load and remove. type BaseKV interface { - Load(key string) (string, error) - MultiLoad(keys []string) ([]string, error) - LoadWithPrefix(key string) ([]string, []string, error) - Save(key, value string) error - MultiSave(kvs map[string]string) error - Remove(key string) error - MultiRemove(keys []string) error - RemoveWithPrefix(key string) error - Has(key string) (bool, error) - HasPrefix(prefix string) (bool, error) + Load(ctx context.Context, key string) (string, error) + MultiLoad(ctx context.Context, keys []string) ([]string, error) + LoadWithPrefix(ctx context.Context, key string) ([]string, []string, error) + Save(ctx context.Context, key, value string) error + MultiSave(ctx context.Context, kvs map[string]string) error + Remove(ctx context.Context, key string) error + MultiRemove(ctx context.Context, keys []string) error + RemoveWithPrefix(ctx context.Context, key string) error + Has(ctx context.Context, key string) (bool, error) + HasPrefix(ctx context.Context, prefix string) (bool, error) Close() } @@ -58,8 +60,8 @@ type BaseKV interface { //go:generate mockery --name=TxnKV --with-expecter type TxnKV interface { BaseKV - MultiSaveAndRemove(saves map[string]string, removals []string, preds ...predicates.Predicate) error - MultiSaveAndRemoveWithPrefix(saves map[string]string, removals []string, preds ...predicates.Predicate) error + MultiSaveAndRemove(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error + MultiSaveAndRemoveWithPrefix(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error } // MetaKv is TxnKV for metadata. It should save data with lease. @@ -68,9 +70,9 @@ type TxnKV interface { type MetaKv interface { TxnKV GetPath(key string) string - LoadWithPrefix(key string) ([]string, []string, error) - CompareVersionAndSwap(key string, version int64, target string) (bool, error) - WalkWithPrefix(prefix string, paginationSize int, fn func([]byte, []byte) error) error + LoadWithPrefix(ctx context.Context, key string) ([]string, []string, error) + CompareVersionAndSwap(ctx context.Context, key string, version int64, target string) (bool, error) + WalkWithPrefix(ctx context.Context, prefix string, paginationSize int, fn func([]byte, []byte) error) error } // WatchKV is watchable MetaKv. As of today(2023/06/24), it's coupled with etcd. @@ -78,19 +80,19 @@ type MetaKv interface { //go:generate mockery --name=WatchKV --with-expecter type WatchKV interface { MetaKv - Watch(key string) clientv3.WatchChan - WatchWithPrefix(key string) clientv3.WatchChan - WatchWithRevision(key string, revision int64) clientv3.WatchChan + Watch(ctx context.Context, key string) clientv3.WatchChan + WatchWithPrefix(ctx context.Context, key string) clientv3.WatchChan + WatchWithRevision(ctx context.Context, key string, revision int64) clientv3.WatchChan } // SnapShotKV is TxnKV for snapshot data. It must save timestamp. // //go:generate mockery --name=SnapShotKV --with-expecter type SnapShotKV interface { - Save(key string, value string, ts typeutil.Timestamp) error - Load(key string, ts typeutil.Timestamp) (string, error) - MultiSave(kvs map[string]string, ts typeutil.Timestamp) error - LoadWithPrefix(key string, ts typeutil.Timestamp) ([]string, []string, error) - MultiSaveAndRemove(saves map[string]string, removals []string, ts typeutil.Timestamp) error - MultiSaveAndRemoveWithPrefix(saves map[string]string, removals []string, ts typeutil.Timestamp) error + Save(ctx context.Context, key string, value string, ts typeutil.Timestamp) error + Load(ctx context.Context, key string, ts typeutil.Timestamp) (string, error) + MultiSave(ctx context.Context, kvs map[string]string, ts typeutil.Timestamp) error + LoadWithPrefix(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) + MultiSaveAndRemove(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error + MultiSaveAndRemoveWithPrefix(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error } diff --git a/pkg/kv/rocksdb/rocksdb_kv.go b/pkg/kv/rocksdb/rocksdb_kv.go index b81e7ead385f2..3b156c414f008 100644 --- a/pkg/kv/rocksdb/rocksdb_kv.go +++ b/pkg/kv/rocksdb/rocksdb_kv.go @@ -17,6 +17,7 @@ package rocksdbkv import ( + "context" "fmt" "github.com/cockroachdb/errors" @@ -97,7 +98,7 @@ func (kv *RocksdbKV) GetName() string { } // Load returns the value of specified key -func (kv *RocksdbKV) Load(key string) (string, error) { +func (kv *RocksdbKV) Load(ctx context.Context, key string) (string, error) { if kv.DB == nil { return "", fmt.Errorf("rocksdb instance is nil when load %s", key) } @@ -114,7 +115,7 @@ func (kv *RocksdbKV) Load(key string) (string, error) { return string(value.Data()), nil } -func (kv *RocksdbKV) LoadBytes(key string) ([]byte, error) { +func (kv *RocksdbKV) LoadBytes(ctx context.Context, key string) ([]byte, error) { if kv.DB == nil { return nil, fmt.Errorf("rocksdb instance is nil when load %s", key) } @@ -139,7 +140,7 @@ func (kv *RocksdbKV) LoadBytes(key string) ([]byte, error) { // LoadWithPrefix returns a batch values of keys with a prefix // if prefix is "", then load every thing from the database -func (kv *RocksdbKV) LoadWithPrefix(prefix string) ([]string, []string, error) { +func (kv *RocksdbKV) LoadWithPrefix(ctx context.Context, prefix string) ([]string, []string, error) { if kv.DB == nil { return nil, nil, fmt.Errorf("rocksdb instance is nil when load %s", prefix) } @@ -164,7 +165,7 @@ func (kv *RocksdbKV) LoadWithPrefix(prefix string) ([]string, []string, error) { return keys, values, nil } -func (kv *RocksdbKV) Has(key string) (bool, error) { +func (kv *RocksdbKV) Has(ctx context.Context, key string) (bool, error) { if kv.DB == nil { return false, fmt.Errorf("rocksdb instance is nil when check if has %s", key) } @@ -180,7 +181,7 @@ func (kv *RocksdbKV) Has(key string) (bool, error) { return value.Size() != 0, nil } -func (kv *RocksdbKV) HasPrefix(prefix string) (bool, error) { +func (kv *RocksdbKV) HasPrefix(ctx context.Context, prefix string) (bool, error) { if kv.DB == nil { return false, fmt.Errorf("rocksdb instance is nil when check if has prefix %s", prefix) } @@ -194,7 +195,7 @@ func (kv *RocksdbKV) HasPrefix(prefix string) (bool, error) { return iter.Valid(), nil } -func (kv *RocksdbKV) LoadBytesWithPrefix(prefix string) ([]string, [][]byte, error) { +func (kv *RocksdbKV) LoadBytesWithPrefix(ctx context.Context, prefix string) ([]string, [][]byte, error) { if kv.DB == nil { return nil, nil, fmt.Errorf("rocksdb instance is nil when load %s", prefix) } @@ -229,7 +230,7 @@ func (kv *RocksdbKV) LoadBytesWithPrefix(prefix string) ([]string, [][]byte, err } // MultiLoad load a batch of values by keys -func (kv *RocksdbKV) MultiLoad(keys []string) ([]string, error) { +func (kv *RocksdbKV) MultiLoad(ctx context.Context, keys []string) ([]string, error) { if kv.DB == nil { return nil, errors.New("rocksdb instance is nil when do MultiLoad") } @@ -253,7 +254,7 @@ func (kv *RocksdbKV) MultiLoad(keys []string) ([]string, error) { return values, nil } -func (kv *RocksdbKV) MultiLoadBytes(keys []string) ([][]byte, error) { +func (kv *RocksdbKV) MultiLoadBytes(ctx context.Context, keys []string) ([][]byte, error) { if kv.DB == nil { return nil, errors.New("rocksdb instance is nil when do MultiLoad") } @@ -281,7 +282,7 @@ func (kv *RocksdbKV) MultiLoadBytes(keys []string) ([][]byte, error) { } // Save a pair of key-value -func (kv *RocksdbKV) Save(key, value string) error { +func (kv *RocksdbKV) Save(ctx context.Context, key, value string) error { if kv.DB == nil { return errors.New("rocksdb instance is nil when do save") } @@ -295,7 +296,7 @@ func (kv *RocksdbKV) Save(key, value string) error { return kv.DB.Put(kv.WriteOptions, []byte(key), []byte(value)) } -func (kv *RocksdbKV) SaveBytes(key string, value []byte) error { +func (kv *RocksdbKV) SaveBytes(ctx context.Context, key string, value []byte) error { if kv.DB == nil { return errors.New("rocksdb instance is nil when do save") } @@ -310,7 +311,7 @@ func (kv *RocksdbKV) SaveBytes(key string, value []byte) error { } // MultiSave a batch of key-values -func (kv *RocksdbKV) MultiSave(kvs map[string]string) error { +func (kv *RocksdbKV) MultiSave(ctx context.Context, kvs map[string]string) error { if kv.DB == nil { return errors.New("rocksdb instance is nil when do MultiSave") } @@ -325,7 +326,7 @@ func (kv *RocksdbKV) MultiSave(kvs map[string]string) error { return kv.DB.Write(kv.WriteOptions, writeBatch) } -func (kv *RocksdbKV) MultiSaveBytes(kvs map[string][]byte) error { +func (kv *RocksdbKV) MultiSaveBytes(ctx context.Context, kvs map[string][]byte) error { if kv.DB == nil { return errors.New("rocksdb instance is nil when do MultiSave") } @@ -342,7 +343,7 @@ func (kv *RocksdbKV) MultiSaveBytes(kvs map[string][]byte) error { // RemoveWithPrefix removes a batch of key-values with specified prefix // If prefix is "", then all data in the rocksdb kv will be deleted -func (kv *RocksdbKV) RemoveWithPrefix(prefix string) error { +func (kv *RocksdbKV) RemoveWithPrefix(ctx context.Context, prefix string) error { if kv.DB == nil { return errors.New("rocksdb instance is nil when do RemoveWithPrefix") } @@ -355,17 +356,17 @@ func (kv *RocksdbKV) RemoveWithPrefix(prefix string) error { // seek to the last key iter.SeekToLast() if iter.Valid() { - return kv.DeleteRange(prefix, typeutil.AddOne(string(iter.Key().Data()))) + return kv.DeleteRange(ctx, prefix, typeutil.AddOne(string(iter.Key().Data()))) } // nothing in the range, skip return nil } prefixEnd := typeutil.AddOne(prefix) - return kv.DeleteRange(prefix, prefixEnd) + return kv.DeleteRange(ctx, prefix, prefixEnd) } // Remove is used to remove a pair of key-value -func (kv *RocksdbKV) Remove(key string) error { +func (kv *RocksdbKV) Remove(ctx context.Context, key string) error { if kv.DB == nil { return errors.New("rocksdb instance is nil when do Remove") } @@ -377,7 +378,7 @@ func (kv *RocksdbKV) Remove(key string) error { } // MultiRemove is used to remove a batch of key-values -func (kv *RocksdbKV) MultiRemove(keys []string) error { +func (kv *RocksdbKV) MultiRemove(ctx context.Context, keys []string) error { if kv.DB == nil { return errors.New("rocksdb instance is nil when do MultiRemove") } @@ -391,7 +392,7 @@ func (kv *RocksdbKV) MultiRemove(keys []string) error { } // MultiSaveAndRemove provides a transaction to execute a batch of operations -func (kv *RocksdbKV) MultiSaveAndRemove(saves map[string]string, removals []string, preds ...predicates.Predicate) error { +func (kv *RocksdbKV) MultiSaveAndRemove(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error { if len(preds) > 0 { return merr.WrapErrServiceUnavailable("predicates not supported") } @@ -411,7 +412,7 @@ func (kv *RocksdbKV) MultiSaveAndRemove(saves map[string]string, removals []stri } // DeleteRange remove a batch of key-values from startKey to endKey -func (kv *RocksdbKV) DeleteRange(startKey, endKey string) error { +func (kv *RocksdbKV) DeleteRange(ctx context.Context, startKey, endKey string) error { if kv.DB == nil { return errors.New("Rocksdb instance is nil when do DeleteRange") } @@ -426,7 +427,7 @@ func (kv *RocksdbKV) DeleteRange(startKey, endKey string) error { } // MultiSaveAndRemoveWithPrefix is used to execute a batch operators with the same prefix -func (kv *RocksdbKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals []string, preds ...predicates.Predicate) error { +func (kv *RocksdbKV) MultiSaveAndRemoveWithPrefix(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error { if len(preds) > 0 { return merr.WrapErrServiceUnavailable("predicates not supported") } diff --git a/pkg/kv/rocksdb/rocksdb_kv_test.go b/pkg/kv/rocksdb/rocksdb_kv_test.go index 14c05a0309339..aeecb0a681d5f 100644 --- a/pkg/kv/rocksdb/rocksdb_kv_test.go +++ b/pkg/kv/rocksdb/rocksdb_kv_test.go @@ -17,6 +17,7 @@ package rocksdbkv_test import ( + "context" "os" "strconv" "sync" @@ -39,27 +40,27 @@ func TestRocksdbKV(t *testing.T) { defer os.RemoveAll(name) defer rocksdbKV.Close() // Need to call RemoveWithPrefix - defer rocksdbKV.RemoveWithPrefix("") + defer rocksdbKV.RemoveWithPrefix(context.TODO(), "") - err = rocksdbKV.Save("abc", "123") + err = rocksdbKV.Save(context.TODO(), "abc", "123") assert.NoError(t, err) - err = rocksdbKV.SaveBytes("abcd", []byte("1234")) + err = rocksdbKV.SaveBytes(context.TODO(), "abcd", []byte("1234")) assert.NoError(t, err) - val, err := rocksdbKV.Load("abc") + val, err := rocksdbKV.Load(context.TODO(), "abc") assert.NoError(t, err) assert.Equal(t, val, "123") - value, err := rocksdbKV.LoadBytes("abc") + value, err := rocksdbKV.LoadBytes(context.TODO(), "abc") assert.NoError(t, err) assert.Equal(t, value, []byte("123")) - _, err = rocksdbKV.Load("") + _, err = rocksdbKV.Load(context.TODO(), "") assert.Error(t, err) - _, err = rocksdbKV.LoadBytes("") + _, err = rocksdbKV.LoadBytes(context.TODO(), "") assert.Error(t, err) - keys, vals, err := rocksdbKV.LoadWithPrefix("abc") + keys, vals, err := rocksdbKV.LoadWithPrefix(context.TODO(), "abc") assert.NoError(t, err) assert.Equal(t, len(keys), len(vals)) assert.Equal(t, len(keys), 2) @@ -70,7 +71,7 @@ func TestRocksdbKV(t *testing.T) { assert.Equal(t, vals[1], "1234") var values [][]byte - keys, values, err = rocksdbKV.LoadBytesWithPrefix("abc") + keys, values, err = rocksdbKV.LoadBytesWithPrefix(context.TODO(), "abc") assert.NoError(t, err) assert.Equal(t, len(keys), len(vals)) assert.Equal(t, len(keys), 2) @@ -80,27 +81,27 @@ func TestRocksdbKV(t *testing.T) { assert.Equal(t, values[0], []byte("123")) assert.Equal(t, values[1], []byte("1234")) - err = rocksdbKV.Save("key_1", "123") + err = rocksdbKV.Save(context.TODO(), "key_1", "123") assert.NoError(t, err) - err = rocksdbKV.Save("key_2", "456") + err = rocksdbKV.Save(context.TODO(), "key_2", "456") assert.NoError(t, err) - err = rocksdbKV.Save("key_3", "789") + err = rocksdbKV.Save(context.TODO(), "key_3", "789") assert.NoError(t, err) keys = []string{"key_1", "key_2"} - vals, err = rocksdbKV.MultiLoad(keys) + vals, err = rocksdbKV.MultiLoad(context.TODO(), keys) assert.NoError(t, err) assert.Equal(t, len(vals), len(keys)) assert.Equal(t, vals[0], "123") assert.Equal(t, vals[1], "456") - values, err = rocksdbKV.MultiLoadBytes(keys) + values, err = rocksdbKV.MultiLoadBytes(context.TODO(), keys) assert.NoError(t, err) assert.Equal(t, len(values), len(keys)) assert.Equal(t, values[0], []byte("123")) assert.Equal(t, values[1], []byte("456")) - err = rocksdbKV.MultiRemove(keys) + err = rocksdbKV.MultiRemove(context.TODO(), keys) assert.NoError(t, err) saves := map[string]string{ @@ -109,10 +110,10 @@ func TestRocksdbKV(t *testing.T) { "s_3": "333", } removals := []string{"key_3"} - err = rocksdbKV.MultiSaveAndRemove(saves, removals) + err = rocksdbKV.MultiSaveAndRemove(context.TODO(), saves, removals) assert.NoError(t, err) - err = rocksdbKV.DeleteRange("s_1", "s_3") + err = rocksdbKV.DeleteRange(context.TODO(), "s_1", "s_3") assert.NoError(t, err) } @@ -125,17 +126,17 @@ func TestRocksdbKV_Prefix(t *testing.T) { defer os.RemoveAll(name) defer rocksdbKV.Close() // Need to call RemoveWithPrefix - defer rocksdbKV.RemoveWithPrefix("") + defer rocksdbKV.RemoveWithPrefix(context.TODO(), "") kvs := map[string]string{ "abcd": "123", "abdd": "1234", "abddqqq": "1234555", } - err = rocksdbKV.MultiSave(kvs) + err = rocksdbKV.MultiSave(context.TODO(), kvs) assert.NoError(t, err) - keys, vals, err := rocksdbKV.LoadWithPrefix("abc") + keys, vals, err := rocksdbKV.LoadWithPrefix(context.TODO(), "abc") assert.NoError(t, err) assert.Equal(t, len(keys), 1) @@ -145,50 +146,50 @@ func TestRocksdbKV_Prefix(t *testing.T) { bytesKvs := map[string][]byte{} for k, v := range kvs { - rocksdbKV.Remove(k) + rocksdbKV.Remove(context.TODO(), k) bytesKvs[k] = []byte(v) } - err = rocksdbKV.MultiSaveBytes(bytesKvs) + err = rocksdbKV.MultiSaveBytes(context.TODO(), bytesKvs) assert.NoError(t, err) var values [][]byte - keys, values, err = rocksdbKV.LoadBytesWithPrefix("abc") + keys, values, err = rocksdbKV.LoadBytesWithPrefix(context.TODO(), "abc") assert.NoError(t, err) assert.Equal(t, len(keys), 1) assert.Equal(t, len(values), 1) assert.Equal(t, keys[0], "abcd") assert.Equal(t, values[0], []byte("123")) - err = rocksdbKV.RemoveWithPrefix("abc") + err = rocksdbKV.RemoveWithPrefix(context.TODO(), "abc") assert.NoError(t, err) - val, err := rocksdbKV.Load("abc") + val, err := rocksdbKV.Load(context.TODO(), "abc") assert.NoError(t, err) assert.Equal(t, len(val), 0) - val, err = rocksdbKV.Load("abdd") + val, err = rocksdbKV.Load(context.TODO(), "abdd") assert.NoError(t, err) assert.Equal(t, val, "1234") - val, err = rocksdbKV.Load("abddqqq") + val, err = rocksdbKV.Load(context.TODO(), "abddqqq") assert.NoError(t, err) assert.Equal(t, val, "1234555") // test remove "" - err = rocksdbKV.RemoveWithPrefix("") + err = rocksdbKV.RemoveWithPrefix(context.TODO(), "") assert.NoError(t, err) // test remove from an empty cf - err = rocksdbKV.RemoveWithPrefix("") + err = rocksdbKV.RemoveWithPrefix(context.TODO(), "") assert.NoError(t, err) - val, err = rocksdbKV.Load("abddqqq") + val, err = rocksdbKV.Load(context.TODO(), "abddqqq") assert.NoError(t, err) assert.Equal(t, len(val), 0) // test we can still save after drop - err = rocksdbKV.Save("abcd", "123") + err = rocksdbKV.Save(context.TODO(), "abcd", "123") assert.NoError(t, err) - val, err = rocksdbKV.Load("abcd") + val, err = rocksdbKV.Load(context.TODO(), "abcd") assert.NoError(t, err) assert.Equal(t, val, "123") } @@ -202,24 +203,24 @@ func TestRocksdbKV_Txn(t *testing.T) { defer os.RemoveAll(name) defer rocksdbKV.Close() // Need to call RemoveWithPrefix - defer rocksdbKV.RemoveWithPrefix("") + defer rocksdbKV.RemoveWithPrefix(context.TODO(), "") kvs := map[string]string{ "abcd": "123", "abdd": "1234", "abddqqq": "1234555", } - err = rocksdbKV.MultiSave(kvs) + err = rocksdbKV.MultiSave(context.TODO(), kvs) assert.NoError(t, err) - keys, vals, err := rocksdbKV.LoadWithPrefix("") + keys, vals, err := rocksdbKV.LoadWithPrefix(context.TODO(), "") assert.NoError(t, err) assert.Equal(t, len(keys), 3) assert.Equal(t, len(vals), 3) - err = rocksdbKV.MultiSave(kvs) + err = rocksdbKV.MultiSave(context.TODO(), kvs) assert.NoError(t, err) - keys, vals, err = rocksdbKV.LoadWithPrefix("") + keys, vals, err = rocksdbKV.LoadWithPrefix(context.TODO(), "") assert.NoError(t, err) assert.Equal(t, len(keys), 3) assert.Equal(t, len(vals), 3) @@ -229,8 +230,8 @@ func TestRocksdbKV_Txn(t *testing.T) { kvs2 := map[string]string{ "abfad": "12345", } - rocksdbKV.MultiSaveAndRemoveWithPrefix(kvs2, removePrefix) - keys, vals, err = rocksdbKV.LoadWithPrefix("") + rocksdbKV.MultiSaveAndRemoveWithPrefix(context.TODO(), kvs2, removePrefix) + keys, vals, err = rocksdbKV.LoadWithPrefix(context.TODO(), "") assert.NoError(t, err) assert.Equal(t, len(keys), 1) assert.Equal(t, len(vals), 1) @@ -242,7 +243,7 @@ func TestRocksdbKV_Goroutines(t *testing.T) { assert.NoError(t, err) defer os.RemoveAll(name) defer rocksdbkv.Close() - defer rocksdbkv.RemoveWithPrefix("") + defer rocksdbkv.RemoveWithPrefix(context.TODO(), "") var wg sync.WaitGroup for i := 0; i < 10; i++ { @@ -251,10 +252,10 @@ func TestRocksdbKV_Goroutines(t *testing.T) { defer wg.Done() key := "key_" + strconv.Itoa(i) val := "val_" + strconv.Itoa(i) - err := rocksdbkv.Save(key, val) + err := rocksdbkv.Save(context.TODO(), key, val) assert.NoError(t, err) - getVal, err := rocksdbkv.Load(key) + getVal, err := rocksdbkv.Load(context.TODO(), key) assert.NoError(t, err) assert.Equal(t, getVal, val) }(i) @@ -268,32 +269,32 @@ func TestRocksdbKV_DummyDB(t *testing.T) { assert.NoError(t, err) defer os.RemoveAll(name) defer rocksdbkv.Close() - defer rocksdbkv.RemoveWithPrefix("") + defer rocksdbkv.RemoveWithPrefix(context.TODO(), "") rocksdbkv.DB = nil - _, err = rocksdbkv.Load("") + _, err = rocksdbkv.Load(context.TODO(), "") assert.Error(t, err) - _, _, err = rocksdbkv.LoadWithPrefix("") + _, _, err = rocksdbkv.LoadWithPrefix(context.TODO(), "") assert.Error(t, err) - _, err = rocksdbkv.MultiLoad(nil) + _, err = rocksdbkv.MultiLoad(context.TODO(), nil) assert.Error(t, err) - err = rocksdbkv.Save("", "") + err = rocksdbkv.Save(context.TODO(), "", "") assert.Error(t, err) - err = rocksdbkv.MultiSave(nil) + err = rocksdbkv.MultiSave(context.TODO(), nil) assert.Error(t, err) - err = rocksdbkv.RemoveWithPrefix("") + err = rocksdbkv.RemoveWithPrefix(context.TODO(), "") assert.Error(t, err) - err = rocksdbkv.Remove("") + err = rocksdbkv.Remove(context.TODO(), "") assert.Error(t, err) - err = rocksdbkv.MultiRemove(nil) + err = rocksdbkv.MultiRemove(context.TODO(), nil) assert.Error(t, err) - err = rocksdbkv.MultiSaveAndRemove(nil, nil) + err = rocksdbkv.MultiSaveAndRemove(context.TODO(), nil, nil) assert.Error(t, err) - err = rocksdbkv.DeleteRange("", "") + err = rocksdbkv.DeleteRange(context.TODO(), "", "") assert.Error(t, err) rocksdbkv.ReadOptions = nil - _, err = rocksdbkv.Load("dummy") + _, err = rocksdbkv.Load(context.TODO(), "dummy") assert.Error(t, err) } @@ -303,20 +304,20 @@ func TestRocksdbKV_CornerCase(t *testing.T) { assert.NoError(t, err) defer os.RemoveAll(name) defer rocksdbkv.Close() - defer rocksdbkv.RemoveWithPrefix("") - _, err = rocksdbkv.Load("") + defer rocksdbkv.RemoveWithPrefix(context.TODO(), "") + _, err = rocksdbkv.Load(context.TODO(), "") assert.Error(t, err) - keys, values, err := rocksdbkv.LoadWithPrefix("") + keys, values, err := rocksdbkv.LoadWithPrefix(context.TODO(), "") assert.NoError(t, err) assert.Equal(t, len(keys), 0) assert.Equal(t, len(values), 0) - err = rocksdbkv.Save("", "") + err = rocksdbkv.Save(context.TODO(), "", "") assert.Error(t, err) - err = rocksdbkv.Save("test", "") + err = rocksdbkv.Save(context.TODO(), "test", "") assert.Error(t, err) - err = rocksdbkv.Remove("") + err = rocksdbkv.Remove(context.TODO(), "") assert.Error(t, err) - err = rocksdbkv.DeleteRange("a", "a") + err = rocksdbkv.DeleteRange(context.TODO(), "a", "a") assert.Error(t, err) } @@ -325,19 +326,19 @@ func TestHas(t *testing.T) { db, err := rocksdbkv.NewRocksdbKV(dir) assert.NoError(t, err) defer db.Close() - defer db.RemoveWithPrefix("") + defer db.RemoveWithPrefix(context.TODO(), "") - has, err := db.Has("key1") + has, err := db.Has(context.TODO(), "key1") assert.NoError(t, err) assert.False(t, has) - err = db.Save("key1", "value1") + err = db.Save(context.TODO(), "key1", "value1") assert.NoError(t, err) - has, err = db.Has("key1") + has, err = db.Has(context.TODO(), "key1") assert.NoError(t, err) assert.True(t, has) - err = db.Remove("key1") + err = db.Remove(context.TODO(), "key1") assert.NoError(t, err) - has, err = db.Has("key1") + has, err = db.Has(context.TODO(), "key1") assert.NoError(t, err) assert.False(t, has) } @@ -347,23 +348,23 @@ func TestHasPrefix(t *testing.T) { db, err := rocksdbkv.NewRocksdbKV(dir) assert.NoError(t, err) defer db.Close() - defer db.RemoveWithPrefix("") + defer db.RemoveWithPrefix(context.TODO(), "") - has, err := db.HasPrefix("key") + has, err := db.HasPrefix(context.TODO(), "key") assert.NoError(t, err) assert.False(t, has) - err = db.Save("key1", "value1") + err = db.Save(context.TODO(), "key1", "value1") assert.NoError(t, err) - has, err = db.HasPrefix("key") + has, err = db.HasPrefix(context.TODO(), "key") assert.NoError(t, err) assert.True(t, has) - err = db.Remove("key1") + err = db.Remove(context.TODO(), "key1") assert.NoError(t, err) - has, err = db.HasPrefix("key") + has, err = db.HasPrefix(context.TODO(), "key") assert.NoError(t, err) assert.False(t, has) } @@ -374,13 +375,13 @@ func TestPredicates(t *testing.T) { require.NoError(t, err) defer db.Close() - defer db.RemoveWithPrefix("") + defer db.RemoveWithPrefix(context.TODO(), "") - err = db.MultiSaveAndRemove(map[string]string{}, []string{}, predicates.ValueEqual("a", "b")) + err = db.MultiSaveAndRemove(context.TODO(), map[string]string{}, []string{}, predicates.ValueEqual("a", "b")) assert.Error(t, err) assert.ErrorIs(t, err, merr.ErrServiceUnavailable) - err = db.MultiSaveAndRemoveWithPrefix(map[string]string{}, []string{}, predicates.ValueEqual("a", "b")) + err = db.MultiSaveAndRemoveWithPrefix(context.TODO(), map[string]string{}, []string{}, predicates.ValueEqual("a", "b")) assert.Error(t, err) assert.ErrorIs(t, err, merr.ErrServiceUnavailable) } diff --git a/pkg/mocks/mock_kv/mock_MetaKv.go b/pkg/mocks/mock_kv/mock_MetaKv.go index 52702fd8cfb27..683ef85f18496 100644 --- a/pkg/mocks/mock_kv/mock_MetaKv.go +++ b/pkg/mocks/mock_kv/mock_MetaKv.go @@ -3,8 +3,11 @@ package mock_kv import ( - predicates "github.com/milvus-io/milvus/pkg/kv/predicates" + context "context" + mock "github.com/stretchr/testify/mock" + + predicates "github.com/milvus-io/milvus/pkg/kv/predicates" ) // MockMetaKv is an autogenerated mock type for the MetaKv type @@ -52,9 +55,9 @@ func (_c *MockMetaKv_Close_Call) RunAndReturn(run func()) *MockMetaKv_Close_Call return _c } -// CompareVersionAndSwap provides a mock function with given fields: key, version, target -func (_m *MockMetaKv) CompareVersionAndSwap(key string, version int64, target string) (bool, error) { - ret := _m.Called(key, version, target) +// CompareVersionAndSwap provides a mock function with given fields: ctx, key, version, target +func (_m *MockMetaKv) CompareVersionAndSwap(ctx context.Context, key string, version int64, target string) (bool, error) { + ret := _m.Called(ctx, key, version, target) if len(ret) == 0 { panic("no return value specified for CompareVersionAndSwap") @@ -62,17 +65,17 @@ func (_m *MockMetaKv) CompareVersionAndSwap(key string, version int64, target st var r0 bool var r1 error - if rf, ok := ret.Get(0).(func(string, int64, string) (bool, error)); ok { - return rf(key, version, target) + if rf, ok := ret.Get(0).(func(context.Context, string, int64, string) (bool, error)); ok { + return rf(ctx, key, version, target) } - if rf, ok := ret.Get(0).(func(string, int64, string) bool); ok { - r0 = rf(key, version, target) + if rf, ok := ret.Get(0).(func(context.Context, string, int64, string) bool); ok { + r0 = rf(ctx, key, version, target) } else { r0 = ret.Get(0).(bool) } - if rf, ok := ret.Get(1).(func(string, int64, string) error); ok { - r1 = rf(key, version, target) + if rf, ok := ret.Get(1).(func(context.Context, string, int64, string) error); ok { + r1 = rf(ctx, key, version, target) } else { r1 = ret.Error(1) } @@ -86,16 +89,17 @@ type MockMetaKv_CompareVersionAndSwap_Call struct { } // CompareVersionAndSwap is a helper method to define mock.On call +// - ctx context.Context // - key string // - version int64 // - target string -func (_e *MockMetaKv_Expecter) CompareVersionAndSwap(key interface{}, version interface{}, target interface{}) *MockMetaKv_CompareVersionAndSwap_Call { - return &MockMetaKv_CompareVersionAndSwap_Call{Call: _e.mock.On("CompareVersionAndSwap", key, version, target)} +func (_e *MockMetaKv_Expecter) CompareVersionAndSwap(ctx interface{}, key interface{}, version interface{}, target interface{}) *MockMetaKv_CompareVersionAndSwap_Call { + return &MockMetaKv_CompareVersionAndSwap_Call{Call: _e.mock.On("CompareVersionAndSwap", ctx, key, version, target)} } -func (_c *MockMetaKv_CompareVersionAndSwap_Call) Run(run func(key string, version int64, target string)) *MockMetaKv_CompareVersionAndSwap_Call { +func (_c *MockMetaKv_CompareVersionAndSwap_Call) Run(run func(ctx context.Context, key string, version int64, target string)) *MockMetaKv_CompareVersionAndSwap_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string), args[1].(int64), args[2].(string)) + run(args[0].(context.Context), args[1].(string), args[2].(int64), args[3].(string)) }) return _c } @@ -105,7 +109,7 @@ func (_c *MockMetaKv_CompareVersionAndSwap_Call) Return(_a0 bool, _a1 error) *Mo return _c } -func (_c *MockMetaKv_CompareVersionAndSwap_Call) RunAndReturn(run func(string, int64, string) (bool, error)) *MockMetaKv_CompareVersionAndSwap_Call { +func (_c *MockMetaKv_CompareVersionAndSwap_Call) RunAndReturn(run func(context.Context, string, int64, string) (bool, error)) *MockMetaKv_CompareVersionAndSwap_Call { _c.Call.Return(run) return _c } @@ -156,9 +160,9 @@ func (_c *MockMetaKv_GetPath_Call) RunAndReturn(run func(string) string) *MockMe return _c } -// Has provides a mock function with given fields: key -func (_m *MockMetaKv) Has(key string) (bool, error) { - ret := _m.Called(key) +// Has provides a mock function with given fields: ctx, key +func (_m *MockMetaKv) Has(ctx context.Context, key string) (bool, error) { + ret := _m.Called(ctx, key) if len(ret) == 0 { panic("no return value specified for Has") @@ -166,17 +170,17 @@ func (_m *MockMetaKv) Has(key string) (bool, error) { var r0 bool var r1 error - if rf, ok := ret.Get(0).(func(string) (bool, error)); ok { - return rf(key) + if rf, ok := ret.Get(0).(func(context.Context, string) (bool, error)); ok { + return rf(ctx, key) } - if rf, ok := ret.Get(0).(func(string) bool); ok { - r0 = rf(key) + if rf, ok := ret.Get(0).(func(context.Context, string) bool); ok { + r0 = rf(ctx, key) } else { r0 = ret.Get(0).(bool) } - if rf, ok := ret.Get(1).(func(string) error); ok { - r1 = rf(key) + if rf, ok := ret.Get(1).(func(context.Context, string) error); ok { + r1 = rf(ctx, key) } else { r1 = ret.Error(1) } @@ -190,14 +194,15 @@ type MockMetaKv_Has_Call struct { } // Has is a helper method to define mock.On call +// - ctx context.Context // - key string -func (_e *MockMetaKv_Expecter) Has(key interface{}) *MockMetaKv_Has_Call { - return &MockMetaKv_Has_Call{Call: _e.mock.On("Has", key)} +func (_e *MockMetaKv_Expecter) Has(ctx interface{}, key interface{}) *MockMetaKv_Has_Call { + return &MockMetaKv_Has_Call{Call: _e.mock.On("Has", ctx, key)} } -func (_c *MockMetaKv_Has_Call) Run(run func(key string)) *MockMetaKv_Has_Call { +func (_c *MockMetaKv_Has_Call) Run(run func(ctx context.Context, key string)) *MockMetaKv_Has_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string)) + run(args[0].(context.Context), args[1].(string)) }) return _c } @@ -207,14 +212,14 @@ func (_c *MockMetaKv_Has_Call) Return(_a0 bool, _a1 error) *MockMetaKv_Has_Call return _c } -func (_c *MockMetaKv_Has_Call) RunAndReturn(run func(string) (bool, error)) *MockMetaKv_Has_Call { +func (_c *MockMetaKv_Has_Call) RunAndReturn(run func(context.Context, string) (bool, error)) *MockMetaKv_Has_Call { _c.Call.Return(run) return _c } -// HasPrefix provides a mock function with given fields: prefix -func (_m *MockMetaKv) HasPrefix(prefix string) (bool, error) { - ret := _m.Called(prefix) +// HasPrefix provides a mock function with given fields: ctx, prefix +func (_m *MockMetaKv) HasPrefix(ctx context.Context, prefix string) (bool, error) { + ret := _m.Called(ctx, prefix) if len(ret) == 0 { panic("no return value specified for HasPrefix") @@ -222,17 +227,17 @@ func (_m *MockMetaKv) HasPrefix(prefix string) (bool, error) { var r0 bool var r1 error - if rf, ok := ret.Get(0).(func(string) (bool, error)); ok { - return rf(prefix) + if rf, ok := ret.Get(0).(func(context.Context, string) (bool, error)); ok { + return rf(ctx, prefix) } - if rf, ok := ret.Get(0).(func(string) bool); ok { - r0 = rf(prefix) + if rf, ok := ret.Get(0).(func(context.Context, string) bool); ok { + r0 = rf(ctx, prefix) } else { r0 = ret.Get(0).(bool) } - if rf, ok := ret.Get(1).(func(string) error); ok { - r1 = rf(prefix) + if rf, ok := ret.Get(1).(func(context.Context, string) error); ok { + r1 = rf(ctx, prefix) } else { r1 = ret.Error(1) } @@ -246,14 +251,15 @@ type MockMetaKv_HasPrefix_Call struct { } // HasPrefix is a helper method to define mock.On call +// - ctx context.Context // - prefix string -func (_e *MockMetaKv_Expecter) HasPrefix(prefix interface{}) *MockMetaKv_HasPrefix_Call { - return &MockMetaKv_HasPrefix_Call{Call: _e.mock.On("HasPrefix", prefix)} +func (_e *MockMetaKv_Expecter) HasPrefix(ctx interface{}, prefix interface{}) *MockMetaKv_HasPrefix_Call { + return &MockMetaKv_HasPrefix_Call{Call: _e.mock.On("HasPrefix", ctx, prefix)} } -func (_c *MockMetaKv_HasPrefix_Call) Run(run func(prefix string)) *MockMetaKv_HasPrefix_Call { +func (_c *MockMetaKv_HasPrefix_Call) Run(run func(ctx context.Context, prefix string)) *MockMetaKv_HasPrefix_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string)) + run(args[0].(context.Context), args[1].(string)) }) return _c } @@ -263,14 +269,14 @@ func (_c *MockMetaKv_HasPrefix_Call) Return(_a0 bool, _a1 error) *MockMetaKv_Has return _c } -func (_c *MockMetaKv_HasPrefix_Call) RunAndReturn(run func(string) (bool, error)) *MockMetaKv_HasPrefix_Call { +func (_c *MockMetaKv_HasPrefix_Call) RunAndReturn(run func(context.Context, string) (bool, error)) *MockMetaKv_HasPrefix_Call { _c.Call.Return(run) return _c } -// Load provides a mock function with given fields: key -func (_m *MockMetaKv) Load(key string) (string, error) { - ret := _m.Called(key) +// Load provides a mock function with given fields: ctx, key +func (_m *MockMetaKv) Load(ctx context.Context, key string) (string, error) { + ret := _m.Called(ctx, key) if len(ret) == 0 { panic("no return value specified for Load") @@ -278,17 +284,17 @@ func (_m *MockMetaKv) Load(key string) (string, error) { var r0 string var r1 error - if rf, ok := ret.Get(0).(func(string) (string, error)); ok { - return rf(key) + if rf, ok := ret.Get(0).(func(context.Context, string) (string, error)); ok { + return rf(ctx, key) } - if rf, ok := ret.Get(0).(func(string) string); ok { - r0 = rf(key) + if rf, ok := ret.Get(0).(func(context.Context, string) string); ok { + r0 = rf(ctx, key) } else { r0 = ret.Get(0).(string) } - if rf, ok := ret.Get(1).(func(string) error); ok { - r1 = rf(key) + if rf, ok := ret.Get(1).(func(context.Context, string) error); ok { + r1 = rf(ctx, key) } else { r1 = ret.Error(1) } @@ -302,14 +308,15 @@ type MockMetaKv_Load_Call struct { } // Load is a helper method to define mock.On call +// - ctx context.Context // - key string -func (_e *MockMetaKv_Expecter) Load(key interface{}) *MockMetaKv_Load_Call { - return &MockMetaKv_Load_Call{Call: _e.mock.On("Load", key)} +func (_e *MockMetaKv_Expecter) Load(ctx interface{}, key interface{}) *MockMetaKv_Load_Call { + return &MockMetaKv_Load_Call{Call: _e.mock.On("Load", ctx, key)} } -func (_c *MockMetaKv_Load_Call) Run(run func(key string)) *MockMetaKv_Load_Call { +func (_c *MockMetaKv_Load_Call) Run(run func(ctx context.Context, key string)) *MockMetaKv_Load_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string)) + run(args[0].(context.Context), args[1].(string)) }) return _c } @@ -319,14 +326,14 @@ func (_c *MockMetaKv_Load_Call) Return(_a0 string, _a1 error) *MockMetaKv_Load_C return _c } -func (_c *MockMetaKv_Load_Call) RunAndReturn(run func(string) (string, error)) *MockMetaKv_Load_Call { +func (_c *MockMetaKv_Load_Call) RunAndReturn(run func(context.Context, string) (string, error)) *MockMetaKv_Load_Call { _c.Call.Return(run) return _c } -// LoadWithPrefix provides a mock function with given fields: key -func (_m *MockMetaKv) LoadWithPrefix(key string) ([]string, []string, error) { - ret := _m.Called(key) +// LoadWithPrefix provides a mock function with given fields: ctx, key +func (_m *MockMetaKv) LoadWithPrefix(ctx context.Context, key string) ([]string, []string, error) { + ret := _m.Called(ctx, key) if len(ret) == 0 { panic("no return value specified for LoadWithPrefix") @@ -335,27 +342,27 @@ func (_m *MockMetaKv) LoadWithPrefix(key string) ([]string, []string, error) { var r0 []string var r1 []string var r2 error - if rf, ok := ret.Get(0).(func(string) ([]string, []string, error)); ok { - return rf(key) + if rf, ok := ret.Get(0).(func(context.Context, string) ([]string, []string, error)); ok { + return rf(ctx, key) } - if rf, ok := ret.Get(0).(func(string) []string); ok { - r0 = rf(key) + if rf, ok := ret.Get(0).(func(context.Context, string) []string); ok { + r0 = rf(ctx, key) } else { if ret.Get(0) != nil { r0 = ret.Get(0).([]string) } } - if rf, ok := ret.Get(1).(func(string) []string); ok { - r1 = rf(key) + if rf, ok := ret.Get(1).(func(context.Context, string) []string); ok { + r1 = rf(ctx, key) } else { if ret.Get(1) != nil { r1 = ret.Get(1).([]string) } } - if rf, ok := ret.Get(2).(func(string) error); ok { - r2 = rf(key) + if rf, ok := ret.Get(2).(func(context.Context, string) error); ok { + r2 = rf(ctx, key) } else { r2 = ret.Error(2) } @@ -369,14 +376,15 @@ type MockMetaKv_LoadWithPrefix_Call struct { } // LoadWithPrefix is a helper method to define mock.On call +// - ctx context.Context // - key string -func (_e *MockMetaKv_Expecter) LoadWithPrefix(key interface{}) *MockMetaKv_LoadWithPrefix_Call { - return &MockMetaKv_LoadWithPrefix_Call{Call: _e.mock.On("LoadWithPrefix", key)} +func (_e *MockMetaKv_Expecter) LoadWithPrefix(ctx interface{}, key interface{}) *MockMetaKv_LoadWithPrefix_Call { + return &MockMetaKv_LoadWithPrefix_Call{Call: _e.mock.On("LoadWithPrefix", ctx, key)} } -func (_c *MockMetaKv_LoadWithPrefix_Call) Run(run func(key string)) *MockMetaKv_LoadWithPrefix_Call { +func (_c *MockMetaKv_LoadWithPrefix_Call) Run(run func(ctx context.Context, key string)) *MockMetaKv_LoadWithPrefix_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string)) + run(args[0].(context.Context), args[1].(string)) }) return _c } @@ -386,14 +394,14 @@ func (_c *MockMetaKv_LoadWithPrefix_Call) Return(_a0 []string, _a1 []string, _a2 return _c } -func (_c *MockMetaKv_LoadWithPrefix_Call) RunAndReturn(run func(string) ([]string, []string, error)) *MockMetaKv_LoadWithPrefix_Call { +func (_c *MockMetaKv_LoadWithPrefix_Call) RunAndReturn(run func(context.Context, string) ([]string, []string, error)) *MockMetaKv_LoadWithPrefix_Call { _c.Call.Return(run) return _c } -// MultiLoad provides a mock function with given fields: keys -func (_m *MockMetaKv) MultiLoad(keys []string) ([]string, error) { - ret := _m.Called(keys) +// MultiLoad provides a mock function with given fields: ctx, keys +func (_m *MockMetaKv) MultiLoad(ctx context.Context, keys []string) ([]string, error) { + ret := _m.Called(ctx, keys) if len(ret) == 0 { panic("no return value specified for MultiLoad") @@ -401,19 +409,19 @@ func (_m *MockMetaKv) MultiLoad(keys []string) ([]string, error) { var r0 []string var r1 error - if rf, ok := ret.Get(0).(func([]string) ([]string, error)); ok { - return rf(keys) + if rf, ok := ret.Get(0).(func(context.Context, []string) ([]string, error)); ok { + return rf(ctx, keys) } - if rf, ok := ret.Get(0).(func([]string) []string); ok { - r0 = rf(keys) + if rf, ok := ret.Get(0).(func(context.Context, []string) []string); ok { + r0 = rf(ctx, keys) } else { if ret.Get(0) != nil { r0 = ret.Get(0).([]string) } } - if rf, ok := ret.Get(1).(func([]string) error); ok { - r1 = rf(keys) + if rf, ok := ret.Get(1).(func(context.Context, []string) error); ok { + r1 = rf(ctx, keys) } else { r1 = ret.Error(1) } @@ -427,14 +435,15 @@ type MockMetaKv_MultiLoad_Call struct { } // MultiLoad is a helper method to define mock.On call +// - ctx context.Context // - keys []string -func (_e *MockMetaKv_Expecter) MultiLoad(keys interface{}) *MockMetaKv_MultiLoad_Call { - return &MockMetaKv_MultiLoad_Call{Call: _e.mock.On("MultiLoad", keys)} +func (_e *MockMetaKv_Expecter) MultiLoad(ctx interface{}, keys interface{}) *MockMetaKv_MultiLoad_Call { + return &MockMetaKv_MultiLoad_Call{Call: _e.mock.On("MultiLoad", ctx, keys)} } -func (_c *MockMetaKv_MultiLoad_Call) Run(run func(keys []string)) *MockMetaKv_MultiLoad_Call { +func (_c *MockMetaKv_MultiLoad_Call) Run(run func(ctx context.Context, keys []string)) *MockMetaKv_MultiLoad_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].([]string)) + run(args[0].(context.Context), args[1].([]string)) }) return _c } @@ -444,22 +453,22 @@ func (_c *MockMetaKv_MultiLoad_Call) Return(_a0 []string, _a1 error) *MockMetaKv return _c } -func (_c *MockMetaKv_MultiLoad_Call) RunAndReturn(run func([]string) ([]string, error)) *MockMetaKv_MultiLoad_Call { +func (_c *MockMetaKv_MultiLoad_Call) RunAndReturn(run func(context.Context, []string) ([]string, error)) *MockMetaKv_MultiLoad_Call { _c.Call.Return(run) return _c } -// MultiRemove provides a mock function with given fields: keys -func (_m *MockMetaKv) MultiRemove(keys []string) error { - ret := _m.Called(keys) +// MultiRemove provides a mock function with given fields: ctx, keys +func (_m *MockMetaKv) MultiRemove(ctx context.Context, keys []string) error { + ret := _m.Called(ctx, keys) if len(ret) == 0 { panic("no return value specified for MultiRemove") } var r0 error - if rf, ok := ret.Get(0).(func([]string) error); ok { - r0 = rf(keys) + if rf, ok := ret.Get(0).(func(context.Context, []string) error); ok { + r0 = rf(ctx, keys) } else { r0 = ret.Error(0) } @@ -473,14 +482,15 @@ type MockMetaKv_MultiRemove_Call struct { } // MultiRemove is a helper method to define mock.On call +// - ctx context.Context // - keys []string -func (_e *MockMetaKv_Expecter) MultiRemove(keys interface{}) *MockMetaKv_MultiRemove_Call { - return &MockMetaKv_MultiRemove_Call{Call: _e.mock.On("MultiRemove", keys)} +func (_e *MockMetaKv_Expecter) MultiRemove(ctx interface{}, keys interface{}) *MockMetaKv_MultiRemove_Call { + return &MockMetaKv_MultiRemove_Call{Call: _e.mock.On("MultiRemove", ctx, keys)} } -func (_c *MockMetaKv_MultiRemove_Call) Run(run func(keys []string)) *MockMetaKv_MultiRemove_Call { +func (_c *MockMetaKv_MultiRemove_Call) Run(run func(ctx context.Context, keys []string)) *MockMetaKv_MultiRemove_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].([]string)) + run(args[0].(context.Context), args[1].([]string)) }) return _c } @@ -490,22 +500,22 @@ func (_c *MockMetaKv_MultiRemove_Call) Return(_a0 error) *MockMetaKv_MultiRemove return _c } -func (_c *MockMetaKv_MultiRemove_Call) RunAndReturn(run func([]string) error) *MockMetaKv_MultiRemove_Call { +func (_c *MockMetaKv_MultiRemove_Call) RunAndReturn(run func(context.Context, []string) error) *MockMetaKv_MultiRemove_Call { _c.Call.Return(run) return _c } -// MultiSave provides a mock function with given fields: kvs -func (_m *MockMetaKv) MultiSave(kvs map[string]string) error { - ret := _m.Called(kvs) +// MultiSave provides a mock function with given fields: ctx, kvs +func (_m *MockMetaKv) MultiSave(ctx context.Context, kvs map[string]string) error { + ret := _m.Called(ctx, kvs) if len(ret) == 0 { panic("no return value specified for MultiSave") } var r0 error - if rf, ok := ret.Get(0).(func(map[string]string) error); ok { - r0 = rf(kvs) + if rf, ok := ret.Get(0).(func(context.Context, map[string]string) error); ok { + r0 = rf(ctx, kvs) } else { r0 = ret.Error(0) } @@ -519,14 +529,15 @@ type MockMetaKv_MultiSave_Call struct { } // MultiSave is a helper method to define mock.On call +// - ctx context.Context // - kvs map[string]string -func (_e *MockMetaKv_Expecter) MultiSave(kvs interface{}) *MockMetaKv_MultiSave_Call { - return &MockMetaKv_MultiSave_Call{Call: _e.mock.On("MultiSave", kvs)} +func (_e *MockMetaKv_Expecter) MultiSave(ctx interface{}, kvs interface{}) *MockMetaKv_MultiSave_Call { + return &MockMetaKv_MultiSave_Call{Call: _e.mock.On("MultiSave", ctx, kvs)} } -func (_c *MockMetaKv_MultiSave_Call) Run(run func(kvs map[string]string)) *MockMetaKv_MultiSave_Call { +func (_c *MockMetaKv_MultiSave_Call) Run(run func(ctx context.Context, kvs map[string]string)) *MockMetaKv_MultiSave_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(map[string]string)) + run(args[0].(context.Context), args[1].(map[string]string)) }) return _c } @@ -536,19 +547,19 @@ func (_c *MockMetaKv_MultiSave_Call) Return(_a0 error) *MockMetaKv_MultiSave_Cal return _c } -func (_c *MockMetaKv_MultiSave_Call) RunAndReturn(run func(map[string]string) error) *MockMetaKv_MultiSave_Call { +func (_c *MockMetaKv_MultiSave_Call) RunAndReturn(run func(context.Context, map[string]string) error) *MockMetaKv_MultiSave_Call { _c.Call.Return(run) return _c } -// MultiSaveAndRemove provides a mock function with given fields: saves, removals, preds -func (_m *MockMetaKv) MultiSaveAndRemove(saves map[string]string, removals []string, preds ...predicates.Predicate) error { +// MultiSaveAndRemove provides a mock function with given fields: ctx, saves, removals, preds +func (_m *MockMetaKv) MultiSaveAndRemove(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error { _va := make([]interface{}, len(preds)) for _i := range preds { _va[_i] = preds[_i] } var _ca []interface{} - _ca = append(_ca, saves, removals) + _ca = append(_ca, ctx, saves, removals) _ca = append(_ca, _va...) ret := _m.Called(_ca...) @@ -557,8 +568,8 @@ func (_m *MockMetaKv) MultiSaveAndRemove(saves map[string]string, removals []str } var r0 error - if rf, ok := ret.Get(0).(func(map[string]string, []string, ...predicates.Predicate) error); ok { - r0 = rf(saves, removals, preds...) + if rf, ok := ret.Get(0).(func(context.Context, map[string]string, []string, ...predicates.Predicate) error); ok { + r0 = rf(ctx, saves, removals, preds...) } else { r0 = ret.Error(0) } @@ -572,23 +583,24 @@ type MockMetaKv_MultiSaveAndRemove_Call struct { } // MultiSaveAndRemove is a helper method to define mock.On call +// - ctx context.Context // - saves map[string]string // - removals []string // - preds ...predicates.Predicate -func (_e *MockMetaKv_Expecter) MultiSaveAndRemove(saves interface{}, removals interface{}, preds ...interface{}) *MockMetaKv_MultiSaveAndRemove_Call { +func (_e *MockMetaKv_Expecter) MultiSaveAndRemove(ctx interface{}, saves interface{}, removals interface{}, preds ...interface{}) *MockMetaKv_MultiSaveAndRemove_Call { return &MockMetaKv_MultiSaveAndRemove_Call{Call: _e.mock.On("MultiSaveAndRemove", - append([]interface{}{saves, removals}, preds...)...)} + append([]interface{}{ctx, saves, removals}, preds...)...)} } -func (_c *MockMetaKv_MultiSaveAndRemove_Call) Run(run func(saves map[string]string, removals []string, preds ...predicates.Predicate)) *MockMetaKv_MultiSaveAndRemove_Call { +func (_c *MockMetaKv_MultiSaveAndRemove_Call) Run(run func(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate)) *MockMetaKv_MultiSaveAndRemove_Call { _c.Call.Run(func(args mock.Arguments) { - variadicArgs := make([]predicates.Predicate, len(args)-2) - for i, a := range args[2:] { + variadicArgs := make([]predicates.Predicate, len(args)-3) + for i, a := range args[3:] { if a != nil { variadicArgs[i] = a.(predicates.Predicate) } } - run(args[0].(map[string]string), args[1].([]string), variadicArgs...) + run(args[0].(context.Context), args[1].(map[string]string), args[2].([]string), variadicArgs...) }) return _c } @@ -598,19 +610,19 @@ func (_c *MockMetaKv_MultiSaveAndRemove_Call) Return(_a0 error) *MockMetaKv_Mult return _c } -func (_c *MockMetaKv_MultiSaveAndRemove_Call) RunAndReturn(run func(map[string]string, []string, ...predicates.Predicate) error) *MockMetaKv_MultiSaveAndRemove_Call { +func (_c *MockMetaKv_MultiSaveAndRemove_Call) RunAndReturn(run func(context.Context, map[string]string, []string, ...predicates.Predicate) error) *MockMetaKv_MultiSaveAndRemove_Call { _c.Call.Return(run) return _c } -// MultiSaveAndRemoveWithPrefix provides a mock function with given fields: saves, removals, preds -func (_m *MockMetaKv) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals []string, preds ...predicates.Predicate) error { +// MultiSaveAndRemoveWithPrefix provides a mock function with given fields: ctx, saves, removals, preds +func (_m *MockMetaKv) MultiSaveAndRemoveWithPrefix(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error { _va := make([]interface{}, len(preds)) for _i := range preds { _va[_i] = preds[_i] } var _ca []interface{} - _ca = append(_ca, saves, removals) + _ca = append(_ca, ctx, saves, removals) _ca = append(_ca, _va...) ret := _m.Called(_ca...) @@ -619,8 +631,8 @@ func (_m *MockMetaKv) MultiSaveAndRemoveWithPrefix(saves map[string]string, remo } var r0 error - if rf, ok := ret.Get(0).(func(map[string]string, []string, ...predicates.Predicate) error); ok { - r0 = rf(saves, removals, preds...) + if rf, ok := ret.Get(0).(func(context.Context, map[string]string, []string, ...predicates.Predicate) error); ok { + r0 = rf(ctx, saves, removals, preds...) } else { r0 = ret.Error(0) } @@ -634,23 +646,24 @@ type MockMetaKv_MultiSaveAndRemoveWithPrefix_Call struct { } // MultiSaveAndRemoveWithPrefix is a helper method to define mock.On call +// - ctx context.Context // - saves map[string]string // - removals []string // - preds ...predicates.Predicate -func (_e *MockMetaKv_Expecter) MultiSaveAndRemoveWithPrefix(saves interface{}, removals interface{}, preds ...interface{}) *MockMetaKv_MultiSaveAndRemoveWithPrefix_Call { +func (_e *MockMetaKv_Expecter) MultiSaveAndRemoveWithPrefix(ctx interface{}, saves interface{}, removals interface{}, preds ...interface{}) *MockMetaKv_MultiSaveAndRemoveWithPrefix_Call { return &MockMetaKv_MultiSaveAndRemoveWithPrefix_Call{Call: _e.mock.On("MultiSaveAndRemoveWithPrefix", - append([]interface{}{saves, removals}, preds...)...)} + append([]interface{}{ctx, saves, removals}, preds...)...)} } -func (_c *MockMetaKv_MultiSaveAndRemoveWithPrefix_Call) Run(run func(saves map[string]string, removals []string, preds ...predicates.Predicate)) *MockMetaKv_MultiSaveAndRemoveWithPrefix_Call { +func (_c *MockMetaKv_MultiSaveAndRemoveWithPrefix_Call) Run(run func(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate)) *MockMetaKv_MultiSaveAndRemoveWithPrefix_Call { _c.Call.Run(func(args mock.Arguments) { - variadicArgs := make([]predicates.Predicate, len(args)-2) - for i, a := range args[2:] { + variadicArgs := make([]predicates.Predicate, len(args)-3) + for i, a := range args[3:] { if a != nil { variadicArgs[i] = a.(predicates.Predicate) } } - run(args[0].(map[string]string), args[1].([]string), variadicArgs...) + run(args[0].(context.Context), args[1].(map[string]string), args[2].([]string), variadicArgs...) }) return _c } @@ -660,22 +673,22 @@ func (_c *MockMetaKv_MultiSaveAndRemoveWithPrefix_Call) Return(_a0 error) *MockM return _c } -func (_c *MockMetaKv_MultiSaveAndRemoveWithPrefix_Call) RunAndReturn(run func(map[string]string, []string, ...predicates.Predicate) error) *MockMetaKv_MultiSaveAndRemoveWithPrefix_Call { +func (_c *MockMetaKv_MultiSaveAndRemoveWithPrefix_Call) RunAndReturn(run func(context.Context, map[string]string, []string, ...predicates.Predicate) error) *MockMetaKv_MultiSaveAndRemoveWithPrefix_Call { _c.Call.Return(run) return _c } -// Remove provides a mock function with given fields: key -func (_m *MockMetaKv) Remove(key string) error { - ret := _m.Called(key) +// Remove provides a mock function with given fields: ctx, key +func (_m *MockMetaKv) Remove(ctx context.Context, key string) error { + ret := _m.Called(ctx, key) if len(ret) == 0 { panic("no return value specified for Remove") } var r0 error - if rf, ok := ret.Get(0).(func(string) error); ok { - r0 = rf(key) + if rf, ok := ret.Get(0).(func(context.Context, string) error); ok { + r0 = rf(ctx, key) } else { r0 = ret.Error(0) } @@ -689,14 +702,15 @@ type MockMetaKv_Remove_Call struct { } // Remove is a helper method to define mock.On call +// - ctx context.Context // - key string -func (_e *MockMetaKv_Expecter) Remove(key interface{}) *MockMetaKv_Remove_Call { - return &MockMetaKv_Remove_Call{Call: _e.mock.On("Remove", key)} +func (_e *MockMetaKv_Expecter) Remove(ctx interface{}, key interface{}) *MockMetaKv_Remove_Call { + return &MockMetaKv_Remove_Call{Call: _e.mock.On("Remove", ctx, key)} } -func (_c *MockMetaKv_Remove_Call) Run(run func(key string)) *MockMetaKv_Remove_Call { +func (_c *MockMetaKv_Remove_Call) Run(run func(ctx context.Context, key string)) *MockMetaKv_Remove_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string)) + run(args[0].(context.Context), args[1].(string)) }) return _c } @@ -706,22 +720,22 @@ func (_c *MockMetaKv_Remove_Call) Return(_a0 error) *MockMetaKv_Remove_Call { return _c } -func (_c *MockMetaKv_Remove_Call) RunAndReturn(run func(string) error) *MockMetaKv_Remove_Call { +func (_c *MockMetaKv_Remove_Call) RunAndReturn(run func(context.Context, string) error) *MockMetaKv_Remove_Call { _c.Call.Return(run) return _c } -// RemoveWithPrefix provides a mock function with given fields: key -func (_m *MockMetaKv) RemoveWithPrefix(key string) error { - ret := _m.Called(key) +// RemoveWithPrefix provides a mock function with given fields: ctx, key +func (_m *MockMetaKv) RemoveWithPrefix(ctx context.Context, key string) error { + ret := _m.Called(ctx, key) if len(ret) == 0 { panic("no return value specified for RemoveWithPrefix") } var r0 error - if rf, ok := ret.Get(0).(func(string) error); ok { - r0 = rf(key) + if rf, ok := ret.Get(0).(func(context.Context, string) error); ok { + r0 = rf(ctx, key) } else { r0 = ret.Error(0) } @@ -735,14 +749,15 @@ type MockMetaKv_RemoveWithPrefix_Call struct { } // RemoveWithPrefix is a helper method to define mock.On call +// - ctx context.Context // - key string -func (_e *MockMetaKv_Expecter) RemoveWithPrefix(key interface{}) *MockMetaKv_RemoveWithPrefix_Call { - return &MockMetaKv_RemoveWithPrefix_Call{Call: _e.mock.On("RemoveWithPrefix", key)} +func (_e *MockMetaKv_Expecter) RemoveWithPrefix(ctx interface{}, key interface{}) *MockMetaKv_RemoveWithPrefix_Call { + return &MockMetaKv_RemoveWithPrefix_Call{Call: _e.mock.On("RemoveWithPrefix", ctx, key)} } -func (_c *MockMetaKv_RemoveWithPrefix_Call) Run(run func(key string)) *MockMetaKv_RemoveWithPrefix_Call { +func (_c *MockMetaKv_RemoveWithPrefix_Call) Run(run func(ctx context.Context, key string)) *MockMetaKv_RemoveWithPrefix_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string)) + run(args[0].(context.Context), args[1].(string)) }) return _c } @@ -752,22 +767,22 @@ func (_c *MockMetaKv_RemoveWithPrefix_Call) Return(_a0 error) *MockMetaKv_Remove return _c } -func (_c *MockMetaKv_RemoveWithPrefix_Call) RunAndReturn(run func(string) error) *MockMetaKv_RemoveWithPrefix_Call { +func (_c *MockMetaKv_RemoveWithPrefix_Call) RunAndReturn(run func(context.Context, string) error) *MockMetaKv_RemoveWithPrefix_Call { _c.Call.Return(run) return _c } -// Save provides a mock function with given fields: key, value -func (_m *MockMetaKv) Save(key string, value string) error { - ret := _m.Called(key, value) +// Save provides a mock function with given fields: ctx, key, value +func (_m *MockMetaKv) Save(ctx context.Context, key string, value string) error { + ret := _m.Called(ctx, key, value) if len(ret) == 0 { panic("no return value specified for Save") } var r0 error - if rf, ok := ret.Get(0).(func(string, string) error); ok { - r0 = rf(key, value) + if rf, ok := ret.Get(0).(func(context.Context, string, string) error); ok { + r0 = rf(ctx, key, value) } else { r0 = ret.Error(0) } @@ -781,15 +796,16 @@ type MockMetaKv_Save_Call struct { } // Save is a helper method to define mock.On call +// - ctx context.Context // - key string // - value string -func (_e *MockMetaKv_Expecter) Save(key interface{}, value interface{}) *MockMetaKv_Save_Call { - return &MockMetaKv_Save_Call{Call: _e.mock.On("Save", key, value)} +func (_e *MockMetaKv_Expecter) Save(ctx interface{}, key interface{}, value interface{}) *MockMetaKv_Save_Call { + return &MockMetaKv_Save_Call{Call: _e.mock.On("Save", ctx, key, value)} } -func (_c *MockMetaKv_Save_Call) Run(run func(key string, value string)) *MockMetaKv_Save_Call { +func (_c *MockMetaKv_Save_Call) Run(run func(ctx context.Context, key string, value string)) *MockMetaKv_Save_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string), args[1].(string)) + run(args[0].(context.Context), args[1].(string), args[2].(string)) }) return _c } @@ -799,22 +815,22 @@ func (_c *MockMetaKv_Save_Call) Return(_a0 error) *MockMetaKv_Save_Call { return _c } -func (_c *MockMetaKv_Save_Call) RunAndReturn(run func(string, string) error) *MockMetaKv_Save_Call { +func (_c *MockMetaKv_Save_Call) RunAndReturn(run func(context.Context, string, string) error) *MockMetaKv_Save_Call { _c.Call.Return(run) return _c } -// WalkWithPrefix provides a mock function with given fields: prefix, paginationSize, fn -func (_m *MockMetaKv) WalkWithPrefix(prefix string, paginationSize int, fn func([]byte, []byte) error) error { - ret := _m.Called(prefix, paginationSize, fn) +// WalkWithPrefix provides a mock function with given fields: ctx, prefix, paginationSize, fn +func (_m *MockMetaKv) WalkWithPrefix(ctx context.Context, prefix string, paginationSize int, fn func([]byte, []byte) error) error { + ret := _m.Called(ctx, prefix, paginationSize, fn) if len(ret) == 0 { panic("no return value specified for WalkWithPrefix") } var r0 error - if rf, ok := ret.Get(0).(func(string, int, func([]byte, []byte) error) error); ok { - r0 = rf(prefix, paginationSize, fn) + if rf, ok := ret.Get(0).(func(context.Context, string, int, func([]byte, []byte) error) error); ok { + r0 = rf(ctx, prefix, paginationSize, fn) } else { r0 = ret.Error(0) } @@ -828,16 +844,17 @@ type MockMetaKv_WalkWithPrefix_Call struct { } // WalkWithPrefix is a helper method to define mock.On call +// - ctx context.Context // - prefix string // - paginationSize int // - fn func([]byte , []byte) error -func (_e *MockMetaKv_Expecter) WalkWithPrefix(prefix interface{}, paginationSize interface{}, fn interface{}) *MockMetaKv_WalkWithPrefix_Call { - return &MockMetaKv_WalkWithPrefix_Call{Call: _e.mock.On("WalkWithPrefix", prefix, paginationSize, fn)} +func (_e *MockMetaKv_Expecter) WalkWithPrefix(ctx interface{}, prefix interface{}, paginationSize interface{}, fn interface{}) *MockMetaKv_WalkWithPrefix_Call { + return &MockMetaKv_WalkWithPrefix_Call{Call: _e.mock.On("WalkWithPrefix", ctx, prefix, paginationSize, fn)} } -func (_c *MockMetaKv_WalkWithPrefix_Call) Run(run func(prefix string, paginationSize int, fn func([]byte, []byte) error)) *MockMetaKv_WalkWithPrefix_Call { +func (_c *MockMetaKv_WalkWithPrefix_Call) Run(run func(ctx context.Context, prefix string, paginationSize int, fn func([]byte, []byte) error)) *MockMetaKv_WalkWithPrefix_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(string), args[1].(int), args[2].(func([]byte, []byte) error)) + run(args[0].(context.Context), args[1].(string), args[2].(int), args[3].(func([]byte, []byte) error)) }) return _c } @@ -847,7 +864,7 @@ func (_c *MockMetaKv_WalkWithPrefix_Call) Return(_a0 error) *MockMetaKv_WalkWith return _c } -func (_c *MockMetaKv_WalkWithPrefix_Call) RunAndReturn(run func(string, int, func([]byte, []byte) error) error) *MockMetaKv_WalkWithPrefix_Call { +func (_c *MockMetaKv_WalkWithPrefix_Call) RunAndReturn(run func(context.Context, string, int, func([]byte, []byte) error) error) *MockMetaKv_WalkWithPrefix_Call { _c.Call.Return(run) return _c } diff --git a/pkg/mq/mqimpl/rocksmq/server/rocksmq_impl.go b/pkg/mq/mqimpl/rocksmq/server/rocksmq_impl.go index 195604b5eeb7f..43554e1620d2a 100644 --- a/pkg/mq/mqimpl/rocksmq/server/rocksmq_impl.go +++ b/pkg/mq/mqimpl/rocksmq/server/rocksmq_impl.go @@ -12,6 +12,7 @@ package server import ( + "context" "fmt" "path" "strconv" @@ -356,7 +357,7 @@ func (rmq *rocksmq) Info() bool { } pageTsSizeKey := constructKey(PageTsTitle, topic) - pages, _, err := rmq.kv.LoadWithPrefix(pageTsSizeKey) + pages, _, err := rmq.kv.LoadWithPrefix(context.TODO(), pageTsSizeKey) if err != nil { log.Error("Rocksmq get page num failed", zap.String("topic", topic)) rtn = false @@ -364,7 +365,7 @@ func (rmq *rocksmq) Info() bool { } msgSizeKey := MessageSizeTitle + topic - msgSizeVal, err := rmq.kv.Load(msgSizeKey) + msgSizeVal, err := rmq.kv.Load(context.TODO(), msgSizeKey) if err != nil { log.Error("Rocksmq get last page size failed", zap.String("topic", topic)) rtn = false @@ -405,7 +406,7 @@ func (rmq *rocksmq) CreateTopic(topicName string) error { // topicIDKey is the only identifier of a topic topicIDKey := TopicIDTitle + topicName - val, err := rmq.kv.Load(topicIDKey) + val, err := rmq.kv.Load(context.TODO(), topicIDKey) if err != nil { return err } @@ -429,7 +430,7 @@ func (rmq *rocksmq) CreateTopic(topicName string) error { // Initialize topic id to its creating time, we don't really use it for now nowTs := strconv.FormatInt(time.Now().Unix(), 10) kvs[topicIDKey] = nowTs - if err = rmq.kv.MultiSave(kvs); err != nil { + if err = rmq.kv.MultiSave(context.TODO(), kvs); err != nil { return retry.Unrecoverable(err) } @@ -459,28 +460,28 @@ func (rmq *rocksmq) DestroyTopic(topicName string) error { // clean the topic data it self fixTopicName := topicName + "/" - err := rmq.kv.RemoveWithPrefix(fixTopicName) + err := rmq.kv.RemoveWithPrefix(context.TODO(), fixTopicName) if err != nil { return err } // clean page size info pageMsgSizeKey := constructKey(PageMsgSizeTitle, topicName) - err = rmq.kv.RemoveWithPrefix(pageMsgSizeKey) + err = rmq.kv.RemoveWithPrefix(context.TODO(), pageMsgSizeKey) if err != nil { return err } // clean page ts info pageMsgTsKey := constructKey(PageTsTitle, topicName) - err = rmq.kv.RemoveWithPrefix(pageMsgTsKey) + err = rmq.kv.RemoveWithPrefix(context.TODO(), pageMsgTsKey) if err != nil { return err } // cleaned acked ts info ackedTsKey := constructKey(AckedTsTitle, topicName) - err = rmq.kv.RemoveWithPrefix(ackedTsKey) + err = rmq.kv.RemoveWithPrefix(context.TODO(), ackedTsKey) if err != nil { return err } @@ -492,7 +493,7 @@ func (rmq *rocksmq) DestroyTopic(topicName string) error { var removedKeys []string removedKeys = append(removedKeys, topicIDKey, msgSizeKey) // Batch remove, atomic operation - err = rmq.kv.MultiRemove(removedKeys) + err = rmq.kv.MultiRemove(context.TODO(), removedKeys) if err != nil { return err } @@ -702,7 +703,7 @@ func (rmq *rocksmq) Produce(topicName string, messages []ProducerMessage) ([]Uni func (rmq *rocksmq) updatePageInfo(topicName string, msgIDs []UniqueID, msgSizes map[UniqueID]int64) error { params := paramtable.Get() msgSizeKey := MessageSizeTitle + topicName - msgSizeVal, err := rmq.kv.Load(msgSizeKey) + msgSizeVal, err := rmq.kv.Load(context.TODO(), msgSizeKey) if err != nil { return err } @@ -731,7 +732,7 @@ func (rmq *rocksmq) updatePageInfo(topicName string, msgIDs []UniqueID, msgSizes } } mutateBuffer[msgSizeKey] = strconv.FormatInt(curMsgSize, 10) - err = rmq.kv.MultiSave(mutateBuffer) + err = rmq.kv.MultiSave(context.TODO(), mutateBuffer) return err } @@ -1125,7 +1126,7 @@ func (rmq *rocksmq) updateAckedInfo(topicName, groupName string, firstID UniqueI ackedTsKvs[pageAckedTsKey] = nowTs } } - err := rmq.kv.MultiSave(ackedTsKvs) + err := rmq.kv.MultiSave(context.TODO(), ackedTsKvs) if err != nil { return err } diff --git a/pkg/mq/mqimpl/rocksmq/server/rocksmq_retention.go b/pkg/mq/mqimpl/rocksmq/server/rocksmq_retention.go index 98ce3d03bcc22..fe58c8cc6433d 100644 --- a/pkg/mq/mqimpl/rocksmq/server/rocksmq_retention.go +++ b/pkg/mq/mqimpl/rocksmq/server/rocksmq_retention.go @@ -12,6 +12,7 @@ package server import ( + "context" "fmt" "path" "strconv" @@ -55,7 +56,7 @@ func initRetentionInfo(kv *rocksdbkv.RocksdbKV, db *gorocksdb.DB) (*retentionInf closeWg: sync.WaitGroup{}, } // Get topic from topic begin id - topicKeys, _, err := ri.kv.LoadWithPrefix(TopicIDTitle) + topicKeys, _, err := ri.kv.LoadWithPrefix(context.TODO(), TopicIDTitle) if err != nil { return nil, err } @@ -164,7 +165,7 @@ func (ri *retentionInfo) expiredCleanUp(topic string) error { return err } ackedTsKey := fixedAckedTsKey + "/" + strconv.FormatInt(pageID, 10) - ackedTsVal, err := ri.kv.Load(ackedTsKey) + ackedTsVal, err := ri.kv.Load(context.TODO(), ackedTsKey) if err != nil { return err } @@ -265,7 +266,7 @@ func (ri *retentionInfo) calculateTopicAckedSize(topic string) (int64, error) { // check if page is acked ackedTsKey := fixedAckedTsKey + "/" + strconv.FormatInt(pageID, 10) - ackedTsVal, err := ri.kv.Load(ackedTsKey) + ackedTsVal, err := ri.kv.Load(context.TODO(), ackedTsKey) if err != nil { return -1, err } diff --git a/pkg/mq/mqimpl/rocksmq/server/rocksmq_retention_test.go b/pkg/mq/mqimpl/rocksmq/server/rocksmq_retention_test.go index fcb7b143f7585..1266d4068c550 100644 --- a/pkg/mq/mqimpl/rocksmq/server/rocksmq_retention_test.go +++ b/pkg/mq/mqimpl/rocksmq/server/rocksmq_retention_test.go @@ -12,6 +12,7 @@ package server import ( + "context" "os" "strconv" "testing" @@ -97,24 +98,24 @@ func TestRmqRetention_Basic(t *testing.T) { // test acked size acked ts and other meta are updated as expect msgSizeKey := MessageSizeTitle + topicName - msgSizeVal, err := rmq.kv.Load(msgSizeKey) + msgSizeVal, err := rmq.kv.Load(context.TODO(), msgSizeKey) assert.NoError(t, err) assert.Equal(t, msgSizeVal, "0") pageMsgSizeKey := constructKey(PageMsgSizeTitle, topicName) - keys, values, err := rmq.kv.LoadWithPrefix(pageMsgSizeKey) + keys, values, err := rmq.kv.LoadWithPrefix(context.TODO(), pageMsgSizeKey) assert.NoError(t, err) assert.Equal(t, len(keys), 0) assert.Equal(t, len(values), 0) pageTsSizeKey := constructKey(PageTsTitle, topicName) - keys, values, err = rmq.kv.LoadWithPrefix(pageTsSizeKey) + keys, values, err = rmq.kv.LoadWithPrefix(context.TODO(), pageTsSizeKey) assert.NoError(t, err) assert.Equal(t, len(keys), 0) assert.Equal(t, len(values), 0) aclTsSizeKey := constructKey(AckedTsTitle, topicName) - keys, values, err = rmq.kv.LoadWithPrefix(aclTsSizeKey) + keys, values, err = rmq.kv.LoadWithPrefix(context.TODO(), aclTsSizeKey) assert.NoError(t, err) assert.Equal(t, len(keys), 0) assert.Equal(t, len(values), 0) @@ -183,7 +184,7 @@ func TestRmqRetention_NotConsumed(t *testing.T) { id := cMsgs[0].MsgID aclTsSizeKey := constructKey(AckedTsTitle, topicName) - keys, values, err := rmq.kv.LoadWithPrefix(aclTsSizeKey) + keys, values, err := rmq.kv.LoadWithPrefix(context.TODO(), aclTsSizeKey) assert.NoError(t, err) assert.Equal(t, len(keys), 2) assert.Equal(t, len(values), 2) @@ -202,26 +203,26 @@ func TestRmqRetention_NotConsumed(t *testing.T) { // test acked size acked ts and other meta are updated as expect msgSizeKey := MessageSizeTitle + topicName - msgSizeVal, err := rmq.kv.Load(msgSizeKey) + msgSizeVal, err := rmq.kv.Load(context.TODO(), msgSizeKey) assert.NoError(t, err) assert.Equal(t, msgSizeVal, "0") // should only clean 2 pages pageMsgSizeKey := constructKey(PageMsgSizeTitle, topicName) - keys, values, err = rmq.kv.LoadWithPrefix(pageMsgSizeKey) + keys, values, err = rmq.kv.LoadWithPrefix(context.TODO(), pageMsgSizeKey) assert.NoError(t, err) assert.Equal(t, len(keys), 48) assert.Equal(t, len(values), 48) pageTsSizeKey := constructKey(PageTsTitle, topicName) - keys, values, err = rmq.kv.LoadWithPrefix(pageTsSizeKey) + keys, values, err = rmq.kv.LoadWithPrefix(context.TODO(), pageTsSizeKey) assert.NoError(t, err) assert.Equal(t, len(keys), 48) assert.Equal(t, len(values), 48) aclTsSizeKey = constructKey(AckedTsTitle, topicName) - keys, values, err = rmq.kv.LoadWithPrefix(aclTsSizeKey) + keys, values, err = rmq.kv.LoadWithPrefix(context.TODO(), aclTsSizeKey) assert.NoError(t, err) assert.Equal(t, len(keys), 0) assert.Equal(t, len(values), 0) @@ -317,46 +318,46 @@ func TestRmqRetention_MultipleTopic(t *testing.T) { // test acked size acked ts and other meta are updated as expect msgSizeKey := MessageSizeTitle + topicName - msgSizeVal, err := rmq.kv.Load(msgSizeKey) + msgSizeVal, err := rmq.kv.Load(context.TODO(), msgSizeKey) assert.NoError(t, err) assert.Equal(t, msgSizeVal, "0") // 100 page left, each entity is a page pageMsgSizeKey := constructKey(PageMsgSizeTitle, "topic_a") - keys, values, err := rmq.kv.LoadWithPrefix(pageMsgSizeKey) + keys, values, err := rmq.kv.LoadWithPrefix(context.TODO(), pageMsgSizeKey) assert.NoError(t, err) assert.Equal(t, len(keys), 0) assert.Equal(t, len(values), 0) pageTsSizeKey := constructKey(PageTsTitle, "topic_a") - keys, values, err = rmq.kv.LoadWithPrefix(pageTsSizeKey) + keys, values, err = rmq.kv.LoadWithPrefix(context.TODO(), pageTsSizeKey) assert.NoError(t, err) assert.Equal(t, len(keys), 0) assert.Equal(t, len(values), 0) aclTsSizeKey := constructKey(AckedTsTitle, "topic_a") - keys, values, err = rmq.kv.LoadWithPrefix(aclTsSizeKey) + keys, values, err = rmq.kv.LoadWithPrefix(context.TODO(), aclTsSizeKey) assert.NoError(t, err) assert.Equal(t, len(keys), 0) assert.Equal(t, len(values), 0) // for topic B, nothing has been cleadn pageMsgSizeKey = constructKey(PageMsgSizeTitle, "topic_b") - keys, values, err = rmq.kv.LoadWithPrefix(pageMsgSizeKey) + keys, values, err = rmq.kv.LoadWithPrefix(context.TODO(), pageMsgSizeKey) assert.NoError(t, err) assert.Equal(t, len(keys), 50) assert.Equal(t, len(values), 50) pageTsSizeKey = constructKey(PageTsTitle, "topic_b") - keys, values, err = rmq.kv.LoadWithPrefix(pageTsSizeKey) + keys, values, err = rmq.kv.LoadWithPrefix(context.TODO(), pageTsSizeKey) assert.NoError(t, err) assert.Equal(t, len(keys), 50) assert.Equal(t, len(values), 50) aclTsSizeKey = constructKey(AckedTsTitle, "topic_b") - keys, values, err = rmq.kv.LoadWithPrefix(aclTsSizeKey) + keys, values, err = rmq.kv.LoadWithPrefix(context.TODO(), aclTsSizeKey) assert.NoError(t, err) assert.Equal(t, len(keys), 0) assert.Equal(t, len(values), 0) @@ -539,26 +540,26 @@ func TestRmqRetention_PageTimeExpire(t *testing.T) { // test acked size acked ts and other meta are updated as expect msgSizeKey := MessageSizeTitle + topicName - msgSizeVal, err := rmq.kv.Load(msgSizeKey) + msgSizeVal, err := rmq.kv.Load(context.TODO(), msgSizeKey) assert.NoError(t, err) assert.Equal(t, msgSizeVal, "0") // 100 page left, each entity is a page pageMsgSizeKey := constructKey(PageMsgSizeTitle, topicName) - keys, values, err := rmq.kv.LoadWithPrefix(pageMsgSizeKey) + keys, values, err := rmq.kv.LoadWithPrefix(context.TODO(), pageMsgSizeKey) assert.NoError(t, err) assert.Equal(t, len(keys), 100) assert.Equal(t, len(values), 100) pageTsSizeKey := constructKey(PageTsTitle, topicName) - keys, values, err = rmq.kv.LoadWithPrefix(pageTsSizeKey) + keys, values, err = rmq.kv.LoadWithPrefix(context.TODO(), pageTsSizeKey) assert.NoError(t, err) assert.Equal(t, len(keys), 100) assert.Equal(t, len(values), 100) aclTsSizeKey := constructKey(AckedTsTitle, topicName) - keys, values, err = rmq.kv.LoadWithPrefix(aclTsSizeKey) + keys, values, err = rmq.kv.LoadWithPrefix(context.TODO(), aclTsSizeKey) assert.NoError(t, err) assert.Equal(t, len(keys), 100) assert.Equal(t, len(values), 100)