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 490182f37b..cd46326de0 100644 --- a/airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py +++ b/airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py @@ -229,13 +229,15 @@ def stream_slices(self) -> Iterable[StreamSlice]: ) record_data = parent_record.data - try: - partition_value = dpath.get( - record_data, # type: ignore [arg-type] - parent_field, + partition_value = dpath.get( + record_data, # type: ignore [arg-type] + parent_field, + default=None, + ) + if partition_value is None: + self.logger.debug( + f"Skipping {parent_stream.name} record with missing or null parent key '{parent_field}'" ) - except KeyError: - # FIXME a log here would go a long way for debugging emit_slice = False if emit_slice: 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 b69849ebe6..696606d141 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 @@ -1242,30 +1242,48 @@ def test_substream_partition_router_closes_all_partitions_even_when_no_records() assert close_partition_calls[2][0][0] == partition_3 -def test_substream_partition_router_closes_partition_even_when_parent_key_missing(): +@pytest.mark.parametrize( + "parent_key, valid_record, skipped_record, expected_value", + [ + pytest.param( + "id", {"id": "record_1"}, {"other_field": "value"}, "record_1", id="missing_key" + ), + pytest.param("id", {"id": 0}, {"id": None}, 0, id="null_value"), + pytest.param( + "a/id", + {"a": {"id": "record_1"}}, + {"a": {"id": None}}, + "record_1", + id="nested_null_value", + ), + pytest.param( + "a/id", {"a": {"id": "record_1"}}, {"a": None}, "record_1", id="null_intermediate" + ), + ], +) +def test_substream_partition_router_skips_parent_record_without_partition_value( + parent_key, valid_record, skipped_record, expected_value +): """ - 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. + A parent record whose parent_key is missing or null produces no slice, but the parent + cursor still observes it and closes its partition so parent state keeps progressing. """ 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]), + _build_records_for_slice([valid_record], 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]), + _build_records_for_slice([skipped_record], mock_slices[1]), ) mock_cursor = Mock() @@ -1279,7 +1297,7 @@ def test_substream_partition_router_closes_partition_even_when_parent_key_missin "first_stream", cursor=mock_cursor, ), - parent_key="id", + parent_key=parent_key, partition_field="partition_field", parameters={}, config={}, @@ -1291,13 +1309,11 @@ def test_substream_partition_router_closes_partition_even_when_parent_key_missin 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"}}, + {"partition_field": expected_value, "parent_slice": {"slice": "first"}}, ] - # Both partitions should be closed, even though the second one had a KeyError + assert mock_cursor.observe.call_count == 2 assert mock_cursor.close_partition.call_count == 2 close_partition_calls = mock_cursor.close_partition.call_args_list