Wip, TODO: Solve patching of service_locator from Crawlee

Author: Pijukatel

pyproject.toml

@@ -36,7 +36,7 @@ keywords = [
 dependencies = [
     "apify-client>=2.0.0,<3.0.0",
     "apify-shared>=2.0.0,<3.0.0",
-    "crawlee @ git+https://github.com/apify/crawlee-python.git@storage-clients-and-configurations",
+    "crawlee @ git+https://github.com/apify/crawlee-python.git@storage-clients-and-configurations-2",
     "cachetools>=5.5.0",
     "cryptography>=42.0.0",
     # TODO: ensure compatibility with the latest version of lazy-object-proxy

src/apify/_actor.py

@@ -13,6 +13,7 @@
 
 from apify_client import ApifyClientAsync
 from apify_shared.consts import ActorEnvVars, ActorExitCodes, ApifyEnvVars
+from crawlee import service_locator
 from crawlee.errors import ServiceConflictError
 from crawlee.events import (
     Event,
@@ -25,7 +26,7 @@
 )
 
 from apify._charging import ChargeResult, ChargingManager, ChargingManagerImplementation
-from apify._configuration import Configuration, service_locator
+from apify._configuration import Configuration
 from apify._consts import EVENT_LISTENERS_TIMEOUT
 from apify._crypto import decrypt_input_secrets, load_private_key
 from apify._models import ActorRun

src/apify/_configuration.py

@@ -1,22 +1,26 @@
 from __future__ import annotations
 
 import traceback
+import types
 from datetime import datetime, timedelta
 from decimal import Decimal
 from logging import getLogger
-from typing import Annotated, Any
+from typing import TYPE_CHECKING, Annotated, Any, cast
 
 from pydantic import AliasChoices, BeforeValidator, Field, model_validator
 from typing_extensions import Self, deprecated
 
 import crawlee
-from crawlee._service_locator import ServiceLocator
+from crawlee import service_locator
 from crawlee._utils.models import timedelta_ms
 from crawlee._utils.urls import validate_http_url
 from crawlee.configuration import Configuration as CrawleeConfiguration
 
 from apify._utils import docs_group
 
+if TYPE_CHECKING:
+    from crawlee._service_locator import ServiceLocator
+
 logger = getLogger(__name__)
 
 
@@ -431,20 +435,21 @@ def get_global_configuration(cls) -> Configuration:
         Mostly for the backwards compatibility. It is recommended to use the `service_locator.get_configuration()`
         instead.
         """
-        return service_locator.get_configuration()
+        return cast('Configuration', service_locator.get_configuration())
 
 
-class ApifyServiceLocator(ServiceLocator):
-    """Same as ServiceLocator from Crawlee, but it always returns Apify Configuration."""
+def _patch_service_locator() -> None:
+    """Patch the Crawlee ServiceLocator to ensure that Apify Configuration is always returned."""
+    old_get_configuration = crawlee.service_locator.get_configuration
 
-    def get_configuration(self) -> Configuration:
+    def get_configuration(self: ServiceLocator) -> Configuration:
         # ApifyServiceLocator can store any children of Crawlee Configuration, but in Apify context it is desired to
         # return Apify Configuration.
         if isinstance(self._configuration, Configuration):
             # If Apify configuration was already stored in service locator, return it.
             return self._configuration
 
-        stored_configuration = super().get_configuration()
+        stored_configuration = old_get_configuration()
         apify_configuration = Configuration()
 
         # Ensure the returned configuration is of type Apify Configuration.
@@ -457,7 +462,8 @@ def get_configuration(self) -> Configuration:
 
         return apify_configuration
 
+    crawlee.service_locator.get_configuration = types.MethodType(get_configuration, crawlee.service_locator)
+
 
 # Ensure that ApifyServiceLocator is used to make sure Apify Configuration is used.
-service_locator = ApifyServiceLocator()
-crawlee.service_locator = service_locator
+_patch_service_locator()

src/apify/storage_clients/_apify/_storage_client.py

@@ -4,12 +4,12 @@
 
 from typing_extensions import override
 
+from crawlee import service_locator
 from crawlee.storage_clients._base import StorageClient
 
 from ._dataset_client import ApifyDatasetClient
 from ._key_value_store_client import ApifyKeyValueStoreClient
 from ._request_queue_client import ApifyRequestQueueClient
-from apify._configuration import service_locator
 from apify._utils import docs_group
 
 if TYPE_CHECKING:
@@ -40,12 +40,13 @@ async def create_dataset_client(
         # Import here to avoid circular imports.
         from apify import Configuration as ApifyConfiguration  # noqa: PLC0415
 
-        if isinstance(self._configuration, ApifyConfiguration):
-            return await ApifyDatasetClient.open(id=id, name=name, configuration=self._configuration)
+        configuration = configuration or ApifyConfiguration.get_global_configuration()
+        if isinstance(configuration, ApifyConfiguration):
+            return await ApifyDatasetClient.open(id=id, name=name, configuration=configuration)
 
         raise TypeError(
             f'Expected "configuration" to be an instance of "apify.Configuration", '
-            f'but got {type(self._configuration).__name__} instead.'
+            f'but got {type(configuration).__name__} instead.'
         )
 
     @override
@@ -59,12 +60,13 @@ async def create_kvs_client(
         # Import here to avoid circular imports.
         from apify import Configuration as ApifyConfiguration  # noqa: PLC0415
 
-        if isinstance(self._configuration, ApifyConfiguration):
-            return await ApifyKeyValueStoreClient.open(id=id, name=name, configuration=self._configuration)
+        configuration = configuration or ApifyConfiguration.get_global_configuration()
+        if isinstance(configuration, ApifyConfiguration):
+            return await ApifyKeyValueStoreClient.open(id=id, name=name, configuration=configuration)
 
         raise TypeError(
             f'Expected "configuration" to be an instance of "apify.Configuration", '
-            f'but got {type(self._configuration).__name__} instead.'
+            f'but got {type(configuration).__name__} instead.'
         )
 
     @override
@@ -78,15 +80,11 @@ async def create_rq_client(
         # Import here to avoid circular imports.
         from apify import Configuration as ApifyConfiguration  # noqa: PLC0415
 
-        if isinstance(self._configuration, ApifyConfiguration):
-            return await ApifyRequestQueueClient.open(id=id, name=name, configuration=self._configuration)
+        configuration = configuration or ApifyConfiguration.get_global_configuration()
+        if isinstance(configuration, ApifyConfiguration):
+            return await ApifyRequestQueueClient.open(id=id, name=name, configuration=configuration)
 
         raise TypeError(
             f'Expected "configuration" to be an instance of "apify.Configuration", '
-            f'but got {type(self._configuration).__name__} instead.'
+            f'but got {type(configuration).__name__} instead.'
         )
-
-    @override
-    def create_client(self, configuration: Configuration) -> ApifyStorageClient:
-        """Create a storage client from an existing storage client potentially just replacing the configuration."""
-        return ApifyStorageClient(configuration)

tests/integration/conftest.py

@@ -58,7 +58,7 @@ def _prepare_test_env() -> None:
         service_locator._configuration = None
         service_locator._event_manager = None
         service_locator._storage_client = None
-        service_locator._storage_instance_manager = None
+        service_locator.storage_instance_manager.clear_cache()
 
         # Verify that the test environment was set up correctly.
         assert os.environ.get(ApifyEnvVars.LOCAL_STORAGE_DIR) == str(tmp_path)

tests/unit/actor/test_configuration.py

@@ -1,12 +1,14 @@
+from pathlib import Path
+
 import pytest
 
+from crawlee import service_locator
 from crawlee.configuration import Configuration as CrawleeConfiguration
 from crawlee.crawlers import BasicCrawler
 from crawlee.errors import ServiceConflictError
 
 from apify import Actor
 from apify import Configuration as ApifyConfiguration
-from apify._configuration import service_locator
 
 
 @pytest.mark.parametrize(
@@ -105,9 +107,9 @@ async def test_crawler_implicit_configuration() -> None:
     """Test that crawler and Actor use implicit service_locator based configuration unless explicit configuration
     was passed to them."""
     async with Actor():
-        crawler_1 = BasicCrawler()
+        crawler = BasicCrawler()
 
-    assert service_locator.get_configuration() is crawler_1._service_locator.get_configuration()
+    assert service_locator.get_configuration() is crawler._service_locator.get_configuration()
 
 
 async def test_crawlers_own_configuration() -> None:
@@ -132,30 +134,54 @@ async def test_crawler_global_configuration() -> None:
     service_locator.set_configuration(config_global)
 
     async with Actor():
-        crawler_1 = BasicCrawler()
+        crawler = BasicCrawler()
 
     assert service_locator.get_configuration() is config_global
-    assert crawler_1._service_locator.get_configuration() is config_global
+    assert crawler._service_locator.get_configuration() is config_global
 
 
-async def test_storage_retrieved_is_different_with_different_config() -> None:
+async def test_storage_retrieved_is_different_with_different_config(tmp_path: Path) -> None:
     """Test that retrieving storage depends on used configuration."""
+    dir_1 = tmp_path / 'dir_1'
+    dir_2 = tmp_path / 'dir_2'
     config_actor = ApifyConfiguration()
-    apify_crawler_1 = ApifyConfiguration()
+    config_actor.storage_dir = str(dir_1)
+    apify_crawler = ApifyConfiguration()
+    apify_crawler.storage_dir = str(dir_2)
 
     async with Actor(configuration=config_actor):
         actor_kvs = await Actor.open_key_value_store()
-        crawler_1 = BasicCrawler(configuration=apify_crawler_1)
-        crawler_kvs = await crawler_1.get_key_value_store()
+        crawler = BasicCrawler(configuration=apify_crawler)
+        crawler_kvs = await crawler.get_key_value_store()
 
     assert actor_kvs is not crawler_kvs
 
 
+async def test_storage_retrieved_is_same_with_equivalent_config() -> None:
+    """Test that retrieving storage depends on used configuration. If two same configuration(even if they are different
+    instances) are used it returns same storage."""
+    config_actor = ApifyConfiguration()
+    apify_crawler = ApifyConfiguration()
+
+    async with Actor(configuration=config_actor):
+        actor_kvs = await Actor.open_key_value_store()
+        crawler = BasicCrawler(configuration=apify_crawler)
+        crawler_kvs = await crawler.get_key_value_store()
+
+    assert actor_kvs is crawler_kvs
+
+
 async def test_storage_retrieved_is_same_with_same_config() -> None:
     """Test that retrieving storage is same if same configuration is used."""
     async with Actor():
         actor_kvs = await Actor.open_key_value_store()
-        crawler_1 = BasicCrawler()
-        crawler_kvs = await crawler_1.get_key_value_store()
+        crawler = BasicCrawler()
+        crawler_kvs = await crawler.get_key_value_store()
 
     assert actor_kvs is crawler_kvs
+
+
+async def test_crawler_uses_apify_config() -> None:
+    """Test that crawler is using ApifyConfiguration in SDK context."""
+    crawler = BasicCrawler()
+    assert isinstance(crawler._service_locator.get_configuration(), ApifyConfiguration)

tests/unit/conftest.py

@@ -13,9 +13,9 @@
 
 from apify_client import ApifyClientAsync
 from apify_shared.consts import ApifyEnvVars
+from crawlee import service_locator
 
 import apify._actor
-from apify._configuration import service_locator
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Iterator
@@ -49,7 +49,7 @@ def _prepare_test_env() -> None:
         service_locator._configuration = None
         service_locator._event_manager = None
         service_locator._storage_client = None
-        service_locator._storage_instance_manager = None
+        service_locator.storage_instance_manager.clear_cache()
 
         # Verify that the test environment was set up correctly.
         assert os.environ.get(ApifyEnvVars.LOCAL_STORAGE_DIR) == str(tmp_path)