Revert "CR"

achals · achals · commit 852c27efd9c7 · 2022-03-22T14:08:19.000-07:00
This reverts commit 989d6c55752b54d2dce8a65271d6963e13b99d1b.

Signed-off-by: Achal Shah &lt;achals@gmail.com&gt;
diff --git a/sdk/python/feast/data_source.py b/sdk/python/feast/data_source.py
@@ -134,6 +134,18 @@ def to_proto(self) -> DataSourceProto.KinesisOptions:
         return kinesis_options_proto
 
 
+_DATA_SOURCE_OPTIONS = {
+    DataSourceProto.SourceType.BATCH_FILE: "feast.infra.offline_stores.file_source.FileSource",
+    DataSourceProto.SourceType.BATCH_BIGQUERY: "feast.infra.offline_stores.bigquery_source.BigQuerySource",
+    DataSourceProto.SourceType.BATCH_REDSHIFT: "feast.infra.offline_stores.redshift_source.RedshiftSource",
+    DataSourceProto.SourceType.BATCH_SNOWFLAKE: "feast.infra.offline_stores.snowflake_source.SnowflakeSource",
+    DataSourceProto.SourceType.STREAM_KAFKA: "feast.data_source.KafkaSource",
+    DataSourceProto.SourceType.STREAM_KINESIS: "feast.data_source.KinesisSource",
+    DataSourceProto.SourceType.REQUEST_SOURCE: "feast.data_source.RequestDataSource",
+    DataSourceProto.SourceType.PUSH_SOURCE: "feast.data_source.PushSource",
+}
+
+
 class DataSource(ABC):
     """
     DataSource that can be used to source features.
@@ -210,13 +222,20 @@ def from_proto(data_source: DataSourceProto) -> Any:
         Raises:
             ValueError: The type of DataSource could not be identified.
         """
+        data_source_type = data_source.type
+        if not data_source_type or (
+            data_source_type
+            not in list(_DATA_SOURCE_OPTIONS.keys())
+            + [DataSourceProto.SourceType.CUSTOM_SOURCE]
+        ):
+            raise ValueError("Could not identify the source type being added.")
 
-        if data_source.data_source_class_type:
+        if data_source_type == DataSourceProto.SourceType.CUSTOM_SOURCE:
             cls = get_data_source_class_from_type(data_source.data_source_class_type)
             return cls.from_proto(data_source)
-        raise ValueError(
-            f"Could not identify the type for data source: {data_source.name}."
-        )
+
+        cls = get_data_source_class_from_type(_DATA_SOURCE_OPTIONS[data_source_type])
+        return cls.from_proto(data_source)
 
     @abstractmethod
     def to_proto(self) -> DataSourceProto:
diff --git a/sdk/python/feast/registry.py b/sdk/python/feast/registry.py
@@ -320,9 +320,6 @@ def apply_data_source(
                 del registry.data_sources[idx]
         data_source_proto = data_source.to_proto()
         data_source_proto.project = project
-        data_source_proto.data_source_class_type = (
-            f"{data_source.__class__.__module__}.{data_source.__class__.__name__}"
-        )
         registry.data_sources.append(data_source_proto)
         if commit:
             self.commit()
