fix: handle GroupingPartitionRouter at call sites instead of in get_partition_router()

Co-Authored-By: unknown <>

Author: devin-ai-integration[bot]

airbyte_cdk/sources/concurrent_source/concurrent_read_processor.py

@@ -13,6 +13,9 @@
 )
 from airbyte_cdk.sources.concurrent_source.stream_thread_exception import StreamThreadException
 from airbyte_cdk.sources.concurrent_source.thread_pool_manager import ThreadPoolManager
+from airbyte_cdk.sources.declarative.partition_routers.grouping_partition_router import (
+    GroupingPartitionRouter,
+)
 from airbyte_cdk.sources.declarative.partition_routers.substream_partition_router import (
     SubstreamPartitionRouter,
 )
@@ -398,6 +401,8 @@ def _collect_all_parent_stream_names(self, stream_name: str) -> Set[str]:
         partition_router = (
             stream.get_partition_router() if isinstance(stream, DefaultStream) else None
         )
+        if isinstance(partition_router, GroupingPartitionRouter):
+            partition_router = partition_router.underlying_partition_router
 
         if isinstance(partition_router, SubstreamPartitionRouter):
             for parent_config in partition_router.parent_stream_configs:

airbyte_cdk/sources/declarative/concurrent_declarative_source.py

@@ -76,6 +76,9 @@
 from airbyte_cdk.sources.declarative.parsers.model_to_component_factory import (
     ModelToComponentFactory,
 )
+from airbyte_cdk.sources.declarative.partition_routers.grouping_partition_router import (
+    GroupingPartitionRouter,
+)
 from airbyte_cdk.sources.declarative.partition_routers.substream_partition_router import (
     SubstreamPartitionRouter,
 )
@@ -454,6 +457,8 @@ def _apply_stream_groups(self, streams: List[AbstractStream]) -> None:
             if not isinstance(stream, DefaultStream) or stream.name not in stream_name_to_group:
                 continue
             partition_router = stream.get_partition_router()
+            if isinstance(partition_router, GroupingPartitionRouter):
+                partition_router = partition_router.underlying_partition_router
             if not isinstance(partition_router, SubstreamPartitionRouter):
                 continue
             group_name = stream_name_to_group[stream.name]

airbyte_cdk/sources/streams/concurrent/default_stream.py

@@ -9,9 +9,6 @@
 from airbyte_cdk.sources.declarative.incremental.concurrent_partition_cursor import (
     ConcurrentPerPartitionCursor,
 )
-from airbyte_cdk.sources.declarative.partition_routers.grouping_partition_router import (
-    GroupingPartitionRouter,
-)
 from airbyte_cdk.sources.declarative.partition_routers.partition_router import PartitionRouter
 from airbyte_cdk.sources.declarative.stream_slicers.declarative_partition_generator import (
     StreamSlicerPartitionGenerator,
@@ -116,20 +113,13 @@ def block_simultaneous_read(self, value: str) -> None:
         self._block_simultaneous_read = value
 
     def get_partition_router(self) -> PartitionRouter | None:
-        """Return the partition router for this stream, or None if not available.
-
-        If the router is a GroupingPartitionRouter, unwraps it to return the
-        underlying router so callers can inspect parent stream relationships.
-        """
+        """Return the partition router for this stream, or None if not available."""
         if not isinstance(self._stream_partition_generator, StreamSlicerPartitionGenerator):
             return None
         stream_slicer = self._stream_partition_generator._stream_slicer
         if not isinstance(stream_slicer, ConcurrentPerPartitionCursor):
             return None
-        router = stream_slicer._partition_router
-        if isinstance(router, GroupingPartitionRouter):
-            return router.underlying_partition_router
-        return router
+        return stream_slicer._partition_router
 
     def check_availability(self) -> StreamAvailability:
         """