diff --git a/docs/upgrading/upgrading_to_v2.md b/docs/upgrading/upgrading_to_v2.md index f7e0bd036b..2c7b25d9e8 100644 --- a/docs/upgrading/upgrading_to_v2.md +++ b/docs/upgrading/upgrading_to_v2.md @@ -343,6 +343,12 @@ class MyHttpClient(HttpClient): `HttpResponse` is now an abstract base class instead of a protocol, so a custom response subclasses it. It has three new abstract members, `url`, `headers_set_cookie` and `aclose`, and works as an async context manager that closes the response. +## `get_pending_count` on request loaders + +`RequestLoader` has a new abstract `get_pending_count` method, which returns the number of requests that haven't been handled yet. If you have a custom request loader or manager, implement it. `RequestList` reads its source lazily, so like `get_total_count`, its `get_pending_count` leaves out the requests it hasn't fetched yet. + +`RequestManagerTandem.get_total_count` no longer counts a request twice once it moves from the loader to the manager. In v1, it added up the totals of both, so the crawler's status message showed too many pages in total. + ## Logs changed If you configure loggers by name or match log messages, update them for the following changes. diff --git a/src/crawlee/request_loaders/_request_list.py b/src/crawlee/request_loaders/_request_list.py index 18d2f71d85..fa6e2324d7 100644 --- a/src/crawlee/request_loaders/_request_list.py +++ b/src/crawlee/request_loaders/_request_list.py @@ -159,6 +159,11 @@ async def get_handled_count(self) -> int: async def get_total_count(self) -> int: return self._assumed_total_count + @override + async def get_pending_count(self) -> int: + # The source is read lazily, so like the total count, this leaves out the requests not fetched yet. + return self._assumed_total_count - self._handled_count + @override async def check_readiness(self) -> RequestLoaderStatus: await self._ensure_next_request() diff --git a/src/crawlee/request_loaders/_request_loader.py b/src/crawlee/request_loaders/_request_loader.py index 747610bdf5..2bf7d5a23d 100644 --- a/src/crawlee/request_loaders/_request_loader.py +++ b/src/crawlee/request_loaders/_request_loader.py @@ -24,7 +24,7 @@ class RequestLoader(ABC): Key responsibilities: - Fetching the next request to be processed. - Marking requests as successfully handled after processing. - - Managing state information such as the total and handled request counts. + - Managing state information such as the total, pending and handled request counts. """ @abstractmethod @@ -35,6 +35,10 @@ async def get_handled_count(self) -> int: async def get_total_count(self) -> int: """Get an offline approximation of the total number of requests in the loader (i.e. pending + handled).""" + @abstractmethod + async def get_pending_count(self) -> int: + """Get an approximation of the number of requests in the loader that haven't been handled yet.""" + @abstractmethod async def check_readiness(self) -> RequestSourceStatus: """Report whether the loader has a request to hand over, is waiting for one, or is done. diff --git a/src/crawlee/request_loaders/_request_manager_tandem.py b/src/crawlee/request_loaders/_request_manager_tandem.py index fb79f55374..4c9aee5d95 100644 --- a/src/crawlee/request_loaders/_request_manager_tandem.py +++ b/src/crawlee/request_loaders/_request_manager_tandem.py @@ -42,7 +42,14 @@ async def get_handled_count(self) -> int: @override async def get_total_count(self) -> int: - return (await self._read_only_loader.get_total_count()) + (await self._read_write_manager.get_total_count()) + # A request moved into the manager counts as handled in the loader, so only the loader's pending ones are added. + loader_pending = await self._read_only_loader.get_pending_count() + return loader_pending + (await self._read_write_manager.get_total_count()) + + @override + async def get_pending_count(self) -> int: + loader_pending = await self._read_only_loader.get_pending_count() + return loader_pending + (await self._read_write_manager.get_pending_count()) @override async def check_readiness(self) -> RequestSourceStatus: diff --git a/src/crawlee/request_loaders/_sitemap_request_loader.py b/src/crawlee/request_loaders/_sitemap_request_loader.py index 9f51c458ca..5ebfc38331 100644 --- a/src/crawlee/request_loaders/_sitemap_request_loader.py +++ b/src/crawlee/request_loaders/_sitemap_request_loader.py @@ -209,6 +209,12 @@ async def get_handled_count(self) -> int: state = await self._get_state() return state.handled_count + @override + async def get_pending_count(self) -> int: + """Return the number of URLs found so far that haven't been handled yet.""" + state = await self._get_state() + return len(state.url_queue) + len(state.in_progress) + @override async def check_readiness(self) -> RequestLoaderStatus: state = await self._get_state() diff --git a/src/crawlee/request_loaders/_throttling_request_manager.py b/src/crawlee/request_loaders/_throttling_request_manager.py index 750d11e56b..6913488dbf 100644 --- a/src/crawlee/request_loaders/_throttling_request_manager.py +++ b/src/crawlee/request_loaders/_throttling_request_manager.py @@ -305,6 +305,15 @@ async def get_total_count(self) -> int: ) return sum(counts) - self._migrated_from_inner + @override + async def get_pending_count(self) -> int: + await self._ensure_sub_managers() + counts = await asyncio.gather( + self._inner.get_pending_count(), *(sm.get_pending_count() for sm in self._sub_managers.values()) + ) + # `inner` counts a migrated request as handled, so it is pending only in its sub-manager and needs no fix-up. + return sum(counts) + @override async def check_readiness(self) -> RequestSourceStatus: """Report whether anything can be dispatched now, and otherwise when it can be or why it never will. diff --git a/src/crawlee/storages/_request_queue.py b/src/crawlee/storages/_request_queue.py index a2f4ca6ad6..b70a61c6ff 100644 --- a/src/crawlee/storages/_request_queue.py +++ b/src/crawlee/storages/_request_queue.py @@ -116,6 +116,11 @@ async def get_total_count(self) -> int: metadata = await self._backend.get_metadata() return metadata.total_request_count + @override + async def get_pending_count(self) -> int: + metadata = await self._backend.get_metadata() + return metadata.total_request_count - metadata.handled_request_count + @override @classmethod async def open( diff --git a/tests/unit/request_loaders/test_request_list.py b/tests/unit/request_loaders/test_request_list.py index 36f510b84c..e1102b2c16 100644 --- a/tests/unit/request_loaders/test_request_list.py +++ b/tests/unit/request_loaders/test_request_list.py @@ -36,6 +36,19 @@ async def generator() -> AsyncGenerator[str]: assert await request_list.get_handled_count() == 3 +async def test_pending_count_counts_fetched() -> None: + """Test that the pending count covers only the fetched requests not handled yet, since the source is lazy.""" + request_list = RequestList(['https://a.placeholder.com', 'https://b.placeholder.com', 'https://c.placeholder.com']) + + first = await request_list.fetch_next_request() + second = await request_list.fetch_next_request() + assert first is not None + assert second is not None + await request_list.mark_request_as_handled(first) + + assert await request_list.get_pending_count() == 1 + + async def test_readiness_during_traversal() -> None: """Test that the list waits for fetched requests and finishes once they're all handled.""" request_list = RequestList(['https://a.placeholder.com', 'https://b.placeholder.com', 'https://c.placeholder.com']) diff --git a/tests/unit/request_loaders/test_sitemap_request_loader.py b/tests/unit/request_loaders/test_sitemap_request_loader.py index 7bc1b2be03..60b07925ad 100644 --- a/tests/unit/request_loaders/test_sitemap_request_loader.py +++ b/tests/unit/request_loaders/test_sitemap_request_loader.py @@ -229,6 +229,23 @@ async def test_fetch_none_while_in_progress(server_url: URL, http_client: HttpCl assert (await sitemap_loader.check_readiness()).status == 'waiting' +async def test_pending_count_during_traversal(server_url: URL, http_client: HttpClient) -> None: + """Test that the pending count covers the queued URLs and the ones in progress.""" + sitemap_url = (server_url / 'sitemap.xml').with_query( + base64=encode_base64(get_basic_sitemap(url=server_url).encode()) + ) + sitemap_loader = SitemapRequestLoader([str(sitemap_url)], http_client=http_client, enqueue_strategy='all') + + first = await sitemap_loader.fetch_next_request() + second = await sitemap_loader.fetch_next_request() + assert first is not None + assert second is not None + await sitemap_loader.mark_request_as_handled(first) + + assert await poll_until_condition(sitemap_loader.get_total_count, lambda count: count == 5) == 5 + assert await sitemap_loader.get_pending_count() == 4 + + async def test_abort_sitemap_loading(server_url: URL, http_client: HttpClient) -> None: sitemap_url = (server_url / 'sitemap.xml').with_query( base64=encode_base64(get_basic_sitemap(url=server_url).encode()) diff --git a/tests/unit/storages/test_request_manager_tandem.py b/tests/unit/storages/test_request_manager_tandem.py index 1e7c9bc081..e45a9792d4 100644 --- a/tests/unit/storages/test_request_manager_tandem.py +++ b/tests/unit/storages/test_request_manager_tandem.py @@ -10,6 +10,7 @@ from crawlee.request_loaders import ( FinishedStatus, ReadyStatus, + RequestList, RequestLoader, RequestManager, RequestManagerTandem, @@ -157,3 +158,28 @@ async def test_transfers_only_when_loader_ready() -> None: assert request is not None assert request.url == 'https://a.placeholder.com' loader.fetch_next_request.assert_not_awaited() + + +async def test_total_count_counts_moved_once() -> None: + """Test that a request moved from the loader to the manager counts once in the total.""" + tandem = RequestManagerTandem( + RequestList(['https://a.placeholder.com', 'https://b.placeholder.com']), await RequestQueue.open() + ) + + while (await tandem.check_readiness()).status != 'finished': + request = await tandem.fetch_next_request() + assert request is not None + await tandem.mark_request_as_handled(request) + + assert await tandem.get_total_count() == 2 + assert await tandem.get_pending_count() == 0 + + +async def test_pending_count_sums_sources() -> None: + """Test that the pending count adds up the pending requests of the loader and the manager.""" + loader = create_autospec(RequestLoader, instance=True, spec_set=True) + loader.get_pending_count.return_value = 2 + manager = create_autospec(RequestManager, instance=True, spec_set=True) + manager.get_pending_count.return_value = 3 + + assert await RequestManagerTandem(loader, manager).get_pending_count() == 5 diff --git a/tests/unit/storages/test_request_queue.py b/tests/unit/storages/test_request_queue.py index 6e0ec5bdf7..9b553f164b 100644 --- a/tests/unit/storages/test_request_queue.py +++ b/tests/unit/storages/test_request_queue.py @@ -628,6 +628,17 @@ async def test_add_requests_wait_for_all( assert await rq.get_total_count() == 15 +async def test_get_pending_count(rq: RequestQueue) -> None: + """Test that the pending count leaves out the handled requests.""" + await rq.add_requests(['https://example.com/1', 'https://example.com/2', 'https://example.com/3']) + + request = await rq.fetch_next_request() + assert request is not None + await rq.mark_request_as_handled(request) + + assert await rq.get_pending_count() == 2 + + async def test_mark_non_existent_request_as_handled(rq: RequestQueue) -> None: """Test marking a non-existent request as handled.""" # Create a request that hasn't been added to the queue diff --git a/tests/unit/test_throttling_request_manager.py b/tests/unit/test_throttling_request_manager.py index 1ef7e072f7..9ef3896c2e 100644 --- a/tests/unit/test_throttling_request_manager.py +++ b/tests/unit/test_throttling_request_manager.py @@ -891,6 +891,32 @@ async def test_get_total_count_aggregates(manager: ThrottlingRequestManager[Requ assert await manager.get_total_count() == 2 +async def test_get_pending_count_aggregates(manager: ThrottlingRequestManager[RequestQueue]) -> None: + """Test that the pending count sums inner and every sub-manager, without the handled requests.""" + await manager.add_request(f'https://{THROTTLED_DOMAIN}/page1') + await manager.add_request(f'https://{NON_THROTTLED_DOMAIN}/page1') + + request = await manager.fetch_next_request() + assert request is not None + await manager.mark_request_as_handled(request) + + assert await manager.get_pending_count() == 1 + + +async def test_pending_count_after_migration( + manager: ThrottlingRequestManager[RequestQueue], + inner_queue: RequestQueue, +) -> None: + """Test that a request migrated from inner to its domain's sub-manager counts once as pending.""" + await inner_queue.add_requests([f'https://{THROTTLED_DOMAIN}/page{i}' for i in range(2)]) + manager.record_domain_delay(f'https://{THROTTLED_DOMAIN}/', retry_after=timedelta(seconds=60)) + + assert await manager.fetch_next_request() is None + + assert await inner_queue.get_handled_count() == 2 + assert await manager.get_pending_count() == 2 + + async def test_readiness_aggregates(manager: ThrottlingRequestManager[RequestQueue]) -> None: """Test that the readiness follows a request through the sub-manager it's routed to.""" assert (await manager.check_readiness()).status == 'finished'