feat: Sourcer changes: * deprecate partitions_handler, * Add support for active_partitions_handler and total_partitions, * Maintain backward compatibility

vaibhavtiwari33 · vaibhavtiwari33 · commit 9ff135646fb2 · 2026-03-26T19:25:07.000-04:00
Signed-off-by: Vaibhav Tiwari &lt;vaibhav.tiwari33@gmail.com&gt;
diff --git a/packages/pynumaflow/pynumaflow/proto/accumulator/accumulator_pb2_grpc.py b/packages/pynumaflow/pynumaflow/proto/accumulator/accumulator_pb2_grpc.py
@@ -6,7 +6,7 @@
 from google.protobuf import empty_pb2 as google_dot_protobuf_dot_empty__pb2
 from pynumaflow.proto.accumulator import accumulator_pb2 as pynumaflow_dot_proto_dot_accumulator_dot_accumulator__pb2
 
-GRPC_GENERATED_VERSION = '1.75.0'
+GRPC_GENERATED_VERSION = '1.78.0'
 GRPC_VERSION = grpc.__version__
 _version_not_supported = False
 
@@ -19,7 +19,7 @@
 if _version_not_supported:
     raise RuntimeError(
         f'The grpc package installed is at version {GRPC_VERSION},'
-        + f' but the generated code in pynumaflow/proto/accumulator/accumulator_pb2_grpc.py depends on'
+        + ' but the generated code in pynumaflow/proto/accumulator/accumulator_pb2_grpc.py depends on'
         + f' grpcio>={GRPC_GENERATED_VERSION}.'
         + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}'
         + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.'
diff --git a/packages/pynumaflow/pynumaflow/proto/common/metadata_pb2_grpc.py b/packages/pynumaflow/pynumaflow/proto/common/metadata_pb2_grpc.py
@@ -4,7 +4,7 @@
 import warnings
 
 
-GRPC_GENERATED_VERSION = '1.75.0'
+GRPC_GENERATED_VERSION = '1.78.0'
 GRPC_VERSION = grpc.__version__
 _version_not_supported = False
 
@@ -17,7 +17,7 @@
 if _version_not_supported:
     raise RuntimeError(
         f'The grpc package installed is at version {GRPC_VERSION},'
-        + f' but the generated code in pynumaflow/proto/common/metadata_pb2_grpc.py depends on'
+        + ' but the generated code in pynumaflow/proto/common/metadata_pb2_grpc.py depends on'
         + f' grpcio>={GRPC_GENERATED_VERSION}.'
         + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}'
         + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.'
diff --git a/packages/pynumaflow/pynumaflow/proto/mapper/map_pb2_grpc.py b/packages/pynumaflow/pynumaflow/proto/mapper/map_pb2_grpc.py
@@ -6,7 +6,7 @@
 from google.protobuf import empty_pb2 as google_dot_protobuf_dot_empty__pb2
 from pynumaflow.proto.mapper import map_pb2 as pynumaflow_dot_proto_dot_mapper_dot_map__pb2
 
-GRPC_GENERATED_VERSION = '1.75.0'
+GRPC_GENERATED_VERSION = '1.78.0'
 GRPC_VERSION = grpc.__version__
 _version_not_supported = False
 
@@ -19,7 +19,7 @@
 if _version_not_supported:
     raise RuntimeError(
         f'The grpc package installed is at version {GRPC_VERSION},'
-        + f' but the generated code in pynumaflow/proto/mapper/map_pb2_grpc.py depends on'
+        + ' but the generated code in pynumaflow/proto/mapper/map_pb2_grpc.py depends on'
         + f' grpcio>={GRPC_GENERATED_VERSION}.'
         + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}'
         + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.'
diff --git a/packages/pynumaflow/pynumaflow/proto/reducer/reduce_pb2_grpc.py b/packages/pynumaflow/pynumaflow/proto/reducer/reduce_pb2_grpc.py
@@ -6,7 +6,7 @@
 from google.protobuf import empty_pb2 as google_dot_protobuf_dot_empty__pb2
 from pynumaflow.proto.reducer import reduce_pb2 as pynumaflow_dot_proto_dot_reducer_dot_reduce__pb2
 
-GRPC_GENERATED_VERSION = '1.75.0'
+GRPC_GENERATED_VERSION = '1.78.0'
 GRPC_VERSION = grpc.__version__
 _version_not_supported = False
 
@@ -19,7 +19,7 @@
 if _version_not_supported:
     raise RuntimeError(
         f'The grpc package installed is at version {GRPC_VERSION},'
-        + f' but the generated code in pynumaflow/proto/reducer/reduce_pb2_grpc.py depends on'
+        + ' but the generated code in pynumaflow/proto/reducer/reduce_pb2_grpc.py depends on'
         + f' grpcio>={GRPC_GENERATED_VERSION}.'
         + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}'
         + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.'
