Use a per-thread requests.Session instead of module-level sessions

test/test_wbi_helpers.py

@@ -3,11 +3,13 @@
 mapping), search, merge, SPARQL and the various pure helper functions.
 """
 import logging
+import threading
 from types import MappingProxyType
 
 import pytest
 import requests
 
+from wikibaseintegrator import wbi_helpers
 from wikibaseintegrator.wbi_config import config as wbi_config
 from wikibaseintegrator.wbi_exceptions import AnonymousEditNotAllowedError, MaxRetriesReachedException, ModificationFailed, MWApiError, NonExistentEntityError, SaveFailed
 from wikibaseintegrator.wbi_helpers import (check_constraints, download_entity_ttl, execute_sparql_query, format2wbi, format_amount, fulltext_search, generate_entity_instances,
@@ -91,6 +93,41 @@ def test_format_must_be_json(self):
             mediawiki_api_call('POST', mediawiki_api_url='https://example.org/w/api.php', data={'format': 'xml'})
 
 
+class TestDefaultSession:
+    """Requests made without a Login share a session per thread, never across threads."""
+
+    def test_session_is_reused_within_a_thread(self):
+        assert wbi_helpers._get_default_session() is wbi_helpers._get_default_session()
+
+    def test_session_is_not_shared_across_threads(self):
+        sessions = []
+        thread = threading.Thread(target=lambda: sessions.append(wbi_helpers._get_default_session()))
+        thread.start()
+        thread.join()
+
+        assert sessions[0] is not wbi_helpers._get_default_session()
+
+    def test_anonymous_call_uses_default_session(self, requests_mock, monkeypatch):
+        url = 'https://example.org/w/api.php'
+        requests_mock.post(url, json={'success': 1})
+        session = requests.Session()
+        monkeypatch.setattr(wbi_helpers, '_get_default_session', lambda: session)
+        sent = []
+        monkeypatch.setattr(session, 'request', lambda **kwargs: sent.append(kwargs) or requests.Session.request(session, **kwargs))
+
+        mediawiki_api_call('POST', mediawiki_api_url=url, data={'action': 'query'})
+        assert len(sent) == 1
+
+    @pytest.mark.parametrize('name', ['default_session', 'helpers_session'])
+    def test_former_module_sessions_are_deprecated(self, name):
+        with pytest.warns(DeprecationWarning):
+            assert getattr(wbi_helpers, name) is wbi_helpers._get_default_session()
+
+    def test_unknown_attribute_still_raises(self):
+        with pytest.raises(AttributeError):
+            getattr(wbi_helpers, 'does_not_exist')
+
+
 class TestTimeout:
     def test_default_timeout_is_applied(self, wikibase, requests_mock):
         wbi_config['TIMEOUT'] = (3, 33)

wikibaseintegrator/wbi_helpers.py

@@ -7,6 +7,7 @@
 import json
 import logging
 import re
+import threading
 import warnings
 from time import sleep
 from typing import TYPE_CHECKING, Any
@@ -29,7 +30,27 @@
 
 log = logging.getLogger(__name__)
 
-helpers_session = requests.Session()
+# Sessions used for requests made without a Login instance, one per thread since requests.Session isn't thread-safe.
+_thread_local = threading.local()
+
+
+def _get_default_session() -> Session:
+    """
+    Return the session used for requests made without a Login instance, created on first use in each thread.
+    Connections are kept alive between calls in a same thread, but no state is shared across threads.
+    """
+    session = getattr(_thread_local, 'session', None)
+    if session is None:
+        session = _thread_local.session = requests.Session()
+    return session
+
+
+def __getattr__(name: str) -> Any:
+    # Backward compatibility for the former module-level sessions
+    if name in ('default_session', 'helpers_session'):
+        warnings.warn(f"wbi_helpers.{name} is deprecated, the session is now per-thread and private.", DeprecationWarning, stacklevel=2)
+        return _get_default_session()
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
 
 
 class BColors:
@@ -50,9 +71,6 @@ class BColors:
     UNDERLINE = '\033[4m'
 
 
-# Session used for all anonymous requests
-default_session = requests.Session()
-
 
 # MediaWiki error codes meaning the server no longer considers the current session authenticated
 # (e.g. its session store evicted/expired the session, see #902). A CSRF token fetched right before the
@@ -68,7 +86,7 @@ def mediawiki_api_call(method: str, mediawiki_api_url: str | None = None, sessio
 
     :param method: 'GET' or 'POST'
     :param mediawiki_api_url:
-    :param session: If a session is passed, it will be used. Otherwise, a new requests session is created
+    :param session: If a session is passed, it will be used. Otherwise, the current thread's default session is used
     :param login: If provided and the API reports that the session is no longer authenticated (see
                   SESSION_LOST_ERROR_CODES), it is used to fully re-authenticate before retrying.
     :param max_retries: If api request fails due to rate limiting, maxlag, or readonly mode, retry up to `max_retries` times
@@ -93,7 +111,7 @@ def mediawiki_api_call(method: str, mediawiki_api_url: str | None = None, sessio
         kwargs['timeout'] = config['TIMEOUT']
 
     response = None
-    session = session if session else default_session
+    session = session or _get_default_session()
     for n in range(max_retries):
         try:
             response = session.request(method=method, url=mediawiki_api_url, **kwargs)
@@ -300,7 +318,7 @@ def execute_sparql_query(query: str, prefix: str | None = None, endpoint: str |
 
     for _ in range(max_retries):
         try:
-            response = helpers_session.post(sparql_endpoint_url, data=body, headers=request_headers, auth=auth, timeout=config['TIMEOUT'])
+            response = _get_default_session().post(sparql_endpoint_url, data=body, headers=request_headers, auth=auth, timeout=config['TIMEOUT'])
         except requests.exceptions.ConnectionError as e:
             log.exception("Connection error: %s. Sleeping for %d seconds.", e, retry_after)
             sleep(retry_after)
@@ -1114,7 +1132,7 @@ def download_entity_ttl(entity: str, wikibase_url: str | None = None, user_agent
         'User-Agent': get_user_agent(user_agent)
     }
 
-    response = helpers_session.get(wikibase_url + '/entity/' + entity + '.ttl', headers=headers, timeout=config['TIMEOUT'])
+    response = _get_default_session().get(wikibase_url + '/entity/' + entity + '.ttl', headers=headers, timeout=config['TIMEOUT'])
     response.raise_for_status()
     results = response.text