Add HasPrefix in kv (#25462)

Signed-off-by: sunby <bingyi.sun@zilliz.com>
Co-authored-by: sunby <bingyi.sun@zilliz.com>

internal/kv/etcd/embed_etcd_kv.go

@@ -151,6 +151,21 @@ func (kv *EmbedEtcdKV) Has(key string) (bool, error) {
 	return resp.Count != 0, nil
 }
 
+func (kv *EmbedEtcdKV) HasPrefix(prefix string) (bool, error) {
+	prefix = path.Join(kv.rootPath, prefix)
+	log.Debug("HasPrefix", zap.String("prefix", prefix))
+
+	ctx, cancel := context.WithTimeout(context.TODO(), RequestTimeout)
+	defer cancel()
+
+	resp, err := kv.client.Get(ctx, prefix, clientv3.WithPrefix(), clientv3.WithCountOnly(), clientv3.WithLimit(1))
+	if err != nil {
+		return false, err
+	}
+
+	return resp.Count != 0, nil
+}
+
 // LoadBytesWithPrefix returns all the keys and values with the given key prefix
 func (kv *EmbedEtcdKV) LoadBytesWithPrefix(key string) ([]string, [][]byte, error) {
 	key = path.Join(kv.rootPath, key)

internal/kv/etcd/embed_etcd_kv_test.go

@@ -865,4 +865,31 @@ func TestEmbedEtcd(te *testing.T) {
 		assert.NoError(t, err)
 		assert.False(t, has)
 	})
+
+	te.Run("test has prefix", func(t *testing.T) {
+		rootPath := "/etcd/test/root/hasprefix"
+		kv, err := embed_etcd_kv.NewMetaKvFactory(rootPath, &param.EtcdCfg)
+		assert.NoError(t, err)
+
+		defer kv.Close()
+		defer kv.RemoveWithPrefix("")
+
+		has, err := kv.HasPrefix("key")
+		assert.NoError(t, err)
+		assert.False(t, has)
+
+		err = kv.Save("key1", "value1")
+		assert.NoError(t, err)
+
+		has, err = kv.HasPrefix("key")
+		assert.NoError(t, err)
+		assert.True(t, has)
+
+		err = kv.Remove("key1")
+		assert.NoError(t, err)
+
+		has, err = kv.HasPrefix("key")
+		assert.NoError(t, err)
+		assert.False(t, has)
+	})
 }

internal/kv/etcd/etcd_kv.go

@@ -135,6 +135,21 @@ func (kv *etcdKV) Has(key string) (bool, error) {
 	return resp.Count != 0, nil
 }
 
+func (kv *etcdKV) HasPrefix(prefix string) (bool, error) {
+	start := time.Now()
+	prefix = path.Join(kv.rootPath, prefix)
+	ctx, cancel := context.WithTimeout(context.TODO(), RequestTimeout)
+	defer cancel()
+
+	resp, err := kv.getEtcdMeta(ctx, prefix, clientv3.WithPrefix(), clientv3.WithLimit(1), clientv3.WithCountOnly())
+	if err != nil {
+		return false, err
+	}
+
+	CheckElapseAndWarn(start, "Slow etcd operation has", zap.String("prefix", prefix))
+	return resp.Count != 0, nil
+}
+
 // LoadBytesWithPrefix returns all the keys and values with the given key prefix.
 func (kv *etcdKV) LoadBytesWithPrefix(key string) ([]string, [][]byte, error) {
 	start := time.Now()

internal/kv/etcd/etcd_kv_test.go

@@ -862,3 +862,41 @@ func TestHas(t *testing.T) {
 	assert.NoError(t, err)
 	assert.False(t, has)
 }
+
+func TestHasPrefix(t *testing.T) {
+	etcdCli, err := etcd.GetEtcdClient(
+		Params.EtcdCfg.UseEmbedEtcd.GetAsBool(),
+		Params.EtcdCfg.EtcdUseSSL.GetAsBool(),
+		Params.EtcdCfg.Endpoints.GetAsStrings(),
+		Params.EtcdCfg.EtcdTLSCert.GetValue(),
+		Params.EtcdCfg.EtcdTLSKey.GetValue(),
+		Params.EtcdCfg.EtcdTLSCACert.GetValue(),
+		Params.EtcdCfg.EtcdTLSMinVersion.GetValue())
+	defer etcdCli.Close()
+	assert.NoError(t, err)
+	rootPath := "/etcd/test/root/hasprefix"
+	kv := etcdkv.NewEtcdKV(etcdCli, rootPath)
+	err = kv.RemoveWithPrefix("")
+	require.NoError(t, err)
+
+	defer kv.Close()
+	defer kv.RemoveWithPrefix("")
+
+	has, err := kv.HasPrefix("key")
+	assert.NoError(t, err)
+	assert.False(t, has)
+
+	err = kv.Save("key1", "value1")
+	assert.NoError(t, err)
+
+	has, err = kv.HasPrefix("key")
+	assert.NoError(t, err)
+	assert.True(t, has)
+
+	err = kv.Remove("key1")
+	assert.NoError(t, err)
+
+	has, err = kv.HasPrefix("key")
+	assert.NoError(t, err)
+	assert.False(t, has)
+}

internal/kv/kv.go

@@ -48,6 +48,7 @@ type BaseKV interface {
 	MultiRemove(keys []string) error
 	RemoveWithPrefix(key string) error
 	Has(key string) (bool, error)
+	HasPrefix(prefix string) (bool, error)
 	Close()
 }
 

internal/kv/mem/mem_kv.go

@@ -359,3 +359,16 @@ func (kv *MemoryKV) Has(key string) (bool, error) {
 	defer kv.Unlock()
 	return kv.tree.Has(memoryKVItem{key: key}), nil
 }
+
+func (kv *MemoryKV) HasPrefix(prefix string) (bool, error) {
+	kv.Lock()
+	defer kv.Unlock()
+
+	var has bool
+	kv.tree.AscendGreaterOrEqual(memoryKVItem{key: prefix}, func(i btree.Item) bool {
+		has = strings.HasPrefix(i.(memoryKVItem).key, prefix)
+		return false
+	})
+
+	return has, nil
+}

internal/kv/mem/mem_kv_test.go

@@ -220,3 +220,25 @@ func TestHas(t *testing.T) {
 	assert.NoError(t, err)
 	assert.False(t, has)
 }
+
+func TestHasPrefix(t *testing.T) {
+	kv := NewMemoryKV()
+
+	has, err := kv.HasPrefix("key")
+	assert.NoError(t, err)
+	assert.False(t, has)
+
+	err = kv.Save("key1", "value1")
+	assert.NoError(t, err)
+
+	has, err = kv.HasPrefix("key")
+	assert.NoError(t, err)
+	assert.True(t, has)
+
+	err = kv.Remove("key1")
+	assert.NoError(t, err)
+
+	has, err = kv.HasPrefix("key")
+	assert.NoError(t, err)
+	assert.False(t, has)
+}

internal/kv/mocks/MetaKv.go

@@ -197,6 +197,58 @@ func (_c *MetaKv_Has_Call) RunAndReturn(run func(string) (bool, error)) *MetaKv_
 	return _c
 }
 
+// HasPrefix provides a mock function with given fields: prefix
+func (_m *MetaKv) HasPrefix(prefix string) (bool, error) {
+	ret := _m.Called(prefix)
+
+	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(string) bool); ok {
+		r0 = rf(prefix)
+	} else {
+		r0 = ret.Get(0).(bool)
+	}
+
+	if rf, ok := ret.Get(1).(func(string) error); ok {
+		r1 = rf(prefix)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// MetaKv_HasPrefix_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'HasPrefix'
+type MetaKv_HasPrefix_Call struct {
+	*mock.Call
+}
+
+// HasPrefix is a helper method to define mock.On call
+//   - prefix string
+func (_e *MetaKv_Expecter) HasPrefix(prefix interface{}) *MetaKv_HasPrefix_Call {
+	return &MetaKv_HasPrefix_Call{Call: _e.mock.On("HasPrefix", prefix)}
+}
+
+func (_c *MetaKv_HasPrefix_Call) Run(run func(prefix string)) *MetaKv_HasPrefix_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(string))
+	})
+	return _c
+}
+
+func (_c *MetaKv_HasPrefix_Call) Return(_a0 bool, _a1 error) *MetaKv_HasPrefix_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *MetaKv_HasPrefix_Call) RunAndReturn(run func(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)

internal/kv/mocks/TxnKV.go

@@ -101,6 +101,58 @@ func (_c *TxnKV_Has_Call) RunAndReturn(run func(string) (bool, error)) *TxnKV_Ha
 	return _c
 }
 
