Make batch_source optional in PushSource (#5440)

snehsuresh · snehsuresh · commit b76ee666f6ba · 2025-06-14T19:02:08.000-05:00
Signed-off-by: snehsuresh &lt;snehsuresh02@gmail.com&gt;
diff --git a/sdk/python/feast/data_source.py b/sdk/python/feast/data_source.py
@@ -764,13 +764,13 @@ class PushSource(DataSource):
 
     # TODO(adchia): consider adding schema here in case where Feast manages pushing events to the offline store
     # TODO(adchia): consider a "mode" to support pushing raw vs transformed events
-    batch_source: DataSource
+    batch_source: Optional[DataSource] = None
 
     def __init__(
         self,
         *,
         name: str,
-        batch_source: DataSource,
+        batch_source: Optional[DataSource] = None,
         description: Optional[str] = "",
         tags: Optional[Dict[str, str]] = None,
         owner: Optional[str] = "",
@@ -815,8 +815,12 @@ def get_table_column_names_and_types(
 
     @staticmethod
     def from_proto(data_source: DataSourceProto):
-        assert data_source.HasField("batch_source")
-        batch_source = DataSource.from_proto(data_source.batch_source)
+        # assert data_source.HasField("batch_source")
+        batch_source = (
+            DataSource.from_proto(data_source.batch_source)
+            if data_source.HasField("batch_source")
+            else None
+        )
 
         return PushSource(
             name=data_source.name,
@@ -827,19 +831,19 @@ def from_proto(data_source: DataSourceProto):
         )
 
     def to_proto(self) -> DataSourceProto:
-        batch_source_proto = None
-        if self.batch_source:
-            batch_source_proto = self.batch_source.to_proto()
+        # batch_source_proto = None
 
         data_source_proto = DataSourceProto(
             name=self.name,
             type=DataSourceProto.PUSH_SOURCE,
             description=self.description,
             tags=self.tags,
             owner=self.owner,
-            batch_source=batch_source_proto,
         )
 
+        if self.batch_source:
+            data_source_proto.batch_source.MergeFrom(self.batch_source.to_proto())
+
         return data_source_proto
 
     def get_table_query_string(self) -> str:
diff --git a/sdk/python/tests/unit/test_data_sources.py b/sdk/python/tests/unit/test_data_sources.py
@@ -30,6 +30,22 @@ def test_push_with_batch():
     assert push_source.batch_source.name == push_source_unproto.batch_source.name
 
 
+def test_push_source_without_batch_source():
+    # Create PushSource with no batch_source
+    push_source = PushSource(name="test_push_source")
+
+    # Convert to proto
+    push_source_proto = push_source.to_proto()
+
+    # Assert batch_source is not present in proto
+    assert not push_source_proto.HasField("batch_source")
+
+    # Deserialize and check again
+    push_source_unproto = PushSource.from_proto(push_source_proto)
+    assert push_source_unproto.batch_source is None
+    assert push_source_unproto.name == "test_push_source"
+
+
 def test_request_source_primitive_type_to_proto():
     schema = [
         Field(name="f1", dtype=Float32),
