From 566d6456b3b16942485e81eb3fa9ce1da64d336c Mon Sep 17 00:00:00 2001
From: Daniele Martinoli <86618610+dmartinol@users.noreply.github.com>
Date: Tue, 25 Jun 2024 10:08:40 +0200
Subject: [PATCH] Added missing type

Signed-off-by: Daniele Martinoli <86618610+dmartinol@users.noreply.github.com>
---
 sdk/python/feast/diff/registry_diff.py | 2 ++
 sdk/python/feast/feast_object.py       | 5 ++++-
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/sdk/python/feast/diff/registry_diff.py b/sdk/python/feast/diff/registry_diff.py
index b608757496f..9236b087d4b 100644
--- 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
index 2d06d8d669d..d9505dcb9f0 100644
--- 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,
 ]