fix(cdk): fix SubstreamPartitionRouter updating cursor value when no records were read in partition#889
fix(cdk): fix SubstreamPartitionRouter updating cursor value when no records were read in partition#889darynaishchenko wants to merge 8 commits intomainfrom
Conversation
…re read in partition
👋 Greetings, Airbyte Team Member!Here are some helpful tips and reminders for your convenience. 💡 Show Tips and TricksTesting This CDK VersionYou can test this version of the CDK using the following: # Run the CLI from this branch:
uvx 'git+https://github.com/airbytehq/airbyte-python-cdk.git@daryna/fix-substream-partition-router#egg=airbyte-python-cdk[dev]' --help
# Update a connector to use the CDK from this branch ref:
cd airbyte-integrations/connectors/source-example
poe use-cdk-branch daryna/fix-substream-partition-routerPR Slash CommandsAirbyte Maintainers can execute the following slash commands on your PR:
|
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughChanged SubstreamPartitionRouter to add Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Would you like me to generate a small focused patch-level checklist for reviewers (key lines/behaviors to verify), wdyt? 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
No actionable comments were generated in the recent review. 🎉 🧹 Recent nitpick comments
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py (1)
42-49:⚠️ Potential issue | 🔴 CriticalType annotation mismatch causing pipeline failure.
The pipeline is flagging a type error here. Yielding
(None, True)producestuple[None, bool], but the return type expectstuple[T, bool]. Since the sentinel introducesNoneas a possible value, would updating the return type toOptional[T]work for you, wdyt?🔧 Proposed fix
-def iterate_with_last_flag(generator: Iterable[T]) -> Iterable[tuple[T, bool]]: +def iterate_with_last_flag(generator: Iterable[T]) -> Iterable[tuple[Optional[T], bool]]:
🤖 Fix all issues with AI agents
In
`@airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py`:
- Around line 210-211: In stream_slices() in substream_partition_router.py the
check "if partition is None: return" prematurely exits the entire method when a
parent_stream.generate_partitions() yields the (None, True) sentinel; change
this to "break" (or otherwise exit the inner partition loop) so we stop
iterating that parent stream's partitions and continue to the next
parent_stream_config instead of returning from stream_slices(), leaving the rest
of the parent streams to be processed; locate the partition handling inside the
loop that consumes parent_stream.generate_partitions() and replace the return
with a break (or equivalent loop-control) so behavior matches the original
semantics.
airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py
Outdated
Show resolved
Hide resolved
|
/prerelease
|
…titionRouter null handling Co-Authored-By: Daryna Ishchenko <darina.ishchenko17@gmail.com>
…ulti-parent behavior Co-Authored-By: Daryna Ishchenko <darina.ishchenko17@gmail.com>
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py (2)
42-49:⚠️ Potential issue | 🟠 MajorFix mypy type mismatch in
iterate_with_last_flag.
Line 48 yields(None, True), but the signature advertisestuple[T, bool], which is tripping mypy. Could we widen the return type toOptional[T]to match the sentinel, wdyt?🔧 Proposed fix
-def iterate_with_last_flag(generator: Iterable[T]) -> Iterable[tuple[T, bool]]: +def iterate_with_last_flag(generator: Iterable[T]) -> Iterable[tuple[Optional[T], bool]]:
186-187:⚠️ Potential issue | 🟡 MinorUpdate docstring to match empty-slice behavior.
Lines 186–187 still say an empty parent slice emitsparent_record=None, but the code now skips emission. Could we adjust the docstring to reflect the new behavior, wdyt?📝 Suggested doc tweak
- If a parent slice contains no record, emit a slice with parent_record=None. + If a parent slice contains no record, emit no slice.
| assert result == expected_output | ||
|
|
||
|
|
||
| def test_substream_partition_router_no_cursor_update_when_partition_has_no_records(): |
There was a problem hiding this comment.
Please also add a test that validates that all parent stream partitions are properly closed, even if no parent records are produced for a given parent partition.
… when no records Co-Authored-By: Daryna Ishchenko <darina.ishchenko17@gmail.com>
… empty generators Co-Authored-By: Daryna Ishchenko <darina.ishchenko17@gmail.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In
`@airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py`:
- Around line 231-267: The try/except around dpath.get in
substream_partition_router should not use an early `continue` because that skips
the later partition closure logic; instead, introduce a boolean like `skip_slice
= True` when KeyError occurs (and only set `skip_slice = False` when extraction
succeeds), keep `extracted_extra_fields` in a safe default state, and let
execution fall through so the existing `if is_last_record_in_slice:
parent_stream.cursor.close_partition(partition)` / `if is_last_slice:
parent_stream.cursor.ensure_at_least_one_state_emitted()` blocks always run;
finally, guard the `yield StreamSlice(...)` (the emission that uses
`parent_record`, `partition_value`, `extracted_extra_fields`,
`parent_partition`, etc.) so it only occurs when `skip_slice` is False. This
ensures partition closure always executes while still skipping emission for
failed extractions.
| try: | ||
| partition_value = dpath.get( | ||
| record_data, # type: ignore [arg-type] | ||
| parent_field, | ||
| ) | ||
| except KeyError: | ||
| # FIXME a log here would go a long way for debugging | ||
| continue | ||
|
|
||
| # Add extra fields | ||
| extracted_extra_fields = self._extract_extra_fields( | ||
| record_data, extra_fields | ||
| ) | ||
|
|
||
| if parent_stream_config.lazy_read_pointer: | ||
| extracted_extra_fields = { | ||
| "child_response": self._extract_child_response( | ||
| record_data, | ||
| parent_stream_config.lazy_read_pointer, # type: ignore[arg-type] # lazy_read_pointer type handeled in __post_init__ of parent_stream_config | ||
| ), | ||
| **extracted_extra_fields, | ||
| } | ||
|
|
||
| if is_last_record_in_slice: | ||
| parent_stream.cursor.close_partition(partition) | ||
| if is_last_slice: | ||
| parent_stream.cursor.ensure_at_least_one_state_emitted() | ||
|
|
||
| yield StreamSlice( | ||
| partition={ | ||
| partition_field: partition_value, | ||
| "parent_slice": parent_partition or {}, | ||
| }, | ||
| cursor_slice={}, | ||
| extra_fields=extracted_extra_fields, | ||
| ) | ||
| if parent_record is not None: | ||
| yield StreamSlice( | ||
| partition={ | ||
| partition_field: partition_value, | ||
| "parent_slice": parent_partition or {}, | ||
| }, | ||
| cursor_slice={}, | ||
| extra_fields=extracted_extra_fields, | ||
| ) |
There was a problem hiding this comment.
Avoid skipping partition closure when partition-value extraction fails.
The early continue in the extraction-failure path skips close_partition() / ensure_at_least_one_state_emitted() when that record is last, leaving the cursor partition open. Could we switch to a skip_slice flag so closure always runs and only slice emission is skipped, wdyt?
🔧 Possible fix
for parent_record, is_last_record_in_slice in iterate_with_last_flag(
partition.read()
):
- if parent_record is not None:
+ skip_slice = True
+ if parent_record is not None:
# In the previous CDK implementation, state management was done internally by the stream.
# However, this could cause issues when doing availability check for example as the availability
# check would progress the state so state management was moved outside of the read method.
# Hence, we need to call the cursor here.
# Note that we call observe and close_partition before emitting the associated record as the
# ConcurrentPerPartitionCursor will associate a record with the state of the stream after the
# record was consumed.
parent_stream.cursor.observe(parent_record)
parent_partition = (
parent_record.associated_slice.partition
if parent_record.associated_slice
else {}
)
record_data = parent_record.data
try:
partition_value = dpath.get(
record_data, # type: ignore [arg-type]
parent_field,
)
except KeyError:
# FIXME a log here would go a long way for debugging
- continue
-
- # Add extra fields
- extracted_extra_fields = self._extract_extra_fields(
- record_data, extra_fields
- )
-
- if parent_stream_config.lazy_read_pointer:
- extracted_extra_fields = {
- "child_response": self._extract_child_response(
- record_data,
- parent_stream_config.lazy_read_pointer, # type: ignore[arg-type] # lazy_read_pointer type handeled in __post_init__ of parent_stream_config
- ),
- **extracted_extra_fields,
- }
+ skip_slice = True
+ else:
+ skip_slice = False
+
+ # Add extra fields
+ extracted_extra_fields = self._extract_extra_fields(
+ record_data, extra_fields
+ )
+
+ if parent_stream_config.lazy_read_pointer:
+ extracted_extra_fields = {
+ "child_response": self._extract_child_response(
+ record_data,
+ parent_stream_config.lazy_read_pointer, # type: ignore[arg-type] # lazy_read_pointer type handeled in __post_init__ of parent_stream_config
+ ),
+ **extracted_extra_fields,
+ }
if is_last_record_in_slice:
parent_stream.cursor.close_partition(partition)
if is_last_slice:
parent_stream.cursor.ensure_at_least_one_state_emitted()
- if parent_record is not None:
+ if not skip_slice:
yield StreamSlice(
partition={
partition_field: partition_value,
"parent_slice": parent_partition or {},
},
cursor_slice={},
extra_fields=extracted_extra_fields,
)🤖 Prompt for AI Agents
In
`@airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py`
around lines 231 - 267, The try/except around dpath.get in
substream_partition_router should not use an early `continue` because that skips
the later partition closure logic; instead, introduce a boolean like `skip_slice
= True` when KeyError occurs (and only set `skip_slice = False` when extraction
succeeds), keep `extracted_extra_fields` in a safe default state, and let
execution fall through so the existing `if is_last_record_in_slice:
parent_stream.cursor.close_partition(partition)` / `if is_last_slice:
parent_stream.cursor.ensure_at_least_one_state_emitted()` blocks always run;
finally, guard the `yield StreamSlice(...)` (the emission that uses
`parent_record`, `partition_value`, `extracted_extra_fields`,
`parent_partition`, etc.) so it only occurs when `skip_slice` is False. This
ensures partition closure always executes while still skipping emission for
failed extractions.
…_update_with_cursor Co-Authored-By: Daryna Ishchenko <darina.ishchenko17@gmail.com>
2 participants
Summary
This PR fixes an issue where
SubstreamPartitionRouterwould incorrectly update cursor values when no records were read in a partition. The fix adds defensive null-checking throughoutstream_slices()to gracefully handle:iterate_with_last_flagnow yields(None, True)sentinel when the input generator is empty, andstream_slices()breaks out of the partition loop when partition is None (preserving multi-parent behavior)if parent_record is not None)Updates since last revision
if partition is None: returntoif partition is None: breakper review feedback. Usingbreakensures that if one parent stream has no partitions, we continue processing subsequentparent_stream_configsrather than exiting the entire method.test_substream_partition_router_closes_all_partitions_even_when_no_recordsto verify thatcursor.close_partition()is called for ALL partitions, including those with no records. This addresses the review feedback requesting validation of partition lifecycle management.iterate_with_last_flagreturn type fromtuple[T, bool]totuple[T | None, bool]to properly reflect that the function yields(None, True)as a sentinel for empty generators.Review & Testing Checklist for Human
test_substream_slicer_parent_state_update_with_cursoris failing on this branch but passes on main. The test expectslookback_window: 0but getslookback_window: 1, and there's an extrastatekey in the output. This appears to be a pre-existing issue in the PR's original changes that needs investigation.if partition is None: breakexits only the current parent's partition loop. Test with a connector that has multipleparent_stream_configswhere one parent has no partitions to confirm subsequent parents are still processed. (Note: current tests only cover single-parent scenarios)SubstreamPartitionRouterwhere a parent stream returns empty partitions or partitions with no recordspoetry run pytest unit_tests/sources/declarative/partition_routers/test_substream_partition_router.py -vNotes
Summary by CodeRabbit
New Features
Bug Fixes
Tests