diff --git a/sdk/python/feast/feature_view.py b/sdk/python/feast/feature_view.py index e26759ba92..67f9662d31 100644 --- a/sdk/python/feast/feature_view.py +++ b/sdk/python/feast/feature_view.py @@ -17,6 +17,7 @@ from typing import Dict, List, Optional, Tuple, Type from google.protobuf.duration_pb2 import Duration +from google.protobuf.message import Message from typeguard import typechecked from feast import utils @@ -274,7 +275,7 @@ def ensure_valid(self): raise ValueError("Feature view has no entities.") @property - def proto_class(self) -> Type[FeatureViewProto]: + def proto_class(self) -> Type[Message]: return FeatureViewProto def with_join_key_map(self, join_key_map: Dict[str, str]): diff --git a/sdk/python/feast/stream_feature_view.py b/sdk/python/feast/stream_feature_view.py index d3a2164788..6a204d6813 100644 --- a/sdk/python/feast/stream_feature_view.py +++ b/sdk/python/feast/stream_feature_view.py @@ -3,9 +3,10 @@ import warnings from datetime import datetime, timedelta from types import FunctionType -from typing import Dict, List, Optional, Tuple, Union +from typing import Dict, List, Optional, Tuple, Type, Union import dill +from google.protobuf.message import Message from typeguard import typechecked from feast import flags_helper, utils @@ -298,6 +299,10 @@ def __copy__(self): fv.projection = copy.copy(self.projection) return fv + @property + def proto_class(self) -> Type[Message]: + return StreamFeatureViewProto + def stream_feature_view( *, diff --git a/sdk/python/tests/unit/test_feature_views.py b/sdk/python/tests/unit/test_feature_views.py index 379396e5c6..afef332d37 100644 --- a/sdk/python/tests/unit/test_feature_views.py +++ b/sdk/python/tests/unit/test_feature_views.py @@ -10,6 +10,9 @@ from feast.feature_view import FeatureView from feast.field import Field from feast.infra.offline_stores.file_source import FileSource +from feast.protos.feast.core.StreamFeatureView_pb2 import ( + StreamFeatureView as StreamFeatureViewProto, +) from feast.protos.feast.types.Value_pb2 import ValueType from feast.stream_feature_view import StreamFeatureView, stream_feature_view from feast.types import Float32 @@ -277,3 +280,22 @@ def test_hash(): def test_field_types(): with pytest.raises(TypeError): Field(name="name", dtype=ValueType.INT32) + + +def test_stream_feature_view_proto_type(): + stream_source = KafkaSource( + name="kafka", + timestamp_field="event_timestamp", + kafka_bootstrap_servers="", + message_format=AvroFormat(""), + topic="topic", + batch_source=FileSource(path="some path"), + ) + sfv = StreamFeatureView( + name="test stream featureview proto class", + entities=[], + ttl=timedelta(days=30), + source=stream_source, + aggregations=[], + ) + assert sfv.proto_class is StreamFeatureViewProto