-
Notifications
You must be signed in to change notification settings - Fork 787
fix: honor session cookies across HTTP client request paths #2104
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: master
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change | ||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
|
@@ -4,6 +4,7 @@ | |||||||||||
| from contextlib import asynccontextmanager | ||||||||||||
| from logging import getLogger | ||||||||||||
| from typing import TYPE_CHECKING, Any, TypedDict | ||||||||||||
| from urllib.request import Request as UrllibRequest | ||||||||||||
|
|
||||||||||||
| from cachetools import LRUCache | ||||||||||||
| from impit import AsyncClient, Browser, HTTPError, Response, TimeoutException, TransportError | ||||||||||||
|
|
@@ -30,6 +31,9 @@ | |||||||||||
|
|
||||||||||||
| logger = getLogger(__name__) | ||||||||||||
|
|
||||||||||||
| # Cache key: (proxy_url, id(cookie_jar) or None) | ||||||||||||
| _ClientCacheKey = tuple[str | None, int | None] | ||||||||||||
|
|
||||||||||||
|
|
||||||||||||
| class _ClientCacheEntry(TypedDict): | ||||||||||||
| """Type definition for client cache entries.""" | ||||||||||||
|
|
@@ -116,7 +120,43 @@ def __init__( | |||||||||||
|
|
||||||||||||
| self._async_client_kwargs = async_client_kwargs | ||||||||||||
|
|
||||||||||||
| self._client_by_proxy_url = LRUCache[str | None, _ClientCacheEntry](maxsize=10) | ||||||||||||
| self._client_cache = LRUCache[_ClientCacheKey, _ClientCacheEntry](maxsize=10) | ||||||||||||
|
|
||||||||||||
| def _prepare_cookies_and_headers( | ||||||||||||
| self, | ||||||||||||
| *, | ||||||||||||
| session: Session | None, | ||||||||||||
| url: str, | ||||||||||||
| headers: HttpHeaders | dict[str, str] | None, | ||||||||||||
| ) -> tuple[CookieJar | None, HttpHeaders | None]: | ||||||||||||
| """Resolve cookie jar / Cookie header based on `persist_cookies_per_session`. | ||||||||||||
|
|
||||||||||||
| When persistence is enabled, attach the session jar to Impit so response cookies update it. | ||||||||||||
| When persistence is disabled, send existing cookies via the `Cookie` header and keep the | ||||||||||||
| shared client (no jar) so clients stay cached and reusable. | ||||||||||||
| """ | ||||||||||||
| if isinstance(headers, dict) or headers is None: | ||||||||||||
| headers = HttpHeaders(headers or {}) | ||||||||||||
|
|
||||||||||||
| if session is None: | ||||||||||||
| return None, headers or None | ||||||||||||
|
|
||||||||||||
| if self._persist_cookies_per_session: | ||||||||||||
| return session.cookies.jar, headers or None | ||||||||||||
|
|
||||||||||||
| cookie_header = self._get_cookie_header(session.cookies.jar, url, headers) | ||||||||||||
| if cookie_header and 'cookie' not in headers: | ||||||||||||
| headers = headers | HttpHeaders({'Cookie': cookie_header}) | ||||||||||||
|
Comment on lines
+147
to
+149
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
|
||||||||||||
|
|
||||||||||||
| return None, headers or None | ||||||||||||
|
|
||||||||||||
| @staticmethod | ||||||||||||
| def _get_cookie_header(jar: CookieJar, url: str, headers: HttpHeaders | None = None) -> str: | ||||||||||||
| """Build a Cookie request header from a jar without attaching the jar to the client.""" | ||||||||||||
| # UrllibRequest is only used to format Cookie headers via CookieJar; it never opens a connection. | ||||||||||||
| request = UrllibRequest(url, headers=dict(headers) if headers else {}) # noqa: S310 | ||||||||||||
| jar.add_cookie_header(request) | ||||||||||||
| return request.get_header('Cookie') or '' | ||||||||||||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
|
||||||||||||
|
|
||||||||||||
| @override | ||||||||||||
| async def crawl( | ||||||||||||
|
|
@@ -128,14 +168,19 @@ async def crawl( | |||||||||||
| statistics: Statistics | None = None, | ||||||||||||
| timeout: timedelta | None = None, | ||||||||||||
| ) -> HttpCrawlingResult: | ||||||||||||
| client = self._get_client(proxy_info.url if proxy_info else None, session.cookies.jar if session else None) | ||||||||||||
| cookie_jar, headers = self._prepare_cookies_and_headers( | ||||||||||||
| session=session, | ||||||||||||
| url=request.url, | ||||||||||||
| headers=request.headers, | ||||||||||||
| ) | ||||||||||||
| client = self._get_client(proxy_info.url if proxy_info else None, cookie_jar) | ||||||||||||
|
|
||||||||||||
| try: | ||||||||||||
| response = await client.request( | ||||||||||||
| url=request.url, | ||||||||||||
| method=request.method, | ||||||||||||
| content=request.payload, | ||||||||||||
| headers=dict(request.headers) if request.headers else None, | ||||||||||||
| headers=dict(headers) if headers else None, | ||||||||||||
| timeout=timeout.total_seconds() if timeout else None, | ||||||||||||
| ) | ||||||||||||
| except TimeoutException as exc: | ||||||||||||
|
|
@@ -166,10 +211,8 @@ async def send_request( | |||||||||||
| ) -> HttpResponse: | ||||||||||||
| validate_http_url(url) | ||||||||||||
|
|
||||||||||||
| if isinstance(headers, dict) or headers is None: | ||||||||||||
| headers = HttpHeaders(headers or {}) | ||||||||||||
|
|
||||||||||||
| client = self._get_client(proxy_info.url if proxy_info else None, session.cookies.jar if session else None) | ||||||||||||
| cookie_jar, headers = self._prepare_cookies_and_headers(session=session, url=url, headers=headers) | ||||||||||||
| client = self._get_client(proxy_info.url if proxy_info else None, cookie_jar) | ||||||||||||
|
|
||||||||||||
| try: | ||||||||||||
| response = await client.request( | ||||||||||||
|
|
@@ -203,7 +246,8 @@ async def stream( | |||||||||||
| ) -> AsyncGenerator[HttpResponse]: | ||||||||||||
| validate_http_url(url) | ||||||||||||
|
|
||||||||||||
| client = self._get_client(proxy_info.url if proxy_info else None, session.cookies.jar if session else None) | ||||||||||||
| cookie_jar, headers = self._prepare_cookies_and_headers(session=session, url=url, headers=headers) | ||||||||||||
| client = self._get_client(proxy_info.url if proxy_info else None, cookie_jar) | ||||||||||||
|
|
||||||||||||
| try: | ||||||||||||
| response = await client.request( | ||||||||||||
|
|
@@ -222,18 +266,22 @@ async def stream( | |||||||||||
| finally: | ||||||||||||
| response.close() | ||||||||||||
|
|
||||||||||||
| @staticmethod | ||||||||||||
| def _make_cache_key(proxy_url: str | None, cookie_jar: CookieJar | None) -> _ClientCacheKey: | ||||||||||||
| return (proxy_url, id(cookie_jar) if cookie_jar is not None else None) | ||||||||||||
|
|
||||||||||||
| def _get_client(self, proxy_url: str | None, cookie_jar: CookieJar | None) -> AsyncClient: | ||||||||||||
| """Retrieve or create an HTTP client for the given proxy URL. | ||||||||||||
| """Retrieve or create an HTTP client for the given proxy URL and cookie jar. | ||||||||||||
|
|
||||||||||||
| If a client for the specified proxy URL does not exist, create and store a new one. | ||||||||||||
| Clients are cached by `(proxy_url, cookie_jar identity)` so sessions with different jars do not share | ||||||||||||
| a client. When cookie persistence is disabled, cookies are sent via headers and `cookie_jar` is `None`, | ||||||||||||
| so a shared client can be reused for the proxy. | ||||||||||||
| """ | ||||||||||||
| cached_data = self._client_by_proxy_url.get(proxy_url) | ||||||||||||
| if cached_data: | ||||||||||||
| client = cached_data['client'] | ||||||||||||
| client_cookie_jar = cached_data['cookie_jar'] | ||||||||||||
| if client_cookie_jar is cookie_jar: | ||||||||||||
| # If the cookie jar matches, return the existing client. | ||||||||||||
| return client | ||||||||||||
| cache_key = self._make_cache_key(proxy_url, cookie_jar) | ||||||||||||
|
|
||||||||||||
| cached_data = self._client_cache.get(cache_key) | ||||||||||||
| if cached_data and cached_data['cookie_jar'] is cookie_jar: | ||||||||||||
| return cached_data['client'] | ||||||||||||
|
|
||||||||||||
| # Prepare a default kwargs for the new client. | ||||||||||||
| kwargs: dict[str, Any] = { | ||||||||||||
|
|
@@ -249,7 +297,11 @@ def _get_client(self, proxy_url: str | None, cookie_jar: CookieJar | None) -> As | |||||||||||
|
|
||||||||||||
| client = AsyncClient(**kwargs, cookie_jar=cookie_jar) | ||||||||||||
|
|
||||||||||||
| self._client_by_proxy_url[proxy_url] = _ClientCacheEntry(client=client, cookie_jar=cookie_jar) | ||||||||||||
| # Evict the least-recently-used entry explicitly before inserting. | ||||||||||||
| if cache_key not in self._client_cache and len(self._client_cache) >= self._client_cache.maxsize: | ||||||||||||
| self._client_cache.popitem() | ||||||||||||
|
Comment on lines
+301
to
+302
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This is dead code now. The LRU cache automatically evicts the items. |
||||||||||||
|
|
||||||||||||
| self._client_cache[cache_key] = _ClientCacheEntry(client=client, cookie_jar=cookie_jar) | ||||||||||||
|
|
||||||||||||
| return client | ||||||||||||
|
|
||||||||||||
|
|
@@ -270,4 +322,4 @@ def _is_proxy_error(error: HTTPError) -> bool: | |||||||||||
| @override | ||||||||||||
| async def cleanup(self) -> None: | ||||||||||||
| """Clean up resources used by the HTTP client.""" | ||||||||||||
| self._client_by_proxy_url.clear() | ||||||||||||
| self._client_cache.clear() | ||||||||||||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -2,6 +2,7 @@ | |
|
|
||
| import asyncio | ||
| import importlib | ||
| import json | ||
| import os | ||
| import sys | ||
| from typing import TYPE_CHECKING | ||
|
|
@@ -12,8 +13,11 @@ | |
| from pydantic import ValidationError | ||
|
|
||
| from crawlee import Request | ||
| from crawlee._types import HttpHeaders | ||
| from crawlee.errors import ProxyError | ||
| from crawlee.fingerprint_suite import HeaderGenerator | ||
| from crawlee.http_clients import CurlImpersonateHttpClient, HttpClient, HttpxHttpClient, ImpitHttpClient | ||
| from crawlee.sessions import Session | ||
| from crawlee.statistics import Statistics | ||
| from tests.unit.server import generate_file_content | ||
| from tests.unit.server_endpoints import HELLO_WORLD | ||
|
|
@@ -323,3 +327,98 @@ def test_import_error_handled(optional_module_name: str, import_path: str) -> No | |
| sys.modules.pop(mod_name, None) | ||
| with pytest.raises(ImportError): | ||
| importlib.import_module(import_path) | ||
|
|
||
|
|
||
| async def test_send_request_sends_session_cookies(http_client: HttpClient, server_url: URL) -> None: | ||
| """`send_request` must attach existing session cookies (same as `crawl`).""" | ||
| session = Session() | ||
| session.cookies.set('auth', 'token-1', domain=server_url.host or '127.0.0.1', path='/') | ||
|
|
||
| response = await http_client.send_request(str(server_url / 'cookies'), session=session) | ||
| body = json.loads(await response.read()) | ||
|
|
||
| assert body['cookies'] == {'auth': 'token-1'} | ||
|
|
||
|
|
||
| async def test_stream_sends_session_cookies(http_client: HttpClient, server_url: URL) -> None: | ||
| """`stream` must attach existing session cookies (same as `crawl`).""" | ||
| session = Session() | ||
| session.cookies.set('auth', 'token-2', domain=server_url.host or '127.0.0.1', path='/') | ||
|
|
||
| content = b'' | ||
| async with http_client.stream(str(server_url / 'cookies'), session=session) as response: | ||
| async for chunk in response.read_stream(): | ||
| content += chunk | ||
|
|
||
| assert json.loads(content)['cookies'] == {'auth': 'token-2'} | ||
|
|
||
|
|
||
| @pytest.mark.parametrize( | ||
| 'custom_http_client', | ||
| [ | ||
| pytest.param(CurlImpersonateHttpClient(persist_cookies_per_session=False), id='curl'), | ||
| pytest.param(HttpxHttpClient(persist_cookies_per_session=False), id='httpx'), | ||
| pytest.param(ImpitHttpClient(persist_cookies_per_session=False), id='impit'), | ||
| ], | ||
| indirect=['custom_http_client'], | ||
| ) | ||
| async def test_persist_cookies_per_session_false(custom_http_client: HttpClient, server_url: URL) -> None: | ||
| """When persistence is disabled, response Set-Cookie must not update the session jar.""" | ||
| session = Session() | ||
| request = Request.from_url(str(server_url.with_path('set_cookies').extend_query(a=1))) | ||
|
|
||
| await custom_http_client.crawl(request, session=session) | ||
|
|
||
| assert {cookie['name']: cookie['value'] for cookie in session.cookies.get_cookies_as_dicts()} == {} | ||
|
|
||
|
|
||
| @pytest.mark.parametrize( | ||
| 'custom_http_client', | ||
| [ | ||
| pytest.param(CurlImpersonateHttpClient(persist_cookies_per_session=True), id='curl'), | ||
| pytest.param(HttpxHttpClient(persist_cookies_per_session=True), id='httpx'), | ||
| pytest.param(ImpitHttpClient(persist_cookies_per_session=True), id='impit'), | ||
| ], | ||
| indirect=['custom_http_client'], | ||
| ) | ||
| async def test_persist_cookies_per_session_true(custom_http_client: HttpClient, server_url: URL) -> None: | ||
| """When persistence is enabled, response Set-Cookie must update the session jar.""" | ||
| session = Session() | ||
| request = Request.from_url(str(server_url.with_path('set_cookies').extend_query(a=1))) | ||
|
|
||
| await custom_http_client.crawl(request, session=session) | ||
|
|
||
| assert {cookie['name']: cookie['value'] for cookie in session.cookies.get_cookies_as_dicts()} == {'a': '1'} | ||
|
|
||
|
|
||
| async def test_httpx_headers_come_from_single_fingerprint() -> None: | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Tests that verify the behavior of a specific client should be placed in a separate test file for that client |
||
| """Accept and User-Agent must come from the same generated fingerprint profile.""" | ||
| header_generator = HeaderGenerator() | ||
| fingerprint = {'Accept': 'text/html', 'Accept-Language': 'en-US', 'User-Agent': 'TestAgent/1.0'} | ||
|
|
||
| with patch.object(header_generator, 'get_specific_headers', return_value=HttpHeaders(fingerprint)) as mocked: | ||
| client = HttpxHttpClient(header_generator=header_generator) | ||
| combined = client._combine_headers(None) | ||
|
|
||
| mocked.assert_called_once_with(header_names={'Accept', 'Accept-Language', 'User-Agent'}) | ||
| assert combined is not None | ||
| assert combined['accept'] == 'text/html' | ||
| assert combined['accept-language'] == 'en-US' | ||
| assert combined['user-agent'] == 'TestAgent/1.0' | ||
|
|
||
|
|
||
| async def test_impit_cleanup_clears_client_cache(server_url: URL) -> None: | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Tests that verify the behavior of a specific client should be placed in a separate test file for that client |
||
| """`ImpitHttpClient.cleanup` must drop cached clients so the next request creates a fresh one.""" | ||
| client = ImpitHttpClient() | ||
| async with client: | ||
| await client.send_request(str(server_url)) | ||
| assert len(client._client_cache) == 1 | ||
| first_client = next(iter(client._client_cache.values()))['client'] | ||
|
|
||
| await client.cleanup() | ||
| assert len(client._client_cache) == 0 | ||
|
|
||
| await client.send_request(str(server_url)) | ||
| assert len(client._client_cache) == 1 | ||
| second_client = next(iter(client._client_cache.values()))['client'] | ||
| assert second_client is not first_client | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
or Noneis unnecessary after: