fix: Fix registry proto (#2435)

* Fix registry Signed-off-by: Kevin Zhang <kzhang@tecton.ai> * Clean up Signed-off-by: Kevin Zhang <kzhang@tecton.ai> * Clean up Signed-off-by: Kevin Zhang <kzhang@tecton.ai> * Fix issues Signed-off-by: Kevin Zhang <kzhang@tecton.ai> * Remove the allowCache flag Signed-off-by: Kevin Zhang <kzhang@tecton.ai> * Address review issues Signed-off-by: Kevin Zhang <kzhang@tecton.ai> * Fix dumb mistake Signed-off-by: Kevin Zhang <kzhang@tecton.ai>
feast-dev · Mar 23, 2022 · ea6a9b2 · ea6a9b2
1 parent ed5e928
commit ea6a9b2
Show file tree

Hide file tree

Showing 2 changed files with 128 additions and 92 deletions.
diff --git a/go/internal/feast/featurestore.go b/go/internal/feast/featurestore.go
@@ -4,14 +4,15 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"sort"
+	"strings"
+
 	"github.com/feast-dev/feast/go/protos/feast/serving"
 	"github.com/feast-dev/feast/go/protos/feast/types"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/types/known/durationpb"
 	"google.golang.org/protobuf/types/known/timestamppb"
-	"sort"
-	"strings"
 )
 
 type FeatureStore struct {
@@ -92,7 +93,7 @@ func (fs *FeatureStore) GetOnlineFeatures(ctx context.Context, request *serving.
 		return nil, err
 	}
 
-	fvs, requestedFeatureViews, requestedRequestFeatureViews, requestedOnDemandFeatureViews, err := fs.getFeatureViewsToUse(features, true, false)
+	fvs, requestedFeatureViews, requestedRequestFeatureViews, requestedOnDemandFeatureViews, err := fs.getFeatureViewsToUse(features, false)
 
 	if len(requestedRequestFeatureViews)+len(requestedOnDemandFeatureViews) > 0 {
 		return nil, status.Errorf(codes.InvalidArgument, "on demand feature views are currently not supported")
@@ -253,22 +254,31 @@ func (fs *FeatureStore) getFeatureRefs(features *Features) ([]string, error) {
 		retrieving all feature views. Similar argument to FeatureService applies.
 
 */
-func (fs *FeatureStore) getFeatureViewsToUse(features *Features, allowCache, hideDummyEntity bool) (map[string]*FeatureView, map[*FeatureView][]string, []*RequestFeatureView, []*OnDemandFeatureView, error) {
+func (fs *FeatureStore) getFeatureViewsToUse(features *Features, hideDummyEntity bool) (map[string]*FeatureView, map[*FeatureView][]string, []*RequestFeatureView, []*OnDemandFeatureView, error) {
 	fvs := make(map[string]*FeatureView)
 	requestFvs := make(map[string]*RequestFeatureView)
 	odFvs := make(map[string]*OnDemandFeatureView)
 
-	featureViews := fs.listFeatureViews(allowCache, hideDummyEntity)
+	featureViews, err := fs.listFeatureViews(hideDummyEntity)
+	if err != nil {
+		return nil, nil, nil, nil, err
+	}
 	for _, featureView := range featureViews {
 		fvs[featureView.base.name] = featureView
 	}
 
-	requestFeatureViews := fs.registry.listRequestFeatureViews(fs.config.Project)
+	requestFeatureViews, err := fs.registry.listRequestFeatureViews(fs.config.Project)
+	if err != nil {
+		return nil, nil, nil, nil, err
+	}
 	for _, requestFeatureView := range requestFeatureViews {
 		requestFvs[requestFeatureView.base.name] = requestFeatureView
 	}
 
-	onDemandFeatureViews := fs.registry.listOnDemandFeatureViews(fs.config.Project)
+	onDemandFeatureViews, err := fs.registry.listOnDemandFeatureViews(fs.config.Project)
+	if err != nil {
+		return nil, nil, nil, nil, err
+	}
 	for _, onDemandFeatureView := range onDemandFeatureViews {
 		odFvs[onDemandFeatureView.base.name] = onDemandFeatureView
 	}
@@ -347,7 +357,10 @@ func (fs *FeatureStore) getEntityMaps(requestedFeatureViews map[*FeatureView][]s
 	var joinKeyMap map[string]string
 	var featureView *FeatureView
 
-	entities := fs.listEntities(true, false)
+	entities, err := fs.listEntities(false)
+	if err != nil {
+		return nil, err
+	}
 
 	for _, entity := range entities {
 		entityNameToJoinKeyMap[entity.name] = entity.joinKey
@@ -636,26 +649,31 @@ func (fs *FeatureStore) dropUnneededColumns(onlineFeaturesResponse *serving.GetO
 	}
 }
 
-func (fs *FeatureStore) listFeatureViews(allowCache, hideDummyEntity bool) []*FeatureView {
-	featureViews := fs.registry.listFeatureViews(fs.config.Project)
-	for _, featureView := range featureViews {
-		if _, ok := featureView.entities[DUMMY_ENTITY_NAME]; ok && hideDummyEntity {
-			featureView.entities = make(map[string]struct{})
-		}
+func (fs *FeatureStore) listFeatureViews(hideDummyEntity bool) ([]*FeatureView, error) {
+	featureViews, err := fs.registry.listFeatureViews(fs.config.Project)
+	if err != nil {
+		return featureViews, err
 	}
-	return featureViews
+	return featureViews, nil
+}
+
+func (fs *FeatureStore) listRequestFeatureViews() ([]*RequestFeatureView, error) {
+	return fs.registry.listRequestFeatureViews(fs.config.Project)
 }
 
-func (fs *FeatureStore) listEntities(allowCache, hideDummyEntity bool) []*Entity {
+func (fs *FeatureStore) listEntities(hideDummyEntity bool) ([]*Entity, error) {
 
-	allEntities := fs.registry.listEntities(fs.config.Project)
+	allEntities, err := fs.registry.listEntities(fs.config.Project)
+	if err != nil {
+		return allEntities, err
+	}
 	entities := make([]*Entity, 0)
 	for _, entity := range allEntities {
 		if entity.name != DUMMY_ENTITY_NAME || !hideDummyEntity {
 			entities = append(entities, entity)
 		}
 	}
-	return entities
+	return entities, nil
 }
 
 func (fs *FeatureStore) getFvEntityValues(fv *FeatureView,
@@ -820,7 +838,7 @@ func (fs *FeatureStore) groupFeatureRefs(requestedFeatureViews map[*FeatureView]
 	return fvFeatures, nil
 }
 
-func (fs *FeatureStore) getFeatureView(project, featureViewName string, allowCache, hideDummyEntity bool) (*FeatureView, error) {
+func (fs *FeatureStore) getFeatureView(project, featureViewName string, hideDummyEntity bool) (*FeatureView, error) {
 	fv, err := fs.registry.getFeatureView(fs.config.Project, featureViewName)
 	if err != nil {
 		return nil, err