fixing generate public url for KVS records

Author: vdusek

src/apify/storage_clients/_apify/_dataset_client.py

@@ -14,7 +14,6 @@
 
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator
-    from datetime import datetime
 
     from apify_client.clients import DatasetClientAsync
     from crawlee._types import JsonSerializable
@@ -38,31 +37,23 @@ class ApifyDatasetClient(DatasetClient):
     def __init__(
         self,
         *,
-        id: str,
-        name: str | None,
-        created_at: datetime,
-        accessed_at: datetime,
-        modified_at: datetime,
-        item_count: int,
+        metadata: DatasetMetadata,
         api_client: DatasetClientAsync,
+        api_public_base_url: str,
         lock: asyncio.Lock,
     ) -> None:
         """Initialize a new instance.
 
         Preferably use the `ApifyDatasetClient.open` class method to create a new instance.
         """
-        self._metadata = DatasetMetadata(
-            id=id,
-            name=name,
-            created_at=created_at,
-            accessed_at=accessed_at,
-            modified_at=modified_at,
-            item_count=item_count,
-        )
+        self._metadata = metadata
 
         self._api_client = api_client
         """The Apify dataset client for API operations."""
 
+        self._api_public_base_url = api_public_base_url
+        """The public base URL for accessing the key-value store records."""
+
         self._lock = lock
         """A lock to ensure that only one operation is performed at a time."""
 
