diff --git a/src/ralph/backends/data/async_lrs.py b/src/ralph/backends/data/async_lrs.py index 24b363470..1e767777b 100644 --- a/src/ralph/backends/data/async_lrs.py +++ b/src/ralph/backends/data/async_lrs.py @@ -5,7 +5,7 @@ from urllib.parse import ParseResult, parse_qs, urljoin, urlparse from httpx import AsyncClient, HTTPError, HTTPStatusError, RequestError -from pydantic import AnyHttpUrl, parse_obj_as +from pydantic import AnyHttpUrl, PositiveInt, parse_obj_as from ralph.backends.data.base import ( AsyncWritable, @@ -76,8 +76,8 @@ async def read( # noqa: PLR0913 chunk_size: Optional[int] = None, raw_output: bool = False, ignore_errors: bool = False, - prefetch: Optional[int] = None, - max_statements: Optional[int] = None, + prefetch: Optional[PositiveInt] = None, + max_statements: Optional[PositiveInt] = None, ) -> Union[AsyncIterator[bytes], AsyncIterator[dict]]: """Get statements from LRS `target` endpoint. @@ -99,12 +99,8 @@ async def read( # noqa: PLR0913 ignore_errors (bool): If `True`, errors during the read operation are ignored and logged. If `False` (default), a `BackendException` is raised if an error occurs. - prefetch: The number of records to prefetch (queue) while yielding. - If `prefetch` is `None` or `0` it defaults to `1` - no records are - prefetched. - If `prefetch` is less than zero, all records are prefetched. - Caution: setting `prefetch<0` might potentially lead to large amounts - of API calls and to the memory filling up. + prefetch (int): The number of records to prefetch (queue) while yielding. + If `prefetch` is `None` it defaults to `1` - no records are prefetched. max_statements: The maximum number of statements to yield. """ statements = super().read( diff --git a/src/ralph/backends/data/lrs.py b/src/ralph/backends/data/lrs.py index 9eaebfbd0..01f291fb9 100644 --- a/src/ralph/backends/data/lrs.py +++ b/src/ralph/backends/data/lrs.py @@ -5,7 +5,7 @@ from urllib.parse import ParseResult, parse_qs, urljoin, urlparse from httpx import Client, HTTPError, HTTPStatusError, RequestError -from pydantic import AnyHttpUrl, BaseModel, Field, parse_obj_as +from pydantic import AnyHttpUrl, BaseModel, Field, PositiveInt, parse_obj_as from ralph.backends.data.base import ( BaseDataBackend, @@ -117,7 +117,7 @@ def read( # noqa: PLR0913 chunk_size: Optional[int] = None, raw_output: bool = False, ignore_errors: bool = False, - max_statements: Optional[int] = None, + max_statements: Optional[PositiveInt] = None, ) -> Union[Iterator[bytes], Iterator[dict]]: """Get statements from LRS `target` endpoint. diff --git a/tests/backends/data/test_async_lrs.py b/tests/backends/data/test_async_lrs.py index 6dd4377a3..e504b287b 100644 --- a/tests/backends/data/test_async_lrs.py +++ b/tests/backends/data/test_async_lrs.py @@ -708,7 +708,7 @@ async def _simulate_slow_processing(): without_prefetch_duration = time.time() - time_1 time_2 = time.time() - async for _ in backend.read(target=targets[0], chunk_size=chunk_size, prefetch=-1): + async for _ in backend.read(target=targets[0], chunk_size=chunk_size, prefetch=100): await _simulate_slow_processing() prefetch_duration = time.time() - time_2