fix error to search with key instead of name

Author: msistla96

go/internal/feast/featurestore.go

@@ -291,13 +291,15 @@ func (fs *FeatureStore) ListEntities(hideDummyEntity bool) ([]*model.Entity, err
 	return entities, nil
 }
 
-func (fs *FeatureStore) ListEntity(entityName string) (*model.Entity, error) {
+func (fs *FeatureStore) ListEntityByJoinKey(entitykey string) (*model.Entity, error) {
 
-	entity, err := fs.registry.GetEntity(fs.config.Project, entityName)
-	if err != nil {
-		return entity, err
+	entities, _ := fs.ListEntities(false)
+	for _, entity := range entities {
+		if entity.JoinKey == entitykey {
+			return entity, nil
+		}
 	}
-	return entity, nil
+	return nil, fmt.Errorf("entity with join key %s not found", entitykey)
 }
 
 func (fs *FeatureStore) ListOnDemandFeatureViews() ([]*model.OnDemandFeatureView, error) {

go/internal/feast/server/http_server.go

@@ -171,10 +171,10 @@ func logWithSpanContext(span tracer.Span) zerolog.Logger {
 
 /*
 *
-This function is to typecast entity values unmarshalled from JSON to the actual entity type specified in the entity definition.
-This is because of a bug when using an entity defined with ValueType_INT32
+This function fixes the bug when the entity join key type is INT32.
+It typecasts the entity values to int32, as unmarshalled JSON values only return int64.
 */
-func typecastToEntityIntType(val *repeatedValue, entityType prototypes.ValueType_Enum) {
+func typecastToInt32Type(val *repeatedValue, entityType prototypes.ValueType_Enum) {
 	if val.int64Val != nil {
 		if entityType == prototypes.ValueType_INT32 {
 			for _, v := range val.int64Val {
@@ -183,14 +183,6 @@ func typecastToEntityIntType(val *repeatedValue, entityType prototypes.ValueType
 			val.int64Val = nil
 		}
 	}
-	if val.int32Val != nil {
-		if entityType == prototypes.ValueType_INT64 {
-			for _, v := range val.int32Val {
-				val.int64Val = append(val.int64Val, int64(v))
-			}
-			val.int32Val = nil
-		}
-	}
 }
 
 func (s *httpServer) getOnlineFeatures(w http.ResponseWriter, r *http.Request) {
@@ -238,11 +230,15 @@ func (s *httpServer) getOnlineFeatures(w http.ResponseWriter, r *http.Request) {
 	}
 	entitiesProto := make(map[string]*prototypes.RepeatedValue)
 	for key, value := range request.Entities {
-		var entity, err = s.fs.ListEntity(key)
+		var entity, err = s.fs.ListEntityByJoinKey(key)
 		if err != nil {
-
+			logSpanContext.Error().Err(err).Msg("Error getting entity from registry")
+			writeJSONError(w, fmt.Errorf("Error getting entity from registry: %+v", err), http.StatusInternalServerError)
+			return
+		}
+		if value.int64Val != nil {
+			typecastToInt32Type(&value, entity.Dtype)
 		}
-		typecastToEntityIntType(&value, entity.Dtype)
 		entitiesProto[key] = value.ToProto()
 	}
 	requestContextProto := make(map[string]*prototypes.RepeatedValue)