refactor: add get_partition_router() helper to DefaultStream

Replace hasattr chain in ConcurrentReadProcessor._collect_all_parent_stream_names
with DefaultStream.get_partition_router() that safely traverses the internal
partition_generator -> stream_slicer -> partition_router chain using isinstance checks.

Co-Authored-By: unknown <>

Author: devin-ai-integration[bot]

airbyte_cdk/sources/concurrent_source/concurrent_read_processor.py

@@ -371,42 +371,30 @@ def _is_stream_done(self, stream_name: str) -> bool:
         return stream_name in self._streams_done
 
     def _collect_all_parent_stream_names(self, stream_name: str) -> Set[str]:
-        """
-        Recursively collect all parent stream names for a given stream.
-        For example, if we have: epics -> issues -> comments
-        Then for comments, this returns {issues, epics}
+        """Recursively collect all parent stream names for a given stream.
 
-        :param stream_name: The stream to collect parents for
-        :return: Set of all parent stream names (recursively)
+        For example, if we have: epics -> issues -> comments
+        Then for comments, this returns {issues, epics}.
         """
+        from airbyte_cdk.sources.declarative.partition_routers.substream_partition_router import (
+            SubstreamPartitionRouter,
+        )
+        from airbyte_cdk.sources.streams.concurrent.default_stream import DefaultStream
+
         parent_names: Set[str] = set()
         stream = self._stream_name_to_instance.get(stream_name)
 
         if not stream:
             return parent_names
 
-        # Get partition router if it exists (this is where parent streams are defined)
-        partition_router = None
-
-        # Try DefaultStream path first (_stream_partition_generator._stream_slicer._partition_router)
-        if (
-            hasattr(stream, "_stream_partition_generator")
-            and hasattr(stream._stream_partition_generator, "_stream_slicer")
-            and hasattr(stream._stream_partition_generator._stream_slicer, "_partition_router")
-        ):
-            partition_router = stream._stream_partition_generator._stream_slicer._partition_router
-        # Fallback to legacy path (retriever.partition_router) for backward compatibility and test mocks
-        elif hasattr(stream, "retriever") and hasattr(stream.retriever, "partition_router"):
-            partition_router = stream.retriever.partition_router
+        partition_router = (
+            stream.get_partition_router() if isinstance(stream, DefaultStream) else None
+        )
 
-        # SubstreamPartitionRouter has parent_stream_configs
-        if partition_router and hasattr(partition_router, "parent_stream_configs"):
+        if isinstance(partition_router, SubstreamPartitionRouter):
             for parent_config in partition_router.parent_stream_configs:
-                parent_stream = parent_config.stream
-                parent_name = parent_stream.name
+                parent_name = parent_config.stream.name
                 parent_names.add(parent_name)
-
-                # Recursively collect grandparents, great-grandparents, etc.
                 parent_names.update(self._collect_all_parent_stream_names(parent_name))
 
         return parent_names

airbyte_cdk/sources/streams/concurrent/default_stream.py

@@ -3,7 +3,10 @@
 #
 
 from logging import Logger
-from typing import Any, Callable, Iterable, List, Mapping, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, Iterable, List, Mapping, Optional, Union
+
+if TYPE_CHECKING:
+    from airbyte_cdk.sources.declarative.partition_routers.partition_router import PartitionRouter
 
 from airbyte_cdk.models import AirbyteStream, SyncMode
 from airbyte_cdk.sources.streams.concurrent.abstract_stream import AbstractStream
@@ -105,6 +108,22 @@ def block_simultaneous_read(self) -> str:
     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."""
+        from airbyte_cdk.sources.declarative.incremental.concurrent_partition_cursor import (
+            ConcurrentPerPartitionCursor,
+        )
+        from airbyte_cdk.sources.declarative.stream_slicers.declarative_partition_generator import (
+            StreamSlicerPartitionGenerator,
+        )
+
+        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
+        return stream_slicer._partition_router
+
     def check_availability(self) -> StreamAvailability:
         """
         Check stream availability by attempting to read the first record of the stream.