replace result request and session in context by originals

Author: Mantisus

src/crawlee/_types.py

@@ -4,7 +4,7 @@
 from collections.abc import Callable, Iterator, Mapping
 from copy import deepcopy
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Literal, Protocol, TypedDict, TypeVar, cast, overload
+from typing import TYPE_CHECKING, Annotated, Any, Literal, Protocol, TypedDict, TypeVar, cast, overload
 
 from pydantic import ConfigDict, Field, PlainValidator, RootModel
 
@@ -261,11 +261,6 @@ async def get_value(self, key: str, default_value: T | None = None) -> T | None:
 class RequestHandlerRunResult:
     """Record of calls to storage-related context helpers."""
 
-    _REQUEST_SYNC_FIELDS: ClassVar[frozenset[str]] = frozenset({'headers', 'user_data'})
-    _SESSION_SYNC_FIELDS: ClassVar[frozenset[str]] = frozenset(
-        {'_user_data', '_usage_count', '_error_score', '_cookies'}
-    )
-
     def __init__(
         self,
         *,
@@ -279,8 +274,16 @@ def __init__(
         self.key_value_store_changes = dict[tuple[str | None, str | None, str | None], KeyValueStoreChangeRecords]()
 
         # Isolated copies for handler execution
-        self.request = deepcopy(request)
-        self.session = deepcopy(session) if session else None
+        self._request = deepcopy(request)
+        self._session = deepcopy(session) if session else None
+
+    @property
+    def request(self) -> Request:
+        return self._request
+
+    @property
+    def session(self) -> Session | None:
+        return self._session
 
     async def add_requests(
         self,
@@ -331,22 +334,29 @@ async def get_key_value_store(
 
         return self.key_value_store_changes[id, name, alias]
 
-    def sync_request(self, sync_request: Request) -> None:
-        """Sync request state from copies back to originals."""
-        for field in self._REQUEST_SYNC_FIELDS:
-            value = getattr(self.request, field)
-            original_value = getattr(sync_request, field)
-            if value != original_value:
-                object.__setattr__(sync_request, field, value)
-
-    def sync_session(self, sync_session: Session | None = None) -> None:
-        """Sync session state from copies back to originals."""
-        if self.session and sync_session:
-            for field in self._SESSION_SYNC_FIELDS:
+    def apply_request_changes(self, target: Request) -> None:
+        """Apply tracked changes from handler copy to original request."""
+        if self.request.user_data != target.user_data:
+            target.user_data.update(self.request.user_data)
+
+        if self.request.headers != target.headers:
+            target.headers = target.headers | self.request.headers
+
+    def apply_session_changes(self, target: Session | None = None) -> None:
+        """Apply tracked changes from handler copy to original session."""
+        simple_fields: set[str] = {'_usage_count', '_error_score'}
+
+        if self.session and target:
+            if self.session.user_data != target.user_data:
+                target.user_data.update(self.session.user_data)
+
+            if self.session.cookies != target.cookies:
+                target.cookies.set_cookies(self.session.cookies.get_cookies_as_dicts())
+            for field in simple_fields:
                 value = getattr(self.session, field)
-                original_value = getattr(sync_session, field)
+                original_value = getattr(target, field)
                 if value != original_value:
-                    object.__setattr__(sync_session, field, value)
+                    object.__setattr__(target, field, value)
 
 
 @docs_group('Functions')

src/crawlee/crawlers/_basic/_basic_crawler.py

@@ -65,6 +65,7 @@
 from crawlee.storages import Dataset, KeyValueStore, RequestQueue
 
 from ._context_pipeline import ContextPipeline
+from ._context_utils import swaped_context
 from ._logging_utils import (
     get_one_line_error_summary_if_possible,
     reduce_asyncio_timeout_error_to_relevant_traceback_parts,
@@ -1315,8 +1316,6 @@ async def _add_requests(
     async def _commit_request_handler_result(
         self,
         context: BasicCrawlingContext,
-        original_request: Request,
-        original_session: Session | None = None,
     ) -> None:
         """Commit request handler result for the input `context`. Result is taken from `_context_result_map`."""
         result = self._context_result_map[context]
@@ -1329,8 +1328,8 @@ async def _commit_request_handler_result(
 
         await self._commit_key_value_store_changes(result, get_kvs=self.get_key_value_store)
 
-        result.sync_session(sync_session=original_session)
-        result.sync_request(sync_request=original_request)
+        result.apply_session_changes(target=context.session)
+        result.apply_request_changes(target=context.request)
 
     @staticmethod
     async def _commit_key_value_store_changes(
@@ -1419,14 +1418,14 @@ async def __run_task_function(self) -> None:
         try:
             request.state = RequestState.REQUEST_HANDLER
 
-            self._check_request_collision(context.request, context.session)
-
             try:
-                await self._run_request_handler(context=context)
+                with swaped_context(context, request, session):
+                    self._check_request_collision(request, session)
+                    await self._run_request_handler(context=context)
             except asyncio.TimeoutError as e:
                 raise RequestHandlerError(e, context) from e
 
-            await self._commit_request_handler_result(context, original_request=request, original_session=session)
+            await self._commit_request_handler_result(context)
             await wait_for(
                 lambda: request_manager.mark_request_as_handled(request),
                 timeout=self._internal_timeout,
@@ -1438,13 +1437,13 @@ async def __run_task_function(self) -> None:
 
             request.state = RequestState.DONE
 
-            if context.session and context.session.is_usable:
-                context.session.mark_good()
+            if session and session.is_usable:
+                session.mark_good()
 
             self._statistics.record_request_processing_finish(request.unique_key)
 
         except RequestCollisionError as request_error:
-            context.request.no_retry = True
+            request.no_retry = True
             await self._handle_request_error(context, request_error)
 
         except RequestHandlerError as primary_error:
@@ -1459,7 +1458,7 @@ async def __run_task_function(self) -> None:
             await self._handle_request_error(primary_error.crawling_context, primary_error.wrapped_exception)
 
         except SessionError as session_error:
-            if not context.session:
+            if not session:
                 raise RuntimeError('SessionError raised in a crawling context without a session') from session_error
 
             if self._error_handler:
@@ -1469,16 +1468,17 @@ async def __run_task_function(self) -> None:
                 exc_only = ''.join(traceback.format_exception_only(session_error)).strip()
                 self._logger.warning('Encountered "%s", rotating session and retrying...', exc_only)
 
-                context.session.retire()
+                if session:
+                    session.retire()
 
                 # Increment session rotation count.
-                context.request.session_rotation_count = (context.request.session_rotation_count or 0) + 1
+                request.session_rotation_count = (request.session_rotation_count or 0) + 1
 
                 await request_manager.reclaim_request(request)
                 await self._statistics.error_tracker_retry.add(error=session_error, context=context)
             else:
                 await wait_for(
-                    lambda: request_manager.mark_request_as_handled(context.request),
+                    lambda: request_manager.mark_request_as_handled(request),
                     timeout=self._internal_timeout,
                     timeout_message='Marking request as handled timed out after '
                     f'{self._internal_timeout.total_seconds()} seconds',
@@ -1493,7 +1493,7 @@ async def __run_task_function(self) -> None:
             self._logger.debug('The context pipeline was interrupted', exc_info=interrupted_error)
 
             await wait_for(
-                lambda: request_manager.mark_request_as_handled(context.request),
+                lambda: request_manager.mark_request_as_handled(request),
                 timeout=self._internal_timeout,
                 timeout_message='Marking request as handled timed out after '
                 f'{self._internal_timeout.total_seconds()} seconds',

src/crawlee/crawlers/_basic/_context_utils.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+    from crawlee._request import Request
+    from crawlee.sessions import Session
+
+    from ._basic_crawling_context import BasicCrawlingContext
+
+
+@contextmanager
+def swaped_context(
+    context: BasicCrawlingContext,
+    request: Request,
+    session: Session | None,
+) -> Iterator[None]:
+    """Replace context's isolated copies with originals after handler execution."""
+    try:
+        yield
+    finally:
+        # Restore original context state to avoid side effects between different handlers.
+        object.__setattr__(context, 'request', request)
+        object.__setattr__(context, 'session', session)