Merge remote-tracking branch 'origin/master' into no-locking-queue

Author: Pijukatel

src/apify/storage_clients/_apify/_models.py

@@ -5,7 +5,7 @@
 
 from pydantic import BaseModel, ConfigDict, Field
 
-from crawlee.storage_clients.models import KeyValueStoreMetadata
+from crawlee.storage_clients.models import KeyValueStoreMetadata, RequestQueueMetadata
 
 from apify import Request
 from apify._utils import docs_group
@@ -105,3 +105,27 @@ class CachedRequest(BaseModel):
 
     lock_expires_at: datetime | None = None
     """The expiration time of the lock on the request."""
+
+
+class RequestQueueStats(BaseModel):
+    model_config = ConfigDict(populate_by_name=True)
+
+    delete_count: Annotated[int, Field(alias='deleteCount', default=0)]
+    """"The number of request queue deletes."""
+
+    head_item_read_count: Annotated[int, Field(alias='headItemReadCount', default=0)]
+    """The number of request queue head reads."""
+
+    read_count: Annotated[int, Field(alias='readCount', default=0)]
+    """The number of request queue reads."""
+
+    storage_bytes: Annotated[int, Field(alias='storageBytes', default=0)]
+    """Storage size in Bytes."""
+
+    write_count: Annotated[int, Field(alias='writeCount', default=0)]
+    """The number of request queue writes."""
+
+
+class ApifyRequestQueueMetadata(RequestQueueMetadata):
+    stats: Annotated[RequestQueueStats, Field(alias='stats', default_factory=RequestQueueStats)]
+    """Additional statistics about the request queue."""

src/apify/storage_clients/_apify/_request_queue_client.py

@@ -14,6 +14,7 @@
 from crawlee.storage_clients.models import RequestQueueMetadata
 from crawlee.storages import RequestQueue
 
+from ._models import ApifyRequestQueueMetadata, RequestQueueStats
 from ._utils import AliasResolver
 
 if TYPE_CHECKING:
@@ -70,7 +71,7 @@ def __init__(
         """Additional data related to the RequestQueue."""
 
     @override
-    async def get_metadata(self) -> RequestQueueMetadata:
+    async def get_metadata(self) -> ApifyRequestQueueMetadata:
         """Get metadata about the request queue.
 
         Returns:
@@ -81,7 +82,7 @@ async def get_metadata(self) -> RequestQueueMetadata:
         if response is None:
             raise ValueError('Failed to fetch request queue metadata from the API.')
         # Enhance API response by local estimations (API can be delayed few seconds, while local estimation not.)
-        return RequestQueueMetadata(
+        return ApifyRequestQueueMetadata(
             id=response['id'],
             name=response['name'],
             total_request_count=max(response['totalRequestCount'], self._metadata.total_request_count),
@@ -91,6 +92,7 @@ async def get_metadata(self) -> RequestQueueMetadata:
             modified_at=max(response['modifiedAt'], self._metadata.modified_at),
             accessed_at=max(response['accessedAt'], self._metadata.accessed_at),
             had_multiple_clients=response['hadMultipleClients'] or self._metadata.had_multiple_clients,
+            stats=RequestQueueStats.model_validate(response['stats'], by_alias=True),
         )
 
     @classmethod

tests/integration/test_request_queue.py

@@ -2,7 +2,7 @@
 
 import asyncio
 from datetime import datetime, timezone
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 
 import pytest
 
@@ -20,6 +20,7 @@
     from crawlee._types import BasicCrawlingContext
 
     from .conftest import MakeActorFunction, RunActorFunction
+    from apify.storage_clients._apify._models import ApifyRequestQueueMetadata
 
 
 async def test_add_and_fetch_requests(default_request_queue_apify: RequestQueue) -> None:
@@ -1126,14 +1127,9 @@ async def default_handler(context: BasicCrawlingContext) -> None:
     # Check the request queue stats
     await asyncio.sleep(10)  # Wait to be sure that metadata are updated
 
-    # Get raw client, because stats are not exposed in `RequestQueue` class, but are available in raw client
-    # https://github.com/apify/apify-sdk-python/pull/574
-    rq_client = Actor.apify_client.request_queue(request_queue_id=rq.id)
-    _rq = await rq_client.get()
-    assert _rq
-    request_queue_stats = _rq.get('stats', {})
-    Actor.log.info(f'{request_queue_stats=}')
-    assert request_queue_stats['writeCount'] == requests * expected_write_count_per_request
+    metadata = cast('ApifyRequestQueueMetadata', await rq.get_metadata())
+    Actor.log.info(f'{metadata.stats=}')
+    assert metadata.stats.write_count == requests * expected_write_count_per_request
     await rq.drop()
 
 
@@ -1153,13 +1149,8 @@ async def test_cache_initialization(apify_token: str, monkeypatch: pytest.Monkey
 
             # Check that it is correctly in the API
             await asyncio.sleep(10)  # Wait to be sure that metadata are updated
-
-            # Get raw client, because stats are not exposed in `RequestQueue` class, but are available in raw client
-            # https://github.com/apify/apify-sdk-python/pull/574
-            rq_client = Actor.apify_client.request_queue(request_queue_id=rq.id)
-            _rq = await rq_client.get()
-            assert _rq
-            stats_before = _rq.get('stats', {})
+            metadata = cast('ApifyRequestQueueMetadata', await rq.get_metadata())
+            stats_before = metadata.stats
             Actor.log.info(stats_before)
 
             # Clear service locator cache to simulate creating RQ instance from scratch
@@ -1170,15 +1161,37 @@ async def test_cache_initialization(apify_token: str, monkeypatch: pytest.Monkey
             await rq.add_requests(requests)
 
             await asyncio.sleep(10)  # Wait to be sure that metadata are updated
-            _rq = await rq_client.get()
-            assert _rq
-            stats_after = _rq.get('stats', {})
+            metadata = cast('ApifyRequestQueueMetadata', await rq.get_metadata())
+            stats_after = metadata.stats
             Actor.log.info(stats_after)
 
             # Cache was actually initialized, readCount increased
-            assert (stats_after['readCount'] - stats_before['readCount']) == len(requests)
+            assert (stats_after.read_count - stats_before.read_count) == len(requests)
             # Deduplication happened locally, writeCount should be the same
-            assert stats_after['writeCount'] == stats_before['writeCount']
+            assert stats_after.write_count == stats_before.write_count
 
         finally:
             await rq.drop()
+
+
+async def test_request_queue_has_stats(request_queue_force_cloud: RequestQueue) -> None:
+    """Test that Apify based request queue has stats in metadata."""
+
+    add_request_count = 3
+    read_request_count = 2
+
+    await request_queue_force_cloud.add_requests(
+        [Request.from_url(f'http://example.com/{i}') for i in range(add_request_count)]
+    )
+    for _ in range(read_request_count):
+        await request_queue_force_cloud.get_request(Request.from_url('http://example.com/1').unique_key)
+
+    # Wait for stats to become stable
+    await asyncio.sleep(10)
+
+    metadata = await request_queue_force_cloud.get_metadata()
+
+    assert hasattr(metadata, 'stats')
+    apify_metadata = cast('ApifyRequestQueueMetadata', metadata)
+    assert apify_metadata.stats.read_count == read_request_count
+    assert apify_metadata.stats.write_count == add_request_count