From f15007f6881ba7b91c64c3c1edc5770ad96e2b74 Mon Sep 17 00:00:00 2001 From: "Gareth.Y" <91742523+linjiaye929@users.noreply.github.com> Date: Tue, 13 Oct 2026 01:07:20 -0700 Subject: [PATCH 1/2] fix(dataset): invalidate fetched records after deletion --- py/src/braintrust/logger.py | 1 + py/src/braintrust/test_logger.py | 39 ++++++++++++++++++++++++++++++++ 2 files changed, 40 insertions(+) diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py index e8243f070..f352c2c7b 100644 --- a/py/src/braintrust/logger.py +++ b/py/src/braintrust/logger.py @@ -5562,6 +5562,7 @@ def compute_args(): ) self.state.global_bg_logger().log(LazyValue(compute_args, use_mutex=False)) + self._clear_cache() return id def summarize(self, summarize_data: bool = True) -> "DatasetSummary": diff --git a/py/src/braintrust/test_logger.py b/py/src/braintrust/test_logger.py index 19dc3add3..bf5be3198 100644 --- a/py/src/braintrust/test_logger.py +++ b/py/src/braintrust/test_logger.py @@ -3188,6 +3188,45 @@ def test_dataset_fetch_uses_generated_resource(self): ) mock_state.api_conn.assert_not_called() + def test_dataset_delete_invalidates_fetched_records(self): + mock_state = MagicMock() + fetch = mock_state.api_client.return_value.datasets.post_dataset_id_fetch + deleted = {"id": "deleted-row", "expected": "obsolete"} + retained = {"id": "retained-row", "expected": "current"} + fetch.side_effect = [ + {"events": [deleted, retained]}, + {"events": [retained]}, + ] + metadata = logger.ProjectDatasetMetadata( + project=logger.ObjectMetadata(id="test-project-id", name="test-project", full_info={}), + dataset=logger.ObjectMetadata(id="test-dataset-id", name="test-dataset", full_info={}), + ) + dataset = logger.Dataset( + lazy_metadata=LazyValue(lambda: metadata, use_mutex=False), + legacy=False, + state=mock_state, + ) + + self.assertEqual(list(dataset), [deleted, retained]) + self.assertEqual(list(dataset), [deleted, retained]) + self.assertEqual(fetch.call_count, 1) + + self.assertEqual(dataset.delete("deleted-row"), "deleted-row") + background_logger = mock_state.global_bg_logger.return_value + background_logger.log.assert_called_once() + tombstone = background_logger.log.call_args.args[0].get() + self.assertEqual(tombstone["id"], "deleted-row") + self.assertEqual(tombstone["dataset_id"], "test-dataset-id") + self.assertIs(tombstone["_object_delete"], True) + self.assertIsInstance(tombstone["created"], str) + + dataset.flush() + background_logger.flush.assert_called_once() + self.assertEqual(list(dataset), [retained]) + self.assertEqual(fetch.call_count, 2) + self.assertEqual(list(dataset), [retained]) + self.assertEqual(fetch.call_count, 2) + def test_dataset_summary_uses_generated_resource_and_configured_public_url(self): mock_state = MagicMock() mock_state.app_public_url = "https://public.example.com" From 98ea56114999d9204f56cc326b03cfc599a6107d Mon Sep 17 00:00:00 2001 From: "Gareth.Y" <91742523+linjiaye929@users.noreply.github.com> Date: Tue, 13 Oct 2026 02:09:01 -0700 Subject: [PATCH 2/2] fix(dataset): preserve active iterator snapshots during mutation --- py/src/braintrust/logger.py | 6 ++- py/src/braintrust/test_logger.py | 63 +++++++++++++++++++++++++++++++- 2 files changed, 67 insertions(+), 2 deletions(-) diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py index f352c2c7b..0ce39efee 100644 --- a/py/src/braintrust/logger.py +++ b/py/src/braintrust/logger.py @@ -3398,12 +3398,16 @@ class ObjectIterator(Generic[T]): def __init__(self, refetch_fn: Callable[[], Sequence[T]]): self.refetch_fn = refetch_fn self.idx = 0 + self._data: Sequence[T] | None = None def __iter__(self): return self def __next__(self) -> T: - data = self.refetch_fn() + if self._data is None: + # Keep this iteration stable if a mutation invalidates the fetcher's cache. + self._data = self.refetch_fn() + data = self._data if self.idx >= len(data): raise StopIteration value = data[self.idx] diff --git a/py/src/braintrust/test_logger.py b/py/src/braintrust/test_logger.py index bf5be3198..963608956 100644 --- a/py/src/braintrust/test_logger.py +++ b/py/src/braintrust/test_logger.py @@ -2883,6 +2883,7 @@ def test_experiment_fetch_uses_generated_paginated_fetch(self): api_client.experiments.post_experiment_id_fetch.side_effect = [ {"events": [first_event], "cursor": "next-page"}, {"events": [second_event]}, + {"events": [second_event]}, ] metadata = ProjectExperimentMetadata( project=ObjectMetadata(id="test-project-id", name="test-project", full_info={}), @@ -2890,7 +2891,10 @@ def test_experiment_fetch_uses_generated_paginated_fetch(self): ) experiment = Experiment(LazyValue(lambda: metadata, use_mutex=False), state=mock_state) - events = list(experiment.fetch(batch_size=1)) + iterator = experiment.fetch(batch_size=1) + events = [next(iterator)] + experiment._clear_cache() + events.extend(iterator) assert [event["id"] for event in events] == ["first-event", "second-event"] assert api_client.experiments.post_experiment_id_fetch.call_args_list == [ @@ -2898,6 +2902,28 @@ def test_experiment_fetch_uses_generated_paginated_fetch(self): call("test-experiment-id", body={"limit": 1, "cursor": "next-page"}), ] mock_state.api_conn.assert_not_called() + assert list(experiment.fetch(batch_size=1)) == [second_event] + assert api_client.experiments.post_experiment_id_fetch.call_count == 3 + + +@pytest.mark.parametrize("initial_rows", [[], ["first"]]) +def test_object_iterator_keeps_exhausted_snapshot(initial_rows): + fetch = MagicMock(side_effect=[initial_rows, ["new"]]) + iterator = logger.ObjectIterator(fetch) + fetch.assert_not_called() + assert list(iterator) == initial_rows + with pytest.raises(StopIteration): + next(iterator) + fetch.assert_called_once() + + +def test_object_iterator_retries_failed_initial_fetch(): + fetch = MagicMock(side_effect=[RuntimeError("fetch failed"), ["row"]]) + iterator = logger.ObjectIterator(fetch) + with pytest.raises(RuntimeError, match="fetch failed"): + next(iterator) + assert list(iterator) == ["row"] + assert fetch.call_count == 2 class TestLoggerProjectGroup(TestCase): @@ -3227,6 +3253,41 @@ def test_dataset_delete_invalidates_fetched_records(self): self.assertEqual(list(dataset), [retained]) self.assertEqual(fetch.call_count, 2) + def test_dataset_delete_during_iteration_preserves_snapshot(self): + mock_state = MagicMock() + fetch = mock_state.api_client.return_value.datasets.post_dataset_id_fetch + rows = [{"id": str(index), "expected": index} for index in range(4)] + fetch.side_effect = lambda *args, **kwargs: {"events": list(rows)} + metadata = logger.ProjectDatasetMetadata( + project=logger.ObjectMetadata(id="test-project-id", name="test-project", full_info={}), + dataset=logger.ObjectMetadata(id="test-dataset-id", name="test-dataset", full_info={}), + ) + dataset = logger.Dataset( + lazy_metadata=LazyValue(lambda: metadata, use_mutex=False), legacy=False, state=mock_state + ) + background_logger = mock_state.global_bg_logger.return_value + pending_deletes = [] + background_logger.log.side_effect = lambda event: pending_deletes.append(event.get()["id"]) + + def flush_deletes(): + rows[:] = [row for row in rows if row["id"] not in pending_deletes] + pending_deletes.clear() + + background_logger.flush.side_effect = flush_deletes + iterator = iter(dataset) + fetch.assert_not_called() + visited = [] + for row in iterator: + visited.append(row["id"]) + dataset.delete(row["id"]) + dataset.flush() + + self.assertEqual(visited, ["0", "1", "2", "3"]) + self.assertEqual(rows, []) + self.assertEqual(fetch.call_count, 1) + self.assertEqual(list(dataset), []) + self.assertEqual(fetch.call_count, 2) + def test_dataset_summary_uses_generated_resource_and_configured_public_url(self): mock_state = MagicMock() mock_state.app_public_url = "https://public.example.com"