@@ -109,6 +100,13 @@ async def open(
         if not api_url:
             raise ValueError(f'Apify storage client requires a valid API URL in Configuration (api_url={api_url}).')
 
+        api_public_base_url = getattr(configuration, 'api_public_base_url', None)
+        if not api_public_base_url:
+            raise ValueError(
+                'Apify storage client requires a valid API public base URL in Configuration '
+                f'(api_public_base_url={api_public_base_url}).'
+            )
+
         if id and name:
             raise ValueError('Only one of "id" or "name" can be specified, not both.')
 
@@ -144,13 +142,9 @@ async def open(
         metadata = DatasetMetadata.model_validate(await apify_dataset_client.get())
 
         return cls(
-            id=metadata.id,
-            name=metadata.name,
-            created_at=metadata.created_at,
-            accessed_at=metadata.accessed_at,
-            modified_at=metadata.modified_at,
-            item_count=metadata.item_count,
+            metadata=metadata,
             api_client=apify_dataset_client,
+            api_public_base_url=api_public_base_url,
             lock=asyncio.Lock(),
         )
 

src/apify/storage_clients/_apify/_key_value_store_client.py

@@ -16,7 +16,6 @@
 
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator
-    from datetime import datetime
 
     from apify_client.clients import KeyValueStoreClientAsync
     from crawlee.configuration import Configuration
@@ -30,28 +29,22 @@ class ApifyKeyValueStoreClient(KeyValueStoreClient):
     def __init__(
         self,
         *,
-        id: str,
-        name: str | None,
-        created_at: datetime,
-        accessed_at: datetime,
-        modified_at: datetime,
+        metadata: KeyValueStoreMetadata,
         api_client: KeyValueStoreClientAsync,
+        api_public_base_url: str,
         lock: asyncio.Lock,
     ) -> None:
         """Initialize a new instance.
 
         Preferably use the `ApifyKeyValueStoreClient.open` class method to create a new instance.
         """
-        self._metadata = KeyValueStoreMetadata(
-            id=id,
-            name=name,
-            created_at=created_at,
-            accessed_at=accessed_at,
-            modified_at=modified_at,
-        )
+        self._metadata = metadata
 
         self._api_client = api_client
-        """The Apify key-value store client for API operations."""
+        """The Apify KVS client for API operations."""
+
+        self._api_public_base_url = api_public_base_url
+        """The public base URL for accessing the key-value store records."""
 
         self._lock = lock
         """A lock to ensure that only one operation is performed at a time."""
@@ -99,6 +92,13 @@ async def open(
         if not api_url:
             raise ValueError(f'Apify storage client requires a valid API URL in Configuration (api_url={api_url}).')
 
+        api_public_base_url = getattr(configuration, 'api_public_base_url', None)
+        if not api_public_base_url:
+            raise ValueError(
+                'Apify storage client requires a valid API public base URL in Configuration '
+                f'(api_public_base_url={api_public_base_url}).'
+            )
+
         if id and name:
             raise ValueError('Only one of "id" or "name" can be specified, not both.')
 
@@ -134,12 +134,9 @@ async def open(
         metadata = KeyValueStoreMetadata.model_validate(await apify_kvs_client.get())
 
         return cls(
-            id=metadata.id,
-            name=metadata.name,
-            created_at=metadata.created_at,
-            accessed_at=metadata.accessed_at,
-            modified_at=metadata.modified_at,
+            metadata=metadata,
             api_client=apify_kvs_client,
+            api_public_base_url=api_public_base_url,
             lock=asyncio.Lock(),
         )
 
@@ -230,7 +227,7 @@ async def get_public_url(self, key: str) -> str:
             raise ValueError('resource_id cannot be None when generating a public URL')
 
         public_url = (
-            URL(self._api_client.base_url) / 'v2' / 'key-value-stores' / self._api_client.resource_id / 'records' / key
+            URL(self._api_public_base_url) / 'v2' / 'key-value-stores' / self._api_client.resource_id / 'records' / key
         )
 
         if self.metadata.model_extra is not None:

src/apify/storage_clients/_apify/_request_queue_client.py

@@ -38,39 +38,23 @@ class ApifyRequestQueueClient(RequestQueueClient):
     def __init__(
         self,
         *,
-        id: str,
-        name: str | None,
-        created_at: datetime,
-        accessed_at: datetime,
-        modified_at: datetime,
-        had_multiple_clients: bool,
-        handled_request_count: int,
-        pending_request_count: int,
-        stats: dict,
-        total_request_count: int,
+        metadata: RequestQueueMetadata,
         api_client: RequestQueueClientAsync,
+        api_public_base_url: str,
         lock: asyncio.Lock,
     ) -> None:
         """Initialize a new instance.
 
         Preferably use the `ApifyRequestQueueClient.open` class method to create a new instance.
         """
-        self._metadata = RequestQueueMetadata(
-            id=id,
-            name=name,
-            created_at=created_at,
-            accessed_at=accessed_at,
-            modified_at=modified_at,
-            had_multiple_clients=had_multiple_clients,
-            handled_request_count=handled_request_count,
-            pending_request_count=pending_request_count,
-            stats=stats,
-            total_request_count=total_request_count,
-        )
+        self._metadata = metadata
 
         self._api_client = api_client
         """The Apify request queue client for API operations."""
 
+        self._api_public_base_url = api_public_base_url
+        """The public base URL for accessing the key-value store records."""
+
         self._lock = lock
         """A lock to ensure that only one operation is performed at a time."""
 
@@ -130,6 +114,13 @@ async def open(
         if not api_url:
             raise ValueError(f'Apify storage client requires a valid API URL in Configuration (api_url={api_url}).')
 
+        api_public_base_url = getattr(configuration, 'api_public_base_url', None)
+        if not api_public_base_url:
+            raise ValueError(
+                'Apify storage client requires a valid API public base URL in Configuration '
+                f'(api_public_base_url={api_public_base_url}).'
+            )
+
         if id and name:
             raise ValueError('Only one of "id" or "name" can be specified, not both.')
 
@@ -165,17 +156,9 @@ async def open(
         metadata = RequestQueueMetadata.model_validate(await apify_rq_client.get())
 
         return cls(
-            id=metadata.id,
-            name=metadata.name,
-            created_at=metadata.created_at,
-            accessed_at=metadata.accessed_at,
-            modified_at=metadata.modified_at,
-            had_multiple_clients=metadata.had_multiple_clients,
-            handled_request_count=metadata.handled_request_count,
-            pending_request_count=metadata.pending_request_count,
-            stats=metadata.stats,
-            total_request_count=metadata.total_request_count,
+            metadata=metadata,
             api_client=apify_rq_client,
+            api_public_base_url=api_public_base_url,
             lock=asyncio.Lock(),
         )
 

tests/integration/test_actor_key_value_store.py

@@ -217,13 +217,13 @@ async def main() -> None:
             await kvs.set_value(record_key, {'exposedData': 'test'}, 'application/json')
 
             record_url = await kvs.get_public_url(record_key)
-
             signature = create_hmac_signature(url_signing_secret_key, record_key)
-            assert (
-                record_url
-                == f'{public_api_url}/v2/key-value-stores/{default_kvs_id}/records/{record_key}?signature={signature}'
+            expected_record_url = (
+                f'{public_api_url}/v2/key-value-stores/{default_kvs_id}/records/{record_key}?signature={signature}'
             )
 
+            assert record_url == expected_record_url
+
     actor = await make_actor(label='kvs-get-public-url', main_func=main)
     run_result = await run_actor(actor)