fix: Added missing type (#4315)

Added missing type Signed-off-by: Daniele Martinoli <[email protected]>
feast-dev · Jun 25, 2024 · 86af60a · 86af60a
1 parent 21deec8
commit 86af60a
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/sdk/python/feast/diff/registry_diff.py b/sdk/python/feast/diff/registry_diff.py
@@ -20,6 +20,7 @@
     OnDemandFeatureView as OnDemandFeatureViewProto,
 )
 from feast.protos.feast.core.OnDemandFeatureView_pb2 import OnDemandFeatureViewSpec
+from feast.protos.feast.core.SavedDataset_pb2 import SavedDataset as SavedDatasetProto
 from feast.protos.feast.core.StreamFeatureView_pb2 import (
     StreamFeatureView as StreamFeatureViewProto,
 )
@@ -109,6 +110,7 @@ def tag_objects_for_keep_delete_update_add(
     OnDemandFeatureViewProto,
     StreamFeatureViewProto,
     ValidationReferenceProto,
+    SavedDatasetProto,
 )
 
 

diff --git a/sdk/python/feast/feast_object.py b/sdk/python/feast/feast_object.py
@@ -11,11 +11,12 @@
 from .protos.feast.core.FeatureService_pb2 import FeatureServiceSpec
 from .protos.feast.core.FeatureView_pb2 import FeatureViewSpec
 from .protos.feast.core.OnDemandFeatureView_pb2 import OnDemandFeatureViewSpec
+from .protos.feast.core.SavedDataset_pb2 import SavedDatasetSpec
 from .protos.feast.core.StreamFeatureView_pb2 import StreamFeatureViewSpec
 from .protos.feast.core.ValidationProfile_pb2 import (
     ValidationReference as ValidationReferenceProto,
 )
-from .saved_dataset import ValidationReference
+from .saved_dataset import SavedDataset, ValidationReference
 from .stream_feature_view import StreamFeatureView
 
 # Convenience type representing all Feast objects
@@ -28,6 +29,7 @@
     FeatureService,
     DataSource,
     ValidationReference,
+    SavedDataset,
 ]
 
 FeastObjectSpecProto = Union[
@@ -38,4 +40,5 @@
     FeatureServiceSpec,
     DataSourceProto,
     ValidationReferenceProto,
+    SavedDatasetSpec,
 ]