diff --git a/packages/pynumaflow/pynumaflow/proto/sideinput/sideinput_pb2_grpc.py b/packages/pynumaflow/pynumaflow/proto/sideinput/sideinput_pb2_grpc.py
@@ -6,7 +6,7 @@
 from google.protobuf import empty_pb2 as google_dot_protobuf_dot_empty__pb2
 from pynumaflow.proto.sideinput import sideinput_pb2 as pynumaflow_dot_proto_dot_sideinput_dot_sideinput__pb2
 
-GRPC_GENERATED_VERSION = '1.75.0'
+GRPC_GENERATED_VERSION = '1.78.0'
 GRPC_VERSION = grpc.__version__
 _version_not_supported = False
 
@@ -19,7 +19,7 @@
 if _version_not_supported:
     raise RuntimeError(
         f'The grpc package installed is at version {GRPC_VERSION},'
-        + f' but the generated code in pynumaflow/proto/sideinput/sideinput_pb2_grpc.py depends on'
+        + ' but the generated code in pynumaflow/proto/sideinput/sideinput_pb2_grpc.py depends on'
         + f' grpcio>={GRPC_GENERATED_VERSION}.'
         + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}'
         + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.'
diff --git a/packages/pynumaflow/pynumaflow/proto/sinker/sink_pb2_grpc.py b/packages/pynumaflow/pynumaflow/proto/sinker/sink_pb2_grpc.py
@@ -6,7 +6,7 @@
 from google.protobuf import empty_pb2 as google_dot_protobuf_dot_empty__pb2
 from pynumaflow.proto.sinker import sink_pb2 as pynumaflow_dot_proto_dot_sinker_dot_sink__pb2
 
-GRPC_GENERATED_VERSION = '1.75.0'
+GRPC_GENERATED_VERSION = '1.78.0'
 GRPC_VERSION = grpc.__version__
 _version_not_supported = False
 
@@ -19,7 +19,7 @@
 if _version_not_supported:
     raise RuntimeError(
         f'The grpc package installed is at version {GRPC_VERSION},'
-        + f' but the generated code in pynumaflow/proto/sinker/sink_pb2_grpc.py depends on'
+        + ' but the generated code in pynumaflow/proto/sinker/sink_pb2_grpc.py depends on'
         + f' grpcio>={GRPC_GENERATED_VERSION}.'
         + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}'
         + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.'
