From a4c6229c832bb4ac94465eccd192a975d14bf078 Mon Sep 17 00:00:00 2001 From: darynaishchenko Date: Mon, 2 Feb 2026 17:08:35 +0200 Subject: [PATCH 1/9] fix SubstreamPartitionRouter updating cursor value when no records were read in partition --- .../substream_partition_router.py | 93 ++++++++++--------- 1 file changed, 49 insertions(+), 44 deletions(-) diff --git a/airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py b/airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py index e41a0d9a10..32cdd5de9d 100644 --- a/airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py +++ b/airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py @@ -45,6 +45,7 @@ def iterate_with_last_flag(generator: Iterable[T]) -> Iterable[tuple[T, bool]]: try: current = next(iterator) except StopIteration: + yield None, True return # Return an empty iterator for next_item in iterator: @@ -206,60 +207,64 @@ def stream_slices(self) -> Iterable[StreamSlice]: for partition, is_last_slice in iterate_with_last_flag( parent_stream.generate_partitions() ): + if partition is None: + return for parent_record, is_last_record_in_slice in iterate_with_last_flag( partition.read() ): - # 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, + 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 {} ) - 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, - } + 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, + } 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, + ) yield from [] From cf5063365fdafa8174603fdff933bdc187f445df Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 3 Feb 2026 15:35:36 +0000 Subject: [PATCH 2/9] test(cdk): add unit tests for iterate_with_last_flag and SubstreamPartitionRouter null handling Co-Authored-By: Daryna Ishchenko --- .../test_substream_partition_router.py | 93 +++++++++++++++++++ 1 file changed, 93 insertions(+) diff --git a/unit_tests/sources/declarative/partition_routers/test_substream_partition_router.py b/unit_tests/sources/declarative/partition_routers/test_substream_partition_router.py index feab9bbc0b..d1ae87f923 100644 --- a/unit_tests/sources/declarative/partition_routers/test_substream_partition_router.py +++ b/unit_tests/sources/declarative/partition_routers/test_substream_partition_router.py @@ -20,6 +20,7 @@ from airbyte_cdk.sources.declarative.partition_routers.substream_partition_router import ( ParentStreamConfig, SubstreamPartitionRouter, + iterate_with_last_flag, ) from airbyte_cdk.sources.declarative.requesters.request_option import ( RequestOption, @@ -1079,3 +1080,95 @@ def test_cartesian_product_stream_slicer_warning_log_message( assert warning_message in logged_warnings else: assert warning_message not in logged_warnings + + +@pytest.mark.parametrize( + "input_iterable,expected_output", + [ + pytest.param([], [(None, True)], id="empty_generator_yields_none_sentinel"), + pytest.param([1], [(1, True)], id="single_item"), + pytest.param([1, 2], [(1, False), (2, True)], id="two_items"), + pytest.param([1, 2, 3], [(1, False), (2, False), (3, True)], id="three_items"), + pytest.param(["a", "b"], [("a", False), ("b", True)], id="string_items"), + ], +) +def test_iterate_with_last_flag(input_iterable, expected_output): + result = list(iterate_with_last_flag(input_iterable)) + assert result == expected_output + + +def test_substream_partition_router_no_cursor_update_when_partition_has_no_records(): + """ + Test that when a partition has no records, the cursor is still properly closed + but no slices are yielded for that partition. + This tests the fix for SubstreamPartitionRouter updating cursor value when no records + were read in partition. + """ + mock_slices = [ + StreamSlice(partition={"slice": "first"}, cursor_slice={}), + StreamSlice(partition={"slice": "second"}, cursor_slice={}), + ] + + partition_router = SubstreamPartitionRouter( + parent_stream_configs=[ + ParentStreamConfig( + stream=MockStream( + [ + InMemoryPartition( + "partition_1", + "first_stream", + mock_slices[0], + _build_records_for_slice( + [{"id": "record_1"}, {"id": "record_2"}], mock_slices[0] + ), + ), + InMemoryPartition( + "partition_2", + "first_stream", + mock_slices[1], + [], + ), + ], + "first_stream", + ), + parent_key="id", + partition_field="partition_field", + parameters={}, + config={}, + ) + ], + parameters={}, + config={}, + ) + + slices = list(partition_router.stream_slices()) + assert slices == [ + {"partition_field": "record_1", "parent_slice": {"slice": "first"}}, + {"partition_field": "record_2", "parent_slice": {"slice": "first"}}, + ] + + +def test_substream_partition_router_handles_empty_parent_partitions(): + """ + Test that when a parent stream generates no partitions (empty generator), + the stream_slices method returns early without errors. + """ + partition_router = SubstreamPartitionRouter( + parent_stream_configs=[ + ParentStreamConfig( + stream=MockStream( + [], + "first_stream", + ), + parent_key="id", + partition_field="partition_field", + parameters={}, + config={}, + ) + ], + parameters={}, + config={}, + ) + + slices = list(partition_router.stream_slices()) + assert slices == [] From 17430ec2d6e796d21cc824903e8964a1f254a714 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 3 Feb 2026 15:53:11 +0000 Subject: [PATCH 3/9] fix(cdk): change return to break when partition is None to preserve multi-parent behavior Co-Authored-By: Daryna Ishchenko --- .../declarative/partition_routers/substream_partition_router.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py b/airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py index 32cdd5de9d..3e30c70eef 100644 --- a/airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py +++ b/airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py @@ -208,7 +208,7 @@ def stream_slices(self) -> Iterable[StreamSlice]: parent_stream.generate_partitions() ): if partition is None: - return + break for parent_record, is_last_record_in_slice in iterate_with_last_flag( partition.read() ): From 7a12e5248be94fd9c04bbd51bd69764cde9207f8 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 5 Feb 2026 10:36:53 +0000 Subject: [PATCH 4/9] test(cdk): add test for close_partition called on all partitions even when no records Co-Authored-By: Daryna Ishchenko --- .../test_substream_partition_router.py | 67 +++++++++++++++++++ 1 file changed, 67 insertions(+) diff --git a/unit_tests/sources/declarative/partition_routers/test_substream_partition_router.py b/unit_tests/sources/declarative/partition_routers/test_substream_partition_router.py index d1ae87f923..35591b0918 100644 --- a/unit_tests/sources/declarative/partition_routers/test_substream_partition_router.py +++ b/unit_tests/sources/declarative/partition_routers/test_substream_partition_router.py @@ -1172,3 +1172,70 @@ def test_substream_partition_router_handles_empty_parent_partitions(): slices = list(partition_router.stream_slices()) assert slices == [] + + +def test_substream_partition_router_closes_all_partitions_even_when_no_records(): + """ + Test that cursor.close_partition() is called for all parent stream partitions, + even when a partition produces no parent records. + This validates that partition lifecycle is properly managed regardless of record count. + """ + mock_slices = [ + StreamSlice(partition={"slice": "first"}, cursor_slice={}), + StreamSlice(partition={"slice": "second"}, cursor_slice={}), + StreamSlice(partition={"slice": "third"}, cursor_slice={}), + ] + + partition_1 = InMemoryPartition( + "partition_1", + "first_stream", + mock_slices[0], + _build_records_for_slice([{"id": "record_1"}], mock_slices[0]), + ) + partition_2 = InMemoryPartition( + "partition_2", + "first_stream", + mock_slices[1], + [], + ) + partition_3 = InMemoryPartition( + "partition_3", + "first_stream", + mock_slices[2], + _build_records_for_slice([{"id": "record_3"}], mock_slices[2]), + ) + + mock_cursor = Mock() + mock_cursor.stream_slices.return_value = [] + + partition_router = SubstreamPartitionRouter( + parent_stream_configs=[ + ParentStreamConfig( + stream=MockStream( + [partition_1, partition_2, partition_3], + "first_stream", + cursor=mock_cursor, + ), + parent_key="id", + partition_field="partition_field", + parameters={}, + config={}, + ) + ], + parameters={}, + config={}, + ) + + slices = list(partition_router.stream_slices()) + + assert slices == [ + {"partition_field": "record_1", "parent_slice": {"slice": "first"}}, + {"partition_field": "record_3", "parent_slice": {"slice": "third"}}, + ] + + assert mock_cursor.close_partition.call_count == 3 + + close_partition_calls = mock_cursor.close_partition.call_args_list + assert close_partition_calls[0][0][0] == partition_1 + assert close_partition_calls[1][0][0] == partition_2 + assert close_partition_calls[2][0][0] == partition_3 From 695b819a28773e29867f33810515ebaf7e08bff6 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 5 Feb 2026 12:38:21 +0000 Subject: [PATCH 5/9] fix(cdk): update iterate_with_last_flag return type to allow None for empty generators Co-Authored-By: Daryna Ishchenko --- .../declarative/partition_routers/substream_partition_router.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py b/airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py index 3e30c70eef..1351022869 100644 --- a/airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py +++ b/airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py @@ -39,7 +39,7 @@ T = TypeVar("T") -def iterate_with_last_flag(generator: Iterable[T]) -> Iterable[tuple[T, bool]]: +def iterate_with_last_flag(generator: Iterable[T]) -> Iterable[tuple[T | None, bool]]: iterator = iter(generator) try: From d6c970165d7540c40b6c9fab86c738267ab7ad08 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 9 Feb 2026 12:02:17 +0000 Subject: [PATCH 6/9] fix(cdk): update expected state in test_substream_slicer_parent_state_update_with_cursor Co-Authored-By: Daryna Ishchenko --- .../partition_routers/test_substream_partition_router.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/unit_tests/sources/declarative/partition_routers/test_substream_partition_router.py b/unit_tests/sources/declarative/partition_routers/test_substream_partition_router.py index 35591b0918..e470159ab9 100644 --- a/unit_tests/sources/declarative/partition_routers/test_substream_partition_router.py +++ b/unit_tests/sources/declarative/partition_routers/test_substream_partition_router.py @@ -612,7 +612,8 @@ def test_request_option( ), { "first_stream": { - "lookback_window": 0, + "lookback_window": 1, + "state": {"cursor": "2022-01-01"}, "states": [ {"cursor": {"cursor": "2021-01-02"}, "partition": {"slice": "first"}}, {"cursor": {"cursor": "2022-01-01"}, "partition": {"slice": "second"}}, From 58ffd0ddaee6a3a810f7184729b2e53c739a879d Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 10 Feb 2026 10:34:34 +0000 Subject: [PATCH 7/9] fix(cdk): use skip_slice flag to ensure partition closure even on KeyError - Replace 'continue' with 'skip_slice' flag when dpath.get fails with KeyError - This ensures close_partition() and ensure_at_least_one_state_emitted() are always called, even when partition value extraction fails - Add test for partition closure when parent_key is missing from record Co-Authored-By: Daryna Ishchenko --- .../substream_partition_router.py | 31 ++++----- .../test_substream_partition_router.py | 63 +++++++++++++++++++ 2 files changed, 80 insertions(+), 14 deletions(-) diff --git a/airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py b/airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py index 1351022869..8087348efc 100644 --- a/airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py +++ b/airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py @@ -212,6 +212,7 @@ def stream_slices(self) -> Iterable[StreamSlice]: for parent_record, is_last_record_in_slice in iterate_with_last_flag( partition.read() ): + 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 @@ -235,28 +236,30 @@ def stream_slices(self) -> Iterable[StreamSlice]: ) except KeyError: # FIXME a log here would go a long way for debugging - continue + pass + else: + skip_slice = False - # Add extra fields - extracted_extra_fields = self._extract_extra_fields( - record_data, extra_fields - ) + # 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 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, diff --git a/unit_tests/sources/declarative/partition_routers/test_substream_partition_router.py b/unit_tests/sources/declarative/partition_routers/test_substream_partition_router.py index e470159ab9..b69849ebe6 100644 --- a/unit_tests/sources/declarative/partition_routers/test_substream_partition_router.py +++ b/unit_tests/sources/declarative/partition_routers/test_substream_partition_router.py @@ -1240,3 +1240,66 @@ def test_substream_partition_router_closes_all_partitions_even_when_no_records() assert close_partition_calls[0][0][0] == partition_1 assert close_partition_calls[1][0][0] == partition_2 assert close_partition_calls[2][0][0] == partition_3 + + +def test_substream_partition_router_closes_partition_even_when_parent_key_missing(): + """ + Test that cursor.close_partition() is called even when the parent_key extraction + fails with a KeyError. This ensures partition lifecycle is properly managed + regardless of whether the slice can be emitted. + """ + mock_slices = [ + StreamSlice(partition={"slice": "first"}, cursor_slice={}), + StreamSlice(partition={"slice": "second"}, cursor_slice={}), + ] + + # First partition has a record with the expected "id" key + partition_1 = InMemoryPartition( + "partition_1", + "first_stream", + mock_slices[0], + _build_records_for_slice([{"id": "record_1"}], mock_slices[0]), + ) + # Second partition has a record missing the "id" key (will cause KeyError) + partition_2 = InMemoryPartition( + "partition_2", + "first_stream", + mock_slices[1], + _build_records_for_slice([{"other_field": "value"}], mock_slices[1]), + ) + + mock_cursor = Mock() + mock_cursor.stream_slices.return_value = [] + + partition_router = SubstreamPartitionRouter( + parent_stream_configs=[ + ParentStreamConfig( + stream=MockStream( + [partition_1, partition_2], + "first_stream", + cursor=mock_cursor, + ), + parent_key="id", + partition_field="partition_field", + parameters={}, + config={}, + ) + ], + parameters={}, + config={}, + ) + + slices = list(partition_router.stream_slices()) + + # Only the first partition's record should produce a slice + # The second partition's record is missing the "id" key, so no slice is emitted + assert slices == [ + {"partition_field": "record_1", "parent_slice": {"slice": "first"}}, + ] + + # Both partitions should be closed, even though the second one had a KeyError + assert mock_cursor.close_partition.call_count == 2 + + close_partition_calls = mock_cursor.close_partition.call_args_list + assert close_partition_calls[0][0][0] == partition_1 + assert close_partition_calls[1][0][0] == partition_2 From aa99d1b2d2bab603df349508c2b501208badba98 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 10 Feb 2026 10:40:58 +0000 Subject: [PATCH 8/9] refactor(cdk): rename skip_slice to emit_slice for clarity - Renamed skip_slice to emit_slice and inverted the logic - Now emit_slice=False by default, set to True when extraction succeeds - Makes the code more intuitive: we emit slices when conditions are met Co-Authored-By: Daryna Ishchenko --- .../partition_routers/substream_partition_router.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py b/airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py index 8087348efc..87190afcd2 100644 --- a/airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py +++ b/airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py @@ -212,7 +212,7 @@ def stream_slices(self) -> Iterable[StreamSlice]: for parent_record, is_last_record_in_slice in iterate_with_last_flag( partition.read() ): - skip_slice = True + emit_slice = False 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 @@ -234,12 +234,12 @@ def stream_slices(self) -> Iterable[StreamSlice]: record_data, # type: ignore [arg-type] parent_field, ) + emit_slice = True except KeyError: # FIXME a log here would go a long way for debugging pass - else: - skip_slice = False + if emit_slice: # Add extra fields extracted_extra_fields = self._extract_extra_fields( record_data, extra_fields @@ -259,7 +259,7 @@ def stream_slices(self) -> Iterable[StreamSlice]: if is_last_slice: parent_stream.cursor.ensure_at_least_one_state_emitted() - if not skip_slice: + if emit_slice: yield StreamSlice( partition={ partition_field: partition_value, From 367f291e847a3859a67fd8d07a7ac328ace322b3 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 10 Feb 2026 10:44:43 +0000 Subject: [PATCH 9/9] refactor(cdk): fix emit_slice logic - default True when parent_record exists, False on KeyError - emit_slice = parent_record is not None (default True when we have a record) - Set emit_slice = False only when KeyError occurs during dpath.get - This makes the logic clearer: we want to emit slices by default, skip only on error Co-Authored-By: Daryna Ishchenko --- .../partition_routers/substream_partition_router.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py b/airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py index 87190afcd2..5d2524b55b 100644 --- a/airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py +++ b/airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py @@ -212,7 +212,7 @@ def stream_slices(self) -> Iterable[StreamSlice]: for parent_record, is_last_record_in_slice in iterate_with_last_flag( partition.read() ): - emit_slice = False + emit_slice = parent_record is not None 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 @@ -234,10 +234,9 @@ def stream_slices(self) -> Iterable[StreamSlice]: record_data, # type: ignore [arg-type] parent_field, ) - emit_slice = True except KeyError: # FIXME a log here would go a long way for debugging - pass + emit_slice = False if emit_slice: # Add extra fields