Update partition router with process_parent_record

Author: lazebnyi

airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py

@@ -1,10 +1,12 @@
 #
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
+
+
 import copy
 import logging
 from dataclasses import InitVar, dataclass
-from typing import TYPE_CHECKING, Any, Iterable, List, Mapping, MutableMapping, Optional, Union
+from typing import TYPE_CHECKING, Any, Iterable, Tuple, List, Mapping, MutableMapping, Optional, Union
 
 import dpath
 
@@ -131,6 +133,40 @@ def _get_request_option(
                         parent_config.request_option.inject_into_request(params, value, self.config)
         return params
 
+    def process_parent_record(self, parent_record: Union[AirbyteMessage, Record, Mapping], parent_stream_name: str) -> Tuple[Optional[Mapping], Optional[Mapping]]:
+        """
+        Processes and extracts data from a parent record, handling different record types
+        and ensuring only valid types proceed.
+
+        :param parent_record: The parent record to process.
+        :param parent_stream_name: The parent stream name associated with the record.
+        :return: Extracted record data and partition (if applicable).
+        :raises AirbyteTracedException: If the record type is invalid.
+        """
+        if isinstance(parent_record, AirbyteMessage):
+            self.logger.warning(
+                f"Parent stream {parent_stream_name} returns records of type AirbyteMessage. "
+                f"This SubstreamPartitionRouter is not able to checkpoint incremental parent state."
+            )
+            if parent_record.type == MessageType.RECORD:
+                return parent_record.record.data, {}
+            return None, None   # Skip invalid or non-record data
+
+        # Handle Record type
+        if isinstance(parent_record, Record):
+            parent_partition = (
+                parent_record.associated_slice.partition if parent_record.associated_slice else {}
+            )
+            return parent_record.data, parent_partition
+
+        # Validate the record type
+        if not isinstance(parent_record, Mapping):
+            raise AirbyteTracedException(
+                message=f"Parent stream returned records as invalid type {type(parent_record)}"
+            )
+
+        return parent_record, {}
+
     def stream_slices(self) -> Iterable[StreamSlice]:
         """
         Iterate over each parent stream's record and create a StreamSlice for each record.
@@ -163,28 +199,13 @@ def stream_slices(self) -> Iterable[StreamSlice]:
                 # read_stateless() assumes the parent is not concurrent. This is currently okay since the concurrent CDK does
                 # not support either substreams or RFR, but something that needs to be considered once we do
                 for parent_record in parent_stream.read_only_records():
-                    parent_partition = None
-                    # Skip non-records (eg AirbyteLogMessage)
-                    if isinstance(parent_record, AirbyteMessage):
-                        self.logger.warning(
-                            f"Parent stream {parent_stream.name} returns records of type AirbyteMessage. This SubstreamPartitionRouter is not able to checkpoint incremental parent state."
-                        )
-                        if parent_record.type == MessageType.RECORD:
-                            parent_record = parent_record.record.data  # type: ignore[union-attr, assignment]  # record is always a Record
-                        else:
-                            continue
-                    elif isinstance(parent_record, Record):
-                        parent_partition = (
-                            parent_record.associated_slice.partition
-                            if parent_record.associated_slice
-                            else {}
-                        )
-                        parent_record = parent_record.data
-                    elif not isinstance(parent_record, Mapping):
-                        # The parent_record should only take the form of a Record, AirbyteMessage, or Mapping. Anything else is invalid
-                        raise AirbyteTracedException(
-                            message=f"Parent stream returned records as invalid type {type(parent_record)}"
-                        )
+                    # Process the parent record
+                    parent_record, parent_partition = self.process_parent_record(parent_record, parent_stream.name)
+
+                    # Skip invalid or non-record data
+                    if parent_record is None:
+                        continue
+
                     try:
                         partition_value = dpath.get(
                             parent_record,  # type: ignore [arg-type]