diff --git a/packages/pynumaflow/pynumaflow/proto/sourcer/source.proto b/packages/pynumaflow/pynumaflow/proto/sourcer/source.proto
@@ -190,8 +190,10 @@ message PendingResponse {
  */
 message PartitionsResponse {
   message Result {
-    // Required field holding the list of partitions.
+    // Required field holding the list of active partitions.
     repeated int32 partitions = 1;
+    // Total number of partitions in the source.
+    optional int32 total_partitions = 2;
   }
   // Required field holding the result.
   Result result = 1;
diff --git a/packages/pynumaflow/pynumaflow/proto/sourcer/source_pb2.py b/packages/pynumaflow/pynumaflow/proto/sourcer/source_pb2.py
diff --git a/packages/pynumaflow/pynumaflow/proto/sourcer/source_pb2.pyi b/packages/pynumaflow/pynumaflow/proto/sourcer/source_pb2.pyi
@@ -156,10 +156,12 @@ class PendingResponse(_message.Message):
 class PartitionsResponse(_message.Message):
     __slots__ = ("result",)
     class Result(_message.Message):
-        __slots__ = ("partitions",)
+        __slots__ = ("partitions", "total_partitions")
         PARTITIONS_FIELD_NUMBER: _ClassVar[int]
+        TOTAL_PARTITIONS_FIELD_NUMBER: _ClassVar[int]
         partitions: _containers.RepeatedScalarFieldContainer[int]
-        def __init__(self, partitions: _Optional[_Iterable[int]] = ...) -> None: ...
+        total_partitions: int
+        def __init__(self, partitions: _Optional[_Iterable[int]] = ..., total_partitions: _Optional[int] = ...) -> None: ...
     RESULT_FIELD_NUMBER: _ClassVar[int]
     result: PartitionsResponse.Result
     def __init__(self, result: _Optional[_Union[PartitionsResponse.Result, _Mapping]] = ...) -> None: ...
diff --git a/packages/pynumaflow/pynumaflow/proto/sourcer/source_pb2_grpc.py b/packages/pynumaflow/pynumaflow/proto/sourcer/source_pb2_grpc.py
@@ -6,7 +6,7 @@
 from google.protobuf import empty_pb2 as google_dot_protobuf_dot_empty__pb2
 from pynumaflow.proto.sourcer import source_pb2 as pynumaflow_dot_proto_dot_sourcer_dot_source__pb2
 
-GRPC_GENERATED_VERSION = '1.75.0'
+GRPC_GENERATED_VERSION = '1.78.0'
 GRPC_VERSION = grpc.__version__
 _version_not_supported = False
 
@@ -19,7 +19,7 @@
 if _version_not_supported:
     raise RuntimeError(
         f'The grpc package installed is at version {GRPC_VERSION},'
-        + f' but the generated code in pynumaflow/proto/sourcer/source_pb2_grpc.py depends on'
+        + ' but the generated code in pynumaflow/proto/sourcer/source_pb2_grpc.py depends on'
         + f' grpcio>={GRPC_GENERATED_VERSION}.'
         + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}'
         + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.'
diff --git a/packages/pynumaflow/pynumaflow/proto/sourcetransformer/transform_pb2_grpc.py b/packages/pynumaflow/pynumaflow/proto/sourcetransformer/transform_pb2_grpc.py
@@ -6,7 +6,7 @@
 from google.protobuf import empty_pb2 as google_dot_protobuf_dot_empty__pb2
 from pynumaflow.proto.sourcetransformer import transform_pb2 as pynumaflow_dot_proto_dot_sourcetransformer_dot_transform__pb2
 
-GRPC_GENERATED_VERSION = '1.75.0'
+GRPC_GENERATED_VERSION = '1.78.0'
 GRPC_VERSION = grpc.__version__
 _version_not_supported = False
 
@@ -19,7 +19,7 @@
 if _version_not_supported:
     raise RuntimeError(
         f'The grpc package installed is at version {GRPC_VERSION},'
-        + f' but the generated code in pynumaflow/proto/sourcetransformer/transform_pb2_grpc.py depends on'
+        + ' but the generated code in pynumaflow/proto/sourcetransformer/transform_pb2_grpc.py depends on'
         + f' grpcio>={GRPC_GENERATED_VERSION}.'
         + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}'
         + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.'
diff --git a/packages/pynumaflow/pynumaflow/sourcer/_dtypes.py b/packages/pynumaflow/pynumaflow/sourcer/_dtypes.py
@@ -1,4 +1,5 @@
 import os
+import warnings
 from abc import ABCMeta, abstractmethod
 from collections.abc import AsyncIterator, Callable, Iterator
 from dataclasses import dataclass
@@ -240,25 +241,32 @@ def count(self) -> int:
 class PartitionsResponse:
     """
     PartitionsResponse is the response for the partition request.
-    It indicates the number of partitions at the user defined source.
-    A negative count indicates that the partition information is not available.
+    It indicates the active partitions at the user defined source.
 
     Args:
-        count: the number of partitions.
+        partitions: the list of active partitions.
+        total_partitions: the total number of partitions in the source (optional).
     """
 
     _partitions: list[int]
+    _total_partitions: int | None
 
-    def __init__(self, partitions: list[int]):
+    def __init__(self, partitions: list[int], total_partitions: int | None = None):
         if not isinstance(partitions, list):
             raise TypeError(f"Wrong data type: {type(partitions)} for Partition.partitions")
         self._partitions = partitions
+        self._total_partitions = total_partitions
 
     @property
     def partitions(self) -> list[int]:
-        """Returns the list of partitions"""
+        """Returns the list of active partitions"""
         return self._partitions
 
+    @property
+    def total_partitions(self) -> int | None:
+        """Returns the total number of partitions, or None if not available"""
+        return self._total_partitions
+
 
 class Sourcer(metaclass=ABCMeta):
     """
@@ -297,12 +305,36 @@ async def pending_handler(self) -> PendingResponse:
         """
         pass
 
-    @abstractmethod
     async def partitions_handler(self) -> PartitionsResponse:
         """
-        The simple source always returns zero to indicate there is no pending record.
+        .. deprecated::
+            Use :meth:`active_partitions_handler` instead.
+
+        Returns the active partitions associated with the source.
         """
-        pass
+        warnings.warn(
+            "partitions_handler is deprecated, use active_partitions_handler instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return PartitionsResponse(partitions=get_default_partitions())
+
+    async def active_partitions_handler(self) -> PartitionsResponse:
+        """
+        Returns the active partitions associated with the source, used by the platform
+        to determine the partitions to which the watermark should be published.
+        Falls back to partitions_handler() if not overridden.
+        """
+        return await self.partitions_handler()
+
+    async def total_partitions_handler(self) -> int | None:
+        """
+        Returns the total number of partitions in the source.
+        Used by the platform for watermark progression to know when all
+        processors have reported in.
+        Returns None by default, indicating the source does not report total partitions.
+        """
+        return None
 
 
 # Create default partition id from the environment variable "NUMAFLOW_REPLICA"
diff --git a/packages/pynumaflow/pynumaflow/sourcer/servicer/async_servicer.py b/packages/pynumaflow/pynumaflow/sourcer/servicer/async_servicer.py
diff --git a/packages/pynumaflow/tests/source/test_async_source.py b/packages/pynumaflow/tests/source/test_async_source.py
diff --git a/packages/pynumaflow/tests/source/utils.py b/packages/pynumaflow/tests/source/utils.py