+// HasPrefix provides a mock function with given fields: prefix
+func (_m *TxnKV) HasPrefix(prefix string) (bool, error) {
+	ret := _m.Called(prefix)
+
+	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(string) bool); ok {
+		r0 = rf(prefix)
+	} else {
+		r0 = ret.Get(0).(bool)
+	}
+
+	if rf, ok := ret.Get(1).(func(string) error); ok {
+		r1 = rf(prefix)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// TxnKV_HasPrefix_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'HasPrefix'
+type TxnKV_HasPrefix_Call struct {
+	*mock.Call
+}
+
+// HasPrefix is a helper method to define mock.On call
+//   - prefix string
+func (_e *TxnKV_Expecter) HasPrefix(prefix interface{}) *TxnKV_HasPrefix_Call {
+	return &TxnKV_HasPrefix_Call{Call: _e.mock.On("HasPrefix", prefix)}
+}
+
+func (_c *TxnKV_HasPrefix_Call) Run(run func(prefix string)) *TxnKV_HasPrefix_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(string))
+	})
+	return _c
+}
+
+func (_c *TxnKV_HasPrefix_Call) Return(_a0 bool, _a1 error) *TxnKV_HasPrefix_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *TxnKV_HasPrefix_Call) RunAndReturn(run func(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)

internal/kv/mocks/WatchKV.go

@@ -201,6 +201,58 @@ func (_c *WatchKV_Has_Call) RunAndReturn(run func(string) (bool, error)) *WatchK
 	return _c
 }
 
+// HasPrefix provides a mock function with given fields: prefix
+func (_m *WatchKV) HasPrefix(prefix string) (bool, error) {
+	ret := _m.Called(prefix)
+
+	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(string) bool); ok {
+		r0 = rf(prefix)
+	} else {
+		r0 = ret.Get(0).(bool)
+	}
+
+	if rf, ok := ret.Get(1).(func(string) error); ok {
+		r1 = rf(prefix)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// WatchKV_HasPrefix_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'HasPrefix'
+type WatchKV_HasPrefix_Call struct {
+	*mock.Call
+}
+
+// HasPrefix is a helper method to define mock.On call
+//   - prefix string
+func (_e *WatchKV_Expecter) HasPrefix(prefix interface{}) *WatchKV_HasPrefix_Call {
+	return &WatchKV_HasPrefix_Call{Call: _e.mock.On("HasPrefix", prefix)}
+}
+
+func (_c *WatchKV_HasPrefix_Call) Run(run func(prefix string)) *WatchKV_HasPrefix_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(string))
+	})
+	return _c
+}
+
+func (_c *WatchKV_HasPrefix_Call) Return(_a0 bool, _a1 error) *WatchKV_HasPrefix_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *WatchKV_HasPrefix_Call) RunAndReturn(run func(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)

internal/kv/rocksdb/rocksdb_kv.go

@@ -178,6 +178,20 @@ func (kv *RocksdbKV) Has(key string) (bool, error) {
 	return value.Size() != 0, nil
 }
 
+func (kv *RocksdbKV) HasPrefix(prefix string) (bool, error) {
+	if kv.DB == nil {
+		return false, fmt.Errorf("rocksdb instance is nil when check if has prefix %s", prefix)
+	}
+
+	option := gorocksdb.NewDefaultReadOptions()
+	defer option.Destroy()
+	iter := NewRocksIteratorWithUpperBound(kv.DB, typeutil.AddOne(prefix), option)
+	defer iter.Close()
+
+	iter.Seek([]byte(prefix))
+	return iter.Valid(), nil
+}
+
 func (kv *RocksdbKV) LoadBytesWithPrefix(prefix string) ([]string, [][]byte, error) {
 	if kv.DB == nil {
 		return nil, nil, fmt.Errorf("rocksdb instance is nil when load %s", prefix)

internal/kv/rocksdb/rocksdb_kv_test.go

@@ -363,3 +363,29 @@ func TestHas(t *testing.T) {
 	assert.NoError(t, err)
 	assert.False(t, has)
 }
+
+func TestHasPrefix(t *testing.T) {
+	dir := t.TempDir()
+	db, err := rocksdbkv.NewRocksdbKV(dir)
+	assert.NoError(t, err)
+	defer db.Close()
+	defer db.RemoveWithPrefix("")
+
+	has, err := db.HasPrefix("key")
+	assert.NoError(t, err)
+	assert.False(t, has)
+
+	err = db.Save("key1", "value1")
+	assert.NoError(t, err)
+
+	has, err = db.HasPrefix("key")
+	assert.NoError(t, err)
+	assert.True(t, has)
+
+	err = db.Remove("key1")
+	assert.NoError(t, err)
+
+	has, err = db.HasPrefix("key")
+	assert.NoError(t, err)
+	assert.False(t, has)
+}