fix: correct handled/pending counts when reclaiming a previously handled request

Author: vdusek

src/apify/storage_clients/_apify/_request_queue_shared_client.py

@@ -250,7 +250,10 @@ async def reclaim_request(
         """Specific implementation of this method for the RQ shared access mode."""
         # Check if the request was marked as handled and clear it. When reclaiming,
         # we want to put the request back for processing.
-        if request.was_already_handled:
+        # Capture this before clearing `handled_at`, otherwise the computed `was_already_handled` property
+        # would always be False below and the metadata counters would never be adjusted.
+        was_already_handled = request.was_already_handled
+        if was_already_handled:
             request.handled_at = None
 
         # Reclaim with lock to prevent race conditions that could lead to double processing of the same request.
@@ -262,7 +265,7 @@ async def reclaim_request(
 
                 # If the request was previously handled, decrement our handled count since
                 # we're putting it back for processing.
-                if request.was_already_handled and not processed_request.was_already_handled:
+                if was_already_handled and not processed_request.was_already_handled:
                     self.metadata.handled_request_count -= 1
                     self.metadata.pending_request_count += 1
 

src/apify/storage_clients/_apify/_request_queue_single_client.py

@@ -250,7 +250,10 @@ async def reclaim_request(
 
         request_id = unique_key_to_request_id(request.unique_key)
 
-        if request.was_already_handled:
+        # Capture this before clearing `handled_at`, otherwise the computed `was_already_handled` property
+        # would always be False below and the metadata counters would never be adjusted.
+        was_already_handled = request.was_already_handled
+        if was_already_handled:
             request.handled_at = None
 
         try:
@@ -271,7 +274,7 @@ async def reclaim_request(
             processed_request.unique_key = request.unique_key
             # If the request was previously handled, decrement our handled count since
             # we're putting it back for processing.
-            if request.was_already_handled and not processed_request.was_already_handled:
+            if was_already_handled and not processed_request.was_already_handled:
                 self.metadata.handled_request_count -= 1
                 self.metadata.pending_request_count += 1
 

tests/unit/storage_clients/test_apify_request_queue_client.py

@@ -1,25 +1,27 @@
 from __future__ import annotations
 
 from datetime import UTC, datetime
+from typing import TYPE_CHECKING
 from unittest.mock import AsyncMock
 
 import pytest
 
 from apify_client._models import Request as ClientRequest
-from apify_client._models import RequestQueueHead
+from apify_client._models import RequestQueueHead, RequestRegistration
 from crawlee.storage_clients.models import RequestQueueMetadata
 
+from apify import Request
+from apify.storage_clients._apify._request_queue_shared_client import ApifyRequestQueueSharedClient
 from apify.storage_clients._apify._request_queue_single_client import ApifyRequestQueueSingleClient
 from apify.storage_clients._apify._utils import unique_key_to_request_id
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
 
-def _make_single_client(
-    api_client: AsyncMock | None = None,
-) -> tuple[ApifyRequestQueueSingleClient, AsyncMock]:
-    if api_client is None:
-        api_client = AsyncMock()
+
+def _make_metadata() -> RequestQueueMetadata:
     now = datetime.now(tz=UTC)
-    metadata = RequestQueueMetadata(
+    return RequestQueueMetadata(
         id='test-rq-id',
         name='test-rq',
         accessed_at=now,
@@ -30,7 +32,28 @@ def _make_single_client(
         pending_request_count=0,
         total_request_count=0,
     )
-    client = ApifyRequestQueueSingleClient(api_client=api_client, metadata=metadata, cache_size=100)
+
+
+def _make_single_client(
+    api_client: AsyncMock | None = None,
+) -> tuple[ApifyRequestQueueSingleClient, AsyncMock]:
+    if api_client is None:
+        api_client = AsyncMock()
+    client = ApifyRequestQueueSingleClient(api_client=api_client, metadata=_make_metadata(), cache_size=100)
+    return client, api_client
+
+
+def _make_shared_client(
+    api_client: AsyncMock | None = None,
+) -> tuple[ApifyRequestQueueSharedClient, AsyncMock]:
+    if api_client is None:
+        api_client = AsyncMock()
+    client = ApifyRequestQueueSharedClient(
+        api_client=api_client,
+        metadata=_make_metadata(),
+        cache_size=100,
+        metadata_getter=AsyncMock(),
+    )
     return client, api_client
 
 
@@ -136,3 +159,34 @@ async def test_fetch_next_request_skips_already_handled() -> None:
     assert result is None, 'Already-handled request must not be fetched.'
     assert request_id not in client._requests_in_progress, 'Handled request must not be left in progress.'
     assert request_id in client._requests_already_handled, 'Handled request id should be cached for deduplication.'
+
+
+@pytest.mark.parametrize(
+    'make_client',
+    [_make_single_client, _make_shared_client],
+    ids=['single_client', 'shared_client'],
+)
+async def test_reclaim_previously_handled_adjusts_counts(
+    make_client: Callable[[], tuple[ApifyRequestQueueSingleClient | ApifyRequestQueueSharedClient, AsyncMock]],
+) -> None:
+    """Reclaiming a previously handled request must move it from handled back to pending in the metadata."""
+    client, api_client = make_client()
+    client.metadata.handled_request_count = 1
+    client.metadata.pending_request_count = 0
+
+    unique_key = 'https://example.com'
+    request_id = unique_key_to_request_id(unique_key)
+    request = Request.from_url(unique_key, unique_key=unique_key)
+    request.handled_at = datetime.now(tz=UTC)
+
+    # After reclaiming, the platform reports the request as no longer handled.
+    api_client.update_request = AsyncMock(
+        return_value=RequestRegistration.model_validate(
+            {'requestId': request_id, 'wasAlreadyPresent': True, 'wasAlreadyHandled': False}
+        )
+    )
+
+    await client.reclaim_request(request)
+
+    assert client.metadata.handled_request_count == 0, 'Reclaimed request must be removed from the handled count.'
+    assert client.metadata.pending_request_count == 1, 'Reclaimed request must be added back to the pending count.'