Skip to content

Commit 852c27e

Browse files
committed
Revert "CR"
This reverts commit 989d6c55752b54d2dce8a65271d6963e13b99d1b. Signed-off-by: Achal Shah <achals@gmail.com>
1 parent 772cc5c commit 852c27e

File tree

2 files changed

+23
-7
lines changed

2 files changed

+23
-7
lines changed

sdk/python/feast/data_source.py

+23-4
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,18 @@ def to_proto(self) -> DataSourceProto.KinesisOptions:
134134
return kinesis_options_proto
135135

136136

137+
_DATA_SOURCE_OPTIONS = {
138+
DataSourceProto.SourceType.BATCH_FILE: "feast.infra.offline_stores.file_source.FileSource",
139+
DataSourceProto.SourceType.BATCH_BIGQUERY: "feast.infra.offline_stores.bigquery_source.BigQuerySource",
140+
DataSourceProto.SourceType.BATCH_REDSHIFT: "feast.infra.offline_stores.redshift_source.RedshiftSource",
141+
DataSourceProto.SourceType.BATCH_SNOWFLAKE: "feast.infra.offline_stores.snowflake_source.SnowflakeSource",
142+
DataSourceProto.SourceType.STREAM_KAFKA: "feast.data_source.KafkaSource",
143+
DataSourceProto.SourceType.STREAM_KINESIS: "feast.data_source.KinesisSource",
144+
DataSourceProto.SourceType.REQUEST_SOURCE: "feast.data_source.RequestDataSource",
145+
DataSourceProto.SourceType.PUSH_SOURCE: "feast.data_source.PushSource",
146+
}
147+
148+
137149
class DataSource(ABC):
138150
"""
139151
DataSource that can be used to source features.
@@ -210,13 +222,20 @@ def from_proto(data_source: DataSourceProto) -> Any:
210222
Raises:
211223
ValueError: The type of DataSource could not be identified.
212224
"""
225+
data_source_type = data_source.type
226+
if not data_source_type or (
227+
data_source_type
228+
not in list(_DATA_SOURCE_OPTIONS.keys())
229+
+ [DataSourceProto.SourceType.CUSTOM_SOURCE]
230+
):
231+
raise ValueError("Could not identify the source type being added.")
213232

214-
if data_source.data_source_class_type:
233+
if data_source_type == DataSourceProto.SourceType.CUSTOM_SOURCE:
215234
cls = get_data_source_class_from_type(data_source.data_source_class_type)
216235
return cls.from_proto(data_source)
217-
raise ValueError(
218-
f"Could not identify the type for data source: {data_source.name}."
219-
)
236+
237+
cls = get_data_source_class_from_type(_DATA_SOURCE_OPTIONS[data_source_type])
238+
return cls.from_proto(data_source)
220239

221240
@abstractmethod
222241
def to_proto(self) -> DataSourceProto:

sdk/python/feast/registry.py

-3
Original file line numberDiff line numberDiff line change
@@ -320,9 +320,6 @@ def apply_data_source(
320320
del registry.data_sources[idx]
321321
data_source_proto = data_source.to_proto()
322322
data_source_proto.project = project
323-
data_source_proto.data_source_class_type = (
324-
f"{data_source.__class__.__module__}.{data_source.__class__.__name__}"
325-
)
326323
registry.data_sources.append(data_source_proto)
327324
if commit:
328325
self.commit()

0 commit comments

Comments
 (0)