From 883d285b869d370d807c8ad9a5697848deddcb4d Mon Sep 17 00:00:00 2001 From: Anton Burnashev Date: Tue, 2 Apr 2024 17:29:55 +0300 Subject: [PATCH] Ensure total count is an integer in OffsetPaginator --- dlt/sources/helpers/rest_client/paginators.py | 8 ++++++++ .../sources/helpers/rest_client/test_paginators.py | 13 +++++++++++++ 2 files changed, 21 insertions(+) diff --git a/dlt/sources/helpers/rest_client/paginators.py b/dlt/sources/helpers/rest_client/paginators.py index c098ea667f..6d77e844c8 100644 --- a/dlt/sources/helpers/rest_client/paginators.py +++ b/dlt/sources/helpers/rest_client/paginators.py @@ -87,6 +87,14 @@ def update_state(self, response: Response) -> None: f"Total count not found in response for {self.__class__.__name__}" ) + try: + total = int(total) + except ValueError: + raise ValueError( + f"Total count is not an integer in response for {self.__class__.__name__}. " + f"Expected an integer, got {total}" + ) + self.offset += self.limit if self.offset >= total: diff --git a/tests/sources/helpers/rest_client/test_paginators.py b/tests/sources/helpers/rest_client/test_paginators.py index cc4dea65dc..43b6dd7951 100644 --- a/tests/sources/helpers/rest_client/test_paginators.py +++ b/tests/sources/helpers/rest_client/test_paginators.py @@ -75,6 +75,19 @@ def test_update_state(self): paginator.update_state(response) assert paginator.has_next_page is False + def test_update_state_with_string_total(self): + paginator = OffsetPaginator(0, 10) + response = Mock(Response, json=lambda: {"total": "20"}) + paginator.update_state(response) + assert paginator.offset == 10 + assert paginator.has_next_page is True + + def test_update_state_with_invalid_total(self): + paginator = OffsetPaginator(0, 10) + response = Mock(Response, json=lambda: {"total": "invalid"}) + with pytest.raises(ValueError): + paginator.update_state(response) + def test_update_state_without_total(self): paginator = OffsetPaginator(0, 10) response = Mock(Response, json=lambda: {})