Repository: oom-ai/oomstore
Pull request #1071 (2022-01-25 01:46:46)

Jinghan/filter features by group in `ListCachedFeature`
This PR refactors informer method `ListCachedFeature`: add parameter `groupName` and `groupID` so that callers could filter features by group.

related to #1008


Patch:
diff --git a/go.mod b/go.mod
index b0f4def9..00079c27 100644
--- a/go.mod
+++ b/go.mod
@@ -20,7 +20,6 @@ require (
 	github.com/lib/pq v1.10.4
 	github.com/mitchellh/mapstructure v1.4.3
 	github.com/olekukonko/tablewriter v0.0.5
-	github.com/pingcap/errors v0.11.5-0.20201126102027-b0a155152ca3
 	github.com/pingcap/log v0.0.0-20211215031037-e024ba4eb0ee
 	github.com/pkg/errors v0.9.1
 	github.com/snowflakedb/gosnowflake v1.6.5
@@ -30,7 +29,6 @@ require (
 	github.com/tikv/client-go/v2 v2.0.0-rc
 	go.uber.org/zap v1.19.1
 	google.golang.org/api v0.63.0
-	google.golang.org/genproto v0.0.0-20211223182754-3ac035c7e7cb
 	google.golang.org/grpc v1.43.0
 	google.golang.org/protobuf v1.27.1
 	gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b
@@ -93,6 +91,7 @@ require (
 	github.com/matttproud/golang_protobuf_extensions v1.0.1 // indirect
 	github.com/opentracing/opentracing-go v1.2.0 // indirect
 	github.com/pierrec/lz4/v4 v4.1.12 // indirect
+	github.com/pingcap/errors v0.11.5-0.20201126102027-b0a155152ca3 // indirect
 	github.com/pingcap/failpoint v0.0.0-20210918120811-547c13e3eb00 // indirect
 	github.com/pingcap/kvproto v0.0.0-20211229082925-7a8280c36daf // indirect
 	github.com/pkg/browser v0.0.0-20210911075715-681adbf594b8 // indirect
@@ -120,6 +119,7 @@ require (
 	golang.org/x/tools v0.1.5 // indirect
 	golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 // indirect
 	google.golang.org/appengine v1.6.7 // indirect
+	google.golang.org/genproto v0.0.0-20211223182754-3ac035c7e7cb // indirect
 	gopkg.in/inf.v0 v0.9.1 // indirect
 	gopkg.in/natefinch/lumberjack.v2 v2.0.0 // indirect
 	lukechampine.com/uint128 v1.1.1 // indirect
diff --git a/internal/database/metadata/informer/feature.go b/internal/database/metadata/informer/feature.go
index 20c5f15a..bdcde282 100644
--- a/internal/database/metadata/informer/feature.go
+++ b/internal/database/metadata/informer/feature.go
@@ -1,6 +1,7 @@
 package informer
 
 import (
+	"github.com/oom-ai/oomstore/internal/database/metadata"
 	"github.com/oom-ai/oomstore/pkg/oomstore/types"
 )
 
@@ -16,13 +17,13 @@ func (c *FeatureCache) Enrich(groupCache *GroupCache) {
 	}
 }
 
-func (c *FeatureCache) List(fullNames *[]string) types.FeatureList {
+func (c *FeatureCache) List(opt metadata.ListCachedFeatureOpt) types.FeatureList {
 	features := c.FeatureList
 
-	// filter names
-	if fullNames != nil {
+	// filter featureNames
+	if opt.FullNames != nil {
 		var tmp types.FeatureList
-		for _, fullName := range *fullNames {
+		for _, fullName := range *opt.FullNames {
 			if f := features.Find(func(f *types.Feature) bool {
 				return f.FullName() == fullName
 			}); f != nil {
@@ -31,6 +32,17 @@ func (c *FeatureCache) List(fullNames *[]string) types.FeatureList {
 		}
 		features = tmp
 	}
-
+	// filter groupName
+	if opt.GroupName != nil {
+		features = features.Filter(func(f *types.Feature) bool {
+			return f.Group.Name == *opt.GroupName
+		})
+	}
+	// filter groupID
+	if opt.GroupID != nil {
+		features = features.Filter(func(f *types.Feature) bool {
+			return f.Group.ID == *opt.GroupID
+		})
+	}
 	return features
 }
diff --git a/internal/database/metadata/informer/informer.go b/internal/database/metadata/informer/informer.go
index fbf81c96..98f4dcea 100644
--- a/internal/database/metadata/informer/informer.go
+++ b/internal/database/metadata/informer/informer.go
@@ -88,6 +88,6 @@ func (f *Informer) Cache() *Cache {
 	return f.cache.Load().(*Cache)
 }
 
-func (f *Informer) ListCachedFeature(ctx context.Context, featureNames *[]string) types.FeatureList {
-	return f.Cache().Features.List(featureNames).Copy()
+func (f *Informer) ListCachedFeature(ctx context.Context, opt metadata.ListCachedFeatureOpt) types.FeatureList {
+	return f.Cache().Features.List(opt).Copy()
 }
diff --git a/internal/database/metadata/mock_metadata/store.go b/internal/database/metadata/mock_metadata/store.go
index 4179e367..1fcca36c 100644
--- a/internal/database/metadata/mock_metadata/store.go
+++ b/internal/database/metadata/mock_metadata/store.go
@@ -233,17 +233,17 @@ func (mr *MockStoreMockRecorder) GetRevisionBy(ctx, groupID, revision interface{
 }
 
 // ListCachedFeature mocks base method.
-func (m *MockStore) ListCachedFeature(ctx context.Context, fullNames *[]string) types.FeatureList {
+func (m *MockStore) ListCachedFeature(ctx context.Context, opt metadata.ListCachedFeatureOpt) types.FeatureList {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "ListCachedFeature", ctx, fullNames)
+	ret := m.ctrl.Call(m, "ListCachedFeature", ctx, opt)
 	ret0, _ := ret[0].(types.FeatureList)
 	return ret0
 }
 
 // ListCachedFeature indicates an expected call of ListCachedFeature.
-func (mr *MockStoreMockRecorder) ListCachedFeature(ctx, fullNames interface{}) *gomock.Call {
+func (mr *MockStoreMockRecorder) ListCachedFeature(ctx, opt interface{}) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListCachedFeature", reflect.TypeOf((*MockStore)(nil).ListCachedFeature), ctx, fullNames)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListCachedFeature", reflect.TypeOf((*MockStore)(nil).ListCachedFeature), ctx, opt)
 }
 
 // ListEntity mocks base method.
@@ -762,17 +762,17 @@ func (m *MockCacheStore) EXPECT() *MockCacheStoreMockRecorder {
 }
 
 // ListCachedFeature mocks base method.
-func (m *MockCacheStore) ListCachedFeature(ctx context.Context, fullNames *[]string) types.FeatureList {
+func (m *MockCacheStore) ListCachedFeature(ctx context.Context, opt metadata.ListCachedFeatureOpt) types.FeatureList {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "ListCachedFeature", ctx, fullNames)
+	ret := m.ctrl.Call(m, "ListCachedFeature", ctx, opt)
 	ret0, _ := ret[0].(types.FeatureList)
 	return ret0
 }
 
 // ListCachedFeature indicates an expected call of ListCachedFeature.
-func (mr *MockCacheStoreMockRecorder) ListCachedFeature(ctx, fullNames interface{}) *gomock.Call {
+func (mr *MockCacheStoreMockRecorder) ListCachedFeature(ctx, opt interface{}) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListCachedFeature", reflect.TypeOf((*MockCacheStore)(nil).ListCachedFeature), ctx, fullNames)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListCachedFeature", reflect.TypeOf((*MockCacheStore)(nil).ListCachedFeature), ctx, opt)
 }
 
 // Refresh mocks base method.
diff --git a/internal/database/metadata/store.go b/internal/database/metadata/store.go
index b22badb8..af86fba7 100644
--- a/internal/database/metadata/store.go
+++ b/internal/database/metadata/store.go
@@ -16,7 +16,7 @@ type Store interface {
 	io.Closer
 }
 
-// The interface defines the methods that a database backend store must implement.
+// DBStore defines the methods that a database backend store must implement.
 type DBStore interface {
 	// entity
 	CreateEntity(ctx context.Context, opt CreateEntityOpt) (int, error)
@@ -50,11 +50,11 @@ type DBStore interface {
 	WithTransaction(ctx context.Context, fn func(context.Context, DBStore) error) error
 }
 
-// The interface defines methods that a memory backend store must implement.
+// CacheStore defines methods that a memory backend store must implement.
 type CacheStore interface {
-	ListCachedFeature(ctx context.Context, fullNames *[]string) types.FeatureList
+	ListCachedFeature(ctx context.Context, opt ListCachedFeatureOpt) types.FeatureList
 
-	// Pull data from database and update cache.
+	// Refresh pulls data from database and update cache.
 	Refresh() error
 }
 
diff --git a/internal/database/metadata/types.go b/internal/database/metadata/types.go
index 9bcc0478..c9ab48e7 100644
--- a/internal/database/metadata/types.go
+++ b/internal/database/metadata/types.go
@@ -65,3 +65,9 @@ type ListFeatureOpt struct {
 	GroupID    *int
 	FeatureIDs *[]int
 }
+
+type ListCachedFeatureOpt struct {
+	FullNames *[]string
+	GroupName *string
+	GroupID   *int
+}
diff --git a/pkg/oomstore/online_query.go b/pkg/oomstore/online_query.go
index 88fd0d63..d38b6860 100644
--- a/pkg/oomstore/online_query.go
+++ b/pkg/oomstore/online_query.go
@@ -3,6 +3,8 @@ package oomstore
 import (
 	"context"
 
+	"github.com/oom-ai/oomstore/internal/database/metadata"
+
 	"github.com/oom-ai/oomstore/pkg/errdefs"
 
 	"github.com/oom-ai/oomstore/internal/database/dbutil"
@@ -21,7 +23,9 @@ func (s *OomStore) OnlineGet(ctx context.Context, opt types.OnlineGetOpt) (*type
 		FeatureNames:    opt.FeatureNames,
 		FeatureValueMap: make(map[string]interface{}),
 	}
-	features := s.metadata.ListCachedFeature(ctx, &opt.FeatureNames)
+	features := s.metadata.ListCachedFeature(ctx, metadata.ListCachedFeatureOpt{
+		FullNames: &opt.FeatureNames,
+	})
 	if len(features) == 0 {
 		return &rs, nil
 	}
@@ -63,7 +67,9 @@ func (s *OomStore) OnlineMultiGet(ctx context.Context, opt types.OnlineMultiGetO
 		return nil, err
 	}
 	result := make(map[string]*types.FeatureValues)
-	features := s.metadata.ListCachedFeature(ctx, &opt.FeatureNames)
+	features := s.metadata.ListCachedFeature(ctx, metadata.ListCachedFeatureOpt{
+		FullNames: &opt.FeatureNames,
+	})
 	if len(features) == 0 {
 		return result, nil
 	}
diff --git a/pkg/oomstore/push.go b/pkg/oomstore/push.go
index 6789758d..bceb9462 100644
--- a/pkg/oomstore/push.go
+++ b/pkg/oomstore/push.go
@@ -4,6 +4,8 @@ import (
 	"context"
 	"time"
 
+	"github.com/oom-ai/oomstore/internal/database/metadata"
+
 	"github.com/oom-ai/oomstore/pkg/errdefs"
 
 	"github.com/oom-ai/oomstore/internal/database/online"
@@ -12,12 +14,9 @@ import (
 
 // Push inserts stream feature values to online store and offline store
 func (s *OomStore) Push(ctx context.Context, opt types.PushOpt) error {
-	features, err := s.ListFeature(ctx, types.ListFeatureOpt{
+	features := s.metadata.ListCachedFeature(ctx, metadata.ListCachedFeatureOpt{
 		GroupName: &opt.GroupName,
 	})
-	if err != nil {
-		return err
-	}
 	var featureNames []string
 	for name := range opt.FeatureValues {
 		featureNames = append(featureNames, name)
@@ -33,7 +32,7 @@ func (s *OomStore) Push(ctx context.Context, opt types.PushOpt) error {
 		values = append(values, opt.FeatureValues[f.Name])
 	}
 
-	if err = s.online.Push(ctx, online.PushOpt{
+	if err := s.online.Push(ctx, online.PushOpt{
 		Entity:        entity,
 		EntityKey:     opt.EntityKey,
 		GroupID:       group.ID,


Test patch:
diff --git a/internal/database/metadata/informer/informer_test.go b/internal/database/metadata/informer/informer_test.go
index 5bf54421..46733200 100644
--- a/internal/database/metadata/informer/informer_test.go
+++ b/internal/database/metadata/informer/informer_test.go
@@ -5,6 +5,8 @@ import (
 	"testing"
 	"time"
 
+	"github.com/oom-ai/oomstore/internal/database/metadata"
+
 	"github.com/oom-ai/oomstore/internal/database/metadata/informer"
 	"github.com/oom-ai/oomstore/pkg/oomstore/types"
 	"github.com/oom-ai/oomstore/pkg/oomstore/util"
@@ -50,7 +52,9 @@ func TestInformer(t *testing.T) {
 	defer informer.Close()
 
 	name := util.ComposeFullFeatureName("group", "price")
-	features := informer.ListCachedFeature(ctx, &[]string{name})
+	features := informer.ListCachedFeature(ctx, metadata.ListCachedFeatureOpt{
+		FullNames: &[]string{name},
+	})
 	require.Equal(t, 1, len(features))
 	feature := features[0]
 
@@ -62,18 +66,25 @@ func TestInformer(t *testing.T) {
 	assert.Equal(t, 100, feature.Group.ID)
 	assert.Equal(t, 1, feature.Group.Entity.ID)
 	assert.Equal(t, "entity", feature.Group.Entity.Name)
+
+	groupName := "group"
+	featuresByGroup := informer.ListCachedFeature(ctx, metadata.ListCachedFeatureOpt{
+		GroupName: &groupName,
+	})
+	assert.Equal(t, 1, len(featuresByGroup))
+	assert.Equal(t, feature, featuresByGroup[0])
 }
 
 func TestInformerDeepCopy(t *testing.T) {
 	ctx, informer := prepareInformer(t)
 	defer informer.Close()
 
-	features := informer.ListCachedFeature(ctx, nil)
+	features := informer.ListCachedFeature(ctx, metadata.ListCachedFeatureOpt{})
 	require.Equal(t, 1, len(features))
 	// changing this entity should not change the internal state of the informer
 	features[0].Name = "new_price"
 
-	features = informer.ListCachedFeature(ctx, nil)
+	features = informer.ListCachedFeature(ctx, metadata.ListCachedFeatureOpt{})
 	require.Equal(t, 1, len(features))
 	assert.Equal(t, 1, features[0].ID)
 	assert.Equal(t, 100, features[0].GroupID)
diff --git a/internal/database/metadata/test_impl/feature.go b/internal/database/metadata/test_impl/feature.go
index b1ad9252..7c6a4b8b 100644
--- a/internal/database/metadata/test_impl/feature.go
+++ b/internal/database/metadata/test_impl/feature.go
@@ -176,7 +176,7 @@ func TestListCachedFeature(t *testing.T, prepareStore PrepareStoreFn, destroySto
 	_, groupID := prepareEntityAndGroup(t, ctx, store)
 
 	// case 1: no feature to list
-	features := store.ListCachedFeature(ctx, nil)
+	features := store.ListCachedFeature(ctx, metadata.ListCachedFeatureOpt{})
 	assert.Equal(t, 0, features.Len())
 
 	_, err := store.CreateFeature(ctx, metadata.CreateFeatureOpt{
@@ -189,11 +189,13 @@ func TestListCachedFeature(t *testing.T, prepareStore PrepareStoreFn, destroySto
 	require.NoError(t, store.Refresh())
 
 	// case 2: no condition, list all features
-	features = store.ListCachedFeature(ctx, nil)
+	features = store.ListCachedFeature(ctx, metadata.ListCachedFeatureOpt{})
 	assert.Equal(t, 1, features.Len())
 
 	// case 8: list features by FeatureNames
-	features = store.ListCachedFeature(ctx, &[]string{"device_info.phone"})
+	features = store.ListCachedFeature(ctx, metadata.ListCachedFeatureOpt{
+		FullNames: &[]string{"device_info.phone"},
+	})
 	assert.Equal(t, 1, len(features))
 }
 
diff --git a/oomcli/test/test_export.sh b/oomcli/test/test_export.sh
index e540126e..c5633544 100755
--- a/oomcli/test/test_export.sh
+++ b/oomcli/test/test_export.sh
@@ -52,15 +52,6 @@ oomcli_export_batch() {
     assert_eq "$case" "$(sort <<< "$expected")" "$(sort <<< "$actual")"
 }
 
-oomcli_export_batch_and_stream_before_import() {
-    case="export batch and stream feature before import"
-    expected='user,student.name,student.gender,student.age,user-click.last_5_click_posts,user-click.number_of_user_starred_posts
-1,,,,"2,3,4,5,6",11
-'
-    actual=$(oomcli export --feature student.name,student.gender,student.age,user-click.last_5_click_posts,user-click.number_of_user_starred_posts --unix-milli $t2 -o csv)
-    assert_eq "$case" "$(sort <<< "$expected")" "$(sort <<< "$actual")"
-}
-
 oomcli_export_batch_and_stream() {
     case="export batch and stream feature"
     expected='user,student.name,student.gender,student.age,user-click.last_5_click_posts,user-click.number_of_user_starred_posts
@@ -80,7 +71,6 @@ main() {
     oomcli_export_push_feature
     oomcli_export_update_feature
     oomcli_export_batch
-    oomcli_export_batch_and_stream_before_import
     oomcli_export_batch_and_stream
 }
 
diff --git a/pkg/oomstore/online_query_test.go b/pkg/oomstore/online_query_test.go
index 534f2e5c..b47c5cf6 100644
--- a/pkg/oomstore/online_query_test.go
+++ b/pkg/oomstore/online_query_test.go
@@ -5,6 +5,8 @@ import (
 	"fmt"
 	"testing"
 
+	"github.com/oom-ai/oomstore/internal/database/metadata"
+
 	"github.com/golang/mock/gomock"
 	"github.com/oom-ai/oomstore/internal/database/dbutil"
 	"github.com/oom-ai/oomstore/internal/database/metadata/mock_metadata"
@@ -80,7 +82,9 @@ func TestOnlineGet(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.description, func(t *testing.T) {
-			metadataStore.EXPECT().ListCachedFeature(gomock.Any(), &tc.opt.FeatureNames).Return(tc.features)
+			metadataStore.EXPECT().ListCachedFeature(gomock.Any(), metadata.ListCachedFeatureOpt{
+				FullNames: &tc.opt.FeatureNames,
+			}).Return(tc.features)
 			if tc.entityName != nil {
 				onlineStore.EXPECT().Get(gomock.Any(), gomock.Any()).Return(dbutil.RowMap{
 					"price": int64(100),
@@ -174,7 +178,9 @@ func TestOnlineMultiGet(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.description, func(t *testing.T) {
-			metadataStore.EXPECT().ListCachedFeature(gomock.Any(), &tc.opt.FeatureNames).Return(tc.features)
+			metadataStore.EXPECT().ListCachedFeature(gomock.Any(), metadata.ListCachedFeatureOpt{
+				FullNames: &tc.opt.FeatureNames,
+			}).Return(tc.features)
 			if tc.entityName != nil {
 				onlineStore.EXPECT().MultiGet(gomock.Any(), gomock.Any()).Return(map[string]dbutil.RowMap{
 					"1234": {